diff --git a/docs/operators/README.md b/docs/operators/README.md index 6bf8e06b2..d2beba6c8 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -52,3 +52,4 @@ Every operator page should include: - [MiniMax-H3 Full Transformer and Dual Trajectory](h3-full-transformer.md) - [MiniMax-H3 Integration Rows](h3-integration.md) - [MiniMax-H3 Trajectory Artifact](h3-trajectory-artifact.md) +- [MiniMax-H3 Tensor/Sequence/Context Parallel Rows](h3-parallel.md) diff --git a/docs/operators/h3-checkpoint-mapper-vllm.md b/docs/operators/h3-checkpoint-mapper-vllm.md new file mode 100644 index 000000000..050653650 --- /dev/null +++ b/docs/operators/h3-checkpoint-mapper-vllm.md @@ -0,0 +1,32 @@ +# MiniMax-H3 Checkpoint Mapper and vLLM-Omni Rollout Adapter + +## Summary + +RFC #420 integration rows, validated on H100 against the installed vLLM-Omni name tables: + +| Row | Where | What it fixes | +| --- | --- | --- | +| `h3_checkpoint_mapper` | `models/minimax_h3/checkpoint_mapper.py` | Diffusers <-> vLLM-Omni native names and layouts | +| `vllm_runtime_adapter_h3` | `integrations/engines/rollout/vllm/h3_omni.py` | strict fixed-order kernels behind vLLM-Omni's `qkv_proj` / `out_proj` / `fc1` / `fc2` | + +## Checkpoint mapper + +`transformer_blocks.N` -> `blocks.N`, `token_refiner.refiner_blocks.N` -> `token_refiner.blocks.N`; +`to_q/k/v` -> one `attn.qkv_proj` (`[Q;K;V]` rows); `to_out.0` -> `out_proj`; `ff.net.0.proj` +(value-first) -> `mlp.fc1` (gate-first: halves swapped); `ff.net.2` -> `mlp.fc2`; model-level names +from vLLM-Omni's `fasth3` table. All 638 checkpoint tensors map; the round trip is bitwise; +unmapped names and incomplete q/k/v sets fail closed. A test compares the tables with +`vllm_omni.diffusion.models.minimax_h3.fasth3` when importable (passes with vLLM-Omni installed). + +## vLLM-Omni adapter + +`install_h3_rollout(model, collective=None, tp=1)` replaces the four fused linears' forwards +(returning vLLM's `(output, None)`): `qkv_proj` / `fc1` use `h3-projection-gemm-v1` (a fused column +equals the same column alone, so fused == separate bitwise, tested); `out_proj` / `fc2` use the +canonical row-parallel reduction of [`h3-parallel-v1`](h3-parallel.md). Wrong dtype, non-CUDA +weights or a model without those modules raise in strict mode. Attention, rotary, QK norm, +SwiGLU and AdaLN stay vLLM-Omni's and are reported `external`. `handle.remove()` restores. + +## Evidence + +`tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py`. The name-table test compares with `vllm_omni.diffusion.models.minimax_h3.fasth3` when it is importable and is skipped otherwise. diff --git a/docs/operators/h3-parallel.md b/docs/operators/h3-parallel.md new file mode 100644 index 000000000..4ebdef6ed --- /dev/null +++ b/docs/operators/h3-parallel.md @@ -0,0 +1,101 @@ +# MiniMax-H3 Tensor / Sequence / Context Parallel Rows (WS2) + +## Summary + +RFC #420 rows `tp_attention_7168`, `tp_ffn_14336`, `cp_full_attention` (Ulysses), +`cp_metadata_coshard`, `cp_arbitrary_length`, `distributed_packed_io`, and closeouts +`ws2_tp2_tp4_tp8`, `ws2_tp4_cp2`, `ws2_tp4_cp2_sp`. Contract `h3-parallel-v1` +(`rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py`); the block that composes them is +`rl_engine/validation/models/h3_tp.py`. Cross-node reductions use +`rl_engine/distributed/algorithms/hierarchical.py` (`hierarchical_ipc_nccl_tree`). + +Result (H100, 2 nodes x 4 GPUs, block 0 of the pinned checkpoint, 264 packed rows): every rank's +rows of the forward, the input gradient and every weight-gradient shard are **bitwise equal** to the +single-GPU canonical block for TP2, TP4, TP8 (two nodes), TP4 x CP2 (two nodes) and TP4 x CP2 + SP. +Evidence: `reports/experiments/h3-ws2-parallel-h100/*.json`. + +## The canonical reduction + +All cross-rank sums are defined on **8 canonical units** (the divisors of 8 are the TP sizes): + +| Layer | Sharding | Reduction | Rule | +| --- | --- | --- | --- | +| Q/K/V, gate/up forward | output rows (whole heads / halves) | none | the unsharded GEMM's columns (`h3-projection-gemm-v1`) | +| O (7168 -> 5376), down (14336 -> 5376) forward | reduction axis, `8/tp` K-units per rank | K | each K-unit is one fixed-order FP32 GEMM; a rank adds its units as a balanced tree; the collective adds ranks as the same tree; one BF16 rounding | +| O / down backward (`dX`, `dW`) | - | N / tokens | plain fixed-order GEMM, independent of the shard | +| Q/K/V, gate/up backward `dX` | output rows | N | unit `j` = block `j` of every fused weight, summed in the order q,k,v (hidden,gate); tree over units; collective; one rounding | + +The unit of the O projection is 896 K (14 K-tiles) and of down 1792 (28 K-tiles), so a unit never +splits a tile. FP32 partials cross ranks (`DeterministicCollective` fixed tree on a node; +`HierarchicalCollective` across nodes). The row-parallel contract differs from the single-sweep +K reduction of the unsharded O/down GEMM rows (claimed by another contributor): it is the +TP-invariant definition, and TP=1 is that definition, not the sweep. + +## Cross-node collective + +`DeterministicCollective` is CUDA IPC and therefore single node. `HierarchicalCollective(local_size)` +reduces each group of `local_size` ranks with that tree, moves the group partials with NCCL +`all_gather` (exact copies, no NCCL arithmetic) and adds them in one fixed balanced order. For +power-of-two sizes this is the 8-way tree; a unit test checks it equals the single-node tree of the +same world on 4 GPUs for FP32 and BF16 (`local_size` 1 and 2), and the TP8 run above uses it across +two nodes. + +## Context parallelism + +- `SeqShard`: contiguous row slices differing by at most one row, any length >= 0 (including fewer + rows than ranks) for the data movement (`scatter_packed` / `gather_packed`, Ulysses all-to-all): + `cp_arbitrary_length`. The compute paths (`seq_sharded_linear`, the parallel block) need at + least one row per rank; a shorter sequence is refused on every rank alike, before any collective. +- `coshard`: hidden, AdaLN row index, token tags, rotary tables are cut with one `SeqShard` so the + metadata follows its rows: `cp_metadata_coshard`. +- `ulysses_attention`: all-to-all sequence <-> head (ragged split sizes), attention on `H/cp` heads + of the full sequence, all-to-all back; both moves are exact, so rows equal the unsharded + attention's if the attention is head independent (the repo's deterministic attention is). Tested at + lengths 1, 2, 7, 130 and in the block at 261 / 263 rows. +- `scatter_packed` / `gather_packed`: exact packed IO, ragged gather by padding: `distributed_packed_io`. +- `seq_sharded_linear`: weight gradient from gathered rows so it equals the unsharded token sweep. + +## Sequence parallelism + +With `sp=True` the residual stream is sharded by rows across the TP group: norm and gate-residual are +row-local, the column-parallel input is all-gathered (exact) and the row-parallel output is +reduce-scattered with the same tree (`rows % tp == 0` required, checked). + +## Tests and evidence + +`tests/models/minimax_h3/test_h3_parallel.py`: 70 tests on 4 H100 (sharding, canonical folding for +TP 2/4/8, accuracy against FP64, row independence, Ulysses ragged, ten block grids with backward, +hierarchical collective). `tools/validation/models/h3_parallel_run.py` runs a grid under `srun` on +any node count and compares against the single-GPU block. + +| Grid | GPUs / nodes | Collective | Forward | dX | dW shards | +| --- | --- | --- | --- | --- | --- | +| TP2 | 2 / 1 | IPC tree | bitwise | bitwise | bitwise | +| TP4 | 4 / 1 | IPC tree | bitwise | bitwise | bitwise | +| TP8 | 8 / 2 | hierarchical | bitwise | bitwise | bitwise | +| TP4 x CP2 | 8 / 2 | IPC tree (TP in node), NCCL all-to-all (CP across) | bitwise | bitwise | not compared (token axis sharded) | +| TP4 x CP2 + SP | 8 / 2 | same | bitwise | bitwise | not compared | + +Mean time per rank of one forward call (ms), from the `h3-ws2-parallel-h100/*.json` reports: TP2 +33.1, TP4 32.6, TP8 32.4, TP4xCP2 31.5, TP4xCP2+SP 33.3 (tiny 264-row case, launch-bound; not a +scaling claim). Each call also rebuilds the block context, the operators and the weight shards, so +this is the time of a whole `block_case(..., backward=False)` call, not of the kernels alone. + +## Limits + +- Attention, SwiGLU and QK norm inside the block are the repo's interim operators; the parallel + rows only require them to be row/head independent, which the bitwise results confirm. +- Pipeline parallelism is out of scope (`ws2_pipeline_parallel` is not part of this change). +- dW under CP is exact through `seq_sharded_linear` for non-TP linears; the TP x CP block compares + forward and `dX` (CP shards the token axis of the TP weight gradients). +- `HierarchicalCollective` assumes ranks are packed node-major and node counts of 1/2/4/8. + With several TP groups (CP > 1) it is given every TP group's ranks (`group_ranks`, as + `make_grid` does) and every rank creates all their subgroups in one order, as + `torch.distributed.new_group` requires; a TP subgroup without that list is refused. The + hierarchical path with CP > 1 has a CPU test of the creation order only, no GPU run. + +## Hopper evidence chart + +These outputs and figures are not part of this diff; they are on the fork's `pr_assets` branch (links in the PR description). +Per-layout results of the 2 nodes x 4 H100 runs (`reports/experiments/h3-ws2-parallel-h100/*.json`): every layout is bitwise equal to the single-GPU reference +(forward and input gradient); mean per-rank forward time of block 0 is in the figure. diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py b/rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py new file mode 100644 index 000000000..e8bfbbb46 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py @@ -0,0 +1,372 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""MiniMax-H3 tensor-, sequence- and context-parallel building blocks (contract ``h3-parallel-v1``). + +Rows: ``tp_attention_7168``, ``tp_ffn_14336``, ``cp_full_attention`` (Ulysses), +``cp_metadata_coshard``, ``cp_arbitrary_length``, ``distributed_packed_io``. + +Every reduction that crosses ranks is defined on **8 canonical units** so TP 1/2/4/8 add the +same numbers in the same order: + +* *row-parallel* (O and FFN-down, ``K`` = 7168 / 14336): the reduction axis is cut into 8 + equal K-units; a unit is one fixed-order FP32 GEMM (``h3-projection-gemm-v1`` tile, ``K`` + swept ascending). A rank holding ``8/tp`` consecutive units adds them in balanced-tree + order and the collective adds the ranks' partials in the same tree: the result is the + balanced tree over the 8 units for every TP size, rounded to BF16 once. +* *column-parallel backward* (``dX`` of Q/K/V and gate/up, reduction over the output axis): + the same 8 units, where unit ``j`` is output block ``j`` of every weight of the fused + projection, added in the order given (q, k, v / hidden, gate). + +The forward of a column-parallel projection and the backward of a row-parallel one reduce +over axes that do not depend on the sharding (they are plain ``h3-projection-gemm-v1``), so +they are bitwise the unsharded result without any collective. + +Context parallelism (Ulysses) moves data only: sequence <-> head all-to-alls are exact, so +attention (per head, full sequence) is the unsharded attention. Weight gradients under +sequence sharding are computed from the gathered rows, so they equal the unsharded GEMM. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass + +import torch +import torch.distributed as dist + +from rl_engine.backends.shared.triton.gemm import h3_projection as g +from rl_engine.distributed.algorithms.hierarchical import tree_add + +CONTRACT = "h3-parallel-v1" +UNITS = 8 # canonical reduction units; TP sizes are the divisors 1, 2, 4, 8 +TP_SIZES = (1, 2, 4, 8) +HEADS, HEAD_DIM = 56, 128 +INNER = HEADS * HEAD_DIM # 7168 +FFN = 14336 + + +# --- Sharding helpers --------------------------------------------------------------------------- + + +def check_tp(tp: int, rank: int) -> None: + if tp not in TP_SIZES or not 0 <= rank < tp: + raise ValueError( + f"tp must be one of {TP_SIZES} and rank in [0, tp); got tp={tp} rank={rank}" + ) + + +def shard_rows(weight: torch.Tensor, tp: int, rank: int) -> torch.Tensor: + """Output-axis shard of a column-parallel weight ``(N, K)``: rows ``[r*N/tp, (r+1)*N/tp)``.""" + + check_tp(tp, rank) + n = weight.shape[0] + if n % UNITS: + raise ValueError(f"{n} output rows do not split into {UNITS} units") + return weight.narrow(0, rank * n // tp, n // tp).contiguous() + + +def shard_cols(weight: torch.Tensor, tp: int, rank: int) -> torch.Tensor: + """Reduction-axis shard of a row-parallel weight ``(N, K)``: columns of ``8/tp`` K-units.""" + + check_tp(tp, rank) + k = weight.shape[1] + unit = k // UNITS + if unit * UNITS != k or unit % g.CONFIG_BF16["BLOCK_K"]: + raise ValueError(f"K={k} must split into {UNITS} units of a multiple of the GEMM K tile") + return weight.narrow(1, rank * k // tp, k // tp).contiguous() + + +def shard_heads(x: torch.Tensor, tp: int, rank: int, dim: int = -1) -> torch.Tensor: + """Head shard of a ``(..., 7168)`` activation or weight axis (whole 128-wide heads).""" + + check_tp(tp, rank) + return x.narrow(dim, rank * x.shape[dim] // tp, x.shape[dim] // tp).contiguous() + + +def shard_gate_up(weight: torch.Tensor, tp: int, rank: int) -> torch.Tensor: + """Shard of the fused ``(2*FFN, K)`` gate/up weight keeping ``[hidden; gate]`` halves. + + Each half is sharded on its own, so local rows are ``[hidden shard; gate shard]`` and + ``chunk(2)`` still splits them. + """ + + half = weight.shape[0] // 2 + return torch.cat( + [shard_rows(weight[:half], tp, rank), shard_rows(weight[half:], tp, rank)] + ).contiguous() + + +# --- Row-parallel linear (reduction over the sharded axis) --------------------------------------- + + +def _units_partial(x: torch.Tensor, w: torch.Tensor, tp: int) -> torch.Tensor: + """FP32 balanced-tree sum of this rank's ``8/tp`` K-units of ``x @ w.T``.""" + + local = UNITS // tp + unit = x.shape[1] // local + parts = [ + g.gemm( + x[:, i * unit : (i + 1) * unit], + w[:, i * unit : (i + 1) * unit], + b_t=True, + out_dtype=torch.float32, + ) + for i in range(local) + ] + return tree_add(parts) + + +class _RowParallel(torch.autograd.Function): + @staticmethod + def forward(ctx, x, w, collective, tp, scatter): + ctx.save_for_backward(x, w) + ctx.scatter, ctx.collective = scatter, collective + partial = _units_partial(x, w, tp) + if tp > 1: + partial = ( + collective.reduce_scatter(partial) if scatter else collective.all_reduce(partial) + ) + return partial.to(x.dtype) + + @staticmethod + def backward(ctx, grad): + x, w = ctx.saved_tensors + if ctx.scatter: # sequence-parallel: the rows of dY are sharded; gather them (exact) + grad = ctx.collective.all_gather(grad.contiguous()) + grad = grad.contiguous() + dx = g.gemm(grad, w) # reduction over N: independent of the shard + dw = g.gemm(grad, x, a_t=True) # reduction over tokens: independent of the shard + return dx, dw, None, None, None + + +def row_parallel_linear(x, w_shard, collective=None, tp: int = 1, scatter: bool = False): + """``y = x @ W.T`` with ``x (T, K/tp)`` and ``w_shard (N, K/tp)``; ``y`` bitwise TP-invariant. + + ``scatter`` gives this rank's ``T/tp`` rows of ``y`` (sequence parallel) instead of all of it. + """ + + if x.dim() != 2 or w_shard.dim() != 2 or x.shape[1] != w_shard.shape[1]: + raise ValueError("row_parallel_linear needs x (T, K/tp) and w (N, K/tp)") + if tp > 1 and collective is None: + raise ValueError("tp > 1 needs a collective") + return _RowParallel.apply(x.contiguous(), w_shard, collective, tp, scatter) + + +# --- Column-parallel linear (backward reduces over the sharded axis) ------------------------------ + + +class _ColumnParallel(torch.autograd.Function): + @staticmethod + def forward(ctx, x, collective, tp, gather, *weights): + ctx.save_for_backward(x, *weights) + ctx.collective, ctx.tp, ctx.gather = collective, tp, gather + full_x = collective.all_gather(x) if gather else x + ctx.full_x = full_x if gather else None + return tuple(g.gemm(full_x, w, b_t=True) for w in weights) + + @staticmethod + def backward(ctx, *grads): + x, *weights = ctx.saved_tensors + full_x = ctx.full_x if ctx.gather else x + local = UNITS // ctx.tp + grads = [t.contiguous() for t in grads] + units = [] + for j in range(local): + acc = None + for grad, w in zip(grads, weights): + width = grad.shape[1] // local + piece = g.gemm( + grad[:, j * width : (j + 1) * width], + w[j * width : (j + 1) * width], + out_dtype=torch.float32, + ) + acc = piece if acc is None else acc + piece + units.append(acc) + partial = tree_add(units) + if ctx.tp > 1: + partial = ( + ctx.collective.reduce_scatter(partial) + if ctx.gather + else ctx.collective.all_reduce(partial) + ) + dx = partial.to(x.dtype) + dws = tuple(g.gemm(grad, full_x, a_t=True) for grad in grads) + return (dx, None, None, None, *dws) + + +def column_parallel_linear( + x, weights: Sequence[torch.Tensor], collective=None, tp: int = 1, gather: bool = False +): + """``[x @ w.T for w in weights]`` with each ``w`` a row shard ``(N/tp, K)``. + + The forward is the unsharded projection's columns (no collective). With ``gather``, ``x`` is + this rank's rows of a sequence-parallel activation, gathered first (exact), and ``dX`` is + reduce-scattered back to rows. ``dX`` is the canonical 8-unit tree for every ``tp``. + """ + + if tp > 1 and collective is None: + raise ValueError("tp > 1 needs a collective") + return _ColumnParallel.apply(x.contiguous(), collective, tp, gather, *weights) + + +# --- Context parallelism (Ulysses) --------------------------------------------------------------- + + +@dataclass(frozen=True) +class SeqShard: + """This rank's contiguous slice of a packed sequence of ``length`` rows over ``cp`` ranks. + + Slices differ by at most one row (first ``length % cp`` ranks get the extra row). For data + movement (scatter / gather / all-to-all) any length is legal, including ``length < cp`` (some + ranks hold nothing). The compute paths (``seq_sharded_linear``, the parallel block) need a row + on every rank and refuse a shorter sequence on every rank alike + (``require_rows_on_every_rank``), so no rank is left waiting in a collective. + """ + + length: int + cp: int + rank: int + + @property + def sizes(self) -> list[int]: + base, extra = divmod(self.length, self.cp) + return [base + (r < extra) for r in range(self.cp)] + + @property + def lo(self) -> int: + return sum(self.sizes[: self.rank]) + + @property + def hi(self) -> int: + return self.lo + self.sizes[self.rank] + + +def require_rows_on_every_rank(shard: SeqShard) -> None: + """Raise, identically on every rank, when some rank's slice of the sequence is empty.""" + + if min(shard.sizes) == 0: + raise ValueError( + f"{shard.length} rows over cp={shard.cp}: the sequence-sharded compute paths need at " + "least one row per rank (data movement alone accepts fewer)" + ) + + +def coshard(shard: SeqShard, *tensors: torch.Tensor, dim: int = 0) -> list[torch.Tensor]: + """Slice every per-row tensor (hidden, timestep_indices, token_tags, position_ids, ...) alike. + + The metadata that selects AdaLN rows and rotary angles must follow the rows it describes; + slicing them in one call with one ``SeqShard`` is what keeps them aligned. + """ + + out = [] + for t in tensors: + if t.shape[dim] != shard.length: + raise ValueError( + f"tensor has {t.shape[dim]} rows along dim {dim}, expected {shard.length}" + ) + out.append(t.narrow(dim, shard.lo, shard.hi - shard.lo).contiguous()) + return out + + +def scatter_packed(packed: torch.Tensor, shard: SeqShard, dim: int = 0) -> torch.Tensor: + """This rank's rows of a packed tensor that every rank holds in full.""" + + return coshard(shard, packed, dim=dim)[0] + + +def gather_packed(local: torch.Tensor, shard: SeqShard, group=None, dim: int = 0) -> torch.Tensor: + """Reassemble the packed tensor from every rank's rows (exact, any length).""" + + if shard.cp == 1: + return local + sizes = shard.sizes + widest = max(sizes) + moved = local.movedim(dim, 0).contiguous() + pad = moved.new_zeros((widest - moved.shape[0], *moved.shape[1:])) + chunks = [torch.empty_like(torch.cat([moved, pad])) for _ in range(shard.cp)] + dist.all_gather(chunks, torch.cat([moved, pad]), group=group) + return torch.cat([c[:n] for c, n in zip(chunks, sizes)]).movedim(0, dim) + + +class _AllToAll(torch.autograd.Function): + @staticmethod + def forward(ctx, x, in_sizes, out_sizes, group): + ctx.sizes, ctx.group = (in_sizes, out_sizes), group + out = x.new_empty((sum(out_sizes), *x.shape[1:])) + dist.all_to_all_single(out, x.contiguous(), out_sizes, in_sizes, group=group) + return out + + @staticmethod + def backward(ctx, grad): + in_sizes, out_sizes = ctx.sizes + back = grad.new_empty((sum(in_sizes), *grad.shape[1:])) + dist.all_to_all_single(back, grad.contiguous(), in_sizes, out_sizes, group=ctx.group) + return back, None, None, None + + +def ulysses_attention(q, k, v, attention, shard: SeqShard, group=None): + """Full attention of a sequence-sharded ``q, k, v (B, S_local, H, D)`` by head scatter. + + All-to-all makes each rank hold ``H/cp`` heads for the whole sequence, ``attention(q, k, v)`` + (``(B, H', S, D)`` layout, per-head independent) runs, and the result goes back. Both moves + are exact, so the output rows equal the unsharded attention's rows bitwise as long as + ``attention`` is batch/head invariant. + """ + + cp = shard.cp + if cp == 1: + return attention(q, k, v) + batch, local_rows, heads, dim = q.shape + if heads % cp: + raise ValueError(f"{heads} heads do not divide over cp={cp}") + sizes, hp = shard.sizes, heads // cp + if local_rows != sizes[shard.rank]: + raise ValueError(f"expected {sizes[shard.rank]} local rows, got {local_rows}") + + def to_heads(t): + # peer r receives head block r of my rows; I receive every peer's rows of my head block + send = t.reshape(batch, local_rows, cp, hp, dim).permute(2, 1, 0, 3, 4) + send = send.reshape(cp * local_rows, batch, hp, dim) + got = _AllToAll.apply(send, [local_rows] * cp, sizes, group) + return got.permute(1, 0, 2, 3) # (B, S, hp, D), sequence in rank order + + out = attention(to_heads(q), to_heads(k), to_heads(v)) # (B, S, hp, D) + send = out.permute(1, 0, 2, 3).contiguous() + got = _AllToAll.apply(send, sizes, [local_rows] * cp, group) + got = got.reshape(cp, local_rows, batch, hp, dim) + return got.permute(2, 1, 0, 3, 4).reshape(batch, local_rows, heads, dim) + + +# --- Weight gradients under sequence sharding ---------------------------------------------------- + + +class _SeqLinear(torch.autograd.Function): + """Row-local linear on a sequence shard whose ``dW`` is computed from gathered rows.""" + + @staticmethod + def forward(ctx, x, w, shard, group): + ctx.save_for_backward(x, w) + ctx.shard, ctx.group = shard, group + return g.gemm(x, w, b_t=True) + + @staticmethod + def backward(ctx, grad): + x, w = ctx.saved_tensors + grad = grad.contiguous() + dx = g.gemm(grad, w) + full_x = gather_packed(x, ctx.shard, ctx.group) + full_g = gather_packed(grad, ctx.shard, ctx.group) + return dx, g.gemm(full_g, full_x, a_t=True), None, None + + +def seq_sharded_linear(x, w, shard: SeqShard, group=None): + """``x @ W.T`` for this rank's rows; ``dW`` equals the unsharded GEMM (token sweep).""" + + if shard.cp > 1: + require_rows_on_every_rank(shard) + return ( + _SeqLinear.apply(x.contiguous(), w, shard, group) + if shard.cp > 1 + else g.linear_forward(x, w, None) + ) diff --git a/rl_engine/distributed/algorithms/hierarchical.py b/rl_engine/distributed/algorithms/hierarchical.py new file mode 100644 index 000000000..6e7e35620 --- /dev/null +++ b/rl_engine/distributed/algorithms/hierarchical.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Cross-node deterministic collective: node-local IPC fixed tree, then one exact gather. + +``DeterministicCollective`` (``cuda_ipc_fixed_tree``) reads peers' memory through CUDA IPC, +so it is single node. This wrapper (``hierarchical_ipc_nccl_tree``) reaches the same bits on +several nodes without relying on any NCCL reduction order: + +1. every group of ``local_size`` consecutive ranks reduces with the node-local fixed tree; +2. the group partials move with NCCL ``all_gather`` (bit-exact data movement, no arithmetic); +3. each rank adds the group partials in one fixed, balanced order. + +With ``world_size = nodes * local_size`` and ``local_size``, ``nodes`` powers of two, the +tree ``((0+1)+(2+3)) + ((4+5)+(6+7))`` of an 8-way node-local reduction is exactly "the 4-way +tree of each half, then the two halves added". The input dtype is the accumulation dtype, as in +``DeterministicCollective`` (FP32 inputs keep every partial in FP32, which is what the H3 +parallel operators use), so the result is bitwise the single-node tree of the same world size. +""" + +from __future__ import annotations + +import torch +import torch.distributed as dist + +from rl_engine.distributed.algorithms.collectives import DeterministicCollective + +_GROUP_COUNTS = (1, 2, 4, 8) + + +def tree_add(parts: list[torch.Tensor]) -> torch.Tensor: + """Balanced pairwise sum ``((p0+p1)+(p2+p3))+...`` of a power-of-two number of tensors.""" + + if len(parts) & (len(parts) - 1) or not parts: + raise ValueError(f"tree reduction needs a power-of-two part count, got {len(parts)}") + while len(parts) > 1: + parts = [parts[i] + parts[i + 1] for i in range(0, len(parts), 2)] + return parts[0] + + +class HierarchicalCollective: + """TP-invariant all-reduce / reduce-scatter over ``world_size`` ranks on several nodes.""" + + backend_id = "hierarchical_ipc_nccl_tree" + + def __init__( + self, + local_size: int, + group: dist.ProcessGroup | None = None, + device: torch.device | str | int | None = None, + *, + max_size_bytes: int = 64 << 20, + group_ranks: list[list[int]] | None = None, + ) -> None: + """``group`` defaults to the world. ``group_ranks`` lists the global ranks of every group + whose collective is built at the same time (e.g. all TP groups of a TP x CP grid, this + rank's included); it is required when ``group`` does not span the world. + """ + + if not dist.is_initialized(): + raise RuntimeError("torch.distributed must be initialized before collectives") + self.group = group if group is not None else dist.group.WORLD + self.world_size = dist.get_world_size(group=self.group) + self.rank = dist.get_rank(group=self.group) + if self.world_size % local_size or local_size not in _GROUP_COUNTS: + raise ValueError( + f"local_size {local_size} must be 1/2/4/8 and divide world {self.world_size}" + ) + self.local_size = local_size + self.num_groups = self.world_size // local_size + if self.num_groups not in _GROUP_COUNTS: + raise ValueError(f"{self.num_groups} node groups; need 1, 2, 4 or 8") + self.group_index, self.local_rank = divmod(self.rank, local_size) + # ``dist.new_group`` must be entered by every rank of the world, in the same order and + # with the same rank lists. A subgroup does not know its sibling groups, so every rank + # creates the subgroups of every listed group in one rank-independent order and keeps + # its own; a subgroup without the full list is refused rather than left to hang. + ranks = list(dist.get_process_group_ranks(self.group)) # global rank of each member + everyone = [list(g) for g in group_ranks] if group_ranks is not None else [ranks] + if sorted(r for g in everyone for r in g) != list(range(dist.get_world_size())): + raise ValueError( + "group_ranks must list every group built together, covering each world rank once" + ) + if ranks not in everyone or any(len(g) != self.world_size for g in everyone): + raise ValueError("this rank's group must be in group_ranks, all groups the same size") + local_groups: list[dist.ProcessGroup] = [] + cross_groups: list[dist.ProcessGroup] = [] + for members in everyone: + local = [ + dist.new_group(members[g * local_size : (g + 1) * local_size]) + for g in range(self.num_groups) + ] + cross = [ + dist.new_group([members[g * local_size + r] for g in range(self.num_groups)]) + for r in range(local_size) + ] + if members == ranks: + local_groups, cross_groups = local, cross + self._cross = cross_groups[self.local_rank] + self._local = DeterministicCollective( + local_groups[self.group_index], device, max_size_bytes=max_size_bytes + ) + + def all_reduce(self, input: torch.Tensor) -> torch.Tensor: + """The fixed-tree sum of every rank's ``input`` on every rank (new tensor).""" + + if input.dtype not in (torch.float32, torch.bfloat16, torch.float16): + raise TypeError(f"unsupported dtype {input.dtype}") + wide = ( + input.contiguous() + ) # the input dtype is the accumulation dtype, as in the node-local tree + partial = self._local.all_reduce(wide) if self.local_size > 1 else wide.clone() + if self.num_groups > 1: + gathered = [torch.empty_like(partial) for _ in range(self.num_groups)] + dist.all_gather(gathered, partial, group=self._cross) + partial = tree_add(gathered) + return partial + + def reduce_scatter(self, input: torch.Tensor) -> torch.Tensor: + """``all_reduce`` then this rank's equal share of dimension 0.""" + + if input.shape[0] % self.world_size: + raise ValueError("dimension 0 must divide by the world size") + rows = input.shape[0] // self.world_size + return self.all_reduce(input)[self.rank * rows : (self.rank + 1) * rows].contiguous() + + def all_gather(self, input: torch.Tensor) -> torch.Tensor: + """Rank-ordered concatenation along dimension 0 (exact data movement).""" + + out = [torch.empty_like(input) for _ in range(self.world_size)] + dist.all_gather(out, input.contiguous(), group=self.group) + return torch.cat(out) + + def close(self) -> None: + self._local.close() diff --git a/rl_engine/integrations/engines/rollout/vllm/h3_omni.py b/rl_engine/integrations/engines/rollout/vllm/h3_omni.py new file mode 100644 index 000000000..815fa16fa --- /dev/null +++ b/rl_engine/integrations/engines/rollout/vllm/h3_omni.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``vllm_runtime_adapter_h3``: strict H3 projections inside a vLLM-Omni DiT. + +vLLM-Omni's ``MiniMaxH3Attention`` / MLP keep fused linears: ``attn.qkv_proj`` (``[Q;K;V]`` rows), +``attn.out_proj`` (row parallel), ``mlp.fc1`` (gate-first), ``mlp.fc2`` (row parallel); each +returns ``(output, bias)``. ``install_h3_rollout(model)`` replaces those four forwards with the +fixed-order kernels so rollout reads the same bits as the strict training graph: + +* ``qkv_proj`` / ``fc1``: ``h3-projection-gemm-v1`` over the fused rows (a column of a fused + projection equals the same column computed alone, so fused == separate bitwise); +* ``out_proj`` / ``fc2``: the canonical 8-unit row-parallel reduction of ``h3-parallel-v1`` + (``tp`` ranks with a deterministic collective, ``tp=1`` alone), one BF16 rounding. + +H3's four projections have no bias, and the patched forward computes ``x @ W.T`` and returns +``(output, None)``. A module that carries a bias is therefore refused in strict mode (and left +on vLLM-Omni's own forward otherwise) instead of silently losing its bias. + +Attention, rotary and norm stay vLLM-Omni's (their rows are separate); the report lists them as +``external``. Everything is restored by ``handle.remove()``. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field + +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3 import parallel as par +from rl_engine.backends.shared.triton.gemm import h3_projection as g + +LINEAR_NAMES = {"qkv_proj": "column", "fc1": "column", "out_proj": "row", "fc2": "row"} +# The four projections of a transformer block or a token-refiner block, by module path (an outer +# wrapper prefix is allowed). A module elsewhere that happens to share a leaf name is not touched. +_PROJECTION_PATH = re.compile( + r"(?:^|\.)(?:token_refiner\.)?blocks\.\d+\.(?:attn\.(?:qkv_proj|out_proj)|mlp\.(?:fc1|fc2))$" +) +_ABSENT = object() + + +@dataclass +class H3RolloutHandle: + patched: list[str] = field(default_factory=list) + _restore: list = field(default_factory=list) # (module, its own forward or _ABSENT) + + def remove(self) -> None: + """Put back exactly what each module had: its own instance ``forward`` or none.""" + + for module, previous in reversed(self._restore): + if previous is _ABSENT: + module.__dict__.pop("forward", None) + else: + module.forward = previous + self._restore.clear() + + def __enter__(self): + return self + + def __exit__(self, *exc): + self.remove() + + def report(self) -> dict: + return { + "contract": par.CONTRACT, + "patched": list(self.patched), + "external": ["attention", "rotary", "qk_norm", "swiglu", "adaln"], + } + + +def _wrap(module, kind, collective, tp): + def forward(x): + shape = x.shape + flat = x.reshape(-1, shape[-1]).contiguous() + if kind == "column": + out = g.gemm(flat, module.weight, b_t=True) + else: + out = par.row_parallel_linear(flat, module.weight, collective, tp) + return out.view(*shape[:-1], out.shape[-1]), None + + return forward + + +def install_h3_rollout( + model: torch.nn.Module, *, collective=None, tp: int = 1, strict: bool = True +) -> H3RolloutHandle: + """Patch every ``qkv_proj`` / ``out_proj`` / ``fc1`` / ``fc2`` of ``blocks`` and the refiner. + + Only modules at a block projection path are considered. Every candidate is checked before + any forward changes, so a strict refusal leaves the model untouched. + """ + + if tp > 1 and collective is None: + raise ValueError("tp > 1 needs a deterministic collective") + selected = [] + for name, module in model.named_modules(): + if not _PROJECTION_PATH.search(name) or not hasattr(module, "weight"): + continue + weight = module.weight + if not weight.is_cuda or weight.dtype != torch.bfloat16: + if strict: + raise RuntimeError( + f"{name}: strict H3 rollout needs CUDA BF16 weights, got " + f"{weight.dtype} on {weight.device}" + ) + continue + if getattr(module, "bias", None) is not None: + if strict: + raise RuntimeError( + f"{name}: has a bias; the strict H3 rollout kernels compute the bias-free H3 " + "projections only" + ) + continue + selected.append((name, module)) + if strict and not selected: + raise RuntimeError("no H3 projection modules found (qkv_proj/out_proj/fc1/fc2)") + handle = H3RolloutHandle() + for name, module in selected: + kind = LINEAR_NAMES[name.rsplit(".", 1)[-1]] + handle._restore.append((module, module.__dict__.get("forward", _ABSENT))) + module.forward = _wrap(module, kind, collective, tp) + handle.patched.append(name) + return handle diff --git a/rl_engine/models/minimax_h3/checkpoint_mapper.py b/rl_engine/models/minimax_h3/checkpoint_mapper.py new file mode 100644 index 000000000..e2c69a429 --- /dev/null +++ b/rl_engine/models/minimax_h3/checkpoint_mapper.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``h3_checkpoint_mapper``: Diffusers <-> vLLM-Omni native H3 names and layouts. + +The training side (diffusers, FSDP2) and the rollout side (vLLM-Omni) hold the same 638 +tensors under different names and fused layouts. This module is the single, tested mapping, so +the strict operators on both sides read bit-identical weights: + +* ``transformer_blocks.N`` -> ``blocks.N``; + ``token_refiner.refiner_blocks.N`` -> ``token_refiner.blocks.N``; +* ``attn.to_q/to_k/to_v`` -> one fused ``attn.qkv_proj`` laid out ``[Q; K; V]`` (rows), the + layout ``h3_qkv_gemm.forward_fused`` is bitwise-equal to; +* ``attn.to_out.0`` -> ``attn.out_proj``, ``attn.norm_q/k`` -> ``attn.q_norm/k_norm``; +* ``ff.net.0.proj`` (value-first ``[hidden; gate]``) -> ``mlp.fc1`` (gate-first, halves swapped), + ``ff.net.2`` -> ``mlp.fc2``; +* model-level names per vLLM-Omni's ``fasth3`` table (``proj_in`` -> ``video_patch_proj`` ...). + +The tables mirror ``vllm_omni.diffusion.models.minimax_h3.fasth3`` (checked by a test when that +package is importable). Mapping is pure data movement: +``native_to_diffusers(diffusers_to_native(x))`` is bitwise ``x``. + +Two QKV layouts exist on the native side, and both directions take ``qkv_layout``: + +* ``"module"`` (default): ``attn.qkv_proj.weight`` as the loaded module holds it, ``[Q; K; V]`` + rows. This is what ``install_h3_rollout`` reads and what ``load_state_dict`` writes in place; +* ``"checkpoint"``: a serialised native checkpoint, per-head ``[q, k, v]`` groups, which + vLLM-Omni's loader converts to ``[Q; K; V]`` itself. Writing the module layout into a file + for that loader would permute the rows twice. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +import torch + +PLAIN, SWAP_HALVES = "plain", "swap_halves" +QKV_SLOTS = ("q", "k", "v") +QKV_LAYOUTS = ("module", "checkpoint") +HEAD_DIM = 128 +TENSOR_KINDS = ("weight", "bias") + +MODEL_LEVEL = { + "proj_in": "video_patch_proj", + "proj_out": "final_layer.video_out", + "audio_proj_in": "audio_patch_proj", + "audio_proj_out": "final_layer.audio_out", + "context_embedder": "condition_proj", + "time_embedder.linear_1": "time_embedder.proj_in", + "time_embedder.linear_2": "time_embedder.proj_out", + "norm_out.linear": "final_layer.adaln_proj.linear", + "norm_out.norm": "final_layer.norm", + "token_refiner.final_norm": "token_refiner.final_norm", +} +BLOCK_PREFIXES = ( + ("token_refiner.refiner_blocks.", "token_refiner.blocks."), + ("transformer_blocks.", "blocks."), +) +BLOCK_TARGETS = { + "attn.to_q": ("attn.qkv_proj", "q"), + "attn.to_k": ("attn.qkv_proj", "k"), + "attn.to_v": ("attn.qkv_proj", "v"), + "attn.to_out.0": ("attn.out_proj", PLAIN), + "attn.norm_q": ("attn.q_norm", PLAIN), + "attn.norm_k": ("attn.k_norm", PLAIN), + "ff.net.0.proj": ("mlp.fc1", SWAP_HALVES), + "ff.net.2": ("mlp.fc2", PLAIN), + "adaln_proj.linear": ("adaln_proj.linear", PLAIN), + "norm1": ("norm1", PLAIN), + "norm2": ("norm2", PLAIN), +} + + +def resolve(module: str) -> tuple[str, str] | None: + """Diffusers module path -> ``(native module path, layout)`` or ``None``.""" + + if module in MODEL_LEVEL: + return MODEL_LEVEL[module], PLAIN + for source, target in BLOCK_PREFIXES: + if module.startswith(source): + index, _, suffix = module[len(source) :].partition(".") + mapped = BLOCK_TARGETS.get(suffix) + if index.isdigit() and mapped is not None: + return f"{target}{index}.{mapped[0]}", mapped[1] + return None + + +def swap_halves(tensor: torch.Tensor) -> torch.Tensor: + first, second = tensor.chunk(2, dim=0) + return torch.cat((second, first), dim=0) + + +def fuse_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + """``[Q; K; V]`` rows (the serving layout).""" + + return torch.cat([q, k, v], dim=0) + + +def split_qkv(fused: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if fused.shape[0] % 3: + raise ValueError(f"fused QKV rows {fused.shape[0]} are not a multiple of 3") + q, k, v = fused.chunk(3, dim=0) + return q, k, v + + +def grouped_qkv_to_qkv(weight: torch.Tensor, *, heads: int, head_dim: int) -> torch.Tensor: + """Native checkpoint's per-head ``[q,k,v]`` groups -> ``[Q; K; V]`` (vLLM-Omni's loader).""" + + per_group = 3 * head_dim + if weight.shape[0] != heads * per_group: + raise ValueError(f"grouped QKV rows {weight.shape[0]} != {heads * per_group}") + grouped = weight.reshape(heads, per_group, *weight.shape[1:]) + q, k, v = grouped.split(head_dim, dim=1) + return torch.cat([t.reshape(heads * head_dim, *weight.shape[1:]) for t in (q, k, v)]) + + +def qkv_to_grouped(fused: torch.Tensor, *, heads: int, head_dim: int) -> torch.Tensor: + q, k, v = split_qkv(fused) + parts = [t.reshape(heads, head_dim, *fused.shape[1:]) for t in (q, k, v)] + return torch.cat(parts, dim=1).reshape(heads * 3 * head_dim, *fused.shape[1:]) + + +def _check_layout(qkv_layout: str) -> None: + if qkv_layout not in QKV_LAYOUTS: + raise ValueError(f"qkv_layout must be one of {QKV_LAYOUTS}, got {qkv_layout!r}") + + +def _heads(fused: torch.Tensor) -> int: + if fused.shape[0] % (3 * HEAD_DIM): + raise ValueError(f"fused QKV rows {fused.shape[0]} are not whole {HEAD_DIM}-wide heads") + return fused.shape[0] // (3 * HEAD_DIM) + + +def diffusers_to_native( + state: Mapping[str, torch.Tensor], *, qkv_layout: str = "module" +) -> dict[str, torch.Tensor]: + """Diffusers state dict -> native names, fused QKV, gate-first ``fc1``; unmapped keys fail. + + ``qkv_layout`` selects the module layout (``[Q; K; V]``) or the serialised checkpoint + layout (per-head groups) of the fused QKV, see the module docstring. + """ + + _check_layout(qkv_layout) + out: dict[str, torch.Tensor] = {} + qkv: dict[str, dict[str, torch.Tensor]] = {} + for name, tensor in state.items(): + module, _, kind = name.rpartition(".") + target = resolve(module) + if target is None or kind not in TENSOR_KINDS: + raise ValueError(f"unmapped Diffusers H3 tensor: {name}") + native, layout = target + key = f"{native}.{kind}" + if layout in QKV_SLOTS: + qkv.setdefault(key, {})[layout] = tensor + elif layout == SWAP_HALVES: + out[key] = swap_halves(tensor) + else: + out[key] = tensor + for key, parts in qkv.items(): + if set(parts) != set(QKV_SLOTS): + raise ValueError(f"{key} needs q, k and v, got {sorted(parts)}") + fused = fuse_qkv(parts["q"], parts["k"], parts["v"]) + if qkv_layout == "checkpoint": + fused = qkv_to_grouped(fused, heads=_heads(fused), head_dim=HEAD_DIM) + out[key] = fused + return out + + +def _native_block(module: str) -> tuple[str, str, str] | None: + for source, target in BLOCK_PREFIXES: + if module.startswith(target): + index, _, suffix = module[len(target) :].partition(".") + if index.isdigit(): + return source, index, suffix + return None + + +def native_to_diffusers( + state: Mapping[str, torch.Tensor], *, qkv_layout: str = "module" +) -> dict[str, torch.Tensor]: + """Inverse of :func:`diffusers_to_native` (splits fused QKV, swaps ``fc1`` halves back). + + ``qkv_layout`` must name the layout the native QKV tensors are in (``"checkpoint"`` for a + serialised native checkpoint). Tensor kinds other than ``weight`` / ``bias`` are refused. + """ + + _check_layout(qkv_layout) + model = {native: diff for diff, native in MODEL_LEVEL.items()} + block = { + native: (diff, layout) + for diff, (native, layout) in BLOCK_TARGETS.items() + if layout not in QKV_SLOTS + } + out: dict[str, torch.Tensor] = {} + for name, tensor in state.items(): + module, _, kind = name.rpartition(".") + if kind not in TENSOR_KINDS: + raise ValueError(f"unmapped native H3 tensor: {name}") + if module in model: + out[f"{model[module]}.{kind}"] = tensor + continue + parsed = _native_block(module) + if parsed is None: + raise ValueError(f"unmapped native H3 tensor: {name}") + source, index, suffix = parsed + if suffix == "attn.qkv_proj": + if qkv_layout == "checkpoint": + tensor = grouped_qkv_to_qkv(tensor, heads=_heads(tensor), head_dim=HEAD_DIM) + for slot, part in zip(QKV_SLOTS, split_qkv(tensor)): + out[f"{source}{index}.attn.to_{slot}.{kind}"] = part + elif suffix in block: + diff, layout = block[suffix] + out[f"{source}{index}.{diff}.{kind}"] = ( + swap_halves(tensor) if layout == SWAP_HALVES else tensor + ) + else: + raise ValueError(f"unmapped native H3 tensor: {name}") + return out diff --git a/rl_engine/validation/models/h3_block.py b/rl_engine/validation/models/h3_block.py index fef290898..eb3c0c602 100644 --- a/rl_engine/validation/models/h3_block.py +++ b/rl_engine/validation/models/h3_block.py @@ -83,7 +83,7 @@ class CandidateOps: deterministic on every GPU this replay runs on. """ - def __init__(self, registry: KernelRegistry | None = None) -> None: + def __init__(self, registry: KernelRegistry | None = None, attention: str = "auto") -> None: from rl_engine.backends.cuda.activation.swiglu import SwiGLUCudaOp from rl_engine.backends.cuda.attention.deterministic_attn import DeterministicAttentionOp from rl_engine.backends.shared.triton.gemm.det_gemm import TritonDetGemmOp @@ -95,11 +95,27 @@ def __init__(self, registry: KernelRegistry | None = None) -> None: self.gemm = TritonDetGemmOp() self.qkv = registry.get_op("h3_qkv_gemm", device="cuda") self.attention = DeterministicAttentionOp() + self.attention_h3 = None + if attention != "interim": + # h3_full_attention (#553): query-chunked, bounded workspace, same #240 arithmetic + try: + from rl_engine.backends.cuda.attention.h3_full_attention import ( + H3FullAttentionOp, + _extension, + ) + + _extension() + self.attention_h3 = H3FullAttentionOp() + except (ImportError, RuntimeError): + if attention == "h3_full": + raise self.swiglu = SwiGLUCudaOp() def backends(self) -> dict[str, str]: return { - name: f"{type(op).__module__}.{type(op).__name__}" for name, op in vars(self).items() + name: f"{type(op).__module__}.{type(op).__name__}" + for name, op in vars(self).items() + if op is not None } @@ -146,6 +162,9 @@ def _cand_linear(ops: CandidateOps, x: torch.Tensor, weight: torch.Tensor) -> to def _cand_attention(ops: CandidateOps, q, k, v) -> torch.Tensor: + if ops.attention_h3 is not None: # BSHD, bounded workspace + q, k, v = (t.contiguous() for t in (q, k, _heads(v))) + return ops.attention_h3(q, k, v).flatten(2, 3) q, k, v = (t.transpose(1, 2).contiguous() for t in (q, k, _heads(v))) out = ops.attention(q, k, v, causal=False) return out.transpose(1, 2).flatten(2, 3) diff --git a/rl_engine/validation/models/h3_tp.py b/rl_engine/validation/models/h3_tp.py new file mode 100644 index 000000000..d298a09ec --- /dev/null +++ b/rl_engine/validation/models/h3_tp.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""One H3 block under tensor / sequence / context parallelism (RFC #420 WS2 closeouts). + +``run_tp_block`` is the block graph with the attention and FFN projections replaced by the +``h3-parallel-v1`` operators: column-parallel Q/K/V and gate/up, per-head norm, rotary and +attention on the rank's heads, Ulysses head<->sequence all-to-all across the CP group, +row-parallel O and down with the deterministic collective, optional sequence parallelism +(rows of the residual stream sharded across the TP group). With ``tp=cp=1`` it is the +canonical single-GPU block every other layout must equal bitwise. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist + +from rl_engine.backends.cuda.model_specific.minimax_h3 import parallel as par +from rl_engine.validation.models import h3_block as hb + +HIDDEN, HEAD_DIM = hb.HIDDEN, hb.HEAD_DIM + + +@dataclass +class Grid: + tp: int + cp: int + tp_rank: int + cp_rank: int + tp_collective: Any = None + cp_group: Any = None + + def seq_shard(self, length: int) -> par.SeqShard: + return par.SeqShard(length, self.cp, self.cp_rank) + + +def make_grid(tp: int, cp: int, *, local_size: int | None = None) -> Grid: + """Rank ``r`` is ``(cp_rank, tp_rank) = divmod(r, tp)``; TP groups are consecutive ranks. + + ``local_size`` (ranks per node) selects the cross-node collective when the TP group spans + nodes; ``None`` means the TP group is on one node (CUDA IPC fixed tree). + """ + + from rl_engine.distributed.algorithms.collectives import DeterministicCollective + from rl_engine.distributed.algorithms.hierarchical import HierarchicalCollective + + world, rank = dist.get_world_size(), dist.get_rank() + if world != tp * cp: + raise ValueError(f"world {world} != tp {tp} x cp {cp}") + cp_rank, tp_rank = divmod(rank, tp) + tp_ranks = [list(range(c * tp, (c + 1) * tp)) for c in range(cp)] + tp_groups = [dist.new_group(ranks) for ranks in tp_ranks] + cp_groups = [dist.new_group(list(range(t, world, tp))) for t in range(tp)] + collective: DeterministicCollective | HierarchicalCollective | None = None + if tp > 1: + mine = tp_groups[cp_rank] + if local_size is not None and local_size < tp: + collective = HierarchicalCollective( + local_size, + mine, + torch.cuda.current_device(), + max_size_bytes=256 << 20, + group_ranks=tp_ranks, # every rank creates all TP groups' subgroups (CP > 1) + ) + else: + collective = DeterministicCollective( + mine, torch.cuda.current_device(), max_size_bytes=256 << 20 + ) + return Grid(tp, cp, tp_rank, cp_rank, collective, cp_groups[tp_rank] if cp > 1 else None) + + +def _attention_fn(ops): + def attend(q, k, v): # (B, S, H', D) -> (B, S, H', D) + q, k, v = (t.transpose(1, 2).contiguous() for t in (q, k, v)) + return ops.attention(q, k, v, causal=False).transpose(1, 2).contiguous() + + return attend + + +def run_tp_block( + weights: dict[str, torch.Tensor], + hidden: torch.Tensor, + temb: torch.Tensor, + index: torch.Tensor, + rope: tuple[torch.Tensor, torch.Tensor], + ops: hb.CandidateOps, + grid: Grid, + *, + sp: bool = False, + seq_len: int | None = None, + leaves: dict[str, torch.Tensor] | None = None, +) -> torch.Tensor: + """Block 0 for this rank. ``hidden/index/rope`` are the full packed sequence on every rank. + + The rank slices its CP rows (``cp_metadata_coshard``: hidden, AdaLN row index and both + rotary tables with the same ``SeqShard``) and, with ``sp``, its TP row share of those. + Returns this rank's rows of the block output (CP rows, or the SP share of them). With + ``leaves`` the weight shards are tracked and stored there by name so the caller can read + their gradients. + """ + + if hidden.shape[0] != 1: + raise ValueError("the parallel block takes one packed sequence (batch 1)") + tp, w = grid.tp, (lambda name: weights[f"{hb.BLOCK}.{name}"]) + seq = seq_len or hidden.shape[1] + shard = grid.seq_shard(seq) + par.require_rows_on_every_rank(shard) # every rank refuses alike, before any collective + h = par.coshard(shard, hidden, dim=1)[0] + idx = par.coshard(shard, index, dim=0)[0] + cos, sin = par.coshard(shard, rope[0], rope[1], dim=0) + rows = h.shape[1] + if sp: + if rows % tp: + raise ValueError(f"sequence parallelism needs {rows} local rows divisible by tp={tp}") + share = rows // tp + lo, hi = grid.tp_rank * share, (grid.tp_rank + 1) * share + h_res, idx_res = h[:, lo:hi].contiguous(), idx[lo:hi].contiguous() + else: + h_res, idx_res = h, idx + coll = grid.tp_collective + + def leaf(name, tensor): + tensor = tensor.detach().clone() + if leaves is not None: + tensor.requires_grad_(True) + leaves[name] = tensor + return tensor + + table = ops.projection(temb, w("adaln_proj.linear.weight"), w("adaln_proj.linear.bias")) + + def proj_in(x, name_weights, gather): + flat = x.reshape(-1, HIDDEN) + return par.column_parallel_linear(flat, name_weights, coll, tp, gather=gather) + + x = ops.norm.forward_modulated(h_res, w("norm1.weight"), table[0], table[1], idx_res) + q, k, v = proj_in( + x, + [ + leaf(f"attn.to_{n}.weight", par.shard_rows(w(f"attn.to_{n}.weight"), tp, grid.tp_rank)) + for n in "qkv" + ], + sp, + ) + local_heads = par.HEADS // tp + + def heads(t): + return t.view(1, rows, local_heads, HEAD_DIM) + + q = ops.norm(heads(q), w("attn.norm_q.weight")) + k = ops.norm(heads(k), w("attn.norm_k.weight")) + from rl_engine.validation.models import h3_provider as prov + + q, k = prov.provider_apply_rotary(q, cos, sin), prov.provider_apply_rotary(k, cos, sin) + attn = par.ulysses_attention(q, k, heads(v), _attention_fn(ops), shard, grid.cp_group) + o = par.row_parallel_linear( + attn.reshape(rows, local_heads * HEAD_DIM), + leaf("attn.to_out.0.weight", par.shard_cols(w("attn.to_out.0.weight"), tp, grid.tp_rank)), + coll, + tp, + scatter=sp, + ) + h1 = ops.gate(h_res, o.view(1, -1, HIDDEN), table[2], idx_res) + x2 = ops.norm.forward_modulated(h1, w("norm2.weight"), table[3], table[4], idx_res) + fused = w("ff.net.0.proj.weight") + half = fused.shape[0] // 2 + hid, gate = proj_in( + x2, + [ + leaf("ff.hidden", par.shard_rows(fused[:half], tp, grid.tp_rank)), + leaf("ff.gate", par.shard_rows(fused[half:], tp, grid.tp_rank)), + ], + sp, + ) + act = ops.swiglu(gate.contiguous(), hid.contiguous()) + down = par.row_parallel_linear( + act, + leaf("ff.net.2.weight", par.shard_cols(w("ff.net.2.weight"), tp, grid.tp_rank)), + coll, + tp, + scatter=sp, + ) + return ops.gate(h1, down.view(1, -1, HIDDEN), table[5], idx_res) + + +def block_target( + collective, + dummy, + *, + tp: int, + cp: int, + sp: bool = False, + seq_len: int = 0, + backward: bool = False, + local_size: int = -1, + layout: str = "tiny", + seed: int = 0, +): + """``run_world`` target: this rank's rows of the parallel block (and gradients) for one case.""" + + del collective # the grid builds its own TP/CP groups + from rl_engine.validation.models.h3_weights import load_h3_block_weights + + grid = make_grid(tp, cp, local_size=local_size if local_size > 0 else None) + return block_case( + load_h3_block_weights("cuda"), + grid, + sp=sp, + seq_len=seq_len, + backward=backward, + layout=layout, + seed=seed, + ) + + +def block_case(weights, grid: Grid, *, sp=False, seq_len=0, backward=False, layout="tiny", seed=0): + """Run block 0 on ``grid``; the global rows this rank returns are in ``rows``.""" + + ctx = hb.make_block_context(weights, layout=layout, seed=seed) + n = seq_len or ctx["seq_len"] + hidden = ctx["hidden"][:, :n].clone().requires_grad_(backward) + index = ctx["index"][:n] + rope = (ctx["rope"][0][:n], ctx["rope"][1][:n]) + ops = hb.CandidateOps() + leaves: dict[str, torch.Tensor] | None = {} if backward else None + with torch.set_grad_enabled(backward): + out = run_tp_block( + weights, hidden, ctx["temb"], index, rope, ops, grid, sp=sp, seq_len=n, leaves=leaves + ) + shard = grid.seq_shard(n) + lo, hi = shard.lo, shard.hi + if sp: + share = (hi - lo) // grid.tp + lo, hi = lo + grid.tp_rank * share, lo + (grid.tp_rank + 1) * share + result = {"out": out.detach(), "rows": torch.tensor([lo, hi])} + if backward: + gen = torch.Generator().manual_seed(seed + 99) + full_grad = torch.randn(1, n, HIDDEN, generator=gen).to(out.dtype).cuda() + out.backward(full_grad[:, lo:hi]) + result["dx"] = hidden.grad[:, lo:hi].detach() + assert leaves is not None + for name, tensor in leaves.items(): + result["dw." + name] = tensor.grad.detach() + return result diff --git a/tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py b/tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py new file mode 100644 index 000000000..3c25869df --- /dev/null +++ b/tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 integration rows: checkpoint mapper and the vLLM-Omni rollout adapter.""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.models.minimax_h3 import checkpoint_mapper as cm +from rl_engine.models.minimax_h3 import fingerprint as fp + +needs_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +# --- checkpoint mapper ---------------------------------------------------------------------- + + +def _names(): + return list(fp.expected_tensor_table(fp.transformer_config())) + + +def test_every_checkpoint_tensor_maps_and_native_names_are_unique(): + names = _names() + assert len(names) == 638 + state = {n: torch.zeros(6, 2) for n in names} + native = cm.diffusers_to_native(state) + # q/k/v of each of the 52 blocks (50 + 2 refiner) fuse per kind (weight): 3 -> 1 + assert len(native) == 638 - 2 * 52 + assert len(set(native)) == len(native) + + +def test_diffusers_native_roundtrip_is_bitwise(): + names = _names() + gen = torch.Generator().manual_seed(0) + state = {n: torch.randn(6, 2, generator=gen) for n in names} + back = cm.native_to_diffusers(cm.diffusers_to_native(state)) + assert set(back) == set(state) + assert all(torch.equal(back[n], state[n]) for n in state) + + +def test_fc1_halves_are_swapped_and_qkv_is_q_then_k_then_v(): + w = torch.arange(8.0).view(8, 1) + state = { + "transformer_blocks.3.ff.net.0.proj.weight": w, + "transformer_blocks.3.attn.to_q.weight": torch.full((2, 1), 1.0), + "transformer_blocks.3.attn.to_k.weight": torch.full((2, 1), 2.0), + "transformer_blocks.3.attn.to_v.weight": torch.full((2, 1), 3.0), + } + native = cm.diffusers_to_native(state) + assert torch.equal(native["blocks.3.mlp.fc1.weight"], torch.cat([w[4:], w[:4]])) + assert native["blocks.3.attn.qkv_proj.weight"].flatten().tolist() == [1, 1, 2, 2, 3, 3] + + +def test_unmapped_names_fail_closed(): + with pytest.raises(ValueError, match="unmapped"): + cm.diffusers_to_native({"transformer_blocks.0.surprise.weight": torch.zeros(1)}) + with pytest.raises(ValueError, match="unmapped"): + cm.native_to_diffusers({"blocks.0.surprise.weight": torch.zeros(1)}) + with pytest.raises(ValueError, match="q, k and v"): + cm.diffusers_to_native({"transformer_blocks.0.attn.to_q.weight": torch.zeros(1, 1)}) + + +def test_grouped_checkpoint_layout_roundtrip(): + w = torch.randn(4 * 3 * 8, 5) + fused = cm.grouped_qkv_to_qkv(w, heads=4, head_dim=8) + assert torch.equal(cm.qkv_to_grouped(fused, heads=4, head_dim=8), w) + + +def test_tables_match_vllm_omni_when_importable(): + fasth3 = pytest.importorskip("vllm_omni.diffusion.models.minimax_h3.fasth3") + assert dict(fasth3._MODEL_LEVEL_TARGETS) == cm.MODEL_LEVEL + assert {k: v[0] for k, v in fasth3._BLOCK_TARGETS.items() if k in cm.BLOCK_TARGETS} == { + k: v[0] for k, v in cm.BLOCK_TARGETS.items() + } + + +class _Lin(torch.nn.Module): + def __init__(self, n, k): + super().__init__() + self.weight = torch.nn.Parameter( + (torch.randn(n, k) / 30).bfloat16().cuda(), requires_grad=False + ) + + def forward(self, x): + raise AssertionError("must be replaced") + + +@needs_gpu +def test_vllm_adapter_patches_projections_and_restores(): + pytest.importorskip("triton") + from rl_engine.backends.cuda.model_specific.minimax_h3 import parallel as par + from rl_engine.integrations.engines.rollout.vllm.h3_omni import install_h3_rollout + + torch.manual_seed(0) + model = torch.nn.Module() + model.blocks = torch.nn.ModuleList([torch.nn.Module()]) + blk = model.blocks[0] + blk.attn = torch.nn.Module() + blk.attn.qkv_proj, blk.attn.out_proj = _Lin(3 * 1024, 640), _Lin(640, 1024) + blk.mlp = torch.nn.Module() + blk.mlp.fc1, blk.mlp.fc2 = _Lin(2 * 2048, 640), _Lin(640, 2048) + x = torch.randn(2, 17, 640).bfloat16().cuda() + with install_h3_rollout(model) as handle: + qkv, bias = blk.attn.qkv_proj(x) + assert bias is None and qkv.shape == (2, 17, 3 * 1024) + q, k, v = (blk.attn.qkv_proj.weight[i * 1024 : (i + 1) * 1024] for i in range(3)) + sep = [par.g.gemm(x.reshape(-1, 640), w, b_t=True) for w in (q, k, v)] + assert torch.equal(qkv.reshape(-1, 3072), torch.cat(sep, dim=1)) # fused == separate + y, _ = blk.attn.out_proj(torch.randn(2, 17, 1024).bfloat16().cuda()) + assert y.shape == (2, 17, 640) + assert sorted(handle.patched) == [ + "blocks.0.attn.out_proj", + "blocks.0.attn.qkv_proj", + "blocks.0.mlp.fc1", + "blocks.0.mlp.fc2", + ] + with pytest.raises(AssertionError, match="must be replaced"): + blk.attn.qkv_proj(x) + + +@needs_gpu +def test_vllm_adapter_strict_rejects_wrong_dtype_and_empty_models(): + from rl_engine.integrations.engines.rollout.vllm.h3_omni import install_h3_rollout + + with pytest.raises(RuntimeError, match="no H3 projection"): + install_h3_rollout(torch.nn.Linear(2, 2).cuda()) + holder = _model_with(attn=("qkv_proj", torch.nn.Linear(8, 8).cuda())) # float32 + with pytest.raises(RuntimeError, match="BF16"): + install_h3_rollout(holder) + # the patched forward has no bias term: refuse, never drop it + biased = _model_with(mlp=("fc2", torch.nn.Linear(8, 8, bias=True).cuda().bfloat16())) + with pytest.raises(RuntimeError, match="has a bias"): + install_h3_rollout(biased) + fc2 = biased.blocks[0].mlp.fc2 + with install_h3_rollout(biased, strict=False) as handle: + assert handle.patched == [] and "forward" not in fc2.__dict__ + + +def _model_with(**parts): + """``blocks.0..`` holding the given module, as in vLLM-Omni's DiT.""" + + model = torch.nn.Module() + model.blocks = torch.nn.ModuleList([torch.nn.Module()]) + for group, (name, module) in parts.items(): + setattr(model.blocks[0], group, torch.nn.Module()) + setattr(getattr(model.blocks[0], group), name, module) + return model + + +@needs_gpu +def test_vllm_adapter_patches_block_projections_only_and_restores_exactly(): + from rl_engine.integrations.engines.rollout.vllm.h3_omni import install_h3_rollout + + model = _model_with(attn=("qkv_proj", _Lin(3 * 128, 64)), mlp=("fc1", _Lin(256, 64))) + model.audio_head = torch.nn.Module() + model.audio_head.fc1 = _Lin(64, 64) # same leaf name, not a block projection + own = model.blocks[0].mlp.fc1.forward = lambda x: ("own forward", x) + with install_h3_rollout(model) as handle: + assert sorted(handle.patched) == ["blocks.0.attn.qkv_proj", "blocks.0.mlp.fc1"] + assert "forward" not in model.audio_head.fc1.__dict__ + assert model.blocks[0].mlp.fc1.forward is own # the instance override is back + assert "forward" not in model.blocks[0].attn.qkv_proj.__dict__ + + +@needs_gpu +def test_vllm_adapter_strict_refusal_leaves_every_module_unpatched(): + from rl_engine.integrations.engines.rollout.vllm.h3_omni import install_h3_rollout + + model = _model_with(attn=("qkv_proj", _Lin(3 * 128, 64))) + model.blocks.append(torch.nn.Module()) + model.blocks[1].attn = torch.nn.Module() + model.blocks[1].attn.qkv_proj = torch.nn.Linear(64, 3 * 128).cuda() # float32: refused + with pytest.raises(RuntimeError, match="BF16"): + install_h3_rollout(model) + assert "forward" not in model.blocks[0].attn.qkv_proj.__dict__ + + +def test_qkv_layout_is_explicit_module_or_serialised_checkpoint(): + """The module layout is ``[Q; K; V]``; a serialised native checkpoint holds per-head groups + (vLLM-Omni's loader converts them). Each direction round-trips in its own layout.""" + + gen = torch.Generator().manual_seed(1) + heads, d = 2, cm.HEAD_DIM + state = { + f"transformer_blocks.0.attn.to_{s}.weight": torch.randn(heads * d, 3, generator=gen) + for s in "qkv" + } + module = cm.diffusers_to_native(state)["blocks.0.attn.qkv_proj.weight"] + ckpt = cm.diffusers_to_native(state, qkv_layout="checkpoint")["blocks.0.attn.qkv_proj.weight"] + assert torch.equal( + module, torch.cat([state[f"transformer_blocks.0.attn.to_{s}.weight"] for s in "qkv"]) + ) + assert torch.equal(ckpt, cm.qkv_to_grouped(module, heads=heads, head_dim=d)) + assert not torch.equal(ckpt, module) # writing the module layout for the loader is wrong + for layout, tensor in (("module", module), ("checkpoint", ckpt)): + back = cm.native_to_diffusers({"blocks.0.attn.qkv_proj.weight": tensor}, qkv_layout=layout) + assert all(torch.equal(back[n], state[n]) for n in state) + with pytest.raises(ValueError, match="qkv_layout"): + cm.diffusers_to_native(state, qkv_layout="grouped") + + +def test_native_tensor_kinds_other_than_weight_and_bias_are_refused(): + with pytest.raises(ValueError, match="unmapped"): + cm.native_to_diffusers({"video_patch_proj.running_mean": torch.zeros(1)}) + with pytest.raises(ValueError, match="unmapped"): + cm.native_to_diffusers({"blocks.0.mlp.fc2.scale": torch.zeros(1)}) + + +def test_checkpoint_qkv_layout_is_what_the_vllm_omni_loader_reorders(): + """vLLM-Omni's ``load_weights`` reorders a native ``attn.qkv_proj.weight`` from per-head groups + to ``[Q; K; V]``: the ``"checkpoint"`` layout must come out as the module layout.""" + + loader = pytest.importorskip("vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer") + gen = torch.Generator().manual_seed(2) + heads = 3 + state = { + f"transformer_blocks.0.attn.to_{s}.weight": torch.randn( + heads * cm.HEAD_DIM, 4, generator=gen + ) + for s in "qkv" + } + key = "blocks.0.attn.qkv_proj.weight" + module = cm.diffusers_to_native(state)[key] + ckpt = cm.diffusers_to_native(state, qkv_layout="checkpoint")[key] + loaded = loader._reorder_grouped_qkv_to_qkv( + ckpt, num_query_groups=heads, heads_per_group=1, head_dim=cm.HEAD_DIM + ) + assert torch.equal(loaded, module) diff --git a/tests/models/minimax_h3/test_h3_parallel.py b/tests/models/minimax_h3/test_h3_parallel.py new file mode 100644 index 000000000..a3c7a8be1 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_parallel.py @@ -0,0 +1,378 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 WS2 rows: TP attention/FFN, Ulysses CP, metadata co-shard, arbitrary length, packed IO. + +Single-process tests pin the canonical 8-unit reduction (summing per-rank partials in tree +order reproduces the TP=1 result for TP 2/4/8) and the sharding helpers. Multi-process tests +(2-4 GPUs) run the whole block 0 under TP / CP / SP grids and require every rank's rows to be +bitwise equal to the single-GPU canonical block, forward and backward. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3 import parallel as par +from rl_engine.distributed.algorithms.hierarchical import tree_add +from rl_engine.validation.distributed.h3_ws2 import run_world +from rl_engine.validation.models import h3_tp + +needs_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def gpus(n): + return pytest.mark.skipif(torch.cuda.device_count() < n, reason=f"needs {n} GPUs") + + +def _triton(): + pytest.importorskip("triton") + + +class TestSharding: + @pytest.mark.parametrize("tp", [1, 2, 4, 8]) + def test_row_and_column_shards_tile_once(self, tp): + w = torch.arange(56 * 128 * 16, dtype=torch.float32).view(56 * 128, 16) + assert torch.equal(torch.cat([par.shard_rows(w, tp, r) for r in range(tp)]), w) + w2 = torch.arange(8 * 64 * 4, dtype=torch.float32).view(4, 8 * 64) + assert torch.equal(torch.cat([par.shard_cols(w2, tp, r) for r in range(tp)], 1), w2) + + @pytest.mark.parametrize("tp", [3, 5, 6, 16]) + def test_illegal_tp_is_rejected(self, tp): + with pytest.raises(ValueError): + par.shard_rows(torch.zeros(56 * 128, 4), tp, 0) + + def test_shards_are_whole_heads(self): + for tp in (1, 2, 4, 8): + assert (par.HEADS // tp) * par.HEAD_DIM == par.INNER // tp + + def test_gate_up_shards_keep_both_halves(self): + w = torch.arange(2 * 64 * 3, dtype=torch.float32).view(128, 3) + s = par.shard_gate_up(w, 4, 1) + hidden, gate = s.chunk(2) + assert torch.equal(hidden, w[:64][16:32]) and torch.equal(gate, w[64:][16:32]) + + @pytest.mark.parametrize("length", [1, 2, 3, 7, 261, 264, 925]) + @pytest.mark.parametrize("cp", [1, 2, 3, 4]) + def test_seq_shards_tile_any_length(self, length, cp): + spans = [ + (par.SeqShard(length, cp, r).lo, par.SeqShard(length, cp, r).hi) for r in range(cp) + ] + assert spans[0][0] == 0 and spans[-1][1] == length + assert all(a[1] == b[0] for a, b in zip(spans, spans[1:])) + sizes = [hi - lo for lo, hi in spans] + assert max(sizes) - min(sizes) <= 1 + + def test_coshard_slices_every_per_row_tensor_alike(self): + n = 11 + idx, tags, rows = torch.arange(n), torch.arange(n) * 3, torch.arange(n * 2).view(n, 2) + for r in range(3): + s = par.SeqShard(n, 3, r) + a, b, c = par.coshard(s, idx, tags, rows) + assert torch.equal(a * 3, b) and torch.equal(c[:, 0] // 2, a) + + def test_coshard_rejects_wrong_length(self): + with pytest.raises(ValueError, match="rows"): + par.coshard(par.SeqShard(5, 2, 0), torch.zeros(4)) + + +@needs_gpu +class TestCanonicalReduction: + """Per-rank partials added in tree order are bitwise the TP=1 result (no collective needed).""" + + @pytest.mark.parametrize("k", [7168, 14336]) + @pytest.mark.parametrize("tp", [2, 4, 8]) + def test_row_parallel_partials_fold_to_tp1(self, k, tp): + _triton() + gen = torch.Generator().manual_seed(1) + x = torch.randn(33, k, generator=gen).bfloat16().cuda() + w = (torch.randn(5376, k, generator=gen) / k**0.5).bfloat16().cuda() + full = par._units_partial(x, w, 1) + parts = [ + par._units_partial(par.shard_heads(x, tp, r, 1), par.shard_cols(w, tp, r), tp) + for r in range(tp) + ] + assert torch.equal(tree_add(parts), full) + assert torch.equal(par.row_parallel_linear(x, w), full.bfloat16()) + + def test_row_parallel_close_to_fp64(self): + _triton() + gen = torch.Generator().manual_seed(2) + x = torch.randn(64, 7168, generator=gen).bfloat16().cuda() + w = (torch.randn(5376, 7168, generator=gen) / 7168**0.5).bfloat16().cuda() + got = par.row_parallel_linear(x, w).double() + ref = x.double() @ w.double().t() + provider = torch.nn.functional.linear(x, w).double() + err = (got - ref).norm() / ref.norm() + assert err <= 1.25 * (provider - ref).norm() / ref.norm() + 1e-3 + + def test_each_row_is_independent_of_the_others(self): + _triton() + gen = torch.Generator().manual_seed(3) + x = torch.randn(70, 7168, generator=gen).bfloat16().cuda() + w = (torch.randn(5376, 7168, generator=gen) / 7168**0.5).bfloat16().cuda() + whole = par.row_parallel_linear(x, w) + for rows in (slice(0, 1), slice(5, 6), slice(30, 31), slice(0, 33)): + assert torch.equal(par.row_parallel_linear(x[rows].contiguous(), w), whole[rows]) + + @pytest.mark.parametrize("tp", [2, 4, 8]) + def test_column_parallel_backward_folds_to_tp1(self, tp): + _triton() + gen = torch.Generator().manual_seed(4) + k, n = 5376, 7168 + ws = [(torch.randn(n, k, generator=gen) / k**0.5).bfloat16().cuda() for _ in range(3)] + grads = [torch.randn(40, n, generator=gen).bfloat16().cuda() for _ in range(3)] + + def unit_partials(weights, gs, local): + units = [] + for j in range(local): + acc = None + for grad, w in zip(gs, weights): + width = grad.shape[1] // local + piece = par.g.gemm( + grad[:, j * width : (j + 1) * width], + w[j * width : (j + 1) * width], + out_dtype=torch.float32, + ) + acc = piece if acc is None else acc + piece + units.append(acc) + return tree_add(units) + + full = unit_partials(ws, grads, 8) + ranks = [ + unit_partials( + [par.shard_rows(w, tp, r) for w in ws], + [par.shard_heads(g_, tp, r, 1) for g_ in grads], + 8 // tp, + ) + for r in range(tp) + ] + assert torch.equal(tree_add(ranks), full) + + def test_column_parallel_forward_is_the_unsharded_columns(self): + _triton() + gen = torch.Generator().manual_seed(5) + x = torch.randn(40, 5376, generator=gen).bfloat16().cuda() + w = (torch.randn(7168, 5376, generator=gen) / 5376**0.5).bfloat16().cuda() + (full,) = par.column_parallel_linear(x, [w]) + for tp in (2, 4, 8): + for r in range(tp): + (part,) = par.column_parallel_linear(x, [par.shard_rows(w, tp, r)]) + assert torch.equal(part, par.shard_heads(full, tp, r, 1)) + + +@needs_gpu +class TestPackedIO: + @pytest.mark.parametrize("length", [1, 5, 261]) + def test_scatter_then_gather_is_identity_single_rank(self, length): + t = torch.randn(length, 7).cuda() + shard = par.SeqShard(length, 1, 0) + assert torch.equal(par.gather_packed(par.scatter_packed(t, shard), shard), t) + + +# --- Multi-process grids -------------------------------------------------------------------- + + +def _reference(weights, seq_len=0, backward=False): + return h3_tp.block_case(weights, h3_tp.Grid(1, 1, 0, 0), seq_len=seq_len, backward=backward) + + +def _assemble(ranks, tp, cp, key="out"): + """Global rows from the ranks; replicated ranks must agree bitwise.""" + + pieces: dict[tuple[int, int], torch.Tensor] = {} + for r in ranks: + span = tuple(int(v) for v in r["rows"]) + if span in pieces: + assert torch.equal(pieces[span], r[key]), f"ranks disagree on rows {span}" + pieces[span] = r[key] + return torch.cat([pieces[s] for s in sorted(pieces)], dim=1), sorted(pieces) + + +GRIDS = [ + pytest.param(2, 1, False, 0, id="tp2", marks=gpus(2)), + pytest.param(4, 1, False, 0, id="tp4", marks=gpus(4)), + pytest.param(1, 2, False, 0, id="cp2", marks=gpus(2)), + pytest.param(2, 2, False, 0, id="tp2_cp2", marks=gpus(4)), + pytest.param(2, 1, True, 0, id="tp2_sp", marks=gpus(2)), + pytest.param(4, 1, True, 0, id="tp4_sp", marks=gpus(4)), + pytest.param(2, 2, True, 0, id="tp2_cp2_sp", marks=gpus(4)), + pytest.param(1, 2, False, 261, id="cp2_len261", marks=gpus(2)), + pytest.param(2, 2, False, 261, id="tp2_cp2_len261", marks=gpus(4)), + pytest.param(4, 1, False, 263, id="tp4_len263", marks=gpus(4)), +] + + +@needs_gpu +@pytest.mark.parametrize(("tp", "cp", "sp", "length"), GRIDS) +def test_parallel_block_equals_single_gpu_bitwise(tp, cp, sp, length, h3_block_weights_cuda): + # the fixture skips without extracted weights; every rank loads the same pinned files + _triton() + ranks = run_world( + tp * cp, + h3_tp.block_target, + {"dummy": torch.zeros(1)}, + tp=tp, + cp=cp, + sp=sp, + seq_len=length, + backward=True, + ) + ref = _reference(h3_block_weights_cuda, length, backward=True) + out, spans = _assemble(ranks, tp, cp) + assert out.shape == ref["out"].shape, (spans, out.shape) + assert torch.equal(out.cpu(), ref["out"].cpu()), "forward differs from the single-GPU block" + dx, _ = _assemble(ranks, tp, cp, "dx") + assert torch.equal(dx.cpu(), ref["dx"].cpu()), "input gradient differs" + if cp == 1: # weight gradients: shards of the unsharded gradient + for name, axis in ( + ("attn.to_out.0.weight", 1), + ("ff.net.2.weight", 1), + ("attn.to_q.weight", 0), + ("attn.to_k.weight", 0), + ("attn.to_v.weight", 0), + ("ff.hidden", 0), + ("ff.gate", 0), + ): + if sp: + continue # SP shards sequence rows; dW is reduced over rows below + full = ref["dw." + name].cpu() + shards = [r["dw." + name].cpu() for r in sorted(ranks, key=lambda r: r["rank"])] + assert torch.equal(torch.cat(shards, dim=axis), full), f"dW {name} differs" + + +@needs_gpu +@gpus(2) +def test_ulysses_matches_unsharded_attention_for_ragged_lengths(): + _triton() + for length in (1, 2, 7, 130): + ranks = run_world(2, _ulysses_target, {"dummy": torch.zeros(1)}, length=length) + gathered = torch.cat([r["out"] for r in ranks], dim=1) + assert torch.equal(gathered.cpu(), ranks[0]["ref"].cpu()), f"length {length}" + + +def _ulysses_target(collective, dummy, *, length): + import torch.distributed as dist + + from rl_engine.backends.cuda.attention.deterministic_attn import DeterministicAttentionOp + + del collective, dummy + rank = dist.get_rank() + gen = torch.Generator().manual_seed(7) + q, k, v = (torch.randn(1, length, 8, 128, generator=gen).bfloat16().cuda() for _ in range(3)) + op = DeterministicAttentionOp() + + def attend(q_, k_, v_): + t = lambda x: x.transpose(1, 2).contiguous() # noqa: E731 + return op(t(q_), t(k_), t(v_), causal=False).transpose(1, 2).contiguous() + + shard = par.SeqShard(length, 2, rank) + local = [par.scatter_packed(t, shard, dim=1) for t in (q, k, v)] + out = par.ulysses_attention(*local, attend, shard, dist.group.WORLD) + return {"out": out, "ref": attend(q, k, v)} + + +# --- Hierarchical (cross-node design) collective, exercised inside one node ---------------- + + +def _hier_target(collective, dummy, *, local_size, dtype): + import torch.distributed as dist + + from rl_engine.distributed.algorithms.hierarchical import HierarchicalCollective + + del dummy + h = HierarchicalCollective( + local_size, device=torch.cuda.current_device(), max_size_bytes=64 << 20 + ) + gen = torch.Generator().manual_seed(100 + dist.get_rank()) + x = ( + (torch.randn(257, 129, generator=gen) * 10 ** torch.randint(-3, 4, (257, 1), generator=gen)) + .to(dtype) + .cuda() + ) + flat = collective.all_reduce(x.clone()) + mine = h.all_reduce(x) + rs = h.reduce_scatter(x[:256].contiguous()) + ag = h.all_gather(x[:2].contiguous()) + return {"flat": flat, "hier": mine, "rs": rs, "ag": ag, "x": x} + + +@needs_gpu +@gpus(4) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("local_size", [1, 2]) +def test_hierarchical_equals_single_node_fixed_tree(local_size, dtype): + ranks = run_world( + 4, _hier_target, {"dummy": torch.zeros(1)}, local_size=local_size, dtype=dtype + ) + for r in ranks: + assert torch.equal(r["flat"], r["hier"]), f"rank {r['rank']}: hierarchical != fixed tree" + assert torch.equal(ranks[0]["rs"], ranks[0]["hier"][:64]) + assert torch.equal(ranks[3]["rs"], ranks[3]["hier"][192:256]) + assert torch.equal(ranks[1]["ag"], torch.cat([r["x"][:2] for r in ranks])) + + +def _fake_dist(rank: int, world: int, calls: list): + """Just enough of ``torch.distributed`` for one rank to build a ``HierarchicalCollective``.""" + + from types import SimpleNamespace + + everything = tuple(range(world)) + + def new_group(ranks): + calls.append(tuple(ranks)) + return tuple(ranks) + + return SimpleNamespace( + is_initialized=lambda: True, + group=SimpleNamespace(WORLD=everything), + get_world_size=lambda group=None: len(everything if group is None else group), + get_rank=lambda group=None: (everything if group is None else group).index(rank), + get_process_group_ranks=lambda group: list(group), + new_group=new_group, + ProcessGroup=tuple, + ) + + +@pytest.mark.parametrize(("tp", "cp", "local_size"), [(2, 2, 1), (4, 2, 2), (8, 2, 4), (8, 1, 4)]) +def test_hierarchical_subgroups_are_created_in_one_global_order(monkeypatch, tp, cp, local_size): + """``dist.new_group`` is entered by every rank with the same lists in the same order, also + when several TP groups (CP > 1) each build a hierarchical collective.""" + + from types import SimpleNamespace + + import rl_engine.distributed.algorithms.hierarchical as hier + + world = tp * cp + tp_ranks = [list(range(c * tp, (c + 1) * tp)) for c in range(cp)] + monkeypatch.setattr(hier, "DeterministicCollective", lambda g, d, **_: SimpleNamespace(g=g)) + sequences = [] + for rank in range(world): + calls: list = [] + monkeypatch.setattr(hier, "dist", _fake_dist(rank, world, calls)) + mine = tuple(tp_ranks[rank // tp]) + h = hier.HierarchicalCollective(local_size, mine, group_ranks=tp_ranks) + sequences.append(calls) + assert rank in h._local.g and rank in h._cross # its own subgroups, of its own TP group + assert set(h._local.g) | set(h._cross) <= set(mine) + assert len(h._local.g) == local_size and len(h._cross) == tp // local_size + assert all(seq == sequences[0] for seq in sequences) + if cp > 1: # a TP subgroup alone would make the ranks disagree; it is refused instead + monkeypatch.setattr(hier, "dist", _fake_dist(0, world, [])) + with pytest.raises(ValueError, match="group_ranks"): + hier.HierarchicalCollective(local_size, tuple(tp_ranks[0])) + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_compute_paths_refuse_a_sequence_shorter_than_cp_on_every_rank(rank): + """``seq_len < cp`` is fine for data movement but not for the compute paths: every rank + (also the one that holds a row) refuses before any GEMM or collective.""" + + shard = par.SeqShard(1, 2, rank) + with pytest.raises(ValueError, match="at least one row per rank"): + par.seq_sharded_linear(torch.ones(shard.hi - shard.lo, 4), torch.ones(3, 4), shard) + grid = h3_tp.Grid(1, 2, 0, rank) + with pytest.raises(ValueError, match="at least one row per rank"): + h3_tp.run_tp_block({}, torch.ones(1, 1, 8), None, torch.zeros(1), (None, None), None, grid) + par.require_rows_on_every_rank(par.SeqShard(2, 2, rank)) # one row each: accepted diff --git a/tools/validation/models/h3_full_speed.py b/tools/validation/models/h3_full_speed.py new file mode 100644 index 000000000..7a7682ca4 --- /dev/null +++ b/tools/validation/models/h3_full_speed.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Time one full-model forward + backward step, provider against the strict RL-Kernel graph. + + python tools/validation/models/h3_full_speed.py --checkpoint \ + --sizes 12,925,4253,8349 --out speed.json +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import time +from pathlib import Path + +import torch + +SIZES = { + "12": None, + "925": dict(text=77, frames=3, side=32, audio=40), + "4253": dict(text=77, frames=1, side=128, audio=40), + "8349": dict(text=77, frames=2, side=128, audio=40), +} + + +def sized_inputs(device, *, text, frames, side, audio): + from rl_engine.validation.models.h3_cases import h3_fl2va_layout + + layout = h3_fl2va_layout( + num_text_tokens=text, + num_latent_frames=frames, + latent_height=side, + latent_width=side, + num_audio_latents=audio, + device=device, + ) + tags = layout["token_tags"] + video, text_idx, audio_idx = ((tags == t).nonzero().flatten() for t in (0, 1, 2)) + gen = torch.Generator().manual_seed(420539) + + def rnd(*shape): + return torch.randn(shape, generator=gen).to(device).requires_grad_() + + return dict( + **layout, + hidden_states=rnd(1, len(video), 96), + audio_hidden_states=rnd(1, len(audio_idx), 32), + encoder_hidden_states=rnd(1, len(text_idx), 5120), + timestep=torch.tensor([0.0, 500.0], device=device), + video_indices=video, + text_indices=text_idx, + audio_indices=audio_idx, + return_dict=False, + ) + + +def timed(model, inputs, candidate, repeats, attention="auto"): + from rl_engine.validation.models.h3_full_model import FullModelTrace, full_model_execution + + grad_inputs = ("hidden_states", "audio_hidden_states", "encoder_hidden_states") + backends = {} + ops = None + if candidate: # built once and outside every timed region: registry lookups, CUDA init + from rl_engine.validation.models.h3_block import CandidateOps + + ops = CandidateOps(attention=attention) + + def once(): + fresh = { + k: (v.detach().clone().requires_grad_() if k in grad_inputs else v) + for k, v in inputs.items() + } + trace = FullModelTrace() + with full_model_execution( + model, candidate=candidate, trace=trace, ops=ops, strict_boundaries=candidate + ): + # timed: forward and backward of the bound model only, not the binding setup + torch.cuda.synchronize() + t0 = time.perf_counter() + with torch.autograd.graph.save_on_cpu(pin_memory=True): + video, audio = model(**fresh) + torch.cuda.synchronize() + t1 = time.perf_counter() + (video.float().square().mean() + audio.float().square().mean()).backward() + torch.cuda.synchronize() + backends.update(getattr(trace, "actual_backends", None) or {}) + return t1 - t0, time.perf_counter() - t0 + + torch.cuda.reset_peak_memory_stats() + once() + runs = [once() for _ in range(repeats)] + return dict( + forward_s=statistics.median(r[0] for r in runs), + total_s=statistics.median(r[1] for r in runs), + peak_gib=torch.cuda.max_memory_allocated() / 2**30, + attention_backend=backends.get("attention_h3") or backends.get("attention"), + ) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--checkpoint", type=Path, required=True) + p.add_argument("--sizes", default="12,925,4253,8349") + p.add_argument("--attention", default="auto", choices=("auto", "h3_full", "interim")) + p.add_argument("--repeats", type=int, default=3) + p.add_argument("--out", type=Path, required=True) + args = p.parse_args() + from diffusers import MiniMaxH3Transformer3DModel + + from tools.validation.models.h3_full_model import packed_inputs + + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + model = MiniMaxH3Transformer3DModel.from_pretrained( + str(args.checkpoint / "transformer"), + torch_dtype=torch.bfloat16, + local_files_only=True, + low_cpu_mem_usage=True, + ).to("cuda") + model.eval().enable_gradient_checkpointing() + for q in model.parameters(): + q.register_post_accumulate_grad_hook(lambda t: setattr(t, "grad", None)) + result = dict(device=torch.cuda.get_device_name(0), torch=torch.__version__, sizes={}) + for size in args.sizes.split(","): + inputs = ( + packed_inputs("cuda") if SIZES[size] is None else sized_inputs("cuda", **SIZES[size]) + ) + entry = dict(tokens=int(inputs["position_ids"].shape[0])) + for name, candidate in (("provider", False), ("strict", True)): + try: + entry[name] = timed(model, inputs, candidate, args.repeats, args.attention) + except torch.OutOfMemoryError: + torch.cuda.empty_cache() + entry[name] = dict(error="CUDA out of memory") + if "total_s" in entry["provider"] and "total_s" in entry["strict"]: + entry["ratio"] = entry["strict"]["total_s"] / entry["provider"]["total_s"] + entry["fwd_ratio"] = entry["strict"]["forward_s"] / entry["provider"]["forward_s"] + result["sizes"][size] = entry + print(size, json.dumps(entry), flush=True) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(result, indent=1)) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_mapper_vllm_report.py b/tools/validation/models/h3_mapper_vllm_report.py new file mode 100644 index 000000000..1b3e50fbe --- /dev/null +++ b/tools/validation/models/h3_mapper_vllm_report.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Measurements of the RFC #420 ``h3_checkpoint_mapper`` and ``vllm_runtime_adapter_h3`` rows on +the real checkpoint: JSON report plus a PNG. + + python tools/validation/models/h3_mapper_vllm_report.py --checkpoint \ + --out report.json --png figure.png +""" + +from __future__ import annotations + +import argparse +import json +import re +import time +from collections import defaultdict +from pathlib import Path + +import torch +import torch.nn.functional as F +from safetensors import safe_open + +from rl_engine.backends.shared.triton.gemm import h3_projection as g +from rl_engine.integrations.engines.rollout.vllm.h3_omni import install_h3_rollout +from rl_engine.models.minimax_h3 import checkpoint_mapper as cm + +GROUP = re.compile(r"^((?:token_refiner\.refiner_blocks|transformer_blocks)\.\d+)\.") + + +def rel(a, b): + a, b = a.double(), b.double() + return ((a - b).norm() / b.norm()).item() + + +def mapping_report(root): + index = {} + for f in sorted(root.glob("*.safetensors")): + with safe_open(f, "pt") as h: + for k in h.keys(): + index[k] = f + groups = defaultdict(list) + for name in index: + m = GROUP.match(name) + groups[m.group(1) if m else "model-level"].append(name) + tensors = native_count = 0 + bitwise, t_fwd, t_back = True, 0.0, 0.0 + native_names = set() + for names in groups.values(): + state = {} + for n in names: + with safe_open(index[n], "pt") as h: + state[n] = h.get_tensor(n) + t0 = time.perf_counter() + native = cm.diffusers_to_native(state) + t1 = time.perf_counter() + back = cm.native_to_diffusers(native) + t_back += time.perf_counter() - t1 + t_fwd += t1 - t0 + tensors += len(state) + native_count += len(native) + native_names |= set(native) + bitwise &= set(back) == set(state) and all( + torch.equal(back[n], state[n]) and back[n].dtype == state[n].dtype for n in state + ) + return ( + dict( + checkpoint_tensors=tensors, + native_tensors=native_count, + unique_native_names=len(native_names), + groups=len(groups), + roundtrip_bitwise=bool(bitwise), + map_seconds=t_fwd, + inverse_seconds=t_back, + ), + index, + ) + + +def load(index, name): + with safe_open(index[name], "pt") as h: + return h.get_tensor(name).cuda() + + +def timeit(fn, n=20): + for _ in range(3): + fn() + torch.cuda.synchronize() + s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + s.record() + for _ in range(n): + fn() + e.record() + torch.cuda.synchronize() + return s.elapsed_time(e) / n + + +class Lin(torch.nn.Module): + def __init__(self, w): + super().__init__() + self.weight = torch.nn.Parameter(w, requires_grad=False) + + def forward(self, x): + return F.linear(x, self.weight), None + + +def layer_rows(index, tokens=(70, 925)): + rows = [] + for blk in (0, 25, 49): + p = f"transformer_blocks.{blk}." + q, k, v = (load(index, p + f"attn.to_{s}.weight") for s in "qkv") + o = load(index, p + "attn.to_out.0.weight") + f1 = load(index, p + "ff.net.0.proj.weight") + f2 = load(index, p + "ff.net.2.weight") + native = cm.diffusers_to_native( + { + p + "attn.to_q.weight": q, + p + "attn.to_k.weight": k, + p + "attn.to_v.weight": v, + p + "attn.to_out.0.weight": o, + p + "ff.net.0.proj.weight": f1, + p + "ff.net.2.weight": f2, + } + ) + model = torch.nn.Module() + model.blocks = torch.nn.ModuleList([torch.nn.Module()]) + b = model.blocks[0] + b.attn, b.mlp = torch.nn.Module(), torch.nn.Module() + b.attn.qkv_proj = Lin(native["blocks.%d.attn.qkv_proj.weight" % blk]) + b.attn.out_proj = Lin(native["blocks.%d.attn.out_proj.weight" % blk]) + b.mlp.fc1 = Lin(native["blocks.%d.mlp.fc1.weight" % blk]) + b.mlp.fc2 = Lin(native["blocks.%d.mlp.fc2.weight" % blk]) + for t in tokens: + gen = torch.Generator().manual_seed(blk * 1000 + t) + hidden = q.shape[1] + x = torch.randn(1, t, hidden, generator=gen).bfloat16().cuda() + xo = torch.randn(1, t, o.shape[1], generator=gen).bfloat16().cuda() + x2 = torch.randn(1, t, f2.shape[1], generator=gen).bfloat16().cuda() + ref = { + "qkv_proj": (x, b.attn.qkv_proj), + "out_proj": (xo, b.attn.out_proj), + "fc1": (x, b.mlp.fc1), + "fc2": (x2, b.mlp.fc2), + } + torch_out = {n: m(xi)[0] for n, (xi, m) in ref.items()} + torch_ms = {n: timeit(lambda m=m, xi=xi: m(xi)) for n, (xi, m) in ref.items()} + with install_h3_rollout(model): + out = {n: m(xi)[0] for n, (xi, m) in ref.items()} + ms = {n: timeit(lambda m=m, xi=xi: m(xi)) for n, (xi, m) in ref.items()} + again = {n: m(xi)[0] for n, (xi, m) in ref.items()} + flat = x.reshape(-1, hidden) + sep = torch.cat([g.gemm(flat, w, b_t=True) for w in (q, k, v)], dim=1) + half = f1.shape[0] // 2 + sep_fc1 = g.gemm(flat, f1, b_t=True) + fc1_native = out["fc1"].reshape(t, -1) + swapped = torch.cat([sep_fc1[:, half:], sep_fc1[:, :half]], dim=1) + for n, (xi, m) in ref.items(): + exact = xi.reshape(-1, xi.shape[-1]).double() @ m.weight.double().t() + rows.append( + dict( + block=blk, + tokens=t, + linear=n, + repeatable=bool(torch.equal(out[n], again[n])), + rel_err_vs_fp64=rel(out[n].reshape(t, -1), exact), + torch_rel_err_vs_fp64=rel(torch_out[n].reshape(t, -1), exact), + fused_equals_separate=( + bool(torch.equal(out[n].reshape(t, -1), sep)) + if n == "qkv_proj" + else None + ), + fc1_equals_diffusers_swapped=( + bool(torch.equal(fc1_native, swapped)) if n == "fc1" else None + ), + adapter_ms=ms[n], + torch_ms=torch_ms[n], + slowdown=ms[n] / torch_ms[n], + ) + ) + return rows + + +def plot(report, path): + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + rows = report["layers"] + names = ("qkv_proj", "out_proj", "fc1", "fc2") + fig, ax = plt.subplots(1, 3, figsize=(18, 5)) + m = report["mapping"] + ax[0].bar( + ["checkpoint\ntensors", "native\ntensors", "round trip\nbitwise equal"], + [ + m["checkpoint_tensors"], + m["native_tensors"], + m["checkpoint_tensors"] if m["roundtrip_bitwise"] else 0, + ], + ) + ax[0].set_title("Checkpoint mapper (real 33B checkpoint)") + w = 0.38 + for j, key in enumerate(("rel_err_vs_fp64", "torch_rel_err_vs_fp64")): + vals = [max(r[key] for r in rows if r["linear"] == n) for n in names] + ax[1].bar( + [i + (j - 0.5) * w for i in range(4)], + vals, + w, + label=("adapter" if j == 0 else "torch F.linear"), + ) + ax[1].set_xticks(range(4), names) + ax[1].set_yscale("log") + ax[1].set_title("Relative error vs FP64 (worst of 3 blocks x 2 sizes)") + ax[1].legend() + for t, c in ((70, "C0"), (925, "C1")): + vals = [ + sum(r["slowdown"] for r in rows if r["linear"] == n and r["tokens"] == t) / 3 + for n in names + ] + ax[2].plot(names, vals, "o-", color=c, label=f"{t} tokens") + ax[2].axhline(1, color="gray", lw=1) + ax[2].set_title("Adapter time / torch F.linear time (mean over blocks)") + ax[2].legend() + fig.tight_layout() + fig.savefig(path, dpi=130) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--checkpoint", type=Path, required=True) + p.add_argument("--out", type=Path, required=True) + p.add_argument("--png", type=Path) + args = p.parse_args() + torch.backends.cuda.matmul.allow_tf32 = False + mapping, index = mapping_report(args.checkpoint / "transformer") + report = dict( + device=torch.cuda.get_device_name(0), + torch=torch.__version__, + mapping=mapping, + layers=layer_rows(index), + ) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=1)) + if args.png: + plot(report, args.png) + print(json.dumps(mapping, indent=1)) + for r in report["layers"]: + print(r) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_parallel_run.py b/tools/validation/models/h3_parallel_run.py new file mode 100644 index 000000000..db5576bb2 --- /dev/null +++ b/tools/validation/models/h3_parallel_run.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Run the H3 parallel block on a TP x CP grid under ``srun`` (one task per GPU, any node count). + + srun -N2 --ntasks-per-node=4 python tools/validation/models/h3_parallel_run.py \\ + --tp 8 --cp 1 --out evidence/ws2_tp8 + +Every rank writes its rows (and gradients) to ``--out/rank{r}.pt``; rank 0 then recomputes the +single-GPU canonical block and compares the assembled forward, input gradient and weight-gradient +shards bitwise, and writes ``report.json``. ``--local-size`` is the number of GPUs per node for a +TP group that spans nodes (selects ``hierarchical_ipc_nccl_tree``); a TP group inside one node +uses the CUDA-IPC fixed tree. +""" + +from __future__ import annotations + +import argparse +import json +import os +import socket +import time +from pathlib import Path + +import torch +import torch.distributed as dist + +from rl_engine.validation.models import h3_tp +from rl_engine.validation.models.h3_weights import load_h3_block_weights + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tp", type=int, required=True) + parser.add_argument("--cp", type=int, default=1) + parser.add_argument("--sp", action="store_true") + parser.add_argument("--seq-len", type=int, default=0) + parser.add_argument("--layout", default="tiny") + parser.add_argument("--local-size", type=int, default=4, help="GPUs per node") + parser.add_argument("--iters", type=int, default=5) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + + rank = int(os.environ["SLURM_PROCID"]) + world = int(os.environ["SLURM_NTASKS"]) + local = int(os.environ["SLURM_LOCALID"]) + torch.cuda.set_device(local) + dist.init_process_group( + "nccl", + init_method=f"tcp://{os.environ['MASTER_ADDR']}:{os.environ.get('MASTER_PORT', '29511')}", + rank=rank, + world_size=world, + device_id=torch.device("cuda", local), + ) + if world != args.tp * args.cp: + raise SystemExit(f"world {world} != tp {args.tp} x cp {args.cp}") + probe = torch.ones(1, device="cuda") * rank + dist.all_reduce(probe) # NCCL connectivity across nodes + spans_nodes = args.tp > args.local_size + grid = h3_tp.make_grid(args.tp, args.cp, local_size=args.local_size if spans_nodes else None) + weights = load_h3_block_weights("cuda") + kwargs = dict(sp=args.sp, seq_len=args.seq_len, layout=args.layout) + result = h3_tp.block_case(weights, grid, backward=True, **kwargs) + # forward_s: time of a whole forward-only block_case call. It includes rebuilding the block + # context, the operators and the weight shards on every call (kept so the recorded reports + # stay comparable); it is not the time of the kernels alone. + torch.cuda.synchronize() + dist.barrier() + begin = time.perf_counter() + for _ in range(args.iters): + h3_tp.block_case(weights, grid, backward=False, **kwargs) + torch.cuda.synchronize() + seconds = (time.perf_counter() - begin) / args.iters + args.out.mkdir(parents=True, exist_ok=True) + torch.save({k: v.cpu() for k, v in result.items()}, args.out / f"rank{rank}.pt") + dist.barrier() + timings = [None] * world + dist.all_gather_object( + timings, {"rank": rank, "host": socket.gethostname(), "forward_s": seconds} + ) + if rank == 0: + ref = h3_tp.block_case( + weights, h3_tp.Grid(1, 1, 0, 0), backward=True, **{**kwargs, "sp": False} + ) + ranks = [torch.load(args.out / f"rank{r}.pt") for r in range(world)] + + def assemble(key): + pieces = {} + for r in ranks: + span = tuple(int(v) for v in r["rows"]) + if span in pieces and not torch.equal(pieces[span], r[key]): + return None + pieces[span] = r[key] + return torch.cat([pieces[s] for s in sorted(pieces)], dim=1) + + out, dx = assemble("out"), assemble("dx") + checks = { + "forward_bitwise": out is not None and torch.equal(out, ref["out"].cpu()), + "input_grad_bitwise": dx is not None and torch.equal(dx, ref["dx"].cpu()), + } + if args.cp == 1 and not args.sp: + for name, axis in ( + ("attn.to_out.0.weight", 1), + ("ff.net.2.weight", 1), + ("attn.to_q.weight", 0), + ("attn.to_k.weight", 0), + ("attn.to_v.weight", 0), + ("ff.hidden", 0), + ("ff.gate", 0), + ): + shards = [r["dw." + name] for r in ranks] + checks["dw." + name] = torch.equal( + torch.cat(shards, dim=axis), ref["dw." + name].cpu() + ) + report = dict( + tp=args.tp, + cp=args.cp, + sp=args.sp, + world=world, + hosts=sorted({t["host"] for t in timings}), + collective="hierarchical_ipc_nccl_tree" if spans_nodes else "cuda_ipc_fixed_tree", + seq_len=int(ref["out"].shape[1]) if not args.seq_len else args.seq_len, + checks=checks, + all_bitwise=all(checks.values()), + forward_seconds_per_rank=timings, + torch=torch.__version__, + device=torch.cuda.get_device_name(), + ) + (args.out / "report.json").write_text(json.dumps(report, indent=1) + "\n") + print( + json.dumps( + { + k: report[k] + for k in ("tp", "cp", "sp", "hosts", "collective", "checks", "all_bitwise") + }, + indent=1, + ) + ) + dist.barrier() + if rank == 0: + time.sleep(3) # keep the rendezvous store up until every rank has left + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_ws2_charts.py b/tools/validation/models/h3_ws2_charts.py new file mode 100644 index 000000000..ae9261524 --- /dev/null +++ b/tools/validation/models/h3_ws2_charts.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Chart of the WS2 evidence: bitwise agreement and per-rank forward time per parallel layout.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 + + +def main(evidence: Path, out: Path) -> None: + # h3_parallel_run.py writes /report.json; published evidence keeps one + # .json per layout. Read either layout, and refuse an empty chart. + paths = sorted(evidence.glob("*/report.json")) or sorted(evidence.glob("*.json")) + loaded = [json.loads(p.read_text()) for p in paths] + reports = [r for r in loaded if isinstance(r, dict) and "forward_seconds_per_rank" in r] + if not reports: + raise SystemExit(f"no WS2 reports under {evidence} (*/report.json or .json)") + labels = [] + times = [] + colors = [] + for r in reports: + labels.append(f"TP{r['tp']}xCP{r['cp']}{'+SP' if r.get('sp') else ''}") + ranks = r["forward_seconds_per_rank"] + times.append(sum(x["forward_s"] for x in ranks) / len(ranks)) + colors.append("tab:green" if r["all_bitwise"] else "tab:red") + fig, ax = plt.subplots(figsize=(8, 4)) + ax.bar(labels, times, color=colors) + ax.set( + ylabel="mean forward s per rank (H100)", title="WS2 layouts: green = bitwise == single GPU" + ) + for i, r in enumerate(reports): + ax.text( + i, times[i], f"{r['world']} GPU\n{len(set(r['hosts']))} node", ha="center", va="bottom" + ) + out.parent.mkdir(parents=True, exist_ok=True) + fig.tight_layout() + fig.savefig(out, dpi=120) + + +if __name__ == "__main__": + main(Path(sys.argv[1]), Path(sys.argv[2]))