From 1baed6e2d916384f0c404f3a6db16b29fd9a3f39 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 18:47:03 +0000 Subject: [PATCH 01/12] feat: support dense canonical multi-adapter LoRA Keep independent adapter ownership through canonical root slicing and checkpoint recomputation. Reject segmented routed experts. Co-Authored-By: Claude Opus 5.5 --- src/art/megatron/canonical/operators.py | 8 +- .../kernels/cute_grouped_lora_quack.py | 6 +- src/art/megatron/lora.py | 25 ++ src/art/megatron/segmented_lora.py | 202 ++++++++++++ .../megatron/lora/test_segmented_lora.py | 289 ++++++++++++++++++ tests/unit/test_canonical_numerics.py | 45 ++- tests/unit/test_segmented_lora.py | 239 +++++++++++++++ 7 files changed, 806 insertions(+), 8 deletions(-) create mode 100644 src/art/megatron/segmented_lora.py create mode 100644 tests/integration/megatron/lora/test_segmented_lora.py create mode 100644 tests/unit/test_segmented_lora.py diff --git a/src/art/megatron/canonical/operators.py b/src/art/megatron/canonical/operators.py index caf860bc8..27c853be9 100644 --- a/src/art/megatron/canonical/operators.py +++ b/src/art/megatron/canonical/operators.py @@ -6,6 +6,7 @@ from torch.nn.attention import SDPBackend, sdpa_kernel import torch.nn.functional as F +from ..segmented_lora import slice_lora_segments from .layout import NumericalPolicy, current_layout @@ -39,9 +40,10 @@ def forward(x: torch.Tensor, *args: Any, **kwargs: Any) -> Any: or len(layout.lengths) == 1 ): return original(x, *args, **kwargs) - pieces = [ - original(x[part].contiguous(), *args, **kwargs) for part in layout.slices() - ] + pieces = [] + for part in layout.slices(): + with slice_lora_segments(part.start, part.stop): + pieces.append(original(x[part].contiguous(), *args, **kwargs)) return _join(pieces, layout.padded_lengths[0]) return forward diff --git a/src/art/megatron/kernels/cute_grouped_lora_quack.py b/src/art/megatron/kernels/cute_grouped_lora_quack.py index 804b64e58..36339ee4d 100644 --- a/src/art/megatron/kernels/cute_grouped_lora_quack.py +++ b/src/art/megatron/kernels/cute_grouped_lora_quack.py @@ -494,9 +494,9 @@ def backward(ctx, *grad_outputs: Any): actual_rank = ctx.actual_rank scale = ctx.scale grad_out = cast(torch.Tensor, grad_outputs[0]) - assert grad_out.stride(-1) == 1, ( - "QuACK grouped LoRA backward requires grad_out stride(-1) == 1" - ) + # Scalar reductions and strided views need the GEMM's unit inner stride. + if grad_out.stride(-1) != 1: + grad_out = grad_out.contiguous() grad_tmp = _varlen_quack_gemm( grad_out, diff --git a/src/art/megatron/lora.py b/src/art/megatron/lora.py index ce3349974..09f6ef07d 100644 --- a/src/art/megatron/lora.py +++ b/src/art/megatron/lora.py @@ -46,6 +46,10 @@ default_lora_rank_for_handler, ) from .routed_experts import capture_routes +from .segmented_lora import ( + _CURRENT_LORA_SEGMENTS, + segmented_lora_factors, +) _LAYER_BLOCK_RE = re.compile(r"^(?P.*\.layers\.\d+)\.") @@ -168,15 +172,18 @@ def _collect_compile_garbage() -> None: def _with_captured_lora_slot(function: _F) -> _F: context = _CURRENT_LORA_SLOT.get() + segments = _CURRENT_LORA_SEGMENTS.get() function = capture_routes(capture_canonical(function)) @functools.wraps(function) def wrapped(*args: Any, **kwargs: Any) -> Any: _collect_compile_garbage() token = _CURRENT_LORA_SLOT.set(context) + segment_token = _CURRENT_LORA_SEGMENTS.set(segments) try: result = function(*args, **kwargs) finally: + _CURRENT_LORA_SEGMENTS.reset(segment_token) _CURRENT_LORA_SLOT.reset(token) # A compile inside this call has unwound; collect before its backward. # Failed calls leave the mark for the next call, keeping their error. @@ -1052,6 +1059,8 @@ def sharded_lora_grad_dict(self) -> dict[str, torch.Tensor]: def active_lora_tensors( self, ) -> tuple[torch.Tensor, torch.Tensor, float] | None: + if self.is_expert and _CURRENT_LORA_SEGMENTS.get() is not None: + raise NotImplementedError("segmented routed experts are not supported") active = self._active_lora_storage() return None if active is None else self._expert_views(*active) @@ -1085,6 +1094,22 @@ def _expert_views( def forward( self, x: torch.Tensor, tokens_per_expert: list[int] | torch.Tensor | None = None ) -> torch.Tensor: + segments = _CURRENT_LORA_SEGMENTS.get() + if segments is not None: + if self.is_expert or tokens_per_expert is not None: + raise NotImplementedError("segmented routed experts are not supported") + factors: list[tuple[torch.Tensor, torch.Tensor, float] | None] = [] + for ref in segments.refs: + if ref is None: + factors.append((self.A_T, self.B_T, self.scale)) + else: + slot = self._slot(ref) if ref.name is not None else None + factors.append( + (slot.A_T, slot.B_T, slot.scale) if slot is not None else None + ) + return segmented_lora_factors( + x, factors, segments, template_a=self.A_T, template_b=self.B_T + ) compact_experts = ( tokens_per_expert is not None and self.shared_factor is not None diff --git a/src/art/megatron/segmented_lora.py b/src/art/megatron/segmented_lora.py new file mode 100644 index 000000000..806c24801 --- /dev/null +++ b/src/art/megatron/segmented_lora.py @@ -0,0 +1,202 @@ +"""Adapter-sorted, rank-local dense LoRA, with ordinary slot parameters. + +The planner owns row order. This module never permutes model activations: attention +and GDN plans must describe the same order, including padding and SP gathers. +""" + +from collections.abc import Iterator, Sequence +from contextlib import contextmanager +import contextvars +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +import torch.nn.functional as F + +if TYPE_CHECKING: + from .lora import LoRASlotRef + + +@dataclass(frozen=True) +class LoRASegments: + """One contiguous span per adapter, in flattened (sequence, batch) order. + + ``None`` selects the module's default parameters; a base slot or a missing + module contributes zero. Counts live on CPU and are validated once, before + entering the hot path. A segment with zero rows is allowed. The caller + must pin referenced slot versions until backward/recompute completes. + """ + + refs: tuple["LoRASlotRef | None", ...] + counts: tuple[int, ...] + grouped: bool = False + + def __post_init__(self) -> None: + object.__setattr__(self, "refs", tuple(self.refs)) + object.__setattr__(self, "counts", tuple(self.counts)) + if not self.refs or len(self.refs) != len(self.counts): + raise ValueError("segments require equally sized, nonempty refs/counts") + if any(type(n) is not int or n < 0 for n in self.counts): + raise ValueError("segment counts must be nonnegative Python integers") + if len(set(self.refs)) != len(self.refs): + raise ValueError("an adapter must occupy exactly one sorted segment") + + @property + def rows(self) -> int: + return sum(self.counts) + + def slice(self, start: int, stop: int) -> "LoRASegments": + """Intersect a canonical root slice with the original adapter spans.""" + if not 0 <= start <= stop <= self.rows: + raise ValueError("segment slice must lie within the original rows") + offset = 0 + counts = [] + for count in self.counts: + counts.append(max(0, min(stop, offset + count) - max(start, offset))) + offset += count + return LoRASegments(self.refs, tuple(counts), grouped=self.grouped) + + def adapter_indices(self, device: torch.device | str = "cpu") -> torch.Tensor: + """Materialize the per-token adapter table for a planner/debugger.""" + return torch.repeat_interleave( + torch.arange(len(self.refs), device=device), + torch.tensor(self.counts, device=device), + ) + + @classmethod + def from_adapter_indices( + cls, + refs: Sequence["LoRASlotRef | None"], + indices: torch.Tensor, + ) -> "LoRASegments": + # This is a plan-construction API; a CUDA input intentionally synchronizes. + if indices.ndim != 1 or indices.dtype not in {torch.int32, torch.int64}: + raise ValueError("adapter indices must be a 1D int32/int64 table") + ids = indices.detach().cpu().to(torch.int64) + if ids.numel() and ( + int(ids.min()) < 0 + or int(ids.max()) >= len(refs) + or bool((ids[1:] < ids[:-1]).any()) + ): + raise ValueError("rows must be sorted by an in-range adapter index") + return cls( + tuple(refs), tuple(torch.bincount(ids, minlength=len(refs)).tolist()) + ) + + +_CURRENT_LORA_SEGMENTS: contextvars.ContextVar[LoRASegments | None] = ( + contextvars.ContextVar("art_megatron_current_lora_segments", default=None) +) + + +@contextmanager +def use_lora_segments(segments: LoRASegments | None) -> Iterator[None]: + token = _CURRENT_LORA_SEGMENTS.set(segments) + try: + yield + finally: + _CURRENT_LORA_SEGMENTS.reset(token) + + +@contextmanager +def slice_lora_segments(start: int, stop: int) -> Iterator[None]: + """Keep adapter ownership when an operator executes one root at a time.""" + segments = _CURRENT_LORA_SEGMENTS.get() + with use_lora_segments(None if segments is None else segments.slice(start, stop)): + yield + + +def stack_lora_factors( + factors: Sequence[tuple[torch.Tensor, torch.Tensor, float] | None], + *, + template_a: torch.Tensor, + template_b: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pad and stack without detaching: gradients return to each owning slot. + + Scaling B makes QuACK's scalar scale=1 correct for heterogeneous alphas. + These transient stacks must not be cached across backward/optimizer steps. + """ + rank = max( + (int(pair[0].shape[-1]) for pair in factors if pair is not None), default=1 + ) + rank = max(8, (rank + 7) // 8 * 8) + a_stack, b_stack = [], [] + for pair in factors: + if pair is None: + a_stack.append(template_a.new_zeros(template_a.shape[-2], rank)) + b_stack.append(template_b.new_zeros(rank, template_b.shape[-1])) + else: + a, b, scale = pair + a_stack.append(F.pad(a, (0, rank - a.shape[-1]))) + b_stack.append(F.pad(b * scale, (0, 0, 0, rank - b.shape[-2]))) + return torch.stack(a_stack), torch.stack(b_stack) + + +def segmented_lora_factors( + x: torch.Tensor, + factors: Sequence[tuple[torch.Tensor, torch.Tensor, float] | None], + segments: LoRASegments, + *, + template_a: torch.Tensor, + template_b: torch.Tensor, +) -> torch.Tensor: + """Keep sequential rank/scale geometry unless grouped math is requested. + + Both modes use the original slot Parameters and build a fresh graph. The + enclosing executor must isolate attention/GDN roots and pin slot versions. + """ + if x.numel() // x.shape[-1] != segments.rows or len(factors) != len(segments.refs): + raise ValueError("projection rows/factors disagree with the segment table") + if segments.grouped: + a, b = stack_lora_factors(factors, template_a=template_a, template_b=template_b) + return segmented_lora(x, a, b, segments) + outputs = [] + for chunk, pair in zip( + x.reshape(-1, x.shape[-1]).split(segments.counts), factors, strict=True + ): + if pair is None: + outputs.append(chunk.new_zeros((chunk.shape[0], template_b.shape[-1]))) + else: + a, b, scale = pair + output = (chunk @ a) @ b + outputs.append(output if scale == 1.0 else output * scale) + return torch.cat(outputs).reshape(*x.shape[:-1], template_b.shape[-1]) + + +def segmented_lora( + x: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + segments: LoRASegments, +) -> torch.Tensor: + """Use QuACK for CUDA bf16/fp16 and Torch otherwise; rounding can differ.""" + shape = x.shape + flat = x.reshape(-1, shape[-1]) + if flat.shape[0] != segments.rows: + raise ValueError( + f"segment table has {segments.rows} rows, projection has {flat.shape[0]}; " + "supply metadata for this rank's actual layout (including SP/padding)" + ) + if a.shape[0] != len(segments.refs) or b.shape[0] != len(segments.refs): + raise ValueError("stacked factors must have one group per segment") + if flat.shape[0] == 0: + # Keep zero gradients attached to every participating parameter. + out = flat.new_zeros(0, b.shape[-1]) + (a.sum() + b.sum() + flat.sum()) * 0 + elif ( + x.is_cuda + and x.dtype in {torch.bfloat16, torch.float16} + and torch.cuda.get_device_capability(x.device)[0] == 9 + ): + from .kernels.cute_grouped_lora_quack import quack_grouped_lora + + out = quack_grouped_lora(flat, a, b, list(segments.counts)) + else: + out = torch.cat( + [ + (chunk @ a[i]) @ b[i] + for i, chunk in enumerate(torch.split(flat, list(segments.counts))) + ], + dim=0, + ) + return out.reshape(*shape[:-1], b.shape[-1]) diff --git a/tests/integration/megatron/lora/test_segmented_lora.py b/tests/integration/megatron/lora/test_segmented_lora.py new file mode 100644 index 000000000..232a3997b --- /dev/null +++ b/tests/integration/megatron/lora/test_segmented_lora.py @@ -0,0 +1,289 @@ +"""Slot ownership, row layout, numerical parity, and checkpoint-context oracles.""" + +from dataclasses import replace +import json +import tempfile +from typing import Any, cast + +from megatron.core import parallel_state as ps +import pytest +import torch +import torch.distributed as dist + +from art.megatron.lora import LoRA, LoRASlotRef, use_lora_slot +from art.megatron.segmented_lora import LoRASegments, use_lora_segments + + +@pytest.fixture(scope="module", autouse=True) +def parallel_state(): + if not torch.cuda.is_available(): + pytest.skip("one CUDA device is required") + torch.cuda.set_device(0) + with tempfile.TemporaryDirectory() as init_dir: + dist.init_process_group( + "nccl", init_method=f"file://{init_dir}/store", rank=0, world_size=1 + ) + ps.initialize_model_parallel() + yield + ps.destroy_model_parallel() + dist.destroy_process_group() + + +def make_case(dtype=torch.float32): + torch.manual_seed(74) + lora = LoRA("dense", 64, 96, 8, 8, dtype, torch.device("cuda")) + refs = tuple(LoRASlotRef("checkpoint", f"adapter-{i}") for i in range(4)) + for i, (ref, rank) in enumerate(zip(refs, (1, 3, 13, 64), strict=True)): + lora.load_lora_slot( + ref, + { + "dense.lora_A.weight": torch.randn(rank, 64, device="cuda", dtype=dtype) + * 0.02, + "dense.lora_B.weight": torch.randn(96, rank, device="cuda", dtype=dtype) + * 0.02, + }, + alpha=rank * (0.75 + i * 0.125), + requires_grad=True, + ) + # Include an empty adapter, a base segment and a module missing this target. + refs += (LoRASlotRef("checkpoint", None), LoRASlotRef("lora", "absent")) + segments = LoRASegments(refs, (11, 0, 17, 23, 3, 5)) + return lora, segments + + +def sequential(lora, x, segments): + chunks = x.split(segments.counts) + outputs = [] + with use_lora_segments(None): + for ref, chunk in zip(segments.refs, chunks, strict=True): + with use_lora_slot(ref): + outputs.append(lora(chunk)) + return torch.cat(outputs) + + +@pytest.mark.parametrize( + "checkpoint_kind", ("none", "torch", "torch_reentrant", "megatron") +) +def test_segmented_fp32_outputs_grads_and_recompute(checkpoint_kind): + torch.backends.cuda.matmul.allow_tf32 = False + lora, segments = make_case() + x = torch.randn(segments.rows, 1, 64, device="cuda", requires_grad=True) + dy = torch.randn(segments.rows, 1, 96, device="cuda") + expected = sequential(lora, x, segments) + expected.backward(dy) + dx = cast(torch.Tensor, x.grad).clone() + grads = { + name: p.grad.clone() + for name, p in lora.named_parameters() + if p.grad is not None + } + lora.zero_grad(set_to_none=True) + x.grad = None + with use_lora_segments(segments): + if checkpoint_kind.startswith("torch"): + # Access the patched module attribute, as the real runtime does. + from torch.utils import checkpoint as checkpoint_module + + actual = checkpoint_module.checkpoint( + lora, x, use_reentrant=checkpoint_kind == "torch_reentrant" + ) + elif checkpoint_kind == "megatron": + from megatron.core.tensor_parallel import checkpoint + + actual = checkpoint(lora, False, x) + else: + actual = lora(x) + # Backward must use the captured table even under a conflicting ambient table. + wrong = LoRASegments((segments.refs[0],), (segments.rows,)) + with use_lora_segments(wrong), use_lora_slot(segments.refs[1]): + actual.backward(dy) + torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5) + torch.testing.assert_close(x.grad, dx, atol=1e-6, rtol=1e-5) + for name, p in lora.named_parameters(): + if name in grads: + torch.testing.assert_close(p.grad, grads[name], atol=1e-6, rtol=1e-5) + else: + assert p.grad is None, name + assert lora.A_T.grad is None and lora.B_T.grad is None + print( + "ORACLE " + + json.dumps( + { + "dtype": "float32", + "checkpoint": checkpoint_kind, + "output_max_abs_error": float( + (actual.detach() - expected.detach()).abs().max() + ), + "input_grad_max_abs_error": float( + (cast(torch.Tensor, x.grad) - dx).abs().max() + ), + "param_grad_max_abs_error": max( + float((p.grad - grads[name]).abs().max()) + for name, p in lora.named_parameters() + if name in grads + ), + "atol": 1e-6, + "rtol": 1e-5, + "passed": True, + } + ) + ) + + +def test_segment_table_validation_and_zero_rows(): + lora, segments = make_case() + indices = segments.adapter_indices() + assert LoRASegments.from_adapter_indices(segments.refs, indices) == segments + with pytest.raises(ValueError, match="sorted"): + LoRASegments.from_adapter_indices(segments.refs, indices.flip(0)) + with pytest.raises(ValueError, match="exactly one"): + LoRASegments((segments.refs[0], segments.refs[0]), (1, 1)) + with use_lora_segments(segments), pytest.raises(ValueError, match="row"): + lora(torch.zeros(1, 64, device="cuda")) + lora.zero_grad(set_to_none=True) + empty = LoRASegments(segments.refs, (0,) * len(segments.refs)) + with use_lora_segments(empty): + lora(torch.empty(0, 64, device="cuda")).sum().backward() + for ref in segments.refs[:4]: + for p in lora.lora_slot_params(ref): + assert p.grad is not None and not p.grad.any() + + +def test_segmented_bf16_quack_matches_sequential(): + lora, segments = make_case(torch.bfloat16) + segments = replace(segments, grouped=True) + x = torch.randn( + segments.rows, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + dy = torch.randn(segments.rows, 96, device="cuda", dtype=torch.bfloat16) + expected = sequential(lora, x, segments) + expected.backward(dy) + dx = cast(torch.Tensor, x.grad).clone() + grads = { + name: p.grad.clone() + for name, p in lora.named_parameters() + if p.grad is not None + } + lora.zero_grad(set_to_none=True) + x.grad = None + with use_lora_segments(segments): + actual = lora(x) + actual.backward(dy) + assert x.grad is not None + + # Keep strict FP32 checks above. BF16 is qualified with a per-adapter + # relative-L2 envelope; pointwise ratios around cancellation are unstable. + def rms_error(a, b): + return float( + (a.detach().float() - b.detach().float()).norm() + / b.detach().float().norm().clamp_min(1e-30) + ) + + def assert_bf16_close(a, b): + assert torch.isfinite(a).all() + assert rms_error(a, b) < 0.01 + + for a, b in zip( + actual.split(segments.counts), expected.split(segments.counts), strict=True + ): + assert_bf16_close(a, b) + for a, b in zip( + x.grad.split(segments.counts), dx.split(segments.counts), strict=True + ): + assert_bf16_close(a, b) + for name, p in lora.named_parameters(): + if name in grads: + assert p.grad is not None + assert_bf16_close(p.grad, grads[name]) + else: + assert p.grad is None + + print( + "ORACLE " + + json.dumps( + { + "dtype": "bfloat16", + "checkpoint": "none", + "output_max_abs_error": float( + (actual.float() - expected.float()).abs().max() + ), + "input_grad_max_abs_error": float( + (x.grad.float() - dx.float()).abs().max() + ), + "param_grad_max_abs_error": max( + float((p.grad.float() - grads[name].float()).abs().max()) + for name, p in lora.named_parameters() + if name in grads + ), + "passed_bf16_gate": True, + "bf16_relative_l2_limit": 0.01, + "output_relative_l2": rms_error(actual, expected), + "input_grad_relative_l2": rms_error(x.grad, dx), + "param_grad_max_relative_l2": max( + rms_error(p.grad, grads[name]) + for name, p in lora.named_parameters() + if name in grads + ), + "passed_fp32_tolerance": bool( + torch.allclose( + actual.float(), expected.float(), atol=1e-6, rtol=1e-5 + ) + ), + } + ) + ) + + # A scalar reduction supplies stride-zero gradients to this new dense caller. + lora.zero_grad(set_to_none=True) + x.grad = None + with use_lora_segments(segments): + output = lora(x) + output.sum().backward() + assert x.grad is not None and torch.isfinite(x.grad).all() + for ref in segments.refs[:4]: + for p in lora.lora_slot_params(ref): + assert p.grad is not None and torch.isfinite(p.grad).all() + + +def test_routed_experts_reject_adapter_segments(): + lora = LoRA("experts.{expert}", 64, 96, 8, 8, torch.float32, torch.device("cuda")) + segments = LoRASegments((None,), (3,)) + with use_lora_segments(segments): + with pytest.raises(NotImplementedError, match="expert, adapter"): + lora(torch.zeros(3, 64, device="cuda"), [3]) + with pytest.raises(NotImplementedError, match="expert, adapter"): + lora.active_lora_tensors() + + +def test_empty_shared_fc1_attaches_owning_adapter_gradients(): + from art.megatron.lora import SharedExpertsLinearFC1LoRA + + # The empty path must not invoke the base projection at all. + base = torch.nn.Linear(64, 192, bias=False, device="cuda") + wrapped = SharedExpertsLinearFC1LoRA("shared", cast(Any, base), rank=8, alpha=8) + refs = (LoRASlotRef("checkpoint", "first"), LoRASlotRef("checkpoint", "second")) + for lora in (wrapped.gate_lora, wrapped.up_lora): + for ref in refs: + lora.load_lora_slot( + ref, + { + f"{lora.adapter_model_prefix}.lora_A.weight": torch.randn( + 3, 64, device="cuda" + ), + f"{lora.adapter_model_prefix}.lora_B.weight": torch.randn( + 96, 3, device="cuda" + ), + }, + requires_grad=True, + ) + with use_lora_segments(LoRASegments(refs, (0, 0))): + y, bias = wrapped(torch.empty(0, 1, 64, device="cuda")) + y.sum().backward() + assert y.shape == (0, 1, 192) and bias is None + for lora in (wrapped.gate_lora, wrapped.up_lora): + assert lora.A_T.grad is None and lora.B_T.grad is None + for ref in refs: + assert all( + p.grad is not None and not p.grad.any() + for p in lora.lora_slot_params(ref) + ) diff --git a/tests/unit/test_canonical_numerics.py b/tests/unit/test_canonical_numerics.py index 008ffcb5c..75900edff 100644 --- a/tests/unit/test_canonical_numerics.py +++ b/tests/unit/test_canonical_numerics.py @@ -1,5 +1,5 @@ from types import SimpleNamespace -from typing import Any +from typing import Any, cast import pytest import torch @@ -144,7 +144,9 @@ def test_attention_rejects_partial_root_before_launching_kernel(): def test_canonical_gdn_never_selects_uniform_bucket_algorithm(): from art.megatron.gdn.operator import _bucket_has_uniform_lengths - bucket: Any = SimpleNamespace(length=128, lengths_cpu=torch.tensor([128, 128])) + bucket = cast( + Any, SimpleNamespace(length=128, lengths_cpu=torch.tensor([128, 128])) + ) assert _bucket_has_uniform_lengths(bucket) with execution(RootLayout((128, 128))): assert not _bucket_has_uniform_lengths(bucket) @@ -282,3 +284,42 @@ class Projection(te.RMSNorm): wrapper = projection.forward assert install(model) == {"0": "row"} assert projection.forward is wrapper + + +def test_segmented_slots_follow_canonical_root_slices_through_backward(): + from art.megatron.segmented_lora import ( + _CURRENT_LORA_SEGMENTS, + LoRASegments, + segmented_lora_factors, + use_lora_segments, + ) + + factors = [ + ( + torch.randn(3, 2, requires_grad=True), + torch.randn(2, 4, requires_grad=True), + 1.0, + ) + for _ in range(2) + ] + segments = LoRASegments((None, cast(Any, "second")), (128, 256)) + seen = [] + + def original(x): + current = _CURRENT_LORA_SEGMENTS.get() + assert current is not None + seen.append(current.counts) + return segmented_lora_factors( + x, factors, current, template_a=factors[0][0], template_b=factors[0][1] + ) + + x = torch.randn(384, 3) + with use_lora_segments(segments), execution(RootLayout((3, 129))): + actual = _row_forward(original)(x) + actual[:3].sum().backward() + assert seen == [(128, 0), (0, 256)] + assert torch.equal(actual[:128], x[:128] @ factors[0][0] @ factors[0][1]) + assert torch.equal(actual[128:], x[128:] @ factors[1][0] @ factors[1][1]) + assert not cast(torch.Tensor, factors[1][0].grad).any() + assert not cast(torch.Tensor, factors[1][1].grad).any() + assert _CURRENT_LORA_SEGMENTS.get() is None diff --git a/tests/unit/test_segmented_lora.py b/tests/unit/test_segmented_lora.py new file mode 100644 index 000000000..013edd7b7 --- /dev/null +++ b/tests/unit/test_segmented_lora.py @@ -0,0 +1,239 @@ +"""Offline checks of rank/scale geometry and independent gradient ownership.""" + +from dataclasses import replace +from typing import Any, cast + +import pytest +import torch + +from art.megatron.segmented_lora import ( + _CURRENT_LORA_SEGMENTS, + LoRASegments, + segmented_lora_factors, + use_lora_segments, +) + + +@pytest.mark.parametrize("grouped", [False, True]) +@pytest.mark.parametrize("adapters", [1, 4, 16, 64]) +def test_outputs_losses_and_original_parameter_gradients(adapters, grouped): + torch.manual_seed(913) + counts = tuple(0 if i % 7 == 1 else 3 + i % 5 for i in range(adapters)) + (5,) + refs = tuple(cast(Any, f"run-{i}") for i in range(adapters)) + (None,) + segments = LoRASegments(refs, counts, grouped=grouped) + x = torch.randn(sum(counts), 11, requires_grad=True) + factors = [ + ( + torch.nn.Parameter(torch.randn(11, i + 1) * 0.02), + torch.nn.Parameter(torch.randn(i + 1, 13) * 0.02), + 1.375 / (i + 1), + ) + for i in range(adapters) + ] + chunks = x.split(counts) + expected = torch.cat( + [(chunks[i] @ a) @ b * scale for i, (a, b, scale) in enumerate(factors)] + + [chunks[-1].new_zeros(counts[-1], 13)] + ) + + def losses(output): + # Unequal run lengths, separate request normalization. Empty requests + # still produce connected zero gradients for their installed slots. + return torch.stack( + [ + chunk.square().sum() / max(n, 1) + for chunk, n in zip(output.split(counts), counts, strict=True) + ] + ) + + losses(expected).sum().backward() + assert x.grad is not None + dx = x.grad.clone() + grads = [ + (cast(torch.Tensor, a.grad).clone(), cast(torch.Tensor, b.grad).clone()) + for a, b, _ in factors + ] + x.grad = None + for a, b, _ in factors: + a.grad = b.grad = None + actual = segmented_lora_factors( + x, + factors + [None], + segments, + template_a=factors[0][0], + template_b=factors[0][1], + ) + losses(actual).sum().backward() + atol, rtol = (1e-6, 1e-5) if grouped else (0, 0) + torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol) + torch.testing.assert_close(losses(actual), losses(expected), atol=atol, rtol=rtol) + torch.testing.assert_close(x.grad, dx, atol=atol, rtol=rtol) + for (a, b, _), (ga, gb) in zip(factors, grads, strict=True): + torch.testing.assert_close(a.grad, ga, atol=atol, rtol=rtol) + torch.testing.assert_close(b.grad, gb, atol=atol, rtol=rtol) + + +@pytest.mark.parametrize("grouped", [False, True]) +def test_masked_loss_has_no_other_adapter_gradients(grouped): + a = [torch.randn(7, r, requires_grad=True) for r in (1, 8)] + b = [torch.randn(r, 5, requires_grad=True) for r in (1, 8)] + segments = LoRASegments(cast(Any, ("a", "b")), (3, 9), grouped) + output = segmented_lora_factors( + torch.randn(12, 7), + list(zip(a, b, (1.0, 0.5), strict=True)), + segments, + template_a=a[0], + template_b=b[0], + ) + output[:3].square().sum().backward() + for p in (a[0], b[0]): + assert p.grad is not None and p.grad.any() + for p in (a[1], b[1]): + assert p.grad is not None and not p.grad.any() + + +def test_layout_validation_empty_gradients_and_context_restoration(): + refs = cast(Any, ("first", "second")) + table = LoRASegments(refs, (2, 5)) + assert LoRASegments.from_adapter_indices(refs, table.adapter_indices()) == table + with pytest.raises(ValueError, match="sorted"): + LoRASegments.from_adapter_indices(refs, table.adapter_indices().flip(0)) + for counts in ((True, 1), (-1, 2)): + with pytest.raises(ValueError, match="nonnegative"): + LoRASegments(refs, counts) + with pytest.raises(ValueError, match="one sorted"): + LoRASegments(cast(Any, ("a", "a")), (1, 1)) + a, b = torch.randn(7, 3, requires_grad=True), torch.randn(3, 5, requires_grad=True) + for grouped in (False, True): + a.grad = b.grad = None + empty = replace(table, counts=(0, 0), grouped=grouped) + result = segmented_lora_factors( + torch.empty(0, 7), + [(a, b, 1.0), None], + empty, + template_a=a, + template_b=b, + ) + result.sum().backward() + assert a.grad is not None and not a.grad.any() + assert b.grad is not None and not b.grad.any() + with pytest.raises(ValueError, match="rows"): + segmented_lora_factors( + torch.empty(1, 7), [(a, b, 1.0), None], table, template_a=a, template_b=b + ) + with pytest.raises(RuntimeError), use_lora_segments(table): + assert _CURRENT_LORA_SEGMENTS.get() is table + raise RuntimeError("forward failed") + assert _CURRENT_LORA_SEGMENTS.get() is None + + +def test_canonical_root_slices_keep_original_adapter_ownership(): + from art.megatron.segmented_lora import slice_lora_segments + + table = LoRASegments(cast(Any, ("a", "b")), (256, 128)) + with use_lora_segments(table): + with slice_lora_segments(128, 256): + assert _CURRENT_LORA_SEGMENTS.get() == LoRASegments(table.refs, (128, 0)) + with slice_lora_segments(256, 384): + assert _CURRENT_LORA_SEGMENTS.get() == LoRASegments(table.refs, (0, 128)) + assert _CURRENT_LORA_SEGMENTS.get() is table + with pytest.raises(ValueError, match="within"): + table.slice(0, 385) + + +def test_segment_metadata_owns_constructor_lists(): + counts = [2, 5] + refs = ["a", "b"] + table = LoRASegments(cast(Any, refs), cast(Any, counts)) + counts[0] = 100 + refs.reverse() + assert table.counts == (2, 5) and table.refs == ("a", "b") + + +@pytest.mark.parametrize("reduction", ["sum", "strided"]) +def test_grouped_backward_accepts_noncontiguous_cotangents(monkeypatch, reduction): + pytest.importorskip("quack") + from art.megatron.kernels import cute_grouped_lora_quack as kernel + + # Exercise the actual custom autograd boundary on CPU. Only the CUDA GEMMs + # are replaced, retaining their last-dimension stride contract. + def gemm(values, weights, *, expert_offsets, alpha=1.0, **kwargs): + assert values.stride(-1) == 1 + counts = expert_offsets.diff().tolist() + return ( + torch.cat( + [chunk @ weights[i].T for i, chunk in enumerate(values.split(counts))] + ) + * alpha + ) + + def gemm_k(left, right, *, expert_offsets, alpha=1.0, **kwargs): + return ( + torch.stack( + [ + left[:, start:stop] @ right[:, start:stop].T + for start, stop in zip( + expert_offsets[:-1].tolist(), + expert_offsets[1:].tolist(), + strict=True, + ) + ] + ) + * alpha + ) + + monkeypatch.setattr(kernel, "_varlen_quack_gemm", gemm) + monkeypatch.setattr(kernel, "_varlen_quack_gemm_k", gemm_k) + torch.manual_seed(74) + counts = torch.tensor([2, 0, 3]) + x = torch.randn(5, 4, requires_grad=True) + a = torch.randn(3, 4, 3, requires_grad=True) + b = torch.randn(3, 3, 6, requires_grad=True) + expected = ( + torch.cat( + [(chunk @ a[i]) @ b[i] for i, chunk in enumerate(x.split(counts.tolist()))] + ) + * 1.375 + ) + cotangent = ( + torch.ones(1).expand_as(expected) if reduction == "sum" else torch.randn(6, 5).T + ) + assert cotangent.stride(-1) != 1 + expected.backward(cotangent) + grads = [cast(torch.Tensor, p.grad).clone() for p in (x, a, b)] + for p in (x, a, b): + p.grad = None + actual = kernel._QuackGroupedLoraFn.apply(x, a, b, counts, 1.375, None) + actual.backward(cotangent) + torch.testing.assert_close(actual, expected, atol=2e-5, rtol=1e-5) + for p, grad in zip((x, a, b), grads, strict=True): + torch.testing.assert_close(p.grad, grad, atol=2e-5, rtol=1e-5) + + +def test_expert_views_reject_missing_dispatch_ownership(monkeypatch): + pytest.importorskip("megatron.core") + from art.megatron import lora as lm + from art.megatron.lora import LoRA + + for name in ( + "get_expert_model_parallel_rank", + "get_expert_data_parallel_rank", + "get_data_parallel_rank", + ): + monkeypatch.setattr(lm.ps, name, lambda **_: 0) + monkeypatch.setattr(lm.ps, "get_tensor_model_parallel_world_size", lambda: 1) + lora = LoRA( + "experts.{expert}", + 3, + 4, + 2, + 3.5, + torch.float32, + torch.device("cpu"), + num_local_experts=3, + ) + with use_lora_segments(LoRASegments((None,), (3,))): + with pytest.raises(NotImplementedError, match="not supported"): + lora.active_lora_tensors() + with pytest.raises(NotImplementedError, match="not supported"): + lora(torch.randn(3, 3)) From 643933e88159c5035fccf02448b2564001f22e87 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 20:10:50 +0000 Subject: [PATCH 02/12] fix: preserve default LoRA behavior outside segmented execution Capture grouped-backward opt-in at forward, preserve main stride failures by default, and route empty segmented FC1 dependencies through the owning slots. Align expert rejection assertions with the production guard and cover gated/non-gated and conflicting-context boundaries on CPU. Co-Authored-By: Claude Opus 5.5 --- .../kernels/cute_grouped_lora_quack.py | 12 +- src/art/megatron/lora.py | 7 +- .../megatron/lora/test_segmented_lora.py | 4 +- tests/unit/test_segmented_lora.py | 105 ++++++++++++++++-- 4 files changed, 111 insertions(+), 17 deletions(-) diff --git a/src/art/megatron/kernels/cute_grouped_lora_quack.py b/src/art/megatron/kernels/cute_grouped_lora_quack.py index 36339ee4d..b0f513548 100644 --- a/src/art/megatron/kernels/cute_grouped_lora_quack.py +++ b/src/art/megatron/kernels/cute_grouped_lora_quack.py @@ -14,6 +14,8 @@ import triton import triton.language as tl +from ..segmented_lora import _CURRENT_LORA_SEGMENTS + _PADDED_LOW_RANK_TARGET = 8 @@ -481,6 +483,8 @@ def forward( ctx.effective_rank = effective_rank ctx.scale = scale ctx.has_residual = has_residual + # Capture the opt-in now: backward may run outside this context. + ctx.segmented = _CURRENT_LORA_SEGMENTS.get() is not None return out @staticmethod @@ -494,9 +498,13 @@ def backward(ctx, *grad_outputs: Any): actual_rank = ctx.actual_rank scale = ctx.scale grad_out = cast(torch.Tensor, grad_outputs[0]) - # Scalar reductions and strided views need the GEMM's unit inner stride. - if grad_out.stride(-1) != 1: + # Only segmented callers opt into cotangent normalization. Ordinary + # grouped/expert execution retains main's stride assertion. + if ctx.segmented and grad_out.stride(-1) != 1: grad_out = grad_out.contiguous() + assert grad_out.stride(-1) == 1, ( + "QuACK grouped LoRA backward requires grad_out stride(-1) == 1" + ) grad_tmp = _varlen_quack_gemm( grad_out, diff --git a/src/art/megatron/lora.py b/src/art/megatron/lora.py index 09f6ef07d..413079960 100644 --- a/src/art/megatron/lora.py +++ b/src/art/megatron/lora.py @@ -2168,8 +2168,11 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: (self.up_lora,) if self.non_gated else (self.gate_lora, self.up_lora) ) for lora in loras: - zero = zero + lora.A_T.to(dtype=x.dtype).sum() * 0.0 - zero = zero + lora.B_T.to(dtype=x.dtype).sum() * 0.0 + if _CURRENT_LORA_SEGMENTS.get() is not None: + zero = zero + lora(x).sum() * 0.0 + else: + zero = zero + lora.A_T.to(dtype=x.dtype).sum() * 0.0 + zero = zero + lora.B_T.to(dtype=x.dtype).sum() * 0.0 return zero.expand(*x.shape[:-1], self.out_features).clone(), None base_output, bias_out = self.linear_fc1(x) if isinstance(base_output, tuple): diff --git a/tests/integration/megatron/lora/test_segmented_lora.py b/tests/integration/megatron/lora/test_segmented_lora.py index 232a3997b..ed1e778dd 100644 --- a/tests/integration/megatron/lora/test_segmented_lora.py +++ b/tests/integration/megatron/lora/test_segmented_lora.py @@ -249,9 +249,9 @@ def test_routed_experts_reject_adapter_segments(): lora = LoRA("experts.{expert}", 64, 96, 8, 8, torch.float32, torch.device("cuda")) segments = LoRASegments((None,), (3,)) with use_lora_segments(segments): - with pytest.raises(NotImplementedError, match="expert, adapter"): + with pytest.raises(NotImplementedError, match="not supported"): lora(torch.zeros(3, 64, device="cuda"), [3]) - with pytest.raises(NotImplementedError, match="expert, adapter"): + with pytest.raises(NotImplementedError, match="not supported"): lora.active_lora_tensors() diff --git a/tests/unit/test_segmented_lora.py b/tests/unit/test_segmented_lora.py index 013edd7b7..8b77c5c48 100644 --- a/tests/unit/test_segmented_lora.py +++ b/tests/unit/test_segmented_lora.py @@ -150,22 +150,23 @@ def test_segment_metadata_owns_constructor_lists(): assert table.counts == (2, 5) and table.refs == ("a", "b") -@pytest.mark.parametrize("reduction", ["sum", "strided"]) -def test_grouped_backward_accepts_noncontiguous_cotangents(monkeypatch, reduction): +@pytest.fixture +def grouped_kernel(monkeypatch): pytest.importorskip("quack") from art.megatron.kernels import cute_grouped_lora_quack as kernel # Exercise the actual custom autograd boundary on CPU. Only the CUDA GEMMs # are replaced, retaining their last-dimension stride contract. - def gemm(values, weights, *, expert_offsets, alpha=1.0, **kwargs): + def gemm(values, weights, *, expert_offsets, alpha=1.0, residual=None, **kwargs): assert values.stride(-1) == 1 counts = expert_offsets.diff().tolist() - return ( + output = ( torch.cat( [chunk @ weights[i].T for i, chunk in enumerate(values.split(counts))] ) * alpha ) + return output if residual is None else output + residual def gemm_k(left, right, *, expert_offsets, alpha=1.0, **kwargs): return ( @@ -184,32 +185,114 @@ def gemm_k(left, right, *, expert_offsets, alpha=1.0, **kwargs): monkeypatch.setattr(kernel, "_varlen_quack_gemm", gemm) monkeypatch.setattr(kernel, "_varlen_quack_gemm_k", gemm_k) + return kernel + + +@pytest.mark.parametrize("segmented", [False, True]) +@pytest.mark.parametrize("residual", [False, True]) +@pytest.mark.parametrize("reduction", ["sum", "strided", "contiguous"]) +def test_grouped_backward_preserves_main_contract( + monkeypatch, grouped_kernel, reduction, residual, segmented +): + monkeypatch.delenv("STARK_FUSED_WAVES", raising=False) torch.manual_seed(74) counts = torch.tensor([2, 0, 3]) x = torch.randn(5, 4, requires_grad=True) a = torch.randn(3, 4, 3, requires_grad=True) b = torch.randn(3, 3, 6, requires_grad=True) + base = torch.randn(5, 6, requires_grad=True) if residual else None expected = ( torch.cat( [(chunk @ a[i]) @ b[i] for i, chunk in enumerate(x.split(counts.tolist()))] ) * 1.375 ) + if base is not None: + expected = expected + base cotangent = ( - torch.ones(1).expand_as(expected) if reduction == "sum" else torch.randn(6, 5).T + torch.ones(1).expand_as(expected) + if reduction == "sum" + else torch.randn(6, 5).T + if reduction == "strided" + else torch.randn_like(expected) ) - assert cotangent.stride(-1) != 1 + assert (cotangent.stride(-1) == 1) == (reduction == "contiguous") + params = (x, a, b) if base is None else (x, a, b, base) expected.backward(cotangent) - grads = [cast(torch.Tensor, p.grad).clone() for p in (x, a, b)] - for p in (x, a, b): + grads = [cast(torch.Tensor, p.grad).clone() for p in params] + for p in params: p.grad = None - actual = kernel._QuackGroupedLoraFn.apply(x, a, b, counts, 1.375, None) - actual.backward(cotangent) + table = LoRASegments(cast(Any, ("a", "b", "c")), (2, 0, 3), grouped=True) + with use_lora_segments(table if segmented else None): + actual = grouped_kernel._QuackGroupedLoraFn.apply(x, a, b, counts, 1.375, base) + # Backward must use the forward's opt-in, even under conflicting context. + with use_lora_segments(None if segmented else table): + if not segmented and reduction != "contiguous": + # ART main rejects both scalar and strided cotangents by default. + with pytest.raises( + AssertionError, + match=r"^QuACK grouped LoRA backward requires grad_out stride\(-1\) == 1$", + ): + actual.backward(cotangent) + assert all(p.grad is None for p in params) + return + actual.backward(cotangent) torch.testing.assert_close(actual, expected, atol=2e-5, rtol=1e-5) - for p, grad in zip((x, a, b), grads, strict=True): + for p, grad in zip(params, grads, strict=True): torch.testing.assert_close(p.grad, grad, atol=2e-5, rtol=1e-5) +@pytest.mark.parametrize("non_gated", [False, True]) +@pytest.mark.parametrize("grouped", [False, True]) +@pytest.mark.parametrize("segmented", [False, True]) +def test_empty_shared_fc1_preserves_gradient_ownership( + monkeypatch, non_gated, grouped, segmented +): + pytest.importorskip("megatron.core") + from art.megatron import lora as lm + + monkeypatch.setattr(lm.ps, "get_expert_model_parallel_rank", lambda: 0) + monkeypatch.setattr(lm.ps, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(lm.ps, "get_tensor_model_parallel_rank", lambda: 0) + base = torch.nn.Linear(64, 192, bias=False) + monkeypatch.setattr( + base, "forward", lambda *_: pytest.fail("empty base projection") + ) + wrapped = lm.SharedExpertsLinearFC1LoRA( + "shared", cast(Any, base), rank=8, alpha=8, non_gated=non_gated + ) + refs = tuple(lm.LoRASlotRef("checkpoint", name) for name in ("first", "second")) + loras = (wrapped.up_lora,) if non_gated else (wrapped.gate_lora, wrapped.up_lora) + for lora in loras: + for ref in refs: + lora.load_lora_slot( + ref, + { + f"{lora.adapter_model_prefix}.lora_A.weight": torch.randn(3, 64), + f"{lora.adapter_model_prefix}.lora_B.weight": torch.randn( + lora.out_features, 3 + ), + }, + requires_grad=True, + ) + x = torch.empty(0, 1, 64, requires_grad=True) + table = LoRASegments(refs, (0, 0), grouped=grouped) + # Main's ordinary empty path still attaches the default factors, even with + # a named single-slot context. Only the segmented API changes ownership. + with lm.use_lora_slot(refs[0]), use_lora_segments(table if segmented else None): + output, bias = wrapped(x) + output.sum().backward() + assert output.shape == (0, 1, 192) and bias is None + assert x.grad is not None and x.grad.shape == x.shape + assert base.weight.grad is not None and not base.weight.grad.any() + for lora in loras: + defaults = (lora.A_T, lora.B_T) + selected = tuple(p for ref in refs for p in lora.lora_slot_params(ref)) + owned, untouched = (selected, defaults) if segmented else (defaults, selected) + assert all(p.grad is not None and not p.grad.any() for p in owned) + assert all(p.grad is None for p in untouched) + + def test_expert_views_reject_missing_dispatch_ownership(monkeypatch): pytest.importorskip("megatron.core") from art.megatron import lora as lm From bd6d388e01ed0233781743ab1335047c9b8fd7fc Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 21:03:53 +0000 Subject: [PATCH 03/12] fix: preserve native adapter geometry in opt-in MoE waves Retain original slot ownership through native EP1 dispatch and recomputation. Keep TE row geometry and router aux/z/capacity computations per adapter; reuse compact expert kernels and leave ordinary forwards unchanged. Co-Authored-By: Claude Opus 5.5 --- src/art/megatron/lora.py | 59 ++- src/art/megatron/segmented_expert_lora.py | 377 ++++++++++++++++++ .../lora/test_segmented_expert_lora.py | 255 ++++++++++++ tests/unit/test_segmented_expert_lora.py | 241 +++++++++++ tests/unit/test_segmented_router.py | 108 +++++ 5 files changed, 1033 insertions(+), 7 deletions(-) create mode 100644 src/art/megatron/segmented_expert_lora.py create mode 100644 tests/integration/megatron/lora/test_segmented_expert_lora.py create mode 100644 tests/unit/test_segmented_expert_lora.py create mode 100644 tests/unit/test_segmented_router.py diff --git a/src/art/megatron/lora.py b/src/art/megatron/lora.py index 413079960..dbe454664 100644 --- a/src/art/megatron/lora.py +++ b/src/art/megatron/lora.py @@ -46,6 +46,11 @@ default_lora_rank_for_handler, ) from .routed_experts import capture_routes +from .segmented_expert_lora import ( + _CURRENT_EXPERT_ADAPTER_PLAN, + capture_expert_context, + segmented_expert_lora, +) from .segmented_lora import ( _CURRENT_LORA_SEGMENTS, segmented_lora_factors, @@ -173,16 +178,19 @@ def _collect_compile_garbage() -> None: def _with_captured_lora_slot(function: _F) -> _F: context = _CURRENT_LORA_SLOT.get() segments = _CURRENT_LORA_SEGMENTS.get() - function = capture_routes(capture_canonical(function)) + expert_plan = _CURRENT_EXPERT_ADAPTER_PLAN.get() + function = capture_expert_context(capture_routes(capture_canonical(function))) @functools.wraps(function) def wrapped(*args: Any, **kwargs: Any) -> Any: _collect_compile_garbage() token = _CURRENT_LORA_SLOT.set(context) segment_token = _CURRENT_LORA_SEGMENTS.set(segments) + expert_token = _CURRENT_EXPERT_ADAPTER_PLAN.set(expert_plan) try: result = function(*args, **kwargs) finally: + _CURRENT_EXPERT_ADAPTER_PLAN.reset(expert_token) _CURRENT_LORA_SEGMENTS.reset(segment_token) _CURRENT_LORA_SLOT.reset(token) # A compile inside this call has unwound; collect before its backward. @@ -1095,11 +1103,23 @@ def forward( self, x: torch.Tensor, tokens_per_expert: list[int] | torch.Tensor | None = None ) -> torch.Tensor: segments = _CURRENT_LORA_SEGMENTS.get() - if segments is not None: - if self.is_expert or tokens_per_expert is not None: - raise NotImplementedError("segmented routed experts are not supported") + expert_plan = _CURRENT_EXPERT_ADAPTER_PLAN.get() if self.is_expert else None + if ( + expert_plan is not None + and segments is not None + and expert_plan.refs != segments.refs + ): + raise ValueError("expert plan adapter refs disagree with dense segments") + refs = ( + expert_plan.refs + if expert_plan is not None + else segments.refs + if segments is not None + else None + ) + if refs is not None: factors: list[tuple[torch.Tensor, torch.Tensor, float] | None] = [] - for ref in segments.refs: + for ref in refs: if ref is None: factors.append((self.A_T, self.B_T, self.scale)) else: @@ -1107,6 +1127,19 @@ def forward( factors.append( (slot.A_T, slot.B_T, slot.scale) if slot is not None else None ) + if expert_plan is not None: + return segmented_expert_lora( + x, + factors, + expert_plan, + out_features=self.out_features, + shared_outer_kernel=self.lora_layout.shared_outer_kernel, + ) + if segments is not None: + if self.is_expert or tokens_per_expert is not None: + raise NotImplementedError( + "segmented routed experts are not supported without an (expert, adapter) row layout" + ) return segmented_lora_factors( x, factors, segments, template_a=self.A_T, template_b=self.B_T ) @@ -1366,8 +1399,12 @@ def _expert_grouped_lora_forward( tokens_per_expert: list[int] | torch.Tensor, out_features: int, ) -> torch.Tensor: - if x.shape[0] == 0 and not ( - lora.shared_factor is not None and lora.lora_layout.shared_outer_kernel + if ( + x.shape[0] == 0 + and _CURRENT_EXPERT_ADAPTER_PLAN.get() is None + and not ( + lora.shared_factor is not None and lora.lora_layout.shared_outer_kernel + ) ): return x.new_zeros((x.shape[0], out_features)) return lora(x, tokens_per_expert=tokens_per_expert) @@ -1386,6 +1423,14 @@ def _expert_grouped_lora_dual_forward( x: torch.Tensor, tokens_per_expert: list[int] | torch.Tensor, ) -> torch.Tensor: + if _CURRENT_EXPERT_ADAPTER_PLAN.get() is not None: + return torch.cat( + [ + module.gate_lora(x, tokens_per_expert), + module.up_lora(x, tokens_per_expert), + ], + dim=-1, + ) counts = tokens_per_expert if isinstance(counts, list): counts = torch.tensor(counts, dtype=torch.int64, device="cpu") diff --git a/src/art/megatron/segmented_expert_lora.py b/src/art/megatron/segmented_expert_lora.py new file mode 100644 index 000000000..404bccc1b --- /dev/null +++ b/src/art/megatron/segmented_expert_lora.py @@ -0,0 +1,377 @@ +"""Wave-local adapter ownership for expert-major dispatched rows. + +The plan follows the dispatcher's actual row order, including repeated top-k +rows and padding. Shared-outer factors stay compact; Parameters are never +replaced or expanded across experts. Transport must carry adapter IDs with rows. +""" + +from collections.abc import Iterator, Sequence +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from functools import partial +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from .lora import LoRASlotRef + + +@dataclass(frozen=True) +class ExpertAdapterPlan: + refs: tuple["LoRASlotRef | None", ...] + permutation: torch.Tensor + inverse: torch.Tensor + counts: tuple[tuple[int, ...], ...] # adapter-major, then expert + + @classmethod + def build( + cls, + refs: Sequence["LoRASlotRef | None"], + adapter_ids: torch.Tensor, + expert_counts: Sequence[int] | torch.Tensor, + ) -> "ExpertAdapterPlan": + """Build once after dispatch; CPU metadata validation synchronizes CUDA.""" + if not refs or len(set(refs)) != len(refs): + raise ValueError("expert plan requires unique adapter refs") + if adapter_ids.ndim != 1 or adapter_ids.dtype not in (torch.int32, torch.int64): + raise ValueError("adapter IDs must be a 1D int32/int64 table") + raw = torch.as_tensor(expert_counts) + if ( + raw.ndim != 1 + or raw.dtype not in (torch.int32, torch.int64) + or not raw.numel() + ): + raise ValueError("expert counts must be a nonempty integer vector") + counts = raw.detach().cpu().tolist() + ids = adapter_ids.detach().cpu().tolist() + if any(n < 0 for n in counts) or sum(counts) != len(ids): + raise ValueError("expert counts must cover every dispatched row") + if any(i < 0 or i >= len(refs) for i in ids): + raise ValueError("adapter IDs must select a wave adapter") + experts = torch.repeat_interleave( + torch.arange(len(counts), device=adapter_ids.device), + torch.tensor(counts, device=adapter_ids.device), + output_size=adapter_ids.numel(), + ) + groups = adapter_ids.to(torch.int64) * len(counts) + experts + permutation = groups.argsort(stable=True) + inverse = permutation.argsort() + histogram = torch.bincount(groups.cpu(), minlength=len(refs) * len(counts)) + return cls( + tuple(refs), + permutation, + inverse, + tuple(tuple(row) for row in histogram.reshape(len(refs), -1).tolist()), + ) + + @property + def rows(self) -> int: + return sum(map(sum, self.counts)) + + +_CURRENT_EXPERT_ADAPTER_PLAN: ContextVar[ExpertAdapterPlan | None] = ContextVar( + "art_megatron_expert_adapter_plan", default=None +) + + +@contextmanager +def use_expert_adapter_plan(plan: ExpertAdapterPlan | None) -> Iterator[None]: + token = _CURRENT_EXPERT_ADAPTER_PLAN.set(plan) + try: + yield + finally: + _CURRENT_EXPERT_ADAPTER_PLAN.reset(token) + + +def expert_lora( + x: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + counts: Sequence[int], + *, + scale: float = 1.0, + shared_outer_kernel: bool = True, +) -> torch.Tensor: + """Per-adapter math for per-expert or compact shared-outer A/B factors. + + H200 reuses the ordinary native kernels with the adapter's own row geometry. + Other devices and FP32 use the Torch reference below. + + Shared A projects all of an adapter's rows together before expert B. Shared + B projects all intermediate rows together after expert A. This avoids both + expert-sized shared weights and expert-sized shared gradient temporaries. + """ + experts = len(counts) + if ( + x.ndim != 2 + or a.ndim != 3 + or b.ndim != 3 + or a.shape[0] not in (1, experts) + or b.shape[0] not in (1, experts) + or a.shape[1] != x.shape[1] + or a.shape[2] != b.shape[1] + or sum(counts) != x.shape[0] + or any(n < 0 for n in counts) + ): + raise ValueError("expert rows/counts/factor shapes disagree") + if not shared_outer_kernel: + a = a.expand(experts, -1, -1) + b = b.expand(experts, -1, -1) + if ( + x.shape[0] > 0 + and x.is_cuda + and x.dtype in {torch.bfloat16, torch.float16} + and torch.cuda.get_device_capability(x.device)[0] == 9 + ): + if a.shape[0] == b.shape[0] == experts: + from .kernels.cute_grouped_lora_quack import quack_grouped_lora + + return quack_grouped_lora(x, a, b, list(counts), scale=scale) + if (a.shape[0], b.shape[0]) in {(1, experts), (experts, 1)}: + from .kernels.shared_outer_lora import shared_outer_lora + + return shared_outer_lora( + x, + a, + b, + list(counts), + scale=scale, + shared_factor="A" if a.shape[0] == 1 else "B", + ) + if a.shape[0] == 1: + tmp = x @ a[0] + else: + tmp = torch.cat([chunk @ a[e] for e, chunk in enumerate(x.split(counts))]) + if b.shape[0] == 1: + out = tmp @ b[0] + else: + out = torch.cat([chunk @ b[e] for e, chunk in enumerate(tmp.split(counts))]) + return out if scale == 1.0 else out * scale + + +def segmented_expert_lora( + x: torch.Tensor, + factors: Sequence[tuple[torch.Tensor, torch.Tensor, float] | None], + plan: ExpertAdapterPlan, + *, + out_features: int, + shared_outer_kernel: bool = True, +) -> torch.Tensor: + """One base expert pass can use distinct original adapter parameters.""" + if x.ndim != 2 or x.shape[0] != plan.rows or len(factors) != len(plan.refs): + raise ValueError("projection rows/factors disagree with expert plan") + sorted_x = x.index_select(0, plan.permutation) + outputs = [] + for chunk, pair, counts in zip( + sorted_x.split(tuple(map(sum, plan.counts))), factors, plan.counts, strict=True + ): + if pair is None: + outputs.append(chunk.new_zeros((chunk.shape[0], out_features))) + else: + a, b, scale = pair + if b.shape[-1] != out_features: + raise ValueError("expert output width disagrees with factors") + outputs.append( + expert_lora( + chunk, + a, + b, + counts, + scale=scale, + shared_outer_kernel=shared_outer_kernel, + ) + ) + return torch.cat(outputs).index_select(0, plan.inverse) + + +@dataclass(frozen=True) +class AdapterTokens: + refs: tuple["LoRASlotRef | None", ...] + ids: torch.Tensor + + def __post_init__(self) -> None: + object.__setattr__(self, "refs", tuple(self.refs)) + if not 1 <= len(self.refs) <= 256 or len(set(self.refs)) != len(self.refs): + raise ValueError("transport supports 1..256 distinct adapter refs") + if self.ids.ndim != 1 or self.ids.dtype not in (torch.int32, torch.int64): + raise ValueError("token adapter IDs must be a 1D integer vector") + ids = self.ids.detach().clone() + if ids.numel() and (int(ids.min()) < 0 or int(ids.max()) >= len(self.refs)): + raise ValueError("token adapter ID is outside the wave") + object.__setattr__(self, "ids", ids) + + +_CURRENT_ADAPTER_TOKENS: ContextVar[AdapterTokens | None] = ContextVar( + "art_megatron_adapter_tokens", default=None +) + + +@contextmanager +def multi_adapter_tokens(tokens: AdapterTokens | None) -> Iterator[None]: + """Scope immutable source IDs for native EP1 dispatch and recomputation.""" + token = _CURRENT_ADAPTER_TOKENS.set(tokens) + plan_token = _CURRENT_EXPERT_ADAPTER_PLAN.set(None) + try: + yield + finally: + _CURRENT_EXPERT_ADAPTER_PLAN.reset(plan_token) + _CURRENT_ADAPTER_TOKENS.reset(token) + + +def capture_expert_context(function): + """Restore source IDs and dispatched plans on recomputation worker threads.""" + from functools import wraps + + tokens = _CURRENT_ADAPTER_TOKENS.get() + plan = _CURRENT_EXPERT_ADAPTER_PLAN.get() + + @wraps(function) + def wrapped(*args, **kwargs): + with multi_adapter_tokens(tokens), use_expert_adapter_plan(plan): + return function(*args, **kwargs) + + return wrapped + + +def _ep1_dispatch_postprocess(dispatcher, original, hidden_states, probs): + """Replay EP1's exact native permutation on a detached scalar ID column.""" + from megatron.core.transformer.moe.moe_utils import permute + + tokens = _CURRENT_ADAPTER_TOKENS.get() + if tokens is None: + return original(hidden_states, probs) + if ( + tokens.ids.device != hidden_states.device + or tokens.ids.numel() != dispatcher.routing_map.shape[0] + ): + raise ValueError("source adapter IDs must match native routing rows/device") + result = original(hidden_states, probs) + # TP=EP=1 makes native permutation 2 an identity. Use the same permutation + # implementation/flags as hidden states, including TE's fused row map. + ids, _, _, _, _ = permute( + tokens.ids[:, None].to(hidden_states.dtype), + dispatcher.routing_map, + # TE zeros the payload for zero-probability padding slots. Metadata + # must retain IDs even there, so use nonzero detached dummy scores. + probs=torch.ones_like(dispatcher.routing_map, dtype=torch.float32), + num_out_tokens=dispatcher.num_out_tokens, + fused=dispatcher.config.moe_permute_fusion, + drop_and_pad=False, + ) + _CURRENT_EXPERT_ADAPTER_PLAN.set( + ExpertAdapterPlan.build(tokens.refs, ids[:, 0].to(torch.int64), result[1]) + ) + return result + + +def _router_forward(original, hidden_states, padding_mask=None): + """Keep native gate geometry, capacity and auxiliary losses per adapter.""" + from .segmented_lora import _CURRENT_LORA_SEGMENTS, slice_lora_segments + + segments = _CURRENT_LORA_SEGMENTS.get() + if segments is None: + return original(hidden_states, padding_mask=padding_mask) + if ( + hidden_states.ndim != 3 + or hidden_states.shape[1] != 1 + or hidden_states.shape[0] != segments.rows + or not segments.rows + or (padding_mask is not None and padding_mask.shape != hidden_states.shape[:2]) + ): + raise ValueError("segmented router requires complete batch-one adapter rows") + results = [] + start = 0 + for count in segments.counts: + stop = start + count + if count: + with slice_lora_segments(start, stop): + results.append( + original( + hidden_states[start:stop].contiguous(), + padding_mask=None + if padding_mask is None + else padding_mask[start:stop], + ) + ) + start = stop + return tuple(torch.cat(parts, dim=0) for parts in zip(*results, strict=True)) + + +def _row_forward(original, hidden_states, *args, **kwargs): + """Preserve native row shapes for TE normalization/projection algorithms.""" + from .canonical.operators import _join + from .segmented_lora import _CURRENT_LORA_SEGMENTS, slice_lora_segments + + segments = _CURRENT_LORA_SEGMENTS.get() + if ( + segments is None + or hidden_states.shape[0] != segments.rows + or sum(count > 0 for count in segments.counts) < 2 + ): + return original(hidden_states, *args, **kwargs) + outputs = [] + start = 0 + for count in segments.counts: + stop = start + count + if count: + with slice_lora_segments(start, stop): + outputs.append( + original(hidden_states[start:stop].contiguous(), *args, **kwargs) + ) + start = stop + return _join(outputs, next(count for count in segments.counts if count)) + + +def install_segmented_row_modules(model) -> None: + """Opt into native adapter geometry without changing ordinary forwards.""" + from .canonical.operators import _module_types + from .lora import LoRA + + row_types, _, _ = _module_types() + for module in model.modules(): + if ( + isinstance(module, row_types) + and not isinstance(module, LoRA) + and not getattr(module, "_art_segmented_rows", False) + ): + module.forward = partial(_row_forward, module.forward) + module._art_segmented_rows = True + + +def install_segmented_dispatcher(moe_layer) -> None: + """Opt into native TP1/CP1/EP1 adapter IDs; ordinary dispatch stays native.""" + config = moe_layer.config + if ( + getattr(config, "tensor_model_parallel_size", 1) != 1 + or getattr(config, "context_parallel_size", 1) != 1 + or getattr(config, "expert_model_parallel_size", 1) != 1 + or getattr(config, "fp8", None) + or getattr(config, "fp4", None) + or getattr(config, "moe_router_enable_expert_bias", False) + or getattr(getattr(moe_layer, "router", None), "global_tokens_per_expert", None) + is not None + ): + raise ValueError( + "segmented transport requires stateless TP1/CP1/EP1 BF16/FP16 routing" + ) + dispatcher = moe_layer.token_dispatcher + if getattr(dispatcher, "_art_segmented_ep1", False): + return + from megatron.core.transformer.moe.token_dispatcher import ( + MoEAlltoAllTokenDispatcher, + ) + + if type(dispatcher) is not MoEAlltoAllTokenDispatcher: + raise ValueError("unsupported native segmented dispatcher") + if dispatcher.tp_size != 1 or dispatcher.ep_size != 1 or dispatcher.drop_and_pad: + raise ValueError( + "native segmented dispatch requires TP1/EP1 without capacity padding" + ) + dispatcher.dispatch_postprocess = partial( + _ep1_dispatch_postprocess, dispatcher, dispatcher.dispatch_postprocess + ) + if hasattr(moe_layer, "router"): + moe_layer.router.forward = partial(_router_forward, moe_layer.router.forward) + dispatcher._art_segmented_ep1 = True diff --git a/tests/integration/megatron/lora/test_segmented_expert_lora.py b/tests/integration/megatron/lora/test_segmented_expert_lora.py new file mode 100644 index 000000000..96062c94b --- /dev/null +++ b/tests/integration/megatron/lora/test_segmented_expert_lora.py @@ -0,0 +1,255 @@ +"""Native ART slot/recompute integration for compact shared-outer waves.""" + +import tempfile +from typing import cast + +from megatron.core import parallel_state as ps +import pytest +import torch +import torch.distributed as dist +from torch.utils import checkpoint as checkpoint_module + +from art.megatron.lora import LoRA, LoraLayout, LoRASlot, LoRASlotRef +from art.megatron.segmented_expert_lora import ( + ExpertAdapterPlan, + use_expert_adapter_plan, +) + + +@pytest.fixture(scope="module", autouse=True) +def parallel_state(): + if not torch.cuda.is_available(): + pytest.skip("one CUDA device is required") + torch.cuda.set_device(0) + with tempfile.TemporaryDirectory() as init_dir: + dist.init_process_group( + "nccl", init_method=f"file://{init_dir}/store", rank=0, world_size=1 + ) + ps.initialize_model_parallel() + yield + ps.destroy_model_parallel() + dist.destroy_process_group() + + +def test_native_te_normalization_projection_keeps_adapter_row_geometry(): + import transformer_engine.pytorch as te + + from art.megatron.segmented_expert_lora import install_segmented_row_modules + from art.megatron.segmented_lora import LoRASegments, use_lora_segments + + torch.manual_seed(731) + module = te.LayerNormLinear( + 2048, + 64, + params_dtype=torch.bfloat16, + device="cuda", + normalization="RMSNorm", + return_layernorm_output=True, + ) + for parameter in module.parameters(): + parameter.requires_grad_(False) + x = torch.randn( + 386, 1, 2048, dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + dy = torch.randn(386, 1, 64, dtype=x.dtype, device=x.device) + separate = [module(part.contiguous()) for part in x.split((257, 129))] + expected = torch.cat([result[0] for result in separate]) + expected_norm = torch.cat([result[1] for result in separate]) + expected.backward(dy) + dx = x.grad.clone() + x.grad = None + install_segmented_row_modules(module) + with use_lora_segments( + LoRASegments((None, LoRASlotRef("checkpoint", None)), (257, 129)) + ): + actual, normalized = module(x) + actual.backward(dy) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + torch.testing.assert_close(normalized, expected_norm, atol=0, rtol=0) + torch.testing.assert_close(x.grad, dx, atol=0, rtol=0) + + +@pytest.mark.parametrize("shared", ["A", "B"]) +@pytest.mark.parametrize("reentrant", [False, True]) +def test_slot_gradients_and_recompute_keep_dispatched_adapter_plan(shared, reentrant): + torch.backends.cuda.matmul.allow_tf32 = False + torch.manual_seed(41) + lora = LoRA( + "experts.{expert}", + 16, + 24, + 8, + 8, + torch.float32, + torch.device("cuda"), + num_local_experts=3, + ) + refs = tuple(LoRASlotRef("checkpoint", f"run-{i}") for i in range(3)) + for index, ref in enumerate(refs): + slot = LoRASlot( + ref=ref, + a_t=torch.randn(1 if shared == "A" else 3, 16, 3, device="cuda"), + b_t=torch.randn(1 if shared == "B" else 3, 3, 24, device="cuda"), + alpha=3.0, + a_template=lora.A_T, + b_template=lora.B_T, + requires_grad=True, + ) + key = f"slot_{index}" + lora._slot_keys[ref] = key + lora._slot_modules[key] = slot + ids = torch.tensor([2, 0, 1, 0, 1, 2, 0, 1, 2, 0, 2, 1], device="cuda") + experts = [0] * 7 + [2] * 5 + plan = ExpertAdapterPlan.build(refs, ids, [7, 0, 5]) + x = torch.randn(12, 16, device="cuda", requires_grad=True) + dy = torch.randn(12, 24, device="cuda") + rows = [] + for row, (adapter, expert) in enumerate(zip(ids.tolist(), experts, strict=True)): + slot = lora._slot(refs[adapter]) + assert slot is not None + rows.append( + x[row] + @ slot.A_T[0 if shared == "A" else expert] + @ slot.B_T[0 if shared == "B" else expert] + ) + expected = torch.stack(rows) + expected.backward(dy) + dx = cast(torch.Tensor, x.grad).clone() + gradients = { + name: parameter.grad.clone() + for name, parameter in lora.named_parameters() + if parameter.grad is not None + } + lora.zero_grad(set_to_none=True) + x.grad = None + with use_expert_adapter_plan(plan): + actual = checkpoint_module.checkpoint( + lambda value: lora(value, [7, 0, 5]), x, use_reentrant=reentrant + ) + # A subsequent wave changes both geometry and slot ownership before backward. + other = ExpertAdapterPlan.build(refs, torch.tensor([0], device="cuda"), [0, 1, 0]) + with use_expert_adapter_plan(other): + actual.backward(dy) + torch.testing.assert_close(actual, expected, atol=2e-5, rtol=1e-5) + torch.testing.assert_close(x.grad, dx, atol=2e-5, rtol=1e-5) + for name, parameter in lora.named_parameters(): + if name in gradients: + torch.testing.assert_close( + parameter.grad, gradients[name], atol=2e-5, rtol=1e-5 + ) + else: + assert parameter.grad is None + + +@pytest.mark.parametrize("shared", ["A", "B"]) +def test_native_ep1_dispatch_outputs_input_probability_and_slot_gradients(shared): + from types import SimpleNamespace + + from megatron.core.process_groups_config import ProcessGroupCollection + from megatron.core.transformer.moe.token_dispatcher import ( + MoEAlltoAllTokenDispatcher, + ) + from megatron.core.transformer.transformer_config import TransformerConfig + + from art.megatron.segmented_expert_lora import ( + AdapterTokens, + install_segmented_dispatcher, + multi_adapter_tokens, + ) + + torch.manual_seed(81) + config = TransformerConfig( + num_layers=1, + hidden_size=16, + num_attention_heads=2, + num_moe_experts=3, + moe_router_topk=2, + moe_permute_fusion=True, + moe_token_dispatcher_type="alltoall", + params_dtype=torch.bfloat16, + ) + dispatcher = MoEAlltoAllTokenDispatcher( + 3, + [0, 1, 2], + config, + pg_collection=ProcessGroupCollection.use_mpu_process_groups( + required_pgs=["ep", "expt_tp", "tp_ep"] + ), + ) + layer = SimpleNamespace(config=config, token_dispatcher=dispatcher) + install_segmented_dispatcher(layer) + install_segmented_dispatcher(layer) # Idempotent, preserves the native object. + assert layer.token_dispatcher is dispatcher + prefix = "experts.{expert}." + ("gate_proj" if shared == "A" else "down_proj") + lora = LoRA( + prefix, + 16, + 16, + 8, + 8, + torch.bfloat16, + torch.device("cuda"), + num_local_experts=3, + lora_layout=LoraLayout(expert="shared_outer"), + ) + refs = tuple(LoRASlotRef("checkpoint", f"run-{i}") for i in range(3)) + for index, ref in enumerate(refs): + slot = LoRASlot( + ref=ref, + a_t=torch.randn_like(lora.A_T) * 0.02, + b_t=torch.randn_like(lora.B_T) * 0.02, + alpha=8 * (index + 1), + a_template=lora.A_T, + b_template=lora.B_T, + requires_grad=True, + ) + key = f"slot_{index}" + lora._slot_keys[ref], lora._slot_modules[key] = key, slot + x = torch.randn(7, 1, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True) + probabilities = torch.randn(7, 3, device="cuda", requires_grad=True) + routing = torch.ones(7, 3, device="cuda", dtype=torch.bool) + routing[torch.arange(7), torch.arange(7) % 3] = False + ids = torch.tensor([0, 2, 1, 0, 2, 1, 0], device="cuda") + cotangent = torch.randn_like(x) + + def routed(hidden, routes, probs): + hidden, probs = dispatcher.dispatch_preprocess(hidden, routes, probs) + hidden, probs = dispatcher.token_dispatch(hidden, probs) + hidden, counts, probs = dispatcher.dispatch_postprocess(hidden, probs) + projected = lora(hidden, counts) * probs[:, None].to(hidden.dtype) + projected = dispatcher.combine_preprocess(projected) + projected = dispatcher.token_combine(projected) + return dispatcher.combine_postprocess(projected) + + expected = torch.zeros_like(x) + for index, ref in enumerate(refs): + rows = torch.where(ids == index)[0] + tokens = AdapterTokens( + (ref,), torch.zeros(len(rows), dtype=torch.int64, device="cuda") + ) + with multi_adapter_tokens(tokens): + result = routed(x[rows], routing[rows], probabilities[rows]) + expected = expected.index_copy(0, rows, result) + expected.backward(cotangent) + dx, dp = ( + cast(torch.Tensor, x.grad).clone(), + cast(torch.Tensor, probabilities.grad).clone(), + ) + gradients = { + name: parameter.grad.clone() + for name, parameter in lora.named_parameters() + if parameter.grad is not None + } + lora.zero_grad(set_to_none=True) + x.grad = probabilities.grad = None + with multi_adapter_tokens(AdapterTokens(refs, ids)): + actual = routed(x, routing, probabilities) + actual.backward(cotangent) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + torch.testing.assert_close(x.grad, dx, atol=0, rtol=0) + torch.testing.assert_close(probabilities.grad, dp, atol=0, rtol=0) + for name, parameter in lora.named_parameters(): + if name in gradients: + torch.testing.assert_close(parameter.grad, gradients[name], atol=0, rtol=0) + else: + assert parameter.grad is None diff --git a/tests/unit/test_segmented_expert_lora.py b/tests/unit/test_segmented_expert_lora.py new file mode 100644 index 000000000..c2f17a785 --- /dev/null +++ b/tests/unit/test_segmented_expert_lora.py @@ -0,0 +1,241 @@ +"""Independent per-token oracle for shared-outer and per-expert waves.""" + +from typing import Any, cast + +import pytest +import torch + +from art.megatron.segmented_expert_lora import ( + ExpertAdapterPlan, + segmented_expert_lora, +) + + +def test_native_row_geometry_nested_scopes_and_gradients(): + from art.megatron.segmented_expert_lora import _row_forward + from art.megatron.segmented_lora import ( + _CURRENT_LORA_SEGMENTS, + LoRASegments, + use_lora_segments, + ) + + counts = (3, 0, 5) + segments = LoRASegments(cast(Any, ("a", "idle", "b")), counts) + calls = [] + + def inner(value): + calls.append((len(value), _CURRENT_LORA_SEGMENTS.get())) + return value * len(value), torch.ones(7), value + 1 + + def outer(value): + return _row_forward(inner, value) + + x = torch.randn(sum(counts), 1, 7, requires_grad=True) + with use_lora_segments(segments): + output, bias, normalized = _row_forward(outer, x) + assert _CURRENT_LORA_SEGMENTS.get() is segments + assert [n for n, _ in calls] == [3, 5] + assert [scope.counts for _, scope in calls] == [(3, 0, 0), (0, 0, 5)] + expected = torch.cat((x[:3] * 3, x[3:] * 5)) + torch.testing.assert_close(output, expected, atol=0, rtol=0) + torch.testing.assert_close(normalized, x + 1, atol=0, rtol=0) + assert bias.shape == (7,) + output.sum().backward() + torch.testing.assert_close( + x.grad, torch.cat((torch.full_like(x[:3], 3), torch.full_like(x[3:], 5))) + ) + calls.clear() + _row_forward(outer, x) + assert calls == [(8, None)] + + +@pytest.mark.parametrize("shared", ["A", "B", "neither"]) +@pytest.mark.parametrize("empty", [False, True]) +def test_original_factor_gradients_against_per_token_dense_oracle(shared, empty): + torch.manual_seed(17) + counts = [0, 0, 0] if empty else [7, 0, 5] + ids = torch.tensor( + [] if empty else [2, 0, 1, 0, 1, 2, 0, 1, 2, 0, 2, 1], dtype=torch.int64 + ) + refs = cast(Any, ("first", "second", "third", "idle", "missing")) + plan = ExpertAdapterPlan.build(refs, ids, counts) + x = torch.randn(sum(counts), 5, requires_grad=True) + experts = torch.repeat_interleave(torch.arange(3), torch.tensor(counts)) + factors = [ + ( + torch.randn(1 if shared == "A" else 3, 5, rank, requires_grad=True), + torch.randn(1 if shared == "B" else 3, rank, 7, requires_grad=True), + scale, + ) + for rank, scale in zip((1, 3, 8, 2), (0.75, 1.5, 0.25, 2.0), strict=True) + ] + cotangent = torch.randn(sum(counts), 7) + # Fully dense per-row factors are an independent oracle, not the production + # strategy. The implementation must preserve compact shared leaf gradients. + outputs = [] + for row, (adapter, expert) in enumerate( + zip(ids.tolist(), experts.tolist(), strict=True) + ): + a, b, scale = factors[adapter] + outputs.append( + x[row] + @ a[0 if shared == "A" else expert] + @ b[0 if shared == "B" else expert] + * scale + ) + if not empty: + expected = torch.stack(outputs) + expected.backward(cotangent) + dx = cast(torch.Tensor, x.grad).clone() + grads = [ + (cast(torch.Tensor, a.grad).clone(), cast(torch.Tensor, b.grad).clone()) + for a, b, _ in factors[:3] + ] + for a, b, _ in factors: + a.grad = b.grad = None + x.grad = None + actual = segmented_expert_lora(x, factors + [None], plan, out_features=7) + actual.backward(cotangent) + assert x.grad is not None + if not empty: + torch.testing.assert_close(actual, expected, atol=2e-5, rtol=1e-5) + torch.testing.assert_close(x.grad, dx, atol=2e-5, rtol=1e-5) + for (a, b, _), (ga, gb) in zip(factors[:3], grads, strict=True): + torch.testing.assert_close(a.grad, ga, atol=2e-5, rtol=1e-5) + torch.testing.assert_close(b.grad, gb, atol=2e-5, rtol=1e-5) + for a, b, _ in factors: + assert a.grad is not None and a.grad.shape == a.shape + assert b.grad is not None and b.grad.shape == b.shape + for a, b, _ in factors[3:] if not empty else factors: + assert not cast(torch.Tensor, a.grad).any() + assert not cast(torch.Tensor, b.grad).any() + + +def test_plan_validates_dispatch_metadata_and_preserves_row_order(): + refs = cast(Any, ("a", "b")) + ids = torch.tensor([1, 0, 1, 0]) + plan = ExpertAdapterPlan.build(refs, ids, [1, 3]) + values = torch.arange(4) + assert torch.equal(values[plan.permutation][plan.inverse], values) + assert plan.counts == ((0, 2), (1, 1)) + # External ID mutation cannot change a plan retained by autograd. + ids.fill_(0) + assert plan.counts == ((0, 2), (1, 1)) + for bad_ids, counts in ( + (torch.tensor([2]), [1]), + (torch.tensor([0]), [-1]), + (torch.tensor([0]), [2]), + (torch.tensor([0.0]), [1]), + ): + with pytest.raises(ValueError): + ExpertAdapterPlan.build(refs, bad_ids, counts) + with pytest.raises(ValueError, match="rows"): + segmented_expert_lora(torch.empty(5, 3), [None, None], plan, out_features=2) + + +def test_recompute_captures_original_tokens_and_dispatched_plan(): + from art.megatron.segmented_expert_lora import ( + _CURRENT_ADAPTER_TOKENS, + _CURRENT_EXPERT_ADAPTER_PLAN, + AdapterTokens, + capture_expert_context, + multi_adapter_tokens, + use_expert_adapter_plan, + ) + + refs = cast(Any, ("a", "b")) + tokens = AdapterTokens(refs, torch.tensor([0, 1, 1])) + plan = ExpertAdapterPlan.build(refs, tokens.ids, [1, 2]) + with multi_adapter_tokens(tokens), use_expert_adapter_plan(plan): + captured = capture_expert_context( + lambda: (_CURRENT_ADAPTER_TOKENS.get(), _CURRENT_EXPERT_ADAPTER_PLAN.get()) + ) + other = AdapterTokens(refs, torch.tensor([1])) + with multi_adapter_tokens(other): + assert captured() == (tokens, plan) + assert _CURRENT_ADAPTER_TOKENS.get() is other + assert _CURRENT_ADAPTER_TOKENS.get() is None + + +def test_native_ep1_dispatch_replays_routing_without_modifying_base_result(monkeypatch): + import sys + from types import ModuleType, SimpleNamespace + + from art.megatron.segmented_expert_lora import ( + _CURRENT_EXPERT_ADAPTER_PLAN, + AdapterTokens, + _ep1_dispatch_postprocess, + multi_adapter_tokens, + ) + + routing = torch.tensor( + [[1, 0, 1], [1, 1, 0], [0, 1, 1], [0, 0, 1]], dtype=torch.bool + ) + source_rows = routing.T.nonzero()[:, 1] + calls = [] + + def native_permute( + values, actual_routing, *, probs, num_out_tokens, fused, drop_and_pad + ): + assert actual_routing is routing and fused and not drop_and_pad + assert num_out_tokens == 7 and not values.requires_grad + calls.append("permute IDs") + selected = probs.T[routing.T] + # TE masks payloads at zero-probability slots even with a true route. + payload = values[source_rows].masked_fill(selected[:, None] == 0, 0) + return payload, selected, source_rows, None, None + + module = ModuleType("megatron.core.transformer.moe.moe_utils") + setattr(module, "permute", native_permute) + monkeypatch.setitem(sys.modules, module.__name__, module) + dispatcher = SimpleNamespace( + routing_map=routing, + num_out_tokens=7, + config=SimpleNamespace(moe_permute_fusion=True), + ) + hidden = torch.randn(7, 5, requires_grad=True) + probabilities = torch.randn(7, requires_grad=True) + result = (hidden, torch.tensor([2, 2, 3]), probabilities) + + def original(x, p): + assert x is hidden and p is probabilities + calls.append("native postprocess") + return result + + tokens = AdapterTokens(cast(Any, ("a", "b")), torch.tensor([0, 1, 0, 1])) + with multi_adapter_tokens(tokens): + assert ( + _ep1_dispatch_postprocess(dispatcher, original, hidden, probabilities) + is result + ) + plan = _CURRENT_EXPERT_ADAPTER_PLAN.get() + assert plan is not None and plan.counts == ((1, 1, 2), (1, 1, 1)) + assert calls == ["native postprocess", "permute IDs"] + assert _CURRENT_EXPERT_ADAPTER_PLAN.get() is None + # Installation must preserve ordinary non-wave dispatch on the same model. + calls.clear() + assert ( + _ep1_dispatch_postprocess(dispatcher, original, hidden, probabilities) is result + ) + assert calls == ["native postprocess"] + + +def test_installer_rejects_unsupported_transport_and_parallelism(monkeypatch): + import sys + from types import ModuleType, SimpleNamespace + + from art.megatron.segmented_expert_lora import install_segmented_dispatcher + + module = ModuleType("megatron.core.transformer.moe.token_dispatcher") + setattr( + module, "MoEAlltoAllTokenDispatcher", type("MoEAlltoAllTokenDispatcher", (), {}) + ) + monkeypatch.setitem(sys.modules, module.__name__, module) + dispatcher = SimpleNamespace() + layer = SimpleNamespace(config=SimpleNamespace(), token_dispatcher=dispatcher) + with pytest.raises(ValueError, match="unsupported native"): + install_segmented_dispatcher(layer) + assert layer.token_dispatcher is dispatcher + layer.config.tensor_model_parallel_size = 2 + with pytest.raises(ValueError, match="TP1/CP1/EP1"): + install_segmented_dispatcher(layer) diff --git a/tests/unit/test_segmented_router.py b/tests/unit/test_segmented_router.py new file mode 100644 index 000000000..98a9b9cbc --- /dev/null +++ b/tests/unit/test_segmented_router.py @@ -0,0 +1,108 @@ +"""Native router semantics must remain local to each adapter's token rows.""" + +from types import SimpleNamespace +from typing import Any, cast + +import pytest +import torch +import torch.distributed as dist + +from art.megatron.segmented_expert_lora import install_segmented_dispatcher +from art.megatron.segmented_lora import LoRASegments, use_lora_segments + + +@pytest.fixture(scope="module") +def group(tmp_path_factory): + path = tmp_path_factory.mktemp("router-gloo") / "store" + dist.init_process_group("gloo", init_method=f"file://{path}", rank=0, world_size=1) + yield dist.group.WORLD + dist.destroy_process_group() + + +@pytest.mark.parametrize("counts", [(3, 11), (3, 0, 11)]) +@pytest.mark.parametrize("per_token", [False, True]) +@pytest.mark.parametrize( + "aux,z,capacity", [(0.3, None, None), (0, 0.2, None), (0, None, 0.5)] +) +def test_aux_z_and_capacity_match_independent_adapters( + monkeypatch, group, per_token, aux, z, capacity, counts +): + from megatron.core.process_groups_config import ProcessGroupCollection + from megatron.core.transformer.moe import router as router_module + from megatron.core.transformer.moe import token_dispatcher + from megatron.core.transformer.transformer_config import TransformerConfig + + # Keep all actual routing/loss/autograd math; only GPU gate movement and + # logging are outside this CPU operator fixture. + monkeypatch.setattr( + router_module, "save_to_aux_losses_tracker", lambda *a, **k: None + ) + config = TransformerConfig( + num_layers=1, + hidden_size=16, + num_attention_heads=1, + num_moe_experts=4, + moe_router_topk=2, + moe_router_load_balancing_type="aux_loss" if aux else "global_aux_loss", + moe_aux_loss_coeff=aux, + moe_z_loss_coeff=z, + moe_expert_capacity_factor=capacity, + calculate_per_token_loss=per_token, + add_bias_linear=False, + moe_router_fusion=False, + ) + groups = ProcessGroupCollection() + groups.tp = groups.cp = groups.tp_cp = groups.tp_dp_cp = group + router = router_module.TopKRouter(config, pg_collection=groups) + router.set_layer_number(1) + monkeypatch.setattr(router, "gating", lambda value: value) + + class Dispatcher: + tp_size = ep_size = 1 + drop_and_pad = False + + def dispatch_postprocess(self, hidden, probabilities): + return hidden, probabilities + + monkeypatch.setattr(token_dispatcher, "MoEAlltoAllTokenDispatcher", Dispatcher) + layer = SimpleNamespace(config=config, router=router, token_dispatcher=Dispatcher()) + torch.manual_seed(17) + logits = torch.randn(sum(counts), 1, 4) + logits[:3, :, 0] += 4 + logits[3:, :, 1] += 4 + padding = torch.zeros(sum(counts), 1, dtype=torch.bool) + padding[-2:] = True + + def run(fused): + value = logits.clone().requires_grad_() + if fused: + segments = LoRASegments( + cast(Any, tuple(f"run-{i}" for i in range(len(counts)))), counts + ) + with use_lora_segments(segments): + probabilities, routes = router(value, padding_mask=padding) + else: + results = [ + router(part, padding_mask=mask) + for part, mask in zip( + value.split(counts), padding.split(counts), strict=True + ) + if part.shape[0] + ] + probabilities, routes = ( + torch.cat(parts) for parts in zip(*results, strict=True) + ) + probabilities.sum().backward() + return probabilities, routes, value.grad + + expected = run(False) + install_segmented_dispatcher(layer) + installed = router.forward + install_segmented_dispatcher(layer) + assert router.forward is installed + actual = run(True) + for candidate, reference in zip(actual, expected, strict=True): + torch.testing.assert_close(candidate, reference, atol=0, rtol=0) + # Installation alone never changes the ordinary path. + for candidate, reference in zip(run(False), expected, strict=True): + torch.testing.assert_close(candidate, reference, atol=0, rtol=0) From 349d016cf146bb659d6783178b3059ab2a570214 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 21:15:00 +0000 Subject: [PATCH 04/12] fix: retain native attention tile origins in MoE waves Run each independent root through the existing attention backend with cached root-local masks. The full-model trace localizes the remaining forward drift to layer 3 attention with exact Q/K/V inputs. Preserve native execution outside an explicit segmented scope. Co-Authored-By: Claude Opus 5.5 --- src/art/megatron/segmented_expert_lora.py | 105 ++++++++++++++++++ .../lora/test_segmented_expert_lora.py | 69 ++++++++++++ tests/unit/test_segmented_expert_lora.py | 70 ++++++++++++ 3 files changed, 244 insertions(+) diff --git a/src/art/megatron/segmented_expert_lora.py b/src/art/megatron/segmented_expert_lora.py index 404bccc1b..897bb8a3a 100644 --- a/src/art/megatron/segmented_expert_lora.py +++ b/src/art/megatron/segmented_expert_lora.py @@ -340,6 +340,111 @@ def install_segmented_row_modules(model) -> None: module._art_segmented_rows = True +def _attention_forward( + module, + original, + query, + key, + value, + attention_mask, + attn_mask_type=None, + attention_bias=None, + packed_seq_params=None, +): + """Use the native attention backend with each root's original tile origin.""" + from .segmented_lora import _CURRENT_LORA_SEGMENTS, slice_lora_segments + + segments = _CURRENT_LORA_SEGMENTS.get() + if segments is None or sum(n > 0 for n in segments.counts) < 2: + return original( + query, + key, + value, + attention_mask, + attn_mask_type=attn_mask_type, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + ) + if packed_seq_params is not None or any( + tensor.shape[0] != segments.rows or tensor.shape[1] != 1 + for tensor in (query, key, value) + ): + raise ValueError("segmented attention requires complete batch-one roots") + if not all( + hasattr(attention_bias, name) + for name in ("model_state", "group_ids", "parent_ids") + ): + raise ValueError("segmented attention requires prefix-tree metadata") + from .prefix_tree_state import create_prefix_tree_state + + assert attention_bias is not None + window = getattr(module, "art_sliding_window", None) + cache = attention_bias.model_state.setdefault("_art_segmented_attention", {}) + cache_key = (segments.counts, query.shape[-1], value.shape[-1], window) + if cache_key not in cache: + states = [] + start = 0 + for count in segments.counts: + stop = start + count + if count: + groups = attention_bias.group_ids[:, start:stop] + parents = attention_bias.parent_ids[:, start:stop] + if ( + groups.numel() != count + or not torch.equal(groups, parents) + or not bool((groups == groups.flatten()[0]).all()) + ): + raise ValueError( + "segmented attention requires independent unshared roots" + ) + states.append( + create_prefix_tree_state( + groups, + parents, + input_pos=torch.arange(count).unsqueeze(0), + target_device=query.device, + sliding_windows=() if window is None else (window,), + attention_head_dim=query.shape[-1], + attention_value_head_dim=value.shape[-1], + ) + ) + start = stop + cache[cache_key] = states + outputs = [] + start = 0 + states = iter(cache[cache_key]) + for count in segments.counts: + stop = start + count + if count: + with slice_lora_segments(start, stop): + outputs.append( + original( + query[start:stop].contiguous(), + key[start:stop].contiguous(), + value[start:stop].contiguous(), + attention_mask, + attn_mask_type=attn_mask_type, + attention_bias=next(states), + packed_seq_params=None, + ) + ) + start = stop + return torch.cat(outputs, dim=0) + + +def install_segmented_attention_modules(model) -> None: + """Install root-local native attention without changing ordinary execution.""" + from .canonical.operators import _module_types + + _, attention_types, _ = _module_types() + for module in model.modules(): + if isinstance(module, attention_types) and not getattr( + module, "_art_segmented_attention", False + ): + module.forward = partial(_attention_forward, module, module.forward) + module._art_segmented_attention = True + + def install_segmented_dispatcher(moe_layer) -> None: """Opt into native TP1/CP1/EP1 adapter IDs; ordinary dispatch stays native.""" config = moe_layer.config diff --git a/tests/integration/megatron/lora/test_segmented_expert_lora.py b/tests/integration/megatron/lora/test_segmented_expert_lora.py index 96062c94b..548c72c81 100644 --- a/tests/integration/megatron/lora/test_segmented_expert_lora.py +++ b/tests/integration/megatron/lora/test_segmented_expert_lora.py @@ -56,6 +56,7 @@ def test_native_te_normalization_projection_keeps_adapter_row_geometry(): expected = torch.cat([result[0] for result in separate]) expected_norm = torch.cat([result[1] for result in separate]) expected.backward(dy) + assert x.grad is not None dx = x.grad.clone() x.grad = None install_segmented_row_modules(module) @@ -69,6 +70,74 @@ def test_native_te_normalization_projection_keeps_adapter_row_geometry(): torch.testing.assert_close(x.grad, dx, atol=0, rtol=0) +def test_native_attention_preserves_independent_root_tile_origins(): + from megatron.core.transformer.enums import AttnMaskType + from megatron.core.transformer.transformer_config import TransformerConfig + + from art.megatron.flex_attn.attention import FlexDotProductAttention + from art.megatron.prefix_tree_state import create_prefix_tree_state + from art.megatron.segmented_expert_lora import install_segmented_attention_modules + from art.megatron.segmented_lora import LoRASegments, use_lora_segments + + torch.manual_seed(731) + config = TransformerConfig( + num_layers=1, + hidden_size=512, + num_attention_heads=2, + num_query_groups=1, + kv_channels=256, + hidden_dropout=0, + attention_dropout=0, + params_dtype=torch.bfloat16, + ) + module = FlexDotProductAttention(config, 1, AttnMaskType.causal, "self") + counts = (257, 129) + q = torch.randn(386, 1, 2, 256, dtype=torch.bfloat16, device="cuda") + k, v = [ + torch.randn(386, 1, 1, 256, dtype=q.dtype, device=q.device) for _ in range(2) + ] + mask = torch.empty(0, dtype=torch.bool, device=q.device) + expected = [] + start = 0 + for count in counts: + groups = torch.zeros(1, count, dtype=torch.long) + state = create_prefix_tree_state( + groups, + groups, + input_pos=torch.arange(count).unsqueeze(0), + target_device=q.device, + attention_head_dim=256, + attention_value_head_dim=256, + ) + expected.append( + module( + q[start : start + count].contiguous(), + k[start : start + count].contiguous(), + v[start : start + count].contiguous(), + mask, + attention_bias=state, + ) + ) + start += count + groups = torch.cat( + [torch.full((count,), i, dtype=torch.long) for i, count in enumerate(counts)] + )[None, :] + state = create_prefix_tree_state( + groups, + groups, + input_pos=torch.cat([torch.arange(count) for count in counts])[None, :], + target_device=q.device, + attention_head_dim=256, + attention_value_head_dim=256, + ) + install_segmented_attention_modules(module) + with use_lora_segments( + LoRASegments((None, LoRASlotRef("checkpoint", None)), counts) + ): + actual = module(q, k, v, mask, attention_bias=state) + torch.testing.assert_close(actual, torch.cat(expected), atol=0, rtol=0) + + @pytest.mark.parametrize("shared", ["A", "B"]) @pytest.mark.parametrize("reentrant", [False, True]) def test_slot_gradients_and_recompute_keep_dispatched_adapter_plan(shared, reentrant): diff --git a/tests/unit/test_segmented_expert_lora.py b/tests/unit/test_segmented_expert_lora.py index c2f17a785..6c04a3c9b 100644 --- a/tests/unit/test_segmented_expert_lora.py +++ b/tests/unit/test_segmented_expert_lora.py @@ -49,6 +49,76 @@ def outer(value): assert calls == [(8, None)] +def test_native_attention_root_masks_scope_cache_and_foreign_parent_rejection( + monkeypatch, +): + from types import SimpleNamespace + + from art.megatron import prefix_tree_state + from art.megatron.segmented_expert_lora import _attention_forward + from art.megatron.segmented_lora import ( + _CURRENT_LORA_SEGMENTS, + LoRASegments, + use_lora_segments, + ) + + built, calls = [], [] + groups = torch.tensor([[0, 0, 0, 1, 1, 1, 1, 1]]) + state = SimpleNamespace(group_ids=groups, parent_ids=groups.clone(), model_state={}) + module = SimpleNamespace(art_sliding_window=None) + segments = LoRASegments(cast(Any, ("a", "idle", "b")), (3, 0, 5)) + + def build(groups, parents, **kwargs): + assert torch.equal( + kwargs["input_pos"], torch.arange(groups.numel()).unsqueeze(0) + ) + built.append(groups.clone()) + return SimpleNamespace(group_ids=groups, parent_ids=parents) + + monkeypatch.setattr(prefix_tree_state, "create_prefix_tree_state", build) + + def original(q, k, v, mask, *, attention_bias, **kwargs): + current = _CURRENT_LORA_SEGMENTS.get() + calls.append((len(q), current)) + assert attention_bias.group_ids.numel() == len(q) + if current is not None: + assert sum(n > 0 for n in current.counts) == 1 + assert bool( + ( + attention_bias.group_ids == attention_bias.group_ids.flatten()[0] + ).all() + ) + else: + assert attention_bias is state + return (q + k + v).flatten(2) + + q = torch.randn(8, 1, 2, 4, requires_grad=True) + k, v = torch.randn_like(q), torch.randn_like(q) + expected = (q + k + v).flatten(2) + with use_lora_segments(segments): + actual = _attention_forward( + module, original, q, k, v, None, attention_bias=state + ) + _attention_forward(module, original, q, k, v, None, attention_bias=state) + assert _CURRENT_LORA_SEGMENTS.get() is segments + assert [len(g.flatten()) for g in built] == [3, 5] + assert [n for n, _ in calls] == [3, 5, 3, 5] + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + actual.sum().backward() + torch.testing.assert_close(q.grad, torch.ones_like(q), atol=0, rtol=0) + _attention_forward(module, original, q, k, v, None, attention_bias=state) + assert calls[-1] == (8, None) + state.model_state.clear() + state.parent_ids[:, 3:] = 0 + calls.clear() + with ( + use_lora_segments(segments), + pytest.raises(ValueError, match="independent unshared"), + ): + _attention_forward(module, original, q, k, v, None, attention_bias=state) + assert not calls + + @pytest.mark.parametrize("shared", ["A", "B", "neither"]) @pytest.mark.parametrize("empty", [False, True]) def test_original_factor_gradients_against_per_token_dense_oracle(shared, empty): From f20b8b3b0a5a3444fee21ec3f6d698edc858a9cb Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 21:26:33 +0000 Subject: [PATCH 05/12] fix: reuse forward-preserving deterministic backward for MoE waves Reuse native-repeat d17408e0 lowering and regression tests verbatim. Add an explicit native execution scope that enables the standard determinism opt-in without changing root geometry. Native controls and ART defaults retain their existing execution. Co-Authored-By: Claude Opus 5.5 --- src/art/megatron/canonical/README.md | 6 + src/art/megatron/canonical/layout.py | 11 +- src/art/megatron/flex_attn/compiled.py | 6 +- .../megatron/flex_attn/flash_dlse_patch.py | 38 +++++++ .../megatron/test_native_determinism.py | 107 ++++++++++++++++++ tests/unit/test_canonical_numerics.py | 18 +++ tests/unit/test_flex_deterministic.py | 27 +++++ 7 files changed, 209 insertions(+), 4 deletions(-) create mode 100644 tests/integration/megatron/test_native_determinism.py create mode 100644 tests/unit/test_flex_deterministic.py diff --git a/src/art/megatron/canonical/README.md b/src/art/megatron/canonical/README.md index 817aaad8a..773be8455 100644 --- a/src/art/megatron/canonical/README.md +++ b/src/art/megatron/canonical/README.md @@ -27,3 +27,9 @@ mask the root layout without modifying process flags. Record `POLICY_VERSION` and kernel/runtime versions with a run; retain the same policy between forward, backward and replay. The GPU regression is `tests/integration/megatron/test_canonical_numerics.py`. + +`execution(None, numerical_policy="native", deterministic_native=True)` scopes +the standard PyTorch determinism opt-in while retaining native root geometry. +ART's FLASH lowering preserves its forward and uses deterministic TRITON +backward, converting saved LSE/cotangent units at the existing boundary. The +process-wide flags, locking and thread-quiescence requirements above still apply. diff --git a/src/art/megatron/canonical/layout.py b/src/art/megatron/canonical/layout.py index 67843c284..9dd3df0fa 100644 --- a/src/art/megatron/canonical/layout.py +++ b/src/art/megatron/canonical/layout.py @@ -93,7 +93,10 @@ def wrapped(*args: Any, **kwargs: Any) -> Any: @contextmanager def execution( - layout: RootLayout | None, *, numerical_policy: NumericalPolicy = "canonical" + layout: RootLayout | None, + *, + numerical_policy: NumericalPolicy = "canonical", + deterministic_native: bool = False, ) -> Iterator[None]: """One synchronous forward/recompute/backward scope; no async suspension. @@ -101,7 +104,9 @@ def execution( backward completes. Other TrainerRank/planner/prefetch threads inherit them. Callers must dedicate the process to this policy or quiesce other PyTorch work; this lock only serializes callers of execution(), including native - scopes. Never suspend async work here. Native scopes only mask the layout. + scopes. Never suspend async work here. Native scopes mask the layout; + deterministic_native=True also opts into deterministic backward without + selecting canonical root geometry. """ if numerical_policy not in ("native", "canonical"): raise ValueError(f"unknown numerical policy: {numerical_policy}") @@ -109,7 +114,7 @@ def execution( raise ValueError("canonical execution requires a root layout") with _EXECUTION_LOCK: token = _LAYOUT.set(layout if numerical_policy == "canonical" else None) - if numerical_policy == "native": + if numerical_policy == "native" and not deterministic_native: try: yield finally: diff --git a/src/art/megatron/flex_attn/compiled.py b/src/art/megatron/flex_attn/compiled.py index bc38248e6..ec7caede9 100644 --- a/src/art/megatron/flex_attn/compiled.py +++ b/src/art/megatron/flex_attn/compiled.py @@ -14,9 +14,13 @@ flex_attention, ) -from art.megatron.flex_attn.flash_dlse_patch import apply_flash_flex_dlse_patch +from art.megatron.flex_attn.flash_dlse_patch import ( + apply_deterministic_flash_backward_patch, + apply_flash_flex_dlse_patch, +) apply_flash_flex_dlse_patch() +apply_deterministic_flash_backward_patch() # Integration tests patch this module in-process when they need a non-default diff --git a/src/art/megatron/flex_attn/flash_dlse_patch.py b/src/art/megatron/flex_attn/flash_dlse_patch.py index f80bab880..a7919da70 100644 --- a/src/art/megatron/flex_attn/flash_dlse_patch.py +++ b/src/art/megatron/flex_attn/flash_dlse_patch.py @@ -8,6 +8,7 @@ from __future__ import annotations import inspect +import math from typing import Any, cast import torch @@ -15,6 +16,7 @@ _PATCH_APPLIED = False _TILE_PATCH_APPLIED = False _DQ_POSTPROCESS_PATCH_APPLIED = False +_DETERMINISTIC_BWD_PATCH_APPLIED = False def _sm90_block_sparse_fwd_config(cute_interface: Any, head_dim: int, head_dim_v: int): @@ -583,3 +585,39 @@ def flex_attention_backward_with_flash_dlse(*args, **kwargs): flex_attention_backward_with_flash_dlse ) _PATCH_APPLIED = True + + +def apply_deterministic_flash_backward_patch() -> None: + """Keep FLASH forward; honor PyTorch determinism with Triton backward. + + FLASH saves natural-log LSE, whereas Triton consumes log2 LSE. Convert + both the saved value and its cotangent at the existing lowering boundary. + """ + global _DETERMINISTIC_BWD_PATCH_APPLIED + if _DETERMINISTIC_BWD_PATCH_APPLIED: + return + + import torch._inductor.kernel.flex.flex_attention as flex_attention_mod + from torch._inductor.lowering import lowerings + + operation = torch.ops.higher_order.flex_attention_backward + original = lowerings[operation] + + def backward(*args, **kwargs): + if ( + not kwargs + and torch.are_deterministic_algorithms_enabled() + and args[11].get("BACKEND") == "FLASH" + ): + converted = list(args) + multiply = lowerings[torch.ops.aten.mul.Tensor] + converted[4] = multiply(args[4], math.log2(math.e)) + if args[6] is not None: + converted[6] = multiply(args[6], math.log(2.0)) + converted[11] = {**args[11], "BACKEND": "TRITON"} + return original(*converted) + return original(*args, **kwargs) + + cast(Any, flex_attention_mod).flex_attention_backward = backward + lowerings[operation] = backward + _DETERMINISTIC_BWD_PATCH_APPLIED = True diff --git a/tests/integration/megatron/test_native_determinism.py b/tests/integration/megatron/test_native_determinism.py new file mode 100644 index 000000000..6edd02e66 --- /dev/null +++ b/tests/integration/megatron/test_native_determinism.py @@ -0,0 +1,107 @@ +"""Deterministic backward preserves FLASH forward, including natural-log LSE.""" + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("use_lse", [False, True]) +@pytest.mark.parametrize( + "length,dim,heads,kv_heads", [(256, 128, 8, 2), (183, 256, 16, 2)] +) +def test_masked_gqa_deterministic_backward(use_lse, length, dim, heads, kv_heads): + from torch.nn.attention.flex_attention import AuxRequest, create_block_mask + + from art.megatron.flex_attn.compiled import ( + flash_sparse_block_size_for_head_dim, + flex_backend_for_head_dims, + get_dense_compiled_flex_attention, + normalize_flex_lse, + ) + + original = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + fill = getattr(torch.utils.deterministic, "fill_uninitialized_memory") + generator = torch.Generator(device="cuda").manual_seed(108) + q, k, v, grad = [ + torch.randn(shape, generator=generator, device="cuda", dtype=torch.bfloat16) + for shape in [ + (1, heads, length, dim), + (1, kv_heads, length, dim), + (1, kv_heads, length, dim), + (1, heads, length, dim), + ] + ] + + def causal(batch, head, query, key): + return query >= key + + try: + torch.use_deterministic_algorithms(True) + setattr(torch.utils.deterministic, "fill_uninitialized_memory", True) + backend = flex_backend_for_head_dims( + head_dim=dim, head_dim_v=dim, device=q.device + ) + assert backend == "FLASH" + attention = get_dense_compiled_flex_attention( + backend=backend, head_dim=dim, head_dim_v=dim, device=q.device + ) + block_size = flash_sparse_block_size_for_head_dim( + head_dim=dim, head_dim_v=dim, device=q.device + ) + mask = create_block_mask( + causal, 1, None, length, length, device="cuda", BLOCK_SIZE=block_size + ) + + def forward(inputs): + result = attention( + *inputs, + block_mask=mask, + scale=dim**-0.5, + enable_gqa=True, + return_aux=AuxRequest(lse=True) if use_lse else None, + ) + if use_lse: + out, aux = result + assert aux.lse is not None + return out, normalize_flex_lse(aux.lse, backend=backend) + return (result,) + + torch.use_deterministic_algorithms(False) + baseline = forward((q, k, v)) + torch.use_deterministic_algorithms(True) + runs = [] + for _ in range(3): + inputs = [x.detach().requires_grad_() for x in (q, k, v)] + outputs = forward(inputs) + assert all( + torch.equal(a, b) for a, b in zip(outputs, baseline, strict=True) + ) + upstream = (grad, torch.full_like(outputs[1], 0.31)) if use_lse else (grad,) + runs.append( + ( + *[x.detach() for x in outputs], + *torch.autograd.grad(outputs, inputs, upstream), + ) + ) + for run in runs[1:]: + assert all(torch.equal(a, b) for a, b in zip(runs[0], run, strict=True)) + + rq, rk, rv = [x.double().requires_grad_() for x in (q, k, v)] + scores = ( + rq @ rk.repeat_interleave(heads // kv_heads, 1).transpose(-1, -2) + ) * dim**-0.5 + visible = torch.ones(length, length, device=q.device, dtype=torch.bool).tril() + scores = scores.masked_fill(~visible, -torch.inf) + expected = scores.softmax(-1) @ rv.repeat_interleave(heads // kv_heads, 1) + reference_outputs = (expected, scores.logsumexp(-1)) if use_lse else (expected,) + oracle = torch.autograd.grad( + reference_outputs, (rq, rk, rv), tuple(x.double() for x in upstream) + ) + for actual, reference in zip(runs[0][-3:], oracle, strict=True): + assert torch.isfinite(actual).all() + assert (actual.double() - reference).norm() / reference.norm() <= 0.01 + finally: + setattr(torch.utils.deterministic, "fill_uninitialized_memory", fill) + torch.use_deterministic_algorithms(original, warn_only=warn_only) diff --git a/tests/unit/test_canonical_numerics.py b/tests/unit/test_canonical_numerics.py index 75900edff..95d640139 100644 --- a/tests/unit/test_canonical_numerics.py +++ b/tests/unit/test_canonical_numerics.py @@ -29,6 +29,24 @@ def test_layout_pads_roots_independently_and_masks_labels(): layout.pad(roots[::-1]) +def test_deterministic_native_scope_keeps_layout_native_and_restores_flags(): + initial = torch.are_deterministic_algorithms_enabled() + warning = torch.is_deterministic_algorithms_warn_only_enabled() + fill = getattr(torch.utils.deterministic, "fill_uninitialized_memory") + with pytest.raises(RuntimeError, match="abort"): + with execution(None, numerical_policy="native", deterministic_native=True): + assert current_layout() is None + assert torch.are_deterministic_algorithms_enabled() + assert not torch.is_deterministic_algorithms_warn_only_enabled() + with execution(None, numerical_policy="native"): + assert current_layout() is None + assert torch.are_deterministic_algorithms_enabled() + raise RuntimeError("abort") + assert torch.are_deterministic_algorithms_enabled() == initial + assert torch.is_deterministic_algorithms_warn_only_enabled() == warning + assert getattr(torch.utils.deterministic, "fill_uninitialized_memory") == fill + + @pytest.mark.parametrize("lengths", [(), (0,), (-1,), (1.5,), (True,)]) def test_invalid_root_lengths(lengths): with pytest.raises(ValueError): diff --git a/tests/unit/test_flex_deterministic.py b/tests/unit/test_flex_deterministic.py new file mode 100644 index 000000000..cc31e3b5b --- /dev/null +++ b/tests/unit/test_flex_deterministic.py @@ -0,0 +1,27 @@ +"""The standard PyTorch determinism opt-in must select a supported backend.""" + +import pytest +import torch + +compiled = pytest.importorskip("art.megatron.flex_attn.compiled") + + +@pytest.mark.parametrize("enabled,expected", [(False, "FLASH"), (True, "FLASH")]) +def test_determinism_preserves_forward_backend(monkeypatch, enabled, expected): + monkeypatch.setattr(torch, "are_deterministic_algorithms_enabled", lambda: enabled) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (9, 0)) + monkeypatch.setattr(compiled, "_FORCED_FLEX_BACKEND", "FLASH") + assert ( + compiled.flex_backend_for_head_dims( + head_dim=128, head_dim_v=128, device=torch.device("cuda") + ) + == expected + ) + + +def test_determinism_preserves_flash_mask_geometry(monkeypatch): + monkeypatch.setattr(torch, "are_deterministic_algorithms_enabled", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (9, 0)) + assert compiled.flash_sparse_block_size_for_head_dim( + head_dim=256, head_dim_v=256, device=torch.device("cuda") + ) == (128, 64) From 16605fa0e51fc72aa206919542d17dcf2dc7bc34 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 22:14:52 +0000 Subject: [PATCH 06/12] fix: preserve native GDN gate reductions in opt-in MoE waves Capture independent-root reduction scope on FLA autograd contexts. Keep recurrent GEMMs batched while using batch-one scalar cumsum geometry in forward and backward. Co-Authored-By: Claude Opus 5.5 --- src/art/megatron/segmented_gdn.py | 78 +++++++++++++++++++ .../megatron/test_segmented_gdn.py | 69 ++++++++++++++++ tests/unit/test_segmented_gdn.py | 72 +++++++++++++++++ 3 files changed, 219 insertions(+) create mode 100644 src/art/megatron/segmented_gdn.py create mode 100644 tests/integration/megatron/test_segmented_gdn.py create mode 100644 tests/unit/test_segmented_gdn.py diff --git a/src/art/megatron/segmented_gdn.py b/src/art/megatron/segmented_gdn.py new file mode 100644 index 000000000..6e2149b55 --- /dev/null +++ b/src/art/megatron/segmented_gdn.py @@ -0,0 +1,78 @@ +"""Opt-in native sequence geometry for FLA's scalar gate reductions. + +FLA autotunes cumulative sums by batch size. Different warp counts change the +FP32 reverse sum tree and therefore alpha gradients, even with exact Q/K/V and +outputs. Keep the heavy recurrent kernels batched; only these small reductions +use each independent root's native geometry. Capture the scope on FLA's autograd +context so backward does not depend on the caller's ContextVars. +""" + +from contextlib import contextmanager +import contextvars +from functools import partial +import importlib + +import torch + +_ACTIVE = contextvars.ContextVar("art_native_gdn_reductions", default=False) + + +def _cumsum(original, tensor, chunk_size, **kwargs): + if not _ACTIVE.get(): + return original(tensor, chunk_size, **kwargs) + cu = kwargs.get("cu_seqlens") + if cu is None: + if tensor.shape[0] == 1: + return original(tensor, chunk_size, **kwargs) + return torch.cat( + [ + original(part.contiguous(), chunk_size, **kwargs) + for part in tensor.split(1) + ], + dim=0, + ) + if tensor.shape[0] != 1 or kwargs.get("head_first", False): + raise ValueError("native wave GDN requires batch-one packed gate streams") + boundaries = cu.detach().cpu().tolist() + local = dict(kwargs, cu_seqlens=None, chunk_indices=None) + return torch.cat( + [ + original(tensor[:, start:stop].contiguous(), chunk_size, **local) + for start, stop in zip(boundaries[:-1], boundaries[1:], strict=True) + ], + dim=1, + ) + + +def _forward(original, ctx, *args, **kwargs): + ctx._art_native_gdn_reductions = _ACTIVE.get() + return original(ctx, *args, **kwargs) + + +def _backward(original, ctx, *args, **kwargs): + token = _ACTIVE.set(ctx._art_native_gdn_reductions) + try: + return original(ctx, *args, **kwargs) + finally: + _ACTIVE.reset(token) + + +def _install(module): + function = module.ChunkGatedDeltaRuleFunction + if getattr(function, "_art_native_reductions", False): + return + module.chunk_local_cumsum = partial(_cumsum, module.chunk_local_cumsum) + function.forward = staticmethod(partial(_forward, function.forward)) + function.backward = staticmethod(partial(_backward, function.backward)) + function._art_native_reductions = True + + +@contextmanager +def native_gdn_reductions(): + """Explicit independent-root MoE scope; ordinary forwards stay unchanged.""" + _install(importlib.import_module("fla.ops.gated_delta_rule.chunk")) + token = _ACTIVE.set(True) + try: + yield + finally: + _ACTIVE.reset(token) diff --git a/tests/integration/megatron/test_segmented_gdn.py b/tests/integration/megatron/test_segmented_gdn.py new file mode 100644 index 000000000..2ae21769d --- /dev/null +++ b/tests/integration/megatron/test_segmented_gdn.py @@ -0,0 +1,69 @@ +"""Native FLA gate gradients must keep their per-root reduction geometry.""" + +import pytest +import torch + +from art.megatron.segmented_gdn import native_gdn_reductions + + +@pytest.mark.parametrize("packed", [False, True]) +def test_batched_recurrence_matches_separate_native_gate_gradients(packed): + if not torch.cuda.is_available(): + pytest.skip("one CUDA device required") + from fla.ops.gated_delta_rule import chunk_gated_delta_rule + + torch.manual_seed(731) + n, length, heads, width = 16, 1024, 32, 128 + q, k = [ + torch.nn.functional.normalize( + torch.randn(n, length, heads, width, device="cuda"), dim=-1 + ) + .bfloat16() + .requires_grad_() + for _ in range(2) + ] + v = torch.randn_like(q, requires_grad=True) + g = (-torch.rand(n, length, heads, device="cuda") * 0.1).requires_grad_() + beta = torch.rand( + n, length, heads, device="cuda", dtype=q.dtype, requires_grad=True + ) + dy = torch.randn_like(v) + variables = [q, k, v, g, beta] + expected = torch.cat( + [ + chunk_gated_delta_rule( + q[i : i + 1], + k[i : i + 1], + v[i : i + 1], + g=g[i : i + 1], + beta=beta[i : i + 1], + output_final_state=False, + use_qk_l2norm_in_kernel=False, + )[0] + for i in range(n) + ] + ) + expected.backward(dy) + gradients = [x.grad.clone() for x in variables] + for x in variables: + x.grad = None + args = variables + kwargs = {} + if packed: + args = [x.reshape(1, n * length, *x.shape[2:]) for x in variables] + boundaries = torch.arange(n + 1, dtype=torch.int32) * length + kwargs = {"cu_seqlens": boundaries.cuda(), "cu_seqlens_cpu": boundaries} + with native_gdn_reductions(): + actual = chunk_gated_delta_rule( + *args[:3], + g=args[3], + beta=args[4], + output_final_state=False, + use_qk_l2norm_in_kernel=False, + **kwargs, + )[0].reshape_as(expected) + # Scope has exited before backward: FLA's context must retain it. + actual.backward(dy) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + for x, gradient in zip(variables, gradients, strict=True): + torch.testing.assert_close(x.grad, gradient, atol=0, rtol=0) diff --git a/tests/unit/test_segmented_gdn.py b/tests/unit/test_segmented_gdn.py new file mode 100644 index 000000000..f86d5a06e --- /dev/null +++ b/tests/unit/test_segmented_gdn.py @@ -0,0 +1,72 @@ +from types import SimpleNamespace + +import pytest +import torch + +from art.megatron import segmented_gdn as gdn + + +@pytest.mark.parametrize("packed", [False, True]) +def test_gate_reductions_keep_native_sequence_shapes_and_backward_scope(packed): + calls = [] + + def cumsum(value, chunk_size, **kwargs): + calls.append((tuple(value.shape), kwargs)) + return value.cumsum(dim=1) + + class Function: + @staticmethod + def forward(ctx, value, **kwargs): + ctx.value, ctx.kwargs = value, kwargs + return module.chunk_local_cumsum(value, 64, **kwargs) + + @staticmethod + def backward(ctx): + return module.chunk_local_cumsum(ctx.value, 64, **ctx.kwargs) + + module = SimpleNamespace( + ChunkGatedDeltaRuleFunction=Function, chunk_local_cumsum=cumsum + ) + gdn._install(module) + installed = module.chunk_local_cumsum + gdn._install(module) + assert module.chunk_local_cumsum is installed + values = torch.arange(24.0).reshape(2, 3, 4) + expected = values.cumsum(dim=1) + kwargs = {} + if packed: + values = values.reshape(1, 6, 4) + kwargs = { + "cu_seqlens": torch.tensor([0, 3, 6]), + "chunk_indices": torch.ones(2, 2), + } + expected = expected.reshape_as(values) + ctx = SimpleNamespace() + token = gdn._ACTIVE.set(True) + try: + actual = Function.forward(ctx, values, **kwargs) + finally: + gdn._ACTIVE.reset(token) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert [shape for shape, _ in calls] == [(1, 3, 4), (1, 3, 4)] + calls.clear() + torch.testing.assert_close(Function.backward(ctx), expected, atol=0, rtol=0) + assert [shape for shape, _ in calls] == [(1, 3, 4), (1, 3, 4)] + assert not gdn._ACTIVE.get() + calls.clear() + ordinary = SimpleNamespace() + Function.forward(ordinary, values, **kwargs) + Function.backward(ordinary) + assert [shape for shape, _ in calls] == [tuple(values.shape)] * 2 + + +def test_backward_scope_restores_after_failure(): + ctx = SimpleNamespace(_art_native_gdn_reductions=True) + + def fail(ctx): + assert gdn._ACTIVE.get() + raise RuntimeError("backward failure") + + with pytest.raises(RuntimeError, match="backward failure"): + gdn._backward(fail, ctx) + assert not gdn._ACTIVE.get() From 8d8b17360229a7532255588cf70c86a8695c93a0 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 22:24:26 +0000 Subject: [PATCH 07/12] fix: retain MoE gate reductions through activation recompute Use the existing LoRA checkpoint context capture to carry the opt-in GDN reduction scope onto autograd worker threads. Add thread and real recurrence checkpoint regressions. Co-Authored-By: Claude Opus 5.5 --- src/art/megatron/lora.py | 6 +++- src/art/megatron/segmented_gdn.py | 18 ++++++++++- .../megatron/test_segmented_gdn.py | 30 ++++++++++++++----- tests/unit/test_segmented_gdn.py | 18 +++++++++++ 4 files changed, 63 insertions(+), 9 deletions(-) diff --git a/src/art/megatron/lora.py b/src/art/megatron/lora.py index dbe454664..08992bcfa 100644 --- a/src/art/megatron/lora.py +++ b/src/art/megatron/lora.py @@ -179,7 +179,11 @@ def _with_captured_lora_slot(function: _F) -> _F: context = _CURRENT_LORA_SLOT.get() segments = _CURRENT_LORA_SEGMENTS.get() expert_plan = _CURRENT_EXPERT_ADAPTER_PLAN.get() - function = capture_expert_context(capture_routes(capture_canonical(function))) + from .segmented_gdn import capture as capture_gdn_reductions + + function = capture_gdn_reductions( + capture_expert_context(capture_routes(capture_canonical(function))) + ) @functools.wraps(function) def wrapped(*args: Any, **kwargs: Any) -> Any: diff --git a/src/art/megatron/segmented_gdn.py b/src/art/megatron/segmented_gdn.py index 6e2149b55..60b8fe216 100644 --- a/src/art/megatron/segmented_gdn.py +++ b/src/art/megatron/segmented_gdn.py @@ -9,7 +9,7 @@ from contextlib import contextmanager import contextvars -from functools import partial +from functools import partial, wraps import importlib import torch @@ -76,3 +76,19 @@ def native_gdn_reductions(): yield finally: _ACTIVE.reset(token) + + +def capture(function): + """Preserve this opt-in scope across checkpoint/autograd worker threads.""" + if not _ACTIVE.get(): + return function + + @wraps(function) + def wrapped(*args, **kwargs): + token = _ACTIVE.set(True) + try: + return function(*args, **kwargs) + finally: + _ACTIVE.reset(token) + + return wrapped diff --git a/tests/integration/megatron/test_segmented_gdn.py b/tests/integration/megatron/test_segmented_gdn.py index 2ae21769d..b430c2df5 100644 --- a/tests/integration/megatron/test_segmented_gdn.py +++ b/tests/integration/megatron/test_segmented_gdn.py @@ -7,7 +7,10 @@ @pytest.mark.parametrize("packed", [False, True]) -def test_batched_recurrence_matches_separate_native_gate_gradients(packed): +@pytest.mark.parametrize("checkpointed", [False, True]) +def test_batched_recurrence_matches_separate_native_gate_gradients( + packed, checkpointed +): if not torch.cuda.is_available(): pytest.skip("one CUDA device required") from fla.ops.gated_delta_rule import chunk_gated_delta_rule @@ -53,15 +56,28 @@ def test_batched_recurrence_matches_separate_native_gate_gradients(packed): args = [x.reshape(1, n * length, *x.shape[2:]) for x in variables] boundaries = torch.arange(n + 1, dtype=torch.int32) * length kwargs = {"cu_seqlens": boundaries.cuda(), "cu_seqlens_cpu": boundaries} - with native_gdn_reductions(): - actual = chunk_gated_delta_rule( - *args[:3], - g=args[3], - beta=args[4], + + def run(*values): + return chunk_gated_delta_rule( + *values[:3], + g=values[3], + beta=values[4], output_final_state=False, use_qk_l2norm_in_kernel=False, **kwargs, - )[0].reshape_as(expected) + )[0] + + from torch.utils import checkpoint as checkpoint_module + + from art.megatron.lora import install_lora_checkpoint_context_hooks + + install_lora_checkpoint_context_hooks() + with native_gdn_reductions(): + actual = ( + checkpoint_module.checkpoint(run, *args, use_reentrant=True) + if checkpointed + else run(*args) + ).reshape_as(expected) # Scope has exited before backward: FLA's context must retain it. actual.backward(dy) torch.testing.assert_close(actual, expected, atol=0, rtol=0) diff --git a/tests/unit/test_segmented_gdn.py b/tests/unit/test_segmented_gdn.py index f86d5a06e..b02e1280b 100644 --- a/tests/unit/test_segmented_gdn.py +++ b/tests/unit/test_segmented_gdn.py @@ -70,3 +70,21 @@ def fail(ctx): with pytest.raises(RuntimeError, match="backward failure"): gdn._backward(fail, ctx) assert not gdn._ACTIVE.get() + + +def test_art_checkpoint_capture_restores_reduction_scope_on_worker_thread(): + from concurrent.futures import ThreadPoolExecutor + + from art.megatron.lora import _with_captured_lora_slot + + ordinary = _with_captured_lora_slot(lambda: gdn._ACTIVE.get()) + token = gdn._ACTIVE.set(True) + try: + captured = _with_captured_lora_slot(lambda: gdn._ACTIVE.get()) + finally: + gdn._ACTIVE.reset(token) + with ThreadPoolExecutor(max_workers=1) as pool: + assert pool.submit(captured).result() is True + assert pool.submit(gdn._ACTIVE.get).result() is False + assert pool.submit(ordinary).result() is False + assert not gdn._ACTIVE.get() From 4bba0bf93da54477f5537f9cbe5ced43afcdd63e Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 22:30:54 +0000 Subject: [PATCH 08/12] fix: preserve preexisting native GDN autograd contexts Contexts created before the explicit installer retain native reductions. Flagged MoE contexts follow the identical captured-scope path qualified in 8d8b1736. Co-Authored-By: Claude Opus 5.5 --- src/art/megatron/segmented_gdn.py | 2 +- tests/unit/test_segmented_gdn.py | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/src/art/megatron/segmented_gdn.py b/src/art/megatron/segmented_gdn.py index 60b8fe216..f90352ff6 100644 --- a/src/art/megatron/segmented_gdn.py +++ b/src/art/megatron/segmented_gdn.py @@ -50,7 +50,7 @@ def _forward(original, ctx, *args, **kwargs): def _backward(original, ctx, *args, **kwargs): - token = _ACTIVE.set(ctx._art_native_gdn_reductions) + token = _ACTIVE.set(getattr(ctx, "_art_native_gdn_reductions", False)) try: return original(ctx, *args, **kwargs) finally: diff --git a/tests/unit/test_segmented_gdn.py b/tests/unit/test_segmented_gdn.py index b02e1280b..f66e2bb26 100644 --- a/tests/unit/test_segmented_gdn.py +++ b/tests/unit/test_segmented_gdn.py @@ -88,3 +88,13 @@ def test_art_checkpoint_capture_restores_reduction_scope_on_worker_thread(): assert pool.submit(gdn._ACTIVE.get).result() is False assert pool.submit(ordinary).result() is False assert not gdn._ACTIVE.get() + + +def test_context_created_before_installation_keeps_native_backward(): + ctx = SimpleNamespace() + token = gdn._ACTIVE.set(True) + try: + assert gdn._backward(lambda ctx: gdn._ACTIVE.get(), ctx) is False + assert gdn._ACTIVE.get() is True + finally: + gdn._ACTIVE.reset(token) From 9f026c90f95416cb4b164959fb1c790536db8e95 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 21:57:32 +0000 Subject: [PATCH 09/12] test: cover Hopper tail attention in deterministic backward Include the dense fixture GQA geometry and 65-token tail, with and without LSE gradients. Scope the FLASH-specific regression to SM90 because Blackwell already selects Triton. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 355fccefcec8c5827a2a7be74a65df5724de4f79) --- tests/integration/megatron/test_native_determinism.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/integration/megatron/test_native_determinism.py b/tests/integration/megatron/test_native_determinism.py index 6edd02e66..03e113290 100644 --- a/tests/integration/megatron/test_native_determinism.py +++ b/tests/integration/megatron/test_native_determinism.py @@ -3,12 +3,15 @@ import pytest import torch -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9, + reason="SM90 FLASH attention required", +) @pytest.mark.parametrize("use_lse", [False, True]) @pytest.mark.parametrize( - "length,dim,heads,kv_heads", [(256, 128, 8, 2), (183, 256, 16, 2)] + "length,dim,heads,kv_heads", [(256, 128, 8, 2), (183, 256, 16, 2), (65, 256, 24, 4)] ) def test_masked_gqa_deterministic_backward(use_lse, length, dim, heads, kv_heads): from torch.nn.attention.flex_attention import AuxRequest, create_block_mask From ca75d38cb09e44ad45a7c59f3f7d386d5c0bf966 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 22:03:56 +0000 Subject: [PATCH 10/12] test: isolate compiled attention policy cases Reset Dynamo between independent model geometries and reject its unfused debug fallback. Shared test specialization history exhausted the compile cache, changed LSE conventions and violated the unchanged FP64 gate; do not widen that tolerance. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 3b4517eca0179058fac520ac19ee425c148a9142) --- tests/integration/megatron/test_native_determinism.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/integration/megatron/test_native_determinism.py b/tests/integration/megatron/test_native_determinism.py index 03e113290..718dda5a2 100644 --- a/tests/integration/megatron/test_native_determinism.py +++ b/tests/integration/megatron/test_native_determinism.py @@ -9,6 +9,17 @@ ) +@pytest.fixture(autouse=True) +def fresh_compilation_history(): + # Each case represents a fresh trainer with one fixed numerical policy. + # Shared specialization history can exhaust Dynamo's cache and run the + # unfused debug fallback, whose LSE convention is different from FLASH. + torch._dynamo.reset() + yield + torch._dynamo.reset() + + +@pytest.mark.filterwarnings("error:flex_attention called without torch.compile") @pytest.mark.parametrize("use_lse", [False, True]) @pytest.mark.parametrize( "length,dim,heads,kv_heads", [(256, 128, 8, 2), (183, 256, 16, 2), (65, 256, 24, 4)] From f47f793c621a5942cbb51db5c82d7985a7edfe01 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 21 Oct 2026 22:34:23 +0000 Subject: [PATCH 11/12] fix: reject uncompiled FLASH fallback in deterministic mode Keep the deterministic opt-in on its compiled backward lowering when Dynamo exhausts specialization capacity. Preserve default eager behavior and cover the cache transition with a failing-first CPU regression. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit e23d04821f29e4ac9f03f07d9e4bc4f1b709855d) --- src/art/megatron/flex_attn/compiled.py | 17 ++++++++++++++++ tests/unit/test_flex_deterministic.py | 28 ++++++++++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/src/art/megatron/flex_attn/compiled.py b/src/art/megatron/flex_attn/compiled.py index ec7caede9..357b775d0 100644 --- a/src/art/megatron/flex_attn/compiled.py +++ b/src/art/megatron/flex_attn/compiled.py @@ -113,6 +113,20 @@ def flash_sparse_block_size_for_head_dim( return (128, 64) +def _require_compiled_deterministic_flash(backend: object) -> None: + if ( + backend == "FLASH" + and torch.are_deterministic_algorithms_enabled() + and not torch.compiler.is_compiling() + ): + raise RuntimeError( + "Deterministic attention requires compiled FLASH execution; " + "the uncompiled fallback bypasses its backward lowering and uses " + "a different LSE convention. Start a fresh trainer or increase " + "torch._dynamo.config.recompile_limit before compiling." + ) + + def _forced_flex_attention_dense( q, k, @@ -123,6 +137,7 @@ def _forced_flex_attention_dense( enable_gqa, return_aux: AuxRequest | None = None, ): + _require_compiled_deterministic_flash(_FORCED_FLEX_BACKEND) return flex_attention( q, k, @@ -149,6 +164,7 @@ def _sparse_flex_attention_with_options(kernel_options: FlexKernelOptions) -> An resolved_options["OUTPUT_MAX"] = False def _flex_attention(q, k, v, *, block_mask, scale, enable_gqa): + _require_compiled_deterministic_flash(resolved_options.get("BACKEND")) del enable_gqa for tensor in (q, k, v): torch._dynamo.mark_static(tensor, -3) @@ -214,6 +230,7 @@ def _flex_attention( enable_gqa, return_aux: AuxRequest | None = None, ): + _require_compiled_deterministic_flash(kernel_options.get("BACKEND")) return flex_attention( q, k, diff --git a/tests/unit/test_flex_deterministic.py b/tests/unit/test_flex_deterministic.py index cc31e3b5b..1560b832b 100644 --- a/tests/unit/test_flex_deterministic.py +++ b/tests/unit/test_flex_deterministic.py @@ -25,3 +25,31 @@ def test_determinism_preserves_flash_mask_geometry(monkeypatch): assert compiled.flash_sparse_block_size_for_head_dim( head_dim=256, head_dim_v=256, device=torch.device("cuda") ) == (128, 64) + + +@pytest.mark.parametrize("deterministic", [False, True]) +def test_deterministic_flash_rejects_uncompiled_cache_fallback( + monkeypatch, deterministic +): + original = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + monkeypatch.setattr(compiled, "flex_attention", lambda q, *args, **kwargs: q + 1) + torch._dynamo.reset() + try: + torch.use_deterministic_algorithms(deterministic) + attention = torch.compile( + compiled._forced_flex_attention_dense, backend="eager", dynamic=False + ) + with torch._dynamo.config.patch(recompile_limit=1): + options = dict(block_mask=None, scale=1.0, enable_gqa=False) + q = torch.zeros(2) + assert torch.equal(attention(q, q, q, **options), q + 1) + q = torch.zeros(3) + if deterministic: + with pytest.raises(RuntimeError, match="compiled FLASH"): + attention(q, q, q, **options) + else: + assert torch.equal(attention(q, q, q, **options), q + 1) + finally: + torch._dynamo.reset() + torch.use_deterministic_algorithms(original, warn_only=warn_only) From fab384edee89850310f215f0169b6efbd65a3e88 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 22 Oct 2026 04:07:16 +0000 Subject: [PATCH 12/12] ci: run Megatron-dependent MoE wave unit tests in the Megatron step The router, expert-plan and GDN capture tests import Megatron, which the backend unit-test environment lacks. Run those three files with the other Megatron lightweight tests and exclude them from the backend run. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/prek.yml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml index b270b0645..e4f55032c 100644 --- a/.github/workflows/prek.yml +++ b/.github/workflows/prek.yml @@ -252,6 +252,9 @@ jobs: tests/unit/test_trainer_rank_split.py \ tests/unit/test_megatron_compile_garbage.py \ tests/unit/test_canonical_numerics.py \ + tests/unit/test_segmented_expert_lora.py \ + tests/unit/test_segmented_gdn.py \ + tests/unit/test_segmented_router.py \ tests/unit/test_lora_layouts.py \ tests/unit/test_trainer_rank_forward_profile.py \ tests/unit/test_trainer_rank_cache_recovery.py::test_dense_cp_exact_demand_fits_after_recovery \ @@ -302,6 +305,9 @@ jobs: --ignore=tests/unit/test_trainer_rank_head_memory.py \ --ignore=tests/unit/test_grouped_planner_replay.py \ --ignore=tests/unit/test_canonical_numerics.py \ + --ignore=tests/unit/test_segmented_expert_lora.py \ + --ignore=tests/unit/test_segmented_gdn.py \ + --ignore=tests/unit/test_segmented_router.py \ --ignore=tests/unit/test_lora_layouts.py \ --ignore=tests/unit/test_trainer_rank_forward_profile.py \ --ignore=tests/unit/test_hybridep_planner_replay.py \