diff --git a/.gitignore b/.gitignore index 5556d1d5a4a..71fff8b9fee 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ -__pycache__ +__pycache__/ +*.pyc *.so build .coverage_* @@ -16,10 +17,11 @@ onelogger.err runs/ /test_cases/ **/dist/ +AGENTS.md # Sphinx documentation docs/_build docs/apidocs # Git worktrees -.worktrees/ \ No newline at end of file +.worktrees/ diff --git a/docs/gdn_cuda_optimization_reproduction.md b/docs/gdn_cuda_optimization_reproduction.md new file mode 100644 index 00000000000..539f0f7784c --- /dev/null +++ b/docs/gdn_cuda_optimization_reproduction.md @@ -0,0 +1,204 @@ +# GDN CUDA Optimization Reproduction + +This note covers the current GatedDeltaNet CUDA optimization test flow for +Megatron-LM on B200/H100. The optimized kernels are provided by +the `mcore_gdn_opt` Python package; Megatron can optionally route its +gated-delta-rule calls through that package without modifying FLA source. + +## Install + +Use a GPU container that already has `mcore_gdn_opt` and its CUDA extensions +installed, or build + install the package from the internal repository before +running Megatron-LM. + +> **Use commit `12605c5` or later.** Earlier kernel commits produce wrong or +> NaN gradients — see [Kernel version history](#kernel-version-history-dhu-backward) +> below. `12605c5` is the first commit that is both numerically correct and +> NaN-free at training scale. + +```bash +git clone https://gitlab-master.nvidia.com/bhsueh/mcore_gdn_opt.git +cd mcore_gdn_opt +git checkout 12605c515bdbc1f239df991a9cea570f7e79d234 + +# IMPORTANT: use the PINNED submodule commits. Do NOT use +# `git submodule update --remote` — it tracks the branch tip of +# `gated_delta_rule_bwd` and pulls an OLDER, buggy DHU kernel. +git submodule update --init --recursive + +# The nested `gated_delta_rule_bwd/third_party/cutlass` submodule is sometimes +# left unpopulated; point it at the top-level cutlass so the SM100 DHU kernel +# can build. +if [ ! -e third_party/gated_delta_rule_bwd/third_party/cutlass/include ]; then + rm -rf third_party/gated_delta_rule_bwd/third_party/cutlass + ln -s "$PWD/third_party/cutlass" third_party/gated_delta_rule_bwd/third_party/cutlass +fi +export CUTLASS_PATH="${PWD}/third_party/cutlass" + +# FLA must come from the pinned submodule (the wrapper imports +# `fla.ops.gated_delta_rule.chunk_fwd`, present only in that FLA commit). +python -m pip install --no-build-isolation --no-deps \ + third_party/gated_delta_rule_bwd/third_party/fla + +# Build + install the CUDA extensions (editable) and the package. +python -m pip install -e third_party/gated_delta_rule_bwd --no-build-isolation # chunk_delta_h_bwd_sm100 (fwd_h / wy_bwd / dhu) +python -m pip install -e chunk_bwd_kernel_dqkwg --no-build-isolation # dqkwg +python -m pip install -e chunk_gated_delta_rule_fwd --no-build-isolation # fwd_h (forward state recompute) +python -m pip install -e chunk_delta_fused_fwd_bwd --no-build-isolation # fused dv_dhu (bundled since 5661cfa) +python -m pip install -e . --no-build-isolation +``` + +Do not use `PYTHONPATH` or ad-hoc `sys.modules` injection for these tests. The +package and CUDA extensions should be installed in editable mode. + +### Building portable wheels (install into a fixed container without rebuilding) + +The CUDA compile is slow (~30–40 min). Build wheels once, then force-reinstall +them into a running container (the wheels are pinned to e.g. py3.12 / torch2.12 / +SM100a aarch64): + +```bash +for d in third_party/gated_delta_rule_bwd chunk_bwd_kernel_dqkwg \ + chunk_gated_delta_rule_fwd chunk_delta_fused_fwd_bwd .; do + python -m pip wheel --no-build-isolation --no-deps -w ./wheels "./$d" +done +pip install --no-deps --force-reinstall ./wheels/*.whl +``` + +`--force-reinstall` cleanly replaces the image's editable installs with the wheel +`.so` (verified: the imported `.so` becomes byte-identical to the wheel's). + +> **Fused `dv_dhu` (DV_DHU=1) install caveat.** The compiled +> `chunk_delta_fused_fwd_bwd_cuda` extension installs as a *top-level* module, but +> its wrapper `chunk_delta_fused_fwd_bwd/__init__.py` imports it as a *submodule* +> (`from . import chunk_delta_fused_fwd_bwd_cuda`), so a plain wheel install fails +> with `ImportError: chunk_delta_fused_fwd_bwd_cuda extension not found`. Until the +> packaging is fixed, copy the `.so` into the package dir after installing: +> ```bash +> SO=$(python -c "import chunk_delta_fused_fwd_bwd_cuda as m; print(m.__file__)") +> PKG=$(python -c "import chunk_delta_fused_fwd_bwd, os; print(os.path.dirname(chunk_delta_fused_fwd_bwd.__file__))") +> cp "$SO" "$PKG/" +> ``` +> The standalone `dhu` path (`DV_DHU=0, ENABLE_DHU=1`) and the editable `-e` install +> do not need this. + +### Kernel version history (DHU backward) + +| `mcore_gdn_opt` | submodule `gated_delta_rule_bwd` | status | +|---|---|---| +| `3a371f5` | `949c959` | backward grads wrong (DHU dk/dv ~1.4×, WY dβ ~4.6×) | +| `297386a` | `f2351d3` | accuracy fixed, but DHU decay-reciprocal **underflow → NaN** on real inputs | +| `5661cfa` | `2e6d892` | still NaN: the standalone cute DHU kernel emits NaN under very-negative `g` | +| **`12605c5`** | — | **FIXED** — finite and numerically matches FLA; training-safe | + +Validation of `12605c5`: Qwen3.5-VL 397B proxy, 8×GB200, 20 steps, `DV_DHU=0` +(so the standalone cute `chunk_gated_delta_rule_bwd_dhu_cute` kernel is the one +running). grad norm finite for all steps (0 NaN iterations); loss / grad-norm +track the Triton(FLA) baseline to ~4 decimals (iter1 `13.24028` / `15.12`, +iter20 `0.073` / `0.506`). + +## Runtime Flags + +| Case | Flags | +|---|---| +| Triton baseline | `MCORE_GDN_USE_OPT_WRAPPER=0` | +| wrapper auto | `MCORE_GDN_USE_OPT_WRAPPER=1 MCORE_GDN_OPT_BACKEND=auto` | +| `wy_bwd` only | `MCORE_GDN_USE_OPT_WRAPPER=1 MCORE_GDN_OPT_BACKEND=cuda` with other optimized stages disabled | +| `dhu` only | `MCORE_GDN_USE_OPT_WRAPPER=1 MCORE_GDN_OPT_BACKEND=cuda` with other optimized stages disabled | +| `dqkwg` only | `MCORE_GDN_USE_OPT_WRAPPER=1 MCORE_GDN_OPT_BACKEND=cuda` with other optimized stages disabled | +| all three separate | `wy_bwd+dhu+dqkwg` enabled, `fwd_h` and `dv_dhu` disabled | +| all four | `fwd_h+wy_bwd+dhu+dqkwg` enabled, `dv_dhu` disabled | +| `fwd_h+wy_bwd+fused_dv_dhu+dqkwg` | `fwd_h+wy_bwd+dv_dhu+dqkwg` enabled, standalone `dhu` disabled | + +The `dhu_dqkwg` wrapper is intentionally not exposed as a benchmark scenario +because the single-kernel DHU+DQKWG path is not implemented. The remaining +optimized scenarios use standalone `dhu`/`dqkwg` or the real fused `dv_dhu` +kernel. + +## GDN-Only Direct Test + +This bypasses the full GPT layer and measures a direct `GatedDeltaNet` +forward/backward. It checks output, input grad, and parameter grads against the +Triton baseline. + +```bash +python -m tests.unit_tests.ssm.bench_gdn_cuda_opt \ + --dtype bf16 \ + --loss sum \ + --scenarios baseline,separate,all_four,fwd_h_wy_dv_dhu_dqkwg \ + --warmup 5 --repeats 20 --rounds 3 +``` + +Use `--loss square_mean` to reproduce the earlier loss used during debugging, +and add `--fail-on-accuracy` when the command should return non-zero on any +accuracy mismatch. + +Latest B200 spot check for +`B=2,T=8192,H=64,D=128,bf16,loss=sum,warmup=3,repeats=10,rounds=3` +on Megatron-LM `c42dc298a`, `mcore_gdn_opt@9121702`, and +`gated_delta_rule_bwd@949c959`: + +| Scenario | Accuracy vs Triton | Mean us | Speedup | +|---|---:|---:|---:| +| Triton baseline | PASS | 15220.135 | 1.000x | +| CUDA all three separate | FAIL | 13420.346 | 1.134x | +| CUDA all four | FAIL | 12875.528 | 1.182x | + +For this direct `loss=sum` GDN-only check, the optimized scenarios still fail +the strict gradient comparison against the Triton baseline. The current +production workflow is validated with `loss=square_mean`; the latest B200 full +workflow validation passed all requested scenarios and measured `CUDA all four` +at `12895.830 us` (`1.182x`) and `CUDA fwd_h+wy_bwd+fused_dv_dhu+dqkwg` at +`12732.651 us` (`1.197x`). Fresh logs: +`third_party/gdn_doc_loss_sum_20260528_205336.log` and +`third_party/gdn_full_validation_cb51345_20260528_204219.log`. + +## E2E Pytest + +This runs the focused GDN CUDA optimization pytest path. It checks correctness +by default and can print the benchmark table when `MCORE_GDN_UNIT_TEST_PERF=1`. + +```bash +MCORE_GDN_UNIT_TEST_SCENARIOS=baseline,fwd_h_wy_dv_dhu_dqkwg \ +pytest -s tests/unit_tests/ssm/test_gated_delta_net_cuda_opt.py::test_gated_delta_net_cuda_opt_correctness_and_optional_perf -k bf16 +``` + +To generate the E2E benchmark table with NVTX labels: + +```bash +MCORE_GDN_UNIT_TEST_SCENARIOS=baseline,wy,dhu,dqkwg,separate,all_four,fwd_h_wy_dv_dhu_dqkwg \ +MCORE_GDN_UNIT_TEST_PERF=1 \ +MCORE_GDN_UNIT_TEST_WARMUP=5 \ +MCORE_GDN_UNIT_TEST_REPEATS=20 \ +MCORE_GDN_UNIT_TEST_ROUNDS=3 \ +pytest -s tests/unit_tests/ssm/test_gated_delta_net_cuda_opt.py::test_gated_delta_net_cuda_opt_correctness_and_optional_perf -k bf16 +``` + +Latest B200 full workflow validation for `loss=square_mean` passed correctness +for wrapper forced FLA, wrapper auto, wrapper forced CUDA, `CUDA all four`, and +`CUDA fwd_h+wy_bwd+fused_dv_dhu+dqkwg`. Observed speedups were `1.198x` for +wrapper auto, `1.182x` for `CUDA all four`, and `1.197x` for +`CUDA fwd_h+wy_bwd+fused_dv_dhu+dqkwg` versus the Triton baseline. + +## Nsight Systems + +Use the E2E pytest command above under `nsys profile`. The benchmark emits NVTX +labels in this format: + +```text +gdn_only/_/round_/iter_ +``` + +Example: + +```bash +MCORE_GDN_UNIT_TEST_SCENARIOS=baseline,fwd_h_wy_dv_dhu_dqkwg \ +MCORE_GDN_UNIT_TEST_PERF=1 \ +MCORE_GDN_UNIT_TEST_WARMUP=5 \ +MCORE_GDN_UNIT_TEST_REPEATS=20 \ +MCORE_GDN_UNIT_TEST_ROUNDS=3 \ +nsys profile -f true -o gdn_e2e_b200 \ + pytest -s tests/unit_tests/ssm/test_gated_delta_net_cuda_opt.py::test_gated_delta_net_cuda_opt_correctness_and_optional_perf -k bf16 +``` + +Keep profiler outputs (`*.nsys-rep`, `*.sqlite`, `*.qdrep`) out of commits. diff --git a/examples/multimodal_dev/README.md b/examples/multimodal_dev/README.md new file mode 100644 index 00000000000..bdc34414da9 --- /dev/null +++ b/examples/multimodal_dev/README.md @@ -0,0 +1,162 @@ +# multimodal_dev — Standalone Multimodal Training + +Standalone, model-agnostic training entry point for multimodal +vision-language models built on Megatron-Core (FSDP + EP). + +## Directory Structure + +``` +multimodal_dev/ +├── pretrain_multimodal.py # Training entry point (model-agnostic) +├── forward_step.py # Forward step, TP broadcast, loss computation +├── arguments.py # Multimodal CLI arguments +├── data/ +│ └── mock.py # Mock dataset for end-to-end testing +├── models/ +│ ├── __init__.py # MODEL_REGISTRY — central model registry +│ ├── base.py # MultimodalModel base class (vision encoder + GPTModel) +│ └── qwen35_vl/ # Qwen3.5-VL architecture +│ ├── factory.py # Factory functions for pretrain entry point +│ ├── model.py # Qwen35VLModel (MRoPE, vision encoder wiring) +│ ├── configuration.py # TransformerConfig builders and constants +│ ├── specs.py # Layer spec builders (hybrid attention, ViT) +│ ├── mrope.py # 3D MRoPE position ID computation +│ └── vision_encoder.py# ViT encoder (patch embed, merger, RoPE) +└── scripts/ # Launch scripts (torchrun, Slurm) +``` + +## Quick Start + +```bash +torchrun --nproc_per_node=8 multimodal_dev/pretrain_multimodal.py \ + --model-arch qwen35_vl \ + --dataset-provider mock \ + ... # other Megatron args (--num-layers, --hidden-size, etc.) +``` + +## Architecture + +`pretrain_multimodal.py` is **model-agnostic**. All model-specific logic +is delegated to factory functions registered in `MODEL_REGISTRY` +(`models/__init__.py`). The entry point handles only generic concerns: + +- Building `language_config` from Megatron CLI args +- Constructing `vision_config` via the registry +- Applying vision recompute and dtype propagation +- Routing to model and dataset factories + +The `forward_step` is also model-agnostic — it uses the model's +`compute_position_ids()` method polymorphically and passes a standard +batch dict. + +## Adding a New Model Architecture + +Adding a new model (e.g. `llava_next`) requires **no changes** to +`pretrain_multimodal.py` or `forward_step.py`. Follow these steps: + +### Step 1 — Create the model package + +``` +multimodal_dev/models/llava_next/ +├── __init__.py +├── factory.py # Required: factory functions +├── configuration.py # Vision/language TransformerConfig builders +├── model.py # Model class (subclass MultimodalModel) +├── specs.py # Layer spec builders +└── vision_encoder.py # Vision encoder (if custom) +``` + +### Step 2 — Implement factory functions + +Create `factory.py` with up to three functions: + +```python +# models/llava_next/factory.py + +def post_language_config(language_config, args): + """(Optional) Mutate language_config with model-specific fields.""" + # e.g. language_config.some_field = value + pass + +def set_vision_flops_metadata(args, language_config, vision_config): + """(Optional) Set vision FLOPs metadata on args.""" + args.count_vision_model_flops = True + args.vision_flops_variant = "llava_next" + # ... set dimension fields for FLOPs calculation + +def build_model(args, language_config, vision_config, **kwargs): + """(Required) Build and return the complete model instance.""" + from .model import LlavaNextModel + from .specs import get_llava_next_language_spec + + language_spec = get_llava_next_language_spec( + config=language_config, + vp_stage=kwargs.get("vp_stage", None), + pp_rank=None, + ) + return LlavaNextModel( + language_config=language_config, + language_spec=language_spec, + vision_config=vision_config, + # ... model-specific args + ) +``` + +### Step 3 — Register in `MODEL_REGISTRY` + +Add an entry in `models/__init__.py`: + +```python +from multimodal_dev.models.llava_next.configuration import ( + get_llava_next_vision_config, +) +from multimodal_dev.models.llava_next.factory import ( + build_model as _build_llava_next_model, + post_language_config as _llava_next_post_language_config, + set_vision_flops_metadata as _llava_next_vision_flops, +) + +MODEL_REGISTRY["llava_next"] = { + "model_factory_fn": _build_llava_next_model, # required + "vision_config_fn": get_llava_next_vision_config, # required + "post_language_config_fn": _llava_next_post_language_config, # optional + "vision_flops_fn": _llava_next_vision_flops, # optional + "dataset_providers": { # optional + "mock": "multimodal_dev.data.llava_mock.train_valid_test_datasets_provider", + }, +} +``` + +### Step 4 — (Optional) Add a dataset provider + +Create a dataset module under `data/` if the model needs custom data +preprocessing. The provider function signature is: + +```python +def train_valid_test_datasets_provider(train_val_test_num_samples): + """Return (train_dataset, val_dataset, test_dataset).""" + ... +``` + +Register it in the `dataset_providers` dict of the registry entry. +Providers can be either direct callables or dotted import path strings +(resolved lazily at runtime). + +### Step 5 — Launch + +```bash +torchrun --nproc_per_node=8 multimodal_dev/pretrain_multimodal.py \ + --model-arch llava_next \ + --dataset-provider mock \ + ... +``` + +## Registry Entry Reference + +| Field | Required | Signature | +|-------|----------|-----------| +| `model_factory_fn` | Yes | `(args, language_config, vision_config, **kwargs) -> MegatronModule` | +| `vision_config_fn` | Yes | `(num_layers_override=None) -> TransformerConfig` | +| `post_language_config_fn` | No | `(language_config, args) -> None` | +| `vision_flops_fn` | No | `(args, language_config, vision_config) -> None` | +| `dataset_providers` | No | `Dict[str, str \| callable]` | diff --git a/examples/multimodal_dev/__init__.py b/examples/multimodal_dev/__init__.py new file mode 100644 index 00000000000..26496bfed70 --- /dev/null +++ b/examples/multimodal_dev/__init__.py @@ -0,0 +1 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. diff --git a/examples/multimodal_dev/arguments.py b/examples/multimodal_dev/arguments.py new file mode 100644 index 00000000000..35655831bfb --- /dev/null +++ b/examples/multimodal_dev/arguments.py @@ -0,0 +1,100 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Extra CLI arguments for multimodal_dev standalone training.""" + + +def add_multimodal_args(parser): + """Add multimodal-specific arguments to the Megatron argument parser.""" + group = parser.add_argument_group( + "Multimodal", "Multimodal model arguments", + ) + + group.add_argument( + "--model-arch", + type=str, + default="qwen35_vl", + help="Model architecture. Available: qwen35_vl", + ) + group.add_argument( + "--model-variant", + type=str, + default="proxy", + help="Model variant (size). E.g. proxy, 9b, 397b_a17b", + ) + group.add_argument( + "--dataset-provider", + type=str, + default="mock", + help="Dataset provider: mock", + ) + group.add_argument( + "--image-token-id", + type=int, + default=248056, + help="Token ID for image placeholder tokens", + ) + group.add_argument( + "--image-size", + type=int, + default=224, + help="Image size (height and width) for mock data", + ) + group.add_argument( + "--total-seq-length", + type=int, + default=1024, + help="Total sequence length for mock data", + ) + group.add_argument( + "--image-seq-length", + type=int, + default=256, + help="Number of image tokens in mock data", + ) + group.add_argument( + "--vision-num-layers", + type=int, + default=None, + help=( + "Override for vision backbone depth. " + "Useful for proxy perf runs." + ), + ) + group.add_argument( + "--hf-processor-path", + type=str, + default=None, + help=( + "HuggingFace processor path for real VLM datasets " + "(e.g. Qwen/Qwen2.5-VL-7B-Instruct)" + ), + ) + group.add_argument( + "--recompute-vision", + action="store_true", + default=False, + help=( + "Enable full activation recomputation for vision encoder layers. " + "Uses uniform method and recomputes every layer. " + "Independent of the decoder --recompute-* flags." + ), + ) + group.add_argument( + "--use-packed-sequence", + action="store_true", + default=False, + help=( + "Pack variable-length sequences into THD format to eliminate " + "padding waste." + ), + ) + group.add_argument( + "--use-vanilla-collate-fn", + action="store_true", + default=False, + help=( + "Use vanilla collate function to collate the data." + ), + ) + + return parser diff --git a/examples/multimodal_dev/data/__init__.py b/examples/multimodal_dev/data/__init__.py new file mode 100644 index 00000000000..26496bfed70 --- /dev/null +++ b/examples/multimodal_dev/data/__init__.py @@ -0,0 +1 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. diff --git a/examples/multimodal_dev/data/mock.py b/examples/multimodal_dev/data/mock.py new file mode 100644 index 00000000000..0975b132013 --- /dev/null +++ b/examples/multimodal_dev/data/mock.py @@ -0,0 +1,192 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Mock dataset for multimodal_dev end-to-end testing. + +Generates synthetic image + text data. Each sample has random text +tokens with image-token placeholders, random pixel values sized for the +vision encoder, 3D MRoPE position IDs, and shifted labels. +""" + +import torch +from torch.utils.data import Dataset + +from examples.multimodal_dev.models.qwen35_vl.configuration import ( + QWEN35_VL_IMAGE_TOKEN_ID, + QWEN35_VL_VIDEO_TOKEN_ID, + QWEN35_VL_VISION_START_TOKEN_ID, +) +from examples.multimodal_dev.models.qwen35_vl.mrope import get_rope_index + + +class MockQwen35VLDataset(Dataset): + """Synthetic Qwen3.5-VL training samples. + + Args: + num_samples: Number of samples. + seq_length: Total sequence length (text + image tokens). + image_seq_length: Number of image tokens per sample. + vocab_size: Vocabulary size for random text tokens. + image_token_id: Token ID for image placeholders. + video_token_id: Token ID for video placeholders. + vision_start_token_id: Token ID marking start of a vision region. + image_size: Image height and width in pixels. + patch_size: Spatial patch size. + temporal_patch_size: Temporal patch size. + spatial_merge_size: Spatial merge factor. + """ + + def __init__( + self, + num_samples: int = 1000, + seq_length: int = 1024, + image_seq_length: int = 256, + vocab_size: int = 248320, + image_token_id: int = QWEN35_VL_IMAGE_TOKEN_ID, + video_token_id: int = QWEN35_VL_VIDEO_TOKEN_ID, + vision_start_token_id: int = QWEN35_VL_VISION_START_TOKEN_ID, + image_size: int = 224, + patch_size: int = 16, + temporal_patch_size: int = 2, + spatial_merge_size: int = 2, + ): + self.num_samples = num_samples + self.seq_length = seq_length + self.vocab_size = vocab_size + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.vision_start_token_id = vision_start_token_id + self.image_size = image_size + self.patch_size = patch_size + self.temporal_patch_size = temporal_patch_size + self.spatial_merge_size = spatial_merge_size + + h_patches = image_size // patch_size + w_patches = image_size // patch_size + t_patches = temporal_patch_size + self.grid_thw = torch.tensor([[t_patches, h_patches, w_patches]]) + + self.num_merged_tokens = ( + t_patches + * (h_patches // spatial_merge_size) + * (w_patches // spatial_merge_size) + ) + self.image_seq_length = min( + image_seq_length, self.num_merged_tokens, + ) + self.total_patches = t_patches * h_patches * w_patches + + def __len__(self): + return self.num_samples + + def __getitem__(self, idx): + # Reserve 1 slot for the vision_start sentinel before image tokens. + text_length = self.seq_length - self.image_seq_length - 1 + text_tokens = torch.randint( + 1, self.vocab_size, (text_length,), dtype=torch.long, + ) + special_ids = { + self.image_token_id, + self.video_token_id, + self.vision_start_token_id, + } + for sid in special_ids: + text_tokens[text_tokens == sid] = 1 + + prefix_len = text_length // 2 + suffix_len = text_length - prefix_len + input_ids = torch.cat([ + text_tokens[:prefix_len], + torch.tensor( + [self.vision_start_token_id], dtype=torch.long, + ), + torch.full( + (self.image_seq_length,), + self.image_token_id, + dtype=torch.long, + ), + text_tokens[prefix_len: prefix_len + suffix_len], + ]) + + labels = input_ids.clone() + labels[:-1] = input_ids[1:] + labels[-1] = 0 + + loss_mask = (input_ids != self.image_token_id).float() + loss_mask[-1] = 0 + + pixel_dim = ( + 3 + * self.temporal_patch_size + * self.patch_size + * self.patch_size + ) + pixel_values = torch.randn(self.total_patches, pixel_dim) + + image_grid_thw = self.grid_thw.clone() + + position_ids, _ = get_rope_index( + spatial_merge_size=self.spatial_merge_size, + image_token_id=self.image_token_id, + video_token_id=self.video_token_id, + vision_start_token_id=self.vision_start_token_id, + input_ids=input_ids.unsqueeze(0), + image_grid_thw=image_grid_thw, + ) + position_ids = position_ids.squeeze(1) + + return { + "input_ids": input_ids, + "labels": labels, + "loss_mask": loss_mask, + "cu_seqlens": torch.tensor([0, self.seq_length], dtype=torch.int32), + "cu_seqlens_padded": torch.tensor( + [0, self.seq_length], dtype=torch.int32, + ), + "max_seqlen": torch.tensor(self.seq_length, dtype=torch.int32), + "position_ids": position_ids, + "pixel_values": pixel_values, + "image_grid_thw": image_grid_thw, + } + + +def mock_collate_fn(batch): + """Collate: handles position_ids ``[3, S]`` stacking.""" + result = {} + keys = batch[0].keys() + for key in keys: + tensors = [sample[key] for sample in batch] + if key == "position_ids": + result[key] = torch.stack(tensors, dim=1) + elif key == "image_grid_thw": + result[key] = torch.cat(tensors, dim=0) + elif key == "pixel_values": + result[key] = torch.cat(tensors, dim=0) + else: + result[key] = torch.stack(tensors, dim=0) + return result + + +def train_valid_test_datasets_provider(train_val_test_num_samples): + """Provide mock train / val / test datasets.""" + from megatron.training import get_args + + args = get_args() + kwargs = dict( + seq_length=getattr(args, "total_seq_length", 1024), + image_seq_length=getattr(args, "image_seq_length", 256), + vocab_size=getattr(args, "padded_vocab_size", 248320), + image_token_id=getattr(args, "image_token_id", 248056), + image_size=getattr(args, "image_size", 224), + ) + + train_ds = MockQwen35VLDataset( + num_samples=train_val_test_num_samples[0], **kwargs, + ) + val_ds = MockQwen35VLDataset( + num_samples=train_val_test_num_samples[1], **kwargs, + ) + test_ds = MockQwen35VLDataset( + num_samples=train_val_test_num_samples[2], **kwargs, + ) + + return train_ds, val_ds, test_ds diff --git a/examples/multimodal_dev/data/vlm_dataset.py b/examples/multimodal_dev/data/vlm_dataset.py new file mode 100644 index 00000000000..3a493087252 --- /dev/null +++ b/examples/multimodal_dev/data/vlm_dataset.py @@ -0,0 +1,371 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Simple VLM dataset for multimodal_dev training. + +Single-turn image-text dataset using a HuggingFace ``AutoProcessor`` for +tokenization and image preprocessing. Currently supports CORD-V2 (receipt +OCR). No multi-turn support — each sample is one image + question → +answer pair. + +Each image is preprocessed via ``qwen_vl_utils.process_vision_info`` and +fed to the processor with Qwen-VL's recommended ``min_pixels`` / +``max_pixels`` budget, so the per-sample patch grid varies with aspect +ratio. The run script must therefore pass ``--use-vanilla-collate-fn`` +(which the example launcher does) so the dataloader does not try to stack +variable-shape tensors. + +Usage:: + + torchrun ... pretrain_multimodal.py \\ + --model-arch qwen35_vl --dataset-provider cord_v2 \\ + --hf-processor-path Qwen/Qwen3.5-397B-A17B \\ + --total-seq-length 4096 --use-vanilla-collate-fn +""" + +import json +import logging +import random +from typing import Dict, List, Optional + +import torch +from torch.utils.data import Dataset + +try: + from qwen_vl_utils import process_vision_info + HAVE_QWEN_VL_UTILS = True +except ImportError: + HAVE_QWEN_VL_UTILS = False + +logger = logging.getLogger(__name__) + +# Qwen-VL recommended pixel-budget range; lets the processor pick a +# per-image patch grid that respects aspect ratio. +_QWEN_VL_MIN_PIXELS = 256 * 28 * 28 # 200_704 +_QWEN_VL_MAX_PIXELS = 1280 * 28 * 28 # 1_003_520 + + +# --------------------------------------------------------------------------- +# CORD-V2 helpers +# --------------------------------------------------------------------------- + +def _json2token(obj, sort_json_key=True): + """Convert a JSON object to a token-sequence string (Donut format).""" + if isinstance(obj, dict): + if len(obj) == 1 and "text_sequence" in obj: + return obj["text_sequence"] + output = "" + keys = sorted(obj.keys(), reverse=True) if sort_json_key else obj.keys() + for k in keys: + output += f"" + _json2token(obj[k], sort_json_key) + f"" + return output + if isinstance(obj, list): + return "".join(_json2token(item, sort_json_key) for item in obj) + return str(obj) + + +def load_cord_v2(split="train"): + """Load CORD-V2 and return a list of ``{image, question, answer}`` dicts.""" + from datasets import load_dataset + + ds = load_dataset("naver-clova-ix/cord-v2", split=split) + rng = random.Random(42) + examples = [] + for ex in ds: + gt = json.loads(ex["ground_truth"]) + gt_jsons = gt.get("gt_parses") or [gt["gt_parse"]] + text = rng.choice( + [_json2token(g, sort_json_key=True) for g in gt_jsons] + ) + examples.append( + {"image": ex["image"], "question": "Describe this image.", "answer": text} + ) + return examples + + +# --------------------------------------------------------------------------- +# Dataset +# --------------------------------------------------------------------------- + +class CordV2VLMDataset(Dataset): + """Single-turn VLM dataset backed by CORD-V2. + + Each sample is tokenized by the HF ``AutoProcessor`` and the image is + handed to the processor with Qwen-VL's dynamic-resolution budget + (``min_pixels`` / ``max_pixels``); the per-image patch grid varies with + aspect ratio. + + Args: + examples: Output of :func:`load_cord_v2`. + processor: ``AutoProcessor`` instance. + seq_length: End-truncate ``input_ids`` at this length. + image_token_id: Token ID for image placeholders. + target_length: Virtual dataset length (repeats examples if needed). + + NOTE: + For the Qwen3.5-VL processor, the temporal patch dimension is + always 2 (the processor duplicates a single frame so the 3D conv + behaves like a 2D conv on one image) — ``image_grid_thw`` therefore + has shape ``[num_images, 3]`` with ``T=2`` per image. + ``pixel_values`` has shape ``[total_patches, 3 * T * P * P]`` where + ``P`` is the processor's patch size. + """ + + def __init__( + self, + examples: List[Dict], + processor, + seq_length: int = 2048, + image_token_id: Optional[int] = None, + target_length: Optional[int] = None, + ): + if not HAVE_QWEN_VL_UTILS: + raise ImportError( + "qwen_vl_utils is required for Qwen3.5-VL preprocessing. " + "Install with `pip install qwen-vl-utils`.", + ) + self.examples = examples + self.processor = processor + self.seq_length = seq_length + self._length = target_length if target_length else len(examples) + tok = processor.tokenizer + # Falling back to 0 is unsafe: token 0 is a real vocab token in many + # tokenizers (incl. Qwen) and would be silently masked. Prefer EOS, + # and require at least one of pad/eos to be set. + if tok.pad_token_id is not None: + self.pad_token_id = int(tok.pad_token_id) + elif tok.eos_token_id is not None: + self.pad_token_id = int(tok.eos_token_id) + else: + raise ValueError( + "Tokenizer has neither pad_token_id nor eos_token_id; " + "cannot derive a safe pad id for loss masking.", + ) + + # Resolve image token ID. Vision embeddings are scattered into + # positions equal to this id by the model, so a wrong id silently + # breaks training — fail loudly rather than return None. + if image_token_id is not None: + self.image_token_id = int(image_token_id) + else: + vocab = tok.get_vocab() + for candidate in ("<|image_pad|>", "<|placeholder|>"): + if candidate in vocab: + self.image_token_id = int(vocab[candidate]) + break + else: + raise ValueError( + "Could not resolve image token id from tokenizer " + f"({type(tok).__name__}); pass --image-token-id " + "explicitly.", + ) + + # Structural tokens that must never appear as a loss target: + # pad, image, plus everything the tokenizer registered as special + # (im_start/im_end, vision_start/vision_end, video_pad, endoftext...). + # Mirrors megatron-bridge's extract_skipped_token_ids convention. + skipped: set = set(int(x) for x in (tok.all_special_ids or [])) + skipped.add(self.pad_token_id) + skipped.add(self.image_token_id) + self.skipped_token_ids = torch.tensor( + sorted(skipped), dtype=torch.long, + ) + + def __len__(self) -> int: + return self._length + + def _mark_assistant_span( + self, + input_ids_list: List[int], + asst_text: str, + loss_mask: torch.Tensor, + ) -> bool: + """Find ``asst_text`` as a contiguous token span in ``input_ids_list`` + and set ``loss_mask`` to 1 over those positions. + + Substring tokenization is sensitive to surrounding whitespace and + BPE merge boundaries, so we try a few common variants. Returns + True if a span was found. + """ + tokenizer = self.processor.tokenizer + n = len(input_ids_list) + variants = ( + asst_text, + asst_text + "\n", + asst_text.strip(), + asst_text.strip() + "\n", + ) + for variant in variants: + span_tokens = tokenizer( + variant, add_special_tokens=False, + )["input_ids"] + m = len(span_tokens) + if m == 0 or m > n: + continue + # Backward search: rightmost match = the actual assistant turn. + for start in range(n - m, -1, -1): + if input_ids_list[start : start + m] == span_tokens: + loss_mask[start : start + m] = 1.0 + return True + return False + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + example = self.examples[idx % len(self.examples)] + + # Conversation schema must include the actual image object inside + # the content so the chat template + process_vision_info can extract + # it (matches megatron-bridge's qwen2_5_collate_fn convention; also + # used by the Qwen3-VL processor). + conversation = [ + { + "role": "user", + "content": [ + {"type": "image", "image": example["image"]}, + {"type": "text", "text": example["question"]}, + ], + }, + { + "role": "assistant", + "content": [{"type": "text", "text": example["answer"]}], + }, + ] + + text = self.processor.apply_chat_template( + conversation, tokenize=False, add_generation_prompt=False, + ) + images, _ = process_vision_info(conversation) + batch = self.processor( + text=[text], + images=images, + return_tensors="pt", + min_pixels=_QWEN_VL_MIN_PIXELS, + max_pixels=_QWEN_VL_MAX_PIXELS, + ) + + input_ids = batch["input_ids"].squeeze(0) + pixel_values = batch["pixel_values"].to(torch.bfloat16) + image_grid_thw = batch["image_grid_thw"] # [num_images, 3] + + # End-truncate so the model never sees more than seq_length tokens. + # Qwen-VL chat template puts the image at the user-turn start and the + # assistant answer trails at the end, so end-truncation preserves the + # image_pad block for normal-sized images; if the image grid alone + # already exceeds seq_length, the model will fail loudly at the + # masked_scatter step. + if input_ids.shape[0] > self.seq_length: + logger.warning( + "Sample idx=%d has %d tokens > seq_length=%d; truncating.", + idx, input_ids.shape[0], self.seq_length, + ) + input_ids = input_ids[: self.seq_length] + + # SFT loss mask: start fully masked, then unmask only the assistant + # answer span found via substring token search (mirrors + # megatron-bridge's create_multiturn_loss_mask_by_search). The user + # turn, chat-template tags, and image tokens stay masked. + loss_mask = torch.zeros_like(input_ids, dtype=torch.float32) + found = self._mark_assistant_span( + input_ids.tolist(), example["answer"], loss_mask, + ) + if not found: + logger.warning( + "Assistant span not located for example idx=%d; " + "loss_mask will be all-zero for this sample.", + idx, + ) + + # Shifted next-token labels: labels[i] is the target for position i. + labels = input_ids.clone() + labels[:-1] = input_ids[1:] + labels[-1] = -100 + + # Mask structural tokens on the *labels* (the prediction targets), + # not on input_ids — matches the next-token timeline. + labels[torch.isin(labels, self.skipped_token_ids)] = -100 + + # Shift loss_mask left by one so position i decides whether to learn + # input_ids[i] -> labels[i] (== input_ids[i+1]). Last position is + # never trained (no next token to predict). + loss_mask = torch.cat( + [loss_mask[1:], torch.zeros(1, dtype=loss_mask.dtype)], + ) + + # Enforce label = -100 wherever we won't compute loss. + labels[loss_mask == 0] = -100 + + return { + "input_ids": input_ids, + "labels": labels, + "loss_mask": loss_mask, + "pixel_values": pixel_values, + "image_grid_thw": image_grid_thw, + } + + +# --------------------------------------------------------------------------- +# Megatron dataset provider interface +# --------------------------------------------------------------------------- + +def train_valid_test_datasets_provider(train_val_test_num_samples): + """Provide CORD-V2 train / val / test datasets. + + Requires ``--hf-processor-path`` to point to a HuggingFace VL model + (e.g. ``Qwen/Qwen3.5-397B-A17B``) whose processor handles tokenization + and image preprocessing. + """ + from transformers import AutoProcessor + + from megatron.training import get_args + + args = get_args() + + processor_path = getattr(args, "hf_processor_path", None) + if processor_path is None: + raise ValueError( + "cord_v2 dataset requires --hf-processor-path " + "(e.g. Qwen/Qwen3.5-397B-A17B)" + ) + processor = AutoProcessor.from_pretrained( + processor_path, trust_remote_code=True, + ) + + seq_length = ( + getattr(args, "total_seq_length", None) + or getattr(args, "seq_length", 2048) + ) + image_token_id = getattr(args, "image_token_id", None) + + # Load real data + train_examples = load_cord_v2(split="train") + val_examples = load_cord_v2(split="validation") + test_examples = load_cord_v2(split="test") + + def _make(examples, num_samples): + return CordV2VLMDataset( + examples=examples, + processor=processor, + seq_length=seq_length, + image_token_id=image_token_id, + target_length=num_samples, + ) + + # MegatronPretrainingSampler asserts total_samples > 0, so val/test + # datasets must have non-zero length even when eval is disabled. + train_ds = _make(train_examples, train_val_test_num_samples[0]) + val_ds = _make(val_examples, max(train_val_test_num_samples[1], 1)) + test_ds = _make(test_examples, max(train_val_test_num_samples[2], 1)) + + return train_ds, val_ds, test_ds + + +if __name__ == "__main__": + from transformers import AutoProcessor + processor = AutoProcessor.from_pretrained( + "Qwen/Qwen3.5-397B-A17B", trust_remote_code=True, + ) + examples = load_cord_v2(split="train") + dataset = CordV2VLMDataset( + examples=examples, + processor=processor, + image_token_id=248056, + ) + print(dataset[0]) diff --git a/examples/multimodal_dev/forward_step.py b/examples/multimodal_dev/forward_step.py new file mode 100644 index 00000000000..cc141b3e16d --- /dev/null +++ b/examples/multimodal_dev/forward_step.py @@ -0,0 +1,444 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Forward step, TP broadcast, and loss for multimodal_dev training.""" + +import math +from functools import partial +from itertools import accumulate +from typing import Any, Dict, Iterator, Optional + +import torch +import torch.nn.functional as F + +from megatron.core import mpu +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.parallel_state import ( + get_tensor_model_parallel_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_src_rank, +) +from megatron.training import get_args + +# ------------------------------------------------------------------- +# dtype <-> int mapping for cross-rank broadcast +# ------------------------------------------------------------------- + +_DTYPE_MAP = { + torch.float32: 0, + torch.float16: 1, + torch.bfloat16: 2, + torch.int64: 3, + torch.int32: 4, + torch.bool: 5, +} +_ID_MAP = {v: k for k, v in _DTYPE_MAP.items()} + + +def _dtype_to_id(dtype): + return _DTYPE_MAP.get(dtype, 0) + + +def _id_to_dtype(id_val): + return _ID_MAP.get(id_val, torch.float32) + + +# ------------------------------------------------------------------- +# Tensor broadcast helper +# ------------------------------------------------------------------- + + +def _broadcast_tensor(tensor, src, group, device): + """Broadcast a single tensor from *src* to all ranks in *group*.""" + ndim = torch.tensor( + [len(tensor.shape) if tensor is not None else 0], dtype=torch.long, device=device + ) + torch.distributed.broadcast(ndim, src, group=group) + + if ndim.item() == 0: + return None + + if tensor is not None: + shape_tensor = torch.tensor(list(tensor.shape), dtype=torch.long, device=device) + dtype_id = torch.tensor([_dtype_to_id(tensor.dtype)], dtype=torch.long, device=device) + else: + shape_tensor = torch.zeros(ndim.item(), dtype=torch.long, device=device) + dtype_id = torch.zeros(1, dtype=torch.long, device=device) + + torch.distributed.broadcast(shape_tensor, src, group=group) + torch.distributed.broadcast(dtype_id, src, group=group) + + dtype = _id_to_dtype(dtype_id.item()) + shape = tuple(shape_tensor.tolist()) + + if tensor is None: + tensor = torch.empty(shape, dtype=dtype, device=device) + torch.distributed.broadcast(tensor, src, group=group) + return tensor + + +# ------------------------------------------------------------------- +# Batch broadcast across TP ranks +# ------------------------------------------------------------------- + + +def broadcast_data_batch(data, device="cuda"): + """Broadcast a data-batch dict from TP rank 0 to all TP ranks.""" + src = get_tensor_model_parallel_src_rank() + group = get_tensor_model_parallel_group() + + if data is None: + data = {} + + if get_tensor_model_parallel_rank() == 0: + keys = list(data.keys()) + key_str = ",".join(keys) + key_bytes = key_str.encode("utf-8") + key_len = torch.tensor([len(key_bytes)], dtype=torch.long, device=device) + else: + key_len = torch.zeros(1, dtype=torch.long, device=device) + keys = [] + + torch.distributed.broadcast(key_len, src, group=group) + + if get_tensor_model_parallel_rank() == 0: + key_tensor = torch.tensor(list(key_bytes), dtype=torch.uint8, device=device) + else: + key_tensor = torch.zeros(key_len.item(), dtype=torch.uint8, device=device) + + torch.distributed.broadcast(key_tensor, src, group=group) + + if get_tensor_model_parallel_rank() != 0: + key_str = bytes(key_tensor.cpu().tolist()).decode("utf-8") + keys = key_str.split(",") if key_str else [] + + result = {} + for key in keys: + tensor = data.get(key, None) if data else None + if tensor is not None and isinstance(tensor, torch.Tensor): + tensor = tensor.to(device) + result[key] = _broadcast_tensor( + tensor if isinstance(tensor, torch.Tensor) else None, src, group, device + ) + + return result + + +# ------------------------------------------------------------------- +# THD (packed sequence) helpers +# ------------------------------------------------------------------- + + +def _build_packed_seq_params(seq_lengths: torch.Tensor, device: torch.device) -> PackedSeqParams: + """Build ``PackedSeqParams`` from per-sample valid sequence lengths. + + Args: + seq_lengths: ``[B]`` valid token counts per sample. + device: Target device for cu_seqlens tensors. + + Returns: + A ``PackedSeqParams`` instance with ``qkv_format='thd'``. + """ + if not isinstance(seq_lengths, torch.Tensor): + seq_lengths = torch.tensor(seq_lengths) + lengths_t = seq_lengths.to(device=device, dtype=torch.int32) + cu_seqlens = torch.zeros(lengths_t.numel() + 1, dtype=torch.int32, device=device) + torch.cumsum(lengths_t, dim=0, out=cu_seqlens[1:]) + max_seqlen = int(lengths_t.max().item()) + return _build_packed_seq_params_from_cu_seqlens(cu_seqlens=cu_seqlens, max_seqlen=max_seqlen) + + +def _build_packed_seq_params_from_cu_seqlens( + cu_seqlens: torch.Tensor, max_seqlen: int +) -> PackedSeqParams: + """Build ``PackedSeqParams`` from packed cumulative sequence lengths. + + ``cu_seqlens`` must already be on the target compute device. + """ + cs = cu_seqlens.to(dtype=torch.int32) + total_tokens = int(cs[-1].item()) + return PackedSeqParams( + cu_seqlens_q=cs, + cu_seqlens_kv=cs, + cu_seqlens_q_padded=cs, + cu_seqlens_kv_padded=cs, + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + qkv_format='thd', + total_tokens=total_tokens, + ) + + +def pack_or_pad_batch( + batch: Optional[list[Dict[str, Any]]], + use_packed_sequence: bool = False, + seq_length: Optional[int] = None, + device="cuda", +) -> Dict[str, Any]: + """Pack or pad a ``[B, S]`` batch into ``[1, T]`` THD or ``[B, S]`` BSHD. + + Must be invoked on every TP rank. On the TP source rank ``batch`` is + the per-sample dict list from the dataset; on other TP ranks ``batch`` + may be ``None`` (the function relies on the trailing TP broadcast to + distribute results). All metadata needed to reconstruct + ``PackedSeqParams`` (``cu_seqlens``, ``cu_seqlens_padded``, + ``max_seqlen``, ``total_tokens``) is broadcast alongside the data, so + every rank can build an identical ``PackedSeqParams`` on its own. + """ + tp_size = mpu.get_tensor_model_parallel_world_size() + cp_size = mpu.get_context_parallel_world_size() + is_src = mpu.get_tensor_model_parallel_rank() == 0 + + # SP is an explicit runtime option; TP>1 does not imply SP is enabled. + # get_args() itself raises in test contexts where megatron globals are + # not initialised. + try: + has_sp = bool(getattr(get_args(), "sequence_parallel", False)) + except AssertionError: + has_sp = False + + if cp_size > 1: + divisible_by = (tp_size * cp_size * 2) if has_sp else (cp_size * 2) + else: + divisible_by = tp_size if has_sp else 1 + + if use_packed_sequence: + packed_batch: Dict[str, Any] = {} + + if is_src: + assert batch is not None, "source TP rank must provide a batch" + input_ids_list, labels_list, loss_mask_list = [], [], [] + pixel_values_list, image_grid_thw_list = [], [] + seqlens_list, seqlens_padded_list = [], [] + + for sample in batch: + seqlen = sample["input_ids"].shape[0] + assert ( + sample["labels"].shape == sample["input_ids"].shape == sample["loss_mask"].shape + ), "labels, input_ids, and loss_mask must have the same shape" + target_len = math.ceil(seqlen / divisible_by) * divisible_by + input_ids_list.append(F.pad(sample["input_ids"], (0, target_len - seqlen), value=0)) + labels_list.append(F.pad(sample["labels"], (0, target_len - seqlen), value=-100)) + loss_mask_list.append(F.pad(sample["loss_mask"], (0, target_len - seqlen), value=0)) + seqlens_list.append(seqlen) + seqlens_padded_list.append(target_len) + pixel_values_list.append(sample["pixel_values"]) + image_grid_thw_list.append(sample["image_grid_thw"]) + + cu_seqlens = list(accumulate(seqlens_list, initial=0)) + cu_seqlens_padded = list(accumulate(seqlens_padded_list, initial=0)) + + packed_batch["input_ids"] = torch.concat(input_ids_list, dim=0).unsqueeze(0) + packed_batch["labels"] = torch.concat(labels_list, dim=0).unsqueeze(0) + packed_batch["loss_mask"] = torch.concat(loss_mask_list, dim=0).unsqueeze(0) + packed_batch["pixel_values"] = torch.concat(pixel_values_list) + packed_batch["image_grid_thw"] = torch.concat(image_grid_thw_list) + # cu_seqlens / cu_seqlens_padded need to reach non-source TP ranks + # so each rank can build an identical PackedSeqParams. + packed_batch["cu_seqlens"] = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + packed_batch["cu_seqlens_padded"] = torch.tensor( + cu_seqlens_padded, dtype=torch.int32, device=device + ) + + packed_batch = broadcast_data_batch(packed_batch, device=device) + + cu_seqlens_t = packed_batch.pop("cu_seqlens") + cu_seqlens_padded_t = packed_batch.pop("cu_seqlens_padded") + # Derive max_seqlen / total_tokens from the (broadcast) cu_seqlens — + # no extra collective needed. + max_seqlen_q = int((cu_seqlens_padded_t[1:] - cu_seqlens_padded_t[:-1]).max().item()) + total_tokens = int(cu_seqlens_padded_t[-1].item()) + + packed_batch["packed_seq_params"] = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens_t, + cu_seqlens_kv=cu_seqlens_t, + cu_seqlens_q_padded=cu_seqlens_padded_t, + cu_seqlens_kv_padded=cu_seqlens_padded_t, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_q, + total_tokens=total_tokens, + ) + return packed_batch + + # ---------- padded (BSHD) branch ---------- + assert seq_length is not None, "seq_length must be provided when use_packed_sequence is False" + padded_batch: Dict[str, Any] = {} + + if is_src: + assert batch is not None, "source TP rank must provide a batch" + max_seqlens = max(x["input_ids"].shape[0] for x in batch) + target_seqlens = min(max_seqlens, seq_length) + # Round target seqlen up to the parallelism alignment factor so the + # batched tensor is divisible for CP (+SP) splitting downstream. + if divisible_by > 1: + target_seqlens = math.ceil(target_seqlens / divisible_by) * divisible_by + + for sample in batch: + sample["input_ids"] = F.pad( + sample["input_ids"], (0, target_seqlens - sample["input_ids"].shape[0]), value=0 + ) + sample["labels"] = F.pad( + sample["labels"], (0, target_seqlens - sample["labels"].shape[0]), value=-100 + ) + sample["loss_mask"] = F.pad( + sample["loss_mask"], (0, target_seqlens - sample["loss_mask"].shape[0]), value=0 + ) + + padded_batch["input_ids"] = torch.concat( + [x["input_ids"].unsqueeze(0) for x in batch], dim=0 + ) + padded_batch["labels"] = torch.concat([x["labels"].unsqueeze(0) for x in batch], dim=0) + padded_batch["loss_mask"] = torch.concat( + [x["loss_mask"].unsqueeze(0) for x in batch], dim=0 + ) + padded_batch["pixel_values"] = torch.concat([x["pixel_values"] for x in batch]) + padded_batch["image_grid_thw"] = torch.concat([x["image_grid_thw"] for x in batch]) + + return broadcast_data_batch(padded_batch, device=device) + + +# ------------------------------------------------------------------- +# get_batch +# ------------------------------------------------------------------- + + +def get_batch(data_iterator: Iterator[list[Dict[str, Any]]]): + """Get a batch from *data_iterator* and broadcast across TP ranks.""" + device = "cuda" + args = get_args() + + if get_tensor_model_parallel_rank() == 0: + try: + data = next(data_iterator) + has_data = torch.tensor([1], dtype=torch.uint8, device=device) + except StopIteration: + has_data = torch.tensor([0], dtype=torch.uint8, device=device) + data = None + else: + has_data = torch.empty(1, dtype=torch.uint8, device=device) + data = None + + src = get_tensor_model_parallel_src_rank() + group = get_tensor_model_parallel_group() + torch.distributed.broadcast(has_data, src, group=group) + + if has_data.item() == 0: + return None + + # Because broadcast will not broadcast packed_seq_params, we move it into pack_or_pad_batch + batch = pack_or_pad_batch(data, args.use_packed_sequence, args.seq_length, device=device) + + # Fix shapes produced by default_collate. + if "position_ids" in batch and batch["position_ids"] is not None: + p = batch["position_ids"] + if p.dim() == 3 and p.shape[1] == 3: + batch["position_ids"] = p.permute(1, 0, 2).contiguous() + + if "pixel_values" in batch and batch["pixel_values"] is not None: + pv = batch["pixel_values"] + if pv.dim() == 3: + B, P, D = pv.shape + batch["pixel_values"] = pv.reshape(B * P, D) + + if "image_grid_thw" in batch and batch["image_grid_thw"] is not None: + g = batch["image_grid_thw"] + if g.dim() == 3: + batch["image_grid_thw"] = g.squeeze(1) + + return batch + + +# ------------------------------------------------------------------- +# Loss +# ------------------------------------------------------------------- + + +def loss_func(loss_mask, output_tensor): + """Compute masked language model loss.""" + losses = output_tensor.float() + loss_mask = loss_mask.contiguous().view(-1).float() + + total_tokens = loss_mask.sum().clone().detach().to(torch.int) + total_loss = torch.sum(losses.view(-1) * loss_mask) + reporting_loss = torch.cat([total_loss.clone().detach().view(1), total_tokens.view(1)]) + + return (total_loss, total_tokens, {"lm loss": reporting_loss}) + + +# ------------------------------------------------------------------- +# Forward step +# ------------------------------------------------------------------- + + +def forward_step(data_iterator, model, return_schedule_plan: bool = False): + """Forward step for multimodal_dev training. + + When ``return_schedule_plan=True`` (EP A2A overlap path via combined_1f1b), + delegate to ``model.build_schedule_plan(...)`` instead of running the + full forward — the inner schedule plan handles decoder layers. + """ + batch = get_batch(data_iterator) + + if batch is None: + return None, None + + pixel_values = batch.get("pixel_values", None) + if ( + pixel_values is not None + and pixel_values.is_floating_point() + and pixel_values.dtype == torch.float32 + ): + pixel_values = pixel_values.bfloat16() + + if return_schedule_plan: + from megatron.training import get_args as _get_args + + _args = _get_args() + assert _args.overlap_moe_expert_parallel_comm, ( + "overlap_moe_expert_parallel_comm must be enabled to return the schedule plan" + ) + schedule_plan = model.build_schedule_plan( + input_ids=batch["input_ids"], + position_ids=batch.get("position_ids"), + attention_mask=batch.get("attention_mask", None), + labels=batch.get("labels", None), + loss_mask=batch.get("loss_mask", None), + pixel_values=pixel_values, + image_grid_thw=batch.get("image_grid_thw", None), + packed_seq_params=batch.get("packed_seq_params", None), + ) + + loss_mask = batch.get("loss_mask", None) + if loss_mask is None: + loss_mask = torch.ones_like(batch["input_ids"], dtype=torch.float) + from examples.multimodal_dev.models.base import MultimodalModel + + loss_mask = MultimodalModel.cp_split_loss_mask( + loss_mask, batch.get("packed_seq_params", None) + ) + return schedule_plan, partial(loss_func, loss_mask) + + # We don't provide position_ids, now. Let model handle it itself. + output_tensor = model( + input_ids=batch["input_ids"], + position_ids=batch.get("position_ids"), + attention_mask=batch.get("attention_mask", None), + labels=batch.get("labels", None), + loss_mask=batch.get("loss_mask", None), + pixel_values=pixel_values, + image_grid_thw=batch.get("image_grid_thw", None), + packed_seq_params=batch.get("packed_seq_params", None), + ) + + loss_mask = batch.get("loss_mask", None) + if loss_mask is None: + loss_mask = torch.ones_like(batch["input_ids"], dtype=torch.float) + + # Slice loss_mask the same way the model sliced its inputs, so the + # mask aligns with the CP-shard output. Delegated to MultimodalModel + # so the slicing rule lives in one place. + from examples.multimodal_dev.models.base import MultimodalModel + + loss_mask = MultimodalModel.cp_split_loss_mask(loss_mask, batch.get("packed_seq_params", None)) + + return output_tensor, partial(loss_func, loss_mask) diff --git a/examples/multimodal_dev/models/__init__.py b/examples/multimodal_dev/models/__init__.py new file mode 100644 index 00000000000..e8ed05f1ca2 --- /dev/null +++ b/examples/multimodal_dev/models/__init__.py @@ -0,0 +1,62 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Model registry for multimodal_dev training. + +Maps ``--model-arch`` to a set of factory functions that fully encapsulate +model-specific logic. The training entry point (``pretrain_multimodal.py``) +remains model-agnostic — adding a new architecture only requires a new +registry entry (and its backing module) without touching the entry point. + +Registry entry fields +--------------------- +``model_factory_fn`` *(required)* + ``(args, language_config, vision_config, **kwargs) -> MegatronModule`` + Builds and returns the complete model instance. + +``vision_config_fn`` *(required)* + ``(num_layers_override=None, variant=None) -> TransformerConfig`` + Returns the vision encoder TransformerConfig. + +``post_language_config_fn`` *(optional)* + ``(language_config, args) -> None`` + Mutates the language TransformerConfig in-place with model-specific + fields (e.g. ``mrope_section``). + +``vision_flops_fn`` *(optional)* + ``(args, language_config, vision_config) -> None`` + Sets vision FLOPs metadata on ``args`` for training throughput logging. + +``dataset_providers`` *(optional)* + ``Dict[str, str | callable]`` + Maps ``--dataset-provider`` names to callables (or dotted import paths + resolved lazily) with signature + ``(train_val_test_num_samples) -> (train_ds, val_ds, test_ds)``. +""" + +from examples.multimodal_dev.models.qwen35_vl.configuration import get_qwen35_vl_vision_config +from examples.multimodal_dev.models.qwen35_vl.factory import build_model as _build_qwen35_vl_model +from examples.multimodal_dev.models.qwen35_vl.factory import ( + post_language_config as _qwen35_vl_post_language_config, +) +from examples.multimodal_dev.models.qwen35_vl.factory import ( + set_vision_flops_metadata as _qwen35_vl_vision_flops, +) + +MODEL_REGISTRY = { + "qwen35_vl": { + "model_factory_fn": _build_qwen35_vl_model, + "vision_config_fn": get_qwen35_vl_vision_config, + "post_language_config_fn": _qwen35_vl_post_language_config, + "vision_flops_fn": _qwen35_vl_vision_flops, + "dataset_providers": { + "mock": ( + "examples.multimodal_dev.data.mock" + ".train_valid_test_datasets_provider" + ), + "cord_v2": ( + "examples.multimodal_dev.data.vlm_dataset" + ".train_valid_test_datasets_provider" + ), + }, + }, +} diff --git a/examples/multimodal_dev/models/base.py b/examples/multimodal_dev/models/base.py new file mode 100644 index 00000000000..d0597f540ef --- /dev/null +++ b/examples/multimodal_dev/models/base.py @@ -0,0 +1,430 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Base multimodal model for FSDP + EP training. + +Composes a vision encoder and a ``GPTModel`` language decoder. Designed +for FSDP + EP: always builds the **full** model on every rank (no PP +flags). PP support is only available through the MIMO ``MimoModel`` +assembly path. + +Subclasses override ``compute_position_ids()`` for model-specific +position encoding (e.g. MRoPE for Qwen3.5-VL). +""" + +import contextlib +from typing import Optional + +import torch +from torch import Tensor + +from megatron.core import parallel_state, tensor_parallel +from megatron.core.models.gpt import GPTModel +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_config import TransformerConfig + + +def _cp_split_tensor(tensor, seq_dim, cp_size, cp_rank): + """Zigzag-split *tensor* along *seq_dim* for context parallelism (BSHD). + + Splits the sequence into ``2 * cp_size`` equal chunks, then selects + chunks ``[cp_rank, 2*cp_size - cp_rank - 1]`` and concatenates them. + This mirrors ``megatron.core.utils.get_batch_on_this_cp_rank``. + """ + S = tensor.shape[seq_dim] + assert S % (2 * cp_size) == 0, f"seq_len {S} not divisible by 2*cp_size={2 * cp_size}" + tensor = tensor.view( + *tensor.shape[:seq_dim], 2 * cp_size, S // (2 * cp_size), *tensor.shape[seq_dim + 1 :] + ) + index = torch.zeros(2, dtype=torch.int64, device=tensor.device) + index[0] = cp_rank + index[1] = 2 * cp_size - cp_rank - 1 + tensor = tensor.index_select(seq_dim, index) + tensor = tensor.view(*tensor.shape[:seq_dim], -1, *tensor.shape[seq_dim + 2 :]) + return tensor + + +class _NoCPGroup: + """Dummy size-1 process group used to bypass BSHD-style CP slicing + for THD MRoPE call sites that do not pass ``packed_seq=True``. + """ + + def size(self): + """Pretend this group has exactly one rank.""" + return 1 + + def rank(self): + """This rank's id within the fake group is always 0.""" + return 0 + + +_NO_CP_GROUP = _NoCPGroup() + +# Note: reported ``mtp_1 loss`` drifts ~1.3% from the CP=1 baseline under +# THD+CP. Megatron-Core's logging averages per-rank pre-divided ratios +# with op=AVG, and per-rank num_tokens are unequal after MTP rolling. +# Gradients are correct; only the *logged* value drifts. + + +def _thd_cp_partition_index(cu_seqlens_padded, total_tokens, cp_size, cp_rank): + """Per-rank token index for THD + CP via TE's + ``thd_get_partitioned_indices``. Cast to int64 so the result can be + used directly with ``index_select`` regardless of TE's return dtype. + """ + from transformer_engine.pytorch import cpp_extensions as tex + + idx = tex.thd_get_partitioned_indices(cu_seqlens_padded, total_tokens, cp_size, cp_rank) + return idx.long() + + +class MultimodalModel(MegatronModule): + """Base class for multimodal vision-language models. + + Composes a pre-constructed vision encoder and a ``GPTModel`` language + decoder. Designed for FSDP + EP; always builds the full model on + every rank. + + Args: + language_config: ``TransformerConfig`` for the language decoder. + language_spec: ``ModuleSpec`` for decoder transformer layers. + vision_encoder: Pre-constructed vision encoder module. + vocab_size: Language model vocabulary size. + max_sequence_length: Maximum sequence length. + image_token_id: Token ID for image placeholder tokens. + position_embedding_type: Position embedding type for the decoder. + rotary_percent: Fraction of hidden dim for RoPE. + rotary_base: Base frequency for RoPE. + mrope_section: MRoPE channel sections. + mtp_block_spec: Optional MTP block spec. + parallel_output: Keep outputs split across TP ranks. + share_embeddings_and_output_weights: Tie input/output embeddings. + """ + + def __init__( + self, + language_config: TransformerConfig, + language_spec: ModuleSpec, + vision_encoder: MegatronModule, + vocab_size: int, + max_sequence_length: int, + image_token_id: int, + position_embedding_type: str = "rope", + rotary_percent: float = 1.0, + rotary_base: int = 10000, + mrope_section: Optional[list] = None, + mtp_block_spec: Optional[ModuleSpec] = None, + parallel_output: bool = True, + share_embeddings_and_output_weights: bool = False, + ): + super().__init__(config=language_config) + + self.image_token_id = image_token_id + + self.vision_model = vision_encoder + self.language_model = GPTModel( + config=language_config, + transformer_layer_spec=language_spec, + vocab_size=vocab_size, + max_sequence_length=max_sequence_length, + pre_process=True, + post_process=True, + parallel_output=parallel_output, + share_embeddings_and_output_weights=(share_embeddings_and_output_weights), + position_embedding_type=position_embedding_type, + rotary_percent=rotary_percent, + rotary_base=rotary_base, + mtp_block_spec=mtp_block_spec, + ) + + def set_input_tensor(self, input_tensor): + """Route input tensors (simplified, no PP routing).""" + if not isinstance(input_tensor, list): + input_tensor = [input_tensor] + assert len(input_tensor) == 1 + self.language_model.set_input_tensor(input_tensor[0]) + + def build_schedule_plan( + self, + input_ids: Tensor, + position_ids: Tensor = None, + attention_mask: Tensor = None, + labels: Tensor = None, + loss_mask: Tensor = None, + pixel_values: Tensor = None, + image_grid_thw: Tensor = None, + packed_seq_params=None, + **kwargs, + ): + """Build a schedule plan for EP A2A overlap on the decoder only. + + Runs the vision encoder + embedding scatter eagerly (these stay on + the main path; not part of the overlap schedule), then delegates the + decoder transformer-layer schedule plan to the inner ``language_model``. + + Vision encoder is intentionally NOT included in the overlap schedule — + per design, only decoder layers participate in the EP A2A overlap. + """ + if position_ids is None: + position_ids = self.compute_position_ids( + input_ids=input_ids, + image_grid_thw=image_grid_thw, + packed_seq_params=packed_seq_params, + ) + + vision_embeddings = None + if self.vision_model is not None and pixel_values is not None: + vision_embeddings = self.vision_model(pixel_values, image_grid_thw) + + text_embeddings = self.language_model.embedding(input_ids=input_ids, position_ids=None) + if vision_embeddings is not None: + decoder_input = self._scatter_vision_embeddings( + input_ids, text_embeddings, vision_embeddings + ) + else: + decoder_input = text_embeddings + + (decoder_input, input_ids, labels, loss_mask, attention_mask, position_ids) = ( + self._cp_split_for_forward( + decoder_input=decoder_input, + input_ids=input_ids, + labels=labels, + loss_mask=loss_mask, + attention_mask=attention_mask, + position_ids=position_ids, + packed_seq_params=packed_seq_params, + ) + ) + + with self._thd_mrope_no_cp_override(packed_seq_params): + return self.language_model.build_schedule_plan( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + decoder_input=decoder_input, + labels=labels, + loss_mask=loss_mask, + packed_seq_params=packed_seq_params, + ) + + def _scatter_vision_embeddings( + self, input_ids: Tensor, text_embeddings: Tensor, vision_embeddings: Tensor + ) -> Tensor: + """Replace image-token positions with vision embeddings. + + Handles sequence parallelism (gather → scatter → re-scatter). + + Args: + input_ids: ``[B, S]`` token IDs. + text_embeddings: ``[S, B, D]`` (or ``[S/TP, B, D]`` with SP). + vision_embeddings: ``[num_visual_tokens, D]``. + + Returns: + Combined embeddings, same shape as *text_embeddings*. + """ + sp = ( + self.config.sequence_parallel + and parallel_state.get_tensor_model_parallel_world_size() > 1 + ) + + if sp: + text_embeddings = tensor_parallel.gather_from_sequence_parallel_region( + text_embeddings, tensor_parallel_output_grad=False + ) + + combined = text_embeddings.transpose(0, 1).contiguous() + image_mask = input_ids == self.image_token_id + mask_expanded = image_mask.unsqueeze(-1).expand_as(combined) + combined = combined.masked_scatter(mask_expanded, vision_embeddings) + combined = combined.transpose(0, 1).contiguous() + + if sp: + combined = tensor_parallel.scatter_to_sequence_parallel_region(combined) + + return combined + + def compute_position_ids( + self, input_ids: Tensor, image_grid_thw: Optional[Tensor] = None, packed_seq_params=None + ) -> Tensor: + """Compute position IDs. Override for MRoPE etc. + + Default: simple sequential positions. ``packed_seq_params`` is + accepted for subclass compatibility (e.g. MRoPE in THD mode). + """ + B, S = input_ids.shape + return torch.arange(S, device=input_ids.device).unsqueeze(0).expand(B, -1) + + def _cp_split_for_forward( + self, + *, + decoder_input, + input_ids, + labels, + loss_mask, + attention_mask, + position_ids, + packed_seq_params, + ): + """Apply CP split to model-forward inputs. + + BSHD path zigzag-splits each tensor along its seq dim. THD path + partitions per-sample via ``tex.thd_get_partitioned_indices`` so + chunks line up with ``cu_seqlens_q_padded`` boundaries. + ``position_ids`` and ``attention_mask`` are NOT split in THD — + MRoPE returns full freqs and TE attention's + ``_apply_rotary_pos_emb_thd`` does the per-sample CP zigzag + itself via ``_get_thd_freqs_on_this_cp_rank``. + """ + cp_size = parallel_state.get_context_parallel_world_size() + if cp_size <= 1: + return (decoder_input, input_ids, labels, loss_mask, attention_mask, position_ids) + cp_rank = parallel_state.get_context_parallel_rank() + + if packed_seq_params is not None: + total_tokens = ( + decoder_input.shape[0] if decoder_input is not None else input_ids.shape[1] + ) + idx = _thd_cp_partition_index( + packed_seq_params.cu_seqlens_q_padded, total_tokens, cp_size, cp_rank + ) + if decoder_input is not None: + decoder_input = decoder_input.index_select(0, idx) + if input_ids is not None: + input_ids = input_ids.index_select(1, idx) + if labels is not None: + labels = labels.index_select(1, idx) + if loss_mask is not None: + loss_mask = loss_mask.index_select(1, idx) + else: + + def _split(t, seq_dim): + return ( + None + if t is None + else _cp_split_tensor(t, seq_dim=seq_dim, cp_size=cp_size, cp_rank=cp_rank) + ) + + decoder_input = _split(decoder_input, 0) + input_ids = _split(input_ids, 1) + labels = _split(labels, 1) + loss_mask = _split(loss_mask, 1) + attention_mask = _split(attention_mask, 1) + + return (decoder_input, input_ids, labels, loss_mask, attention_mask, position_ids) + + @staticmethod + def cp_split_loss_mask(loss_mask, packed_seq_params): + """Slice ``loss_mask`` the same way the model slices its inputs. + + Mirrors the slicing done inside :meth:`_cp_split_for_forward` so + the loss computation outside the model can index a mask aligned + with the model's CP-shard output. Returns ``loss_mask`` unchanged + when ``CP <= 1``. + """ + cp_size = parallel_state.get_context_parallel_world_size() + if cp_size <= 1 or loss_mask is None: + return loss_mask + cp_rank = parallel_state.get_context_parallel_rank() + if packed_seq_params is not None: + idx = _thd_cp_partition_index( + packed_seq_params.cu_seqlens_q_padded, loss_mask.shape[1], cp_size, cp_rank + ) + return loss_mask.index_select(1, idx) + return _cp_split_tensor(loss_mask, seq_dim=1, cp_size=cp_size, cp_rank=cp_rank) + + @contextlib.contextmanager + def _thd_mrope_no_cp_override(self, packed_seq_params): + """Force ``rotary_pos_emb.cp_group`` to size 1 for the wrapped + forward call so MRoPE returns full-length freqs in THD mode. + Attention then applies per-sample CP zigzag itself via + ``_apply_rotary_pos_emb_thd``. Done by direct mutation rather + than via ``packed_seq_params.cp_group`` so MTP's CP-aware roll + (which reads that field) still sees the real CP group. + """ + mrope = ( + getattr(self.language_model, "rotary_pos_emb", None) + if packed_seq_params is not None + and parallel_state.get_context_parallel_world_size() > 1 + else None + ) + saved = getattr(mrope, "cp_group", None) if mrope is not None else None + if mrope is not None: + mrope.cp_group = _NO_CP_GROUP + try: + yield + finally: + if mrope is not None: + mrope.cp_group = saved + + def forward( + self, + input_ids: Tensor, + position_ids: Tensor, + attention_mask: Tensor = None, + labels: Tensor = None, + loss_mask: Tensor = None, + pixel_values: Tensor = None, + image_grid_thw: Tensor = None, + decoder_input: Tensor = None, + packed_seq_params=None, + **kwargs, + ): + """Forward pass. + + Args: + input_ids: ``[B, S]`` token IDs (or ``[1, T]`` in THD mode). + position_ids: ``[3, B, S]`` for MRoPE or ``[B, S]`` + (``[3, 1, T]`` / ``[1, T]`` in THD mode). + attention_mask: ``[B, S]`` attention mask (None in THD). + labels: ``[B, S]`` target token IDs (``[1, T]`` in THD). + loss_mask: ``[B, S]`` mask for loss (``[1, T]`` in THD). + pixel_values: Preprocessed image pixels. + image_grid_thw: ``[num_images, 3]`` grid dimensions. + decoder_input: Pre-computed decoder input (skip embed). + packed_seq_params: ``PackedSeqParams`` for THD attention. + + Returns: + Loss tensor (post_process=True) or hidden states. + """ + if position_ids is None: + position_ids = self.compute_position_ids( + input_ids=input_ids, + image_grid_thw=image_grid_thw, + packed_seq_params=packed_seq_params, + ) + + vision_embeddings = None + if self.vision_model is not None and pixel_values is not None: + vision_embeddings = self.vision_model(pixel_values, image_grid_thw) + + if decoder_input is None and self.language_model is not None: + text_embeddings = self.language_model.embedding(input_ids=input_ids, position_ids=None) + + if vision_embeddings is not None: + decoder_input = self._scatter_vision_embeddings( + input_ids, text_embeddings, vision_embeddings + ) + else: + decoder_input = text_embeddings + + (decoder_input, input_ids, labels, loss_mask, attention_mask, position_ids) = ( + self._cp_split_for_forward( + decoder_input=decoder_input, + input_ids=input_ids, + labels=labels, + loss_mask=loss_mask, + attention_mask=attention_mask, + position_ids=position_ids, + packed_seq_params=packed_seq_params, + ) + ) + + with self._thd_mrope_no_cp_override(packed_seq_params): + return self.language_model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + decoder_input=decoder_input, + labels=labels, + loss_mask=loss_mask, + packed_seq_params=packed_seq_params, + ) diff --git a/examples/multimodal_dev/models/qwen35_vl/__init__.py b/examples/multimodal_dev/models/qwen35_vl/__init__.py new file mode 100644 index 00000000000..1a0bad8b219 --- /dev/null +++ b/examples/multimodal_dev/models/qwen35_vl/__init__.py @@ -0,0 +1,70 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Qwen3.5-VL model components — the single source of truth. + +Both the standalone ``multimodal_dev`` training path and the MIMO path +import from here. +""" + +from examples.multimodal_dev.models.qwen35_vl.configuration import ( + MROPE_SECTION, + QWEN35_VL_IMAGE_TOKEN_ID, + QWEN35_VL_VIDEO_TOKEN_ID, + QWEN35_VL_VISION_END_TOKEN_ID, + QWEN35_VL_VISION_START_TOKEN_ID, + QWEN35_VL_VOCAB_SIZE, + ROTARY_BASE, + ROTARY_PERCENT, + VISION_KWARGS, + get_qwen35_vl_language_config, + get_qwen35_vl_vision_config, +) +from examples.multimodal_dev.models.qwen35_vl.factory import ( + build_model, + post_language_config, + set_vision_flops_metadata, +) +from examples.multimodal_dev.models.qwen35_vl.model import Qwen35VLModel +from examples.multimodal_dev.models.qwen35_vl.mrope import get_rope_index +from examples.multimodal_dev.models.qwen35_vl.specs import ( + get_qwen35_vl_language_spec, + get_qwen35_vl_vision_spec, +) +from examples.multimodal_dev.models.qwen35_vl.vision_encoder import ( + Qwen35VLPatchEmbed, + Qwen35VLPatchMerger, + Qwen35VLVisionEncoder, + Qwen35VLVisionRotaryEmbedding, +) + +__all__ = [ + # Model class + "Qwen35VLModel", + # Factory functions + "build_model", + "post_language_config", + "set_vision_flops_metadata", + # Vision encoder + "Qwen35VLVisionEncoder", + "Qwen35VLPatchEmbed", + "Qwen35VLPatchMerger", + "Qwen35VLVisionRotaryEmbedding", + # Config helpers + "get_qwen35_vl_vision_config", + "get_qwen35_vl_language_config", + # Spec helpers + "get_qwen35_vl_language_spec", + "get_qwen35_vl_vision_spec", + # MRoPE + "get_rope_index", + # Constants + "QWEN35_VL_IMAGE_TOKEN_ID", + "QWEN35_VL_VIDEO_TOKEN_ID", + "QWEN35_VL_VISION_START_TOKEN_ID", + "QWEN35_VL_VISION_END_TOKEN_ID", + "QWEN35_VL_VOCAB_SIZE", + "ROTARY_BASE", + "ROTARY_PERCENT", + "MROPE_SECTION", + "VISION_KWARGS", +] diff --git a/examples/multimodal_dev/models/qwen35_vl/configuration.py b/examples/multimodal_dev/models/qwen35_vl/configuration.py new file mode 100644 index 00000000000..8f65fe25a02 --- /dev/null +++ b/examples/multimodal_dev/models/qwen35_vl/configuration.py @@ -0,0 +1,368 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Configuration helpers for Qwen3.5-VL vision-language model. + +Provides TransformerConfig builders for the vision encoder and all language +decoder variants. Both the standalone ``multimodal_dev`` training path and the +MIMO path import from here — this is the single source of truth. + +Supported language variants (HuggingFace Qwen3.5 series): + ``0.8b`` Dense 0.8B + ``2b`` Dense 2B + ``4b`` Dense 4B + ``9b`` Dense 9B + ``27b`` Dense 27B + ``35b_a3b`` MoE 35B-A3B (256 experts, top-8) + ``122b_a10b`` MoE 122B-A10B (256 experts, top-8) + ``397b_a17b`` MoE 397B-A17B (512 experts, top-10) + ``35b_a3b_light`` Reduced 35B-A3B for testing + ``proxy`` Reduced proxy based on 397B for single-node testing +""" + +from typing import Optional + +import torch + +from megatron.core.transformer.transformer_config import TransformerConfig + +# --------------------------------------------------------------------------- +# Public constants +# --------------------------------------------------------------------------- + +QWEN35_VL_IMAGE_TOKEN_ID: int = 248056 +QWEN35_VL_VIDEO_TOKEN_ID: int = 248057 +QWEN35_VL_VISION_START_TOKEN_ID: int = 248053 +QWEN35_VL_VISION_END_TOKEN_ID: int = 248054 +QWEN35_VL_VOCAB_SIZE: int = 248320 + +ROTARY_BASE: int = 10_000_000 +ROTARY_PERCENT: float = 0.25 +MROPE_SECTION: list = [11, 11, 10] + +# --------------------------------------------------------------------------- +# Vision config +# --------------------------------------------------------------------------- + +VISION_KWARGS = { + "in_channels": 3, + "patch_size": 16, + "temporal_patch_size": 2, + "spatial_merge_size": 2, + "out_hidden_size": 3584, + "max_num_positions": 2304, +} + +# Three distinct vision encoder architectures in the Qwen3.5 family. +_VISION_SMALL = { + "num_layers": 12, "hidden_size": 768, "num_attention_heads": 12, + "kv_channels": 64, "ffn_hidden_size": 3072, +} +_VISION_MEDIUM = { + "num_layers": 24, "hidden_size": 1024, "num_attention_heads": 16, + "kv_channels": 64, "ffn_hidden_size": 4096, +} +_VISION_LARGE = { + "num_layers": 27, "hidden_size": 1152, "num_attention_heads": 16, + "kv_channels": 72, "ffn_hidden_size": 4304, +} + +# Per-variant vision config. ``out_hidden_size`` equals the language model's +# hidden_size and controls the merger projection output dimension. +_VISION_VARIANT_CONFIGS = { + "0.8b": {**_VISION_SMALL, "out_hidden_size": 1024}, + "2b": {**_VISION_MEDIUM, "out_hidden_size": 2048}, + "4b": {**_VISION_MEDIUM, "out_hidden_size": 2560}, + "9b": {**_VISION_LARGE, "out_hidden_size": 4096}, + "27b": {**_VISION_LARGE, "out_hidden_size": 5120}, + "35b_a3b": {**_VISION_LARGE, "out_hidden_size": 2048}, + "122b_a10b": {**_VISION_LARGE, "out_hidden_size": 3072}, + "397b_a17b": {**_VISION_LARGE, "out_hidden_size": 4096}, +} + +# Fallback for proxy/unknown variants (large ViT, generic out_hidden_size). +_VISION_DEFAULT = {**_VISION_LARGE, "out_hidden_size": 3584} + + +def get_qwen35_vl_vision_config( + num_layers_override: Optional[int] = None, + variant: Optional[str] = None, +) -> TransformerConfig: + """TransformerConfig for the Qwen3.5-VL vision encoder. + + Three ViT architectures are used across the family: + - Small (0.8b): depth 12, 768-dim, 12 heads + - Medium (2b, 4b): depth 24, 1024-dim, 16 heads + - Large (9b, 27b, MoE variants): depth 27, 1152-dim, 16 heads + + Args: + num_layers_override: Override vision backbone depth for proxy runs. + variant: Language model variant name. When set, selects the + matching vision config from ``_VISION_VARIANT_CONFIGS`` if one + exists; otherwise the default large-ViT config is used. + """ + vcfg = _VISION_VARIANT_CONFIGS.get(variant, _VISION_DEFAULT) + num_layers = vcfg["num_layers"] + if num_layers_override is not None: + num_layers = num_layers_override + + vision_head_dim = vcfg["kv_channels"] + assert vision_head_dim % 4 == 0, ( + "Qwen3.5-VL vision RoPE expects the per-head dimension to split " + f"evenly across row/column frequencies, got {vision_head_dim}" + ) + vision_rope_axis_dim = vision_head_dim // 4 + + return TransformerConfig( + num_layers=num_layers, + hidden_size=vcfg["hidden_size"], + num_attention_heads=vcfg["num_attention_heads"], + kv_channels=vcfg["kv_channels"], + ffn_hidden_size=vcfg["ffn_hidden_size"], + hidden_dropout=0.0, + attention_dropout=0.0, + layernorm_epsilon=1e-6, + normalization="LayerNorm", + gated_linear_unit=False, + activation_func=lambda x: torch.nn.functional.gelu(x, approximate="tanh"), + bias_activation_fusion=False, + apply_query_key_layer_scaling=False, + apply_rope_fusion=False, + # Vision RoPE is 2D row/column RoPE. Represent it as sectioned raw + # mRoPE with a zero temporal section so the fused mRoPE dispatcher can + # reuse the same Triton kernel when rope fusion is enabled. + mrope_section=[0, vision_rope_axis_dim, vision_rope_axis_dim], + mrope_interleaved=False, + rotary_interleaved=False, + bf16=False, + ) + + +# --------------------------------------------------------------------------- +# Language config variants +# --------------------------------------------------------------------------- + +_VARIANT_CONFIGS = { + "0.8b": { + "num_layers": 24, + "hidden_size": 1024, + "ffn_hidden_size": 3584, + "num_attention_heads": 8, + "num_query_groups": 2, + "kv_channels": 256, + "linear_num_value_heads": 16, + "num_moe_experts": None, + "moe_router_topk": None, + "moe_ffn_hidden_size": None, + "moe_shared_expert_intermediate_size": None, + }, + "2b": { + "num_layers": 24, + "hidden_size": 2048, + "ffn_hidden_size": 6144, + "num_attention_heads": 8, + "num_query_groups": 2, + "kv_channels": 256, + "linear_num_value_heads": 16, + "num_moe_experts": None, + "moe_router_topk": None, + "moe_ffn_hidden_size": None, + "moe_shared_expert_intermediate_size": None, + }, + "4b": { + "num_layers": 32, + "hidden_size": 2560, + "ffn_hidden_size": 9216, + "num_attention_heads": 16, + "num_query_groups": 4, + "kv_channels": 256, + "linear_num_value_heads": 32, + "num_moe_experts": None, + "moe_router_topk": None, + "moe_ffn_hidden_size": None, + "moe_shared_expert_intermediate_size": None, + }, + "9b": { + "num_layers": 32, + "hidden_size": 4096, + "ffn_hidden_size": 12288, + "num_attention_heads": 16, + "num_query_groups": 4, + "kv_channels": 256, + "linear_num_value_heads": 32, + "num_moe_experts": None, + "moe_router_topk": None, + "moe_ffn_hidden_size": None, + "moe_shared_expert_intermediate_size": None, + }, + "27b": { + "num_layers": 64, + "hidden_size": 5120, + "ffn_hidden_size": 17408, + "num_attention_heads": 24, + "num_query_groups": 4, + "kv_channels": 256, + "linear_num_value_heads": 48, + "num_moe_experts": None, + "moe_router_topk": None, + "moe_ffn_hidden_size": None, + "moe_shared_expert_intermediate_size": None, + }, + "35b_a3b": { + "num_layers": 40, + "hidden_size": 2048, + "ffn_hidden_size": 4096, + "num_attention_heads": 16, + "num_query_groups": 2, + "kv_channels": 256, + "linear_num_value_heads": 32, + "num_moe_experts": 256, + "moe_router_topk": 8, + "moe_ffn_hidden_size": 512, + "moe_shared_expert_intermediate_size": 512, + }, + "35b_a3b_light": { + "num_layers": 20, + "hidden_size": 2048, + "ffn_hidden_size": 4096, + "num_attention_heads": 16, + "num_query_groups": 2, + "kv_channels": 256, + "linear_num_value_heads": 32, + "num_moe_experts": 256, + "moe_router_topk": 8, + "moe_ffn_hidden_size": 512, + "moe_shared_expert_intermediate_size": 512, + }, + "122b_a10b": { + "num_layers": 48, + "hidden_size": 3072, + "ffn_hidden_size": 8192, + "num_attention_heads": 32, + "num_query_groups": 2, + "kv_channels": 256, + "linear_num_value_heads": 64, + "num_moe_experts": 256, + "moe_router_topk": 8, + "moe_ffn_hidden_size": 1024, + "moe_shared_expert_intermediate_size": 1024, + }, + "397b_a17b": { + "num_layers": 60, + "hidden_size": 4096, + "ffn_hidden_size": 10240, + "num_attention_heads": 32, + "num_query_groups": 2, + "kv_channels": 256, + "linear_num_value_heads": 64, + "num_moe_experts": 512, + "moe_router_topk": 10, + "moe_ffn_hidden_size": 1024, + "moe_shared_expert_intermediate_size": 1024, + }, + "proxy": { + "num_layers": 4, + "hidden_size": 4096, + "ffn_hidden_size": 10240, + "num_attention_heads": 32, + "num_query_groups": 2, + "kv_channels": 256, + "linear_num_value_heads": 64, + "num_moe_experts": 16, + "moe_router_topk": 2, + "moe_ffn_hidden_size": 1024, + "moe_shared_expert_intermediate_size": 1024, + }, +} + + +def get_qwen35_vl_language_config( + variant: str = "proxy", + **overrides, +) -> TransformerConfig: + """TransformerConfig for the Qwen3.5-VL language decoder. + + The ``397b_a17b`` variant reproduces the MIMO + ``get_qwen35_language_model_config()`` output exactly. + + Args: + variant: One of ``0.8b``, ``2b``, ``4b``, ``9b``, ``27b``, + ``35b_a3b``, ``122b_a10b``, ``397b_a17b``, + ``35b_a3b_light``, ``proxy``. + **overrides: Override any TransformerConfig field. + + Returns: + Fully-populated TransformerConfig. + """ + if variant not in _VARIANT_CONFIGS: + raise ValueError( + f"Unknown variant '{variant}'. " + f"Choose from {list(_VARIANT_CONFIGS.keys())}" + ) + + v = _VARIANT_CONFIGS[variant] + + kwargs = dict( + # Architecture + num_layers=v["num_layers"], + hidden_size=v["hidden_size"], + ffn_hidden_size=v["ffn_hidden_size"], + num_attention_heads=v["num_attention_heads"], + num_query_groups=v["num_query_groups"], + kv_channels=v["kv_channels"], + # Normalization & activation + normalization="RMSNorm", + layernorm_epsilon=1e-6, + layernorm_zero_centered_gamma=True, + apply_residual_connection_post_layernorm=False, + gated_linear_unit=True, + activation_func=torch.nn.functional.silu, + # MRoPE section (interleaved T/H/W layout, Qwen3.5-VL style) + mrope_section=list(MROPE_SECTION), + mrope_interleaved=True, + rotary_interleaved=False, + # Attention + qk_layernorm=True, + attention_output_gate=True, + attention_dropout=0.0, + hidden_dropout=0.0, + add_bias_linear=False, + # Hybrid attention (GatedDeltaNet) + experimental_attention_variant="gated_delta_net", + linear_attention_freq=4, + linear_conv_kernel_dim=4, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=v["linear_num_value_heads"], + # Kernel / TE fusions + bias_activation_fusion=True, + masked_softmax_fusion=True, + persist_layer_norm=True, + bias_dropout_fusion=True, + apply_rope_fusion=False, + # Precision + bf16=True, + ) + + # MoE config (only for MoE variants) + if v["num_moe_experts"] is not None: + kwargs.update( + num_moe_experts=v["num_moe_experts"], + moe_router_topk=v["moe_router_topk"], + moe_ffn_hidden_size=v["moe_ffn_hidden_size"], + moe_shared_expert_intermediate_size=v[ + "moe_shared_expert_intermediate_size" + ], + moe_shared_expert_gate=True, + moe_layer_freq=1, + moe_router_pre_softmax=False, + moe_router_load_balancing_type="global_aux_loss", + moe_permute_fusion=True, + moe_aux_loss_coeff=1e-3, + moe_grouped_gemm=True, + moe_token_dispatcher_type="alltoall", + moe_router_dtype="fp32", + ) + + kwargs.update(overrides) + return TransformerConfig(**kwargs) diff --git a/examples/multimodal_dev/models/qwen35_vl/factory.py b/examples/multimodal_dev/models/qwen35_vl/factory.py new file mode 100644 index 00000000000..3064bc5b7f4 --- /dev/null +++ b/examples/multimodal_dev/models/qwen35_vl/factory.py @@ -0,0 +1,101 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Factory functions for Qwen3.5-VL model construction. + +Encapsulates all Qwen3.5-VL-specific logic needed by ``pretrain_multimodal.py`` +so that the training entry point remains model-agnostic. +""" + +from examples.multimodal_dev.models.qwen35_vl.configuration import ( + MROPE_SECTION, + VISION_KWARGS, +) + + +def post_language_config(language_config, args): + """Apply Qwen3.5-VL-specific settings to the language TransformerConfig. + + Called after ``core_transformer_config_from_args`` to inject model-specific + fields that cannot be expressed via CLI args alone. + """ + language_config.mrope_section = list(MROPE_SECTION) + language_config.mrope_interleaved = True + + +def set_vision_flops_metadata(args, language_config, vision_config): + """Expose Qwen3.5-VL vision-model dimensions for FLOPs estimation.""" + args.count_vision_model_flops = True + args.vision_flops_variant = "qwen35_vl_v2" + args.vision_num_layers = vision_config.num_layers + args.vision_hidden_size = vision_config.hidden_size + args.vision_ffn_hidden_size = vision_config.ffn_hidden_size + args.vision_num_attention_heads = vision_config.num_attention_heads + args.vision_kv_channels = vision_config.kv_channels + args.vision_in_channels = VISION_KWARGS["in_channels"] + args.vision_patch_size = VISION_KWARGS["patch_size"] + args.vision_temporal_patch_size = VISION_KWARGS["temporal_patch_size"] + args.vision_spatial_merge_size = VISION_KWARGS["spatial_merge_size"] + args.vision_out_hidden_size = language_config.hidden_size + + +def build_model(args, language_config, vision_config, **kwargs): + """Build a complete Qwen3.5-VL model instance. + + Handles language spec construction, optional MTP block spec, and + model instantiation with Qwen3.5-VL-specific parameters. + + Args: + args: Megatron parsed arguments. + language_config: ``TransformerConfig`` for the language decoder + (already post-processed by :func:`post_language_config`). + vision_config: ``TransformerConfig`` for the vision encoder. + **kwargs: Extra keyword arguments (e.g. ``vp_stage``). + + Returns: + A :class:`Qwen35VLModel` instance. + """ + from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_mtp_block_spec, + ) + + from examples.multimodal_dev.models.qwen35_vl.model import Qwen35VLModel + from examples.multimodal_dev.models.qwen35_vl.specs import ( + get_qwen35_vl_language_spec, + ) + + language_spec = get_qwen35_vl_language_spec( + config=language_config, + vp_stage=kwargs.get("vp_stage", None), + pp_rank=None, + ) + + mtp_block_spec = None + if getattr(args, "mtp_num_layers", None): + mtp_block_spec = get_gpt_mtp_block_spec( + config=language_config, + spec=language_spec, + use_transformer_engine=( + args.transformer_impl == "transformer_engine" + ), + vp_stage=kwargs.get("vp_stage", None), + pp_rank=None, + ) + + # When --untie-embeddings-and-output-weights is NOT passed, Megatron + # defaults to tied embeddings (share_embeddings_and_output_weights=True). + # The 0.8B variant uses tied embeddings, while larger variants untie them. + share_embeddings = not getattr( + args, "untie_embeddings_and_output_weights", False + ) + + return Qwen35VLModel( + language_config=language_config, + language_spec=language_spec, + vision_config=vision_config, + vocab_size=args.padded_vocab_size, + max_sequence_length=args.max_position_embeddings, + image_token_id=getattr(args, "image_token_id", 248056), + mtp_block_spec=mtp_block_spec, + parallel_output=True, + share_embeddings_and_output_weights=share_embeddings, + ) diff --git a/examples/multimodal_dev/models/qwen35_vl/model.py b/examples/multimodal_dev/models/qwen35_vl/model.py new file mode 100644 index 00000000000..a8fdaf67d33 --- /dev/null +++ b/examples/multimodal_dev/models/qwen35_vl/model.py @@ -0,0 +1,129 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Qwen3.5-VL multimodal model for standalone FSDP + EP training. + +Composes a Megatron-native Qwen3.5 vision encoder with a ``GPTModel`` +language decoder using MRoPE and hybrid GatedDeltaNet / full-attention +layers. +""" + +from typing import Optional + +from torch import Tensor + +from examples.multimodal_dev.models.base import MultimodalModel +from examples.multimodal_dev.models.qwen35_vl.configuration import ( + QWEN35_VL_IMAGE_TOKEN_ID, + QWEN35_VL_VIDEO_TOKEN_ID, + QWEN35_VL_VISION_START_TOKEN_ID, + QWEN35_VL_VOCAB_SIZE, + ROTARY_BASE, + ROTARY_PERCENT, + VISION_KWARGS, +) +from examples.multimodal_dev.models.qwen35_vl.mrope import get_rope_index +from examples.multimodal_dev.models.qwen35_vl.specs import get_qwen35_vl_vision_spec +from examples.multimodal_dev.models.qwen35_vl.vision_encoder import Qwen35VLVisionEncoder +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_config import TransformerConfig + + +class Qwen35VLModel(MultimodalModel): + """Qwen3.5-VL multimodal model. + + Args: + language_config: ``TransformerConfig`` for the language decoder. + language_spec: ``ModuleSpec`` for language decoder layers. + vision_config: ``TransformerConfig`` for the vision encoder. + vision_spec: ``ModuleSpec`` for vision encoder layers. + vocab_size: Vocabulary size. + max_sequence_length: Maximum sequence length. + image_token_id: Token ID for image placeholders. + spatial_merge_size: Vision encoder spatial merge factor. + mtp_block_spec: Optional MTP block spec. + parallel_output: Keep outputs split across TP. + share_embeddings_and_output_weights: Tie embeddings. + """ + + def __init__( + self, + language_config: TransformerConfig, + language_spec: ModuleSpec, + vision_config: TransformerConfig, + vision_spec: ModuleSpec = None, + vocab_size: int = QWEN35_VL_VOCAB_SIZE, + max_sequence_length: int = 262144, + image_token_id: int = QWEN35_VL_IMAGE_TOKEN_ID, + video_token_id: int = QWEN35_VL_VIDEO_TOKEN_ID, + vision_start_token_id: int = QWEN35_VL_VISION_START_TOKEN_ID, + spatial_merge_size: int = 2, + mtp_block_spec: ModuleSpec = None, + parallel_output: bool = True, + share_embeddings_and_output_weights: bool = False, + ): + if vision_spec is None: + vision_spec = get_qwen35_vl_vision_spec() + + self.video_token_id = video_token_id + self.vision_start_token_id = vision_start_token_id + self.spatial_merge_size = spatial_merge_size + + vkw = dict(VISION_KWARGS) + vkw["spatial_merge_size"] = spatial_merge_size + vkw["out_hidden_size"] = language_config.hidden_size + + vision_encoder = Qwen35VLVisionEncoder( + config=vision_config, + transformer_layer_spec=vision_spec, + in_channels=vkw["in_channels"], + patch_size=vkw["patch_size"], + temporal_patch_size=vkw["temporal_patch_size"], + spatial_merge_size=vkw["spatial_merge_size"], + out_hidden_size=vkw["out_hidden_size"], + max_num_positions=vkw["max_num_positions"], + ) + + super().__init__( + language_config=language_config, + language_spec=language_spec, + vision_encoder=vision_encoder, + vocab_size=vocab_size, + max_sequence_length=max_sequence_length, + image_token_id=image_token_id, + position_embedding_type="mrope", + rotary_percent=ROTARY_PERCENT, + rotary_base=ROTARY_BASE, + mrope_section=language_config.mrope_section, + mtp_block_spec=mtp_block_spec, + parallel_output=parallel_output, + share_embeddings_and_output_weights=( + share_embeddings_and_output_weights + ), + ) + + def compute_position_ids( + self, + input_ids: Tensor, + image_grid_thw: Optional[Tensor] = None, + packed_seq_params=None, + ) -> Tensor: + """Compute 3D MRoPE position IDs for Qwen3.5-VL. + + In THD mode ``input_ids`` is ``[1, T]`` and ``packed_seq_params`` + supplies per-segment boundaries; positions restart at 0 per + segment. In BSHD mode ``input_ids`` is ``[B, S]`` and + ``packed_seq_params`` should be ``None``. + + Returns: + ``[3, B, S]`` position IDs for MRoPE (``[3, 1, T]`` in THD). + """ + position_ids, _ = get_rope_index( + spatial_merge_size=self.spatial_merge_size, + image_token_id=self.image_token_id, + video_token_id=self.video_token_id, + vision_start_token_id=self.vision_start_token_id, + input_ids=input_ids, + image_grid_thw=image_grid_thw, + packed_seq_params=packed_seq_params, + ) + return position_ids diff --git a/examples/multimodal_dev/models/qwen35_vl/mrope.py b/examples/multimodal_dev/models/qwen35_vl/mrope.py new file mode 100644 index 00000000000..763070929be --- /dev/null +++ b/examples/multimodal_dev/models/qwen35_vl/mrope.py @@ -0,0 +1,378 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""MRoPE (Multimodal Rotary Position Embedding) position ID computation. + +Computes 3D position IDs for Qwen3.5-VL: for text tokens all three +dimensions share sequential positions; for image/video tokens the three +dimensions encode (temporal, height, width) in the merged spatial grid. + +Supports two input layouts: + +* **BSHD** — ``input_ids`` is ``[B, S]``; each row is an independent + sample (possibly padded) and ``attention_mask`` marks valid tokens. +* **THD** — ``input_ids`` is ``[1, T]``, a concatenation of ``N`` + sub-sequences. ``packed_seq_params.cu_seqlens_q_padded`` gives the + physical segment boundaries in the packed tensor and + ``cu_seqlens_q`` gives the valid (unpadded) token count inside each + segment. Position IDs restart at 0 at every segment boundary; image + / video grid rows are consumed in packed order across segments. + +Ported from Megatron-Bridge ``get_rope_index`` (which itself is adapted +from HF ``Qwen3VLForConditionalGeneration.get_rope_index``). The inner +loop iterates over vision occurrences, not individual tokens. +""" + +from typing import Optional + +import torch +from torch import Tensor + +from megatron.core.packed_seq_params import PackedSeqParams + + +def _build_sample_mrope_positions( + sample_input_ids: Tensor, + image_grid_thw: Optional[Tensor], + video_grid_thw: Optional[Tensor], + image_index: int, + video_index: int, + spatial_merge_size: int, + image_token_id: int, + video_token_id: int, + vision_start_token_id: int, +) -> tuple[Tensor, int, int]: + """Compute MRoPE position IDs for a single sub-sequence. + + Walks vision occurrences in ``sample_input_ids`` and produces a + ``[3, L]`` position tensor whose values start at 0. Advances + ``image_index`` / ``video_index`` through ``image_grid_thw`` / + ``video_grid_thw`` so callers can keep a running cursor across + multiple sub-sequences. + """ + vision_start_indices = torch.argwhere( + sample_input_ids == vision_start_token_id, + ).squeeze(1) + vision_tokens = sample_input_ids[vision_start_indices + 1] + image_nums = int((vision_tokens == image_token_id).sum()) + video_nums = int((vision_tokens == video_token_id).sum()) + input_tokens = sample_input_ids.tolist() + llm_pos_ids_list: list = [] + st = 0 + remain_images, remain_videos = image_nums, video_nums + + for _ in range(image_nums + video_nums): + if image_token_id in input_tokens and remain_images > 0: + ed_image = input_tokens.index(image_token_id, st) + else: + ed_image = len(input_tokens) + 1 + if video_token_id in input_tokens and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + + if ed_image < ed_video: + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + image_index += 1 + remain_images -= 1 + ed = ed_image + else: + t, h, w = ( + video_grid_thw[video_index][0], + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + video_index += 1 + remain_videos -= 1 + ed = ed_video + + llm_grid_t, llm_grid_h, llm_grid_w = ( + t.item(), + h.item() // spatial_merge_size, + w.item() // spatial_merge_size, + ) + text_len = ed - st + + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if llm_pos_ids_list + else 0 + ) + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + + st_idx + ) + + t_index = ( + torch.arange(llm_grid_t) + .view(-1, 1) + .expand(-1, llm_grid_h * llm_grid_w) + .flatten() + ) + h_index = ( + torch.arange(llm_grid_h) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + + text_len + + st_idx + ) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + + if st < len(input_tokens): + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if llm_pos_ids_list + else 0 + ) + text_len = len(input_tokens) - st + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + + st_idx + ) + + if llm_pos_ids_list: + positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + else: + positions = torch.zeros( + 3, 0, + dtype=sample_input_ids.dtype, + device=sample_input_ids.device, + ) + return positions, image_index, video_index + + +def get_rope_index( + spatial_merge_size: int, + image_token_id: int, + video_token_id: int, + vision_start_token_id: int, + input_ids: Optional[Tensor] = None, + image_grid_thw: Optional[Tensor] = None, + video_grid_thw: Optional[Tensor] = None, + attention_mask: Optional[Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, +) -> tuple[Tensor, Tensor]: + """Compute 3D MRoPE position IDs for Qwen3-VL / Qwen3.5-VL. + + Qwen3-VL uses timestamps rather than absolute time position IDs. + + For text tokens all three dimensions share sequential positions. + For vision tokens the three dimensions encode (temporal, height, + width) in the merged spatial grid. + + Args: + spatial_merge_size: Merge factor for spatial dimensions. + image_token_id: Token ID for image placeholders. + video_token_id: Token ID for video placeholders. + vision_start_token_id: Token ID marking start of a vision region. + input_ids: ``[B, S]`` in BSHD or ``[1, T]`` in THD. + image_grid_thw: ``[num_images, 3]`` per-image + ``(temporal, height, width)`` in patch-grid units. Rows are + consumed in the order their image tokens appear in + ``input_ids`` (packed order across segments in THD). + video_grid_thw: ``[num_videos, 3]`` per-video grid dimensions. + attention_mask: ``[B, S]`` mask (1 = keep, 0 = pad). BSHD only. + packed_seq_params: When provided, selects the THD branch and + supplies segment boundaries via ``cu_seqlens_q`` (valid + lengths) and ``cu_seqlens_q_padded`` (packed layout). + + Returns: + ``(position_ids, mrope_position_deltas)`` where *position_ids* + has shape ``[3, B, S]`` (``[3, 1, T]`` in THD). + """ + if video_grid_thw is not None: + video_grid_thw = torch.repeat_interleave( + video_grid_thw, video_grid_thw[:, 0], dim=0, + ) + video_grid_thw[:, 0] = 1 + + # ----------------------------------------------------------------- + # THD (packed) branch + # ----------------------------------------------------------------- + if packed_seq_params is not None and input_ids is not None: + cu_seqlens = packed_seq_params.cu_seqlens_q + cu_seqlens_padded = getattr( + packed_seq_params, "cu_seqlens_q_padded", None, + ) + if cu_seqlens_padded is None: + cu_seqlens_padded = cu_seqlens + + assert ( + input_ids.dim() == 2 and input_ids.shape[0] == 1 + ), "THD get_rope_index expects input_ids shape [1, T]" + + total_tokens = input_ids.shape[1] + device = input_ids.device + + # Padding slots default to 1 (matches BSHD convention where + # masked positions get filled with 1). + position_ids = torch.ones( + 3, 1, total_tokens, + dtype=input_ids.dtype, device=device, + ) + deltas: list = [] + image_index = 0 + video_index = 0 + num_segs = cu_seqlens.numel() - 1 + + for k in range(num_segs): + seg_start = int(cu_seqlens_padded[k].item()) + valid_len = int( + cu_seqlens[k + 1].item() - cu_seqlens[k].item() + ) + valid_end = seg_start + valid_len + + if valid_len == 0: + deltas.append(0) + continue + + sample_input_ids = input_ids[0, seg_start:valid_end] + + if ( + image_grid_thw is not None + or video_grid_thw is not None + ): + ( + positions, + image_index, + video_index, + ) = _build_sample_mrope_positions( + sample_input_ids=sample_input_ids, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + image_index=image_index, + video_index=video_index, + spatial_merge_size=spatial_merge_size, + image_token_id=image_token_id, + video_token_id=video_token_id, + vision_start_token_id=vision_start_token_id, + ) + else: + positions = ( + torch.arange(valid_len, device=device) + .view(1, -1) + .expand(3, -1) + ) + + position_ids[:, 0, seg_start:valid_end] = positions.to( + device=device, dtype=position_ids.dtype, + ) + + if positions.numel() > 0: + deltas.append( + int(positions.max().item()) + 1 - valid_len + ) + else: + deltas.append(0) + + mrope_position_deltas = torch.tensor( + deltas, device=device, + ).unsqueeze(1) + return position_ids, mrope_position_deltas + + # ----------------------------------------------------------------- + # BSHD branch with vision + # ----------------------------------------------------------------- + if input_ids is not None and ( + image_grid_thw is not None or video_grid_thw is not None + ): + total_input_ids = input_ids + if attention_mask is None: + attention_mask = torch.ones_like(total_input_ids) + elif attention_mask.dim() > 2: + attention_mask = attention_mask.any(dim=-1) + if attention_mask.dim() == 3: + attention_mask = attention_mask.squeeze(1) + attention_mask = attention_mask.to(dtype=total_input_ids.dtype) + + position_ids = torch.ones( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + mrope_position_deltas = [] + image_index, video_index = 0, 0 + attention_mask = attention_mask.to(total_input_ids.device) + + for i, sample_input_ids in enumerate(total_input_ids): + sample_input_ids = sample_input_ids[attention_mask[i] == 1] + ( + llm_positions, + image_index, + video_index, + ) = _build_sample_mrope_positions( + sample_input_ids=sample_input_ids, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + image_index=image_index, + video_index=video_index, + spatial_merge_size=spatial_merge_size, + image_token_id=image_token_id, + video_token_id=video_token_id, + vision_start_token_id=vision_start_token_id, + ) + position_ids[ + ..., i, attention_mask[i] == 1 + ] = llm_positions.to(position_ids.device) + mrope_position_deltas.append( + llm_positions.max() + 1 - len(total_input_ids[i]), + ) + + mrope_position_deltas = torch.tensor( + mrope_position_deltas, device=total_input_ids.device, + ).unsqueeze(1) + return position_ids, mrope_position_deltas + + # ----------------------------------------------------------------- + # Text-only fallback + # ----------------------------------------------------------------- + if attention_mask is not None: + if attention_mask.dim() > 2: + attention_mask = attention_mask.any(dim=-1) + if attention_mask.dim() == 3: + attention_mask = attention_mask.squeeze(1) + attention_mask = attention_mask.to(dtype=torch.long) + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + position_ids = ( + position_ids.unsqueeze(0) + .expand(3, -1, -1) + .to(attention_mask.device) + ) + max_position_ids = ( + position_ids.max(0, keepdim=False)[0] + .max(-1, keepdim=True)[0] + ) + mrope_position_deltas = ( + max_position_ids + 1 - attention_mask.shape[-1] + ) + else: + position_ids = ( + torch.arange( + input_ids.shape[1], device=input_ids.device, + ) + .view(1, 1, -1) + .expand(3, input_ids.shape[0], -1) + ) + mrope_position_deltas = torch.zeros( + [input_ids.shape[0], 1], + device=input_ids.device, + dtype=input_ids.dtype, + ) + + return position_ids, mrope_position_deltas diff --git a/examples/multimodal_dev/models/qwen35_vl/specs.py b/examples/multimodal_dev/models/qwen35_vl/specs.py new file mode 100644 index 00000000000..eac6d543a04 --- /dev/null +++ b/examples/multimodal_dev/models/qwen35_vl/specs.py @@ -0,0 +1,162 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Layer spec helpers for Qwen3.5-VL vision encoder and language decoder. + +Provides ModuleSpec builders that define the transformer layer composition. +Both the standalone and MIMO training paths import from here. +""" + +from typing import Optional + +from examples.multimodal_dev.models.base import _NO_CP_GROUP +from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_transformer_block_with_experimental_attention_variant_spec, +) +from megatron.core.models.vision.vit_layer_specs import get_vit_layer_with_transformer_engine_spec +from megatron.core.transformer.attention import SelfAttention +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_block import TransformerBlockSubmodules +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import nvtx_range_pop, nvtx_range_push + + +def _apply_rope_fp32(t, freqs, config, cu_seqlens=None, mscale=1.0, cp_group=None): + """Apply rotary positional embedding in fp32, then cast back to original dtype. + + Mirrors ``Qwen3VLSelfAttention.apply_rotary_pos_emb_absolute`` in Megatron-Bridge + with ``apply_rotary_pos_emb_in_fp32=True``. + """ + from megatron.core.models.common.embeddings import rope_utils + from megatron.core.models.common.embeddings.rope_utils import apply_rotary_pos_emb + + orig_dtype = t.dtype + if ( + cu_seqlens is not None + and getattr(config, "apply_rope_fusion", False) + and getattr(config, "mrope_section", None) is not None + and getattr(config, "rotary_interleaved", False) is False + and getattr(config, "multi_latent_attention", False) is False + and mscale == 1.0 + and t.dim() == 3 + and freqs.dim() == 4 + and freqs.shape[0] == 3 + and cp_group is not None + and rope_utils.fused_apply_mrope_thd is not None + and rope_utils.get_fused_mrope_thd_unavailable_reason is not None + ): + unavailable_reason = rope_utils.get_fused_mrope_thd_unavailable_reason( + t, + cu_seqlens, + freqs, + rotary_interleaved=config.rotary_interleaved, + cp_size=cp_group.size(), + cp_rank=cp_group.rank(), + ) + if unavailable_reason is None: + return rope_utils.fused_apply_mrope_thd( + t, + cu_seqlens, + freqs, + config.mrope_section, + interleaved_mrope=config.mrope_interleaved, + rotary_interleaved=config.rotary_interleaved, + cp_size=cp_group.size(), + cp_rank=cp_group.rank(), + fp32_compute=True, + ) + + t_fp32 = t.float() + out = apply_rotary_pos_emb( + t_fp32, + freqs, + config=config, + cu_seqlens=cu_seqlens, + mscale=mscale, + cp_group=cp_group, + mla_rotary_interleaved=getattr(config, 'multi_latent_attention', False), + ) + return out.to(orig_dtype) + + +def _apply_rope_fp32_no_cp(t, freqs, config, cu_seqlens=None, mscale=1.0, cp_group=None): + """Same as ``_apply_rope_fp32`` but forces CP-size=1. + + The vision encoder uses THD packed sequences for variable-resolution + images. When the language model uses CP>1, the global CP group would + incorrectly split the vision seqlens. This wrapper substitutes a + trivial group so the vision RoPE sees the full packed sequence. + """ + range_name = "qwen35_vl.vision_encoder.rope_apply" + nvtx_range_push(range_name) + try: + return _apply_rope_fp32( + t, + freqs, + config, + cu_seqlens, + mscale, + cp_group=_NO_CP_GROUP, + ) + finally: + nvtx_range_pop(range_name) + + +class Qwen35VLVisionSelfAttention(SelfAttention): + """ViT self-attention with RoPE applied in fp32. + + Matches Bridge's ``Qwen3VLSelfAttention`` behaviour when + ``apply_rotary_pos_emb_in_fp32=True``: query and key are cast to float32 + before the rotary multiply and cast back to bf16 afterwards. The + monkey-patch approach avoids duplicating the 300-line ``SelfAttention.forward`` + while keeping the change local to this class. + """ + + def forward(self, *args, **kwargs): + import megatron.core.transformer.attention as _attn_mod + + _orig = _attn_mod.apply_rotary_pos_emb + _attn_mod.apply_rotary_pos_emb = _apply_rope_fp32_no_cp + try: + return super().forward(*args, **kwargs) + finally: + _attn_mod.apply_rotary_pos_emb = _orig + + +def get_qwen35_vl_language_spec( + config: TransformerConfig, + vp_stage: Optional[int] = None, + pp_rank: Optional[int] = None, +) -> TransformerBlockSubmodules: + """Transformer block spec for the Qwen3.5-VL language decoder. + + Uses the experimental attention variant infrastructure to build hybrid + GatedDeltaNet + full-attention layers with optional MoE interleaving. + + Args: + config: Language decoder TransformerConfig. + vp_stage: Virtual pipeline stage. + pp_rank: Pipeline parallel rank. + + Returns: + TransformerBlockSubmodules with per-layer specs. + """ + return get_transformer_block_with_experimental_attention_variant_spec( + config=config, + vp_stage=vp_stage, + pp_rank=pp_rank, + ) + + +def get_qwen35_vl_vision_spec() -> ModuleSpec: + """ModuleSpec for vision encoder transformer layers. + + Uses ``TEDotProductAttention`` which supports packed-sequence (THD) + attention via ``PackedSeqParams`` for variable-length images. + + ``Qwen35VLVisionSelfAttention`` replaces the default ``SelfAttention`` so + that RoPE is applied in fp32, matching Bridge's + ``apply_rotary_pos_emb_in_fp32=True`` behaviour. + """ + spec = get_vit_layer_with_transformer_engine_spec() + spec.submodules.self_attention.module = Qwen35VLVisionSelfAttention + return spec diff --git a/examples/multimodal_dev/models/qwen35_vl/vision_encoder.py b/examples/multimodal_dev/models/qwen35_vl/vision_encoder.py new file mode 100644 index 00000000000..1dc00221141 --- /dev/null +++ b/examples/multimodal_dev/models/qwen35_vl/vision_encoder.py @@ -0,0 +1,609 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Megatron-native Qwen3.5-VL vision encoder. + +Architecture (matches HF ``Qwen3VLVisionModel`` exactly): + + PatchEmbed (Conv3d) + → learned position embedding (bilinear interpolation) + → 2D Vision RoPE + → TransformerBlock × N (with PackedSeqParams / THD attention) + → PatchMerger (per-token LN → spatial merge → MLP) + +Key design choices: + * ``Conv3d`` patch embedding is replicated across TP ranks (no MCore + equivalent for 3D convolutions). + * ``PatchMerger`` MLP uses ``ColumnParallelLinear`` / ``RowParallelLinear`` + for TP sharding. + * Inherits from ``VisionModule``. + * Expects pixel values in block-merge order (as produced by the HF + processor) so the merger's simple reshape is correct. +""" + +from typing import List, Optional + +import torch +import torch.nn.functional as F +from torch import Tensor + +from megatron.core.extensions.transformer_engine import TENorm +from megatron.core.models.common.vision_module.vision_module import VisionModule +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_block import TransformerBlock +from megatron.core.transformer.transformer_config import TransformerConfig + +# ------------------------------------------------------------------- +# PatchEmbed — Conv3d (replicated, no TP sharding) +# ------------------------------------------------------------------- + +class Qwen35VLPatchEmbed(MegatronModule): + """3D convolution patch embedding matching HF ``Qwen3VLVisionPatchEmbed``. + + Uses ``nn.Conv3d`` with kernel = stride = ``[temporal_patch_size, + patch_size, patch_size]`` and ``bias=True``. The module is replicated + across TP ranks (no MCore equivalent for 3D conv). + + Args: + config: TransformerConfig (used by MegatronModule base). + in_channels: Number of input channels (3 for RGB). + hidden_size: Output embedding dimension. + patch_size: Spatial patch size. + temporal_patch_size: Temporal patch size. + """ + + def __init__( + self, + config: TransformerConfig, + in_channels: int = 3, + hidden_size: int = 1152, + patch_size: int = 16, + temporal_patch_size: int = 2, + ): + super().__init__(config=config) + self.patch_size = patch_size + self.temporal_patch_size = temporal_patch_size + self.in_channels = in_channels + self.hidden_size = hidden_size + + kernel = [temporal_patch_size, patch_size, patch_size] + self.proj = torch.nn.Conv3d( + in_channels, + hidden_size, + kernel_size=kernel, + stride=kernel, + bias=True, + ) + + def forward(self, pixel_values: Tensor) -> Tensor: + """Forward pass. + + Args: + pixel_values: ``[total_patches, C * T * pH * pW]`` + pre-extracted flat patches. + + Returns: + Patch embeddings ``[total_patches, hidden_size]``. + """ + target_dtype = self.proj.weight.dtype + pixel_values = pixel_values.view( + -1, + self.in_channels, + self.temporal_patch_size, + self.patch_size, + self.patch_size, + ) + return self.proj(pixel_values.to(dtype=target_dtype)).view( + -1, self.hidden_size + ) + + +# ------------------------------------------------------------------- +# VisionRotaryEmbedding — 1D frequency table +# ------------------------------------------------------------------- + +class Qwen35VLVisionRotaryEmbedding(MegatronModule): + """1D rotary position frequency table for the vision transformer. + + Generates RoPE frequencies for integer positions ``0 .. seqlen-1``. + The encoder maps 2D (row, col) positions to embeddings via table + lookup. Matches HF ``Qwen3VLVisionRotaryEmbedding``. + + Args: + dim: Frequency dimension (``head_dim // 2``). + theta: RoPE base frequency. + config: Optional TransformerConfig for MegatronModule base. + """ + + def __init__( + self, + dim: int, + theta: float = 10000.0, + config: Optional[TransformerConfig] = None, + ): + super().__init__(config=config) + self.dim = dim + self.theta = theta + inv_freq = 1.0 / ( + theta + ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def _get_inv_freq(self, device: torch.device) -> Tensor: + """Return ``inv_freq`` in float32 on *device*. + + Always recomputes in float32 regardless of the buffer's stored dtype. + This matches Bridge's lazy-init behaviour where ``inv_freq`` is + constructed fresh (in float32) on the first forward call, after any + ``model.bfloat16()`` cast has already occurred. + """ + return 1.0 / ( + self.theta + ** ( + torch.arange( + 0, self.dim, 2, + dtype=torch.float32, device=device, + ) + / self.dim + ) + ) + + def forward( + self, + seqlen: int, + device: Optional[torch.device] = None, + ) -> Tensor: + """Frequency lookup table for positions ``0 .. seqlen-1``. + + Args: + seqlen: Number of positions. + device: Runtime device (required for meta-init safety). + + Returns: + ``[seqlen, dim // 2]`` frequencies. + """ + if device is None: + if self.inv_freq.device.type != "meta": + device = self.inv_freq.device + else: + device = torch.device( + "cuda", torch.cuda.current_device() + ) + inv_freq = self._get_inv_freq(device) + seq = torch.arange(seqlen, device=device, dtype=inv_freq.dtype) + return torch.outer(seq, inv_freq) + + +# ------------------------------------------------------------------- +# PatchMerger — per-token LN, spatial merge, TP-sharded MLP +# ------------------------------------------------------------------- + +class Qwen35VLPatchMerger(MegatronModule): + """Spatial patch merger matching HF ``Qwen3VLVisionPatchMerger``. + + Per-token ``LayerNorm`` on ``hidden_size`` → reshape to merge + ``spatial_merge_size ** 2`` adjacent patches → two-layer MLP + (``ColumnParallelLinear`` → GELU → ``RowParallelLinear``). + + MLP dimensions: ``merge_dim → merge_dim → out_hidden_size`` + where ``merge_dim = hidden_size * spatial_merge_size ** 2``. + + Args: + config: TransformerConfig (provides TP settings, init_method). + hidden_size: Per-token hidden size from the ViT. + out_hidden_size: Output dimension (language model hidden_size). + spatial_merge_size: Merge factor per spatial dimension. + """ + + def __init__( + self, + config: TransformerConfig, + hidden_size: int = 1152, + out_hidden_size: int = 3584, + spatial_merge_size: int = 2, + ): + super().__init__(config=config) + self.spatial_merge_size = spatial_merge_size + self.merge_dim = hidden_size * (spatial_merge_size ** 2) + merge_dim = self.merge_dim + + self.patch_norm = TENorm(config=config, hidden_size=hidden_size, eps=1e-6) + self.linear_fc1 = build_module( + ColumnParallelLinear, + merge_dim, + merge_dim, + config=config, + init_method=config.init_method, + bias=True, + gather_output=False, + ) + self.linear_fc2 = build_module( + RowParallelLinear, + merge_dim, + out_hidden_size, + config=config, + init_method=config.output_layer_init_method, + bias=True, + input_is_parallel=True, + skip_bias_add=False, + ) + + def forward(self, hidden_states: Tensor) -> Tensor: + """Merge patches spatially. + + Args: + hidden_states: ``[total_patches, hidden_size]`` in block-merge + order from the ViT transformer blocks. + + Returns: + ``[total_merged_patches, out_hidden_size]``. + """ + hidden_states = self.patch_norm(hidden_states) + merged = hidden_states.view(-1, self.merge_dim) + merged, _ = self.linear_fc1(merged) + # NOTE: Official HuggingFace uses default approximate='none' in Qwen3VLVisionPatchMerger. + merged = torch.nn.functional.gelu(merged, approximate="tanh") + merged, _ = self.linear_fc2(merged) + return merged + + +# ------------------------------------------------------------------- +# Qwen35VLVisionEncoder — top-level encoder module +# ------------------------------------------------------------------- + +class Qwen35VLVisionEncoder(VisionModule): + """Megatron-native Qwen3.5-VL vision encoder. + + Processes image / video inputs through: + + 1. ``Qwen35VLPatchEmbed`` (Conv3d) + 2. Learned ``nn.Embedding`` position table with bilinear interpolation + 3. 2D Vision RoPE from ``(row, col)`` patch positions + 4. ``TransformerBlock`` × N with ``PackedSeqParams`` (THD attention) + 5. ``Qwen35VLPatchMerger`` + + Output dimension matches the language model ``hidden_size``. + + Args: + config: Vision ``TransformerConfig``. + transformer_layer_spec: ``ModuleSpec`` for ViT layers. + in_channels: Image channels (3 for RGB). + patch_size: Spatial patch size. + temporal_patch_size: Temporal patch size. + spatial_merge_size: Spatial merge factor. + out_hidden_size: Output dim (language decoder hidden_size). + max_num_positions: Size of the learned position table. + """ + + def __init__( + self, + config: TransformerConfig, + transformer_layer_spec: ModuleSpec = None, + in_channels: int = 3, + patch_size: int = 16, + temporal_patch_size: int = 2, + spatial_merge_size: int = 2, + out_hidden_size: int = 3584, + max_num_positions: int = 2304, + ): + super().__init__(config=config) + + self.hidden_size = config.hidden_size + self.spatial_merge_size = spatial_merge_size + + # --- Patch embedding (Conv3d) --- + self.patch_embed = Qwen35VLPatchEmbed( + config=config, + in_channels=in_channels, + hidden_size=config.hidden_size, + patch_size=patch_size, + temporal_patch_size=temporal_patch_size, + ) + + # --- Learned position embedding with bilinear interpolation --- + self.pos_embed = torch.nn.Embedding( + max_num_positions, config.hidden_size, + ) + self.num_grid_per_side = int(max_num_positions ** 0.5) + + # --- Vision rotary embeddings --- + head_dim = config.hidden_size // config.num_attention_heads + self.rot_pos_emb = Qwen35VLVisionRotaryEmbedding( + head_dim // 2, config=config, + ) + + # --- Transformer blocks --- + if transformer_layer_spec is None: + from examples.multimodal_dev.models.qwen35_vl.specs import get_qwen35_vl_vision_spec + transformer_layer_spec = get_qwen35_vl_vision_spec() + + self.decoder = TransformerBlock( + config=config, + spec=transformer_layer_spec, + pre_process=True, + post_process=True, + post_layer_norm=False, + ) + + # --- Patch merger --- + self.merger = Qwen35VLPatchMerger( + config=config, + hidden_size=config.hidden_size, + out_hidden_size=out_hidden_size, + spatial_merge_size=spatial_merge_size, + ) + + # --------------------------------------------------------------- + # Learned position embedding with bilinear interpolation + # --------------------------------------------------------------- + + def _fast_pos_embed_interpolate( + self, grid_thw: Tensor, + ) -> Tensor: + """Bilinear interpolation of the learned 2D position table. + + Matches HF ``Qwen3VLVisionModel.fast_pos_embed_interpolate``. + + Args: + grid_thw: ``[num_images, 3]`` (T, H, W) in patch-grid units. + + Returns: + ``[total_patches, hidden_size]`` position embeddings in + block-merge order. + """ + grid_thw_list = grid_thw.tolist() + grid_ts = [int(row[0]) for row in grid_thw_list] + grid_hs = [int(row[1]) for row in grid_thw_list] + grid_ws = [int(row[2]) for row in grid_thw_list] + device = self.pos_embed.weight.device + n = self.num_grid_per_side + + idx_list: List[List[int]] = [[] for _ in range(4)] + weight_list: List[List[float]] = [[] for _ in range(4)] + + for t, h, w in grid_thw_list: + t, h, w = int(t), int(h), int(w) + h_idxs = torch.linspace(0, n - 1, h) + w_idxs = torch.linspace(0, n - 1, w) + + h_floor = h_idxs.int() + w_floor = w_idxs.int() + h_ceil = (h_floor + 1).clip(max=n - 1) + w_ceil = (w_floor + 1).clip(max=n - 1) + + dh = h_idxs - h_floor.float() + dw = w_idxs - w_floor.float() + + base_h = h_floor * n + base_h_ceil = h_ceil * n + + indices = [ + (base_h[None].T + w_floor[None]).flatten(), + (base_h[None].T + w_ceil[None]).flatten(), + (base_h_ceil[None].T + w_floor[None]).flatten(), + (base_h_ceil[None].T + w_ceil[None]).flatten(), + ] + weights = [ + ((1 - dh)[None].T * (1 - dw)[None]).flatten(), + ((1 - dh)[None].T * dw[None]).flatten(), + (dh[None].T * (1 - dw)[None]).flatten(), + (dh[None].T * dw[None]).flatten(), + ] + + for i in range(4): + idx_list[i].extend(indices[i].tolist()) + weight_list[i].extend(weights[i].tolist()) + + idx_tensor = torch.tensor( + idx_list, dtype=torch.long, device=device, + ) + weight_tensor = torch.tensor( + weight_list, + dtype=self.pos_embed.weight.dtype, + device=device, + ) + pos_embeds = ( + self.pos_embed(idx_tensor).to(device) + * weight_tensor[:, :, None] + ) + patch_pos_embeds = ( + pos_embeds[0] + pos_embeds[1] + + pos_embeds[2] + pos_embeds[3] + ) + + patch_pos_embeds = patch_pos_embeds.split( + [h * w for h, w in zip(grid_hs, grid_ws)] + ) + + merge = self.spatial_merge_size + result = [] + for pe, t, h, w in zip( + patch_pos_embeds, grid_ts, grid_hs, grid_ws, + ): + pe = pe.repeat(t, 1) + pe = ( + pe.view( + t, h // merge, merge, w // merge, merge, -1, + ) + .permute(0, 1, 3, 2, 4, 5) + .flatten(0, 4) + ) + result.append(pe) + + return torch.cat(result) + + # --------------------------------------------------------------- + # 2D Vision RoPE + # --------------------------------------------------------------- + + def _compute_rotary_pos_emb(self, grid_thw: Tensor) -> Tensor: + """Compute 2D Vision RoPE for all patches in block-merge order. + + Matches HF ``Qwen3VLVisionModel.rot_pos_emb``. + + Args: + grid_thw: ``[num_images, 3]`` (T, H, W) per image. + + Returns: + Raw sectioned frequencies ``[3, 1, total_patches, head_dim // 2]`` + when ``config.mrope_section`` is set. Otherwise returns the legacy + ``[total_patches, head_dim // 2]`` row/column frequency tensor. + """ + merge = self.spatial_merge_size + grid_thw_list = grid_thw.tolist() + + max_hw = max(max(int(h), int(w)) for _, h, w in grid_thw_list) + freq_table = self.rot_pos_emb( + max_hw, device=grid_thw.device, + ) + device = freq_table.device + + total_tokens = sum( + int(t) * int(h) * int(w) for t, h, w in grid_thw_list + ) + pos_ids = torch.empty( + (total_tokens, 2), dtype=torch.long, device=device, + ) + + offset = 0 + for num_frames, height, width in grid_thw_list: + num_frames = int(num_frames) + height = int(height) + width = int(width) + merged_h = height // merge + merged_w = width // merge + + block_rows = torch.arange(merged_h, device=device) + block_cols = torch.arange(merged_w, device=device) + intra_row = torch.arange(merge, device=device) + intra_col = torch.arange(merge, device=device) + + row_idx = ( + block_rows[:, None, None, None] * merge + + intra_row[None, None, :, None] + ) + col_idx = ( + block_cols[None, :, None, None] * merge + + intra_col[None, None, None, :] + ) + + row_idx = row_idx.expand( + merged_h, merged_w, merge, merge, + ).reshape(-1) + col_idx = col_idx.expand( + merged_h, merged_w, merge, merge, + ).reshape(-1) + + coords = torch.stack((row_idx, col_idx), dim=-1) + if num_frames > 1: + coords = coords.repeat(num_frames, 1) + + n_tokens = coords.shape[0] + pos_ids[offset: offset + n_tokens] = coords + offset += n_tokens + + embeddings = freq_table[pos_ids] + embeddings = embeddings.flatten(1) + + mrope_section = getattr(self.config, "mrope_section", None) + if mrope_section is None: + return embeddings + + sec_t, sec_h, sec_w = (int(section) for section in mrope_section) + if sec_t != 0 or sec_h + sec_w != embeddings.shape[-1]: + raise ValueError( + "Qwen3.5-VL vision RoPE expects mrope_section " + f"[0, row_dim, col_dim] summing to {embeddings.shape[-1]}, " + f"got {mrope_section}" + ) + + raw_freqs = embeddings.new_zeros( + 3, 1, embeddings.shape[0], embeddings.shape[1], + ) + raw_freqs[1, 0, :, :sec_h] = embeddings[:, :sec_h] + raw_freqs[2, 0, :, sec_h : sec_h + sec_w] = embeddings[ + :, sec_h : sec_h + sec_w + ] + return raw_freqs + + # --------------------------------------------------------------- + # PackedSeqParams for variable-length attention + # --------------------------------------------------------------- + + @staticmethod + def _build_packed_seq_params(grid_thw: Tensor) -> PackedSeqParams: + """Build ``PackedSeqParams`` from grid dimensions. + + Each temporal frame of each image forms a separate sub-sequence + in the packed THD layout, matching HF's ``cu_seqlens`` computation. + + Args: + grid_thw: ``[num_images, 3]``. + + Returns: + ``PackedSeqParams`` for ``TransformerBlock``. + """ + cu_seqlens = torch.repeat_interleave( + grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0], + ).cumsum(dim=0, dtype=torch.int32) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + max_seqlen = int( + (grid_thw[:, 1] * grid_thw[:, 2]).max().item() + ) + + return PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + ) + + # --------------------------------------------------------------- + # Forward + # --------------------------------------------------------------- + + def forward( + self, + pixel_values: Tensor, + grid_thw: Tensor, + ) -> Tensor: + """Encode images / video frames. + + Args: + pixel_values: ``[total_patches, C * T * pH * pW]`` + pre-extracted flat patches in block-merge order. + grid_thw: ``[num_images, 3]`` (T, H, W) in patch-grid units. + + Returns: + ``[total_merged_patches, out_hidden_size]`` visual embeddings. + """ + # 1. Patch embedding (Conv3d) + hidden_states = self.patch_embed(pixel_values) + + # 2. Learned position embedding (bilinear interpolation) + pos_embeds = self._fast_pos_embed_interpolate(grid_thw) + hidden_states = hidden_states + pos_embeds + + # 3. 2D Vision RoPE + rot_freqs = self._compute_rotary_pos_emb(grid_thw) + if getattr(self.config, "mrope_section", None) is None: + emb = torch.cat((rot_freqs, rot_freqs), dim=-1) + rot_freqs = emb.unsqueeze(1).unsqueeze(1) + + # 4. Transformer blocks with PackedSeqParams + packed_seq_params = self._build_packed_seq_params(grid_thw) + hidden_states = hidden_states.unsqueeze(1) + hidden_states = self.decoder( + hidden_states=hidden_states, + attention_mask=None, + rotary_pos_emb=rot_freqs, + packed_seq_params=packed_seq_params, + ) + hidden_states = hidden_states.squeeze(1) + + # 5. Patch merger + return self.merger(hidden_states) diff --git a/examples/multimodal_dev/pretrain_multimodal.py b/examples/multimodal_dev/pretrain_multimodal.py new file mode 100644 index 00000000000..3792f05f9f1 --- /dev/null +++ b/examples/multimodal_dev/pretrain_multimodal.py @@ -0,0 +1,159 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Standalone entry point for multimodal_dev model training (FSDP + EP). + +This entry point is **model-agnostic**. All model-specific logic (layer +specs, model construction, FLOPs metadata, dataset generation) is +delegated to factory functions registered in +:data:`multimodal_dev.models.MODEL_REGISTRY`. + +Adding a new architecture only requires: + +1. Creating a new model package under ``multimodal_dev/models//`` + with the appropriate factory functions. +2. Registering an entry in ``MODEL_REGISTRY``. + +No changes to this file are necessary. + +Usage:: + + torchrun --nproc_per_node=8 multimodal_dev/pretrain_multimodal.py \\ + --model-arch qwen35_vl \\ + --dataset-provider mock \\ + ... (other megatron args) +""" + +import importlib +import os +import sys + +sys.path.insert( + 0, + os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")), +) + +from examples.multimodal_dev.arguments import add_multimodal_args +from examples.multimodal_dev.forward_step import forward_step +from megatron.core.enums import ModelType +from megatron.training import get_args, pretrain +from megatron.training.argument_utils import pretrain_cfg_container_from_args +from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args + + +def model_provider( + pre_process: bool = True, + post_process: bool = True, + **kwargs, +): + """Build a multimodal model from ``--model-arch``. + + The language ``TransformerConfig`` is built from CLI args so that + parallelism settings, precision, and fusion flags are inherited. + Model-specific post-processing and construction are delegated to the + registry factory functions. + """ + args = get_args() + model_arch = getattr(args, "model_arch", "qwen35_vl") + + from examples.multimodal_dev.models import MODEL_REGISTRY + + if model_arch not in MODEL_REGISTRY: + raise ValueError( + f"Unknown model arch '{model_arch}'. " + f"Available: {list(MODEL_REGISTRY.keys())}" + ) + + registry = MODEL_REGISTRY[model_arch] + + # --- language config (generic + model-specific post-processing) --- + language_config = core_transformer_config_from_args(args) + post_language_config_fn = registry.get("post_language_config_fn") + if post_language_config_fn is not None: + post_language_config_fn(language_config, args) + + # --- vision config --- + vision_config = registry["vision_config_fn"]( + num_layers_override=getattr(args, "vision_num_layers", None), + variant=getattr(args, "model_variant", None), + ) + vision_config.bf16 = language_config.bf16 + vision_config.fp16 = language_config.fp16 + vision_config.apply_rope_fusion = language_config.apply_rope_fusion + + if getattr(args, "recompute_vision", False): + vision_config.recompute_granularity = "full" + vision_config.recompute_method = "uniform" + vision_config.recompute_num_layers = 1 + + # --- vision FLOPs metadata --- + vision_flops_fn = registry.get("vision_flops_fn") + if vision_flops_fn is not None: + vision_flops_fn(args, language_config, vision_config) + + # --- build model (fully delegated to the arch factory) --- + model = registry["model_factory_fn"]( + args=args, + language_config=language_config, + vision_config=vision_config, + **kwargs, + ) + + return model + + +def _resolve_provider_fn(provider_fn): + """Resolve a provider that may be a dotted import path string.""" + if isinstance(provider_fn, str): + module_path, func_name = provider_fn.rsplit(".", 1) + provider_fn = getattr( + importlib.import_module(module_path), func_name, + ) + return provider_fn + + +def datasets_provider(train_val_test_num_samples): + """Dataset provider dispatcher. + + Routes to the dataset factory registered for the current + ``(--model-arch, --dataset-provider)`` combination. + """ + args = get_args() + model_arch = getattr(args, "model_arch", "qwen35_vl") + provider = getattr(args, "dataset_provider", "mock") + + from examples.multimodal_dev.models import MODEL_REGISTRY + + if model_arch not in MODEL_REGISTRY: + raise ValueError( + f"Unknown model arch '{model_arch}'. " + f"Available: {list(MODEL_REGISTRY.keys())}" + ) + + registry = MODEL_REGISTRY[model_arch] + available = registry.get("dataset_providers", {}) + + if provider not in available: + raise ValueError( + f"Unknown dataset provider '{provider}' for arch " + f"'{model_arch}'. Available: {list(available.keys())}" + ) + + provider_fn = _resolve_provider_fn(available[provider]) + return provider_fn(train_val_test_num_samples) + + +if __name__ == "__main__": + datasets_provider.is_distributed = True + + args = parse_and_validate_args( + extra_args_provider=add_multimodal_args, + args_defaults={}, + ) + full_config = pretrain_cfg_container_from_args(args) + pretrain( + full_config, + datasets_provider, + model_provider, + ModelType.encoder_or_decoder, + forward_step, + ) diff --git a/examples/multimodal_dev/scripts/run_qwen35_vl.sh b/examples/multimodal_dev/scripts/run_qwen35_vl.sh new file mode 100755 index 00000000000..44c1fb5e2a5 --- /dev/null +++ b/examples/multimodal_dev/scripts/run_qwen35_vl.sh @@ -0,0 +1,559 @@ +#!/bin/bash + +# Launch script for Qwen3.5-VL training via multimodal_dev (FSDP + EP). +# +# Usage (from the Megatron-LM repo root): +# ./examples/multimodal_dev/scripts/run_qwen35_vl.sh +# +# Environment variables: +# MODEL_VARIANT: proxy (default), 0.8b, 2b, 4b, 9b, 27b, 35b_a3b, 122b_a10b, 397b_a17b, 35b_a3b_light +# CKPT_LOAD: path to a pre-converted checkpoint to load (enables --load + --finetune) +# CKPT_FORMAT: checkpoint format override (e.g. torch_dist); auto-detected when empty +# TP, EP, PP: parallelism sizes +# MBS, GBS: micro/global batch sizes +# NUM_LAYERS, NUM_EXPERTS: override for proxy testing +# MTP_NUM_LAYERS: number of MTP layers (default: 1, set 0 to disable) +# LINEAR_ATTENTION_FREQ: every Nth decoder layer uses standard attention (default: 4; set 1 to force all standard attention) +# DATASET_PROVIDER: cord_v2 (default) or mock +# TOKENIZER_TYPE: HuggingFaceTokenizer (default) or NullTokenizer +# NO_ROPE_FUSION: set to 1 to pass --no-rope-fusion for baseline profiling +# SAVE_CHECKPOINTS: set to 0 to skip checkpoint saves in short profiling runs +# LAUNCHER: torchrun (default) or python +# TORCHRUN_PYTHON: Python executable for LAUNCHER=torchrun (default: python) +# PROFILE: set to 1 to enable Nsight Systems profiling (default: 0) +# NVTX_RANGES: set to 1 to emit Megatron custom NVTX ranges when PROFILE=1 (default: 1) +# PROFILE_STEP_START/PROFILE_STEP_END: profiled iteration window (default: 4-5) + +# example script: +# DRY_RUN=0 MODEL_VARIANT=proxy USE_PACKED_SEQUENCE=1 bash ./examples/multimodal_dev/scripts/run_qwen35_vl.sh + +set -euo pipefail + +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export NCCL_IB_SL=1 +export NVTE_FUSED_ATTN=1 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True + +DRY_RUN=${DRY_RUN:-1} +GPUS_PER_NODE=${GPUS_PER_NODE:-8} +if [ -n "${SLURM_JOB_NUM_NODES:-}" ]; then + NUM_NODES="$SLURM_JOB_NUM_NODES" +else + NUM_NODES=${NNODES:-1} +fi +PROFILE=${PROFILE:-0} +NVTX_RANGES=${NVTX_RANGES:-1} +PROFILE_STEP_START=${PROFILE_STEP_START:-4} +PROFILE_STEP_END=${PROFILE_STEP_END:-5} +PROFILE_RANKS=${PROFILE_RANKS:-0} +LAUNCHER=${LAUNCHER:-torchrun} +TORCHRUN_PYTHON=${TORCHRUN_PYTHON:-python} +NO_ROPE_FUSION=${NO_ROPE_FUSION:-0} + +MODEL_VARIANT=${MODEL_VARIANT:-proxy} +VISION_NUM_LAYERS=${VISION_NUM_LAYERS:-} + +# Batch sizes +MBS=${MBS:-2} +GBS=${GBS:-16} +MTP_NUM_LAYERS=${MTP_NUM_LAYERS:-1} +LINEAR_ATTENTION_FREQ=${LINEAR_ATTENTION_FREQ:-4} + +# Parallelism +TP=${TP:-1} +EP=${EP:-2} +PP=${PP:-1} +CP=${CP:-1} + +# Variant-aware architecture defaults. +# The model provider builds configs from the variant dict in +# multimodal_dev/models/qwen35_vl/configuration.py, but Megatron also +# uses these CLI args internally (PP splits, param counting). +case "$MODEL_VARIANT" in + 0.8b) + NUM_LAYERS=${NUM_LAYERS:-24} + NUM_EXPERTS=${NUM_EXPERTS:-0} + HIDDEN_SIZE=1024 + FFN_HIDDEN_SIZE=3584 + NUM_ATTN_HEADS=8 + NUM_QUERY_GROUPS=2 + LINEAR_NUM_VALUE_HEADS=16 + VISION_NUM_LAYERS=${VISION_NUM_LAYERS:-12} + ;; + 2b) + NUM_LAYERS=${NUM_LAYERS:-24} + NUM_EXPERTS=${NUM_EXPERTS:-0} + HIDDEN_SIZE=2048 + FFN_HIDDEN_SIZE=6144 + NUM_ATTN_HEADS=8 + NUM_QUERY_GROUPS=2 + LINEAR_NUM_VALUE_HEADS=16 + VISION_NUM_LAYERS=${VISION_NUM_LAYERS:-24} + ;; + 4b) + NUM_LAYERS=${NUM_LAYERS:-32} + NUM_EXPERTS=${NUM_EXPERTS:-0} + HIDDEN_SIZE=2560 + FFN_HIDDEN_SIZE=9216 + NUM_ATTN_HEADS=16 + NUM_QUERY_GROUPS=4 + LINEAR_NUM_VALUE_HEADS=32 + VISION_NUM_LAYERS=${VISION_NUM_LAYERS:-24} + ;; + proxy) + NUM_LAYERS=${NUM_LAYERS:-4} + NUM_EXPERTS=${NUM_EXPERTS:-16} + HIDDEN_SIZE=4096 + FFN_HIDDEN_SIZE=10240 + NUM_ATTN_HEADS=32 + NUM_QUERY_GROUPS=2 + LINEAR_NUM_VALUE_HEADS=64 + VISION_NUM_LAYERS=${VISION_NUM_LAYERS:-2} + ;; + 9b) + NUM_LAYERS=${NUM_LAYERS:-32} + NUM_EXPERTS=${NUM_EXPERTS:-0} + HIDDEN_SIZE=4096 + FFN_HIDDEN_SIZE=12288 + NUM_ATTN_HEADS=16 + NUM_QUERY_GROUPS=4 + LINEAR_NUM_VALUE_HEADS=32 + VISION_NUM_LAYERS=${VISION_NUM_LAYERS:-27} + ;; + 27b) + NUM_LAYERS=${NUM_LAYERS:-64} + NUM_EXPERTS=${NUM_EXPERTS:-0} + HIDDEN_SIZE=5120 + FFN_HIDDEN_SIZE=17408 + NUM_ATTN_HEADS=24 + NUM_QUERY_GROUPS=4 + LINEAR_NUM_VALUE_HEADS=48 + VISION_NUM_LAYERS=${VISION_NUM_LAYERS:-27} + ;; + 35b_a3b) + NUM_LAYERS=${NUM_LAYERS:-40} + NUM_EXPERTS=${NUM_EXPERTS:-256} + HIDDEN_SIZE=2048 + FFN_HIDDEN_SIZE=4096 + NUM_ATTN_HEADS=16 + NUM_QUERY_GROUPS=2 + LINEAR_NUM_VALUE_HEADS=32 + VISION_NUM_LAYERS=${VISION_NUM_LAYERS:-27} + ;; + 35b_a3b_light) + NUM_LAYERS=${NUM_LAYERS:-12} + NUM_EXPERTS=${NUM_EXPERTS:-128} + HIDDEN_SIZE=2048 + FFN_HIDDEN_SIZE=4096 + NUM_ATTN_HEADS=16 + NUM_QUERY_GROUPS=2 + LINEAR_NUM_VALUE_HEADS=32 + VISION_NUM_LAYERS=${VISION_NUM_LAYERS:-7} + ;; + 122b_a10b) + NUM_LAYERS=${NUM_LAYERS:-48} + NUM_EXPERTS=${NUM_EXPERTS:-256} + HIDDEN_SIZE=3072 + FFN_HIDDEN_SIZE=8192 + NUM_ATTN_HEADS=32 + NUM_QUERY_GROUPS=2 + LINEAR_NUM_VALUE_HEADS=64 + VISION_NUM_LAYERS=${VISION_NUM_LAYERS:-27} + ;; + 397b_a17b) + NUM_LAYERS=${NUM_LAYERS:-60} + NUM_EXPERTS=${NUM_EXPERTS:-512} + HIDDEN_SIZE=4096 + FFN_HIDDEN_SIZE=10240 + NUM_ATTN_HEADS=32 + NUM_QUERY_GROUPS=2 + LINEAR_NUM_VALUE_HEADS=64 + VISION_NUM_LAYERS=${VISION_NUM_LAYERS:-27} + ;; + *) + : "${NUM_LAYERS:?NUM_LAYERS must be set for MODEL_VARIANT=$MODEL_VARIANT}" + : "${NUM_EXPERTS:?NUM_EXPERTS must be set for MODEL_VARIANT=$MODEL_VARIANT}" + : "${HIDDEN_SIZE:?HIDDEN_SIZE must be set for MODEL_VARIANT=$MODEL_VARIANT}" + : "${FFN_HIDDEN_SIZE:?FFN_HIDDEN_SIZE must be set for MODEL_VARIANT=$MODEL_VARIANT}" + : "${NUM_ATTN_HEADS:?NUM_ATTN_HEADS must be set for MODEL_VARIANT=$MODEL_VARIANT}" + : "${NUM_QUERY_GROUPS:?NUM_QUERY_GROUPS must be set for MODEL_VARIANT=$MODEL_VARIANT}" + : "${LINEAR_NUM_VALUE_HEADS:?LINEAR_NUM_VALUE_HEADS must be set for MODEL_VARIANT=$MODEL_VARIANT}" + VISION_NUM_LAYERS=${VISION_NUM_LAYERS:-27} + ;; +esac +SEQ_LEN=${SEQ_LEN:-4096} + +WANDB_PROJECT=${WANDB_PROJECT:-'qwen35-vl-0524'} +EXP_NAME="qwen35vl_${MODEL_VARIANT}_tp${TP}_ep${EP}_pp${PP}_cp${CP}" + +RECOMPUTE_VISION=${RECOMPUTE_VISION:-0} +if [ "$RECOMPUTE_VISION" -eq 1 ]; then + EXP_NAME+="_recompute_encoder" +fi +RECOMPUTE=${RECOMPUTE:-0} +if [ "$RECOMPUTE" -eq 1 ]; then + EXP_NAME+="_recompute_decoder" +fi + +USE_PACKED_SEQUENCE=${USE_PACKED_SEQUENCE:-0} +if [ "$USE_PACKED_SEQUENCE" -eq 1 ]; then + EXP_NAME+="_thd" +fi +if [ "$NO_ROPE_FUSION" -eq 1 ]; then + EXP_NAME+="_no_rope_fusion" +fi + +MEGATRON_LM_PATH="${MEGATRON_LM_PATH:-$(cd "$(dirname "$0")/../../.." && pwd)}" +ROOT_DIR="${ROOT_DIR:-${MEGATRON_LM_PATH}/local/}" +CHECKPOINT_STORE_PATH="${ROOT_DIR}${EXP_NAME}" +mkdir -p "$CHECKPOINT_STORE_PATH" + +TENSORBOARD_LOGS_PATH="${TENSORBOARD_LOGS_PATH:-${MEGATRON_LM_PATH}/logs}" +mkdir -p "$TENSORBOARD_LOGS_PATH" + +DISTRIBUTED_ARGS=( + --nproc_per_node "$GPUS_PER_NODE" + --nnodes "$NUM_NODES" +) + +if [ "$NUM_NODES" -gt 1 ]; then + DISTRIBUTED_ARGS+=( + --master_addr "${MASTER_ADDR:-localhost}" + --master_port "${MASTER_PORT:-6000}" + ) +fi + +# --- Parallelism --- +MODEL_PARALLEL_ARGS=( + --tensor-model-parallel-size "$TP" + --pipeline-model-parallel-size "$PP" + --expert-model-parallel-size "$EP" + --context-parallel-size "$CP" + --cp-comm-type "a2a" + --expert-tensor-parallel-size 1 + --use-distributed-optimizer + --sequence-parallel +) + +# --- Training --- +TRAINING_ARGS=( + --micro-batch-size "$MBS" + --global-batch-size "$GBS" + --train-iters "${TRAIN_ITERS:-500}" + --adam-beta1 0.9 + --adam-beta2 0.95 + --lr 1.2e-4 + --min-lr 1.2e-5 + --lr-decay-style cosine + --lr-warmup-iters 100 + --lr-decay-iters 2000 + --weight-decay 0.1 + --clip-grad 1.0 + --bf16 + --use-mcore-models + --transformer-impl transformer_engine + --cross-entropy-loss-fusion + --cross-entropy-fusion-impl te + --enable-experimental + --manual-gc + --manual-gc-interval 50 + --sft + --use-flash-attn + # --attention-backend flash + --calculate-per-token-loss +) +if [ "$MTP_NUM_LAYERS" -gt 0 ]; then + TRAINING_ARGS+=( + --mtp-num-layers "$MTP_NUM_LAYERS" + --mtp-loss-scaling-factor 0.1 + ) +fi + +PROFILE_ARGS=() +NSYS_CMD=() +if [ "$PROFILE" = "1" ]; then + PROFILE_ARGS=( + --profile + --profile-step-start "$PROFILE_STEP_START" + --profile-step-end "$PROFILE_STEP_END" + --profile-ranks "$PROFILE_RANKS" + ) + if [ "$NVTX_RANGES" -eq 1 ]; then + PROFILE_ARGS+=( --nvtx-ranges ) + fi + + NSYS_OUTPUT_DIR="${CHECKPOINT_STORE_PATH}/nsys" + mkdir -p "$NSYS_OUTPUT_DIR" + NSYS_CMD=( + nsys profile + --sample=none + --cpuctxsw=none + --trace=cuda,nvtx,cublas,cudnn + --force-overwrite=true + --capture-range=cudaProfilerApi + --capture-range-end=stop + -o "${NSYS_OUTPUT_DIR}/${EXP_NAME}_$(date +%Y%m%d_%H%M%S)" + ) +fi + +# --- Logging & Checkpointing --- +SAVE_CHECKPOINTS=${SAVE_CHECKPOINTS:-1} +SAVE_INTERVAL=${SAVE_INTERVAL:-500} +EVAL_AND_LOGGING_ARGS=( + --log-interval 1 + --eval-interval 500 + --eval-iters 10 + --tensorboard-dir "$TENSORBOARD_LOGS_PATH" + --wandb-project "$WANDB_PROJECT" + --wandb-exp-name "$EXP_NAME" + --wandb-save-dir "$CHECKPOINT_STORE_PATH" + --log-throughput + --log-timers-to-tensorboard + --log-params-norm +) +if [ "$SAVE_CHECKPOINTS" -eq 1 ]; then + EVAL_AND_LOGGING_ARGS+=( + --save-interval "$SAVE_INTERVAL" + --save "$CHECKPOINT_STORE_PATH" + ) +fi + +# --- Tokenizer --- +TOKENIZER_MODEL=${TOKENIZER_MODEL:-Qwen/Qwen3.5-397B-A17B} +TOKENIZER_TYPE=${TOKENIZER_TYPE:-HuggingFaceTokenizer} +VOCAB_SIZE=${VOCAB_SIZE:-248320} +TOKENIZER_ARGS=( + --tokenizer-type "$TOKENIZER_TYPE" +) +if [ "$TOKENIZER_TYPE" = "NullTokenizer" ]; then + TOKENIZER_ARGS+=( --vocab-size "$VOCAB_SIZE" ) +else + TOKENIZER_ARGS+=( --tokenizer-model "$TOKENIZER_MODEL" ) +fi + +# --- Multimodal-specific --- +DATASET_PROVIDER=${DATASET_PROVIDER:-cord_v2} +HF_PROCESSOR_PATH=${HF_PROCESSOR_PATH-Qwen/Qwen3.5-397B-A17B} +IMAGE_SEQ_LENGTH=${IMAGE_SEQ_LENGTH:-256} +MULTIMODAL_ARGS=( + --model-arch qwen35_vl + --model-variant "$MODEL_VARIANT" + --dataset-provider "$DATASET_PROVIDER" + --use-vanilla-collate-fn + --image-token-id 248056 + --image-size 224 + --total-seq-length "$SEQ_LEN" + --image-seq-length "$IMAGE_SEQ_LENGTH" + --vision-num-layers "$VISION_NUM_LAYERS" +) +if [ -n "$HF_PROCESSOR_PATH" ]; then + MULTIMODAL_ARGS+=( --hf-processor-path "$HF_PROCESSOR_PATH" ) +fi + +if [ "$USE_PACKED_SEQUENCE" -eq 1 ]; then + MULTIMODAL_ARGS+=( --use-packed-sequence ) +fi + +# --- Qwen3.5 Decoder Architecture (variant-specific dims set above) --- +# These must match examples/multimodal_dev/models/qwen35_vl/configuration.py +GPT_MODEL_ARGS=( + --num-layers "$NUM_LAYERS" + --hidden-size "$HIDDEN_SIZE" + --ffn-hidden-size "$FFN_HIDDEN_SIZE" + --num-attention-heads "$NUM_ATTN_HEADS" + --group-query-attention + --num-query-groups "$NUM_QUERY_GROUPS" + --kv-channels 256 + --max-position-embeddings 262144 + --seq-length "$SEQ_LEN" + --normalization RMSNorm + --apply-layernorm-1p + --norm-epsilon 1e-06 + --swiglu + --disable-bias-linear + --position-embedding-type rope + --rotary-percent 0.25 + --rotary-base 10000000 + --rotary-seq-len-interpolation-factor 1 + --qk-layernorm + --attention-output-gate + --attention-dropout 0.0 + --hidden-dropout 0.0 + --experimental-attention-variant gated_delta_net + --linear-attention-freq "$LINEAR_ATTENTION_FREQ" + --linear-conv-kernel-dim 4 + --linear-key-head-dim 128 + --linear-value-head-dim 128 + --linear-num-key-heads 16 + --linear-num-value-heads "$LINEAR_NUM_VALUE_HEADS" + --make-vocab-size-divisible-by 485 + --moe-router-force-load-balancing +) +if [ "$NO_ROPE_FUSION" -eq 1 ]; then + GPT_MODEL_ARGS+=( --no-rope-fusion ) +fi + +# --- Tied / untied embeddings --- +# 0.8B, 2B, 4B use tied embeddings; all other variants untie them. +case "$MODEL_VARIANT" in + 0.8b|2b|4b) ;; + *) GPT_MODEL_ARGS+=( --untie-embeddings-and-output-weights ) ;; +esac + +# --- MoE args (MoE variants only) --- +MOE_ARGS=() +case "$MODEL_VARIANT" in + proxy) + MOE_TOPK=2; MOE_FFN_HIDDEN=1024; MOE_SHARED_HIDDEN=1024 + ;; + 35b_a3b|35b_a3b_light) + MOE_TOPK=8; MOE_FFN_HIDDEN=512; MOE_SHARED_HIDDEN=512 + ;; + 122b_a10b) + MOE_TOPK=8; MOE_FFN_HIDDEN=1024; MOE_SHARED_HIDDEN=1024 + ;; + 397b_a17b) + MOE_TOPK=10; MOE_FFN_HIDDEN=1024; MOE_SHARED_HIDDEN=1024 + ;; + 0.8b|2b|4b|9b|27b) + ;; +esac +if [ "${NUM_EXPERTS:-0}" -gt 0 ]; then + MOE_ARGS=( + --num-experts "$NUM_EXPERTS" + --moe-ffn-hidden-size "$MOE_FFN_HIDDEN" + --moe-shared-expert-intermediate-size "$MOE_SHARED_HIDDEN" + --moe-shared-expert-gate + --moe-router-load-balancing-type aux_loss + --moe-router-topk "$MOE_TOPK" + --moe-grouped-gemm + --moe-aux-loss-coeff 1e-3 + --moe-token-dispatcher-type alltoall + --moe-router-dtype fp32 + --moe-permute-fusion + --moe-router-fusion + ) +fi + +# --- Recompute --- +if [ "$RECOMPUTE" -eq 1 ]; then + RECOMPUTE_ARGS=( + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + ) + # RECOMPUTE_ARGS=( + # --recompute-granularity selective + # --recompute-modules moe_act shared_experts layernorm moe + # ) +else + RECOMPUTE_ARGS=() +fi +if [ "$RECOMPUTE_VISION" -eq 1 ]; then + RECOMPUTE_ARGS+=( --recompute-vision ) +fi + +# --- Checkpoint loading --- +# CKPT_LOAD: path to checkpoint directory +# CKPT_FORMAT: override checkpoint format (default: auto-detect) +# CKPT_RESUME: set to 1 to resume training (keep iteration, optimizer, rng); +# default 0 = finetune mode (reset iteration, skip optim/rng) +CKPT_LOAD=${CKPT_LOAD:-} +CKPT_FORMAT=${CKPT_FORMAT:-} +CKPT_RESUME=${CKPT_RESUME:-0} +CKPT_OVERRIDE_SCHEDULER=${CKPT_OVERRIDE_SCHEDULER:-0} +CKPT_ARGS=() +if [ -n "$CKPT_LOAD" ]; then + CKPT_ARGS+=( --load "$CKPT_LOAD" ) + if [ "$CKPT_RESUME" -eq 0 ]; then + CKPT_ARGS+=( --finetune --no-load-optim --no-load-rng ) + fi + if [ -n "$CKPT_FORMAT" ]; then + CKPT_ARGS+=( --ckpt-format "$CKPT_FORMAT" ) + fi + if [ "$CKPT_OVERRIDE_SCHEDULER" -eq 1 ]; then + CKPT_ARGS+=( --override-opt-param-scheduler ) + fi +fi + +# --- FSDP --- +USE_FSDP=${USE_FSDP:-1} +if [ "$USE_FSDP" -eq 1 ]; then + FSDP_ARGS=( + --use-megatron-fsdp + --data-parallel-sharding-strategy optim_grads_params + --init-model-with-meta-device + --use-distributed-optimizer + --ckpt-format fsdp_dtensor + ) + export CUDA_DEVICE_MAX_CONNECTIONS=8 +else + FSDP_ARGS=() +fi + +echo "================================================================" +echo "Qwen3.5-VL Multimodal Training (multimodal_dev)" +echo " Variant: $MODEL_VARIANT" +echo " Vision layers: $VISION_NUM_LAYERS" +echo " GPUs per node: $GPUS_PER_NODE" +echo " Num nodes: $NUM_NODES" +echo " TP=$TP EP=$EP PP=$PP CP=$CP" +echo " MBS=$MBS GBS=$GBS" +echo " MTP layers: $MTP_NUM_LAYERS" +echo " Linear attn freq: $LINEAR_ATTENTION_FREQ" +echo " Launcher: $LAUNCHER" +if [ "$LAUNCHER" = "torchrun" ]; then + echo " Torchrun py: $TORCHRUN_PYTHON" +fi +echo " FSDP: $USE_FSDP" +echo " PROFILE: $PROFILE" +echo " RoPE fusion: $([ "$NO_ROPE_FUSION" -eq 1 ] && echo off || echo on)" +echo " Dataset: $DATASET_PROVIDER" +echo " Tokenizer: $TOKENIZER_TYPE" +echo " Checkpoints: $([ "$SAVE_CHECKPOINTS" -eq 1 ] && echo on || echo off)" +if [ -n "$CKPT_LOAD" ]; then + echo " CKPT_LOAD: $CKPT_LOAD" + echo " CKPT_FORMAT: ${CKPT_FORMAT:-auto}" + echo " CKPT_RESUME: $CKPT_RESUME" +fi +if [ "$PROFILE" = "1" ]; then + echo " Profile steps: ${PROFILE_STEP_START}-${PROFILE_STEP_END}" + echo " Profile ranks: $PROFILE_RANKS" + echo " NVTX ranges: $([ "$NVTX_RANGES" -eq 1 ] && echo on || echo off)" +fi +echo "================================================================" + +if [ "$LAUNCHER" = "python" ]; then + LAUNCH_CMD=( python $MEGATRON_LM_PATH/examples/multimodal_dev/pretrain_multimodal.py ) +elif [ "$LAUNCHER" = "torchrun" ]; then + LAUNCH_CMD=( + "$TORCHRUN_PYTHON" -m torch.distributed.run + "${DISTRIBUTED_ARGS[@]}" + $MEGATRON_LM_PATH/examples/multimodal_dev/pretrain_multimodal.py + ) +else + echo "Unsupported LAUNCHER=$LAUNCHER (expected torchrun or python)" >&2 + exit 1 +fi + +cmd=( "${NSYS_CMD[@]}" "${LAUNCH_CMD[@]}" \ + "${TRAINING_ARGS[@]}" \ + "${PROFILE_ARGS[@]}" \ + "${MODEL_PARALLEL_ARGS[@]}" \ + "${EVAL_AND_LOGGING_ARGS[@]}" \ + "${TOKENIZER_ARGS[@]}" \ + "${MULTIMODAL_ARGS[@]}" \ + "${GPT_MODEL_ARGS[@]}" \ + "${MOE_ARGS[@]}" \ + "${RECOMPUTE_ARGS[@]}" \ + "${FSDP_ARGS[@]}" \ + "${CKPT_ARGS[@]}" ) + +echo "${cmd[@]}" + +if [ "$DRY_RUN" -eq 1 ]; then + echo "=== DRY RUN ===" + exit 0 +else + "${cmd[@]}" +fi diff --git a/examples/multimodal_dev/tests/__init__.py b/examples/multimodal_dev/tests/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/examples/multimodal_dev/tests/_helpers.py b/examples/multimodal_dev/tests/_helpers.py new file mode 100644 index 00000000000..b0c69207f3a --- /dev/null +++ b/examples/multimodal_dev/tests/_helpers.py @@ -0,0 +1,19 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Test helpers shared across the multimodal_dev test suite.""" + + +def grad_norm(model): + """L2 norm of all populated parameter gradients on this rank.""" + total = 0.0 + for p in model.parameters(): + if p.grad is not None: + total += p.grad.data.float().norm(2).item() ** 2 + return total**0.5 + + +def mean_loss(per_token_loss, loss_mask): + """Mean per-token loss over valid (mask>0) positions on this rank.""" + flat = per_token_loss.float().view(-1) + mask = loss_mask.float().view(-1) + return (flat * mask).sum() / mask.sum().clamp(min=1) diff --git a/examples/multimodal_dev/tests/test_cp_correctness.py b/examples/multimodal_dev/tests/test_cp_correctness.py new file mode 100644 index 00000000000..fe156ed54af --- /dev/null +++ b/examples/multimodal_dev/tests/test_cp_correctness.py @@ -0,0 +1,313 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Distributed correctness test for Context Parallelism (CP) support. + +Verifies that CP>1 produces the same (or numerically close) loss as CP=1 +for the Qwen3.5-VL multimodal model by running forward passes with +deterministic data and comparing the per-rank reduced losses. + +Launch with torchrun (N must be >= 2*max_cp_size for zigzag splitting): + + # Test CP=2 on 2 GPUs: + torchrun --nproc_per_node=2 examples/multimodal_dev/tests/test_cp_correctness.py --cp-size 2 + + # Test CP=4 on 4 GPUs: + torchrun --nproc_per_node=4 examples/multimodal_dev/tests/test_cp_correctness.py --cp-size 4 + +The test: + 1. Builds a tiny proxy model (2 layers, no MoE, no vision encoder). + 2. Generates a deterministic batch (same seed on all ranks). + 3. Runs forward with CP=1 (each rank processes the full sequence independently). + 4. Re-initialises model-parallel groups with the target CP size. + 5. Runs forward with CP=target (sequence is split across ranks). + 6. Compares the all-reduced loss values. + +Exit code 0 = PASS, 1 = FAIL. +""" + +import argparse +import os +import sys + +import torch +import torch.distributed as dist + +# Ensure the repo root is on the path so that megatron and examples are importable. +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + + +def _parse_args(): + parser = argparse.ArgumentParser(description="CP correctness test") + parser.add_argument( + "--cp-size", type=int, default=2, + help="Target context-parallel size to compare against CP=1 baseline", + ) + parser.add_argument( + "--seq-len", type=int, default=128, + help="Sequence length (must be divisible by 2*max(cp_size, tp_size*cp_size))", + ) + parser.add_argument( + "--atol", type=float, default=1e-4, + help="Absolute tolerance for loss comparison", + ) + parser.add_argument( + "--rtol", type=float, default=5e-2, + help="Relative tolerance for loss comparison (default 5%%)", + ) + parser.add_argument( + "--seed", type=int, default=42, + help="Random seed for reproducibility", + ) + # Megatron adds extra args; ignore them. + args, _ = parser.parse_known_args() + return args + + +def _init_distributed(): + """Initialise torch.distributed if not already done.""" + if not dist.is_initialized(): + dist.init_process_group(backend="nccl") + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + return local_rank + + +def _init_megatron_parallel(tp_size=1, pp_size=1, cp_size=1, seed=42): + """(Re-)initialise Megatron model-parallel groups and RNG tracker.""" + from megatron.core import parallel_state as ps + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + ps.destroy_model_parallel() + ps.initialize_model_parallel( + tensor_model_parallel_size=tp_size, + pipeline_model_parallel_size=pp_size, + context_parallel_size=cp_size, + ) + model_parallel_cuda_manual_seed(seed) + + +def _make_deterministic_batch(seed, batch_size, seq_len, vocab_size, device): + """Create a deterministic batch identical on all ranks.""" + rng = torch.Generator(device="cpu") + rng.manual_seed(seed) + + input_ids = torch.randint( + 0, vocab_size, (batch_size, seq_len), generator=rng, + ).to(device) + labels = torch.randint( + 0, vocab_size, (batch_size, seq_len), generator=rng, + ).to(device) + loss_mask = torch.ones(batch_size, seq_len, device=device) + # Standard position_ids [B, S] + position_ids = torch.arange(seq_len, device=device).unsqueeze(0).expand(batch_size, -1) + + return { + "input_ids": input_ids, + "labels": labels, + "loss_mask": loss_mask, + "position_ids": position_ids, + } + + +def _build_tiny_model(cp_size, device): + """Build a minimal GPTModel for testing (no vision, no MoE).""" + from megatron.core.models.gpt import GPTModel + from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec + from megatron.core.transformer.spec_utils import ModuleSpec + from megatron.core.transformer.transformer_config import TransformerConfig + + hidden_size = 256 + num_heads = 4 + config = TransformerConfig( + num_layers=2, + hidden_size=hidden_size, + ffn_hidden_size=hidden_size * 4, + num_attention_heads=num_heads, + kv_channels=hidden_size // num_heads, + normalization="RMSNorm", + layernorm_epsilon=1e-6, + gated_linear_unit=True, + activation_func=torch.nn.functional.silu, + bf16=True, + context_parallel_size=cp_size, + add_bias_linear=False, + attention_dropout=0.0, + hidden_dropout=0.0, + sequence_parallel=False, + ) + + spec = get_gpt_layer_with_transformer_engine_spec() + + model = GPTModel( + config=config, + transformer_layer_spec=spec, + vocab_size=1024, + max_sequence_length=4096, + pre_process=True, + post_process=True, + parallel_output=False, + share_embeddings_and_output_weights=True, + position_embedding_type="rope", + rotary_percent=1.0, + rotary_base=10000, + ) + model = model.to(device=device, dtype=torch.bfloat16) + return model, config + + +def _forward_with_cp(model, batch, cp_size): + """Run forward pass, handling CP splitting of the batch. + + When cp_size > 1, splits the batch tensors using the same zigzag + logic as multimodal_dev/models/base.py. + """ + from examples.multimodal_dev.models.base import _cp_split_tensor + from megatron.core import parallel_state as ps + + input_ids = batch["input_ids"].clone() + labels = batch["labels"].clone() + loss_mask = batch["loss_mask"].clone() + position_ids = batch["position_ids"].clone() + + if cp_size > 1: + cp_rank = ps.get_context_parallel_rank() + input_ids = _cp_split_tensor(input_ids, seq_dim=1, cp_size=cp_size, cp_rank=cp_rank) + labels = _cp_split_tensor(labels, seq_dim=1, cp_size=cp_size, cp_rank=cp_rank) + loss_mask = _cp_split_tensor(loss_mask, seq_dim=1, cp_size=cp_size, cp_rank=cp_rank) + # position_ids are NOT split — the RoPE layer handles CP slicing internally. + + with torch.no_grad(): + output = model( + input_ids=input_ids, + position_ids=position_ids, + labels=labels, + attention_mask=None, + ) + + # output is the per-token loss [B, S/CP] + masked_loss = (output.float() * loss_mask.float()).sum() + num_tokens = loss_mask.sum() + + # All-reduce across CP ranks to get global loss + if cp_size > 1: + cp_group = ps.get_context_parallel_group() + dist.all_reduce(masked_loss, group=cp_group) + dist.all_reduce(num_tokens, group=cp_group) + + avg_loss = masked_loss / num_tokens.clamp(min=1) + return avg_loss.item() + + +def main(): + args = _parse_args() + local_rank = _init_distributed() + device = torch.device(f"cuda:{local_rank}") + world_size = dist.get_world_size() + rank = dist.get_rank() + + target_cp = args.cp_size + if world_size < target_cp: + if rank == 0: + print( + f"SKIP: world_size={world_size} < cp_size={target_cp}. " + f"Need at least {target_cp} GPUs.", + flush=True, + ) + dist.destroy_process_group() + sys.exit(0) + if world_size % target_cp != 0: + if rank == 0: + print( + f"SKIP: world_size={world_size} is not divisible by cp_size={target_cp}.", + flush=True, + ) + dist.destroy_process_group() + sys.exit(0) + + vocab_size = 1024 + + # Ensure seq_len is divisible by 2 * target_cp + seq_len = args.seq_len + align = 2 * target_cp + if seq_len % align != 0: + seq_len = ((seq_len + align - 1) // align) * align + if rank == 0: + print(f"Adjusted seq_len to {seq_len} for alignment with CP={target_cp}", flush=True) + + # --- Step 1: CP=1 baseline --- + if rank == 0: + print(f"=== CP=1 baseline (world_size={world_size}) ===", flush=True) + + _init_megatron_parallel(cp_size=1) + + # Set deterministic seed for model init + torch.manual_seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + model_cp1, _ = _build_tiny_model(cp_size=1, device=device) + + batch = _make_deterministic_batch( + seed=args.seed + 1, batch_size=1, seq_len=seq_len, + vocab_size=vocab_size, device=device, + ) + + loss_cp1 = _forward_with_cp(model_cp1, batch, cp_size=1) + + if rank == 0: + print(f" CP=1 loss: {loss_cp1:.6f}", flush=True) + + # Save model state for reuse + state_dict = model_cp1.state_dict() + del model_cp1 + torch.cuda.empty_cache() + + # --- Step 2: CP=target --- + if rank == 0: + print(f"=== CP={target_cp} (world_size={world_size}) ===", flush=True) + + _init_megatron_parallel(cp_size=target_cp) + + torch.manual_seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + model_cpN, _ = _build_tiny_model(cp_size=target_cp, device=device) + + # Load the same weights to ensure identical model + model_cpN.load_state_dict(state_dict, strict=True) + del state_dict + + loss_cpN = _forward_with_cp(model_cpN, batch, cp_size=target_cp) + + if rank == 0: + print(f" CP={target_cp} loss: {loss_cpN:.6f}", flush=True) + + del model_cpN + torch.cuda.empty_cache() + + # --- Step 3: Compare --- + if rank == 0: + diff = abs(loss_cpN - loss_cp1) + rel_diff = diff / max(abs(loss_cp1), 1e-10) + + print(f"\n=== Comparison ===", flush=True) + print(f" CP=1 loss: {loss_cp1:.6f}", flush=True) + print(f" CP={target_cp} loss: {loss_cpN:.6f}", flush=True) + print(f" Absolute diff: {diff:.6e}", flush=True) + print(f" Relative diff: {rel_diff:.6e}", flush=True) + print(f" Tolerance (atol): {args.atol:.6e}", flush=True) + print(f" Tolerance (rtol): {args.rtol:.6e}", flush=True) + + passed = diff <= args.atol + args.rtol * abs(loss_cp1) + if passed: + print(f"\nPASS: CP={target_cp} matches CP=1 baseline", flush=True) + else: + print(f"\nFAIL: CP={target_cp} loss differs from CP=1 beyond tolerance", flush=True) + + dist.barrier() + dist.destroy_process_group() + + if rank == 0 and not passed: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/multimodal_dev/tests/test_cp_support.py b/examples/multimodal_dev/tests/test_cp_support.py new file mode 100644 index 00000000000..d46b5d8ef71 --- /dev/null +++ b/examples/multimodal_dev/tests/test_cp_support.py @@ -0,0 +1,347 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for Context Parallelism (CP) support in multimodal_dev. + +Tests cover: + 1. _cp_split_tensor — zigzag split correctness, reconstruction, and edge cases + 2. _NoCPGroup — dummy process group behaviour + 3. _thd_cp_partition_index — TE-based per-sample THD CP partitioning + 4. Cross-validation against megatron.core.utils.get_batch_on_this_cp_rank + +Run with: pytest examples/multimodal_dev/tests/test_cp_support.py -v +""" + +import pytest +import torch + +from examples.multimodal_dev.models.base import _cp_split_tensor, _NoCPGroup + + +class TestCpSplitTensor: + """Tests for zigzag CP splitting.""" + + def test_basic_2d_cp2(self): + """[B, S] tensor with CP=2 splits and reconstructs correctly.""" + B, S = 2, 16 + t = torch.arange(B * S).reshape(B, S) + cp_size = 2 + + chunks = [] + for rank in range(cp_size): + chunks.append(_cp_split_tensor(t, seq_dim=1, cp_size=cp_size, cp_rank=rank)) + + # Each rank gets S / CP = 8 tokens + for c in chunks: + assert c.shape == (B, S // cp_size) + + # Reconstruct: rank 0 gets chunks [0, 3], rank 1 gets chunks [1, 2] + # Original split into 4 chunks of size 4: + # chunk0=[0..3], chunk1=[4..7], chunk2=[8..11], chunk3=[12..15] + # rank0 = [chunk0, chunk3] = [0..3, 12..15] + # rank1 = [chunk1, chunk2] = [4..7, 8..11] + assert torch.equal(chunks[0][0], torch.tensor([0, 1, 2, 3, 12, 13, 14, 15])) + assert torch.equal(chunks[1][0], torch.tensor([4, 5, 6, 7, 8, 9, 10, 11])) + + def test_3d_mrope_cp2(self): + """[3, B, S] MRoPE tensor with CP=2.""" + B, S = 1, 8 + cp_size = 2 + t = torch.arange(3 * B * S).reshape(3, B, S) + + chunk = _cp_split_tensor(t, seq_dim=2, cp_size=cp_size, cp_rank=0) + assert chunk.shape == (3, B, S // cp_size) + + # All 3 MRoPE components should be split consistently + for d in range(3): + original_row = t[d, 0] # [S] + # With S=8, CP=2: 4 chunks of size 2 + # rank0 gets chunks [0, 3] = positions [0,1, 6,7] + expected = torch.cat([original_row[0:2], original_row[6:8]]) + assert torch.equal(chunk[d, 0], expected) + + def test_sbh_decoder_input(self): + """[S, B, H] decoder input split along dim=0.""" + S, B, H = 16, 2, 4 + cp_size = 2 + t = torch.randn(S, B, H) + + chunk = _cp_split_tensor(t, seq_dim=0, cp_size=cp_size, cp_rank=0) + assert chunk.shape == (S // cp_size, B, H) + + def test_cp4(self): + """CP=4 zigzag pattern.""" + S = 32 + cp_size = 4 + t = torch.arange(S).unsqueeze(0) # [1, 32] + + all_chunks = [] + for rank in range(cp_size): + c = _cp_split_tensor(t, seq_dim=1, cp_size=cp_size, cp_rank=rank) + all_chunks.append(c) + assert c.shape == (1, S // cp_size) + + # All tokens should appear exactly once across ranks + combined = torch.cat(all_chunks, dim=1) + assert torch.equal(combined.sort(dim=1).values, t.sort(dim=1).values) + + def test_cp8(self): + """CP=8 zigzag pattern — all tokens appear exactly once.""" + S = 64 + cp_size = 8 + t = torch.arange(S).unsqueeze(0) # [1, 64] + + all_chunks = [] + for rank in range(cp_size): + c = _cp_split_tensor(t, seq_dim=1, cp_size=cp_size, cp_rank=rank) + all_chunks.append(c) + assert c.shape == (1, S // cp_size) + + combined = torch.cat(all_chunks, dim=1) + assert torch.equal(combined.sort(dim=1).values, t.sort(dim=1).values) + + def test_not_divisible_raises(self): + """Should raise when seq_len not divisible by 2*cp_size.""" + t = torch.randn(2, 10) # S=10, not divisible by 4 + with pytest.raises(AssertionError): + _cp_split_tensor(t, seq_dim=1, cp_size=2, cp_rank=0) + + def test_zigzag_symmetry(self): + """rank 0 and rank (cp_size-1) should get mirror chunks.""" + S = 16 + cp_size = 2 + t = torch.arange(S).unsqueeze(0) # [1, 16] + + c0 = _cp_split_tensor(t, seq_dim=1, cp_size=cp_size, cp_rank=0) + c1 = _cp_split_tensor(t, seq_dim=1, cp_size=cp_size, cp_rank=1) + + # rank0 gets chunks [0, 3], rank1 gets chunks [1, 2] + # chunk0=[0..3], chunk3=[12..15] -> rank0 gets [0..3, 12..15] + # chunk1=[4..7], chunk2=[8..11] -> rank1 gets [4..7, 8..11] + # rank0's first half is earliest, rank1's first half is next + assert c0[0, 0].item() < c1[0, 0].item() # rank0 starts earlier + + def test_matches_megatron_core(self): + """Cross-validate against megatron.core.utils.get_batch_on_this_cp_rank logic. + + We simulate the core function's logic (seq_dim=1, attention_mask seq_dim=2) + and compare. + """ + B, S = 2, 32 + cp_size = 4 + + input_ids = torch.arange(B * S).reshape(B, S) + labels = torch.arange(B * S).reshape(B, S) + 1000 + + for cp_rank in range(cp_size): + # Our implementation + our_ids = _cp_split_tensor(input_ids, seq_dim=1, cp_size=cp_size, cp_rank=cp_rank) + our_labels = _cp_split_tensor(labels, seq_dim=1, cp_size=cp_size, cp_rank=cp_rank) + + # Simulate megatron core logic inline + def core_split(val, seq_dim): + val = val.view( + *val.shape[0:seq_dim], + 2 * cp_size, + val.shape[seq_dim] // (2 * cp_size), + *val.shape[(seq_dim + 1):], + ) + index = torch.zeros(2, dtype=torch.int64, device=val.device) + index[0].fill_(cp_rank) + index[1].fill_(2 * cp_size - cp_rank - 1) + val = val.index_select(seq_dim, index) + val = val.view(*val.shape[0:seq_dim], -1, *val.shape[(seq_dim + 2):]) + return val + + ref_ids = core_split(input_ids.clone(), seq_dim=1) + ref_labels = core_split(labels.clone(), seq_dim=1) + + assert torch.equal(our_ids, ref_ids), f"input_ids mismatch at rank {cp_rank}" + assert torch.equal(our_labels, ref_labels), f"labels mismatch at rank {cp_rank}" + + def test_batch_dim_preserved(self): + """Batch dimension must be unchanged after split.""" + B, S = 4, 32 + cp_size = 4 + t = torch.randn(B, S) + + for rank in range(cp_size): + c = _cp_split_tensor(t, seq_dim=1, cp_size=cp_size, cp_rank=rank) + assert c.shape[0] == B + + +class TestNoCPGroup: + """Tests for the dummy CP group used by the vision encoder.""" + + def test_size_is_one(self): + g = _NoCPGroup() + assert g.size() == 1 + + def test_rank_is_zero(self): + g = _NoCPGroup() + assert g.rank() == 0 + + +try: + from transformer_engine.pytorch import cpp_extensions as _tex # noqa: F401 + + _HAS_TE = True +except Exception: + _HAS_TE = False + + +@pytest.mark.skipif(not _HAS_TE, reason="TransformerEngine not installed") +class TestThdCpPartition: + """Verify TE-based per-sample THD + CP partition matches THD semantics. + + Each packed sub-sample of length ``s_i`` (where ``s_i % (2*cp_size) == 0``) + is split into ``2*cp_size`` zigzag chunks per sample; rank ``r`` gets + chunks ``[r, 2*cp_size - r - 1]`` of every sample. The union across + ranks must cover every token position exactly once. + """ + + @staticmethod + def _make_padded_packed(seqlens, divisor): + """Concatenate per-sample dummy tokens after padding each sample to a + multiple of *divisor*. Returns ``(input_ids[1, T], cu_seqlens_padded)``. + """ + import math + padded = [math.ceil(s / divisor) * divisor for s in seqlens] + chunks = [] + next_id = 1 + for s, p in zip(seqlens, padded): + chunks.append(torch.arange(next_id, next_id + s, dtype=torch.int64)) + chunks.append(torch.zeros(p - s, dtype=torch.int64)) # padding + next_id += s + input_ids = torch.cat(chunks, dim=0).unsqueeze(0) # [1, T] + cu_seqlens_padded = torch.tensor( + [0] + list(torch.tensor(padded).cumsum(0).tolist()), + dtype=torch.int32, + ) + return input_ids, cu_seqlens_padded + + def _ensure_cuda(self, x): + return x.cuda() if torch.cuda.is_available() else x + + def test_partition_covers_all_positions_cp2(self): + from examples.multimodal_dev.models.base import _thd_cp_partition_index + + cp_size = 2 + seqlens = [5, 7, 3] # valid lengths + input_ids, cu_seqlens_padded = self._make_padded_packed( + seqlens, divisor=2 * cp_size, + ) + input_ids = self._ensure_cuda(input_ids) + cu_seqlens_padded = self._ensure_cuda(cu_seqlens_padded) + T = input_ids.shape[1] + + # Union of per-rank indices must be all positions exactly once. + seen = torch.zeros(T, dtype=torch.long, device=input_ids.device) + for cp_rank in range(cp_size): + idx = _thd_cp_partition_index( + cu_seqlens_padded, T, cp_size, cp_rank, + ) + assert idx.numel() == T // cp_size, ( + f"rank {cp_rank}: expected {T // cp_size} tokens, got {idx.numel()}" + ) + seen.scatter_add_( + 0, idx.long(), torch.ones_like(idx, dtype=seen.dtype), + ) + assert torch.all(seen == 1), ( + f"Position coverage broken: counts={seen.tolist()}" + ) + + def test_index_select_aligns_inputs_and_position_ids_cp2(self): + """input_ids, loss_mask, and (3, 1, T) position_ids index_select with + the same partition index produce shape-consistent per-rank tensors.""" + from examples.multimodal_dev.models.base import _thd_cp_partition_index + + cp_size = 2 + seqlens = [8, 4] + input_ids, cu_seqlens_padded = self._make_padded_packed( + seqlens, divisor=2 * cp_size, + ) + input_ids = self._ensure_cuda(input_ids) + cu_seqlens_padded = self._ensure_cuda(cu_seqlens_padded) + T = input_ids.shape[1] + labels = input_ids + 1000 + loss_mask = (input_ids != 0).float() + position_ids = ( + torch.arange(T, device=input_ids.device) + .unsqueeze(0).unsqueeze(0).expand(3, 1, T).contiguous() + ) + H = 4 + decoder_input = ( + torch.arange(T * H, dtype=torch.float32, device=input_ids.device) + .view(T, 1, H) + ) + + for cp_rank in range(cp_size): + idx = _thd_cp_partition_index( + cu_seqlens_padded, T, cp_size, cp_rank, + ) + ii = input_ids.index_select(1, idx) + ll = labels.index_select(1, idx) + lm = loss_mask.index_select(1, idx) + pi = position_ids.index_select(2, idx) + di = decoder_input.index_select(0, idx) + + assert ii.shape == (1, T // cp_size) + assert ll.shape == (1, T // cp_size) + assert lm.shape == (1, T // cp_size) + assert pi.shape == (3, 1, T // cp_size) + assert di.shape == (T // cp_size, 1, H) + # Sliced position_ids is just the partition index itself + # (since position_ids was arange(T) over all positions). + assert torch.equal(pi[0, 0], idx.to(pi.dtype)) + # All MRoPE rows agree. + assert torch.equal(pi[1, 0], pi[0, 0]) + assert torch.equal(pi[2, 0], pi[0, 0]) + + def test_partition_cp4_three_samples(self): + from examples.multimodal_dev.models.base import _thd_cp_partition_index + + cp_size = 4 + seqlens = [12, 4, 8] + input_ids, cu_seqlens_padded = self._make_padded_packed( + seqlens, divisor=2 * cp_size, + ) + input_ids = self._ensure_cuda(input_ids) + cu_seqlens_padded = self._ensure_cuda(cu_seqlens_padded) + T = input_ids.shape[1] + + seen = torch.zeros(T, dtype=torch.long, device=input_ids.device) + for cp_rank in range(cp_size): + idx = _thd_cp_partition_index( + cu_seqlens_padded, T, cp_size, cp_rank, + ) + assert idx.numel() == T // cp_size + seen.scatter_add_( + 0, idx.long(), torch.ones_like(idx, dtype=seen.dtype), + ) + assert torch.all(seen == 1) + + def test_loss_mask_zero_kept_per_rank(self): + """Pad-token positions (loss_mask=0) survive as 0 on whichever rank + they land — sanity check that we don't accidentally discard them.""" + from examples.multimodal_dev.models.base import _thd_cp_partition_index + + cp_size = 2 + seqlens = [5, 3] + input_ids, cu_seqlens_padded = self._make_padded_packed( + seqlens, divisor=2 * cp_size, + ) + input_ids = self._ensure_cuda(input_ids) + cu_seqlens_padded = self._ensure_cuda(cu_seqlens_padded) + T = input_ids.shape[1] + loss_mask = (input_ids != 0).float() + total_zeros = (loss_mask == 0).sum().item() + + zeros_seen = 0 + for cp_rank in range(cp_size): + idx = _thd_cp_partition_index( + cu_seqlens_padded, T, cp_size, cp_rank, + ) + zeros_seen += ( + loss_mask.index_select(1, idx) == 0 + ).sum().item() + assert zeros_seen == total_zeros diff --git a/examples/multimodal_dev/tests/test_cp_thd_correctness.py b/examples/multimodal_dev/tests/test_cp_thd_correctness.py new file mode 100644 index 00000000000..e815a948474 --- /dev/null +++ b/examples/multimodal_dev/tests/test_cp_thd_correctness.py @@ -0,0 +1,455 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# This is a stdout-reporting standalone script; `print` is intentional. +# pylint: disable=bad-builtin + +"""CP=1 vs CP=4 correctness test for THD and BSHD packing. + +Runs the production forward path (:class:`MultimodalModel`) twice in a +single ``torchrun`` invocation: + + Phase 1 — CP=1 baseline. All 4 ranks initialise with TP=1, CP=1 + (DP=4 implicit). Each rank computes the full sequence, + producing identical loss / grad_norm on every rank; rank 0's + value is the baseline. + + Phase 2 — CP=4. After ``destroy_model_parallel`` + ``initialize_model_parallel(CP=4)`` + the 4 ranks form a single CP group. The model's internal + ``_cp_split_for_forward`` slices inputs per rank; per-rank + loss / gradients are aggregated via AllReduce on the CP group. + +We compare CP=1 and CP=4 results for both BSHD and THD packing modes, +asserting that loss and grad_norm match within tolerance. + +Run with:: + + PYTHONPATH=. torchrun --nproc-per-node 4 \\ + examples/multimodal_dev/tests/test_cp_thd_correctness.py +""" + +import argparse +import os +import sys + +import torch + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +from examples.multimodal_dev.forward_step import pack_or_pad_batch +from examples.multimodal_dev.models.base import ( + MultimodalModel, + _cp_split_tensor, + _thd_cp_partition_index, +) +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec +from megatron.core.parallel_state import get_context_parallel_group, get_context_parallel_rank +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.transformer_config import TransformerConfig +from tests.unit_tests.test_utilities import Utils + +# =================================================================== +# Stub vision encoder +# =================================================================== + + +class _StubVisionEncoder(MegatronModule): + """Vision encoder placeholder. The vision branch is skipped in + :meth:`MultimodalModel.forward` whenever ``pixel_values is None``, so + this module is never called — it only satisfies the constructor's + ``vision_encoder: MegatronModule`` requirement. + """ + + def __init__(self, config): + """Initialise the stub with the given TransformerConfig.""" + super().__init__(config=config) + + def forward(self, pixel_values, image_grid_thw): + """Never called when ``pixel_values=None``; raises if it ever is.""" + raise RuntimeError("vision branch should not run when pixel_values=None") + + +# =================================================================== +# Model builder +# =================================================================== + + +def _build_model(config, vocab_size, max_seq_len, image_token_id): + spec = get_gpt_layer_with_transformer_engine_spec() + vision = _StubVisionEncoder(config) + model = MultimodalModel( + language_config=config, + language_spec=spec, + vision_encoder=vision, + vocab_size=vocab_size, + max_sequence_length=max_seq_len, + image_token_id=image_token_id, + position_embedding_type="rope", + parallel_output=False, + ) + model.cuda() + return model + + +def _make_config( + num_layers, hidden_size, ffn_hidden_size, num_heads, num_kv_heads, context_parallel_size +): + return TransformerConfig( + num_layers=num_layers, + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + num_attention_heads=num_heads, + num_query_groups=num_kv_heads, + bf16=True, + params_dtype=torch.bfloat16, + pipeline_dtype=torch.bfloat16, + hidden_dropout=0.0, + attention_dropout=0.0, + tensor_model_parallel_size=1, + context_parallel_size=context_parallel_size, + sequence_parallel=False, + ) + + +# =================================================================== +# Loss / grad-norm aggregation +# =================================================================== + + +def _global_loss(output, rank_loss_mask, cp_size): + """Mean per-token loss over all CP shards (matches CP=1 mean exactly).""" + num = (output.float().view(-1) * rank_loss_mask.float().view(-1)).sum() + den = rank_loss_mask.float().view(-1).sum().clamp(min=1) + if cp_size > 1: + group = get_context_parallel_group() + torch.distributed.all_reduce(num, group=group) + torch.distributed.all_reduce(den, group=group) + return (num / den).item() + + +def _global_grad_norm(model, cp_size): + """Global L2 grad norm. For CP>1, AllReduce(SUM) gradients across CP + then divide by ``cp_size`` so each rank holds the CP-mean gradient + (matching CP=1's behaviour, where backward on the per-batch mean loss + yields exactly that gradient). + """ + if cp_size > 1: + group = get_context_parallel_group() + for p in model.parameters(): + if p.grad is not None: + torch.distributed.all_reduce(p.grad, group=group) + p.grad /= cp_size + + sq = 0.0 + for p in model.parameters(): + if p.grad is not None: + sq += p.grad.float().norm(2).item() ** 2 + return sq**0.5 + + +# =================================================================== +# Data — identical across all ranks (deterministic generator) +# =================================================================== + + +def _make_data(B, S, vocab_size, image_token_id, seed): + """Same input on every rank thanks to the seeded generator.""" + g = torch.Generator(device="cuda") + g.manual_seed(seed) + input_ids = torch.randint(0, vocab_size, (B, S), generator=g, device="cuda") + # Ensure no accidental image tokens (we never run the vision branch). + input_ids = torch.where(input_ids == image_token_id, (input_ids + 1) % vocab_size, input_ids) + labels = torch.randint(0, vocab_size, (B, S), generator=g, device="cuda") + loss_mask = torch.ones(B, S, device="cuda") + position_ids = torch.arange(S, device="cuda").unsqueeze(0).expand(B, -1).contiguous() + return input_ids, labels, loss_mask, position_ids + + +# =================================================================== +# One BSHD or THD forward+backward, returning (loss, grad_norm) +# =================================================================== + + +def _run_bshd(model, B, S, vocab_size, image_token_id, cp_size, seed): + input_ids, labels, loss_mask, position_ids = _make_data(B, S, vocab_size, image_token_id, seed) + + output = model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=None, + labels=labels, + loss_mask=loss_mask, + pixel_values=None, + image_grid_thw=None, + packed_seq_params=None, + ) + + # Slice loss_mask the same way forward_step does for BSHD + CP. + rank_loss_mask = loss_mask + if cp_size > 1: + rank_loss_mask = _cp_split_tensor( + rank_loss_mask, seq_dim=1, cp_size=cp_size, cp_rank=get_context_parallel_rank() + ) + + loss_val = _global_loss(output, rank_loss_mask, cp_size) + + # Backward on the LOCAL mean loss (each rank's contribution + # equal-weighted; SUM-then-divide across CP recovers CP=1's gradient). + local = ( + output.float().view(-1) * rank_loss_mask.float().view(-1) + ).sum() / rank_loss_mask.float().view(-1).sum().clamp(min=1) + model.zero_grad() + local.backward() + + gn = _global_grad_norm(model, cp_size) + return loss_val, gn + + +def _run_thd(model, B, S, vocab_size, image_token_id, cp_size, seed): + input_ids, labels, loss_mask, _ = _make_data(B, S, vocab_size, image_token_id, seed) + + # Build the per-sample dict list and pack to [1, T]. + samples = [] + for i in range(B): + samples.append( + { + "input_ids": input_ids[i].clone(), + "labels": labels[i].clone(), + "loss_mask": loss_mask[i].clone(), + # No vision; empty tensors satisfy pack_or_pad_batch. + "pixel_values": torch.zeros(0, 1, device="cuda"), + "image_grid_thw": torch.empty(0, 3, dtype=torch.long, device="cuda"), + } + ) + packed = pack_or_pad_batch(samples, use_packed_sequence=True, device="cuda") + psp = packed.pop("packed_seq_params") + + # THD position_ids: per-sample restart at 0. Each sample has length S + # (equal-length data), so this is arange(S) repeated B times. + thd_pos = ( + torch.cat([torch.arange(S, device="cuda") for _ in range(B)]).unsqueeze(0).contiguous() + ) + + output = model( + input_ids=packed["input_ids"], + position_ids=thd_pos, + attention_mask=None, + labels=packed["labels"], + loss_mask=packed["loss_mask"], + pixel_values=None, + image_grid_thw=None, + packed_seq_params=psp, + ) + + rank_loss_mask = packed["loss_mask"] + if cp_size > 1: + T = rank_loss_mask.shape[1] + idx = _thd_cp_partition_index( + psp.cu_seqlens_q_padded, T, cp_size, get_context_parallel_rank() + ) + rank_loss_mask = rank_loss_mask.index_select(1, idx) + + loss_val = _global_loss(output, rank_loss_mask, cp_size) + + local = ( + output.float().view(-1) * rank_loss_mask.float().view(-1) + ).sum() / rank_loss_mask.float().view(-1).sum().clamp(min=1) + model.zero_grad() + local.backward() + + gn = _global_grad_norm(model, cp_size) + return loss_val, gn + + +# =================================================================== +# State-dict roundtrip — keep weights identical across phases +# =================================================================== + + +def _cpu_state_dict(model): + """Snapshot of model.state_dict() detached to CPU (kept in memory). + + Some entries (TransformerEngine ``_extra_state``) are non-tensor or + ``None``; pass them through untouched. + """ + snap = {} + for k, v in model.state_dict().items(): + if isinstance(v, torch.Tensor): + snap[k] = v.detach().to("cpu").clone() + else: + snap[k] = v + return snap + + +def _restore_state_dict(model, snapshot): + """Load a saved snapshot back into a freshly built model.""" + payload = {k: (v.to("cuda") if isinstance(v, torch.Tensor) else v) for k, v in snapshot.items()} + model.load_state_dict(payload) + + +# =================================================================== +# Main +# =================================================================== + + +def _is_rank0(): + return not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 + + +def _print_banner(title): + if _is_rank0(): + print(f"\n{'=' * 60}") + print(f" {title}") + print(f"{'=' * 60}") + + +def _print_compare(label, baseline, trial, atol, rtol): + """Print a CP=1 vs CP=4 comparison line and return whether it passed.""" + if not _is_rank0(): + return True + + abs_diff = abs(baseline - trial) + rel_diff = abs_diff / max(abs(baseline), 1e-8) + ok = abs_diff < atol or rel_diff < rtol + flag = "PASS" if ok else "FAIL" + print( + f" {label:<30s} CP=1: {baseline:.8f} CP=4: {trial:.8f}" + f" abs={abs_diff:.2e} rel={rel_diff:.2e} [{flag}]" + ) + return ok + + +def main(): + """Run CP=1 baseline + CP=4 trial and compare losses / grad_norms.""" + parser = argparse.ArgumentParser() + parser.add_argument("--batch-size", type=int, default=2) + # Must be divisible by 2*cp_size (=8 for CP=4 zigzag). + parser.add_argument("--seq-len", type=int, default=64) + parser.add_argument("--vocab-size", type=int, default=1024) + parser.add_argument("--hidden-size", type=int, default=256) + parser.add_argument("--num-layers", type=int, default=2) + parser.add_argument("--num-heads", type=int, default=4) + parser.add_argument("--num-kv-heads", type=int, default=2) + parser.add_argument("--ffn-hidden-size", type=int, default=512) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--atol-loss", type=float, default=1e-3) + parser.add_argument("--rtol-grad", type=float, default=5e-3) + parser.add_argument("--data-seed", type=int, default=123) + args = parser.parse_args() + + image_token_id = 0 # never appears in input (data filters this id out) + + # ---------------------------------------------------------------- + # Phase 1: CP=1 baseline + # ---------------------------------------------------------------- + _print_banner("Phase 1 — building CP=1 baseline (TP=1, CP=1, DP=4)") + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=1) + model_parallel_cuda_manual_seed(args.seed) + + config_cp1 = _make_config( + args.num_layers, + args.hidden_size, + args.ffn_hidden_size, + args.num_heads, + args.num_kv_heads, + context_parallel_size=1, + ) + torch.manual_seed(args.seed) + model_cp1 = _build_model(config_cp1, args.vocab_size, args.seq_len, image_token_id) + + bshd_loss_cp1, bshd_gn_cp1 = _run_bshd( + model_cp1, + args.batch_size, + args.seq_len, + args.vocab_size, + image_token_id, + cp_size=1, + seed=args.data_seed, + ) + thd_loss_cp1, thd_gn_cp1 = _run_thd( + model_cp1, + args.batch_size, + args.seq_len, + args.vocab_size, + image_token_id, + cp_size=1, + seed=args.data_seed, + ) + + # Snapshot weights *before* the optimizer would have touched them. + # (We've zeroed grads but never stepped; weights at this point are + # the just-initialised baseline.) + weights_snapshot = _cpu_state_dict(model_cp1) + del model_cp1 + torch.cuda.empty_cache() + + # ---------------------------------------------------------------- + # Phase 2: CP=4 + # ---------------------------------------------------------------- + _print_banner("Phase 2 — re-initialising for CP=4 (TP=1, CP=4)") + Utils.destroy_model_parallel() + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=4) + model_parallel_cuda_manual_seed(args.seed) + + config_cp4 = _make_config( + args.num_layers, + args.hidden_size, + args.ffn_hidden_size, + args.num_heads, + args.num_kv_heads, + context_parallel_size=4, + ) + torch.manual_seed(args.seed) + model_cp4 = _build_model(config_cp4, args.vocab_size, args.seq_len, image_token_id) + _restore_state_dict(model_cp4, weights_snapshot) + + bshd_loss_cp4, bshd_gn_cp4 = _run_bshd( + model_cp4, + args.batch_size, + args.seq_len, + args.vocab_size, + image_token_id, + cp_size=4, + seed=args.data_seed, + ) + thd_loss_cp4, thd_gn_cp4 = _run_thd( + model_cp4, + args.batch_size, + args.seq_len, + args.vocab_size, + image_token_id, + cp_size=4, + seed=args.data_seed, + ) + + # ---------------------------------------------------------------- + # Compare + # ---------------------------------------------------------------- + _print_banner("Results — CP=1 vs CP=4") + all_ok = True + all_ok &= _print_compare( + "BSHD loss", bshd_loss_cp1, bshd_loss_cp4, args.atol_loss, args.rtol_grad + ) + all_ok &= _print_compare( + "BSHD grad_norm", bshd_gn_cp1, bshd_gn_cp4, args.atol_loss, args.rtol_grad + ) + all_ok &= _print_compare( + "THD loss", thd_loss_cp1, thd_loss_cp4, args.atol_loss, args.rtol_grad + ) + all_ok &= _print_compare( + "THD grad_norm", thd_gn_cp1, thd_gn_cp4, args.atol_loss, args.rtol_grad + ) + + _print_banner("Summary") + if _is_rank0(): + print(f" {'ALL TESTS PASSED' if all_ok else 'SOME TESTS FAILED'}") + print(f"{'=' * 60}\n") + + Utils.destroy_model_parallel() + if not all_ok: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/multimodal_dev/tests/test_mrope_parity.py b/examples/multimodal_dev/tests/test_mrope_parity.py new file mode 100644 index 00000000000..154284a4fc1 --- /dev/null +++ b/examples/multimodal_dev/tests/test_mrope_parity.py @@ -0,0 +1,663 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Parity tests for ``get_rope_index`` (MRoPE position-ID computation). + +Two properties are verified: + +1. **BSHD backwards compatibility** — the refactored ``get_rope_index`` + returns bit-identical ``(position_ids, mrope_position_deltas)`` to + the pre-refactor implementation on padded ``[B, S]`` batches. +2. **THD == BSHD on the valid region** — when the same variable-length + samples are fed through both layouts (BSHD with right-padding; THD + packed with ``cu_seqlens_q`` / ``cu_seqlens_q_padded``), positions at + every valid slot agree. + +The pre-refactor function is pinned inline as ``_old_get_rope_index`` +so this test stays self-contained. Run with:: + + python -m pytest examples/multimodal_dev/tests/test_mrope_parity.py -v + +or directly:: + + python examples/multimodal_dev/tests/test_mrope_parity.py +""" + +import math +import os +import sys +from itertools import accumulate + +import torch +import torch.nn.functional as F + +_REPO_ROOT = os.path.abspath( + os.path.join(os.path.dirname(__file__), "../../.."), +) +# Insert at position 0 unconditionally — other entries on sys.path +# (e.g. a sibling Megatron-LM checkout) have their own ``examples`` +# package that would otherwise shadow ours. +if _REPO_ROOT in sys.path: + sys.path.remove(_REPO_ROOT) +sys.path.insert(0, _REPO_ROOT) + +from megatron.core.packed_seq_params import PackedSeqParams + +from examples.multimodal_dev.models.qwen35_vl.mrope import get_rope_index + +# ----------------------------------------------------------------------------- +# Token-ID constants (match Qwen3.5-VL, but values are arbitrary for this test) +# ----------------------------------------------------------------------------- + +IMAGE_TOKEN_ID = 248056 +VIDEO_TOKEN_ID = 248057 +VISION_START_TOKEN_ID = 248053 +SPATIAL_MERGE_SIZE = 2 + + +# ----------------------------------------------------------------------------- +# Pinned reference implementation (pre-refactor BSHD path) +# ----------------------------------------------------------------------------- + +def _old_get_rope_index( + spatial_merge_size, + image_token_id, + video_token_id, + vision_start_token_id, + input_ids=None, + image_grid_thw=None, + video_grid_thw=None, + attention_mask=None, +): + """Pre-refactor BSHD implementation of ``get_rope_index``. + + Copied verbatim (modulo the broken cu_seqlens branch, which this + parity test does not exercise) so we can diff against the new + implementation on BSHD inputs. + """ + if video_grid_thw is not None: + video_grid_thw = torch.repeat_interleave( + video_grid_thw, video_grid_thw[:, 0], dim=0, + ) + video_grid_thw[:, 0] = 1 + + mrope_position_deltas = [] + + if input_ids is not None and ( + image_grid_thw is not None or video_grid_thw is not None + ): + total_input_ids = input_ids + if attention_mask is None: + attention_mask = torch.ones_like(total_input_ids) + + position_ids = torch.ones( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + image_index, video_index = 0, 0 + attention_mask = attention_mask.to(total_input_ids.device) + + for i, sample_input_ids in enumerate(total_input_ids): + sample_input_ids = sample_input_ids[attention_mask[i] == 1] + vision_start_indices = torch.argwhere( + sample_input_ids == vision_start_token_id, + ).squeeze(1) + vision_tokens = sample_input_ids[vision_start_indices + 1] + image_nums = (vision_tokens == image_token_id).sum() + video_nums = (vision_tokens == video_token_id).sum() + input_tokens = sample_input_ids.tolist() + llm_pos_ids_list = [] + st = 0 + remain_images, remain_videos = image_nums, video_nums + + for _ in range(image_nums + video_nums): + if image_token_id in input_tokens and remain_images > 0: + ed_image = input_tokens.index(image_token_id, st) + else: + ed_image = len(input_tokens) + 1 + if video_token_id in input_tokens and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + + if ed_image < ed_video: + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + image_index += 1 + remain_images -= 1 + ed = ed_image + else: + t, h, w = ( + video_grid_thw[video_index][0], + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + video_index += 1 + remain_videos -= 1 + ed = ed_video + + llm_grid_t, llm_grid_h, llm_grid_w = ( + t.item(), + h.item() // spatial_merge_size, + w.item() // spatial_merge_size, + ) + text_len = ed - st + + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if llm_pos_ids_list + else 0 + ) + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + + st_idx + ) + + t_index = ( + torch.arange(llm_grid_t) + .view(-1, 1) + .expand(-1, llm_grid_h * llm_grid_w) + .flatten() + ) + h_index = ( + torch.arange(llm_grid_h) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + + text_len + + st_idx + ) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + + if st < len(input_tokens): + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if llm_pos_ids_list + else 0 + ) + text_len = len(input_tokens) - st + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + + st_idx + ) + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + position_ids[ + ..., i, attention_mask[i] == 1 + ] = llm_positions.to(position_ids.device) + mrope_position_deltas.append( + llm_positions.max() + 1 - len(total_input_ids[i]), + ) + + mrope_position_deltas = torch.tensor( + mrope_position_deltas, device=total_input_ids.device, + ).unsqueeze(1) + return position_ids, mrope_position_deltas + + # Text-only fallback. + if attention_mask is not None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + position_ids = ( + position_ids.unsqueeze(0) + .expand(3, -1, -1) + .to(attention_mask.device) + ) + max_position_ids = ( + position_ids.max(0, keepdim=False)[0] + .max(-1, keepdim=True)[0] + ) + mrope_position_deltas = ( + max_position_ids + 1 - attention_mask.shape[-1] + ) + else: + position_ids = ( + torch.arange( + input_ids.shape[1], device=input_ids.device, + ) + .view(1, 1, -1) + .expand(3, input_ids.shape[0], -1) + ) + mrope_position_deltas = torch.zeros( + [input_ids.shape[0], 1], + device=input_ids.device, + dtype=input_ids.dtype, + ) + return position_ids, mrope_position_deltas + + +# ----------------------------------------------------------------------------- +# Synthetic-sample builder +# ----------------------------------------------------------------------------- + +def _build_sample( + prefix_text_len, + grids, + suffix_text_len, + text_token=100, +): + """Build one variable-length sample with ``len(grids)`` images. + + Layout per image: ``vision_start_id`` then + ``llm_grid_t * llm_grid_h * llm_grid_w`` ``image_token_id`` slots + (where ``llm_grid_* = grid_* // spatial_merge_size`` for h/w). + Grids use ``t=1``. + + Returns ``(input_ids [L], image_grid_thw [N, 3])``. + """ + tokens = [text_token] * prefix_text_len + grid_rows = [] + for t, h, w in grids: + n_image_tokens = ( + t * (h // SPATIAL_MERGE_SIZE) * (w // SPATIAL_MERGE_SIZE) + ) + tokens.append(VISION_START_TOKEN_ID) + tokens.extend([IMAGE_TOKEN_ID] * n_image_tokens) + grid_rows.append([t, h, w]) + tokens.extend([text_token + 1] * suffix_text_len) + input_ids = torch.tensor(tokens, dtype=torch.int64) + image_grid_thw = torch.tensor(grid_rows, dtype=torch.int64) + return input_ids, image_grid_thw + + +def _sample_bank(): + """A small bank of samples covering text-only, single-image, multi-image.""" + return [ + _build_sample( + prefix_text_len=5, + grids=[(1, 4, 4)], + suffix_text_len=7, + ), + _build_sample( + prefix_text_len=3, + grids=[(1, 2, 2), (1, 4, 6)], + suffix_text_len=4, + ), + _build_sample( + prefix_text_len=10, + grids=[], + suffix_text_len=0, + ), + _build_sample( + prefix_text_len=0, + grids=[(1, 6, 4)], + suffix_text_len=2, + ), + ] + + +# ----------------------------------------------------------------------------- +# Test 1: BSHD backwards compatibility +# ----------------------------------------------------------------------------- + +def test_bshd_matches_old_reference(): + """New ``get_rope_index`` equals the pinned reference on BSHD inputs.""" + samples = _sample_bank() + max_len = max(s.numel() for s, _ in samples) + + input_ids_rows = [] + mask_rows = [] + grid_rows = [] + for tokens, grids in samples: + L = tokens.numel() + padded = F.pad(tokens, (0, max_len - L), value=0) + mask = torch.zeros(max_len, dtype=torch.int64) + mask[:L] = 1 + input_ids_rows.append(padded) + mask_rows.append(mask) + if grids.numel() > 0: + grid_rows.append(grids) + + input_ids = torch.stack(input_ids_rows) # [B, S] + attention_mask = torch.stack(mask_rows) # [B, S] + image_grid_thw = ( + torch.cat(grid_rows, dim=0) if grid_rows else None + ) + + old_pos, old_delta = _old_get_rope_index( + spatial_merge_size=SPATIAL_MERGE_SIZE, + image_token_id=IMAGE_TOKEN_ID, + video_token_id=VIDEO_TOKEN_ID, + vision_start_token_id=VISION_START_TOKEN_ID, + input_ids=input_ids, + image_grid_thw=image_grid_thw, + attention_mask=attention_mask, + ) + new_pos, new_delta = get_rope_index( + spatial_merge_size=SPATIAL_MERGE_SIZE, + image_token_id=IMAGE_TOKEN_ID, + video_token_id=VIDEO_TOKEN_ID, + vision_start_token_id=VISION_START_TOKEN_ID, + input_ids=input_ids, + image_grid_thw=image_grid_thw, + attention_mask=attention_mask, + packed_seq_params=None, + ) + + assert torch.equal(old_pos, new_pos), ( + f"BSHD position_ids differ.\nold:\n{old_pos}\nnew:\n{new_pos}" + ) + assert torch.equal(old_delta, new_delta), ( + f"BSHD mrope_position_deltas differ.\n" + f"old: {old_delta}\nnew: {new_delta}" + ) + + +# ----------------------------------------------------------------------------- +# Test 2: THD positions match BSHD positions on the valid region +# ----------------------------------------------------------------------------- + +def _pack_samples(samples, divisible_by=1): + """Pack ``samples`` into a single ``[1, T]`` tensor the same way + ``pack_or_pad_batch`` does, and build ``PackedSeqParams``. + + Each per-sample tensor is right-padded to a multiple of + ``divisible_by`` before concatenation. ``cu_seqlens_q`` tracks + unpadded lengths; ``cu_seqlens_q_padded`` tracks the packed layout. + """ + padded_chunks = [] + seqlens = [] + seqlens_padded = [] + grid_rows = [] + for tokens, grids in samples: + L = tokens.numel() + target_L = math.ceil(L / divisible_by) * divisible_by + padded_chunks.append(F.pad(tokens, (0, target_L - L), value=0)) + seqlens.append(L) + seqlens_padded.append(target_L) + if grids.numel() > 0: + grid_rows.append(grids) + + packed = torch.cat(padded_chunks, dim=0).unsqueeze(0) # [1, T] + cu_seqlens = torch.tensor( + list(accumulate(seqlens, initial=0)), dtype=torch.int32, + ) + cu_seqlens_padded = torch.tensor( + list(accumulate(seqlens_padded, initial=0)), dtype=torch.int32, + ) + psp = PackedSeqParams( + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=max(seqlens_padded), + max_seqlen_kv=max(seqlens_padded), + ) + image_grid_thw = ( + torch.cat(grid_rows, dim=0) if grid_rows else None + ) + return packed, psp, image_grid_thw, seqlens, seqlens_padded + + +def test_thd_matches_bshd_padded(): + """THD positions at every valid slot equal BSHD positions on the + equivalent right-padded batch. + """ + samples = _sample_bank() + + # BSHD side: right-pad to common max_len. + max_len = max(s.numel() for s, _ in samples) + input_ids_rows = [] + mask_rows = [] + grid_rows = [] + for tokens, grids in samples: + L = tokens.numel() + input_ids_rows.append(F.pad(tokens, (0, max_len - L), value=0)) + m = torch.zeros(max_len, dtype=torch.int64) + m[:L] = 1 + mask_rows.append(m) + if grids.numel() > 0: + grid_rows.append(grids) + bshd_input_ids = torch.stack(input_ids_rows) + bshd_mask = torch.stack(mask_rows) + bshd_grid = torch.cat(grid_rows, dim=0) if grid_rows else None + + bshd_pos, _ = get_rope_index( + spatial_merge_size=SPATIAL_MERGE_SIZE, + image_token_id=IMAGE_TOKEN_ID, + video_token_id=VIDEO_TOKEN_ID, + vision_start_token_id=VISION_START_TOKEN_ID, + input_ids=bshd_input_ids, + image_grid_thw=bshd_grid, + attention_mask=bshd_mask, + ) + # bshd_pos: [3, B, S_pad] + + # THD side: pack with a non-trivial divisor so the padded and + # unpadded cu_seqlens diverge — this exercises the distinction. + for divisible_by in (1, 4): + packed_input_ids, psp, thd_grid, seqlens, seqlens_padded = ( + _pack_samples(samples, divisible_by=divisible_by) + ) + thd_pos, _ = get_rope_index( + spatial_merge_size=SPATIAL_MERGE_SIZE, + image_token_id=IMAGE_TOKEN_ID, + video_token_id=VIDEO_TOKEN_ID, + vision_start_token_id=VISION_START_TOKEN_ID, + input_ids=packed_input_ids, + image_grid_thw=thd_grid, + packed_seq_params=psp, + ) + # thd_pos: [3, 1, T] + assert thd_pos.shape == ( + 3, 1, packed_input_ids.shape[1], + ), f"bad THD shape {thd_pos.shape}" + + seg_starts = list(accumulate(seqlens_padded, initial=0)) + for k, (valid_len, seg_start) in enumerate( + zip(seqlens, seg_starts) + ): + thd_slice = thd_pos[:, 0, seg_start:seg_start + valid_len] + bshd_slice = bshd_pos[:, k, :valid_len] + assert torch.equal(thd_slice, bshd_slice), ( + f"[divisible_by={divisible_by}] segment {k} " + f"(valid_len={valid_len}, seg_start={seg_start}) " + f"disagrees:\nTHD:\n{thd_slice}\nBSHD:\n{bshd_slice}" + ) + + +# ----------------------------------------------------------------------------- +# Test 3: THD with no images (text-only packed) +# ----------------------------------------------------------------------------- + +def test_thd_text_only_restarts_per_segment(): + """Text-only THD: each segment gets a fresh ``[0..valid_len-1]`` range.""" + samples = [ + _build_sample(prefix_text_len=6, grids=[], suffix_text_len=0), + _build_sample(prefix_text_len=11, grids=[], suffix_text_len=0), + _build_sample(prefix_text_len=3, grids=[], suffix_text_len=0), + ] + packed_input_ids, psp, _, seqlens, seqlens_padded = _pack_samples( + samples, divisible_by=4, + ) + thd_pos, _ = get_rope_index( + spatial_merge_size=SPATIAL_MERGE_SIZE, + image_token_id=IMAGE_TOKEN_ID, + video_token_id=VIDEO_TOKEN_ID, + vision_start_token_id=VISION_START_TOKEN_ID, + input_ids=packed_input_ids, + image_grid_thw=None, + packed_seq_params=psp, + ) + + seg_starts = list(accumulate(seqlens_padded, initial=0)) + for valid_len, seg_start in zip(seqlens, seg_starts): + expected = ( + torch.arange(valid_len, dtype=thd_pos.dtype) + .view(1, -1) + .expand(3, -1) + ) + got = thd_pos[:, 0, seg_start:seg_start + valid_len] + assert torch.equal(got, expected), ( + f"text-only segment mismatch at seg_start={seg_start}, " + f"valid_len={valid_len}:\n{got}\nexpected:\n{expected}" + ) + + +# ----------------------------------------------------------------------------- +# Test 4: Explicit two-sequence batch with vision, both in BSHD and THD +# ----------------------------------------------------------------------------- + +def _two_image_samples(): + """Two samples, each with one image — the smallest case that can + expose a bug where segment k > 0 positions leak state from segment + k - 1 (e.g. non-restarted ``st_idx`` or a stale ``image_index``). + """ + return [ + _build_sample( + prefix_text_len=5, + grids=[(1, 4, 4)], # 4 image tokens after spatial merge + suffix_text_len=3, + ), + _build_sample( + prefix_text_len=4, + grids=[(1, 4, 4)], + suffix_text_len=6, + ), + ] + + +def test_bshd_batch_size_2_with_vision(): + """BSHD with ``B == 2``: both rows' positions restart at 0 and match + the pinned reference. + """ + samples = _two_image_samples() + max_len = max(s.numel() for s, _ in samples) + + input_ids_rows, mask_rows, grid_rows = [], [], [] + for tokens, grids in samples: + L = tokens.numel() + input_ids_rows.append(F.pad(tokens, (0, max_len - L), value=0)) + m = torch.zeros(max_len, dtype=torch.int64) + m[:L] = 1 + mask_rows.append(m) + grid_rows.append(grids) + + input_ids = torch.stack(input_ids_rows) # [2, S] + attention_mask = torch.stack(mask_rows) + image_grid_thw = torch.cat(grid_rows, dim=0) + assert input_ids.shape[0] == 2 + + old_pos, old_delta = _old_get_rope_index( + spatial_merge_size=SPATIAL_MERGE_SIZE, + image_token_id=IMAGE_TOKEN_ID, + video_token_id=VIDEO_TOKEN_ID, + vision_start_token_id=VISION_START_TOKEN_ID, + input_ids=input_ids, + image_grid_thw=image_grid_thw, + attention_mask=attention_mask, + ) + new_pos, new_delta = get_rope_index( + spatial_merge_size=SPATIAL_MERGE_SIZE, + image_token_id=IMAGE_TOKEN_ID, + video_token_id=VIDEO_TOKEN_ID, + vision_start_token_id=VISION_START_TOKEN_ID, + input_ids=input_ids, + image_grid_thw=image_grid_thw, + attention_mask=attention_mask, + ) + assert torch.equal(old_pos, new_pos), ( + f"BSHD B=2 position mismatch vs reference.\n" + f"old:\n{old_pos}\nnew:\n{new_pos}" + ) + assert torch.equal(old_delta, new_delta) + + # Both rows must start at position 0. + for i in range(2): + valid_len = int(attention_mask[i].sum().item()) + assert torch.all(new_pos[:, i, 0] == 0), ( + f"row {i} does not start at 0: {new_pos[:, i, 0]}" + ) + # Sanity: positions within the valid region are strictly < valid_len + # would be wrong (MRoPE can skip positions), so just check max. + assert new_pos[:, i, :valid_len].max() < valid_len + + +def test_thd_batch_size_2_with_vision(): + """THD with 2 packed sequences: seg 1 positions restart at 0 and + equal BSHD row 1 on the valid region (bit-identical). + """ + samples = _two_image_samples() + + # BSHD reference. + max_len = max(s.numel() for s, _ in samples) + rows, masks, grids_bshd = [], [], [] + for tokens, grids in samples: + L = tokens.numel() + rows.append(F.pad(tokens, (0, max_len - L), value=0)) + m = torch.zeros(max_len, dtype=torch.int64) + m[:L] = 1 + masks.append(m) + grids_bshd.append(grids) + bshd_input_ids = torch.stack(rows) + bshd_mask = torch.stack(masks) + bshd_grid = torch.cat(grids_bshd, dim=0) + bshd_pos, _ = get_rope_index( + spatial_merge_size=SPATIAL_MERGE_SIZE, + image_token_id=IMAGE_TOKEN_ID, + video_token_id=VIDEO_TOKEN_ID, + vision_start_token_id=VISION_START_TOKEN_ID, + input_ids=bshd_input_ids, + image_grid_thw=bshd_grid, + attention_mask=bshd_mask, + ) + + # THD packed version with a non-trivial divisor so padded and + # unpadded cu_seqlens disagree. + packed_input_ids, psp, thd_grid, seqlens, seqlens_padded = ( + _pack_samples(samples, divisible_by=4) + ) + assert len(seqlens) == 2 + thd_pos, _ = get_rope_index( + spatial_merge_size=SPATIAL_MERGE_SIZE, + image_token_id=IMAGE_TOKEN_ID, + video_token_id=VIDEO_TOKEN_ID, + vision_start_token_id=VISION_START_TOKEN_ID, + input_ids=packed_input_ids, + image_grid_thw=thd_grid, + packed_seq_params=psp, + ) + + seg_starts = list(accumulate(seqlens_padded, initial=0)) + for k, (valid_len, seg_start) in enumerate( + zip(seqlens, seg_starts) + ): + thd_slice = thd_pos[:, 0, seg_start:seg_start + valid_len] + bshd_slice = bshd_pos[:, k, :valid_len] + assert torch.equal(thd_slice, bshd_slice), ( + f"seg {k} THD vs BSHD row {k} mismatch.\n" + f"THD:\n{thd_slice}\nBSHD:\n{bshd_slice}" + ) + # Critical: seg k must start at position 0 (bug 2 check). + assert torch.all(thd_slice[:, 0] == 0), ( + f"seg {k} does not start at 0 — positions leaked from " + f"previous segment: first col = {thd_slice[:, 0]}" + ) + + +if __name__ == "__main__": + test_bshd_matches_old_reference() + print("[ok] test_bshd_matches_old_reference") + test_thd_matches_bshd_padded() + print("[ok] test_thd_matches_bshd_padded") + test_thd_text_only_restarts_per_segment() + print("[ok] test_thd_text_only_restarts_per_segment") + test_bshd_batch_size_2_with_vision() + print("[ok] test_bshd_batch_size_2_with_vision") + test_thd_batch_size_2_with_vision() + print("[ok] test_thd_batch_size_2_with_vision") + print("All parity tests passed.") diff --git a/examples/multimodal_dev/tests/test_thd_correctness.py b/examples/multimodal_dev/tests/test_thd_correctness.py new file mode 100644 index 00000000000..1f92cc2e2f7 --- /dev/null +++ b/examples/multimodal_dev/tests/test_thd_correctness.py @@ -0,0 +1,380 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# This is a stdout-reporting standalone script; `print` is intentional. +# pylint: disable=bad-builtin + +"""BSHD vs THD correctness test for multimodal_dev packed sequence support. + +Validates that packing a [B, S] batch into [1, T] THD format produces +numerically equivalent loss values and gradient norms through a GPTModel. + +The test uses equal-length sequences (no padding) so that BSHD causal +attention and THD cu_seqlens-based causal attention are mathematically +identical. This makes any numerical deviation a real bug rather than an +expected consequence of different padding/masking strategies. + +Usage:: + + # Single GPU (flash attention): + torchrun --nproc_per_node=1 \\ + examples/multimodal_dev/tests/test_thd_correctness.py + + # Override model size: + torchrun --nproc_per_node=1 \\ + examples/multimodal_dev/tests/test_thd_correctness.py \\ + --num-layers 4 --hidden-size 512 --num-heads 8 --num-kv-heads 4 +""" + +import argparse +import os +import sys + +import torch + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +from examples.multimodal_dev.forward_step import pack_or_pad_batch +from examples.multimodal_dev.tests._helpers import grad_norm, mean_loss +from megatron.core.models.gpt import GPTModel +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.transformer_config import TransformerConfig +from tests.unit_tests.test_utilities import Utils + + +def _samples_from_bshd(input_ids, labels, loss_mask, seq_lengths=None): + """Slice ``[B, S]`` tensors into the per-sample dict list that + ``pack_or_pad_batch`` expects. ``seq_lengths`` lets variable-length + samples carry only their valid tokens (no attention_mask needed). + """ + B, S = input_ids.shape + if seq_lengths is None: + seq_lengths = [S] * B + samples = [] + for i, L in enumerate(seq_lengths): + samples.append( + { + "input_ids": input_ids[i, :L].clone(), + "labels": labels[i, :L].clone(), + "loss_mask": loss_mask[i, :L].clone(), + # pack_or_pad_batch requires these keys but the model call + # below ignores them; provide minimal dummies. + "pixel_values": torch.zeros(1, 1, device=input_ids.device), + "image_grid_thw": torch.tensor( + [[2, 1, 1]], dtype=torch.long, device=input_ids.device + ), + } + ) + return samples + + +def _thd_position_ids(seq_lengths, device): + """Build ``[1, T]`` THD position_ids: each segment restarts at 0.""" + return ( + torch.cat([torch.arange(L, device=device) for L in seq_lengths]).unsqueeze(0).contiguous() + ) + + +# =================================================================== +# Helpers +# =================================================================== + + +def _build_model(cfg, vocab_size, max_seq_len): + """Build a small GPTModel for testing.""" + spec = get_gpt_layer_with_transformer_engine_spec() + model = GPTModel( + config=cfg, + transformer_layer_spec=spec, + vocab_size=vocab_size, + max_sequence_length=max_seq_len, + pre_process=True, + post_process=True, + parallel_output=False, + position_embedding_type="rope", + ) + model.cuda() + return model + + +# =================================================================== +# Core test logic +# =================================================================== + + +def run_equal_length_test(model, batch_size, seq_len, vocab_size, seed, atol_loss, rtol_grad): + """Compare BSHD and THD with equal-length sequences (no padding). + + Returns a dict with test metrics for logging. + """ + B, S = batch_size, seq_len + + # Deterministic data generation. + torch.manual_seed(seed + 1) + input_ids = torch.randint(0, vocab_size, (B, S), device="cuda") + labels = torch.randint(0, vocab_size, (B, S), device="cuda") + loss_mask = torch.ones(B, S, device="cuda") + position_ids = torch.arange(S, device="cuda").unsqueeze(0).expand(B, -1).contiguous() + + # ---- BSHD forward / backward ---- + output_bshd = model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=None, + labels=labels, + loss_mask=loss_mask, + ) + bshd_loss = mean_loss(output_bshd, loss_mask) + bshd_loss.backward() + bshd_gn = grad_norm(model) + bshd_lv = bshd_loss.item() + bshd_per_token = output_bshd.detach().float().view(-1).clone() + + model.zero_grad() + + # ---- THD forward / backward ---- + samples = _samples_from_bshd(input_ids, labels, loss_mask) + packed = pack_or_pad_batch(samples, use_packed_sequence=True, device="cuda") + psp = packed.pop("packed_seq_params") + thd_position_ids = _thd_position_ids([S] * B, device="cuda") + + output_thd = model( + input_ids=packed["input_ids"], + position_ids=thd_position_ids, + attention_mask=None, + labels=packed["labels"], + loss_mask=packed["loss_mask"], + packed_seq_params=psp, + ) + thd_loss = mean_loss(output_thd, packed["loss_mask"]) + thd_loss.backward() + thd_gn = grad_norm(model) + thd_lv = thd_loss.item() + thd_per_token = output_thd.detach().float().view(-1).clone() + + model.zero_grad() + + # ---- Comparison ---- + loss_diff = abs(bshd_lv - thd_lv) + grad_diff = abs(bshd_gn - thd_gn) + grad_rel = grad_diff / max(bshd_gn, 1e-8) + token_max_diff = (bshd_per_token - thd_per_token).abs().max().item() + token_mean_diff = (bshd_per_token - thd_per_token).abs().mean().item() + + loss_ok = loss_diff < atol_loss + grad_ok = grad_rel < rtol_grad + + metrics = dict( + bshd_loss=bshd_lv, + thd_loss=thd_lv, + loss_diff=loss_diff, + bshd_grad_norm=bshd_gn, + thd_grad_norm=thd_gn, + grad_diff=grad_diff, + grad_rel=grad_rel, + token_max_diff=token_max_diff, + token_mean_diff=token_mean_diff, + loss_ok=loss_ok, + grad_ok=grad_ok, + ) + return metrics + + +def run_variable_length_smoke_test(model, vocab_size, seed): + """Smoke test: variable-length sequences packed to THD. + + Does NOT compare against BSHD (padding in BSHD changes attention + context). Validates that: + - Packing produces correct shapes + - Forward + backward complete without error + - Loss is finite + - Gradients are finite and non-zero + + Returns a dict with test metrics. + """ + seq_lengths = [128, 96, 112, 80] + S = max(seq_lengths) + B = len(seq_lengths) + + torch.manual_seed(seed + 2) + input_ids = torch.randint(0, vocab_size, (B, S), device="cuda") + labels = torch.randint(0, vocab_size, (B, S), device="cuda") + loss_mask = torch.ones(B, S, device="cuda") + + # Mask out padded positions in loss_mask for the input we hand to + # pack_or_pad_batch (variable-length samples carry only valid tokens). + for i, sl in enumerate(seq_lengths): + loss_mask[i, sl:] = 0.0 + + samples = _samples_from_bshd(input_ids, labels, loss_mask, seq_lengths=seq_lengths) + packed = pack_or_pad_batch(samples, use_packed_sequence=True, device="cuda") + psp = packed.pop("packed_seq_params") + thd_position_ids = _thd_position_ids(seq_lengths, device="cuda") + + T = sum(seq_lengths) + assert packed["input_ids"].shape == ( + 1, + T, + ), f"Expected [1, {T}], got {packed['input_ids'].shape}" + assert packed["labels"].shape == (1, T) + assert packed["loss_mask"].shape == (1, T) + assert psp.cu_seqlens_q.tolist() == [ + 0, + seq_lengths[0], + seq_lengths[0] + seq_lengths[1], + seq_lengths[0] + seq_lengths[1] + seq_lengths[2], + T, + ] + + output = model( + input_ids=packed["input_ids"], + position_ids=thd_position_ids, + attention_mask=None, + labels=packed["labels"], + loss_mask=packed["loss_mask"], + packed_seq_params=psp, + ) + loss = mean_loss(output, packed["loss_mask"]) + loss.backward() + gn = grad_norm(model) + loss_val = loss.item() + + model.zero_grad() + + loss_finite = torch.isfinite(torch.tensor(loss_val)).item() + grad_finite = torch.isfinite(torch.tensor(gn)).item() + grad_nonzero = gn > 0 + + return dict( + loss=loss_val, + grad_norm=gn, + total_tokens=T, + loss_finite=loss_finite, + grad_finite=grad_finite, + grad_nonzero=grad_nonzero, + passed=loss_finite and grad_finite and grad_nonzero, + ) + + +# =================================================================== +# Main +# =================================================================== + + +def _print_banner(title): + print(f"\n{'='*60}") + print(f" {title}") + print(f"{'='*60}") + + +def main(): + """CLI entry: run the equal-length parity test + variable-length smoke test.""" + parser = argparse.ArgumentParser(description="BSHD vs THD correctness test") + parser.add_argument("--batch-size", type=int, default=4) + parser.add_argument("--seq-len", type=int, default=128) + parser.add_argument("--vocab-size", type=int, default=1024) + parser.add_argument("--hidden-size", type=int, default=256) + parser.add_argument("--num-layers", type=int, default=2) + parser.add_argument("--num-heads", type=int, default=4) + parser.add_argument("--num-kv-heads", type=int, default=2) + parser.add_argument("--ffn-hidden-size", type=int, default=512) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--atol-loss", type=float, default=1e-5, help="Absolute tolerance for loss comparison" + ) + parser.add_argument( + "--rtol-grad", type=float, default=1e-3, help="Relative tolerance for grad norm comparison" + ) + args = parser.parse_args() + + Utils.initialize_model_parallel(tensor_model_parallel_size=1) + model_parallel_cuda_manual_seed(args.seed) + + config = TransformerConfig( + num_layers=args.num_layers, + hidden_size=args.hidden_size, + ffn_hidden_size=args.ffn_hidden_size, + num_attention_heads=args.num_heads, + num_query_groups=args.num_kv_heads, + bf16=True, + params_dtype=torch.bfloat16, + pipeline_dtype=torch.bfloat16, + hidden_dropout=0.0, + attention_dropout=0.0, + tensor_model_parallel_size=1, + sequence_parallel=False, + ) + + model = _build_model(config, args.vocab_size, args.seq_len) + + all_passed = True + + # ---------------------------------------------------------------- + # Test 1: equal-length correctness (BSHD vs THD) + # ---------------------------------------------------------------- + _print_banner("Test 1: Equal-length BSHD vs THD correctness") + m = run_equal_length_test( + model=model, + batch_size=args.batch_size, + seq_len=args.seq_len, + vocab_size=args.vocab_size, + seed=args.seed, + atol_loss=args.atol_loss, + rtol_grad=args.rtol_grad, + ) + print( + f" Config: B={args.batch_size}, S={args.seq_len}, " + f"H={args.hidden_size}, L={args.num_layers}, " + f"heads={args.num_heads}/{args.num_kv_heads}" + ) + print(f" BSHD loss: {m['bshd_loss']:.8f}") + print(f" THD loss: {m['thd_loss']:.8f}") + print(f" Loss abs diff: {m['loss_diff']:.2e}") + print(f" BSHD grad norm: {m['bshd_grad_norm']:.8f}") + print(f" THD grad norm: {m['thd_grad_norm']:.8f}") + print(f" Grad norm rel diff: {m['grad_rel']:.2e}") + print(f" Per-token max diff: {m['token_max_diff']:.2e}") + print(f" Per-token mean diff: {m['token_mean_diff']:.2e}") + print( + f" Loss match: {'PASS' if m['loss_ok'] else 'FAIL'} " f"(atol={args.atol_loss})" + ) + print( + f" Grad norm match: {'PASS' if m['grad_ok'] else 'FAIL'} " f"(rtol={args.rtol_grad})" + ) + if not (m["loss_ok"] and m["grad_ok"]): + all_passed = False + + # ---------------------------------------------------------------- + # Test 2: variable-length smoke test (THD only) + # ---------------------------------------------------------------- + _print_banner("Test 2: Variable-length THD smoke test") + v = run_variable_length_smoke_test(model, args.vocab_size, args.seed) + print(f" Seq lengths: [128, 96, 112, 80]") + print(f" Total packed tokens: {v['total_tokens']}") + print(f" Loss: {v['loss']:.8f}") + print(f" Grad norm: {v['grad_norm']:.8f}") + print(f" Loss finite: {'PASS' if v['loss_finite'] else 'FAIL'}") + print(f" Grad finite: {'PASS' if v['grad_finite'] else 'FAIL'}") + print(f" Grad nonzero: {'PASS' if v['grad_nonzero'] else 'FAIL'}") + if not v["passed"]: + all_passed = False + + # ---------------------------------------------------------------- + # Summary + # ---------------------------------------------------------------- + _print_banner("Summary") + if all_passed: + print(" ALL TESTS PASSED") + else: + print(" SOME TESTS FAILED") + print(f"{'='*60}\n") + + Utils.destroy_model_parallel() + + if not all_passed: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/multimodal_dev/tests/test_thd_e2e.py b/examples/multimodal_dev/tests/test_thd_e2e.py new file mode 100644 index 00000000000..c8d9f2e1622 --- /dev/null +++ b/examples/multimodal_dev/tests/test_thd_e2e.py @@ -0,0 +1,314 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Tests for THD / padded batch construction in multimodal_dev. + +Exercises the production data path :func:`pack_or_pad_batch`, which +consumes a list of per-sample dicts produced by the dataset's +``__getitem__`` and produces either a packed THD batch (``[1, T]``) or a +padded BSHD batch (``[B, S]``). + +``pack_or_pad_batch`` ends with a TP-group broadcast, so these tests +require ``torch.distributed`` to be initialised. Run via:: + + torchrun --nproc-per-node 1 -m pytest -q \\ + examples/multimodal_dev/tests/test_thd_e2e.py +""" + +import os +import sys + +import pytest +import torch + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +from examples.multimodal_dev.forward_step import _build_packed_seq_params, pack_or_pad_batch +from tests.unit_tests.test_utilities import Utils + + +@pytest.fixture(scope="module", autouse=True) +def _init_model_parallel(): + """Single-rank TP init so pack_or_pad_batch's TP broadcast is a no-op.""" + Utils.initialize_model_parallel(tensor_model_parallel_size=1) + yield + Utils.destroy_model_parallel() + + +def _make_sample( + seq_len: int, *, base: int = 0, num_patches: int = 4, pixel_dim: int = 8, device: str = "cuda" +): + """Per-sample dict in the shape produced by ``CordV2VLMDataset.__getitem__``. + + 1-D ``input_ids`` / ``labels`` / ``loss_mask`` over the sequence dim; + ``pixel_values`` is ``[num_patches, pixel_dim]``; ``image_grid_thw`` is + ``[1, 3]``. + """ + return { + "input_ids": torch.arange(seq_len, dtype=torch.long, device=device) + base, + "labels": (torch.arange(seq_len, dtype=torch.long, device=device) + base + 100), + "loss_mask": torch.ones(seq_len, dtype=torch.float, device=device), + "pixel_values": torch.full((num_patches, pixel_dim), float(base), device=device), + "image_grid_thw": torch.tensor([[2, 4, 4]], dtype=torch.long, device=device), + } + + +# =================================================================== +# _build_packed_seq_params — pure helper, exercised independently +# =================================================================== + + +class TestBuildPackedSeqParams: + """Tests for ``_build_packed_seq_params``.""" + + def test_basic(self): + """Mixed-length sample build sanity check.""" + params = _build_packed_seq_params(torch.tensor([5, 3, 7], dtype=torch.int32), device="cpu") + assert params.qkv_format == "thd" + assert params.cu_seqlens_q.tolist() == [0, 5, 8, 15] + assert params.cu_seqlens_kv.tolist() == [0, 5, 8, 15] + assert params.max_seqlen_q == 7 + assert params.max_seqlen_kv == 7 + assert params.total_tokens == 15 + assert params.cu_seqlens_q_padded.tolist() == [0, 5, 8, 15] + assert params.cu_seqlens_kv_padded.tolist() == [0, 5, 8, 15] + + def test_equal_lengths(self): + """Equal-length samples produce uniform cu_seqlens.""" + params = _build_packed_seq_params(torch.tensor([4, 4, 4], dtype=torch.int32), device="cpu") + assert params.cu_seqlens_q.tolist() == [0, 4, 8, 12] + assert params.max_seqlen_q == 4 + assert params.total_tokens == 12 + + def test_single_sample(self): + """Single-sample batch round-trips its own length.""" + params = _build_packed_seq_params(torch.tensor([10], dtype=torch.int32), device="cpu") + assert params.cu_seqlens_q.tolist() == [0, 10] + assert params.max_seqlen_q == 10 + assert params.total_tokens == 10 + + def test_dtype_is_int32(self): + """``cu_seqlens_q`` is cast to int32 regardless of input dtype.""" + params = _build_packed_seq_params(torch.tensor([3, 5], dtype=torch.int32), device="cpu") + assert params.cu_seqlens_q.dtype == torch.int32 + + def test_seq_idx_computed(self): + """``__post_init__`` computes ``seq_idx`` for Mamba compatibility.""" + params = _build_packed_seq_params(torch.tensor([3, 2], dtype=torch.int32), device="cpu") + assert params.seq_idx is not None + assert params.seq_idx.shape == (1, 5) + assert params.seq_idx[0].tolist() == [0, 0, 0, 1, 1] + + +# =================================================================== +# pack_or_pad_batch — packed (THD) mode +# =================================================================== + + +class TestPackOrPadBatchPacked: + """``pack_or_pad_batch(..., use_packed_sequence=True)`` produces ``[1, T]``.""" + + def test_equal_lengths(self): + """Two equal-length samples → packed ``[1, 2S]``.""" + S = 8 + batch = [_make_sample(S, base=0), _make_sample(S, base=1000)] + packed = pack_or_pad_batch(batch, use_packed_sequence=True, device="cuda") + + T = 2 * S + assert packed["input_ids"].shape == (1, T) + assert packed["labels"].shape == (1, T) + assert packed["loss_mask"].shape == (1, T) + psp = packed["packed_seq_params"] + assert psp.cu_seqlens_q.tolist() == [0, S, T] + assert psp.cu_seqlens_q_padded.tolist() == [0, S, T] + assert psp.max_seqlen_q == S + assert psp.total_tokens == T + + def test_variable_lengths(self): + """Variable-length samples concatenated end-to-end.""" + lens = [5, 8, 3] + batch = [_make_sample(L, base=i * 1000) for i, L in enumerate(lens)] + packed = pack_or_pad_batch(batch, use_packed_sequence=True, device="cuda") + + T = sum(lens) + assert packed["input_ids"].shape == (1, T) + psp = packed["packed_seq_params"] + assert psp.cu_seqlens_q.tolist() == [0, 5, 13, 16] + assert psp.cu_seqlens_q_padded.tolist() == [0, 5, 13, 16] + assert psp.max_seqlen_q == 8 + assert psp.total_tokens == T + + def test_token_order_preserved(self): + """Sample 0's tokens precede sample 1's tokens in the packed sequence.""" + s0 = _make_sample(3, base=10) + s1 = _make_sample(3, base=40) + packed = pack_or_pad_batch([s0, s1], use_packed_sequence=True, device="cuda") + assert packed["input_ids"][0].tolist() == [10, 11, 12, 40, 41, 42] + + def test_labels_loss_mask_content_preserved(self): + """labels and loss_mask carry through unchanged when divisible_by=1.""" + s = _make_sample(4, base=0) + packed = pack_or_pad_batch([s], use_packed_sequence=True, device="cuda") + assert packed["labels"][0].tolist() == [100, 101, 102, 103] + assert packed["loss_mask"][0].tolist() == [1.0, 1.0, 1.0, 1.0] + + def test_pixel_values_concatenated(self): + """``pixel_values`` are concatenated along the patch dim.""" + s0 = _make_sample(4, base=0, num_patches=4, pixel_dim=8) + s1 = _make_sample(4, base=10, num_patches=6, pixel_dim=8) + packed = pack_or_pad_batch([s0, s1], use_packed_sequence=True, device="cuda") + assert packed["pixel_values"].shape == (10, 8) + assert packed["pixel_values"][:4].eq(0.0).all().item() + assert packed["pixel_values"][4:].eq(10.0).all().item() + + def test_image_grid_thw_concatenated(self): + """``image_grid_thw`` rows are concatenated along the first dim.""" + s0 = _make_sample(4, base=0) + s1 = _make_sample(4, base=10) + packed = pack_or_pad_batch([s0, s1], use_packed_sequence=True, device="cuda") + assert packed["image_grid_thw"].shape == (2, 3) + + def test_single_sample_round_trip(self): + """A single sample packs to its own length.""" + s = _make_sample(7, base=0) + packed = pack_or_pad_batch([s], use_packed_sequence=True, device="cuda") + assert packed["input_ids"].shape == (1, 7) + psp = packed["packed_seq_params"] + assert psp.cu_seqlens_q.tolist() == [0, 7] + assert psp.max_seqlen_q == 7 + assert psp.total_tokens == 7 + + +# =================================================================== +# pack_or_pad_batch — padded (BSHD) mode +# =================================================================== + + +class TestPackOrPadBatchPadded: + """``pack_or_pad_batch(..., use_packed_sequence=False)`` produces ``[B, S]``.""" + + def test_equal_lengths(self): + """Equal-length samples → ``[B, S]`` without further padding.""" + S = 6 + batch = [_make_sample(S, base=0), _make_sample(S, base=10)] + padded = pack_or_pad_batch(batch, use_packed_sequence=False, seq_length=S, device="cuda") + assert padded["input_ids"].shape == (2, S) + assert padded["labels"].shape == (2, S) + assert padded["loss_mask"].shape == (2, S) + # Sample-0 content is preserved verbatim. + assert padded["input_ids"][0].tolist() == list(range(S)) + + def test_pads_short_sample_to_batch_max(self): + """Shorter sample is right-padded to match the batch max length.""" + long_sample = _make_sample(7, base=0) + short_sample = _make_sample(3, base=10) + padded = pack_or_pad_batch( + [long_sample, short_sample], use_packed_sequence=False, seq_length=7, device="cuda" + ) + assert padded["input_ids"].shape == (2, 7) + # Short sample: original [10, 11, 12] then pad zeros. + assert padded["input_ids"][1].tolist() == [10, 11, 12, 0, 0, 0, 0] + # labels pad with -100 (ignore index). + assert padded["labels"][1].tolist() == [110, 111, 112, -100, -100, -100, -100] + # loss_mask pads with 0. + assert padded["loss_mask"][1].tolist() == [1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0] + + def test_seq_length_required(self): + """``seq_length`` must be provided in padded mode.""" + s = _make_sample(4, base=0) + with pytest.raises(AssertionError, match="seq_length"): + pack_or_pad_batch([s], use_packed_sequence=False, seq_length=None, device="cuda") + + def test_pixel_values_concatenated(self): + """``pixel_values`` concat preserves both samples' patches.""" + s0 = _make_sample(4, base=0, num_patches=4, pixel_dim=8) + s1 = _make_sample(4, base=10, num_patches=6, pixel_dim=8) + padded = pack_or_pad_batch([s0, s1], use_packed_sequence=False, seq_length=4, device="cuda") + assert padded["pixel_values"].shape == (10, 8) + assert padded["pixel_values"][:4].eq(0.0).all().item() + assert padded["pixel_values"][4:].eq(10.0).all().item() + + +# =================================================================== +# pack_or_pad_batch — divisible_by = 4 alignment +# =================================================================== + + +class TestPackOrPadBatchDivisibleBy4: + """Per-sample sequence alignment when ``divisible_by = 4``. + + The function computes ``divisible_by`` from the parallel state. With + ``world_size=1`` we cannot stand up a real CP=2 group, so we patch + :func:`mpu.get_context_parallel_world_size` to return 2; the function + then takes the ``cp_size > 1`` branch and yields + ``divisible_by = cp_size * 2 = 4`` (no SP). + """ + + @pytest.fixture + def cp2(self, monkeypatch): + """Patch CP world size to 2 so ``divisible_by = 4``.""" + from examples.multimodal_dev import forward_step + + monkeypatch.setattr(forward_step.mpu, "get_context_parallel_world_size", lambda: 2) + + def test_packed_aligned_samples_no_padding(self, cp2): + """Samples already multiples of 4 → cu_seqlens == cu_seqlens_padded.""" + batch = [_make_sample(8, base=0), _make_sample(4, base=100)] + packed = pack_or_pad_batch(batch, use_packed_sequence=True, device="cuda") + + T = 12 + assert packed["input_ids"].shape == (1, T) + psp = packed["packed_seq_params"] + assert psp.cu_seqlens_q.tolist() == [0, 8, 12] + assert psp.cu_seqlens_q_padded.tolist() == [0, 8, 12] + assert psp.max_seqlen_q == 8 + assert psp.total_tokens == 12 + + def test_packed_misaligned_samples_padded_per_sample(self, cp2): + """Each sample padded up to the nearest multiple of 4.""" + # lens=[5, 8, 3] → padded=[8, 8, 4] → T_padded = 20. + batch = [_make_sample(5, base=0), _make_sample(8, base=100), _make_sample(3, base=200)] + packed = pack_or_pad_batch(batch, use_packed_sequence=True, device="cuda") + + T_padded = 20 + assert packed["input_ids"].shape == (1, T_padded) + psp = packed["packed_seq_params"] + # cu_seqlens reflects real per-sample lengths. + assert psp.cu_seqlens_q.tolist() == [0, 5, 13, 16] + # cu_seqlens_padded reflects per-sample alignment to 4. + assert psp.cu_seqlens_q_padded.tolist() == [0, 8, 16, 20] + # max_seqlen comes from the padded lengths. + assert psp.max_seqlen_q == 8 + assert psp.total_tokens == T_padded + + def test_packed_pad_values(self, cp2): + """Pad slots filled with input_ids=0, labels=-100, loss_mask=0.""" + # Single sample len=3 → target_len=4 → 1 pad slot at position 3. + batch = [_make_sample(3, base=10)] + packed = pack_or_pad_batch(batch, use_packed_sequence=True, device="cuda") + + assert packed["input_ids"].shape == (1, 4) + assert packed["input_ids"][0].tolist() == [10, 11, 12, 0] + assert packed["labels"][0].tolist() == [110, 111, 112, -100] + assert packed["loss_mask"][0].tolist() == [1.0, 1.0, 1.0, 0.0] + psp = packed["packed_seq_params"] + assert psp.cu_seqlens_q.tolist() == [0, 3] + assert psp.cu_seqlens_q_padded.tolist() == [0, 4] + assert psp.max_seqlen_q == 4 + assert psp.total_tokens == 4 + + def test_padded_target_rounded_up_to_multiple_of_4(self, cp2): + """Padded (BSHD) mode: ``target = ceil(min(max, seq_length) / 4) * 4``.""" + # lens=[5, 3], seq_length=10 → min(5, 10) = 5 → ceil(5/4)*4 = 8. + long_sample = _make_sample(5, base=0) + short_sample = _make_sample(3, base=10) + padded = pack_or_pad_batch( + [long_sample, short_sample], use_packed_sequence=False, seq_length=10, device="cuda" + ) + + assert padded["input_ids"].shape == (2, 8) + assert padded["input_ids"][0].tolist() == [0, 1, 2, 3, 4, 0, 0, 0] + assert padded["input_ids"][1].tolist() == [10, 11, 12, 0, 0, 0, 0, 0] + assert padded["labels"][1].tolist() == [110, 111, 112, -100, -100, -100, -100, -100] + assert padded["loss_mask"][1].tolist() == [1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0] diff --git a/examples/multimodal_dev/tests/test_vision_rope_fusion.py b/examples/multimodal_dev/tests/test_vision_rope_fusion.py new file mode 100644 index 00000000000..602115af1fb --- /dev/null +++ b/examples/multimodal_dev/tests/test_vision_rope_fusion.py @@ -0,0 +1,182 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Tests for Qwen3.5-VL vision RoPE fusion dispatch.""" + +import os +import sys +from types import SimpleNamespace + +import pytest +import torch + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +import megatron.core.models.common.embeddings.rope_utils as rope_utils +from examples.multimodal_dev.models.qwen35_vl.configuration import get_qwen35_vl_vision_config +from examples.multimodal_dev.models.qwen35_vl.specs import _apply_rope_fp32_no_cp +from examples.multimodal_dev.models.qwen35_vl.vision_encoder import Qwen35VLVisionEncoder +from megatron.core.fusions.fused_mrope import is_fused_mrope_available, mrope_freqs_to_rotary_emb + + +class _FakeVisionRotaryEmbedding: + def __init__(self, axis_dim): + self.axis_dim = axis_dim + + def __call__(self, seqlen, device=None): + device = torch.device("cpu") if device is None else device + positions = torch.arange(seqlen, device=device, dtype=torch.float32)[:, None] + dims = torch.arange(self.axis_dim, device=device, dtype=torch.float32)[None, :] + return positions * 0.125 + dims * 0.01 + + +def test_vision_config_sets_2d_rope_as_sectioned_raw_mrope(): + config = get_qwen35_vl_vision_config(variant="0.8b") + + assert config.kv_channels == 64 + assert config.mrope_section == [0, 16, 16] + assert config.mrope_interleaved is False + assert config.rotary_interleaved is False + + +def test_vision_raw_mrope_freqs_match_legacy_materialized_rope(): + grid_thw = torch.tensor([[1, 4, 4], [2, 2, 2]], dtype=torch.long) + legacy_encoder = SimpleNamespace( + spatial_merge_size=2, + rot_pos_emb=_FakeVisionRotaryEmbedding(axis_dim=16), + config=SimpleNamespace(mrope_section=None), + ) + fused_encoder = SimpleNamespace( + spatial_merge_size=2, + rot_pos_emb=_FakeVisionRotaryEmbedding(axis_dim=16), + config=SimpleNamespace(mrope_section=[0, 16, 16]), + ) + + legacy_freqs = Qwen35VLVisionEncoder._compute_rotary_pos_emb(legacy_encoder, grid_thw) + raw_freqs = Qwen35VLVisionEncoder._compute_rotary_pos_emb(fused_encoder, grid_thw) + + expected = torch.cat((legacy_freqs, legacy_freqs), dim=-1).unsqueeze(1).unsqueeze(1) + converted = mrope_freqs_to_rotary_emb( + raw_freqs, + [0, 16, 16], + interleaved_mrope=False, + rotary_interleaved=False, + ) + + assert raw_freqs.shape == (3, 1, legacy_freqs.shape[0], legacy_freqs.shape[1]) + torch.testing.assert_close(converted, expected) + + +def test_vision_fp32_wrapper_dispatches_raw_freqs_to_fused_mrope_thd(monkeypatch): + calls = {} + + def fake_fused_apply_mrope_thd( + t, + cu_seqlens, + freqs, + mrope_section, + interleaved_mrope=False, + rotary_interleaved=False, + cp_size=1, + cp_rank=0, + fp32_compute=False, + ): + calls["t_shape"] = tuple(t.shape) + calls["t_dtype"] = t.dtype + calls["cu_seqlens"] = cu_seqlens.tolist() + calls["freqs_shape"] = tuple(freqs.shape) + calls["mrope_section"] = list(mrope_section) + calls["interleaved_mrope"] = interleaved_mrope + calls["rotary_interleaved"] = rotary_interleaved + calls["cp_size"] = cp_size + calls["cp_rank"] = cp_rank + calls["fp32_compute"] = fp32_compute + return t + 1.0 + + monkeypatch.setattr(rope_utils, "fused_apply_mrope_thd", fake_fused_apply_mrope_thd) + monkeypatch.setattr(rope_utils, "get_fused_mrope_thd_unavailable_reason", lambda *args, **kwargs: None) + + config = SimpleNamespace( + apply_rope_fusion=True, + mrope_section=[0, 2, 2], + mrope_interleaved=False, + rotary_interleaved=False, + multi_latent_attention=False, + ) + t = torch.zeros(6, 2, 8, dtype=torch.bfloat16) + freqs = torch.zeros(3, 1, 6, 4, dtype=torch.float32) + cu_seqlens = torch.tensor([0, 3, 6], dtype=torch.int32) + + out = _apply_rope_fp32_no_cp(t, freqs, config, cu_seqlens=cu_seqlens) + + assert out.dtype == torch.bfloat16 + torch.testing.assert_close(out, torch.ones_like(out)) + assert calls == { + "t_shape": (6, 2, 8), + "t_dtype": torch.bfloat16, + "cu_seqlens": [0, 3, 6], + "freqs_shape": (3, 1, 6, 4), + "mrope_section": [0, 2, 2], + "interleaved_mrope": False, + "rotary_interleaved": False, + "cp_size": 1, + "cp_rank": 0, + "fp32_compute": True, + } + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_vision_fused_rope_matches_unfused_forward_backward_cuda(): + generator = torch.Generator(device="cuda").manual_seed(1234) + total_tokens = 64 + num_heads = 2 + head_dim = 72 + half_rotary_dim = head_dim // 2 + section = [0, half_rotary_dim // 2, half_rotary_dim // 2] + cu_seqlens = torch.tensor([0, total_tokens], dtype=torch.int32, device="cuda") + + t_ref = torch.randn( + total_tokens, + num_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + generator=generator, + requires_grad=True, + ) + t_fused = t_ref.detach().clone().requires_grad_(True) + freqs = torch.randn( + 3, + 1, + total_tokens, + half_rotary_dim, + device="cuda", + dtype=torch.float32, + generator=generator, + ) + + ref_config = SimpleNamespace( + apply_rope_fusion=False, + mrope_section=section, + mrope_interleaved=False, + rotary_interleaved=False, + multi_latent_attention=False, + ) + fused_config = SimpleNamespace( + apply_rope_fusion=True, + mrope_section=section, + mrope_interleaved=False, + rotary_interleaved=False, + multi_latent_attention=False, + ) + + ref = _apply_rope_fp32_no_cp(t_ref, freqs, ref_config, cu_seqlens=cu_seqlens) + out = _apply_rope_fp32_no_cp(t_fused, freqs, fused_config, cu_seqlens=cu_seqlens) + torch.testing.assert_close(ref.float(), out.float(), rtol=2.0e-2, atol=5.0e-2) + + grad = torch.randn_like(ref) + ref.backward(grad) + out.backward(grad) + torch.testing.assert_close(t_ref.grad.float(), t_fused.grad.float(), rtol=2.0e-2, atol=5.0e-2) diff --git a/megatron/core/fusions/fused_mega_pre_gated_delta_rule.py b/megatron/core/fusions/fused_mega_pre_gated_delta_rule.py new file mode 100644 index 00000000000..66ce10a99ea --- /dev/null +++ b/megatron/core/fusions/fused_mega_pre_gated_delta_rule.py @@ -0,0 +1,1109 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Mega fused pre-gated-delta-rule kernels. + +This is the "mega" sibling of :mod:`fused_pre_gated_delta_rule`. The streamed +path splits the pre-gated-delta-rule front-end into four separate Triton launch +scopes (QK / V / Z / g-beta) plus an external conv backward, optimized for +per-kernel quality and overlap under CUDA-graph capture. The mega path instead +folds **all forward tasks into a single Triton launch** using a flat logical +task space, trading a little per-kernel efficiency for far fewer host-side +launches. It is the right choice for non-CUDA-graph recipes where launch +overhead is visible in the trace. + +Public contract is identical across unfused / streamed / mega: +``(query, key, value, gate, beta, g)``. + +Design notes: + +* The forward kernel maps ``program_id(0)`` onto a flat row space partitioned + into QK, V, Z, and g/beta ranges; ``program_id(1)`` tiles the sequence axis. + Each program inspects its row id and runs exactly one task body. This keeps + every sub-computation in one launch while still letting Triton schedule + memory-bound (Z copy) and compute-bound (QK/V conv) tiles concurrently on the + SMs. +* Numerics mirror the streamed/unfused reference **bit-for-bit within the unit + test tolerance**: the conv accumulator is rounded through the activation + dtype before SiLU; the SiLU output is rounded again before the L2-norm + reduction; ``g`` uses an fp32 ``log(1+exp(...))`` softplus; ``beta`` uses an + fp32 sigmoid. The QK ``silu(conv(x))`` intermediate is persisted channel-last + exactly as the streamed path saves it, so the backward can be shared. +* The kernel assumes ``key_head_dim == value_head_dim`` so a single + ``HEAD_DIM`` constexpr drives the QK/V/Z channel tiles (true for the GDN + production shapes and the unit tests). This is asserted at the Python entry. + +The backward mirrors the forward: a single fused Triton kernel folds the four +streamed branch backward scopes (QK l2norm/repeat, V layout, Z layout, g/beta +chain rule) into one flat-task launch, then the depthwise conv input/weight +gradients are delegated to the same external ``causal_conv1d_bwd_function`` the +streamed path uses (its hand-tuned C++ remains the conv-backward anchor). That +is two launches total (one fused branch kernel + one external conv backward), +down from the streamed path's five, while staying numerically bit-identical to +the streamed branch kernels. +""" + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl +from torch import Tensor + +# Reuse the streamed module's validated constants and the external conv backward +# binding. Importing is not a modification of that module. +from megatron.core.fusions.fused_pre_gated_delta_rule import ( + _L2NORM_EPS, + _causal_conv1d_bwd_function, + _is_power_of_two, + _resolve_packed_seq_idx, +) + + +# --------------------------------------------------------------------------- +# Forward kernel +# --------------------------------------------------------------------------- + + +def _mega_autotune_configs(): + return [ + triton.Config({"BLOCK_S": 32}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_S": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=3), + triton.Config({"BLOCK_S": 256}, num_warps=8, num_stages=2), + ] + + +@triton.jit +def _mega_seq_bounds(cu_seqlens_ptr, token_offsets, total_tokens, num_packed_seqs): + """Lane-wise packed-sequence [start, end) bounds for flattened THD tokens. + + Local copy of the streamed helper so this kernel never depends on + cross-module ``@triton.jit`` symbol resolution. + """ + + safe_tokens = tl.minimum(token_offsets, total_tokens - 1) + seq_start = token_offsets * 0 + seq_end = token_offsets * 0 + total_tokens + + seq_id = 0 + while seq_id < num_packed_seqs: + start = tl.load(cu_seqlens_ptr + seq_id) + end = tl.load(cu_seqlens_ptr + seq_id + 1) + in_seq = (safe_tokens >= start) & (safe_tokens < end) + seq_start = tl.where(in_seq, start, seq_start) + seq_end = tl.where(in_seq, end, seq_end) + seq_id += 1 + + return seq_start, seq_end + + +@triton.autotune( + configs=_mega_autotune_configs(), + key=["seq_len", "HEAD_DIM", "K_W", "num_key_heads", "num_value_heads", "REPEAT", "HAS_THD"], +) +@triton.jit +def _mega_forward_kernel( + qkvzba_ptr, + weight_ptr, + A_log_ptr, + dt_bias_ptr, + qk_out_ptr, + value_ptr, + gate_ptr, + g_ptr, + beta_ptr, + silu_save_ptr, + cu_seqlens_ptr, + seq_len, + num_packed_seqs, + num_key_heads, + num_value_heads, + qk_channels, + v_channels, + R_qk, + R_v, + R_z, + qkvzba_s_stride, + qkvzba_b_stride, + qkvzba_c_stride, + weight_c_stride, + weight_w_stride, + qk_g_stride, + qk_b_stride, + qk_s_stride, + qk_h_stride, + v_b_stride, + v_s_stride, + v_h_stride, + z_b_stride, + z_s_stride, + z_h_stride, + g_b_stride, + g_s_stride, + g_h_stride, + beta_b_stride, + beta_s_stride, + beta_h_stride, + silu_b_stride, + silu_c_stride, + silu_s_stride, + eps, + HEAD_DIM: tl.constexpr, + K_W: tl.constexpr, + REPEAT: tl.constexpr, + HAS_THD: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """All-in-one forward for the pre-gated-delta-rule front-end. + + Flat task space on ``program_id(0)``: + rows [0, R_qk) -> QK conv+silu+l2norm+repeat + rows [R_qk, R_qk+R_v) -> V conv+silu + rows [R_qk+R_v, +R_z) -> Z copy + rows [.., end) -> g/beta + ``program_id(1)`` tiles the (flattened, for THD) sequence axis. + """ + + pid_row = tl.program_id(0) + pid_s = tl.program_id(1) + + # Common rounding dtype (activation dtype, e.g. bf16). All q/k/v/gate/beta + # outputs share this; g is fp32. + out_ty = qk_out_ptr.dtype.element_ty + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + chan_off = tl.arange(0, HEAD_DIM) + + R_qkv = R_qk + R_v + R_qkvz = R_qkv + R_z + + if pid_row < R_qk: + # ---- QK: depthwise causal conv + silu + l2norm + head repeat ---- + local = pid_row + heads_per_batch = 2 * num_key_heads + batch_id = local // heads_per_batch + lb = local - batch_id * heads_per_batch + group_id = lb // num_key_heads # 0 -> Q, 1 -> K + head_id = lb - group_id * num_key_heads + chan = group_id * qk_channels + head_id * HEAD_DIM + chan_off + + if HAS_THD: + seq_start, seq_end = _mega_seq_bounds( + cu_seqlens_ptr, s_offs, seq_len, num_packed_seqs + ) + + acc = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for i in tl.static_range(K_W): + x_s = s_offs - (K_W - 1) + i + if HAS_THD: + x_mask = s_mask & (x_s >= seq_start) & (x_s < seq_end) + safe_x_s = tl.minimum(tl.maximum(x_s, 0), seq_len - 1) + else: + x_mask = (x_s >= 0) & (x_s < seq_len) + safe_x_s = x_s + x_ptr = ( + qkvzba_ptr + + safe_x_s[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + chan[None, :] * qkvzba_c_stride + ) + x_val = tl.load(x_ptr, mask=x_mask[:, None], other=0.0).to(tl.float32) + w_tap = tl.load( + weight_ptr + chan * weight_c_stride + i * weight_w_stride + ).to(tl.float32) + acc += w_tap[None, :] * x_val + + acc = acc.to(out_ty).to(tl.float32) # F.conv1d rounding + silu_out = acc * tl.sigmoid(acc) + silu_out = silu_out.to(out_ty).to(tl.float32) # round before l2norm + + # Persist silu(conv(x)) for the QK channels, channel-last (b, 2*qk, s). + silu_chan = group_id * qk_channels + head_id * HEAD_DIM + chan_off + silu_ptrs = ( + silu_save_ptr + + batch_id * silu_b_stride + + silu_chan[None, :] * silu_c_stride + + s_offs[:, None] * silu_s_stride + ) + tl.store( + silu_ptrs, + silu_out.to(silu_save_ptr.dtype.element_ty), + mask=s_mask[:, None], + ) + + norm_sq = tl.sum(silu_out * silu_out, axis=1) + rstd = 1.0 / tl.sqrt(norm_sq + eps) + out_typed = (silu_out * rstd[:, None]).to(out_ty) + + for r in tl.static_range(REPEAT): + v_head = head_id * REPEAT + r + write_ptr = ( + qk_out_ptr + + group_id * qk_g_stride + + batch_id * qk_b_stride + + s_offs[:, None] * qk_s_stride + + v_head * qk_h_stride + + chan_off[None, :] + ) + tl.store(write_ptr, out_typed, mask=s_mask[:, None]) + + elif pid_row < R_qkv: + # ---- V: depthwise causal conv + silu (no l2norm, no repeat) ---- + local = pid_row - R_qk + batch_id = local // num_value_heads + head_id = local - batch_id * num_value_heads + chan = 2 * qk_channels + head_id * HEAD_DIM + chan_off + + if HAS_THD: + seq_start, seq_end = _mega_seq_bounds( + cu_seqlens_ptr, s_offs, seq_len, num_packed_seqs + ) + + acc = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for i in tl.static_range(K_W): + x_s = s_offs - (K_W - 1) + i + if HAS_THD: + x_mask = s_mask & (x_s >= seq_start) & (x_s < seq_end) + safe_x_s = tl.minimum(tl.maximum(x_s, 0), seq_len - 1) + else: + x_mask = (x_s >= 0) & (x_s < seq_len) + safe_x_s = x_s + x_ptr = ( + qkvzba_ptr + + safe_x_s[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + chan[None, :] * qkvzba_c_stride + ) + x_val = tl.load(x_ptr, mask=x_mask[:, None], other=0.0).to(tl.float32) + w_tap = tl.load( + weight_ptr + chan * weight_c_stride + i * weight_w_stride + ).to(tl.float32) + acc += w_tap[None, :] * x_val + + acc = acc.to(out_ty).to(tl.float32) + silu_out = acc * tl.sigmoid(acc) + out_typed = silu_out.to(out_ty) + write_ptr = ( + value_ptr + + batch_id * v_b_stride + + s_offs[:, None] * v_s_stride + + head_id * v_h_stride + + chan_off[None, :] + ) + tl.store(write_ptr, out_typed, mask=s_mask[:, None]) + + elif pid_row < R_qkvz: + # ---- Z: copy qkvzba z slice into the final gate layout ---- + local = pid_row - R_qkv + batch_id = local // num_value_heads + head_id = local - batch_id * num_value_heads + z_chan = 2 * qk_channels + v_channels + head_id * HEAD_DIM + chan_off + src_ptr = ( + qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + z_chan[None, :] * qkvzba_c_stride + ) + z_val = tl.load(src_ptr, mask=s_mask[:, None]) + write_ptr = ( + gate_ptr + + batch_id * z_b_stride + + s_offs[:, None] * z_s_stride + + head_id * z_h_stride + + chan_off[None, :] + ) + tl.store(write_ptr, z_val, mask=s_mask[:, None]) + + else: + # ---- g/beta: -exp(A_log)*softplus(alpha+dt_bias) and sigmoid(beta) ---- + local = pid_row - R_qkvz + batch_id = local // num_value_heads + head_id = local - batch_id * num_value_heads + beta_chan = 2 * qk_channels + 2 * v_channels + head_id + alpha_chan = beta_chan + num_value_heads + + alpha_ptr = ( + qkvzba_ptr + + s_offs * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + alpha_chan * qkvzba_c_stride + ) + beta_raw_ptr = ( + qkvzba_ptr + + s_offs * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + beta_chan * qkvzba_c_stride + ) + alpha = tl.load(alpha_ptr, mask=s_mask, other=0.0).to(tl.float32) + beta_raw = tl.load(beta_raw_ptr, mask=s_mask, other=0.0).to(tl.float32) + A_log = tl.load(A_log_ptr + head_id).to(tl.float32) + dt_bias = tl.load(dt_bias_ptr + head_id).to(tl.float32) + + pre = alpha + dt_bias + softplus_val = tl.log(1.0 + tl.exp(pre)) + g = -tl.exp(A_log) * softplus_val + beta_sig = tl.sigmoid(beta_raw) + + g_store_ptr = ( + g_ptr + batch_id * g_b_stride + s_offs * g_s_stride + head_id * g_h_stride + ) + beta_store_ptr = ( + beta_ptr + + batch_id * beta_b_stride + + s_offs * beta_s_stride + + head_id * beta_h_stride + ) + tl.store(g_store_ptr, g.to(g_ptr.dtype.element_ty), mask=s_mask) + tl.store(beta_store_ptr, beta_sig.to(beta_ptr.dtype.element_ty), mask=s_mask) + + +# --------------------------------------------------------------------------- +# Forward orchestration +# --------------------------------------------------------------------------- + + +def _mega_pre_gated_delta_rule_forward( + qkvzba: Tensor, + conv1d_weight: Tensor, + A_log: Tensor, + dt_bias: Tensor, + *, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, + cu_seqlens: Optional[Tensor] = None, +) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + """Single-launch mega forward. + + Returns ``(query, key, value, gate, beta, g, silu_qk_save)``; the last + element is the bf16-rounded QK ``silu(conv(x))`` laid out channel-last, + matching the streamed forward so the shared backward can consume it. + """ + + seq_len, batch, total_channels = qkvzba.shape + is_packed_thd = cu_seqlens is not None + num_packed_seqs = (cu_seqlens.shape[0] - 1) if is_packed_thd else 0 + + assert key_head_dim == value_head_dim, ( + "fused_mega_pre_gated_delta_rule currently requires " + f"key_head_dim == value_head_dim; got {key_head_dim=} {value_head_dim=}." + ) + assert _is_power_of_two(key_head_dim), ( + f"Mega kernel expects key_head_dim to be a power of two; got {key_head_dim=}." + ) + head_dim = key_head_dim + + qk_channels = num_key_heads * key_head_dim + v_channels = num_value_heads * value_head_dim + repeat_factor = num_value_heads // num_key_heads + k_w = conv1d_weight.shape[-1] + + expected_channels = 2 * qk_channels + 2 * v_channels + 2 * num_value_heads + assert total_channels == expected_channels, ( + f"qkvzba last-dim mismatch: got {total_channels}, expected {expected_channels}." + ) + + out_dtype = qkvzba.dtype + device = qkvzba.device + + # Output buffers (identical layouts to the streamed path). + qk_out = torch.empty( + 2, batch, seq_len, num_value_heads, key_head_dim, dtype=out_dtype, device=device + ) + query = qk_out[0] + key = qk_out[1] + value = torch.empty( + batch, seq_len, num_value_heads, value_head_dim, dtype=out_dtype, device=device + ) + gate = torch.empty( + batch, seq_len, num_value_heads, value_head_dim, dtype=out_dtype, device=device + ) + g = torch.empty(batch, seq_len, num_value_heads, dtype=torch.float32, device=device) + beta = torch.empty(batch, seq_len, num_value_heads, dtype=out_dtype, device=device) + + # QK silu(conv(x)) persisted channel-last: (b, 2*qk_channels, s), stride(1)==1. + silu_qk_save = torch.empty( + (batch, seq_len, 2 * qk_channels), dtype=out_dtype, device=device + ).permute(0, 2, 1) + + weight_2d = conv1d_weight.view(conv1d_weight.shape[0], k_w) + + # Flat task-space row partition. + R_qk = batch * 2 * num_key_heads + R_v = batch * num_value_heads + R_z = batch * num_value_heads + R_gb = batch * num_value_heads + num_rows = R_qk + R_v + R_z + R_gb + + cu_seqlens_arg = cu_seqlens if is_packed_thd else qkvzba # dummy when dense + + grid = lambda meta: (num_rows, triton.cdiv(seq_len, meta["BLOCK_S"])) + _mega_forward_kernel[grid]( + qkvzba, + weight_2d, + A_log, + dt_bias, + qk_out, + value, + gate, + g, + beta, + silu_qk_save, + cu_seqlens_arg, + seq_len, + num_packed_seqs, + num_key_heads, + num_value_heads, + qk_channels, + v_channels, + R_qk, + R_v, + R_z, + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + weight_2d.stride(0), + weight_2d.stride(1), + qk_out.stride(0), + qk_out.stride(1), + qk_out.stride(2), + qk_out.stride(3), + value.stride(0), + value.stride(1), + value.stride(2), + gate.stride(0), + gate.stride(1), + gate.stride(2), + g.stride(0), + g.stride(1), + g.stride(2), + beta.stride(0), + beta.stride(1), + beta.stride(2), + silu_qk_save.stride(0), + silu_qk_save.stride(1), + silu_qk_save.stride(2), + _L2NORM_EPS, + HEAD_DIM=head_dim, + K_W=k_w, + REPEAT=repeat_factor, + HAS_THD=is_packed_thd, + ) + + return query, key, value, gate, beta, g, silu_qk_save + + +# --------------------------------------------------------------------------- +# Backward kernel +# --------------------------------------------------------------------------- + + +def _mega_backward_autotune_configs(): + return [ + triton.Config({"BLOCK_S": 32}, num_warps=2, num_stages=2), + triton.Config({"BLOCK_S": 32}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_S": 64}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 256}, num_warps=8, num_stages=2), + ] + + +@triton.autotune( + configs=_mega_backward_autotune_configs(), + key=["seq_len", "HEAD_DIM", "REPEAT", "num_key_heads", "num_value_heads"], + # The g/beta task atomic-adds per-head partials into these accumulators. + # reset_to_zero clears them before each autotune trial so trials don't stack. + reset_to_zero=["d_A_log_ptr", "d_dt_bias_ptr"], +) +@triton.jit +def _mega_backward_kernel( + # inputs + dq_ptr, + dk_ptr, + dv_ptr, + dgate_ptr, + dg_ptr, + dbeta_ptr, + silu_save_ptr, + qkvzba_ptr, + A_log_ptr, + dt_bias_ptr, + # outputs + d_silu_conv_ptr, + d_qkvzba_ptr, + d_A_log_ptr, + d_dt_bias_ptr, + # sizes / layout + seq_len, + num_key_heads, + num_value_heads, + qk_channels, + v_channels, + R_qk, + R_v, + R_z, + eps, + # dq / dk strides (b, s, h, d) + dq_b_stride, + dq_s_stride, + dq_h_stride, + dk_b_stride, + dk_s_stride, + dk_h_stride, + # dv strides (b, s, h, d) + dv_b_stride, + dv_s_stride, + dv_h_stride, + # dgate strides (b, s, h, d) + dgate_b_stride, + dgate_s_stride, + dgate_h_stride, + # dg / dbeta strides (b, s, h) + dg_b_stride, + dg_s_stride, + dg_h_stride, + dbeta_b_stride, + dbeta_s_stride, + dbeta_h_stride, + # silu_save strides (b, 2*qk, s) + silu_b_stride, + silu_c_stride, + silu_s_stride, + # qkvzba / d_qkvzba strides (s, b, C) + qkvzba_s_stride, + qkvzba_b_stride, + qkvzba_c_stride, + # d_silu_conv strides (b, conv_dim, s) + dsc_b_stride, + dsc_c_stride, + dsc_s_stride, + HEAD_DIM: tl.constexpr, + REPEAT: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """All-in-one backward for the QK / V / Z / g-beta branches. + + Mirrors the four streamed branch kernels in one flat task space. Conv input + and weight gradients are NOT produced here; the caller feeds ``d_silu_conv`` + into the external ``causal_conv1d_bwd_function`` exactly as the streamed + path does. Flat task space on ``program_id(0)``: + rows [0, R_qk) -> QK l2norm + repeat backward -> d_silu_conv[Q/K] + rows [R_qk, +R_v) -> V layout copy -> d_silu_conv[V] + rows [.., +R_z) -> Z layout copy -> d_qkvzba[z] + rows [.., end) -> g/beta chain rule -> d_qkvzba[alpha,beta], + atomic d_A_log/d_dt_bias + """ + + pid_row = tl.program_id(0) + pid_s = tl.program_id(1) + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + chan_off = tl.arange(0, HEAD_DIM) + + R_qkv = R_qk + R_v + R_qkvz = R_qkv + R_z + + if pid_row < R_qk: + # ---- QK: repeat-reduce + l2norm backward -> d_silu_conv[Q/K] ---- + local = pid_row + heads_per_batch = 2 * num_key_heads + batch_id = local // heads_per_batch + lb = local - batch_id * heads_per_batch + group_id = lb // num_key_heads + head_id = lb - group_id * num_key_heads + is_query = group_id == 0 + is_key = group_id == 1 + chan = group_id * qk_channels + head_id * HEAD_DIM + chan_off + + d_normed = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for r in tl.static_range(REPEAT): + v_head = head_id * REPEAT + r + dq_ptrs = ( + dq_ptr + + batch_id * dq_b_stride + + s_offs[:, None] * dq_s_stride + + v_head * dq_h_stride + + chan_off[None, :] + ) + dk_ptrs = ( + dk_ptr + + batch_id * dk_b_stride + + s_offs[:, None] * dk_s_stride + + v_head * dk_h_stride + + chan_off[None, :] + ) + d_normed += tl.load(dq_ptrs, mask=s_mask[:, None] & is_query, other=0.0).to(tl.float32) + d_normed += tl.load(dk_ptrs, mask=s_mask[:, None] & is_key, other=0.0).to(tl.float32) + + silu_ptrs = ( + silu_save_ptr + + batch_id * silu_b_stride + + chan[None, :] * silu_c_stride + + s_offs[:, None] * silu_s_stride + ) + silu_bf16 = tl.load(silu_ptrs, mask=s_mask[:, None], other=0.0).to(tl.float32) + + norm_sq = tl.sum(silu_bf16 * silu_bf16, axis=1) + rstd = 1.0 / tl.sqrt(norm_sq + eps) + s_row = tl.sum(d_normed * silu_bf16, axis=1) + rstd3 = rstd * rstd * rstd + d_silu = rstd[:, None] * d_normed - rstd3[:, None] * silu_bf16 * s_row[:, None] + + dsc_ptrs = ( + d_silu_conv_ptr + + batch_id * dsc_b_stride + + chan[None, :] * dsc_c_stride + + s_offs[:, None] * dsc_s_stride + ) + tl.store(dsc_ptrs, d_silu.to(d_silu_conv_ptr.dtype.element_ty), mask=s_mask[:, None]) + + elif pid_row < R_qkv: + # ---- V: relayout dv -> d_silu_conv[V] ---- + local = pid_row - R_qk + batch_id = local // num_value_heads + head_id = local - batch_id * num_value_heads + dv_ptrs = ( + dv_ptr + + batch_id * dv_b_stride + + s_offs[:, None] * dv_s_stride + + head_id * dv_h_stride + + chan_off[None, :] + ) + dv_val = tl.load(dv_ptrs, mask=s_mask[:, None], other=0.0) + dsc_chan = 2 * qk_channels + head_id * HEAD_DIM + chan_off + dsc_ptrs = ( + d_silu_conv_ptr + + batch_id * dsc_b_stride + + dsc_chan[None, :] * dsc_c_stride + + s_offs[:, None] * dsc_s_stride + ) + tl.store(dsc_ptrs, dv_val, mask=s_mask[:, None]) + + elif pid_row < R_qkvz: + # ---- Z: relayout dgate -> d_qkvzba[z] ---- + local = pid_row - R_qkv + batch_id = local // num_value_heads + head_id = local - batch_id * num_value_heads + dgate_ptrs = ( + dgate_ptr + + batch_id * dgate_b_stride + + s_offs[:, None] * dgate_s_stride + + head_id * dgate_h_stride + + chan_off[None, :] + ) + dgate_val = tl.load(dgate_ptrs, mask=s_mask[:, None], other=0.0) + dz_chan = 2 * qk_channels + v_channels + head_id * HEAD_DIM + chan_off + dz_ptrs = ( + d_qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + dz_chan[None, :] * qkvzba_c_stride + ) + tl.store(dz_ptrs, dgate_val, mask=s_mask[:, None]) + + else: + # ---- g/beta: chain rule -> d_qkvzba[alpha,beta] + atomic d_A_log/d_dt_bias ---- + local = pid_row - R_qkvz + batch_id = local // num_value_heads + head_id = local - batch_id * num_value_heads + beta_chan = 2 * qk_channels + 2 * v_channels + head_id + alpha_chan = beta_chan + num_value_heads + + alpha_ptrs = ( + qkvzba_ptr + + s_offs * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + alpha_chan * qkvzba_c_stride + ) + beta_ptrs = ( + qkvzba_ptr + + s_offs * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + beta_chan * qkvzba_c_stride + ) + alpha = tl.load(alpha_ptrs, mask=s_mask, other=0.0).to(tl.float32) + beta_raw = tl.load(beta_ptrs, mask=s_mask, other=0.0).to(tl.float32) + A_log = tl.load(A_log_ptr + head_id).to(tl.float32) + dt_bias = tl.load(dt_bias_ptr + head_id).to(tl.float32) + + pre = alpha + dt_bias + sigmoid_pre = tl.sigmoid(pre) + softplus_pre = tl.log(1.0 + tl.exp(pre)) + exp_A = tl.exp(A_log) + g = -exp_A * softplus_pre + beta_sig = tl.sigmoid(beta_raw) + + dg_ptrs = ( + dg_ptr + batch_id * dg_b_stride + s_offs * dg_s_stride + head_id * dg_h_stride + ) + dbeta_ptrs = ( + dbeta_ptr + + batch_id * dbeta_b_stride + + s_offs * dbeta_s_stride + + head_id * dbeta_h_stride + ) + d_g = tl.load(dg_ptrs, mask=s_mask, other=0.0).to(tl.float32) + d_beta_out = tl.load(dbeta_ptrs, mask=s_mask, other=0.0).to(tl.float32) + + d_alpha = d_g * (-exp_A * sigmoid_pre) + d_beta_raw = d_beta_out * beta_sig * (1.0 - beta_sig) + + d_alpha_ptrs = ( + d_qkvzba_ptr + + s_offs * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + alpha_chan * qkvzba_c_stride + ) + d_beta_ptrs = ( + d_qkvzba_ptr + + s_offs * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + beta_chan * qkvzba_c_stride + ) + tl.store(d_alpha_ptrs, d_alpha.to(d_qkvzba_ptr.dtype.element_ty), mask=s_mask) + tl.store(d_beta_ptrs, d_beta_raw.to(d_qkvzba_ptr.dtype.element_ty), mask=s_mask) + + d_g_masked = tl.where(s_mask, d_g, 0.0) + d_alpha_masked = tl.where(s_mask, d_alpha, 0.0) + d_A_log_partial = tl.sum(d_g_masked * g) + d_dt_bias_partial = tl.sum(d_alpha_masked) + tl.atomic_add(d_A_log_ptr + head_id, d_A_log_partial) + tl.atomic_add(d_dt_bias_ptr + head_id, d_dt_bias_partial) + + +# --------------------------------------------------------------------------- +# Backward orchestration +# --------------------------------------------------------------------------- + + +def _mega_pre_gated_delta_rule_backward( + qkvzba: Tensor, + conv1d_weight: Tensor, + silu_qk_save: Tensor, + dq: Tensor, + dk: Tensor, + dv: Tensor, + dgate: Tensor, + dbeta: Tensor, + dg: Tensor, + A_log: Tensor, + dt_bias: Tensor, + *, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, + seq_idx: Optional[Tensor] = None, +) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + """Two-launch mega backward: one fused branch kernel + external conv bwd. + + Collapses the four streamed branch kernels (QK l2norm/repeat, V layout, Z + layout, g/beta chain rule) into a single Triton launch, then delegates the + depthwise conv input/weight gradients to ``causal_conv1d_bwd_function`` as + the streamed path does. Returns ``(d_qkvzba, d_weight, d_A_log, d_dt_bias)``. + """ + + seq_len, batch, _ = qkvzba.shape + qk_channels = num_key_heads * key_head_dim + v_channels = num_value_heads * value_head_dim + conv_dim = 2 * qk_channels + v_channels + k_w = conv1d_weight.shape[-1] + device = qkvzba.device + head_dim = key_head_dim + + weight_2d = conv1d_weight.view(conv1d_weight.shape[0], k_w) + # Channel-last conv input view (stride(1)==1) — no copy. + qkvzba_conv = qkvzba[:, :, :conv_dim].permute(1, 2, 0) + + # d_silu_conv channel-last (b, conv_dim, s), stride(1)==1. + d_silu_conv = torch.empty( + (batch, seq_len, conv_dim), dtype=qkvzba.dtype, device=device + ).permute(0, 2, 1) + d_qkvzba = torch.empty_like(qkvzba) + d_A_log_fp32 = torch.zeros(num_value_heads, dtype=torch.float32, device=device) + d_dt_bias_fp32 = torch.zeros(num_value_heads, dtype=torch.float32, device=device) + + R_qk = batch * 2 * num_key_heads + R_v = batch * num_value_heads + R_z = batch * num_value_heads + R_gb = batch * num_value_heads + num_rows = R_qk + R_v + R_z + R_gb + + grid = lambda meta: (num_rows, triton.cdiv(seq_len, meta["BLOCK_S"])) + _mega_backward_kernel[grid]( + dq, + dk, + dv, + dgate, + dg, + dbeta, + silu_qk_save, + qkvzba, + A_log, + dt_bias, + d_silu_conv, + d_qkvzba, + d_A_log_fp32, + d_dt_bias_fp32, + seq_len, + num_key_heads, + num_value_heads, + qk_channels, + v_channels, + R_qk, + R_v, + R_z, + _L2NORM_EPS, + dq.stride(0), + dq.stride(1), + dq.stride(2), + dk.stride(0), + dk.stride(1), + dk.stride(2), + dv.stride(0), + dv.stride(1), + dv.stride(2), + dgate.stride(0), + dgate.stride(1), + dgate.stride(2), + dg.stride(0), + dg.stride(1), + dg.stride(2), + dbeta.stride(0), + dbeta.stride(1), + dbeta.stride(2), + silu_qk_save.stride(0), + silu_qk_save.stride(1), + silu_qk_save.stride(2), + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + d_silu_conv.stride(0), + d_silu_conv.stride(1), + d_silu_conv.stride(2), + HEAD_DIM=head_dim, + REPEAT=num_value_heads // num_key_heads, + ) + + # External conv backward: writes d_x into d_qkvzba's conv slice (strided + # view, no copy) and returns d_weight. Same call shape as the streamed path. + seq_stride = qkvzba.stride(0) + batch_stride = qkvzba.stride(1) + d_x_conv_view = d_qkvzba.as_strided( + (batch, conv_dim, seq_len), + (batch_stride, 1, seq_stride), + ) + if _causal_conv1d_bwd_function is None: + raise RuntimeError( + "Fused pre-gated-delta-rule backward requires the 'causal_conv1d' package. " + "Install it, or use pre_gated_delta_rule_impl='unfused'." + ) + _, d_weight_fp32, _, _ = _causal_conv1d_bwd_function( + qkvzba_conv, + weight_2d, + None, # no bias + d_silu_conv, + seq_idx, + None, # initial_states + None, # dfinal_states + d_x_conv_view, # dx pre-allocated into d_qkvzba's conv slice + False, # return_dinitial_states + True, # activation (silu folded into conv bwd) + ) + + d_weight = d_weight_fp32.view(*conv1d_weight.shape).to(conv1d_weight.dtype) + d_A_log = d_A_log_fp32.to(A_log.dtype) + d_dt_bias = d_dt_bias_fp32.to(dt_bias.dtype) + return d_qkvzba, d_weight, d_A_log, d_dt_bias + + +# --------------------------------------------------------------------------- +# Autograd wiring +# --------------------------------------------------------------------------- + + +class _FusedMegaPreGatedDeltaRuleFunction(torch.autograd.Function): + """Autograd entry point for the mega path. + + Forward dispatches to the single-launch mega forward. Backward currently + reuses the streamed conv-backend-delegated backward (which consumes the + same saved ``silu_qk_save`` layout); a dedicated mega backward is layered + in behind this same entry point. + """ + + @staticmethod + def forward( + ctx, + qkvzba, + conv1d_weight, + A_log, + dt_bias, + cu_seqlens, + seq_idx, + num_key_heads, + num_value_heads, + key_head_dim, + value_head_dim, + ): + ctx.num_key_heads = num_key_heads + ctx.num_value_heads = num_value_heads + ctx.key_head_dim = key_head_dim + ctx.value_head_dim = value_head_dim + query, key, value, gate, beta, g, silu_qk_save = ( + _mega_pre_gated_delta_rule_forward( + qkvzba, + conv1d_weight, + A_log, + dt_bias, + num_key_heads=num_key_heads, + num_value_heads=num_value_heads, + key_head_dim=key_head_dim, + value_head_dim=value_head_dim, + cu_seqlens=cu_seqlens, + ) + ) + ctx.has_seq_idx = seq_idx is not None + if ctx.has_seq_idx: + ctx.save_for_backward(qkvzba, conv1d_weight, A_log, dt_bias, silu_qk_save, seq_idx) + else: + ctx.save_for_backward(qkvzba, conv1d_weight, A_log, dt_bias, silu_qk_save) + return query, key, value, gate, beta, g + + @staticmethod + def backward(ctx, dq, dk, dv, dgate, dbeta, dg): + if ctx.has_seq_idx: + qkvzba, conv1d_weight, A_log, dt_bias, silu_qk_save, seq_idx = ctx.saved_tensors + else: + qkvzba, conv1d_weight, A_log, dt_bias, silu_qk_save = ctx.saved_tensors + seq_idx = None + d_qkvzba, d_weight, d_A_log, d_dt_bias = _mega_pre_gated_delta_rule_backward( + qkvzba, + conv1d_weight, + silu_qk_save, + dq, + dk, + dv, + dgate, + dbeta, + dg, + A_log, + dt_bias, + num_key_heads=ctx.num_key_heads, + num_value_heads=ctx.num_value_heads, + key_head_dim=ctx.key_head_dim, + value_head_dim=ctx.value_head_dim, + seq_idx=seq_idx, + ) + return (d_qkvzba, d_weight, d_A_log, d_dt_bias, None, None, None, None, None, None) + + +def fused_mega_pre_gated_delta_rule( + qkvzba: Tensor, + conv1d_weight: Tensor, + conv1d_bias: Optional[Tensor], + A_log: Tensor, + dt_bias: Tensor, + *, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, + use_qk_l2norm: bool = True, + cu_seqlens: Optional[Tensor] = None, + seq_idx: Optional[Tensor] = None, +) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + """Mega fused pre-gated-delta-rule entry point. + + Args: + qkvzba: ``[seq_len, batch, in_proj_dim]`` projection output. + conv1d_weight: ``[conv_dim, 1, k_w]`` depthwise conv weight. + conv1d_bias: Must be ``None`` in the mega path. + A_log: ``[num_value_heads]`` raw decay parameter. + dt_bias: ``[num_value_heads]`` time-step bias. + num_key_heads / num_value_heads / key_head_dim / value_head_dim: GDN + architecture parameters. ``num_value_heads`` must be a multiple of + ``num_key_heads`` and ``key_head_dim == value_head_dim``. + use_qk_l2norm: Must be ``True`` for parity with the streamed path. + cu_seqlens: Optional packed THD cumulative sequence lengths. + seq_idx: Optional precomputed token-to-sequence map for packed THD mode. + + Returns: + ``(query, key, value, gate, beta, g)`` matching the unfused and streamed + fused pre-GDR APIs. + """ + + assert qkvzba.is_cuda, ( + "fused_mega_pre_gated_delta_rule requires CUDA inputs; " + f"got qkvzba.device={qkvzba.device}." + ) + assert conv1d_bias is None, ( + "Conv bias is not supported by fused_mega_pre_gated_delta_rule " + "(production GDN config has none)." + ) + assert use_qk_l2norm, ( + "use_qk_l2norm=False is not supported by fused_mega_pre_gated_delta_rule " + "(the backward closes over the l2norm path)." + ) + assert num_value_heads % num_key_heads == 0, ( + f"{num_value_heads=} must be a multiple of {num_key_heads=}." + ) + assert key_head_dim == value_head_dim, ( + "fused_mega_pre_gated_delta_rule currently requires " + f"key_head_dim == value_head_dim; got {key_head_dim=} {value_head_dim=}." + ) + if cu_seqlens is not None: + assert cu_seqlens.is_cuda, ( + "Packed fused_mega_pre_gated_delta_rule requires CUDA cu_seqlens; " + f"got cu_seqlens.device={cu_seqlens.device}." + ) + assert cu_seqlens.dtype == torch.int32, ( + "Packed fused_mega_pre_gated_delta_rule requires int32 cu_seqlens; " + f"got {cu_seqlens.dtype=}." + ) + assert cu_seqlens.dim() == 1, ( + "Packed fused_mega_pre_gated_delta_rule expects 1-D cu_seqlens; " + f"got {cu_seqlens.shape=}." + ) + assert qkvzba.shape[1] == 1, ( + "Packed THD fused_mega_pre_gated_delta_rule expects batch dimension 1; " + f"got qkvzba.shape={qkvzba.shape}." + ) + assert cu_seqlens.shape[0] >= 2, ( + "Packed fused_mega_pre_gated_delta_rule requires at least one packed sequence; " + f"got {cu_seqlens.shape=}." + ) + assert cu_seqlens[0].item() == 0, ( + "Packed fused_mega_pre_gated_delta_rule requires cu_seqlens[0] == 0, " + f"got {cu_seqlens[0].item()}." + ) + assert cu_seqlens[-1].item() == qkvzba.shape[0], ( + "Packed fused_mega_pre_gated_delta_rule requires cu_seqlens[-1] to match " + f"seq_len, got {cu_seqlens[-1].item()} vs {qkvzba.shape[0]}." + ) + cu_seqlens = cu_seqlens.contiguous() + seq_idx = _resolve_packed_seq_idx(cu_seqlens, seq_idx, qkvzba.shape[0]) + else: + assert seq_idx is None, "seq_idx requires cu_seqlens for packed THD mode." + + return _FusedMegaPreGatedDeltaRuleFunction.apply( + qkvzba, + conv1d_weight, + A_log, + dt_bias, + cu_seqlens, + seq_idx, + num_key_heads, + num_value_heads, + key_head_dim, + value_head_dim, + ) diff --git a/megatron/core/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py index 6eed7581d03..4319af230ff 100644 --- a/megatron/core/fusions/fused_mla_yarn_rope_apply.py +++ b/megatron/core/fusions/fused_mla_yarn_rope_apply.py @@ -41,12 +41,17 @@ def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size): last_cum_seqlen = cur_cum_seqlen seq_idx += 1 if cp_size > 1: - if token_idx < this_seq_len // 2: - token_idx = token_idx + cp_rank * this_seq_len // 2 + first_cp_seg = (this_seq_len + 1) // 2 + second_cp_seg = this_seq_len // 2 + if token_idx < first_cp_seg: + token_idx = token_idx + cp_rank * first_cp_seg else: - token_idx = (token_idx - this_seq_len // 2) + ( - 2 * cp_size - cp_rank - 1 - ) * this_seq_len // 2 + token_idx = ( + token_idx + - first_cp_seg + + cp_size * first_cp_seg + + (cp_size - cp_rank - 1) * second_cp_seg + ) return token_idx diff --git a/megatron/core/fusions/fused_mrope.py b/megatron/core/fusions/fused_mrope.py new file mode 100644 index 00000000000..6ebad4df933 --- /dev/null +++ b/megatron/core/fusions/fused_mrope.py @@ -0,0 +1,871 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Triton fused multimodal RoPE apply. + +The fused path consumes the raw three-axis mRoPE frequencies with shape +``[3, batch, seq, rotary_dim / 2]`` and applies the rotation directly to a BSHD +tensor. It supports both Qwen2-VL section-based mRoPE and Qwen3.5-VL +stride-3 interleaved mRoPE layouts. +""" + +from __future__ import annotations + +from typing import List, Optional +from unittest.mock import MagicMock + +import torch + +from megatron.core.utils import null_decorator + +try: + import triton + import triton.language as tl + + HAVE_TRITON = True +except ImportError: + HAVE_TRITON = False + +if not HAVE_TRITON: + triton = MagicMock() + triton.jit = null_decorator + tl = MagicMock() + + +def _smallest_power_of_2_at_least(x: int) -> int: + block = 1 + while block < x: + block *= 2 + return block + + +def _expected_interleaved_mrope_section(half_rotary_dim: int) -> tuple[int, int, int]: + return ((half_rotary_dim + 2) // 3, (half_rotary_dim + 1) // 3, half_rotary_dim // 3) + + +def _validate_mrope_section( + mrope_section: List[int], half_rotary_dim: int, interleaved_mrope: bool +) -> tuple[int, int, int]: + assert len(mrope_section) == 3, f"mrope_section must have length 3, got {mrope_section}" + + sec_t, sec_h, sec_w = (int(section) for section in mrope_section) + assert ( + min(sec_t, sec_h, sec_w) >= 0 + ), f"mrope_section values must be non-negative, got {mrope_section}" + assert half_rotary_dim > 0, "raw mRoPE rotary dim must be greater than 0" + assert ( + sec_t + sec_h + sec_w == half_rotary_dim + ), f"mrope_section {mrope_section} must sum to rotary_dim / 2 = {half_rotary_dim}" + if interleaved_mrope: + expected = _expected_interleaved_mrope_section(half_rotary_dim) + assert (sec_t, sec_h, sec_w) == expected, ( + f"interleaved mRoPE with rotary_dim / 2 = {half_rotary_dim} requires " + f"mrope_section {list(expected)}, got {mrope_section}" + ) + return sec_t, sec_h, sec_w + + +def _validate_mrope_inputs( + t: torch.Tensor, freqs: torch.Tensor, mrope_section: List[int], interleaved_mrope: bool +) -> tuple[int, int, int, int, int, int, int, int]: + assert t.dim() == 4, f"t must have shape [seq, batch, heads, head_dim], got {t.shape}" + assert freqs.dim() == 4, ( + "raw mRoPE freqs must have shape [3, batch, seq, rotary_dim / 2], " f"got {freqs.shape}" + ) + + seq, batch, heads, head_dim = t.shape + axes, freq_batch, freq_seq, half_rotary_dim = freqs.shape + assert axes == 3, f"raw mRoPE freqs first dimension must be 3, got {axes}" + assert ( + freq_batch == batch and freq_seq == seq + ), f"freqs shape {tuple(freqs.shape)} is incompatible with t shape {tuple(t.shape)}" + + sec_t, sec_h, sec_w = _validate_mrope_section(mrope_section, half_rotary_dim, interleaved_mrope) + + rotary_dim = half_rotary_dim * 2 + assert ( + rotary_dim <= head_dim + ), f"raw mRoPE rotary dim {rotary_dim} exceeds input head dim {head_dim}" + return seq, batch, heads, head_dim, half_rotary_dim, sec_t, sec_h, sec_w + + +def _validate_mrope_thd_inputs( + t: torch.Tensor, + cu_seqlens: torch.Tensor, + freqs: torch.Tensor, + mrope_section: List[int], + interleaved_mrope: bool, + cp_size: int, +) -> tuple[int, int, int, int, int, int, int]: + assert t.dim() == 3, f"t must have shape [tokens, heads, head_dim], got {t.shape}" + assert freqs.dim() == 4, ( + "raw mRoPE freqs must have shape [3, 1, total_seqlen, rotary_dim / 2], " + f"got {freqs.shape}" + ) + assert cu_seqlens.dim() == 1, f"cu_seqlens must be 1D, got {cu_seqlens.shape}" + + tokens, heads, head_dim = t.shape + axes, freq_batch, freq_seq, half_rotary_dim = freqs.shape + assert axes == 3, f"raw mRoPE freqs first dimension must be 3, got {axes}" + assert freq_batch == 1, ( + "raw mRoPE THD freqs must have singleton batch dimension, " f"got {freqs.shape}" + ) + assert freq_seq == tokens * cp_size, ( + "raw mRoPE THD freqs sequence length must match local tokens times cp_size, " + f"got freqs.shape[2]={freq_seq}, tokens={tokens}, cp_size={cp_size}" + ) + + sec_t, sec_h, sec_w = _validate_mrope_section(mrope_section, half_rotary_dim, interleaved_mrope) + rotary_dim = half_rotary_dim * 2 + assert ( + rotary_dim <= head_dim + ), f"raw mRoPE rotary dim {rotary_dim} exceeds input head dim {head_dim}" + return tokens, heads, head_dim, half_rotary_dim, sec_t, sec_h, sec_w + + +def get_fused_mrope_unavailable_reason( + t: Optional[torch.Tensor] = None, + freqs: Optional[torch.Tensor] = None, + rotary_interleaved: bool = False, +) -> Optional[str]: + """Return why fused mRoPE cannot run, or None when it is launchable.""" + if not HAVE_TRITON: + return "Triton is not available" + if rotary_interleaved: + return "rotary_interleaved=True is not supported" + if t is None or freqs is None: + return None + if not t.is_cuda or not freqs.is_cuda: + return "Triton fused mRoPE requires CUDA tensors" + if t.device != freqs.device: + return ( + "Triton fused mRoPE requires t and freqs on the same device, " + f"got {t.device} and {freqs.device}" + ) + if freqs.dtype != torch.float32: + return f"raw mRoPE freqs must be float32, got {freqs.dtype}" + if t.dtype not in (torch.float16, torch.bfloat16, torch.float32): + return f"input dtype {t.dtype} is not supported" + if t.stride(-1) != 1: + return f"input head dimension must be contiguous, got stride {t.stride()}" + try: + capability = torch.cuda.get_device_capability(t.device) + except RuntimeError as exc: + return f"could not query CUDA device capability: {exc}" + if capability < (7, 0): + return f"requires CUDA compute capability >= 7.0, got {capability[0]}.{capability[1]}" + if t.dtype == torch.bfloat16 and capability < (8, 0): + return ( + "requires CUDA compute capability >= 8.0 for bfloat16 inputs, " + f"got {capability[0]}.{capability[1]}" + ) + return None + + +def get_fused_mrope_thd_unavailable_reason( + t: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, + freqs: Optional[torch.Tensor] = None, + rotary_interleaved: bool = False, + cp_size: int = 1, + cp_rank: int = 0, +) -> Optional[str]: + """Return why fused THD mRoPE cannot run, or None when it is launchable.""" + if not HAVE_TRITON: + return "Triton is not available" + if rotary_interleaved: + return "rotary_interleaved=True is not supported" + if cp_size < 1: + return f"cp_size must be positive, got {cp_size}" + if cp_rank < 0 or cp_rank >= cp_size: + return f"cp_rank must be in [0, {cp_size}), got {cp_rank}" + if t is None or cu_seqlens is None or freqs is None: + return None + if t.dim() != 3: + return ( + f"THD fused mRoPE expects t with shape [tokens, heads, head_dim], got {tuple(t.shape)}" + ) + if freqs.dim() != 4: + return ( + "raw mRoPE THD freqs must have shape [3, 1, total_seqlen, rotary_dim / 2], " + f"got {tuple(freqs.shape)}" + ) + if cu_seqlens.dim() != 1: + return f"cu_seqlens must be 1D, got {tuple(cu_seqlens.shape)}" + if not t.is_cuda or not freqs.is_cuda or not cu_seqlens.is_cuda: + return "Triton fused THD mRoPE requires CUDA tensors" + if t.device != freqs.device or t.device != cu_seqlens.device: + return ( + "Triton fused THD mRoPE requires t, freqs, and cu_seqlens on the same device, " + f"got {t.device}, {freqs.device}, and {cu_seqlens.device}" + ) + if freqs.dtype != torch.float32: + return f"raw mRoPE freqs must be float32, got {freqs.dtype}" + if t.dtype not in (torch.float16, torch.bfloat16, torch.float32): + return f"input dtype {t.dtype} is not supported" + if cu_seqlens.dtype not in (torch.int32, torch.int64): + return f"cu_seqlens dtype {cu_seqlens.dtype} is not supported" + if t.stride(-1) != 1: + return f"input head dimension must be contiguous, got stride {t.stride()}" + if freqs.shape[0] != 3 or freqs.shape[1] != 1: + return ( + "raw mRoPE THD freqs must have shape [3, 1, total_seqlen, rotary_dim / 2], " + f"got {tuple(freqs.shape)}" + ) + if cp_size > 1 and freqs.shape[2] % cp_size != 0: + return ( + "raw mRoPE THD freqs sequence length must be divisible by context parallel size, " + f"got freqs.shape[2]={freqs.shape[2]}, cp_size={cp_size}" + ) + if cp_size > 1: + # Guard: each packed sub-sequence length must satisfy seqlen % cp_size == 0. + seq_bounds = cu_seqlens.tolist() + for seq_start, seq_end in zip(seq_bounds[:-1], seq_bounds[1:]): + if (seq_end - seq_start) % cp_size != 0: + return ( + "each packed THD sub-sequence length must be divisible by context " + f"parallel size, got sub-sequence length {seq_end - seq_start} " + f"with cp_size={cp_size}" + ) + if freqs.shape[2] != t.shape[0] * cp_size: + return ( + "raw mRoPE THD freqs sequence length must match local tokens times cp_size, " + f"got freqs.shape[2]={freqs.shape[2]}, tokens={t.shape[0]}, cp_size={cp_size}" + ) + try: + capability = torch.cuda.get_device_capability(t.device) + except RuntimeError as exc: + return f"could not query CUDA device capability: {exc}" + if capability < (7, 0): + return f"requires CUDA compute capability >= 7.0, got {capability[0]}.{capability[1]}" + if t.dtype == torch.bfloat16 and capability < (8, 0): + return ( + "requires CUDA compute capability >= 8.0 for bfloat16 inputs, " + f"got {capability[0]}.{capability[1]}" + ) + return None + + +def can_launch_fused_mrope( + t: Optional[torch.Tensor] = None, + freqs: Optional[torch.Tensor] = None, + rotary_interleaved: bool = False, +) -> bool: + """Return whether the Triton fused mRoPE kernel can be launched.""" + return get_fused_mrope_unavailable_reason(t, freqs, rotary_interleaved) is None + + +def can_launch_fused_mrope_thd( + t: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, + freqs: Optional[torch.Tensor] = None, + rotary_interleaved: bool = False, + cp_size: int = 1, + cp_rank: int = 0, +) -> bool: + """Return whether the Triton fused THD mRoPE kernel can be launched.""" + return ( + get_fused_mrope_thd_unavailable_reason( + t, + cu_seqlens, + freqs, + rotary_interleaved=rotary_interleaved, + cp_size=cp_size, + cp_rank=cp_rank, + ) + is None + ) + + +def mrope_freqs_to_rotary_emb( + freqs: torch.Tensor, + mrope_section: List[int], + interleaved_mrope: bool = False, + rotary_interleaved: bool = False, +) -> torch.Tensor: + """Convert raw mRoPE freqs to the unfused RoPE embedding layout. + + Args: + freqs: Raw mRoPE frequencies with shape ``[3, batch, seq, rotary_dim / 2]``. + mrope_section: Temporal, height, and width channel sections. + interleaved_mrope: Use Qwen3.5-VL stride-3 T/H/W layout when True. Use + Qwen2-VL section layout when False. + rotary_interleaved: Use adjacent-pair RoPE layout when True. This is + available for reference conversion; fused Triton currently supports + split-half layout only. + + Returns: + Tensor with shape ``[seq, batch, 1, rotary_dim]``. + """ + assert freqs.dim() == 4, ( + "raw mRoPE freqs must have shape [3, batch, seq, rotary_dim / 2], " f"got {freqs.shape}" + ) + assert freqs.size(0) == 3, f"raw mRoPE freqs first dimension must be 3, got {freqs.size(0)}" + assert len(mrope_section) == 3, f"mrope_section must have length 3, got {mrope_section}" + + half_rotary_dim = freqs.size(-1) + sec_t, sec_h, sec_w = _validate_mrope_section(mrope_section, half_rotary_dim, interleaved_mrope) + + if interleaved_mrope: + freqs_out = freqs[0].clone() + for dim_idx, offset in enumerate((1, 2), start=1): + length = int(mrope_section[dim_idx]) * 3 + idx = slice(offset, length, 3) + freqs_out[..., idx] = freqs[dim_idx, ..., idx] + if rotary_interleaved: + batch = freqs_out.shape[0] + emb = torch.stack( + (freqs_out.reshape(batch, -1, 1), freqs_out.reshape(batch, -1, 1)), dim=-1 + ) + emb = emb.view(batch, freqs_out.shape[1], -1) + else: + emb = torch.cat((freqs_out, freqs_out), dim=-1) + else: + if rotary_interleaved: + batch = freqs.shape[1] + emb = torch.stack( + (freqs.reshape(3, batch, -1, 1), freqs.reshape(3, batch, -1, 1)), dim=-1 + ).view(3, batch, freqs.shape[2], -1) + mrope_section_doubled = list(mrope_section) * 2 + emb = torch.cat( + [chunk[i % 3] for i, chunk in enumerate(emb.split(mrope_section_doubled, dim=-1))], + dim=-1, + ) + else: + freqs_out = torch.empty_like(freqs[0]) + freqs_out[..., :sec_t] = freqs[0, ..., :sec_t] + freqs_out[..., sec_t : sec_t + sec_h] = freqs[1, ..., sec_t : sec_t + sec_h] + freqs_out[..., sec_t + sec_h :] = freqs[2, ..., sec_t + sec_h :] + emb = torch.cat((freqs_out, freqs_out), dim=-1) + return emb[..., None, :].transpose(0, 1).contiguous() + + +@triton.jit +def _mrope_axis( + k, + SEC_T: tl.constexpr, + SEC_H: tl.constexpr, + SEC_W: tl.constexpr, + INTERLEAVED_MROPE: tl.constexpr, +): + if INTERLEAVED_MROPE: + rem = k % 3 + section_idx = k // 3 + is_h = (rem == 1) & (section_idx < SEC_H) + is_w = (rem == 2) & (section_idx < SEC_W) + return tl.where(is_h, 1, tl.where(is_w, 2, 0)) + + is_h = (k >= SEC_T) & (k < (SEC_T + SEC_H)) + is_w = k >= (SEC_T + SEC_H) + return tl.where(is_h, 1, tl.where(is_w, 2, 0)) + + +@triton.jit +def _fused_mrope_kernel( + T, + FREQS, + OUT, + t_s_seq, + t_s_batch, + t_s_head, + t_s_dim, + f_s_axis, + f_s_batch, + f_s_seq, + f_s_dim, + o_s_seq, + o_s_batch, + o_s_head, + o_s_dim, + HEAD_DIM: tl.constexpr, + HALF_ROTARY_DIM: tl.constexpr, + PASS_DIM: tl.constexpr, + SEC_T: tl.constexpr, + SEC_H: tl.constexpr, + SEC_W: tl.constexpr, + INTERLEAVED_MROPE: tl.constexpr, + ROTARY_INTERLEAVED: tl.constexpr, + INVERSE: tl.constexpr, + BLOCK_HALF: tl.constexpr, + BLOCK_PASS: tl.constexpr, +): + seq_idx = tl.program_id(0) + batch_idx = tl.program_id(1) + head_idx = tl.program_id(2) + + k = tl.arange(0, BLOCK_HALF) + mask = k < HALF_ROTARY_DIM + + axis = _mrope_axis(k, SEC_T, SEC_H, SEC_W, INTERLEAVED_MROPE) + + freqs_offset = axis * f_s_axis + batch_idx * f_s_batch + seq_idx * f_s_seq + k * f_s_dim + freqs = tl.load(FREQS + freqs_offset, mask=mask, other=0.0) + # Match PyTorch pointwise dtype semantics: cast cos/sin before the multiply. + cos_v = tl.cos(freqs).to(OUT.dtype.element_ty) + sin_v = tl.sin(freqs).to(OUT.dtype.element_ty) + if INVERSE: + sin_v = -sin_v + + t_base = T + seq_idx * t_s_seq + batch_idx * t_s_batch + head_idx * t_s_head + out_base = OUT + seq_idx * o_s_seq + batch_idx * o_s_batch + head_idx * o_s_head + + if ROTARY_INTERLEAVED: + lo_offset = (2 * k) * t_s_dim + hi_offset = (2 * k + 1) * t_s_dim + out_lo_offset = (2 * k) * o_s_dim + out_hi_offset = (2 * k + 1) * o_s_dim + else: + lo_offset = k * t_s_dim + hi_offset = (k + HALF_ROTARY_DIM) * t_s_dim + out_lo_offset = k * o_s_dim + out_hi_offset = (k + HALF_ROTARY_DIM) * o_s_dim + + t_lo = tl.load(t_base + lo_offset, mask=mask, other=0.0).to(OUT.dtype.element_ty) + t_hi = tl.load(t_base + hi_offset, mask=mask, other=0.0).to(OUT.dtype.element_ty) + + lo_cos = (t_lo * cos_v).to(OUT.dtype.element_ty) + hi_sin = (t_hi * sin_v).to(OUT.dtype.element_ty) + hi_cos = (t_hi * cos_v).to(OUT.dtype.element_ty) + lo_sin = (t_lo * sin_v).to(OUT.dtype.element_ty) + + out_lo = (lo_cos - hi_sin).to(OUT.dtype.element_ty) + out_hi = (hi_cos + lo_sin).to(OUT.dtype.element_ty) + + tl.store(out_base + out_lo_offset, out_lo, mask=mask) + tl.store(out_base + out_hi_offset, out_hi, mask=mask) + + if PASS_DIM > 0: + pass_idx = tl.arange(0, BLOCK_PASS) + pass_mask = pass_idx < PASS_DIM + src_dim = 2 * HALF_ROTARY_DIM + pass_idx + pass_values = tl.load(t_base + src_dim * t_s_dim, mask=pass_mask, other=0.0) + tl.store(out_base + src_dim * o_s_dim, pass_values, mask=pass_mask) + + +@triton.jit +def _fused_mrope_thd_kernel( + T, + CU_SEQLENS, + FREQS, + OUT, + t_s_token, + t_s_head, + t_s_dim, + cu_s_idx, + f_s_axis, + f_s_seq, + f_s_dim, + o_s_token, + o_s_head, + o_s_dim, + NUM_SEQS, + HEAD_DIM: tl.constexpr, + HALF_ROTARY_DIM: tl.constexpr, + PASS_DIM: tl.constexpr, + SEC_T: tl.constexpr, + SEC_H: tl.constexpr, + SEC_W: tl.constexpr, + INTERLEAVED_MROPE: tl.constexpr, + ROTARY_INTERLEAVED: tl.constexpr, + INVERSE: tl.constexpr, + CP_SIZE: tl.constexpr, + CP_RANK: tl.constexpr, + FP32_COMPUTE: tl.constexpr, + BLOCK_HALF: tl.constexpr, + BLOCK_PASS: tl.constexpr, +): + token_idx = tl.program_id(0) + head_idx = tl.program_id(1) + + freq_seq_idx = token_idx + seq_i = 0 + while seq_i < NUM_SEQS: + global_start = tl.load(CU_SEQLENS + seq_i * cu_s_idx) + global_end = tl.load(CU_SEQLENS + (seq_i + 1) * cu_s_idx) + local_start = global_start // CP_SIZE + local_end = global_end // CP_SIZE + in_seq = (token_idx >= local_start) & (token_idx < local_end) + local_offset = token_idx - local_start + + if CP_SIZE > 1: + local_seq_len = local_end - local_start + first_cp_seg = (local_seq_len + 1) // 2 + second_cp_seg = local_seq_len // 2 + first_freq_idx = global_start + CP_RANK * first_cp_seg + local_offset + second_freq_idx = ( + global_end - (CP_RANK + 1) * second_cp_seg + (local_offset - first_cp_seg) + ) + seq_freq_idx = tl.where(local_offset < first_cp_seg, first_freq_idx, second_freq_idx) + else: + seq_freq_idx = global_start + local_offset + + freq_seq_idx = tl.where(in_seq, seq_freq_idx, freq_seq_idx) + seq_i += 1 + + k = tl.arange(0, BLOCK_HALF) + mask = k < HALF_ROTARY_DIM + axis = _mrope_axis(k, SEC_T, SEC_H, SEC_W, INTERLEAVED_MROPE) + + freqs_offset = axis * f_s_axis + freq_seq_idx * f_s_seq + k * f_s_dim + freqs = tl.load(FREQS + freqs_offset, mask=mask, other=0.0) + if FP32_COMPUTE: + cos_v = tl.cos(freqs) + sin_v = tl.sin(freqs) + else: + cos_v = tl.cos(freqs).to(OUT.dtype.element_ty) + sin_v = tl.sin(freqs).to(OUT.dtype.element_ty) + if INVERSE: + sin_v = -sin_v + + t_base = T + token_idx * t_s_token + head_idx * t_s_head + out_base = OUT + token_idx * o_s_token + head_idx * o_s_head + + if ROTARY_INTERLEAVED: + lo_offset = (2 * k) * t_s_dim + hi_offset = (2 * k + 1) * t_s_dim + out_lo_offset = (2 * k) * o_s_dim + out_hi_offset = (2 * k + 1) * o_s_dim + else: + lo_offset = k * t_s_dim + hi_offset = (k + HALF_ROTARY_DIM) * t_s_dim + out_lo_offset = k * o_s_dim + out_hi_offset = (k + HALF_ROTARY_DIM) * o_s_dim + + if FP32_COMPUTE: + t_lo = tl.load(t_base + lo_offset, mask=mask, other=0.0).to(tl.float32) + t_hi = tl.load(t_base + hi_offset, mask=mask, other=0.0).to(tl.float32) + else: + t_lo = tl.load(t_base + lo_offset, mask=mask, other=0.0).to(OUT.dtype.element_ty) + t_hi = tl.load(t_base + hi_offset, mask=mask, other=0.0).to(OUT.dtype.element_ty) + + if FP32_COMPUTE: + lo_cos = t_lo * cos_v + hi_sin = t_hi * sin_v + hi_cos = t_hi * cos_v + lo_sin = t_lo * sin_v + else: + lo_cos = (t_lo * cos_v).to(OUT.dtype.element_ty) + hi_sin = (t_hi * sin_v).to(OUT.dtype.element_ty) + hi_cos = (t_hi * cos_v).to(OUT.dtype.element_ty) + lo_sin = (t_lo * sin_v).to(OUT.dtype.element_ty) + + if FP32_COMPUTE: + out_lo = lo_cos - hi_sin + out_hi = hi_cos + lo_sin + else: + out_lo = (lo_cos - hi_sin).to(OUT.dtype.element_ty) + out_hi = (hi_cos + lo_sin).to(OUT.dtype.element_ty) + + tl.store(out_base + out_lo_offset, out_lo, mask=mask) + tl.store(out_base + out_hi_offset, out_hi, mask=mask) + + if PASS_DIM > 0: + pass_idx = tl.arange(0, BLOCK_PASS) + pass_mask = pass_idx < PASS_DIM + src_dim = 2 * HALF_ROTARY_DIM + pass_idx + pass_values = tl.load(t_base + src_dim * t_s_dim, mask=pass_mask, other=0.0) + tl.store(out_base + src_dim * o_s_dim, pass_values, mask=pass_mask) + + +def _launch_fused_mrope( + t: torch.Tensor, + freqs: torch.Tensor, + mrope_section: List[int], + interleaved_mrope: bool, + rotary_interleaved: bool, + inverse: bool, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + unavailable_reason = get_fused_mrope_unavailable_reason(t, freqs, rotary_interleaved) + assert unavailable_reason is None, unavailable_reason + + seq, batch, heads, head_dim, half_rotary_dim, sec_t, sec_h, sec_w = _validate_mrope_inputs( + t, freqs, mrope_section, interleaved_mrope + ) + + if out is None: + out = torch.empty_like(t) + else: + assert out.shape == t.shape and out.dtype == t.dtype + assert ( + out.stride(-1) == 1 + ), f"fused mRoPE requires output contiguous head dimension, got {out.stride()}" + + block_half = _smallest_power_of_2_at_least(half_rotary_dim) + pass_dim = head_dim - (2 * half_rotary_dim) + block_pass = _smallest_power_of_2_at_least(max(pass_dim, 1)) + + grid = (seq, batch, heads) + _fused_mrope_kernel[grid]( + t, + freqs, + out, + t.stride(0), + t.stride(1), + t.stride(2), + t.stride(3), + freqs.stride(0), + freqs.stride(1), + freqs.stride(2), + freqs.stride(3), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + HEAD_DIM=head_dim, + HALF_ROTARY_DIM=half_rotary_dim, + PASS_DIM=pass_dim, + SEC_T=sec_t, + SEC_H=sec_h, + SEC_W=sec_w, + INTERLEAVED_MROPE=interleaved_mrope, + ROTARY_INTERLEAVED=rotary_interleaved, + INVERSE=inverse, + BLOCK_HALF=block_half, + BLOCK_PASS=block_pass, + num_warps=4, + ) + return out + + +def _launch_fused_mrope_thd( + t: torch.Tensor, + cu_seqlens: torch.Tensor, + freqs: torch.Tensor, + mrope_section: List[int], + interleaved_mrope: bool, + rotary_interleaved: bool, + inverse: bool, + cp_size: int, + cp_rank: int, + fp32_compute: bool = False, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + unavailable_reason = get_fused_mrope_thd_unavailable_reason( + t, + cu_seqlens, + freqs, + rotary_interleaved=rotary_interleaved, + cp_size=cp_size, + cp_rank=cp_rank, + ) + assert unavailable_reason is None, unavailable_reason + + tokens, heads, head_dim, half_rotary_dim, sec_t, sec_h, sec_w = _validate_mrope_thd_inputs( + t, cu_seqlens, freqs, mrope_section, interleaved_mrope, cp_size + ) + + if out is None: + out = torch.empty_like(t) + else: + assert out.shape == t.shape and out.dtype == t.dtype + assert ( + out.stride(-1) == 1 + ), f"fused THD mRoPE requires output contiguous head dimension, got {out.stride()}" + + block_half = _smallest_power_of_2_at_least(half_rotary_dim) + pass_dim = head_dim - (2 * half_rotary_dim) + block_pass = _smallest_power_of_2_at_least(max(pass_dim, 1)) + num_seqs = cu_seqlens.numel() - 1 + + grid = (tokens, heads) + _fused_mrope_thd_kernel[grid]( + t, + cu_seqlens, + freqs, + out, + t.stride(0), + t.stride(1), + t.stride(2), + cu_seqlens.stride(0), + freqs.stride(0), + freqs.stride(2), + freqs.stride(3), + out.stride(0), + out.stride(1), + out.stride(2), + num_seqs, + HEAD_DIM=head_dim, + HALF_ROTARY_DIM=half_rotary_dim, + PASS_DIM=pass_dim, + SEC_T=sec_t, + SEC_H=sec_h, + SEC_W=sec_w, + INTERLEAVED_MROPE=interleaved_mrope, + ROTARY_INTERLEAVED=rotary_interleaved, + INVERSE=inverse, + CP_SIZE=cp_size, + CP_RANK=cp_rank, + FP32_COMPUTE=fp32_compute, + BLOCK_HALF=block_half, + BLOCK_PASS=block_pass, + num_warps=4, + ) + return out + + +class _FusedMRoPE(torch.autograd.Function): + """Autograd wrapper for fused mRoPE. + + The raw frequency table is generated from position IDs and inverse frequencies, + so gradients are only propagated to the rotated tensor. + """ + + @staticmethod + def forward(ctx, t, freqs, mrope_section, interleaved_mrope, rotary_interleaved): + assert not freqs.requires_grad, "fused mRoPE expects non-gradient raw frequency tensors" + ctx.mrope_section = tuple(int(section) for section in mrope_section) + ctx.interleaved_mrope = bool(interleaved_mrope) + ctx.rotary_interleaved = bool(rotary_interleaved) + ctx.save_for_backward(freqs) + return _launch_fused_mrope( + t, + freqs, + ctx.mrope_section, + ctx.interleaved_mrope, + ctx.rotary_interleaved, + inverse=False, + ) + + @staticmethod + def backward(ctx, grad_output): + (freqs,) = ctx.saved_tensors + grad_input = _launch_fused_mrope( + grad_output.contiguous(), + freqs, + ctx.mrope_section, + ctx.interleaved_mrope, + ctx.rotary_interleaved, + inverse=True, + ) + return grad_input, None, None, None, None + + +class _FusedMRoPETHD(torch.autograd.Function): + """Autograd wrapper for fused THD mRoPE.""" + + @staticmethod + def forward( + ctx, + t, + cu_seqlens, + freqs, + mrope_section, + interleaved_mrope, + rotary_interleaved, + cp_size, + cp_rank, + fp32_compute, + ): + assert not freqs.requires_grad, "fused THD mRoPE expects non-gradient raw frequency tensors" + ctx.mrope_section = tuple(int(section) for section in mrope_section) + ctx.interleaved_mrope = bool(interleaved_mrope) + ctx.rotary_interleaved = bool(rotary_interleaved) + ctx.cp_size = int(cp_size) + ctx.cp_rank = int(cp_rank) + ctx.fp32_compute = bool(fp32_compute) + ctx.save_for_backward(cu_seqlens, freqs) + return _launch_fused_mrope_thd( + t, + cu_seqlens, + freqs, + ctx.mrope_section, + ctx.interleaved_mrope, + ctx.rotary_interleaved, + inverse=False, + cp_size=ctx.cp_size, + cp_rank=ctx.cp_rank, + fp32_compute=ctx.fp32_compute, + ) + + @staticmethod + def backward(ctx, grad_output): + cu_seqlens, freqs = ctx.saved_tensors + grad_input = _launch_fused_mrope_thd( + grad_output.contiguous(), + cu_seqlens, + freqs, + ctx.mrope_section, + ctx.interleaved_mrope, + ctx.rotary_interleaved, + inverse=True, + cp_size=ctx.cp_size, + cp_rank=ctx.cp_rank, + fp32_compute=ctx.fp32_compute, + ) + return grad_input, None, None, None, None, None, None, None, None + + +def fused_apply_mrope( + t: torch.Tensor, + freqs: torch.Tensor, + mrope_section: List[int], + interleaved_mrope: bool = False, + rotary_interleaved: bool = False, +) -> torch.Tensor: + """Apply multimodal RoPE with a fused Triton kernel. + + Args: + t: Input tensor with shape ``[seq, batch, heads, head_dim]``. + freqs: Raw mRoPE frequencies with shape ``[3, batch, seq, rotary_dim / 2]``. + mrope_section: Temporal, height, and width channel sections. + interleaved_mrope: Use Qwen3.5-VL stride-3 T/H/W layout when True. Use + Qwen2-VL section layout when False. + rotary_interleaved: Must be False. The integrated fused mRoPE path + currently supports split-half RoPE layout. + + Returns: + Rotated tensor with the same shape and dtype as ``t``. + """ + return _FusedMRoPE.apply(t, freqs, mrope_section, interleaved_mrope, rotary_interleaved) + + +def fused_apply_mrope_thd( + t: torch.Tensor, + cu_seqlens: torch.Tensor, + freqs: torch.Tensor, + mrope_section: List[int], + interleaved_mrope: bool = False, + rotary_interleaved: bool = False, + cp_size: int = 1, + cp_rank: int = 0, + fp32_compute: bool = False, +) -> torch.Tensor: + """Apply multimodal RoPE to THD-packed tensors with a fused Triton kernel. + + Args: + t: Input tensor with shape ``[total_tokens, heads, head_dim]``. + cu_seqlens: Global cumulative sequence lengths for the packed batch. + freqs: Raw mRoPE frequencies with shape ``[3, 1, total_seqlen, rotary_dim / 2]``. + mrope_section: Temporal, height, and width channel sections. + interleaved_mrope: Use Qwen3.5-VL stride-3 T/H/W layout when True. + rotary_interleaved: Must be False. + cp_size: Context parallel world size for THD token mapping. + cp_rank: Context parallel rank for THD token mapping. + fp32_compute: Apply the rotary math in fp32 and cast directly to output dtype. + + Returns: + Rotated tensor with the same shape and dtype as ``t``. + """ + return _FusedMRoPETHD.apply( + t, + cu_seqlens, + freqs, + mrope_section, + interleaved_mrope, + rotary_interleaved, + cp_size, + cp_rank, + fp32_compute, + ) + + +def is_fused_mrope_available() -> bool: + """Return whether the Triton mRoPE fusion can be used on this host. + + This does not check tensor device, dtype, stride, or CUDA capability. Use + ``can_launch_fused_mrope`` or ``get_fused_mrope_unavailable_reason`` with + tensors before dispatching to the fused kernel. + """ + if not torch.cuda.is_available(): + return False + return can_launch_fused_mrope() diff --git a/megatron/core/fusions/fused_pre_gated_delta_rule.py b/megatron/core/fusions/fused_pre_gated_delta_rule.py new file mode 100644 index 00000000000..72c702ecbb4 --- /dev/null +++ b/megatron/core/fusions/fused_pre_gated_delta_rule.py @@ -0,0 +1,2191 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Fused pre-gated-delta-rule projection kernels. + +The public entry point consumes the dense ``qkvzba`` projection and returns +``query``, ``key``, ``value``, ``gate``, ``beta``, and ``g`` in the layouts +expected by the gated delta rule. The forward path keeps QK, V, Z, and +G/Beta as separate streamed scopes. The backward mirrors those scopes for +layout/l2norm/g-beta work, then delegates depthwise conv gradients to the +``causal_conv1d`` backend. + +Unsupported cases are rejected at the Python entry point: CPU tensors, +conv bias, and ``use_qk_l2norm=False``. Packed THD sequences use separate +QK/V causal-conv kernels so the dense BSHD kernels stay free of packed +metadata and runtime branches. +""" + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + +# The 1.6.1+ ``causal_conv1d`` package exposes the lower-level binding via +# ``causal_conv1d.cpp_functions.causal_conv1d_bwd_function``; older builds +# (still common in some older environments) expose the same +# function under ``causal_conv1d_cuda.causal_conv1d_bwd``. Try both so the +# fast path is taken everywhere the package is installed. +from torch import Tensor + +try: + from causal_conv1d.cpp_functions import ( + causal_conv1d_bwd_function as _causal_conv1d_bwd_function, + ) +except ImportError: + try: + import causal_conv1d_cuda as _causal_conv1d_cuda + + _causal_conv1d_bwd_function = _causal_conv1d_cuda.causal_conv1d_bwd + except ImportError: + # The external causal_conv1d package is optional: only the fused pre-GDR + # backward needs it. Importing this module (and hence GatedDeltaNet) must + # not fail when it is absent; raise a clear error only if the fused path + # is actually exercised. + _causal_conv1d_bwd_function = None + + +_L2NORM_EPS = 1e-6 + +_QK_STREAM_SLOT = 0 +_V_STREAM_SLOT = 2 +_G_BETA_STREAM_SLOT = 3 +_Z_STREAM_SLOT = 4 + +_LAYOUT_BLOCK_S = 64 + + +# --------------------------------------------------------------------------- +# Forward kernels +# --------------------------------------------------------------------------- + + +def _conv_autotune_configs(): + return [ + triton.Config({"BLOCK_S": 16}, num_warps=2, num_stages=2), + triton.Config({"BLOCK_S": 32}, num_warps=2, num_stages=2), + triton.Config({"BLOCK_S": 32}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 32}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_S": 64}, num_warps=2, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=4), + triton.Config({"BLOCK_S": 64}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=3), + triton.Config({"BLOCK_S": 256}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 256}, num_warps=8, num_stages=2), + ] + + +def _g_beta_autotune_configs(): + return [ + triton.Config({"BLOCK_S": 32, "BLOCK_H": 16}, num_warps=2, num_stages=2), + triton.Config({"BLOCK_S": 64, "BLOCK_H": 16}, num_warps=2, num_stages=2), + triton.Config({"BLOCK_S": 64, "BLOCK_H": 32}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 128, "BLOCK_H": 16}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 128, "BLOCK_H": 32}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 128, "BLOCK_H": 64}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 256, "BLOCK_H": 16}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 256, "BLOCK_H": 32}, num_warps=8, num_stages=2), + ] + + +@triton.autotune( + configs=_conv_autotune_configs(), + key=["seq_len", "HEAD_DIM", "K_W", "APPLY_L2", "REPEAT", "NUM_GROUPS"], +) +@triton.jit +def _conv_silu_project_kernel( + qkvzba_ptr, + weight_ptr, + bias_ptr, + out_ptr, + silu_save_ptr, + seq_len, + num_in_heads, + in_channel_offset, + in_group_stride, + silu_save_chan_offset, + silu_save_group_stride, + qkvzba_s_stride, + qkvzba_b_stride, + qkvzba_c_stride, + weight_c_stride, + weight_w_stride, + bias_stride, + out_group_dim_stride, + out_b_stride, + out_s_stride, + out_h_stride, + silu_save_b_stride, + silu_save_c_stride, + silu_save_s_stride, + eps, + HEAD_DIM: tl.constexpr, + K_W: tl.constexpr, + REPEAT: tl.constexpr, + NUM_GROUPS: tl.constexpr, + HAS_BIAS: tl.constexpr, + APPLY_L2: tl.constexpr, + SAVE_SILU: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """Depthwise conv1d + silu + (optional l2norm) + (optional head repeat). + + Grid layout (program_id): + 0: batch * NUM_GROUPS * num_in_heads (flat) + 1: num_seq_blocks + + Args: + in_channel_offset: starting channel index of the first group inside + ``qkvzba``. 0 for QK, ``v_channel_offset`` for V. + in_group_stride: channel distance between logical groups. For QK this + is ``qk_channels`` so group 0 is Q and group 1 is K. For V this is + 0 because ``NUM_GROUPS == 1``. + out_group_dim_stride: output-storage distance between logical groups. QK + passes a grouped output buffer and V passes 0. + """ + + pid_bgh = tl.program_id(0) + pid_s = tl.program_id(1) + + heads_per_batch = num_in_heads * NUM_GROUPS + batch_id = pid_bgh // heads_per_batch + local_bgh = pid_bgh - batch_id * heads_per_batch + group_id = local_bgh // num_in_heads + head_id = local_bgh - group_id * num_in_heads + + chan_off = tl.arange(0, HEAD_DIM) + group_channel_offset = in_channel_offset + group_id * in_group_stride + chan = group_channel_offset + head_id * HEAD_DIM + chan_off + + if HAS_BIAS: + bias = tl.load(bias_ptr + chan * bias_stride).to(tl.float32) + else: + bias = tl.zeros([HEAD_DIM], dtype=tl.float32) + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + + acc = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for i in tl.static_range(K_W): + x_s = s_offs - (K_W - 1) + i + x_mask = (x_s >= 0) & (x_s < seq_len) + x_ptr = ( + qkvzba_ptr + + x_s[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + chan[None, :] * qkvzba_c_stride + ) + x_val = tl.load(x_ptr, mask=x_mask[:, None], other=0.0).to(tl.float32) + w_tap = tl.load( + weight_ptr + chan * weight_c_stride + i * weight_w_stride + ).to(tl.float32) + acc += w_tap[None, :] * x_val + + acc += bias[None, :] + # Mimic the unfused F.conv1d rounding: the reference path stores the conv + # output in the input dtype (bf16) before silu, so do the same here. This + # keeps the fused output bit-aligned with the reference within one ULP. + acc = acc.to(out_ptr.dtype.element_ty).to(tl.float32) + silu_out = acc * tl.sigmoid(acc) + + if APPLY_L2: + # F.silu rounds to the input dtype before l2norm reads it. Round-trip + # via bf16 to match that precision. + silu_out = silu_out.to(out_ptr.dtype.element_ty).to(tl.float32) + if SAVE_SILU: + # Persist only the QK silu output in the channel-last layout + # consumed by the QK l2norm backward. + silu_save_chan = ( + silu_save_chan_offset + + group_id * silu_save_group_stride + + head_id * HEAD_DIM + + chan_off + ) + silu_save_ptrs = ( + silu_save_ptr + + batch_id * silu_save_b_stride + + silu_save_chan[None, :] * silu_save_c_stride + + s_offs[:, None] * silu_save_s_stride + ) + tl.store( + silu_save_ptrs, + silu_out.to(silu_save_ptr.dtype.element_ty), + mask=s_mask[:, None], + ) + norm_sq = tl.sum(silu_out * silu_out, axis=1) + rstd = 1.0 / tl.sqrt(norm_sq + eps) + out = silu_out * rstd[:, None] + else: + # No l2norm follows. The final store→bf16 already does the rounding; + # an intermediate bf16 round-trip would be redundant. + out = silu_out + + out_typed = out.to(out_ptr.dtype.element_ty) + + # Write the same data to ``REPEAT`` adjacent value heads. ``REPEAT == 1`` + # is the no-repeat case (V branch is handled by a separate kernel that + # always has REPEAT == 1, but using the same code here is convenient). + for r in tl.static_range(REPEAT): + v_head = head_id * REPEAT + r + write_ptr = ( + out_ptr + + group_id * out_group_dim_stride + + batch_id * out_b_stride + + s_offs[:, None] * out_s_stride + + v_head * out_h_stride + + chan_off[None, :] + ) + tl.store(write_ptr, out_typed, mask=s_mask[:, None]) + + +@triton.jit +def _thd_seq_bounds(cu_seqlens_ptr, token_offsets, total_tokens, num_packed_seqs): + """Return lane-wise packed sequence bounds for flattened THD tokens.""" + + safe_tokens = tl.minimum(token_offsets, total_tokens - 1) + seq_start = token_offsets * 0 + seq_end = token_offsets * 0 + total_tokens + + seq_id = 0 + while seq_id < num_packed_seqs: + start = tl.load(cu_seqlens_ptr + seq_id) + end = tl.load(cu_seqlens_ptr + seq_id + 1) + in_seq = (safe_tokens >= start) & (safe_tokens < end) + seq_start = tl.where(in_seq, start, seq_start) + seq_end = tl.where(in_seq, end, seq_end) + seq_id += 1 + + return seq_start, seq_end + + +@triton.autotune( + configs=_conv_autotune_configs(), + key=["seq_len", "HEAD_DIM", "K_W", "APPLY_L2", "REPEAT", "NUM_GROUPS"], +) +@triton.jit +def _conv_silu_project_thd_kernel( + qkvzba_ptr, + weight_ptr, + bias_ptr, + out_ptr, + silu_save_ptr, + cu_seqlens_ptr, + seq_len, + num_packed_seqs, + num_in_heads, + in_channel_offset, + in_group_stride, + silu_save_chan_offset, + silu_save_group_stride, + qkvzba_s_stride, + qkvzba_b_stride, + qkvzba_c_stride, + weight_c_stride, + weight_w_stride, + bias_stride, + out_group_dim_stride, + out_b_stride, + out_s_stride, + out_h_stride, + silu_save_b_stride, + silu_save_c_stride, + silu_save_s_stride, + eps, + HEAD_DIM: tl.constexpr, + K_W: tl.constexpr, + REPEAT: tl.constexpr, + NUM_GROUPS: tl.constexpr, + HAS_BIAS: tl.constexpr, + APPLY_L2: tl.constexpr, + SAVE_SILU: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """THD depthwise conv1d + silu + optional l2norm/repeat. + + This is intentionally separate from ``_conv_silu_project_kernel`` so + packed sequence boundary metadata never enters the dense BSHD hot path. + Only the causal-conv loads use ``cu_seqlens``; the following per-token + transforms and stores are identical to the dense path. + """ + + pid_bgh = tl.program_id(0) + pid_s = tl.program_id(1) + + heads_per_batch = num_in_heads * NUM_GROUPS + batch_id = pid_bgh // heads_per_batch + local_bgh = pid_bgh - batch_id * heads_per_batch + group_id = local_bgh // num_in_heads + head_id = local_bgh - group_id * num_in_heads + + chan_off = tl.arange(0, HEAD_DIM) + group_channel_offset = in_channel_offset + group_id * in_group_stride + chan = group_channel_offset + head_id * HEAD_DIM + chan_off + + if HAS_BIAS: + bias = tl.load(bias_ptr + chan * bias_stride).to(tl.float32) + else: + bias = tl.zeros([HEAD_DIM], dtype=tl.float32) + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + seq_start, seq_end = _thd_seq_bounds(cu_seqlens_ptr, s_offs, seq_len, num_packed_seqs) + + acc = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for i in tl.static_range(K_W): + x_s = s_offs - (K_W - 1) + i + x_mask = s_mask & (x_s >= seq_start) & (x_s < seq_end) + safe_x_s = tl.minimum(tl.maximum(x_s, 0), seq_len - 1) + x_ptr = ( + qkvzba_ptr + + safe_x_s[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + chan[None, :] * qkvzba_c_stride + ) + x_val = tl.load(x_ptr, mask=x_mask[:, None], other=0.0).to(tl.float32) + w_tap = tl.load( + weight_ptr + chan * weight_c_stride + i * weight_w_stride + ).to(tl.float32) + acc += w_tap[None, :] * x_val + + acc += bias[None, :] + acc = acc.to(out_ptr.dtype.element_ty).to(tl.float32) + silu_out = acc * tl.sigmoid(acc) + + if APPLY_L2: + silu_out = silu_out.to(out_ptr.dtype.element_ty).to(tl.float32) + if SAVE_SILU: + silu_save_chan = ( + silu_save_chan_offset + + group_id * silu_save_group_stride + + head_id * HEAD_DIM + + chan_off + ) + silu_save_ptrs = ( + silu_save_ptr + + batch_id * silu_save_b_stride + + silu_save_chan[None, :] * silu_save_c_stride + + s_offs[:, None] * silu_save_s_stride + ) + tl.store( + silu_save_ptrs, + silu_out.to(silu_save_ptr.dtype.element_ty), + mask=s_mask[:, None], + ) + norm_sq = tl.sum(silu_out * silu_out, axis=1) + rstd = 1.0 / tl.sqrt(norm_sq + eps) + out = silu_out * rstd[:, None] + else: + out = silu_out + + out_typed = out.to(out_ptr.dtype.element_ty) + + for r in tl.static_range(REPEAT): + v_head = head_id * REPEAT + r + write_ptr = ( + out_ptr + + group_id * out_group_dim_stride + + batch_id * out_b_stride + + s_offs[:, None] * out_s_stride + + v_head * out_h_stride + + chan_off[None, :] + ) + tl.store(write_ptr, out_typed, mask=s_mask[:, None]) + + +@triton.jit +def _copy_z_kernel( + qkvzba_ptr, + gate_ptr, + seq_len, + num_v_heads, + z_channel_offset, + qkvzba_s_stride, + qkvzba_b_stride, + qkvzba_c_stride, + gate_b_stride, + gate_s_stride, + gate_h_stride, + HEAD_DIM: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """Copy the z slice from qkvzba into the final gate layout.""" + + pid_bh = tl.program_id(0) + pid_s = tl.program_id(1) + + batch_id = pid_bh // num_v_heads + head_id = pid_bh - batch_id * num_v_heads + + chan_off = tl.arange(0, HEAD_DIM) + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + + z_chan = z_channel_offset + head_id * HEAD_DIM + chan_off + z_src_ptr = ( + qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + z_chan[None, :] * qkvzba_c_stride + ) + z_val = tl.load(z_src_ptr, mask=s_mask[:, None]) + z_write_ptr = ( + gate_ptr + + batch_id * gate_b_stride + + s_offs[:, None] * gate_s_stride + + head_id * gate_h_stride + + chan_off[None, :] + ) + tl.store(z_write_ptr, z_val, mask=s_mask[:, None]) + + +@triton.autotune(configs=_g_beta_autotune_configs(), key=["seq_len", "num_v_heads"]) +@triton.jit +def _compute_g_and_beta_kernel( + qkvzba_ptr, + A_log_ptr, + dt_bias_ptr, + g_out_ptr, + beta_out_ptr, + seq_len, + num_v_heads, + beta_channel_offset, + alpha_channel_offset, + qkvzba_s_stride, + qkvzba_b_stride, + qkvzba_c_stride, + g_b_stride, + g_s_stride, + g_h_stride, + beta_b_stride, + beta_s_stride, + beta_h_stride, + BLOCK_S: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """Compute ``g = -exp(A_log) * softplus(alpha + dt_bias)`` and ``sigmoid(beta)``.""" + + pid_b = tl.program_id(0) + pid_s = tl.program_id(1) + pid_h = tl.program_id(2) + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + s_mask = s_offs < seq_len + h_mask = h_offs < num_v_heads + mask = s_mask[:, None] & h_mask[None, :] + + alpha_ptr = ( + qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + pid_b * qkvzba_b_stride + + (alpha_channel_offset + h_offs[None, :]) * qkvzba_c_stride + ) + beta_ptr = ( + qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + pid_b * qkvzba_b_stride + + (beta_channel_offset + h_offs[None, :]) * qkvzba_c_stride + ) + + alpha = tl.load(alpha_ptr, mask=mask, other=0.0).to(tl.float32) + beta = tl.load(beta_ptr, mask=mask, other=0.0).to(tl.float32) + + A_log = tl.load(A_log_ptr + h_offs, mask=h_mask, other=0.0).to(tl.float32) + dt_bias = tl.load(dt_bias_ptr + h_offs, mask=h_mask, other=0.0).to(tl.float32) + + pre = alpha + dt_bias[None, :] + # softplus(x) = log(1 + exp(x)); torch's softplus thresholds at x>20 but we + # rely on fp32 evaluation here, which stays well within range for typical + # GDN inputs (the unfused path computes the same expression). + softplus_val = tl.log(1.0 + tl.exp(pre)) + g = -tl.exp(A_log)[None, :] * softplus_val + beta_sig = tl.sigmoid(beta) + + g_ptr = ( + g_out_ptr + + pid_b * g_b_stride + + s_offs[:, None] * g_s_stride + + h_offs[None, :] * g_h_stride + ) + beta_out_ptr_calc = ( + beta_out_ptr + + pid_b * beta_b_stride + + s_offs[:, None] * beta_s_stride + + h_offs[None, :] * beta_h_stride + ) + tl.store(g_ptr, g.to(g_out_ptr.dtype.element_ty), mask=mask) + tl.store(beta_out_ptr_calc, beta_sig.to(beta_out_ptr.dtype.element_ty), mask=mask) + + +# --------------------------------------------------------------------------- +# Backward kernels +# --------------------------------------------------------------------------- + + +@triton.jit +def _conv_silu_l2norm_backward_kernel( + qkvzba_ptr, + weight_ptr, + d_out_ptr, + d_qkvzba_ptr, + d_w_partial_ptr, + seq_len, + num_qk_heads, + in_channel_offset, + eps, + d_out_scale, + qkvzba_s_stride, + qkvzba_b_stride, + qkvzba_c_stride, + weight_c_stride, + weight_w_stride, + d_out_b_stride, + d_out_s_stride, + d_out_h_stride, + d_wp_b_stride, + d_wp_h_stride, + d_wp_s_stride, + d_wp_c_stride, + d_wp_w_stride, + HEAD_DIM: tl.constexpr, + K_W: tl.constexpr, + REPEAT: tl.constexpr, + BLOCK_S: tl.constexpr, + USE_L2NORM: tl.constexpr, + V_HEAD_SHARED: tl.constexpr, +): + """Backward for the Q / K / V branches of ``_conv_silu_project_kernel``. + + ``USE_L2NORM`` is a constexpr branch: ``True`` for the QK branches (with + l2norm) and ``False`` for the V branch. The V case skips the l2norm + intermediates entirely — Triton DCE drops them at compile time. The + ``REPEAT=2`` workaround for the channel-collapse codegen bug still + applies in both branches. + + Forward (no bias, with l2norm, with REPEAT-way head broadcast): + acc = depthwise_conv(qkvzba_qk_slice, weight_qk_slice) + acc_bf16 = acc.to(bf16).to(fp32) # F.conv1d rounding + silu_out = acc_bf16 * sigmoid(acc_bf16) + silu_bf16 = silu_out.to(bf16).to(fp32) # round before l2norm + norm_sq = sum_c silu_bf16^2 + rstd = 1 / sqrt(norm_sq + eps) + out = silu_bf16 * rstd + # out is stored identically to REPEAT adjacent value heads. + + Backward (given d_out for each v_head): + d_qk_out = Σ_{r in REPEAT} d_v_out[head_id * REPEAT + r] + S = Σ_c d_qk_out_c * silu_bf16_c + d_silu_c = rstd * d_qk_out_c - rstd^3 * silu_bf16_c * S + d_acc_c = d_silu_c * silu'(acc_bf16) + d_w[c, i] = Σ_{b, t} d_acc[t, c] * x[t + i - (K_W - 1), c] + d_x[u, c] += Σ_i d_acc[u + (K_W - 1) - i, c] * w[c, i] + + ``d_w`` uses the per-program partial-buffer pattern (see the V backward + kernel). ``d_qkvzba`` uses bf16 atomic_add because the K_W − 1 boundary + input rows cross seq-block programs. + """ + + pid_bh = tl.program_id(0) + pid_s = tl.program_id(1) + + batch_id = pid_bh // num_qk_heads + head_id = pid_bh - batch_id * num_qk_heads + + chan_off = tl.arange(0, HEAD_DIM) + chan = in_channel_offset + head_id * HEAD_DIM + chan_off + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + + # ----- Forward recompute (conv + silu + l2norm) ----- + acc = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for i in tl.static_range(K_W): + x_s = s_offs - (K_W - 1) + i + x_mask = (x_s >= 0) & (x_s < seq_len) + x_ptr = ( + qkvzba_ptr + + x_s[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + chan[None, :] * qkvzba_c_stride + ) + x_val = tl.load(x_ptr, mask=x_mask[:, None], other=0.0).to(tl.float32) + w_tap = tl.load( + weight_ptr + chan * weight_c_stride + i * weight_w_stride + ).to(tl.float32) + acc += w_tap[None, :] * x_val + acc = acc.to(d_qkvzba_ptr.dtype.element_ty).to(tl.float32) + + # ----- Sum d_out across REPEAT v_heads ----- + # For QK: v_head = head_id*REPEAT + r, summing REPEAT distinct heads. + # For V (V_HEAD_SHARED=True): both r iterations load the SAME v_head + # (head_id), so d_qk_out = REPEAT * d_value[head_id]; the host passes + # d_out_scale = 1/REPEAT to recover d_value[head_id]. The duplicate + # load goes through L2, so the kernel-side cost is roughly one load; + # the trick was needed to avoid the REPEAT=1 codegen bug without + # having to allocate a doubled d_value tensor on the host. + d_qk_out = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for r in tl.static_range(REPEAT): + if V_HEAD_SHARED: + v_head = head_id + else: + v_head = head_id * REPEAT + r + d_out_ptrs = ( + d_out_ptr + + batch_id * d_out_b_stride + + s_offs[:, None] * d_out_s_stride + + v_head * d_out_h_stride + + chan_off[None, :] + ) + d_qk_out += tl.load(d_out_ptrs, mask=s_mask[:, None], other=0.0).to(tl.float32) + d_qk_out = d_qk_out * d_out_scale + + # ----- l2norm backward gated by USE_L2NORM constexpr. ----- + if USE_L2NORM: + silu_out = acc * tl.sigmoid(acc) + silu_bf16 = silu_out.to(d_qkvzba_ptr.dtype.element_ty).to(tl.float32) + norm_sq = tl.sum(silu_bf16 * silu_bf16, axis=1) + rstd = 1.0 / tl.sqrt(norm_sq + eps) + s_row = tl.sum(d_qk_out * silu_bf16, axis=1) + rstd3 = rstd * rstd * rstd + d_silu = rstd[:, None] * d_qk_out - rstd3[:, None] * silu_bf16 * s_row[:, None] + else: + d_silu = d_qk_out + + # ----- silu backward ----- + sig_acc = tl.sigmoid(acc) + silu_prime = sig_acc + acc * sig_acc * (1.0 - sig_acc) + d_acc = d_silu * silu_prime + d_acc = tl.where(s_mask[:, None], d_acc, 0.0) + + # ----- d_w via per-program partial, d_x via atomic_add ----- + partial_base = ( + d_w_partial_ptr + + batch_id * d_wp_b_stride + + head_id * d_wp_h_stride + + pid_s * d_wp_s_stride + ) + + for i in tl.static_range(K_W): + x_s = s_offs - (K_W - 1) + i + x_mask_inner = (x_s >= 0) & (x_s < seq_len) + x_ptr = ( + qkvzba_ptr + + x_s[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + chan[None, :] * qkvzba_c_stride + ) + x_val = tl.load(x_ptr, mask=x_mask_inner[:, None], other=0.0).to(tl.float32) + d_w_partial = tl.sum(d_acc * x_val, axis=0) + tl.store( + partial_base + chan_off * d_wp_c_stride + i * d_wp_w_stride, + d_w_partial, + ) + + w_tap = tl.load( + weight_ptr + chan * weight_c_stride + i * weight_w_stride + ).to(tl.float32) + contribution = d_acc * w_tap[None, :] + d_qkvzba_target = ( + d_qkvzba_ptr + + x_s[:, None] * qkvzba_s_stride + + batch_id * qkvzba_b_stride + + chan[None, :] * qkvzba_c_stride + ) + tl.atomic_add( + d_qkvzba_target, + contribution.to(d_qkvzba_ptr.dtype.element_ty), + mask=x_mask_inner[:, None], + ) + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_S": 32}, num_warps=2, num_stages=2), + triton.Config({"BLOCK_S": 32}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_S": 64}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=3), + triton.Config({"BLOCK_S": 256}, num_warps=8, num_stages=2), + ], + key=["seq_len", "HEAD_DIM", "REPEAT"], +) +@triton.jit +def _l2norm_repeat_backward_kernel( + d_qk_out_ptr, # (b, s, num_v_heads, head_dim) — gradient from downstream + silu_bf16_ptr, # (b, conv_dim, s) — silu(conv(x)) recomputed for QK channels + d_silu_bf16_ptr, # (b, conv_dim, s) — output gradient w.r.t. silu(conv(x)) + seq_len, + num_qk_heads, + channel_offset, # 0 for Q, qk_channels for K — indexes into conv_dim + eps, + d_qk_b_stride, + d_qk_s_stride, + d_qk_h_stride, + silu_b_stride, + silu_c_stride, + silu_s_stride, + d_silu_b_stride, + d_silu_c_stride, + d_silu_s_stride, + HEAD_DIM: tl.constexpr, + REPEAT: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """l2norm + REPEAT-way head broadcast backward. + + Forward (per QK head): + silu_bf16 ∈ R^{HEAD_DIM} # silu(conv(x)) rounded to bf16 + norm_sq = Σ_c silu_bf16_c^2 + rstd = 1 / sqrt(norm_sq + eps) + out = silu_bf16 * rstd + # out is broadcast identically to REPEAT adjacent v_heads. + + Backward (given d_qk_out for each v_head): + d_normed = Σ_{r in REPEAT} d_qk_out[head_id * REPEAT + r] + S = Σ_c d_normed_c * silu_bf16_c + d_silu_c = rstd * d_normed_c - rstd^3 * silu_bf16_c * S + + The output ``d_silu_bf16`` is the gradient w.r.t. ``silu(conv(x))`` — + exactly what ``causal_conv1d_bwd_function`` consumes as its ``dout`` + argument when ``activation="silu"`` is in effect on the forward. + """ + + pid_bh = tl.program_id(0) + pid_s = tl.program_id(1) + + batch_id = pid_bh // num_qk_heads + head_id = pid_bh - batch_id * num_qk_heads + + chan_off = tl.arange(0, HEAD_DIM) + chan = channel_offset + head_id * HEAD_DIM + chan_off + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + + # ----- Sum d_qk_out across REPEAT v_heads ----- + d_normed = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for r in tl.static_range(REPEAT): + v_head = head_id * REPEAT + r + d_out_ptrs = ( + d_qk_out_ptr + + batch_id * d_qk_b_stride + + s_offs[:, None] * d_qk_s_stride + + v_head * d_qk_h_stride + + chan_off[None, :] + ) + d_normed += tl.load(d_out_ptrs, mask=s_mask[:, None], other=0.0).to(tl.float32) + + # ----- Load silu_bf16 ----- + silu_ptrs = ( + silu_bf16_ptr + + batch_id * silu_b_stride + + chan[None, :] * silu_c_stride + + s_offs[:, None] * silu_s_stride + ) + silu_bf16 = tl.load(silu_ptrs, mask=s_mask[:, None], other=0.0).to(tl.float32) + + # ----- l2norm backward ----- + norm_sq = tl.sum(silu_bf16 * silu_bf16, axis=1) + rstd = 1.0 / tl.sqrt(norm_sq + eps) + s_row = tl.sum(d_normed * silu_bf16, axis=1) + rstd3 = rstd * rstd * rstd + d_silu = rstd[:, None] * d_normed - rstd3[:, None] * silu_bf16 * s_row[:, None] + + # ----- Store d_silu (same (b, conv_dim, s) layout as silu_bf16_ptr) ----- + d_silu_ptrs = ( + d_silu_bf16_ptr + + batch_id * d_silu_b_stride + + chan[None, :] * d_silu_c_stride + + s_offs[:, None] * d_silu_s_stride + ) + tl.store(d_silu_ptrs, d_silu.to(d_silu_bf16_ptr.dtype.element_ty), mask=s_mask[:, None]) + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_S": 32}, num_warps=2, num_stages=2), + triton.Config({"BLOCK_S": 32}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 64}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_S": 64}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_S": 128}, num_warps=8, num_stages=3), + triton.Config({"BLOCK_S": 256}, num_warps=8, num_stages=2), + ], + key=["seq_len", "HEAD_DIM", "REPEAT"], +) +@triton.jit +def _qk_l2norm_repeat_backward_kernel( + dq_ptr, + dk_ptr, + silu_bf16_ptr, + d_silu_bf16_ptr, + seq_len, + num_qk_heads, + qk_channels, + eps, + dq_b_stride, + dq_s_stride, + dq_h_stride, + dk_b_stride, + dk_s_stride, + dk_h_stride, + silu_b_stride, + silu_c_stride, + silu_s_stride, + d_silu_b_stride, + d_silu_c_stride, + d_silu_s_stride, + HEAD_DIM: tl.constexpr, + REPEAT: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """Merged Q/K l2norm + REPEAT-way head broadcast backward.""" + + pid_bgh = tl.program_id(0) + pid_s = tl.program_id(1) + + heads_per_batch = num_qk_heads * 2 + batch_id = pid_bgh // heads_per_batch + local_bgh = pid_bgh - batch_id * heads_per_batch + group_id = local_bgh // num_qk_heads + head_id = local_bgh - group_id * num_qk_heads + is_query = group_id == 0 + is_key = group_id == 1 + + chan_off = tl.arange(0, HEAD_DIM) + chan = group_id * qk_channels + head_id * HEAD_DIM + chan_off + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + + d_normed = tl.zeros([BLOCK_S, HEAD_DIM], dtype=tl.float32) + for r in tl.static_range(REPEAT): + v_head = head_id * REPEAT + r + dq_ptrs = ( + dq_ptr + + batch_id * dq_b_stride + + s_offs[:, None] * dq_s_stride + + v_head * dq_h_stride + + chan_off[None, :] + ) + dk_ptrs = ( + dk_ptr + + batch_id * dk_b_stride + + s_offs[:, None] * dk_s_stride + + v_head * dk_h_stride + + chan_off[None, :] + ) + d_normed += tl.load( + dq_ptrs, mask=s_mask[:, None] & is_query, other=0.0 + ).to(tl.float32) + d_normed += tl.load( + dk_ptrs, mask=s_mask[:, None] & is_key, other=0.0 + ).to(tl.float32) + + silu_ptrs = ( + silu_bf16_ptr + + batch_id * silu_b_stride + + chan[None, :] * silu_c_stride + + s_offs[:, None] * silu_s_stride + ) + silu_bf16 = tl.load(silu_ptrs, mask=s_mask[:, None], other=0.0).to(tl.float32) + + norm_sq = tl.sum(silu_bf16 * silu_bf16, axis=1) + rstd = 1.0 / tl.sqrt(norm_sq + eps) + s_row = tl.sum(d_normed * silu_bf16, axis=1) + rstd3 = rstd * rstd * rstd + d_silu = rstd[:, None] * d_normed - rstd3[:, None] * silu_bf16 * s_row[:, None] + + d_silu_ptrs = ( + d_silu_bf16_ptr + + batch_id * d_silu_b_stride + + chan[None, :] * d_silu_c_stride + + s_offs[:, None] * d_silu_s_stride + ) + tl.store(d_silu_ptrs, d_silu.to(d_silu_bf16_ptr.dtype.element_ty), mask=s_mask[:, None]) + + +@triton.jit +def _v_layout_to_conv_kernel( + dv_ptr, # (b, s, num_v_heads, value_head_dim) + d_silu_conv_ptr, # (b, conv_dim, s) — write into V channel slice + seq_len, + num_v_heads, + v_channel_offset, # = 2 * qk_channels + dv_b_stride, + dv_s_stride, + dv_h_stride, + d_silu_b_stride, + d_silu_c_stride, + d_silu_s_stride, + HEAD_DIM: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """Write V-branch gradients into the conv-backward layout. + + ``dv`` is the gradient of ``value`` (forward layout + ``(b, s, num_v_heads, value_head_dim)``). The conv backward needs + ``d_silu_conv`` in layout ``(b, conv_dim, s)`` for the V channel + slice. + """ + + pid_bh = tl.program_id(0) + pid_s = tl.program_id(1) + + batch_id = pid_bh // num_v_heads + head_id = pid_bh - batch_id * num_v_heads + + chan_off = tl.arange(0, HEAD_DIM) + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + + # Read dv at (batch, s, head, chan). + dv_ptrs = ( + dv_ptr + + batch_id * dv_b_stride + + s_offs[:, None] * dv_s_stride + + head_id * dv_h_stride + + chan_off[None, :] + ) + dv_val = tl.load(dv_ptrs, mask=s_mask[:, None], other=0.0) + + # Write to d_silu_conv at (batch, v_channel_offset + head*HEAD_DIM + chan, s). + d_silu_chan = v_channel_offset + head_id * HEAD_DIM + chan_off + d_silu_ptrs = ( + d_silu_conv_ptr + + batch_id * d_silu_b_stride + + d_silu_chan[None, :] * d_silu_c_stride + + s_offs[:, None] * d_silu_s_stride + ) + tl.store(d_silu_ptrs, dv_val, mask=s_mask[:, None]) + + +@triton.jit +def _z_layout_to_qkvzba_kernel( + dgate_ptr, # (b, s, num_v_heads, value_head_dim) + d_qkvzba_ptr, # (s, b, total_channels) — write into z channel slice + seq_len, + num_v_heads, + z_channel_offset, # = 2 * qk_channels + v_channels + dgate_b_stride, + dgate_s_stride, + dgate_h_stride, + d_qkvzba_s_stride, + d_qkvzba_b_stride, + d_qkvzba_c_stride, + HEAD_DIM: tl.constexpr, + BLOCK_S: tl.constexpr, +): + """Write gate gradients into the z slice of ``d_qkvzba``. + + ``dgate`` is the autograd-supplied gradient of ``gate`` (= the z + slice of qkvzba in forward) with layout + ``(b, s, num_v_heads, value_head_dim)``. We need to write it into + ``d_qkvzba``'s z slice — layout ``(s, b, total_channels)`` with + channels in ``[z_channel_offset, z_channel_offset + v_channels)``. + """ + + pid_bh = tl.program_id(0) + pid_s = tl.program_id(1) + + batch_id = pid_bh // num_v_heads + head_id = pid_bh - batch_id * num_v_heads + + chan_off = tl.arange(0, HEAD_DIM) + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < seq_len + + # Read dgate at (batch, s, head, chan). + dgate_ptrs = ( + dgate_ptr + + batch_id * dgate_b_stride + + s_offs[:, None] * dgate_s_stride + + head_id * dgate_h_stride + + chan_off[None, :] + ) + dgate_val = tl.load(dgate_ptrs, mask=s_mask[:, None], other=0.0) + + # Write to d_qkvzba at (s, batch, z_channel_offset + head*HEAD_DIM + chan). + d_qkvzba_chan = z_channel_offset + head_id * HEAD_DIM + chan_off + d_qkvzba_ptrs = ( + d_qkvzba_ptr + + s_offs[:, None] * d_qkvzba_s_stride + + batch_id * d_qkvzba_b_stride + + d_qkvzba_chan[None, :] * d_qkvzba_c_stride + ) + tl.store(d_qkvzba_ptrs, dgate_val, mask=s_mask[:, None]) + + +@triton.autotune( + configs=_g_beta_autotune_configs(), + key=["seq_len", "num_v_heads"], + # Each autotune trial atomic-adds partial sums into these accumulators. + # Without reset_to_zero the trials would stack on top of one another and + # produce values that are ``num_trials`` × the correct result. + reset_to_zero=["d_A_log_ptr", "d_dt_bias_ptr"], +) +@triton.jit +def _g_beta_backward_kernel( + qkvzba_ptr, + A_log_ptr, + dt_bias_ptr, + d_g_ptr, + d_beta_out_ptr, + d_qkvzba_ptr, + d_A_log_ptr, + d_dt_bias_ptr, + seq_len, + num_v_heads, + beta_channel_offset, + alpha_channel_offset, + qkvzba_s_stride, + qkvzba_b_stride, + qkvzba_c_stride, + d_g_b_stride, + d_g_s_stride, + d_g_h_stride, + d_beta_b_stride, + d_beta_s_stride, + d_beta_h_stride, + BLOCK_S: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """Backward for ``_compute_g_and_beta_kernel``. + + Forward: + pre = alpha + dt_bias # fp32 + softplus_pre = log(1 + exp(pre)) + g = -exp(A_log) * softplus_pre + beta_sig = sigmoid(beta_raw) + + Backward (given d_g and d_beta_out): + d_alpha = d_g * (-exp(A_log) * sigmoid(pre)) + d_beta_raw = d_beta_out * beta_sig * (1 - beta_sig) + d_dt_bias[h] = Σ_{b,s} d_alpha[b,s,h] + d_A_log[h] = Σ_{b,s} d_g[b,s,h] * g[b,s,h] + + ``d_alpha`` and ``d_beta_raw`` are written into the matching channel slices + of ``d_qkvzba``. ``d_A_log`` and ``d_dt_bias`` are reduced via per-element + atomic_add to fp32 buffers; the caller casts those to the parameter dtype. + """ + + pid_b = tl.program_id(0) + pid_s = tl.program_id(1) + pid_h = tl.program_id(2) + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + s_mask = s_offs < seq_len + h_mask = h_offs < num_v_heads + mask = s_mask[:, None] & h_mask[None, :] + + # ----- Forward recompute ----- + alpha_ptr = ( + qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + pid_b * qkvzba_b_stride + + (alpha_channel_offset + h_offs[None, :]) * qkvzba_c_stride + ) + beta_ptr = ( + qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + pid_b * qkvzba_b_stride + + (beta_channel_offset + h_offs[None, :]) * qkvzba_c_stride + ) + alpha = tl.load(alpha_ptr, mask=mask, other=0.0).to(tl.float32) + beta_raw = tl.load(beta_ptr, mask=mask, other=0.0).to(tl.float32) + A_log = tl.load(A_log_ptr + h_offs, mask=h_mask, other=0.0).to(tl.float32) + dt_bias = tl.load(dt_bias_ptr + h_offs, mask=h_mask, other=0.0).to(tl.float32) + + pre = alpha + dt_bias[None, :] + sigmoid_pre = tl.sigmoid(pre) + softplus_pre = tl.log(1.0 + tl.exp(pre)) + exp_A = tl.exp(A_log)[None, :] + g = -exp_A * softplus_pre + beta_sig = tl.sigmoid(beta_raw) + + # ----- Load upstream gradients ----- + d_g_ptrs = ( + d_g_ptr + + pid_b * d_g_b_stride + + s_offs[:, None] * d_g_s_stride + + h_offs[None, :] * d_g_h_stride + ) + d_beta_out_ptrs = ( + d_beta_out_ptr + + pid_b * d_beta_b_stride + + s_offs[:, None] * d_beta_s_stride + + h_offs[None, :] * d_beta_h_stride + ) + d_g = tl.load(d_g_ptrs, mask=mask, other=0.0).to(tl.float32) + d_beta_out = tl.load(d_beta_out_ptrs, mask=mask, other=0.0).to(tl.float32) + + # ----- Per-element gradients ----- + d_alpha = d_g * (-exp_A * sigmoid_pre) + d_beta_raw = d_beta_out * beta_sig * (1.0 - beta_sig) + + # ----- (b, s) → h reductions ----- + d_g_masked = tl.where(mask, d_g, 0.0) + d_alpha_masked = tl.where(mask, d_alpha, 0.0) + d_A_log_partial = tl.sum(d_g_masked * g, axis=0) + d_dt_bias_partial = tl.sum(d_alpha_masked, axis=0) + + # ----- Store per-element grads back to d_qkvzba ----- + d_alpha_ptrs = ( + d_qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + pid_b * qkvzba_b_stride + + (alpha_channel_offset + h_offs[None, :]) * qkvzba_c_stride + ) + d_beta_ptrs = ( + d_qkvzba_ptr + + s_offs[:, None] * qkvzba_s_stride + + pid_b * qkvzba_b_stride + + (beta_channel_offset + h_offs[None, :]) * qkvzba_c_stride + ) + tl.store( + d_alpha_ptrs, d_alpha.to(d_qkvzba_ptr.dtype.element_ty), mask=mask + ) + tl.store( + d_beta_ptrs, d_beta_raw.to(d_qkvzba_ptr.dtype.element_ty), mask=mask + ) + + # ----- Atomic-add (b, s) partials into per-head accumulators ----- + tl.atomic_add(d_A_log_ptr + h_offs, d_A_log_partial, mask=h_mask) + tl.atomic_add(d_dt_bias_ptr + h_offs, d_dt_bias_partial, mask=h_mask) + + +# --------------------------------------------------------------------------- +# Python entry points +# --------------------------------------------------------------------------- + + + + +def _is_power_of_two(value: int) -> bool: + return value > 0 and (value & (value - 1)) == 0 + + +_SIDE_STREAMS: dict = {} + + +def _get_side_stream(device: torch.device, slot: int) -> "torch.cuda.Stream": + """Lazily allocate and cache CUDA streams keyed by ``(device, slot)``. + + Reusing streams across calls keeps launches free of stream-creation + overhead, which would otherwise dominate the small kernels. + """ + + key = (device.index if device.index is not None else torch.cuda.current_device(), slot) + stream = _SIDE_STREAMS.get(key) + if stream is None: + stream = torch.cuda.Stream(device=device) + _SIDE_STREAMS[key] = stream + return stream + + +def _triton_l2norm_repeat_backward( + d_qk_out: Tensor, + silu_bf16: Tensor, + d_silu_bf16: Tensor, + *, + is_query: bool, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + eps: float = 1e-6, + stream: Optional["torch.cuda.Stream"] = None, +) -> Tensor: + """l2norm + REPEAT backward. + + ``silu_bf16`` is the (b, conv_dim, s) bf16 tensor produced by re-running + causal_conv1d_fn (forward, no-grad). Output ``d_silu_bf16`` is written + in place; only the matching channel slice (Q or K) is filled in. + """ + + batch = d_qk_out.shape[0] + seq_len = d_qk_out.shape[1] + qk_channels = num_key_heads * key_head_dim + repeat = num_value_heads // num_key_heads + channel_offset = 0 if is_query else qk_channels + + device = d_qk_out.device + + grid = lambda meta: ( + batch * num_key_heads, + triton.cdiv(seq_len, meta["BLOCK_S"]), + ) + + with _launch_context(device, stream): + _l2norm_repeat_backward_kernel[grid]( + d_qk_out, + silu_bf16, + d_silu_bf16, + seq_len, + num_key_heads, + channel_offset, + eps, + d_qk_out.stride(0), + d_qk_out.stride(1), + d_qk_out.stride(2), + silu_bf16.stride(0), + silu_bf16.stride(1), + silu_bf16.stride(2), + d_silu_bf16.stride(0), + d_silu_bf16.stride(1), + d_silu_bf16.stride(2), + HEAD_DIM=key_head_dim, + REPEAT=repeat, + ) + + return d_silu_bf16 + + +def _triton_qk_l2norm_repeat_backward( + dq: Tensor, + dk: Tensor, + silu_bf16: Tensor, + d_silu_bf16: Tensor, + *, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + eps: float = 1e-6, + stream: Optional["torch.cuda.Stream"] = None, +) -> Tensor: + """Merged Q/K l2norm + REPEAT backward launch.""" + + batch = dq.shape[0] + seq_len = dq.shape[1] + qk_channels = num_key_heads * key_head_dim + repeat = num_value_heads // num_key_heads + device = dq.device + + grid = lambda meta: ( + batch * 2 * num_key_heads, + triton.cdiv(seq_len, meta["BLOCK_S"]), + ) + + with _launch_context(device, stream): + _qk_l2norm_repeat_backward_kernel[grid]( + dq, + dk, + silu_bf16, + d_silu_bf16, + seq_len, + num_key_heads, + qk_channels, + eps, + dq.stride(0), + dq.stride(1), + dq.stride(2), + dk.stride(0), + dk.stride(1), + dk.stride(2), + silu_bf16.stride(0), + silu_bf16.stride(1), + silu_bf16.stride(2), + d_silu_bf16.stride(0), + d_silu_bf16.stride(1), + d_silu_bf16.stride(2), + HEAD_DIM=key_head_dim, + REPEAT=repeat, + ) + + return d_silu_bf16 + + +def _triton_v_layout_to_conv( + dv: Tensor, + d_silu_conv: Tensor, + *, + v_channel_offset: int, + num_value_heads: int, + value_head_dim: int, + stream: Optional["torch.cuda.Stream"] = None, +) -> None: + """Write ``dv`` into ``d_silu_conv``'s V channel slice.""" + + batch, seq_len, _, _ = dv.shape + device = dv.device + + BLOCK_S = _LAYOUT_BLOCK_S + num_seq_blocks = triton.cdiv(seq_len, BLOCK_S) + grid = (batch * num_value_heads, num_seq_blocks) + + with _launch_context(device, stream): + _v_layout_to_conv_kernel[grid]( + dv, + d_silu_conv, + seq_len, + num_value_heads, + v_channel_offset, + dv.stride(0), + dv.stride(1), + dv.stride(2), + d_silu_conv.stride(0), + d_silu_conv.stride(1), + d_silu_conv.stride(2), + HEAD_DIM=value_head_dim, + BLOCK_S=BLOCK_S, + num_warps=4, + num_stages=2, + ) + + +def _triton_z_layout_to_qkvzba( + dgate: Tensor, + d_qkvzba: Tensor, + *, + z_channel_offset: int, + num_value_heads: int, + value_head_dim: int, + stream: Optional["torch.cuda.Stream"] = None, +) -> None: + """Write ``dgate`` into ``d_qkvzba``'s z channel slice.""" + + batch, seq_len, _, _ = dgate.shape + device = dgate.device + + BLOCK_S = _LAYOUT_BLOCK_S + num_seq_blocks = triton.cdiv(seq_len, BLOCK_S) + grid = (batch * num_value_heads, num_seq_blocks) + + with _launch_context(device, stream): + _z_layout_to_qkvzba_kernel[grid]( + dgate, + d_qkvzba, + seq_len, + num_value_heads, + z_channel_offset, + dgate.stride(0), + dgate.stride(1), + dgate.stride(2), + d_qkvzba.stride(0), + d_qkvzba.stride(1), + d_qkvzba.stride(2), + HEAD_DIM=value_head_dim, + BLOCK_S=BLOCK_S, + num_warps=4, + num_stages=2, + ) + + +def _triton_g_beta_backward( + qkvzba: Tensor, + A_log: Tensor, + dt_bias: Tensor, + d_g: Tensor, + d_beta_out: Tensor, + *, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, + num_key_heads: int, + d_qkvzba_out: Optional[Tensor] = None, + stream: Optional["torch.cuda.Stream"] = None, +) -> Tuple[Tensor, Tensor, Tensor]: + """Launch ``_g_beta_backward_kernel`` and return its outputs. + + Returns: + ``(d_qkvzba_out, d_A_log, d_dt_bias)``. ``d_qkvzba_out`` only has its + alpha and beta slices filled in; the caller is expected to allocate + the buffer while the other backward kernels fill the rest. + ``d_A_log`` and ``d_dt_bias`` are fp32 and need to be cast back to + the parameter dtype by the caller. + """ + + seq_len, batch, total_channels = qkvzba.shape + qk_channels = num_key_heads * key_head_dim + v_channels = num_value_heads * value_head_dim + beta_channel_offset = 2 * qk_channels + 2 * v_channels + alpha_channel_offset = beta_channel_offset + num_value_heads + + if d_qkvzba_out is None: + d_qkvzba_out = torch.zeros_like(qkvzba) + + device = qkvzba.device + + g_beta_grid = lambda meta: ( + batch, + triton.cdiv(seq_len, meta["BLOCK_S"]), + triton.cdiv(num_value_heads, meta["BLOCK_H"]), + ) + with _launch_context(device, stream): + d_param_grads = torch.empty((2, num_value_heads), dtype=torch.float32, device=device) + d_param_grads.zero_() + d_A_log = d_param_grads[0] + d_dt_bias = d_param_grads[1] + _g_beta_backward_kernel[g_beta_grid]( + qkvzba, + A_log, + dt_bias, + d_g, + d_beta_out, + d_qkvzba_out, + d_A_log, + d_dt_bias, + seq_len, + num_value_heads, + beta_channel_offset, + alpha_channel_offset, + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + d_g.stride(0), + d_g.stride(1), + d_g.stride(2), + d_beta_out.stride(0), + d_beta_out.stride(1), + d_beta_out.stride(2), + ) + return d_qkvzba_out, d_A_log, d_dt_bias + + +class _NullContext: + def __enter__(self): + return None + + def __exit__(self, exc_type, exc_val, exc_tb): + return False + + +def _launch_context( + device: torch.device, + stream: Optional["torch.cuda.Stream"], +): + """Return a CUDA launch context after wiring the optional side stream.""" + + if stream is None: + return _NullContext() + stream.wait_stream(torch.cuda.current_stream(device)) + return torch.cuda.stream(stream) + + +def _wait_for_streams( + dst_stream: "torch.cuda.Stream", + *src_streams: "torch.cuda.Stream", +) -> None: + for stream in src_streams: + dst_stream.wait_stream(stream) + + +def _resolve_packed_seq_idx( + cu_seqlens: Optional[Tensor], + seq_idx: Optional[Tensor], + total_tokens: int, +) -> Optional[Tensor]: + """Return the token-level sequence-id buffer for causal-conv backward.""" + + if cu_seqlens is None: + assert seq_idx is None, "seq_idx requires cu_seqlens for packed THD mode." + return None + + if seq_idx is None: + seq_lengths = cu_seqlens[1:] - cu_seqlens[:-1] + seq_idx = torch.repeat_interleave( + torch.arange(seq_lengths.numel(), device=cu_seqlens.device, dtype=torch.int32), + seq_lengths, + ) + seq_idx = seq_idx.unsqueeze(0) + elif seq_idx.dim() == 1: + seq_idx = seq_idx.unsqueeze(0) + + assert seq_idx.is_cuda, f"Packed seq_idx must be CUDA, got {seq_idx.device}." + assert seq_idx.dtype == torch.int32, f"Packed seq_idx must be int32, got {seq_idx.dtype}." + assert seq_idx.shape == (1, total_tokens), ( + "Packed seq_idx must have shape [1, total_tokens], " + f"got {seq_idx.shape=} and {total_tokens=}." + ) + return seq_idx.contiguous() + + +def _triton_pre_gated_delta_rule_forward( + qkvzba: Tensor, + conv1d_weight: Tensor, + A_log: Tensor, + dt_bias: Tensor, + *, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, + cu_seqlens: Optional[Tensor] = None, +) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + """Triton-backed forward for the pre-gated-delta-rule front-end. + + Returns ``(query, key, value, gate, beta, g, silu_qk_save)``. The last + element is the bf16-rounded ``silu(conv(x))`` for the QK channel range + laid out channel-last so the backward can feed it straight into + ``causal_conv1d_bwd_function`` — see module docstring. + """ + + seq_len, batch, total_channels = qkvzba.shape + is_packed_thd = cu_seqlens is not None + if is_packed_thd: + assert batch == 1, ( + "Packed THD fused_pre_gated_delta_rule expects batch dimension 1; " + f"got {batch=}." + ) + num_packed_seqs = cu_seqlens.shape[0] - 1 + else: + num_packed_seqs = 0 + + qk_channels = num_key_heads * key_head_dim + v_channels = num_value_heads * value_head_dim + repeat_factor = num_value_heads // num_key_heads + k_w = conv1d_weight.shape[-1] + assert _is_power_of_two(key_head_dim), ( + "Triton kernel currently expects key_head_dim to be a power of two; " + f"got {key_head_dim=}." + ) + assert _is_power_of_two(value_head_dim), ( + "Triton kernel currently expects value_head_dim to be a power of two; " + f"got {value_head_dim=}." + ) + + expected_channels = 2 * qk_channels + 2 * v_channels + 2 * num_value_heads + assert total_channels == expected_channels, ( + f"qkvzba last-dim mismatch: got {total_channels}, expected {expected_channels}." + ) + + out_dtype = qkvzba.dtype + device = qkvzba.device + + # Output buffers: contiguous (b, s, h, d) for q/k/v and (b, s, h) for g/beta. + # Q and K share one allocation so the fused-streamed QK kernel can select + # the logical group by pointer stride instead of branching between two + # unrelated base pointers inside Triton. + qk_out = torch.empty( + 2, batch, seq_len, num_value_heads, key_head_dim, dtype=out_dtype, device=device + ) + query = qk_out[0] + key = qk_out[1] + value = torch.empty( + batch, seq_len, num_value_heads, value_head_dim, dtype=out_dtype, device=device + ) + g = torch.empty(batch, seq_len, num_value_heads, dtype=torch.float32, device=device) + beta = torch.empty(batch, seq_len, num_value_heads, dtype=out_dtype, device=device) + + # Conv weight is (conv_dim, 1, K_W); we treat it as (conv_dim, K_W). + weight_2d = conv1d_weight.view(conv1d_weight.shape[0], k_w) + + # No conv bias support: the entry point asserts this. We still pass a + # dummy ``bias_tensor`` to the kernel so the launch signature stays + # stable; ``HAS_BIAS=False`` ensures the kernel never reads it. + bias_tensor = qkvzba + bias_stride = 0 + + # Allocate the gate (z) output buffer that the independent Z kernel will + # populate. Keeping Z separate makes the forward scopes QK / V / Z / + # G-Beta explicit. + gate = torch.empty( + batch, seq_len, num_value_heads, value_head_dim, dtype=out_dtype, device=device + ) + + # Persist the QK silu(conv(x)) intermediate in channel-last layout so the + # backward can feed it directly into the l2norm backward. + silu_qk_save = torch.empty( + (batch, seq_len, 2 * qk_channels), dtype=out_dtype, device=device + ).permute(0, 2, 1) # → (b, 2*qk_c, s) with stride(1)==1 + silu_save_b_stride = silu_qk_save.stride(0) + silu_save_c_stride = silu_qk_save.stride(1) + silu_save_s_stride = silu_qk_save.stride(2) + + # Stream setup. Each side stream handles one of the four sub-computations + # (QK conv+l2norm, V conv, Z copy, g/beta). + main_stream = torch.cuda.current_stream(device=device) + qk_stream = _get_side_stream(device, slot=_QK_STREAM_SLOT) + v_stream = _get_side_stream(device, slot=_V_STREAM_SLOT) + g_beta_stream = _get_side_stream(device, slot=_G_BETA_STREAM_SLOT) + z_stream = _get_side_stream(device, slot=_Z_STREAM_SLOT) + for stream in (qk_stream, v_stream, g_beta_stream, z_stream): + stream.wait_stream(main_stream) + + # --- QK conv + silu + l2norm + repeat --- + qk_grid = lambda meta: ( + batch * 2 * num_key_heads, + triton.cdiv(seq_len, meta["BLOCK_S"]), + ) + with torch.cuda.stream(qk_stream): + if is_packed_thd: + _conv_silu_project_thd_kernel[qk_grid]( + qkvzba, + weight_2d, + bias_tensor, + qk_out, + silu_qk_save, + cu_seqlens, + seq_len, + num_packed_seqs, + num_key_heads, + 0, # QK starts at channel 0; group 1 starts at +qk_channels. + qk_channels, + 0, # silu_save_chan_offset + qk_channels, + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + weight_2d.stride(0), + weight_2d.stride(1), + bias_stride, + qk_out.stride(0), + qk_out.stride(1), + qk_out.stride(2), + qk_out.stride(3), + silu_save_b_stride, + silu_save_c_stride, + silu_save_s_stride, + _L2NORM_EPS, + HEAD_DIM=key_head_dim, + K_W=k_w, + REPEAT=repeat_factor, + NUM_GROUPS=2, + HAS_BIAS=False, + SAVE_SILU=True, + APPLY_L2=True, + ) + else: + _conv_silu_project_kernel[qk_grid]( + qkvzba, + weight_2d, + bias_tensor, + qk_out, + silu_qk_save, + seq_len, + num_key_heads, + 0, # QK starts at channel 0; group 1 starts at +qk_channels. + qk_channels, + 0, # silu_save_chan_offset + qk_channels, + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + weight_2d.stride(0), + weight_2d.stride(1), + bias_stride, + qk_out.stride(0), + qk_out.stride(1), + qk_out.stride(2), + qk_out.stride(3), + silu_save_b_stride, + silu_save_c_stride, + silu_save_s_stride, + _L2NORM_EPS, + HEAD_DIM=key_head_dim, + K_W=k_w, + REPEAT=repeat_factor, + NUM_GROUPS=2, + HAS_BIAS=False, + SAVE_SILU=True, + APPLY_L2=True, + ) + + # --- V conv + silu (no l2norm, no repeat) --- + v_channel_offset = 2 * qk_channels + z_channel_offset = 2 * qk_channels + v_channels + v_grid = lambda meta: (batch * num_value_heads, triton.cdiv(seq_len, meta["BLOCK_S"])) + with torch.cuda.stream(v_stream): + if is_packed_thd: + _conv_silu_project_thd_kernel[v_grid]( + qkvzba, + weight_2d, + bias_tensor, + value, + qkvzba, # silu_save unused (SAVE_SILU=False) + cu_seqlens, + seq_len, + num_packed_seqs, + num_value_heads, + v_channel_offset, + 0, # in_group_stride unused for NUM_GROUPS=1 + 0, # silu_save_chan_offset unused + 0, # silu_save_group_stride unused + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + weight_2d.stride(0), + weight_2d.stride(1), + bias_stride, + 0, # out_group_dim_stride unused for NUM_GROUPS=1 + value.stride(0), + value.stride(1), + value.stride(2), + 0, # silu_save strides unused + 0, + 0, + _L2NORM_EPS, + HEAD_DIM=value_head_dim, + K_W=k_w, + REPEAT=1, + NUM_GROUPS=1, + HAS_BIAS=False, + SAVE_SILU=False, + APPLY_L2=False, + ) + else: + _conv_silu_project_kernel[v_grid]( + qkvzba, + weight_2d, + bias_tensor, + value, + qkvzba, # silu_save unused (SAVE_SILU=False) + seq_len, + num_value_heads, + v_channel_offset, + 0, # in_group_stride unused for NUM_GROUPS=1 + 0, # silu_save_chan_offset unused + 0, # silu_save_group_stride unused + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + weight_2d.stride(0), + weight_2d.stride(1), + bias_stride, + 0, # out_group_dim_stride unused for NUM_GROUPS=1 + value.stride(0), + value.stride(1), + value.stride(2), + 0, # silu_save strides unused + 0, + 0, + _L2NORM_EPS, + HEAD_DIM=value_head_dim, + K_W=k_w, + REPEAT=1, + NUM_GROUPS=1, + HAS_BIAS=False, + SAVE_SILU=False, + APPLY_L2=False, + ) + + # --- Z copy --- + BLOCK_Z_S = _LAYOUT_BLOCK_S + z_grid = (batch * num_value_heads, triton.cdiv(seq_len, BLOCK_Z_S)) + with torch.cuda.stream(z_stream): + _copy_z_kernel[z_grid]( + qkvzba, + gate, + seq_len, + num_value_heads, + z_channel_offset, + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + gate.stride(0), + gate.stride(1), + gate.stride(2), + HEAD_DIM=value_head_dim, + BLOCK_S=BLOCK_Z_S, + num_warps=4, + num_stages=2, + ) + + # --- g and beta --- + beta_channel_offset = 2 * qk_channels + 2 * v_channels + alpha_channel_offset = beta_channel_offset + num_value_heads + g_beta_grid = lambda meta: ( + batch, + triton.cdiv(seq_len, meta["BLOCK_S"]), + triton.cdiv(num_value_heads, meta["BLOCK_H"]), + ) + with torch.cuda.stream(g_beta_stream): + _compute_g_and_beta_kernel[g_beta_grid]( + qkvzba, + A_log, + dt_bias, + g, + beta, + seq_len, + num_value_heads, + beta_channel_offset, + alpha_channel_offset, + qkvzba.stride(0), + qkvzba.stride(1), + qkvzba.stride(2), + g.stride(0), + g.stride(1), + g.stride(2), + beta.stride(0), + beta.stride(1), + beta.stride(2), + ) + + # Re-join the side streams so the caller's stream observes the writes. + _wait_for_streams(main_stream, qk_stream, v_stream, z_stream, g_beta_stream) + + return query, key, value, gate, beta, g, silu_qk_save + + +def _triton_pre_gated_delta_rule_backward( + qkvzba: Tensor, + conv1d_weight: Tensor, + silu_qk_save: Tensor, + dq: Tensor, + dk: Tensor, + dv: Tensor, + dgate: Tensor, + dbeta: Tensor, + dg: Tensor, + A_log: Tensor, + dt_bias: Tensor, + *, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, + seq_idx: Optional[Tensor] = None, +) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + """Triton-backed backward for the pre-gated-delta-rule front-end. + + Mirror of :func:`_triton_pre_gated_delta_rule_forward`. Takes upstream + gradients (``dq``/``dk``/``dv``/``dgate``/``dbeta``/``dg``) plus the + saved forward intermediates and returns input/parameter gradients + ``(d_qkvzba, d_weight, d_A_log, d_dt_bias)``. + + Five Triton kernels + one C++ ``causal_conv1d_bwd_function`` call, + fanned out on five side streams so memory-bound work overlaps while + the conv backward runs on the default stream. See module docstring + for the overall design. + """ + + seq_len, batch, _ = qkvzba.shape + qk_channels = num_key_heads * key_head_dim + v_channels = num_value_heads * value_head_dim + conv_dim = 2 * qk_channels + v_channels + z_offset = 2 * qk_channels + v_channels + k_w = conv1d_weight.shape[-1] + device = qkvzba.device + + # Rebuild the conv input as a NON-contiguous (b, c, s) view of qkvzba. + # ``causal_conv1d_fn`` / ``_bwd_function`` accept inputs where either + # ``stride(1) == 1`` or ``stride(2) == 1``; the permuted view of qkvzba + # satisfies the former (channel stride is 1 in the original (s, b, c) + # layout), so we can skip a 256 MB ``.contiguous()`` copy. + qkvzba_conv = qkvzba[:, :, :conv_dim].permute(1, 2, 0) + weight_2d = conv1d_weight.view(conv1d_weight.shape[0], k_w) + + # ``silu_qk_save`` is the (b, 2*qk_channels, s) bf16 buffer the + # forward wrote ``silu(conv(x))`` into for QK. Reuse it directly as + # the silu input to the l2norm backward. + silu_conv = silu_qk_save + + # Allocate d_silu_conv channel-last (stride(1)==1) — that's what + # ``causal_conv1d_channellast_bwd_kernel`` consumes natively. + d_silu_conv = torch.empty( + (batch, seq_len, conv_dim), dtype=qkvzba.dtype, device=device + ).permute(0, 2, 1) + + # Use the same stream slots as the forward for the matching scopes. + qk_stream = _get_side_stream(device, slot=_QK_STREAM_SLOT) + v_stream = _get_side_stream(device, slot=_V_STREAM_SLOT) + g_beta_stream = _get_side_stream(device, slot=_G_BETA_STREAM_SLOT) + z_stream = _get_side_stream(device, slot=_Z_STREAM_SLOT) + + # Q + K: l2norm + REPEAT backward writes into d_silu_conv's Q/K slices. + _triton_qk_l2norm_repeat_backward( + dq, + dk, + silu_conv, + d_silu_conv, + num_key_heads=num_key_heads, + num_value_heads=num_value_heads, + key_head_dim=key_head_dim, + stream=qk_stream, + ) + + # V: no l2norm and no REPEAT in forward, so d_silu_conv's V slice is + # just dv re-laid-out from (b, s, num_v_heads, value_head_dim) to + # (b, v_channels, s). + _triton_v_layout_to_conv( + dv, + d_silu_conv, + v_channel_offset=2 * qk_channels, + num_value_heads=num_value_heads, + value_head_dim=value_head_dim, + stream=v_stream, + ) + + # g + beta backward fully stores d_qkvzba's alpha + beta slices, plus + # per-head d_A_log / d_dt_bias. Conv and z slices are filled by the + # causal-conv and z kernels, so d_qkvzba does not need a pre-zero. + d_qkvzba = torch.empty_like(qkvzba) + _, d_A_log_fp32, d_dt_bias_fp32 = _triton_g_beta_backward( + qkvzba, + A_log, + dt_bias, + dg, + dbeta, + num_value_heads=num_value_heads, + key_head_dim=key_head_dim, + value_head_dim=value_head_dim, + num_key_heads=num_key_heads, + d_qkvzba_out=d_qkvzba, + stream=g_beta_stream, + ) + + # Z slice gradient: stream dgate into d_qkvzba's z slice. + _triton_z_layout_to_qkvzba( + dgate, + d_qkvzba, + z_channel_offset=z_offset, + num_value_heads=num_value_heads, + value_head_dim=value_head_dim, + stream=z_stream, + ) + + # Join only streams that wrote into d_silu_conv before causal_conv1d_bwd_function. + # g/beta and z write disjoint outputs and can continue overlapping with conv bwd. + default_stream = torch.cuda.current_stream(device) + _wait_for_streams(default_stream, qk_stream, v_stream) + + # Pre-allocate d_x_conv as a strided view INTO d_qkvzba's conv slice. + # d_qkvzba memory layout is (s, b, total_channels) contiguous, so + # element [s, b, c] sits at offset s*b_stride + b*c_stride + c. + # Re-interpreting that storage as (b, conv_dim, s) lets + # causal_conv1d_bwd_function write d_x directly into the right cells. + seq_stride = qkvzba.stride(0) + batch_stride = qkvzba.stride(1) + d_x_conv_view = d_qkvzba.as_strided( + (batch, conv_dim, seq_len), + (batch_stride, 1, seq_stride), + ) + + # Hand-tuned C++ conv backward. Internally folds the silu' factor and + # computes both d_x and d_w in fp32; writes d_x directly into the + # view above. + if _causal_conv1d_bwd_function is None: + raise RuntimeError( + "Fused pre-gated-delta-rule backward requires the 'causal_conv1d' package. " + "Install it, or use pre_gated_delta_rule_impl='unfused'." + ) + _, d_weight_fp32, _, _ = _causal_conv1d_bwd_function( + qkvzba_conv, + weight_2d, + None, # no bias + d_silu_conv, + seq_idx, + None, # initial_states + None, # dfinal_states + d_x_conv_view, # dx pre-allocated into d_qkvzba's conv slice + False, # return_dinitial_states + True, # activation (silu) + ) + + d_weight = d_weight_fp32.view(*conv1d_weight.shape).to(conv1d_weight.dtype) + default_stream.wait_stream(g_beta_stream) + d_A_log = d_A_log_fp32.to(A_log.dtype) + d_dt_bias = d_dt_bias_fp32.to(dt_bias.dtype) + default_stream.wait_stream(z_stream) + + return d_qkvzba, d_weight, d_A_log, d_dt_bias + + +class _FusedPreGatedDeltaRuleFunction(torch.autograd.Function): + """Thin :class:`torch.autograd.Function` wrapper around the fused path. + + Stashes the forward inputs + the saved ``silu_qk_save`` intermediate + in ``ctx`` and dispatches to :func:`_triton_pre_gated_delta_rule_forward` + / :func:`_triton_pre_gated_delta_rule_backward`. The actual kernel + logic lives in those two free functions so it's easy to read and + reuse outside the autograd machinery. + """ + + @staticmethod + def forward( + ctx, + qkvzba, + conv1d_weight, + A_log, + dt_bias, + cu_seqlens, + seq_idx, + num_key_heads, + num_value_heads, + key_head_dim, + value_head_dim, + ): + ctx.num_key_heads = num_key_heads + ctx.num_value_heads = num_value_heads + ctx.key_head_dim = key_head_dim + ctx.value_head_dim = value_head_dim + query, key, value, gate, beta, g, silu_qk_save = ( + _triton_pre_gated_delta_rule_forward( + qkvzba, + conv1d_weight, + A_log, + dt_bias, + num_key_heads=num_key_heads, + num_value_heads=num_value_heads, + key_head_dim=key_head_dim, + value_head_dim=value_head_dim, + cu_seqlens=cu_seqlens, + ) + ) + ctx.has_seq_idx = seq_idx is not None + if ctx.has_seq_idx: + ctx.save_for_backward(qkvzba, conv1d_weight, A_log, dt_bias, silu_qk_save, seq_idx) + else: + ctx.save_for_backward(qkvzba, conv1d_weight, A_log, dt_bias, silu_qk_save) + return query, key, value, gate, beta, g + + @staticmethod + def backward(ctx, dq, dk, dv, dgate, dbeta, dg): + if ctx.has_seq_idx: + qkvzba, conv1d_weight, A_log, dt_bias, silu_qk_save, seq_idx = ctx.saved_tensors + else: + qkvzba, conv1d_weight, A_log, dt_bias, silu_qk_save = ctx.saved_tensors + seq_idx = None + d_qkvzba, d_weight, d_A_log, d_dt_bias = _triton_pre_gated_delta_rule_backward( + qkvzba, + conv1d_weight, + silu_qk_save, + dq, + dk, + dv, + dgate, + dbeta, + dg, + A_log, + dt_bias, + num_key_heads=ctx.num_key_heads, + num_value_heads=ctx.num_value_heads, + key_head_dim=ctx.key_head_dim, + value_head_dim=ctx.value_head_dim, + seq_idx=seq_idx, + ) + # Match forward inputs: (qkvzba, conv1d_weight, A_log, dt_bias, + # cu_seqlens, seq_idx, num_key_heads, num_value_heads, + # key_head_dim, value_head_dim). + # Non-tensor args get None. + return ( + d_qkvzba, + d_weight, + d_A_log, + d_dt_bias, + None, + None, + None, + None, + None, + None, + ) + + +def fused_streamed_pre_gated_delta_rule( + qkvzba: Tensor, + conv1d_weight: Tensor, + conv1d_bias: Optional[Tensor], + A_log: Tensor, + dt_bias: Tensor, + *, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, + use_qk_l2norm: bool = True, + cu_seqlens: Optional[Tensor] = None, + seq_idx: Optional[Tensor] = None, +) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + """Streamed fused pre-gated-delta-rule entry point. + + Args: + qkvzba: ``[seq_len, batch, in_proj_dim]`` projection output. Must be + on CUDA. + conv1d_weight: ``[conv_dim, 1, k_w]`` depthwise conv weight. + conv1d_bias: Must be ``None`` (conv bias is not supported). + A_log: ``[num_value_heads]`` raw decay parameter. + dt_bias: ``[num_value_heads]`` time-step bias. + num_key_heads / num_value_heads / key_head_dim / value_head_dim: GDN + architecture parameters. ``num_value_heads`` must be an integer + multiple of ``num_key_heads``. + use_qk_l2norm: Must be ``True``; the fused backward closes over the + l2norm path. + cu_seqlens: Optional packed THD cumulative sequence lengths. When set, + ``qkvzba`` must have ``batch == 1`` and ``cu_seqlens[-1] == seq_len``. + seq_idx: Optional precomputed token-to-sequence map with shape + ``[1, seq_len]``. Used by causal-conv backward in packed THD mode. + + Returns: + ``(query, key, value, gate, beta, g)`` matching the unfused + :meth:`GatedDeltaNet.pre_gated_delta_rule` API. + """ + + assert qkvzba.is_cuda, ( + "fused_pre_gated_delta_rule requires CUDA inputs; " + f"got qkvzba.device={qkvzba.device}." + ) + assert conv1d_bias is None, ( + "Conv bias is not supported by fused_pre_gated_delta_rule " + "(production GDN config has none)." + ) + assert use_qk_l2norm, ( + "use_qk_l2norm=False is not supported by fused_pre_gated_delta_rule " + "(the backward closes over the l2norm path)." + ) + assert num_value_heads % num_key_heads == 0, ( + f"{num_value_heads=} must be a multiple of {num_key_heads=}." + ) + if cu_seqlens is not None: + assert cu_seqlens.is_cuda, ( + "Packed fused_pre_gated_delta_rule requires CUDA cu_seqlens; " + f"got cu_seqlens.device={cu_seqlens.device}." + ) + assert cu_seqlens.dtype == torch.int32, ( + "Packed fused_pre_gated_delta_rule requires int32 cu_seqlens; " + f"got {cu_seqlens.dtype=}." + ) + assert cu_seqlens.dim() == 1, ( + "Packed fused_pre_gated_delta_rule expects 1-D cu_seqlens; " + f"got {cu_seqlens.shape=}." + ) + assert qkvzba.shape[1] == 1, ( + "Packed THD fused_pre_gated_delta_rule expects batch dimension 1; " + f"got qkvzba.shape={qkvzba.shape}." + ) + assert cu_seqlens.shape[0] >= 2, ( + "Packed fused_pre_gated_delta_rule requires at least one packed sequence; " + f"got {cu_seqlens.shape=}." + ) + assert cu_seqlens[0].item() == 0, ( + "Packed fused_pre_gated_delta_rule requires cu_seqlens[0] == 0, " + f"got {cu_seqlens[0].item()}." + ) + assert torch.all(cu_seqlens[1:] >= cu_seqlens[:-1]).item(), ( + "Packed fused_pre_gated_delta_rule requires monotonically non-decreasing " + f"cu_seqlens, got {cu_seqlens}." + ) + assert cu_seqlens[-1].item() == qkvzba.shape[0], ( + "Packed fused_pre_gated_delta_rule requires cu_seqlens[-1] to match " + f"seq_len, got {cu_seqlens[-1].item()} vs {qkvzba.shape[0]}." + ) + cu_seqlens = cu_seqlens.contiguous() + seq_idx = _resolve_packed_seq_idx(cu_seqlens, seq_idx, qkvzba.shape[0]) + else: + assert seq_idx is None, "seq_idx requires cu_seqlens for packed THD mode." + + return _FusedPreGatedDeltaRuleFunction.apply( + qkvzba, + conv1d_weight, + A_log, + dt_bias, + cu_seqlens, + seq_idx, + num_key_heads, + num_value_heads, + key_head_dim, + value_head_dim, + ) + + +fused_pre_gated_delta_rule = fused_streamed_pre_gated_delta_rule diff --git a/megatron/core/models/common/embeddings/rope_utils.py b/megatron/core/models/common/embeddings/rope_utils.py index c97f738771b..0c89cdbf084 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -16,6 +16,7 @@ from megatron.core import parallel_state logger = logging.getLogger(__name__) +_ROPE_FUSION_FALLBACK_WARNINGS: set[str] = set() try: from megatron.core.extensions.transformer_engine import fused_apply_rotary_pos_emb @@ -29,6 +30,26 @@ fused_apply_rotary_pos_emb_thd = None +try: + from megatron.core.fusions.fused_mrope import ( + can_launch_fused_mrope_thd, + fused_apply_mrope, + fused_apply_mrope_thd, + get_fused_mrope_thd_unavailable_reason, + get_fused_mrope_unavailable_reason, + is_fused_mrope_available, + mrope_freqs_to_rotary_emb, + ) +except ImportError: + can_launch_fused_mrope_thd = None + fused_apply_mrope = None + fused_apply_mrope_thd = None + get_fused_mrope_thd_unavailable_reason = None + get_fused_mrope_unavailable_reason = None + is_fused_mrope_available = None + mrope_freqs_to_rotary_emb = None + + try: from flash_attn.layers.rotary import apply_rotary_emb as apply_rotary_emb_flash except ImportError: @@ -41,10 +62,101 @@ 'apply_rotary_pos_emb_with_cos_sin', 'fused_apply_rotary_pos_emb', 'fused_apply_rotary_pos_emb_thd', + 'can_launch_fused_mrope_thd', + 'fused_apply_mrope', + 'fused_apply_mrope_thd', + 'get_fused_mrope_thd_unavailable_reason', + 'get_fused_mrope_unavailable_reason', + 'is_fused_mrope_available', + 'mrope_freqs_to_rotary_emb', 'get_pos_emb_on_this_cp_rank', ] +def _is_raw_mrope_freqs(t: Tensor, freqs: Tensor, config: TransformerConfig) -> bool: + """Return whether freqs is the raw 3-axis mRoPE tensor for fused apply.""" + if config.mrope_section is None or freqs.dim() != 4 or freqs.shape[0] != 3: + return False + if sum(config.mrope_section) != freqs.shape[-1] or freqs.shape[-1] * 2 > t.shape[-1]: + return False + if t.dim() == 4: + return freqs.shape[1] == t.shape[1] and freqs.shape[2] == t.shape[0] + if t.dim() == 3: + return freqs.shape[1] == 1 + return False + + +def _is_raw_mrope_freqs_thd( + t: Tensor, freqs: Tensor, cu_seqlens: Tensor, config: TransformerConfig, cp_size: int +) -> bool: + """Return whether freqs is raw mRoPE for THD layout, or fail on raw-like bad shapes.""" + if config.mrope_section is None or freqs.dim() != 4 or freqs.shape[0] != 3: + return False + if t.dim() != 3: + raise ValueError( + f"raw mRoPE THD expects t with shape [tokens, heads, head_dim], got {tuple(t.shape)}" + ) + if sum(config.mrope_section) != freqs.shape[-1] or freqs.shape[-1] * 2 > t.shape[-1]: + return False + + if freqs.shape[1] != 1: + raise ValueError( + "raw mRoPE THD freqs must have singleton batch dimension with shape " + f"[3, 1, total_seqlen, rotary_dim / 2], got {tuple(freqs.shape)}" + ) + if cp_size > 1 and freqs.shape[2] % cp_size != 0: + raise ValueError( + "raw mRoPE THD freqs sequence length must be divisible by context parallel size, " + f"got freqs.shape[2]={freqs.shape[2]}, cp_size={cp_size}" + ) + expected_total_seqlen = t.shape[0] * cp_size + if freqs.shape[2] != expected_total_seqlen: + raise ValueError( + "raw mRoPE THD freqs sequence length must match local tokens times cp_size, " + f"got freqs.shape[2]={freqs.shape[2]}, tokens={t.shape[0]}, cp_size={cp_size}" + ) + if cu_seqlens.dim() != 1: + raise ValueError(f"raw mRoPE THD cu_seqlens must be 1D, got {tuple(cu_seqlens.shape)}") + return True + + +def _raw_mrope_freqs_to_emb(freqs: Tensor, config: TransformerConfig) -> Tensor: + assert mrope_freqs_to_rotary_emb is not None, "mRoPE frequency conversion is unavailable." + return mrope_freqs_to_rotary_emb( + freqs, + config.mrope_section, + interleaved_mrope=config.mrope_interleaved, + rotary_interleaved=config.rotary_interleaved, + ) + + +def _warn_rope_fusion_fallback_once(key: str, message: str) -> None: + if key in _ROPE_FUSION_FALLBACK_WARNINGS: + return + _ROPE_FUSION_FALLBACK_WARNINGS.add(key) + warnings.warn(message, stacklevel=2) + + +def _fused_mrope_unavailable_warning_key(reason: str, thd: bool = False) -> str: + prefix = "triton-mrope-thd-unavailable" if thd else "triton-mrope-unavailable" + reason_lower = reason.lower() + if "triton is not available" in reason_lower: + category = "import" + elif "cuda tensors" in reason_lower or "same device" in reason_lower: + category = "device" + elif "dtype" in reason_lower or "float32" in reason_lower: + category = "dtype" + elif "stride" in reason_lower or "contiguous" in reason_lower: + category = "stride" + elif "capability" in reason_lower: + category = "capability" + elif "rotary_interleaved" in reason_lower: + category = "rotary-interleaved" + else: + category = "other" + return f"{prefix}-{category}" + + def get_pos_emb_on_this_cp_rank( pos_emb: Tensor, seq_dim: int, cp_group: torch.distributed.ProcessGroup ) -> Tensor: @@ -181,20 +293,21 @@ def _get_thd_freqs_on_this_cp_rank( compatibility. """ if cp_size > 1: - cp_seg = x.size(0) // 2 + first_cp_seg = (x.size(0) + 1) // 2 + second_cp_seg = x.size(0) // 2 full_seqlen = cp_size * x.size(0) # Apply offset to both forward and backward segments for context parallelism - # offset=0: traditional behavior, freqs[0:cp_seg] and freqs[...] - # offset>0: exact mapping, freqs[offset+0:offset+cp_seg] and freqs[offset+...] + # offset=0: traditional behavior, freqs[0:first_cp_seg] and freqs[...] + # offset>0: exact mapping, freqs[offset+0:offset+first_cp_seg] and freqs[offset+...] return torch.cat( [ - freqs[offset + cp_rank * cp_seg : offset + (cp_rank + 1) * cp_seg], + freqs[offset + cp_rank * first_cp_seg : offset + (cp_rank + 1) * first_cp_seg], freqs[ offset + full_seqlen - - (cp_rank + 1) * cp_seg : offset + - (cp_rank + 1) * second_cp_seg : offset + full_seqlen - - cp_rank * cp_seg + - cp_rank * second_cp_seg ], ] ) @@ -205,6 +318,84 @@ def _get_thd_freqs_on_this_cp_rank( return freqs[offset : offset + x.size(0)] +def _get_thd_raw_mrope_freqs_on_this_cp_rank( + cp_rank: int, cp_size: int, x: Tensor, freqs: Tensor, offset: int = 0 +) -> Tensor: + """Get raw mRoPE frequency slices for this CP rank in THD layout.""" + if cp_size > 1: + first_cp_seg = (x.size(0) + 1) // 2 + second_cp_seg = x.size(0) // 2 + full_seqlen = cp_size * x.size(0) + return torch.cat( + [ + freqs[ + :, :, offset + cp_rank * first_cp_seg : offset + (cp_rank + 1) * first_cp_seg + ], + freqs[ + :, + :, + offset + + full_seqlen + - (cp_rank + 1) * second_cp_seg : offset + + full_seqlen + - cp_rank * second_cp_seg, + ], + ], + dim=2, + ) + else: + return freqs[:, :, offset : offset + x.size(0)] + + +def _get_thd_cp_splits(cu_seqlens: Tensor, cp_size: int) -> tuple[list[int], list[int]]: + """Return global sequence offsets and per-rank sequence lengths for THD CP fallback.""" + cu_seqlens_list = cu_seqlens.tolist() + local_seqlens = [] + for seq_start, seq_end in zip(cu_seqlens_list[:-1], cu_seqlens_list[1:]): + seq_len = seq_end - seq_start + if cp_size > 1 and seq_len % cp_size != 0: + raise ValueError( + "THD sequence lengths must be divisible by context parallel size, " + f"got sequence length {seq_len}, cp_size={cp_size}" + ) + local_seqlens.append(seq_len // cp_size) + return cu_seqlens_list, local_seqlens + + +def _pack_thd_raw_mrope_freqs( + t: Tensor, + cu_seqlens: Tensor, + freqs: Tensor, + cp_group: torch.distributed.ProcessGroup, + total_seqlen: Optional[int] = None, +) -> Tensor: + """Pack raw mRoPE freqs into the same local token order as THD tensor ``t``.""" + cp_size = cp_group.size() + cp_rank = cp_group.rank() + cu_seqlens_list, seqlens = _get_thd_cp_splits(cu_seqlens, cp_size) + sequence_splits = torch.split(t, seqlens) + if total_seqlen is None: + total_seqlen = cu_seqlens_list[-1] + assert freqs.size(2) == total_seqlen, ( + f"raw mRoPE THD freqs sequence length {freqs.size(2)} must match " + f"cu_seqlens[-1] = {total_seqlen}" + ) + + freq_slices = [] + for i, x in enumerate(sequence_splits): + seq_start_offset = cu_seqlens_list[i] + freq_slices.append( + _get_thd_raw_mrope_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs, seq_start_offset) + ) + + packed_freqs = torch.cat(freq_slices, dim=2) + assert packed_freqs.shape[2] == t.shape[0], ( + f"packed raw mRoPE freqs sequence length {packed_freqs.shape[2]} " + f"does not match THD tensor length {t.shape[0]}" + ) + return packed_freqs.contiguous() + + def _apply_rotary_pos_emb_thd( t: Tensor, cu_seqlens: Tensor, @@ -240,21 +431,21 @@ def _apply_rotary_pos_emb_thd( raise ValueError("cp_group must be provided for THD format RoPE") cp_size = cp_group.size() cp_rank = cp_group.rank() - seqlens = ((cu_seqlens[1:] - cu_seqlens[:-1]) // cp_size).tolist() + cu_seqlens_list, seqlens = _get_thd_cp_splits(cu_seqlens, cp_size) # Handle two different frequency tensor formats: - # 1. If freqs.size(0) == cu_seqlens[-1]: freqs contains all positions across all sequences + # 1. If freqs.size(0) == cu_seqlens_list[-1]: freqs contains all positions across all sequences # -> Use offset-based mapping for exact positional correspondence # 2. Otherwise: freqs contains only max sequence length positions # -> Use traditional mapping without offsets (map first :seqlen part) - if freqs.dim() >= 1 and freqs.size(0) == cu_seqlens[-1]: + if freqs.dim() >= 1 and freqs.size(0) == cu_seqlens_list[-1]: # CASE 1: Exact mapping with offsets # Build packed freqs in one pass, then apply once to the whole packed tensor sequence_splits = torch.split(t, seqlens) freq_slices = [] for i, x in enumerate(sequence_splits): # cu_seqlens[i] is the starting offset of this sequence in the original batch - seq_start_offset = cu_seqlens[i].item() + seq_start_offset = cu_seqlens_list[i] freq_slices.append( _get_thd_freqs_on_this_cp_rank(cp_rank, cp_size, x, freqs, seq_start_offset) ) @@ -311,40 +502,254 @@ def apply_rotary_pos_emb( if cp_group is None: cp_group = parallel_state.get_context_parallel_group() + is_raw_mrope_freqs = ( + _is_raw_mrope_freqs(t, freqs, config) + if cu_seqlens is None + else _is_raw_mrope_freqs_thd(t, freqs, cu_seqlens, config, cp_group.size()) + ) + if config.apply_rope_fusion: if cu_seqlens is None: + force_unfused_mrope = False + if is_raw_mrope_freqs: + unavailable_reason = None + can_try_fused_mrope = ( + fused_apply_mrope is not None + and get_fused_mrope_unavailable_reason is not None + and not mla_rotary_interleaved + and not inverse + and mscale == 1.0 + ) + if can_try_fused_mrope: + unavailable_reason = get_fused_mrope_unavailable_reason( + t, freqs, config.rotary_interleaved + ) + use_fused_mrope = can_try_fused_mrope and unavailable_reason is None + if use_fused_mrope: + return fused_apply_mrope( + t, + freqs, + config.mrope_section, + interleaved_mrope=config.mrope_interleaved, + rotary_interleaved=config.rotary_interleaved, + ) + + if unavailable_reason is not None: + _warn_rope_fusion_fallback_once( + _fused_mrope_unavailable_warning_key(unavailable_reason), + f"Triton fused mRoPE is unavailable: {unavailable_reason}. " + "Using unfused implementation.", + ) + force_unfused_mrope = True + unavailable_is_rotary_interleaved = ( + unavailable_reason is not None + and "rotary_interleaved" in unavailable_reason.lower() + ) + if mscale != 1.0: + _warn_rope_fusion_fallback_once( + "triton-mrope-mscale", + f"mscale={mscale} is not supported by Triton fused mRoPE. " + "Using unfused implementation.", + ) + force_unfused_mrope = True + if mla_rotary_interleaved: + _warn_rope_fusion_fallback_once( + "triton-mrope-mla-rotary-interleaved", + "Triton fused mRoPE does not support MLA-style interleaving in RoPE. " + "Using unfused implementation.", + ) + force_unfused_mrope = True + if inverse: + _warn_rope_fusion_fallback_once( + "triton-mrope-inverse", + "inverse RoPE is not supported by Triton fused mRoPE. " + "Using unfused implementation.", + ) + force_unfused_mrope = True + if config.rotary_interleaved and not unavailable_is_rotary_interleaved: + _warn_rope_fusion_fallback_once( + "triton-mrope-rotary-interleaved", + "Triton fused mRoPE currently supports rotary_interleaved=False. " + "Using unfused implementation.", + ) + force_unfused_mrope = True + freqs = _raw_mrope_freqs_to_emb(freqs, config) + is_raw_mrope_freqs = False + if force_unfused_mrope: + return _apply_rotary_pos_emb_bshd( + t, + freqs, + rotary_interleaved=config.rotary_interleaved, + mla_rotary_interleaved=mla_rotary_interleaved, + mscale=mscale, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, + ) + # NOTE: TE backends do not support mRoPE in bshd format when bs > 1. use_unfused = False if config.mrope_section is not None and freqs.shape[1] > 1: # TODO: Add a check in TransformerConfig and remove this unfused implementation. - warnings.warn( - "apply_rope_fusion does not support mRoPE in bshd format when bs > 1. " - "Please set apply_rope_fusion to false. This will become an error in v0.16." + _warn_rope_fusion_fallback_once( + "te-mrope-bshd-batch", + "Transformer Engine fused RoPE does not support mRoPE in bshd format when " + "bs > 1 without raw mRoPE freqs. Using unfused implementation.", ) use_unfused = True if mscale != 1.0: - warnings.warn( + _warn_rope_fusion_fallback_once( + "te-rope-mscale", f"mscale={mscale} is not supported by TE's fused RoPE. " - "Using unfused implementation." + "Using unfused implementation.", ) use_unfused = True if mla_rotary_interleaved: - warnings.warn( - "apply_rope_fusion does not support MLA-style interleaving in RoPE." - "Using unfused implementation." + _warn_rope_fusion_fallback_once( + "te-rope-mla-rotary-interleaved", + "apply_rope_fusion does not support MLA-style interleaving in RoPE. " + "Using unfused implementation.", ) use_unfused = True if inverse: - warnings.warn( + _warn_rope_fusion_fallback_once( + "te-rope-inverse", "inverse RoPE is not supported by TE's fused RoPE. " - "Using unfused implementation." + "Using unfused implementation.", + ) + use_unfused = True + if fused_apply_rotary_pos_emb is None: + _warn_rope_fusion_fallback_once( + "te-rope-unavailable", + "Transformer Engine fused RoPE is unavailable. Using unfused implementation.", ) use_unfused = True if not use_unfused: - assert fused_apply_rotary_pos_emb is not None, "apply_rope_fusion is not available." return fused_apply_rotary_pos_emb(t, freqs, interleaved=config.rotary_interleaved) else: - assert fused_apply_rotary_pos_emb_thd is not None, "apply_rope_fusion is not available." + if is_raw_mrope_freqs: + use_fused_mrope_thd = ( + fused_apply_mrope_thd is not None + and can_launch_fused_mrope_thd is not None + and get_fused_mrope_thd_unavailable_reason is not None + and mscale == 1.0 + and not mla_rotary_interleaved + and not inverse + and not config.rotary_interleaved + ) + if use_fused_mrope_thd: + unavailable_reason = get_fused_mrope_thd_unavailable_reason( + t, + cu_seqlens, + freqs, + rotary_interleaved=config.rotary_interleaved, + cp_size=cp_group.size(), + cp_rank=cp_group.rank(), + ) + if unavailable_reason is None: + return fused_apply_mrope_thd( + t, + cu_seqlens, + freqs, + config.mrope_section, + interleaved_mrope=config.mrope_interleaved, + rotary_interleaved=config.rotary_interleaved, + cp_size=cp_group.size(), + cp_rank=cp_group.rank(), + ) + _warn_rope_fusion_fallback_once( + _fused_mrope_unavailable_warning_key(unavailable_reason, thd=True), + f"Triton fused mRoPE for THD layout is unavailable: " + f"{unavailable_reason}. Using unfused implementation.", + ) + else: + has_unsupported_option = False + if mscale != 1.0: + _warn_rope_fusion_fallback_once( + "triton-mrope-thd-mscale", + f"mscale={mscale} is not supported by Triton fused mRoPE for THD " + "layout. Using unfused implementation.", + ) + has_unsupported_option = True + if mla_rotary_interleaved: + _warn_rope_fusion_fallback_once( + "triton-mrope-thd-mla-rotary-interleaved", + "Triton fused mRoPE for THD layout does not support MLA-style " + "interleaving in RoPE. Using unfused implementation.", + ) + has_unsupported_option = True + if inverse: + _warn_rope_fusion_fallback_once( + "triton-mrope-thd-inverse", + "inverse RoPE is not supported by Triton fused mRoPE for THD layout. " + "Using unfused implementation.", + ) + has_unsupported_option = True + if config.rotary_interleaved: + _warn_rope_fusion_fallback_once( + "triton-mrope-thd-rotary-interleaved", + "Triton fused mRoPE for THD layout currently supports " + "rotary_interleaved=False. Using unfused implementation.", + ) + has_unsupported_option = True + if not has_unsupported_option: + _warn_rope_fusion_fallback_once( + "triton-mrope-thd-unavailable", + "Triton fused mRoPE for THD layout is unavailable. " + "Using unfused implementation.", + ) + freqs = _raw_mrope_freqs_to_emb(freqs, config) + return _apply_rotary_pos_emb_thd( + t, + cu_seqlens, + freqs, + rotary_interleaved=config.rotary_interleaved, + mla_rotary_interleaved=mla_rotary_interleaved, + mscale=mscale, + cp_group=cp_group, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, + ) + use_unfused_thd = False + if mscale != 1.0: + _warn_rope_fusion_fallback_once( + "te-rope-thd-mscale", + f"mscale={mscale} is not supported by TE's fused RoPE for THD layout. " + "Using unfused implementation.", + ) + use_unfused_thd = True + if mla_rotary_interleaved: + _warn_rope_fusion_fallback_once( + "te-rope-thd-mla-rotary-interleaved", + "TE fused RoPE for THD layout does not support MLA-style interleaving " + "in RoPE. Using unfused implementation.", + ) + use_unfused_thd = True + if inverse: + _warn_rope_fusion_fallback_once( + "te-rope-thd-inverse", + "inverse RoPE is not supported by TE's fused RoPE for THD layout. " + "Using unfused implementation.", + ) + use_unfused_thd = True + if fused_apply_rotary_pos_emb_thd is None: + _warn_rope_fusion_fallback_once( + "te-rope-thd-unavailable", + "Transformer Engine fused RoPE for THD layout is unavailable. " + "Using unfused implementation.", + ) + use_unfused_thd = True + if use_unfused_thd: + return _apply_rotary_pos_emb_thd( + t, + cu_seqlens, + freqs, + rotary_interleaved=config.rotary_interleaved, + mla_rotary_interleaved=mla_rotary_interleaved, + mscale=mscale, + cp_group=cp_group, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, + ) return fused_apply_rotary_pos_emb_thd( t, cu_seqlens, @@ -354,6 +759,9 @@ def apply_rotary_pos_emb( interleaved=config.rotary_interleaved, ) # use unfused implementation + if is_raw_mrope_freqs: + freqs = _raw_mrope_freqs_to_emb(freqs, config) + if cu_seqlens is None: return _apply_rotary_pos_emb_bshd( t, diff --git a/megatron/core/models/common/embeddings/rotary_pos_embedding.py b/megatron/core/models/common/embeddings/rotary_pos_embedding.py index 804bdb7c537..e056ffeb9be 100644 --- a/megatron/core/models/common/embeddings/rotary_pos_embedding.py +++ b/megatron/core/models/common/embeddings/rotary_pos_embedding.py @@ -341,6 +341,8 @@ def forward( position_ids: torch.Tensor, mrope_section: List[int], cp_group: Optional[torch.distributed.ProcessGroup] = None, + return_raw_freqs: bool = False, + packed_seq: bool = False, ) -> Tensor: """Forward pass of multimodal RoPE embedding. @@ -350,9 +352,14 @@ def forward( height and width in rope calculation. cp_group (torch.distributed.ProcessGroup, optional): Context parallel group. Defaults to None. + return_raw_freqs (bool, optional): If True, return the raw per-axis frequencies with + shape [3, batchsize, seqlens, dim / 2] for fused mRoPE application. + packed_seq (bool, optional): Whether the sequence uses THD packing. Packed sequences + keep full position frequencies because THD RoPE applies CP partitioning later. Returns: - Tensor: Embeddings after applying RoPE. + Tensor: Embeddings after applying RoPE, or raw per-axis frequencies when + return_raw_freqs is True. """ seq = position_ids.to(device=self.inv_freq.device, dtype=self.inv_freq.dtype) @@ -366,6 +373,13 @@ def forward( # shape (3, bs, seq_length, dim) freqs = (inv_freq_expanded @ seq_expanded).transpose(2, 3) + if cp_group is None: + cp_group = self.cp_group + if return_raw_freqs: + if cp_group is not None and cp_group.size() > 1 and not packed_seq: + freqs = get_pos_emb_on_this_cp_rank(freqs, 2, cp_group) + return freqs.contiguous() + # first part even vector components, second part odd vector components, # 2 * dim in dimension size if self.interleaved_mrope: @@ -376,9 +390,9 @@ def forward( emb = torch.cat((freqs, freqs), dim=-1) # shape (bs, seq_length, 2 * dim) else: bs = freqs.shape[0] - emb = torch.stack((freqs.view(bs, -1, 1), freqs.view(bs, -1, 1)), dim=-1).view( - bs, freqs.shape[1], -1 - ) + emb = torch.stack( + (freqs.reshape(bs, -1, 1), freqs.reshape(bs, -1, 1)), dim=-1 + ).view(bs, freqs.shape[1], -1) else: # Original section-based layout (Qwen2-VL style). if not self.rotary_interleaved: @@ -386,8 +400,8 @@ def forward( else: bs = freqs.shape[1] emb = torch.stack( - (freqs.view(3, bs, -1, 1), freqs.view(3, bs, -1, 1)), dim=-1 - ).view(3, bs, freqs.shape[0], -1) + (freqs.reshape(3, bs, -1, 1), freqs.reshape(3, bs, -1, 1)), dim=-1 + ).view(3, bs, freqs.shape[2], -1) # generate freqs with mrope_section: cycle T/H/W per section chunk mrope_section_doubled = list(mrope_section) * 2 emb = torch.cat( @@ -396,9 +410,7 @@ def forward( # shape (seq_length, bs, 1, 2 * dim) emb = emb[..., None, :].transpose(0, 1).contiguous() - if cp_group is None: - cp_group = self.cp_group - if cp_group is not None and cp_group.size() > 1: + if cp_group is not None and cp_group.size() > 1 and not packed_seq: # slice rotary_pos_emb along sequence dimension and select the parition of the current # CP rank emb = get_pos_emb_on_this_cp_rank(emb, 0, cp_group) diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index b41e1e7afff..67c64d02a7c 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/megatron/core/models/gpt/fine_grained_callables.py @@ -634,6 +634,14 @@ def submodule_moe_forward(node: ScheduleNode, dispatched_tokens: torch.Tensor): # as a gradient hook of expert_output layer.pre_mlp_norm_checkpoint.discard_output_and_register_recompute(expert_output) + # Trigger the shared-expert recompute from expert_output too (output freed in + # submodule_combine_forward). Registering on the same tensor AFTER the pre_mlp_norm + # recompute orders it after pre_mlp_layernorm_output is restored and before the attn node's + # shared-expert backward. + shared_experts_checkpoint = getattr(layer.mlp, "shared_experts_checkpoint", None) + if shared_experts_checkpoint is not None: + shared_experts_checkpoint.register_recompute_hook(expert_output) + return expert_output def submodule_combine_forward(node: ScheduleNode, output: torch.Tensor): @@ -683,6 +691,13 @@ def submodule_combine_forward(node: ScheduleNode, output: torch.Tensor): if not node.is_mtp and final_layernorm and node.is_last_layer: output = final_layernorm(output) output = make_viewless_tensor(inp=output, requires_grad=True, keep_graph=True) + + # postprocess() has consumed the shared-expert output; free its storage now (the recompute + # hook was registered on expert_output in submodule_moe_forward). + shared_experts_checkpoint = getattr(layer.mlp, "shared_experts_checkpoint", None) + if shared_experts_checkpoint is not None: + shared_experts_checkpoint.discard_output() + layer.mlp.shared_experts_checkpoint = None return output @copy_signature(layer._forward_mlp, handle_first_dst_param='preserve') diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index 7bce9d96d2c..7adcc03ecd1 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -147,6 +147,7 @@ def __init__( self.mtp_process = mtp_block_spec is not None and mtp_on_this_rank( self.config, ignore_virtual=False, vp_stage=vp_stage ) + self._fused_mrope_available = False self.fuse_linear_cross_entropy = ( self.config.cross_entropy_loss_fusion @@ -209,6 +210,13 @@ def __init__( assert ( self.mrope_section is not None ), "mrope require mrope_section setting, but we got None from TransformerConfig" + if self.config.apply_rope_fusion and not self.config.rotary_interleaved: + try: + from megatron.core.fusions.fused_mrope import is_fused_mrope_available + + self._fused_mrope_available = is_fused_mrope_available() + except ImportError: + self._fused_mrope_available = False # Cache for RoPE tensors which do not change between iterations. self.rotary_pos_emb_cache = {} @@ -402,10 +410,25 @@ def _preprocess( ) elif self.position_embedding_type == 'mrope' and not self.config.multi_latent_attention: if self.training or not self.config.flash_decode: + packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + use_fused_mrope = False + use_raw_mrope_freqs = ( + self.config.apply_rope_fusion and not self.config.rotary_interleaved + ) + if self.config.fused_single_qkv_rope: + use_raw_mrope_freqs = False + # Inference indexes rotary_pos_emb as seq-major materialized embeddings. + # Raw mRoPE freqs are axis-major and are only safe for the normal decoder path. + if in_inference_mode: + use_raw_mrope_freqs = False + if use_raw_mrope_freqs: + use_fused_mrope = self._fused_mrope_available rotary_pos_emb = self.rotary_pos_emb( position_ids, self.mrope_section, cp_group=packed_seq_params.cp_group if packed_seq_params is not None else None, + return_raw_freqs=use_fused_mrope, + packed_seq=packed_seq, ) else: # Flash decoding uses precomputed cos and sin for RoPE diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py index 863b5d55d9d..f2c9b0b5181 100644 --- a/megatron/core/parallel_state.py +++ b/megatron/core/parallel_state.py @@ -154,6 +154,12 @@ def get_nccl_options(pg_name, nccl_comm_cfgs): nccl_comm_cfgs (dict): nccl communicator configurations When an option (e.g., max_ctas) is not found in the config, use the NCCL default setting. """ + # The fake distributed backend (--fake-process-group) cannot accept + # ProcessGroupNCCL.Options; PyTorch's FakeProcessGroup._create_internal + # rejects them with a TypeError. Return None so callers create the + # fake sub-groups without NCCL-specific options. + if torch.distributed.is_initialized() and torch.distributed.get_backend() == "fake": + return None if pg_name in nccl_comm_cfgs: # When fields in nccl_options.config are not specified, NCCL applies default settings. # The default values for Hopper GPUs are as follows: diff --git a/megatron/core/pipeline_parallel/combined_1f1b.py b/megatron/core/pipeline_parallel/combined_1f1b.py index b1ebbb876ff..ecda8e55a14 100644 --- a/megatron/core/pipeline_parallel/combined_1f1b.py +++ b/megatron/core/pipeline_parallel/combined_1f1b.py @@ -372,9 +372,21 @@ def forward_backward_step(): ) from megatron.core.models.gpt.gpt_model import GPTModel - assert isinstance(unwrapped_model, GPTModel), ( - "The final unwrapped model must be a GPTModel instance " - "since only GPTModel is supported for EP A2A overlapping." + # GPTModel is the canonical model class supporting EP A2A overlap. + # MultimodalModel wraps a GPTModel decoder and exposes its own + # build_schedule_plan that delegates decoder-layer scheduling to the + # inner GPTModel; vision encoder runs eagerly on the main path. + _allowed_for_a2a_overlap = (GPTModel,) + try: + from examples.multimodal_dev.models.base import MultimodalModel + + _allowed_for_a2a_overlap = (GPTModel, MultimodalModel) + except ImportError: + pass + assert isinstance(unwrapped_model, _allowed_for_a2a_overlap), ( + "The final unwrapped model must be a GPTModel or MultimodalModel " + "(decoder-only EP A2A overlap) instance. " + f"Got {type(unwrapped_model).__name__}." ) f_schedule_plan, loss_func = forward_step_func( data_iterator, unwrapped_model, return_schedule_plan=True diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index f9b923632f5..08747f12952 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -6,6 +6,7 @@ # LICENSE file in the root directory of this source tree. import logging +import os from dataclasses import dataclass, replace from typing import List, Optional, Tuple, Union @@ -14,9 +15,16 @@ import torch.nn.functional as F from torch import Tensor +from megatron.core import tensor_parallel from megatron.core.dist_checkpointing import ShardedTensor from megatron.core.dist_checkpointing.mapping import ReplicaId, ShardedTensorFactory from megatron.core.fp8_utils import get_fp8_align_size +from megatron.core.fusions.fused_mega_pre_gated_delta_rule import ( + fused_mega_pre_gated_delta_rule, +) +from megatron.core.fusions.fused_pre_gated_delta_rule import ( + fused_streamed_pre_gated_delta_rule, +) from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.jit import jit_fuser from megatron.core.packed_seq_params import PackedSeqParams @@ -28,6 +36,7 @@ _undo_attention_load_balancing, ) from megatron.core.tensor_parallel import get_cuda_rng_tracker +from megatron.core.tensor_parallel.random import CheckpointManager from megatron.core.transformer import TransformerConfig from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.module import MegatronModule @@ -43,7 +52,14 @@ try: from fla.modules.convolution import causal_conv1d from fla.modules.l2norm import l2norm - from fla.ops.gated_delta_rule import chunk_gated_delta_rule + + if os.environ.get("MCORE_GDN_USE_OPT_WRAPPER", "0") == "1": + try: + from mcore_gdn_opt.gated_delta_rule import chunk_gated_delta_rule + except ImportError: + from fla.ops.gated_delta_rule import chunk_gated_delta_rule + else: + from fla.ops.gated_delta_rule import chunk_gated_delta_rule HAVE_FLA = True except ImportError: @@ -122,6 +138,11 @@ def __init__( self.cp_size = self.pg_collection.cp.size() self.tp_size = self.pg_collection.tp.size() self.sp_size = self.tp_size if config.sequence_parallel else 1 + self.pre_gated_delta_rule_impl = config.pre_gated_delta_rule_impl + if self.pre_gated_delta_rule_impl != "unfused": + assert ( + self.cp_size == 1 + ), "Fused pre_gated_delta_rule does not support context parallelism yet." # Attributes from config self.config = config @@ -219,6 +240,15 @@ def __init__( hidden_size=self.value_head_dim, eps=self.config.layernorm_epsilon, ) + self.recompute_norm_out = False + self.recompute_qkv = False + if self.config.recompute_granularity == "selective": + self.recompute_norm_out = "gdn_norm_out" in self.config.recompute_modules + # gdn_qkv: recompute the whole QKV proj+prep block as a discard-output checkpoint. + self.recompute_qkv = "gdn_qkv" in self.config.recompute_modules + + # Per-forward CheckpointManager for the GDN discard-output recompute (gdn_qkv/gdn_norm_out). + self.gdn_recompute_manager = None self.out_proj = build_module( submodules.out_proj, @@ -332,6 +362,109 @@ def forward( cu_seqlens_q = None cu_seqlens_kv = None + # gdn_qkv (QKV proj+prep) and gdn_norm_out (gated norm) are discard-output checkpoints; the + # QKV output `gate` feeds the gated-norm block, so when both are on the CheckpointManager + # replays them in forward order (qkv -> norm_out) from one grad hook on `out`. + recompute_qkv = self.recompute_qkv and self.training + recompute_norm_out = self.recompute_norm_out and self.training + self.gdn_recompute_manager = ( + CheckpointManager() if (recompute_qkv or recompute_norm_out) else None + ) + + # QKV projection + prep block (in_proj -> CP a2a -> conv1d -> _prepare_qkv -> g/beta). + def _qkv_proj_and_prepare(hidden_states): + return self._compute_qkv_for_gated_delta_rule( + hidden_states, batch, seq_len, cu_seqlens_q, packed_seq_params + ) + + if recompute_qkv: + # Discard the QKV outputs now; regenerate them in backward. Synchronous recompute + # (no async reload), so it is safe with the fla/compiled gated_delta_rule backward. + query, key, value, g, beta, gate = tensor_parallel.CheckpointWithoutOutput( + fp8=(self.config.fp8 or self.config.fp4), + ckpt_manager=self.gdn_recompute_manager, + ).checkpoint(_qkv_proj_and_prepare, hidden_states) + else: + query, key, value, g, beta, gate = _qkv_proj_and_prepare(hidden_states) + + # seq_len was reassigned to the post-CP-a2a sequence length inside the block; recover it + # from a produced tensor so the downstream gated-norm reshape uses the correct value. + seq_len = value.shape[1] + + nvtx_range_push(suffix="gated_delta_rule") + core_attn_out, last_recurrent_state = self.gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + cu_seqlens=cu_seqlens_q, + ) + nvtx_range_pop(suffix="gated_delta_rule") + + def _gated_norm_and_a2a(core_attn_out: torch.Tensor, gate: torch.Tensor): + # RMSNorm + nvtx_range_push(suffix="gated_norm") + norm_out_hp = self._apply_gated_norm(core_attn_out, gate) + nvtx_range_pop(suffix="gated_norm") + + # Transpose: b s x --> s b x + # From bshd back to sbhd format + norm_out_hp = norm_out_hp.reshape(batch, seq_len, -1) + norm_out_hp = norm_out_hp.transpose(0, 1).contiguous() + + # CP all to all: HP to CP + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + unpacked_norm_out = _unpack_sequence(norm_out_hp, cu_seqlens_q, dim=0) + outputs = [] + for norm_out_i in unpacked_norm_out: + norm_out_i = tensor_a2a_hp2cp( + norm_out_i, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp + ) + outputs.append(norm_out_i) + norm_out = torch.cat(outputs, dim=0) + else: + norm_out = tensor_a2a_hp2cp( + norm_out_hp, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp + ) + + return norm_out + + if recompute_norm_out: + norm_out = tensor_parallel.CheckpointWithoutOutput( + ckpt_manager=self.gdn_recompute_manager + ).checkpoint(_gated_norm_and_a2a, core_attn_out, gate) + else: + norm_out = _gated_norm_and_a2a(core_attn_out, gate) + + # Output projection + nvtx_range_push(suffix="out_proj") + out, out_bias = self.out_proj(norm_out) + nvtx_range_pop(suffix="out_proj") + + # Discard the checkpointed outputs (now consumed) and register the unified recompute hook on + # `out` — its grad is computed first in backward, before the backwards that need them. + if self.gdn_recompute_manager is not None: + self.gdn_recompute_manager.discard_all_outputs_and_register_unified_recompute(out) + self.gdn_recompute_manager = None + + return out, out_bias + + def _compute_qkv_for_gated_delta_rule( + self, hidden_states, batch, seq_len, cu_seqlens_q, packed_seq_params + ): + """QKV projection + preparation block for the gated delta rule. + + Runs in_proj, CP all-to-all, conv1d, _prepare_qkv and g/beta, producing the tensors consumed + by ``self.gated_delta_rule`` plus the ``gate`` for the gated norm. Extracted so it can be + checkpointed when ``recompute_modules`` contains ``"gdn_qkv"``. + + Returns: + Tuple of (query, key, value, g, beta, gate). + """ # Input projection nvtx_range_push(suffix="in_proj") qkvzba, _ = self.in_proj(hidden_states) @@ -374,6 +507,43 @@ def forward( ], ) + # Fused pre-gated-delta-rule path: a single fused kernel replaces the conv1d -> + # _prepare_qkv -> g/beta block below. The fused wrappers consume the post-CP-a2a + # qkvzba (s b x) directly and return (query, key, value, gate, beta, g); reorder to the + # (query, key, value, g, beta, gate) layout this method returns. + if self.pre_gated_delta_rule_impl != "unfused": + seq_idx = ( + packed_seq_params.seq_idx + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + else None + ) + if self.pre_gated_delta_rule_impl == "fused_streamed": + nvtx_range_push(suffix="fused_streamed_pre_gated_delta_rule") + query, key, value, gate, beta, g = self._fused_streamed_pre_gated_delta_rule( + qkvzba, cu_seqlens_q=cu_seqlens_q, seq_idx=seq_idx + ) + nvtx_range_pop(suffix="fused_streamed_pre_gated_delta_rule") + else: + assert self.pre_gated_delta_rule_impl == "fused_mega" + nvtx_range_push(suffix="fused_mega_pre_gated_delta_rule") + query, key, value, gate, beta, g = self._fused_mega_pre_gated_delta_rule( + qkvzba, cu_seqlens_q=cu_seqlens_q, seq_idx=seq_idx + ) + nvtx_range_pop(suffix="fused_mega_pre_gated_delta_rule") + return query, key, value, g, beta, gate + + # Unfused path: conv1d -> _prepare_qkv -> g/beta on the post-CP-a2a qkvzba. + query, key, value, gate, beta, g = self.pre_gated_delta_rule( + qkvzba, batch, seq_len, cu_seqlens_q=cu_seqlens_q + ) + return query, key, value, g, beta, gate + + def pre_gated_delta_rule(self, qkvzba, batch, seq_len, cu_seqlens_q=None): + """Unfused pre-gated-delta-rule on the post-CP-a2a qkvzba. + + Runs the split -> conv1d -> _prepare_qkv -> g/beta block (the reference path that the + fused wrappers replace). Returns (query, key, value, gate, beta, g). + """ # Transpose: s b x --> b s x # From sbhd to bshd format qkvzba = qkvzba.transpose(0, 1) @@ -459,51 +629,45 @@ def forward( g, beta = self._compute_g_and_beta(A_log_local_cp, dt_bias_local_cp, alpha, beta) nvtx_range_pop(suffix="g_and_beta") - nvtx_range_push(suffix="gated_delta_rule") - core_attn_out, last_recurrent_state = self.gated_delta_rule( - query, - key, - value, - g=g, - beta=beta, - initial_state=None, - output_final_state=False, - use_qk_l2norm_in_kernel=False, - cu_seqlens=cu_seqlens_q, - ) - nvtx_range_pop(suffix="gated_delta_rule") - - # RMSNorm - nvtx_range_push(suffix="gated_norm") - norm_out = self._apply_gated_norm(core_attn_out, gate) - nvtx_range_pop(suffix="gated_norm") + return query, key, value, gate, beta, g - # Transpose: b s x --> s b x - # From bshd back to sbhd format - norm_out = norm_out.reshape(batch, seq_len, -1) - norm_out = norm_out.transpose(0, 1).contiguous() + def _fused_streamed_pre_gated_delta_rule(self, qkvzba, cu_seqlens_q=None, seq_idx=None): + """Call the streamed fused pre-GDR wrapper. Returns (query, key, value, gate, beta, g).""" - # CP all to all: HP to CP - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': - unpacked_norm_out = _unpack_sequence(norm_out, cu_seqlens_q, dim=0) - outputs = [] - for norm_out_i in unpacked_norm_out: - norm_out_i = tensor_a2a_hp2cp( - norm_out_i, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp - ) - outputs.append(norm_out_i) - norm_out = torch.cat(outputs, dim=0) - else: - norm_out = tensor_a2a_hp2cp( - norm_out, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp - ) + assert self.cp_size == 1, "Fused pre_gated_delta_rule does not support CP yet." + return fused_streamed_pre_gated_delta_rule( + qkvzba, + self.conv1d.weight, + self.conv1d.bias if self.conv_bias else None, + self.A_log, + self.dt_bias, + num_key_heads=self.qk_dim_local_tp // self.key_head_dim, + num_value_heads=self.v_dim_local_tp // self.value_head_dim, + key_head_dim=self.key_head_dim, + value_head_dim=self.value_head_dim, + use_qk_l2norm=self.use_qk_l2norm, + cu_seqlens=cu_seqlens_q, + seq_idx=seq_idx, + ) - # Output projection - nvtx_range_push(suffix="out_proj") - out, out_bias = self.out_proj(norm_out) - nvtx_range_pop(suffix="out_proj") + def _fused_mega_pre_gated_delta_rule(self, qkvzba, cu_seqlens_q=None, seq_idx=None): + """Call the mega fused pre-GDR wrapper. Returns (query, key, value, gate, beta, g).""" - return out, out_bias + assert self.cp_size == 1, "Fused pre_gated_delta_rule does not support CP yet." + return fused_mega_pre_gated_delta_rule( + qkvzba, + self.conv1d.weight, + self.conv1d.bias if self.conv_bias else None, + self.A_log, + self.dt_bias, + num_key_heads=self.qk_dim_local_tp // self.key_head_dim, + num_value_heads=self.v_dim_local_tp // self.value_head_dim, + key_head_dim=self.key_head_dim, + value_head_dim=self.value_head_dim, + use_qk_l2norm=self.use_qk_l2norm, + cu_seqlens=cu_seqlens_q, + seq_idx=seq_idx, + ) @jit_fuser def _apply_gated_norm(self, x, gate): diff --git a/megatron/core/tensor_parallel/random.py b/megatron/core/tensor_parallel/random.py index 4516fe10d88..0f1bdbd2da8 100644 --- a/megatron/core/tensor_parallel/random.py +++ b/megatron/core/tensor_parallel/random.py @@ -914,30 +914,43 @@ def _recompute(self, _): self.outputs = None self.ctx = None - def discard_output_and_register_recompute(self, hook_tensor): - """ - Release the output tensor storages and register the recompute function as a grad hook of - the hook_tensor. + def discard_output(self): + """Free the output storages (metadata kept for backward). - Note: the caller should make sure that the output tensors are no longer used - in the forward pass and the gradient of the hook_tensor is computed before the recomputed - tensors are used. + Pair with :meth:`register_recompute_hook` when the output is freed in a different place + than where the recompute hook is registered; otherwise use + :meth:`discard_output_and_register_recompute`. """ - # When ckpt_manager is set, this is a no-op. - # Manager handles all discarding and hook registration uniformly. from megatron.core.transformer.cuda_graphs import is_graph_warmup if self.ckpt_manager is not None or is_graph_warmup(): return - - # use resize to release the output tensor memory and still keep the metadata in the tensors. - # the metadata is still needed for backward + # resize keeps tensor metadata (needed for backward) while releasing the memory. for output in self.outputs: output.untyped_storage().resize_(0) - # register the recomputation as a backward hook, when the the gradient of the hook_tensor - # is computed, the recomputation will be triggered. The hook_tensor should be selected - # carefully to ensure that the tensors are recomputed before it is used by other backward - # computations. + def register_recompute_hook(self, hook_tensor): + """Trigger the recompute from ``hook_tensor``'s grad hook. + + ``hook_tensor`` must have its grad computed before the discarded outputs are needed in + backward (and, if the recompute reads other discarded activations, after those are + restored). + """ + from megatron.core.transformer.cuda_graphs import is_graph_warmup + + if self.ckpt_manager is not None or is_graph_warmup(): + return if hook_tensor.requires_grad: hook_tensor.register_hook(self._recompute) + + def discard_output_and_register_recompute(self, hook_tensor): + """ + Release the output tensor storages and register the recompute function as a grad hook of + the hook_tensor. + + Note: the caller should make sure that the output tensors are no longer used + in the forward pass and the gradient of the hook_tensor is computed before the recomputed + tensors are used. + """ + self.discard_output() + self.register_recompute_hook(hook_tensor) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 31e06a84b48..e9eb8122d57 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -449,6 +449,7 @@ def _build_per_layer_rotary_pos_emb(self, rotary_base: float) -> None: rotary_interleaved=self.config.rotary_interleaved, seq_len_interpolation_factor=seq_len_interpolation_factor, rotary_base=rotary_base, + interleaved_mrope=self.config.mrope_interleaved, ) self.mrope_section = self.config.mrope_section assert ( diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py index e3d19e88d7f..24fd9754a7e 100644 --- a/megatron/core/transformer/moe/fused_a2a.py +++ b/megatron/core/transformer/moe/fused_a2a.py @@ -399,14 +399,18 @@ def forward( # DeepEP calculates tx_depth = 3 * num_tokens + 1. # InfiniBand strictly asserts tx_depth < 65536. tx_depth = 3 * num_tokens + 1 - if tx_depth >= 65536: - raise ValueError( - f"HybridEP RDMA Queue Pair depth ({tx_depth}) exceeds the InfiniBand " - f"hardware limit of 65535. This occurs because the total tokens per rank " - f"({num_tokens}) too high. Reduce sequence length or micro-batch size, " - f"or increase Tensor Parallelism (TP) / Context Parallelism (CP) to reduce " - f"the number of tokens processed per rank." - ) + # PATCH (jinliangl, per expert guidance): false-positive guard, commented out. + # The check incorrectly fires on valid HybridEP configurations (e.g. MBS=6, seq=4096 + # -> 24,576 tokens/rank -> tx_depth=73,729). HybridEP runs cleanly past this threshold + # in practice; re-enable only if a real RDMA QP failure is reproduced. + # if tx_depth >= 65536: + # raise ValueError( + # f"HybridEP RDMA Queue Pair depth ({tx_depth}) exceeds the InfiniBand " + # f"hardware limit of 65535. This occurs because the total tokens per rank " + # f"({num_tokens}) too high. Reduce sequence length or micro-batch size, " + # f"or increase Tensor Parallelism (TP) / Context Parallelism (CP) to reduce " + # f"the number of tokens processed per rank." + # ) fp8_dispatch = False # Currently, we do not support fp8 dispatch init_hybrid_ep_buffer( group, diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index 12ca77e184b..ed5efd34565 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -243,6 +243,18 @@ def __init__( config.recompute_granularity == 'selective' and "shared_experts" in config.recompute_modules ) + # Discard the shared-expert OUTPUT (CheckpointWithoutOutput) instead of a standard + # checkpoint that keeps it. The recompute is wired up where ordering is correct (after any + # pre_mlp_layernorm recompute) by the caller: the fine-grained callables for A2A overlap, + # or TransformerLayer._forward_post_mlp otherwise. Disabled under MoE cudagraph partial + # capture, where the output is a graph output and must keep its storage. + self.shared_experts_recompute_discard_output = ( + self.shared_experts_recompute + and not bool(getattr(config, "cuda_graph_modules", None)) + ) + # The active CheckpointWithoutOutput, handed to the caller that frees the output and + # registers the recompute hook (None when not using discard-output recompute). + self.shared_experts_checkpoint = None self.tp_group = pg_collection.tp self.tp_ep_group = pg_collection.tp_ep @@ -530,7 +542,19 @@ def shared_experts_compute(self, hidden_states: torch.Tensor): shared_expert_output = None if self.use_shared_expert and not self.shared_expert_overlap: # Compute the shared expert separately when not overlapped with communication. - if self.shared_experts_recompute: + if self.shared_experts_recompute_discard_output and self.training: + # Recompute-with-discarded-output: run the shared expert under no_grad, then free + # its output in postprocess() and regenerate it (with its backward graph) from a + # grad hook. CheckpointWithoutOutput handles fp8/fp4 internally via its fp8 flag. + self.shared_experts_checkpoint = tensor_parallel.CheckpointWithoutOutput( + fp8=(self.config.fp8 or self.config.fp4) + ) + shared_expert_output = self.shared_experts_checkpoint.checkpoint( + apply_module(self.shared_experts), hidden_states + ) + elif self.shared_experts_recompute: + # Standard checkpoint fallback (e.g. MoE cudagraph partial capture or eval): keep + # the output, recompute only the intermediates. if self.config.fp8 or self.config.fp4: shared_expert_output = te_checkpoint( apply_module(self.shared_experts), diff --git a/megatron/core/transformer/moe/shared_experts.py b/megatron/core/transformer/moe/shared_experts.py index db36c8ec701..68d05d6992e 100644 --- a/megatron/core/transformer/moe/shared_experts.py +++ b/megatron/core/transformer/moe/shared_experts.py @@ -28,6 +28,8 @@ is_te_min_version, is_torch_min_version, make_sharded_tensor_for_checkpoint, + nvtx_range_pop, + nvtx_range_push, ) if HAVE_TE: @@ -183,11 +185,13 @@ def __init__( def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: """Forward function""" + nvtx_range_push("SharedExpert.forward") output, _ = super().forward(hidden_states) 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 + nvtx_range_pop("SharedExpert.forward") return output def _reset_parameters(self): @@ -233,6 +237,7 @@ def pre_forward_comm(self, input, wait_current_stream=True): if wait_current_stream: self.wait_current_stream() with torch.cuda.stream(self.stream): + nvtx_range_push("SharedExpert.pre_forward_comm") if self.use_shared_expert_gate: logits = torch.nn.functional.linear(input, self.gate_weight) self.gate_score = torch.nn.functional.sigmoid(logits) @@ -243,6 +248,7 @@ def pre_forward_comm(self, input, wait_current_stream=True): else: self.cached_fc1_input = copy_to_tensor_model_parallel_region(input) set_tensor_grad_fn_sequence_sr(self.cached_fc1_input, torch.iinfo(torch.int).max) + nvtx_range_pop("SharedExpert.pre_forward_comm") @overlap_state_check( SharedExpertState.PRE_FORWARD_COMM_DONE, SharedExpertState.FC1_FORWARD_DONE @@ -254,6 +260,7 @@ def linear_fc1_forward_and_act(self, overlapped_comm_output=None): It is only useful when --moe-shared-expert-overlap is set and may be changed. """ with torch.cuda.stream(self.stream): + nvtx_range_push("SharedExpert.linear_fc1_forward_and_act") # [s, b, 4 * h/p] intermediate_parallel, bias_parallel = apply_module(self.linear_fc1)( self.cached_fc1_input @@ -301,6 +308,7 @@ def glu(x): intermediate_parallel = self.activation_func(intermediate_parallel) self.cached_fc2_input = intermediate_parallel + nvtx_range_pop("SharedExpert.linear_fc1_forward_and_act") # Tensor sequence number is used to control the backward order. # Decrease the sequence number of the expert output to make the comm launched first # in the backward order. @@ -320,9 +328,11 @@ def linear_fc2_forward(self, overlapped_comm_output=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): + nvtx_range_push("SharedExpert.linear_fc2_forward") # [s, b, h] self.cached_fc2_output, _ = apply_module(self.linear_fc2)(self.cached_fc2_input) self.cached_fc2_input = None + nvtx_range_pop("SharedExpert.linear_fc2_forward") @overlap_state_check( SharedExpertState.FC2_FORWARD_DONE, SharedExpertState.POST_FORWARD_COMM_DONE @@ -334,6 +344,7 @@ def post_forward_comm(self): It is only useful when --moe-shared-expert-overlap is set and may be changed. """ with torch.cuda.stream(self.stream): + nvtx_range_push("SharedExpert.post_forward_comm") if self.config.sequence_parallel: self.cached_output = reduce_scatter_to_sequence_parallel_region( self.cached_fc2_output @@ -344,6 +355,7 @@ def post_forward_comm(self): ) self.cached_fc2_output = None set_tensor_grad_fn_sequence_sr(self.cached_output, torch.iinfo(torch.int).max) + nvtx_range_pop("SharedExpert.post_forward_comm") @overlap_state_check(SharedExpertState.POST_FORWARD_COMM_DONE, SharedExpertState.IDLE) def get_output(self): @@ -353,6 +365,7 @@ def get_output(self): It is only useful when --moe-shared-expert-overlap is set and may be changed. """ with torch.cuda.stream(self.stream): + nvtx_range_push("SharedExpert.get_output") if self.use_shared_expert_gate: assert self.gate_score is not None output = self.cached_output * self.gate_score @@ -360,6 +373,7 @@ def get_output(self): else: output = self.cached_output self.cached_output = None + nvtx_range_pop("SharedExpert.get_output") torch.cuda.current_stream().wait_stream(self.stream) return output diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index b85f157e3da..084902cc48d 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -365,6 +365,9 @@ class TransformerConfig(ModelParallelConfig): linear_num_value_heads: Optional[int] = 32 """Number of value and gate heads for the gated delta net.""" + pre_gated_delta_rule_impl: Literal["unfused", "fused_streamed", "fused_mega"] = "unfused" + """Pre-gated-delta-rule implementation for GatedDeltaNet.""" + #################### # initialization #################### @@ -543,7 +546,7 @@ class TransformerConfig(ModelParallelConfig): recompute_modules: Optional[List[str]] = None """The submodules to recompute. choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", - "shared_experts", "mhc". + "shared_experts", "mhc", "gdn_norm_out". default: ["core_attn"]. "core_attn": recompute the core attention part of the transformer layer. "moe_act": recompute the MoE MLP activation function. @@ -555,8 +558,9 @@ class TransformerConfig(ModelParallelConfig): "mhc": recompute HyperConnection intermediate activations via CheckpointWithoutOutput + CheckpointManager. Requires enable_hyper_connections=True. Cannot be used with "mlp". - "moe_act", "layernorm", "mla_up_proj", and "mhc" use output-discarding checkpointing, - "core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing. + "gdn_norm_out": recompute the GatedDeltaNet output norm and HP-to-CP all-to-all. + "moe_act", "layernorm", "mla_up_proj", "mhc", and "gdn_norm_out" use output-discarding + checkpointing, "core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing. """ #################### @@ -1385,6 +1389,21 @@ def __post_init__(self): self.experimental_attention_variant = self.linear_attention_type self.linear_attention_type = None + valid_pre_gdr_impls = ("unfused", "fused_streamed", "fused_mega") + if self.pre_gated_delta_rule_impl not in valid_pre_gdr_impls: + raise ValueError( + "pre_gated_delta_rule_impl must be one of " + f"{valid_pre_gdr_impls}, got {self.pre_gated_delta_rule_impl!r}." + ) + if ( + self.pre_gated_delta_rule_impl != "unfused" + and self.experimental_attention_variant != "gated_delta_net" + ): + raise ValueError( + "pre_gated_delta_rule_impl can select a fused path only when " + "experimental_attention_variant='gated_delta_net'." + ) + if self.experimental_attention_variant in ["gated_delta_net"]: assert ( self.linear_attention_freq is not None @@ -1761,6 +1780,8 @@ def __post_init__(self): "moe", "shared_experts", "mhc", + "gdn_norm_out", + "gdn_qkv", } invalid_modules = set(self.recompute_modules) - allowed_modules assert not invalid_modules, ( @@ -1779,6 +1800,24 @@ def __post_init__(self): "multi_latent_attention." ) + if ( + "gdn_norm_out" in self.recompute_modules + and self.experimental_attention_variant != "gated_delta_net" + ): + raise ValueError( + "gdn_norm_out in recompute_modules is only supported with " + "experimental_attention_variant='gated_delta_net'." + ) + + if ( + "gdn_qkv" in self.recompute_modules + and self.experimental_attention_variant != "gated_delta_net" + ): + raise ValueError( + "gdn_qkv in recompute_modules is only supported with " + "experimental_attention_variant='gated_delta_net'." + ) + if "core_attn" in self.recompute_modules: warnings.warn( "If you are using transformer_engine as the transformer implementation, " @@ -2221,7 +2260,18 @@ def __post_init__(self): "It is experimental and may change in future versions." ) else: - if self.rotary_interleaved: + fused_mrope_available = False + # Triton fused mRoPE supports split-half RoPE only. Keep rotary_interleaved + # configs on the TE validation path so the TE >= 2.3 check still applies. + if self.mrope_section is not None and not self.rotary_interleaved: + try: + from megatron.core.fusions.fused_mrope import is_fused_mrope_available + + fused_mrope_available = is_fused_mrope_available() + except ImportError: + fused_mrope_available = False + + if self.rotary_interleaved and not fused_mrope_available: if not is_te_min_version("2.3.0"): raise ValueError( "rotary_interleaved does not work with apply_rope_fusion for " @@ -2233,9 +2283,14 @@ def __post_init__(self): fused_apply_rotary_pos_emb_thd, ) - if fused_apply_rotary_pos_emb is None and fused_apply_rotary_pos_emb_thd is None: + if ( + fused_apply_rotary_pos_emb is None + and fused_apply_rotary_pos_emb_thd is None + and not fused_mrope_available + ): raise ValueError( - "apply_rope_fusion is not available. Please install TE >= 1.4." + "apply_rope_fusion is not available. Please install TE >= 1.4 " + "or Triton for fused mRoPE." ) if self.fused_single_qkv_rope: diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index d2e090de232..e3422381d78 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -904,6 +904,19 @@ def _forward_post_mlp( mlp_output_with_bias[0] ) + # Shared-expert discard-output recompute (non-overlap path; the A2A-overlap path uses the + # fine-grained callables and never reaches here). The output was consumed by the MoE + # postprocess add, so free it and register the recompute on mlp_output_with_bias[0] AFTER + # the pre_mlp_norm recompute above (same hook tensor) — this orders it after the shared + # expert's input pre_mlp_layernorm_output is restored and before its backward. + if self.is_moe_layer: + shared_experts_checkpoint = getattr(self.mlp, "shared_experts_checkpoint", None) + if shared_experts_checkpoint is not None: + shared_experts_checkpoint.discard_output_and_register_recompute( + mlp_output_with_bias[0] + ) + self.mlp.shared_experts_checkpoint = None + # TODO: could we move `bias_dropout_add_exec_handler` itself # inside the module provided in the `bias_dropout_add_spec` module? nvtx_range_push(suffix="mlp_bda") diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 853973b92cd..75ef03b5c59 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1630,7 +1630,7 @@ def validate_args(args, defaults={}): # Legacy RoPE arguments if args.use_rotary_position_embeddings: args.position_embedding_type = 'rope' - if args.position_embedding_type != 'rope': + if args.position_embedding_type not in ('rope', 'mrope'): args.apply_rope_fusion = False # Would just need to add 'NoPE' as a position_embedding_type to support this, but for now @@ -4576,6 +4576,10 @@ def _add_mla_args(parser): def _add_experimental_attention_variant_args(parser): group = parser.add_argument_group(title="experimental_attention_variant") + # NOTE: --pre-gated-delta-rule-impl is auto-generated from the + # TransformerConfig.pre_gated_delta_rule_impl field by ArgumentGroupFactory + # (see _add_transformer_engine_args / build_group), so it must NOT be + # registered manually here — doing so raises an argparse conflict. # Linear attention group.add_argument( '--linear-attention-freq', diff --git a/megatron/training/datasets/data_samplers.py b/megatron/training/datasets/data_samplers.py index 430bd8b85da..8b14b975aba 100644 --- a/megatron/training/datasets/data_samplers.py +++ b/megatron/training/datasets/data_samplers.py @@ -106,7 +106,7 @@ def close_nvidia_fds(): maybe_worker_init_fn = worker_init_fn if args.num_workers > 0 else None # Torch dataloader. - if args.dynamic_context_parallel: + if args.dynamic_context_parallel or getattr(args, "use_vanilla_collate_fn", False): extra_kwargs = {"collate_fn": lambda x: x} else: extra_kwargs = {} diff --git a/megatron/training/training.py b/megatron/training/training.py index 34e65337703..a195aa007de 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -1253,6 +1253,25 @@ def pretrain( # Model, optimizer, and learning rate. timers('model-and-optimizer-setup', log_level=0).start(barrier=True) + + # Enable CUDA memory event recording BEFORE model/optimizer init, so the + # snapshot at dump time contains full allocation event traces (frames + + # device_traces), not just segment state. Gated on the existing + # `--record-memory-history` flag. + if args.record_memory_history and ( + is_last_rank() or torch.distributed.get_backend() == 'fake' + ): + try: + torch.cuda.memory._record_memory_history( + enabled='all', + context='all', + stacks='python', + max_entries=100000, + ) + print_rank_0("[memory_snapshot] enabled torch.cuda.memory event recording (mode=all)") + except Exception as _e: # noqa: BLE001 + print_rank_0(f"[memory_snapshot] _record_memory_history failed: {_e}") + model, optimizer, opt_param_scheduler = setup_model_and_optimizer( model_provider, model_type, checkpointing_context=checkpointing_context ) diff --git a/tests/unit_tests/a2a_overlap/test_fsdp_1f1b_overlap.py b/tests/unit_tests/a2a_overlap/test_fsdp_1f1b_overlap.py index 4e5ddc7eb02..c7005c72884 100644 --- a/tests/unit_tests/a2a_overlap/test_fsdp_1f1b_overlap.py +++ b/tests/unit_tests/a2a_overlap/test_fsdp_1f1b_overlap.py @@ -74,9 +74,12 @@ def test_fsdp_1f1b_training_step( [[], ["attn_norm", "core_attn", "attn_proj", "mlp_norm", "expert_fc1", "moe_act"]], ) def test_fsdp_1f1b_memory_opt(self, recompute_modules, offload_modules): + # Configure shared experts so recompute_modules=[..., "shared_experts"] actually + # exercises the shared-experts discard-output recompute (a no-op without them). self._run_test_helper( dispatcher_type="alltoall", sharding_strategy="optim_grads_params", + shared_expert_intermediate_size=512, recompute_modules=recompute_modules, offload_modules=offload_modules, ) diff --git a/tests/unit_tests/fusions/test_fused_mrope.py b/tests/unit_tests/fusions/test_fused_mrope.py new file mode 100644 index 00000000000..b033f6b9bed --- /dev/null +++ b/tests/unit_tests/fusions/test_fused_mrope.py @@ -0,0 +1,1311 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import warnings +from types import SimpleNamespace + +import pytest +import torch + +import megatron.core.models.common.embeddings.rope_utils as rope_utils +from megatron.core import parallel_state +from megatron.core.fusions.fused_mrope import ( + fused_apply_mrope, + fused_apply_mrope_thd, + get_fused_mrope_thd_unavailable_reason, + get_fused_mrope_unavailable_reason, + is_fused_mrope_available, + mrope_freqs_to_rotary_emb, +) +from megatron.core.models.common.embeddings import apply_rotary_pos_emb +from megatron.core.models.common.embeddings.rope_utils import ( + _ROPE_FUSION_FALLBACK_WARNINGS, + _apply_rotary_pos_emb_bshd, + _apply_rotary_pos_emb_thd, +) +from megatron.core.models.common.embeddings.rotary_pos_embedding import MultimodalRotaryEmbedding +from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.transformer.transformer_config import TransformerConfig +from tests.unit_tests.test_utilities import Utils + + +class FakeCPGroup: + def __init__(self, size=1, rank=0): + self._size = size + self._rank = rank + + def size(self): + return self._size + + def rank(self): + return self._rank + + +class FakeDynamicInferenceContext: + def is_dynamic_batching(self): + return True + + def is_static_batching(self): + return False + + +class FakeStaticInferenceContext: + def is_dynamic_batching(self): + return False + + def is_static_batching(self): + return True + + +@pytest.fixture(autouse=True) +def clear_rope_fusion_fallback_warnings(): + _ROPE_FUSION_FALLBACK_WARNINGS.clear() + yield + _ROPE_FUSION_FALLBACK_WARNINGS.clear() + + +def _dtype_tols(dtype): + if dtype == torch.bfloat16: + return dict(rtol=2.0e-2, atol=5.0e-2) + if dtype == torch.float16: + return dict(rtol=3.0e-3, atol=1.0e-2) + return dict(rtol=1.0e-6, atol=1.0e-6) + + +def _make_inputs( + dtype=torch.bfloat16, + requires_grad=False, + head_dim=20, + rotary_dim=16, + mrope_section=None, + interleaved_mrope=False, + batch=2, +): + seq = 32 + heads = 3 + if mrope_section is None: + mrope_section = [3, 3, 2] if interleaved_mrope else [2, 3, 3] + + generator = torch.Generator(device="cuda").manual_seed(1234) + t = torch.randn( + seq, + batch, + heads, + head_dim, + dtype=dtype, + device="cuda", + generator=generator, + requires_grad=requires_grad, + ) + freqs = torch.randn( + 3, batch, seq, rotary_dim // 2, dtype=torch.float32, device="cuda", generator=generator + ) + return t, freqs, mrope_section + + +def _make_position_ids(seq, batch): + base = torch.arange(seq, device="cuda", dtype=torch.long) + batch_offsets = torch.arange(batch, device="cuda", dtype=torch.long) + return ( + torch.stack((base, base * 2 + 3, base * 3 + 5), dim=0)[:, None, :] + + batch_offsets[None, :, None] + ).contiguous() + + +def _make_thd_inputs( + dtype=torch.bfloat16, + requires_grad=False, + interleaved_mrope=False, + cp_size=1, + padded_seq_lens=(12, 16), + head_dim=20, + rotary_dim=16, + mrope_section=None, +): + total_seq = sum(padded_seq_lens) + local_seq = total_seq // cp_size + heads = 3 + if mrope_section is None: + mrope_section = [3, 3, 2] if interleaved_mrope else [2, 3, 3] + + generator = torch.Generator(device="cuda").manual_seed(5678) + t = torch.randn( + local_seq, + heads, + head_dim, + dtype=dtype, + device="cuda", + generator=generator, + requires_grad=requires_grad, + ) + freqs = torch.randn( + 3, 1, total_seq, rotary_dim // 2, dtype=torch.float32, device="cuda", generator=generator + ) + cu_seqlens = torch.tensor([0, padded_seq_lens[0], total_seq], dtype=torch.int32, device="cuda") + return t, freqs, cu_seqlens, mrope_section + + +def _make_mrope_config( + num_attention_heads, mrope_section, interleaved_mrope=False, rotary_interleaved=False +): + return TransformerConfig( + num_attention_heads=num_attention_heads, + num_layers=1, + apply_rope_fusion=True, + mrope_section=mrope_section, + mrope_interleaved=interleaved_mrope, + rotary_interleaved=rotary_interleaved, + ) + + +def _fallback_warnings(recorded_warnings): + return [ + warning + for warning in recorded_warnings + if issubclass(warning.category, UserWarning) + and "Using unfused implementation" in str(warning.message) + ] + + +def _thd_cp_freq_indices(cu_seqlens_cpu, cp_size, cp_rank): + indices = [] + for global_start, global_end in zip(cu_seqlens_cpu[:-1], cu_seqlens_cpu[1:]): + local_seq_len = (global_end - global_start) // cp_size + first_cp_seg = (local_seq_len + 1) // 2 + second_cp_seg = local_seq_len // 2 + indices.extend( + range( + global_start + cp_rank * first_cp_seg, global_start + (cp_rank + 1) * first_cp_seg + ) + ) + indices.extend( + range(global_end - (cp_rank + 1) * second_cp_seg, global_end - cp_rank * second_cp_seg) + ) + return indices + + +def _assert_thd_cp_freq_index_coverage(cu_seqlens_cpu, cp_size): + expected = [] + actual = [] + for global_start, global_end in zip(cu_seqlens_cpu[:-1], cu_seqlens_cpu[1:]): + expected.extend(range(global_start, global_end)) + for cp_rank in range(cp_size): + actual.extend(_thd_cp_freq_indices(cu_seqlens_cpu, cp_size, cp_rank)) + assert sorted(actual) == expected + assert len(set(actual)) == len(actual) + + +@pytest.mark.parametrize("use_packed_seq", [False, True]) +def test_gpt_mrope_eval_requests_raw_freqs_when_fusion_available(use_packed_seq): + captured_kwargs = {} + + def fake_rotary_pos_emb(*args, **kwargs): + captured_kwargs.update(kwargs) + return "raw-mrope-freqs" + + model = SimpleNamespace( + training=False, + pre_process=False, + mtp_process=False, + position_embedding_type="mrope", + config=SimpleNamespace( + multi_latent_attention=False, + flash_decode=False, + apply_rope_fusion=True, + rotary_interleaved=False, + cuda_graph_impl=None, + fused_single_qkv_rope=False, + ), + rotary_pos_emb=fake_rotary_pos_emb, + mrope_section=[2, 3, 3], + _fused_mrope_available=True, + ) + packed_seq_params = ( + SimpleNamespace(qkv_format="thd", cp_group=FakeCPGroup()) if use_packed_seq else None + ) + + output = GPTModel._preprocess( + model, + input_ids=torch.zeros(1, 4, dtype=torch.long), + position_ids=torch.zeros(3, 1, 4, dtype=torch.long), + decoder_input=torch.zeros(4, 1, 12), + packed_seq_params=packed_seq_params, + ) + + assert output[1] == "raw-mrope-freqs" + assert captured_kwargs["return_raw_freqs"] is True + assert captured_kwargs["packed_seq"] is use_packed_seq + + +def test_gpt_mrope_eval_keeps_materialized_freqs_with_fused_single_qkv_rope(): + captured_kwargs = {} + + def fake_rotary_pos_emb(*args, **kwargs): + captured_kwargs.update(kwargs) + return "materialized-mrope-freqs" + + model = SimpleNamespace( + training=False, + pre_process=False, + mtp_process=False, + position_embedding_type="mrope", + config=SimpleNamespace( + multi_latent_attention=False, + flash_decode=False, + apply_rope_fusion=True, + rotary_interleaved=False, + cuda_graph_impl=None, + fused_single_qkv_rope=True, + ), + rotary_pos_emb=fake_rotary_pos_emb, + mrope_section=[2, 3, 3], + _fused_mrope_available=True, + ) + + output = GPTModel._preprocess( + model, + input_ids=torch.zeros(1, 4, dtype=torch.long), + position_ids=torch.zeros(3, 1, 4, dtype=torch.long), + decoder_input=torch.zeros(4, 1, 12), + ) + + assert output[1] == "materialized-mrope-freqs" + assert captured_kwargs["return_raw_freqs"] is False + + +def test_gpt_mrope_dynamic_inference_keeps_materialized_freqs(): + captured_kwargs = {} + + def fake_rotary_pos_emb(*args, **kwargs): + captured_kwargs.update(kwargs) + return "materialized-mrope-freqs" + + model = SimpleNamespace( + training=False, + pre_process=False, + mtp_process=False, + position_embedding_type="mrope", + config=SimpleNamespace( + multi_latent_attention=False, + flash_decode=False, + apply_rope_fusion=True, + rotary_interleaved=False, + cuda_graph_impl=None, + fused_single_qkv_rope=False, + ), + rotary_pos_emb=fake_rotary_pos_emb, + mrope_section=[2, 3, 3], + _fused_mrope_available=True, + ) + + output = GPTModel._preprocess( + model, + input_ids=torch.zeros(1, 4, dtype=torch.long), + position_ids=torch.zeros(3, 1, 4, dtype=torch.long), + decoder_input=torch.zeros(4, 1, 12), + inference_context=FakeDynamicInferenceContext(), + ) + + assert output[1] == "materialized-mrope-freqs" + assert captured_kwargs["return_raw_freqs"] is False + + +def test_gpt_mrope_static_inference_keeps_materialized_freqs(): + captured_kwargs = {} + + def fake_rotary_pos_emb(*args, **kwargs): + captured_kwargs.update(kwargs) + return "materialized-mrope-freqs" + + model = SimpleNamespace( + training=False, + pre_process=False, + mtp_process=False, + position_embedding_type="mrope", + config=SimpleNamespace( + multi_latent_attention=False, + flash_decode=False, + apply_rope_fusion=True, + rotary_interleaved=False, + cuda_graph_impl=None, + fused_single_qkv_rope=False, + ), + rotary_pos_emb=fake_rotary_pos_emb, + mrope_section=[2, 3, 3], + _fused_mrope_available=True, + ) + + output = GPTModel._preprocess( + model, + input_ids=torch.zeros(1, 4, dtype=torch.long), + position_ids=torch.zeros(3, 1, 4, dtype=torch.long), + decoder_input=torch.zeros(4, 1, 12), + inference_context=FakeStaticInferenceContext(), + ) + + assert output[1] == "materialized-mrope-freqs" + assert captured_kwargs["return_raw_freqs"] is False + + +def test_is_fused_mrope_available_requires_cuda(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + assert not is_fused_mrope_available() + + +def test_transformer_config_rejects_fused_mrope_without_cuda_or_te(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(rope_utils, "fused_apply_rotary_pos_emb", None) + monkeypatch.setattr(rope_utils, "fused_apply_rotary_pos_emb_thd", None) + + with pytest.raises(ValueError, match="apply_rope_fusion is not available"): + TransformerConfig( + num_attention_heads=1, num_layers=1, apply_rope_fusion=True, mrope_section=[1, 1, 1] + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("interleaved_mrope", [False, True]) +@pytest.mark.parametrize("head_dim", [16, 20]) +def test_fused_mrope_matches_unfused_forward_backward(interleaved_mrope, head_dim): + t_ref, freqs, mrope_section = _make_inputs( + requires_grad=True, head_dim=head_dim, interleaved_mrope=interleaved_mrope + ) + t_fused = t_ref.detach().clone().requires_grad_(True) + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + ref = _apply_rotary_pos_emb_bshd(t_ref, emb, rotary_interleaved=False) + out = fused_apply_mrope( + t_fused, freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + + tols = _dtype_tols(t_ref.dtype) + torch.testing.assert_close(ref.float(), out.float(), **tols) + + grad = torch.randn_like(ref) + ref.backward(grad) + out.backward(grad) + torch.testing.assert_close(t_ref.grad.float(), t_fused.grad.float(), **tols) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("interleaved_mrope", [False, True]) +def test_apply_rotary_pos_emb_bshd_eval_uses_triton_without_te(interleaved_mrope, monkeypatch): + t, freqs, mrope_section = _make_inputs(interleaved_mrope=interleaved_mrope, batch=1) + config = _make_mrope_config(t.shape[2], mrope_section, interleaved_mrope) + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + ref = _apply_rotary_pos_emb_bshd(t, emb, rotary_interleaved=False) + + fused_calls = 0 + orig_fused_apply_mrope = rope_utils.fused_apply_mrope + + def wrapped_fused_apply_mrope(*args, **kwargs): + nonlocal fused_calls + fused_calls += 1 + return orig_fused_apply_mrope(*args, **kwargs) + + monkeypatch.setattr(rope_utils, "fused_apply_rotary_pos_emb", None) + monkeypatch.setattr(rope_utils, "fused_apply_mrope", wrapped_fused_apply_mrope) + with torch.no_grad(), warnings.catch_warnings(record=True) as recorded_warnings: + warnings.simplefilter("always") + out = apply_rotary_pos_emb(t, freqs, config, cp_group=FakeCPGroup()) + + assert fused_calls == 1 + assert not _fallback_warnings(recorded_warnings) + assert not out.requires_grad + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize( + "fallback_kwargs, warning_match", + [ + ({"mscale": 1.25}, "mscale=1.25 is not supported by Triton fused mRoPE"), + ({"inverse": True}, "inverse RoPE is not supported by Triton fused mRoPE"), + ], +) +def test_apply_rotary_pos_emb_raw_mrope_fallbacks_match_unfused(fallback_kwargs, warning_match): + t, freqs, mrope_section = _make_inputs() + config = TransformerConfig( + num_attention_heads=t.shape[2], + num_layers=1, + apply_rope_fusion=True, + mrope_section=mrope_section, + ) + + with pytest.warns(UserWarning, match=warning_match): + out = apply_rotary_pos_emb(t, freqs, config, cp_group=FakeCPGroup(), **fallback_kwargs) + + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + ref = _apply_rotary_pos_emb_bshd(t, emb, rotary_interleaved=False, **fallback_kwargs) + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + with warnings.catch_warnings(record=True) as repeated_warnings: + warnings.simplefilter("always") + out_again = apply_rotary_pos_emb( + t, freqs, config, cp_group=FakeCPGroup(), **fallback_kwargs + ) + assert not repeated_warnings + torch.testing.assert_close(ref.float(), out_again.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize( + "fallback_kwargs, config_kwargs, expected_warning_key, warning_text", + [ + ( + {"mscale": 1.25}, + {}, + "triton-mrope-mscale", + "mscale=1.25 is not supported by Triton fused mRoPE", + ), + ( + {"inverse": True}, + {}, + "triton-mrope-inverse", + "inverse RoPE is not supported by Triton fused mRoPE", + ), + ( + {"mla_rotary_interleaved": True}, + {}, + "triton-mrope-mla-rotary-interleaved", + "does not support MLA-style interleaving", + ), + ( + {}, + {"rotary_interleaved": True}, + "triton-mrope-unavailable-rotary-interleaved", + "rotary_interleaved=True is not supported", + ), + ], +) +def test_apply_rotary_pos_emb_raw_mrope_fallback_emits_single_warning( + fallback_kwargs, config_kwargs, expected_warning_key, warning_text +): + t, freqs, mrope_section = _make_inputs() + config = _make_mrope_config(t.shape[2], mrope_section, **config_kwargs) + + with warnings.catch_warnings(record=True) as recorded_warnings: + warnings.simplefilter("always") + out = apply_rotary_pos_emb(t, freqs, config, cp_group=FakeCPGroup(), **fallback_kwargs) + + fallback_warnings = _fallback_warnings(recorded_warnings) + assert len(fallback_warnings) == 1 + assert warning_text in str(fallback_warnings[0].message) + assert _ROPE_FUSION_FALLBACK_WARNINGS == {expected_warning_key} + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, rotary_interleaved=config.rotary_interleaved + ) + ref = _apply_rotary_pos_emb_bshd( + t, emb, rotary_interleaved=config.rotary_interleaved, **fallback_kwargs + ) + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +def test_interleaved_mrope_rejects_inconsistent_sections(): + freqs = torch.randn(3, 2, 8, 8, dtype=torch.float32) + + with pytest.raises(AssertionError, match="interleaved mRoPE"): + mrope_freqs_to_rotary_emb(freqs, [2, 3, 3], interleaved_mrope=True) + + +def test_raw_mrope_cpu_falls_back_to_unfused(): + t = torch.randn(8, 1, 3, 20, dtype=torch.float32) + freqs = torch.randn(3, 1, 8, 8, dtype=torch.float32) + mrope_section = [2, 3, 3] + config = SimpleNamespace( + apply_rope_fusion=True, + mrope_section=mrope_section, + mrope_interleaved=False, + rotary_interleaved=False, + ) + + unavailable_reason = get_fused_mrope_unavailable_reason(t, freqs) + assert unavailable_reason is not None + with pytest.warns( + UserWarning, match="(CUDA tensors|Triton is not available).*Using unfused implementation" + ): + out = apply_rotary_pos_emb(t, freqs, config, cp_group=FakeCPGroup()) + assert _ROPE_FUSION_FALLBACK_WARNINGS in ( + {"triton-mrope-unavailable-device"}, + {"triton-mrope-unavailable-import"}, + ) + + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + ref = _apply_rotary_pos_emb_bshd(t, emb, rotary_interleaved=False) + torch.testing.assert_close(ref, out) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_raw_mrope_unsupported_dtype_falls_back_to_unfused(): + t, freqs, mrope_section = _make_inputs(dtype=torch.float64) + config = TransformerConfig( + num_attention_heads=t.shape[2], + num_layers=1, + apply_rope_fusion=True, + mrope_section=mrope_section, + ) + + assert "dtype" in get_fused_mrope_unavailable_reason(t, freqs) + with pytest.warns(UserWarning, match="dtype.*Using unfused implementation"): + out = apply_rotary_pos_emb(t, freqs, config, cp_group=FakeCPGroup()) + assert _ROPE_FUSION_FALLBACK_WARNINGS == {"triton-mrope-unavailable-dtype"} + + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + ref = _apply_rotary_pos_emb_bshd(t, emb, rotary_interleaved=False) + torch.testing.assert_close(ref, out) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("interleaved_mrope", [False, True]) +def test_apply_rotary_pos_emb_dispatches_raw_mrope(interleaved_mrope): + t, freqs, mrope_section = _make_inputs(interleaved_mrope=interleaved_mrope) + config = TransformerConfig( + num_attention_heads=t.shape[2], + num_layers=1, + apply_rope_fusion=True, + mrope_section=mrope_section, + mrope_interleaved=interleaved_mrope, + ) + + out = apply_rotary_pos_emb(t, freqs, config, cp_group=FakeCPGroup()) + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + ref = _apply_rotary_pos_emb_bshd(t, emb, rotary_interleaved=False) + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_raw_mrope_unsupported_freq_dtype_warning_key_is_dtype(): + t, freqs, mrope_section = _make_inputs() + freqs = freqs.to(torch.float16) + config = TransformerConfig( + num_attention_heads=t.shape[2], + num_layers=1, + apply_rope_fusion=True, + mrope_section=mrope_section, + ) + + assert "float32" in get_fused_mrope_unavailable_reason(t, freqs) + with pytest.warns(UserWarning, match="float32.*Using unfused implementation"): + out = apply_rotary_pos_emb(t, freqs, config, cp_group=FakeCPGroup()) + assert _ROPE_FUSION_FALLBACK_WARNINGS == {"triton-mrope-unavailable-dtype"} + + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + ref = _apply_rotary_pos_emb_bshd(t, emb, rotary_interleaved=False) + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +def test_apply_rotary_pos_emb_raw_mrope_checks_triton_availability_once(monkeypatch): + t = torch.randn(4, 1, 2, 8, dtype=torch.float32) + freqs = torch.randn(3, 1, 4, 4, dtype=torch.float32) + config = SimpleNamespace( + apply_rope_fusion=True, + mrope_section=[1, 1, 2], + mrope_interleaved=False, + rotary_interleaved=False, + ) + + calls = 0 + + def fake_unavailable_reason(*args, **kwargs): + nonlocal calls + calls += 1 + return None + + monkeypatch.setattr(rope_utils, "get_fused_mrope_unavailable_reason", fake_unavailable_reason) + monkeypatch.setattr(rope_utils, "fused_apply_mrope", lambda *args, **kwargs: t + 1) + + out = apply_rotary_pos_emb(t, freqs, config, cp_group=FakeCPGroup()) + + assert calls == 1 + torch.testing.assert_close(out, t + 1) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("layout", ["bshd", "thd"]) +def test_materialized_mrope_falls_back_without_te_fused_rope(monkeypatch, layout): + if layout == "bshd": + t, freqs, mrope_section = _make_inputs(batch=1) + cu_seqlens = None + monkeypatch.setattr(rope_utils, "fused_apply_rotary_pos_emb", None) + else: + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs() + monkeypatch.setattr(rope_utils, "fused_apply_rotary_pos_emb_thd", None) + + config = _make_mrope_config(t.shape[-2], mrope_section) + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + + with pytest.warns(UserWarning, match="Transformer Engine fused RoPE.*unavailable"): + out = apply_rotary_pos_emb(t, emb, config, cu_seqlens, cp_group=FakeCPGroup()) + + if layout == "bshd": + ref = _apply_rotary_pos_emb_bshd(t, emb, rotary_interleaved=False) + else: + ref = _apply_rotary_pos_emb_thd(t, cu_seqlens, emb, cp_group=FakeCPGroup()) + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize( + "fallback_kwargs, expected_warning_key, warning_text", + [ + ( + {"mscale": 1.25}, + "te-rope-thd-mscale", + "mscale=1.25 is not supported by TE's fused RoPE for THD layout", + ), + ( + {"inverse": True}, + "te-rope-thd-inverse", + "inverse RoPE is not supported by TE's fused RoPE for THD layout", + ), + ( + {"mla_rotary_interleaved": True}, + "te-rope-thd-mla-rotary-interleaved", + "does not support MLA-style interleaving", + ), + ], +) +def test_materialized_thd_mrope_option_fallbacks_do_not_call_te( + monkeypatch, fallback_kwargs, expected_warning_key, warning_text +): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs() + config = _make_mrope_config(t.shape[1], mrope_section) + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + + def unexpected_te_thd_call(*args, **kwargs): + raise AssertionError("TE THD fused RoPE should not be called") + + monkeypatch.setattr(rope_utils, "fused_apply_rotary_pos_emb_thd", unexpected_te_thd_call) + with warnings.catch_warnings(record=True) as recorded_warnings: + warnings.simplefilter("always") + out = apply_rotary_pos_emb( + t, emb, config, cu_seqlens, cp_group=FakeCPGroup(), **fallback_kwargs + ) + + fallback_warnings = _fallback_warnings(recorded_warnings) + assert len(fallback_warnings) == 1 + assert warning_text in str(fallback_warnings[0].message) + assert _ROPE_FUSION_FALLBACK_WARNINGS == {expected_warning_key} + + ref = _apply_rotary_pos_emb_thd(t, cu_seqlens, emb, cp_group=FakeCPGroup(), **fallback_kwargs) + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("interleaved_mrope", [False, True]) +@pytest.mark.parametrize("cp_size, cp_rank", [(1, 0), (2, 0), (2, 1)]) +def test_fused_mrope_thd_matches_unfused_forward_backward( + interleaved_mrope, cp_size, cp_rank, monkeypatch +): + t_ref, freqs, cu_seqlens, mrope_section = _make_thd_inputs( + requires_grad=True, interleaved_mrope=interleaved_mrope, cp_size=cp_size + ) + t_fused = t_ref.detach().clone().requires_grad_(True) + cp_group = FakeCPGroup(size=cp_size, rank=cp_rank) + config = TransformerConfig( + num_attention_heads=t_ref.shape[1], + num_layers=1, + context_parallel_size=cp_size, + apply_rope_fusion=True, + mrope_section=mrope_section, + mrope_interleaved=interleaved_mrope, + ) + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + ref = _apply_rotary_pos_emb_thd(t_ref, cu_seqlens, emb, cp_group=cp_group) + + fused_calls = 0 + orig_fused_apply_mrope_thd = rope_utils.fused_apply_mrope_thd + + def wrapped_fused_apply_mrope_thd(*args, **kwargs): + nonlocal fused_calls + fused_calls += 1 + return orig_fused_apply_mrope_thd(*args, **kwargs) + + def unexpected_pack(*args, **kwargs): + raise AssertionError("raw THD mRoPE fusion should not materialize packed freqs") + + monkeypatch.setattr(rope_utils, "fused_apply_mrope_thd", wrapped_fused_apply_mrope_thd) + monkeypatch.setattr(rope_utils, "_pack_thd_raw_mrope_freqs", unexpected_pack) + out = apply_rotary_pos_emb(t_fused, freqs, config, cu_seqlens, cp_group=cp_group) + assert fused_calls == 1 + + tols = _dtype_tols(t_ref.dtype) + torch.testing.assert_close(ref.float(), out.float(), **tols) + + grad = torch.randn_like(ref) + ref.backward(grad) + out.backward(grad) + torch.testing.assert_close(t_ref.grad.float(), t_fused.grad.float(), **tols) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("interleaved_mrope", [False, True]) +def test_apply_rotary_pos_emb_thd_eval_uses_triton_without_te(interleaved_mrope, monkeypatch): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs(interleaved_mrope=interleaved_mrope) + config = _make_mrope_config(t.shape[1], mrope_section, interleaved_mrope) + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + ref = _apply_rotary_pos_emb_thd(t, cu_seqlens, emb, cp_group=FakeCPGroup()) + + fused_calls = 0 + orig_fused_apply_mrope_thd = rope_utils.fused_apply_mrope_thd + + def wrapped_fused_apply_mrope_thd(*args, **kwargs): + nonlocal fused_calls + fused_calls += 1 + return orig_fused_apply_mrope_thd(*args, **kwargs) + + def unexpected_pack(*args, **kwargs): + raise AssertionError("raw THD mRoPE fusion should not materialize packed freqs") + + monkeypatch.setattr(rope_utils, "fused_apply_rotary_pos_emb_thd", None) + monkeypatch.setattr(rope_utils, "fused_apply_mrope_thd", wrapped_fused_apply_mrope_thd) + monkeypatch.setattr(rope_utils, "_pack_thd_raw_mrope_freqs", unexpected_pack) + with torch.no_grad(), warnings.catch_warnings(record=True) as recorded_warnings: + warnings.simplefilter("always") + out = apply_rotary_pos_emb(t, freqs, config, cu_seqlens, cp_group=FakeCPGroup()) + + assert fused_calls == 1 + assert not _fallback_warnings(recorded_warnings) + assert not out.requires_grad + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_apply_rotary_pos_emb_thd_fused_dispatch_does_not_read_cuda_scalars(monkeypatch): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs() + config = _make_mrope_config(t.shape[1], mrope_section) + + def unexpected_item(_tensor): + raise AssertionError("fused raw THD mRoPE dispatch should not call Tensor.item()") + + monkeypatch.setattr(torch.Tensor, "item", unexpected_item) + out = apply_rotary_pos_emb(t, freqs, config, cu_seqlens, cp_group=FakeCPGroup()) + + assert out.shape == t.shape + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize( + "fallback_kwargs, config_kwargs, expected_warning_key, warning_text", + [ + ( + {"mscale": 1.25}, + {}, + "triton-mrope-thd-mscale", + "mscale=1.25 is not supported by Triton fused mRoPE for THD layout", + ), + ( + {"inverse": True}, + {}, + "triton-mrope-thd-inverse", + "inverse RoPE is not supported by Triton fused mRoPE for THD layout", + ), + ( + {"mla_rotary_interleaved": True}, + {}, + "triton-mrope-thd-mla-rotary-interleaved", + "does not support MLA-style interleaving", + ), + ( + {}, + {"rotary_interleaved": True}, + "triton-mrope-thd-rotary-interleaved", + "currently supports rotary_interleaved=False", + ), + ], +) +def test_apply_rotary_pos_emb_thd_raw_mrope_fallback_emits_option_warning( + fallback_kwargs, config_kwargs, expected_warning_key, warning_text +): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs() + config = _make_mrope_config(t.shape[1], mrope_section, **config_kwargs) + + with warnings.catch_warnings(record=True) as recorded_warnings: + warnings.simplefilter("always") + out = apply_rotary_pos_emb( + t, freqs, config, cu_seqlens, cp_group=FakeCPGroup(), **fallback_kwargs + ) + + fallback_warnings = _fallback_warnings(recorded_warnings) + assert len(fallback_warnings) == 1 + assert warning_text in str(fallback_warnings[0].message) + assert _ROPE_FUSION_FALLBACK_WARNINGS == {expected_warning_key} + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, rotary_interleaved=config.rotary_interleaved + ) + ref = _apply_rotary_pos_emb_thd( + t, + cu_seqlens, + emb, + rotary_interleaved=config.rotary_interleaved, + cp_group=FakeCPGroup(), + **fallback_kwargs, + ) + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_thd_raw_mrope_rejects_sequence_length_mismatch(): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs() + config = _make_mrope_config(t.shape[1], mrope_section) + bad_freqs = freqs[:, :, :-1, :].contiguous() + + with pytest.raises(ValueError, match="sequence length must match local tokens"): + apply_rotary_pos_emb(t, bad_freqs, config, cu_seqlens, cp_group=FakeCPGroup()) + + +def test_thd_raw_mrope_rejects_global_sequence_length_not_divisible_by_cp(): + t = torch.randn(2, 3, 20, dtype=torch.float32) + freqs = torch.randn(3, 1, 5, 8, dtype=torch.float32) + cu_seqlens = torch.tensor([0, 5], dtype=torch.int32) + config = SimpleNamespace( + apply_rope_fusion=True, + mrope_section=[2, 3, 3], + mrope_interleaved=False, + rotary_interleaved=False, + ) + + with pytest.raises(ValueError, match="divisible by context parallel size"): + apply_rotary_pos_emb(t, freqs, config, cu_seqlens, cp_group=FakeCPGroup(size=2)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_thd_raw_mrope_unavailable_reason_rejects_global_sequence_length_not_divisible_by_cp(): + t = torch.randn(3, 3, 20, dtype=torch.bfloat16, device="cuda") + freqs = torch.randn(3, 1, 5, 8, dtype=torch.float32, device="cuda") + cu_seqlens = torch.tensor([0, 5], dtype=torch.int32, device="cuda") + + assert "divisible by context parallel size" in get_fused_mrope_thd_unavailable_reason( + t, cu_seqlens, freqs, cp_size=2, cp_rank=0 + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("cp_rank", [0, 1]) +def test_thd_raw_mrope_cp_odd_local_sequence_lengths_match_manual_reference(cp_rank): + cp_size = 2 + t_ref, freqs, cu_seqlens, mrope_section = _make_thd_inputs( + requires_grad=True, cp_size=cp_size, padded_seq_lens=(10, 14) + ) + t_fused = t_ref.detach().clone().requires_grad_(True) + config = _make_mrope_config(t_ref.shape[1], mrope_section) + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + + cu_seqlens_cpu = cu_seqlens.cpu().tolist() + _assert_thd_cp_freq_index_coverage(cu_seqlens_cpu, cp_size) + packed_freqs = emb[_thd_cp_freq_indices(cu_seqlens_cpu, cp_size, cp_rank)] + + ref = _apply_rotary_pos_emb_bshd(t_ref.unsqueeze(1), packed_freqs).squeeze(1) + out = apply_rotary_pos_emb( + t_fused, freqs, config, cu_seqlens, cp_group=FakeCPGroup(size=cp_size, rank=cp_rank) + ) + + tols = _dtype_tols(t_ref.dtype) + torch.testing.assert_close(ref.float(), out.float(), **tols) + + grad = torch.randn_like(ref) + ref.backward(grad) + out.backward(grad) + torch.testing.assert_close(t_ref.grad.float(), t_fused.grad.float(), **tols) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("cp_rank", [0, 1]) +def test_thd_raw_mrope_fallback_supports_odd_local_sequence_lengths(cp_rank): + cp_size = 2 + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs( + cp_size=cp_size, padded_seq_lens=(10, 14) + ) + config = _make_mrope_config(t.shape[1], mrope_section) + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + cu_seqlens_cpu = cu_seqlens.cpu().tolist() + packed_freqs = emb[_thd_cp_freq_indices(cu_seqlens_cpu, cp_size, cp_rank)] + + with pytest.warns(UserWarning, match="mscale=1.25.*Using unfused implementation"): + out = apply_rotary_pos_emb( + t, + freqs, + config, + cu_seqlens, + mscale=1.25, + cp_group=FakeCPGroup(size=cp_size, rank=cp_rank), + ) + + ref = _apply_rotary_pos_emb_bshd(t.unsqueeze(1), packed_freqs, mscale=1.25).squeeze(1) + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_thd_raw_mrope_rejects_batch_dimension_greater_than_one(): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs() + config = _make_mrope_config(t.shape[1], mrope_section) + bad_freqs = freqs.expand(-1, 2, -1, -1).contiguous() + + with pytest.raises(ValueError, match="singleton batch dimension"): + apply_rotary_pos_emb(t, bad_freqs, config, cu_seqlens, cp_group=FakeCPGroup()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_thd_raw_mrope_rejects_non_thd_tensor_shape(): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs() + config = _make_mrope_config(t.shape[1], mrope_section) + + with pytest.raises(ValueError, match="raw mRoPE THD expects t"): + apply_rotary_pos_emb( + t[..., :8].unsqueeze(1), freqs, config, cu_seqlens, cp_group=FakeCPGroup() + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_fused_mrope_thd_public_api_matches_unfused(): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs() + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + + ref = _apply_rotary_pos_emb_thd(t, cu_seqlens, emb, cp_group=FakeCPGroup()) + assert ( + get_fused_mrope_thd_unavailable_reason(t, cu_seqlens, freqs, cp_size=1, cp_rank=0) is None + ) + out = fused_apply_mrope_thd(t, cu_seqlens, freqs, mrope_section) + + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("padded_seq_lens", [(28,), (8, 10, 10)]) +def test_fused_mrope_thd_matches_unfused_for_different_sequence_counts(padded_seq_lens): + t, freqs, cu_seqlens, mrope_section = _make_thd_inputs(padded_seq_lens=padded_seq_lens) + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + + ref = _apply_rotary_pos_emb_thd(t, cu_seqlens, emb, cp_group=FakeCPGroup()) + out = fused_apply_mrope_thd(t, cu_seqlens, freqs, mrope_section) + + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +def test_fused_mrope_thd_fp32_compute_matches_explicit_cast_forward_backward(): + t_ref, freqs, cu_seqlens, mrope_section = _make_thd_inputs(requires_grad=True) + t_fused = t_ref.detach().clone().requires_grad_(True) + emb = mrope_freqs_to_rotary_emb(freqs, mrope_section, rotary_interleaved=False) + + ref = _apply_rotary_pos_emb_thd(t_ref.float(), cu_seqlens, emb, cp_group=FakeCPGroup()).to( + t_ref.dtype + ) + out = fused_apply_mrope_thd(t_fused, cu_seqlens, freqs, mrope_section, fp32_compute=True) + + torch.testing.assert_close(ref.float(), out.float(), **_dtype_tols(t_ref.dtype)) + + grad = torch.randn_like(ref) + ref.backward(grad) + out.backward(grad) + torch.testing.assert_close(t_ref.grad.float(), t_fused.grad.float(), **_dtype_tols(t_ref.dtype)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize("return_raw_freqs", [False, True]) +def test_mrope_packed_seq_keeps_global_freqs_with_context_parallel(return_raw_freqs): + class FakeCPGroup2: + def size(self): + return 2 + + def rank(self): + return 0 + + seq = 16 + batch = 1 + head_dim = 20 + rotary_dim = 16 + mrope_section = [2, 3, 3] + cp_group = FakeCPGroup2() + position_ids = _make_position_ids(seq, batch) + rope = MultimodalRotaryEmbedding( + head_dim, rotary_percent=rotary_dim / head_dim, cp_group=cp_group + ) + + unpacked_freqs = rope( + position_ids, + mrope_section, + cp_group=cp_group, + return_raw_freqs=return_raw_freqs, + packed_seq=False, + ) + packed_freqs = rope( + position_ids, + mrope_section, + cp_group=cp_group, + return_raw_freqs=return_raw_freqs, + packed_seq=True, + ) + + seq_dim = 2 if return_raw_freqs else 0 + assert unpacked_freqs.shape[seq_dim] == seq // cp_group.size() + assert packed_freqs.shape[seq_dim] == seq + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.skipif(Utils.world_size < 2, reason="CP test requires at least 2 distributed ranks") +@pytest.mark.parametrize("interleaved_mrope", [False, True]) +def test_raw_mrope_fusion_matches_unfused_with_context_parallel(interleaved_mrope): + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=2) + try: + cp_group = parallel_state.get_context_parallel_group() + seq = 32 + batch = 2 + heads = 3 + head_dim = 20 + rotary_dim = 16 + mrope_section = [3, 3, 2] if interleaved_mrope else [2, 3, 3] + position_ids = _make_position_ids(seq, batch) + + rope = MultimodalRotaryEmbedding( + head_dim, + rotary_percent=rotary_dim / head_dim, + cp_group=cp_group, + interleaved_mrope=interleaved_mrope, + ) + raw_freqs = rope(position_ids, mrope_section, cp_group=cp_group, return_raw_freqs=True) + materialized_emb = rope(position_ids, mrope_section, cp_group=cp_group) + raw_freqs_emb = mrope_freqs_to_rotary_emb( + raw_freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + torch.testing.assert_close(raw_freqs_emb, materialized_emb) + + local_seq = seq // cp_group.size() + assert raw_freqs.shape == (3, batch, local_seq, rotary_dim // 2) + assert materialized_emb.shape == (local_seq, batch, 1, rotary_dim) + + generator = torch.Generator(device="cuda").manual_seed(4321) + t_ref = torch.randn( + local_seq, + batch, + heads, + head_dim, + dtype=torch.bfloat16, + device="cuda", + generator=generator, + requires_grad=True, + ) + t_fused = t_ref.detach().clone().requires_grad_(True) + + config = TransformerConfig( + num_attention_heads=heads, + num_layers=1, + context_parallel_size=cp_group.size(), + apply_rope_fusion=True, + mrope_section=mrope_section, + mrope_interleaved=interleaved_mrope, + ) + + ref = _apply_rotary_pos_emb_bshd(t_ref, materialized_emb, rotary_interleaved=False) + out = apply_rotary_pos_emb(t_fused, raw_freqs, config, cp_group=cp_group) + tols = _dtype_tols(t_ref.dtype) + torch.testing.assert_close(ref.float(), out.float(), **tols) + + grad = torch.randn_like(ref) + ref.backward(grad) + out.backward(grad) + torch.testing.assert_close(t_ref.grad.float(), t_fused.grad.float(), **tols) + finally: + Utils.destroy_model_parallel() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.skipif( + Utils.world_size < 2, reason="THD CP test requires at least 2 distributed ranks" +) +@pytest.mark.parametrize("interleaved_mrope", [False, True]) +def test_raw_mrope_thd_fusion_matches_unfused_with_context_parallel(interleaved_mrope): + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=2) + try: + cp_group = parallel_state.get_context_parallel_group() + t_ref, _, cu_seqlens, mrope_section = _make_thd_inputs( + requires_grad=True, interleaved_mrope=interleaved_mrope, cp_size=cp_group.size() + ) + t_fused = t_ref.detach().clone().requires_grad_(True) + config = TransformerConfig( + num_attention_heads=t_ref.shape[1], + num_layers=1, + context_parallel_size=cp_group.size(), + apply_rope_fusion=True, + mrope_section=mrope_section, + mrope_interleaved=interleaved_mrope, + ) + total_seq = int(cu_seqlens[-1].item()) + position_ids = _make_position_ids(total_seq, 1) + rope = MultimodalRotaryEmbedding( + t_ref.shape[-1], + rotary_percent=16 / t_ref.shape[-1], + cp_group=cp_group, + interleaved_mrope=interleaved_mrope, + ) + freqs = rope( + position_ids, mrope_section, cp_group=cp_group, return_raw_freqs=True, packed_seq=True + ) + emb = rope(position_ids, mrope_section, cp_group=cp_group, packed_seq=True) + + raw_freqs_emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + assert freqs.shape == (3, 1, total_seq, 8) + assert emb.shape == (total_seq, 1, 1, 16) + torch.testing.assert_close(raw_freqs_emb, emb) + + ref = _apply_rotary_pos_emb_thd(t_ref, cu_seqlens, emb, cp_group=cp_group) + out = apply_rotary_pos_emb(t_fused, freqs, config, cu_seqlens, cp_group=cp_group) + tols = _dtype_tols(t_ref.dtype) + torch.testing.assert_close(ref.float(), out.float(), **tols) + + grad = torch.randn_like(ref) + ref.backward(grad) + out.backward(grad) + torch.testing.assert_close(t_ref.grad.float(), t_fused.grad.float(), **tols) + finally: + Utils.destroy_model_parallel() + + +# --------------------------------------------------------------------------- +# Real Qwen3.5-VL deployment shapes. +# +# The parametrized tests above use head_dim=16/20 with rotary_dim=16 (~80% of +# channels rotated). The real Qwen3.5-VL config is head_dim=256 with +# rotary_percent=0.25 -> rotary_dim=64 (only 25% rotated, 75% pass-through) and +# mrope_section=[11,11,10] (interleaved). Exercise those exact shapes so a kernel +# regression in the large-pass-through / large-section regime is caught. +# --------------------------------------------------------------------------- + +# (head_dim, rotary_dim, mrope_section, interleaved_mrope) +_REAL_BSHD_SHAPES = [ + (256, 64, [11, 11, 10], True), # Qwen3.5-VL LLM decoder (75% pass-through) + (256, 64, [10, 11, 11], False), # same, section (non-interleaved) layout + (256, 256, [43, 43, 42], True), # full rotary (no pass-through) +] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("head_dim,rotary_dim,mrope_section,interleaved_mrope", _REAL_BSHD_SHAPES) +def test_fused_mrope_matches_unfused_real_shapes( + head_dim, rotary_dim, mrope_section, interleaved_mrope +): + t_ref, freqs, mrope_section = _make_inputs( + requires_grad=True, + head_dim=head_dim, + rotary_dim=rotary_dim, + mrope_section=mrope_section, + interleaved_mrope=interleaved_mrope, + ) + t_fused = t_ref.detach().clone().requires_grad_(True) + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + ref = _apply_rotary_pos_emb_bshd(t_ref, emb, rotary_interleaved=False) + out = fused_apply_mrope( + t_fused, freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + + tols = _dtype_tols(t_ref.dtype) + torch.testing.assert_close(ref.float(), out.float(), **tols) + + grad = torch.randn_like(ref) + ref.backward(grad) + out.backward(grad) + torch.testing.assert_close(t_ref.grad.float(), t_fused.grad.float(), **tols) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not is_fused_mrope_available(), reason="Triton fused mRoPE not available") +@pytest.mark.parametrize("interleaved_mrope", [False, True]) +def test_fused_mrope_thd_matches_unfused_real_shapes(interleaved_mrope): + # Real Qwen3.5-VL head_dim=256, rotary_dim=64 in THD packed layout. + section = [11, 11, 10] if interleaved_mrope else [10, 11, 11] + t_ref, freqs, cu_seqlens, mrope_section = _make_thd_inputs( + requires_grad=True, + interleaved_mrope=interleaved_mrope, + head_dim=256, + rotary_dim=64, + mrope_section=section, + ) + t_fused = t_ref.detach().clone().requires_grad_(True) + cp_group = FakeCPGroup(size=1, rank=0) + + emb = mrope_freqs_to_rotary_emb( + freqs, mrope_section, interleaved_mrope=interleaved_mrope, rotary_interleaved=False + ) + ref = _apply_rotary_pos_emb_thd(t_ref, cu_seqlens, emb, cp_group=cp_group) + out = fused_apply_mrope_thd( + t_fused, cu_seqlens, freqs, mrope_section, + interleaved_mrope=interleaved_mrope, rotary_interleaved=False, cp_size=1, cp_rank=0, + ) + + tols = _dtype_tols(t_ref.dtype) + torch.testing.assert_close(ref.float(), out.float(), **tols) + + grad = torch.randn_like(ref) + ref.backward(grad) + out.backward(grad) + torch.testing.assert_close(t_ref.grad.float(), t_fused.grad.float(), **tols) + + +def test_thd_unavailable_reason_rejects_non_cp_divisible_subsequence(): + # Per-sequence CP divisibility: total length is divisible by cp_size but an + # individual packed sub-sequence is not. The fused THD launch path + # (apply_rotary_pos_emb -> fused_apply_mrope_thd) must reject this so it falls + # back to the unfused path (which splits per-sequence correctly), instead of + # silently computing wrong CP token indices. + cp_size = 2 + # sub-sequence lengths 10 and 14 -> both even (OK); 9 and 15 -> total 24 even + # but each odd (must be rejected). + cu_seqlens = torch.tensor([0, 9, 24], dtype=torch.int32, device="cuda") + local_tokens = 24 // cp_size + t = torch.randn(local_tokens, 3, 20, dtype=torch.bfloat16, device="cuda") + freqs = torch.randn(3, 1, 24, 8, dtype=torch.float32, device="cuda") + reason = get_fused_mrope_thd_unavailable_reason( + t, cu_seqlens, freqs, rotary_interleaved=False, cp_size=cp_size, cp_rank=0 + ) + assert reason is not None and "sub-sequence" in reason, reason + + # Control: all sub-sequences divisible by cp_size -> launchable (reason None). + cu_ok = torch.tensor([0, 10, 24], dtype=torch.int32, device="cuda") + reason_ok = get_fused_mrope_thd_unavailable_reason( + t, cu_ok, freqs, rotary_interleaved=False, cp_size=cp_size, cp_rank=0 + ) + assert reason_ok is None, reason_ok diff --git a/tests/unit_tests/ssm/bench_gdn_cuda_opt.py b/tests/unit_tests/ssm/bench_gdn_cuda_opt.py new file mode 100644 index 00000000000..9668f18a4d1 --- /dev/null +++ b/tests/unit_tests/ssm/bench_gdn_cuda_opt.py @@ -0,0 +1,459 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Direct GatedDeltaNet CUDA optimization correctness and performance runner. + +This runner intentionally uses installed packages and normal project imports. +Install `mcore_gdn_opt` and FLA in editable mode before running it. +""" + +import argparse +import importlib.util +import os +import statistics +from contextlib import nullcontext +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +from megatron.core import parallel_state +from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_experimental_attention_variant_module_spec, +) +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer import TransformerConfig +from tests.unit_tests.test_utilities import Utils + +FLAGS = ( + "MCORE_GDN_USE_OPT_WRAPPER", + "MCORE_GDN_OPT_BACKEND", + "MCORE_GDN_OPT_WARN_FALLBACK", + "MCORE_GDN_OPT_ENABLE_FWD_H", + "MCORE_GDN_OPT_ENABLE_WY_BWD", + "MCORE_GDN_OPT_ENABLE_DV_DHU", + "MCORE_GDN_OPT_ENABLE_DHU", + "MCORE_GDN_OPT_ENABLE_DQKWG", + "FLA_CUTE_FWD_H", + "CHUNK_DELTA_FWD_USE_BWD_PORT", + "FLA_CUTE_WY_BWD", + "FLA_CUTE_BWD_DV_DHU", + "FLA_CUTE_BWD_DHU", + "FLA_CUTE_BWD_DQKWG", +) + + +SCENARIOS = { + "baseline": ("Triton baseline", {}), + "wrapper_fla": ( + "MCore wrapper forced FLA", + {"MCORE_GDN_USE_OPT_WRAPPER": "1", "MCORE_GDN_OPT_BACKEND": "fla"}, + ), + "wrapper_auto": ( + "MCore wrapper auto", + {"MCORE_GDN_USE_OPT_WRAPPER": "1", "MCORE_GDN_OPT_BACKEND": "auto"}, + ), + "wrapper_cuda": ( + "MCore wrapper forced CUDA", + {"MCORE_GDN_USE_OPT_WRAPPER": "1", "MCORE_GDN_OPT_BACKEND": "cuda"}, + ), + "wy": ( + "CUDA wy_bwd", + { + "MCORE_GDN_USE_OPT_WRAPPER": "1", + "MCORE_GDN_OPT_BACKEND": "cuda", + "MCORE_GDN_OPT_ENABLE_FWD_H": "0", + "MCORE_GDN_OPT_ENABLE_DV_DHU": "0", + "MCORE_GDN_OPT_ENABLE_DHU": "0", + "MCORE_GDN_OPT_ENABLE_DQKWG": "0", + }, + ), + "dv_dhu": ( + "CUDA dv_local+delta_h fused", + { + "MCORE_GDN_USE_OPT_WRAPPER": "1", + "MCORE_GDN_OPT_BACKEND": "cuda", + "MCORE_GDN_OPT_ENABLE_FWD_H": "0", + "MCORE_GDN_OPT_ENABLE_WY_BWD": "0", + "MCORE_GDN_OPT_ENABLE_DHU": "0", + "MCORE_GDN_OPT_ENABLE_DQKWG": "0", + }, + ), + "dhu": ( + "CUDA delta_h", + { + "MCORE_GDN_USE_OPT_WRAPPER": "1", + "MCORE_GDN_OPT_BACKEND": "cuda", + "MCORE_GDN_OPT_ENABLE_FWD_H": "0", + "MCORE_GDN_OPT_ENABLE_WY_BWD": "0", + "MCORE_GDN_OPT_ENABLE_DV_DHU": "0", + "MCORE_GDN_OPT_ENABLE_DQKWG": "0", + }, + ), + "dqkwg": ( + "CUDA dqkwg", + { + "MCORE_GDN_USE_OPT_WRAPPER": "1", + "MCORE_GDN_OPT_BACKEND": "cuda", + "MCORE_GDN_OPT_ENABLE_FWD_H": "0", + "MCORE_GDN_OPT_ENABLE_WY_BWD": "0", + "MCORE_GDN_OPT_ENABLE_DV_DHU": "0", + "MCORE_GDN_OPT_ENABLE_DHU": "0", + }, + ), + "separate": ( + "CUDA all three separate", + { + "MCORE_GDN_USE_OPT_WRAPPER": "1", + "MCORE_GDN_OPT_BACKEND": "cuda", + "MCORE_GDN_OPT_ENABLE_FWD_H": "0", + "MCORE_GDN_OPT_ENABLE_DV_DHU": "0", + }, + ), + "dv_dhu_dqkwg": ( + "CUDA fused_dv_dhu+dqkwg", + { + "MCORE_GDN_USE_OPT_WRAPPER": "1", + "MCORE_GDN_OPT_BACKEND": "cuda", + "MCORE_GDN_OPT_ENABLE_FWD_H": "0", + "MCORE_GDN_OPT_ENABLE_WY_BWD": "0", + "MCORE_GDN_OPT_ENABLE_DHU": "0", + }, + ), + "all_four": ( + "CUDA fwd_h+wy_bwd+dhu+dqkwg", + { + "MCORE_GDN_USE_OPT_WRAPPER": "1", + "MCORE_GDN_OPT_BACKEND": "cuda", + "MCORE_GDN_OPT_ENABLE_DV_DHU": "0", + }, + ), + "fwd_h_wy_dv_dhu_dqkwg": ( + "CUDA fwd_h+wy_bwd+fused_dv_dhu+dqkwg", + { + "MCORE_GDN_USE_OPT_WRAPPER": "1", + "MCORE_GDN_OPT_BACKEND": "cuda", + "MCORE_GDN_OPT_ENABLE_DHU": "0", + }, + ), +} + + +@dataclass +class AccuracyRow: + name: str + status: str + output_max_abs: float + input_grad_max_abs: float + worst_param: str + worst_param_max_abs: float + + +@dataclass +class PerfRow: + name: str + mean_us: float + median_us: float + min_us: float + max_us: float + speedup: float + + +def set_env(overrides): + for flag in FLAGS: + os.environ.pop(flag, None) + if "MCORE_GDN_USE_OPT_WRAPPER" not in overrides: + os.environ["MCORE_GDN_USE_OPT_WRAPPER"] = "0" + if "MCORE_GDN_OPT_BACKEND" not in overrides: + os.environ["MCORE_GDN_OPT_BACKEND"] = "fla" + os.environ.update(overrides) + + +def set_model_dispatch(model): + if os.environ.get("MCORE_GDN_USE_OPT_WRAPPER", "0") == "1": + from mcore_gdn_opt.gated_delta_rule import chunk_gated_delta_rule + else: + from fla.ops.gated_delta_rule import chunk_gated_delta_rule + + model.gated_delta_rule = chunk_gated_delta_rule + + +def validate_dispatch_sources(scenario_items): + if any("MCORE_GDN_OPT_BACKEND" in env for _, (_, env) in scenario_items): + for module_name in ( + "mcore_gdn_opt.gated_delta_rule.chunk", + "mcore_gdn_opt.gated_delta_rule.backward", + ): + spec = importlib.util.find_spec(module_name) + if spec is None or spec.origin is None: + raise RuntimeError(f"cannot locate required mcore_gdn_opt module {module_name!r}") + print( + f"MCORE_GDN_OPT_DISPATCH_SOURCE module={module_name} path={spec.origin}", flush=True + ) + + +def nvtx_range(label, enabled=True): + if enabled and torch.cuda.is_available(): + return torch.cuda.nvtx.range(label) + return nullcontext() + + +def scenario_label(index, name): + safe_name = name.replace(" ", "_").replace("+", "plus").replace("/", "_") + return f"gdn_only/{index:02d}_{safe_name}" + + +def make_model(dtype): + from megatron.core.ssm.gated_delta_net import GatedDeltaNet + + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, context_parallel_size=1 + ) + model_parallel_cuda_manual_seed(123) + pg_collection = ProcessGroupCollection( + tp=parallel_state.get_tensor_model_parallel_group(), + cp=parallel_state.get_context_parallel_group(), + ) + cfg = TransformerConfig( + hidden_size=128, + linear_conv_kernel_dim=2, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_num_key_heads=64, + linear_num_value_heads=64, + num_layers=1, + normalization="RMSNorm", + use_cpu_initialization=True, + layernorm_zero_centered_gamma=True, + num_attention_heads=64, + activation_func=F.silu, + bf16=(dtype == torch.bfloat16), + fp16=(dtype == torch.float16), + experimental_attention_variant="gated_delta_net", + linear_attention_freq=[1], + transformer_impl="transformer_engine", + ) + submodules = get_experimental_attention_variant_module_spec(config=cfg).submodules + return ( + GatedDeltaNet( + cfg, + submodules=submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ) + .cuda() + .to(dtype) + ) + + +def zero_grads(model): + model.zero_grad(set_to_none=True) + + +def compute_loss(output, loss): + if loss == "sum": + return output.float().sum() + if loss == "square_mean": + return output.float().square().mean() + raise ValueError(f"unknown loss: {loss}") + + +def run_once(model, x, env, loss, nvtx_label=None, use_nvtx=True): + set_env(env) + set_model_dispatch(model) + print( + "RUN_ONCE " + f"label={nvtx_label or 'none'} " + f"use_wrapper={os.environ.get('MCORE_GDN_USE_OPT_WRAPPER', '')} " + f"backend={os.environ.get('MCORE_GDN_OPT_BACKEND', '')}", + flush=True, + ) + zero_grads(model) + inp = x.detach().clone().requires_grad_(True) + with nvtx_range(nvtx_label, enabled=use_nvtx and nvtx_label is not None): + out, _ = model(inp, attention_mask=None) + compute_loss(out, loss).backward() + torch.cuda.synchronize() + grads = { + name: param.grad.detach().float().clone().cpu() + for name, param in model.named_parameters() + if param.grad is not None + } + return out.detach().float().clone().cpu(), inp.grad.detach().float().clone().cpu(), grads + + +def diff_max_abs(actual, expected): + return float((actual - expected).abs().max().item()) + + +def allclose(actual, expected, atol, rtol): + return bool(torch.isfinite(actual).all().item()) and bool( + torch.allclose(actual, expected, atol=atol, rtol=rtol) + ) + + +def check_accuracy(model, x, scenario_items, loss, atol, rtol, use_nvtx=True): + base_name, base_env = SCENARIOS["baseline"] + base_out, base_grad, base_params = run_once( + model, x, base_env, loss, "gdn_only/00_accuracy_reference/Triton_baseline", use_nvtx + ) + rows = [] + for scenario_idx, (_, (name, env)) in enumerate(scenario_items, start=1): + label = f"{scenario_label(scenario_idx, name)}/accuracy" + out, grad, params = run_once(model, x, env, loss, label, use_nvtx) + output_ok = allclose(out, base_out, atol, rtol) + grad_ok = allclose(grad, base_grad, atol, rtol) + worst_param = "" + worst_param_abs = 0.0 + params_ok = True + for param_name, expected in base_params.items(): + actual = params[param_name] + params_ok = params_ok and allclose(actual, expected, atol, rtol) + param_abs = diff_max_abs(actual, expected) + if param_abs > worst_param_abs: + worst_param = param_name + worst_param_abs = param_abs + rows.append( + AccuracyRow( + name=name, + status="PASS" if output_ok and grad_ok and params_ok else "FAIL", + output_max_abs=diff_max_abs(out, base_out), + input_grad_max_abs=diff_max_abs(grad, base_grad), + worst_param=worst_param, + worst_param_max_abs=worst_param_abs, + ) + ) + return rows + + +def fwd_bwd(model, x, env, loss, nvtx_label=None, use_nvtx=True): + set_env(env) + set_model_dispatch(model) + zero_grads(model) + inp = x.detach().requires_grad_(True) + with nvtx_range(nvtx_label, enabled=use_nvtx and nvtx_label is not None): + out, _ = model(inp, attention_mask=None) + compute_loss(out, loss).backward() + + +def benchmark(model, x, scenario_items, loss, warmup, repeats, rounds, use_nvtx=True): + rows = [] + baseline_us = None + for scenario_idx, (_, (name, env)) in enumerate(scenario_items, start=1): + base_label = scenario_label(scenario_idx, name) + for warmup_idx in range(warmup): + fwd_bwd(model, x, env, loss, f"{base_label}/warmup_{warmup_idx:02d}", use_nvtx) + torch.cuda.synchronize() + samples = [] + for round_idx in range(rounds): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + with nvtx_range( + f"{base_label}/round_{round_idx:02d}/measured_{repeats}iters", enabled=use_nvtx + ): + start.record() + for iter_idx in range(repeats): + fwd_bwd( + model, + x, + env, + loss, + f"{base_label}/round_{round_idx:02d}/iter_{iter_idx:02d}", + use_nvtx, + ) + end.record() + torch.cuda.synchronize() + samples.append(start.elapsed_time(end) * 1000.0 / repeats) + mean_us = statistics.mean(samples) + if baseline_us is None: + baseline_us = mean_us + rows.append( + PerfRow( + name=name, + mean_us=mean_us, + median_us=statistics.median(samples), + min_us=min(samples), + max_us=max(samples), + speedup=baseline_us / mean_us, + ) + ) + return rows + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dtype", choices=("bf16", "fp16"), default="bf16") + parser.add_argument("--loss", choices=("sum", "square_mean"), default="square_mean") + parser.add_argument("--scenarios", default="baseline,separate,all_four,fwd_h_wy_dv_dhu_dqkwg") + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--repeats", type=int, default=20) + parser.add_argument("--rounds", type=int, default=3) + parser.add_argument("--atol", type=float, default=5e-3) + parser.add_argument("--rtol", type=float, default=5e-3) + parser.add_argument("--fail-on-accuracy", action="store_true") + parser.add_argument("--no-nvtx", dest="use_nvtx", action="store_false", default=True) + return parser.parse_args() + + +def main(): + args = parse_args() + keys = [key.strip() for key in args.scenarios.split(",") if key.strip()] + if "baseline" not in keys: + keys.insert(0, "baseline") + unknown = [key for key in keys if key not in SCENARIOS] + if unknown: + raise ValueError(f"unknown scenarios: {unknown}; choices={sorted(SCENARIOS)}") + scenario_items = [(key, SCENARIOS[key]) for key in keys] + validate_dispatch_sources(scenario_items) + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 + + torch.manual_seed(123) + set_env({}) + print( + f"DEVICE {torch.cuda.get_device_name(0)} SHAPE B=2 T=8192 H=64 D=128 " + f"dtype={args.dtype} loss={args.loss}" + ) + try: + model = make_model(dtype).eval() + x = torch.randn(8192, 2, 128, device="cuda", dtype=dtype) + accuracy_rows = check_accuracy( + model, x, scenario_items, args.loss, args.atol, args.rtol, args.use_nvtx + ) + for row in accuracy_rows: + print( + f"ACCURACY name={row.name!r} status={row.status} " + f"output_max_abs={row.output_max_abs:.9f} " + f"input_grad_max_abs={row.input_grad_max_abs:.9f} " + f"worst_param={row.worst_param} " + f"worst_param_max_abs={row.worst_param_max_abs:.9f}" + ) + perf_rows = benchmark( + model, + x, + scenario_items, + args.loss, + args.warmup, + args.repeats, + args.rounds, + args.use_nvtx, + ) + for row in perf_rows: + print( + f"PERF name={row.name!r} mean_us={row.mean_us:.3f} " + f"median_us={row.median_us:.3f} min_us={row.min_us:.3f} " + f"max_us={row.max_us:.3f} speedup_vs_baseline={row.speedup:.3f}" + ) + if args.fail_on_accuracy and any(row.status != "PASS" for row in accuracy_rows): + raise SystemExit(1) + finally: + set_env({}) + Utils.destroy_model_parallel() + + +if __name__ == "__main__": + main() diff --git a/tests/unit_tests/ssm/test_bench_gdn_cuda_opt_scenarios.py b/tests/unit_tests/ssm/test_bench_gdn_cuda_opt_scenarios.py new file mode 100644 index 00000000000..0f57245ecdf --- /dev/null +++ b/tests/unit_tests/ssm/test_bench_gdn_cuda_opt_scenarios.py @@ -0,0 +1,56 @@ +import ast +from pathlib import Path + +BENCH = Path(__file__).with_name("bench_gdn_cuda_opt.py") + + +def _literal_assignment(name): + tree = ast.parse(BENCH.read_text()) + for node in tree.body: + if isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name) and target.id == name: + return ast.literal_eval(node.value) + raise AssertionError(f"{name} assignment not found") + + +def test_optimized_scenarios_route_through_mcore_wrapper(): + scenarios = _literal_assignment("SCENARIOS") + optimized = [ + "wy", + "dv_dhu", + "dhu", + "dqkwg", + "separate", + "dv_dhu_dqkwg", + "all_four", + "fwd_h_wy_dv_dhu_dqkwg", + ] + + for key in optimized: + env = scenarios[key][1] + assert env["MCORE_GDN_USE_OPT_WRAPPER"] == "1", key + assert env["MCORE_GDN_OPT_BACKEND"] == "cuda", key + assert not any(flag.startswith("FLA_CUTE_") for flag in env), key + + +def test_benchmark_does_not_require_patched_fla_sources(): + text = BENCH.read_text() + + assert "patched flash-linear-attention" not in text + assert "FLA_DISPATCH_SOURCE" not in text + + +def test_benchmark_does_not_expose_dhu_dqkwg_wrapper_path(): + scenarios = _literal_assignment("SCENARIOS") + flags = _literal_assignment("FLAGS") + forbidden = { + "MCORE_GDN_OPT_ENABLE_DHU_DQKWG", + "FLA_CUTE_BWD_DHU_DQKWG", + "FLA_CUTE_BWD_DHU_DQKWG_KERNEL", + "FLA_CUTE_BWD_DHU_DQKWG_DIRECT", + } + + assert forbidden.isdisjoint(flags) + for key, (_label, env) in scenarios.items(): + assert forbidden.isdisjoint(env), key diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index 3eb02442fe9..038a11d63f1 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -1,5 +1,6 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import copy from functools import partial from unittest import mock @@ -16,6 +17,7 @@ get_transformer_block_with_experimental_attention_variant_spec, ) from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.gated_delta_net import GatedDeltaNet from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed @@ -45,6 +47,49 @@ HAVE_FLA = False +def _make_gdn_config(**overrides): + config_kwargs = { + "hidden_size": 128, + "linear_conv_kernel_dim": 2, + "linear_key_head_dim": 32, + "linear_value_head_dim": 32, + "linear_num_key_heads": 4, + "linear_num_value_heads": 8, + "num_layers": 1, + "normalization": "RMSNorm", + "use_cpu_initialization": True, + "layernorm_zero_centered_gamma": True, + "num_attention_heads": 8, + "activation_func": F.silu, + "bf16": True, + "experimental_attention_variant": "gated_delta_net", + "linear_attention_freq": [1], + "transformer_impl": "transformer_engine", + } + config_kwargs.update(overrides) + return TransformerConfig(**config_kwargs) + + +@pytest.mark.parametrize("pre_gated_delta_rule_impl", ["unfused", "fused_streamed", "fused_mega"]) +def test_pre_gated_delta_rule_impl_accepts_gdn_modes(pre_gated_delta_rule_impl): + config = _make_gdn_config(pre_gated_delta_rule_impl=pre_gated_delta_rule_impl) + assert config.pre_gated_delta_rule_impl == pre_gated_delta_rule_impl + + +def test_pre_gated_delta_rule_impl_rejects_invalid_value(): + with pytest.raises(ValueError, match="pre_gated_delta_rule_impl must be one of"): + _make_gdn_config(pre_gated_delta_rule_impl="fused") + + +def test_pre_gated_delta_rule_impl_requires_gdn_variant(): + with pytest.raises(ValueError, match="experimental_attention_variant='gated_delta_net'"): + _make_gdn_config( + experimental_attention_variant=None, + linear_attention_freq=None, + pre_gated_delta_rule_impl="fused_streamed", + ) + + @pytest.mark.parametrize( ("tp_size", "sp", "cp_size"), [(1, False, 1), (2, False, 1), (2, True, 1), (1, False, 2), (2, False, 2), (2, True, 2)], @@ -142,6 +187,171 @@ def test_gpu_forward(self): output.dtype == hidden_states.dtype ), f"Output dtype {output.dtype=} mismatch with {hidden_states.dtype=}" + def test_selective_recompute_norm_out(self): + tp_group = parallel_state.get_tensor_model_parallel_group() + cp_group = parallel_state.get_context_parallel_group() + pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group) + + def build_gdn(config): + gdn_submodules = get_experimental_attention_variant_module_spec( + config=config + ).submodules + gdn = GatedDeltaNet( + config, + submodules=gdn_submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ) + return gdn.cuda().bfloat16() + + def run(gdn, hidden_states): + output, _ = gdn(hidden_states, None) + output.float().sum().backward() + grads = { + name: param.grad.detach() + for name, param in gdn.named_parameters() + if param.grad is not None + } + input_grad = hidden_states.grad.detach().clone() + return output.detach(), grads, input_grad + + micro_batch_size = 2 + seq_length = 64 + base_config = copy.deepcopy(self.transformer_config) + rec_config = copy.deepcopy(self.transformer_config) + rec_config.recompute_granularity = "selective" + rec_config.recompute_modules = ["gdn_norm_out"] + + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + hidden_states = torch.randn( + ( + seq_length // self.sp_size // self.cp_size, + micro_batch_size, + self.gdn.config.hidden_size, + ), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + requires_grad=True, + ) + + # --- Baseline (no recompute) --- + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + base_gdn = build_gdn(base_config) + assert base_gdn.recompute_norm_out is False + base_output, base_grads, base_input_grad = run(base_gdn, hidden_states) + hidden_states.grad = None + del base_gdn + torch.cuda.empty_cache() + + # --- Recompute --- + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + rec_gdn = build_gdn(rec_config) + assert rec_gdn.recompute_norm_out is True + rec_output, rec_grads, rec_input_grad = run(rec_gdn, hidden_states) + + rank = torch.distributed.get_rank() + assert torch.equal(rec_output, base_output), f"Output not identical ({rank=})" + assert torch.equal(rec_input_grad, base_input_grad), f"Input grad not identical ({rank=})" + assert set(rec_grads.keys()) == set(base_grads.keys()) + for name in base_grads: + assert torch.equal( + rec_grads[name], base_grads[name] + ), f"Grad not identical for {name} ({rank=})" + + def test_selective_recompute_gdn_qkv(self): + """gdn_qkv discard-output recompute must be numerically exact. + + recompute_modules=["gdn_qkv"] recomputes the whole QKV projection + + preparation block (in_proj -> CP a2a -> conv1d -> _prepare_qkv -> g/beta) + as a discard-output checkpoint. Output, parameter grads and input grad + must match the no-recompute baseline bit-for-bit. + """ + tp_group = parallel_state.get_tensor_model_parallel_group() + cp_group = parallel_state.get_context_parallel_group() + pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group) + + def build_gdn(config): + gdn_submodules = get_experimental_attention_variant_module_spec( + config=config + ).submodules + gdn = GatedDeltaNet( + config, + submodules=gdn_submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ) + return gdn.cuda().bfloat16() + + def run(gdn, hidden_states): + output, _ = gdn(hidden_states, None) + output.float().sum().backward() + grads = { + name: param.grad.detach() + for name, param in gdn.named_parameters() + if param.grad is not None + } + input_grad = hidden_states.grad.detach().clone() + return output.detach(), grads, input_grad + + micro_batch_size = 2 + seq_length = 64 + base_config = copy.deepcopy(self.transformer_config) + rec_config = copy.deepcopy(self.transformer_config) + rec_config.recompute_granularity = "selective" + rec_config.recompute_modules = ["gdn_qkv"] + + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + hidden_states = torch.randn( + ( + seq_length // self.sp_size // self.cp_size, + micro_batch_size, + self.gdn.config.hidden_size, + ), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + requires_grad=True, + ) + + # --- Baseline (no recompute) --- + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + base_gdn = build_gdn(base_config) + assert base_gdn.recompute_qkv is False + base_output, base_grads, base_input_grad = run(base_gdn, hidden_states) + hidden_states.grad = None + del base_gdn + torch.cuda.empty_cache() + + # --- Recompute --- + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + rec_gdn = build_gdn(rec_config) + assert rec_gdn.recompute_qkv is True + rec_output, rec_grads, rec_input_grad = run(rec_gdn, hidden_states) + + rank = torch.distributed.get_rank() + assert torch.equal(rec_output, base_output), f"Output not identical ({rank=})" + assert torch.equal(rec_input_grad, base_input_grad), f"Input grad not identical ({rank=})" + assert set(rec_grads.keys()) == set(base_grads.keys()) + for name in base_grads: + assert torch.equal( + rec_grads[name], base_grads[name] + ), f"Grad not identical for {name} ({rank=})" + def test_jit_compiled_helpers(self): import torch._dynamo @@ -307,6 +517,462 @@ def test_gpu_forward_thd_padding_correctness(self): self.gdn(hidden_states_thd, None, packed_seq_params=actual_mismatch_params) +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +@pytest.mark.internal +class TestFusedPreGatedDeltaRule: + + @pytest.fixture(scope='function', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=1, + ) + model_parallel_cuda_manual_seed(123) + + tp_group = parallel_state.get_tensor_model_parallel_group() + cp_group = parallel_state.get_context_parallel_group() + self.pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group) + + self.unfused_gdn = self._build_gdn(pre_gated_delta_rule_impl="unfused") + self.fused_gdn = self._build_gdn(pre_gated_delta_rule_impl="fused_streamed") + self.fused_gdn.load_state_dict(self.unfused_gdn.state_dict()) + + def teardown_method(self): + Utils.destroy_model_parallel() + + def _build_gdn( + self, + pre_gated_delta_rule_impl: str, + *, + deterministic_mode: bool = True, + conv_kernel_dim: int = 2, + ): + transformer_config = TransformerConfig( + hidden_size=256, + linear_conv_kernel_dim=conv_kernel_dim, + linear_key_head_dim=64, + linear_value_head_dim=64, + linear_num_key_heads=4, + linear_num_value_heads=8, + num_layers=1, + normalization="RMSNorm", + use_cpu_initialization=True, + layernorm_zero_centered_gamma=True, + num_attention_heads=8, + activation_func=F.silu, + bf16=True, + tensor_model_parallel_size=1, + context_parallel_size=1, + experimental_attention_variant="gated_delta_net", + linear_attention_freq=[1], + transformer_impl="transformer_engine", + deterministic_mode=deterministic_mode, + pre_gated_delta_rule_impl=pre_gated_delta_rule_impl, + ) + gdn_submodules = get_experimental_attention_variant_module_spec( + config=transformer_config + ).submodules + gdn = GatedDeltaNet( + transformer_config, + submodules=gdn_submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=self.pg_collection, + ) + return gdn.cuda().bfloat16() + + def _packed_pre_gated_delta_rule_reference(self, gdn, qkvzba, cu_seqlens): + """Run the dense torch reference independently on each packed sequence.""" + + segment_outputs = [[] for _ in range(6)] + for start, end in zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist()): + outputs = gdn.pre_gated_delta_rule(qkvzba[start:end], batch=1, seq_len=end - start) + for output_list, output in zip(segment_outputs, outputs): + output_list.append(output) + return tuple(torch.cat(outputs, dim=1) for outputs in segment_outputs) + + def _assert_pre_gated_delta_rule_outputs_close( + self, + fused_outputs, + unfused_outputs, + *, + atol: float, + rtol: float, + output_tolerances=None, + ): + """Compare named pre-GDR outputs with optional per-output tolerances.""" + + output_names = ("query", "key", "value", "gate", "beta", "g") + output_tolerances = output_tolerances or {} + for name, fused, unfused in zip(output_names, fused_outputs, unfused_outputs): + output_atol, output_rtol = output_tolerances.get(name, (atol, rtol)) + torch.testing.assert_close( + fused, + unfused, + atol=output_atol, + rtol=output_rtol, + msg=lambda msg, output_name=name: f"{output_name} mismatch: {msg}", + ) + + def test_fused_and_unfused_forward_match(self): + hidden_states = torch.randn( + (32, 2, self.unfused_gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + with torch.no_grad(): + unfused_output, unfused_bias = self.unfused_gdn(hidden_states, None) + fused_output, fused_bias = self.fused_gdn(hidden_states, None) + + torch.testing.assert_close(fused_output, unfused_output, atol=1e-3, rtol=1e-3) + assert fused_bias == unfused_bias + + @pytest.mark.parametrize("pre_gated_delta_rule_impl", ["fused_streamed", "fused_mega"]) + def test_fused_and_unfused_forward_thd_match(self, pre_gated_delta_rule_impl): + unfused_gdn = self._build_gdn( + pre_gated_delta_rule_impl="unfused", + deterministic_mode=False, + conv_kernel_dim=4, + ) + fused_gdn = self._build_gdn( + pre_gated_delta_rule_impl=pre_gated_delta_rule_impl, + deterministic_mode=False, + conv_kernel_dim=4, + ) + fused_gdn.load_state_dict(unfused_gdn.state_dict()) + + hidden_states = torch.randn( + (32, 1, unfused_gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + cu_seqlens = torch.tensor([0, 1, 4, 11, 32], device=torch.cuda.current_device(), dtype=torch.int32) + packed_seq_params = PackedSeqParams( + qkv_format='thd', + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=21, + max_seqlen_kv=21, + total_tokens=hidden_states.shape[0], + ) + assert packed_seq_params.seq_idx is not None + + with torch.no_grad(): + unfused_output, unfused_bias = unfused_gdn( + hidden_states, None, packed_seq_params=packed_seq_params + ) + fused_output, fused_bias = fused_gdn( + hidden_states, None, packed_seq_params=packed_seq_params + ) + + torch.testing.assert_close(fused_output, unfused_output, atol=2e-3, rtol=2e-3) + assert fused_bias == unfused_bias + + def test_fused_and_unfused_forward_thd_padding_match(self): + unfused_gdn = self._build_gdn( + pre_gated_delta_rule_impl="unfused", + deterministic_mode=False, + conv_kernel_dim=4, + ) + fused_gdn = self._build_gdn( + pre_gated_delta_rule_impl="fused_streamed", + deterministic_mode=False, + conv_kernel_dim=4, + ) + fused_gdn.load_state_dict(unfused_gdn.state_dict()) + + hidden_states = torch.randn( + (12, 1, unfused_gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + cu_seqlens = torch.tensor([0, 1, 4, 9], device=torch.cuda.current_device(), dtype=torch.int32) + cu_seqlens_padded = torch.tensor( + [0, 2, 6, 12], device=torch.cuda.current_device(), dtype=torch.int32 + ) + packed_seq_params = PackedSeqParams( + qkv_format='thd', + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=6, + max_seqlen_kv=6, + total_tokens=hidden_states.shape[0], + ) + assert packed_seq_params.seq_idx is not None + + with torch.no_grad(): + unfused_output, unfused_bias = unfused_gdn( + hidden_states, None, packed_seq_params=packed_seq_params + ) + fused_output, fused_bias = fused_gdn( + hidden_states, None, packed_seq_params=packed_seq_params + ) + + torch.testing.assert_close(fused_output, unfused_output, atol=2e-3, rtol=2e-3) + assert fused_bias == unfused_bias + + def test_fused_and_unfused_pre_gated_delta_rule_match(self): + batch = 2 + seq_len = 32 + hidden_states = torch.randn( + (seq_len, batch, self.unfused_gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + with torch.no_grad(): + qkvzba, _ = self.unfused_gdn.in_proj(hidden_states) + unfused_outputs = self.unfused_gdn.pre_gated_delta_rule(qkvzba, batch, seq_len) + fused_outputs = self.fused_gdn._fused_streamed_pre_gated_delta_rule(qkvzba) + + self._assert_pre_gated_delta_rule_outputs_close( + fused_outputs, + unfused_outputs, + atol=1e-3, + rtol=1e-3, + # g uses Triton exp/log softplus in the fused path and torch softplus + # in the reference path, so its direct intermediate parity needs a + # slightly looser relative tolerance than the layout/conv outputs. + output_tolerances={"g": (1e-3, 3e-3)}, + ) + + def test_fused_and_unfused_packed_pre_gated_delta_rule_forward_match(self): + reference_gdn = self._build_gdn( + pre_gated_delta_rule_impl="unfused", + deterministic_mode=True, + conv_kernel_dim=4, + ) + fused_gdn = self._build_gdn( + pre_gated_delta_rule_impl="fused_streamed", + deterministic_mode=False, + conv_kernel_dim=4, + ) + fused_gdn.load_state_dict(reference_gdn.state_dict()) + + batch = 1 + cu_seqlens = torch.tensor([0, 1, 4, 6, 11], device=torch.cuda.current_device(), dtype=torch.int32) + seq_len = cu_seqlens[-1].item() + qkvzba = torch.randn( + (seq_len, batch, reference_gdn.in_proj_dim), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + with torch.no_grad(): + unfused_outputs = self._packed_pre_gated_delta_rule_reference( + reference_gdn, qkvzba, cu_seqlens + ) + fused_outputs = fused_gdn._fused_streamed_pre_gated_delta_rule( + qkvzba, cu_seqlens_q=cu_seqlens + ) + + self._assert_pre_gated_delta_rule_outputs_close( + fused_outputs, unfused_outputs, atol=2e-3, rtol=2e-3 + ) + + def test_fused_and_unfused_packed_pre_gated_delta_rule_backward_match(self): + reference_gdn = self._build_gdn( + pre_gated_delta_rule_impl="unfused", + deterministic_mode=True, + conv_kernel_dim=4, + ) + fused_gdn = self._build_gdn( + pre_gated_delta_rule_impl="fused_streamed", + deterministic_mode=False, + conv_kernel_dim=4, + ) + fused_gdn.load_state_dict(reference_gdn.state_dict()) + + batch = 1 + cu_seqlens = torch.tensor([0, 1, 4, 6, 11], device=torch.cuda.current_device(), dtype=torch.int32) + seq_len = cu_seqlens[-1].item() + qkvzba = torch.randn( + (seq_len, batch, reference_gdn.in_proj_dim), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + qkvzba_unfused = qkvzba.detach().clone().requires_grad_(True) + qkvzba_fused = qkvzba.detach().clone().requires_grad_(True) + + reference_gdn.zero_grad(set_to_none=True) + fused_gdn.zero_grad(set_to_none=True) + + unfused_outputs = self._packed_pre_gated_delta_rule_reference( + reference_gdn, qkvzba_unfused, cu_seqlens + ) + fused_outputs = fused_gdn._fused_streamed_pre_gated_delta_rule( + qkvzba_fused, cu_seqlens_q=cu_seqlens + ) + grad_outputs = [torch.randn_like(output.float()) for output in unfused_outputs] + + unfused_loss = sum( + (output.float() * grad).sum() for output, grad in zip(unfused_outputs, grad_outputs) + ) + fused_loss = sum( + (output.float() * grad).sum() for output, grad in zip(fused_outputs, grad_outputs) + ) + unfused_loss.backward() + fused_loss.backward() + + torch.testing.assert_close(qkvzba_fused.grad, qkvzba_unfused.grad, atol=3e-2, rtol=3e-2) + torch.testing.assert_close( + fused_gdn.conv1d.weight.grad, + reference_gdn.conv1d.weight.grad, + atol=3e-2, + rtol=3e-2, + ) + torch.testing.assert_close(fused_gdn.A_log.grad, reference_gdn.A_log.grad, atol=3e-2, rtol=3e-2) + torch.testing.assert_close( + fused_gdn.dt_bias.grad, reference_gdn.dt_bias.grad, atol=3e-2, rtol=3e-2 + ) + + def test_fused_packed_conv_forward_boundary_isolation(self): + from megatron.core.fusions.fused_pre_gated_delta_rule import ( + fused_streamed_pre_gated_delta_rule, + ) + + seq_len = 5 + boundary = 3 + num_key_heads = 1 + # Keep qkvzba.stride(0) aligned for causal_conv1d's channel-last + # backward guard; the boundary condition under test is independent + # of the value-head repeat factor. + num_value_heads = 4 + key_head_dim = 32 + value_head_dim = 32 + conv_width = 4 + qk_channels = num_key_heads * key_head_dim + v_channels = num_value_heads * value_head_dim + v_offset = 2 * qk_channels + k_offset = qk_channels + total_channels = 2 * qk_channels + 2 * v_channels + 2 * num_value_heads + device = torch.cuda.current_device() + + qkvzba = torch.zeros((seq_len, 1, total_channels), device=device, dtype=torch.bfloat16) + qkvzba[boundary - 1, 0, :qk_channels] = 10.0 + qkvzba[boundary - 1, 0, k_offset : k_offset + qk_channels] = 10.0 + qkvzba[boundary - 1, 0, v_offset : v_offset + v_channels] = 10.0 + conv_weight = torch.zeros((2 * qk_channels + v_channels, 1, conv_width), device=device) + conv_weight[:qk_channels, 0, conv_width - 2] = 1.0 + conv_weight[k_offset : k_offset + qk_channels, 0, conv_width - 2] = 1.0 + conv_weight[v_offset : v_offset + v_channels, 0, conv_width - 2] = 1.0 + A_log = torch.zeros((num_value_heads,), device=device, dtype=torch.bfloat16) + dt_bias = torch.zeros((num_value_heads,), device=device, dtype=torch.bfloat16) + cu_seqlens = torch.tensor([0, boundary, seq_len], device=device, dtype=torch.int32) + + query, key, value, _, _, _ = fused_streamed_pre_gated_delta_rule( + qkvzba, + conv_weight.to(torch.bfloat16), + None, + A_log, + dt_bias, + num_key_heads=num_key_heads, + num_value_heads=num_value_heads, + key_head_dim=key_head_dim, + value_head_dim=value_head_dim, + cu_seqlens=cu_seqlens, + ) + + torch.testing.assert_close( + query[0, boundary], + torch.zeros_like(query[0, boundary]), + atol=0.0, + rtol=0.0, + ) + torch.testing.assert_close( + key[0, boundary], + torch.zeros_like(key[0, boundary]), + atol=0.0, + rtol=0.0, + ) + torch.testing.assert_close( + value[0, boundary], + torch.zeros_like(value[0, boundary]), + atol=0.0, + rtol=0.0, + ) + + def test_fused_packed_conv_backward_boundary_isolation(self): + from megatron.core.fusions.fused_pre_gated_delta_rule import ( + fused_streamed_pre_gated_delta_rule, + ) + + seq_len = 5 + boundary = 3 + num_key_heads = 1 + # Keep qkvzba.stride(0) aligned for causal_conv1d's channel-last + # backward guard; the boundary condition under test is independent + # of the value-head repeat factor. + num_value_heads = 4 + key_head_dim = 32 + value_head_dim = 32 + conv_width = 4 + qk_channels = num_key_heads * key_head_dim + v_channels = num_value_heads * value_head_dim + v_offset = 2 * qk_channels + k_offset = qk_channels + total_channels = 2 * qk_channels + 2 * v_channels + 2 * num_value_heads + device = torch.cuda.current_device() + + qkvzba = torch.zeros( + (seq_len, 1, total_channels), device=device, dtype=torch.bfloat16, requires_grad=True + ) + conv_weight = torch.zeros( + (2 * qk_channels + v_channels, 1, conv_width), + device=device, + dtype=torch.bfloat16, + requires_grad=True, + ) + with torch.no_grad(): + conv_weight[:qk_channels, 0, conv_width - 2] = 1.0 + conv_weight[k_offset : k_offset + qk_channels, 0, conv_width - 2] = 1.0 + conv_weight[v_offset : v_offset + v_channels, 0, conv_width - 2] = 1.0 + A_log = torch.zeros((num_value_heads,), device=device, dtype=torch.bfloat16, requires_grad=True) + dt_bias = torch.zeros( + (num_value_heads,), device=device, dtype=torch.bfloat16, requires_grad=True + ) + cu_seqlens = torch.tensor([0, boundary, seq_len], device=device, dtype=torch.int32) + + query, key, value, gate, beta, g = fused_streamed_pre_gated_delta_rule( + qkvzba, + conv_weight, + None, + A_log, + dt_bias, + num_key_heads=num_key_heads, + num_value_heads=num_value_heads, + key_head_dim=key_head_dim, + value_head_dim=value_head_dim, + cu_seqlens=cu_seqlens, + ) + + loss = ( + query[0, boundary].float().sum() + + key[0, boundary].float().sum() + + value[0, boundary].float().sum() + ) + loss = loss + 0.0 * ( + gate.float().sum() + + beta.float().sum() + + g.float().sum() + ) + loss.backward() + leaked_q_grad = qkvzba.grad[boundary - 1, 0, :qk_channels] + leaked_k_grad = qkvzba.grad[boundary - 1, 0, k_offset : k_offset + qk_channels] + leaked_grad = qkvzba.grad[boundary - 1, 0, v_offset : v_offset + v_channels] + torch.testing.assert_close(leaked_q_grad, torch.zeros_like(leaked_q_grad), atol=0.0, rtol=0.0) + torch.testing.assert_close(leaked_k_grad, torch.zeros_like(leaked_k_grad), atol=0.0, rtol=0.0) + torch.testing.assert_close(leaked_grad, torch.zeros_like(leaked_grad), atol=0.0, rtol=0.0) + + @pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") @pytest.mark.internal class TestGDNCuSeqlensResolve: diff --git a/tests/unit_tests/ssm/test_gated_delta_net_cuda_opt.py b/tests/unit_tests/ssm/test_gated_delta_net_cuda_opt.py new file mode 100644 index 00000000000..5eb354bddc1 --- /dev/null +++ b/tests/unit_tests/ssm/test_gated_delta_net_cuda_opt.py @@ -0,0 +1,88 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Focused GatedDeltaNet CUDA optimization coverage. + +This keeps the optimized-kernel correctness and optional perf check separate +from the generic GatedDeltaNet unit tests. +""" + +import os + +import pytest +import torch + +from tests.unit_tests.ssm import bench_gdn_cuda_opt as runner + +try: + import fla # noqa: F401 + + HAVE_FLA = True +except ImportError: + HAVE_FLA = False + + +def _scenario_items(): + keys = [ + key.strip() + for key in os.environ.get( + "MCORE_GDN_UNIT_TEST_SCENARIOS", "baseline,fwd_h_wy_dv_dhu_dqkwg" + ).split(",") + if key.strip() + ] + if "baseline" not in keys: + keys.insert(0, "baseline") + unknown = [key for key in keys if key not in runner.SCENARIOS] + if unknown: + raise ValueError(f"unknown GDN CUDA opt scenarios: {unknown}") + return [(key, runner.SCENARIOS[key]) for key in keys] + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available.") +@pytest.mark.internal +def test_gated_delta_net_cuda_opt_correctness_and_optional_perf(dtype): + scenario_items = _scenario_items() + runner.validate_dispatch_sources(scenario_items) + + torch.manual_seed(123) + runner.set_env({}) + try: + model = runner.make_model(dtype).eval() + seq_len = int(os.environ.get("MCORE_GDN_UNIT_TEST_T", "8192")) + batch = int(os.environ.get("MCORE_GDN_UNIT_TEST_B", "2")) + x = torch.randn(seq_len, batch, 128, device="cuda", dtype=dtype) + + atol = float(os.environ.get("MCORE_GDN_UNIT_TEST_ATOL", "5e-3")) + rtol = float(os.environ.get("MCORE_GDN_UNIT_TEST_RTOL", "5e-3")) + loss = os.environ.get("MCORE_GDN_UNIT_TEST_LOSS", "sum") + accuracy_rows = runner.check_accuracy( + model, x, scenario_items, loss=loss, atol=atol, rtol=rtol, use_nvtx=False + ) + failed = [row for row in accuracy_rows if row.status != "PASS"] + assert not failed, "\n".join( + f"{row.name}: output={row.output_max_abs:.9f} " + f"input_grad={row.input_grad_max_abs:.9f} " + f"{row.worst_param}={row.worst_param_max_abs:.9f}" + for row in failed + ) + + if os.environ.get("MCORE_GDN_UNIT_TEST_PERF", "0") == "1": + perf_rows = runner.benchmark( + model, + x, + scenario_items, + loss=loss, + warmup=int(os.environ.get("MCORE_GDN_UNIT_TEST_WARMUP", "5")), + repeats=int(os.environ.get("MCORE_GDN_UNIT_TEST_REPEATS", "20")), + rounds=int(os.environ.get("MCORE_GDN_UNIT_TEST_ROUNDS", "3")), + use_nvtx=True, + ) + for row in perf_rows: + print( + f"PERF {row.name}: mean_us={row.mean_us:.3f} " + f"speedup_vs_baseline={row.speedup:.3f}" + ) + finally: + runner.set_env({}) + runner.Utils.destroy_model_parallel()