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 \ 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..357b775d0 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 @@ -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, 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/src/art/megatron/lora.py b/src/art/megatron/lora.py index 413079960..08992bcfa 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,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,11 +1107,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 +1131,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 +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") diff --git a/src/art/megatron/segmented_expert_lora.py b/src/art/megatron/segmented_expert_lora.py new file mode 100644 index 000000000..897bb8a3a --- /dev/null +++ b/src/art/megatron/segmented_expert_lora.py @@ -0,0 +1,482 @@ +"""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 _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 + 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/src/art/megatron/segmented_gdn.py b/src/art/megatron/segmented_gdn.py new file mode 100644 index 000000000..f90352ff6 --- /dev/null +++ b/src/art/megatron/segmented_gdn.py @@ -0,0 +1,94 @@ +"""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, wraps +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(getattr(ctx, "_art_native_gdn_reductions", False)) + 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) + + +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/lora/test_segmented_expert_lora.py b/tests/integration/megatron/lora/test_segmented_expert_lora.py new file mode 100644 index 000000000..548c72c81 --- /dev/null +++ b/tests/integration/megatron/lora/test_segmented_expert_lora.py @@ -0,0 +1,324 @@ +"""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) + assert x.grad is not None + 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) + + +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): + 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/integration/megatron/test_native_determinism.py b/tests/integration/megatron/test_native_determinism.py new file mode 100644 index 000000000..718dda5a2 --- /dev/null +++ b/tests/integration/megatron/test_native_determinism.py @@ -0,0 +1,121 @@ +"""Deterministic backward preserves FLASH forward, including natural-log LSE.""" + +import pytest +import torch + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9, + reason="SM90 FLASH attention required", +) + + +@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)] +) +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/integration/megatron/test_segmented_gdn.py b/tests/integration/megatron/test_segmented_gdn.py new file mode 100644 index 000000000..b430c2df5 --- /dev/null +++ b/tests/integration/megatron/test_segmented_gdn.py @@ -0,0 +1,85 @@ +"""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]) +@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 + + 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} + + 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] + + 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) + 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_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..1560b832b --- /dev/null +++ b/tests/unit/test_flex_deterministic.py @@ -0,0 +1,55 @@ +"""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) + + +@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) diff --git a/tests/unit/test_segmented_expert_lora.py b/tests/unit/test_segmented_expert_lora.py new file mode 100644 index 000000000..6c04a3c9b --- /dev/null +++ b/tests/unit/test_segmented_expert_lora.py @@ -0,0 +1,311 @@ +"""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)] + + +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): + 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_gdn.py b/tests/unit/test_segmented_gdn.py new file mode 100644 index 000000000..f66e2bb26 --- /dev/null +++ b/tests/unit/test_segmented_gdn.py @@ -0,0 +1,100 @@ +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() + + +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() + + +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) 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)