diff --git a/docs/index.md b/docs/index.md index 11337315588..133317dc449 100644 --- a/docs/index.md +++ b/docs/index.md @@ -67,6 +67,7 @@ models/index user-guide/features/moe user-guide/features/context_parallel +user-guide/features/mamba_state_passing_context_parallel user-guide/features/megatron_fsdp user-guide/features/dist_optimizer user-guide/features/optimizer_cpu_offload diff --git a/docs/user-guide/features/index.md b/docs/user-guide/features/index.md index cb2e895afdc..ce0fb95938b 100644 --- a/docs/user-guide/features/index.md +++ b/docs/user-guide/features/index.md @@ -18,6 +18,7 @@ cuda_graph fine_grained_activation_offloading moe context_parallel +mamba_state_passing_context_parallel megatron_fsdp dist_optimizer optimizer_cpu_offload diff --git a/docs/user-guide/features/mamba_state_passing_context_parallel.md b/docs/user-guide/features/mamba_state_passing_context_parallel.md new file mode 100644 index 00000000000..f57622ea049 --- /dev/null +++ b/docs/user-guide/features/mamba_state_passing_context_parallel.md @@ -0,0 +1,279 @@ + + +# Mamba2 State-Passing Context Parallelism + +Sequence-sharded, state-passing context parallelism (CP) for Mamba2 training. +The production implementation lives in +[`megatron/core/ssm/ops/ssd_state_passing_cp.py`](../../../megatron/core/ssm/ops/ssd_state_passing_cp.py) +and is opt-in; the existing all-to-all Mamba CP path remains the default. + +## Overview + +The existing Mamba CP path redistributes activations with an all-to-all (A2A) so +that every rank holds the full sequence for a subset of heads and groups. +State-passing CP instead keeps the sequence shard local and communicates only +what the causal boundary requires: the convolution halo and the SSD state +summary. + +| Path | Sequence handling | Dominant communication | +|---|---|---| +| No CP | one rank processes the whole sequence | none | +| Mamba A2A CP | full sequence, head/group shard | activation A2A proportional to sequence length | +| State-passing CP | local sequence shard is kept | conv halo and SSD boundary all-gather | + +Approximate rank-local payloads: + +```text +A2A: + O(local_sequence * hidden) + +State-passing boundary: + Conv: O(batch * channels * (d_conv - 1)) + SSD : O(batch * heads * headdim * dstate) + O(batch * heads) for decay +``` + +The boundary collective is an all-gather. Its size does not scale with sequence +length, but the received volume grows with the CP size. + +## Configuration + +- `--use-mamba-state-passing-cp` enables the path. +- `--mamba-state-passing-cp-load-balancing {none,permute_p2p,permute_a2a,virtual}` + selects how the standard balanced CP layout is handled. + +`none` expects an already-contiguous causal shard and is only valid for direct +calls into the state-passing kernels. Standard Megatron CP hands the mixer a +front/back balanced layout, so `MambaMixer` requires `permute_p2p`, +`permute_a2a`, or `virtual`. + +`MambaMixer` dispatch: + +```text +hidden_states + -> in_proj -> zxBCdt + -> MambaMixer._use_mamba_state_passing_cp() + | + +-- cp_size == 1 + | -> standard fused Mamba + | + +-- state-passing OFF + | -> MambaContextParallel.pre_conv_ssm() + | -> fused Mamba + | -> MambaContextParallel.post_conv_ssm() + | + `-- state-passing ON + -> MambaStatePassingCPAdapter.forward() + -> out_proj +``` + +## Load-balancing modes + +Megatron's balanced CP assigns one front chunk and one back chunk to each rank: + +```text +CP=3 balanced ownership: + rank 0: [1, 6] + rank 1: [2, 5] + rank 2: [3, 4] + +causal order: + [1, 2, 3, 4, 5, 6] +``` + +### Permute + +`permute_p2p` and `permute_a2a` physically exchange the balanced activations for +a contiguous causal shard: + +```text +balanced: + rank 0 [1,6], rank 1 [2,5], rank 2 [3,4] + +contiguous: + rank 0 [1,2], rank 1 [3,4], rank 2 [5,6] +``` + +`undo_state_passing_cp_load_balancing()` runs on the forward input and +`redo_state_passing_cp_load_balancing()` on the output. A custom autograd +function restores gradient ownership with the inverse permutation in the +backward pass. + +- `permute_p2p` sends and receives the required remote chunks point-to-point. +- `permute_a2a` exchanges them in one unequal-split `all_to_all_single`. + +Both backends produce the same layout and the same mathematics; only the +communication implementation differs. + +### Relation to `context_parallel_layout` + +Megatron Core already models the balanced/contiguous distinction as +`CpPartitionMode = Literal["zigzag", "contiguous"]` in +`megatron.core.context_parallel_layout`, and Gated DeltaNet uses it at its module +entry point through `convert_module_input_tensors_cp_partition_mode()` when it +runs its chunkwise (sequence-sharded) CP path. `permute_p2p` and `permute_a2a` +perform the same conversion, so the two overlap. + +They are implemented separately here for two reasons: the permutation is driven +from inside the fused Conv+SSD autograd function rather than at the module entry +point, and `permute_p2p` adds a point-to-point backend that the shared helper +does not provide. Consolidating onto the shared helper is intended follow-up +work. + +### Virtual + +`virtual` moves no activations. Each physical rank's front and back chunks are +interpreted as two independent virtual causal segments, and the input is viewed +as an interleaved virtual batch: + +```text +[local_L, batch, ...] -> [2 * batch, local_L / 2, ...] +``` + +Boundary routing then selects the predecessor and successor virtual segment out +of the rank-major gather result. + +## Production flow + +The path fuses the convolution and the SSD scan into a single custom autograd +function, `MambaSplitConv1dScanCombinedStatePassingCPFn`. Its public API keeps +the upstream `MambaSplitConv1dScanCombinedFn` argument order and appends +`state_passing_cp_group` and `state_passing_cp_virtual`. + +### Forward + +```text +zxBCdt [local_L, batch, packed] + -> load-balancing handling + -> virtual: view as [2 * batch, local_L / 2, packed] + -> split z / xBC / dt + +Conv + -> all-gather local tail + -> select predecessor tail + -> causal_conv1d(initial_states=predecessor_tail) + -> split x / B / C + +SSD + -> dt cumsum and local chunk contribution + -> packed summary [S_ext, a_block] + -> async all-gather(summary) || local CB computation + -> exclusive causal boundary scan + -> S_in + -> local state passing and chunk scan + +Output + -> virtual unpack or permutation redo + -> [local_L, batch, d_inner] +``` + +The SSD transform of one causal segment is affine: + +```text +S_out = a_block * S_in + S_ext +``` + +- `S_ext` is the final state the local segment produces from a zero initial state. +- `a_block` is the per-head decay across the whole segment. +- `S_in` is the exclusive prefix composition of the preceding segment summaries. + +`_state_passing_summary_fwd_kernel()` writes `S_ext` and the decay straight into +an FP32 packed buffer. `_state_passing_boundary_scan_kernel()` handles both the +forward and the backward causal order through the compile-time `VIRTUAL_CP` and +`REVERSE` flags. + +### Convolution boundary + +The causal convolution needs the last `d_conv - 1` tokens of the preceding +segment: + +```text +forward: + local tail -> all-gather -> predecessor tail + -> causal_conv1d_fwd_function(initial_states=...) + +backward: + causal_conv1d_bwd_function(return_dinitial_states=True) + -> all-gather d(initial_state) + -> select successor gradient + -> in-place add into the local dx tail +``` + +### Backward + +```text +re-split xBC from zxbcdt + -> recompute the causal conv forward + -> recompute dt cumsum, CB, chunk contribution, recurrence states + -> local chunk dstates + -> local dS_in summary + -> async all-gather(dS_in) || local dC / dCB / ddA computation + -> exclusive reverse boundary scan + -> dS_ext from the successor + -> one _state_passing_bwd(dfinal_states=dS_ext) + -> dx / dB / dC / ddt / dA / dD / dz + -> causal conv backward + -> apply the reverse conv boundary gradient to the dx tail +``` + +The post-conv gradient the SSD backward produces is written into a preallocated +`dzxBCdt` view and consumed directly by the convolution backward; there is no +separate autograd accumulation to re-join the split gradients. + +### Saved tensors and dtypes + +The forward context keeps the upstream tensor order and appends +`state_passing_conv_initial_states`, `state_passing_initial_states`, and +`state_passing_gathered_decays`. Quantities that scale with sequence length +(`xBC_conv`, the split `x`/`B`/`C`, `dt_proc`, the chunk states, and `CB`) are +recomputed in the backward pass instead of being saved. + +| Value | dtype | +|---|---| +| `S_ext`, `a_block` | FP32 | +| `S_in`, recurrence states | FP32 | +| `dS_in`, `dS_ext` | FP32 | +| scan/BMM boundary tensors | upstream input dtype | + +## Supported scope + +Supported: + +- the training memory-efficient path +- fixed-length, non-packed sequences +- the standard balanced CP input layout +- `permute_p2p`, `permute_a2a`, and `virtual` +- BF16 activations with FP32 boundary state + +Not supported yet: + +- inference +- `seq_idx` and packed sequences +- hybrid, dynamic, or otherwise variable-length CP +- an externally supplied SSD `initial_states` +- `--mamba-training-ssm-states-dtype` +- fused RMSNorm and output projection inside the custom function + +The mixer adapter runs RMSNorm and the output projection outside the custom +function, so both remain available on the full `MambaMixer` training path. + +Shape constraints: each virtual segment length must be a multiple of the SSD +`chunk_size`, and each convolution segment must hold at least `d_conv - 1` +tokens. + +## Tests + +- [`tests/unit_tests/ssm/ops/test_ssd_state_passing_cp.py`](../../../tests/unit_tests/ssm/ops/test_ssd_state_passing_cp.py) + checks the load-balancing permutation for exactness and the convolution and + SSD kernels against a full-sequence reference. +- [`tests/unit_tests/ssm/test_mamba_mixer_state_passing_cp.py`](../../../tests/unit_tests/ssm/test_mamba_mixer_state_passing_cp.py) + checks every mode end to end against a full-sequence mixer and against the + A2A CP path. +- [`tests/unit_tests/ssm/test_mamba_state_passing_cp_cuda_graph.py`](../../../tests/unit_tests/ssm/test_mamba_state_passing_cp_cuda_graph.py) + covers the Megatron local and Transformer Engine CUDA Graph backends. diff --git a/examples/mamba_state_passing_context_parallel/README.md b/examples/mamba_state_passing_context_parallel/README.md new file mode 100644 index 00000000000..a1986370bec --- /dev/null +++ b/examples/mamba_state_passing_context_parallel/README.md @@ -0,0 +1,88 @@ + + +# Mamba2 State-Passing Context Parallelism Benchmark + +Compares the Mamba2 context-parallel paths against each other so the +`--mamba-state-passing-cp-load-balancing` mode can be chosen from measurements +rather than by guessing. See +[the feature guide](../../docs/user-guide/features/mamba_state_passing_context_parallel.md) +for what the paths do. + +## What is measured + +Only the post-projection part of `MambaMixer` — the causal convolution and the +SSD scan — is timed. That is where the CP paths differ; `in_proj`, the output +projection, and RMSNorm are identical across paths and would dilute the +comparison, so they are excluded. + +Four paths are measured: + +| Path | Description | +|---|---| +| `a2a` | The existing Mamba CP path: activations are redistributed so every rank holds the whole sequence for a subset of heads and groups | +| `permute-p2p` | State passing; the balanced shard is exchanged for a contiguous causal shard point-to-point | +| `permute-a2a` | State passing; the same exchange in one unequal-split `all_to_all_single` | +| `virtual` | State passing with no activation exchange; each balanced half is treated as an independent causal segment | + +## Running it + +One process per GPU; the world size is the CP size. No tensor-parallel group is +created — `--tp-size` only divides the head and group counts so that a TP-local +shard shape can be measured with one rank per CP position. + +```bash +# CP=4, 32K sequence +PYTHONPATH=. torchrun --standalone --nproc_per_node=4 \ + examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py \ + --L 32768 --iters 20 + +# Sweep sequence lengths and batch sizes in one run +PYTHONPATH=. torchrun --standalone --nproc_per_node=8 \ + examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py \ + --sequence-lengths 32768 131072 --batch-sizes 1 3 + +# Forward only, and just the two paths worth comparing for a decision +PYTHONPATH=. torchrun --standalone --nproc_per_node=4 \ + examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py \ + --forward-only --paths a2a virtual + +# TP2-local shard shape under CP4 +PYTHONPATH=. torchrun --standalone --nproc_per_node=4 \ + examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py \ + --tp-size 2 --L 32768 --batch 3 +``` + +Defaults are Nemotron-3 Nano's mixer shape: 64 heads of 64, 8 groups, state 128, +SSD chunk 128, `d_conv` 4, BF16 activations. + +## Timing method + +Each iteration is timed with CUDA events and reduced across ranks with `MAX`, +because a CP path is only as fast as its slowest rank. Rank alignment happens on +a barrier outside the timed interval, so waiting for stragglers is not counted +twice. Warmup iterations are discarded, and peak allocated memory is reported +alongside latency. + +Output is one human-readable line per path plus a machine-readable `RESULT,...` +line per path for scripted collection: + +```text +Mamba CP paths (fwd+bwd) L=32768 batch=1 local_L=8192 cp=4 ... + a2a: 12.34 ms +/- 0.12 p50= 12.30 peak= 1.23 GiB + permute-p2p: 9.87 ms +/- 0.09 p50= 9.85 peak= 0.98 GiB +p50 speedup vs a2a: permute-p2p=1.249x, ... +RESULT,cp=4,L=32768,batch=1,tp_size=1,path=a2a,mean_ms=...,p50_ms=...,peak_gib=... +``` + +## Constraints + +The sequence length must divide the CP size, each rank's local length must be +even (balanced CP gives each rank two chunks), and each local half must be a +multiple of the SSD chunk size. The script asserts all three before measuring. diff --git a/examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py b/examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py new file mode 100644 index 00000000000..e842de89af8 --- /dev/null +++ b/examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py @@ -0,0 +1,415 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Benchmark the Mamba2 context-parallel paths against each other. + +Measures the post-projection part of ``MambaMixer`` -- the causal convolution +and the SSD scan, which is where the CP paths differ -- for the existing +all-to-all path and for each state-passing load-balancing mode. The projection +layers are excluded on purpose: they are identical across paths and would dilute +the comparison. + +Run from the repository root, one process per GPU: + + PYTHONPATH=. torchrun --standalone --nproc_per_node=4 \\ + examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py \\ + --L 32768 --iters 20 + +Sweep several shapes in one run: + + PYTHONPATH=. torchrun --standalone --nproc_per_node=8 \\ + examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py \\ + --sequence-lengths 32768 131072 --batch-sizes 1 3 + +The world size is the CP size; no tensor-parallel group is created. ``--tp-size`` +only divides the head and group counts so a TP-local shard shape can be measured +on a single rank per CP position. +""" + +import argparse +import statistics +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Callable, Dict + +import torch +import torch.distributed as dist +import torch.nn as nn +from einops import rearrange +from mamba_ssm.ops.triton.ssd_combined import mamba_split_conv1d_scan_combined + +from megatron.core.ssm.mamba_context_parallel import MambaContextParallel +from megatron.core.ssm.ops.ssd_state_passing_cp import MambaStatePassingCPAdapter + +CP_PATH_NAMES = ("a2a", "permute-p2p", "permute-a2a", "virtual") + + +@dataclass(frozen=True) +class MambaCoreShape: + """TP-local Mamba mixer shape driving the measured kernels.""" + + nheads: int = 64 + headdim: int = 64 + ngroups: int = 8 + d_state: int = 128 + d_conv: int = 4 + chunk_size: int = 128 + + @property + def d_inner(self) -> int: + """Inner (gated) width of the mixer.""" + return self.nheads * self.headdim + + @property + def projected_width(self) -> int: + """Width of the ``in_proj`` output (z, x, B, C, dt).""" + return 2 * self.d_inner + 2 * self.ngroups * self.d_state + self.nheads + + @property + def conv_dim(self) -> int: + """Number of channels the depthwise causal convolution operates on.""" + return self.d_inner + 2 * self.ngroups * self.d_state + + +@dataclass(frozen=True) +class BenchmarkMetrics: + """Per-iteration latency statistics and peak memory for one CP path.""" + + mean_ms: float + stdev_ms: float + p50_ms: float + peak_gib: float + + +@dataclass +class MambaCPInputs: + """Rank-local activation and output gradient for one measured configuration.""" + + projected_input: torch.Tensor + grad_output: torch.Tensor + + +class MambaCPWorkload(nn.Module): + """Minimal post-projection mixer satisfying the contract both CP paths expect. + + ``MambaContextParallel`` and ``MambaStatePassingCPAdapter`` read a specific + set of attributes off the mixer. Reproducing just those keeps the benchmark + independent of the surrounding model, spec, and process-group plumbing. + """ + + def __init__(self, shape: MambaCoreShape, device: torch.device, cp_group: dist.ProcessGroup): + super().__init__() + self.shape = shape + dtype = torch.bfloat16 + generator = torch.Generator(device=device).manual_seed(4321) + + self.nheads_local_tp = shape.nheads + self.d_inner_local_tp = shape.d_inner + self.ngroups_local_tp = shape.ngroups + self.d_state = shape.d_state + self.d_conv = shape.d_conv + self.headdim = shape.headdim + self.chunk_size = shape.chunk_size + self.activation = "silu" + self.act = nn.SiLU() + # RMSNorm and the output projection are outside the measured region. + self.rmsnorm = False + self.norm_before_gate = False + self.D_has_hdim = False + self.config = SimpleNamespace(mamba_state_passing_cp_load_balancing="permute_p2p") + + def randn(*sizes, dtype=dtype, scale=1.0): + return torch.randn(*sizes, device=device, dtype=dtype, generator=generator) * scale + + self.conv1d_weight = nn.Parameter(randn(shape.conv_dim, 1, self.d_conv, scale=0.02)) + self.conv1d_bias = nn.Parameter(randn(shape.conv_dim, scale=0.02)) + self.dt_bias = nn.Parameter(randn(self.nheads_local_tp)) + self.A_log = nn.Parameter(randn(self.nheads_local_tp, dtype=torch.float32)) + self.D = nn.Parameter(randn(self.nheads_local_tp, dtype=torch.float32)) + + self.cp = MambaContextParallel( + cp_group=cp_group, + d_inner_local_tp=self.d_inner_local_tp, + nheads_local_tp=self.nheads_local_tp, + ngroups_local_tp=self.ngroups_local_tp, + d_state=self.d_state, + conv1d_weight_cp1=self.conv1d_weight, + conv1d_bias_cp1=self.conv1d_bias, + conv1d_padding=self.d_conv - 1, + dt_bias_cp1=self.dt_bias, + A_log_cp1=self.A_log, + D_cp1=self.D, + D_has_hdim=self.D_has_hdim, + ) + self.state_passing_cp_adapter = MambaStatePassingCPAdapter(self) + + +def parse_args(): + """Parse benchmark shape, sweep, and iteration-count arguments.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--L", dest="sequence_length", type=int, default=32768) + parser.add_argument("--batch", dest="batch_size", type=int, default=1) + parser.add_argument("--sequence-lengths", type=int, nargs="+") + parser.add_argument("--batch-sizes", type=int, nargs="+") + parser.add_argument("--nheads", type=int, default=64) + parser.add_argument("--headdim", type=int, default=64) + parser.add_argument("--ngroups", type=int, default=8) + parser.add_argument("--d-state", type=int, default=128) + parser.add_argument("--d-conv", type=int, default=4) + parser.add_argument("--chunk", type=int, default=128) + parser.add_argument( + "--tp-size", + type=int, + default=1, + help=( + "emulate a tensor-parallel local Mamba shard by dividing --nheads and " + "--ngroups; no TP process group or TP communication is created" + ), + ) + parser.add_argument("--iters", type=int, default=20) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--forward-only", action="store_true") + parser.add_argument("--paths", nargs="+", choices=CP_PATH_NAMES, default=CP_PATH_NAMES) + return parser.parse_args() + + +def shape_from_args(args) -> MambaCoreShape: + """Build the TP-local shape the benchmark measures.""" + assert args.tp_size > 0 + assert args.nheads % args.tp_size == 0, "global nheads must divide TP size" + assert args.ngroups % args.tp_size == 0, "global ngroups must divide TP size" + return MambaCoreShape( + nheads=args.nheads // args.tp_size, + headdim=args.headdim, + ngroups=args.ngroups // args.tp_size, + d_state=args.d_state, + d_conv=args.d_conv, + chunk_size=args.chunk, + ) + + +def validate_shape(shape: MambaCoreShape, sequence_length: int, cp_size: int) -> int: + """Check the CP and SSD alignment constraints and return the local length.""" + assert shape.nheads % shape.ngroups == 0 + assert sequence_length % cp_size == 0, "global sequence length must divide CP size" + local_length = sequence_length // cp_size + assert local_length % 2 == 0, "load-balanced CP requires two local chunks" + assert ( + local_length // 2 + ) % shape.chunk_size == 0, "each local load-balanced chunk must align with SSD chunk size" + return local_length + + +def create_inputs( + workload: MambaCPWorkload, + sequence_length: int, + batch_size: int, + rank: int, + device: torch.device, + cp_size: int, +) -> MambaCPInputs: + """Allocate the rank-local projected activation and output gradient.""" + local_length = validate_shape(workload.shape, sequence_length, cp_size) + generator = torch.Generator(device=device).manual_seed(1234 + rank) + projected_input = torch.randn( + local_length, + batch_size, + workload.shape.projected_width, + device=device, + dtype=torch.bfloat16, + generator=generator, + requires_grad=True, + ) + grad_output = torch.randn( + local_length, + batch_size, + workload.d_inner_local_tp, + device=device, + dtype=torch.bfloat16, + generator=generator, + ) + return MambaCPInputs(projected_input=projected_input, grad_output=grad_output) + + +def build_cp_path_runners( + workload: MambaCPWorkload, inputs: MambaCPInputs +) -> Dict[str, Callable[[], torch.Tensor]]: + """Return one output-producing closure per measured CP path.""" + + def a2a(): + projected = workload.cp.pre_conv_ssm(inputs.projected_input) + projected = rearrange(projected, "l b d -> b l d").contiguous() + output = mamba_split_conv1d_scan_combined( + projected, + rearrange(workload.cp.get_conv1d_weight(), "d 1 w -> d w"), + workload.cp.get_conv1d_bias(), + workload.cp.get_dt_bias().float(), + -torch.exp(workload.cp.get_A_log().float()), + D=workload.cp.get_D(), + chunk_size=workload.chunk_size, + activation=workload.activation, + headdim=workload.headdim, + ngroups=workload.cp.ngroups_local_tpcp, + norm_before_gate=workload.norm_before_gate, + ) + output = rearrange(output, "b l d -> l b d").contiguous() + return workload.cp.post_conv_ssm(output) + + def state_passing(load_balancing): + workload.config.mamba_state_passing_cp_load_balancing = load_balancing + return workload.state_passing_cp_adapter.forward(inputs.projected_input) + + return { + "a2a": a2a, + "permute-p2p": lambda: state_passing("permute_p2p"), + "permute-a2a": lambda: state_passing("permute_a2a"), + "virtual": lambda: state_passing("virtual"), + } + + +def benchmark_cp_path( + name: str, + runner: Callable[[], torch.Tensor], + *, + warmup: int, + iterations: int, + forward_only: bool, + workload: MambaCPWorkload, + inputs: MambaCPInputs, + group: dist.ProcessGroup, + device: torch.device, + rank: int, +) -> BenchmarkMetrics: + """Time one CP path, reporting the per-iteration maximum across ranks. + + A CP path is only as fast as its slowest rank, so each iteration is reduced + with MAX rather than averaged. Rank alignment happens on a barrier outside + the timed interval so that waiting for stragglers is not counted twice. + """ + times = [] + peaks = [] + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + torch.cuda.synchronize(device) + dist.barrier(group=group) + for iteration in range(warmup + iterations): + inputs.projected_input.grad = None + workload.zero_grad(set_to_none=True) + dist.barrier(group=group) + torch.cuda.reset_peak_memory_stats(device) + start.record() + output = runner() + if not forward_only: + torch.autograd.backward(output, inputs.grad_output) + end.record() + end.synchronize() + + metrics = torch.tensor( + [start.elapsed_time(end), torch.cuda.max_memory_allocated(device) / 1024**3], + device=device, + dtype=torch.float64, + ) + dist.all_reduce(metrics, op=dist.ReduceOp.MAX, group=group) + if iteration >= warmup: + times.append(metrics[0].item()) + peaks.append(metrics[1].item()) + del output + torch.cuda.synchronize(device) + dist.barrier(group=group) + + result = BenchmarkMetrics( + mean_ms=statistics.mean(times), + stdev_ms=statistics.stdev(times) if len(times) > 1 else 0.0, + p50_ms=statistics.median(times), + peak_gib=max(peaks), + ) + if rank == 0: + print( + f"{name:>18}: {result.mean_ms:8.2f} ms +/- {result.stdev_ms:6.2f} " + f"p50={result.p50_ms:8.2f} peak={result.peak_gib:6.2f} GiB", + flush=True, + ) + return result + + +def run_configuration(args, workload, context, sequence_length, batch_size): + """Benchmark every requested path for one (sequence length, batch size).""" + rank, cp_size, device, group = context + local_length = validate_shape(workload.shape, sequence_length, cp_size) + inputs = create_inputs(workload, sequence_length, batch_size, rank, device, cp_size) + runners = build_cp_path_runners(workload, inputs) + + if rank == 0: + mode = "fwd" if args.forward_only else "fwd+bwd" + print( + f"Mamba CP paths ({mode}) L={sequence_length} batch={batch_size} " + f"local_L={local_length} cp={cp_size} tp={args.tp_size}(local-shape-only) " + f"d_inner={workload.d_inner_local_tp} nheads={workload.nheads_local_tp} " + f"ngroups={workload.ngroups_local_tp}", + flush=True, + ) + + results = { + name: benchmark_cp_path( + name, + runners[name], + warmup=args.warmup, + iterations=args.iters, + forward_only=args.forward_only, + workload=workload, + inputs=inputs, + group=group, + device=device, + rank=rank, + ) + for name in args.paths + } + if rank != 0: + return + + if "a2a" in results: + baseline = results["a2a"].p50_ms + speedups = ", ".join( + f"{name}={baseline / metrics.p50_ms:.3f}x" + for name, metrics in results.items() + if name != "a2a" + ) + print(f"p50 speedup vs a2a: {speedups}", flush=True) + + for name, metrics in results.items(): + print( + f"RESULT,cp={cp_size},L={sequence_length},batch={batch_size}," + f"tp_size={args.tp_size},path={name},mean_ms={metrics.mean_ms:.6f}," + f"stdev_ms={metrics.stdev_ms:.6f},p50_ms={metrics.p50_ms:.6f}," + f"peak_gib={metrics.peak_gib:.6f}", + flush=True, + ) + print("Times are per-iteration max latency across ranks; lower is better.", flush=True) + + +def main(): + """Initialize one process per GPU and sweep the requested configurations.""" + args = parse_args() + assert args.iters > 0 and args.warmup >= 0 + + local_rank = dist.get_node_local_rank() + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("nccl", device_id=device) + group = dist.group.WORLD + context = (dist.get_rank(), dist.get_world_size(), device, group) + + workload = MambaCPWorkload(shape_from_args(args), device, group) + for sequence_length in args.sequence_lengths or [args.sequence_length]: + for batch_size in args.batch_sizes or [args.batch_size]: + run_configuration(args, workload, context, sequence_length, batch_size) + torch.cuda.empty_cache() + torch.cuda.synchronize(device) + dist.barrier(group=group) + + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py index c8b3ef583fe..1d3587a9e29 100644 --- a/megatron/core/ssm/mamba_mixer.py +++ b/megatron/core/ssm/mamba_mixer.py @@ -47,6 +47,14 @@ from .mamba_context_parallel import MambaContextParallel +try: + from .ops.ssd_state_passing_cp import MambaStatePassingCPAdapter + + HAVE_MAMBA_STATE_PASSING_CP = True +except ImportError: + MambaStatePassingCPAdapter = None + HAVE_MAMBA_STATE_PASSING_CP = False + try: from causal_conv1d import causal_conv1d_fn from causal_conv1d.causal_conv1d_varlen import causal_conv1d_varlen_states @@ -221,6 +229,10 @@ def __init__( "mamba_training_ssm_states_dtype is set, but the installed mamba_ssm does " "not accept the `state_dtype` argument. Upgrade mamba_ssm or unset the option." ) + self.use_mamba_state_passing_cp = self.config.use_mamba_state_passing_cp + self.mamba_state_passing_cp_load_balancing = ( + self.config.mamba_state_passing_cp_load_balancing + ) self.d_state = self.config.mamba_state_dim self.headdim = self.config.mamba_head_dim self.ngroups = self.config.mamba_num_groups @@ -459,6 +471,9 @@ def __init__( D_has_hdim=self.D_has_hdim, ) self.tp_group = pg_collection.tp + self.state_passing_cp_adapter = ( + MambaStatePassingCPAdapter(self) if HAVE_MAMBA_STATE_PASSING_CP else None + ) def forward( self, @@ -499,9 +514,16 @@ def forward( zxBCdt, _ = self.in_proj(hidden_states) - zxBCdt = self.cp.pre_conv_ssm(zxBCdt, packed_seq_params) + use_state_passing_cp = self._use_mamba_state_passing_cp( + in_inference_mode, packed_seq_params + ) + if not use_state_passing_cp: + zxBCdt = self.cp.pre_conv_ssm(zxBCdt, packed_seq_params) - if in_inference_mode or not self.use_mem_eff_path: + if use_state_passing_cp: + assert ssm_state is None + y = self.state_passing_cp_adapter.forward(zxBCdt) + elif in_inference_mode or not self.use_mem_eff_path: # TODO(ksanthanam): Consider deprecating this path for training assert packed_seq_params is None, ( "Training with packed sequences is not supported " @@ -727,6 +749,30 @@ def _decode( return out, out_bias + def _use_mamba_state_passing_cp( + self, in_inference_mode: bool, packed_seq_params: Optional[PackedSeqParams] + ) -> bool: + if not self.use_mamba_state_passing_cp or self.cp.cp_size == 1: + return False + assert HAVE_MAMBA_STATE_PASSING_CP, "state-passing Mamba CP helper is unavailable" + assert self.use_mem_eff_path, "state-passing Mamba CP is only wired for training fast path" + assert not in_inference_mode, "state-passing Mamba CP is not supported for inference yet" + assert ( + packed_seq_params is None + ), "state-passing Mamba CP does not support packed sequences yet" + assert ( + not self.config.hybrid_context_parallel and not self.config.dynamic_context_parallel + ), "state-passing Mamba CP does not support hybrid/dynamic (variable-length) CP yet" + assert ( + self.config.mamba_training_ssm_states_dtype is None + ), "state-passing Mamba CP does not support --mamba-training-ssm-states-dtype yet" + assert self.mamba_state_passing_cp_load_balancing != "none", ( + "standard Megatron CP always produces front/back load-balanced sequence shards; " + "use 'permute_p2p', 'permute_a2a', or 'virtual'. Mode 'none' is only valid for " + "direct state-passing calls whose input is already contiguous." + ) + return True + def _ssm_training( self, zxBCdt: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None ) -> torch.Tensor: diff --git a/megatron/core/ssm/ops/ssd_state_passing_cp.py b/megatron/core/ssm/ops/ssd_state_passing_cp.py new file mode 100644 index 00000000000..af7a716df6f --- /dev/null +++ b/megatron/core/ssm/ops/ssd_state_passing_cp.py @@ -0,0 +1,1572 @@ +"""Context-parallel state passing for Mamba2 SSD. + +The production autograd path combines causal Conv1d and the local SSD scan, then +exchanges only FP32 state summaries and scalar block decays. Backward recomputes +the conv output, applies the reverse state-passing CP boundary scan, and writes SSD gradients +directly into the buffers consumed by causal Conv1d backward. + +Verified by tests/unit_tests/ssm/ops/test_ssd_state_passing_cp.py. + +Note on the ``permute_p2p`` / ``permute_a2a`` load-balancing modes: converting +Megatron's balanced (zigzag) CP layout to a contiguous causal shard overlaps with +``megatron.core.context_parallel_layout``, whose ``CpPartitionMode`` models the +same zigzag/contiguous pair and which GDN already uses via +``convert_module_input_tensors_cp_partition_mode``. The permutation here is kept +self-contained because it is driven from inside the fused Conv+SSD autograd +function rather than at the module entry point, and because it adds a P2P backend +that the shared helper does not have. Consolidating onto the shared helper is +intended follow-up work, not a decision that the shared path is unsuitable. +""" + +from typing import List, Optional, Tuple + +import torch +import torch.distributed as dist +import triton +import triton.language as tl +from einops import rearrange +from mamba_ssm.ops.triton.ssd_combined import ( + _bmm_chunk_bwd, + _bmm_chunk_fwd, + _chunk_cumsum_bwd, + _chunk_cumsum_fwd, + _chunk_scan_bwd_dC, + _chunk_scan_bwd_dcb, + _chunk_scan_bwd_ddAcs_stable, + _chunk_scan_bwd_dstates, + _chunk_scan_bwd_dz, + _chunk_scan_chunk_state_bwd_dx, + _chunk_scan_fwd, + _chunk_state_bwd_db, + _chunk_state_fwd, + _state_passing_bwd, + _state_passing_fwd, + causal_conv1d_bwd_function, + causal_conv1d_fwd_function, + ensure_stride, +) +from mamba_ssm.utils.torch import custom_bwd, custom_fwd + + +def _all_gather_stack(x, group, async_op=False): + """Gather directly into a stacked contiguous tensor.""" + world = dist.get_world_size(group) + x = x.contiguous() + gathered = torch.empty((world, *x.shape), device=x.device, dtype=x.dtype) + work = dist.all_gather_into_tensor(gathered, x, group=group, async_op=async_op) + return gathered, work + + +@triton.jit +def _state_passing_summary_fwd_kernel( + # Pointers to matrices + states_ptr, + dA_cs_ptr, + summary_ptr, + # Matrix dimensions + dim: tl.constexpr, + nchunks: tl.constexpr, + nheads: tl.constexpr, + state_numel: tl.constexpr, + # Strides + stride_states_batch: tl.constexpr, + stride_states_chunk: tl.constexpr, + stride_states_head: tl.constexpr, + stride_states_dim: tl.constexpr, + stride_dA_cs_batch: tl.constexpr, + stride_dA_cs_head: tl.constexpr, + stride_dA_cs_chunk: tl.constexpr, + # Meta-parameters + BATCH_STRIDE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + pid_m = tl.program_id(axis=0) + batch_idx = pid_b * BATCH_STRIDE + states_ptr += batch_idx * stride_states_batch + pid_h * stride_states_head + dA_cs_ptr += batch_idx * stride_dA_cs_batch + pid_h * stride_dA_cs_head + + offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + states_ptrs = states_ptr + offs_m * stride_states_dim + state = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) + total_decay = 0.0 + for _ in range(nchunks): + new_state = tl.load(states_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + dA_cs = tl.load(dA_cs_ptr).to(tl.float32) + state = tl.exp(dA_cs) * state + new_state + total_decay += dA_cs + states_ptrs += stride_states_chunk + dA_cs_ptr += stride_dA_cs_chunk + + state_offset = (batch_idx * nheads + pid_h) * dim + tl.store(summary_ptr + state_offset + offs_m, state, mask=offs_m < dim) + decay_offset = state_numel + batch_idx * nheads + pid_h + tl.store(summary_ptr + decay_offset, tl.exp(total_decay), mask=pid_m == 0) + + +def _state_passing_summary_fwd(states, dA_chunk_cumsum, *, active_batch=None, batch_stride=1): + """Pack final states and block decays into one contiguous FP32 payload. + + active_batch and batch_stride select the interleaved front slots used by the + virtual state-passing CP rank-0 summary optimization. + """ + assert states.dtype == torch.float32 + assert dA_chunk_cumsum.dtype == torch.float32 + batch, nchunks, nheads, dim = states.shape + active_batch = batch if active_batch is None else active_batch + assert active_batch > 0 + assert (active_batch - 1) * batch_stride < batch + + state_numel = batch * nheads * dim + summary_numel = state_numel + batch * nheads + if active_batch < batch: + summary = torch.zeros(summary_numel, device=states.device, dtype=torch.float32) + else: + summary = torch.empty(summary_numel, device=states.device, dtype=torch.float32) + + block_size = 256 + grid = (triton.cdiv(dim, block_size), active_batch, nheads) + _state_passing_summary_fwd_kernel[grid]( + states, + dA_chunk_cumsum, + summary, + dim=dim, + nchunks=nchunks, + nheads=nheads, + state_numel=state_numel, + stride_states_batch=states.stride(0), + stride_states_chunk=states.stride(1), + stride_states_head=states.stride(2), + stride_states_dim=states.stride(3), + stride_dA_cs_batch=dA_chunk_cumsum.stride(0), + stride_dA_cs_head=dA_chunk_cumsum.stride(1), + stride_dA_cs_chunk=dA_chunk_cumsum.stride(2), + BATCH_STRIDE=batch_stride, + BLOCK_SIZE=block_size, + ) + return summary, state_numel + + +@triton.jit +def _state_passing_summary_bwd_kernel( + # Pointers to matrices + dout_ptr, + dA_cs_ptr, + dinitstates_ptr, + # Matrix dimensions + dim: tl.constexpr, + nchunks: tl.constexpr, + # Strides + stride_dout_batch: tl.constexpr, + stride_dout_chunk: tl.constexpr, + stride_dout_head: tl.constexpr, + stride_dout_dim: tl.constexpr, + stride_dA_cs_batch: tl.constexpr, + stride_dA_cs_head: tl.constexpr, + stride_dA_cs_chunk: tl.constexpr, + stride_dinitstates_batch: tl.constexpr, + stride_dinitstates_head: tl.constexpr, + stride_dinitstates_dim: tl.constexpr, + # Meta-parameters + BLOCK_SIZE: tl.constexpr, +): + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + pid_m = tl.program_id(axis=0) + dout_ptr += ( + pid_b * stride_dout_batch + pid_h * stride_dout_head + (nchunks - 1) * stride_dout_chunk + ) + dA_cs_ptr += ( + pid_b * stride_dA_cs_batch + pid_h * stride_dA_cs_head + (nchunks - 1) * stride_dA_cs_chunk + ) + + offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + dout_ptrs = dout_ptr + offs_m * stride_dout_dim + dstates = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) + for _ in range(nchunks): + dout = tl.load(dout_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + dA_cs = tl.load(dA_cs_ptr).to(tl.float32) + dstates = tl.exp(dA_cs) * dstates + dout + dout_ptrs -= stride_dout_chunk + dA_cs_ptr -= stride_dA_cs_chunk + + dinitstates_ptr += pid_b * stride_dinitstates_batch + pid_h * stride_dinitstates_head + tl.store(dinitstates_ptr + offs_m * stride_dinitstates_dim, dstates, mask=offs_m < dim) + + +def _state_passing_summary_bwd(dout, dA_chunk_cumsum): + """Compute only the gradient entering the first local chunk.""" + assert dout.dtype == torch.float32 + assert dA_chunk_cumsum.dtype == torch.float32 + batch, nchunks, nheads, dim = dout.shape + assert dA_chunk_cumsum.shape == (batch, nheads, nchunks) + + dinitstates = torch.empty((batch, nheads, dim), device=dout.device, dtype=torch.float32) + block_size = 256 + grid = (triton.cdiv(dim, block_size), batch, nheads) + _state_passing_summary_bwd_kernel[grid]( + dout, + dA_chunk_cumsum, + dinitstates, + dim=dim, + nchunks=nchunks, + stride_dout_batch=dout.stride(0), + stride_dout_chunk=dout.stride(1), + stride_dout_head=dout.stride(2), + stride_dout_dim=dout.stride(3), + stride_dA_cs_batch=dA_chunk_cumsum.stride(0), + stride_dA_cs_head=dA_chunk_cumsum.stride(1), + stride_dA_cs_chunk=dA_chunk_cumsum.stride(2), + stride_dinitstates_batch=dinitstates.stride(0), + stride_dinitstates_head=dinitstates.stride(1), + stride_dinitstates_dim=dinitstates.stride(2), + BLOCK_SIZE=block_size, + ) + return dinitstates + + +@triton.jit +def _state_passing_boundary_scan_kernel( + # Pointers to matrices + boundary_values_ptr, + block_decays_ptr, + boundary_output_ptr, + # Matrix dimensions + dim: tl.constexpr, + world: tl.constexpr, + rank: tl.constexpr, + # Strides + stride_values_rank: tl.constexpr, + stride_values_batch: tl.constexpr, + stride_values_head: tl.constexpr, + stride_decays_rank: tl.constexpr, + stride_decays_batch: tl.constexpr, + stride_decays_head: tl.constexpr, + stride_output_batch: tl.constexpr, + stride_output_head: tl.constexpr, + # Meta-parameters + VIRTUAL_STATE_PASSING_CP: tl.constexpr, + REVERSE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + pid_m = tl.program_id(axis=0) + offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + state = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) + + nsegments = 2 * world if VIRTUAL_STATE_PASSING_CP else world + for step in range(nsegments): + segment = nsegments - 1 - step if REVERSE else step + if VIRTUAL_STATE_PASSING_CP: + if segment == rank: + tl.store( + boundary_output_ptr + + (2 * pid_b) * stride_output_batch + + pid_h * stride_output_head + + offs_m, + state, + mask=offs_m < dim, + ) + if segment == nsegments - 1 - rank: + tl.store( + boundary_output_ptr + + (2 * pid_b + 1) * stride_output_batch + + pid_h * stride_output_head + + offs_m, + state, + mask=offs_m < dim, + ) + source_rank = segment if segment < world else nsegments - 1 - segment + source_batch = 2 * pid_b if segment < world else 2 * pid_b + 1 + include = True + else: + source_rank = segment + source_batch = pid_b + include = source_rank > rank if REVERSE else source_rank < rank + + if include: + new_state = tl.load( + boundary_values_ptr + + source_rank * stride_values_rank + + source_batch * stride_values_batch + + pid_h * stride_values_head + + offs_m, + mask=offs_m < dim, + other=0.0, + ).to(tl.float32) + scale = tl.load( + block_decays_ptr + + source_rank * stride_decays_rank + + source_batch * stride_decays_batch + + pid_h * stride_decays_head + ).to(tl.float32) + state = scale * state + new_state + + if not VIRTUAL_STATE_PASSING_CP: + tl.store( + boundary_output_ptr + pid_b * stride_output_batch + pid_h * stride_output_head + offs_m, + state, + mask=offs_m < dim, + ) + + +def _state_passing_boundary_scan( + boundary_values, block_decays, rank, state_passing_cp_virtual, reverse +): + """Scan gathered summaries and return this rank's exclusive causal boundary.""" + assert boundary_values.dtype == torch.float32 + assert block_decays.dtype == torch.float32 + world, batch, nheads = boundary_values.shape[:3] + state_shape = boundary_values.shape[3:] + boundary_values_flat = boundary_values.view(world, batch, nheads, -1) + dim = boundary_values_flat.shape[-1] + boundary_output = torch.empty( + batch, nheads, dim, device=boundary_values.device, dtype=torch.float32 + ) + + block_size = 256 + if state_passing_cp_virtual: + assert batch % 2 == 0 + scan_batch = batch // 2 + else: + scan_batch = batch + grid = (triton.cdiv(dim, block_size), scan_batch, nheads) + _state_passing_boundary_scan_kernel[grid]( + boundary_values_flat, + block_decays, + boundary_output, + dim=dim, + world=world, + rank=rank, + stride_values_rank=boundary_values_flat.stride(0), + stride_values_batch=boundary_values_flat.stride(1), + stride_values_head=boundary_values_flat.stride(2), + stride_decays_rank=block_decays.stride(0), + stride_decays_batch=block_decays.stride(1), + stride_decays_head=block_decays.stride(2), + stride_output_batch=boundary_output.stride(0), + stride_output_head=boundary_output.stride(1), + VIRTUAL_STATE_PASSING_CP=state_passing_cp_virtual, + REVERSE=reverse, + BLOCK_SIZE=block_size, + ) + return boundary_output.view(batch, nheads, *state_shape) + + +def _mamba_chunk_scan_combined_state_passing_cp_fwd( + x, + dt, + A, + B, + C, + chunk_size, + D=None, + z=None, + dt_bias=None, + initial_states=None, + seq_idx=None, + cu_seqlens=None, + dt_softplus=False, + dt_limit=(0.0, float("inf")), + state_passing_cp_group=None, + state_passing_cp_virtual=False, +): + """Mamba combined forward with state-passing CP summary and boundary exchange.""" + assert state_passing_cp_group is not None + batch, seqlen, nheads, headdim = x.shape + _, _, ngroups, dstate = B.shape + assert nheads % ngroups == 0 + assert B.shape == (batch, seqlen, ngroups, dstate) + assert x.shape == (batch, seqlen, nheads, headdim) + assert dt.shape == (batch, seqlen, nheads) + assert A.shape == (nheads,) + assert C.shape == B.shape + if z is not None: + assert z.shape == x.shape + if D is not None: + assert D.shape == (nheads, headdim) or D.shape == (nheads,) + if seq_idx is not None: + assert seq_idx.shape == (batch, seqlen) + if B.stride(-1) != 1: + B = B.contiguous() + if C.stride(-1) != 1: + C = C.contiguous() + if x.stride(-1) != 1 and x.stride(1) != 1: + x = x.contiguous() + if z is not None and z.stride(-1) != 1 and z.stride(1) != 1: + z = z.contiguous() + if D is not None and D.stride(-1) != 1: + D = D.contiguous() + if initial_states is not None: + assert initial_states.shape == (batch, nheads, headdim, dstate) + + assert ( + initial_states is None + ), "external initial_states are not implemented for state-passing CP" + assert seq_idx is None, "seq_idx is not implemented for state-passing CP" + assert cu_seqlens is None, "cu_seqlens is not implemented for state-passing CP" + state_passing_cp_size = dist.get_world_size(state_passing_cp_group) + rank = dist.get_rank(state_passing_cp_group) + if state_passing_cp_virtual: + assert ( + batch % 2 == 0 + ), "virtual state-passing CP expects front/back packed on the batch axis" + virtual_batch = batch // 2 if state_passing_cp_virtual else None + + dA_cumsum, dt = _chunk_cumsum_fwd( + dt, A, chunk_size, dt_bias=dt_bias, dt_softplus=dt_softplus, dt_limit=dt_limit + ) + states = _chunk_state_fwd(B, x, dt, dA_cumsum, seq_idx=seq_idx, states_in_fp32=True) + state_passing_initial_states = None + state_passing_gathered_decays = None + if state_passing_cp_size > 1: + state_dim = headdim * dstate + if not state_passing_cp_virtual and rank == state_passing_cp_size - 1: + # Collective shapes are fixed even though no later rank consumes this summary. + state_numel = batch * nheads * state_dim + packed_summary = torch.zeros( + state_numel + batch * nheads, device=x.device, dtype=torch.float32 + ) + else: + if state_passing_cp_virtual and rank == 0: + # Rank 0's causal-last segment has no successor. + active_batch, batch_stride = virtual_batch, 2 + else: + active_batch, batch_stride = None, 1 + packed_summary, state_numel = _state_passing_summary_fwd( + rearrange(states, "... p n -> ... (p n)"), + dA_cumsum[:, :, :, -1], + active_batch=active_batch, + batch_stride=batch_stride, + ) + + gathered_summary, gather_work = _all_gather_stack( + packed_summary, state_passing_cp_group, async_op=True + ) + CB = _bmm_chunk_fwd(C, B, chunk_size, seq_idx=seq_idx, output_dtype=torch.float32) + gather_work.wait() + gathered_states = gathered_summary[:, :state_numel].view( + state_passing_cp_size, batch, nheads, headdim, dstate + ) + state_passing_gathered_decays = gathered_summary[:, state_numel:].view( + state_passing_cp_size, batch, nheads + ) + state_passing_initial_states = _state_passing_boundary_scan( + gathered_states, + state_passing_gathered_decays, + rank, + state_passing_cp_virtual=state_passing_cp_virtual, + reverse=False, + ) + state_passing_gathered_decays = state_passing_gathered_decays.contiguous() + + states, final_states = _state_passing_fwd( + rearrange(states, "... p n -> ... (p n)"), + dA_cumsum[:, :, :, -1], + initial_states=( + rearrange(state_passing_initial_states, "... p n -> ... (p n)") + if state_passing_initial_states is not None + else None + ), + seq_idx=seq_idx, + chunk_size=chunk_size, + out_dtype=C.dtype, + ) + states, final_states = [ + rearrange(t, "... (p n) -> ... p n", n=dstate) for t in [states, final_states] + ] + if state_passing_cp_size == 1: + CB = _bmm_chunk_fwd(C, B, chunk_size, seq_idx=seq_idx, output_dtype=torch.float32) + # CB comes from the state-passing branch above when cp_size > 1, and from the + # branch just above otherwise; the two cases are exhaustive. + # pylint: disable-next=possibly-used-before-assignment + out, out_x = _chunk_scan_fwd(CB, x, dt, dA_cumsum, C, states, D=D, z=z, seq_idx=seq_idx) + return ( + out, + out_x, + dt, + dA_cumsum, + states, + final_states, + state_passing_initial_states, + state_passing_gathered_decays, + ) + + +def _mamba_chunk_scan_combined_state_passing_cp_bwd( + dout, + x, + dt, + A, + B, + C, + out, + chunk_size, + D=None, + z=None, + dt_bias=None, + initial_states=None, + dfinal_states=None, + seq_idx=None, + dt_softplus=False, + dt_limit=(0.0, float("inf")), + dx=None, + ddt=None, + dB=None, + dC=None, + dz=None, + recompute_output=False, + state_passing_cp_group=None, + state_passing_initial_states=None, + state_passing_gathered_decays=None, + state_passing_cp_virtual=False, +): + """Mamba combined backward with a state-passing CP reverse-boundary handoff.""" + assert state_passing_cp_group is not None + if dout.stride(-1) != 1: + dout = dout.contiguous() + batch, seqlen, nheads, headdim = x.shape + _, _, ngroups, dstate = B.shape + assert dout.shape == (batch, seqlen, nheads, headdim) + assert dt.shape == (batch, seqlen, nheads) + assert A.shape == (nheads,) + assert nheads % ngroups == 0 + assert B.shape == (batch, seqlen, ngroups, dstate) + assert C.shape == B.shape + assert out.shape == x.shape + if initial_states is not None: + assert initial_states.shape == (batch, nheads, headdim, dstate) + if seq_idx is not None: + assert seq_idx.shape == (batch, seqlen) + if dx is not None: + assert dx.shape == x.shape + if dB is not None: + assert dB.shape == B.shape + dB_given = dB + else: + dB_given = torch.empty_like(B) + if dC is not None: + assert dC.shape == C.shape + dC_given = dC + else: + dC_given = torch.empty_like(C) + if dz is not None: + assert z is not None + assert dz.shape == z.shape + if ddt is not None: + assert ddt.shape == dt.shape + ddt_given = ddt + else: + ddt_given = torch.empty_like(dt) + + assert ( + initial_states is None + ), "external initial_states are not implemented for state-passing CP" + assert dfinal_states is None, "dfinal_states are not implemented for state-passing CP" + assert seq_idx is None, "seq_idx is not implemented for state-passing CP" + state_passing_cp_size = dist.get_world_size(state_passing_cp_group) + rank = dist.get_rank(state_passing_cp_group) + + dt_in = dt.clone() + dA_cumsum, dt = _chunk_cumsum_fwd( + dt_in, A, chunk_size, dt_bias=dt_bias, dt_softplus=dt_softplus, dt_limit=dt_limit + ) + CB = _bmm_chunk_fwd(C, B, chunk_size, seq_idx=seq_idx, output_dtype=torch.float32) + states = _chunk_state_fwd(B, x, dt, dA_cumsum, seq_idx=seq_idx, states_in_fp32=True) + states, _ = _state_passing_fwd( + rearrange(states, "... p n -> ... (p n)"), + dA_cumsum[:, :, :, -1], + initial_states=( + rearrange(state_passing_initial_states, "... p n -> ... (p n)") + if state_passing_initial_states is not None + else None + ), + seq_idx=seq_idx, + chunk_size=chunk_size, + ) + states = rearrange(states, "... (p n) -> ... p n", n=dstate) + + if z is not None: + dz, dout, dD, *rest = _chunk_scan_bwd_dz( + x, + z, + out, + dout, + chunk_size=chunk_size, + has_ddAcs=False, + D=D, + dz=dz, + recompute_output=recompute_output, + ) + outz = rest[0] if recompute_output else out + else: + dz = None + outz = out + + dstates = _chunk_scan_bwd_dstates(C, dA_cumsum, dout, seq_idx=seq_idx, dtype=states.dtype) + + # Everything in this block is state-passing CP-specific. The local kernels run while the + # reverse-boundary all-gather is in flight. + dC_local = dCB = ddA = None + final_state_grad = None + if state_passing_cp_size > 1: + assert ( + state_passing_initial_states is not None and state_passing_gathered_decays is not None + ) + initial_state_grad = _state_passing_summary_bwd( + rearrange(dstates, "... p n -> ... (p n)"), dA_cumsum[:, :, :, -1] + ) + initial_state_grad = rearrange(initial_state_grad, "... (p n) -> ... p n", n=dstate) + gathered_initial_state_grads, gather_work = _all_gather_stack( + initial_state_grad, state_passing_cp_group, async_op=True + ) + + states_for_scan = states.to(x.dtype) + dC_local, ddA_cumsum_prev = _chunk_scan_bwd_dC( + states_for_scan, dA_cumsum, dout, seq_idx=seq_idx, C=C, ngroups=ngroups + ) + del states_for_scan + dCB = _chunk_scan_bwd_dcb(x, dt, dA_cumsum, dout, seq_idx=seq_idx, ngroups=ngroups).to( + CB.dtype + ) + ddA = _chunk_scan_bwd_ddAcs_stable(x, dt, dA_cumsum, dout, CB) + gather_work.wait() + final_state_grad = _state_passing_boundary_scan( + gathered_initial_state_grads, + state_passing_gathered_decays, + rank, + state_passing_cp_virtual=state_passing_cp_virtual, + reverse=True, + ) + + state_bwd_outputs = _state_passing_bwd( + rearrange(states, "... p n -> ... (p n)"), + dA_cumsum[:, :, :, -1], + rearrange(dstates, "... p n -> ... (p n)"), + dfinal_states=( + rearrange(final_state_grad, "... p n -> ... (p n)") + if final_state_grad is not None + else None + ), + seq_idx=seq_idx, + has_initial_states=state_passing_initial_states is not None, + dstates_dtype=x.dtype, + # dC already consumed the converted states in the overlapped state-passing CP path. + states_dtype=x.dtype if dC_local is None else None, + chunk_size=chunk_size, + ) + dstates, ddA_chunk_cumsum, _ = state_bwd_outputs[:3] + dstates = rearrange(dstates, "... (p n) -> ... p n", n=dstate) + + if dC_local is None: + states = rearrange(state_bwd_outputs[3], "... (p n) -> ... p n", n=dstate) + dC_local, ddA_cumsum_prev = _chunk_scan_bwd_dC( + states.to(x.dtype), dA_cumsum, dout, seq_idx=seq_idx, C=C, ngroups=ngroups + ) + dCB = _chunk_scan_bwd_dcb(x, dt, dA_cumsum, dout, seq_idx=seq_idx, ngroups=ngroups).to( + CB.dtype + ) + ddA = _chunk_scan_bwd_ddAcs_stable(x, dt, dA_cumsum, dout, CB) + + dx, ddt, dD_from_x = _chunk_scan_chunk_state_bwd_dx( + x, dt, dA_cumsum, B, CB, dout, dstates, D=D, seq_idx=seq_idx, dx=dx + ) + dB, ddA_next = _chunk_state_bwd_db( + x, dt, dA_cumsum, dstates, seq_idx=seq_idx, B=B, ngroups=ngroups + ) + _bmm_chunk_bwd(C, dCB, residual=dB, out=dB_given) + _bmm_chunk_bwd(B, rearrange(dCB, "... l s -> ... s l"), residual=dC_local, out=dC_given) + if z is None: + dD = dD_from_x + + # ddA_cumsum_prev is set by whichever of the two exhaustive dC_local branches + # ran, and dD by whichever of the two exhaustive `z` branches ran. + # pylint: disable-next=possibly-used-before-assignment + ddA_cumsum_prev[..., -1] += ddA_chunk_cumsum + ddA_prev = ddA_cumsum_prev.flip([-1]).cumsum(dim=-1).flip([-1]) + ddA += ddA_next + ddA_prev + ddt_given, dA, ddt_bias = _chunk_cumsum_bwd( + ddA, + ddt, + dt_in, + A, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + dt_limit=dt_limit, + ddt=ddt_given, + ) + + dinitial_states = None + # pylint: disable-next=possibly-used-before-assignment + return_vals = (dx, ddt_given, dA, dB_given, dC_given, dD, dz, ddt_bias, dinitial_states) + return return_vals if not recompute_output else (*return_vals, outz) + + +# Mixer integration, conv boundary exchange, and state-passing CP layout handling. +def _pack_state_passing_cp_virtual_segments(x: torch.Tensor) -> torch.Tensor: + """View balanced chunks as interleaved front/back segments on the batch axis.""" + assert x.shape[1] % 2 == 0, "virtual state-passing CP requires two equal local sequence chunks" + batch = x.shape[0] + half = x.shape[1] // 2 + return x.reshape(2 * batch, half, *x.shape[2:]) + + +def _unpack_state_passing_cp_virtual_segments(x: torch.Tensor) -> torch.Tensor: + """View interleaved front/back segments as the balanced local sequence layout.""" + assert x.shape[0] % 2 == 0 + batch = x.shape[0] // 2 + return x.reshape(batch, 2 * x.shape[1], *x.shape[2:]) + + +def _state_passing_cp_chunk_owner_slot( + chunk_id: int, state_passing_cp_size: int, layout: str +) -> Tuple[int, int]: + if layout == "contiguous": + return chunk_id // 2, chunk_id % 2 + if chunk_id < state_passing_cp_size: + return chunk_id, 0 + return 2 * state_passing_cp_size - 1 - chunk_id, 1 + + +def _state_passing_cp_local_chunk_ids( + rank: int, state_passing_cp_size: int, layout: str +) -> Tuple[int, int]: + if layout == "contiguous": + return 2 * rank, 2 * rank + 1 + return rank, 2 * state_passing_cp_size - 1 - rank + + +def permute_state_passing_cp_sequence_chunks( + x: torch.Tensor, + state_passing_cp_group: torch.distributed.ProcessGroup, + undo_load_balancing: bool, + backend: str = "p2p", +) -> torch.Tensor: + """Exchange sequence chunks between the balanced and contiguous CP layouts. + + ``undo_load_balancing=True`` converts Megatron's balanced front/back layout + into a contiguous causal shard; ``False`` restores the balanced layout. The + backward pass applies the inverse permutation. + + This mirrors the zigzag/contiguous conversion in + ``megatron.core.context_parallel_layout``; see this module's docstring for why + it is currently implemented separately. + """ + if dist.get_world_size(state_passing_cp_group) == 1: + return x + if backend not in ("p2p", "a2a"): + raise ValueError(f"Unsupported state-passing CP sequence permutation backend {backend!r}") + return _StatePassingCPSequenceChunkPermutationFn.apply( + x, state_passing_cp_group, undo_load_balancing, backend + ) + + +def undo_state_passing_cp_load_balancing( + x: torch.Tensor, state_passing_cp_group: torch.distributed.ProcessGroup, backend: str = "p2p" +) -> torch.Tensor: + """Convert a balanced CP shard into a contiguous causal shard.""" + return permute_state_passing_cp_sequence_chunks(x, state_passing_cp_group, True, backend) + + +def redo_state_passing_cp_load_balancing( + x: torch.Tensor, state_passing_cp_group: torch.distributed.ProcessGroup, backend: str = "p2p" +) -> torch.Tensor: + """Convert a contiguous causal shard back into the balanced CP layout.""" + return permute_state_passing_cp_sequence_chunks(x, state_passing_cp_group, False, backend) + + +class _StatePassingCPSequenceChunkPermutationFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x, + state_passing_cp_group: torch.distributed.ProcessGroup, + undo_load_balancing: bool, + backend: str, + ): + ctx.state_passing_cp_group = state_passing_cp_group + ctx.undo_load_balancing = undo_load_balancing + ctx.backend = backend + return _permute_state_passing_cp_sequence_chunks( + x, state_passing_cp_group, undo_load_balancing, backend + ) + + @staticmethod + def backward(ctx, grad_output): + grad_x = _permute_state_passing_cp_sequence_chunks( + grad_output, ctx.state_passing_cp_group, not ctx.undo_load_balancing, ctx.backend + ) + return grad_x, None, None, None + + +def _permute_state_passing_cp_sequence_chunks( + x: torch.Tensor, + state_passing_cp_group: torch.distributed.ProcessGroup, + undo_load_balancing: bool, + backend: str = "p2p", +) -> torch.Tensor: + if backend == "a2a": + return _permute_state_passing_cp_sequence_chunks_a2a( + x, state_passing_cp_group, undo_load_balancing + ) + if backend != "p2p": + raise ValueError(f"Unsupported state-passing CP sequence permutation backend {backend!r}") + return _permute_state_passing_cp_sequence_chunks_p2p( + x, state_passing_cp_group, undo_load_balancing + ) + + +def _permute_state_passing_cp_sequence_chunks_p2p( + x: torch.Tensor, + state_passing_cp_group: torch.distributed.ProcessGroup, + undo_load_balancing: bool, +) -> torch.Tensor: + state_passing_cp_size = dist.get_world_size(state_passing_cp_group) + rank = dist.get_rank(state_passing_cp_group) + assert ( + x.size(0) % 2 == 0 + ), "State-passing CP load-balanced sequence shard must contain two chunks" + chunk_len = x.size(0) // 2 + input_layout = "balanced" if undo_load_balancing else "contiguous" + output_layout = "contiguous" if undo_load_balancing else "balanced" + input_chunks = _state_passing_cp_local_chunk_ids(rank, state_passing_cp_size, input_layout) + output_chunks = _state_passing_cp_local_chunk_ids(rank, state_passing_cp_size, output_layout) + + # P2P receive buffers must be non-overlapping and dense. In backward, + # grad_output can inherit a non-standard stride from the downstream layout. + out = torch.empty_like(x, memory_format=torch.contiguous_format) + ops = [] + sends = [] + + for out_slot, chunk_id in enumerate(output_chunks): + src_rank, src_slot = _state_passing_cp_chunk_owner_slot( + chunk_id, state_passing_cp_size, input_layout + ) + out_slice = out[out_slot * chunk_len : (out_slot + 1) * chunk_len] + if src_rank == rank: + src = x[src_slot * chunk_len : (src_slot + 1) * chunk_len] + out_slice.copy_(src) + else: + ops.append( + dist.P2POp(dist.irecv, out_slice, group=state_passing_cp_group, group_peer=src_rank) + ) + + for in_slot, chunk_id in enumerate(input_chunks): + dst_rank, _ = _state_passing_cp_chunk_owner_slot( + chunk_id, state_passing_cp_size, output_layout + ) + if dst_rank != rank: + send = x[in_slot * chunk_len : (in_slot + 1) * chunk_len].contiguous() + sends.append(send) + ops.append( + dist.P2POp(dist.isend, send, group=state_passing_cp_group, group_peer=dst_rank) + ) + + for work in dist.batch_isend_irecv(ops): + work.wait() + return out + + +def _permute_state_passing_cp_sequence_chunks_a2a( + x: torch.Tensor, + state_passing_cp_group: torch.distributed.ProcessGroup, + undo_load_balancing: bool, +) -> torch.Tensor: + state_passing_cp_size = dist.get_world_size(state_passing_cp_group) + rank = dist.get_rank(state_passing_cp_group) + assert ( + x.size(0) % 2 == 0 + ), "State-passing CP load-balanced sequence shard must contain two chunks" + chunk_len = x.size(0) // 2 + input_layout = "balanced" if undo_load_balancing else "contiguous" + output_layout = "contiguous" if undo_load_balancing else "balanced" + input_chunks = _state_passing_cp_local_chunk_ids(rank, state_passing_cp_size, input_layout) + + local_destinations = [ + _state_passing_cp_chunk_owner_slot(chunk_id, state_passing_cp_size, output_layout) + for chunk_id in input_chunks + ] + send_slot_order = sorted(range(2), key=lambda slot: local_destinations[slot]) + local_chunks = (x[:chunk_len], x[chunk_len:]) + send_buf = torch.cat([local_chunks[slot] for slot in send_slot_order], dim=0).contiguous() + + input_split_chunks = [0] * state_passing_cp_size + for dst_rank, _ in local_destinations: + input_split_chunks[dst_rank] += 1 + + output_split_chunks = [0] * state_passing_cp_size + recv_slots_by_source = [[] for _ in range(state_passing_cp_size)] + for src_rank in range(state_passing_cp_size): + src_chunks = _state_passing_cp_local_chunk_ids( + src_rank, state_passing_cp_size, input_layout + ) + src_destinations = [ + _state_passing_cp_chunk_owner_slot(chunk_id, state_passing_cp_size, output_layout) + for chunk_id in src_chunks + ] + for src_slot in sorted(range(2), key=lambda slot: src_destinations[slot]): + dst_rank, dst_slot = src_destinations[src_slot] + if dst_rank == rank: + output_split_chunks[src_rank] += 1 + recv_slots_by_source[src_rank].append(dst_slot) + + input_split_sizes = [count * chunk_len for count in input_split_chunks] + output_split_sizes = [count * chunk_len for count in output_split_chunks] + recv_buf = torch.empty_like(x, memory_format=torch.contiguous_format) + dist.all_to_all_single( + recv_buf, + send_buf, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=state_passing_cp_group, + ) + + target_slots: List[Optional[torch.Tensor]] = [None, None] + offset = 0 + for src_rank in range(state_passing_cp_size): + for dst_slot in recv_slots_by_source[src_rank]: + target_slots[dst_slot] = recv_buf[offset : offset + chunk_len] + offset += chunk_len + assert all( + slot is not None for slot in target_slots + ), "Incomplete state-passing CP sequence A2A reassembly" + return torch.cat(target_slots, dim=0) + + +@triton.jit +def _route_virtual_state_passing_cp_conv_boundary_kernel( + gathered, + output, + batch: tl.constexpr, + channels: tl.constexpr, + halo: tl.constexpr, + rank: tl.constexpr, + world: tl.constexpr, + stride_g_rank: tl.constexpr, + stride_g_batch: tl.constexpr, + stride_g_channel: tl.constexpr, + stride_g_halo: tl.constexpr, + stride_o_batch: tl.constexpr, + stride_o_channel: tl.constexpr, + stride_o_halo: tl.constexpr, + REVERSE: tl.constexpr, + BLOCK_C: tl.constexpr, + BLOCK_H: tl.constexpr, +): + pid_b = tl.program_id(0) + pid_c = tl.program_id(1) + slot = pid_b % 2 + local_batch = pid_b // 2 + + if not REVERSE: + if slot == 0: + valid = rank > 0 + source_rank = rank - 1 + source_batch = 2 * local_batch + else: + valid = True + source_rank = rank + 1 + source_batch = 2 * local_batch + 1 + if rank == world - 1: + source_rank = rank + source_batch = 2 * local_batch + else: + if slot == 0: + valid = True + source_rank = rank + 1 + source_batch = 2 * local_batch + if rank == world - 1: + source_rank = rank + source_batch = 2 * local_batch + 1 + else: + valid = rank > 0 + source_rank = rank - 1 + source_batch = 2 * local_batch + 1 + + offsets_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) + offsets_h = tl.arange(0, BLOCK_H) + mask = (offsets_c[:, None] < channels) & (offsets_h[None, :] < halo) + values = tl.load( + gathered + + source_rank * stride_g_rank + + source_batch * stride_g_batch + + offsets_c[:, None] * stride_g_channel + + offsets_h[None, :] * stride_g_halo, + mask=mask & valid, + other=0.0, + ) + tl.store( + output + + pid_b * stride_o_batch + + offsets_c[:, None] * stride_o_channel + + offsets_h[None, :] * stride_o_halo, + values, + mask=mask, + ) + + +def _route_virtual_state_passing_cp_conv_boundary( + gathered: torch.Tensor, + batch: int, + rank: int, + world: int, + reverse: bool, + channel_last_output: bool, +) -> torch.Tensor: + channels, halo = gathered.shape[-2:] + if channel_last_output: + output = torch.empty( + 2 * batch, halo, channels, device=gathered.device, dtype=gathered.dtype + ).transpose(1, 2) + else: + output = torch.empty( + 2 * batch, channels, halo, device=gathered.device, dtype=gathered.dtype + ) + block_c = 128 + block_h = triton.next_power_of_2(halo) + _route_virtual_state_passing_cp_conv_boundary_kernel[ + (2 * batch, triton.cdiv(channels, block_c)) + ]( + gathered, + output, + batch=batch, + channels=channels, + halo=halo, + rank=rank, + world=world, + stride_g_rank=gathered.stride(0), + stride_g_batch=gathered.stride(1), + stride_g_channel=gathered.stride(2), + stride_g_halo=gathered.stride(3), + stride_o_batch=output.stride(0), + stride_o_channel=output.stride(1), + stride_o_halo=output.stride(2), + REVERSE=reverse, + BLOCK_C=block_c, + BLOCK_H=block_h, + ) + return output + + +def _gather_causal_conv1d_state_passing_cp_boundary( + boundary: torch.Tensor, + state_passing_cp_group: torch.distributed.ProcessGroup, + state_passing_cp_virtual: bool, + reverse: bool, + channel_last_output: bool, +) -> torch.Tensor: + """Gather a predecessor state or successor gradient for causal conv.""" + world = dist.get_world_size(state_passing_cp_group) + rank = dist.get_rank(state_passing_cp_group) + batch, channels, halo = boundary.shape + gathered = boundary.new_empty(world * batch, channels, halo) + dist.all_gather_into_tensor(gathered, boundary.contiguous(), group=state_passing_cp_group) + gathered = gathered.view(world, batch, channels, halo) + + if state_passing_cp_virtual: + assert batch % 2 == 0 + return _route_virtual_state_passing_cp_conv_boundary( + gathered, + batch // 2, + rank, + world, + reverse=reverse, + channel_last_output=channel_last_output, + ) + + if channel_last_output: + output = torch.empty( + batch, halo, channels, device=boundary.device, dtype=boundary.dtype + ).transpose(1, 2) + else: + output = torch.empty_like(boundary, memory_format=torch.contiguous_format) + source_rank = rank + 1 if reverse else rank - 1 + if 0 <= source_rank < world: + output.copy_(gathered[source_rank]) + else: + output.zero_() + return output + + +def _causal_conv1d_state_passing_cp_fwd( + x, + weight, + bias, + seq_idx=None, + initial_states=None, + final_states_out=None, + activation=None, + state_passing_cp_group=None, + state_passing_cp_virtual=False, +): + """Run causal Conv1d using a gathered state-passing CP predecessor state.""" + assert causal_conv1d_fwd_function is not None + assert state_passing_cp_group is not None + assert seq_idx is None, "seq_idx is not implemented for state-passing CP" + assert final_states_out is None + halo = weight.shape[-1] - 1 + assert halo > 0 and x.size(-1) >= halo + state_passing_conv_initial_states = initial_states + if state_passing_conv_initial_states is None: + state_passing_conv_initial_states = _gather_causal_conv1d_state_passing_cp_boundary( + x[..., -halo:], + state_passing_cp_group, + state_passing_cp_virtual=state_passing_cp_virtual, + reverse=False, + channel_last_output=True, + ) + out = causal_conv1d_fwd_function( + x, + weight, + bias, + seq_idx, + state_passing_conv_initial_states, + final_states_out, + activation in ["silu", "swish"], + ) + return out, state_passing_conv_initial_states + + +def _causal_conv1d_state_passing_cp_bwd( + x, + weight, + bias, + dout, + seq_idx=None, + initial_states=None, + dfinal_states=None, + dx=None, + return_dinitial_states=False, + activation=None, + state_passing_cp_group=None, + state_passing_cp_virtual=False, +): + """Run causal Conv1d backward and route its state-passing CP halo gradient.""" + assert causal_conv1d_bwd_function is not None + assert state_passing_cp_group is not None + assert seq_idx is None, "seq_idx is not implemented for state-passing CP" + assert dfinal_states is None + assert not return_dinitial_states, "the state-passing boundary gradient is consumed internally" + if dout.stride(2) != 1 and dout.stride(1) != 1: + dout = dout.contiguous() + dx, dweight, dbias, dinitial_states = causal_conv1d_bwd_function( + x, + weight, + bias, + dout, + seq_idx, + initial_states, + dfinal_states, + dx, + True, + activation in ["silu", "swish"], + ) + halo = weight.shape[-1] - 1 + grad_tail = _gather_causal_conv1d_state_passing_cp_boundary( + dinitial_states, + state_passing_cp_group, + state_passing_cp_virtual=state_passing_cp_virtual, + reverse=True, + channel_last_output=False, + ) + dx[..., -halo:].add_(grad_tail) + return dx, dweight, dbias + + +class MambaSplitConv1dScanCombinedStatePassingCPFn(torch.autograd.Function): + """State-passing CP counterpart of ``MambaSplitConv1dScanCombinedFn``.""" + + @staticmethod + @custom_fwd + def forward( + ctx, + zxbcdt, + conv1d_weight, + conv1d_bias, + dt_bias, + A, + D, + chunk_size, + initial_states=None, + seq_idx=None, + dt_limit=(0.0, float("inf")), + return_final_states=False, + activation="silu", + rmsnorm_weight=None, + rmsnorm_eps=1e-6, + outproj_weight=None, + outproj_bias=None, + headdim=None, + ngroups=1, + norm_before_gate=True, + state_passing_cp_group=None, + state_passing_cp_virtual=False, + ): + assert activation in [None, "silu", "swish"] + assert state_passing_cp_group is not None + assert ( + initial_states is None + ), "external initial_states are not implemented for state-passing CP" + assert seq_idx is None, "seq_idx packed input is not implemented for state-passing CP" + assert ( + not return_final_states + ), "return_final_states is not implemented for state-passing CP" + assert rmsnorm_weight is None, "fused RMSNorm is not implemented for state-passing CP" + assert outproj_weight is None and outproj_bias is None + if D.dim() == 1: + assert headdim is not None + (nheads,) = D.shape + else: + nheads, headdim = D.shape + batch, seqlen, _ = zxbcdt.shape + dim = nheads * headdim + assert nheads % ngroups == 0 + dstate = (conv1d_weight.shape[0] - dim) // ngroups // 2 + d_nonssm = (zxbcdt.shape[-1] - 2 * dim - 2 * ngroups * dstate - nheads) // 2 + assert d_nonssm >= 0 + assert d_nonssm == 0, "non-SSM channels are not implemented for state-passing CP" + assert zxbcdt.shape == ( + batch, + seqlen, + 2 * d_nonssm + 2 * dim + 2 * ngroups * dstate + nheads, + ) + assert dt_bias.shape == (nheads,) + assert A.shape == (nheads,) + zx0, z, xBC, dt = torch.split( + zxbcdt, [2 * d_nonssm, dim, dim + 2 * ngroups * dstate, nheads], dim=-1 + ) + seq_idx = seq_idx.contiguous() if seq_idx is not None else None + xBC_conv, state_passing_conv_initial_states = _causal_conv1d_state_passing_cp_fwd( + rearrange(ensure_stride(xBC), "b s d -> b d s"), + conv1d_weight, + conv1d_bias, + seq_idx, + None, + None, + activation, + state_passing_cp_group, + state_passing_cp_virtual, + ) + xBC_conv = rearrange(xBC_conv, "b d s -> b s d") + x, B, C = torch.split(xBC_conv, [dim, ngroups * dstate, ngroups * dstate], dim=-1) + x = rearrange(x, "b l (h p) -> b l h p", h=nheads) + B = rearrange(B, "b l (g n) -> b l g n", g=ngroups) + C = rearrange(C, "b l (g n) -> b l g n", g=ngroups) + z = rearrange(z, "b l (h p) -> b l h p", h=nheads) if z is not None else None + + ( + out, + out_x, + dt_out, + dA_cumsum, + states, + final_states, + state_passing_initial_states, + state_passing_gathered_decays, + ) = _mamba_chunk_scan_combined_state_passing_cp_fwd( + x, + dt, + A, + B, + C, + chunk_size=chunk_size, + D=D, + z=z, + dt_bias=dt_bias, + initial_states=initial_states, + seq_idx=seq_idx, + dt_softplus=True, + dt_limit=dt_limit, + state_passing_cp_group=state_passing_cp_group, + state_passing_cp_virtual=state_passing_cp_virtual, + ) + out = rearrange(out, "b s h p -> b s (h p)") + rstd = None + + ctx.save_for_backward( + zxbcdt, + conv1d_weight, + conv1d_bias, + out_x, + A, + D, + dt_bias, + initial_states, + seq_idx, + rmsnorm_weight, + rstd, + outproj_weight, + outproj_bias, + state_passing_conv_initial_states, + state_passing_initial_states, + state_passing_gathered_decays, + ) + ctx.dt_limit = dt_limit + ctx.return_final_states = return_final_states + ctx.activation = activation + ctx.rmsnorm_eps = rmsnorm_eps + ctx.norm_before_gate = norm_before_gate + ctx.chunk_size = chunk_size + ctx.headdim = headdim + ctx.ngroups = ngroups + ctx.state_passing_cp_group = state_passing_cp_group + ctx.state_passing_cp_virtual = state_passing_cp_virtual + return out if not return_final_states else (out, final_states) + + @staticmethod + @custom_bwd + def backward(ctx, dout, *args): + ( + zxbcdt, + conv1d_weight, + conv1d_bias, + out, + A, + D, + dt_bias, + initial_states, + seq_idx, + rmsnorm_weight, + rstd, + outproj_weight, + outproj_bias, + state_passing_conv_initial_states, + state_passing_initial_states, + state_passing_gathered_decays, + ) = ctx.saved_tensors + dfinal_states = args[0] if ctx.return_final_states else None + headdim = ctx.headdim + nheads = D.shape[0] + dim = nheads * headdim + assert nheads % ctx.ngroups == 0 + dstate = (conv1d_weight.shape[0] - dim) // ctx.ngroups // 2 + d_nonssm = (zxbcdt.shape[-1] - 2 * dim - 2 * ctx.ngroups * dstate - nheads) // 2 + assert d_nonssm >= 0 + assert d_nonssm == 0 + recompute_output = outproj_weight is not None + + zx0, z, xBC, dt = torch.split( + zxbcdt, [2 * d_nonssm, dim, dim + 2 * ctx.ngroups * dstate, nheads], dim=-1 + ) + # Recompute x, B, C. + xBC_conv, _ = _causal_conv1d_state_passing_cp_fwd( + rearrange(ensure_stride(xBC), "b s d -> b d s"), + conv1d_weight, + conv1d_bias, + seq_idx, + state_passing_conv_initial_states, + None, + ctx.activation, + ctx.state_passing_cp_group, + ctx.state_passing_cp_virtual, + ) + xBC_conv = rearrange(xBC_conv, "b d s -> b s d") + x, B, C = torch.split(xBC_conv, [dim, ctx.ngroups * dstate, ctx.ngroups * dstate], dim=-1) + x = rearrange(x, "b l (h p) -> b l h p", h=nheads) + B = rearrange(B, "b l (g n) -> b l g n", g=ctx.ngroups) + C = rearrange(C, "b l (g n) -> b l g n", g=ctx.ngroups) + + dzxbcdt = torch.empty_like(zxbcdt) + dzx0, dz, dxBC_given, ddt_given = torch.split( + dzxbcdt, [2 * d_nonssm, dim, dim + 2 * ctx.ngroups * dstate, nheads], dim=-1 + ) + dxBC = torch.empty_like(xBC) + dx, dB, dC = torch.split(dxBC, [dim, ctx.ngroups * dstate, ctx.ngroups * dstate], dim=-1) + z = rearrange(z, "b l (h p) -> b l h p", h=nheads) + dx = rearrange(dx, "b l (h p) -> b l h p", h=nheads) + dB = rearrange(dB, "b l (g n) -> b l g n", g=ctx.ngroups) + dC = rearrange(dC, "b l (g n) -> b l g n", g=ctx.ngroups) + dout = rearrange(dout, "b s (h p) -> b s h p", p=headdim) + dz = rearrange(dz, "b l (h p) -> b l h p", h=nheads) + + dx, ddt, dA, dB, dC, dD, dz, ddt_bias, dinitial_states = ( + _mamba_chunk_scan_combined_state_passing_cp_bwd( + dout, + x, + dt, + A, + B, + C, + out, + ctx.chunk_size, + D=D, + z=z, + dt_bias=dt_bias, + initial_states=initial_states, + dfinal_states=dfinal_states, + seq_idx=seq_idx, + dt_softplus=True, + dt_limit=ctx.dt_limit, + dx=dx, + ddt=ddt_given, + dB=dB, + dC=dC, + dz=dz, + recompute_output=recompute_output, + state_passing_cp_group=ctx.state_passing_cp_group, + state_passing_initial_states=state_passing_initial_states, + state_passing_gathered_decays=state_passing_gathered_decays, + state_passing_cp_virtual=ctx.state_passing_cp_virtual, + ) + ) + + dxBC_given_update, dweight, dbias = _causal_conv1d_state_passing_cp_bwd( + rearrange(ensure_stride(xBC), "b s d -> b d s"), + conv1d_weight, + conv1d_bias, + rearrange(ensure_stride(dxBC), "b s d -> b d s"), + seq_idx, + state_passing_conv_initial_states, + None, + rearrange(ensure_stride(dxBC_given), "b s d -> b d s"), + False, + ctx.activation, + ctx.state_passing_cp_group, + ctx.state_passing_cp_virtual, + ) + dxBC_given_update = rearrange(dxBC_given_update, "b d s -> b s d") + if dxBC_given.stride() != dxBC_given_update.stride(): + dxBC_given.copy_(dxBC_given_update) + else: + dxBC_given = dxBC_given_update + + drmsnorm_weight = None + doutproj_weight = None + doutproj_bias = None + return ( + dzxbcdt, + dweight, + dbias, + ddt_bias, + dA, + dD, + None, + dinitial_states, + None, + None, + None, + None, + drmsnorm_weight, + None, + doutproj_weight, + doutproj_bias, + None, + None, + None, + None, + None, + ) + + +def mamba_split_conv1d_scan_combined_state_passing_cp( + zxbcdt, + conv1d_weight, + conv1d_bias, + dt_bias, + A, + D, + chunk_size, + initial_states=None, + seq_idx=None, + dt_limit=(0.0, float("inf")), + return_final_states=False, + activation="silu", + rmsnorm_weight=None, + rmsnorm_eps=1e-6, + outproj_weight=None, + outproj_bias=None, + headdim=None, + ngroups=1, + norm_before_gate=True, + state_passing_cp_group=None, + state_passing_cp_virtual=False, +): + """State-passing CP counterpart of ``mamba_split_conv1d_scan_combined``.""" + return MambaSplitConv1dScanCombinedStatePassingCPFn.apply( + zxbcdt, + conv1d_weight, + conv1d_bias, + dt_bias, + A, + D, + chunk_size, + initial_states, + seq_idx, + dt_limit, + return_final_states, + activation, + rmsnorm_weight, + rmsnorm_eps, + outproj_weight, + outproj_bias, + headdim, + ngroups, + norm_before_gate, + state_passing_cp_group, + state_passing_cp_virtual, + ) + + +class MambaStatePassingCPAdapter: + """Adapter between ``MambaMixer`` and the state-passing CP production path.""" + + def __init__(self, mixer): + self.mixer = mixer + + def forward(self, zxBCdt: torch.Tensor) -> torch.Tensor: + """Run the fused state-passing Conv+SSD path on a projected activation.""" + mixer = self.mixer + state_passing_cp_load_balancing = mixer.config.mamba_state_passing_cp_load_balancing + assert state_passing_cp_load_balancing in ("none", "permute_p2p", "permute_a2a", "virtual") + if state_passing_cp_load_balancing.startswith("permute_"): + state_passing_cp_permute_backend = state_passing_cp_load_balancing.removeprefix( + "permute_" + ) + zxBCdt = undo_state_passing_cp_load_balancing( + zxBCdt, mixer.cp.cp_group, backend=state_passing_cp_permute_backend + ) + zxBCdt = rearrange(zxBCdt, "l b d -> b l d").contiguous() + state_passing_cp_virtual = state_passing_cp_load_balancing == "virtual" + if state_passing_cp_virtual: + zxBCdt = _pack_state_passing_cp_virtual_segments(zxBCdt) + A = -torch.exp(mixer.A_log.float()) + D = ( + rearrange(mixer.D.float(), "(h p) -> h p", p=mixer.headdim) + if mixer.D_has_hdim + else mixer.D + ) + + assert ( + causal_conv1d_fwd_function is not None and causal_conv1d_bwd_function is not None + ), "Mamba state-passing CP requires causal-conv1d" + assert mixer.activation in ["silu", "swish"] + y = mamba_split_conv1d_scan_combined_state_passing_cp( + zxBCdt, + rearrange(mixer.conv1d_weight, "d 1 w -> d w"), + mixer.conv1d_bias, + mixer.dt_bias.float(), + A, + D, + mixer.chunk_size, + activation=mixer.activation, + headdim=None if mixer.D_has_hdim else mixer.headdim, + ngroups=mixer.ngroups_local_tp, + norm_before_gate=mixer.norm_before_gate, + state_passing_cp_group=mixer.cp.cp_group, + state_passing_cp_virtual=state_passing_cp_virtual, + ) + + if mixer.rmsnorm: + # Match MambaMixer's fused path: materialize the BF16 gated scan + # output before applying RMSNorm. + y = mixer.norm(y) + if state_passing_cp_virtual: + y = _unpack_state_passing_cp_virtual_segments(y) + y = rearrange(y, "b l d -> l b d").contiguous() + if state_passing_cp_load_balancing.startswith("permute_"): + y = redo_state_passing_cp_load_balancing( + y, mixer.cp.cp_group, backend=state_passing_cp_permute_backend + ) + return y diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index d6afac69c48..6c09fd57181 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1362,6 +1362,29 @@ class TransformerConfig(ModelParallelConfig): ) """Controls usage of the memory efficient path for Mamba layers.""" + use_mamba_state_passing_cp: bool = field( + default=False, metadata={"argparse_meta": {"arg_names": ["--use-mamba-state-passing-cp"]}} + ) + """Use experimental sequence-sharded state-passing context parallel for Mamba training.""" + + mamba_state_passing_cp_load_balancing: Literal[ + "none", "permute_p2p", "permute_a2a", "virtual" + ] = field( + default="permute_p2p", + metadata={ + "argparse_meta": { + "arg_names": ["--mamba-state-passing-cp-load-balancing"], + "choices": ["none", "permute_p2p", "permute_a2a", "virtual"], + } + }, + ) + """How state-passing Mamba CP handles the standard balanced CP sequence layout. + + ``none`` expects contiguous causal shards, ``permute_p2p`` and ``permute_a2a`` + exchange activation chunks with the selected communication method, and ``virtual`` + treats each balanced front/back chunk as an independent virtual rank. + """ + mlp_chunks_for_prefill: int = 1 """The number of chunks along the sequence dimension to use for MLP computation during prefill.""" diff --git a/tests/unit_tests/ssm/mamba_state_passing_cp_utils.py b/tests/unit_tests/ssm/mamba_state_passing_cp_utils.py new file mode 100644 index 00000000000..0e8357b25a9 --- /dev/null +++ b/tests/unit_tests/ssm/mamba_state_passing_cp_utils.py @@ -0,0 +1,276 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Shared helpers for Mamba state-passing context-parallel unit tests.""" + +from dataclasses import dataclass, replace + +import torch +import torch.distributed as dist + +from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.mamba_block import MambaStackSubmodules +from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules +from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules +from megatron.core.transformer import TransformerConfig + +# Relative RMS error tolerance used for BF16 comparisons against a +# full-sequence reference. +DEFAULT_ERROR_RATIO = 3e-3 + +# Load-balancing modes exercised by the state-passing CP tests. +STATE_PASSING_CP_MODES = ("permute_p2p", "permute_a2a", "virtual") + + +def balanced_cp_chunk_ids(rank: int, cp_size: int) -> tuple[int, int]: + """Return the front/back chunk indices Megatron's balanced CP assigns to ``rank``.""" + assert cp_size > 0 and 0 <= rank < cp_size + return rank, 2 * cp_size - 1 - rank + + +def select_balanced_cp_shard( + tensor: torch.Tensor, rank: int, cp_size: int, *, sequence_dim: int = 0 +) -> torch.Tensor: + """Select Megatron's front/back balanced CP shard without packing the batch axis.""" + assert tensor.shape[sequence_dim] % (2 * cp_size) == 0 + chunks = torch.chunk(tensor, 2 * cp_size, dim=sequence_dim) + front, back = balanced_cp_chunk_ids(rank, cp_size) + return torch.cat((chunks[front], chunks[back]), dim=sequence_dim).contiguous() + + +def select_contiguous_cp_shard( + tensor: torch.Tensor, rank: int, cp_size: int, *, sequence_dim: int = 0 +) -> torch.Tensor: + """Select the two causally adjacent chunks owned by a contiguous CP rank.""" + assert tensor.shape[sequence_dim] % (2 * cp_size) == 0 + chunks = torch.chunk(tensor, 2 * cp_size, dim=sequence_dim) + return torch.cat((chunks[2 * rank], chunks[2 * rank + 1]), dim=sequence_dim).contiguous() + + +def select_state_passing_cp_shard( + tensor: torch.Tensor, + rank: int, + cp_size: int, + *, + virtual: bool, + batch_dim: int = 0, + sequence_dim: int = 1, +) -> torch.Tensor: + """Select a contiguous shard, or pack balanced front/back chunks on the batch axis. + + The state-passing kernels consume either a contiguous causal shard or, in + ``virtual`` mode, an interleaved view where each balanced front/back chunk + becomes an independent virtual batch entry. + """ + assert cp_size > 0 and 0 <= rank < cp_size + assert tensor.shape[sequence_dim] % cp_size == 0 + if not virtual: + local_length = tensor.shape[sequence_dim] // cp_size + return tensor.narrow(sequence_dim, rank * local_length, local_length).contiguous() + + assert batch_dim != sequence_dim + moved = tensor.movedim((batch_dim, sequence_dim), (0, 1)) + assert moved.shape[1] % (2 * cp_size) == 0 + chunks = torch.chunk(moved, 2 * cp_size, dim=1) + front, back = balanced_cp_chunk_ids(rank, cp_size) + packed = torch.stack((chunks[front], chunks[back]), dim=1).flatten(0, 1).contiguous() + return packed.movedim((0, 1), (batch_dim, sequence_dim)) + + +def relative_rms_error(actual: torch.Tensor, expected: torch.Tensor) -> float: + """Return the RMS error of ``actual`` relative to the RMS magnitude of ``expected``. + + A relative criterion is used instead of elementwise ``allclose`` because the + state-passing path reassociates the SSD scan across CP boundaries, which + changes BF16 rounding without changing the mathematical result. + """ + actual_float = actual.float() + expected_float = expected.float() + assert torch.isfinite(actual_float).all(), "state-passing result contains non-finite values" + assert torch.isfinite(expected_float).all(), "reference contains non-finite values" + difference = actual_float - expected_float + if difference.abs().max().item() == 0.0: + return 0.0 + expected_rms = expected_float.square().mean().sqrt().item() + return difference.square().mean().sqrt().item() / (expected_rms + 1e-8) + + +def assert_all_close_rms( + checks: dict[str, tuple[torch.Tensor, torch.Tensor]], + group: dist.ProcessGroup, + ratio: float = DEFAULT_ERROR_RATIO, +) -> None: + """Compare every (actual, expected) pair and fail identically on all CP ranks. + + Every comparison is evaluated before anything is asserted, and the verdict is + all-reduced. Asserting per comparison would let one rank leave the test while + the others are still inside a collective, which deadlocks the whole run + instead of reporting a failure. + """ + errors = {name: relative_rms_error(*pair) for name, pair in checks.items()} + names = list(errors) + local = torch.tensor([errors[name] for name in names], device=torch.cuda.current_device()) + dist.all_reduce(local, op=dist.ReduceOp.MAX, group=group) + failures = [ + f"{name}: relative RMS error {error:.3e} exceeds {ratio:.3e}" + for name, error in zip(names, local.tolist()) + if not error < ratio + ] + assert not failures, "state-passing CP mismatch on at least one rank:\n" + "\n".join(failures) + + +def broadcast_module_parameters(module: torch.nn.Module, group: dist.ProcessGroup) -> None: + """Make every CP rank hold the same weights. + + The CP path combines activations produced from every rank's parameters, so a + per-rank difference in initialization would invalidate any comparison against + a rank-local full-sequence reference. + """ + source = dist.get_global_rank(group, 0) + for tensor in list(module.parameters()) + list(module.buffers()): + dist.broadcast(tensor.data, src=source, group=group) + + +def collect_parameter_grads( + module: torch.nn.Module, reduce_group: dist.ProcessGroup | None = None +) -> dict[str, torch.Tensor]: + """Collect parameter grads as FP32, optionally summing replicated grads across CP.""" + gradients = {} + for name, parameter in module.named_parameters(): + if parameter.grad is None: + continue + gradient = parameter.grad.detach().float().clone() + if reduce_group is not None: + dist.all_reduce(gradient, group=reduce_group) + gradients[name] = gradient + return gradients + + +@dataclass(frozen=True) +class MambaModelShape: + """TP-local Mamba mixer shape used to build test models.""" + + hidden_size: int = 2688 + nheads: int = 64 + head_dim: int = 64 + state_dim: int = 128 + ngroups: int = 8 + chunk_size: int = 128 + d_conv: int = 4 + + @property + def d_inner(self) -> int: + """Inner (gated) width of the mixer.""" + return self.nheads * self.head_dim + + @property + def conv_dim(self) -> int: + """Number of channels the depthwise causal convolution operates on.""" + return self.d_inner + 2 * self.ngroups * self.state_dim + + @property + def projected_width(self) -> int: + """Width of the ``in_proj`` output (z, x, B, C, dt).""" + return 2 * self.d_inner + 2 * self.ngroups * self.state_dim + self.nheads + + +def build_mamba_config( + shape: MambaModelShape, + cp_size: int, + *, + load_balancing: str, + use_state_passing_cp: bool = True, + cuda_graph_impl: str = "none", + cuda_graph_warmup_steps: int = 2, +) -> TransformerConfig: + """Build a single-layer BF16 Mamba config for state-passing CP tests.""" + return TransformerConfig( + hidden_size=shape.hidden_size, + num_layers=1, + num_attention_heads=1, + bf16=True, + params_dtype=torch.bfloat16, + pipeline_dtype=torch.bfloat16, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=cp_size, + mamba_num_heads=shape.nheads, + mamba_head_dim=shape.head_dim, + mamba_state_dim=shape.state_dim, + mamba_num_groups=shape.ngroups, + use_cpu_initialization=True, + use_mamba_mem_eff_path=True, + use_mamba_state_passing_cp=use_state_passing_cp, + mamba_state_passing_cp_load_balancing=load_balancing, + cuda_graph_impl=cuda_graph_impl, + cuda_graph_scope=[], + cuda_graph_warmup_steps=cuda_graph_warmup_steps, + # Both CUDA Graph backends require the TE RNG tracker: the default + # tracker refuses to be captured. + use_te_rng_tracker=cuda_graph_impl != "none", + hidden_dropout=0.0, + ) + + +def _mamba_layer_spec(): + assert isinstance(mamba_stack_spec.submodules, MambaStackSubmodules) + spec = mamba_stack_spec.submodules.mamba_layer + assert isinstance(spec.submodules, MambaLayerSubmodules) + return spec + + +def build_mamba_mixer( + shape: MambaModelShape, cp_group: dist.ProcessGroup, tp_group: dist.ProcessGroup, **kwargs +) -> MambaMixer: + """Build a ``MambaMixer`` bound to ``cp_group`` and move it to the current device.""" + config = build_mamba_config(shape, cp_group.size(), **kwargs) + mixer_spec = _mamba_layer_spec().submodules.mixer + assert isinstance(mixer_spec.submodules, MambaMixerSubmodules) + mixer = MambaMixer( + config, + mixer_spec.submodules, + shape.hidden_size, + d_conv=shape.d_conv, + chunk_size=shape.chunk_size, + layer_number=1, + pg_collection=ProcessGroupCollection(tp=tp_group, cp=cp_group), + ).to(torch.cuda.current_device()) + mixer.train() + assert mixer.conv1d_weight.shape[0] == shape.conv_dim + assert mixer.in_proj.weight.shape[0] == shape.projected_width + return mixer + + +def build_mamba_layer( + shape: MambaModelShape, cp_group: dist.ProcessGroup, tp_group: dist.ProcessGroup, **kwargs +) -> tuple[MambaLayer, TransformerConfig]: + """Build a ``MambaLayer`` bound to ``cp_group`` and move it to the current device.""" + config = build_mamba_config(shape, cp_group.size(), **kwargs) + spec = _mamba_layer_spec() + mixer_spec = replace( + spec.submodules.mixer, + params={ + **spec.submodules.mixer.params, + "d_conv": shape.d_conv, + "chunk_size": shape.chunk_size, + }, + ) + layer = MambaLayer( + config, + replace(spec.submodules, mixer=mixer_spec), + layer_number=1, + pg_collection=ProcessGroupCollection(tp=tp_group, cp=cp_group), + ).to(torch.cuda.current_device()) + layer.train() + return layer, config + + +def set_state_passing_cp_mode(mixer: MambaMixer, mode: str | None) -> None: + """Switch a built mixer between the A2A CP path (``None``) and a state-passing mode.""" + if mode is None: + mixer.use_mamba_state_passing_cp = False + return + assert mode in STATE_PASSING_CP_MODES + mixer.use_mamba_state_passing_cp = True + mixer.mamba_state_passing_cp_load_balancing = mode + mixer.config.mamba_state_passing_cp_load_balancing = mode diff --git a/tests/unit_tests/ssm/ops/test_ssd_state_passing_cp.py b/tests/unit_tests/ssm/ops/test_ssd_state_passing_cp.py new file mode 100644 index 00000000000..2ca8c4bd132 --- /dev/null +++ b/tests/unit_tests/ssm/ops/test_ssd_state_passing_cp.py @@ -0,0 +1,305 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Kernel-level tests for the Mamba2 state-passing context-parallel path. + +Each test builds a full-sequence reference on every rank, shards it the way +Megatron's context parallelism would, and checks that the state-passing kernels +reproduce the reference shard for both the forward output and every gradient. + +The mixer dimensions are Nemotron-3 Nano's (64 heads of 64, 8 groups, state 128, +chunk 128); the batch size and sequence length are what the parametrization +varies, because those are what the CP sharding and the boundary exchange +actually depend on. +""" + +import pytest +import torch + +from tests.unit_tests.ssm.mamba_state_passing_cp_utils import ( + MambaModelShape, + assert_all_close_rms, + select_balanced_cp_shard, + select_contiguous_cp_shard, + select_state_passing_cp_shard, +) +from tests.unit_tests.test_utilities import Utils + +try: + from causal_conv1d import causal_conv1d_fn + from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined + + from megatron.core.ssm.ops.ssd_state_passing_cp import ( + _causal_conv1d_state_passing_cp_bwd, + _causal_conv1d_state_passing_cp_fwd, + _mamba_chunk_scan_combined_state_passing_cp_bwd, + _mamba_chunk_scan_combined_state_passing_cp_fwd, + redo_state_passing_cp_load_balancing, + undo_state_passing_cp_load_balancing, + ) + + HAVE_STATE_PASSING_CP = True +except ImportError: + HAVE_STATE_PASSING_CP = False + +pytestmark = [ + pytest.mark.internal, + pytest.mark.skipif( + not HAVE_STATE_PASSING_CP, reason="mamba_ssm and causal_conv1d are required" + ), + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required"), +] + +SHAPE = MambaModelShape() + +# (batch, sequence_length). Every causal segment must align with the SSD chunk +# size, and ``virtual`` mode halves the segment length, so the sequence length +# must be a multiple of ``4 * cp_size * chunk_size``. +BATCH_AND_SEQUENCE_LENGTHS = ((1, 2048), (3, 4096)) + + +@pytest.fixture +def cp_context(request): + """Initialize a context-parallel group of the requested size and tear it down.""" + cp_size = request.param + if Utils.world_size % cp_size != 0: + pytest.skip(f"world size {Utils.world_size} is not a multiple of cp_size {cp_size}") + Utils.initialize_model_parallel(context_parallel_size=cp_size) + from megatron.core import parallel_state + + group = parallel_state.get_context_parallel_group() + yield group, group.rank(), cp_size, torch.device(torch.cuda.current_device()) + Utils.destroy_model_parallel() + + +@pytest.mark.parametrize("cp_context", (2, 4), indirect=True) +@pytest.mark.parametrize("backend", ("p2p", "a2a")) +def test_load_balancing_permutation_round_trip(cp_context, backend): + """The balanced-to-contiguous permutation and its backward must be exact.""" + group, rank, cp_size, device = cp_context + sequence_length, hidden_size = 2048, SHAPE.hidden_size + + total_chunks = 2 * cp_size + chunk_length = sequence_length // total_chunks + full = torch.arange(sequence_length * hidden_size, device=device, dtype=torch.float32).view( + sequence_length, 1, hidden_size + ) + grad_full = torch.empty_like(full) + for chunk_id in range(total_chunks): + grad_full[chunk_id * chunk_length : (chunk_id + 1) * chunk_length].fill_(chunk_id + 1) + + balanced = select_balanced_cp_shard(full, rank, cp_size).clone().requires_grad_(True) + contiguous = undo_state_passing_cp_load_balancing(balanced, group, backend=backend) + restored = redo_state_passing_cp_load_balancing(contiguous, group, backend=backend) + (contiguous * select_contiguous_cp_shard(grad_full, rank, cp_size)).sum().backward() + + torch.testing.assert_close( + contiguous, select_contiguous_cp_shard(full, rank, cp_size), rtol=0, atol=0 + ) + torch.testing.assert_close(restored, balanced.detach(), rtol=0, atol=0) + torch.testing.assert_close( + balanced.grad, select_balanced_cp_shard(grad_full, rank, cp_size), rtol=0, atol=0 + ) + + +@pytest.mark.parametrize("cp_context", (2, 4), indirect=True) +@pytest.mark.parametrize("virtual", (False, True)) +@pytest.mark.parametrize("batch,sequence_length", BATCH_AND_SEQUENCE_LENGTHS) +def test_causal_conv1d_matches_full_sequence(cp_context, virtual, batch, sequence_length): + """The halo-exchanging causal conv must match a full-sequence causal_conv1d.""" + group, rank, cp_size, device = cp_context + channels = SHAPE.conv_dim + dtype = torch.bfloat16 + + generator = torch.Generator(device=device).manual_seed(1234) + x = torch.randn( + batch, channels, sequence_length, device=device, dtype=dtype, generator=generator + ) + weight = torch.randn(channels, SHAPE.d_conv, device=device, dtype=dtype, generator=generator) + bias = torch.randn(channels, device=device, dtype=dtype, generator=generator) + # Every input must come from the explicit seeded generator: the per-rank RNG + # that ``randn_like`` would use differs across ranks, which would give each + # rank a different "full-sequence" reference. + grad_output = torch.randn( + batch, channels, sequence_length, device=device, dtype=dtype, generator=generator + ) + + x_reference = x.clone().requires_grad_(True) + weight_reference = weight.clone().requires_grad_(True) + bias_reference = bias.clone().requires_grad_(True) + y_reference = causal_conv1d_fn(x_reference, weight_reference, bias_reference, activation="silu") + (y_reference * grad_output).sum().backward() + + def local_view(tensor): + return select_state_passing_cp_shard( + tensor, rank, cp_size, virtual=virtual, batch_dim=0, sequence_dim=2 + ) + + # causal_conv1d requires a channel-last-in-memory layout. + x_local = local_view(x).transpose(1, 2).contiguous().transpose(1, 2) + y_local, conv_initial_states = _causal_conv1d_state_passing_cp_fwd( + x_local, + weight.clone(), + bias.clone(), + activation="silu", + state_passing_cp_group=group, + state_passing_cp_virtual=virtual, + ) + dx_local, dweight_local, dbias_local = _causal_conv1d_state_passing_cp_bwd( + x_local, + weight.clone(), + bias.clone(), + local_view(grad_output), + initial_states=conv_initial_states, + activation="silu", + state_passing_cp_group=group, + state_passing_cp_virtual=virtual, + ) + + # Weight and bias are replicated across CP, so their gradients are summed. + dweight = dweight_local.detach().float() + dbias = dbias_local.detach().float() + torch.distributed.all_reduce(dweight, group=group) + torch.distributed.all_reduce(dbias, group=group) + + assert_all_close_rms( + { + "conv output": (y_local, local_view(y_reference)), + "conv dx": (dx_local, local_view(x_reference.grad)), + "conv dweight": (dweight, weight_reference.grad), + "conv dbias": (dbias, bias_reference.grad), + }, + group, + ) + + +@pytest.mark.parametrize("cp_context", (2, 4), indirect=True) +@pytest.mark.parametrize("virtual", (False, True)) +@pytest.mark.parametrize("batch,sequence_length", BATCH_AND_SEQUENCE_LENGTHS) +def test_chunk_scan_combined_matches_full_sequence(cp_context, virtual, batch, sequence_length): + """The boundary-scanning SSD kernel must match a full-sequence chunk scan.""" + group, rank, cp_size, device = cp_context + nheads, headdim, ngroups = SHAPE.nheads, SHAPE.head_dim, SHAPE.ngroups + dstate, chunk_size = SHAPE.state_dim, SHAPE.chunk_size + segment_length = sequence_length // (cp_size * (2 if virtual else 1)) + assert ( + segment_length % chunk_size == 0 + ), f"each causal segment ({segment_length}) must align with the SSD chunk size" + dtype = torch.bfloat16 + kernel_kwargs = dict(dt_softplus=True, dt_limit=(0.0, float("inf"))) + + generator = torch.Generator(device=device).manual_seed(1234) + x = torch.randn( + batch, sequence_length, nheads, headdim, device=device, dtype=dtype, generator=generator + ) + dt = ( + torch.rand(batch, sequence_length, nheads, device=device, dtype=dtype, generator=generator) + * 0.5 + + 0.01 + ) + A_log = torch.log( + torch.rand(nheads, device=device, dtype=torch.float32, generator=generator) + 0.1 + ) + B = torch.randn( + batch, sequence_length, ngroups, dstate, device=device, dtype=dtype, generator=generator + ) + C = torch.randn( + batch, sequence_length, ngroups, dstate, device=device, dtype=dtype, generator=generator + ) + # Every input must come from the explicit seeded generator: the per-rank RNG + # that ``randn_like`` would use differs across ranks, which would give each + # rank a different "full-sequence" reference. + z = torch.randn( + batch, sequence_length, nheads, headdim, device=device, dtype=dtype, generator=generator + ) + D = torch.randn(nheads, device=device, dtype=torch.float32, generator=generator) + dt_bias = torch.randn(nheads, device=device, dtype=dtype, generator=generator) * 0.1 + grad_output = torch.randn( + batch, sequence_length, nheads, headdim, device=device, dtype=dtype, generator=generator + ) + + def local_view(tensor): + return select_state_passing_cp_shard( + tensor, rank, cp_size, virtual=virtual, batch_dim=0, sequence_dim=1 + ) + + reference_inputs = [tensor.clone().requires_grad_(True) for tensor in (x, dt, B, C, z)] + A_log_reference = A_log.clone().requires_grad_(True) + A_reference = -torch.exp(A_log_reference) + A_reference.retain_grad() + D_reference = D.clone().requires_grad_(True) + dt_bias_reference = dt_bias.clone().requires_grad_(True) + y_reference = mamba_chunk_scan_combined( + reference_inputs[0], + reference_inputs[1], + A_reference, + reference_inputs[2], + reference_inputs[3], + chunk_size, + D=D_reference, + z=reference_inputs[4], + dt_bias=dt_bias_reference, + **kernel_kwargs, + ) + (y_reference * grad_output).sum().backward() + + A = -torch.exp(A_log) + cp_inputs = [local_view(tensor).clone().requires_grad_(True) for tensor in (x, dt, B, C, z)] + y_cp, out_x, _, _, _, _, initial_states, gathered_decays = ( + _mamba_chunk_scan_combined_state_passing_cp_fwd( + cp_inputs[0], + cp_inputs[1], + A, + cp_inputs[2], + cp_inputs[3], + chunk_size, + D=D, + z=cp_inputs[4], + dt_bias=dt_bias, + state_passing_cp_group=group, + state_passing_cp_virtual=virtual, + **kernel_kwargs, + ) + ) + cp_grads = _mamba_chunk_scan_combined_state_passing_cp_bwd( + local_view(grad_output), + cp_inputs[0], + cp_inputs[1], + A, + cp_inputs[2], + cp_inputs[3], + out_x, + chunk_size, + D=D, + z=cp_inputs[4], + dt_bias=dt_bias, + state_passing_initial_states=initial_states, + state_passing_gathered_decays=gathered_decays, + state_passing_cp_group=group, + state_passing_cp_virtual=virtual, + **kernel_kwargs, + ) + dx, ddt, dA, dB, dC, dD, dz, ddt_bias = cp_grads[:8] + + # A, D, and dt_bias are replicated across CP, so their gradients are summed. + reduced = [] + for gradient in (dA, dD, ddt_bias): + value = gradient.detach().float() + torch.distributed.all_reduce(value, group=group) + reduced.append(value) + dA_reduced, dD_reduced, ddt_bias_reduced = reduced + + assert_all_close_rms( + { + "ssd output": (y_cp, local_view(y_reference)), + "ssd dx": (dx, local_view(reference_inputs[0].grad)), + "ssd ddt": (ddt, local_view(reference_inputs[1].grad)), + "ssd dB": (dB, local_view(reference_inputs[2].grad)), + "ssd dC": (dC, local_view(reference_inputs[3].grad)), + "ssd dz": (dz, local_view(reference_inputs[4].grad)), + "ssd dA": (dA_reduced, A_reference.grad), + "ssd dA_log": (dA_reduced * A, A_log_reference.grad), + "ssd dD": (dD_reduced, D_reference.grad), + "ssd ddt_bias": (ddt_bias_reduced, dt_bias_reference.grad), + }, + group, + ) diff --git a/tests/unit_tests/ssm/test_mamba_mixer_state_passing_cp.py b/tests/unit_tests/ssm/test_mamba_mixer_state_passing_cp.py new file mode 100644 index 00000000000..0f8d44cd47b --- /dev/null +++ b/tests/unit_tests/ssm/test_mamba_mixer_state_passing_cp.py @@ -0,0 +1,138 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""End-to-end tests for ``MambaMixer`` with state-passing context parallelism. + +Each state-passing load-balancing mode is compared against a single-rank +full-sequence ``MambaMixer`` and against the existing all-to-all CP path, for +the output, the input gradient, and every parameter gradient. +""" + +import pytest +import torch + +from megatron.core import parallel_state +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from tests.unit_tests.ssm.mamba_state_passing_cp_utils import ( + STATE_PASSING_CP_MODES, + MambaModelShape, + assert_all_close_rms, + broadcast_module_parameters, + build_mamba_mixer, + collect_parameter_grads, + select_balanced_cp_shard, + set_state_passing_cp_mode, +) +from tests.unit_tests.test_utilities import Utils + +try: + import causal_conv1d # noqa: F401 + import mamba_ssm # noqa: F401 + + from megatron.core.ssm.mamba_mixer import HAVE_MAMBA_STATE_PASSING_CP +except ImportError: + HAVE_MAMBA_STATE_PASSING_CP = False + +pytestmark = [ + pytest.mark.internal, + pytest.mark.skipif( + not HAVE_MAMBA_STATE_PASSING_CP, reason="mamba_ssm and causal_conv1d are required" + ), + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required"), +] + + +def _run_mixer(mixer, mode, hidden_states, grad_output, parameter_reduce_group=None): + """Run one forward/backward of ``mixer`` in ``mode`` and collect all gradients.""" + mixer.zero_grad(set_to_none=True) + local_input = hidden_states.detach().clone().requires_grad_(True) + set_state_passing_cp_mode(mixer, mode) + output, output_bias = mixer(local_input) + assert output_bias is None, "this test expects a bias-free output projection" + torch.autograd.backward(output, grad_output) + return ( + output.detach(), + local_input.grad.detach(), + collect_parameter_grads(mixer, parameter_reduce_group), + ) + + +@pytest.mark.parametrize("cp_size", (2, 4)) +@pytest.mark.parametrize("batch_size,sequence_length", ((1, 2048), (3, 4096))) +def test_mixer_state_passing_cp_matches_full_sequence_and_a2a(cp_size, batch_size, sequence_length): + """All state-passing modes must match the full-sequence and A2A CP references.""" + if Utils.world_size % cp_size != 0: + pytest.skip(f"world size {Utils.world_size} is not a multiple of cp_size {cp_size}") + Utils.initialize_model_parallel(context_parallel_size=cp_size) + try: + torch.manual_seed(1234) + model_parallel_cuda_manual_seed(1234, force_reset_rng=True) + cp_group = parallel_state.get_context_parallel_group() + tp_group = parallel_state.get_tensor_model_parallel_group() + assert tp_group.size() == 1, "the full-sequence reference needs a size-1 CP group" + rank = cp_group.rank() + device = torch.device(torch.cuda.current_device()) + + shape = MambaModelShape() + # ``virtual`` mode halves the segment length, and every causal segment + # must align with the SSD chunk size. + assert sequence_length % (4 * cp_size * shape.chunk_size) == 0 + + mixer = build_mamba_mixer(shape, cp_group, tp_group, load_balancing="permute_p2p") + broadcast_module_parameters(mixer, cp_group) + # A size-1 CP group makes this mixer process the whole sequence. + reference_mixer = build_mamba_mixer( + shape, tp_group, tp_group, load_balancing="permute_p2p", use_state_passing_cp=False + ) + reference_mixer.load_state_dict(mixer.state_dict()) + + generator = torch.Generator(device=device).manual_seed(5678) + global_input = torch.randn( + sequence_length, + batch_size, + shape.hidden_size, + device=device, + dtype=torch.bfloat16, + generator=generator, + ) + global_grad_output = torch.randn( + global_input.shape, device=device, dtype=torch.bfloat16, generator=generator + ) + local_input = select_balanced_cp_shard(global_input, rank, cp_size) + local_grad_output = select_balanced_cp_shard(global_grad_output, rank, cp_size) + + reference_output, reference_input_grad, reference_parameter_grads = _run_mixer( + reference_mixer, None, global_input, global_grad_output + ) + expected_output = select_balanced_cp_shard(reference_output, rank, cp_size) + expected_input_grad = select_balanced_cp_shard(reference_input_grad, rank, cp_size) + + # The existing all-to-all CP path is the second reference: it is the + # numerical behaviour production already ships. + a2a_output, a2a_input_grad, a2a_parameter_grads = _run_mixer( + mixer, None, local_input, local_grad_output, cp_group + ) + checks = { + "a2a output vs full sequence": (a2a_output, expected_output), + "a2a input grad vs full sequence": (a2a_input_grad, expected_input_grad), + } + for mode in STATE_PASSING_CP_MODES: + output, input_grad, parameter_grads = _run_mixer( + mixer, mode, local_input, local_grad_output, cp_group + ) + checks[f"{mode} output vs full sequence"] = (output, expected_output) + checks[f"{mode} input grad vs full sequence"] = (input_grad, expected_input_grad) + checks[f"{mode} output vs a2a"] = (output, a2a_output) + checks[f"{mode} input grad vs a2a"] = (input_grad, a2a_input_grad) + assert parameter_grads.keys() == reference_parameter_grads.keys() + for name, gradient in parameter_grads.items(): + checks[f"{mode} {name} vs full sequence"] = ( + gradient, + reference_parameter_grads[name], + ) + checks[f"{mode} {name} vs a2a"] = (gradient, a2a_parameter_grads[name]) + + # Asserted in one place so that a mismatch fails every rank identically + # instead of leaving some ranks inside a collective. + assert_all_close_rms(checks, cp_group) + finally: + Utils.destroy_model_parallel() diff --git a/tests/unit_tests/ssm/test_mamba_state_passing_cp_cuda_graph.py b/tests/unit_tests/ssm/test_mamba_state_passing_cp_cuda_graph.py new file mode 100644 index 00000000000..f8d89de28bb --- /dev/null +++ b/tests/unit_tests/ssm/test_mamba_state_passing_cp_cuda_graph.py @@ -0,0 +1,267 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""CUDA Graph tests for Mamba state-passing context parallelism. + +Both graph backends are covered: Megatron's local CUDA Graph lifecycle +(``cuda_graph_impl="local"``) and the Transformer Engine helper +(``cuda_graph_impl="transformer_engine"``). Each test records a graph, replays +it, checks bitwise parity with eager execution on the first replay, and then +changes the input to confirm that the static input buffer is actually updated. +""" + +import gc + +import pytest +import torch +import torch.distributed as dist +import torch.nn as nn + +from megatron.core import parallel_state +from megatron.core.num_microbatches_calculator import ( + init_num_microbatches_calculator, + unset_num_microbatches_calculator, +) +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.cuda_graphs import ( + TECudaGraphHelper, + create_cudagraphs, + delete_cuda_graphs, +) +from tests.unit_tests.ssm.mamba_state_passing_cp_utils import ( + STATE_PASSING_CP_MODES, + MambaModelShape, + build_mamba_layer, + select_balanced_cp_shard, +) +from tests.unit_tests.test_utilities import Utils + +try: + import causal_conv1d # noqa: F401 + import mamba_ssm # noqa: F401 + + from megatron.core.ssm.mamba_mixer import HAVE_MAMBA_STATE_PASSING_CP +except ImportError: + HAVE_MAMBA_STATE_PASSING_CP = False + +pytestmark = [ + pytest.mark.internal, + pytest.mark.skipif( + not HAVE_MAMBA_STATE_PASSING_CP, reason="mamba_ssm and causal_conv1d are required" + ), + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required"), +] + +CP_SIZE = 2 +INPUT_GRAD_RTOL = 1e-3 +INPUT_GRAD_ATOL = 3e-6 + + +class _DecoderShell(nn.Module): + """Minimal decoder interface used by ``TECudaGraphHelper`` discovery.""" + + def __init__(self, layer): + super().__init__() + self.layers = nn.ModuleList([layer]) + + +class _ModelChunkShell(nn.Module): + """Minimal training model chunk required by ``TECudaGraphHelper``.""" + + def __init__(self, layer, config): + super().__init__() + self.config = config + self.decoder = _DecoderShell(layer) + + def zero_grad_buffer(self): + """Match the training loop's interface for clearing gradient buffers.""" + self.zero_grad(set_to_none=True) + + +def _unwrap_output(output): + """Normalize graph replay's optional one-tensor tuple to a tensor.""" + if isinstance(output, tuple): + assert len(output) == 1 + return output[0] + return output + + +def _assert_module_grads_finite(module): + for name, parameter in module.named_parameters(): + assert ( + parameter.grad is None or torch.isfinite(parameter.grad).all() + ), f"non-finite gradient for {name}" + + +def _setup(mode, cuda_graph_impl, warmup_steps): + """Initialize CP, build a graphed MambaLayer, and return it with its local input.""" + Utils.initialize_model_parallel(context_parallel_size=CP_SIZE) + torch.manual_seed(1234) + model_parallel_cuda_manual_seed(1234, te_rng_tracker=True, force_reset_rng=True) + cp_group = parallel_state.get_context_parallel_group() + tp_group = parallel_state.get_tensor_model_parallel_group() + device = torch.device(torch.cuda.current_device()) + + shape = MambaModelShape() + sequence_length = shape.chunk_size * 2 * CP_SIZE + batch_size = 1 + layer, config = build_mamba_layer( + shape, + cp_group, + tp_group, + load_balancing=mode, + cuda_graph_impl=cuda_graph_impl, + cuda_graph_warmup_steps=warmup_steps, + ) + generator = torch.Generator(device=device).manual_seed(5678) + global_input = torch.randn( + sequence_length, + batch_size, + shape.hidden_size, + device=device, + dtype=torch.bfloat16, + generator=generator, + ) + local_input = select_balanced_cp_shard(global_input, cp_group.rank(), CP_SIZE) + + # Cudagraph backward capture accumulates wgrads into ``main_grad``, which + # DDP normally provides; this test drives a bare module, so create them here + # (the same setup tests/unit_tests/transformer/test_cuda_graphs.py uses). + for parameter in layer.parameters(): + parameter.main_grad = torch.zeros_like(parameter) + + return layer, config, cp_group, local_input, sequence_length, batch_size + + +def _replay_and_check(layer, local_input, eager_output, eager_input_grad): + """Replay the graph twice, checking eager parity and static-buffer refresh.""" + replay_outputs = [] + for replay_index in range(2): + layer.zero_grad(set_to_none=True) + replay_input = (local_input * (1.0 + 0.01 * replay_index)).detach().requires_grad_(True) + replay_output = _unwrap_output(layer(hidden_states=replay_input, attention_mask=None)) + replay_output.float().square().mean().backward() + torch.cuda.synchronize() + + assert torch.isfinite(replay_output).all() + assert torch.isfinite(replay_input.grad).all() + _assert_module_grads_finite(layer) + if replay_index == 0: + torch.testing.assert_close(replay_output, eager_output, rtol=0, atol=0) + torch.testing.assert_close( + replay_input.grad, eager_input_grad, rtol=INPUT_GRAD_RTOL, atol=INPUT_GRAD_ATOL + ) + replay_outputs.append(replay_output.detach().clone()) + + assert not torch.equal( + replay_outputs[0], replay_outputs[1] + ), "CUDA Graph replay did not consume the changed input" + + +@pytest.mark.parametrize("mode", STATE_PASSING_CP_MODES) +def test_local_cuda_graph(mode): + """Megatron's local CUDA Graph lifecycle must capture and replay the CP path.""" + if Utils.world_size % CP_SIZE != 0: + pytest.skip(f"world size {Utils.world_size} is not a multiple of {CP_SIZE}") + layer, _, cp_group, local_input, _, _ = _setup(mode, "local", warmup_steps=2) + runner = None + try: + # A normal Megatron call records the runners and establishes eager references. + recorded_input = local_input.detach().clone().requires_grad_(True) + recorded_output = layer(hidden_states=recorded_input, attention_mask=None) + recorded_output.float().square().mean().backward() + eager_output = recorded_output.detach().clone() + eager_input_grad = recorded_input.grad.detach().clone() + torch.cuda.synchronize() + dist.barrier(group=cp_group) + + assert len(layer.cudagraph_manager.cudagraph_runners) == 1 + runner = layer.cudagraph_manager.cudagraph_runners[0] + assert runner.fwd_graph_recorded and runner.bwd_graph_recorded + + create_cudagraphs() + torch.cuda.synchronize() + dist.barrier(group=cp_group) + assert runner.cudagraph_created + assert runner.fwd_graph is not None and runner.bwd_graph is not None + + _replay_and_check(layer, local_input, eager_output, eager_input_grad) + dist.barrier(group=cp_group) + finally: + # Release the graph execs and the global cudagraph record before + # destroying the NCCL process groups. Leaving the global record in place + # makes the next test see `cudagraph_created` with no matching runners. + torch.cuda.synchronize() + delete_cuda_graphs() + del layer, runner + gc.collect() + torch.cuda.empty_cache() + Utils.destroy_model_parallel() + + +@pytest.mark.parametrize("mode", STATE_PASSING_CP_MODES) +def test_te_cuda_graph(mode): + """The Transformer Engine graph helper must capture and replay the CP path.""" + if Utils.world_size % CP_SIZE != 0: + pytest.skip(f"world size {Utils.world_size} is not a multiple of {CP_SIZE}") + warmup_steps = 2 + layer, config, cp_group, local_input, sequence_length, batch_size = _setup( + mode, "transformer_engine", warmup_steps + ) + helper = None + try: + init_num_microbatches_calculator( + rank=torch.distributed.get_rank(), + rampup_batch_size=None, + global_batch_size=batch_size, + micro_batch_size=batch_size, + data_parallel_size=1, + decrease_batch_size_if_needed=False, + ) + model_chunk = _ModelChunkShell(layer, config).to(torch.cuda.current_device()) + + eager_output = None + eager_input_grad = None + for _ in range(warmup_steps): + layer.zero_grad(set_to_none=True) + eager_input = local_input.detach().clone().requires_grad_(True) + eager_result = layer(hidden_states=eager_input, attention_mask=None) + eager_result.float().square().mean().backward() + torch.cuda.synchronize() + eager_output = eager_result.detach().clone() + eager_input_grad = eager_input.grad.detach().clone() + + # Training releases the previous iteration's autograd graph before TE + # capture. Keeping these tensors alive makes an old AccumulateGrad node + # retain its original stream, which can invalidate capture. + del eager_result, eager_input + layer.zero_grad(set_to_none=True) + gc.collect() + torch.cuda.synchronize() + dist.barrier(group=cp_group) + + helper = TECudaGraphHelper( + model=[model_chunk], + config=config, + seq_length=sequence_length, + micro_batch_size=batch_size, + optimizers=[], + ) + assert helper.flattened_callables == [layer] + helper.create_cudagraphs() + torch.cuda.synchronize() + dist.barrier(group=cp_group) + assert helper.capture_finished() and helper.graphs_created() + assert len(layer.cuda_graphs) == 1 + + _replay_and_check(layer, local_input, eager_output, eager_input_grad) + dist.barrier(group=cp_group) + finally: + if helper is not None: + helper.delete_cuda_graphs() + delete_cuda_graphs() + del layer + torch.cuda.synchronize() + gc.collect() + torch.cuda.empty_cache() + unset_num_microbatches_calculator() + Utils.destroy_model_parallel()