feat: support opt-in multi-adapter MoE LoRA waves

.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 \

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.

src/art/megatron/canonical/layout.py

@@ -93,23 +93,28 @@ 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.
 
     Deterministic algorithms and disabled debug fills are PROCESS-WIDE until
     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}")
     if numerical_policy == "canonical" and layout is None:
         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:

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
@@ -109,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,
@@ -119,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,
@@ -145,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)
@@ -210,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,

src/art/megatron/flex_attn/flash_dlse_patch.py

@@ -8,13 +8,15 @@
 from __future__ import annotations
 
 import inspect
+import math
 from typing import Any, cast
 
 import torch
 
 _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

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,23 @@ 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()
+    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:
         _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,18 +1107,43 @@ 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:
                     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
                     )
+            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 +1403,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 +1427,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")