diff --git a/.github/workflows/pr-test-cpu.yml b/.github/workflows/pr-test-cpu.yml index 50074134..8e4e10b4 100644 --- a/.github/workflows/pr-test-cpu.yml +++ b/.github/workflows/pr-test-cpu.yml @@ -78,6 +78,7 @@ jobs: tests/distributed/test_rope_class_b_fsdp_transport.py tests/distributed/test_torch_parallelize_policies.py tests/distributed/test_parallel_plan_meta_slice.py + tests/distributed/test_mixed_dtype_fsdp_split.py # Context/sequence parallel, pipeline parallel, and the core # parallel-state plumbing every other dimension builds on. @@ -93,6 +94,7 @@ jobs: tests/distributed/test_dsv4_exact_cp_attention_layout.py tests/distributed/test_sync_padding.py tests/distributed/test_pipeline_parallel.py + tests/distributed/test_pipeline_model_copy.py tests/distributed/test_pp_*.py tests/distributed/test_parallel_state.py tests/distributed/test_loss_metric_reductions.py diff --git a/README.md b/README.md index 4a276494..bb541b1b 100644 --- a/README.md +++ b/README.md @@ -84,7 +84,12 @@ The default install already includes `xorl-client` from its public repository. T pip install -e submodules/xorl-client ``` -The default profile is a single combined environment: `pyproject.toml` pins the PyTorch 2.11/CUDA 13 stack (Triton 3.6.0, FlashAttention 4) that the checked-in xorl-sglang revision — including its compiled `sglang-kernel` extension — is built against, so XoRL, xorl-client, and xorl-sglang all install into the one `uv sync` environment. The xorl-sglang submodule must be checked out for the install to resolve. +The default profile is a single combined environment: `pyproject.toml` pins +PyTorch 2.12.1+cu132, Triton 3.7.1, FlashAttention 4, and the matching DeepEP +and `sglang-kernel` release wheels used by the checked-in xorl-sglang revision. +XoRL, xorl-client, and xorl-sglang therefore install into one `uv sync` +environment. The xorl-sglang submodule must be checked out for the install to +resolve. See the [installation guide](https://togethercomputer.github.io/xorl/getting-started/installation/) for full setup including optional dependencies (DeepEP, Flash Attention). diff --git a/docs/src/content/docs/getting-started/installation.md b/docs/src/content/docs/getting-started/installation.md index 99e5d058..4ffe3828 100644 --- a/docs/src/content/docs/getting-started/installation.md +++ b/docs/src/content/docs/getting-started/installation.md @@ -13,9 +13,9 @@ XoRL ships a single combined dependency profile: | Manifest | PyTorch / CUDA runtime | Triton | Attention stack | Use it for | |---|---|---|---|---| -| `pyproject.toml` | 2.11.0 / CUDA 13 | 3.6.0 | FlashAttention 4 (`4.0.0b19`) | Local training, the XoRL training server, and the pinned xorl-sglang submodule, all in one environment | +| `pyproject.toml` | 2.12.1 / CUDA 13.2 | 3.7.1 | FlashAttention 4 (`4.0.0b19`) | Local training, the XoRL training server, the pinned xorl-sglang submodule, and DeepEP, all in one environment | -The PyTorch 2.11 pins match the checked-in xorl-sglang package metadata, so its compiled `sglang-kernel` extension loads in the same environment. Do not upgrade or mix the pinned Torch, Triton, or attention packages independently. +The PyTorch 2.12.1+cu132 pins match the checked-in xorl-sglang package metadata and the pinned DeepEP wheel. Do not upgrade or mix the pinned Torch, Triton, DeepEP, or attention packages independently. ## Clone the repo @@ -66,19 +66,40 @@ The default XoRL dependency set already installs `xorl-client` from its public r pip install -e submodules/xorl-client ``` -xorl-sglang installs into the same environment as XoRL: the default profile pins the PyTorch 2.11 stack its compiled `sglang-kernel` extension is built against, and the install steps above already include it (uv via `[tool.uv.sources]`, conda via the explicit editable install). +xorl-sglang installs into the same environment as XoRL: the default profile pins the PyTorch 2.12.1+cu132 stack its compiled extensions and DeepEP wheel target, and the install steps above already include it (uv via `[tool.uv.sources]`, conda via the explicit editable install). ## Verify Installation ```bash python -c "import torch, triton, xorl, sglang; print(torch.__version__, triton.__version__, xorl.__version__)" python -c "from flash_attn.cute import flash_attn_func; print('FlashAttention 4 ok')" -python -c "import sgl_kernel; print('sglang-kernel ok')" +python - <<'PY' +import torch +from sgl_kernel import moe_sum_reduce + +x = torch.arange(2 * 4 * 16, device="cuda", dtype=torch.bfloat16).reshape(2, 4, 16) +out = torch.empty((2, 16), device="cuda", dtype=torch.bfloat16) +moe_sum_reduce(x, out, 1.0) +torch.cuda.synchronize() +torch.testing.assert_close(out, x.float().sum(dim=1).to(torch.bfloat16), rtol=0, atol=0) +print("sglang-kernel MoE GPU operation ok") +PY ``` -## DeepEP Install (Optional) +## DeepEP Backend -DeepEP is a GPU-resident MoE dispatch backend. It uses high-speed GPU interconnects within a node and NVSHMEM/GPUDirect RDMA for supported multi-node deployments. It is only required when using `ep_dispatch: deepep`; the default `ep_dispatch: alltoall` works without it. Install it from [DeepSeek's DeepEP repository](https://github.com/deepseek-ai/DeepEP), then verify it separately with `python -c "import deep_ep; print('DeepEP ok')"`. +DeepEP is a GPU-resident MoE dispatch backend. It uses high-speed GPU interconnects within a node and NVSHMEM/GPUDirect RDMA for supported multi-node deployments. It is only required when using `ep_dispatch: deepep`; the default `ep_dispatch: alltoall` works without it. + +The default XoRL profile installs the pinned wheel for Python 3.12 and PyTorch 2.12.1+cu132. Verify it after the main installation: + +```bash +python -c "import deep_ep; print('DeepEP ok')" +``` + +The wheel URLs, source revisions, and hashes are recorded in +`vendor/deepep-release.lock.json` and +`vendor/sglang-kernel-release.lock.json`. Ordinary DeepEP retains stock reduction. +Setting `deepep_native_exact=true` selects deterministic hierarchical combine. ### Multi-node prerequisites diff --git a/docs/src/content/docs/moe/deepep.mdx b/docs/src/content/docs/moe/deepep.mdx index 5eb592df..5a3c95cc 100644 --- a/docs/src/content/docs/moe/deepep.mdx +++ b/docs/src/content/docs/moe/deepep.mdx @@ -14,9 +14,14 @@ DeepEP is a GPU-resident expert-parallel dispatch backend. It uses GPU interconn ## Installation ```bash -pip install deep_ep-*.whl # use a wheel compatible with the selected CUDA/PyTorch profile +pip install -e . -e "submodules/xorl-sglang/python[all]" ``` +The default dependency profile installs the wheel pinned by +`vendor/deepep-release.lock.json` for Python 3.12 and PyTorch 2.12.1+cu132. +Deterministic hierarchical combine remains an explicit exactness-mode +selection; the installed wheel retains stock DeepEP reduction as its default. + Verify: ```python import deep_ep diff --git a/pyproject.toml b/pyproject.toml index 1cf7006c..f96ebc45 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,14 +49,15 @@ dependencies = [ # P2P / Mooncake weight sync "mooncake-transfer-engine==0.3.9", "xorl-client @ git+https://github.com/togethercomputer/xorl-client.git@2a3a60a783c98e2a8ff722bad06dab18caee350c", - # PyTorch 2.11 with its CUDA 13 runtime dependencies. These versions match - # the checked-in xorl-sglang package metadata, so xorl, xorl-client, and the - # xorl-sglang submodule (including its compiled sglang-kernel extension) all - # run in this one environment. - "torch==2.11.0", - "torchvision==0.26.0", - "triton==3.6.0", + # Four-model zero-K3 qualification used this exact PyTorch/CUDA profile. + # Keep these pins aligned with the checked-in xorl-sglang metadata and the + # DeepEP release lock below. + "torch @ https://download.pytorch.org/whl/cu132/torch-2.12.1%2Bcu132-cp312-cp312-manylinux_2_28_x86_64.whl", + "torchvision @ https://download.pytorch.org/whl/cu132/torchvision-0.27.1%2Bcu132-cp312-cp312-manylinux_2_28_x86_64.whl", + "triton @ https://download.pytorch.org/whl/triton-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", "flash-attn-4==4.0.0b19", + "deep-ep @ https://github.com/togethercomputer/xorl-wheels/releases/download/deepep_sglang_kernel_torch212_cu132_sm90_723b8b3/deep_ep-1.2.1%2B65538ab.xorl.c85744ca7250-cp312-cp312-linux_x86_64.whl#sha256=f30485d585a4cd935f44ffadd276d7ea8603a919cff60302bbb40ded344cdccb", + "sglang-kernel @ https://github.com/togethercomputer/xorl-wheels/releases/download/deepep_sglang_kernel_torch212_cu132_sm90_723b8b3/sglang_kernel-0.4.5%2Bxorl.torch212.cu132.sm90-cp312-cp312-linux_x86_64.whl#sha256=f02e35414c18fd311ce29b7d27c9a07678f3932a9be3cabfec8d322f7b738f21", # Resolved to the checked-in fork by [tool.uv.sources] below. "sglang[all]", # TileLang kernels — used by DeepSeek-V4 sparse MLA, DSA indexer, FP8 diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index a4d9e046..783eac57 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -528,6 +528,13 @@ class ModelArguments: default=False, metadata={"help": "Enable async combine for DeepEP (overlap combine with next layer's compute)."}, ) + deepep_native_exact: bool = field( + default=False, + metadata={ + "help": "Use the versioned real-dispatch DeepEP exact program: BF16 rank leaves " + "and the deterministic hierarchical receiver fold. Requires a frozen router." + }, + ) alltoall_combine_hidden_chunk_size: int = field( default=0, metadata={ @@ -1235,7 +1242,10 @@ def moe_recomputed(self) -> bool: Used to decide whether routing replay is needed with EP: replay is only required when the MoE forward (including EP all-to-all) is recomputed. """ - return self.gradient_checkpointing_method in (None, "recompute_full_layer") + return self.enable_gradient_checkpointing and self.gradient_checkpointing_method in ( + None, + "recompute_full_layer", + ) enable_full_shard: bool = field( default=True, @@ -1865,6 +1875,14 @@ class LoRAArguments: default=False, metadata={"help": "Enable LoRA fine-tuning"}, ) + lora_serving_mode: Optional[Literal["merged", "separate"]] = field( + default=None, + metadata={ + "help": "Exact train/serve LoRA contract. 'merged' publishes W+sBA and " + "serves without an active adapter; 'separate' publishes A/B factors and " + "serves through active-LoRA kernels. Required with deepep_native_exact LoRA." + }, + ) lora_rank: int = field( default=16, metadata={"help": "LoRA rank"}, @@ -2378,11 +2396,32 @@ class Arguments: def __post_init__(self): from xorl.qarl import qarl_unsupported_scope_reason # noqa: PLC0415 + if self.model.deepep_native_exact and self.train.expert_parallel_size <= 1: + raise ValueError("model.deepep_native_exact requires train.expert_parallel_size > 1; EP1 bypasses DeepEP") + + if self.lora.lora_serving_mode not in {None, "merged", "separate"}: + raise ValueError("lora.lora_serving_mode must be 'merged' or 'separate'") + if self.model.deepep_native_exact and self.lora.enable_lora and self.lora.lora_serving_mode is None: + raise ValueError("Exact LoRA requires explicit lora.lora_serving_mode='merged' or 'separate'") + if not self.lora.enable_lora and self.lora.lora_serving_mode is not None: + raise ValueError("lora.lora_serving_mode requires lora.enable_lora=True") + + if ( + self.model.deepep_native_exact + and self.train.enable_gradient_checkpointing + and self.train.gradient_checkpointing_method == "recompute_full_layer" + ): + # Native exact owns live DeepEP dispatch/combine and recomputes its + # router independently. Checkpoint only the pre-dispatch trunk so + # backward never enters the process-wide routing-replay program. + self.train.gradient_checkpointing_method = "recompute_before_dispatch" + if self.train.enable_fp8_training and (self.lora.enable_lora or self.lora.enable_qlora): raise ValueError("enable_fp8_training is a full-weight mode and cannot be combined with LoRA or QLoRA") if self.train.enable_qarl and (self.lora.enable_lora or self.lora.enable_qlora): raise ValueError("enable_qarl is a full-weight mode and cannot be combined with LoRA or QLoRA") if self.lora.block_fp8_qlora_training: + exact_active_lora = self.model.ep_dispatch == "alltoall" requirements = { "lora.enable_lora": (self.lora.enable_lora, True), "lora.enable_qlora": (self.lora.enable_qlora, True), @@ -2390,8 +2429,10 @@ def __post_init__(self): "lora.quant_group_size": (self.lora.quant_group_size, 128), "lora.moe_hybrid_shared_lora": (self.lora.moe_hybrid_shared_lora, True), "model.moe_implementation": (self.model.moe_implementation, "triton"), - "model.ep_dispatch": (self.model.ep_dispatch, "deepep"), - "model.freeze_router": (self.model.freeze_router, True), + "model.ep_dispatch": ( + self.model.ep_dispatch, + "alltoall" if exact_active_lora else "deepep", + ), "model.merge_qkv": (self.model.merge_qkv, True), } mismatches = [ @@ -2401,6 +2442,13 @@ def __post_init__(self): ] if mismatches: raise ValueError("GLM-5.2 block-FP8 QLoRA rejects unsupported configuration: " + ", ".join(mismatches)) + if exact_active_lora: + if self.model.train_router == self.model.freeze_router: + raise ValueError( + "GLM-5.2 exact block-FP8 QLoRA requires train_router and freeze_router to be complementary" + ) + elif self.model.train_router or not self.model.freeze_router: + raise ValueError("GLM-5.2 non-exact block-FP8 QLoRA requires train_router=False and freeze_router=True") if self.lora.lora_target_modules is not None or self.lora.lora_target_manifest is not None: raise ValueError("GLM-5.2 block-FP8 QLoRA uses its complete deterministic target set") if self.lora.exclude_modules is not None: diff --git a/src/xorl/distributed/canonical_moe.py b/src/xorl/distributed/canonical_moe.py index d730ba76..ad879436 100644 --- a/src/xorl/distributed/canonical_moe.py +++ b/src/xorl/distributed/canonical_moe.py @@ -1069,12 +1069,9 @@ def _canonical_moe_reduce( element_size=contribution.tensor.element_size(), ) if transport is CanonicalMoETransport.PACKED_EP16_V2: - # Dense-v1 chunks the 16x-expanded owner slots to bound its allocation. - # Packed-v2's full-capacity send is already only one payload tensor, and - # coalescing here is required because GLM gathers CP shards in - # source-grouped order: an arbitrary subrange need not contain a - # balanced number of logical owners even though the complete capacity - # does. Keep one equal-split A2A over the complete logical row set. + # Dense-v1 chunks the contributor-expanded owner slots to bound its + # allocation. The sparse transports send at most one payload tensor; + # coalescing also preserves the complete owner-row layout. assert effective_chunk_rows == contribution.metadata.capacity tensor = _CanonicalMoEReduce.apply( diff --git a/src/xorl/distributed/moe/__init__.py b/src/xorl/distributed/moe/__init__.py index a3158036..ff5bf9de 100644 --- a/src/xorl/distributed/moe/__init__.py +++ b/src/xorl/distributed/moe/__init__.py @@ -34,7 +34,9 @@ def __getattr__(name): "DEEPEP_AVAILABLE", "token_pre_dispatch", "token_pre_dispatch_no_permute", + "token_pre_dispatch_native", "tokens_post_combine", + "tokens_post_combine_native", "get_default_buffer", "destroy_default_buffer", ): @@ -44,8 +46,10 @@ def __getattr__(name): destroy_default_buffer, get_default_buffer, token_pre_dispatch, + token_pre_dispatch_native, token_pre_dispatch_no_permute, tokens_post_combine, + tokens_post_combine_native, ) globals().update( @@ -54,7 +58,9 @@ def __getattr__(name): "DEEPEP_AVAILABLE": DEEPEP_AVAILABLE, "token_pre_dispatch": token_pre_dispatch, "token_pre_dispatch_no_permute": token_pre_dispatch_no_permute, + "token_pre_dispatch_native": token_pre_dispatch_native, "tokens_post_combine": tokens_post_combine, + "tokens_post_combine_native": tokens_post_combine_native, "get_default_buffer": get_default_buffer, "destroy_default_buffer": destroy_default_buffer, } @@ -76,7 +82,9 @@ def __getattr__(name): "DEEPEP_AVAILABLE", "token_pre_dispatch", "token_pre_dispatch_no_permute", + "token_pre_dispatch_native", "tokens_post_combine", + "tokens_post_combine_native", "get_default_buffer", "destroy_default_buffer", ] diff --git a/src/xorl/distributed/moe/deepep.py b/src/xorl/distributed/moe/deepep.py index e3509958..8a8c57ff 100644 --- a/src/xorl/distributed/moe/deepep.py +++ b/src/xorl/distributed/moe/deepep.py @@ -11,6 +11,7 @@ """ import os as _os +import time as _time from dataclasses import dataclass from typing import List, Optional, Tuple @@ -60,6 +61,7 @@ def preflight_internode_transport( hidden_dim: int, buffer_size_gb: float = 2.0, num_sms: int = 20, + buffer_hidden_bytes: Optional[int] = None, ) -> None: """Probe DeepEP's internode transport with a tiny dispatch+combine before training. @@ -97,7 +99,7 @@ def preflight_internode_transport( buffer = get_default_buffer(ep_group=ep_group, buffer_size_gb=buffer_size_gb, num_sms=num_sms) try: - buffer.init_buffer(hidden_bytes=hidden_dim * 2) + buffer.init_buffer(hidden_bytes=hidden_dim * 2 if buffer_hidden_bytes is None else int(buffer_hidden_bytes)) recv_x, _, _, _, handle = dispatch_no_grad(buffer, x, routing_weights, selected_experts, num_experts=world) out = combine_no_grad(buffer, recv_x.contiguous(), handle) except Exception as exc: @@ -131,9 +133,12 @@ def get_hidden_bytes(x: torch.Tensor) -> int: # Deferred combine sync # --------------------------------------------------------------------------- _pending_combine_event: Optional["EventOverlap"] = None +_ALLOW_UNSAFE_ASYNC_COMBINE: Optional[bool] = None def _allow_unsafe_async_combine() -> bool: + if _ALLOW_UNSAFE_ASYNC_COMBINE is not None: + return _ALLOW_UNSAFE_ASYNC_COMBINE return _os.environ.get("XORL_DEEPEP_UNSAFE_ASYNC_COMBINE", "").strip().lower() in {"1", "true", "yes", "on"} @@ -270,6 +275,52 @@ class DispatchContext: dtype: torch.dtype device: torch.device hidden_dim: int + # DeepEP normal mode returns rank-local expert ids (``-1`` marks a route + # owned by another rank) and the transported FP32 routing weights. The + # generic grouped-GEMM path only needs their expert-sorted projections, + # but DSV4's native Marlin runner consumes the receive layout directly so + # its local top-k reduction is the same numerical program as serving. + recv_topk_idx: Optional[torch.Tensor] = None + recv_topk_weights: Optional[torch.Tensor] = None + # Monotonic process-local identity for this original DeepEP handle. Every + # EP rank constructs handles in the same forward order; the optional trace + # uses this value to diagnose rank-asymmetric backward scheduling without + # adding a collective or touching the numerical program. + call_id: int = -1 + + +_deepep_call_counter = 0 + + +def _next_deepep_call_id() -> int: + global _deepep_call_counter + call_id = _deepep_call_counter + _deepep_call_counter += 1 + return call_id + + +def _trace_deepep_boundary( + call_id: int, + boundary: str, + state: str, + logical_rank: int = -1, + trace_label: str | None = None, +) -> None: + """Append one opt-in liveness record without synchronizing any device.""" + + trace_dir = _os.environ.get("XORL_DEEPEP_BOUNDARY_TRACE_DIR", "").strip() + if not trace_dir: + return + rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else -1 + _os.makedirs(trace_dir, exist_ok=True) + label = "none" if trace_label is None else str(trace_label).replace(" ", "_") + line = ( + f"{_time.monotonic_ns()} rank={rank} call_id={call_id} " + f"boundary={boundary} state={state} logical_rank={logical_rank} label={label}\n" + ) + with open(_os.path.join(trace_dir, f"rank{rank:05d}.log"), "a", encoding="utf-8") as handle: + handle.write(line) + handle.flush() # --------------------------------------------------------------------------- @@ -473,6 +524,8 @@ def forward( buffer: "DeepEPBuffer", num_experts: int, ): + call_id = _next_deepep_call_id() + _trace_deepep_boundary(call_id, "original_dispatch_forward", "enter") buffer.init_buffer(hidden_bytes=get_hidden_bytes(x)) topk_idx_deepep = topk_idx.to(deep_ep.topk_idx_t) topk_weights_f32 = topk_weights.to(torch.float32) @@ -523,6 +576,7 @@ def forward( ctx.buffer = buffer ctx.handle = handle ctx.input_dtype = x.dtype + ctx.call_id = call_id ctx.num_recv_tokens = num_recv_tokens ctx.hidden_dim = x.shape[1] @@ -535,7 +589,11 @@ def forward( dtype=x.dtype, device=x.device, hidden_dim=x.shape[1], + recv_topk_idx=recv_topk_idx, + recv_topk_weights=recv_topk_weights, + call_id=call_id, ) + _trace_deepep_boundary(call_id, "original_dispatch_forward", "exit") return expert_input, cumsum, dispatch_ctx @staticmethod @@ -590,7 +648,18 @@ def forward( topk_weights: torch.Tensor, buffer: "DeepEPBuffer", num_experts: int, + complete_backward_device_boundary: bool, + backward_trace_label: str | None, + backward_layer_dependency: torch.Tensor | None, + backward_shared_dependency: torch.Tensor | None, ): + call_id = _next_deepep_call_id() + _trace_deepep_boundary( + call_id, + "original_dispatch_forward", + "enter", + trace_label=backward_trace_label, + ) buffer.init_buffer(hidden_bytes=get_hidden_bytes(x)) topk_idx_deepep = topk_idx.to(deep_ep.topk_idx_t) topk_weights_f32 = topk_weights.to(torch.float32) @@ -635,6 +704,31 @@ def forward( ctx.buffer = buffer ctx.handle = handle ctx.input_dtype = x.dtype + ctx.call_id = call_id + ctx.complete_backward_device_boundary = bool(complete_backward_device_boundary) + ctx.backward_trace_label = backward_trace_label + # These two ignored forward operands are ordering edges, not value + # operands. Returning explicit zero gradients for them only after the + # terminal reverse-combine completes prevents their FSDP backward + # branches from entering c10d while DeepEP still owns the device. + ctx.backward_layer_dependency_meta = ( + None + if backward_layer_dependency is None + else ( + tuple(backward_layer_dependency.shape), + backward_layer_dependency.dtype, + backward_layer_dependency.device, + ) + ) + ctx.backward_shared_dependency_meta = ( + None + if backward_shared_dependency is None + else ( + tuple(backward_shared_dependency.shape), + backward_shared_dependency.dtype, + backward_shared_dependency.device, + ) + ) dispatch_ctx = DispatchContext( handle=handle, @@ -645,6 +739,15 @@ def forward( dtype=x.dtype, device=x.device, hidden_dim=x.shape[1], + recv_topk_idx=recv_topk_idx, + recv_topk_weights=recv_topk_weights, + call_id=call_id, + ) + _trace_deepep_boundary( + call_id, + "original_dispatch_forward", + "exit", + trace_label=backward_trace_label, ) return recv_x, cumsum, dispatch_ctx @@ -652,10 +755,16 @@ def forward( def backward(ctx, grad_recv_x, grad_cumsum, grad_dispatch_ctx): del grad_cumsum, grad_dispatch_ctx if grad_recv_x is None: - return None, None, None, None, None + return None, None, None, None, None, None, None, None, None buffer = ctx.buffer handle = ctx.handle + _trace_deepep_boundary( + ctx.call_id, + "input_reverse_combine", + "enter", + trace_label=ctx.backward_trace_label, + ) previous_event = EventOverlap(EventHandle()) grad_x, _, event = buffer.buffer.combine( x=grad_recv_x.contiguous(), @@ -667,10 +776,41 @@ def backward(ctx, grad_recv_x, grad_cumsum, grad_dispatch_ctx): ) event.current_stream_wait() grad_x.record_stream(torch.cuda.current_stream()) + _trace_deepep_boundary( + ctx.call_id, + "input_reverse_combine", + "api_return", + trace_label=ctx.backward_trace_label, + ) + if ctx.complete_backward_device_boundary: + # GLM's shared-expert FSDP root and transformer residual can use + # c10d private streams. Do not release either autograd edge while + # this layer's final normal-mode DeepEP reverse-combine can still + # be spinning on its own private stream: that cross-communicator + # overlap forms a device cycle despite identical host order on all + # ranks. This boundary changes scheduling only; no value is cast, + # communicated again, or added to a nonzero gradient. + torch.cuda.current_stream(grad_x.device).synchronize() + _trace_deepep_boundary( + ctx.call_id, + "input_reverse_combine", + "device_complete", + trace_label=ctx.backward_trace_label, + ) if grad_x.dtype != ctx.input_dtype: grad_x = grad_x.to(ctx.input_dtype) - return grad_x, None, None, None, None + dependency_grads = [] + for metadata in ( + getattr(ctx, "backward_layer_dependency_meta", None), + getattr(ctx, "backward_shared_dependency_meta", None), + ): + if metadata is None: + dependency_grads.append(None) + continue + shape, dtype, device = metadata + dependency_grads.append(torch.zeros(shape, dtype=dtype, device=device)) + return grad_x, None, None, None, None, None, None, *dependency_grads class _FusedUnpermuteAndCombine(torch.autograd.Function): @@ -791,6 +931,64 @@ def backward(ctx, grad_output): return grad_expert_output, None, None, None +class _FusedNativeReceiveCombine(torch.autograd.Function): + """Handle-based combine for a runner that consumes DeepEP receive rows. + + Unlike :class:`_FusedUnpermuteAndCombine`, this boundary performs no + trainer-owned scatter-add. The native MoE runner has already reduced the + local top-k slots for each received row, so DeepEP's original dispatch + handle is the complete inverse movement required by serving. + """ + + @staticmethod + def forward( + ctx, + recv_output: torch.Tensor, + buffer: "DeepEPBuffer", + dispatch_ctx: DispatchContext, + async_combine: bool, + ): + previous_event = EventOverlap(EventHandle()) + combined_x, _, event = buffer.buffer.combine( + x=recv_output.contiguous(), + handle=dispatch_ctx.handle, + config=buffer.combine_config, + previous_event=previous_event, + async_finish=True, + allocate_on_comm_stream=True, + ) + if async_combine: + _store_pending_event(event) + else: + event.current_stream_wait() + combined_x.record_stream(torch.cuda.current_stream()) + + ctx.buffer = buffer + ctx.handle = dispatch_ctx.handle + ctx.input_dtype = recv_output.dtype + return combined_x + + @staticmethod + def backward(ctx, grad_output): + if grad_output is None: + return None, None, None, None + + previous_event = EventOverlap(EventHandle()) + grad_recv, _, _, _, _, event = ctx.buffer.buffer.dispatch( + x=grad_output.contiguous(), + handle=ctx.handle, + config=ctx.buffer.dispatch_config, + previous_event=previous_event, + async_finish=True, + allocate_on_comm_stream=True, + ) + event.current_stream_wait() + grad_recv.record_stream(torch.cuda.current_stream()) + if grad_recv.dtype != ctx.input_dtype: + grad_recv = grad_recv.to(ctx.input_dtype) + return grad_recv, None, None, None + + # --------------------------------------------------------------------------- # No-grad dispatch/combine (for inference or profiling) # --------------------------------------------------------------------------- @@ -893,6 +1091,11 @@ def token_pre_dispatch_no_permute( selected_experts: torch.Tensor, num_experts: int, num_local_experts: int = 0, + *, + complete_backward_device_boundary: bool = False, + backward_trace_label: str | None = None, + backward_layer_dependency: torch.Tensor | None = None, + backward_shared_dependency: torch.Tensor | None = None, ) -> Tuple[torch.Tensor, torch.Tensor, DispatchContext]: """Dispatch tokens with DeepEP but leave them in recv order. @@ -907,10 +1110,51 @@ def token_pre_dispatch_no_permute( routing_weights, buffer, num_experts, + complete_backward_device_boundary, + backward_trace_label, + backward_layer_dependency, + backward_shared_dependency, ) return recv_x, cumsum, ctx +def token_pre_dispatch_native( + buffer: DeepEPBuffer, + hidden_states: torch.Tensor, + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, + num_experts: int, + *, + complete_backward_device_boundary: bool = False, + backward_trace_label: str | None = None, + backward_layer_dependency: torch.Tensor | None = None, + backward_shared_dependency: torch.Tensor | None = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, DispatchContext]: + """Dispatch without permutation and expose transported local top-k data. + + This is the native-runner boundary: expert ids and weights come from the + actual DeepEP receive, never from replayed router state. Routing weights + are intentionally non-differentiable here; the admitted DSV4 exact program + freezes its router while hidden-state and active-LoRA gradients cross the + dispatch/combine autograd boundaries. + """ + + recv_x, _cumsum, ctx = token_pre_dispatch_no_permute( + buffer=buffer, + hidden_states=hidden_states, + routing_weights=routing_weights, + selected_experts=selected_experts, + num_experts=num_experts, + complete_backward_device_boundary=complete_backward_device_boundary, + backward_trace_label=backward_trace_label, + backward_layer_dependency=backward_layer_dependency, + backward_shared_dependency=backward_shared_dependency, + ) + if ctx.recv_topk_idx is None or ctx.recv_topk_weights is None: + raise RuntimeError("DeepEP native dispatch did not preserve receive top-k metadata") + return recv_x, ctx.recv_topk_idx, ctx.recv_topk_weights, ctx + + def tokens_post_combine( buffer: DeepEPBuffer, expert_output: torch.Tensor, @@ -938,6 +1182,19 @@ def tokens_post_combine( return _FusedUnpermuteAndCombine.apply(expert_output, buffer, ctx, async_combine) +def tokens_post_combine_native( + buffer: DeepEPBuffer, + recv_output: torch.Tensor, + ctx: DispatchContext, + async_combine: bool = False, +) -> torch.Tensor: + """Combine native runner output in the original DeepEP receive layout.""" + + if async_combine and not _allow_unsafe_async_combine(): + async_combine = False + return _FusedNativeReceiveCombine.apply(recv_output, buffer, ctx, async_combine) + + # --------------------------------------------------------------------------- # Profiled versions # --------------------------------------------------------------------------- diff --git a/src/xorl/distributed/moe/deepep_native_exact.py b/src/xorl/distributed/moe/deepep_native_exact.py new file mode 100644 index 00000000..71f47b01 --- /dev/null +++ b/src/xorl/distributed/moe/deepep_native_exact.py @@ -0,0 +1,612 @@ +"""Native DeepEP transport with explicit BF16-leaf reduction contracts. + +This is the reusable post-expert half of the native zero-K3 program. The +original DeepEP dispatch handle is retained from the real top-k dispatch. The +stock fused ``no_combine=False`` MoE program applies routing weights and sums +the owner-local top-k slots in its normal FP32 accumulator, then stores one +BF16 rank leaf. Exact mode sends those leaves through DeepEP's one-call +deterministic hierarchical receiver tree. Rank-serial implementations belong +only in tests and benchmarks as executable oracles. + +The numerical boundaries are intentional: + +* expert kernels may use their normal internal accumulators; +* the fused local combine stores one BF16 rank leaf; +* every value crossing DeepEP is BF16; +* the selected fixed fold promotes BF16 leaves for its FP64 reduction nodes + and casts only at the specified BF16 leaf/consumer boundaries. + +This is not the older synthetic post-expert rank-leaf dispatcher. It uses the +handle produced by the actual pre-expert top-k dispatch and is consequently a +native DeepEP execution path. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Callable + +import torch +import torch.distributed as dist + + +DEEPEP_DETERMINISTIC_PROTOCOL = "deepep_deterministic_hierarchical_bf16_v2" +DEEPEP_LOW_LATENCY_DETERMINISTIC_PROTOCOL = DEEPEP_DETERMINISTIC_PROTOCOL +DEEPEP_DETERMINISTIC_SUPPORTED_EP_SIZES = frozenset({2, 4, 8, 16}) +logger = logging.getLogger(__name__) +_engagement_logged = False + + +class DeepEPNativeExactError(RuntimeError): + """The real DeepEP dispatch receipt violates the native exact contract.""" + + +def validate_native_combine_geometry(ep_size: int) -> None: + """Validate topology before dispatch can mutate DeepEP buffer state.""" + + if ep_size not in DEEPEP_DETERMINISTIC_SUPPORTED_EP_SIZES: + raise DeepEPNativeExactError( + "DeepEP deterministic combine supports EP sizes " + f"{sorted(DEEPEP_DETERMINISTIC_SUPPORTED_EP_SIZES)}, got EP{ep_size}" + ) + + +@dataclass(frozen=True) +class NativeDeepEPGeometry: + ep_size: int + ep_rank: int + hidden_size: int + + def __post_init__(self) -> None: + if self.ep_size <= 0: + raise DeepEPNativeExactError("native DeepEP requires a positive EP size") + if not 0 <= self.ep_rank < self.ep_size: + raise DeepEPNativeExactError(f"native DeepEP EP rank {self.ep_rank} is outside [0, {self.ep_size})") + if self.hidden_size <= 0: + raise DeepEPNativeExactError("native DeepEP requires a positive hidden size") + + @property + def wire_width(self) -> int: + return self.hidden_size + + @property + def wire_hidden_bytes(self) -> int: + # The wire type is always BF16, independent of the model parameter type. + return self.wire_width * torch.tensor([], dtype=torch.bfloat16).element_size() + + +def resolve_native_deepep_geometry(ep_group, hidden_size: int) -> NativeDeepEPGeometry: + """Resolve physical group rank as the immutable logical leaf ordinal.""" + + if not dist.is_initialized(): + raise DeepEPNativeExactError("native DeepEP requires initialized torch.distributed") + return NativeDeepEPGeometry( + ep_size=dist.get_world_size(ep_group), + ep_rank=dist.get_rank(ep_group), + hidden_size=int(hidden_size), + ) + + +def validate_native_receive_metadata( + recv_output: torch.Tensor, + dispatch_ctx, + *, + num_local_experts: int, +) -> None: + """Fail closed on the actual normal-mode DeepEP receive receipt. + + A received row exists only because at least one route belongs to this rank, + so every non-empty row must name a valid local expert. ``-1`` remains the + required marker for the other top-k slots. Empty receive batches are a + valid load-balancing outcome. + """ + + if recv_output.ndim != 2: + raise DeepEPNativeExactError( + f"native DeepEP runner output must be [recv_rows, hidden], got {tuple(recv_output.shape)}" + ) + if recv_output.dtype is not torch.bfloat16: + raise DeepEPNativeExactError( + f"native DeepEP rank leaves must be BF16 before communication, got {recv_output.dtype}" + ) + if not recv_output.is_contiguous(): + raise DeepEPNativeExactError("native DeepEP runner output must be contiguous") + if int(dispatch_ctx.num_recv_tokens) != recv_output.shape[0]: + raise DeepEPNativeExactError( + "native DeepEP runner row count does not match its dispatch handle: " + f"{recv_output.shape[0]} != {dispatch_ctx.num_recv_tokens}" + ) + if int(dispatch_ctx.hidden_dim) != recv_output.shape[1]: + raise DeepEPNativeExactError( + "native DeepEP runner hidden width does not match its dispatch handle: " + f"{recv_output.shape[1]} != {dispatch_ctx.hidden_dim}" + ) + if num_local_experts <= 0: + raise DeepEPNativeExactError("native DeepEP requires a positive local expert count") + + recv_ids = dispatch_ctx.recv_topk_idx + recv_weights = dispatch_ctx.recv_topk_weights + if recv_ids is None or recv_weights is None: + raise DeepEPNativeExactError("native DeepEP dispatch did not retain receive top-k metadata") + if recv_ids.ndim != 2 or recv_weights.shape != recv_ids.shape: + raise DeepEPNativeExactError("native DeepEP receive ids and weights must have the same [recv_rows, topk] shape") + if recv_ids.shape[0] != recv_output.shape[0]: + raise DeepEPNativeExactError("native DeepEP receive metadata row count changed after dispatch") + if recv_ids.dtype not in (torch.int32, torch.int64): + raise DeepEPNativeExactError(f"native DeepEP receive ids must be integral, got {recv_ids.dtype}") + if recv_weights.dtype is not torch.float32: + raise DeepEPNativeExactError( + f"native DeepEP receive routing weights must be FP32 metadata, got {recv_weights.dtype}" + ) + + valid = recv_ids >= 0 + if recv_ids.numel() and bool(torch.any(recv_ids < -1)): + raise DeepEPNativeExactError("native DeepEP receive ids contain a marker below -1") + if bool(torch.any(recv_ids[valid] >= num_local_experts)): + raise DeepEPNativeExactError("native DeepEP delivered a route outside this rank's local expert slice") + if recv_ids.shape[0] and bool(torch.any(~valid.any(dim=1))): + raise DeepEPNativeExactError("native DeepEP delivered a receive row with no local route") + if recv_weights.numel() and not bool(torch.isfinite(recv_weights).all()): + raise DeepEPNativeExactError("native DeepEP receive routing weights are not finite") + + +def adapt_native_runner_metadata( + recv_topk_ids: torch.Tensor, + recv_topk_weights: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Adapt real DeepEP metadata to the shared serving-runner ABI.""" + + if recv_topk_ids.ndim != 2 or recv_topk_weights.shape != recv_topk_ids.shape: + raise DeepEPNativeExactError("native DeepEP runner metadata must share [recv_rows, topk] shape") + if recv_topk_ids.dtype not in (torch.int32, torch.int64): + raise DeepEPNativeExactError(f"native DeepEP runner expert ids must be integral, got {recv_topk_ids.dtype}") + if recv_topk_weights.dtype is not torch.float32: + raise DeepEPNativeExactError(f"native DeepEP runner weights must remain FP32, got {recv_topk_weights.dtype}") + return ( + recv_topk_ids.to(torch.int32).contiguous(), + recv_topk_weights.contiguous(), + ) + + +def reduce_native_runner_routes_to_bf16( + route_output: torch.Tensor, + recv_topk_ids: torch.Tensor, + recv_topk_weights: torch.Tensor, +) -> torch.Tensor: + """Reduce unweighted rank-local routes in FP32, then store a BF16 leaf. + + This compatibility helper is for a runner that explicitly returns one BF16 + expert result per receive-row/top-k slot. The selected native exact program + instead uses the fused ``no_combine=False`` BF16 local leaf directly. + """ + + if route_output.ndim != 3: + raise DeepEPNativeExactError("native DeepEP no-combine runner output must be [recv_rows, topk, hidden]") + if route_output.dtype is not torch.bfloat16: + raise DeepEPNativeExactError(f"native DeepEP runner routes must be BF16, got {route_output.dtype}") + if recv_topk_ids.shape != route_output.shape[:2] or recv_topk_weights.shape != recv_topk_ids.shape: + raise DeepEPNativeExactError( + "native DeepEP runner routes and receive metadata have different row/top-k geometry" + ) + if recv_topk_ids.dtype is not torch.int32 or recv_topk_weights.dtype is not torch.float32: + raise DeepEPNativeExactError("native DeepEP runner reduction requires int32 ids and FP32 routing weights") + valid = recv_topk_ids >= 0 + weighted_fp32 = torch.where( + valid.unsqueeze(-1), + route_output.to(torch.float32) * recv_topk_weights.unsqueeze(-1), + torch.zeros((), dtype=torch.float32, device=route_output.device), + ) + return weighted_fp32.sum(dim=1).to(torch.bfloat16).contiguous() + + +def native_zero_row_runner_routes( + recv_hidden: torch.Tensor, + recv_topk_ids: torch.Tensor, +) -> torch.Tensor: + """Construct the no-combine runner result for an empty receive batch.""" + + if recv_hidden.ndim != 2 or recv_hidden.shape[0] != 0 or recv_hidden.dtype is not torch.bfloat16: + raise DeepEPNativeExactError("native DeepEP zero-row bypass requires empty BF16 receive rows") + if recv_topk_ids.ndim != 2 or recv_topk_ids.shape[0] != 0: + raise DeepEPNativeExactError("native DeepEP zero-row metadata must be empty [0, topk]") + return recv_hidden.new_empty((0, recv_topk_ids.shape[1], recv_hidden.shape[1])) + + +def native_exact_router_topk( + router_logits: torch.Tensor, + *, + top_k: int, + renormalize: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build independent fixed-order FP32 routing metadata for DeepEP.""" + + if router_logits.ndim != 2 or router_logits.dtype is not torch.float32: + raise DeepEPNativeExactError("native DeepEP router logits must be FP32 [tokens, experts] metadata") + if top_k <= 0 or top_k > router_logits.shape[1]: + raise DeepEPNativeExactError("native DeepEP top-k is outside the expert geometry") + + from xorl.ops.batch_invariant_ops import bi_router_topk_weights # noqa: PLC0415 + + scores = torch.softmax(router_logits, dim=1, dtype=torch.float32) + weights, expert_ids = torch.topk(scores, top_k, dim=-1) + weights = bi_router_topk_weights(weights, renormalize, torch.bfloat16) + return weights.to(torch.float32).contiguous(), expert_ids.contiguous() + + +def canonicalize_native_routing_metadata(routing_weights: torch.Tensor) -> torch.Tensor: + """Preserve routing coefficients in DeepEP's required FP32 metadata ABI. + + Routing coefficients are kernel metadata, not expert-value wire payloads. + Rounding an FP32 coefficient through BF16 here changes the fused + ``no_combine=False`` rank leaf before its declared BF16 storage boundary + and disagrees with serving, which supplies the original FP32 coefficient + to the same fused kernel. Expert outputs and rank leaves remain BF16. + """ + + if routing_weights.dtype not in (torch.bfloat16, torch.float32): + raise DeepEPNativeExactError( + f"native DeepEP routing coefficients must be BF16 or FP32, got {routing_weights.dtype}" + ) + if routing_weights.requires_grad: + raise DeepEPNativeExactError("native DeepEP v1 requires a frozen router and frozen routing coefficients") + return routing_weights.to(torch.float32).contiguous() + + +def _flatten_native_route_metadata( + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, + *, + row_count: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Flatten route metadata without an ambiguous ``reshape(0, -1)``. + + Idle/padded distributed ranks legitimately contribute zero token rows but + must still enter DeepEP's collectives. The top-k width is part of the + metadata contract, so preserve that known width explicitly instead of + asking PyTorch to infer it from zero elements. + """ + + if routing_weights.ndim < 2 or selected_experts.ndim < 2: + raise DeepEPNativeExactError("native DeepEP route metadata must end in an explicit top-k dimension") + top_k = int(routing_weights.shape[-1]) + if top_k <= 0 or int(selected_experts.shape[-1]) != top_k: + raise DeepEPNativeExactError("native DeepEP selected experts and routing weights have different top-k geometry") + expected = int(row_count) * top_k + if routing_weights.numel() != expected or selected_experts.numel() != expected: + raise DeepEPNativeExactError("native DeepEP route metadata does not cover every flattened token row") + return ( + routing_weights.reshape(row_count, top_k).contiguous(), + selected_experts.reshape(row_count, top_k).contiguous(), + ) + + +def native_dispatch_runner_combine( + hidden_states: torch.Tensor, + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, + *, + ep_group, + num_experts: int, + num_local_experts: int, + buffer_size_gb: float, + num_sms: int, + runner: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor], + backward_layer_dependency: torch.Tensor | None = None, + backward_shared_dependency: torch.Tensor | None = None, + backward_trace_label: str | None = None, + complete_backward_device_boundary: bool = False, +) -> torch.Tensor: + """Own the complete reusable real-dispatch native DeepEP program. + + Model and LoRA adapters supply only their fused ``no_combine=False`` local + expert runner. This shared layer owns value/metadata validation, real + top-k dispatch, runner ABI localization, BF16 rank leaves, handle-based + combines, FP64 folding, and the reverse collectives installed by the + combine autograd function. A route cube is rejected so the selected exact + program cannot silently fall back to the superseded external reduction. + """ + + from xorl.distributed.moe.deepep import ( # noqa: PLC0415 + get_default_buffer, + token_pre_dispatch_native, + ) + + if ep_group is None: + raise DeepEPNativeExactError("native DeepEP requires a real EP process group") + if hidden_states.ndim < 2: + raise DeepEPNativeExactError("native DeepEP hidden states must end in a hidden dimension") + if hidden_states.dtype is not torch.bfloat16: + raise DeepEPNativeExactError(f"native DeepEP dispatch values must be BF16, got {hidden_states.dtype}") + if routing_weights.dtype is not torch.float32: + raise DeepEPNativeExactError(f"native DeepEP routing metadata must be FP32, got {routing_weights.dtype}") + if routing_weights.requires_grad: + raise DeepEPNativeExactError("native DeepEP v1 requires frozen FP32 routing metadata") + original_shape = hidden_states.shape + hidden_flat = hidden_states.reshape(-1, hidden_states.shape[-1]).contiguous() + routing_flat, selected_flat = _flatten_native_route_metadata( + routing_weights, + selected_experts, + row_count=hidden_flat.shape[0], + ) + if selected_flat.shape != routing_flat.shape: + raise DeepEPNativeExactError("native DeepEP selected experts and routing weights have different top-k geometry") + + geometry = resolve_native_deepep_geometry(ep_group, hidden_flat.shape[1]) + validate_native_combine_geometry(geometry.ep_size) + if num_local_experts * geometry.ep_size != int(num_experts): + raise DeepEPNativeExactError( + "native DeepEP requires contiguous complete expert ownership: " + f"{num_local_experts} local * {geometry.ep_size} ranks != {num_experts} experts" + ) + buffer = get_default_buffer( + ep_group=ep_group, + buffer_size_gb=buffer_size_gb, + num_sms=num_sms, + ) + buffer.init_buffer(hidden_bytes=geometry.wire_hidden_bytes) + recv_hidden, recv_local_ids, recv_weights, dispatch_ctx = token_pre_dispatch_native( + buffer=buffer, + hidden_states=hidden_flat, + routing_weights=routing_flat, + selected_experts=selected_flat, + num_experts=num_experts, + complete_backward_device_boundary=complete_backward_device_boundary, + backward_trace_label=backward_trace_label, + backward_layer_dependency=backward_layer_dependency, + backward_shared_dependency=backward_shared_dependency, + ) + if recv_hidden.dtype is not torch.bfloat16: + raise DeepEPNativeExactError( + f"native DeepEP dispatch returned {recv_hidden.dtype}; communication must stay BF16" + ) + recv_local_ids, recv_weights = adapt_native_runner_metadata( + recv_local_ids, + recv_weights, + ) + runner_output = runner(recv_hidden, recv_weights, recv_local_ids) + recv_leaf = runner_output + if recv_leaf.shape != recv_hidden.shape or recv_leaf.dtype is not torch.bfloat16: + raise DeepEPNativeExactError( + "native DeepEP fused no_combine=False runner must return one BF16 leaf per receive row" + ) + folded = native_receive_combine_and_fold( + recv_leaf.contiguous(), + buffer=buffer, + dispatch_ctx=dispatch_ctx, + ep_group=ep_group, + num_local_experts=num_local_experts, + backward_layer_dependency=backward_layer_dependency, + backward_trace_label=backward_trace_label, + ) + return folded.reshape(original_shape) + + +def reduce_expert_rows_to_bf16_leaf(expert_output: torch.Tensor, dispatch_ctx) -> torch.Tensor: + """Reduce expert-order route rows locally, then store one BF16 rank leaf. + + This helper is for grouped expert runners that emit one row per valid + route. Native receive-order runners should pass their already-reduced + BF16 output directly to :func:`native_receive_combine_and_fold`. + """ + + if expert_output.ndim != 2: + raise DeepEPNativeExactError("expert-order native output must be two-dimensional") + if expert_output.shape[0] != dispatch_ctx.permuted_indices.numel(): + raise DeepEPNativeExactError("expert-order output does not cover every valid DeepEP route") + if expert_output.shape[1] != int(dispatch_ctx.hidden_dim): + raise DeepEPNativeExactError("expert-order output hidden width changed after DeepEP dispatch") + if dispatch_ctx.permuted_indices.numel() and ( + int(dispatch_ctx.permuted_indices.min()) < 0 + or int(dispatch_ctx.permuted_indices.max()) >= int(dispatch_ctx.num_recv_tokens) + ): + raise DeepEPNativeExactError("DeepEP expert-order row index is outside the receive batch") + + # Local route arithmetic is deliberately wider than the BF16 wire/storage + # boundary. Out-of-place index_add keeps the operation differentiable. + leaf_fp32 = torch.zeros( + (int(dispatch_ctx.num_recv_tokens), int(dispatch_ctx.hidden_dim)), + dtype=torch.float32, + device=expert_output.device, + ) + if expert_output.shape[0]: + leaf_fp32 = leaf_fp32.index_add( + 0, + dispatch_ctx.permuted_indices.to(torch.long), + expert_output.to(torch.float32), + ) + return leaf_fp32.to(torch.bfloat16).contiguous() + + +class _DeepEPDeterministicCombineBF16(torch.autograd.Function): + """One topology-generic deterministic combine and one reverse dispatch.""" + + @staticmethod + def forward( + ctx, + local_leaf: torch.Tensor, + buffer, + dispatch_ctx, + geometry: NativeDeepEPGeometry, + backward_layer_dependency: torch.Tensor | None, + backward_trace_label: str | None, + ): + del backward_layer_dependency + if local_leaf.dtype is not torch.bfloat16 or not local_leaf.is_contiguous(): + raise DeepEPNativeExactError("DeepEP deterministic combine requires contiguous BF16 payload") + if local_leaf.ndim != 2 or local_leaf.shape[1] != geometry.hidden_size: + raise DeepEPNativeExactError("DeepEP deterministic combine received the wrong local-leaf geometry") + try: + from deep_ep import ReductionMode # noqa: PLC0415 + from deep_ep.utils import EventHandle, EventOverlap # noqa: PLC0415 + except (ImportError, AttributeError) as exc: + raise DeepEPNativeExactError("installed DeepEP lacks ReductionMode.DETERMINISTIC") from exc + from xorl.distributed.moe.deepep import _trace_deepep_boundary # noqa: PLC0415 + + reduction_mode = getattr( + ReductionMode, + "DETERMINISTIC", + None, + ) + if reduction_mode is None: + raise DeepEPNativeExactError("installed DeepEP lacks ReductionMode.DETERMINISTIC") + + call_id = int(dispatch_ctx.call_id) + _trace_deepep_boundary( + call_id, + "deterministic_combine_forward", + "enter", + trace_label=backward_trace_label, + ) + previous_event = EventOverlap(EventHandle()) + combined, combined_weights, event = buffer.buffer.combine( + x=local_leaf, + handle=dispatch_ctx.handle, + config=buffer.combine_config, + reduction_mode=reduction_mode, + previous_event=previous_event, + async_finish=True, + allocate_on_comm_stream=True, + ) + event.current_stream_wait() + combined.record_stream(torch.cuda.current_stream()) + if combined_weights is not None: + raise DeepEPNativeExactError("DeepEP deterministic value combine unexpectedly returned routing metadata") + if combined.dtype is not torch.bfloat16: + raise DeepEPNativeExactError(f"DeepEP deterministic combine widened BF16 to {combined.dtype}") + + ctx.buffer = buffer + ctx.handle = dispatch_ctx.handle + ctx.geometry = geometry + ctx.call_id = call_id + ctx.backward_trace_label = backward_trace_label + _trace_deepep_boundary( + call_id, + "deterministic_combine_forward", + "exit", + trace_label=backward_trace_label, + ) + return combined + + @staticmethod + def backward(ctx, grad_output): + if grad_output is None: + return None, None, None, None, None, None + from deep_ep.utils import EventHandle, EventOverlap # noqa: PLC0415 + + from xorl.distributed.moe.deepep import _trace_deepep_boundary # noqa: PLC0415 + + geometry = ctx.geometry + if grad_output.ndim != 2 or grad_output.shape[1] != geometry.hidden_size: + raise DeepEPNativeExactError("DeepEP deterministic backward received the wrong output geometry") + _trace_deepep_boundary( + ctx.call_id, + "output_reverse_dispatch", + "enter", + trace_label=ctx.backward_trace_label, + ) + grad_wire = grad_output.to(torch.bfloat16).contiguous() + previous_event = EventOverlap(EventHandle()) + grad_recv, _, _, _, _, event = ctx.buffer.buffer.dispatch( + x=grad_wire, + handle=ctx.handle, + config=ctx.buffer.dispatch_config, + previous_event=previous_event, + async_finish=True, + allocate_on_comm_stream=True, + ) + event.current_stream_wait() + grad_recv.record_stream(torch.cuda.current_stream()) + if grad_recv.dtype is not torch.bfloat16: + raise DeepEPNativeExactError(f"DeepEP deterministic backward widened BF16 to {grad_recv.dtype}") + _trace_deepep_boundary( + ctx.call_id, + "output_reverse_dispatch", + "exit", + trace_label=ctx.backward_trace_label, + ) + return grad_recv, None, None, None, None, None + + +def native_receive_combine_and_fold( + recv_output: torch.Tensor, + *, + buffer, + dispatch_ctx, + ep_group, + num_local_experts: int, + backward_layer_dependency: torch.Tensor | None = None, + backward_trace_label: str | None = None, +) -> torch.Tensor: + """Transport native receive-order BF16 leaves with deterministic combine.""" + + geometry = resolve_native_deepep_geometry(ep_group, recv_output.shape[1]) + validate_native_combine_geometry(geometry.ep_size) + validate_native_receive_metadata( + recv_output, + dispatch_ctx, + num_local_experts=num_local_experts, + ) + combined = _DeepEPDeterministicCombineBF16.apply( + recv_output, + buffer, + dispatch_ctx, + geometry, + backward_layer_dependency, + backward_trace_label, + ) + + global _engagement_logged + if not _engagement_logged: + logger.info( + "Native DeepEP exact combine ENGAGED: protocol=%s ep_size=%d " + "wire_dtype=bf16 fold=%s wire_width=%d combine_calls=1 " + "backward_schedule=single_reverse_dispatch_v1", + DEEPEP_DETERMINISTIC_PROTOCOL, + geometry.ep_size, + "receiver_fp64_tree8_bf16_node_leaf_fp64_node_fold", + geometry.wire_width, + ) + _engagement_logged = True + return combined + + +def native_expert_combine_and_fold( + expert_output: torch.Tensor, + *, + buffer, + dispatch_ctx, + ep_group, + num_local_experts: int, +) -> torch.Tensor: + """Adapter for expert-order runners using deterministic combine.""" + + local_leaf = reduce_expert_rows_to_bf16_leaf(expert_output, dispatch_ctx) + return native_receive_combine_and_fold( + local_leaf, + buffer=buffer, + dispatch_ctx=dispatch_ctx, + ep_group=ep_group, + num_local_experts=num_local_experts, + ) + + +__all__ = [ + "DEEPEP_DETERMINISTIC_PROTOCOL", + "DEEPEP_LOW_LATENCY_DETERMINISTIC_PROTOCOL", + "DEEPEP_DETERMINISTIC_SUPPORTED_EP_SIZES", + "DeepEPNativeExactError", + "NativeDeepEPGeometry", + "adapt_native_runner_metadata", + "canonicalize_native_routing_metadata", + "native_dispatch_runner_combine", + "native_exact_router_topk", + "native_expert_combine_and_fold", + "native_receive_combine_and_fold", + "native_zero_row_runner_routes", + "reduce_expert_rows_to_bf16_leaf", + "reduce_native_runner_routes_to_bf16", + "resolve_native_deepep_geometry", + "validate_native_receive_metadata", + "validate_native_combine_geometry", +] diff --git a/src/xorl/distributed/pipeline_parallel.py b/src/xorl/distributed/pipeline_parallel.py index 8335b06c..c3baad0a 100644 --- a/src/xorl/distributed/pipeline_parallel.py +++ b/src/xorl/distributed/pipeline_parallel.py @@ -68,6 +68,37 @@ ] +def _deepcopy_pipeline_model(model: nn.Module) -> nn.Module: + """Clone a model while preserving its rank-local process-group handles. + + Exact attention and vocabulary-head modules bind their CP/TP groups during + construction. ``ProcessGroup`` objects are immutable rank-local runtime + handles and cannot be pickled, so a plain ``copy.deepcopy`` fails before + PP pruning. Seeding the deepcopy memo keeps only those handles shared; + parameters, buffers, modules, and ordinary Python state are still cloned. + """ + + memo: dict[int, object] = {} + for module in model.modules(): + for value in vars(module).values(): + if isinstance(value, torch.distributed.ProcessGroup): + memo[id(value)] = value + cloned = copy.deepcopy(model, memo) + + # ``nn.Parameter.__deepcopy__`` clones tensor storage and requires-grad but + # intentionally drops the parameter's Python ``__dict__``. XoRL records + # numerical/ownership contracts there (for example ``_keep_fp32`` and + # ``spec_info``), so silently losing it changes the PP program. Reattach + # the rank-local metadata to the already-cloned parameter objects by FQN. + source_parameters = dict(model.named_parameters(remove_duplicate=False)) + cloned_parameters = dict(cloned.named_parameters(remove_duplicate=False)) + if source_parameters.keys() != cloned_parameters.keys(): + raise RuntimeError("pipeline deepcopy changed the model parameter inventory") + for name, source_parameter in source_parameters.items(): + cloned_parameters[name].__dict__.update(source_parameter.__dict__) + return cloned + + def generate_llm_fqn_from_layer_ranges( layer_ranges, *, @@ -640,7 +671,7 @@ def pipeline_module_split( stages, model_parts = [], [] for stage_idx in stage_ids: module_names = module_names_per_stage[stage_idx] - model = copy.deepcopy(whole_model) + model = _deepcopy_pipeline_model(whole_model) fqns_to_keep = set(module_names) | always_keep # Recursive pruning handles nested HF model structures @@ -707,7 +738,29 @@ def build_pp_stage( dummy forward would otherwise run the intra-stage CP collectives, which deadlock against the cross-stage shape-exchange P2P. Output shapes are derivable from (mbs, seq_len, hidden/vocab), so the forward is unnecessary. + + Manual ``PipelineStage`` treats the supplied examples as the authoritative + static autograd metadata. Activation examples are usually created with + ``torch.empty(..., device="meta")``, whose default ``requires_grad=False`` + would make received activations leaf tensors that never accumulate the + gradient to send to the preceding stage. Mark the differentiable sides of + each inter-stage boundary explicitly; root token IDs remain integral and + non-differentiable. """ + if stage_index > 0 and input_args is not None: + input_args = tuple( + value.requires_grad_(True) + if isinstance(value, torch.Tensor) and (value.is_floating_point() or value.is_complex()) + else value + for value in input_args + ) + if stage_index < num_stages - 1 and output_args is not None: + output_args = tuple( + value.requires_grad_(True) + if isinstance(value, torch.Tensor) and (value.is_floating_point() or value.is_complex()) + else value + for value in output_args + ) return PipelineStage( model_part, stage_index, diff --git a/src/xorl/distributed/torch_parallelize.py b/src/xorl/distributed/torch_parallelize.py index 5211cd7c..cb6b1b05 100644 --- a/src/xorl/distributed/torch_parallelize.py +++ b/src/xorl/distributed/torch_parallelize.py @@ -170,6 +170,65 @@ def _fsdp_kwargs_for_module(fsdp_kwargs: dict, module: nn.Module) -> dict: return module_kwargs +def _fully_shard_declared_mixed_dtype_unit( + module: nn.Module, + *, + compute_kwargs: dict, + full_precision_kwargs: dict, +) -> list[nn.Module]: + """Split a declared composite into dtype-uniform FSDP groups. + + A declaring module owns only the named full-precision parameters directly; + its child modules own the lower-precision compute parameters. Grouping the + children in one FSDP call preserves communication coalescing, while wrapping + the parent afterwards claims only its direct FP32 leaves. Both groups stay + sharded and receive normal FSDP gradient reduction, and parameter FQNs do not + change. + + Returns representatives in forward-prefetch order. Modules passed together + to ``fully_shard`` share one FSDP state, so one child is sufficient to + identify the compute group. + """ + + declared_names = tuple(getattr(module, "fsdp_full_precision_parameter_names", ())) + if not declared_names: + return [] + if len(set(declared_names)) != len(declared_names): + raise ValueError( + f"{type(module).__name__}.fsdp_full_precision_parameter_names contains duplicates: {declared_names!r}" + ) + + direct_parameters = dict(module.named_parameters(recurse=False)) + missing = [name for name in declared_names if name not in direct_parameters] + undeclared = [name for name in direct_parameters if name not in declared_names] + if missing or undeclared: + raise ValueError( + f"{type(module).__name__} must own exactly its declared full-precision parameters directly; " + f"missing={missing}, undeclared={undeclared}" + ) + + full_precision_parameters = [direct_parameters[name] for name in declared_names] + full_precision_dtypes = {parameter.dtype for parameter in full_precision_parameters} + if full_precision_dtypes != {torch.float32}: + raise TypeError( + f"{type(module).__name__} declared full-precision parameters must be FP32; got {full_precision_dtypes}" + ) + + children = list(module.children()) + compute_parameters = [parameter for child in children for parameter in child.parameters()] + if not children or not compute_parameters: + raise ValueError(f"{type(module).__name__} declared a mixed-dtype FSDP split without child compute parameters") + compute_dtypes = {parameter.dtype for parameter in compute_parameters if parameter.is_floating_point()} + if compute_dtypes != {torch.bfloat16}: + raise TypeError( + f"{type(module).__name__} declared compute parameters must be uniformly BF16; got {compute_dtypes}" + ) + + fully_shard(children, **compute_kwargs) + fully_shard(module, **full_precision_kwargs) + return [module, children[0]] + + def _expert_fsdp_kwargs_for_module(expert_fsdp_kwargs: dict, experts_mod: nn.Module) -> dict: """Keep frozen byte-packed expert state sharded without numerically casting it.""" @@ -496,6 +555,12 @@ def _experts_shard_placement_fn(param): # (e.g., some models requires MoE TopK gate layer to have parameters in higher FP32 precision in forward). fsdp_wrapped_experts: List["nn.Module"] = [] preserve_dsv4_fp32 = bool(getattr(model.config, "_dsv4_flash_exact_mode", False)) + split_qwen35_full_weight_gdn = bool( + getattr(model.config, "_qwen35_exact_contract", False) + and not kwargs.get("enable_lora", False) + and not kwargs.get("enable_qlora", False) + ) + split_qwen35_gdn_count = 0 for layer_fqn, layer_mod, experts_mod in layer_pairs: # register all the FSDPModule inside this decoder layer for the convenience of manual prefetching configuration layer_mod._fsdp_modules = [] @@ -515,6 +580,32 @@ def _experts_shard_placement_fn(param): fully_shard(sub_mod, **fsdp_kwargs_without_mp) layer_mod._fsdp_modules.append(sub_mod) + if split_qwen35_full_weight_gdn: + declared_units = [ + submodule + for submodule in layer_mod.modules() + if getattr(submodule, "fsdp_full_precision_parameter_names", ()) + ] + if len(declared_units) > 1: + raise RuntimeError( + f"{layer_fqn} contains multiple declared mixed-dtype FSDP units; " + "nested or repeated declarations are not admitted" + ) + if declared_units: + compute_kwargs = dict(fsdp_kwargs) + if enable_mixed_precision: + compute_kwargs["mp_policy"] = decoder_mp_policy + layer_mod._fsdp_modules.extend( + reversed( + _fully_shard_declared_mixed_dtype_unit( + declared_units[0], + compute_kwargs=compute_kwargs, + full_precision_kwargs=fsdp_kwargs_without_mp, + ) + ) + ) + split_qwen35_gdn_count += 1 + # shard everything else in the decoder layer # If experts_mod has _skip_fsdp, exclude its params from the parent FSDP unit # (each EP rank has different local expert LoRA params — they should not be all-gathered globally) @@ -545,6 +636,17 @@ def _experts_shard_placement_fn(param): fully_shard(layer_mod, **layer_fsdp_kwargs) layer_mod._fsdp_modules.append(layer_mod) logger.debug_rank0(f"{layer_fqn=}, {layer_mod._fsdp_modules=}") + if split_qwen35_full_weight_gdn: + expected_gdn_count = sum(bool(getattr(layer_mod, "linear_attn", None)) for _, layer_mod, _ in layer_pairs) + if split_qwen35_gdn_count != expected_gdn_count: + raise RuntimeError( + "Exact full-weight Qwen3.5 FSDP dtype split did not cover every GDN layer: " + f"wrapped={split_qwen35_gdn_count}, expected={expected_gdn_count}" + ) + logger.info_rank0( + "Exact full-weight Qwen3.5 FSDP dtype split engaged: " + f"{split_qwen35_gdn_count} GDN units with BF16 compute parameters and FP32 A_log/dt_bias" + ) # Torchtitan optimization: group norm + lm_head into a single FSDP unit # with reshard_after_forward=False. When norm.forward() runs inside # the base model, FSDP all-gathers both norm and lm_head weights. @@ -866,6 +968,19 @@ def _build_ep_param_groups(model: "nn.Module") -> None: ) +def refresh_ep_param_groups(model: "nn.Module") -> None: + """Rebind EP parameter groups after a model-wide materialization. + + ``Module.to_empty`` may replace Parameter objects. Any optimizer or EP + ownership groups built before that transition consequently refer to stale + meta tensors even though ``model.parameters()`` is fully materialized. + Rebuild from the preserved FQN placement contract before constructing the + post-restore optimizer. + """ + + _build_ep_param_groups(model) + + def build_parallelize_model( model: "nn.Module", weights_path: Optional[str] = None, diff --git a/src/xorl/models/auto.py b/src/xorl/models/auto.py index 1085225a..a382ccf8 100644 --- a/src/xorl/models/auto.py +++ b/src/xorl/models/auto.py @@ -27,14 +27,18 @@ from .layers.normalization import set_rmsnorm_mode from .layers.rope import set_rope_class_b, set_rope_native from .loader import ModelLoader, get_loader +from .registry import ModelRegistry from .transformers.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config -from .transformers.deepseek_v3.support import validate_deepseek_v3_router_settings +from .transformers.deepseek_v3.support import ( + validate_deepseek_v3_router_settings, +) from .transformers.deepseek_v4.configuration_deepseek_v4 import DeepseekV4Config from .transformers.deepseek_v4.exact_contract import ( is_dsv4_flash_config, validate_dsv4_flash_adapter_program, validate_dsv4_flash_official_geometry, ) +from .transformers.deepseek_v4.moe_program import resolve_dsv4_moe_numerical_program from .transformers.glm4_moe.configuration_glm4_moe import Glm4MoeConfig from .transformers.glm5.configuration_glm5 import Glm5Config from .transformers.glm5.exact_lora_contract import glm52_exact_lora_scaling @@ -176,6 +180,29 @@ def _get_architectures(config: "PretrainedConfig") -> set[str]: return {architectures} +def resolve_deepep_native_exact_capability(config: "PretrainedConfig") -> Optional[dict[str, Any]]: + """Resolve the exact DeepEP contract declared by the registered model class.""" + + architectures = getattr(config, "architectures", None) + if isinstance(architectures, str): + architectures = [architectures] + for architecture in architectures or (): + if architecture not in ModelRegistry.supported_models: + continue + model_cls = ModelRegistry.get_model_cls_from_model_arch(architecture) + capability = getattr(model_cls, "deepep_native_exact_capability", None) + if capability is None: + return None + if ( + not capability.get("produces_local_leaf") + or capability.get("wire_dtype") != "bf16" + or not capability.get("uses_dispatch_handle") + ): + raise ValueError(f"{architecture} declares an incomplete native DeepEP local-leaf contract: {capability!r}") + return dict(capability) + return None + + def _is_gpt_oss_config(config: "PretrainedConfig") -> bool: return getattr(config, "model_type", None) == "gpt_oss" or "GptOssForCausalLM" in _get_architectures(config) @@ -495,14 +522,16 @@ def _validate_exact_qwen35_moe_program( moe_implementation: Optional[str], ep_dispatch: str, deepep_async_combine: bool, + deepep_native_exact: bool = False, ) -> None: if not (_is_exact_qwen35(config) and _is_qwen35_moe(config)): return incompatible = [] if moe_implementation not in (None, "triton"): incompatible.append(f"moe_implementation={moe_implementation!r} (requires 'triton')") - if ep_dispatch != "alltoall": - incompatible.append(f"ep_dispatch={ep_dispatch!r} (requires 'alltoall')") + required_dispatch = "deepep" if deepep_native_exact else "alltoall" + if ep_dispatch != required_dispatch: + incompatible.append(f"ep_dispatch={ep_dispatch!r} (requires {required_dispatch!r})") if deepep_async_combine: incompatible.append("deepep_async_combine=True (requires False)") if incompatible: @@ -800,6 +829,7 @@ def build_foundation_model( deepep_buffer_size_gb: float = 2.0, deepep_num_sms: int = 20, deepep_async_combine: bool = False, + deepep_native_exact: bool = False, alltoall_combine_hidden_chunk_size: int = 0, router_fp32: Optional[bool] = None, lm_head_fp32: Optional[bool] = None, @@ -816,6 +846,7 @@ def build_foundation_model( flash_attention_deterministic: bool = False, server_training: bool = False, enable_lora: bool = False, + lora_serving_mode: Optional[Literal["merged", "separate"]] = None, block_fp8_qlora_training: bool = False, glm52_fullparam_fp8_training: bool = False, lora_rank: Optional[int] = None, @@ -853,7 +884,9 @@ def build_foundation_model( raise ValueError( "glm52_fullparam_fp8_training and block_fp8_qlora_training are mutually exclusive training lanes" ) - exact_active_lora = bool(server_training and glm52_model and block_fp8_qlora_training and ep_dispatch == "alltoall") + exact_active_lora = bool( + server_training and glm52_model and block_fp8_qlora_training and ep_dispatch in {"alltoall", "deepep"} + ) if exact_active_lora: glm52_exact_lora_scaling(lora_rank, lora_alpha) # Training lanes select the same exact value family through their own @@ -883,6 +916,11 @@ def build_foundation_model( dsv4_flash_exact = bool(server_training and is_dsv4_flash_config(config)) config._dsv4_flash_exact_mode = dsv4_flash_exact config._dsv4_flash_exact_active_lora = bool(dsv4_flash_exact and enable_lora) + config._dsv4_moe_numerical_program = resolve_dsv4_moe_numerical_program( + exact=dsv4_flash_exact, + ep_dispatch=ep_dispatch, + deepep_native_exact=deepep_native_exact, + ) if dsv4_flash_exact: validate_dsv4_flash_official_geometry(config) if not enable_lora: @@ -906,14 +944,15 @@ def build_foundation_model( incompatible = [] if moe_implementation not in (None, "triton"): incompatible.append(f"moe_implementation={moe_implementation!r} (requires 'triton')") - if ep_dispatch != "alltoall": - incompatible.append(f"ep_dispatch={ep_dispatch!r} (requires 'alltoall')") if deepep_async_combine: incompatible.append("deepep_async_combine=True (requires False)") if incompatible: raise ValueError( "The DSV4-Flash exact RCA lane rejects incompatible trainer runtime choices: " + ", ".join(incompatible) ) + logger.info_rank0( + f"DSV4 MoE numerical program: {config._dsv4_moe_numerical_program} (ep_dispatch={ep_dispatch})" + ) # Family-neutral exact-contract keys, stamped once at model resolution so # downstream contract sites key off the resolved program rather than @@ -934,6 +973,7 @@ def build_foundation_model( moe_implementation=moe_implementation, ep_dispatch=ep_dispatch, deepep_async_combine=deepep_async_combine, + deepep_native_exact=deepep_native_exact, ) numerical_program = resolve_model_numerical_program( config, @@ -997,12 +1037,48 @@ def build_foundation_model( "Set train_router=False or use ep_dispatch='alltoall'." ) + deepep_capability = resolve_deepep_native_exact_capability(config) + if deepep_native_exact: + incompatible = [] + if deepep_capability is None: + incompatible.append("a model-declared BF16 native DeepEP local-leaf capability") + if ep_dispatch != "deepep": + incompatible.append("ep_dispatch='deepep'") + if train_router: + incompatible.append("train_router=False") + if deepep_async_combine: + incompatible.append("deepep_async_combine=False") + if torch_dtype != "bfloat16": + incompatible.append("torch_dtype='bfloat16'") + if moe_implementation not in (None, "triton"): + incompatible.append("moe_implementation='triton'") + if enable_lora and lora_serving_mode is not None: + supported_lora_modes = set(deepep_capability.get("lora_serving_modes", ())) if deepep_capability else set() + if lora_serving_mode not in supported_lora_modes: + incompatible.append( + f"lora_serving_mode in {sorted(supported_lora_modes)!r} (got {lora_serving_mode!r})" + ) + if incompatible: + raise ValueError( + "deepep_native_exact rejects this unqualified configuration; requires " + ", ".join(incompatible) + ) + + if lora_serving_mode not in {None, "merged", "separate"}: + raise ValueError("lora_serving_mode must be 'merged' or 'separate'") + if deepep_native_exact and enable_lora and lora_serving_mode is None: + raise ValueError("Exact LoRA requires explicit lora_serving_mode='merged' or 'separate'") + if not enable_lora and lora_serving_mode is not None: + raise ValueError("lora_serving_mode requires enable_lora=True") + config._ep_dispatch = ep_dispatch config.train_router = train_router config.record_routing_weights = record_routing_weights config._deepep_buffer_size_gb = deepep_buffer_size_gb config._deepep_num_sms = deepep_num_sms config._deepep_async_combine = deepep_async_combine + config._deepep_native_exact = bool(deepep_native_exact) + config._deepep_native_exact_capability = deepep_capability + config._lora_serving_mode = lora_serving_mode config._alltoall_combine_hidden_chunk_size = alltoall_combine_hidden_chunk_size config._router_fp32 = router_fp32 config._lm_head_fp32 = lm_head_fp32 @@ -1030,6 +1106,13 @@ def build_foundation_model( # otherwise wedges the gang at the first MoE dispatch, minutes later. ep_state = get_parallel_state() if ep_state.ep_enabled: + if deepep_native_exact: + supported_ep_sizes = set(deepep_capability.get("supported_ep_sizes", ())) + if ep_state.ep_size not in supported_ep_sizes: + raise ValueError( + "deepep_native_exact model capability supports EP sizes " + f"{sorted(supported_ep_sizes)}, got EP{ep_state.ep_size}" + ) from ..distributed.moe.deepep import preflight_internode_transport # noqa: PLC0415 preflight_internode_transport( @@ -1037,10 +1120,12 @@ def build_foundation_model( hidden_dim=getattr(config, "hidden_size", 0) or 2048, buffer_size_gb=deepep_buffer_size_gb, num_sms=deepep_num_sms, + buffer_hidden_bytes=((getattr(config, "hidden_size", 0) or 2048) * 2 if deepep_native_exact else None), ) logger.info_rank0( f"DeepEP dispatch enabled (buffer={deepep_buffer_size_gb} GB, " - f"num_sms={deepep_num_sms}, async_combine={deepep_async_combine})" + f"num_sms={deepep_num_sms}, async_combine={deepep_async_combine}, " + f"native_exact={deepep_native_exact})" ) # Validate attention implementation for packed sequences with FlashAttention kwargs diff --git a/src/xorl/models/layers/attention/multi_head_attention.py b/src/xorl/models/layers/attention/multi_head_attention.py index b158182a..540ae6ec 100644 --- a/src/xorl/models/layers/attention/multi_head_attention.py +++ b/src/xorl/models/layers/attention/multi_head_attention.py @@ -51,6 +51,13 @@ def __init__(self, config, layer_idx: int): # Overridable hooks # ------------------------------------------------------------------ # + def _capture_diagnostic_component(self, name: str, value: torch.Tensor) -> None: + """Expose exact-attention operands to the runner's cold diagnostic path.""" + + capture = self.__dict__.get("_diagnostic_capture_component") + if callable(capture): + capture(name, value) + def _init_sliding_window(self, config): """Override in subclasses for model-specific sliding window logic.""" return getattr(config, "sliding_window", None) @@ -80,6 +87,7 @@ def _project_qkv( input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) + self._capture_diagnostic_component("attention_input", hidden_states) if hasattr(self, "qkv_proj"): qkv = project_fused_linear_with_lora( self, @@ -88,6 +96,7 @@ def _project_qkv( projection_names=("q_proj", "k_proj", "v_proj"), projection_sizes=(self.q_dim, self.kv_dim, self.kv_dim), ) + self._capture_diagnostic_component("qkv", qkv) q, k, v = qkv.split([self.q_dim, self.kv_dim, self.kv_dim], dim=-1) else: q = self.q_proj(hidden_states) @@ -95,12 +104,19 @@ def _project_qkv( v = self.v_proj(hidden_states) q = q.view(hidden_shape) k = k.view(hidden_shape) + self._capture_diagnostic_component("q_pre_qk_norm", q) + self._capture_diagnostic_component("k_pre_qk_norm", k) if self._use_qk_norm: q = self.q_norm(q) k = self.k_norm(k) + self._capture_diagnostic_component("q_post_qk_norm", q) + self._capture_diagnostic_component("k_post_qk_norm", k) v = v.view(hidden_shape) + self._capture_diagnostic_component("v", v) cos, sin = position_embeddings + self._capture_diagnostic_component("rope_cos", cos) + self._capture_diagnostic_component("rope_sin", sin) q, k = apply_rotary_pos_emb( q, k, @@ -114,6 +130,8 @@ def _project_qkv( q = q.to(torch.bfloat16) k = k.to(torch.bfloat16) + self._capture_diagnostic_component("q", q) + self._capture_diagnostic_component("k", k) return q, k, v def _project_output(self, attn_output: torch.Tensor) -> torch.Tensor: @@ -121,8 +139,11 @@ def _project_output(self, attn_output: torch.Tensor) -> torch.Tensor: Override for different attention variants (e.g. Multi-head Latent Attention). """ + self._capture_diagnostic_component("attn_output", attn_output) attn_output = attn_output.reshape(*attn_output.shape[:-2], -1).contiguous() - return self.o_proj(attn_output) + output = self.o_proj(attn_output) + self._capture_diagnostic_component("o_proj_output", output) + return output def _append_past_key_values( self, diff --git a/src/xorl/models/layers/moe/dsv4_native_combine.py b/src/xorl/models/layers/moe/dsv4_native_combine.py index c2932cce..eabf796c 100644 --- a/src/xorl/models/layers/moe/dsv4_native_combine.py +++ b/src/xorl/models/layers/moe/dsv4_native_combine.py @@ -164,6 +164,7 @@ def exchange_variable_and_canonical_fold( if partial.shape[0] != total_rows: raise ValueError(f"Partial rows {partial.shape[0]} do not match live row total {total_rows}") local_rows = row_counts[source_ordinal] + exchanged = _AllToAll.apply( group, partial.contiguous(), diff --git a/src/xorl/models/layers/moe/ep_native_combine.py b/src/xorl/models/layers/moe/ep_native_combine.py index d67aa0db..5e2d5958 100644 --- a/src/xorl/models/layers/moe/ep_native_combine.py +++ b/src/xorl/models/layers/moe/ep_native_combine.py @@ -42,10 +42,15 @@ from __future__ import annotations +import logging + import torch import torch.distributed as dist +logger = logging.getLogger(__name__) + + def validate_native_ep_combine_size(ep_size: int) -> None: """Require one complete positive contributor group. @@ -66,11 +71,27 @@ class _AllGatherSumBackward(torch.autograd.Function): Every rank's partial depends on every gathered row, so the gather's backward must SUM the per-rank contributions (unlike the CP helpers, whose backward slices). The reduce-scatter is stock NCCL — backward carries no - bitwise contract. + bitwise contract. ``backward_dependency`` is deliberately unused in the + forward and has no gradient. When supplied, its autograd edge makes the + dependency's producer wait until this gather's complete consumer branch + (including the reduce-scatter) has completed in backward. GLM-5.2 uses + that ordering edge plus an explicit CUDA completion boundary to keep c10d + and DeepEP collectives in the same device order on every rank without + adding arithmetic to either branch. """ @staticmethod - def forward(ctx, x: torch.Tensor, group, padded_rows: int) -> torch.Tensor: + def forward( + ctx, + x: torch.Tensor, + group, + padded_rows: int, + backward_dependency: torch.Tensor | None, + backward_trace_label: str | None, + ) -> torch.Tensor: + ctx.complete_dependency_boundary = backward_dependency is not None + ctx.backward_trace_label = backward_trace_label + del backward_dependency ctx.group = group ctx.local_rows = x.shape[0] ctx.padded_rows = int(padded_rows) @@ -81,17 +102,61 @@ def forward(ctx, x: torch.Tensor, group, padded_rows: int) -> torch.Tensor: x = _pad_rows(x, ctx.padded_rows).contiguous() world_size = dist.get_world_size(group) out = torch.empty((world_size * ctx.padded_rows, *x.shape[1:]), dtype=x.dtype, device=x.device) + from xorl.distributed.moe.deepep import _trace_deepep_boundary # noqa: PLC0415 + + _trace_deepep_boundary( + -1, + "shared_all_gather_forward", + "enter", + trace_label=ctx.backward_trace_label, + ) dist.all_gather_into_tensor(out, x, group=group) + _trace_deepep_boundary( + -1, + "shared_all_gather_forward", + "exit", + trace_label=ctx.backward_trace_label, + ) return out @staticmethod def backward(ctx, grad_output): + from xorl.distributed.moe.deepep import _trace_deepep_boundary # noqa: PLC0415 + grad_output = grad_output.contiguous() grad_local_padded = torch.empty( (ctx.padded_rows, *grad_output.shape[1:]), dtype=grad_output.dtype, device=grad_output.device ) + _trace_deepep_boundary( + -1, + "shared_reduce_scatter_backward", + "enter", + trace_label=ctx.backward_trace_label, + ) dist.reduce_scatter_tensor(grad_local_padded, grad_output, op=dist.ReduceOp.SUM, group=ctx.group) - return grad_local_padded[: ctx.local_rows], None, None + _trace_deepep_boundary( + -1, + "shared_reduce_scatter_backward", + "api_return", + trace_label=ctx.backward_trace_label, + ) + if ctx.complete_dependency_boundary and grad_local_padded.is_cuda: + # ProcessGroupNCCL and normal-mode DeepEP use independent private + # communication streams. Merely ordering their autograd nodes (or + # recording a current-stream event) does not prove that the NCCL + # stream has completed before DeepEP starts spinning in its global + # barrier. Complete the shared branch here; only then does this + # Function return the undefined dependency gradient that makes the + # routed DeepEP producer eligible. This is synchronization only: + # no value is read, communicated again, or added to any gradient. + torch.cuda.current_stream(grad_local_padded.device).synchronize() + _trace_deepep_boundary( + -1, + "shared_reduce_scatter_backward", + "device_complete", + trace_label=ctx.backward_trace_label, + ) + return grad_local_padded[: ctx.local_rows], None, None, None, None def _pad_rows(x: torch.Tensor, padded_rows: int, *, value: int | float = 0) -> torch.Tensor: @@ -109,17 +174,35 @@ def max_rows_for_ep_combine(local_rows: int, device: torch.device, group) -> int return int(rows.item()) -def gather_tokens_for_ep_combine(x: torch.Tensor, group, padded_rows: int | None = None) -> torch.Tensor: +def gather_tokens_for_ep_combine( + x: torch.Tensor, + group, + padded_rows: int | None = None, + *, + backward_dependency: torch.Tensor | None = None, + backward_trace_label: str | None = None, +) -> torch.Tensor: """Autograd token gather with EP-uniform row padding. Dispatcher DP slices can contain different packed sequence lengths. Native EP still needs every expert rank to see every slice, so shorter ranks are padded to the negotiated maximum before the equal-count NCCL gather. The caller slices the final folded result back to its original row count. + + ``backward_dependency`` changes only backward scheduling/synchronization: + this gather's CUDA work completes before the dependency's producer + backward is eligible. No dependency value is read, communicated, or + added to the gradient. """ if padded_rows is None: padded_rows = max_rows_for_ep_combine(x.shape[0], x.device, group) - return _AllGatherSumBackward.apply(x, group, padded_rows) + return _AllGatherSumBackward.apply( + x, + group, + padded_rows, + backward_dependency, + backward_trace_label, + ) def gather_ids_for_ep_combine(ids: torch.Tensor, group, padded_rows: int | None = None) -> torch.Tensor: @@ -190,7 +273,11 @@ def sglang_fused_gate_sigmoid_mul_add( return _SGLangFusedGateSigmoidMulAdd.apply(hidden_states, gate_weight, shared_output, routed_output) -def exchange_and_canonical_fold(partial: torch.Tensor, group, ep_size: int) -> torch.Tensor: +def exchange_and_canonical_fold( + partial: torch.Tensor, + group, + ep_size: int, +) -> torch.Tensor: """RAW BF16 partial exchange + the serving canonical FP64 contributor fold. ``partial`` is this rank's [n*T, H] contribution for ALL gathered tokens. @@ -215,6 +302,7 @@ def exchange_and_canonical_fold(partial: torch.Tensor, group, ep_size: int) -> t raise TypeError("Native EP canonical combine requires BF16 partials") if partial.shape[0] % ep_size: raise ValueError("Native EP canonical combine rows must be divisible by EP size") + exchanged = _AllToAll.apply(group, partial.contiguous(), None, None) # [n*T, H], segment s from rank s rows = exchanged.shape[0] // ep_size logical_sources = exchanged.reshape(ep_size, rows, *exchanged.shape[1:]) diff --git a/src/xorl/models/layers/moe/experts.py b/src/xorl/models/layers/moe/experts.py index 0291c8d7..40dfe0dc 100644 --- a/src/xorl/models/layers/moe/experts.py +++ b/src/xorl/models/layers/moe/experts.py @@ -425,19 +425,25 @@ def forward( moe_index_compute, ) - output = _sglang_fused_experts_kernel_call( - hidden_flat, - gate_up_proj, - down_proj, - routing_flat, - selected_flat, - fused_experts_impl, - activation, - swiglu_limit, - gate_up_bias=None, - weight_cache=weight_cache, - filter_expert=filter_expert, - ) + if hidden_flat.shape[0] == 0 and filter_expert: + # DeepEP may deliver no rows to an expert rank for a small batch. + # The fused runner has no M=0 launch; the combined local leaf is + # the empty BF16 [rows, hidden] tensor. + output = hidden_flat.clone() + else: + output = _sglang_fused_experts_kernel_call( + hidden_flat, + gate_up_proj, + down_proj, + routing_flat, + selected_flat, + fused_experts_impl, + activation, + swiglu_limit, + gate_up_bias=None, + weight_cache=weight_cache, + filter_expert=filter_expert, + ) # Save the xorl scatter bookkeeping alongside the inputs (mirrors the # stock TritonMoeExpertsFunction contract): moe_index_compute uses # relaxed atomics, so the intra-expert row permutation is only @@ -526,7 +532,7 @@ def backward(ctx, grad_output): # Every slot masked (no local pairs): the kernel produced zeros, so # all grads are exactly zero (materialized to keep grad reduction # uniform across EP ranks). - return ( + grads = ( torch.zeros_like(hidden_states), # hidden_flat torch.zeros_like(gate_weights) if ctx.needs_input_grad[1] else None, # routing_flat None, # selected_flat @@ -540,6 +546,7 @@ def backward(ctx, grad_output): None, # weight_cache None, # filter_expert ) + return grads[: len(ctx.needs_input_grad)] # Recompute the cheap xorl-forward intermediates the stock backward # consumes (scatter + gate/up grouped GEMM) from the bookkeeping saved @@ -669,7 +676,7 @@ def backward(ctx, grad_output): num_slots = int(scatter_index.shape[1]) grad_hidden_states = grad_pair_rows.reshape(hidden_states.shape[0], num_slots, -1).sum(dim=1) - return ( + grads = ( grad_hidden_states, # hidden_flat grad_gate_weight, # routing_flat None, # selected_flat @@ -683,6 +690,7 @@ def backward(ctx, grad_output): None, # weight_cache None, # filter_expert ) + return grads[: len(ctx.needs_input_grad)] def _sglang_fused_experts_ep_kernel_call( @@ -1607,6 +1615,11 @@ def __init__( self.deepep_buffer_size_gb: float = 2.0 self.deepep_num_sms: int = 20 self.deepep_async_combine: bool = False + # Versioned real-dispatch exact path. Model adapters may only turn this + # on after construction; the shared layer owns dispatch receipt + # validation, local serving-kernel execution, BF16 transport, FP64 fold, + # and backward communication. + self.deepep_native_exact: bool = False self.fp8_training_enabled: bool = False self.fp8_training_grouped_backend: str = "triton_grouped" self.fp8_training_block_size: int = 128 @@ -2423,6 +2436,8 @@ def sglang_fused_experts_forward( hidden_states: torch.Tensor, routing_weights: torch.Tensor, selected_experts: torch.Tensor, + *, + local_expert_ids: bool = False, ) -> torch.Tensor: """K3 parity mode: run SGLang's serving MoE kernel (``fused_experts_impl``). @@ -2479,6 +2494,11 @@ def sglang_fused_experts_forward( weight_cache = {} self._sglang_fused_weight_cache = weight_cache + # Real DeepEP normal dispatch returns rank-local ids with ``-1`` in + # slots owned by other ranks. The same serving kernel accepts that + # layout through filter_expert=True; its backward must histogram only + # the local weight slice, not the model-global expert count. + kernel_num_experts = int(self.gate_up_proj.shape[0]) if local_expert_ids else int(self.num_experts) needs_grad = torch.is_grad_enabled() and ( hidden_flat.requires_grad or routing_flat.requires_grad @@ -2505,22 +2525,27 @@ def sglang_fused_experts_forward( activation, self.hidden_act, self.swiglu_limit, - self.num_experts, + kernel_num_experts, weight_cache, + local_expert_ids, ) else: - output = _sglang_fused_experts_kernel_call( - hidden_flat, - self.gate_up_proj, - self.down_proj, - routing_flat, - selected_flat, - fused_experts_impl, - activation, - self.swiglu_limit, - self.gate_up_bias, - weight_cache=weight_cache, - ) + if local_expert_ids and hidden_flat.shape[0] == 0: + output = hidden_flat.clone() + else: + output = _sglang_fused_experts_kernel_call( + hidden_flat, + self.gate_up_proj, + self.down_proj, + routing_flat, + selected_flat, + fused_experts_impl, + activation, + self.swiglu_limit, + self.gate_up_bias, + weight_cache=weight_cache, + filter_expert=local_expert_ids, + ) return output.reshape(original_shape) def sglang_fused_experts_auto_supported(self) -> bool: @@ -2974,6 +2999,14 @@ def _ep_forward( or ``"deepep"``). Compute backend by ``self.moe_implementation``. """ + if self.deepep_native_exact: + return self._deepep_native_exact_forward( + hidden_states, + routing_weights, + selected_experts, + parallel_state, + ) + if self.moe_implementation not in EP_EXPERT_COMPUTE: raise ValueError( f"moe_implementation={self.moe_implementation!r} does not support " @@ -3222,6 +3255,54 @@ def _ep_forward( ) return result + def _deepep_native_exact_forward( + self, + hidden_states: torch.Tensor, + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, + parallel_state, + ) -> torch.Tensor: + """Real top-k DeepEP dispatch + local serving runner + canonical fold. + + The model adapter only opts into this shared program. All numerical + and communication boundaries live here so Qwen/DeepSeek adapters do + not grow independent implementations. + """ + + from xorl.distributed.moe.deepep_native_exact import ( # noqa: PLC0415 + canonicalize_native_routing_metadata, + native_dispatch_runner_combine, + ) + + if self.ep_dispatch != "deepep": + raise RuntimeError("deepep_native_exact requires ep_dispatch='deepep'") + if self.deepep_async_combine: + raise RuntimeError("deepep_native_exact owns the immediate FP64 fold and rejects async combine") + if self.fp8_training_enabled: + raise RuntimeError("deepep_native_exact currently admits BF16 expert execution only") + if hidden_states.dtype is not torch.bfloat16: + raise RuntimeError(f"deepep_native_exact requires BF16 dispatch values, got {hidden_states.dtype}") + routing_weights = canonicalize_native_routing_metadata(routing_weights) + if self.gate_up_bias is not None or self.down_bias is not None or not self.gated: + raise NotImplementedError("deepep_native_exact currently admits bias-free gated experts") + num_local_experts = int(self.gate_up_proj.shape[0]) + return native_dispatch_runner_combine( + hidden_states, + routing_weights, + selected_experts, + ep_group=parallel_state.ep_group, + num_experts=int(self.num_experts), + num_local_experts=num_local_experts, + buffer_size_gb=self.deepep_buffer_size_gb, + num_sms=self.deepep_num_sms, + runner=lambda hidden, weights, ids: self.sglang_fused_experts_forward( + hidden, + weights, + ids, + local_expert_ids=True, + ), + ) + def _emit_deepep_parity_diagnostic( self, *, diff --git a/src/xorl/models/layers/moe/lora.py b/src/xorl/models/layers/moe/lora.py index eda4aff7..5b5dd00d 100644 --- a/src/xorl/models/layers/moe/lora.py +++ b/src/xorl/models/layers/moe/lora.py @@ -57,6 +57,72 @@ logger = logging.get_logger(__name__) +class _SglangNativeLoRAHooksTrainFunction(torch.autograd.Function): + """Keep the literal serving-hook value path separate from its VJP. + + The forward is SGLang's ordinary base/hook/activation/hook/combine + sequence. Backward reuses XoRL's existing trainable fused-MoE surrogate, + so changing the value path does not discard hidden, router, or adapter + gradients. + """ + + @staticmethod + def forward( + ctx, + hidden: torch.Tensor, + routing: torch.Tensor, + local_ids: torch.Tensor, + gate_A: torch.Tensor, + gate_B: torch.Tensor, + up_A: torch.Tensor, + up_B: torch.Tensor, + down_A: torch.Tensor, + down_B: torch.Tensor, + module, + ) -> torch.Tensor: + effective = tuple( + factor.to(torch.bfloat16).contiguous() for factor in (gate_A, gate_B, up_A, up_B, down_A, down_B) + ) + output = module._sglang_native_lora_hook_value( + hidden, + routing, + local_ids, + *effective, + ) + expected = (hidden.shape[0], module.hidden_dim) + if output.dtype is not torch.bfloat16 or tuple(output.shape) != expected: + raise RuntimeError( + "Native Qwen MoE-LoRA hook output contract mismatch: " + f"got {output.dtype} {tuple(output.shape)}, expected torch.bfloat16 {expected}" + ) + ctx.module = module + ctx.save_for_backward( + hidden.detach(), + routing.detach(), + local_ids, + gate_A, + gate_B, + up_A, + up_B, + down_A, + down_B, + ) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + hidden, routing, local_ids, *factors = ctx.saved_tensors + gradients = ctx.module._sglang_native_lora_hook_surrogate_vjp( + hidden, + routing, + local_ids, + *factors, + grad_output=grad_output, + needs_input_grad=ctx.needs_input_grad, + ) + return gradients[0], gradients[1], None, *gradients[2:], None + + @dataclass class MoELoRAConfig: """Configuration for MoE LoRA adapters.""" @@ -155,7 +221,9 @@ def __init__( swiglu_limit: float = 0.0, ): super().__init__() + self.num_global_experts = int(num_experts) self.num_experts = num_local_experts if num_local_experts is not None else num_experts + self.num_local_experts = int(self.num_experts) self.hidden_dim = hidden_dim self.intermediate_size = intermediate_size self.hidden_act = hidden_act @@ -245,6 +313,8 @@ def __init__( self.deepep_buffer_size_gb: float = 2.0 self.deepep_num_sms: int = 20 self.deepep_async_combine: bool = False + self.deepep_native_exact: bool = False + self.lora_serving_mode: str | None = None self.alltoall_combine_hidden_chunk_size: int = 0 @property @@ -626,15 +696,66 @@ def sglang_ep_native_routed_partial( routing_flat: torch.Tensor, local_ids: torch.Tensor, ) -> torch.Tensor: - """LoRA-aware local partial for the native EP canonical-combine lane. - - The native combine enters through ``Module.__call__`` so FSDP - materializes this rank's expert slice. Fold the active LoRA factors - into that slice with the same canonical arithmetic used by weight - sync, then run the masked serving kernel (``-1`` means non-local). - Forward bits therefore remain the serving bytes while the custom - folded-weight autograd sends gradients into the low-rank factors. + """Literal SGLang active-LoRA local leaf for native exact EP. + + The base weights stay unmerged. SGLang's ordinary MoE-LoRA hooks add + gate/up deltas before SwiGLU and down deltas before its fused + ``no_combine=False`` route fold. A zero LoRA-B therefore still + constructs and executes the same hooks as a live adapter; only their + numerical delta is zero. """ + from .experts import moe_sglang_fused_experts_weight_mode # noqa: PLC0415 + + if not lora_merged_forward_enabled(self): + raise NotImplementedError( + "Native EP combine on LoRA-adapted experts requires the exact LoRA execution contract" + ) + mode = getattr(self, "lora_serving_mode", None) + if mode == "merged": + return self._sglang_ep_native_merged_partial( + hidden_flat, + routing_flat, + local_ids, + ) + if mode != "separate": + raise ValueError(f"Unknown LoRA serving mode {mode!r}") + if not self.lora_config.hybrid_shared: + raise NotImplementedError("Native exact MoE-LoRA hooks currently require hybrid shared-outer factors") + if self.hidden_act not in {"silu", "gelu_tanh"} or self.swiglu_limit != 0.0: + raise NotImplementedError("LoRA native EP combine supports gated silu/gelu_tanh without swiglu_limit only") + if moe_sglang_fused_experts_weight_mode() == "cached": + raise NotImplementedError( + "Native exact MoE-LoRA hooks do not compose with WEIGHT_MODE=cached; use strided/transient." + ) + factors = tuple( + value + for projection in ("gate_proj", "up_proj", "down_proj") + for value in self._active_lora_views(projection) + ) + if self._merged_lora_needs_grad(hidden_flat, routing_flat): + return _SglangNativeLoRAHooksTrainFunction.apply( + hidden_flat, + routing_flat, + local_ids, + *factors, + self, + ) + effective = tuple(factor.to(torch.bfloat16).contiguous() for factor in factors) + return self._sglang_native_lora_hook_value( + hidden_flat, + routing_flat, + local_ids, + *effective, + ) + + def _sglang_ep_native_merged_partial( + self, + hidden_flat: torch.Tensor, + routing_flat: torch.Tensor, + local_ids: torch.Tensor, + ) -> torch.Tensor: + """Existing canonical folded-weight native local leaf.""" + from .experts import ( # noqa: PLC0415 MoEExperts, _sglang_fused_experts_kernel_call, @@ -642,21 +763,13 @@ def sglang_ep_native_routed_partial( moe_sglang_fused_experts_weight_mode, ) - if not lora_merged_forward_enabled(self): - raise NotImplementedError( - "Native EP combine on LoRA-adapted experts requires canonical merged-LoRA execution" - ) - if self.hidden_act not in {"silu", "gelu_tanh"} or self.swiglu_limit != 0.0: - raise NotImplementedError("LoRA native EP combine supports gated silu/gelu_tanh without swiglu_limit only") if moe_sglang_fused_experts_weight_mode() == "cached": raise NotImplementedError( "Canonical merged-LoRA execution does not compose with WEIGHT_MODE=cached; use strided/transient." ) - fused_experts_impl = MoEExperts._load_sglang_fused_experts_impl() activation = "gelu" if self.hidden_act == "gelu_tanh" else self.hidden_act e_local = int(self.gate_up_proj.shape[0]) - if self._merged_lora_needs_grad(hidden_flat, routing_flat): gate_up_w, down_w = self._merged_trainable_weights() return _SglangFusedExpertsTrainFunction.apply( @@ -673,8 +786,9 @@ def sglang_ep_native_routed_partial( None, True, ) - gate_up_f, down_f = self._merged_weights() + if hidden_flat.shape[0] == 0: + return hidden_flat.clone() return _sglang_fused_experts_kernel_call( hidden_flat, gate_up_f, @@ -689,6 +803,212 @@ def sglang_ep_native_routed_partial( filter_expert=True, ) + def _sglang_native_lora_physical_buffers( + self, + gate_A: torch.Tensor, + gate_B: torch.Tensor, + up_A: torch.Tensor, + up_B: torch.Tensor, + down_A: torch.Tensor, + down_B: torch.Tensor, + ) -> dict[str, torch.Tensor]: + """Build the one-adapter SGLang shared-outer memory-pool views.""" + + scaling = self._active_scaling() + return { + "gate_up_lora_a_weights": torch.cat((gate_A.transpose(1, 2), up_A.transpose(1, 2)), dim=1) + .unsqueeze(0) + .contiguous(), + "gate_up_lora_b_weights": (scaling * torch.cat((gate_B.transpose(1, 2), up_B.transpose(1, 2)), dim=1)) + .unsqueeze(0) + .to(torch.bfloat16) + .contiguous(), + "down_lora_a_weights": down_A.transpose(1, 2).unsqueeze(0).contiguous(), + "down_lora_b_weights": (scaling * down_B.transpose(1, 2)).unsqueeze(0).to(torch.bfloat16).contiguous(), + } + + def _sglang_native_lora_info(self, rows: int, physical: dict[str, torch.Tensor]): + try: + from sglang.srt.lora.lora_moe_runners import LoRAInfo # noqa: PLC0415 + except Exception as exc: + raise RuntimeError("Pinned SGLang MoE-LoRA hooks are required") from exc + + device = physical["gate_up_lora_a_weights"].device + return LoRAInfo( + **physical, + seg_indptr=torch.tensor([0, rows], dtype=torch.int32, device=device), + req_to_lora=torch.zeros(1, dtype=torch.int32, device=device), + lora_ranks=torch.tensor([self.active_r], dtype=torch.int32, device=device), + adapter_enabled=torch.ones(1, dtype=torch.int32, device=device), + token_lora_mapping=torch.zeros(rows, dtype=torch.int32, device=device), + max_lora_rank=self.active_r, + num_experts=self.num_global_experts, + has_active_lora=True, + single_adapter_id=0, + experts_shared_outer_loras=True, + cg_buffers=None, + fully_sharded=False, + tp_size=1, + tp_rank=0, + hidden_size=self.hidden_dim, + lora_use_virtual_experts=False, + ) + + def _sglang_native_lora_hook_value( + self, + hidden: torch.Tensor, + routing: torch.Tensor, + local_ids: torch.Tensor, + gate_A: torch.Tensor, + gate_B: torch.Tensor, + up_A: torch.Tensor, + up_B: torch.Tensor, + down_A: torch.Tensor, + down_B: torch.Tensor, + ) -> torch.Tensor: + """Run SGLang's literal base/hook/activation/hook/fold sequence.""" + + from .experts import MoEExperts # noqa: PLC0415 + + try: + from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import ( # noqa: PLC0415 + _fused_moe_kernel_sequence, + _prepare_fused_moe_run, + ) + from sglang.srt.lora.lora_moe_runners import build_lora_hooks # noqa: PLC0415 + except Exception as exc: + raise RuntimeError("Pinned SGLang MoE-LoRA hook runner is required") from exc + + if hidden.shape[0] == 0: + return hidden.clone() + physical = self._sglang_native_lora_physical_buffers(gate_A, gate_B, up_A, up_B, down_A, down_B) + MoEExperts._ensure_sglang_server_args() + w1 = self.gate_up_proj.transpose(1, 2) + w2 = self.down_proj.transpose(1, 2) + local_ids = local_ids.contiguous() + routing = routing.to(torch.float32).contiguous() + config, down_config, down_tma, sorted_ids, expert_ids, padded = _prepare_fused_moe_run( + hidden, + w1, + w2, + local_ids, + use_fp8_w8a8=False, + use_int8_w8a8=False, + use_int8_w8a16=False, + use_int4_w4a16=False, + per_channel_quant=False, + block_shape=None, + ) + hooks = build_lora_hooks( + hidden, + self._sglang_native_lora_info(hidden.shape[0], physical), + local_ids, + mul_routed_weight=True, + ) + return _fused_moe_kernel_sequence( + hidden, + w1, + w2, + routing, + local_ids, + sorted_ids, + expert_ids, + padded, + config, + down_config, + down_tma, + b1=None, + b2=None, + use_fp8_w8a8=False, + use_int8_w8a8=False, + use_int8_w8a16=False, + use_int4_w4a16=False, + per_channel_quant=False, + w1_scale=None, + w2_scale=None, + w1_zp=None, + w2_zp=None, + a1_scale=None, + a2_scale=None, + block_shape=None, + activation=("gelu" if self.hidden_act == "gelu_tanh" else self.hidden_act), + is_gated=True, + no_combine=False, + inplace=False, + apply_router_weight_on_input=False, + routed_scaling_factor=None, + gemm1_alpha=None, + gemm1_limit=None, + filter_expert=True, + hooks=hooks, + swiglu_limit=None, + gate_up_interleaved=False, + a1_q=None, + ) + + def _sglang_native_lora_hook_surrogate_vjp( + self, + hidden: torch.Tensor, + routing: torch.Tensor, + local_ids: torch.Tensor, + gate_A: torch.Tensor, + gate_B: torch.Tensor, + up_A: torch.Tensor, + up_B: torch.Tensor, + down_A: torch.Tensor, + down_B: torch.Tensor, + *, + grad_output: torch.Tensor, + needs_input_grad: tuple[bool, ...], + ) -> tuple[torch.Tensor | None, ...]: + """VJP through the existing folded-weight training surrogate.""" + + from .experts import ( # noqa: PLC0415 + MoEExperts, + _SglangFusedExpertsTrainFunction, + ) + + needs = (needs_input_grad[0], needs_input_grad[1], *needs_input_grad[3:9]) + if not any(needs): + return (None,) * 8 + values = (hidden, routing, gate_A, gate_B, up_A, up_B, down_A, down_B) + with torch.enable_grad(): + references = [value.detach().requires_grad_(needed) for value, needed in zip(values, needs, strict=True)] + hidden_ref, routing_ref, gate_A_ref, gate_B_ref, up_A_ref, up_B_ref, down_A_ref, down_B_ref = references + scaling = self._active_scaling() + inter = self.intermediate_size + gate = canonical_lora_fold_gkn(self.gate_up_proj[..., :inter], gate_A_ref, gate_B_ref, scaling) + up = canonical_lora_fold_gkn(self.gate_up_proj[..., inter:], up_A_ref, up_B_ref, scaling) + down = canonical_lora_fold_gkn(self.down_proj, down_A_ref, down_B_ref, scaling) + output = _SglangFusedExpertsTrainFunction.apply( + hidden_ref, + routing_ref, + local_ids, + torch.cat((gate, up), dim=-1), + down, + MoEExperts._load_sglang_fused_experts_impl(), + ("gelu" if self.hidden_act == "gelu_tanh" else self.hidden_act), + self.hidden_act, + self.swiglu_limit, + # ParallelPlan/FSDP may shard the expert tensors after this + # module is constructed. In that case the construction-time + # attribute still names the global expert count, while the + # physical weight leading dimension is the authoritative + # local grouped-GEMM count. + int(self.gate_up_proj.shape[0]), + None, + True, + ) + requested = [value for value, needed in zip(references, needs, strict=True) if needed] + computed = torch.autograd.grad( + output, + requested, + grad_outputs=grad_output, + allow_unused=False, + ) + iterator = iter(computed) + return tuple(next(iterator) if needed else None for needed in needs) + def forward( self, hidden_states: torch.Tensor, @@ -698,6 +1018,7 @@ def forward( sglang_ep_native_local_ids: torch.Tensor = None, dsv4_exact_native: bool = False, dsv4_exact_lora_live: bool = True, + dsv4_exact_return_routes: bool = False, ) -> torch.Tensor: """Forward pass with LoRA. @@ -717,6 +1038,7 @@ def forward( selected_experts, self, lora_live=dsv4_exact_lora_live, + return_routes=dsv4_exact_return_routes, ) if sglang_ep_native_local_ids is not None: @@ -777,6 +1099,30 @@ def _ep_forward( """ from .experts import _moe_sglang_fused_experts_env_state # noqa: PLC0415 + if self.deepep_native_exact: + from xorl.distributed.moe.deepep_native_exact import ( # noqa: PLC0415 + canonicalize_native_routing_metadata, + native_dispatch_runner_combine, + ) + + if self.ep_dispatch != "deepep" or self.deepep_async_combine: + raise RuntimeError("LoRA native DeepEP exact requires synchronous ep_dispatch='deepep'") + # The exact router stores coefficients at BF16 precision. DeepEP + # exposes its routing-metadata ABI as FP32, so widen those stored + # BF16 values without introducing any additional information. + routing_weights = canonicalize_native_routing_metadata(routing_weights) + return native_dispatch_runner_combine( + hidden_states, + routing_weights, + selected_experts, + ep_group=parallel_state.ep_group, + num_experts=self.num_global_experts, + num_local_experts=int(self.gate_up_proj.shape[0]), + buffer_size_gb=self.deepep_buffer_size_gb, + num_sms=self.deepep_num_sms, + runner=self.sglang_ep_native_routed_partial, + ) + explicit_sglang_fused = _moe_sglang_fused_experts_env_state() if explicit_sglang_fused is True: if not lora_merged_forward_enabled(self): @@ -969,6 +1315,8 @@ def from_module(cls, module: nn.Module, r: int, lora_alpha: int, **kwargs): lora_experts.deepep_buffer_size_gb = getattr(module, "deepep_buffer_size_gb", 2.0) lora_experts.deepep_num_sms = getattr(module, "deepep_num_sms", 20) lora_experts.deepep_async_combine = getattr(module, "deepep_async_combine", False) + lora_experts.deepep_native_exact = getattr(module, "deepep_native_exact", False) + lora_experts.lora_serving_mode = getattr(module, "lora_serving_mode", None) lora_experts.alltoall_combine_hidden_chunk_size = getattr(module, "alltoall_combine_hidden_chunk_size", 0) lora_experts.swiglu_limit = float(getattr(module, "swiglu_limit", 0.0)) if hasattr(module, "expert_lora_semantics"): @@ -1036,6 +1384,8 @@ def inject_lora_into_experts( lora_experts.deepep_buffer_size_gb = getattr(block.experts, "deepep_buffer_size_gb", 2.0) lora_experts.deepep_num_sms = getattr(block.experts, "deepep_num_sms", 20) lora_experts.deepep_async_combine = getattr(block.experts, "deepep_async_combine", False) + lora_experts.deepep_native_exact = getattr(block.experts, "deepep_native_exact", False) + lora_experts.lora_serving_mode = getattr(block.experts, "lora_serving_mode", None) lora_experts.alltoall_combine_hidden_chunk_size = getattr(block.experts, "alltoall_combine_hidden_chunk_size", 0) lora_experts.swiglu_limit = float(getattr(block.experts, "swiglu_limit", 0.0)) if hasattr(block.experts, "expert_lora_semantics"): diff --git a/src/xorl/models/layers/moe/moe_block.py b/src/xorl/models/layers/moe/moe_block.py index 3d0f95e9..b73c5eec 100644 --- a/src/xorl/models/layers/moe/moe_block.py +++ b/src/xorl/models/layers/moe/moe_block.py @@ -141,6 +141,7 @@ def __init__( activation_native: bool = False, swiglu_limit: float = 0.0, exact_batch_invariant_router: bool = False, + exact_router_weights_fp32: bool = False, ): super().__init__() self.num_experts = num_experts @@ -151,6 +152,7 @@ def __init__( self.train_router = train_router self.record_routing_weights = record_routing_weights self.swiglu_limit = float(swiglu_limit) + self.deepep_native_exact = False # Gate linear — directly on this module for checkpoint path ``mlp.gate.weight`` self.gate = nn.Linear(hidden_size, num_experts, bias=False) @@ -161,6 +163,7 @@ def __init__( top_k, norm_topk_prob, exact_batch_invariant=exact_batch_invariant_router, + exact_weights_fp32=exact_router_weights_fp32, ) self._exact_batch_invariant_router = exact_batch_invariant_router @@ -185,7 +188,11 @@ def __init__( def supports_routing_replay(self) -> bool: """Whether this block's route program can replay cached decisions.""" - return True + return not bool( + self.deepep_native_exact + or getattr(getattr(self, "config", None), "_deepep_native_exact", False) + or getattr(self.experts, "deepep_native_exact", False) + ) def _capture_diagnostic_component(self, name: str, tensor: torch.Tensor) -> None: capture = getattr(self, "_diagnostic_capture_component", None) @@ -335,12 +342,21 @@ def route(self, hidden_states: torch.Tensor): - router_logits: ``(num_tokens, num_experts)`` """ # Route (optionally upcast to fp32 for numerical alignment with SGLang) + native_deepep_exact = bool( + self.deepep_native_exact + or getattr(getattr(self, "config", None), "_deepep_native_exact", False) + or getattr(self.experts, "deepep_native_exact", False) + ) router_fp32 = ( getattr(self, "config", None) is not None and getattr(self.config, "_router_fp32", False) or _router_fp32_layers_enabled(getattr(self, "layer_idx", None)) ) - if self._exact_batch_invariant_router or _moe_bi_router_enabled(getattr(self, "config", None)): + if ( + native_deepep_exact + or self._exact_batch_invariant_router + or _moe_bi_router_enabled(getattr(self, "config", None)) + ): router_logits = self._bi_router_logits(hidden_states) elif router_fp32 and not hasattr(self.gate, "fp8_block_size"): router_logits = F.linear(hidden_states.float(), self.gate.weight.float()) @@ -361,6 +377,11 @@ def route(self, hidden_states: torch.Tensor): stage = get_replay_stage() replay = self._routing_replay + if native_deepep_exact and stage is not None: + raise RuntimeError( + "deepep_native_exact requires independently recomputed routing; routing replay is forbidden" + ) + if stage is not None and replay is not None: if stage == "record": # Determine expert selection without creating autograd nodes @@ -398,11 +419,26 @@ def route(self, hidden_states: torch.Tensor): if cached_weights is not None: routing_weights = cached_weights.to(hidden_states.dtype) else: - # No replay active: use standard router - routing_weights, selected_experts = self.router(router_logits, hidden_states.dtype) + if native_deepep_exact: + from xorl.distributed.moe.deepep_native_exact import ( # noqa: PLC0415 + native_exact_router_topk, + ) + + routing_weights, selected_experts = native_exact_router_topk( + router_logits, + top_k=self.top_k, + renormalize=self.router.norm_topk_prob, + ) + else: + # No replay active: use the model router. + routing_weights, selected_experts = self.router(router_logits, hidden_states.dtype) forced_selected_experts = getattr(self, "_diagnostic_forced_selected_experts", None) if forced_selected_experts is not None: + if native_deepep_exact: + raise RuntimeError( + "deepep_native_exact requires independently recomputed routing; forced routing is forbidden" + ) forced_selected_experts = forced_selected_experts.to( device=selected_experts.device, dtype=selected_experts.dtype, diff --git a/src/xorl/models/layers/moe/router.py b/src/xorl/models/layers/moe/router.py index b10bc38f..617f1716 100644 --- a/src/xorl/models/layers/moe/router.py +++ b/src/xorl/models/layers/moe/router.py @@ -128,6 +128,7 @@ def __init__( topk_method: str | None = None, routed_scaling_factor: float | None = None, exact_batch_invariant: bool = False, + exact_weights_fp32: bool = False, exact_sqrtsoftplus_serving: bool = False, ): super().__init__() @@ -152,6 +153,7 @@ def __init__( self.topk_method = topk_method self.routed_scaling_factor = routed_scaling_factor self._exact_batch_invariant = exact_batch_invariant + self._exact_weights_fp32 = exact_weights_fp32 self._exact_sqrtsoftplus_serving = exact_sqrtsoftplus_serving # Exact model programs are structural and must not be redirected by # process-wide diagnostic environment variables. @@ -251,7 +253,11 @@ def _forward_softmax(self, router_logits: torch.Tensor, input_dtype: torch.dtype # sum(dim=-1) reduction is build-dependent; see bi_router_topk_weights). from xorl.ops.batch_invariant_ops import bi_router_topk_weights # noqa: PLC0415 - routing_weights = bi_router_topk_weights(routing_weights, self.norm_topk_prob, input_dtype) + routing_weights = bi_router_topk_weights( + routing_weights, + self.norm_topk_prob, + torch.float32 if self._exact_weights_fp32 else input_dtype, + ) else: if self.norm_topk_prob: routing_weights /= routing_weights.sum(dim=-1, keepdim=True) diff --git a/src/xorl/models/transformers/deepseek_v4/modeling_deepseek_v4.py b/src/xorl/models/transformers/deepseek_v4/modeling_deepseek_v4.py index e4ba4858..1330cb51 100644 --- a/src/xorl/models/transformers/deepseek_v4/modeling_deepseek_v4.py +++ b/src/xorl/models/transformers/deepseek_v4/modeling_deepseek_v4.py @@ -29,6 +29,7 @@ import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F +from transformers.utils import logging from xorl.distributed.canonical_moe import LogicalRowOwnership from xorl.distributed.parallel_state import get_parallel_state @@ -42,6 +43,9 @@ from xorl.models.layers.normalization import RMSNorm from xorl.models.module_utils import DEFAULT_GRADIENT_CHECKPOINTING_METHOD from xorl.models.outputs import MoeCausalLMOutput, MoeModelOutput +from xorl.models.transformers.deepseek_v4.moe_program import ( + DSV4_DEEPEP_NATIVE_EXACT_V1, +) from xorl.ops.dsv4.attention_core import dense_attn_torch, sparse_attn_tilelang, sparse_attn_torch from xorl.ops.dsv4.compressor import DeepSeekV4Compressor from xorl.ops.dsv4.cp_utils import ( @@ -54,7 +58,7 @@ get_q_positions_for_cp, get_window_topk_idxs_cp, ) -from xorl.ops.dsv4.hyper_connection import DeepSeekV4HyperConnectionUtil +from xorl.ops.dsv4.hyper_connection import DeepSeekV4HyperConnectionUtil, ExactMhcReplaySegment from xorl.ops.dsv4.qat import fp8_simulate_qat from xorl.ops.dsv4.rope import apply_rotary_emb, wrapped_precompute_freqs_cis from xorl.ops.dsv4.utils import dsv4_kv_qat_enabled @@ -73,6 +77,124 @@ # value. _ATTN_IMPL_ENV = "XORL_DSV4_SPARSE_ATTN_IMPL" _ATTN_IMPL_CHOICES = {"tilelang", "sparse", "dense"} +logger = logging.get_logger(__name__) + + +def _build_serving_mhc_segments( + *, + compute_rows: int, + sample_lengths: list[int], + sampler_prefill_lengths: torch.Tensor | None, +) -> tuple[int, ...] | None: + """Map packed requests to serving's one-prefill-then-M=1 row program.""" + + if sampler_prefill_lengths is None: + return None + if sampler_prefill_lengths.ndim != 1 or sampler_prefill_lengths.dtype not in (torch.int32, torch.int64): + raise ValueError("DSV4 sampler_prefill_lengths must be a rank-one integer tensor") + prefill_lengths = [int(value) for value in sampler_prefill_lengths.detach().cpu().tolist()] + if not sample_lengths: + # Dispatcher-created exact-mode dummy ranks own no request rows. Their + # padded activations are excluded before attention/MoE ownership, so no + # serving segmentation is needed for their dead storage. + return None + if len(prefill_lengths) != len(sample_lengths): + raise ValueError( + "DSV4 serving MHC replay needs one prefill boundary per packed request: " + f"prefills={prefill_lengths} sample_lengths={sample_lengths}" + ) + live_rows = sum(sample_lengths) + if live_rows > compute_rows: + raise ValueError(f"DSV4 packed requests contain {live_rows} live rows but only {compute_rows} compute rows") + + segments: list[int] = [] + for sample_length, prefill_length in zip(sample_lengths, prefill_lengths, strict=True): + if sample_length <= 0 or prefill_length <= 0 or prefill_length > sample_length: + raise ValueError( + "DSV4 serving MHC replay requires 0 < prefill <= sample length, got " + f"prefill={prefill_length} sample_length={sample_length}" + ) + segments.append(prefill_length) + segments.extend([1] * (sample_length - prefill_length)) + if live_rows < compute_rows: + segments.append(compute_rows - live_rows) + return tuple(segments) + + +def _build_cp_serving_mhc_segments( + *, + layout: Dsv4ExactCPLayout, + sampler_prefill_lengths: torch.Tensor | None, +) -> tuple[ExactMhcReplaySegment, ...] | None: + """Build global-serving-size MHC calls projected onto local CP rows.""" + + if sampler_prefill_lengths is None or layout.local_live_count == 0: + return None + if sampler_prefill_lengths.ndim != 1 or sampler_prefill_lengths.dtype not in (torch.int32, torch.int64): + raise ValueError("DSV4 sampler_prefill_lengths must be a rank-one integer tensor") + prefill_lengths = [int(value) for value in sampler_prefill_lengths.detach().cpu().tolist()] + if tuple(layout.request_ids) != tuple(range(len(prefill_lengths))): + raise ValueError( + "DSV4 CP serving MHC replay needs one boundary for every packed request id: " + f"request_ids={layout.request_ids} prefills={prefill_lengths}" + ) + + global_request_lengths = { + request_id: int(rows.numel()) + for request_id, rows in zip(layout.request_ids, layout.global_request_row_indices, strict=True) + } + for request_id, prefill_length in enumerate(prefill_lengths): + sample_length = global_request_lengths[request_id] + if prefill_length <= 0 or prefill_length > sample_length: + raise ValueError( + "DSV4 serving MHC replay requires 0 < prefill <= sample length, got " + f"request={request_id} prefill={prefill_length} sample_length={sample_length}" + ) + + segments: list[ExactMhcReplaySegment] = [] + for request_id, local_rows_tensor in zip( + layout.request_ids, + layout.local_request_row_indices, + strict=True, + ): + local_rows = [int(value) for value in local_rows_tensor.detach().cpu().tolist()] + local_positions = [ + int(value) + for value in layout.local_request_positions.index_select(0, local_rows_tensor).detach().cpu().tolist() + ] + prefill_rows = tuple( + row + for row, position in zip(local_rows, local_positions, strict=True) + if position < prefill_lengths[request_id] + ) + prefill_positions = tuple(position for position in local_positions if position < prefill_lengths[request_id]) + if prefill_rows: + segments.append( + ExactMhcReplaySegment( + launch_rows=prefill_lengths[request_id], + source_rows=prefill_rows, + launch_positions=prefill_positions, + ) + ) + for row, position in zip(local_rows, local_positions, strict=True): + if position >= prefill_lengths[request_id]: + segments.append( + ExactMhcReplaySegment( + launch_rows=1, + source_rows=(row,), + launch_positions=(0,), + ) + ) + if layout.local_live_count < layout.compute_rows: + padding_rows = tuple(range(layout.local_live_count, layout.compute_rows)) + segments.append( + ExactMhcReplaySegment( + launch_rows=len(padding_rows), + source_rows=padding_rows, + launch_positions=tuple(range(len(padding_rows))), + ) + ) + return tuple(segments) def _move_preserved_param( @@ -355,6 +477,61 @@ def _capture_diagnostic_component(self, name: str, value: torch.Tensor) -> None: if callable(capture): capture(name, value) + def _maybe_capture_exact_attention_component( + self, + name: str, + value: torch.Tensor, + positions: torch.Tensor, + ) -> None: + """Persist one raw trainer attention boundary for operator localization.""" + + capture_dir = os.environ.get("XORL_DSV4_TRAINER_ATTENTION_CAPTURE_DIR", "").strip() + raw_layer = os.environ.get("XORL_DSV4_ATTENTION_CAPTURE_LAYER", "").strip() + raw_position = os.environ.get("XORL_DSV4_ATTENTION_CAPTURE_POSITION", "").strip() + raw_rank = os.environ.get("XORL_DSV4_ATTENTION_CAPTURE_GLOBAL_RANK", "0").strip() + if not capture_dir or not raw_layer or not raw_position: + return + capture_layers = {int(layer.strip()) for layer in raw_layer.split(",") if layer.strip()} + if self.layer_id not in capture_layers: + return + global_rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + if global_rank != int(raw_rank): + return + capture_key = (self.layer_id, name) + captured = self.__dict__.setdefault("_xorl_dsv4_captured_attention_components", set()) + if capture_key in captured: + return + flat_positions = positions.detach().reshape(-1) + target_position = int(raw_position) + matching = (flat_positions == target_position).nonzero(as_tuple=True)[0] + rows = value.reshape(-1, *value.shape[2:]) if value.ndim >= 3 else value + if matching.numel() == 0: + return + if matching.numel() != 1 or rows.shape[0] != flat_positions.numel(): + raise RuntimeError( + f"DSV4 trainer attention component {name} cannot align value and positions: " + f"value={tuple(value.shape)} positions={tuple(flat_positions.shape)} " + f"matches={matching.numel()}" + ) + captured.add(capture_key) + os.makedirs(capture_dir, exist_ok=True) + output_path = os.path.join( + capture_dir, + f"rank{global_rank:05d}.layer{self.layer_id:03d}.{name}.position{target_position:05d}.pt", + ) + torch.save( + { + "schema": "xorl.dsv4_trainer_attention_component.v1", + "global_rank": global_rank, + "layer": self.layer_id, + "component": name, + "position": target_position, + "source_shape": tuple(value.shape), + "value": rows.index_select(0, matching).detach().cpu(), + }, + output_path, + ) + def forward( self, hidden_states: torch.Tensor, @@ -427,6 +604,18 @@ def validate_lora_metadata(where: str) -> None: q_positions = get_q_positions_for_cp( seqlen_local, cp_size=self.cp_size, cp_group=self.cp_group, device=x.device ) + capture_positions = ( + ( + torch.arange( + carry_offset, + carry_offset + seqlen_local, + dtype=torch.int64, + device=x.device, + ) + ) + if carry_offset is not None + else q_positions + ) # SGLang's DSV4 CP path shards queries for execution, but explicitly # all-gathers and reranges BF16 KV plus compressor scores back into # logical token order before the cache/compressor kernels. The trainer @@ -482,6 +671,7 @@ def validate_lora_metadata(where: str) -> None: apply_rotary_emb(q[..., -rd:], freqs_cis) validate_lora_metadata("Q projection and RoPE") self._capture_diagnostic_component("q", q) + self._maybe_capture_exact_attention_component("q_pre_attention", q, capture_positions) # ---------------- KV (single shared stream) ---------------- kv_pre_norm = self.wkv(x) @@ -529,29 +719,89 @@ def validate_lora_metadata(where: str) -> None: request_x = x.index_select(1, local_rows) if ratio else None if ratio == 0: - from xorl.ops.dsv4.exact_attention import exact_c0_attention # noqa: PLC0415 + from xorl.ops.dsv4.exact_attention import ( # noqa: PLC0415 + exact_attention_layer, + exact_c0_attention, + ) + + with exact_attention_layer(self.layer_id): + request_o = exact_c0_attention( + request_q, + request_kv, + self.kv_norm.weight, + self.attn_sink, + request_freqs, + self.eps, + self.softmax_scale, + query_positions=request_positions, + kv_preprocessed=False, + ) + else: + from xorl.ops.dsv4.exact_attention import ( # noqa: PLC0415 + exact_attention_layer, + exact_compressed_attention, + ) - request_o = exact_c0_attention( - request_q, - request_kv, + with exact_attention_layer(self.layer_id): + request_o = exact_compressed_attention( + request_q, + request_kv, + request_x, + self.kv_norm.weight, + self.attn_sink, + request_freqs, + self.compressor.wkv.weight, + self.compressor.wgate.weight, + self.compressor.ape, + self.compressor.norm.weight, + self.eps, + self.softmax_scale, + ratio, + query_positions=request_positions, + kv_preprocessed=False, + ) + o = o.index_copy(1, local_rows, request_o) + else: + exact_kv_pre_norm = kv_pre_norm + exact_x = x + if self.cp_size > 1: + exact_kv_pre_norm = all_gather_cp(exact_kv_pre_norm, dim=1, cp_group=self.cp_group) + if ratio: + exact_x = all_gather_cp(exact_x, dim=1, cp_group=self.cp_group) + if ratio == 0: + from xorl.ops.dsv4.exact_attention import ( # noqa: PLC0415 + exact_attention_layer, + exact_c0_attention, + ) + + with exact_attention_layer(self.layer_id): + o = exact_c0_attention( + q, + exact_kv_pre_norm, self.kv_norm.weight, self.attn_sink, - request_freqs, + exact_kv_freqs_cis, self.eps, self.softmax_scale, - query_positions=request_positions, + carry_state=carry_state, + position_offset=carry_offset or 0, + query_positions=exact_query_positions, kv_preprocessed=False, ) - else: - from xorl.ops.dsv4.exact_attention import exact_compressed_attention # noqa: PLC0415 + else: + from xorl.ops.dsv4.exact_attention import ( # noqa: PLC0415 + exact_attention_layer, + exact_compressed_attention, + ) - request_o = exact_compressed_attention( - request_q, - request_kv, - request_x, + with exact_attention_layer(self.layer_id): + o = exact_compressed_attention( + q, + exact_kv_pre_norm, + exact_x, self.kv_norm.weight, self.attn_sink, - request_freqs, + exact_kv_freqs_cis, self.compressor.wkv.weight, self.compressor.wgate.weight, self.compressor.ape, @@ -559,55 +809,11 @@ def validate_lora_metadata(where: str) -> None: self.eps, self.softmax_scale, ratio, - query_positions=request_positions, + carry_state=carry_state, + position_offset=carry_offset or 0, + query_positions=exact_query_positions, kv_preprocessed=False, ) - o = o.index_copy(1, local_rows, request_o) - else: - exact_kv_pre_norm = kv_pre_norm - exact_x = x - if self.cp_size > 1: - exact_kv_pre_norm = all_gather_cp(exact_kv_pre_norm, dim=1, cp_group=self.cp_group) - if ratio: - exact_x = all_gather_cp(exact_x, dim=1, cp_group=self.cp_group) - if ratio == 0: - from xorl.ops.dsv4.exact_attention import exact_c0_attention # noqa: PLC0415 - - o = exact_c0_attention( - q, - exact_kv_pre_norm, - self.kv_norm.weight, - self.attn_sink, - exact_kv_freqs_cis, - self.eps, - self.softmax_scale, - carry_state=carry_state, - position_offset=carry_offset or 0, - query_positions=exact_query_positions, - kv_preprocessed=False, - ) - else: - from xorl.ops.dsv4.exact_attention import exact_compressed_attention # noqa: PLC0415 - - o = exact_compressed_attention( - q, - exact_kv_pre_norm, - exact_x, - self.kv_norm.weight, - self.attn_sink, - exact_kv_freqs_cis, - self.compressor.wkv.weight, - self.compressor.wgate.weight, - self.compressor.ape, - self.compressor.norm.weight, - self.eps, - self.softmax_scale, - ratio, - carry_state=carry_state, - position_offset=carry_offset or 0, - query_positions=exact_query_positions, - kv_preprocessed=False, - ) kv_vanilla = None else: kv_vanilla = self.kv_norm(kv_pre_norm) # [B, S, D] @@ -666,6 +872,7 @@ def validate_lora_metadata(where: str) -> None: o = dense_attn_torch(q, kv, attn_sink, topk_idxs, self.softmax_scale) validate_lora_metadata("attention core") self._capture_diagnostic_component("attn_output", o) + self._maybe_capture_exact_attention_component("attention_core_output", o, capture_positions) # Inverse RoPE on the rope slice of the output. if self._exact_attention: @@ -676,10 +883,12 @@ def validate_lora_metadata(where: str) -> None: apply_rotary_emb(o[..., -rd:], freqs_cis, inverse=True) validate_lora_metadata("inverse RoPE") self._capture_diagnostic_component("attn_output_gated", o) + self._maybe_capture_exact_attention_component("attention_inverse_output", o, capture_positions) # ---------------- Grouped output projection ---------------- # o : [B, S, H, D] -> [B, S, n_local_groups, n_heads*D / n_local_groups] o = o.view(bsz, seqlen_local, self.n_local_groups, -1) + self._maybe_capture_exact_attention_component("wo_a_input", o, capture_positions) native_wo_a_payload = self.wo_a._modules.get("native_base_payload") if native_wo_a_payload is not None: from .native_payload import dsv4_native_grouped_wo_a # noqa: PLC0415 @@ -716,8 +925,11 @@ def validate_lora_metadata(where: str) -> None: delta = torch.einsum("bsgr,gor->bsgo", mid, lora_B_g) * self.wo_a.scaling o_base = o_base + delta.to(o_base.dtype) self._capture_diagnostic_component("o_proj_output", o_base) + self._maybe_capture_exact_attention_component("wo_a_output", o_base, capture_positions) - return self.wo_b(o_base.flatten(2)) + output = self.wo_b(o_base.flatten(2)) + self._maybe_capture_exact_attention_component("attention_output", output, capture_positions) + return output class DeepseekV4MLP(nn.Module): @@ -796,12 +1008,24 @@ def __init__( self.is_hash_layer = layer_id < int(config.num_hash_layers) self.routed_scaling_factor = float(config.routed_scaling_factor) self._dsv4_exact_native = bool(getattr(config, "_dsv4_flash_exact_mode", False)) + self._dsv4_moe_numerical_program = getattr( + config, + "_dsv4_moe_numerical_program", + DSV4_DEEPEP_NATIVE_EXACT_V1 if self._dsv4_exact_native else None, + ) + self._dsv4_deepep_native = self._dsv4_moe_numerical_program == DSV4_DEEPEP_NATIVE_EXACT_V1 self.experts.ep_dispatch = getattr(config, "_ep_dispatch", "alltoall") self.experts.deepep_buffer_size_gb = getattr(config, "_deepep_buffer_size_gb", 2.0) self.experts.deepep_num_sms = getattr(config, "_deepep_num_sms", 20) self.experts.deepep_async_combine = getattr(config, "_deepep_async_combine", False) self.experts.alltoall_combine_hidden_chunk_size = getattr(config, "_alltoall_combine_hidden_chunk_size", 0) self.experts.expert_lora_semantics = DSV4_CLAMPED_SWIGLU_LORA_PROGRAM + if self._dsv4_deepep_native and self.experts.ep_dispatch != "deepep": + raise RuntimeError( + f"{DSV4_DEEPEP_NATIVE_EXACT_V1} requires ep_dispatch='deepep'; got {self.experts.ep_dispatch!r}" + ) + if self._dsv4_exact_native and self._dsv4_moe_numerical_program != DSV4_DEEPEP_NATIVE_EXACT_V1: + raise RuntimeError(f"Unknown exact DSV4 MoE numerical program: {self._dsv4_moe_numerical_program!r}") # Replace the parent's softmax router with our V4 one. self.router = TopKRouter( @@ -848,6 +1072,57 @@ def _capture_diagnostic_component(self, name: str, value: torch.Tensor) -> None: if callable(capture): capture(name, value) + def _maybe_capture_native_moe_output( + self, + name: str, + value: torch.Tensor, + input_ids: torch.Tensor | None, + ) -> None: + """Persist one owner-token value around the routed/shared MoE join.""" + + capture_dir = os.environ.get("XORL_DSV4_TRAINER_MOE_CAPTURE_DIR", "").strip() + raw_layer = os.environ.get("XORL_DSV4_MOE_CAPTURE_LAYER", "").strip() + raw_token = os.environ.get("XORL_DSV4_MOE_CAPTURE_TOKEN_ID", "").strip() + if not capture_dir or not raw_layer or not raw_token or input_ids is None: + return + capture_layers = {int(layer.strip()) for layer in raw_layer.split(",") if layer.strip()} + if self.layer_id not in capture_layers: + return + capture_key = (self.layer_id, name) + captured = self.__dict__.setdefault("_xorl_dsv4_captured_moe_outputs", set()) + if capture_key in captured: + return + flat_ids = input_ids.detach().reshape(-1) + target_token = int(raw_token) + matching = (flat_ids == target_token).nonzero(as_tuple=True)[0] + if matching.numel() == 0: + return + if matching.numel() != 1 or value.shape[0] != flat_ids.numel(): + raise RuntimeError( + f"DSV4 trainer MoE capture {name} cannot align value and ids: " + f"value={tuple(value.shape)} ids={tuple(flat_ids.shape)} " + f"matches={matching.numel()}" + ) + captured.add(capture_key) + global_rank = dist.get_rank() if dist.is_initialized() else 0 + os.makedirs(capture_dir, exist_ok=True) + torch.save( + { + "schema": "xorl.dsv4_trainer_moe_output.v1", + "global_rank": global_rank, + "layer": self.layer_id, + "component": name, + "token_id": target_token, + "row_index": int(matching.item()), + "source_shape": tuple(value.shape), + "value": value.index_select(0, matching).detach().cpu(), + }, + os.path.join( + capture_dir, + f"rank{global_rank:05d}.layer{self.layer_id:03d}.{name}.token{target_token:06d}.pt", + ), + ) + def route(self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None): """V4-specific route: passes ``expert_bias`` / ``tid2eid`` / ``input_ids`` into the router, and integrates routing replay for gradient-checkpoint @@ -1059,6 +1334,8 @@ def forward( Returns: ``(output [batch, seqlen, hidden_size], router_logits [N, num_experts])``. """ + if self._dsv4_deepep_native: + return self._forward_deepep_native_exact(hidden_states, input_ids, live_token_count) if self._dsv4_exact_native: return self._forward_exact_native(hidden_states, input_ids, live_token_count) @@ -1078,6 +1355,153 @@ def forward( final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) return final_hidden_states, router_logits + def _forward_deepep_native_exact( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor | None, + live_token_count: int | None, + ): + """Thin DSV4 adapter for the shared native DeepEP exact program. + + DSV4 owns only routing, contiguous expert ownership, its MXFP4-Marlin + active-LoRA runner, and the replicated shared expert. The reusable + layer owns dispatch receipts, runner-output validation, FP32 local + route reduction, BF16 communication leaves, handle combine, FP64 fold, + and reverse communication. + """ + + from xorl.distributed.moe.deepep_native_exact import ( # noqa: PLC0415 + canonicalize_native_routing_metadata, + native_dispatch_runner_combine, + ) + from xorl.distributed.parallel_state import get_parallel_state # noqa: PLC0415 + from xorl.models.transformers.deepseek_v4.native_payload import ( # noqa: PLC0415 + dsv4_join_routed_shared_partial, + dsv4_native_shared_expert_tp_partial, + ) + + batch_size, sequence_length, hidden_dim = hidden_states.shape + local_hidden = hidden_states.reshape(-1, hidden_dim) + local_rows = local_hidden.shape[0] + live_token_count = local_rows if live_token_count is None else int(live_token_count) + if not 0 <= live_token_count <= local_rows: + raise ValueError( + f"Native DSV4 DeepEP live token count {live_token_count} is outside packed row count {local_rows}" + ) + parallel_state = get_parallel_state() + if int(parallel_state.tp_size) != 1: + raise RuntimeError(f"{DSV4_DEEPEP_NATIVE_EXACT_V1} requires stage-local body TP1") + if not parallel_state.ep_enabled or int(parallel_state.ep_size) != 8: + raise RuntimeError(f"{DSV4_DEEPEP_NATIVE_EXACT_V1} requires EP8") + if parallel_state.ep_group is None: + raise RuntimeError(f"{DSV4_DEEPEP_NATIVE_EXACT_V1} requires a stage-local EP process group") + if self.train_router: + raise RuntimeError(f"{DSV4_DEEPEP_NATIVE_EXACT_V1} requires a frozen router") + if self.shared_experts is None: + raise RuntimeError(f"{DSV4_DEEPEP_NATIVE_EXACT_V1} requires the official shared expert") + + live_hidden = local_hidden[:live_token_count] + live_ids = None if input_ids is None else input_ids.reshape(-1)[:live_token_count] + if live_token_count: + routing_weights, selected_experts, router_logits = self.route( + live_hidden, + input_ids=live_ids, + ) + else: + router_logits = local_hidden.new_zeros((0, self.num_experts), dtype=torch.float32) + routing_weights = local_hidden.new_zeros((0, self.top_k), dtype=torch.float32) + selected_experts = torch.empty( + (0, self.top_k), + dtype=torch.int32, + device=local_hidden.device, + ) + self._capture_diagnostic_component("moe_native_deepep_input", live_hidden) + self._capture_diagnostic_component("moe_native_deepep_router_logits", router_logits) + self._capture_diagnostic_component("moe_native_deepep_topk_ids", selected_experts) + routing_weights = canonicalize_native_routing_metadata(routing_weights) + self._capture_diagnostic_component("moe_native_deepep_topk_weights", routing_weights) + self._maybe_capture_native_moe_output("topk_ids", selected_experts, live_ids) + self._maybe_capture_native_moe_output("topk_weights", routing_weights, live_ids) + + payload = self.experts.native_mxfp4_payload + local_experts = int(payload.w13_weight.shape[0]) + if local_experts * int(parallel_state.ep_size) != self.num_experts: + raise RuntimeError( + f"Native DSV4 DeepEP expert ownership is not contiguous EP8: " + f"local={local_experts}, global={self.num_experts}" + ) + local_start = int(parallel_state.ep_rank) * local_experts + + def run_local_leaf(recv_hidden, recv_weights, recv_local_ids): + valid = recv_local_ids >= 0 + recv_global_ids = torch.where( + valid, + recv_local_ids.to(torch.int64) + local_start, + torch.full_like(recv_local_ids, -1, dtype=torch.int64), + ).to(torch.int32) + self._capture_diagnostic_component("moe_native_deepep_recv_hidden", recv_hidden) + self._capture_diagnostic_component("moe_native_deepep_recv_local_ids", recv_local_ids) + self._capture_diagnostic_component("moe_native_deepep_recv_weights", recv_weights) + local_leaf = self.experts( + recv_hidden, + recv_weights, + recv_global_ids, + dsv4_exact_native=True, + dsv4_exact_lora_live=True, + dsv4_exact_return_routes=False, + ) + self._capture_diagnostic_component("moe_native_deepep_recv_leaf", local_leaf) + return local_leaf + + routed_local = native_dispatch_runner_combine( + live_hidden, + routing_weights, + selected_experts, + ep_group=parallel_state.ep_group, + num_experts=self.num_experts, + num_local_experts=local_experts, + buffer_size_gb=self.experts.deepep_buffer_size_gb, + num_sms=self.experts.deepep_num_sms, + runner=run_local_leaf, + ) + self._maybe_capture_native_moe_output("routed_output", routed_local, live_ids) + self._capture_diagnostic_component("moe_native_deepep_combined_routed", routed_local) + if routed_local.shape != live_hidden.shape: + raise RuntimeError( + "Native DSV4 DeepEP combine returned the wrong owner shape: " + f"got {tuple(routed_local.shape)}, expected {tuple(live_hidden.shape)}" + ) + + shared_local = ( + dsv4_native_shared_expert_tp_partial( + live_hidden, + self.shared_experts, + tp_rank=0, + tp_size=1, + diagnostic_capture=self._capture_diagnostic_component, + lora_live=True, + ) + if live_token_count + else live_hidden * 0.0 + ) + self._maybe_capture_native_moe_output("shared_output", shared_local, live_ids) + combined = dsv4_join_routed_shared_partial( + routed_local, + shared_local, + routed_scaling_factor=self.routed_scaling_factor, + ) + self._maybe_capture_native_moe_output("joined_output", combined, live_ids) + self._capture_diagnostic_component("moe_native_deepep_output", combined) + padded = torch.cat((combined, local_hidden[live_token_count:] * 0.0), dim=0) + padded_logits = torch.cat( + ( + router_logits, + router_logits.new_zeros((local_rows - live_token_count, self.num_experts)), + ), + dim=0, + ) + return padded.reshape(batch_size, sequence_length, hidden_dim), padded_logits + def _forward_exact_native( self, hidden_states: torch.Tensor, @@ -1307,6 +1731,75 @@ def _capture_diagnostic_component(self, name: str, value: torch.Tensor) -> None: if callable(capture): capture(name, value) + def _maybe_capture_exact_component( + self, + name: str, + value: torch.Tensor, + input_ids: torch.Tensor | None, + live_token_count: int | None, + ) -> None: + """Persist a raw owner-token component for cross-engine localization.""" + + capture_dir = os.environ.get("XORL_DSV4_TRAINER_COMPONENT_CAPTURE_DIR", "").strip() + raw_layer = os.environ.get("XORL_DSV4_COMPONENT_CAPTURE_LAYER", "").strip() + raw_token = os.environ.get("XORL_DSV4_COMPONENT_CAPTURE_TOKEN_ID", "").strip() + if not capture_dir or not raw_layer or not raw_token or input_ids is None: + return + capture_layers = {int(layer.strip()) for layer in raw_layer.split(",") if layer.strip()} + if self.layer_id not in capture_layers: + return + capture_key = (self.layer_id, name) + captured = self.__dict__.setdefault("_xorl_dsv4_captured_components", set()) + if capture_key in captured: + return + if value.ndim == 4: + rows = value.reshape(-1, value.shape[-2], value.shape[-1]) + elif value.ndim == 3: + rows = value.reshape(-1, value.shape[-1]) + else: + raise RuntimeError(f"DSV4 trainer component {name} has unsupported shape {tuple(value.shape)}") + live_rows = rows.shape[0] if live_token_count is None else int(live_token_count) + if not 0 <= live_rows <= rows.shape[0]: + raise RuntimeError( + f"DSV4 trainer component {name} has invalid live row count {live_rows} for {rows.shape[0]} storage rows" + ) + if live_rows == 0: + return + flat_input_ids = input_ids.detach().reshape(-1) + if flat_input_ids.numel() < live_rows: + raise RuntimeError( + f"DSV4 trainer component {name} has fewer token ids than live rows: " + f"ids={flat_input_ids.numel()} rows={live_rows}" + ) + live_input_ids = flat_input_ids[:live_rows] + target_token = int(raw_token) + matching = (live_input_ids == target_token).nonzero(as_tuple=True)[0] + if matching.numel() == 0: + return + if matching.numel() != 1: + raise RuntimeError( + f"DSV4 trainer component {name} requires one token {target_token}, got {matching.numel()}" + ) + captured.add(capture_key) + global_rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + os.makedirs(capture_dir, exist_ok=True) + output_path = os.path.join( + capture_dir, + f"rank{global_rank:05d}.layer{self.layer_id:03d}.{name}.token{target_token:06d}.pt", + ) + torch.save( + { + "schema": "xorl.dsv4_trainer_component.v1", + "global_rank": global_rank, + "layer": self.layer_id, + "component": name, + "token_id": target_token, + "source_shape": tuple(value.shape), + "value": rows.index_select(0, matching).detach().cpu(), + }, + output_path, + ) + def forward( self, hidden_states_4d: torch.Tensor, @@ -1315,6 +1808,7 @@ def forward( live_token_count: int | None = None, exact_cp_layout: Dsv4ExactCPLayout | None = None, decode_carry_offset: int | None = None, + serving_mhc_segments: tuple[int | ExactMhcReplaySegment, ...] | None = None, ) -> torch.Tensor: """Run the layer. @@ -1326,6 +1820,7 @@ def forward( Returns: ``[batch, seqlen, hc_mult, hidden_size]``. """ + self._maybe_capture_exact_component("layer_input", hidden_states_4d, input_ids, live_token_count) # ---- Attention sublayer ---- if self._exact_mhc: h3d, post, comb = self.hc_util.layer_pre_norm_exact( @@ -1334,6 +1829,7 @@ def forward( self.hc_attn_scale, self.hc_attn_base, self.input_layernorm.weight, + serving_segments=serving_mhc_segments, ) else: h3d, post, comb = self.hc_util.layer_pre( @@ -1341,13 +1837,16 @@ def forward( ) h3d = self.input_layernorm(h3d) self._capture_diagnostic_component("input_norm", h3d) + self._maybe_capture_exact_component("attention_input", h3d, input_ids, live_token_count) attn_out = self.self_attn( h3d, exact_cp_layout=exact_cp_layout, decode_carry_offset=decode_carry_offset, ) + self._maybe_capture_exact_component("attention_output", attn_out, input_ids, live_token_count) post_fn = self.hc_util.layer_post_exact if self._exact_mhc else self.hc_util.layer_post hidden_states_4d = post_fn(attn_out, hidden_states_4d, post, comb) + self._maybe_capture_exact_component("post_attention_residual", hidden_states_4d, input_ids, live_token_count) # ---- MoE / FFN sublayer ---- if self._exact_mhc: @@ -1357,6 +1856,7 @@ def forward( self.hc_ffn_scale, self.hc_ffn_base, self.post_attention_layernorm.weight, + serving_segments=serving_mhc_segments, ) else: h3d, post, comb = self.hc_util.layer_pre( @@ -1364,11 +1864,13 @@ def forward( ) h3d = self.post_attention_layernorm(h3d) self._capture_diagnostic_component("post_attention_norm", h3d) + self._maybe_capture_exact_component("ffn_input", h3d, input_ids, live_token_count) ffn_out, router_logits = self.mlp( h3d, input_ids=input_ids, live_token_count=live_token_count, ) + self._maybe_capture_exact_component("ffn_output", ffn_out, input_ids, live_token_count) if os.environ.get("XORL_DSV4_DIAGNOSTIC_BASE_MARLIN") == "1": from .native_payload import _validate_all_single_adapter_batch_infos # noqa: PLC0415 @@ -1377,6 +1879,7 @@ def forward( where=f"layer {self.layer_id} MoE module return", ) hidden_states_4d = post_fn(ffn_out, hidden_states_4d, post, comb) + self._maybe_capture_exact_component("layer_output", hidden_states_4d, input_ids, live_token_count) if os.environ.get("XORL_DSV4_DIAGNOSTIC_BASE_MARLIN") == "1": _validate_all_single_adapter_batch_infos( h3d.device.index, @@ -1694,6 +2197,82 @@ def _capture_diagnostic_component(self, name: str, value: torch.Tensor) -> None: if callable(capture): capture(name, value) + def _maybe_capture_exact_residual_row( + self, + *, + layer_id: int, + hidden_states: torch.Tensor, + position_ids: torch.Tensor | None, + input_ids: torch.Tensor | None, + ) -> None: + """Persist one trainer residual row for sampler/trainer localization. + + This cold-path hook is inert unless both trainer-specific capture + environment variables are set. Position selection happens before any + token-diagnostic unpacking so packed rows cannot be misattributed. + """ + + capture_dir = os.environ.get("XORL_DSV4_TRAINER_LAYER_CAPTURE_DIR", "").strip() + raw_position = os.environ.get("XORL_DSV4_TRAINER_LAYER_CAPTURE_POSITION", "").strip() + if not capture_dir or not raw_position or position_ids is None: + return + target_position = int(raw_position) + captured = self.__dict__.setdefault("_xorl_dsv4_captured_layer_positions", set()) + capture_key = (target_position, int(layer_id)) + if capture_key in captured: + return + if hidden_states.ndim != 4: + raise RuntimeError( + "DSV4 trainer layer capture expects the four-dimensional MHC residual: " + f"hidden={tuple(hidden_states.shape)}" + ) + hidden_rows = hidden_states.reshape( + -1, + hidden_states.shape[-2], + hidden_states.shape[-1], + ) + storage_positions = position_ids.detach().reshape(-1) + compute_row_count = hidden_rows.shape[0] + if storage_positions.numel() < compute_row_count: + raise RuntimeError( + "DSV4 trainer layer capture has fewer positions than live residual rows: " + f"hidden={tuple(hidden_states.shape)} positions={tuple(storage_positions.shape)}" + ) + # Exact packed execution compacts live rows into the storage prefix and + # leaves the remaining positions as transport padding. Select only the + # positions that still have corresponding residual rows. + flat_positions = storage_positions[:compute_row_count] + matching = (flat_positions == target_position).nonzero(as_tuple=True)[0] + if matching.numel() != 1: + raise RuntimeError( + "DSV4 trainer layer capture requires exactly one matching packed row: " + f"position={target_position} matches={matching.numel()}" + ) + row = hidden_rows.index_select(0, matching).detach().cpu() + captured.add(capture_key) + global_rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + token_id = None + if input_ids is not None and input_ids.numel() >= compute_row_count: + live_input_ids = input_ids.detach().reshape(-1)[:compute_row_count] + token_id = int(live_input_ids.index_select(0, matching).item()) + os.makedirs(capture_dir, exist_ok=True) + output_path = os.path.join( + capture_dir, + f"rank{global_rank:05d}.layer{int(layer_id):03d}.position{target_position:05d}.pt", + ) + torch.save( + { + "schema": "xorl.dsv4_trainer_layer_output.v1", + "global_rank": global_rank, + "layer": int(layer_id), + "position": target_position, + "token_id": token_id, + "source_shape": tuple(hidden_states.shape), + "hidden": row, + }, + output_path, + ) + def forward( self, input_ids: torch.Tensor | None = None, @@ -1711,12 +2290,14 @@ def forward( """ pp_stage_is_first = kwargs.pop("_pp_stage_is_first", None) pp_stage_is_last = kwargs.pop("_pp_stage_is_last", None) + output_hidden_states = bool(kwargs.pop("output_hidden_states", False)) if (pp_stage_is_first is None) != (pp_stage_is_last is None): raise ValueError("DeepSeek-V4 pipeline stage flags must be supplied together") is_pipeline_stage = pp_stage_is_first is not None incoming_hyperconnection_state = bool(is_pipeline_stage and not pp_stage_is_first) decode_cache_carry = False carry_position_offset = 0 + serving_mhc_segments = None if getattr(self.config, "_dsv4_flash_exact_mode", False): decode_cache_carry = bool(kwargs.pop("decode_cache_carry", False)) if decode_cache_carry: @@ -1755,6 +2336,7 @@ def forward( cp_live_mask = kwargs.pop("_cp_live_mask", None) sample_lengths = kwargs.pop("_r3_sample_lengths", None) num_samples = kwargs.pop("num_samples", None) + sampler_prefill_lengths = kwargs.pop("sampler_prefill_lengths", None) if isinstance(sample_lengths, torch.Tensor): sample_lengths = sample_lengths.detach().cpu().reshape(-1).tolist() if not decode_cache_carry: @@ -1845,8 +2427,38 @@ def forward( ) compact_ids = input_ids.index_select(1, exact_cp_layout.local_storage_indices) input_ids = F.pad(compact_ids, (0, compute_rows - live_token_count), value=0) + if int(parallel_state.cp_size) != 1: + serving_mhc_segments = _build_cp_serving_mhc_segments( + layout=exact_cp_layout, + sampler_prefill_lengths=sampler_prefill_lengths, + ) + else: + serving_mhc_segments = _build_serving_mhc_segments( + compute_rows=int(hidden_states.shape[1]), + sample_lengths=( + [int(length) for length in sample_lengths] + if sample_lengths is not None + else ([int(hidden_states.shape[1])] if num_samples is None or int(num_samples) > 0 else []) + ), + sampler_prefill_lengths=sampler_prefill_lengths, + ) + if serving_mhc_segments is not None and not getattr(self, "_serving_mhc_engagement_logged", False): + logger.info( + "Exact DSV4 serving-segment MHC ENGAGED: requests=%s prefill_lengths=%s " + "compute_rows=%s segment_count=%s", + sample_lengths, + sampler_prefill_lengths.detach().cpu().tolist(), + hidden_states.shape[1], + len(serving_mhc_segments), + ) + self._serving_mhc_engagement_logged = True else: - live_token_count = packed_sequence_length + # Decode-cache replay keeps every EP rank in the same sequence + # of collectives, but idle ranks must contribute zero rows just + # as they do in serving. The scorer supplies ``num_samples=0`` + # on dispatcher-created dummy batches. + live_token_count = 0 if num_samples is not None and int(num_samples) == 0 else packed_sequence_length + capture_position_ids = position_ids del position_ids del kwargs @@ -1857,6 +2469,13 @@ def forward( h4d = ( hidden_states if incoming_hyperconnection_state else self.hc_util.block_expand(hidden_states) ) # [B, S, hc_mult, H] + self._maybe_capture_exact_residual_row( + layer_id=-1, + hidden_states=h4d, + position_ids=capture_position_ids, + input_ids=input_ids, + ) + all_hidden_states = [h4d] if output_hidden_states else None use_outer_checkpoint = ( self.gradient_checkpointing @@ -1875,6 +2494,7 @@ def forward( live_token_count=live_token_count, exact_cp_layout=exact_cp_layout, decode_carry_offset=(carry_position_offset if decode_cache_carry else None), + serving_mhc_segments=serving_mhc_segments, ) else: layer_outputs = layer( @@ -1884,6 +2504,7 @@ def forward( live_token_count=live_token_count, exact_cp_layout=exact_cp_layout, decode_carry_offset=(carry_position_offset if decode_cache_carry else None), + serving_mhc_segments=serving_mhc_segments, ) if os.environ.get("XORL_DSV4_DIAGNOSTIC_BASE_MARLIN") == "1": @@ -1907,6 +2528,14 @@ def forward( all_router_logits.append(router_logits) else: h4d = layer_outputs + self._maybe_capture_exact_residual_row( + layer_id=layer.layer_id, + hidden_states=h4d, + position_ids=capture_position_ids, + input_ids=input_ids, + ) + if all_hidden_states is not None: + all_hidden_states.append(h4d) if is_pipeline_stage and not pp_stage_is_last: # Preserve the live multi-stream residual state across the PP wire. @@ -1924,6 +2553,7 @@ def forward( h4d = F.pad(h4d, (0, 0, 0, 0, 0, transport_padding)) return MoeModelOutput( last_hidden_state=h4d, + hidden_states=tuple(all_hidden_states) if all_hidden_states is not None else None, router_logits=tuple(all_router_logits) if all_router_logits is not None else None, ) @@ -1937,6 +2567,7 @@ def forward( h3d = storage_h3d.index_copy(1, exact_cp_layout.local_storage_indices, live_h3d) return MoeModelOutput( last_hidden_state=h3d, + hidden_states=tuple(all_hidden_states) if all_hidden_states is not None else None, router_logits=tuple(all_router_logits) if all_router_logits is not None else None, ) @@ -1944,6 +2575,15 @@ def forward( class DeepseekV4ForCausalLM(DeepseekV4PreTrainedModel): """Causal-LM head wrapping :class:`DeepseekV4Model`.""" + deepep_native_exact_capability = { + "produces_local_leaf": True, + "wire_dtype": "bf16", + "uses_dispatch_handle": True, + "supported_ep_sizes": (8,), + "local_leaf_program": "dsv4_mxfp4_marlin", + "lora_serving_modes": ("separate",), + } + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} # Keep the generic alias while the exact DSV4 PP path uses the explicit # original-ID contract below. @@ -2032,7 +2672,11 @@ def forward( output_router_logits=output_router_logits, **kwargs, ) - return MoeCausalLMOutput(last_hidden_state=outputs.last_hidden_state, router_logits=outputs.router_logits) + return MoeCausalLMOutput( + last_hidden_state=outputs.last_hidden_state, + hidden_states=outputs.hidden_states, + router_logits=outputs.router_logits, + ) def get_parallel_plan(self): """EP plan for the routed-expert weights. Consumed by diff --git a/src/xorl/models/transformers/deepseek_v4/moe_program.py b/src/xorl/models/transformers/deepseek_v4/moe_program.py new file mode 100644 index 00000000..fe13f642 --- /dev/null +++ b/src/xorl/models/transformers/deepseek_v4/moe_program.py @@ -0,0 +1,21 @@ +"""Named DeepSeek-V4 production MoE numerical program.""" + +DSV4_DEEPEP_NATIVE_EXACT_V1 = "deepep_native_exact_v1" + + +def resolve_dsv4_moe_numerical_program(*, exact: bool, ep_dispatch: str, deepep_native_exact: bool) -> str | None: + if not exact: + return None + if not deepep_native_exact or ep_dispatch != "deepep": + raise ValueError( + "Exact DeepSeek-V4 server training requires deepep_native_exact=true " + "with ep_dispatch='deepep'; the retired post-expert diagnostic is not a " + "production fallback" + ) + return DSV4_DEEPEP_NATIVE_EXACT_V1 + + +__all__ = [ + "DSV4_DEEPEP_NATIVE_EXACT_V1", + "resolve_dsv4_moe_numerical_program", +] diff --git a/src/xorl/models/transformers/deepseek_v4/native_payload.py b/src/xorl/models/transformers/deepseek_v4/native_payload.py index 0c93edac..89949234 100644 --- a/src/xorl/models/transformers/deepseek_v4/native_payload.py +++ b/src/xorl/models/transformers/deepseek_v4/native_payload.py @@ -800,6 +800,7 @@ def _build_dsv4_moe_runner_config( hidden_size: int, intermediate_size: int, top_k: int, + no_combine: bool = False, ) -> object: """Freeze the trainer-side serving runner contract for exact DSV4.""" return config_cls( @@ -815,6 +816,7 @@ def _build_dsv4_moe_runner_config( routed_scaling_factor=1.5, swiglu_limit=10.0, inplace=False, + no_combine=no_combine, dsv4_exact_mode=True, ) @@ -830,6 +832,8 @@ def _dsv4_native_mxfp4_forward( down_a: torch.Tensor, down_b: torch.Tensor, experts, + *, + no_combine: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig # noqa: PLC0415 from sglang.srt.layers.moe.moe_runner.runner import MoeRunner # noqa: PLC0415 @@ -886,6 +890,7 @@ def _dsv4_native_mxfp4_forward( hidden_size=payload.hidden_size, intermediate_size=payload.intermediate_size, top_k=selected_experts.shape[1], + no_combine=no_combine, ) dispatch = StandardDispatchOutput( hidden_states=run_hidden_states, @@ -922,6 +927,7 @@ def _dsv4_native_mxfp4_forward( clamp_limit=10.0, expert_map=expert_map, global_num_experts=256, + no_combine=no_combine, ) # This branch is an explicitly opt-in diagnostic control. Surface an # asynchronous Marlin fault at its own call boundary instead of at an @@ -973,6 +979,7 @@ def forward( down_b, experts, lora_live, + return_routes, ): factors = (gate_a, gate_b, up_a, up_b, down_a, down_b) # The gather-aware exact lane sets ``lora_live`` on every EP rank so @@ -985,9 +992,11 @@ def forward( selected_experts, *effective, experts, + no_combine=return_routes, ) ctx.experts = experts ctx.lora_live = lora_live + ctx.return_routes = return_routes ctx.save_for_backward( hidden_states.detach(), routing_weights.detach(), @@ -1020,6 +1029,7 @@ def backward(ctx, grad_output): ) = ctx.saved_tensors payload = ctx.experts.native_mxfp4_payload lora_live = ctx.lora_live + return_routes = ctx.return_routes need_x, need_r = ctx.needs_input_grad[:2] factor_needs = ctx.needs_input_grad[3:9] grad_x = torch.zeros_like(hidden_states) if need_x else None @@ -1038,7 +1048,7 @@ def backward(ctx, grad_output): continue token_ids, slot_ids = positions[:, 0], positions[:, 1] x = hidden_states[token_ids].detach().requires_grad_(need_x) - weight = routing_weights[token_ids, slot_ids].detach().requires_grad_(need_r) + weight = routing_weights[token_ids, slot_ids].detach().requires_grad_(need_r and not return_routes) effective = [ (master[expert_idx].detach() if lora_live else torch.zeros_like(master[expert_idx])) .to(torch.bfloat16) @@ -1063,13 +1073,18 @@ def backward(ctx, grad_output): activated = torch.nn.functional.silu(gate) * up down = torch.nn.functional.linear(activated, w2) down = down + (activated @ effective[4]) @ effective[5] - weighted = (down.float() * weight.float().unsqueeze(-1)).to(torch.bfloat16) + if return_routes: + expert_value = down.to(torch.bfloat16) + expert_grad_output = grad_output[token_ids, slot_ids] + else: + expert_value = (down.float() * weight.float().unsqueeze(-1)).to(torch.bfloat16) + expert_grad_output = grad_output[token_ids] targets = [] target_roles = [] if need_x: targets.append(x) target_roles.append(("x", None)) - if need_r: + if need_r and not return_routes: targets.append(weight) target_roles.append(("r", None)) for factor_idx, (factor, needed) in enumerate(zip(effective, factor_targets)): @@ -1077,9 +1092,9 @@ def backward(ctx, grad_output): targets.append(factor) target_roles.append(("factor", factor_idx)) grads = torch.autograd.grad( - weighted, + expert_value, targets, - grad_outputs=grad_output[token_ids], + grad_outputs=expert_grad_output, allow_unused=False, ) for role, grad in zip(target_roles, grads): @@ -1096,9 +1111,61 @@ def backward(ctx, grad_output): *factor_grads, None, None, + None, ) +class _Dsv4NativeMxfp4ZeroRows(torch.autograd.Function): + """Keep empty DeepEP receives in the active-LoRA autograd program.""" + + @staticmethod + def forward( + ctx, + hidden_states, + routing_weights, + selected_experts, + gate_a, + gate_b, + up_a, + up_b, + down_a, + down_b, + experts, + lora_live, + return_routes, + ): + del experts, lora_live + ctx.save_for_backward( + hidden_states, + routing_weights, + gate_a, + gate_b, + up_a, + up_b, + down_a, + down_b, + ) + if return_routes: + return hidden_states.new_zeros( + hidden_states.shape[0], + selected_experts.shape[1], + hidden_states.shape[1], + ) + return hidden_states.new_zeros(hidden_states.shape) + + @staticmethod + def backward(ctx, grad_output): + del grad_output + hidden_states, routing_weights, *factors = ctx.saved_tensors + grad_hidden = torch.zeros_like(hidden_states) if ctx.needs_input_grad[0] else None + grad_routing = torch.zeros_like(routing_weights) if ctx.needs_input_grad[1] else None + factor_grads = [ + torch.zeros_like(factor) if ctx.needs_input_grad[index] else None + for index, factor in enumerate(factors, start=3) + ] + return grad_hidden, grad_routing, None, *factor_grads, None, None, None + + def dsv4_native_mxfp4_routed_partial( hidden_states: torch.Tensor, routing_weights: torch.Tensor, @@ -1106,12 +1173,15 @@ def dsv4_native_mxfp4_routed_partial( experts, *, lora_live: bool = True, + return_routes: bool = False, ) -> torch.Tensor: """Literal local MXFP4-Marlin routed partial with trainable rank-one LoRA. Gather-aware serving installs rank-major LoRA metadata on every EP rank, so the exact caller keeps ``lora_live`` enabled for every local expert-bank partial. Passing ``False`` remains the explicit base-only program. + ``return_routes=True`` exposes unweighted BF16 ``[rows, topk, hidden]`` + expert values; the shared native layer then owns FP32 weighting/reduction. """ if (experts.active_r, experts.active_lora_alpha, experts._active_scaling()) != ( @@ -1123,6 +1193,16 @@ def dsv4_native_mxfp4_routed_partial( factors = [] for projection in ("gate_proj", "up_proj", "down_proj"): factors.extend(experts._active_lora_views(projection)) + if hidden_states.shape[0] == 0: + return _Dsv4NativeMxfp4ZeroRows.apply( + hidden_states, + routing_weights, + selected_experts, + *factors, + experts, + lora_live, + return_routes, + ) return _Dsv4NativeMxfp4RoutedFunction.apply( hidden_states, routing_weights, @@ -1130,6 +1210,7 @@ def dsv4_native_mxfp4_routed_partial( *factors, experts, lora_live, + return_routes, ) diff --git a/src/xorl/models/transformers/glm5/exact_fullparam_experts.py b/src/xorl/models/transformers/glm5/exact_fullparam_experts.py index 7817e553..1b26ffe4 100644 --- a/src/xorl/models/transformers/glm5/exact_fullparam_experts.py +++ b/src/xorl/models/transformers/glm5/exact_fullparam_experts.py @@ -591,23 +591,10 @@ def _sampler_value( ): raise RuntimeError("GLM-5.2 full-param expert caches and activations must share one CUDA device") - from xorl.models.layers.moe.experts import MoEExperts # noqa: PLC0415 - from xorl.ops.moe.sglang_fused_moe_strided import fused_experts_impl_strided # noqa: PLC0415 - - MoEExperts._ensure_sglang_server_args() - return fused_experts_impl_strided( - hidden.contiguous(), - self.gate_up_proj.transpose(1, 2), - self.down_proj.transpose(1, 2), + return self._sglang_ep_native_routed_value( + hidden, routing, local_ids, - activation="silu", - is_gated=True, - use_fp8_w8a8=True, - w1_scale=self.gate_up_weight_scale_inv.transpose(1, 2), - w2_scale=self.down_weight_scale_inv.transpose(1, 2), - block_shape=[128, 128], - filter_expert=True, routed_scaling_factor=routed_scaling_factor, ) diff --git a/src/xorl/models/transformers/glm5/exact_fullparam_fp8.py b/src/xorl/models/transformers/glm5/exact_fullparam_fp8.py index fcd6e405..b3643fc5 100644 --- a/src/xorl/models/transformers/glm5/exact_fullparam_fp8.py +++ b/src/xorl/models/transformers/glm5/exact_fullparam_fp8.py @@ -940,7 +940,7 @@ def backward(ctx, grad_weights: Tensor): return grad_logits.to(router_logits.dtype), None, None, None, None -def glm52_fullparam_routing_weights_with_grad( +def glm52_routing_weights_with_grad( router_logits: Tensor, serving_weights: Tensor, selected_experts: Tensor, @@ -954,7 +954,7 @@ def glm52_fullparam_routing_weights_with_grad( if not _routing_surrogate_engagement_logged: _routing_surrogate_engagement_logged = True logger.info( - "GLM-5.2 full-param routing-weight surrogate engaged: contract=%s " + "GLM-5.2 exact routing-weight surrogate engaged: contract=%s " "(forward = serving values verbatim; backward = analytic regather vjp " "sigmoid-gather%s x%.6g into the router logits)", GLM52_FULLPARAM_ROUTING_SURROGATE_CONTRACT_VERSION, @@ -970,6 +970,11 @@ def glm52_fullparam_routing_weights_with_grad( ) +# Compatibility name for the full-parameter admission/tests. The surrogate +# is contract-owned by exact routing rather than by one optimizer mode. +glm52_fullparam_routing_weights_with_grad = glm52_routing_weights_with_grad + + __all__ = [ "GLM52_EXACT_FULLPARAM_ROUTER_CONTRACT_VERSION", "GLM52_EXACT_TP1_FULLPARAM_FP8_CONTRACT_VERSION", @@ -978,6 +983,7 @@ def glm52_fullparam_routing_weights_with_grad( "Glm52ExactFullParamRouterWeight", "Glm52ExactTP1BlockFP8FullParamLinear", "Glm52FullParamDenseMLP", + "glm52_routing_weights_with_grad", "glm52_fullparam_routing_weights_with_grad", "quantize_expert_masters_to_serving_bytes", "quantize_master_to_serving_bytes", diff --git a/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py b/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py index 2eb74b5c..e28089f0 100644 --- a/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py +++ b/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py @@ -777,6 +777,7 @@ def _validate_operands( temperature: Tensor | None = None, *, require_cuda: bool, + require_factor_grad: bool = True, ) -> None: for name, value in ( ("hidden_states", hidden_states), @@ -828,7 +829,7 @@ def _validate_operands( ) if local_weight.requires_grad: raise RuntimeError("GLM-5.2 exact LM-head base weight must remain frozen") - if not lora_A.requires_grad or not local_lora_B.requires_grad: + if require_factor_grad and (not lora_A.requires_grad or not local_lora_B.requires_grad): raise RuntimeError("GLM-5.2 exact LM-head A and B factor masters must both be trainable") expected_strides = { @@ -1230,6 +1231,7 @@ def forward( token_ids, temperature, require_cuda=True, + require_factor_grad=torch.is_grad_enabled(), ) self._validate_tp_group() return _Glm52ExactTP16LmHeadFunction.apply( @@ -1263,6 +1265,7 @@ def distributed_selected_logprob( local_token_ids, local_temperature, require_cuda=True, + require_factor_grad=torch.is_grad_enabled(), ) self._validate_tp_group() return _Glm52ExactDistributedTP16LmHeadFunction.apply( @@ -1315,6 +1318,8 @@ def glm52_exact_lm_head_per_token_ce( if not is_glm52_exact_tp16_lm_head(lm_head): raise TypeError("glm52_exact_lm_head_per_token_ce requires the constructed exact GLM-5.2 lm_head") + if not lm_head.lora_A.requires_grad or not lm_head.lora_B.requires_grad: + raise RuntimeError("GLM-5.2 exact LM-head A and B logical factor masters must both be trainable") if ce_mode != "bi_fused": raise NotImplementedError("The GLM-5.2 exact active-LoRA lm_head requires ce_mode='bi_fused'") if not lm_head_fp32: diff --git a/src/xorl/models/transformers/glm5/exact_routed_experts_qlora.py b/src/xorl/models/transformers/glm5/exact_routed_experts_qlora.py index eebeb2de..857898ac 100644 --- a/src/xorl/models/transformers/glm5/exact_routed_experts_qlora.py +++ b/src/xorl/models/transformers/glm5/exact_routed_experts_qlora.py @@ -127,6 +127,8 @@ def forward( @staticmethod def backward(ctx, grad_output: Tensor): + if getattr(ctx.module, "_require_routing_grad", False) and not ctx.needs_input_grad[1]: + raise RuntimeError("GLM-5.2 exact routed expert backward received a detached routing tensor") saved = ctx.saved_tensors hidden, routing, local_ids = saved[:3] effective = saved[3:9] @@ -188,10 +190,10 @@ def expert_adapter_gradient_contract(self) -> ExpertAdapterGradientContract: or self.moe_tp_size != 1 or self.active_r != self.r or self.active_lora_alpha != self.lora_alpha - or self.ep_dispatch != "alltoall" + or self.ep_dispatch not in {"alltoall", "deepep"} or self.hybrid_shared is not True ): - raise ValueError("exact EP16 alltoall routed lane geometry no longer matches its declared contract") + raise ValueError("exact EP16 routed lane geometry no longer matches its declared contract") roles = ("gate_proj", "up_proj", "down_proj") return ExpertAdapterGradientContract( backend=replace( @@ -247,8 +249,8 @@ def __init__( ) if moe_tp_size != 1: raise ValueError(f"GLM-5.2 routed experts admit only effective MoE-TP1, got TP{moe_tp_size}") - if ep_dispatch != "alltoall": - raise ValueError(f"GLM-5.2 exact routed experts reject ep_dispatch={ep_dispatch!r}; DeepEP is not admitted") + if ep_dispatch not in {"alltoall", "deepep"}: + raise ValueError(f"GLM-5.2 exact routed experts reject ep_dispatch={ep_dispatch!r}") if not isinstance(ep_rank, int) or isinstance(ep_rank, bool) or not 0 <= ep_rank < ep_size: raise ValueError(f"GLM-5.2 routed expert owner must be in [0, 15], got {ep_rank!r}") @@ -268,7 +270,7 @@ def __init__( self.ep_rank = ep_rank self.expert_offset = ep_rank * self.num_local_experts self.moe_tp_size = 1 - self.ep_dispatch = "alltoall" + self.ep_dispatch = ep_dispatch self.r = self.active_r = r self.lora_alpha = self.active_lora_alpha = lora_alpha self.scaling = scaling @@ -387,8 +389,8 @@ def _owner_local_factor(self, factor: Tensor, *, name: str) -> Tensor: ) def _validate_runtime_contract(self, hidden: Tensor, routing: Tensor, local_ids: Tensor) -> None: - if self.ep_dispatch == "deepep" or self.ep_dispatch != "alltoall": - raise RuntimeError("GLM-5.2 exact routed experts require canonical alltoall and reject DeepEP") + if self.ep_dispatch not in {"alltoall", "deepep"}: + raise RuntimeError("GLM-5.2 exact routed experts require an admitted canonical leaf transport") if (self.num_experts, self.ep_size, self.num_local_experts, self.moe_tp_size) != (256, 16, 16, 1): raise RuntimeError("GLM-5.2 exact routed EP16/MoE-TP1 topology was mutated") if ( diff --git a/src/xorl/models/transformers/glm5/exact_shared_expert_qlora.py b/src/xorl/models/transformers/glm5/exact_shared_expert_qlora.py index 0925ba83..722d63e8 100644 --- a/src/xorl/models/transformers/glm5/exact_shared_expert_qlora.py +++ b/src/xorl/models/transformers/glm5/exact_shared_expert_qlora.py @@ -232,6 +232,118 @@ def backward(ctx, grad_output: Tensor): return (*gradients, None, None) +class _Glm52ExactTP16SharedExpertAllContributorsFunction(torch.autograd.Function): + """Emit every physical TP16 leaf from one FSDP-root invocation. + + Native DeepEP already returns the routed result on the owning CP rank. A + second distributed shared-expert program is unnecessary: while this root + is materialized, the owner can execute the same sixteen physical sampler + shards over its local rows and canonical-fold their BF16 leaves locally. + Keeping the loop inside one custom Function is important because the + shared-expert root, rather than its projection children, owns the FSDP + unshard/reshard boundary. + """ + + @staticmethod + def forward( + ctx, + input: Tensor, + gate_A: Tensor, + gate_B: Tensor, + up_A: Tensor, + up_B: Tensor, + down_A: Tensor, + down_B: Tensor, + module, + ) -> Tensor: + effective = tuple( + factor.to(torch.bfloat16).contiguous() for factor in (gate_A, gate_B, up_A, up_B, down_A, down_B) + ) + values = [ + module._exact_forward_value(input, *effective, contributor_ordinal=ordinal) + for ordinal in range(module.tp_size) + ] + output = torch.stack( + [value.output.view(*input.shape[:-1], module.hidden_size) for value in values], + dim=0, + ) + expected_shape = (module.tp_size, *input.shape[:-1], module.hidden_size) + if output.dtype is not torch.bfloat16: + raise TypeError(f"GLM-5.2 exact TP16 shared-expert leaves have {output.dtype}, expected torch.bfloat16") + if tuple(output.shape) != expected_shape: + raise RuntimeError( + f"GLM-5.2 exact TP16 shared-expert leaf shape {tuple(output.shape)} does not match {expected_shape}" + ) + + ctx.module = module + # Masters retain normal autograd version checks across the optimizer + # boundary; exact intermediates are stacked in immutable TP ordinal + # order for the staged surrogate VJP. + ctx.save_for_backward( + input.detach(), + *effective, + gate_A, + gate_B, + up_A, + up_B, + down_A, + down_B, + torch.stack([value.gate_up for value in values], dim=0), + torch.stack([value.activated for value in values], dim=0), + ) + return output + + @staticmethod + def backward(ctx, grad_output: Tensor): + ( + input, + effective_gate_A, + effective_gate_B, + effective_up_A, + effective_up_B, + effective_down_A, + effective_down_B, + _gate_A_master, + _gate_B_master, + _up_A_master, + _up_B_master, + _down_A_master, + _down_B_master, + exact_gate_up, + exact_activated, + ) = ctx.saved_tensors + accumulated: list[Tensor | None] = [None] * 7 + for ordinal in range(ctx.module.tp_size): + gradients = ctx.module._surrogate_vjp( + input, + effective_gate_A, + effective_gate_B, + effective_up_A, + effective_up_B, + effective_down_A, + effective_down_B, + exact_gate_up[ordinal], + exact_activated[ordinal], + grad_output[ordinal], + contributor_ordinal=ordinal, + needs_input_grad=ctx.needs_input_grad[:7], + ) + for index, gradient in enumerate(gradients): + if gradient is None: + continue + if accumulated[index] is None: + accumulated[index] = gradient.float() + else: + accumulated[index] = accumulated[index] + gradient.float() + + # Input is the only low-precision differentiable operand. Contributor + # VJPs accumulate in FP32 and cross that BF16 boundary exactly once; + # all six logical factor gradients remain FP32 masters. + if accumulated[0] is not None: + accumulated[0] = accumulated[0].to(input.dtype) + return (*accumulated, None) + + class Glm52ExactTP16SharedExpertBlockFP8QLoRA(nn.Module): """Produce one unreduced sampler-equivalent GLM-5.2 shared-expert partial.""" @@ -713,7 +825,29 @@ def _surrogate_vjp( return grad_input, grad_gate_A, grad_gate_B, grad_up_A, grad_up_B, grad_down_A, grad_down_B - def forward(self, input: Tensor, *, contributor_ordinal: int) -> Tensor: + def forward( + self, + input: Tensor, + *, + contributor_ordinal: int | None = None, + all_contributors: bool = False, + ) -> Tensor: + if all_contributors: + if contributor_ordinal is not None: + raise ValueError("GLM-5.2 shared expert cannot select one contributor while all_contributors=True") + self._validate_engaged_contract(input, 0) + return _Glm52ExactTP16SharedExpertAllContributorsFunction.apply( + input, + self.gate_proj.lora_A, + self.gate_proj.lora_B, + self.up_proj.lora_A, + self.up_proj.lora_B, + self.down_proj.lora_A, + self.down_proj.lora_B, + self, + ) + if contributor_ordinal is None: + raise TypeError("GLM-5.2 shared expert requires contributor_ordinal unless all_contributors=True") ordinal = self._validate_engaged_contract(input, contributor_ordinal) return _Glm52ExactTP16SharedExpertFunction.apply( input, diff --git a/src/xorl/models/transformers/glm5/indexer.py b/src/xorl/models/transformers/glm5/indexer.py index 19effa47..53943889 100644 --- a/src/xorl/models/transformers/glm5/indexer.py +++ b/src/xorl/models/transformers/glm5/indexer.py @@ -103,6 +103,7 @@ def _mix_sampler_index_k_preparation( eps: float, position_embeddings: tuple[torch.Tensor, torch.Tensor], sampler_prefill_lengths: torch.Tensor | None, + sample_lengths: list[int] | tuple[int, ...] | torch.Tensor | None = None, *, query_offset: int, interleaved: bool, @@ -112,45 +113,107 @@ def _mix_sampler_index_k_preparation( if sampler_prefill_lengths is None: return split_prepared_key - if sampler_prefill_lengths.ndim != 1 or sampler_prefill_lengths.numel() != raw_key.shape[0]: - raise ValueError( - "GLM-5.2 sampler_prefill_lengths must contain one boundary per batch row, " - f"got {tuple(sampler_prefill_lengths.shape)} for batch={raw_key.shape[0]}" - ) + if sampler_prefill_lengths.ndim != 1: + raise ValueError("GLM-5.2 sampler_prefill_lengths must be a rank-one tensor") if sampler_prefill_lengths.dtype not in (torch.int32, torch.int64): raise TypeError("GLM-5.2 sampler_prefill_lengths must be an integer tensor") if query_offset < 0: raise ValueError(f"GLM-5.2 exact indexer query_offset must be nonnegative, got {query_offset}") + if isinstance(sample_lengths, torch.Tensor): + sample_lengths = sample_lengths.detach().cpu().reshape(-1).tolist() + if sample_lengths is not None: + sample_lengths = [int(length) for length in sample_lengths] + if raw_key.shape[0] != 1: + raise ValueError( + f"GLM-5.2 packed sampler metadata requires one packed batch row, got batch={raw_key.shape[0]}" + ) + if len(sample_lengths) != sampler_prefill_lengths.numel(): + raise ValueError( + "GLM-5.2 sampler metadata requires one prefill boundary per packed request: " + f"prefills={sampler_prefill_lengths.numel()} sample_lengths={sample_lengths}" + ) + if any(length <= 0 for length in sample_lengths): + raise ValueError(f"GLM-5.2 packed request lengths must be positive, got {sample_lengths}") + elif sampler_prefill_lengths.numel() != raw_key.shape[0]: + raise ValueError( + "GLM-5.2 sampler_prefill_lengths must contain one boundary per batch row when " + "packed request lengths are absent, " + f"got {tuple(sampler_prefill_lengths.shape)} for batch={raw_key.shape[0]}" + ) + cos, sin = position_embeddings - mixed_rows = [] local_length = raw_key.shape[1] - for batch_index in range(raw_key.shape[0]): - prefill_length = int(sampler_prefill_lengths[batch_index].item()) - if prefill_length <= 0: - raise ValueError(f"GLM-5.2 sampler prefill length must be positive, got {prefill_length}") - suffix_start = min(max(prefill_length - query_offset, 0), local_length) - if suffix_start == local_length: - mixed_rows.append(split_prepared_key[batch_index : batch_index + 1]) - continue + + def _mix_slice(batch_index: int, local_start: int, local_end: int, prefill_end: int) -> torch.Tensor: + split_slice = split_prepared_key[batch_index : batch_index + 1, local_start:local_end] + suffix_start = min(max(prefill_end - query_offset, local_start), local_end) + if suffix_start == local_end: + return split_slice fused_suffix = _fused_sampler_index_k_prepare( - raw_key[batch_index : batch_index + 1, suffix_start:], + raw_key[batch_index : batch_index + 1, suffix_start:local_end], norm_weight, norm_bias, eps, ( - cos[batch_index : batch_index + 1, suffix_start:], - sin[batch_index : batch_index + 1, suffix_start:], + cos[batch_index : batch_index + 1, suffix_start:local_end], + sin[batch_index : batch_index + 1, suffix_start:local_end], ), interleaved=interleaved, _native_kernel_for_testing=_native_kernel_for_testing, ) - mixed_rows.append( - torch.cat( - (split_prepared_key[batch_index : batch_index + 1, :suffix_start], fused_suffix), - dim=1, + return torch.cat((split_slice[:, : suffix_start - local_start], fused_suffix), dim=1) + + if sample_lengths is not None: + mixed_slices = [] + request_start = 0 + local_global_start = query_offset + local_global_end = query_offset + local_length + for sample_length, prefill_value in zip( + sample_lengths, + sampler_prefill_lengths.detach().cpu().tolist(), + strict=True, + ): + prefill_length = int(prefill_value) + request_end = request_start + sample_length + if prefill_length <= 0 or prefill_length > sample_length: + raise ValueError( + "GLM-5.2 sampler prefill length must satisfy 0 < prefill <= request length, " + f"got prefill={prefill_length} request_length={sample_length}" + ) + overlap_start = max(request_start, local_global_start) + overlap_end = min(request_end, local_global_end) + if overlap_start < overlap_end: + local_start = overlap_start - local_global_start + local_end = overlap_end - local_global_start + mixed_slices.append( + _mix_slice( + 0, + local_start, + local_end, + request_start + prefill_length, + ) + ) + request_start = request_end + if request_start < local_global_end: + padding_start = max(request_start - local_global_start, 0) + mixed_slices.append(split_prepared_key[:, padding_start:]) + if not mixed_slices: + return split_prepared_key + mixed = torch.cat(mixed_slices, dim=1) + if mixed.shape[1] != local_length: + raise ValueError( + "GLM-5.2 packed request lengths do not cover the local indexer rows: " + f"sample_lengths={sample_lengths} query_offset={query_offset} local_length={local_length}" ) - ) + return mixed + + mixed_rows = [] + for batch_index in range(raw_key.shape[0]): + prefill_length = int(sampler_prefill_lengths[batch_index].item()) + if prefill_length <= 0: + raise ValueError(f"GLM-5.2 sampler prefill length must be positive, got {prefill_length}") + mixed_rows.append(_mix_slice(batch_index, 0, local_length, prefill_length)) return torch.cat(mixed_rows, dim=0) @@ -251,6 +314,7 @@ def project( position_embeddings: tuple[torch.Tensor, torch.Tensor], *, sampler_prefill_lengths: torch.Tensor | None = None, + sample_lengths: list[int] | tuple[int, ...] | torch.Tensor | None = None, query_offset: int = 0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Compute index_query, index_key, head_weights. @@ -347,6 +411,7 @@ def project( self.k_norm.eps, position_embeddings, sampler_prefill_lengths, + sample_lengths, query_offset=query_offset, interleaved=getattr(self.config, "indexer_rope_interleave", True), ) @@ -650,6 +715,7 @@ def forward( position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: torch.Tensor | None = None, sampler_prefill_lengths: torch.Tensor | None = None, + sample_lengths: list[int] | tuple[int, ...] | torch.Tensor | None = None, ) -> torch.Tensor: """One-shot indexer: project + score + top-k. @@ -665,6 +731,7 @@ def forward( q_compressed, position_embeddings, sampler_prefill_lengths=sampler_prefill_lengths, + sample_lengths=sample_lengths, ) return self.select_topk(index_q, index_k, head_weights, attention_mask) diff --git a/src/xorl/models/transformers/glm5/modeling_glm5.py b/src/xorl/models/transformers/glm5/modeling_glm5.py index e1720171..ccd502cf 100644 --- a/src/xorl/models/transformers/glm5/modeling_glm5.py +++ b/src/xorl/models/transformers/glm5/modeling_glm5.py @@ -15,6 +15,7 @@ LogicalRowOwnership, OutputDistribution, ParallelPlan, + canonical_moe_fold_fp64_v3, canonical_moe_leaf_fp32_v1, canonical_moe_reduce_cp_sharded_v3, canonical_moe_reduce_fp64_v3, @@ -38,6 +39,7 @@ from xorl.models.layers.attention.backend import ATTENTION_FUNCTIONS from xorl.models.layers.attention.backend.eager import eager_attention_forward from xorl.models.layers.moe import MoEBlock +from xorl.models.layers.moe.backend import zero_token_lora_output from xorl.models.layers.moe.experts import MoEExperts from xorl.models.layers.moe.moe_block import _BIRouterGemm, _moe_bi_router_enabled from xorl.models.layers.moe.routing_replay import get_replay_stage @@ -76,6 +78,8 @@ logger = logging.get_logger(__name__) GLM52_LOCAL_PARTIAL_POLICY = "glm52_routed_final_scaled_then_shared_ep_slice_fp32_then_bf16_v3" +GLM52_NATIVE_SHARED_PARTIAL_POLICY = "glm52_native_deepep_shared_tp_slice_fp32_then_bf16_v1" +_glm52_native_deepep_engagement_logged = False def _glm52_serving_grouped_topk( @@ -155,6 +159,46 @@ def __init__(self, config: Glm5Config): self.hidden_size = config.hidden_size self.weight = nn.Parameter(torch.empty(config.n_routed_experts, config.hidden_size)) self.register_buffer("e_score_correction_bias", torch.zeros(config.n_routed_experts)) + self._gradient_evidence_parameter_id: int | None = None + self._gradient_evidence_hook_handle = None + self._gradient_evidence: dict[str, torch.Tensor | int] | None = None + + def reset_gradient_evidence(self) -> None: + """Arm a hook on the current post-FSDP router Parameter.""" + + self._gradient_evidence = None + if not self.weight.requires_grad: + return + if self._gradient_evidence_parameter_id == id(self.weight): + return + self._gradient_evidence_hook_handle = self.weight.register_hook(self._capture_gradient_evidence) + self._gradient_evidence_parameter_id = id(self.weight) + + def _capture_gradient_evidence(self, gradient: torch.Tensor) -> torch.Tensor: + local = gradient + to_local = getattr(local, "to_local", None) + if callable(to_local): + local = to_local() + wait = getattr(local, "wait", None) + if callable(wait): + local = wait() + values = local.detach().float() + finite = torch.isfinite(values) + finite_values = values.masked_select(finite) + evidence = { + "sum_sq": torch.sum(finite_values * finite_values, dtype=torch.float64), + "nonfinite": torch.count_nonzero(~finite), + "nonzero": torch.count_nonzero(finite_values), + "elements": values.numel(), + } + previous = self._gradient_evidence + if previous is not None: + evidence["sum_sq"] = evidence["sum_sq"] + previous["sum_sq"] + evidence["nonfinite"] = evidence["nonfinite"] + previous["nonfinite"] + evidence["nonzero"] = evidence["nonzero"] + previous["nonzero"] + evidence["elements"] = int(evidence["elements"]) + int(previous["elements"]) + self._gradient_evidence = evidence + return gradient def _apply(self, fn, recurse: bool = True): correction_bias = self._buffers.pop("e_score_correction_bias") @@ -596,6 +640,7 @@ def forward_sparse( attention_mask: torch.Tensor | None, index_share_context: IndexShareContext | None = None, sampler_prefill_lengths: torch.Tensor | None = None, + _r3_sample_lengths: list[int] | tuple[int, ...] | torch.Tensor | None = None, **_kwargs, ) -> tuple[torch.Tensor, None]: ps = get_parallel_state() @@ -609,6 +654,7 @@ def forward_sparse( position_embeddings, attention_mask, sampler_prefill_lengths=sampler_prefill_lengths, + sample_lengths=_r3_sample_lengths, ), index_share_context, ) @@ -637,6 +683,7 @@ def compute_local_indices() -> torch.Tensor: q_compressed, position_embeddings, sampler_prefill_lengths=sampler_prefill_lengths, + sample_lengths=_r3_sample_lengths, query_offset=query_offset, ) full_index_k = self._gather_ulysses_sequence_no_grad(local_index_k, group) @@ -781,6 +828,7 @@ def __init__(self, config: Glm5Config, *, layer_idx: int | None = None): self.experts.deepep_buffer_size_gb = getattr(config, "_deepep_buffer_size_gb", 2.0) self.experts.deepep_num_sms = getattr(config, "_deepep_num_sms", 20) self.experts.deepep_async_combine = getattr(config, "_deepep_async_combine", False) + self.deepep_native_exact = bool(getattr(config, "_deepep_native_exact", False)) self.n_group = config.n_group self.topk_group = config.topk_group self.norm_topk_prob = config.norm_topk_prob @@ -804,6 +852,15 @@ def __init__(self, config: Glm5Config, *, layer_idx: int | None = None): selectable_experts = self.topk_group * (self.num_experts // self.n_group) if not 1 <= self.top_k <= selectable_experts: raise ValueError("GLM-5.2 canonical routing top_k exceeds the selected expert groups") + if self.deepep_native_exact: + if self.canonical_contract_version is None: + raise RuntimeError("GLM-5.2 native DeepEP requires the canonical exact contract") + if self.experts.ep_dispatch != "deepep": + raise RuntimeError("GLM-5.2 native DeepEP requires ep_dispatch='deepep'") + if self.train_router: + raise RuntimeError("GLM-5.2 native DeepEP v1 requires a frozen router") + if self.experts.deepep_async_combine: + raise RuntimeError("GLM-5.2 native DeepEP owns the immediate FP64 fold") def route(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: flat_hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) @@ -812,6 +869,8 @@ def route(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor if self.canonical_contract_version is not None and not _moe_bi_router_enabled(self.config): raise RuntimeError("GLM-5.2 canonical MoE requires its model-level exact router declaration") + if self.train_router and torch.is_grad_enabled() and not self.gate.weight.requires_grad: + raise RuntimeError("GLM-5.2 trainable router weight lost requires_grad before the BI router GEMM") router_logits = self.gate(flat_hidden_states) if stage is not None and replay is not None: @@ -871,26 +930,28 @@ def route(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor ) ep_dispatch = getattr(self.experts, "ep_dispatch", "alltoall") - if self.train_router and ep_dispatch == "deepep": + if self.canonical_contract_version is None and self.train_router and ep_dispatch == "deepep": raise AssertionError( "train_router=True is not supported with ep_dispatch='deepep'. " "DeepEP cannot propagate gradients through routing weights. " "Set train_router=False or switch to ep_dispatch='alltoall'." ) - if getattr(self.gate, "_glm52_exact_fullparam_component", False): + if getattr(self.gate, "_glm52_exact_fullparam_component", False) or ( + self.canonical_contract_version is not None and self.train_router + ): # Full-param mode trains routers: their only gradient # path is the combine multiply, so the routing weights must carry # a gradient into the router logits. The serving top-k programs # are gradient-opaque, so the values are kept verbatim and the # gradient is supplied by the trainer-owned regather surrogate. - if ep_dispatch == "deepep": + if self.canonical_contract_version is None and ep_dispatch == "deepep": raise AssertionError("GLM-5.2 full-param router training is not supported with ep_dispatch='deepep'") from xorl.models.transformers.glm5.exact_fullparam_fp8 import ( # noqa: PLC0415 - glm52_fullparam_routing_weights_with_grad, + glm52_routing_weights_with_grad, ) canonical = self.canonical_contract_version is not None - routing_weights = glm52_fullparam_routing_weights_with_grad( + routing_weights = glm52_routing_weights_with_grad( router_logits, routing_weights, selected_experts, @@ -900,9 +961,14 @@ def route(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor renormalize=True if canonical else bool(self.norm_topk_prob), scale=1.0 if canonical else float(self.routed_scaling_factor), ) + if self.train_router and torch.is_grad_enabled() and not routing_weights.requires_grad: + raise RuntimeError("GLM-5.2 router surrogate failed to attach routing-weight autograd") elif not self.train_router: routing_weights = routing_weights.detach() + self._capture_diagnostic_component("moe_router_logits", router_logits) + self._capture_diagnostic_component("moe_topk_ids", selected_experts) + self._capture_diagnostic_component("moe_topk_weights", routing_weights) return routing_weights, selected_experts, router_logits def _route_tokens_to_experts( @@ -969,16 +1035,21 @@ def forward_experts_with_shared( routing_weights: torch.Tensor, selected_experts: torch.Tensor, absolute_positions: torch.Tensor | None = None, + backward_layer_dependency: torch.Tensor | None = None, ) -> torch.Tensor: + self._capture_diagnostic_component("moe_input", hidden_states) if self.canonical_contract_version is not None: if absolute_positions is None: raise RuntimeError("GLM-5.2 canonical MoE execution requires explicit absolute positions") - return self._canonical_ep_forward( + output = self._canonical_ep_forward( hidden_states, routing_weights, selected_experts, absolute_positions, + backward_layer_dependency, ) + self._capture_diagnostic_component("moe_experts_output", output) + return output residuals = hidden_states batch_size, sequence_length, hidden_dim = hidden_states.shape @@ -992,7 +1063,9 @@ def forward_experts_with_shared( expert_output = expert_output.view(batch_size, sequence_length, hidden_dim) shared_output = self.shared_experts(residuals) sync_pending_combine() - return expert_output + shared_output + output = expert_output + shared_output + self._capture_diagnostic_component("moe_experts_output", output) + return output def _canonical_expert_slice(self, ep_rank: int, ep_size: int) -> tuple[int, int]: """Resolve this rank's contiguous expert slice for the canonical dispatch. @@ -1054,7 +1127,7 @@ def _canonical_routed_local_partial( self, gathered: torch.Tensor, gathered_routing: torch.Tensor, - gathered_ids: torch.Tensor, + gathered_ids: torch.Tensor | None, local_ids: torch.Tensor, ) -> torch.Tensor: if not isinstance(self.experts, Glm52NativeBlockFP8Experts): @@ -1063,7 +1136,7 @@ def _canonical_routed_local_partial( "sglang_ep_native_local_ids": local_ids, "routed_scaling_factor": self.routed_scaling_factor, } - if isinstance(self.experts, Glm52ExactEP16BlockFP8QLoRARoutedExperts): + if isinstance(self.experts, Glm52ExactEP16BlockFP8QLoRARoutedExperts) and gathered_ids is not None: kwargs["selected_experts"] = gathered_ids return self.experts(gathered, gathered_routing, **kwargs).to(torch.bfloat16) @@ -1116,12 +1189,164 @@ def _canonical_shared_local_partial( self.shared_experts.down_proj.weight[:, shard_start:shard_end], ).to(torch.bfloat16) + def _native_shared_local_fold( + self, + flat: torch.Tensor, + valid_rows: torch.Tensor, + ) -> torch.Tensor: + """Execute and FP64-fold all sampler TP16 shared leaves locally.""" + + if not isinstance(self.shared_experts, Glm52ExactTP16SharedExpertBlockFP8QLoRA): + raise RuntimeError( + "GLM-5.2 native DeepEP requires the exact active-LoRA shared-expert root " + "to produce all TP16 leaves locally" + ) + # Invoke the root through nn.Module.__call__: its FSDP hooks must + # materialize all six logical FP32 factor masters exactly once around + # the complete TP16 leaf program. + shared_leaves = self.shared_experts(flat, all_contributors=True) + self._capture_diagnostic_component("moe_native_shared_down", shared_leaves) + shared_folded = canonical_moe_fold_fp64_v3(shared_leaves) + shared_folded = torch.where( + valid_rows[:, None], + shared_folded, + torch.zeros_like(shared_folded), + ) + self._capture_diagnostic_component("moe_native_shared_folded", shared_folded) + return shared_folded + + def _native_deepep_routed_local( + self, + flat: torch.Tensor, + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, + valid_rows: torch.Tensor, + *, + ep_rank: int, + ep_size: int, + ep_group, + backward_layer_dependency: torch.Tensor | None = None, + backward_shared_dependency: torch.Tensor | None = None, + ) -> torch.Tensor: + """Run only GLM's routed branch through the original DeepEP handle. + + GLM's TP16 shared expert is a separate arithmetic owner. It remains on + the canonical BF16-leaf/FP64-fold path below; folding it into the + routed payload would either replicate the shared branch or silently + change its TP16 reduction tree. + """ + + from xorl.distributed.moe.deepep_native_exact import ( # noqa: PLC0415 + canonicalize_native_routing_metadata, + native_dispatch_runner_combine, + ) + + if self.train_router: + raise RuntimeError("GLM-5.2 native DeepEP v1 requires a frozen router") + live_indices = torch.nonzero(valid_rows, as_tuple=False).flatten() + live_hidden = flat.index_select(0, live_indices) + live_routing = canonicalize_native_routing_metadata(routing_weights.index_select(0, live_indices)) + live_ids = selected_experts.index_select(0, live_indices) + num_local_experts, expert_start = self._canonical_expert_slice(ep_rank, ep_size) + + def run_local_leaf(recv_hidden, recv_weights, recv_local_ids): + # A rank may own no route for this dispatch. The fused GLM expert + # runner intentionally rejects zero rows, while DeepEP considers + # an empty receive batch valid. Keep the empty BF16 leaf connected + # to every active factor: bypassing the expert module here would + # leave ``parameter.grad is None`` on only the empty owner and make + # the public adapter-gradient collective rank-asymmetric. + if recv_hidden.shape[0] == 0: + factors = tuple(getattr(self.experts, name) for name in self.experts.logical_factor_names) + return zero_token_lora_output( + recv_hidden, + recv_hidden.shape[1], + *factors, + ) + + # The diagnostic seam is absent in normal training. When a + # component capture is installed, retain the real DeepEP receive + # receipt and only the expert rows it actually selected. This + # distinguishes transport/id remapping from a loaded-weight or + # fused-kernel mismatch without dumping the full 256-expert bank. + diagnostic_capture = getattr(self, "_diagnostic_capture_component", None) + if callable(diagnostic_capture): + self._capture_diagnostic_component("moe_native_recv_hidden", recv_hidden) + self._capture_diagnostic_component("moe_native_recv_weights", recv_weights) + self._capture_diagnostic_component("moe_native_recv_local_ids", recv_local_ids) + self._capture_diagnostic_component( + "moe_native_expert_start", + torch.tensor([expert_start], dtype=torch.int64, device=recv_hidden.device), + ) + local_ids = torch.unique(recv_local_ids[recv_local_ids >= 0]).tolist() + for local_id in local_ids: + local_id = int(local_id) + self._capture_diagnostic_component( + f"moe_native_gate_up_packed_local_{local_id}", + self.experts.gate_up_packed_weight_f32[local_id], + ) + self._capture_diagnostic_component( + f"moe_native_gate_up_scale_local_{local_id}", + self.experts.gate_up_weight_scale_inv[local_id], + ) + self._capture_diagnostic_component( + f"moe_native_down_packed_local_{local_id}", + self.experts.down_packed_weight_f32[local_id], + ) + self._capture_diagnostic_component( + f"moe_native_down_scale_local_{local_id}", + self.experts.down_weight_scale_inv[local_id], + ) + + # DeepEP already maps every valid slot into this owner's local + # [0, num_local_experts) range and marks all other top-k slots -1. + # Do not reconstruct global IDs here: the QLoRA bank's optional + # global-ID argument has a stricter no-sentinel contract intended + # for the all-gather path. The native path is fully identified by + # the validated local IDs and does not need that redundant input. + local_leaf = self._canonical_routed_local_partial( + recv_hidden, + recv_weights, + None, + recv_local_ids, + ) + self._capture_diagnostic_component("moe_native_recv_leaf", local_leaf) + return local_leaf + + live_routed = native_dispatch_runner_combine( + live_hidden, + live_routing, + live_ids, + ep_group=ep_group, + num_experts=self.num_experts, + num_local_experts=num_local_experts, + buffer_size_gb=self.experts.deepep_buffer_size_gb, + num_sms=self.experts.deepep_num_sms, + runner=run_local_leaf, + backward_layer_dependency=backward_layer_dependency, + backward_shared_dependency=backward_shared_dependency, + backward_trace_label=f"glm52_layer_{getattr(self, 'layer_idx', 'unknown')}", + complete_backward_device_boundary=True, + ) + routed = flat * 0.0 + if live_indices.numel(): + routed = routed.index_copy(0, live_indices, live_routed) + else: + # The empty result still owns DeepEP reverse-dispatch and expert + # factor autograd edges. Keep them in the graph so every EP rank + # enters the same backward collectives as its non-empty peers. + routed = routed + live_routed.sum() * 0.0 + routed = routed.to(torch.bfloat16) + self._capture_diagnostic_component("moe_native_routed", routed) + return routed + def _canonical_ep_forward( self, hidden_states: torch.Tensor, routing_weights: torch.Tensor, selected_experts: torch.Tensor, absolute_positions: torch.Tensor, + backward_layer_dependency: torch.Tensor | None = None, ) -> torch.Tensor: from xorl.models.layers.moe.ep_native_combine import ( # noqa: PLC0415 gather_ids_for_ep_combine, @@ -1158,8 +1383,6 @@ def _canonical_ep_forward( ) elif ps.sp_group is not None: raise RuntimeError("GLM-5.2 CP1 rows must not have a sequence-parallel process group") - if self.experts.ep_dispatch == "deepep": - raise RuntimeError("GLM-5.2 canonical MoE canonical path does not support DeepEP") if hidden_states.dtype is not torch.bfloat16: raise TypeError("GLM-5.2 canonical local partials require BF16 hidden states") if self.config.hidden_act != "silu": @@ -1175,32 +1398,92 @@ def _canonical_ep_forward( ) group = ps.ep_group - padded_rows = max_rows_for_ep_combine(local_rows, flat.device, group) - gathered = gather_tokens_for_ep_combine(flat, group, padded_rows) - gathered_routing = gather_tokens_for_ep_combine(routing_weights.reshape(local_rows, -1), group, padded_rows) - gathered_ids = gather_ids_for_ep_combine(selected_experts.reshape(local_rows, -1), group, padded_rows) - gathered_positions = gather_ids_for_ep_combine(local_positions[:, None], group, padded_rows).squeeze(-1) local_valid = LogicalRowOwnership.valid_positions(local_positions).to(torch.int32)[:, None] - gathered_valid = gather_ids_for_ep_combine(local_valid, group, padded_rows).squeeze(-1) > 0 - ep_rank = dist.get_rank(group) if ep_rank != ownership.source_ordinal: raise RuntimeError( "GLM-5.2 EP rank order does not match the DP-major/CP-minor row layout: " f"ep_rank={ep_rank}, source_ordinal={ownership.source_ordinal}" ) - local_experts, expert_start = self._canonical_expert_slice(ep_rank, ps.ep_size) - local_ids = torch.where( - (gathered_ids >= expert_start) & (gathered_ids < expert_start + local_experts), - gathered_ids - expert_start, - gathered_ids.new_full((), -1), - ).to(torch.int32) - routed = self._canonical_routed_local_partial(gathered, gathered_routing, gathered_ids, local_ids) + + native_routed_local = None + if self.deepep_native_exact: + # Build the local shared value first, then make it (and the + # transformer residual) ordering operands of the routed dispatch. + # Their values do not enter DeepEP. In backward, both edges are + # released only after the routed terminal reverse-combine reaches + # device completion, so shared-root FSDP cannot overlap it. + shared_folded = self._native_shared_local_fold( + flat, + local_valid.squeeze(-1).to(torch.bool), + ) + native_routed_local = self._native_deepep_routed_local( + flat, + routing_weights.reshape(local_rows, -1), + selected_experts.reshape(local_rows, -1), + local_valid.squeeze(-1) > 0, + ep_rank=ep_rank, + ep_size=ps.ep_size, + ep_group=group, + backward_layer_dependency=backward_layer_dependency, + backward_shared_dependency=shared_folded, + ) + local_canonical = canonical_moe_leaf_fp32_v1( + shared_folded, + native_routed_local, + ) + self._capture_diagnostic_component("moe_native_combined", local_canonical) + global _glm52_native_deepep_engagement_logged + if not _glm52_native_deepep_engagement_logged: + logger.info( + "GLM-5.2 trainer native DeepEP split ENGAGED: " + "routed=native_deepep_original_handle wire_dtype=bf16 " + "routed_fold=canonical_moe_fold_fp64_v3 " + "shared=local_tp16_bf16_leaves_fp64_v3 " + "shared_internal_collectives=none " + "join=canonical_moe_leaf_fp32_v1 " + "backward=routed_terminal_before_shared_and_residual_v6" + ) + _glm52_native_deepep_engagement_logged = True + return local_canonical.reshape(batch_size, sequence_length, hidden_dim) + + padded_rows = max_rows_for_ep_combine(local_rows, flat.device, group) + + gathered = gather_tokens_for_ep_combine( + flat, + group, + padded_rows, + ) + gathered_routing = None + gathered_ids = None + if not self.deepep_native_exact: + gathered_routing = gather_tokens_for_ep_combine(routing_weights.reshape(local_rows, -1), group, padded_rows) + if self.train_router and torch.is_grad_enabled() and not gathered_routing.requires_grad: + raise RuntimeError("GLM-5.2 EP routing gather severed router autograd") + gathered_ids = gather_ids_for_ep_combine(selected_experts.reshape(local_rows, -1), group, padded_rows) + gathered_positions = gather_ids_for_ep_combine(local_positions[:, None], group, padded_rows).squeeze(-1) + gathered_valid = gather_ids_for_ep_combine(local_valid, group, padded_rows).squeeze(-1) > 0 + + if self.deepep_native_exact: + routed = gathered * 0.0 + else: + assert gathered_routing is not None and gathered_ids is not None + local_experts, expert_start = self._canonical_expert_slice(ep_rank, ps.ep_size) + local_ids = torch.where( + (gathered_ids >= expert_start) & (gathered_ids < expert_start + local_experts), + gathered_ids - expert_start, + gathered_ids.new_full((), -1), + ).to(torch.int32) + self.experts._require_routing_grad = bool(self.train_router and torch.is_grad_enabled()) + routed = self._canonical_routed_local_partial(gathered, gathered_routing, gathered_ids, local_ids) + if self.train_router and torch.is_grad_enabled() and not routed.requires_grad: + raise RuntimeError("GLM-5.2 exact routed expert output severed router autograd") shared = self._canonical_shared_local_partial( gathered, contributor_ordinal=ep_rank, contributor_count=ps.ep_size, ) + self._capture_diagnostic_component("moe_native_shared_down", shared) local_partial = canonical_moe_leaf_fp32_v1(shared, routed) capacity = int(getattr(self.config, "_glm52_canonical_moe_capacity", local_partial.shape[0])) @@ -1231,7 +1514,11 @@ def _canonical_ep_forward( capacity=capacity, valid_rows=int(gathered_valid.sum().item()), ) - contribution = LocalMoEContribution(local_partial, metadata, GLM52_LOCAL_PARTIAL_POLICY) + contribution = LocalMoEContribution( + local_partial, + metadata, + (GLM52_NATIVE_SHARED_PARTIAL_POLICY if self.deepep_native_exact else GLM52_LOCAL_PARTIAL_POLICY), + ) ep_mesh = ps.ep_fsdp_device_mesh ep_dim = tuple(ep_mesh.mesh_dim_names).index("ep") combine_groups = tuple( @@ -1310,6 +1597,7 @@ def forward( hidden_states: torch.Tensor, *, absolute_positions: torch.Tensor | None = None, + backward_layer_dependency: torch.Tensor | None = None, ): routing_weights, selected_experts, router_logits = self.route(hidden_states) hidden_states = self.forward_experts_with_shared( @@ -1317,6 +1605,7 @@ def forward( routing_weights, selected_experts, absolute_positions, + backward_layer_dependency, ) return hidden_states, router_logits @@ -1455,6 +1744,7 @@ def forward( routing_weights, selected_experts, local_absolute_positions, + residual, ) elif _selective: hidden_states, residual, self_attn_weights = self._gradient_checkpointing_func( @@ -1467,7 +1757,14 @@ def forward( ) hidden_states = self.mlp( hidden_states, - **({"absolute_positions": local_absolute_positions} if isinstance(self.mlp, Glm5MoEBlock) else {}), + **( + { + "absolute_positions": local_absolute_positions, + "backward_layer_dependency": residual, + } + if isinstance(self.mlp, Glm5MoEBlock) + else {} + ), ) if isinstance(hidden_states, tuple): hidden_states, router_logits = hidden_states @@ -1483,7 +1780,14 @@ def forward( ) hidden_states = self.mlp( hidden_states, - **({"absolute_positions": local_absolute_positions} if isinstance(self.mlp, Glm5MoEBlock) else {}), + **( + { + "absolute_positions": local_absolute_positions, + "backward_layer_dependency": residual, + } + if isinstance(self.mlp, Glm5MoEBlock) + else {} + ), ) if isinstance(hidden_states, tuple): hidden_states, router_logits = hidden_states @@ -1791,6 +2095,14 @@ def forward( class Glm5ForCausalLM(Glm5PreTrainedModel): + deepep_native_exact_capability = { + "produces_local_leaf": True, + "wire_dtype": "bf16", + "uses_dispatch_handle": True, + "supported_ep_sizes": (2, 4, 8, 16), + "local_leaf_program": "glm52_scaled", + "lora_serving_modes": ("separate",), + } _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _tp_plan = parallelize.MODEL_TP_PLAN diff --git a/src/xorl/models/transformers/glm5/native_fp8.py b/src/xorl/models/transformers/glm5/native_fp8.py index 81cc9842..3eafd32b 100644 --- a/src/xorl/models/transformers/glm5/native_fp8.py +++ b/src/xorl/models/transformers/glm5/native_fp8.py @@ -34,6 +34,57 @@ ) +def reduce_glm52_no_combine_routes( + routed_values: torch.Tensor, + routing_weights: torch.Tensor, + local_ids: torch.Tensor, + *, + routed_scaling_factor: float, +) -> torch.Tensor: + """Reduce communicated BF16 route rows in GLM-5.2's FP32 local program.""" + + if routed_values.ndim != 3 or routed_values.dtype is not torch.bfloat16: + raise TypeError( + "GLM-5.2 no-combine values must be BF16 [rows, topk, hidden], " + f"got {routed_values.dtype} {tuple(routed_values.shape)}" + ) + route_shape = tuple(routed_values.shape[:2]) + if routing_weights.dtype is not torch.float32 or tuple(routing_weights.shape) != route_shape: + raise TypeError( + "GLM-5.2 no-combine routing weights must be FP32 [rows, topk], " + f"got {routing_weights.dtype} {tuple(routing_weights.shape)}" + ) + if local_ids.dtype is not torch.int32 or tuple(local_ids.shape) != route_shape: + raise TypeError( + f"GLM-5.2 no-combine local ids must be int32 [rows, topk], got {local_ids.dtype} {tuple(local_ids.shape)}" + ) + if not bool(torch.all(torch.isfinite(routing_weights))): + raise ValueError("GLM-5.2 no-combine routing weights contain non-finite values") + routed_scaling_factor = float(routed_scaling_factor) + if not math.isfinite(routed_scaling_factor) or routed_scaling_factor <= 0: + raise ValueError("GLM-5.2 routed scaling factor must be finite and positive") + + # -1 routes are absent on this owner. Mask before multiplying so a stale + # or diagnostic payload in a filtered slot can never enter the reduction. + valid_routes = local_ids >= 0 + safe_values = torch.where( + valid_routes.unsqueeze(-1), + routed_values, + torch.zeros((), dtype=torch.bfloat16, device=routed_values.device), + ) + weights = torch.where( + valid_routes, + routing_weights, + torch.zeros((), dtype=torch.float32, device=routing_weights.device), + ) + reduced = torch.sum( + safe_values.to(torch.float32) * weights.unsqueeze(-1), + dim=1, + dtype=torch.float32, + ) + return reduced.mul(routed_scaling_factor).to(torch.bfloat16) + + def validate_glm52_native_fp8_config(quantization_config: Mapping | None) -> dict: """Validate and normalize the exact official block-FP8 contract.""" @@ -409,7 +460,7 @@ def _sglang_ep_native_routed_value( from xorl.ops.moe.sglang_fused_moe_strided import fused_experts_impl_strided # noqa: PLC0415 MoEExperts._ensure_sglang_server_args() - output = fused_experts_impl_strided( + return fused_experts_impl_strided( hidden_flat.contiguous(), self.gate_up_proj.transpose(1, 2), self.down_proj.transpose(1, 2), @@ -422,9 +473,9 @@ def _sglang_ep_native_routed_value( w2_scale=self.down_weight_scale_inv.transpose(1, 2), block_shape=[128, 128], filter_expert=True, + no_combine=False, routed_scaling_factor=routed_scaling_factor, ) - return output def _dequantized_cached_experts(self) -> tuple[torch.Tensor, torch.Tensor]: """Dequantize the stored bytes into [E, N, K] BF16 banks (LoRA-surrogate treatment).""" diff --git a/src/xorl/models/transformers/glm5/qlora.py b/src/xorl/models/transformers/glm5/qlora.py index 0786392a..863b7d3f 100644 --- a/src/xorl/models/transformers/glm5/qlora.py +++ b/src/xorl/models/transformers/glm5/qlora.py @@ -11,6 +11,7 @@ from xorl.distributed.parallel_state import get_parallel_state from xorl.lora.modules.linear import LoraLinear +from xorl.models.exact_contract import glm52_exact_active_lora_enabled from xorl.models.layers.moe.experts import MoEExperts from xorl.models.transformers.glm5.exact_absorbed_kv_b_qlora import ( Glm52ExactTP1AbsorbedKvBBlockFP8QLoRA, @@ -177,9 +178,9 @@ def _validate_official_config(config) -> dict: exact_lm_head_component, ) ) - required_ep_dispatch = "alltoall" if exact_component_enabled else "deepep" - if getattr(config, "_ep_dispatch", None) != required_ep_dispatch: - raise ValueError(f"GLM-5.2 block-FP8 QLoRA requires ep_dispatch={required_ep_dispatch!r}") + admitted_ep_dispatches = {"alltoall", "deepep"} if exact_component_enabled else {"deepep"} + if getattr(config, "_ep_dispatch", None) not in admitted_ep_dispatches: + raise ValueError(f"GLM-5.2 block-FP8 QLoRA requires ep_dispatch in {sorted(admitted_ep_dispatches)!r}") if exact_dense_component and getattr(config, "hidden_act", None) != "silu": raise ValueError("GLM-5.2 exact active-LoRA dense component requires hidden_act='silu'") if exact_attention_component and not getattr(config, "_sparse_mla_enabled", False): @@ -450,6 +451,25 @@ def _replace_exact_routed_target( hidden_act=config.hidden_act, device=original.gate_up_proj.device, ) + # The replacement happens after ``Glm5MoEBlock`` has attached runtime + # DeepEP configuration to the original expert bank. Preserve those + # non-parameter attributes explicitly: composable FSDP changes the + # replacement's dynamic type but does not synthesize missing attributes. + replacement.deepep_buffer_size_gb = getattr( + original, + "deepep_buffer_size_gb", + getattr(config, "_deepep_buffer_size_gb", 2.0), + ) + replacement.deepep_num_sms = getattr( + original, + "deepep_num_sms", + getattr(config, "_deepep_num_sms", 20), + ) + replacement.deepep_async_combine = getattr( + original, + "deepep_async_combine", + getattr(config, "_deepep_async_combine", False), + ) replacement._source_fqn = target.name replacement._source_quant_format = "block_fp8" _set_submodule(model, target.name, replacement) @@ -705,12 +725,22 @@ def _validate_constructed_model(model: nn.Module, inventory: Glm52AdapterInvento f"GLM-5.2 QLoRA indexer weights_proj for layer {layer_idx} must remain an ordinary BF16 linear" ) + expected_router_names = ( + {f"model.layers.{layer_idx}.mlp.gate.weight" for layer_idx in range(3, 78)} + if bool(getattr(model.config, "train_router", False)) + else set() + ) + expected_trainable = inventory.factor_names | expected_router_names trainable = {name for name, parameter in model.named_parameters() if parameter.requires_grad} - if trainable != inventory.factor_names: + if trainable != expected_trainable: raise RuntimeError( - "GLM-5.2 QLoRA trainable factor set mismatch: " - f"missing={sorted(inventory.factor_names - trainable)} extra={sorted(trainable - inventory.factor_names)}" + "GLM-5.2 QLoRA trainable parameter set mismatch: " + f"missing={sorted(expected_trainable - trainable)} extra={sorted(trainable - expected_trainable)}" ) + for name in expected_router_names: + parameter = dict(model.named_parameters())[name] + if parameter.dtype is not torch.bfloat16: + raise TypeError(f"GLM-5.2 exact QLoRA router {name} must retain its BF16 serving weight") trainable_parameters = [parameter for parameter in model.parameters() if parameter.requires_grad] if len({id(parameter) for parameter in trainable_parameters}) != len(trainable_parameters): raise RuntimeError("GLM-5.2 QLoRA trainable factor set contains aliased Parameter identities") @@ -870,6 +900,19 @@ def prepare_glm52_block_fp8_qlora( # native module to any differentiable policy-bearing path. replace_glm52_native_fp8_modules(model, quantization_config) + # Exact active-LoRA router training deliberately optimizes the checkpoint + # BF16 gate weight itself. The BI router consumes those literal bytes in + # both engines; the analytic routing-weight surrogate supplies the gradient + # without changing the decision-time forward values. + if bool(getattr(config, "train_router", False)): + if not glm52_exact_active_lora_enabled(config): + raise ValueError("GLM-5.2 QLoRA router training requires the complete exact active-LoRA contract") + for layer_idx in range(3, 78): + gate = model.get_submodule(f"model.layers.{layer_idx}.mlp.gate") + if not isinstance(getattr(gate, "weight", None), nn.Parameter): + raise TypeError(f"GLM-5.2 layer {layer_idx} router gate has no trainable weight Parameter") + gate.weight.requires_grad_(True) + inventory = Glm52AdapterInventory(targets=targets, factors=_build_factor_inventory(model, targets)) _validate_constructed_model(model, inventory) model._glm52_adapter_inventory = inventory diff --git a/src/xorl/models/transformers/glm5/support.py b/src/xorl/models/transformers/glm5/support.py index bb3f081b..5974f361 100644 --- a/src/xorl/models/transformers/glm5/support.py +++ b/src/xorl/models/transformers/glm5/support.py @@ -1,5 +1,9 @@ """GLM-5 support helpers shared by model construction and layers.""" +from collections.abc import Iterable + +import torch.nn as nn + from xorl.models.exact_contract import glm52_exact_active_lora_enabled @@ -41,8 +45,36 @@ def validate_glm5_sequence_parallel(config, *, parallel_state=None, cp_enabled: def validate_glm5_router_settings(config, *, train_router: bool) -> None: if not is_glm5_config(config): return - if train_router: - raise ValueError("GLM-5 does not support train_router=True in xorl yet.") + if train_router and not glm52_exact_active_lora_enabled(config): + raise ValueError( + "GLM-5 train_router=True requires the complete exact active-LoRA contract; " + "generic and partial GLM router training remain unsupported." + ) + + +def validate_glm52_local_router_inventory( + model_parts: nn.Module | Iterable[nn.Module], + *, + retained_router_count: int, +) -> int: + """Require one retained router only for each sparse layer owned locally.""" + + parts = (model_parts,) if isinstance(model_parts, nn.Module) else tuple(model_parts) + seen_modules: set[int] = set() + expected_router_count = 0 + for part in parts: + for module in part.modules(): + if id(module) in seen_modules: + continue + seen_modules.add(id(module)) + if module.__class__.__name__ == "Glm5MoEBlock": + expected_router_count += 1 + if retained_router_count != expected_router_count: + raise RuntimeError( + "GLM-5.2 exact QLoRA post-FSDP router inventory does not match the locally owned sparse layers: " + f"retained={retained_router_count}, expected={expected_router_count}" + ) + return expected_router_count def validate_glm5_training_mode( @@ -91,7 +123,6 @@ def validate_glm5_training_mode( "quant_format": (quant_format, "block_fp8"), "quant_group_size": (quant_group_size, 128), "moe_implementation": (moe_implementation, "triton"), - "ep_dispatch": (ep_dispatch, "alltoall" if exact_active_lora else "deepep"), "moe_hybrid_shared_lora": (moe_hybrid_shared_lora, True), } mismatches = [ @@ -101,10 +132,20 @@ def validate_glm5_training_mode( ] if mismatches: raise ValueError("GLM-5.2 block-FP8 QLoRA rejects unsupported configuration: " + ", ".join(mismatches)) + admitted_dispatches = {"alltoall", "deepep"} if exact_active_lora else {"deepep"} + if ep_dispatch not in admitted_dispatches: + raise ValueError( + "GLM-5.2 block-FP8 QLoRA rejects unsupported configuration: " + f"ep_dispatch={ep_dispatch!r} (requires one of {sorted(admitted_dispatches)!r})" + ) elif enable_qlora: raise ValueError("GLM-5 QLoRA requires the explicit block_fp8_qlora_training mode") - if not freeze_router: - raise ValueError("GLM-5 requires freeze_router=True.") + train_router = bool(getattr(config, "train_router", False)) + if train_router != (not freeze_router): + raise ValueError( + "GLM-5 requires train_router and freeze_router to be complementary; " + f"got train_router={train_router}, freeze_router={freeze_router}." + ) if not merge_qkv: raise ValueError("GLM-5 does not support merge_qkv=False yet.") @@ -126,6 +167,7 @@ def glm5_default_lora_targets(*, train_attn: bool, train_mlp: bool, train_unembe "glm5_default_lora_targets", "is_glm5_config", "validate_glm5_router_settings", + "validate_glm52_local_router_inventory", "validate_glm5_sequence_parallel", "validate_glm5_training_mode", ] diff --git a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py index 3645d3ab..1601e246 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -404,11 +404,16 @@ def __init__(self, config, moe_implementation="triton", layer_idx: int | None = ) self.config = config self.layer_idx = layer_idx - self._native_ep_combine = bool(getattr(config, "_qwen35_exact_contract", False)) + self.deepep_native_exact = bool(getattr(config, "_deepep_native_exact", False)) + self._native_ep_combine = bool( + getattr(config, "_qwen35_exact_contract", False) and not self.deepep_native_exact + ) self.experts.ep_dispatch = getattr(config, "_ep_dispatch", "alltoall") self.experts.deepep_buffer_size_gb = getattr(config, "_deepep_buffer_size_gb", 2.0) self.experts.deepep_num_sms = getattr(config, "_deepep_num_sms", 20) self.experts.deepep_async_combine = getattr(config, "_deepep_async_combine", False) + self.experts.deepep_native_exact = self.deepep_native_exact + self.experts.lora_serving_mode = getattr(config, "_lora_serving_mode", None) self.experts.alltoall_combine_hidden_chunk_size = getattr(config, "_alltoall_combine_hidden_chunk_size", 0) self.shared_expert = Qwen3_5MoeMLP(config, intermediate_size=config.shared_expert_intermediate_size) self.shared_expert_gate = nn.Linear(config.hidden_size, 1, bias=False) @@ -558,6 +563,10 @@ def forward(self, hidden_states: torch.Tensor): routing_weights, selected_experts, router_logits = self.route(hidden_states.view(-1, hidden_dim)) out = self._ep_combine_native(hidden_states, routing_weights, selected_experts) return out, router_logits + # Native DeepEP owns the routed-expert transport and canonical fold. + # Qwen3.5's replicated shared expert remains model-specific and is + # joined once, after that fold, on the original token owner. This is + # the same thin composition used by serving's DeepEP path. expert_output, router_logits = super().forward(hidden_states) return expert_output + self._shared_expert(hidden_states), router_logits @@ -918,6 +927,14 @@ def forward( class Qwen3_5MoeForCausalLM(Qwen3_5MoePreTrainedModel): + deepep_native_exact_capability = { + "produces_local_leaf": True, + "wire_dtype": "bf16", + "uses_dispatch_handle": True, + "supported_ep_sizes": (2, 4, 8, 16), + "local_leaf_program": "fused_no_combine_false", + "lora_serving_modes": ("merged", "separate"), + } _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} _tp_plan = parallelize.MODEL_TP_PLAN diff --git a/src/xorl/models/transformers/qwen3_5_shared.py b/src/xorl/models/transformers/qwen3_5_shared.py index 6c168564..5f7eb4a9 100644 --- a/src/xorl/models/transformers/qwen3_5_shared.py +++ b/src/xorl/models/transformers/qwen3_5_shared.py @@ -77,7 +77,12 @@ def _apply_qwen35_gdn_exact(model: torch.nn.Module) -> dict[str, int]: # switch), so propagate it to every injected adapter before the trunk # wrapper validates and composes with those modules. if isinstance(module, LoraModule): + lora_mode = getattr(config, "_lora_serving_mode", None) + # Active serving reconstructs this same canonical folded trunk + # program from A/B; publication mode does not select a different + # trainer arithmetic path. module.exact_merged_forward = True + module.lora_serving_mode = lora_mode if hasattr(module, "rmsnorm_family"): norm_modules.append(module) if module.rmsnorm_family != rmsnorm_family: @@ -87,7 +92,14 @@ def _apply_qwen35_gdn_exact(model: torch.nn.Module) -> dict[str, int]: f"{type(module).__qualname__}." ) if hasattr(module, "_native_ep_combine"): - module._native_ep_combine = is_moe + # The legacy exact Qwen3.5 implementation owns an all-to-all + # exchange in the model block. Shared + # native DeepEP owns dispatch, original-handle combines, and the + # canonical fold inside the routed-expert layer instead. Preserve + # that ownership decision after LoRA injection; re-enabling the + # model-local exchange here would silently route a native launch + # back through the oracle during the final pre-FSDP hook. + module._native_ep_combine = bool(is_moe and not getattr(config, "_deepep_native_exact", False)) if hasattr(module, "_exact_batch_invariant_router"): module._exact_batch_invariant_router = is_moe module.router._exact_batch_invariant = is_moe diff --git a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py index e228a2c7..b9d72033 100644 --- a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py +++ b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py @@ -369,12 +369,17 @@ def __init__(self, config): moe_implementation="eager", train_router=getattr(config, "train_router", False), activation_native=getattr(config, "_activation_native", False), + exact_batch_invariant_router=bool(getattr(config, "_deepep_native_exact", False)), + exact_router_weights_fp32=bool(getattr(config, "_deepep_native_exact", False)), ) self.config = config + self.deepep_native_exact = bool(getattr(config, "_deepep_native_exact", False)) self.experts.ep_dispatch = getattr(config, "_ep_dispatch", "alltoall") self.experts.deepep_buffer_size_gb = getattr(config, "_deepep_buffer_size_gb", 2.0) self.experts.deepep_num_sms = getattr(config, "_deepep_num_sms", 20) self.experts.deepep_async_combine = getattr(config, "_deepep_async_combine", False) + self.experts.deepep_native_exact = getattr(config, "_deepep_native_exact", False) + self.experts.lora_serving_mode = getattr(config, "_lora_serving_mode", None) self.experts.alltoall_combine_hidden_chunk_size = getattr(config, "_alltoall_combine_hidden_chunk_size", 0) @@ -392,12 +397,17 @@ def __init__(self, config): moe_implementation="triton", train_router=getattr(config, "train_router", False), activation_native=getattr(config, "_activation_native", False), + exact_batch_invariant_router=bool(getattr(config, "_deepep_native_exact", False)), + exact_router_weights_fp32=bool(getattr(config, "_deepep_native_exact", False)), ) self.config = config + self.deepep_native_exact = bool(getattr(config, "_deepep_native_exact", False)) self.experts.ep_dispatch = getattr(config, "_ep_dispatch", "alltoall") self.experts.deepep_buffer_size_gb = getattr(config, "_deepep_buffer_size_gb", 2.0) self.experts.deepep_num_sms = getattr(config, "_deepep_num_sms", 20) self.experts.deepep_async_combine = getattr(config, "_deepep_async_combine", False) + self.experts.deepep_native_exact = getattr(config, "_deepep_native_exact", False) + self.experts.lora_serving_mode = getattr(config, "_lora_serving_mode", None) self.experts.alltoall_combine_hidden_chunk_size = getattr(config, "_alltoall_combine_hidden_chunk_size", 0) @@ -415,12 +425,17 @@ def __init__(self, config): moe_implementation="quack", train_router=getattr(config, "train_router", False), activation_native=getattr(config, "_activation_native", False), + exact_batch_invariant_router=bool(getattr(config, "_deepep_native_exact", False)), + exact_router_weights_fp32=bool(getattr(config, "_deepep_native_exact", False)), ) self.config = config + self.deepep_native_exact = bool(getattr(config, "_deepep_native_exact", False)) self.experts.ep_dispatch = getattr(config, "_ep_dispatch", "alltoall") self.experts.deepep_buffer_size_gb = getattr(config, "_deepep_buffer_size_gb", 2.0) self.experts.deepep_num_sms = getattr(config, "_deepep_num_sms", 20) self.experts.deepep_async_combine = getattr(config, "_deepep_async_combine", False) + self.experts.deepep_native_exact = getattr(config, "_deepep_native_exact", False) + self.experts.lora_serving_mode = getattr(config, "_lora_serving_mode", None) self.experts.alltoall_combine_hidden_chunk_size = getattr(config, "_alltoall_combine_hidden_chunk_size", 0) @@ -438,12 +453,17 @@ def __init__(self, config): moe_implementation="native", train_router=getattr(config, "train_router", False), activation_native=getattr(config, "_activation_native", False), + exact_batch_invariant_router=bool(getattr(config, "_deepep_native_exact", False)), + exact_router_weights_fp32=bool(getattr(config, "_deepep_native_exact", False)), ) self.config = config + self.deepep_native_exact = bool(getattr(config, "_deepep_native_exact", False)) self.experts.ep_dispatch = getattr(config, "_ep_dispatch", "alltoall") self.experts.deepep_buffer_size_gb = getattr(config, "_deepep_buffer_size_gb", 2.0) self.experts.deepep_num_sms = getattr(config, "_deepep_num_sms", 20) self.experts.deepep_async_combine = getattr(config, "_deepep_async_combine", False) + self.experts.deepep_native_exact = getattr(config, "_deepep_native_exact", False) + self.experts.lora_serving_mode = getattr(config, "_lora_serving_mode", None) self.experts.alltoall_combine_hidden_chunk_size = getattr(config, "_alltoall_combine_hidden_chunk_size", 0) @@ -883,7 +903,38 @@ def forward( class KwargsForCausalLM(AttentionKwargs): ... +def _apply_qwen3_deepep_native_exact(model: nn.Module) -> dict[str, int]: + """Install the Qwen3-local arithmetic behind its exact DeepEP capability.""" + + from xorl.lora.modules.base import LoraModule # noqa: PLC0415 + from xorl.ops.batch_invariant_ops import ( # noqa: PLC0415 + wrap_trunk_linears_batch_invariant, + ) + from xorl.ops.bi_families_v2 import ( # noqa: PLC0415 + _select_qwen3_dense_families_v2, + ) + + lora_mode = getattr(model.config, "_lora_serving_mode", None) + for module in model.modules(): + if isinstance(module, LoraModule): + # Both publication modes train through the same canonical folded + # trunk weight. In separate mode, active sampler LoRA reproduces + # that fold from the transported factors. + module.exact_merged_forward = True + module.lora_serving_mode = lora_mode + _select_qwen3_dense_families_v2() + return wrap_trunk_linears_batch_invariant(model) + + class Qwen3MoeForCausalLM(Qwen3MoePreTrainedModel): + deepep_native_exact_capability = { + "produces_local_leaf": True, + "wire_dtype": "bf16", + "uses_dispatch_handle": True, + "supported_ep_sizes": (2, 4, 8, 16), + "local_leaf_program": "fused_no_combine_false", + "lora_serving_modes": ("merged", "separate"), + } _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} @@ -905,6 +956,9 @@ def unfuse_for_tp(self): """Unfuse fused projections for tensor parallelism compatibility.""" parallelize.unfuse_for_tp(self) + def _apply_deepep_native_exact(self) -> dict[str, int]: + return _apply_qwen3_deepep_native_exact(self) + def get_input_embeddings(self): return self.model.embed_tokens diff --git a/src/xorl/ops/dsv4/exact_attention.py b/src/xorl/ops/dsv4/exact_attention.py index 7d53a2e6..79b457bd 100644 --- a/src/xorl/ops/dsv4/exact_attention.py +++ b/src/xorl/ops/dsv4/exact_attention.py @@ -3,11 +3,183 @@ from __future__ import annotations import os +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Any, Iterator import torch from torch import Tensor +_ACTIVE_ATTENTION_LAYER: ContextVar[int | None] = ContextVar( + "xorl_dsv4_exact_attention_layer", + default=None, +) +_CAPTURED_OPERATOR_INPUTS: set[tuple[int, int, int]] = set() + + +@contextmanager +def exact_attention_layer(layer_id: int) -> Iterator[None]: + """Associate a literal attention invocation with its model layer. + + The context is used only by the opt-in operator-localization capture. It + does not alter the serving-value forward or its VJP. + """ + + token = _ACTIVE_ATTENTION_LAYER.set(int(layer_id)) + try: + yield + finally: + _ACTIVE_ATTENTION_LAYER.reset(token) + + +def _scheduler_snapshot(metadata: Any) -> dict[str, Any]: + config = getattr(metadata, "config", None) + config_snapshot = None + if config is not None: + config_snapshot = { + name: getattr(config, name) + for name in ( + "b", + "s_q", + "h_q", + "page_block_size", + "h_k", + "causal", + "is_fp8_kvcache", + "topk", + "extra_page_block_size", + "extra_topk", + ) + } + return { + "have_initialized": bool(getattr(metadata, "have_initialized", False)), + "config": config_snapshot, + "tile_scheduler_metadata": ( + metadata.tile_scheduler_metadata.detach().cpu().clone() + if isinstance(getattr(metadata, "tile_scheduler_metadata", None), Tensor) + else None + ), + "num_splits": ( + metadata.num_splits.detach().cpu().clone() + if isinstance(getattr(metadata, "num_splits", None), Tensor) + else None + ), + } + + +def _referenced_cache_rows( + cache: Tensor, + indices: Tensor, + topk_length: Tensor, +) -> tuple[Tensor, Tensor]: + """Gather packed FlashMLA rows in the semantic sparse-decode order. + + FlashMLA does not store each 584-byte token contiguously. A page contains + all 576-byte FP8-nope/BF16-RoPE payloads first, followed by an 8-byte scale + row for every token. Reinterpreting the kernel view as flat 584-byte rows + crosses those planes and manufactures cache differences. + """ + + length = int(topk_length.reshape(-1)[0].item()) + references = indices.reshape(-1, indices.shape[-1])[0, :length] + page_size = cache.shape[1] + valid = references[(references >= 0) & (references < cache.shape[0] * page_size)].to(torch.int64) + if not valid.numel(): + return references.detach().cpu().clone(), torch.empty((0, _SERVING_SLOT_BYTES), dtype=torch.uint8) + + cache_bytes = cache.detach().view(torch.uint8) + page_stride = cache_bytes.stride(0) + page_storage = cache_bytes.as_strided( + (cache.shape[0], page_stride), + (page_stride, 1), + ) + pages = torch.div(valid, page_size, rounding_mode="floor") + offsets = torch.remainder(valid, page_size) + payload_offsets = offsets[:, None] * 576 + torch.arange(576, device=cache.device) + scale_offsets = page_size * 576 + offsets[:, None] * 8 + torch.arange(8, device=cache.device) + rows = torch.cat( + ( + page_storage[pages[:, None], payload_offsets], + page_storage[pages[:, None], scale_offsets], + ), + dim=1, + ) + return references.detach().cpu().clone(), rows.detach().cpu().clone() + + +def _maybe_capture_operator_inputs( + *, + position: int, + ratio: int, + q: Tensor, + attn_sink: Tensor, + k_cache: Tensor, + indices: Tensor, + topk_length: Tensor, + extra_k_cache: Tensor | None, + extra_indices: Tensor | None, + extra_topk_length: Tensor | None, + scheduler_before: dict[str, Any], + scheduler_after: Any, + output: Tensor, +) -> None: + capture_dir = os.environ.get("XORL_DSV4_TRAINER_OPERATOR_CAPTURE_DIR", "").strip() + raw_layer = os.environ.get("XORL_DSV4_OPERATOR_CAPTURE_LAYER", "").strip() + raw_positions = os.environ.get("XORL_DSV4_OPERATOR_CAPTURE_POSITIONS", "").strip() + layer_id = _ACTIVE_ATTENTION_LAYER.get() + if not capture_dir or not raw_layer or not raw_positions or layer_id is None: + return + if layer_id != int(raw_layer): + return + target_positions = {int(item.strip()) for item in raw_positions.split(",") if item.strip()} + if position not in target_positions: + return + global_rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + capture_key = (global_rank, layer_id, position) + if capture_key in _CAPTURED_OPERATOR_INPUTS: + return + + raw_references, raw_rows = _referenced_cache_rows(k_cache, indices, topk_length) + extra_references = None + extra_rows = None + if extra_k_cache is not None and extra_indices is not None and extra_topk_length is not None: + extra_references, extra_rows = _referenced_cache_rows( + extra_k_cache, + extra_indices, + extra_topk_length, + ) + _CAPTURED_OPERATOR_INPUTS.add(capture_key) + os.makedirs(capture_dir, exist_ok=True) + torch.save( + { + "schema": "xorl.dsv4_trainer_flashmla_operator_capture.v2", + "engine": "trainer", + "global_rank": global_rank, + "layer": layer_id, + "position": position, + "ratio": ratio, + "q": q.detach().cpu().clone(), + "attn_sink": attn_sink.detach().cpu().clone(), + "raw_indices": indices.detach().cpu().clone(), + "raw_topk_length": topk_length.detach().cpu().clone(), + "raw_references": raw_references, + "raw_referenced_cache_rows": raw_rows, + "extra_indices": extra_indices.detach().cpu().clone() if extra_indices is not None else None, + "extra_topk_length": (extra_topk_length.detach().cpu().clone() if extra_topk_length is not None else None), + "extra_references": extra_references, + "extra_referenced_cache_rows": extra_rows, + "scheduler_before": scheduler_before, + "scheduler_after": _scheduler_snapshot(scheduler_after), + "output": output.detach().cpu().clone(), + }, + os.path.join( + capture_dir, + f"rank{global_rank:05d}.layer{layer_id:03d}.position{position:05d}.pt", + ), + ) + + def _validate_dsv4_lora_metadata(tensor: Tensor, *, where: str) -> None: if os.environ.get("XORL_DSV4_DIAGNOSTIC_BASE_MARLIN") != "1": return @@ -22,11 +194,24 @@ def _positions(batch_size: int, sequence_length: int, device: torch.device, offs return torch.arange(offset, offset + sequence_length, dtype=torch.int64, device=device).repeat(batch_size) -# Serving's SWA KV pool pages tokens in blocks of the sliding window (the -# DeepseekV4AttnBackend asserts ``swa_page_size == SWA_WINDOW == 128``). The -# decode kernel's tile schedule keys on this page-block size, so the carried -# raw cache must use it for byte parity. -_SERVING_SWA_PAGE_SIZE = 128 +# Serving keeps a logical 128-token SWA window inside a 256-token physical KV +# page. ``_build_dsv4_kv_pool`` requires the CUDA SWA pool page size to be 256, +# while ``DeepseekV4AttnBackend`` independently fixes ``SWA_WINDOW`` at 128. +# FlashMLA's tile schedule keys on the physical cache shape, so the trainer's +# carried raw cache must preserve the 256-token page geometry even though every +# query exposes at most 128 logical SWA indices. +_SERVING_SWA_PAGE_SIZE = 256 +# The qualified eager sampler uses ``max_total_tokens=8192`` with DP2/CP4 and +# materializes these per-rank DSV4 pools (recorded in its startup log as +# swa_size=768, c4_size=2048, c128_size=64). FlashMLA receives the complete +# cache tensors and its dynamic split schedule is sensitive to their physical +# extent, so a short trainer replay must retain the same pool shape rather than +# shrinking each cache to the request length. +_SERVING_CACHE_CAPACITY_BY_PAGE_SIZE = { + 256: 768, + 64: 2048, + 2: 64, +} # One cache slot is 584 bytes: 448 FP8 nope + 128 BF16 rope + 8 scale bytes. _SERVING_SLOT_BYTES = 584 @@ -65,13 +250,19 @@ def _flashmla_page_bytes(page_size: int) -> int: def _ensure_paged_kvcache(cache: Tensor | None, num_slots: int, page_size: int, device: torch.device) -> Tensor: """Grow (never shrink) a paged FlashMLA FP8 cache to hold ``num_slots``. + Serving's paged allocator reserves physical page zero for dummy/padded + writes and allocates real request tokens from page one. Preserve that + address space here because FlashMLA's tile scheduler consumes the physical + indices, not just their page-relative offsets. + Pages are zero-filled like serving's pool allocation (``create_buffer`` uses ``torch.zeros``): the decode kernel's masked lanes still read cache bytes behind invalid indices, and non-zero garbage there perturbs the online-softmax max by a few ULPs. """ - num_pages = max((num_slots + page_size - 1) // page_size, 1) + num_slots = max(num_slots, _SERVING_CACHE_CAPACITY_BY_PAGE_SIZE.get(page_size, num_slots)) + num_pages = 1 + (num_slots + page_size - 1) // page_size page_bytes = _flashmla_page_bytes(page_size) if cache is None: return torch.zeros((num_pages, page_bytes), dtype=torch.uint8, device=device) @@ -93,7 +284,13 @@ def _window_indices_for_positions(positions: Tensor) -> Tensor: def _paged_cache_kernel_view(cache: Tensor, page_size: int) -> Tensor: """View a paged FP8 cache the way serving hands it to the decode kernel.""" - return cache[:, : page_size * _SERVING_SLOT_BYTES].view(cache.shape[0], page_size, 1, _SERVING_SLOT_BYTES) + if cache.dtype != torch.uint8: + raise TypeError(f"packed DSV4 cache storage must be uint8, got {cache.dtype}") + return ( + cache[:, : page_size * _SERVING_SLOT_BYTES] + .view(torch.float8_e4m3fn) + .view(cache.shape[0], page_size, 1, _SERVING_SLOT_BYTES) + ) def _serving_decode_attention( @@ -126,6 +323,7 @@ def _serving_decode_attention( # by a few ULPs on some heads. window_offsets = position - torch.arange(128, dtype=torch.int32, device=device) swa_indices = window_offsets.masked_fill(window_offsets < 0, -1).view(1, 1, 128) + swa_indices = torch.where(swa_indices >= 0, swa_indices + _SERVING_SWA_PAGE_SIZE, swa_indices) swa_topk_length = torch.clamp(positions + 1, max=128).to(torch.int32) extra_kwargs = {} if position < 0 or position >= carry_state.num_tokens: @@ -153,19 +351,27 @@ def _serving_decode_attention( ) extra_indices = torch.full((1, 1, width), -1, dtype=torch.int32, device=device) if blocks: - extra_indices[0, 0, :blocks] = torch.arange(blocks, dtype=torch.int32, device=device) + extra_page_size = 256 // ratio + extra_indices[0, 0, :blocks] = extra_page_size + torch.arange( + blocks, + dtype=torch.int32, + device=device, + ) extra_kwargs = { "extra_k_cache": _paged_cache_kernel_view(carry_state.compressed_kvcache, 256 // ratio), "extra_indices_in_kvcache": extra_indices, "extra_topk_length": torch.tensor([max(blocks, 1)], dtype=torch.int32, device=device), } + scheduler = get_mla_metadata()[0] + scheduler_before = _scheduler_snapshot(scheduler) + k_cache = _paged_cache_kernel_view(carry_state.kvcache, _SERVING_SWA_PAGE_SIZE) output, _ = flash_mla_with_kvcache( q=q.contiguous().view(1, 1, 64, 512), - k_cache=_paged_cache_kernel_view(carry_state.kvcache, _SERVING_SWA_PAGE_SIZE), + k_cache=k_cache, head_dim_v=512, block_table=None, cache_seqlens=None, - tile_scheduler_metadata=get_mla_metadata()[0], + tile_scheduler_metadata=scheduler, softmax_scale=softmax_scale, is_fp8_kvcache=True, indices=swa_indices, @@ -173,6 +379,30 @@ def _serving_decode_attention( topk_length=swa_topk_length, **extra_kwargs, ) + num_splits = getattr(scheduler, "num_splits", None) + if isinstance(num_splits, Tensor) and int(num_splits.reshape(-1)[-1].item()) > 2: + raise RuntimeError( + "DSV4 exact trainer attention requires the qualified MODEL1 fixed-K2 " + "FlashMLA runtime used by serving; the loaded operator selected " + f"num_splits={num_splits.detach().cpu().tolist()} at position={position}, " + f"C{ratio}. Put the qualified DSV4 runtime overlay before the generic " + "sglang-kernel package on PYTHONPATH." + ) + _maybe_capture_operator_inputs( + position=position, + ratio=ratio, + q=q.contiguous().view(1, 1, 64, 512), + attn_sink=attn_sink, + k_cache=k_cache, + indices=swa_indices, + topk_length=swa_topk_length, + extra_k_cache=extra_kwargs.get("extra_k_cache"), + extra_indices=extra_kwargs.get("extra_indices_in_kvcache"), + extra_topk_length=extra_kwargs.get("extra_topk_length"), + scheduler_before=scheduler_before, + scheduler_after=scheduler, + output=output, + ) return output.view(1, 1, 64, 512) @@ -253,7 +483,7 @@ def _store_raw_kv_carry( eps=eps, freqs_cis=freqs_cis, positions=positions, - out_loc=positions.to(torch.int32), + out_loc=positions.to(torch.int32) + page_size, kvcache=carry_state.kvcache, page_size=page_size, ) @@ -264,7 +494,7 @@ def _store_raw_kv_carry( dequantize_k_cache_paged, ) - all_locs = torch.arange(total_tokens, dtype=torch.int32, device=kv_input.device) + all_locs = page_size + torch.arange(total_tokens, dtype=torch.int32, device=kv_input.device) return dequantize_k_cache_paged(carry_state.kvcache, all_locs, page_size).view(total_tokens, 512) @@ -294,7 +524,7 @@ def _store_preprocessed_kv_carry( fused_store_cache( input=kv.contiguous().view(-1, 512), cache=carry_state.kvcache, - indices=positions.to(torch.int32), + indices=positions.to(torch.int32) + page_size, page_size=page_size, type="flashmla", ) @@ -305,7 +535,7 @@ def _store_preprocessed_kv_carry( dequantize_k_cache_paged, ) - all_locs = torch.arange(total_tokens, dtype=torch.int32, device=kv.device) + all_locs = page_size + torch.arange(total_tokens, dtype=torch.int32, device=kv.device) return dequantize_k_cache_paged(carry_state.kvcache, all_locs, page_size).view(total_tokens, 512) @@ -347,8 +577,8 @@ def backward(ctx, grad_output: Tensor): q = q_input.detach().requires_grad_(True) q_float = q.float() normalized = q_float * torch.rsqrt(q_float.square().mean(-1, keepdim=True) + ctx.eps) - surrogate = _apply_rope_torch(normalized.to(q.dtype), freqs_cis, positions, inverse=False) - grad_q = torch.autograd.grad(surrogate, q, grad_output, create_graph=False)[0] + vjp_replay = _apply_rope_torch(normalized.to(q.dtype), freqs_cis, positions, inverse=False) + grad_q = torch.autograd.grad(vjp_replay, q, grad_output, create_graph=False)[0] return grad_q, None, None, None @@ -388,8 +618,8 @@ def backward(ctx, grad_output: Tensor): kv_input, kv_norm_weight, freqs_cis, positions = ctx.saved_tensors with torch.enable_grad(): kv = kv_input.detach().requires_grad_(True) - surrogate = _kv_norm_rope_torch(kv, kv_norm_weight, freqs_cis, positions, ctx.eps) - grad_kv = torch.autograd.grad(surrogate, kv, grad_output, create_graph=False)[0] + vjp_replay = _kv_norm_rope_torch(kv, kv_norm_weight, freqs_cis, positions, ctx.eps) + grad_kv = torch.autograd.grad(vjp_replay, kv, grad_output, create_graph=False)[0] return grad_kv, None, None, None @@ -426,11 +656,9 @@ def _native_swa_kv(kv_input: Tensor, kv_norm_weight: Tensor, freqs_cis: Tensor, positions = _positions(batch_size, sequence_length, kv_input.device) kv_flat = kv_input.contiguous().view(-1, 512) num_tokens = kv_flat.shape[0] - page_size = 256 - page_bytes = ((584 * page_size + 575) // 576) * 576 - num_pages = (num_tokens + page_size - 1) // page_size - kvcache = torch.empty((num_pages, page_bytes), dtype=torch.uint8, device=kv_input.device) - out_loc = torch.arange(num_tokens, dtype=torch.int32, device=kv_input.device) + page_size = _SERVING_SWA_PAGE_SIZE + kvcache = _ensure_paged_kvcache(None, num_tokens, page_size, kv_input.device) + out_loc = page_size + torch.arange(num_tokens, dtype=torch.int32, device=kv_input.device) fused_k_norm_rope_flashmla( kv=kv_flat, kv_weight=kv_norm_weight, @@ -459,13 +687,8 @@ def _native_swa_kv_from_bf16(kv: Tensor) -> Tensor: kv_flat = kv.contiguous().view(-1, 512) num_tokens = kv_flat.shape[0] page_size = _SERVING_SWA_PAGE_SIZE - num_pages = max((num_tokens + page_size - 1) // page_size, 1) - kvcache = torch.zeros( - (num_pages, _flashmla_page_bytes(page_size)), - dtype=torch.uint8, - device=kv.device, - ) - out_loc = torch.arange(num_tokens, dtype=torch.int32, device=kv.device) + kvcache = _ensure_paged_kvcache(None, num_tokens, page_size, kv.device) + out_loc = page_size + torch.arange(num_tokens, dtype=torch.int32, device=kv.device) fused_store_cache( input=kv_flat, cache=kvcache, @@ -581,7 +804,7 @@ def _serving_compressed_decode_step( ) at_block_boundary = seq_len % ratio == 0 out_loc = torch.tensor( - [seq_len // ratio - 1 if at_block_boundary else 0], + [page_size + seq_len // ratio - 1 if at_block_boundary else 0], dtype=torch.int64, device=x.device, ) @@ -602,7 +825,11 @@ def _serving_compressed_decode_step( carry_state.num_compressed = seq_len // ratio if carry_state.num_compressed == 0: return x.new_empty((0, 512)) - compressed_locs = torch.arange(carry_state.num_compressed, dtype=torch.int32, device=x.device) + compressed_locs = page_size + torch.arange( + carry_state.num_compressed, + dtype=torch.int32, + device=x.device, + ) output = dequantize_k_cache_paged(carry_state.compressed_kvcache, compressed_locs, page_size).view( carry_state.num_compressed, 512 ) @@ -688,13 +915,11 @@ def _serving_compressed_kv( carry_state.num_compressed = num_compressed kvcache = carry_state.compressed_kvcache else: - page_bytes = _flashmla_page_bytes(page_size) - num_pages = (num_compressed + page_size - 1) // page_size - kvcache = torch.empty((num_pages, page_bytes), dtype=torch.uint8, device=x.device) + kvcache = _ensure_paged_kvcache(None, num_compressed, page_size, x.device) out_loc = torch.zeros(sequence_length, dtype=torch.int64, device=x.device) if num_compressed: endpoints = torch.arange(ratio - 1, sequence_length, ratio, device=x.device) - out_loc[endpoints] = torch.arange(num_compressed, dtype=torch.int64, device=x.device) + out_loc[endpoints] = page_size + torch.arange(num_compressed, dtype=torch.int64, device=x.device) compress_norm_rope_store( compressed, plan, @@ -714,7 +939,7 @@ def _serving_compressed_kv( dequantize_k_cache_paged, ) - compressed_locs = torch.arange(num_compressed, dtype=torch.int32, device=x.device) + compressed_locs = page_size + torch.arange(num_compressed, dtype=torch.int32, device=x.device) output = dequantize_k_cache_paged(kvcache, compressed_locs, page_size).view(num_compressed, 512) _validate_dsv4_lora_metadata(x, where=f"C{ratio} compressor cache dequantize") return output @@ -778,7 +1003,7 @@ def _hybrid_indices_for_positions(positions: Tensor, ratio: int, compressed_capa return indices, lengths -def _compress_surrogate( +def _compress_vjp_replay( x: Tensor, wkv_weight: Tensor, wgate_weight: Tensor, @@ -788,7 +1013,7 @@ def _compress_surrogate( eps: float, ratio: int, ) -> Tensor: - """Differentiable compressor surrogate for the literal serving VJP.""" + """Recompute the compressor formula used by the literal serving VJP.""" batch_size, sequence_length, _ = x.shape groups = sequence_length // ratio @@ -953,9 +1178,9 @@ def backward(ctx, grad_output: Tensor): kv_input = kv_saved.detach().requires_grad_(True) kv_segment = _kv_norm_rope_torch(kv_input, weight, freqs_cis, positions, ctx.eps) kv = torch.cat((prefix.to(kv_segment.dtype).unsqueeze(0), kv_segment), dim=1) - surrogate = sparse_attn_torch(q, kv, attn_sink, indices, ctx.softmax_scale) + vjp_replay = sparse_attn_torch(q, kv, attn_sink, indices, ctx.softmax_scale) grad_q, grad_kv = torch.autograd.grad( - surrogate, + vjp_replay, (q, kv_input), grad_output, create_graph=False, @@ -973,9 +1198,9 @@ def backward(ctx, grad_output: Tensor): if ctx.kv_preprocessed else _kv_norm_rope_torch(kv_input, weight, freqs_cis, kv_positions, ctx.eps) ) - surrogate = sparse_attn_torch(q, kv, attn_sink, indices, ctx.softmax_scale) + vjp_replay = sparse_attn_torch(q, kv, attn_sink, indices, ctx.softmax_scale) grad_q, grad_kv = torch.autograd.grad( - surrogate, + vjp_replay, (q, kv_input), grad_output, create_graph=False, @@ -1244,7 +1469,7 @@ def backward(ctx, grad_output: Tensor): kv_input = kv_saved.detach().requires_grad_(True) kv_segment = _kv_norm_rope_torch(kv_input, kv_norm_weight, freqs_cis, positions, ctx.eps) kv = torch.cat((kv_constant.to(kv_segment.dtype).unsqueeze(0), kv_segment), dim=1) - surrogate = sparse_attn_torch( + vjp_replay = sparse_attn_torch( q, kv, attn_sink, @@ -1252,7 +1477,7 @@ def backward(ctx, grad_output: Tensor): ctx.softmax_scale, ) grad_q, grad_kv = torch.autograd.grad( - surrogate, + vjp_replay, (q, kv_input), grad_output, create_graph=False, @@ -1305,7 +1530,7 @@ def backward(ctx, grad_output: Tensor): if ctx.kv_preprocessed else _kv_norm_rope_torch(kv_input, kv_norm_weight, freqs_cis, kv_positions, ctx.eps) ) - compressed = _compress_surrogate( + compressed = _compress_vjp_replay( x, compressor_wkv_weight, compressor_wgate_weight, @@ -1324,7 +1549,7 @@ def backward(ctx, grad_output: Tensor): dim=0, ) kv = torch.cat((compressed, vanilla[0]), dim=0).unsqueeze(0) - surrogate = sparse_attn_torch( + vjp_replay = sparse_attn_torch( q, kv, attn_sink, @@ -1332,7 +1557,7 @@ def backward(ctx, grad_output: Tensor): ctx.softmax_scale, ) grad_q, grad_kv, grad_x = torch.autograd.grad( - surrogate, + vjp_replay, (q, kv_input, x), grad_output, create_graph=False, diff --git a/src/xorl/ops/dsv4/hyper_connection.py b/src/xorl/ops/dsv4/hyper_connection.py index 1068e181..03b32c0e 100644 --- a/src/xorl/ops/dsv4/hyper_connection.py +++ b/src/xorl/ops/dsv4/hyper_connection.py @@ -13,6 +13,7 @@ """ import os +from dataclasses import dataclass import einops import torch @@ -30,6 +31,22 @@ _DEFAULT_HC_CHUNK_TOKENS = 1024 +@dataclass(frozen=True) +class ExactMhcReplaySegment: + """One serving-sized MHC launch projected onto this rank's local rows. + + ``launch_rows`` preserves the serving launch geometry (and therefore its + TF32 split count). ``source_rows`` names compact trainer rows, while + ``launch_positions`` places those rows at their positions in the serving + launch. Rows owned by other CP ranks are zero placeholders; MHC pre-norm + has no cross-token arithmetic, so they cannot affect selected local rows. + """ + + launch_rows: int + source_rows: tuple[int, ...] + launch_positions: tuple[int, ...] + + def _exact_mhc_pre_norm_forward( residual: Tensor, hc_fn: Tensor, @@ -425,7 +442,102 @@ def layer_pre_norm_exact( hc_scale: Tensor, hc_base: Tensor, norm_weight: Tensor, + serving_segments: tuple[int | ExactMhcReplaySegment, ...] | None = None, ) -> tuple[Tensor, Tensor, Tensor]: + if serving_segments is not None: + if hidden_states.ndim != 4: + raise ValueError( + "DSV4 serving-segment MHC replay requires [B, S, hc, H] residuals, " + f"got {tuple(hidden_states.shape)}" + ) + compute_rows = hidden_states.shape[1] + layer_inputs = [] + posts = [] + combs = [] + source_order: list[int] = [] + start = 0 + for segment in serving_segments: + if isinstance(segment, int): + if segment <= 0: + raise ValueError(f"DSV4 serving MHC segments must be positive, got {serving_segments}") + end = start + segment + source_rows = tuple(range(start, end)) + launch_positions = tuple(range(segment)) + launch_residual = hidden_states[:, start:end] + start = end + else: + if segment.launch_rows <= 0: + raise ValueError(f"DSV4 serving MHC launch rows must be positive, got {segment}") + source_rows = segment.source_rows + launch_positions = segment.launch_positions + if not source_rows or len(source_rows) != len(launch_positions): + raise ValueError( + "DSV4 CP serving MHC segments require equally sized nonempty source and launch rows, " + f"got {segment}" + ) + if len(set(source_rows)) != len(source_rows) or any( + row < 0 or row >= compute_rows for row in source_rows + ): + raise ValueError(f"DSV4 CP serving MHC source rows are invalid: {segment}") + if len(set(launch_positions)) != len(launch_positions) or any( + row < 0 or row >= segment.launch_rows for row in launch_positions + ): + raise ValueError(f"DSV4 CP serving MHC launch positions are invalid: {segment}") + source_index = torch.tensor(source_rows, dtype=torch.long, device=hidden_states.device) + launch_index = torch.tensor(launch_positions, dtype=torch.long, device=hidden_states.device) + local_residual = hidden_states.index_select(1, source_index) + launch_residual = hidden_states.new_zeros( + hidden_states.shape[0], + segment.launch_rows, + *hidden_states.shape[2:], + ).index_copy(1, launch_index, local_residual) + + layer_input, post, comb = _ExactMhcPreNorm.apply( + launch_residual, + hc_fn, + hc_scale, + hc_base, + norm_weight, + self.norm_eps, + self.hc_eps, + self.hc_sinkhorn_iters, + ) + if not isinstance(segment, int): + launch_index = torch.tensor( + launch_positions, + dtype=torch.long, + device=hidden_states.device, + ) + layer_input = layer_input.index_select(1, launch_index) + post = post.index_select(1, launch_index) + comb = comb.index_select(1, launch_index) + layer_inputs.append(layer_input) + posts.append(post) + combs.append(comb) + source_order.extend(source_rows) + + if sorted(source_order) != list(range(compute_rows)): + raise ValueError( + "DSV4 serving MHC segments must cover the compute rows exactly once: " + f"source_rows={source_order} rows={compute_rows}" + ) + layer_input = torch.cat(layer_inputs, dim=1) + post = torch.cat(posts, dim=1) + comb = torch.cat(combs, dim=1) + if source_order != list(range(compute_rows)): + restore_order = torch.tensor( + sorted(range(compute_rows), key=source_order.__getitem__), + dtype=torch.long, + device=hidden_states.device, + ) + layer_input = layer_input.index_select(1, restore_order) + post = post.index_select(1, restore_order) + comb = comb.index_select(1, restore_order) + return ( + layer_input, + post, + comb, + ) return _ExactMhcPreNorm.apply( hidden_states, hc_fn, diff --git a/src/xorl/ops/linear_attention/layers/gated_deltanet.py b/src/xorl/ops/linear_attention/layers/gated_deltanet.py index 5201a803..bbd3aa8e 100644 --- a/src/xorl/ops/linear_attention/layers/gated_deltanet.py +++ b/src/xorl/ops/linear_attention/layers/gated_deltanet.py @@ -44,6 +44,13 @@ def _sglang_compatible_beta_gate(b_input: torch.Tensor) -> torch.Tensor: class GatedDeltaNet(nn.Module): + # SGLang evaluates the decay/time-step terms from FP32 parameter leaves, + # while every projection in this module remains a BF16 compute parameter. + # FSDP2 requires one original dtype per trainable parameter group, so the + # parallelizer uses this declaration to give the direct parameters their + # own tiny full-precision FSDP unit without changing checkpoint FQNs. + fsdp_full_precision_parameter_names = ("A_log", "dt_bias") + def __init__( self, hidden_size: int = 2048, diff --git a/src/xorl/optim/optimizer.py b/src/xorl/optim/optimizer.py index 426d9efe..ea962fa8 100644 --- a/src/xorl/optim/optimizer.py +++ b/src/xorl/optim/optimizer.py @@ -70,13 +70,32 @@ def _make_param_groups_for_subset( decayed = [p for p in params if name_by_param.get(p) in decay_param_names] undecayed = [p for p in params if name_by_param.get(p) not in decay_param_names] groups: List[Dict[str, Any]] = [] - if decayed: - groups.append({"params": decayed, "weight_decay": weight_decay}) - if undecayed: - groups.append({"params": undecayed, "weight_decay": 0.0}) + for candidates, group_weight_decay in ((decayed, weight_decay), (undecayed, 0.0)): + for homogeneous in _split_dtensor_parameter_groups(candidates): + groups.append({"params": homogeneous, "weight_decay": group_weight_decay}) return groups +def _split_dtensor_parameter_groups(params: Iterable[torch.nn.Parameter]) -> List[List[torch.nn.Parameter]]: + """Keep fused optimizer calls homogeneous in Tensor representation. + + PyTorch's fused AdamW dispatch cannot accept ordinary ``Tensor`` and + ``DTensor`` parameters in the same call. Exact MoE QLoRA legitimately + combines unwrapped expert factors with FSDP-managed projection/router + factors, so split only mixed lists while preserving the historical single + group for homogeneous callers. + """ + + params = list(params) + if not params: + return [] + dtensors = [parameter for parameter in params if isinstance(parameter, DTensor)] + local_tensors = [parameter for parameter in params if not isinstance(parameter, DTensor)] + if not dtensors or not local_tensors: + return [params] + return [dtensors, local_tensors] + + # adapted from https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer_pt_utils.py#L1123 def get_parameter_names(model, forbidden_layer_types, forbidden_param_names): forbidden_layer_types = [] if forbidden_layer_types is None else forbidden_layer_types diff --git a/src/xorl/server/api_server/training_ops.py b/src/xorl/server/api_server/training_ops.py index 0b95b494..f3bcd0cd 100644 --- a/src/xorl/server/api_server/training_ops.py +++ b/src/xorl/server/api_server/training_ops.py @@ -315,7 +315,7 @@ async def forward_backward(self, request: ForwardBackwardRequest) -> ForwardBack # Debug: Log what we got from the engine logger.debug(f"API Server: Received result from engine, keys: {list(result.keys())}") - loss_metrics = {k: v for k, v in result.items() if k.startswith(("is_", "opd_"))} + loss_metrics = {k: v for k, v in result.items() if k.startswith(("is_", "opd_", "router_grad_", "dsv4_"))} if loss_metrics: logger.debug(f"API Server: loss metrics present in result: {list(loss_metrics.keys())}") else: @@ -338,7 +338,7 @@ async def forward_backward(self, request: ForwardBackwardRequest) -> ForwardBack # Add loss-specific metrics if present (already have name:reduction format) for key, value in result.items(): - if key.startswith(("is_", "opd_")): + if key.startswith(("is_", "opd_", "router_grad_", "router_update_", "dsv4_")): # Ensure colon format for tinker compatibility metrics[key if ":" in key else f"{key}:mean"] = value elif key in ( @@ -347,6 +347,10 @@ async def forward_backward(self, request: ForwardBackwardRequest) -> ForwardBack "teacher_hidden_cache_write_s", ): metrics[key] = value + + for key, value in result.items(): + if key.startswith("router_grad_"): + metrics[f"is_{key}:mean"] = value elif ( key.startswith("executor_") or key in PROFILE_TIMING_METRIC_KEYS @@ -440,7 +444,7 @@ async def forward(self, request: ForwardRequest) -> ForwardResponse: "execution_time": result.get("execution_time", 0.0), } for key, value in result.items(): - if key.startswith(("is_", "opd_")): + if key.startswith(("is_", "opd_", "router_grad_", "dsv4_")): metrics[key if ":" in key else f"{key}:mean"] = value elif key in ( "teacher_prefill_tokens", @@ -558,6 +562,9 @@ async def optim_step(self, request: OptimStepRequest) -> OptimStepResponse: for key in ("optim_step_time", "optim_empty_cache_skipped", "glm52_fullparam_publish"): if key in result: metrics[key] = result[key] + for key, value in result.items(): + if key.startswith("router_update_"): + metrics[key] = value return OptimStepResponse( metrics=metrics, diff --git a/src/xorl/server/api_server/weights.py b/src/xorl/server/api_server/weights.py index 979d0b06..6ba34819 100644 --- a/src/xorl/server/api_server/weights.py +++ b/src/xorl/server/api_server/weights.py @@ -716,6 +716,7 @@ async def save_weights_for_sampler(self, request: SaveWeightsForSamplerRequest) model_config = self.model_configs.get(request.model_id, {}) lora_config = model_config.get("lora_config") or {} is_lora = _model_config_is_lora(model_config, default=self.default_session_spec is not None) + lora_serving_mode = lora_config.get("lora_serving_mode") if is_lora: merge_lora_interval = int( (getattr(self, "server_lora_config", {}) or {}).get( @@ -727,7 +728,20 @@ async def save_weights_for_sampler(self, request: SaveWeightsForSamplerRequest) else: merge_lora_interval = 0 - if is_lora and merge_lora_interval == 0: + if lora_serving_mode == "separate" and merge_lora_interval: + raise ValueError("lora_serving_mode='separate' cannot publish periodically merged base weights") + if lora_serving_mode == "merged": + raise ValueError( + "lora_serving_mode='merged' must publish the live canonical W+sBA " + "payload through weight synchronization; save_weights_for_sampler " + "does not materialize a folded full-model snapshot" + ) + publish_separate = is_lora and ( + lora_serving_mode == "separate" or (lora_serving_mode is None and merge_lora_interval == 0) + ) + publish_merged = is_lora and (lora_serving_mode is None and merge_lora_interval > 0) + + if publish_separate: # LoRA with no merge: base weights unchanged, save adapter only engine_request = OrchestratorRequest( operation="save_lora_only", @@ -756,7 +770,7 @@ async def save_weights_for_sampler(self, request: SaveWeightsForSamplerRequest) ) # For LoRA with merge_interval > 0, also save LoRA weights for training recovery - if is_lora and merge_lora_interval > 0: + if publish_merged: lora_save_path = os.path.join(save_path, "lora") lora_request = OrchestratorRequest( operation="save_lora_only", @@ -769,7 +783,7 @@ async def save_weights_for_sampler(self, request: SaveWeightsForSamplerRequest) # Extract results result = output.outputs[0] if output.outputs else {} - saved_path = result.get("lora_path", save_path) if (is_lora and merge_lora_interval == 0) else save_path + saved_path = result.get("lora_path", save_path) if publish_separate else save_path # Validate model_id for xorl_client URI model_id = validate_model_id(request.model_id) @@ -777,9 +791,9 @@ async def save_weights_for_sampler(self, request: SaveWeightsForSamplerRequest) # Build xorl:// URI for the saved checkpoint xorl_uri = self._to_xorl_uri(model_id, request.name, "sampler_weights") - if is_lora and merge_lora_interval == 0: + if publish_separate: save_format = "PEFT LoRA" - elif is_lora and merge_lora_interval > 0: + elif publish_merged: save_format = "safetensors (full weights) + PEFT LoRA" else: save_format = "safetensors (full weights)" diff --git a/src/xorl/server/orchestrator/packing.py b/src/xorl/server/orchestrator/packing.py index c9559a7d..75eed6f4 100644 --- a/src/xorl/server/orchestrator/packing.py +++ b/src/xorl/server/orchestrator/packing.py @@ -334,6 +334,61 @@ def apply_advantages_to_labels( return labels_tensor.tolist() +def derive_sampler_prefill_length( + raw_target_tokens: List[int], + *, + weights: Optional[List[float]], + advantages: Optional[List[float]], + sample_idx: int, +) -> int: + """Recover the sampler's prefill/decode boundary before loss masking. + + Shifted RL payloads may retain ordinary token IDs in every target slot, so + ``IGNORE_INDEX`` alone is not a structural action marker. Binary ``weights`` + are authoritative when present. Advantages also carry the prompt mask when + they contain a live action. An all-zero trajectory may fall back only to an + explicit raw ``IGNORE_INDEX`` prompt mask; ordinary target IDs are + ambiguous and fail closed. + """ + + seq_len = len(raw_target_tokens) + if weights is not None: + if len(weights) != seq_len: + raise ValueError( + f"Sample {sample_idx}: weights length ({len(weights)}) doesn't match target_tokens length ({seq_len})" + ) + first_action = next( + (index for index, weight in enumerate(weights) if weight != 0), + None, + ) + return first_action + 1 if first_action is not None else seq_len + + if advantages is not None: + if len(advantages) != seq_len: + raise ValueError( + f"Sample {sample_idx}: advantages length ({len(advantages)}) " + f"doesn't match target_tokens length ({seq_len})" + ) + first_action = next( + (index for index, advantage in enumerate(advantages) if advantage != 0), + None, + ) + if first_action is not None: + return first_action + 1 + if all(target != IGNORE_INDEX for target in raw_target_tokens): + raise ValueError( + f"Sample {sample_idx}: cannot infer the sampler prefill boundary " + "from ordinary target tokens and all-zero advantages; provide " + "an explicit weights/action mask" + ) + + first_sampled_target = next( + (index for index, target in enumerate(raw_target_tokens) if target != IGNORE_INDEX), + None, + ) + return first_sampled_target + 1 if first_sampled_target is not None else seq_len + + def apply_loss_masks_to_target_tokens( target_tokens: Any, *, @@ -673,9 +728,28 @@ def pack( # Phase 3: build a packed micro-batch per (non-empty) bin. micro_batches: List[Dict[str, Any]] = [] datum_order: List[int] = [] + homogeneous_bins: List[List[_MeasuredSample]] = [] for bin_samples in bins: if not bin_samples: continue + # A shifted RL request needs per-request sampler boundaries for + # exact DSV4 MHC replay, whereas an HF-format request intentionally + # has no such side channel. Never combine the two arithmetic + # contracts in one packed row. Split only at format transitions so + # arrival order and each strategy's bin assignment stay stable. + current_run: List[_MeasuredSample] = [] + current_has_boundary: Optional[bool] = None + for sample in bin_samples: + has_boundary = self._has_sampler_boundary(sample.datum) + if current_run and has_boundary != current_has_boundary: + homogeneous_bins.append(current_run) + current_run = [] + current_run.append(sample) + current_has_boundary = has_boundary + if current_run: + homogeneous_bins.append(current_run) + + for bin_samples in homogeneous_bins: current_batch = self._create_empty_packed_batch(request_id, batch_id=len(micro_batches)) for sample in bin_samples: self._add_sample_to_packed_batch(current_batch, sample.datum, sample.orig_idx) @@ -710,6 +784,19 @@ def _extract_input_ids(datum: Dict[str, Any]) -> Optional[List[int]]: input_ids = input_ids.tolist() if hasattr(input_ids, "tolist") else list(input_ids) return input_ids + @classmethod + def _has_sampler_boundary(cls, datum: Dict[str, Any]) -> bool: + """Return whether ``datum`` is an already-shifted RL request.""" + + input_ids = cls._extract_input_ids(datum) + loss_inputs = datum.get("loss_fn_inputs") + target_tokens = loss_inputs.get("target_tokens") if isinstance(loss_inputs, dict) else None + if "target_tokens" in datum: + target_tokens = datum["target_tokens"] + if input_ids is None or target_tokens is None: + return False + return len(input_ids) == len(target_tokens) + def _measure_and_filter( self, datum_list: List[Dict[str, Any]], @@ -952,6 +1039,7 @@ def _create_empty_packed_batch(self, request_id: str, batch_id: int) -> Dict[str "batch_id": batch_id, "_num_samples": 0, # Internal counter, removed in finalization "_r3_sample_lengths": [], + "_sampler_prefill_lengths_complete": True, } def _add_sample_to_packed_batch( @@ -1035,6 +1123,24 @@ def _add_sample_to_packed_batch( # HF format: len(input_ids) == len(labels) but need shifting is_already_shifted = "target_tokens" in flattened_datum and len(input_ids) == len(labels) + sampler_prefill_length = None + if is_already_shifted: + # Preserve the sampler's arithmetic boundary before loss masks can + # replace generated targets with IGNORE_INDEX. The exact DSV4 + # trainer uses this per-request metadata to replay one prefill + # segment followed by M=1 decode segments inside a packed row. + raw_target_tokens = flattened_datum["target_tokens"] + if not isinstance(raw_target_tokens, list): + raw_target_tokens = ( + raw_target_tokens.tolist() if hasattr(raw_target_tokens, "tolist") else list(raw_target_tokens) + ) + sampler_prefill_length = derive_sampler_prefill_length( + raw_target_tokens, + weights=weights, + advantages=advantages, + sample_idx=sample_idx, + ) + if not is_already_shifted and len(input_ids) == len(labels): # HF format: shift tokens here # input_ids[:-1] predicts labels[1:] @@ -1072,6 +1178,16 @@ def _add_sample_to_packed_batch( seq_len = len(input_ids) batch["_r3_sample_lengths"].append(seq_len) + if sampler_prefill_length is None: + # MHC replay metadata is useful only when it describes every + # request in the packed row. A mixed HF/shifted batch therefore + # carries no boundary side channel rather than a shorter one. + batch["_sampler_prefill_lengths_complete"] = False + batch.pop("sampler_prefill_lengths", None) + elif batch["_sampler_prefill_lengths_complete"]: + # One scalar boundary per packed request. Unlike token-aligned + # fields this list must remain rank one after finalization. + batch.setdefault("sampler_prefill_lengths", []).append(sampler_prefill_length) # Concatenate input_ids to flat list batch["input_ids"].extend(input_ids) @@ -1227,6 +1343,7 @@ def _finalize_packed_batch(self, batch: Dict[str, Any]) -> None: # Store num_samples for reference and remove internal counter num_samples = batch.pop("_num_samples") batch["num_samples"] = num_samples + batch.pop("_sampler_prefill_lengths_complete", None) # Drop OPRD internal accumulators (ints) so they never reach the model forward # as stray kwargs; the offsets they tracked are already baked into the fields. batch.pop("_oprd_cum_teacher_len", None) @@ -1242,6 +1359,7 @@ def _finalize_packed_batch(self, batch: Dict[str, Any]) -> None: "batch_id", "num_samples", "_r3_sample_lengths", + "sampler_prefill_lengths", ]: if isinstance(value, list) and len(value) == len(batch["input_ids"][0]): batch[key] = [value] @@ -1268,6 +1386,7 @@ def _finalize_packed_batch(self, batch: Dict[str, Any]) -> None: "num_samples", "_shifted", "_r3_sample_lengths", + "sampler_prefill_lengths", ): continue if isinstance(value, list) and len(value) == 1 and isinstance(value[0], list): @@ -1389,12 +1508,12 @@ def _add_sample_to_batch( raw_target_tokens = ( raw_target_tokens.tolist() if hasattr(raw_target_tokens, "tolist") else list(raw_target_tokens) ) - first_sampled_target = next( - (index for index, target in enumerate(raw_target_tokens) if target != IGNORE_INDEX), - None, + sampler_prefill_length = derive_sampler_prefill_length( + raw_target_tokens, + weights=weights, + advantages=advantages, + sample_idx=sample_idx, ) - if first_sampled_target is not None: - sampler_prefill_length = first_sampled_target + 1 if labels and not is_already_shifted and len(input_ids) == len(labels): logger.warning( "Sample %s has labels with the same length as input_ids; treating it as HF-format data " diff --git a/src/xorl/server/orchestrator/request_processor.py b/src/xorl/server/orchestrator/request_processor.py index 4829f3e9..bf1a5fba 100644 --- a/src/xorl/server/orchestrator/request_processor.py +++ b/src/xorl/server/orchestrator/request_processor.py @@ -931,9 +931,9 @@ async def _execute_model_pass( } ) - # Add loss-specific metrics (IS/KL divergence, OPD KL stats, ratio stats, etc.) + # Add loss-specific and exact-gradient qualification metrics. for key in result: - if key.startswith(("is_", "opd_")): + if key.startswith(("is_", "opd_", "router_grad_", "router_update_", "dsv4_")): output_dict[key] = result[key] elif key.startswith(FORWARD_BACKWARD_RESULT_PREFIXES): output_dict[key] = result[key] @@ -962,7 +962,11 @@ async def _execute_model_pass( # Unpack per-token outputs if present (tinker API compatibility) if "packed_logprobs" in result and "packed_position_ids" in result: - output_dict["per_sample_outputs"] = self._unpack_per_sample_outputs(result, batches) + packed_outputs = self._unpack_per_sample_outputs(result, batches) + output_dict["per_sample_outputs"] = self._restore_datum_order( + packed_outputs, + datum_order, + ) output = OrchestratorOutputs( request_id=request.request_id, @@ -1049,6 +1053,22 @@ def _teacher_sort_key(datum: Dict[str, Any]) -> int: teacher_id = teacher_ids[0] return int(teacher_id) if teacher_id is not None else 0 + @staticmethod + def _restore_datum_order(outputs: list, datum_order: list[int]) -> list: + """Restore packer-ordered per-sample outputs to surviving input order.""" + + if len(outputs) != len(datum_order): + raise RuntimeError( + "Per-sample output count does not match the packer datum order: " + f"outputs={len(outputs)}, datum_order={len(datum_order)}" + ) + if len(set(datum_order)) != len(datum_order): + raise RuntimeError(f"Packer datum order contains duplicate input indices: {datum_order}") + if any(not isinstance(index, int) or index < 0 for index in datum_order): + raise RuntimeError(f"Packer datum order contains invalid input indices: {datum_order}") + + return [output for _, output in sorted(zip(datum_order, outputs, strict=True))] + @staticmethod def _unpack_per_sample_outputs(result: Dict, batches: list) -> list: """Unpack packed per-token outputs into per-sample lists. @@ -1228,6 +1248,9 @@ def build_output(result): for key in ("optim_step_time", "optim_empty_cache_skipped", "glm52_fullparam_publish"): if key in result: output_dict[key] = result[key] + for key, value in result.items(): + if key.startswith("router_update_"): + output_dict[key] = value if result.get("auto_loaded"): output_dict["auto_loaded"] = True output_dict["auto_load_path"] = result.get("auto_load_path") diff --git a/src/xorl/server/runner/adapters/manager.py b/src/xorl/server/runner/adapters/manager.py index 2aea8674..b78568fa 100644 --- a/src/xorl/server/runner/adapters/manager.py +++ b/src/xorl/server/runner/adapters/manager.py @@ -209,6 +209,25 @@ def get_parallel_plan(self): raise AttributeError("Local model parts do not expose get_parallel_plan") +@dataclass(frozen=True) +class LiveModelLoRAPublisher: + """Minimal PP topology retained after a single-tenant manager detaches.""" + + model: Any + pipeline_parallel_size: int + adapter_process_group: Any + + def materialize_live_model_logical_state_dict(self, *, destination_rank: int = 0) -> Dict[str, torch.Tensor]: + from xorl.lora.utils import get_lora_state_dict # noqa: PLC0415 + + return _gather_pipeline_stage_state( + get_lora_state_dict(self.model), + pipeline_parallel_size=self.pipeline_parallel_size, + adapter_process_group=self.adapter_process_group, + destination_rank=destination_rank, + ) + + def _parameter_layout_tensor(param: Any) -> Any: """Return the tensor carrying a Parameter's static layout metadata.""" @@ -316,6 +335,46 @@ def _optimizer_shard_rank_world() -> Tuple[int, int]: return 0, 1 +def _gather_pipeline_stage_state( + local_stage_state: Dict[str, torch.Tensor], + *, + pipeline_parallel_size: int, + adapter_process_group: Any, + destination_rank: int, +) -> Dict[str, torch.Tensor]: + """Collect one live logical state per PP stage on ``destination_rank``.""" + + rank, world = _optimizer_shard_rank_world() + if world == 1 or pipeline_parallel_size <= 1: + return local_stage_state if rank == destination_rank else {} + + if adapter_process_group is None: + raise RuntimeError("Pipeline adapter publication requires the live stage-local owner group") + stage_ranks = tuple(int(member) for member in torch.distributed.get_process_group_ranks(adapter_process_group)) + if not stage_ranks or rank not in stage_ranks: + raise RuntimeError("Current rank is absent from its pipeline-stage adapter owner group") + stage_payload = local_stage_state if rank == stage_ranks[0] else None + gathered = [None] * world if rank == destination_rank else None + torch.distributed.gather_object( + stage_payload, + object_gather_list=gathered, + dst=destination_rank, + ) + if rank != destination_rank: + return {} + + merged: Dict[str, torch.Tensor] = {} + assert gathered is not None + for payload in gathered: + if payload is None: + continue + for name, tensor in payload.items(): + if name in merged: + raise RuntimeError(f"Pipeline adapter publication produced duplicate parameter {name!r}") + merged[name] = tensor + return merged + + def _adapter_layout_world_identity(adapter_state: "AdapterState") -> Tuple[int, Tuple[int, ...]]: """Return the stage-local layout world captured during registration.""" @@ -3121,43 +3180,54 @@ def get_layout(self, model_id: str, fqn: str) -> AdapterTensorLayout: return layout raise KeyError(f"No adapter layout for {fqn!r}") from None + def _gather_pipeline_stage_state( + self, + local_stage_state: Dict[str, torch.Tensor], + *, + destination_rank: int, + ) -> Dict[str, torch.Tensor]: + """Collect one live logical state per PP stage on ``destination_rank``.""" + + return _gather_pipeline_stage_state( + local_stage_state, + pipeline_parallel_size=self._pipeline_parallel_size, + adapter_process_group=self._adapter_process_group, + destination_rank=destination_rank, + ) + + def make_live_model_lora_publisher(self) -> LiveModelLoRAPublisher: + """Detach only the live model view and PP publication topology.""" + + return LiveModelLoRAPublisher( + model=self.model, + pipeline_parallel_size=self._pipeline_parallel_size, + adapter_process_group=self._adapter_process_group, + ) + def materialize_logical_state_dict(self, model_id: str, *, destination_rank: int = 0) -> Dict[str, torch.Tensor]: """Collectively reconstruct every PP stage's active logical weights.""" self.prepare_forward(model_id) + return self.materialize_live_model_logical_state_dict(destination_rank=destination_rank) + + def materialize_live_model_logical_state_dict(self, *, destination_rank: int = 0) -> Dict[str, torch.Tensor]: + """Publish current live factors without restoring compact adapter slots. + + The exact GLM-5.2 train-router lane detaches the multi-adapter optimizer + after materializing its single default adapter. Its shared optimizer + then updates the live model parameters, so calling ``prepare_forward`` + during publication would overwrite those updates with stale compact + slots. This path treats the live model as source of truth while using + the manager's PP owner topology. + """ + from xorl.lora.utils import get_lora_state_dict # noqa: PLC0415 local_stage_state = get_lora_state_dict(self.model) - rank, world = _optimizer_shard_rank_world() - if world == 1 or self._pipeline_parallel_size <= 1: - return local_stage_state if rank == destination_rank else {} - - stage_group = self._adapter_process_group - if stage_group is None: - raise RuntimeError("Pipeline adapter publication requires the live stage-local owner group") - stage_ranks = tuple(int(member) for member in torch.distributed.get_process_group_ranks(stage_group)) - if not stage_ranks or rank not in stage_ranks: - raise RuntimeError("Current rank is absent from its pipeline-stage adapter owner group") - stage_payload = local_stage_state if rank == stage_ranks[0] else None - gathered = [None] * world if rank == destination_rank else None - torch.distributed.gather_object( - stage_payload, - object_gather_list=gathered, - dst=destination_rank, + return self._gather_pipeline_stage_state( + local_stage_state, + destination_rank=destination_rank, ) - if rank != destination_rank: - return {} - - merged: Dict[str, torch.Tensor] = {} - assert gathered is not None - for payload in gathered: - if payload is None: - continue - for name, tensor in payload.items(): - if name in merged: - raise RuntimeError(f"Pipeline adapter publication produced duplicate parameter {name!r}") - merged[name] = tensor - return merged def load_logical_state_dict(self, model_id: str, state_dict: Dict[str, torch.Tensor]) -> None: """Pack full active logical tensors into this rank's local adapter slots.""" diff --git a/src/xorl/server/runner/checkpoint/manager.py b/src/xorl/server/runner/checkpoint/manager.py index a50e6ac3..c49f7f21 100644 --- a/src/xorl/server/runner/checkpoint/manager.py +++ b/src/xorl/server/runner/checkpoint/manager.py @@ -16,7 +16,7 @@ import shutil import time from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Dict, Optional +from typing import Any, Dict, Optional, Sequence import torch import torch.distributed as dist @@ -36,12 +36,19 @@ contains_dsv4_exact_active_lora_component, contains_glm52_exact_active_lora_component, contains_glm52_fullparam_component, + glm52_exact_active_lora_enabled, ) from xorl.server.runner.adapters.manager import ( + LocalModelPartsView, adapter_gradient_ownership_checkpoint_metadata, save_adapter_optimizer_shards, ) from xorl.server.session_spec import write_session_spec +from xorl.server.weight_sync.glm52_router_bundle import ( + gather_glm52_router_weights_across_ranks, + mark_adapter_config_with_glm52_router_bundle, + save_glm52_router_bundle, +) from xorl.utils import helper from xorl.utils.device import get_device_type @@ -67,8 +74,13 @@ def __init__( rank: int, local_rank: int, adapter_manager=None, + model_parts: Optional[Sequence[nn.Module]] = None, ): self.model = model + self.model_parts = tuple(model_parts) if model_parts is not None else (model,) + self._local_model_view = ( + self.model_parts[0] if len(self.model_parts) == 1 else LocalModelPartsView(self.model_parts) + ) self.optimizer = optimizer self.Checkpointer = checkpointer self.lora_config = lora_config @@ -77,6 +89,10 @@ def __init__( self.rank = rank self.local_rank = local_rank self._adapter_manager = adapter_manager + # The exact single-tenant GLM lane detaches adapter switching after + # promotion, but retains a lightweight PP topology publisher for the + # optimizer-owned live model factors. + self._detached_adapter_publisher = None # These will be set/updated by ModelRunner self.global_step = 0 @@ -89,6 +105,17 @@ def __init__( # Internal helpers # ------------------------------------------------------------------ + def _model_view_for_local_parts(self): + return getattr(self, "_local_model_view", self.model) + + def _has_glm52_exact_active_lora(self) -> bool: + # Config is replicated across PP stages, whereas value components are + # intentionally absent from dense-only stages. Use the replicated + # contract stamp to keep publication collectives rank-symmetric. + return glm52_exact_active_lora_enabled(getattr(self.model, "config", None)) or ( + contains_glm52_exact_active_lora_component(self._model_view_for_local_parts()) + ) + def _get_lora_save_config(self): """Get target_modules and lora_alpha for PEFT-format saving. @@ -259,9 +286,12 @@ def _write_adapter_training_artifacts( def _gather_adapter_lora_params(self, model_id: str) -> Dict[str, torch.Tensor]: """Collectively reconstruct full active logical tensors on rank 0.""" - if self._adapter_manager is None: - return get_lora_state_dict(self.model) - return self._adapter_manager.materialize_logical_state_dict(model_id, destination_rank=0) + if self._adapter_manager is not None: + return self._adapter_manager.materialize_logical_state_dict(model_id, destination_rank=0) + publisher = getattr(self, "_detached_adapter_publisher", None) + if publisher is not None: + return publisher.materialize_live_model_logical_state_dict(destination_rank=0) + return get_lora_state_dict(self.model) def _adapter_publication_error(self, model_id: str, *, strict: bool) -> Optional[str]: if self._adapter_manager is None: @@ -295,7 +325,7 @@ def _require_factor_only_exact_active_lora(self, operation: str) -> None: f"{operation} cannot materialize a merged/full-weight snapshot. Export all 948 factors " "as dsv4_expert_banks and load a fresh sampler adapter version." ) - if contains_glm52_exact_active_lora_component(self.model): + if self._has_glm52_exact_active_lora(): raise RuntimeError( "GLM-5.2 exact active-LoRA composites require factor-only adapter publication; " f"{operation} cannot materialize a merged/full-weight snapshot. Export the complete 1,700-factor " @@ -379,11 +409,24 @@ def _save_lora_weights(self, save_path: str, model_id: str, *, preserve_lora_dty # Ensure adapter weights are synced to model if self._adapter_manager is not None: self._adapter_manager.switch_adapter(model_id, auto_register=True) + publication_step = self._adapter_manager.get_global_step(model_id) + else: + publication_step = self.global_step # Always reconstruct logical tensors through the model's EP+FSDP path; # adapter optimizer slots are rank-local and are never PEFT payloads. logger.info(f"Rank {self.rank}: Using collective topology-aware LoRA weight gathering") lora_state_dict = self._gather_adapter_lora_params(model_id) + exact_glm52_router_state = None + # The router is part of GLM-5.2's trainer/sampler numerical identity + # even when it is frozen. ``train_router`` controls gradients, not + # publication: an exact active-LoRA sampler must receive the trainer's + # router values at every published weight version. + if self._has_glm52_exact_active_lora(): + exact_glm52_router_state = gather_glm52_router_weights_across_ranks( + self._model_view_for_local_parts(), + destination_rank=0, + ) # Only rank 0 writes files if self.rank == 0: @@ -413,11 +456,25 @@ def _save_lora_weights(self, save_path: str, model_id: str, *, preserve_lora_dty lora_export_format=lora_export_format, preserve_lora_dtype=preserve_lora_dtype, ) + if exact_glm52_router_state is not None: + manifest = save_glm52_router_bundle( + save_path, + exact_glm52_router_state, + weight_step=publication_step, + expected_layer_ids=list( + range( + int(getattr(self.model.config, "first_k_dense_replace")), + int(getattr(self.model.config, "num_hidden_layers")), + ) + ), + ) + mark_adapter_config_with_glm52_router_bundle(save_path, manifest) if adapter_session_spec is not None: write_session_spec(save_path, adapter_session_spec) # Cleanup del lora_state_dict + del exact_glm52_router_state gc.collect() torch.cuda.empty_cache() diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index d34654c5..ce36cc7c 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -61,6 +61,7 @@ stage_ids_for_rank, ) from xorl.distributed.sequence_parallel.data import gather_outputs +from xorl.distributed.torch_parallelize import refresh_ep_param_groups from xorl.lora import LoraLinear from xorl.lora.expert_adapter_contract import ( ExpertAdapterFactorOwnership, @@ -69,9 +70,13 @@ ) from xorl.lora.fold import invalidate_lora_merged_weight_caches, lora_merged_forward_enabled from xorl.models import resolve_cross_entropy_mode +from xorl.models.exact_contract import glm52_exact_active_lora_enabled from xorl.models.layers.moe.routing_replay import set_replay_stage from xorl.models.transformers.deepseek_v3.support import deepseek_v3_default_lora_targets -from xorl.models.transformers.deepseek_v4.exact_contract import DSV4_FLASH_REQUIRED_TARGET_MODULES +from xorl.models.transformers.deepseek_v4.exact_contract import ( + DSV4_FLASH_LOGICAL_FACTOR_COUNT, + DSV4_FLASH_REQUIRED_TARGET_MODULES, +) from xorl.models.transformers.glm5.index_share import IndexShareMode from xorl.models.transformers.glm5.support import glm5_default_lora_targets from xorl.ops.batch_invariant_ops import enable_batch_invariant_mode, get_batch_invariant_ops @@ -697,6 +702,11 @@ def __init__( self._initialize_checkpointer() self._checkpoint_mgr = self._build_checkpoint_manager() self._restore_initial_base_checkpoint() + self._glm52_exact_router_single_tenant = bool( + self.lora_config.get("enable_lora", False) + and bool(getattr(self.model.config, "train_router", False)) + and glm52_exact_active_lora_enabled(self.model.config) + ) if enable_full_determinism: # Enabling deterministic algorithms before Kimi DCP/meta materialization # makes startup pathologically slow; training and adapter init happen below. @@ -730,7 +740,24 @@ def __init__( lora_config=self.lora_config, ) self._initialize_default_lora_adapter() - logger.info("Multi-adapter manager initialized with default adapter") + if self._glm52_exact_router_single_tenant: + # The manager is used once to materialize the deterministic + # default adapter into the post-DCP model. Router weights are + # shared base parameters, so retaining per-adapter optimizer + # state would update LoRA factors while silently omitting the + # routers. From this point the ordinary shared optimizer owns + # both sets of model Parameters as one single-tenant unit. + # Adapter registration owns compact slot tensors; copy the + # deterministic default state into the live model before the + # manager is detached. DCP ``to_empty`` may also have replaced + # every Parameter object, so both the EP groups and optimizer + # constructed before restore must be rebound to current model + # identities. The pre-restore optimizer has never stepped and + # the admitted base checkpoint intentionally has no optimizer + # payload, so discarding it here loses no state. + self._promote_exact_glm52_default_adapter_to_shared_optimizer() + else: + logger.info("Multi-adapter manager initialized with default adapter") # Initialize tokenizer for sampling (only on rank 0) if self.rank == 0: @@ -808,6 +835,17 @@ def register_session( initialize_fresh: bool = True, ) -> Dict[str, Any]: """Register a normalized session runtime spec on this worker.""" + if getattr(self, "_glm52_exact_router_single_tenant", False) and self._adapter_manager is None: + self._validate_single_tenant(model_id) + if model_id != "default": + raise ValueError("Exact GLM-5.2 train_router QLoRA is single-tenant and admits only model_id='default'") + return { + "model_id": model_id, + "registered": True, + "materialized": True, + "message": "Exact GLM-5.2 router and adapter factors use the shared single-tenant optimizer.", + } + if not self.lora_enabled: # Full-weight mode remains effectively single-tenant; keep the API # tolerant of create_model but don't install heterogeneous runtime state. @@ -869,6 +907,25 @@ def _initialize_default_lora_adapter(self) -> None: self._adapter_manager.current_adapter_id = "default" self._checkpoint_mgr._adapter_manager = self._adapter_manager + def _promote_exact_glm52_default_adapter_to_shared_optimizer(self) -> None: + """Bind current post-DCP model parameters to the exact single-tenant optimizer.""" + + if self._adapter_manager is None or not self._adapter_manager.has_adapter("default"): + raise RuntimeError("Exact GLM-5.2 shared optimizer requires a materialized default adapter") + self._adapter_manager.prepare_forward("default") + if get_parallel_state().ep_enabled: + refresh_ep_param_groups(self.model) + self._initialize_optimizer() + self._checkpoint_mgr.optimizer = self.optimizer + self._checkpoint_mgr._detached_adapter_publisher = self._adapter_manager.make_live_model_lora_publisher() + self._adapter_manager = None + self._checkpoint_mgr._adapter_manager = None + self._active_session_id = "default" + logger.info( + "Exact GLM-5.2 train_router lane selected single-tenant shared optimizer " + "after deterministic default-adapter materialization and post-DCP parameter rebinding" + ) + @staticmethod def _adapter_gradient_hash(payload: Any) -> str: return hashlib.sha256(json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()).hexdigest() @@ -1373,7 +1430,7 @@ def _validate_single_tenant(self, model_id: str) -> None: Raises: ValueError: If a different session is already active. """ - if self.lora_enabled: + if self.lora_enabled and not getattr(self, "_glm52_exact_router_single_tenant", False): return # Multi-tenant allowed for LoRA mode if self._active_session_id is None: @@ -1381,7 +1438,7 @@ def _validate_single_tenant(self, model_id: str) -> None: logger.info(f"Full-weights session started: {model_id}") elif self._active_session_id != model_id: raise ValueError( - f"Full-weights mode is single-tenant. Active session: {self._active_session_id}, " + f"Shared-optimizer mode is single-tenant. Active session: {self._active_session_id}, " f"requested: {model_id}. Call /api/v1/kill_session first to start a new session." ) @@ -1575,10 +1632,12 @@ def _initialize_model(self): deepep_buffer_size_gb=self.model_config.get("deepep_buffer_size_gb", 2.0), deepep_num_sms=self.model_config.get("deepep_num_sms", 20), deepep_async_combine=self.model_config.get("deepep_async_combine", False), + deepep_native_exact=self.model_config.get("deepep_native_exact", False), alltoall_combine_hidden_chunk_size=self.model_config.get("alltoall_combine_hidden_chunk_size", 0), init_device=self.train_config.get("init_device", "cpu"), merge_qkv=self.model_config.get("merge_qkv", True), enable_lora=lora_enabled, + lora_serving_mode=self.lora_config.get("lora_serving_mode"), lora_rank=self.lora_config.get("max_lora_rank", self.lora_config.get("lora_rank", 32)), lora_alpha=self.lora_config.get("lora_alpha", 16), lora_b_init_std=self.lora_config.get("lora_b_init_std", 0.0), @@ -1984,6 +2043,7 @@ def _build_checkpoint_manager(self, adapter_manager=None) -> CheckpointManager: rank=self.rank, local_rank=self.local_rank, adapter_manager=adapter_manager, + model_parts=self.model_parts if self.pp_enabled else None, ) def _load_initial_checkpoint(self) -> None: @@ -2575,12 +2635,18 @@ def remove(self) -> None: "moe_native_gathered_routing": 81, "moe_native_gathered_ids": 82, "moe_native_local_ids": 83, + "moe_native_recv_hidden": 80, + "moe_native_recv_weights": 81, + "moe_native_recv_local_ids": 82, + "moe_native_expert_start": 83, + "moe_native_recv_leaf": 84, "moe_native_routed": 84, "moe_native_shared_gate_value": 85, "moe_native_shared_gate_up": 86, "moe_native_shared_act": 87, "moe_native_shared_down": 88, "moe_native_local_partial": 89, + "moe_native_shared_folded": 89, "moe_native_combined": 90, "gdn_q_input": 60, "gdn_k_input": 61, @@ -2625,7 +2691,18 @@ def capture(layer_idx: int, name: str, value: Any) -> None: "capture_index": len(captures), "layer": layer_idx, "name": name, - "order": component_order[name], + "order": ( + 83 + if name.startswith( + ( + "moe_native_gate_up_packed_local_", + "moe_native_gate_up_scale_local_", + "moe_native_down_packed_local_", + "moe_native_down_scale_local_", + ) + ) + else component_order[name] + ), "tensor": snapshot, } ) @@ -3193,7 +3270,15 @@ def _compute_token_diagnostics( for _ in range(valid_indices.shape[0]) ] for layer_index, layer_hidden in enumerate(all_hidden_states): - layer_flat = layer_hidden.reshape(-1, layer_hidden.shape[-1]) + if layer_hidden.numel() % labels_flat.numel() != 0: + raise ValueError( + "diagnostic hidden-state rows do not align with labels: " + f"hidden={tuple(layer_hidden.shape)} labels={tuple(labels.shape)}" + ) + # DSV4's residual stream is [B, S, hc_mult, H] between + # decoder layers. Treat hc_mult*H as one feature vector so + # each label position still maps to exactly one hidden row. + layer_flat = layer_hidden.reshape(labels_flat.numel(), -1) rows = layer_flat[valid_indices].float() row_mean = rows.mean(dim=-1) row_std = rows.std(dim=-1, unbiased=False) @@ -3276,6 +3361,71 @@ def _compute_token_diagnostics( return diagnostics + @staticmethod + def _compute_hidden_state_diagnostics( + *, + all_hidden_states: tuple[torch.Tensor, ...], + labels: torch.Tensor, + hidden_sample_count: int = 8, + hidden_sample_indices: Any = None, + ) -> dict[str, Any]: + """Summarize hidden rows without requiring a materialized global lm head.""" + + labels_flat = labels.reshape(-1) + valid_indices = (labels_flat != IGNORE_INDEX).nonzero(as_tuple=True)[0] + if valid_indices.numel() == 0: + return {"valid_positions": [], "hidden_state_summaries": []} + if not all_hidden_states: + raise ValueError("hidden-state diagnostics require at least one hidden-state tensor") + sample_indices = ModelRunner._build_diagnostic_sample_indices( + hidden_dim=all_hidden_states[-1].shape[-1], + hidden_sample_count=hidden_sample_count, + hidden_sample_indices=hidden_sample_indices, + device=all_hidden_states[-1].device, + ) + per_token_summaries = [ + { + "layer_count": len(all_hidden_states), + "sample_indices": sample_indices.cpu().tolist(), + "layers": [], + } + for _ in range(valid_indices.shape[0]) + ] + for layer_index, layer_hidden in enumerate(all_hidden_states): + if layer_hidden.numel() % labels_flat.numel() != 0: + raise ValueError( + "diagnostic hidden-state rows do not align with labels: " + f"hidden={tuple(layer_hidden.shape)} labels={tuple(labels.shape)}" + ) + layer_flat = layer_hidden.reshape(labels_flat.numel(), -1) + rows = layer_flat[valid_indices].float() + row_mean = rows.mean(dim=-1) + row_std = rows.std(dim=-1, unbiased=False) + row_rms = torch.sqrt(torch.mean(rows * rows, dim=-1)) + row_max_abs = rows.abs().amax(dim=-1) + row_min = rows.amin(dim=-1) + row_max = rows.amax(dim=-1) + sampled_values = ( + rows[:, sample_indices] if sample_indices.numel() > 0 else rows.new_empty((rows.shape[0], 0)) + ) + for token_index, summary in enumerate(per_token_summaries): + summary["layers"].append( + { + "index": layer_index, + "mean": float(row_mean[token_index].item()), + "std": float(row_std[token_index].item()), + "rms": float(row_rms[token_index].item()), + "max_abs": float(row_max_abs[token_index].item()), + "min": float(row_min[token_index].item()), + "max": float(row_max[token_index].item()), + "sample_values": sampled_values[token_index].cpu().tolist(), + } + ) + return { + "valid_positions": valid_indices.cpu().tolist(), + "hidden_state_summaries": per_token_summaries, + } + @staticmethod def _teacher_cache_dtype(dtype_name: str) -> torch.dtype: mapping = { @@ -6033,7 +6183,7 @@ def _compute_decode_cache_micro_batch_loss( layer_hidden.new_zeros( input_ids.shape[0], input_ids.shape[1], - layer_hidden.shape[-1], + *layer_hidden.shape[2:], ) for layer_hidden in segment_hidden_states ] @@ -6043,7 +6193,7 @@ def _compute_decode_cache_micro_batch_loss( f"{len(segment_hidden_states)} != {len(all_hidden_states)}" ) for layer_index, layer_hidden in enumerate(segment_hidden_states): - all_hidden_states[layer_index][:, start:end, :] = layer_hidden + all_hidden_states[layer_index][:, start:end, ...] = layer_hidden finally: self._clear_diagnostic_decode_cache(self.model) if was_training: @@ -6084,7 +6234,17 @@ def _compute_decode_cache_micro_batch_loss( "diagnostic_decode_cache diagnostic_topk requires return_per_token=True; skipping token diagnostics" ) elif diagnostic_topk > 0 and loss_tp_group is not None: - logger.warning("diagnostic_decode_cache diagnostic_topk is not supported with vocab-parallel lm_head") + logger.warning( + "diagnostic_decode_cache top-k logits are unavailable with vocab-parallel lm_head; " + "emitting hidden-state summaries only" + ) + if diagnostic_hidden_states: + per_token_outputs["token_diagnostics"] = self._compute_hidden_state_diagnostics( + all_hidden_states=(tuple(all_hidden_states) if all_hidden_states is not None else (hidden_states,)), + labels=labels, + hidden_sample_count=diagnostic_hidden_sample_count, + hidden_sample_indices=diagnostic_hidden_sample_indices, + ) elif diagnostic_topk > 0: token_diagnostics = self._compute_token_diagnostics( hidden_states=hidden_states, @@ -6206,6 +6366,12 @@ def _profile_elapsed_ms(start: float) -> float: force_weights=diagnostic_moe_routing_reference_weights, ) + if "exact_value_logprobs" in micro_batch: + raise ValueError( + "exact_value_logprobs is not a valid training input: exact policy values must be " + "independently recomputed by the differentiated trainer program" + ) + if bool(params.get("diagnostic_decode_cache", False)): if loss_fn not in ["causallm_loss", "cross_entropy"]: raise ValueError("diagnostic_decode_cache only supports causallm_loss/cross_entropy") @@ -6364,7 +6530,21 @@ def _profile_elapsed_ms(start: float) -> float: "diagnostic_topk requires return_per_token=True for per-sample unpacking; skipping diagnostics" ) elif diagnostic_topk > 0 and loss_tp_group is not None: - logger.warning("diagnostic_topk is not supported with vocab-parallel lm_head; skipping diagnostics") + logger.warning( + "diagnostic top-k logits are unavailable with vocab-parallel lm_head; " + "emitting hidden-state summaries only" + ) + if diagnostic_hidden_states: + per_token_outputs["token_diagnostics"] = self._compute_hidden_state_diagnostics( + all_hidden_states=( + tuple(diagnostic_all_hidden_states) + if diagnostic_all_hidden_states is not None + else (hidden_states,) + ), + labels=labels, + hidden_sample_count=diagnostic_hidden_sample_count, + hidden_sample_indices=diagnostic_hidden_sample_indices, + ) elif diagnostic_topk > 0: token_diagnostics = self._compute_token_diagnostics( hidden_states=hidden_states, @@ -6495,7 +6675,6 @@ def _profile_elapsed_ms(start: float) -> float: per_token_outputs["loss"] = _result.per_token_loss is_metrics = dict(_result.metrics or {}) is_metrics.setdefault("valid_tokens", int((target_tokens != IGNORE_INDEX).sum().item())) - elif loss_fn == "policy_loss": target_tokens = micro_batch.get("target_tokens", micro_batch.get("labels")) old_logprobs = micro_batch["logprobs"] @@ -7388,6 +7567,27 @@ def _compute_pp_terminal_objective( key: value.to(terminal_hidden.device, non_blocking=True) if isinstance(value, torch.Tensor) else value for key, value in metadata.micro_batch.items() } + terminal_token_shape = terminal_hidden.shape[:-1] + for key in ( + "labels", + "target_tokens", + "logprobs", + "old_logprobs", + "ref_logprobs", + "rollout_logprobs", + "advantages", + "weights", + "logprob_temperatures", + "logprob_top_ks", + "logprob_top_ps", + "logprob_min_ps", + ): + value = micro_batch.get(key) + if isinstance(value, torch.Tensor) and value.shape != terminal_token_shape: + raise ValueError( + f"Physical PP terminal field {key!r} has shape {tuple(value.shape)}, " + f"but terminal hidden rows have shape {tuple(terminal_token_shape)}" + ) last_part = self._pp_last_stage_part() if last_part is None or getattr(last_part, "lm_head", None) is None: raise RuntimeError("Physical PP terminal objective requires the local terminal lm_head") @@ -8190,6 +8390,8 @@ def _forward_backward_impl( """ self._check_not_sleeping("forward_backward") params = loss_fn_params or {} + if bool(getattr(self.model.config, "train_router", False)): + self._reset_glm52_router_gradient_evidence() # Defragment GPU memory at the top of every forward_backward call by default. # After weight-sync + optim_step from the previous step the CUDA @@ -8485,6 +8687,12 @@ def _forward_backward_impl( stage="after_efsdp_reduce_scatter", ) + if bool(getattr(self.model.config, "_dsv4_flash_exact_active_lora", False)): + result.update(self._collect_dsv4_adapter_gradient_metrics(model_id)) + + if bool(getattr(self.model.config, "train_router", False)): + result.update(self._collect_glm52_router_gradient_metrics()) + # Get step counter (use adapter manager if available, else global) if self._adapter_manager is not None: current_step = self._adapter_manager.get_adapter_state(model_id).global_forward_backward_step @@ -8549,6 +8757,293 @@ def _forward_backward_impl( return result + def _collect_dsv4_adapter_gradient_metrics(self, model_id: str) -> Dict[str, Any]: + """Inspect the authoritative DSV4 adapter numerators before commit. + + Multi-adapter LoRA does not optimize ``Parameter.grad`` directly. The + adapter ownership finalizer first validates and packs each local shard + into staged FP32 numerator tensors; those are the values consumed by + ``LoRAAdapterManager.optim_step`` after the distributed completion + rendezvous. Qualification must therefore measure this staged image, + while separately proving that every DSV4 router remains frozen and has + no gradient. + """ + + if self._adapter_manager is None: + raise RuntimeError("Exact DSV4 active-LoRA gradient evidence requires the adapter manager") + state = self._adapter_manager.get_adapter_state(model_id) + plan = state.gradient_ownership_plan + scratch = state.gradient_scratch + if plan is None or not scratch.capture_open or not scratch.capture_staged: + raise RuntimeError("Exact DSV4 gradient evidence requires a fully staged adapter capture") + + planned_names = tuple(item.fqn for item in plan.parameters) + non_lora_planned = [name for name in planned_names if "lora_A" not in name and "lora_B" not in name] + if len(planned_names) != DSV4_FLASH_LOGICAL_FACTOR_COUNT or non_lora_planned: + raise RuntimeError( + "Exact DSV4 optimizer ownership must contain exactly the complete " + f"{DSV4_FLASH_LOGICAL_FACTOR_COUNT}-factor LoRA inventory; " + f"actual={len(planned_names)} non_lora={non_lora_planned[:8]}" + ) + staged_names = set(scratch.staged_parameter_fqns) + if staged_names != set(scratch.staged_numerators): + raise RuntimeError("Exact DSV4 staged gradient names and numerator tensors disagree") + if staged_names - set(planned_names): + raise RuntimeError("Exact DSV4 staged gradients contain parameters outside the ownership plan") + + local_sum_sq = 0.0 + local_elements = local_nonzero = local_nonfinite = 0 + metric_device: torch.device | None = None + for numerator in scratch.staged_numerators.values(): + values = numerator.detach().float() + metric_device = values.device + local_elements += values.numel() + finite = torch.isfinite(values) + local_nonfinite += int((~finite).sum().item()) + finite_values = values[finite] + local_nonzero += int(torch.count_nonzero(finite_values).item()) + local_sum_sq += float(torch.sum(finite_values * finite_values, dtype=torch.float64).item()) + + router_count = router_requires_grad = router_grad_present = 0 + router_grad_nonzero = router_grad_nonfinite = 0 + trainable_non_lora = 0 + seen_parameters: set[int] = set() + model_parts = self.model_parts if self.pp_enabled else [self.model] + for part in model_parts: + for name, parameter in part.named_parameters(): + if id(parameter) in seen_parameters: + continue + seen_parameters.add(id(parameter)) + is_lora = "lora_A" in name or "lora_B" in name + if parameter.requires_grad and not is_lora: + trainable_non_lora += 1 + if not name.endswith("mlp.gate.weight"): + continue + router_count += 1 + router_requires_grad += int(parameter.requires_grad) + gradient = parameter.grad + if gradient is None: + continue + router_grad_present += 1 + to_local = getattr(gradient, "to_local", None) + if callable(to_local): + gradient = to_local() + wait = getattr(gradient, "wait", None) + if callable(wait): + gradient = wait() + values = gradient.detach().float() + metric_device = metric_device or values.device + finite = torch.isfinite(values) + router_grad_nonfinite += int((~finite).sum().item()) + router_grad_nonzero += int(torch.count_nonzero(values[finite]).item()) + + device = metric_device or torch.device(get_device_type()) + summed = torch.tensor( + [ + local_sum_sq, + local_elements, + local_nonzero, + local_nonfinite, + len(staged_names), + trainable_non_lora, + router_requires_grad, + router_grad_present, + router_grad_nonzero, + router_grad_nonfinite, + ], + dtype=torch.float64, + device=device, + ) + extrema = torch.tensor( + [len(planned_names), len(staged_names), router_count], + dtype=torch.float64, + device=device, + ) + minimum = extrema.clone() + maximum = extrema.clone() + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(summed, op=dist.ReduceOp.SUM) + dist.all_reduce(minimum, op=dist.ReduceOp.MIN) + dist.all_reduce(maximum, op=dist.ReduceOp.MAX) + ( + sum_sq, + elements, + nonzero, + nonfinite, + staged_total, + trainable_non_lora_total, + router_requires_grad_total, + router_grad_present_total, + router_grad_nonzero_total, + router_grad_nonfinite_total, + ) = summed.cpu().tolist() + planned_min, staged_min, router_min = minimum.cpu().tolist() + planned_max, staged_max, router_max = maximum.cpu().tolist() + metrics = { + "dsv4_grad_norm": math.sqrt(sum_sq), + "dsv4_grad_element_count": int(elements), + "dsv4_grad_nonzero_count": int(nonzero), + "dsv4_grad_nonfinite_count": int(nonfinite), + "dsv4_grad_staged_factor_count_total": int(staged_total), + "dsv4_grad_planned_factor_count_min": int(planned_min), + "dsv4_grad_planned_factor_count_max": int(planned_max), + "dsv4_grad_staged_factor_count_min": int(staged_min), + "dsv4_grad_staged_factor_count_max": int(staged_max), + "dsv4_grad_trainable_non_lora_count": int(trainable_non_lora_total), + "dsv4_router_tensor_count_min": int(router_min), + "dsv4_router_tensor_count_max": int(router_max), + "dsv4_router_requires_grad_count": int(router_requires_grad_total), + "dsv4_router_grad_present_count": int(router_grad_present_total), + "dsv4_router_grad_nonzero_count": int(router_grad_nonzero_total), + "dsv4_router_grad_nonfinite_count": int(router_grad_nonfinite_total), + } + logger.info("DSV4 authoritative adapter gradient evidence: %s", metrics) + return metrics + + def _collect_glm52_router_gradient_metrics(self) -> Dict[str, Any]: + """Return rank-aggregated evidence for exact GLM router gradients. + + This intentionally observes the post-backward, post-eFSDP-reduction + gradient without mutating it. The metrics make a cluster gate prove + that all 75 sparse-layer gates participated and that their gradients + are finite and nonzero; a global aggregate prevents rank-0-only false + confidence. + """ + + local_tensors = local_missing = local_nonfinite = 0 + local_elements = local_nonzero = 0 + local_sum_sq = 0.0 + metric_device: torch.device | None = None + seen_parameters: set[int] = set() + for module in self.model.modules(): + if type(module).__name__ != "Glm5TopkRouter": + continue + parameter = getattr(module, "weight", None) + if not isinstance(parameter, torch.Tensor) or id(parameter) in seen_parameters: + continue + seen_parameters.add(id(parameter)) + local_tensors += 1 + evidence = getattr(module, "_gradient_evidence", None) + if evidence is not None: + local_sum_sq += float(evidence["sum_sq"].item()) + local_nonfinite += int(evidence["nonfinite"].item()) + local_nonzero += int(evidence["nonzero"].item()) + local_elements += int(evidence["elements"]) + metric_device = evidence["sum_sq"].device + continue + gradient = parameter.grad + if gradient is None: + local_missing += 1 + continue + to_local = getattr(gradient, "to_local", None) + if callable(to_local): + gradient = to_local() + wait = getattr(gradient, "wait", None) + if callable(wait): + gradient = wait() + values = gradient.detach().float() + metric_device = values.device + local_elements += values.numel() + finite = torch.isfinite(values) + local_nonfinite += int((~finite).sum().item()) + finite_values = values[finite] + local_nonzero += int(torch.count_nonzero(finite_values).item()) + local_sum_sq += float(torch.sum(finite_values * finite_values, dtype=torch.float64).item()) + + device = metric_device or torch.device(get_device_type()) + aggregate = torch.tensor( + [local_sum_sq, local_tensors, local_missing, local_nonfinite, local_elements, local_nonzero], + dtype=torch.float64, + device=device, + ) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(aggregate, op=dist.ReduceOp.SUM) + sum_sq, tensors, missing, nonfinite, elements, nonzero = aggregate.cpu().tolist() + metrics = { + "router_grad_norm": math.sqrt(sum_sq), + "router_grad_tensor_count": int(tensors), + "router_grad_missing_count": int(missing), + "router_grad_nonfinite_count": int(nonfinite), + "router_grad_element_count": int(elements), + "router_grad_nonzero_count": int(nonzero), + } + logger.info("GLM-5.2 router gradient evidence: %s", metrics) + return metrics + + def _reset_glm52_router_gradient_evidence(self) -> None: + count = 0 + for module in self.model.modules(): + reset = getattr(module, "reset_gradient_evidence", None) + if type(module).__name__ == "Glm5TopkRouter" and callable(reset): + reset() + count += 1 + if count == 0: + raise RuntimeError("train_router=True but no GLM-5.2 router modules were found to arm gradient evidence") + + def _snapshot_glm52_router_weights_for_step(self) -> list[tuple[nn.Module, torch.Tensor]]: + """Capture definitive pre-step router bytes for the single-tenant gate.""" + + optimizer_parameter_ids = { + id(parameter) for group in self.optimizer.param_groups for parameter in group.get("params", ()) + } + snapshots: list[tuple[nn.Module, torch.Tensor]] = [] + for module in self.model.modules(): + if type(module).__name__ != "Glm5TopkRouter": + continue + parameter = module.weight + if id(parameter) not in optimizer_parameter_ids: + raise RuntimeError("A trainable GLM-5.2 router is absent from the shared optimizer") + gradient = parameter.grad + if gradient is None: + raise RuntimeError("A trainable GLM-5.2 router has no gradient at the optimizer boundary") + value = parameter.detach() + to_local = getattr(value, "to_local", None) + if callable(to_local): + value = to_local() + snapshots.append((module, value.to(device="cpu", copy=True))) + if not snapshots: + raise RuntimeError("No GLM-5.2 router weights were found at the shared optimizer boundary") + return snapshots + + def _collect_glm52_router_update_metrics( + self, + snapshots: list[tuple[nn.Module, torch.Tensor]], + ) -> Dict[str, Any]: + """Compare post-step router bytes with the pre-step CPU receipts.""" + + local_tensors = len(snapshots) + local_changed_tensors = 0 + local_changed_elements = 0 + metric_device: torch.device | None = None + for module, before in snapshots: + after = module.weight.detach() + to_local = getattr(after, "to_local", None) + if callable(to_local): + after = to_local() + metric_device = after.device + after_cpu = after.to(device="cpu") + changed = torch.ne(after_cpu, before) + changed_elements = int(torch.count_nonzero(changed).item()) + local_changed_elements += changed_elements + local_changed_tensors += int(changed_elements > 0) + + device = metric_device or torch.device(get_device_type()) + aggregate = torch.tensor( + [local_tensors, local_changed_tensors, local_changed_elements], + dtype=torch.int64, + device=device, + ) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(aggregate, op=dist.ReduceOp.SUM) + tensors, changed_tensors, changed_elements = aggregate.cpu().tolist() + metrics = { + "router_update_tensor_count": int(tensors), + "router_update_changed_tensor_count": int(changed_tensors), + "router_update_changed_element_count": int(changed_elements), + } + logger.info("GLM-5.2 router optimizer movement evidence: %s", metrics) + return metrics + def _reshard_exact_forward_only_lm_head(self) -> None: """Restore FP32 masters after the external no-grad vocab projection.""" @@ -8708,6 +9203,7 @@ def optim_step( adapter_mutated = False glm52_fullparam_mutation_started = False glm52_fullparam_publish_receipt = None + glm52_router_update_metrics = None capture_config = dict(sparse_delta_capture or {}) capture_snapshots: dict[str, torch.Tensor] | None = None capture_snapshot_s = 0.0 @@ -8815,6 +9311,11 @@ def optim_step( pp_enabled=self.pp_enabled, pp_group=ps.pp_group if self.pp_enabled else None, ) + router_step_snapshots = ( + self._snapshot_glm52_router_weights_for_step() + if getattr(self, "_glm52_exact_router_single_tenant", False) + else None + ) # Once an optimizer kernel has been launched, any exception may # follow a partial mutation. The exact full-parameter lane must @@ -8827,6 +9328,8 @@ def optim_step( # Python reaches zero_grad/empty_cache. Synchronize before releasing # grad storage to avoid allocator reuse while those kernels are live. synchronize() + if router_step_snapshots is not None: + glm52_router_update_metrics = self._collect_glm52_router_update_metrics(router_step_snapshots) for part in self.model_parts if self.pp_enabled else [self.model]: invalidate_lora_merged_weight_caches(part) try: @@ -8889,6 +9392,8 @@ def optim_step( } if glm52_fullparam_publish_receipt is not None: result["glm52_fullparam_publish"] = glm52_fullparam_publish_receipt + if glm52_router_update_metrics is not None: + result.update(glm52_router_update_metrics) if capture_snapshots is not None: capture_result = write_sparse_source_delta_rank( model=self.model, diff --git a/src/xorl/server/runner/utils/batch_utils.py b/src/xorl/server/runner/utils/batch_utils.py index 05dd0232..fd3cf0ca 100644 --- a/src/xorl/server/runner/utils/batch_utils.py +++ b/src/xorl/server/runner/utils/batch_utils.py @@ -28,6 +28,7 @@ "packed_row_source_num_samples", "packed_row_source_request_ids", "packed_row_source_token_spans", + "sampler_prefill_lengths", "_r3_sample_lengths", "_shifted", "cu_seq_lens_q", @@ -82,11 +83,17 @@ def can_batch_packed_rows(rows: list[Dict[str, Any]]) -> tuple[bool, set[str]]: first_keys = packed_row_sequence_keys(rows[0]) if first_keys is None: return False, set() + first_has_sampler_boundaries = _complete_sampler_prefill_lengths(rows[0]) is not None scalar_keys = set(rows[0]) - first_keys - PACKED_ROW_BATCH_METADATA_KEYS for row in rows[1:]: row_keys = packed_row_sequence_keys(row) if row_keys != first_keys: return False, set() + if (_complete_sampler_prefill_lengths(row) is not None) != first_has_sampler_boundaries: + # Do not recombine an HF-format row with an already-shifted RL row. + # Exact DSV4 relies on complete boundaries to replay the sampler's + # per-request prefill/decode MHC arithmetic. + return False, set() if set(row) - row_keys - PACKED_ROW_BATCH_METADATA_KEYS != scalar_keys: return False, set() for key in scalar_keys: @@ -114,6 +121,23 @@ def _coerce_source_list(value: Any) -> list[Any] | None: return value +def _complete_sampler_prefill_lengths(row: Dict[str, Any]) -> list[int] | None: + """Return row boundaries only when every packed request has one.""" + + sample_lengths = _coerce_source_list(row.get("_r3_sample_lengths")) + prefill_lengths = _coerce_source_list(row.get("sampler_prefill_lengths")) + if not sample_lengths or not prefill_lengths or len(prefill_lengths) != len(sample_lengths): + return None + normalized_samples = [_coerce_int(length, 0) for length in sample_lengths] + normalized_prefills = [_coerce_int(length, 0) for length in prefill_lengths] + if any( + sample_length <= 0 or prefill_length <= 0 or prefill_length > sample_length + for sample_length, prefill_length in zip(normalized_samples, normalized_prefills, strict=True) + ): + return None + return normalized_prefills + + def _row_token_len(row: Dict[str, Any]) -> int: input_rows = row.get("input_ids") if isinstance(input_rows, torch.Tensor): @@ -206,6 +230,9 @@ def merge_packed_row_group(rows: list[Dict[str, Any]], batch_id: int, sequence_k "packed_row_source_group_size": len(source_batch_ids), "_r3_sample_lengths": [length for row in rows for length in row.get("_r3_sample_lengths", [])], } + boundary_rows = [_complete_sampler_prefill_lengths(row) for row in rows] + if all(boundaries is not None for boundaries in boundary_rows): + merged["sampler_prefill_lengths"] = [length for boundaries in boundary_rows for length in boundaries] if "_shifted" in rows[0]: merged["_shifted"] = all(bool(row.get("_shifted", False)) for row in rows) diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index 544b114d..faf58772 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -127,6 +127,13 @@ class ServerArguments: default=False, metadata={"help": "Enable async combine for DeepEP (overlap combine with next layer's compute)."} ) + deepep_native_exact: bool = field( + default=False, + metadata={ + "help": "Use the versioned real-dispatch DeepEP BF16 transport plus deterministic hierarchical fold." + }, + ) + alltoall_combine_hidden_chunk_size: int = field( default=0, metadata={ @@ -1045,6 +1052,13 @@ class ServerArguments: # ======================================================================== enable_lora: bool = field(default=False, metadata={"help": "Enable LoRA adapters for training"}) + lora_serving_mode: Optional[Literal["merged", "separate"]] = field( + default=None, + metadata={ + "help": "Exact train/serve LoRA contract. 'merged' publishes W+sBA with no " + "sampler adapter; 'separate' publishes A/B factors for active-LoRA serving." + }, + ) lora_rank: int = field(default=32, metadata={"help": "LoRA rank (r parameter)"}) @@ -1224,6 +1238,25 @@ def __post_init__(self): """Validate and set defaults.""" from xorl.fp8_training.config_compat import normalize_fp8_training_config # noqa: PLC0415 + if self.deepep_native_exact and self.expert_parallel_size <= 1: + raise ValueError("deepep_native_exact requires expert_parallel_size > 1; EP1 bypasses DeepEP") + + if self.lora_serving_mode not in {None, "merged", "separate"}: + raise ValueError("lora_serving_mode must be 'merged' or 'separate'") + if self.deepep_native_exact and self.enable_lora and self.lora_serving_mode is None: + raise ValueError("Exact LoRA requires explicit lora_serving_mode='merged' or 'separate'") + if not self.enable_lora and self.lora_serving_mode is not None: + raise ValueError("lora_serving_mode requires enable_lora=True") + + if ( + self.deepep_native_exact + and self.enable_gradient_checkpointing + and self.gradient_checkpointing_method in (None, "recompute_full_layer") + ): + # Native exact dispatch/combine cannot be re-entered under the + # routing-replay stage installed by full-layer checkpointing. + self.gradient_checkpointing_method = "recompute_before_dispatch" + if isinstance(self.max_grad_norm, bool): raise ValueError("max_grad_norm must be a finite number; use a value <= 0 to disable clipping") try: @@ -1350,7 +1383,6 @@ def __post_init__(self): "moe_hybrid_shared_lora": (self.moe_hybrid_shared_lora, True), "moe_implementation": (self.moe_implementation, "triton"), "ep_dispatch": (self.ep_dispatch, "alltoall" if exact_active_lora else "deepep"), - "freeze_router": (self.freeze_router, True), "merge_qkv": (self.merge_qkv, True), "lora_export_format": (self.lora_export_format, "sglang_shared_outer"), } @@ -1367,6 +1399,13 @@ def __post_init__(self): ] if mismatches: raise ValueError("GLM-5.2 block-FP8 QLoRA rejects unsupported configuration: " + ", ".join(mismatches)) + if exact_active_lora: + if self.train_router == self.freeze_router: + raise ValueError( + "GLM-5.2 exact block-FP8 QLoRA requires train_router and freeze_router to be complementary" + ) + elif self.train_router or not self.freeze_router: + raise ValueError("GLM-5.2 non-exact block-FP8 QLoRA requires train_router=False and freeze_router=True") if self.lora_target_modules is not None or self.lora_target_manifest is not None: raise ValueError("GLM-5.2 block-FP8 QLoRA uses its complete deterministic target set") if self.qlora_exclude_modules is not None: @@ -1485,6 +1524,7 @@ def to_config_dict(self) -> Dict[str, Any]: "deepep_buffer_size_gb": self.deepep_buffer_size_gb, "deepep_num_sms": self.deepep_num_sms, "deepep_async_combine": self.deepep_async_combine, + "deepep_native_exact": self.deepep_native_exact, "alltoall_combine_hidden_chunk_size": self.alltoall_combine_hidden_chunk_size, "foundation": self.foundation, "encoders": self.encoders, @@ -1621,6 +1661,7 @@ def to_config_dict(self) -> Dict[str, Any]: }, "lora": { "enable_lora": self.enable_lora, + "lora_serving_mode": self.lora_serving_mode, "lora_rank": self.lora_rank, "max_lora_rank": self.max_lora_rank, "lora_alpha": self.lora_alpha, diff --git a/src/xorl/server/weight_sync/glm52_router_bundle.py b/src/xorl/server/weight_sync/glm52_router_bundle.py new file mode 100644 index 00000000..0de70893 --- /dev/null +++ b/src/xorl/server/weight_sync/glm52_router_bundle.py @@ -0,0 +1,157 @@ +"""Atomic router sidecar for exact GLM-5.2 active-LoRA publication.""" + +from __future__ import annotations + +import hashlib +import json +import re +from pathlib import Path + +import torch +from safetensors.torch import save_file +from torch.distributed._tensor import DTensor + + +GLM52_ROUTER_BUNDLE_SCHEMA = "xorl.glm52_router_bundle.v1" +# Keep non-LoRA state outside the adapter root. SGLang's generic adapter +# loader enumerates root-level ``*.safetensors`` files, so placing routers +# beside ``adapter_model.safetensors`` makes their ``layer.N.weight`` keys look +# like malformed LoRA factors before the dedicated receiver can own them. +GLM52_ROUTER_TENSORS = "xorl_router/xorl_glm52_router.safetensors" +GLM52_ROUTER_MANIFEST = "xorl_glm52_router.json" +_ROUTER_MODULE = re.compile(r"(?:^|\.)layers\.(\d+)\.mlp\.gate$") + + +def gather_glm52_router_weights( + model: object, + *, + destination_rank: int | None = 0, +) -> dict[str, torch.Tensor]: + """Collectively reconstruct every GLM router weight on one destination.""" + + state: dict[str, torch.Tensor] = {} + rank = ( + torch.distributed.get_rank() if torch.distributed.is_available() and torch.distributed.is_initialized() else 0 + ) + for name, module in model.named_modules(): + if module.__class__.__name__ != "Glm5TopkRouter": + continue + match = _ROUTER_MODULE.search(name) + if match is None: + raise RuntimeError(f"Cannot derive GLM-5.2 layer id from router module {name!r}") + tensor = module.weight.detach() + retain_tensor = rank == destination_rank if destination_rank is not None else True + if isinstance(tensor, DTensor): + device_mesh = tensor.device_mesh + retain_tensor = ( + retain_tensor + if destination_rank is not None + else all(device_mesh.get_local_rank(mesh_dim) == 0 for mesh_dim in range(device_mesh.ndim)) + ) + tensor = tensor.full_tensor() + if retain_tensor: + key = f"layer.{int(match.group(1))}.weight" + if key in state: + raise RuntimeError(f"Duplicate GLM-5.2 router sidecar key {key!r}") + state[key] = tensor.to(device="cpu", dtype=torch.bfloat16).contiguous() + return state + + +def _merge_glm52_router_states(states: list[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]: + merged: dict[str, torch.Tensor] = {} + for state in states: + for key, tensor in state.items(): + previous = merged.get(key) + if previous is not None: + if not torch.equal(previous, tensor): + raise RuntimeError(f"Conflicting GLM-5.2 router sidecar values for {key!r}") + continue + merged[key] = tensor + return merged + + +def gather_glm52_router_weights_across_ranks( + model: object, + *, + destination_rank: int = 0, +) -> dict[str, torch.Tensor]: + """Reconstruct stage-local routers, then gather all PP stages to one rank.""" + + distributed = torch.distributed.is_available() and torch.distributed.is_initialized() + if not distributed or torch.distributed.get_world_size() == 1: + return gather_glm52_router_weights(model, destination_rank=destination_rank) + + local_state = gather_glm52_router_weights(model, destination_rank=None) + rank = torch.distributed.get_rank() + gathered_states = [None] * torch.distributed.get_world_size() if rank == destination_rank else None + torch.distributed.gather_object(local_state, gathered_states, dst=destination_rank) + if rank != destination_rank: + return {} + return _merge_glm52_router_states([state for state in gathered_states if state]) + + +def save_glm52_router_bundle( + directory: str | Path, + state: dict[str, torch.Tensor], + *, + weight_step: int, + expected_layer_ids: list[int] | None = None, +) -> dict[str, object]: + """Validate and durably save one complete router sidecar.""" + + if not state: + raise RuntimeError("Refusing to publish an empty GLM-5.2 router sidecar") + for key, tensor in state.items(): + if not re.fullmatch(r"layer\.\d+\.weight", key): + raise ValueError(f"Invalid GLM-5.2 router sidecar key {key!r}") + if tensor.dtype is not torch.bfloat16 or tensor.ndim != 2: + raise ValueError(f"GLM-5.2 router {key!r} must be a BF16 matrix") + layer_ids = sorted(int(key.split(".")[1]) for key in state) + if expected_layer_ids is not None and layer_ids != list(expected_layer_ids): + raise RuntimeError( + f"Incomplete GLM-5.2 router sidecar: actual layer ids={layer_ids}, expected={list(expected_layer_ids)}" + ) + + directory = Path(directory) + directory.mkdir(parents=True, exist_ok=True) + tensor_path = directory / GLM52_ROUTER_TENSORS + tensor_path.parent.mkdir(parents=True, exist_ok=True) + save_file(dict(sorted(state.items())), tensor_path) + digest = hashlib.sha256(tensor_path.read_bytes()).hexdigest() + manifest: dict[str, object] = { + "schema": GLM52_ROUTER_BUNDLE_SCHEMA, + "tensor_file": GLM52_ROUTER_TENSORS, + "sha256": digest, + "router_count": len(state), + "layer_ids": layer_ids, + "weight_step": int(weight_step), + } + (directory / GLM52_ROUTER_MANIFEST).write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n") + return manifest + + +def mark_adapter_config_with_glm52_router_bundle(directory: str | Path, manifest: dict[str, object]) -> None: + """Bind the adapter config to its mandatory router sidecar.""" + + config_path = Path(directory) / "adapter_config.json" + config = json.loads(config_path.read_text()) + config["_xorl_glm52_router_bundle"] = { + "schema": manifest["schema"], + "tensor_file": manifest["tensor_file"], + "sha256": manifest["sha256"], + "router_count": manifest["router_count"], + "layer_ids": manifest["layer_ids"], + "weight_step": manifest["weight_step"], + } + config_path.write_text(json.dumps(config, indent=2, sort_keys=True) + "\n") + + +__all__ = [ + "GLM52_ROUTER_BUNDLE_SCHEMA", + "GLM52_ROUTER_MANIFEST", + "GLM52_ROUTER_TENSORS", + "gather_glm52_router_weights", + "gather_glm52_router_weights_across_ranks", + "mark_adapter_config_with_glm52_router_bundle", + "save_glm52_router_bundle", +] diff --git a/src/xorl/server/weight_sync/handler.py b/src/xorl/server/weight_sync/handler.py index e2733cdb..0bea033c 100644 --- a/src/xorl/server/weight_sync/handler.py +++ b/src/xorl/server/weight_sync/handler.py @@ -467,6 +467,17 @@ def _clear_sync_abort(self, abort_path: str) -> None: logger.debug("Rank %d: [WeightSync] failed to clear abort marker %s: %s", self.rank, abort_path, e) def _prepare_lora_adapter_for_sync(self, model_id: Optional[str]) -> Optional[str]: + lora_config = getattr(self.trainer, "lora_config", {}) or {} + lora_serving_mode = lora_config.get("lora_serving_mode") + if lora_serving_mode == "separate": + raise RuntimeError( + "lora_serving_mode='separate' publishes A/B factors and cannot use " + "merged full-weight synchronization; publish the adapter checkpoint " + "and load it as an active sampler adapter" + ) + if lora_serving_mode not in {None, "merged"}: + raise ValueError(f"Unknown lora_serving_mode {lora_serving_mode!r}") + adapter_manager = getattr(self.trainer, "adapter_manager", None) if adapter_manager is None: return None @@ -1538,17 +1549,15 @@ def _add_rank_phase(name: str, start: float) -> None: if _ws_timings: _t_ep_collect = time.perf_counter() - # EP MoE prefixes to skip in extraction - ep_moe_prefixes = set() - for ctx in ep_moe_contexts: - p = ctx["prefix"] - if mod_name != "(root)": - if p == mod_name: - ep_moe_prefixes.add("") - elif p.startswith(mod_name + "."): - ep_moe_prefixes.add(p[len(mod_name) + 1 :]) - else: - ep_moe_prefixes.add(p) + # EP contexts serve two distinct purposes: they identify + # parameters the ordinary dense extractor must skip, and + # (usually) carry expert tensors for the EP transfer path. + # Separate active-LoRA uses skip-only contexts because its + # frozen expert bases must be omitted from both paths. + ep_moe_prefixes, ep_moe_contexts = self._split_ep_moe_contexts_for_sync( + ep_moe_contexts, + mod_name, + ) if _stage_leader or _extract_dense_on_sender: t_phase = time.perf_counter() @@ -2178,6 +2187,30 @@ def _qlora_collective_ops( # EP MoE data collection (all ranks, during unshard) # ======================================================================== + @staticmethod + def _split_ep_moe_contexts_for_sync( + contexts: List[Dict[str, Any]], + mod_name: str, + ) -> Tuple[set, List[Dict[str, Any]]]: + """Return dense-extractor skip prefixes and transferable EP contexts.""" + + prefixes = set() + transferable = [] + for ctx in contexts: + prefix = ctx["prefix"] + if mod_name != "(root)": + if prefix == mod_name: + prefixes.add("") + elif prefix.startswith(mod_name + "."): + prefixes.add(prefix[len(mod_name) + 1 :]) + else: + prefixes.add(prefix) + + if ctx.get("type") != "frozen_active_lora_base": + transferable.append(ctx) + + return prefixes, transferable + def _collect_ep_moe_data( self, fsdp_mod, @@ -2218,6 +2251,39 @@ def _collect_ep_moe_data( if not isinstance(mod, (MoEExperts, MoEExpertsLoRA)): continue + if mname: + full_prefix = f"{mod_name}.{mname}" if mod_name != "(root)" else mname + else: + full_prefix = mod_name + + if isinstance(mod, MoEExpertsLoRA): + lora_serving_mode = getattr(mod, "lora_serving_mode", None) + if lora_serving_mode == "separate": + # The active-LoRA sampler starts from the same immutable base + # checkpoint as the trainer. Routed-expert base parameters are + # frozen by LoRA training, so publishing them is redundant. It + # is also invalid after SGLang wraps FusedMoE for active LoRA: + # the base parameters then live below ``experts.base_layer`` + # while the generic Qwen online loader targets + # ``experts.w13_weight`` / ``experts.w2_weight``. + # + # Keep a skip-only context so the ordinary dense extractor + # also omits these parameters. This context is removed + # before the EP transfer loop by + # _split_ep_moe_contexts_for_sync. + contexts.append( + { + "type": "frozen_active_lora_base", + "prefix": full_prefix, + "local_experts": None, + } + ) + continue + if lora_serving_mode != "merged": + raise ValueError( + f"Unknown LoRA serving mode {lora_serving_mode!r} while collecting {mod_name}.{mname}" + ) + # Get expert params — after unshard they may be plain tensors or DTensors gate_up = getattr(mod, "gate_up_proj", None) if isinstance(gate_up, torch.nn.Parameter): @@ -2234,11 +2300,6 @@ def _collect_ep_moe_data( gated = getattr(mod, "gated", True) proj_names = ("gate_proj", "up_proj", "down_proj") if gated else ("up_proj", "down_proj") - if mname: - full_prefix = f"{mod_name}.{mname}" if mod_name != "(root)" else mname - else: - full_prefix = mod_name - if not collect_tensors: contexts.append( { diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index cf224676..48c189f0 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -6,7 +6,7 @@ """ from dataclasses import dataclass, field -from typing import Any, Callable, List, Optional, Set +from typing import Any, Callable, List, Literal, Optional, Set import torch import torch.nn as nn @@ -23,7 +23,11 @@ freeze_deepseek_v3_router_parameters, validate_deepseek_v3_training_mode, ) -from xorl.models.transformers.glm5.support import is_glm5_config, validate_glm5_training_mode +from xorl.models.transformers.glm5.support import ( + is_glm5_config, + validate_glm5_training_mode, + validate_glm52_local_router_inventory, +) from xorl.qlora import ( detect_prequantized_block_fp8, detect_prequantized_nvfp4, @@ -205,11 +209,13 @@ def build_training_model( deepep_buffer_size_gb: float = 2.0, deepep_num_sms: int = 20, deepep_async_combine: bool = False, + deepep_native_exact: bool = False, alltoall_combine_hidden_chunk_size: int = 0, init_device: str = "meta", merge_qkv: bool = True, # --- LoRA --- enable_lora: bool = False, + lora_serving_mode: Optional[Literal["merged", "separate"]] = None, lora_rank: int = 32, lora_alpha: int = 16, lora_b_init_std: float = 0.0, @@ -318,6 +324,21 @@ def build_training_model( Returns a :class:`TrainingModelResult` with model, config, PP state, etc. """ + if deepep_native_exact and expert_parallel_size <= 1: + raise ValueError( + "deepep_native_exact requires expert_parallel_size > 1; EP1 bypasses " + "DeepEP and cannot satisfy the native exact training contract" + ) + + if ( + deepep_native_exact + and enable_gradient_checkpointing + and gradient_checkpointing_method in (None, "recompute_full_layer") + ): + # Keep the contract safe for direct API callers that do not construct + # a ServerArguments object first. + gradient_checkpointing_method = "recompute_before_dispatch" + if server_training and pp_virtual_stages > 1: raise NotImplementedError( "Server training does not yet support pipeline_parallel_virtual_stages > 1: " @@ -391,6 +412,7 @@ def build_training_model( deepep_buffer_size_gb=deepep_buffer_size_gb, deepep_num_sms=deepep_num_sms, deepep_async_combine=deepep_async_combine, + deepep_native_exact=deepep_native_exact, alltoall_combine_hidden_chunk_size=alltoall_combine_hidden_chunk_size, router_fp32=router_fp32, lm_head_fp32=lm_head_fp32, @@ -405,6 +427,7 @@ def build_training_model( flash_attention_deterministic=flash_attention_deterministic, server_training=server_training, enable_lora=enable_lora, + lora_serving_mode=lora_serving_mode, block_fp8_qlora_training=block_fp8_qlora_training, glm52_fullparam_fp8_training=glm52_fullparam_fp8_training, lora_rank=lora_rank, @@ -636,6 +659,18 @@ def build_training_model( "Exact dense Qwen3 trainer path reinstalled after projection replacement" + (f": {pattern}" if pattern else " (existing wrappers retained)") ) + elif getattr(model.config, "_deepep_native_exact", False): + apply_native_exact = getattr(model, "_apply_deepep_native_exact", None) + if callable(apply_native_exact): + wrapped = apply_native_exact() + else: + wrapped = {} + pattern = ", ".join(f"{name}x{count}" for name, count in sorted(wrapped.items())) + if wrapped: + logger.info_rank0( + "Model-declared native DeepEP exact trainer program: " + f"wrapped {sum(wrapped.values())} trunk linears" + (f"; {pattern}" if pattern else "") + ) apply_exact_qwen = getattr(model, "_apply_qwen35_gdn_exact", None) if server_training and callable(apply_exact_qwen): @@ -694,6 +729,8 @@ def build_training_model( moe_grad_reduce_mode=moe_grad_reduce_mode, fsdp_sharded_lm_head_loss=fsdp_sharded_lm_head_loss, fsdp_reduce_dtype=fsdp_reduce_dtype, + enable_lora=enable_lora, + enable_qlora=enable_qlora, skip_param_upcast=should_skip_generic_param_upcast( enable_lora=enable_lora, enable_qlora=enable_qlora, @@ -730,10 +767,26 @@ def build_training_model( # ------------------------------------------------------------------ if enable_qlora: # After QLoRA quantization, freeze everything except LoRA + keep_exact_glm52_routers = bool( + getattr(model.config, "train_router", False) and glm52_exact_active_lora_enabled(model.config) + ) + retained_router_count = 0 for part in all_parts: for name, param in part.named_parameters(): - if "lora_A" not in name and "lora_B" not in name: + is_exact_glm52_router = keep_exact_glm52_routers and name.endswith("mlp.gate.weight") + if is_exact_glm52_router: + param.requires_grad_(True) + retained_router_count += 1 + elif "lora_A" not in name and "lora_B" not in name: param.requires_grad = False + if keep_exact_glm52_routers: + validate_glm52_local_router_inventory( + all_parts, + retained_router_count=retained_router_count, + ) + logger.info_rank0( + f"Retained {retained_router_count} post-FSDP BF16 router weights for exact GLM-5.2 QLoRA training" + ) helper.print_device_mem_info("VRAM usage after QLoRA quantization") elif enable_lora: for part in all_parts: diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index a07879e9..8479af4e 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -49,7 +49,7 @@ save_model_weights, ) from xorl.models.checkpoint_handlers.buffers import get_prequantized_exclude_modules -from xorl.models.exact_contract import exact_gdn_cp_alignment_required +from xorl.models.exact_contract import exact_gdn_cp_alignment_required, glm52_exact_active_lora_enabled from xorl.models.layers.moe.aux_loss import LoadBalancingBuffer, global_load_balancing_loss_func from xorl.models.layers.moe.routing_replay import RoutingReplay, set_replay_stage from xorl.models.module_utils import compute_loss @@ -58,7 +58,10 @@ validate_deepseek_v3_training_mode, ) from xorl.models.transformers.glm5.index_share import IndexShareMode -from xorl.models.transformers.glm5.support import validate_glm5_training_mode +from xorl.models.transformers.glm5.support import ( + validate_glm5_training_mode, + validate_glm52_local_router_inventory, +) from xorl.optim import build_lr_scheduler, build_optimizer from xorl.qlora import ( detect_prequantized_block_fp8, @@ -706,6 +709,7 @@ def _build_model(self) -> None: deepep_buffer_size_gb=args.model.deepep_buffer_size_gb, deepep_num_sms=args.model.deepep_num_sms, deepep_async_combine=args.model.deepep_async_combine, + deepep_native_exact=getattr(args.model, "deepep_native_exact", False), router_fp32=args.model.router_fp32, lm_head_fp32=args.model.lm_head_fp32, alltoall_combine_hidden_chunk_size=args.model.alltoall_combine_hidden_chunk_size, @@ -721,6 +725,7 @@ def _build_model(self) -> None: block_fp8_qlora_training=getattr(args.lora, "block_fp8_qlora_training", False), lora_rank=args.lora.lora_rank, lora_alpha=args.lora.lora_alpha, + lora_serving_mode=args.lora.lora_serving_mode, init_device=args.train.init_device, pipeline_parallel_virtual_stages=args.train.pipeline_parallel_virtual_stages, pipeline_parallel_input_weight=args.train.pipeline_parallel_input_weight, @@ -1040,6 +1045,30 @@ def _parallelize(self) -> None: if "lora_A" not in name and "lora_B" not in name: param.requires_grad = False + # The complete GLM-5.2 exact QLoRA lane trains the checkpoint-native + # BF16 sparse routers in addition to the adapter factors. Deferred + # QLoRA freezing runs after FSDP and historically froze these weights + # again, leaving the optimizer inventory looking plausible while the + # grad-enabled router GEMM received a non-differentiable tensor. + if ( + args.lora.enable_qlora + and bool(getattr(self.model.config, "train_router", False)) + and glm52_exact_active_lora_enabled(self.model.config) + ): + retained_router_count = 0 + for part in self._all_model_parts(): + for name, param in part.named_parameters(): + if name.endswith("mlp.gate.weight"): + param.requires_grad_(True) + retained_router_count += 1 + validate_glm52_local_router_inventory( + self._all_model_parts(), + retained_router_count=retained_router_count, + ) + logger.info_rank0( + f"Retained {retained_router_count} post-FSDP BF16 router weights for exact GLM-5.2 QLoRA training" + ) + if args.model.freeze_router: frozen = 0 for part in self._all_model_parts(): diff --git a/src/xorl/trainers/training_utils.py b/src/xorl/trainers/training_utils.py index 6c17f5dd..78b35314 100644 --- a/src/xorl/trainers/training_utils.py +++ b/src/xorl/trainers/training_utils.py @@ -610,10 +610,26 @@ def pad_micro_batches_for_pp( if pad_to_multiple_of > 1 and target_sharded % pad_to_multiple_of != 0: target_sharded = ((target_sharded + pad_to_multiple_of - 1) // pad_to_multiple_of) * pad_to_multiple_of + # Every field here is storage-row aligned with input_ids after the + # sequence-shard collator. PP padding must extend the whole objective, + # not only the model inputs: otherwise the terminal hidden state and the + # loss tensors describe different programs. _PAD_VALUES = { - "input_ids": 0, "labels": IGNORE_INDEX, - "attention_mask": 0, + "target_tokens": IGNORE_INDEX, + "logprobs": 0.0, + "old_logprobs": 0.0, + "ref_logprobs": 0.0, + "rollout_logprobs": 0.0, + "advantages": 0.0, + "weights": 0.0, + "teacher_ids": 0, + "teacher_cache_indices": 0, + "teacher_cache_local_indices": 0, + "teacher_weights": 0.0, + "hidden_match_weights": 0.0, + "opd_region_ids": 0, + "opd_sample_ok": 0, # Exact sampling-transform metadata must pad with the mathematical # identity so PP's fixed communication shape cannot change scoring. "logprob_temperatures": 1.0, @@ -634,11 +650,35 @@ def pad_micro_batches_for_pp( for key, pad_value in _PAD_VALUES.items(): if key in mb and isinstance(mb[key], torch.Tensor): + if mb[key].shape[-1] != ids_len: + raise ValueError( + f"PP token-aligned field {key!r} has {mb[key].shape[-1]} rows, but input_ids has {ids_len}" + ) mb[key] = F.pad(mb[key], (0, pad_tokens), value=pad_value) - if "position_ids" in mb and isinstance(mb["position_ids"], torch.Tensor): - scale = mb["position_ids"].shape[-1] // ids_len if ids_len > 0 else 1 - mb["position_ids"] = F.pad(mb["position_ids"], (0, pad_tokens * scale), value=0) + teacher_hidden = mb.get("teacher_hidden_states") + if isinstance(teacher_hidden, torch.Tensor): + if teacher_hidden.ndim != 3 or teacher_hidden.shape[1] != ids_len: + raise ValueError( + "PP token-aligned field 'teacher_hidden_states' must have shape " + f"[batch, {ids_len}, hidden], got {tuple(teacher_hidden.shape)}" + ) + mb["teacher_hidden_states"] = F.pad(teacher_hidden, (0, 0, 0, pad_tokens), value=0.0) + + # The sequence-shard collator retains these in the full CP domain, + # unlike token objectives and input_ids, which are CP-local. + for key in ("position_ids", "attention_mask"): + value = mb.get(key) + if isinstance(value, torch.Tensor): + if ids_len <= 0 or value.shape[-1] % ids_len != 0: + raise ValueError( + f"PP full-domain field {key!r} has {value.shape[-1]} rows, " + f"which is not an integer multiple of input_ids rows {ids_len}" + ) + scale = value.shape[-1] // ids_len + mb[key] = F.pad(value, (0, pad_tokens * scale), value=0) + + mb["input_ids"] = F.pad(mb["input_ids"], (0, pad_tokens), value=0) for key in ("cu_seq_lens_q", "cu_seq_lens_k"): if key in mb and isinstance(mb[key], torch.Tensor): diff --git a/submodules/xorl-sglang b/submodules/xorl-sglang index 4e8684b4..2a878e34 160000 --- a/submodules/xorl-sglang +++ b/submodules/xorl-sglang @@ -1 +1 @@ -Subproject commit 4e8684b482c29caa7fee86ed363431e2313ace48 +Subproject commit 2a878e34ea92fbc468967bec6a97af67c120de04 diff --git a/tests/distributed/test_deepep_async_combine_guard.py b/tests/distributed/test_deepep_async_combine_guard.py index a11a2d5e..e6ab5f48 100644 --- a/tests/distributed/test_deepep_async_combine_guard.py +++ b/tests/distributed/test_deepep_async_combine_guard.py @@ -53,3 +53,154 @@ def fake_apply(expert_output, buffer, ctx, async_combine): assert result is expert_output assert captured["async_combine"] is True + + +def test_native_dispatch_exposes_transported_receive_metadata(monkeypatch): + recv_x = torch.ones(2, 4) + recv_ids = torch.tensor([[0, -1], [1, 0]], dtype=torch.int32) + recv_weights = torch.tensor([[0.5, 0.0], [0.75, 0.25]], dtype=torch.float32) + ctx = SimpleNamespace(recv_topk_idx=recv_ids, recv_topk_weights=recv_weights) + monkeypatch.setattr( + deepep, + "token_pre_dispatch_no_permute", + lambda **_kwargs: (recv_x, torch.tensor([1, 3]), ctx), + ) + + got_x, got_ids, got_weights, got_ctx = deepep.token_pre_dispatch_native( + buffer=None, + hidden_states=torch.zeros(1, 4), + routing_weights=torch.ones(1, 2), + selected_experts=torch.zeros(1, 2, dtype=torch.int32), + num_experts=4, + ) + + assert got_x is recv_x + assert got_ids is recv_ids + assert got_weights is recv_weights + assert got_ctx is ctx + + +def test_native_combine_keeps_receive_layout_and_forces_safe_sync(monkeypatch): + captured = {} + + def fake_apply(recv_output, buffer, ctx, async_combine): + captured.update(buffer=buffer, ctx=ctx, async_combine=async_combine) + return recv_output + + monkeypatch.delenv("XORL_DEEPEP_UNSAFE_ASYNC_COMBINE", raising=False) + monkeypatch.setattr(deepep._FusedNativeReceiveCombine, "apply", staticmethod(fake_apply)) + recv_output = torch.ones(2, 4) + ctx = SimpleNamespace() + result = deepep.tokens_post_combine_native( + buffer="buffer", + recv_output=recv_output, + ctx=ctx, + async_combine=True, + ) + + assert result is recv_output + assert captured == {"buffer": "buffer", "ctx": ctx, "async_combine": False} + + +def test_native_dispatch_backward_can_require_device_completion(monkeypatch): + events = [] + + class FakeEvent: + def current_stream_wait(self): + events.append("event_wait") + + class FakeGrad: + dtype = torch.bfloat16 + device = torch.device("cuda:0") + + def record_stream(self, _stream): + events.append("record_stream") + + def to(self, _dtype): + raise AssertionError("the matching BF16 gradient must not be cast") + + class FakeDeepEP: + def combine(self, **_kwargs): + events.append("combine") + return FakeGrad(), None, FakeEvent() + + class FakeStream: + def synchronize(self): + events.append("device_complete") + + monkeypatch.setattr(deepep, "EventHandle", lambda: object()) + monkeypatch.setattr(deepep, "EventOverlap", lambda _handle: object()) + monkeypatch.setattr(deepep.torch.cuda, "current_stream", lambda _device=None: FakeStream()) + ctx = SimpleNamespace( + buffer=SimpleNamespace(buffer=FakeDeepEP(), combine_config=object()), + handle=object(), + input_dtype=torch.bfloat16, + call_id=17, + complete_backward_device_boundary=True, + backward_trace_label="glm52_layer_7", + backward_layer_dependency_meta=((1,), torch.float32, torch.device("cpu")), + backward_shared_dependency_meta=((2, 2), torch.bfloat16, torch.device("cpu")), + ) + + result = deepep._FusedDispatchNoPermute.backward( + ctx, + torch.ones(2, 4, dtype=torch.bfloat16), + None, + None, + ) + + assert events == ["combine", "event_wait", "record_stream", "device_complete"] + assert isinstance(result[0], FakeGrad) + assert result[1:7] == (None,) * 6 + assert torch.equal(result[7], torch.zeros(1, dtype=torch.float32)) + assert torch.equal(result[8], torch.zeros((2, 2), dtype=torch.bfloat16)) + + +def test_terminal_dispatch_dependency_holds_shared_and_residual_backward(): + events = [] + + class SharedBoundary(torch.autograd.Function): + @staticmethod + def forward(ctx, value): + return value.clone() + + @staticmethod + def backward(ctx, grad_output): + events.append("shared") + return grad_output + + class ResidualBoundary(torch.autograd.Function): + @staticmethod + def forward(ctx, value): + return value.clone() + + @staticmethod + def backward(ctx, grad_output): + events.append("residual") + return grad_output + + class TerminalDispatchBoundary(torch.autograd.Function): + @staticmethod + def forward(ctx, value, layer_dependency, shared_dependency): + ctx.layer_shape = layer_dependency.shape + ctx.shared_shape = shared_dependency.shape + return value.clone() + + @staticmethod + def backward(ctx, grad_output): + events.append("routed_terminal") + return ( + grad_output, + torch.zeros(ctx.layer_shape), + torch.zeros(ctx.shared_shape), + ) + + value = torch.ones((2, 2), requires_grad=True) + shared = SharedBoundary.apply(value) + residual = ResidualBoundary.apply(value) + routed = TerminalDispatchBoundary.apply(value, residual, shared) + (routed.sum() + shared.sum() + residual.sum()).backward() + + assert events[0] == "routed_terminal" + assert set(events[1:]) == {"shared", "residual"} + assert torch.equal(value.grad, torch.full_like(value, 3.0)) diff --git a/tests/distributed/test_deepep_canonical_combine.py b/tests/distributed/test_deepep_canonical_combine.py index ae2ba0d7..e4cc992a 100644 --- a/tests/distributed/test_deepep_canonical_combine.py +++ b/tests/distributed/test_deepep_canonical_combine.py @@ -543,11 +543,17 @@ def test_canonical_combine_over_real_deepep_normal(): def test_canonical_combine_over_real_deepep_q397b_geometry(): - """Exercise 512 experts and top-10 routing (K < EP at EP32).""" + """Exercise 512 experts and top-10 routing (K < EP at EP32). + + DeepEP's normal dispatch supports at most 128 local experts, so the + 512-expert geometry cannot be projected below EP4. Keep the ordinary + 256-expert gate above runnable at EP2, but require a valid topology for + this model-specific gate instead of failing inside the CUDA kernel. + """ pytest.importorskip("deep_ep") pytest.importorskip("nvidia.nvshmem") _launch( - _required_gpus(), + max(_required_gpus(), 4), {"XORL_TEST_DEEPEP_CC_EXPERTS": "512", "XORL_TEST_DEEPEP_CC_TOPK": "10"}, ) diff --git a/tests/distributed/test_deepep_native_exact.py b/tests/distributed/test_deepep_native_exact.py new file mode 100644 index 00000000..c16d4587 --- /dev/null +++ b/tests/distributed/test_deepep_native_exact.py @@ -0,0 +1,282 @@ +from types import SimpleNamespace + +import pytest +import torch + +from xorl.distributed.moe import deepep_native_exact as native_exact_module +from xorl.distributed.moe.deepep_native_exact import ( + DeepEPNativeExactError, + NativeDeepEPGeometry, + _flatten_native_route_metadata, + adapt_native_runner_metadata, + canonicalize_native_routing_metadata, + reduce_expert_rows_to_bf16_leaf, + validate_native_receive_metadata, +) + + +def _dispatch_ctx(*, rows=3, hidden=2, ids=None, weights=None, indices=None): + if ids is None: + ids = torch.tensor([[0, -1], [1, -1], [0, 1]], dtype=torch.int64)[:rows] + if weights is None: + weights = torch.ones_like(ids, dtype=torch.float32) + if indices is None: + indices = torch.tensor([0, 1, 2, 2], dtype=torch.long) + return SimpleNamespace( + num_recv_tokens=rows, + hidden_dim=hidden, + recv_topk_idx=ids, + recv_topk_weights=weights, + permuted_indices=indices, + ) + + +def test_expert_rows_reduce_in_fp32_then_store_one_bf16_leaf(): + # 256.0 + 1.0 - 256.0 distinguishes FP32 local accumulation from a + # left-associated BF16 fold, which loses the unit contribution. + expert_output = torch.tensor( + [[256.0, 1.0], [1.0, 2.0], [-256.0, 4.0]], + dtype=torch.bfloat16, + requires_grad=True, + ) + ctx = _dispatch_ctx( + rows=2, + hidden=2, + ids=torch.tensor([[0, 1], [1, -1]], dtype=torch.int64), + indices=torch.tensor([0, 0, 0], dtype=torch.long), + ) + + leaf = reduce_expert_rows_to_bf16_leaf(expert_output, ctx) + + assert leaf.dtype is torch.bfloat16 + assert torch.equal(leaf, torch.tensor([[1.0, 7.0], [0.0, 0.0]], dtype=torch.bfloat16)) + leaf.float().sum().backward() + assert torch.equal(expert_output.grad, torch.ones_like(expert_output)) + + +def test_native_receive_metadata_accepts_empty_receive_batch(): + output = torch.empty((0, 4), dtype=torch.bfloat16).contiguous() + ctx = _dispatch_ctx( + rows=0, + hidden=4, + ids=torch.empty((0, 2), dtype=torch.int64), + weights=torch.empty((0, 2), dtype=torch.float32), + indices=torch.empty(0, dtype=torch.long), + ) + + validate_native_receive_metadata(output, ctx, num_local_experts=2) + + +def test_native_route_metadata_preserves_topk_on_empty_rank(): + routing = torch.empty((0, 8), dtype=torch.float32) + experts = torch.empty((0, 8), dtype=torch.int64) + + routing_flat, experts_flat = _flatten_native_route_metadata( + routing, + experts, + row_count=0, + ) + + assert routing_flat.shape == (0, 8) + assert experts_flat.shape == (0, 8) + assert routing_flat.is_contiguous() + assert experts_flat.is_contiguous() + + +def test_native_runner_metadata_is_int32_ids_and_fp32_weights(): + ids = torch.tensor([[0, -1], [1, 0]], dtype=torch.int64) + weights = torch.tensor([[0.75, 0.0], [0.5, 0.25]], dtype=torch.float32) + + runner_ids, runner_weights = adapt_native_runner_metadata(ids, weights) + + assert runner_ids.dtype is torch.int32 + assert torch.equal(runner_ids.to(torch.int64), ids) + assert runner_weights.dtype is torch.float32 + assert torch.equal(runner_weights, weights) + + +def test_native_routing_metadata_preserves_sampler_fp32_coefficients(): + weights = torch.tensor( + [[0.31519484519958496, 0.09776496887207031]], + dtype=torch.float32, + ) + + canonical = canonicalize_native_routing_metadata(weights) + + assert canonical.dtype is torch.float32 + assert canonical.is_contiguous() + assert torch.equal(canonical, weights) + assert not torch.equal(canonical, weights.to(torch.bfloat16).to(torch.float32)) + + +@pytest.mark.parametrize( + ("ids", "weights", "match"), + [ + (torch.tensor([[-1, -1]]), torch.ones((1, 2)), "no local route"), + (torch.tensor([[2, -1]]), torch.ones((1, 2)), "outside this rank"), + (torch.tensor([[-2, -1]]), torch.ones((1, 2)), "below -1"), + (torch.tensor([[0, -1]]), torch.tensor([[float("nan"), 0.0]]), "not finite"), + ], +) +def test_native_receive_metadata_fails_closed(ids, weights, match): + output = torch.zeros((1, 4), dtype=torch.bfloat16).contiguous() + ctx = _dispatch_ctx(rows=1, hidden=4, ids=ids, weights=weights) + + with pytest.raises(DeepEPNativeExactError, match=match): + validate_native_receive_metadata(output, ctx, num_local_experts=2) + + +def test_native_receive_metadata_rejects_wider_wire_value(): + output = torch.zeros((1, 4), dtype=torch.float32).contiguous() + ctx = _dispatch_ctx( + rows=1, + hidden=4, + ids=torch.tensor([[0, -1]], dtype=torch.int64), + weights=torch.ones((1, 2), dtype=torch.float32), + ) + + with pytest.raises(DeepEPNativeExactError, match="must be BF16"): + validate_native_receive_metadata(output, ctx, num_local_experts=2) + + +def test_normal_default_is_deterministic_and_admits_ep16_one_call(monkeypatch): + monkeypatch.setattr( + native_exact_module, + "resolve_native_deepep_geometry", + lambda _group, hidden: NativeDeepEPGeometry(ep_size=16, ep_rank=0, hidden_size=hidden), + ) + monkeypatch.setattr( + native_exact_module, + "validate_native_receive_metadata", + lambda *_args, **_kwargs: None, + ) + + calls = [] + + def fake_apply( + recv_output, + _buffer, + _dispatch_ctx, + geometry, + _backward_layer_dependency, + _backward_trace_label, + ): + calls.append(geometry) + return recv_output.clone() + + monkeypatch.setattr( + native_exact_module._DeepEPDeterministicCombineBF16, + "apply", + staticmethod(fake_apply), + ) + recv_output = torch.zeros((1, 4), dtype=torch.bfloat16).contiguous() + + combined = native_exact_module.native_receive_combine_and_fold( + recv_output, + buffer=object(), + dispatch_ctx=object(), + ep_group=object(), + num_local_experts=1, + ) + + assert torch.equal(combined, recv_output) + assert len(calls) == 1 + assert calls[0].ep_size == 16 + + +def test_expert_order_adapter_uses_only_deterministic_receive(monkeypatch): + leaf = torch.zeros((1, 4), dtype=torch.bfloat16) + calls = [] + monkeypatch.setattr( + native_exact_module, + "reduce_expert_rows_to_bf16_leaf", + lambda _output, _dispatch_ctx: leaf, + ) + + def fake_receive(recv_output, **kwargs): + calls.append(kwargs) + return recv_output + + monkeypatch.setattr( + native_exact_module, + "native_receive_combine_and_fold", + fake_receive, + ) + common = dict( + buffer=object(), + dispatch_ctx=object(), + ep_group=object(), + num_local_experts=1, + ) + + assert native_exact_module.native_expert_combine_and_fold(leaf, **common) is leaf + assert len(calls) == 1 + + +def test_normal_deterministic_rejects_unsupported_ep_before_kernel(monkeypatch): + monkeypatch.setattr( + native_exact_module, + "resolve_native_deepep_geometry", + lambda _group, hidden: NativeDeepEPGeometry(ep_size=24, ep_rank=0, hidden_size=hidden), + ) + monkeypatch.setattr( + native_exact_module, + "validate_native_receive_metadata", + lambda *_args, **_kwargs: None, + ) + called = False + + def unexpected_apply(*_args, **_kwargs): + nonlocal called + called = True + raise AssertionError("deterministic kernel must not run") + + monkeypatch.setattr( + native_exact_module._DeepEPDeterministicCombineBF16, + "apply", + staticmethod(unexpected_apply), + ) + + with pytest.raises(DeepEPNativeExactError, match=r"EP sizes.*EP24"): + native_exact_module.native_receive_combine_and_fold( + torch.zeros((1, 4), dtype=torch.bfloat16).contiguous(), + buffer=object(), + dispatch_ctx=object(), + ep_group=object(), + num_local_experts=1, + ) + + assert not called + + +def test_normal_deterministic_rejects_unsupported_ep_before_dispatch(monkeypatch): + from xorl.distributed.moe import deepep as deepep_module + + monkeypatch.setattr( + native_exact_module, + "resolve_native_deepep_geometry", + lambda _group, hidden: NativeDeepEPGeometry(ep_size=24, ep_rank=0, hidden_size=hidden), + ) + called = False + + def unexpected_buffer(**_kwargs): + nonlocal called + called = True + raise AssertionError("DeepEP buffer must not be acquired") + + monkeypatch.setattr(deepep_module, "get_default_buffer", unexpected_buffer) + + with pytest.raises(DeepEPNativeExactError, match=r"EP sizes.*EP24"): + native_exact_module.native_dispatch_runner_combine( + torch.zeros((1, 4), dtype=torch.bfloat16), + torch.ones((1, 1), dtype=torch.float32), + torch.zeros((1, 1), dtype=torch.int64), + ep_group=object(), + num_experts=24, + num_local_experts=1, + buffer_size_gb=1.0, + num_sms=1, + runner=lambda hidden, _weights, _ids: hidden, + ) + + assert not called diff --git a/tests/distributed/test_deepep_native_exact_real.py b/tests/distributed/test_deepep_native_exact_real.py new file mode 100644 index 00000000..648b3104 --- /dev/null +++ b/tests/distributed/test_deepep_native_exact_real.py @@ -0,0 +1,194 @@ +"""Real-GPU component gate for the original-handle native exact transport.""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist + + +_WORKER_ENV = "XORL_TEST_DEEPEP_NATIVE_EXACT_WORKER" +_WORLD_ENV = "XORL_TEST_DEEPEP_NATIVE_EXACT_WORLD" + + +def _hierarchical_fold(leaves: torch.Tensor) -> torch.Tensor: + """Independent Tree8/BF16-node/ascending-FP64-node reference.""" + node_leaves = [] + for begin in range(0, leaves.shape[0], 8): + node = leaves[begin : begin + 8] + if node.shape[0] < 8: + node = torch.cat( + ( + node, + torch.zeros( + (8 - node.shape[0], *node.shape[1:]), + dtype=node.dtype, + device=node.device, + ), + ), + dim=0, + ) + p01 = node[0].double() + node[1].double() + p23 = node[2].double() + node[3].double() + p45 = node[4].double() + node[5].double() + p67 = node[6].double() + node[7].double() + node_leaves.append(((p01 + p23) + (p45 + p67)).to(torch.bfloat16)) + value = node_leaves[0].double() + for node_leaf in node_leaves[1:]: + value = value + node_leaf.double() + return value.to(torch.bfloat16) + + +def _worker_main() -> int: + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl") + rank = dist.get_rank() + world = dist.get_world_size() + device = torch.device("cuda", local_rank) + assert world in (2, 4, 8, 16) + + from xorl.distributed.moe import deepep as deepep_module + from xorl.distributed.moe.deepep import ( + get_default_buffer, + token_pre_dispatch_native, + ) + from xorl.distributed.moe.deepep_native_exact import ( + NativeDeepEPGeometry, + native_receive_combine_and_fold, + ) + + rows, hidden = 3, 2048 + generator = torch.Generator(device="cpu").manual_seed(271828 + rank) + hidden_states = torch.randn((rows, hidden), generator=generator).to(torch.bfloat16).to(device).requires_grad_(True) + # Every source token has exactly one route to every physical expert rank. + # DeepEP therefore performs its real top-k dispatch, while the expected + # return leaf from rank r is unambiguously x + (r+1). + selected = torch.arange(world, dtype=torch.int64, device=device).expand(rows, world).contiguous() + weights = torch.ones((rows, world), dtype=torch.float32, device=device) + + geometry = NativeDeepEPGeometry(ep_size=world, ep_rank=rank, hidden_size=hidden) + buffer = get_default_buffer(ep_group=dist.group.WORLD, buffer_size_gb=2.0, num_sms=20) + buffer.init_buffer(hidden_bytes=geometry.wire_hidden_bytes) + recv_hidden, recv_ids, recv_weights, dispatch_ctx = token_pre_dispatch_native( + buffer=buffer, + hidden_states=hidden_states, + routing_weights=weights, + selected_experts=selected, + num_experts=world, + ) + assert recv_hidden.dtype is torch.bfloat16 + assert recv_ids.dtype in (torch.int32, torch.int64) + assert recv_weights.dtype is torch.float32 + local_leaf = (recv_hidden + float(rank + 1)).to(torch.bfloat16).contiguous() + output = native_receive_combine_and_fold( + local_leaf, + buffer=buffer, + dispatch_ctx=dispatch_ctx, + ep_group=dist.group.WORLD, + num_local_experts=1, + ) + + expected_leaves = torch.stack( + [(hidden_states.detach() + float(source + 1)).to(torch.bfloat16) for source in range(world)] + ) + expected = _hierarchical_fold(expected_leaves) + assert output.dtype is torch.bfloat16 + assert torch.equal(output.view(torch.int16), expected.view(torch.int16)), ( + f"rank {rank}: original-handle BF16 segmented combine differs from the explicit FP64 fold" + ) + + # The complete logical-rank loop must be one autograd node. Independent + # sibling nodes have no cross-rank execution-order guarantee and can enter + # DeepEP's reverse-dispatch barriers in different logical-rank orders. + pending = [output.grad_fn] + seen = set() + ordered_combine_nodes = 0 + while pending: + node = pending.pop() + if node is None or node in seen: + continue + # Retain each Python wrapper. Saving only id(node) permits an + # already-released wrapper's address to be reused while traversing. + seen.add(node) + expected_node = "_DeepEPDeterministicCombineBF16Backward" + if type(node).__name__ == expected_node: + ordered_combine_nodes += 1 + pending.extend(next_node for next_node, _index in node.next_functions) + assert ordered_combine_nodes == 1, ( + f"rank {rank}: expected one deterministic DeepEP autograd boundary, found {ordered_combine_nodes}" + ) + + output.float().sum().backward() + torch.cuda.synchronize() + assert hidden_states.grad is not None + assert torch.equal( + hidden_states.grad, + torch.full_like(hidden_states, float(world)), + ), f"rank {rank}: original-handle native exact backward did not reverse dispatch/combine" + + trace_dir = os.environ.get("XORL_DEEPEP_BOUNDARY_TRACE_DIR", "").strip() + if trace_dir: + trace_lines = Path(trace_dir, f"rank{rank:05d}.log").read_text().splitlines() + reverse_dispatch_enters = [ + line for line in trace_lines if "boundary=output_reverse_dispatch state=enter" in line + ] + expected_reverse_dispatches = 1 + assert len(reverse_dispatch_enters) == expected_reverse_dispatches, ( + f"rank {rank}: expected {expected_reverse_dispatches} reverse dispatches, " + f"observed {len(reverse_dispatch_enters)}" + ) + + widths = [None] * world + dist.all_gather_object(widths, geometry.wire_width) + assert widths == [hidden] * world + if rank == 0: + print( + "deepep_native_exact_real_gate_ok " + f"world={world} dispatch_width={hidden} wire_width={geometry.wire_width} " + "combine_mode=deterministic combine_calls=1 wire_dtype=bf16 " + "fold=hierarchical_tree8_bf16_node_fp64_node " + "reverse_dispatches=1 backward=ok", + flush=True, + ) + if deepep_module._default_buffer is not None: + deepep_module._default_buffer.destroy_buffer() + dist.destroy_process_group() + return 0 + + +if __name__ == "__main__" and os.environ.get(_WORKER_ENV) == "1": + sys.exit(_worker_main()) + + +pytestmark = [pytest.mark.distributed, pytest.mark.gpu] + + +def test_original_dispatch_handle_accepts_native_exact_combine_program(): + pytest.importorskip("deep_ep") + import deep_ep + + if not hasattr(deep_ep, "ReductionMode"): + pytest.skip("installed DeepEP lacks the deterministic reduction program") + from distributed_utils import gpu_count, run_distributed_script + + world = int(os.environ.get(_WORLD_ENV, "8")) + if gpu_count() < world: + pytest.skip(f"requires {world} GPUs, found {gpu_count()}") + result = run_distributed_script( + __file__, + num_gpus=world, + timeout=300, + extra_env={_WORKER_ENV: "1"}, + ) + if not result.success: + print("--- distributed worker stdout ---", file=sys.stderr) + print(result.stdout, file=sys.stderr) + print("--- distributed worker stderr ---", file=sys.stderr) + print(result.stderr, file=sys.stderr) + result.assert_success(f"native exact original-handle gate (world {world})") + assert "deepep_native_exact_real_gate_ok" in result.stdout diff --git a/tests/distributed/test_ep_gradient_reduction_contract.py b/tests/distributed/test_ep_gradient_reduction_contract.py index e5521ffb..e1136ed7 100644 --- a/tests/distributed/test_ep_gradient_reduction_contract.py +++ b/tests/distributed/test_ep_gradient_reduction_contract.py @@ -12,7 +12,7 @@ from xorl.distributed.ep_gradients import synchronize_replicated_gradient_parameters from xorl.distributed.gradient_reduction import GradientReductionDomain -from xorl.distributed.torch_parallelize import _build_ep_param_groups +from xorl.distributed.torch_parallelize import _build_ep_param_groups, refresh_ep_param_groups from xorl.models.layers.moe.backend import ep_lora_gradient_reduction_domain @@ -36,6 +36,31 @@ def test_gradient_reduction_domain_admission_policy(): _build_ep_param_groups(model) +def test_refresh_ep_param_groups_rebinds_replaced_parameter_identity(): + model = nn.Module() + shared = nn.Module() + shared._skip_fsdp = True + original = nn.Parameter(torch.empty(1, device="meta")) + shared.weight = original + model.shared = shared + model._fqn2spec_info = { + "shared.weight": SimpleNamespace( + placement=Replicate(), + gradient_reduction=GradientReductionDomain.EP_SUM, + ) + } + + _build_ep_param_groups(model) + assert model._ep_param_groups["ep_replicated_gradient_sync"] == [original] + + replacement = nn.Parameter(torch.ones(1)) + shared.weight = replacement + refresh_ep_param_groups(model) + + assert model._ep_param_groups["ep_replicated_gradient_sync"] == [replacement] + assert all(parameter is not original for parameter in model._ep_param_groups["ep_replicated_gradient_sync"]) + + def test_real_two_rank_backend_contracts(): from tests.distributed.distributed_utils import run_distributed_script diff --git a/tests/distributed/test_mixed_dtype_fsdp_split.py b/tests/distributed/test_mixed_dtype_fsdp_split.py new file mode 100644 index 00000000..ef328438 --- /dev/null +++ b/tests/distributed/test_mixed_dtype_fsdp_split.py @@ -0,0 +1,125 @@ +"""Real two-rank FSDP2 coverage for a composite with BF16 compute and FP32 state.""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +from torch import nn +from torch.distributed._composable.fsdp import fully_shard +from torch.distributed.tensor import DTensor + +from xorl.distributed.parallel_state import get_parallel_state, init_parallel_state +from xorl.distributed.torch_parallelize import ( + _bf16_mixed_precision_policy, + _fully_shard_declared_mixed_dtype_unit, +) +from xorl.utils.device import get_nccl_backend + + +THIS_DIR = Path(__file__).resolve().parent +if str(THIS_DIR) not in sys.path: + sys.path.insert(0, str(THIS_DIR)) + +from distributed_utils import run_distributed_script, skip_if_gpu_count_less_than # noqa: E402 + + +pytestmark = [pytest.mark.gpu, pytest.mark.distributed] + + +class _MixedComposite(nn.Module): + fsdp_full_precision_parameter_names = ("A_log", "dt_bias") + + def __init__(self) -> None: + super().__init__() + self.A_log = nn.Parameter(torch.linspace(-0.4, 0.2, 4, dtype=torch.float32)) + self.dt_bias = nn.Parameter(torch.linspace(0.1, 0.4, 4, dtype=torch.float32)) + self.left = nn.Linear(4, 4, bias=False, dtype=torch.bfloat16) + self.right = nn.Linear(4, 4, bias=False, dtype=torch.bfloat16) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + hidden = self.right(torch.nn.functional.silu(self.left(inputs))) + state = (self.A_log.exp() + self.dt_bias).to(torch.bfloat16) + return hidden + state + + +class _Model(nn.Module): + def __init__(self) -> None: + super().__init__() + self.mixed = _MixedComposite() + self.out = nn.Linear(4, 1, bias=False, dtype=torch.bfloat16) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + return self.out(self.mixed(inputs)) + + +def _build(device: torch.device) -> _Model: + torch.manual_seed(1234) + return _Model().to(device).train() + + +def _run_split() -> None: + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend=get_nccl_backend()) + try: + world_size = dist.get_world_size() + assert world_size == 2 + device = torch.device("cuda", local_rank) + init_parallel_state(dp_size=world_size, dp_shard_size=world_size, device_type="cuda") + mesh = get_parallel_state().dp_shard_mesh + + reference = _build(device) + sharded = _build(device) + _fully_shard_declared_mixed_dtype_unit( + sharded.mixed, + compute_kwargs={"mesh": mesh, "mp_policy": _bf16_mixed_precision_policy()}, + full_precision_kwargs={"mesh": mesh}, + ) + fully_shard(sharded, mesh=mesh, mp_policy=_bf16_mixed_precision_policy()) + + inputs = torch.arange(8, device=device, dtype=torch.bfloat16).reshape(2, 4) + local_rank + reference_loss = reference(inputs).float().square().mean() + sharded_loss = sharded(inputs).float().square().mean() + reference_loss.backward() + sharded_loss.backward() + + for parameter in reference.parameters(): + assert parameter.grad is not None + dist.all_reduce(parameter.grad, op=dist.ReduceOp.SUM) + parameter.grad.div_(world_size) + + reference_parameters = dict(reference.named_parameters()) + sharded_parameters = dict(sharded.named_parameters()) + assert set(reference_parameters) == set(sharded_parameters) + for name, reference_parameter in reference_parameters.items(): + sharded_parameter = sharded_parameters[name] + assert isinstance(sharded_parameter, DTensor), name + assert isinstance(sharded_parameter.grad, DTensor), name + full_gradient = sharded_parameter.grad.full_tensor() + torch.testing.assert_close(full_gradient, reference_parameter.grad, rtol=0, atol=0) + assert full_gradient.dtype == reference_parameter.dtype + + assert sharded_parameters["mixed.A_log"].dtype == torch.float32 + assert sharded_parameters["mixed.dt_bias"].dtype == torch.float32 + assert sharded_parameters["mixed.left.weight"].dtype == torch.bfloat16 + finally: + dist.destroy_process_group() + + +@skip_if_gpu_count_less_than(2) +def test_real_two_rank_mixed_dtype_fsdp_split() -> None: + if os.environ.get("RUN_MIXED_DTYPE_FSDP_SPLIT") == "1": + _run_split() + return + result = run_distributed_script( + __file__, + num_gpus=2, + timeout=180, + extra_env={"RUN_MIXED_DTYPE_FSDP_SPLIT": "1"}, + ) + result.assert_success("mixed-dtype composite should form two sharded FSDP groups") diff --git a/tests/distributed/test_pipeline_model_copy.py b/tests/distributed/test_pipeline_model_copy.py new file mode 100644 index 00000000..f4711b4a --- /dev/null +++ b/tests/distributed/test_pipeline_model_copy.py @@ -0,0 +1,31 @@ +import torch +import torch.distributed as dist +from torch import nn + +from xorl.distributed.pipeline_parallel import _deepcopy_pipeline_model + + +def test_pipeline_model_copy_preserves_process_groups_and_parameter_metadata(tmp_path) -> None: + owns_group = not dist.is_initialized() + if owns_group: + dist.init_process_group( + "gloo", + rank=0, + world_size=1, + init_method=f"file://{tmp_path / 'gloo-init'}", + ) + try: + group = dist.group.WORLD + holder = nn.Linear(2, 2) + holder.cp_group = group + holder.weight._keep_fp32 = True + + holder_copy = _deepcopy_pipeline_model(holder) + + assert holder_copy.cp_group is group + assert holder_copy.weight is not holder.weight + assert torch.equal(holder_copy.weight, holder.weight) + assert holder_copy.weight._keep_fp32 is True + finally: + if owns_group: + dist.destroy_process_group() diff --git a/tests/distributed/test_torch_parallelize_policies.py b/tests/distributed/test_torch_parallelize_policies.py index 312b7c54..5353ba2f 100644 --- a/tests/distributed/test_torch_parallelize_policies.py +++ b/tests/distributed/test_torch_parallelize_policies.py @@ -4,12 +4,14 @@ import torch from torch import nn +import xorl.distributed.torch_parallelize as torch_parallelize from xorl.distributed.torch_parallelize import ( _coerce_optional_bool_config, _configure_manual_fsdp_prefetch, _exact_lm_head_replicated_params, _expert_mixed_precision_policy, _fsdp_kwargs_for_module, + _fully_shard_declared_mixed_dtype_unit, _resolve_fsdp_reduce_dtype, _sequence_parallel_fully_folded_into_fsdp, _topmost_modules_matching, @@ -66,6 +68,57 @@ def test_explicit_full_precision_module_drops_only_its_fsdp_mp_policy() -> None: assert original["mp_policy"] is policy +def test_declared_mixed_dtype_unit_forms_two_sharded_groups_without_renaming(monkeypatch) -> None: + class _DeclaredComposite(nn.Module): + fsdp_full_precision_parameter_names = ("A_log", "dt_bias") + + def __init__(self) -> None: + super().__init__() + self.A_log = nn.Parameter(torch.zeros(4, dtype=torch.float32)) + self.dt_bias = nn.Parameter(torch.ones(4, dtype=torch.float32)) + self.q_proj = nn.Linear(4, 4, bias=False, dtype=torch.bfloat16) + self.o_proj = nn.Linear(4, 4, bias=False, dtype=torch.bfloat16) + + module = _DeclaredComposite() + parameter_names_before = tuple(dict(module.named_parameters())) + calls = [] + + def fake_fully_shard(target, **kwargs) -> None: + calls.append((target, kwargs)) + + monkeypatch.setattr(torch_parallelize, "fully_shard", fake_fully_shard) + compute_kwargs = {"mesh": "mesh", "mp_policy": "bf16"} + full_precision_kwargs = {"mesh": "mesh"} + + representatives = _fully_shard_declared_mixed_dtype_unit( + module, + compute_kwargs=compute_kwargs, + full_precision_kwargs=full_precision_kwargs, + ) + + assert calls == [([module.q_proj, module.o_proj], compute_kwargs), (module, full_precision_kwargs)] + assert representatives == [module, module.q_proj] + assert tuple(dict(module.named_parameters())) == parameter_names_before + + +def test_declared_mixed_dtype_unit_rejects_non_bf16_compute_parameters(monkeypatch) -> None: + class _BadComposite(nn.Module): + fsdp_full_precision_parameter_names = ("state",) + + def __init__(self) -> None: + super().__init__() + self.state = nn.Parameter(torch.zeros(4, dtype=torch.float32)) + self.proj = nn.Linear(4, 4, bias=False, dtype=torch.float32) + + monkeypatch.setattr(torch_parallelize, "fully_shard", lambda *_args, **_kwargs: None) + with pytest.raises(TypeError, match="compute parameters must be uniformly BF16"): + _fully_shard_declared_mixed_dtype_unit( + _BadComposite(), + compute_kwargs={}, + full_precision_kwargs={}, + ) + + def test_dsv4_exact_lm_head_replicates_only_fp32_a() -> None: head = nn.Module() head.lora_A = nn.Parameter(torch.empty(1, 8, dtype=torch.float32)) diff --git a/tests/models/test_deepep_native_exact_shared_layer.py b/tests/models/test_deepep_native_exact_shared_layer.py new file mode 100644 index 00000000..5f581452 --- /dev/null +++ b/tests/models/test_deepep_native_exact_shared_layer.py @@ -0,0 +1,343 @@ +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from xorl.distributed.moe.deepep_native_exact import ( + DeepEPNativeExactError, + NativeDeepEPGeometry, + canonicalize_native_routing_metadata, + native_dispatch_runner_combine, + native_exact_router_topk, + native_zero_row_runner_routes, + reduce_native_runner_routes_to_bf16, +) +from xorl.models.layers.moe.experts import MoEExperts +from xorl.models.layers.moe.lora import MoEExpertsLoRA, MoELoRAConfig +from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import ( + Qwen3MoeSparseTritonMoeBlock, +) + + +class _FakeBuffer: + def __init__(self): + self.hidden_bytes = None + + def init_buffer(self, *, hidden_bytes): + self.hidden_bytes = hidden_bytes + + +def _local_experts(): + experts = MoEExperts( + num_experts=4, + hidden_dim=4, + intermediate_size=2, + hidden_act="silu", + moe_implementation="triton", + ) + # Model parallelization leaves only the contiguous local expert slice. + experts.gate_up_proj = nn.Parameter(torch.empty(2, 4, 4, dtype=torch.bfloat16)) + experts.down_proj = nn.Parameter(torch.empty(2, 2, 4, dtype=torch.bfloat16)) + experts.ep_dispatch = "deepep" + experts.deepep_native_exact = True + return experts + + +def test_qwen3_native_adapter_selects_structural_fp32_exact_router(): + config = SimpleNamespace( + hidden_size=4, + num_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=2, + hidden_act="silu", + norm_topk_prob=True, + train_router=False, + _activation_native=False, + _ep_dispatch="deepep", + _deepep_native_exact=True, + _lora_serving_mode="separate", + ) + + block = Qwen3MoeSparseTritonMoeBlock(config) + + assert block._exact_batch_invariant_router is True + assert block.router._exact_batch_invariant is True + assert block.router._exact_weights_fp32 is True + assert block.deepep_native_exact is True + assert block.experts.lora_serving_mode == "separate" + + +def test_shared_layer_owns_real_dispatch_runner_and_fold(monkeypatch): + experts = _local_experts() + buffer = _FakeBuffer() + dispatch_ctx = SimpleNamespace(num_recv_tokens=2, hidden_dim=4) + recv_hidden = torch.arange(8, dtype=torch.bfloat16).reshape(2, 4).contiguous() + recv_ids = torch.tensor([[0, -1], [1, 0]], dtype=torch.int64) + recv_weights = torch.tensor([[0.75, 0.0], [0.5, 0.25]], dtype=torch.float32) + folded = torch.full((3, 4), 7.0, dtype=torch.bfloat16) + calls = {} + + import xorl.distributed.moe.deepep as deepep + import xorl.distributed.moe.deepep_native_exact as native_exact + + monkeypatch.setattr(deepep, "get_default_buffer", lambda **kwargs: buffer) + + def fake_dispatch(**kwargs): + calls["dispatch"] = kwargs + return recv_hidden, recv_ids, recv_weights, dispatch_ctx + + monkeypatch.setattr(deepep, "token_pre_dispatch_native", fake_dispatch) + monkeypatch.setattr( + native_exact, + "resolve_native_deepep_geometry", + lambda group, hidden: NativeDeepEPGeometry(ep_size=2, ep_rank=1, hidden_size=hidden), + ) + + def fake_runner(hidden, weights, ids, *, local_expert_ids): + calls["runner"] = (hidden, weights, ids, local_expert_ids) + return hidden.clone() + + monkeypatch.setattr(experts, "sglang_fused_experts_forward", fake_runner) + + def fake_fold(recv_output, **kwargs): + calls["fold"] = (recv_output, kwargs) + return folded + + monkeypatch.setattr(native_exact, "native_receive_combine_and_fold", fake_fold) + + hidden = torch.zeros((3, 4), dtype=torch.bfloat16) + routing = torch.full((3, 2), 0.5, dtype=torch.float32) + selected = torch.tensor([[0, 2], [1, 3], [0, 1]], dtype=torch.int64) + parallel = SimpleNamespace(ep_group=object()) + + result = experts._deepep_native_exact_forward(hidden, routing, selected, parallel) + + assert torch.equal(result, folded) + assert buffer.hidden_bytes == 4 * 2 # original H * sizeof(BF16); repeated once per rank + assert torch.equal(calls["dispatch"]["hidden_states"], hidden) + assert calls["runner"][0] is recv_hidden + assert calls["runner"][1] is recv_weights + assert calls["runner"][2].dtype is torch.int32 + assert torch.equal(calls["runner"][2].to(torch.int64), recv_ids) + assert calls["runner"][3] is True + assert calls["fold"][1]["num_local_experts"] == 2 + + +def test_shared_layer_rejects_superseded_route_cube(monkeypatch): + buffer = _FakeBuffer() + dispatch_ctx = SimpleNamespace(num_recv_tokens=2, hidden_dim=2) + recv_hidden = torch.arange(4, dtype=torch.bfloat16).reshape(2, 2).contiguous() + recv_ids = torch.tensor([[0, -1], [1, 0]], dtype=torch.int64) + recv_weights = torch.tensor([[0.75, 123.0], [0.5, 0.25]], dtype=torch.float32) + routes = torch.tensor( + [[[2.0, 4.0], [99.0, 99.0]], [[3.0, 5.0], [7.0, 11.0]]], + dtype=torch.bfloat16, + ) + import xorl.distributed.moe.deepep as deepep + import xorl.distributed.moe.deepep_native_exact as native_exact + + monkeypatch.setattr(deepep, "get_default_buffer", lambda **_kwargs: buffer) + monkeypatch.setattr( + deepep, + "token_pre_dispatch_native", + lambda **_kwargs: (recv_hidden, recv_ids, recv_weights, dispatch_ctx), + ) + monkeypatch.setattr( + native_exact, + "resolve_native_deepep_geometry", + lambda _group, hidden: NativeDeepEPGeometry(ep_size=2, ep_rank=0, hidden_size=hidden), + ) + + with pytest.raises(RuntimeError, match="no_combine=False runner"): + native_dispatch_runner_combine( + torch.zeros((2, 2), dtype=torch.bfloat16), + torch.full((2, 2), 0.5, dtype=torch.float32), + torch.tensor([[0, 2], [1, 3]], dtype=torch.int64), + ep_group=object(), + num_experts=4, + num_local_experts=2, + buffer_size_gb=1.0, + num_sms=8, + runner=lambda hidden, weights, ids: routes, + ) + + +def test_shared_layer_rejects_trainable_router_metadata(): + experts = _local_experts() + hidden = torch.zeros((1, 4), dtype=torch.bfloat16) + routing = torch.ones((1, 1), dtype=torch.float32, requires_grad=True) + selected = torch.zeros((1, 1), dtype=torch.int64) + + with pytest.raises(RuntimeError, match="frozen router"): + experts._deepep_native_exact_forward( + hidden, + routing, + selected, + SimpleNamespace(ep_group=object()), + ) + + +def test_shared_runner_reduces_bf16_routes_with_fp32_metadata(): + routes = torch.tensor( + [ + [[1.0, 2.0], [3.0, 4.0], [99.0, 99.0]], + [[5.0, 6.0], [99.0, 99.0], [7.0, 8.0]], + ], + dtype=torch.bfloat16, + ) + ids = torch.tensor([[0, 1, -1], [1, -1, 0]], dtype=torch.int32) + weights = torch.tensor( + [[0.25, 0.5, 123.0], [0.75, 123.0, 0.125]], + dtype=torch.float32, + ) + + leaf = reduce_native_runner_routes_to_bf16(routes, ids, weights) + expected = ( + torch.where( + (ids >= 0).unsqueeze(-1), + routes.to(torch.float32) * weights.unsqueeze(-1), + torch.zeros((), dtype=torch.float32), + ) + .sum(dim=1) + .to(torch.bfloat16) + ) + + assert leaf.dtype is torch.bfloat16 + assert leaf.is_contiguous() + assert torch.equal(leaf, expected) + + +def test_shared_runner_empty_receive_avoids_kernel_reduction(): + hidden = torch.empty((0, 4), dtype=torch.bfloat16) + ids = torch.empty((0, 3), dtype=torch.int32) + weights = torch.empty((0, 3), dtype=torch.float32) + + routes = native_zero_row_runner_routes(hidden, ids) + leaf = reduce_native_runner_routes_to_bf16(routes, ids, weights) + + assert routes.shape == (0, 3, 4) + assert leaf.shape == hidden.shape + assert leaf.dtype is torch.bfloat16 + + +def test_shared_native_router_builds_fp32_metadata_with_fixed_order_renorm(): + logits = torch.tensor([[1.0, 3.0, 2.0, -1.0]], dtype=torch.float32) + + weights, ids = native_exact_router_topk(logits, top_k=2, renormalize=True) + + assert weights.dtype is torch.float32 + assert ids.dtype is torch.int64 + assert ids.tolist() == [[1, 2]] + scores = torch.softmax(logits, dim=1).gather(1, ids) + expected = (scores / (scores[:, 0] + scores[:, 1]).unsqueeze(-1)).to(torch.bfloat16).to(torch.float32) + assert torch.equal(weights, expected) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_shared_native_routing_metadata_preserves_available_information(dtype): + source = torch.tensor([[0.1234567, 0.8765433]], dtype=dtype) + + metadata = canonicalize_native_routing_metadata(source) + + assert metadata.dtype is torch.float32 + assert torch.equal(metadata, source.to(torch.float32)) + + +def test_native_route_rejects_forced_routing(monkeypatch): + block = Qwen3MoeSparseTritonMoeBlock( + SimpleNamespace( + hidden_size=4, + num_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=2, + hidden_act="silu", + norm_topk_prob=True, + train_router=False, + _activation_native=False, + _ep_dispatch="deepep", + _deepep_native_exact=True, + ) + ).to(torch.bfloat16) + block._diagnostic_forced_selected_experts = torch.zeros((1, 2), dtype=torch.int64) + monkeypatch.setattr(block, "_bi_router_logits", lambda _hidden: torch.zeros((1, 4), dtype=torch.float32)) + + with pytest.raises(RuntimeError, match="forced routing is forbidden"): + block.route(torch.zeros((1, 4), dtype=torch.bfloat16)) + + +def test_parent_native_marker_survives_wrapped_expert_attribute_loss(monkeypatch): + block = Qwen3MoeSparseTritonMoeBlock( + SimpleNamespace( + hidden_size=4, + num_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=2, + hidden_act="silu", + norm_topk_prob=True, + train_router=False, + _activation_native=False, + _ep_dispatch="deepep", + _deepep_native_exact=True, + ) + ).to(torch.bfloat16) + block.experts.deepep_native_exact = False + monkeypatch.setattr( + block, + "_bi_router_logits", + lambda _hidden: torch.tensor([[1.0, 3.0, 2.0, -1.0]], dtype=torch.float32), + ) + + weights, ids, _ = block.route(torch.zeros((1, 4), dtype=torch.bfloat16)) + + assert weights.dtype is torch.float32 + assert ids.tolist() == [[1, 2]] + + +def test_shared_runner_reduction_rejects_non_bf16_routes(): + with pytest.raises(DeepEPNativeExactError, match="routes must be BF16"): + reduce_native_runner_routes_to_bf16( + torch.empty((1, 1, 4), dtype=torch.float32), + torch.zeros((1, 1), dtype=torch.int32), + torch.ones((1, 1), dtype=torch.float32), + ) + + +def test_lora_adapter_delegates_context_to_shared_native_program(monkeypatch): + experts = MoEExpertsLoRA( + num_experts=4, + num_local_experts=2, + hidden_dim=4, + intermediate_size=2, + moe_implementation="triton", + lora_config=MoELoRAConfig(r=2, lora_alpha=2), + ) + experts.ep_dispatch = "deepep" + experts.deepep_native_exact = True + experts.lora_serving_mode = "separate" + hidden = torch.zeros((1, 4), dtype=torch.bfloat16) + routing = torch.ones((1, 1), dtype=torch.float32) + selected = torch.zeros((1, 1), dtype=torch.int64) + expected = torch.full_like(hidden, 3) + calls = {} + + import xorl.distributed.moe.deepep_native_exact as native_exact + + def fake_program(*args, **kwargs): + calls["args"] = args + calls["kwargs"] = kwargs + return expected + + monkeypatch.setattr(native_exact, "native_dispatch_runner_combine", fake_program) + + result = experts._ep_forward( + hidden, + routing, + selected, + SimpleNamespace(ep_group=object()), + ) + + assert result is expected + assert calls["kwargs"]["num_experts"] == 4 + assert calls["kwargs"]["num_local_experts"] == 2 + assert calls["kwargs"]["runner"].__self__ is experts diff --git a/tests/models/test_dsv4_attention.py b/tests/models/test_dsv4_attention.py index 4921718c..b2bbf6dc 100644 --- a/tests/models/test_dsv4_attention.py +++ b/tests/models/test_dsv4_attention.py @@ -708,6 +708,7 @@ def block_head(hidden, *_args): cfg = _tiny_config(compress_ratios=[0]) cfg._dsv4_flash_exact_mode = True + cfg._ep_dispatch = "deepep" model = DeepseekV4Model(cfg, moe_implementation="eager") recorder = _RecordingLayer() model.layers = nn.ModuleList([recorder]) diff --git a/tests/models/test_dsv4_exact_mhc.py b/tests/models/test_dsv4_exact_mhc.py new file mode 100644 index 00000000..01c10735 --- /dev/null +++ b/tests/models/test_dsv4_exact_mhc.py @@ -0,0 +1,132 @@ +from types import SimpleNamespace + +import pytest +import torch + +from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import ( + _build_cp_serving_mhc_segments, + _build_serving_mhc_segments, +) +from xorl.ops.dsv4.cp_utils import Dsv4ExactCPLayout +from xorl.ops.dsv4.hyper_connection import DeepSeekV4HyperConnectionUtil + + +pytestmark = pytest.mark.cpu + + +def test_serving_mhc_segments_encode_prefill_then_m1_and_padding(): + segments = _build_serving_mhc_segments( + compute_rows=10, + sample_lengths=[4, 3], + sampler_prefill_lengths=torch.tensor([2, 3]), + ) + assert segments == (2, 1, 1, 3, 3) + + +def test_serving_mhc_segments_reject_invalid_boundary(): + with pytest.raises(ValueError, match="0 < prefill <= sample length"): + _build_serving_mhc_segments( + compute_rows=4, + sample_lengths=[4], + sampler_prefill_lengths=torch.tensor([5]), + ) + + +def _cp_layout(): + return Dsv4ExactCPLayout( + local_storage_indices=torch.arange(4), + local_logical_rows=torch.tensor([2, 3, 4, 5, -1, -1]), + local_request_ids=torch.tensor([0, 0, 1, 1, -1, -1]), + local_request_positions=torch.tensor([2, 3, 0, 1, 0, 0]), + local_live_count=4, + compute_rows=6, + gather_order=torch.arange(8), + global_logical_rows=torch.arange(8), + global_request_ids=torch.tensor([0, 0, 0, 0, 1, 1, 1, 1]), + global_request_positions=torch.tensor([0, 1, 2, 3, 0, 1, 2, 3]), + request_ids=(0, 1), + local_request_row_indices=(torch.tensor([0, 1]), torch.tensor([2, 3])), + global_request_row_indices=(torch.arange(4), torch.arange(4, 8)), + ) + + +def test_cp_serving_mhc_segments_preserve_global_prefill_launch_sizes(): + segments = _build_cp_serving_mhc_segments( + layout=_cp_layout(), + sampler_prefill_lengths=torch.tensor([3, 2]), + ) + + assert [segment.launch_rows for segment in segments] == [3, 1, 2, 2] + assert [segment.source_rows for segment in segments] == [(0,), (1,), (2, 3), (4, 5)] + assert [segment.launch_positions for segment in segments] == [(2,), (0,), (0, 1), (0, 1)] + + +def test_exact_mhc_replay_invokes_one_prefill_then_m1(monkeypatch): + calls = [] + + def fake_apply(residual, *_args): + calls.append(residual.shape[1]) + shape = residual.shape[:-2] + return ( + residual[..., 0, :], + torch.zeros(*shape, 4), + torch.zeros(*shape, 4, 4), + ) + + monkeypatch.setattr("xorl.ops.dsv4.hyper_connection._ExactMhcPreNorm.apply", fake_apply) + util = DeepSeekV4HyperConnectionUtil( + SimpleNamespace(rms_norm_eps=1e-6, hc_mult=4, hc_sinkhorn_iters=20, hc_eps=1e-6) + ) + residual = torch.randn(1, 5, 4, 8, requires_grad=True) + layer_input, post, comb = util.layer_pre_norm_exact( + residual, + torch.empty(24, 32), + torch.empty(3), + torch.empty(24), + torch.empty(8), + serving_segments=(3, 1, 1), + ) + + assert calls == [3, 1, 1] + assert layer_input.shape == (1, 5, 8) + assert post.shape == (1, 5, 4) + assert comb.shape == (1, 5, 4, 4) + layer_input.sum().backward() + assert residual.grad is not None + + +def test_exact_cp_mhc_replay_uses_global_m_and_selects_local_rows(monkeypatch): + calls = [] + + def fake_apply(residual, *_args): + calls.append(residual.shape[1]) + shape = residual.shape[:-2] + return ( + residual[..., 0, :], + torch.zeros(*shape, 4), + torch.zeros(*shape, 4, 4), + ) + + monkeypatch.setattr("xorl.ops.dsv4.hyper_connection._ExactMhcPreNorm.apply", fake_apply) + util = DeepSeekV4HyperConnectionUtil( + SimpleNamespace(rms_norm_eps=1e-6, hc_mult=4, hc_sinkhorn_iters=20, hc_eps=1e-6) + ) + residual = torch.randn(1, 6, 4, 8, requires_grad=True) + segments = _build_cp_serving_mhc_segments( + layout=_cp_layout(), + sampler_prefill_lengths=torch.tensor([3, 2]), + ) + + layer_input, _post, _comb = util.layer_pre_norm_exact( + residual, + torch.empty(24, 32), + torch.empty(3), + torch.empty(24), + torch.empty(8), + serving_segments=segments, + ) + + assert calls == [3, 1, 2, 2] + assert torch.equal(layer_input, residual[..., 0, :]) + layer_input.sum().backward() + assert residual.grad is not None diff --git a/tests/models/test_dsv4_model.py b/tests/models/test_dsv4_model.py index 1019d6bd..ade4c69f 100644 --- a/tests/models/test_dsv4_model.py +++ b/tests/models/test_dsv4_model.py @@ -308,6 +308,20 @@ def test_for_causal_lm_forward_backward(): assert grads[name].grad is None, f"{name} unexpectedly received a gradient" +def test_for_causal_lm_propagates_hidden_states(): + from xorl.models.transformers.deepseek_v4 import DeepseekV4ForCausalLM # noqa: PLC0415 + + cfg = _tiny_config(num_hidden_layers=2, compress_ratios=[0, 0]) + model = _make_model(cfg, DeepseekV4ForCausalLM) + input_ids = torch.randint(0, cfg.vocab_size, (1, cfg.sliding_window), dtype=torch.long) + + outputs = model(input_ids=input_ids, output_hidden_states=True) + + assert outputs.hidden_states is not None + assert len(outputs.hidden_states) == cfg.num_hidden_layers + 1 + assert outputs.hidden_states[0].shape == (1, cfg.sliding_window, cfg.hc_mult, cfg.hidden_size) + + def test_for_causal_lm_with_hash_layer(): """First layer is hash-routed; verify input_ids threads through correctly.""" from xorl.models.transformers.deepseek_v4 import DeepseekV4ForCausalLM # noqa: PLC0415 @@ -460,3 +474,96 @@ def fake_checkpoint(func, *args, **kwargs): assert out.shape == (1, cfg.sliding_window, cfg.hidden_size) assert len(calls) == cfg.num_hidden_layers assert all(call_kwargs["input_ids"] is input_ids for _, _, call_kwargs in calls) + + +def test_exact_residual_capture_selects_live_position_before_packed_padding(tmp_path, monkeypatch): + from xorl.models.transformers.deepseek_v4 import DeepseekV4Model # noqa: PLC0415 + + model = DeepseekV4Model.__new__(DeepseekV4Model) + nn.Module.__init__(model) + monkeypatch.setenv("XORL_DSV4_TRAINER_LAYER_CAPTURE_DIR", str(tmp_path)) + monkeypatch.setenv("XORL_DSV4_TRAINER_LAYER_CAPTURE_POSITION", "64") + hidden = torch.arange(3 * 2 * 4, dtype=torch.bfloat16).reshape(1, 3, 2, 4) + positions = torch.tensor([[63, 64, 65, 0, 0]]) + input_ids = torch.tensor([[11, 107413, 13, 0, 0]]) + + model._maybe_capture_exact_residual_row( + layer_id=-1, + hidden_states=hidden, + position_ids=positions, + input_ids=input_ids, + ) + + [capture_path] = list(tmp_path.glob("*.pt")) + payload = torch.load(capture_path, map_location="cpu", weights_only=True) + assert payload["schema"] == "xorl.dsv4_trainer_layer_output.v1" + assert payload["layer"] == -1 + assert payload["position"] == 64 + assert payload["token_id"] == 107413 + torch.testing.assert_close(payload["hidden"], hidden[:, 1]) + + # A later decode request no longer contains position 64. Once the full + # forward captured this layer, the one-shot hook must remain inert. + model._maybe_capture_exact_residual_row( + layer_id=-1, + hidden_states=hidden[:, :1], + position_ids=torch.tensor([[66]]), + input_ids=torch.tensor([[17]]), + ) + assert list(tmp_path.glob("*.pt")) == [capture_path] + + +def test_exact_component_capture_selects_owner_token_and_skips_dummy_rank(tmp_path, monkeypatch): + from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import ( # noqa: PLC0415 + DeepseekV4DecoderLayer, + ) + + monkeypatch.setenv("XORL_DSV4_TRAINER_COMPONENT_CAPTURE_DIR", str(tmp_path)) + monkeypatch.setenv("XORL_DSV4_COMPONENT_CAPTURE_LAYER", "2") + monkeypatch.setenv("XORL_DSV4_COMPONENT_CAPTURE_TOKEN_ID", "107413") + hidden = torch.arange(3 * 4 * 5, dtype=torch.bfloat16).reshape(1, 3, 4, 5) + input_ids = torch.tensor([[11, 107413, 13, 0, 0]]) + + owner = DeepseekV4DecoderLayer.__new__(DeepseekV4DecoderLayer) + nn.Module.__init__(owner) + owner.layer_id = 2 + owner._maybe_capture_exact_component("layer_input", hidden, input_ids, live_token_count=3) + + dummy = DeepseekV4DecoderLayer.__new__(DeepseekV4DecoderLayer) + nn.Module.__init__(dummy) + dummy.layer_id = 2 + dummy._maybe_capture_exact_component("layer_input", hidden, input_ids, live_token_count=0) + + [capture_path] = list(tmp_path.glob("*.pt")) + payload = torch.load(capture_path, map_location="cpu", weights_only=True) + assert payload["schema"] == "xorl.dsv4_trainer_component.v1" + assert payload["component"] == "layer_input" + assert payload["token_id"] == 107413 + torch.testing.assert_close(payload["value"], hidden[:, 1]) + + +def test_exact_attention_component_capture_selects_position(tmp_path, monkeypatch): + from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import ( # noqa: PLC0415 + DeepSeekV4Attention, + ) + + monkeypatch.setenv("XORL_DSV4_TRAINER_ATTENTION_CAPTURE_DIR", str(tmp_path)) + monkeypatch.setenv("XORL_DSV4_ATTENTION_CAPTURE_LAYER", "2") + monkeypatch.setenv("XORL_DSV4_ATTENTION_CAPTURE_POSITION", "64") + attention = DeepSeekV4Attention.__new__(DeepSeekV4Attention) + nn.Module.__init__(attention) + attention.layer_id = 2 + value = torch.arange(3 * 4 * 5, dtype=torch.bfloat16).reshape(1, 3, 4, 5) + + attention._maybe_capture_exact_attention_component( + "q_pre_attention", + value, + torch.tensor([63, 64, 65]), + ) + + [capture_path] = list(tmp_path.glob("*.pt")) + payload = torch.load(capture_path, map_location="cpu", weights_only=True) + assert payload["schema"] == "xorl.dsv4_trainer_attention_component.v1" + assert payload["component"] == "q_pre_attention" + assert payload["position"] == 64 + torch.testing.assert_close(payload["value"], value[:, 1]) diff --git a/tests/models/test_dsv4_moe.py b/tests/models/test_dsv4_moe.py index 91dfb7c5..e9dd6d78 100644 --- a/tests/models/test_dsv4_moe.py +++ b/tests/models/test_dsv4_moe.py @@ -296,6 +296,7 @@ def test_exact_dsv4_checkpointing_does_not_attach_or_activate_routing_replay(): cfg = _tiny_config(num_hash_layers=1) cfg._dsv4_flash_exact_mode = True + cfg._ep_dispatch = "deepep" block = DeepseekV4MoE(cfg, layer_id=0) block._routing_replay = RoutingReplay() # Simulate stale state from a prior enable call. container = nn.Module() @@ -417,6 +418,7 @@ def _shared_partial(hidden_states, _module, **kwargs): shared_experts=object(), routed_scaling_factor=1.5, _capture_diagnostic_component=lambda *_args: None, + _maybe_capture_native_moe_output=lambda *_args: None, route=lambda hidden_states, input_ids=None: ( torch.ones(hidden_states.shape[0], 6), torch.zeros(hidden_states.shape[0], 6, dtype=torch.int32), @@ -482,6 +484,243 @@ def _shared_partial(hidden_states, _module, **kwargs): assert observed == {"routed": True, "shared": True} +def test_dsv4_moe_program_resolves_native_deepep_without_feature_menu(): + from xorl.models.transformers.deepseek_v4.moe_program import ( + DSV4_DEEPEP_NATIVE_EXACT_V1, + resolve_dsv4_moe_numerical_program, + ) + + assert ( + resolve_dsv4_moe_numerical_program( + exact=True, + ep_dispatch="deepep", + deepep_native_exact=True, + ) + == DSV4_DEEPEP_NATIVE_EXACT_V1 + ) + assert ( + resolve_dsv4_moe_numerical_program( + exact=False, + ep_dispatch="alltoall", + deepep_native_exact=False, + ) + is None + ) + with pytest.raises(ValueError, match="retired post-expert diagnostic"): + resolve_dsv4_moe_numerical_program( + exact=True, + ep_dispatch="alltoall", + deepep_native_exact=False, + ) + + +def test_dsv4_native_deepep_program_rejects_non_deepep_dispatch(): + from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepseekV4MoE + from xorl.models.transformers.deepseek_v4.moe_program import DSV4_DEEPEP_NATIVE_EXACT_V1 + + cfg = _tiny_config() + cfg._dsv4_flash_exact_mode = True + cfg._dsv4_moe_numerical_program = DSV4_DEEPEP_NATIVE_EXACT_V1 + cfg._ep_dispatch = "alltoall" + with pytest.raises(RuntimeError, match="requires ep_dispatch='deepep'"): + DeepseekV4MoE(cfg, layer_id=0) + + +def test_dsv4_native_deepep_consumes_transport_receive_layout(monkeypatch): + from types import SimpleNamespace + + import xorl.distributed.moe.deepep_native_exact as native_exact_module + import xorl.distributed.parallel_state as parallel_state_module + import xorl.models.transformers.deepseek_v4.native_payload as payload_module + from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepseekV4MoE + + observed = {} + + class _RecordingExperts: + deepep_buffer_size_gb = 2.0 + deepep_num_sms = 20 + native_mxfp4_payload = SimpleNamespace(w13_weight=torch.empty(32, 1, 1)) + + def __call__(self, hidden_states, routing_weights, selected_experts, **kwargs): + observed["recv_hidden"] = hidden_states.clone() + observed["recv_weights"] = routing_weights.clone() + observed["recv_global_ids"] = selected_experts.clone() + observed["expert_kwargs"] = kwargs + return hidden_states + 3 + + parallel_state = SimpleNamespace( + ep_enabled=True, + ep_group=object(), + ep_size=8, + ep_rank=3, + tp_size=1, + ) + monkeypatch.setattr(parallel_state_module, "get_parallel_state", lambda: parallel_state) + recv_hidden = torch.tensor([[5.0, 6.0, 7.0, 8.0]]) + recv_local_ids = torch.tensor([[0, -1]], dtype=torch.int32) + recv_weights = torch.tensor([[0.75, 0.0]], dtype=torch.float32) + + def _native_program(hidden_states, routing_weights, selected_experts, **kwargs): + observed["program_args"] = (hidden_states, routing_weights, selected_experts) + observed["program_kwargs"] = kwargs + local_leaf = kwargs["runner"](recv_hidden, recv_weights, recv_local_ids) + observed["local_leaf"] = local_leaf + return local_leaf - 1 + + monkeypatch.setattr( + native_exact_module, + "native_dispatch_runner_combine", + _native_program, + ) + + def _shared(hidden_states, _module, **kwargs): + observed["shared_kwargs"] = kwargs + return torch.zeros_like(hidden_states) + + monkeypatch.setattr(payload_module, "dsv4_native_shared_expert_tp_partial", _shared) + monkeypatch.setattr( + payload_module, + "dsv4_join_routed_shared_partial", + lambda routed, shared, **_kwargs: routed + shared, + ) + + fake_moe = SimpleNamespace( + is_hash_layer=False, + layer_id=0, + num_experts=256, + top_k=2, + train_router=False, + experts=_RecordingExperts(), + shared_experts=object(), + routed_scaling_factor=1.5, + _capture_diagnostic_component=lambda *_args: None, + _maybe_capture_native_moe_output=lambda *_args: None, + route=lambda hidden_states, input_ids=None: ( + torch.tensor([[0.75, 0.25]], dtype=torch.float32), + torch.tensor([[96, 129]], dtype=torch.int32), + torch.zeros(hidden_states.shape[0], 256), + ), + ) + hidden = torch.tensor([[[1.0, 2.0, 3.0, 4.0], [9.0, 9.0, 9.0, 9.0]]]) + output, logits = DeepseekV4MoE._forward_deepep_native_exact( + fake_moe, + hidden, + torch.tensor([[7, 8]]), + live_token_count=1, + ) + + assert torch.equal(observed["program_args"][0], hidden.reshape(-1, 4)[:1]) + assert "runner_output_layout" not in observed["program_kwargs"] + assert torch.equal(observed["recv_global_ids"], torch.tensor([[96, -1]], dtype=torch.int32)) + assert observed["expert_kwargs"] == { + "dsv4_exact_native": True, + "dsv4_exact_lora_live": True, + "dsv4_exact_return_routes": False, + } + assert observed["local_leaf"].shape == (1, 4) + assert observed["shared_kwargs"]["tp_rank"] == 0 + assert observed["shared_kwargs"]["tp_size"] == 1 + assert torch.equal(output[0, 0], recv_hidden[0] + 2) + assert torch.equal(output[0, 1], torch.zeros(4)) + assert logits.shape == (2, 256) + + +def test_dsv4_native_expert_adapter_exposes_unweighted_route_layout(monkeypatch): + from types import SimpleNamespace + + import xorl.models.transformers.deepseek_v4.native_payload as payload_module + + factors = { + projection: ( + torch.ones((1, 1, 1), requires_grad=True), + torch.ones((1, 1, 1), requires_grad=True), + ) + for projection in ("gate_proj", "up_proj", "down_proj") + } + + class _Experts: + active_r = 1 + active_lora_alpha = 1 + native_mxfp4_payload = SimpleNamespace(w13_weight=torch.empty((1, 1, 1))) + + @staticmethod + def _active_scaling(): + return 1.0 + + @staticmethod + def _active_lora_views(projection): + return factors[projection] + + observed = {} + + def fake_forward( + hidden_states, + routing_weights, + selected_experts, + *args, + no_combine=False, + ): + observed["no_combine"] = no_combine + observed["weights"] = routing_weights + experts = args[-1] + assert isinstance(experts, _Experts) + routes = hidden_states.unsqueeze(1).expand(-1, selected_experts.shape[1], -1).contiguous() + return routes, selected_experts.to(torch.int32) + + monkeypatch.setattr(payload_module, "_dsv4_native_mxfp4_forward", fake_forward) + hidden = torch.tensor([[1.0, 2.0]], dtype=torch.bfloat16) + weights = torch.tensor([[0.75, 0.25]], dtype=torch.float32) + selected = torch.tensor([[0, -1]], dtype=torch.int32) + + routes = payload_module.dsv4_native_mxfp4_routed_partial( + hidden, + weights, + selected, + _Experts(), + return_routes=True, + ) + + assert observed["no_combine"] is True + assert observed["weights"] is weights + assert routes.shape == (1, 2, 2) + assert routes.dtype is torch.bfloat16 + + +def test_dsv4_native_expert_adapter_defines_empty_route_layout(): + from types import SimpleNamespace + + import xorl.models.transformers.deepseek_v4.native_payload as payload_module + + factors = { + projection: (torch.empty((1, 1, 1)), torch.empty((1, 1, 1))) + for projection in ("gate_proj", "up_proj", "down_proj") + } + + class _Experts: + active_r = 1 + active_lora_alpha = 1 + native_mxfp4_payload = SimpleNamespace(w13_weight=torch.empty((1, 1, 1))) + + @staticmethod + def _active_scaling(): + return 1.0 + + @staticmethod + def _active_lora_views(projection): + return factors[projection] + + routes = payload_module.dsv4_native_mxfp4_routed_partial( + torch.empty((0, 4), dtype=torch.bfloat16), + torch.empty((0, 6), dtype=torch.float32), + torch.empty((0, 6), dtype=torch.int32), + _Experts(), + return_routes=True, + ) + + assert routes.shape == (0, 6, 4) + assert routes.dtype is torch.bfloat16 + + # --------------------------------------------------------------------------- # Routing-replay × hash-routed layer # --------------------------------------------------------------------------- diff --git a/tests/models/test_glm52_contract.py b/tests/models/test_glm52_contract.py index 73ae4bb8..315662af 100644 --- a/tests/models/test_glm52_contract.py +++ b/tests/models/test_glm52_contract.py @@ -1,3 +1,4 @@ +import math from types import MethodType, SimpleNamespace import pytest @@ -912,6 +913,73 @@ def mark_fused_suffix(key, _weight, _bias, _eps, _cache, _positions): assert torch.equal(mixed[:, 4:], torch.full_like(mixed[:, 4:], 11)) +@pytest.mark.cpu +def test_sampler_index_k_preparation_is_segment_aware_for_packed_requests(): + split = torch.full((1, 9, 128), 7, dtype=torch.bfloat16) + raw = torch.zeros_like(split) + weight = torch.ones((128,), dtype=torch.float32) + bias = torch.zeros((128,), dtype=torch.float32) + cos = torch.ones((1, 9, 64), dtype=torch.float32) + sin = torch.zeros_like(cos) + calls = [] + + def mark_fused_suffix(key, _weight, _bias, _eps, _cache, _positions): + calls.append(key.shape[0]) + return torch.full_like(key, 11) + + mixed = _mix_sampler_index_k_preparation( + split, + raw, + weight, + bias, + 1e-6, + (cos, sin), + torch.tensor([3, 2], dtype=torch.int64), + [5, 4], + query_offset=0, + interleaved=True, + _native_kernel_for_testing=mark_fused_suffix, + ) + + assert calls == [2, 2] + assert torch.equal(mixed[:, :3], split[:, :3]) + assert torch.equal(mixed[:, 3:5], torch.full_like(mixed[:, 3:5], 11)) + assert torch.equal(mixed[:, 5:7], split[:, 5:7]) + assert torch.equal(mixed[:, 7:], torch.full_like(mixed[:, 7:], 11)) + + +@pytest.mark.cpu +def test_sampler_index_k_preparation_intersects_packed_segments_with_cp_shard(): + split = torch.full((1, 4, 128), 7, dtype=torch.bfloat16) + raw = torch.zeros_like(split) + cos = torch.ones((1, 4, 64), dtype=torch.float32) + sin = torch.zeros_like(cos) + calls = [] + + def mark_fused_suffix(key, *_args): + calls.append(key.shape[0]) + return torch.full_like(key, 11) + + mixed = _mix_sampler_index_k_preparation( + split, + raw, + torch.ones((128,), dtype=torch.float32), + torch.zeros((128,), dtype=torch.float32), + 1e-6, + (cos, sin), + torch.tensor([3, 2], dtype=torch.int64), + [5, 4], + query_offset=4, + interleaved=True, + _native_kernel_for_testing=mark_fused_suffix, + ) + + assert calls == [1, 1] + assert torch.equal(mixed[:, :1], torch.full_like(mixed[:, :1], 11)) + assert torch.equal(mixed[:, 1:3], split[:, 1:3]) + assert torch.equal(mixed[:, 3:], torch.full_like(mixed[:, 3:], 11)) + + @pytest.mark.cpu def test_sampler_index_k_preparation_maps_4096_boundary_across_cp16(): local_length = 260 @@ -1144,6 +1212,7 @@ def test_correction_bias_stays_fp32_and_checkpoint_ingestion_fails_closed(): def test_canonical_moe_checkpoint_replay_preserves_serving_routing_bytes_and_router_gradients(monkeypatch): config = _small_glm_config() config._glm52_exact_contract = True + config.train_router = True config.routed_scaling_factor = 3.25 block = Glm5MoEBlock(config, layer_idx=1) container = nn.Module() @@ -1157,7 +1226,6 @@ def test_canonical_moe_checkpoint_replay_preserves_serving_routing_bytes_and_rou handler, ) monkeypatch.setattr(block._routing_replay, "_target_device", lambda: torch.device("cpu")) - block.gate._glm52_exact_fullparam_component = True with torch.no_grad(): block.gate.weight.copy_(torch.linspace(-0.1, 0.1, block.gate.weight.numel()).reshape_as(block.gate.weight)) @@ -1194,6 +1262,17 @@ def serving_grouped_topk( original_hidden = torch.randn(2, config.hidden_size, dtype=torch.bfloat16) recompute_hidden = torch.randn_like(original_hidden) try: + # Exact replay qualification begins with a forward-only request. A + # trainable router must preserve the serving bytes in that no-grad + # phase without requiring an autograd edge that PyTorch suppresses by + # contract. + with torch.no_grad(): + forward_only_weights, forward_only_ids, _ = block.route(original_hidden) + assert not forward_only_weights.requires_grad + assert torch.equal(forward_only_weights.view(torch.uint8), serving_weights.view(torch.uint8)) + assert torch.equal(forward_only_ids, serving_ids) + topk_calls.clear() + set_replay_stage("record") recorded_weights, recorded_ids, _ = block.route(original_hidden) assert len(topk_calls) == 1 @@ -1239,6 +1318,26 @@ def serving_grouped_topk( RoutingReplay.clear_all() +@pytest.mark.cpu +def test_model_runner_collects_glm_router_gradient_evidence_by_module_identity(): + config = _small_glm_config() + model = nn.Module() + model.config = config + model.router = Glm5TopkRouter(config) + runner = ModelRunner.__new__(ModelRunner) + runner.model = model + runner._reset_glm52_router_gradient_evidence() + model.router(torch.ones(1, config.hidden_size)).sum().backward() + + metrics = runner._collect_glm52_router_gradient_metrics() + + assert metrics["router_grad_tensor_count"] == 1 + assert metrics["router_grad_missing_count"] == 0 + assert metrics["router_grad_nonfinite_count"] == 0 + assert metrics["router_grad_nonzero_count"] == model.router.weight.numel() + assert metrics["router_grad_norm"] == pytest.approx(math.sqrt(model.router.weight.numel())) + + @pytest.mark.cpu def test_canonical_moe_replay_fails_closed_without_literal_serving_weights(monkeypatch): config = _small_glm_config() @@ -1413,9 +1512,11 @@ def canonical_forward( routing_weights, selected_experts, absolute_positions, + backward_layer_dependency=None, *, _layer_id=layer_id, ): + del backward_layer_dependency partials = _semantic_rank_partials(self, hidden_states, routing_weights, selected_experts) rows = hidden_states.shape[0] * hidden_states.shape[1] metadata = CanonicalMoEGraphMetadata.build( diff --git a/tests/models/test_glm52_deepep_exact_modeling.py b/tests/models/test_glm52_deepep_exact_modeling.py new file mode 100644 index 00000000..3d147b14 --- /dev/null +++ b/tests/models/test_glm52_deepep_exact_modeling.py @@ -0,0 +1,229 @@ +from __future__ import annotations + +from types import MethodType, SimpleNamespace + +import pytest +import torch +from torch import nn + +from xorl.distributed.canonical_moe import canonical_moe_fold_fp64_v3 +from xorl.models.transformers.glm5.exact_routed_experts_qlora import ( + Glm52ExactEP16BlockFP8QLoRARoutedExperts, +) +from xorl.models.transformers.glm5.exact_shared_expert_qlora import ( + Glm52ExactTP16SharedExpertBlockFP8QLoRA, +) +from xorl.models.transformers.glm5.modeling_glm5 import Glm5MoEBlock + + +def _empty_block() -> Glm5MoEBlock: + block = Glm5MoEBlock.__new__(Glm5MoEBlock) + nn.Module.__init__(block) + block.routed_scaling_factor = 2.5 + return block + + +def test_glm_moe_route_exposes_final_independently_computed_routing_diagnostics() -> None: + block = _empty_block() + block.canonical_contract_version = None + block.config = SimpleNamespace() + block.train_router = False + block._routing_replay = None + block.gate = nn.Linear(4, 3, bias=False) + block.experts = nn.Identity() + block.experts.ep_dispatch = "alltoall" + weights = torch.tensor([[0.75, 0.25]], dtype=torch.float32) + ids = torch.tensor([[2, 0]], dtype=torch.int32) + block._route_tokens_to_experts = MethodType( + lambda self, router_logits, input_dtype, **kwargs: (weights, ids), + block, + ) + captured = {} + block._diagnostic_capture_component = lambda name, value: captured.setdefault(name, value) + + routed_weights, routed_ids, router_logits = block.route(torch.ones((1, 4), dtype=torch.float32)) + + assert captured["moe_router_logits"] is router_logits + assert captured["moe_topk_ids"] is routed_ids + assert captured["moe_topk_weights"] is routed_weights + + +def test_glm_moe_canonical_boundary_captures_input_and_combined_output() -> None: + block = _empty_block() + block.canonical_contract_version = "test" + expected = torch.full((1, 2, 4), 0.5, dtype=torch.bfloat16) + block._canonical_ep_forward = MethodType(lambda self, *args: expected, block) + captured = {} + block._diagnostic_capture_component = lambda name, value: captured.setdefault(name, value) + hidden = torch.zeros_like(expected) + + output = block.forward_experts_with_shared( + hidden, + torch.ones((2, 1), dtype=torch.float32), + torch.zeros((2, 1), dtype=torch.int32), + torch.arange(2), + ) + + assert output is expected + assert captured["moe_input"] is hidden + assert captured["moe_experts_output"] is expected + + +def test_canonical_routed_boundary_accepts_deepep_local_ids_without_global_ids() -> None: + block = _empty_block() + experts = Glm52ExactEP16BlockFP8QLoRARoutedExperts(128, 128, ep_rank=7, device="cpu") + captured = {} + + def forward(self, hidden, routing, selected_experts=None, **kwargs): + captured.update( + selected_experts=selected_experts, + local_ids=kwargs["sglang_ep_native_local_ids"], + ) + return torch.ones_like(hidden) + + experts.forward = MethodType(forward, experts) + block.experts = experts + hidden = torch.zeros((3, 128), dtype=torch.bfloat16) + routing = torch.arange(24, dtype=torch.float32).reshape(3, 8).div_(32) + local_ids = torch.tensor( + [[0, -1, -1, 3, -1, -1, -1, -1]] * 3, + dtype=torch.int32, + ) + + output = block._canonical_routed_local_partial(hidden, routing, None, local_ids) + + assert torch.equal(output, torch.ones_like(hidden)) + assert captured["selected_experts"] is None + assert captured["local_ids"] is local_ids + + +def test_native_deepep_routed_boundary_preserves_empty_receive_batch_and_structural_factor_gradients( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from xorl.distributed.moe import deepep_native_exact + + block = _empty_block() + experts = Glm52ExactEP16BlockFP8QLoRARoutedExperts(128, 128, ep_rank=7, device="cpu") + experts.ep_dispatch = "deepep" + experts.deepep_buffer_size_gb = 1.0 + experts.deepep_num_sms = 24 + + def fail_if_called(*args, **kwargs): + raise AssertionError("the fused expert runner must not receive an empty DeepEP batch") + + experts.forward = fail_if_called + block.experts = experts + block.train_router = False + block.num_experts = 256 + captured = {} + program_kwargs = {} + block._diagnostic_capture_component = lambda name, value: captured.setdefault(name, value) + + def fake_program(hidden, routing, selected, **kwargs): + program_kwargs.update(kwargs) + empty_hidden = hidden.new_empty((0, hidden.shape[1])) + empty_weights = routing.new_empty((0, routing.shape[1])) + empty_ids = selected.new_empty((0, selected.shape[1]), dtype=torch.int32) + empty_leaf = kwargs["runner"](empty_hidden, empty_weights, empty_ids) + assert empty_leaf.dtype is torch.bfloat16 + assert empty_leaf.shape == empty_hidden.shape + assert empty_leaf.is_contiguous() + return torch.zeros_like(hidden) + empty_leaf.sum() + + monkeypatch.setattr(deepep_native_exact, "native_dispatch_runner_combine", fake_program) + hidden = torch.zeros((2, 128), dtype=torch.bfloat16) + routing = torch.full((2, 8), 0.125, dtype=torch.float32) + selected = torch.arange(16, dtype=torch.int32).reshape(2, 8) + layer_dependency = torch.ones_like(hidden, requires_grad=True) + shared_dependency = torch.full_like(hidden, 2.0, requires_grad=True) + + output = block._native_deepep_routed_local( + hidden, + routing, + selected, + torch.ones(2, dtype=torch.bool), + ep_rank=7, + ep_size=16, + ep_group=object(), + backward_layer_dependency=layer_dependency, + backward_shared_dependency=shared_dependency, + ) + + assert torch.equal(output, torch.zeros_like(hidden)) + assert captured["moe_native_routed"] is output + assert program_kwargs["complete_backward_device_boundary"] is True + assert program_kwargs["backward_trace_label"] == "glm52_layer_unknown" + assert program_kwargs["backward_layer_dependency"] is layer_dependency + assert program_kwargs["backward_shared_dependency"] is shared_dependency + output.float().sum().backward() + for name in experts.logical_factor_names: + gradient = getattr(experts, name).grad + assert gradient is not None + assert gradient.dtype is torch.float32 + assert torch.equal(gradient, torch.zeros_like(gradient)) + + +def test_native_deepep_empty_source_rank_preserves_combine_autograd_edge( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from xorl.distributed.moe import deepep_native_exact + + block = _empty_block() + block.train_router = False + block.num_experts = 256 + block.experts = SimpleNamespace( + num_experts=256, + deepep_buffer_size_gb=1.0, + deepep_num_sms=24, + ) + block._diagnostic_capture_component = lambda *_args: None + combine_edge = torch.tensor(1.0, dtype=torch.bfloat16, requires_grad=True) + + def fake_program(hidden, _routing, _selected, **_kwargs): + assert hidden.shape == (0, 128) + return hidden + combine_edge + + monkeypatch.setattr(deepep_native_exact, "native_dispatch_runner_combine", fake_program) + hidden = torch.zeros((2, 128), dtype=torch.bfloat16, requires_grad=True) + output = block._native_deepep_routed_local( + hidden, + torch.full((2, 8), 0.125, dtype=torch.float32), + torch.arange(16, dtype=torch.int32).reshape(2, 8), + torch.zeros(2, dtype=torch.bool), + ep_rank=0, + ep_size=16, + ep_group=object(), + ) + + output.float().sum().backward() + + assert combine_edge.grad is not None + assert combine_edge.grad.item() == 0.0 + + +def test_native_shared_boundary_executes_all_tp16_leaves_locally_and_masks_padding() -> None: + block = _empty_block() + shared = Glm52ExactTP16SharedExpertBlockFP8QLoRA(device="meta") + captured = {} + + def forward(self, hidden, *, contributor_ordinal=None, all_contributors=False): + assert all_contributors is True + assert contributor_ordinal is None + captured["hidden"] = hidden + ordinals = torch.arange(16, dtype=torch.float32) - 7.5 + return ordinals[:, None, None].expand(16, *hidden.shape).to(torch.bfloat16).contiguous() + + shared.forward = MethodType(forward, shared) + block.shared_experts = shared + block._diagnostic_capture_component = lambda name, value: captured.setdefault(name, value) + hidden = torch.zeros((3, 6144), dtype=torch.bfloat16) + valid = torch.tensor([True, False, True]) + + actual = block._native_shared_local_fold(hidden, valid) + expected = canonical_moe_fold_fp64_v3(captured["moe_native_shared_down"]) + expected[1].zero_() + + assert captured["hidden"] is hidden + assert captured["moe_native_shared_down"].shape == (16, 3, 6144) + assert torch.equal(actual, expected) + assert captured["moe_native_shared_folded"] is actual diff --git a/tests/models/test_glm52_exact_attention_construction.py b/tests/models/test_glm52_exact_attention_construction.py index 6e79634c..4faf8672 100644 --- a/tests/models/test_glm52_exact_attention_construction.py +++ b/tests/models/test_glm52_exact_attention_construction.py @@ -94,7 +94,7 @@ def test_glm52_exact_attention_component_rejects_nonpositive_rank_or_alpha_befor ("override", "message"), ( ({"_glm52_exact_active_lora_dense_component": False}, "exact active-LoRA dense component"), - ({"_ep_dispatch": "deepep"}, "ep_dispatch='alltoall'"), + ({"_ep_dispatch": "unsupported"}, "ep_dispatch in"), ({"_sparse_mla_enabled": False}, "requires sparse_mla_enabled=true"), ), ) diff --git a/tests/models/test_glm52_exact_lm_head_loss_integration.py b/tests/models/test_glm52_exact_lm_head_loss_integration.py index 3df4f334..51e7e5a5 100644 --- a/tests/models/test_glm52_exact_lm_head_loss_integration.py +++ b/tests/models/test_glm52_exact_lm_head_loss_integration.py @@ -11,7 +11,10 @@ import xorl.models.transformers.glm5.exact_lm_head_qlora as exact_lm_head_impl from xorl.distributed.torch_parallelize import _exact_lm_head_replicated_params from xorl.models.module_utils import get_lm_head_weight -from xorl.models.transformers.glm5.exact_lm_head_qlora import Glm52ExactTP16LmHeadLoraLinear +from xorl.models.transformers.glm5.exact_lm_head_qlora import ( + Glm52ExactTP16LmHeadLoraLinear, + Glm52ExactTP16LmHeadSelectedLogprob, +) from xorl.ops.loss.per_token_ce import compute_per_token_ce from xorl.server.runner.model_runner import ModelRunner from xorl.trainers.training_utils import make_pp_loss_fn @@ -187,6 +190,34 @@ def test_exact_head_fsdp_ignores_only_replicated_a() -> None: _exact_lm_head_replicated_params(lm_head) +def test_exact_head_rejects_frozen_logical_factor_before_local_view_conversion() -> None: + lm_head = _tiny_exact_head() + lm_head._glm52_exact_selected_logprob = Glm52ExactTP16LmHeadSelectedLogprob( + tp_rank=0, + vocab_start=0, + vocab_end=9_680, + padded_vocab_start=0, + padded_vocab_end=9_680, + ) + lm_head.lora_B.requires_grad_(False) + + with pytest.raises(RuntimeError, match="logical factor masters must both be trainable"): + exact_lm_head_impl.glm52_exact_lm_head_per_token_ce( + torch.zeros((1, 4), dtype=torch.bfloat16), + lm_head.weight, + torch.tensor([0], dtype=torch.int64), + lm_head=lm_head, + ignore_index=-100, + ce_mode="bi_fused", + lm_head_fp32=True, + logprob_temperature=1.0, + logprob_top_ks=None, + logprob_top_ps=None, + logprob_min_ps=None, + tp_group=None, + ) + + def test_pp_exact_head_loss_matches_dispatcher_value_and_gradients(monkeypatch: pytest.MonkeyPatch) -> None: def _differentiable_exact(hidden, weight, labels, *, lm_head, ignore_index, **_kwargs): safe_labels = labels.clamp_min(0) diff --git a/tests/models/test_glm52_exact_lm_head_qlora.py b/tests/models/test_glm52_exact_lm_head_qlora.py index bcc9e274..a88867c0 100644 --- a/tests/models/test_glm52_exact_lm_head_qlora.py +++ b/tests/models/test_glm52_exact_lm_head_qlora.py @@ -157,6 +157,28 @@ def _assert_operand_contract_is_official_local_bf16_rank_one_and_stride_exact() token_ids, require_cuda=False, ) + weight.requires_grad_(False) + + detached_A = lora_A.detach() + detached_B = lora_B.detach() + with pytest.raises(RuntimeError, match="factor masters must both be trainable"): + component._validate_operands( + hidden, + weight, + detached_A, + detached_B, + token_ids, + require_cuda=False, + ) + component._validate_operands( + hidden, + weight, + detached_A, + detached_B, + token_ids, + require_cuda=False, + require_factor_grad=False, + ) def _assert_cpu_rejection_happens_before_sglang_import_or_group_use() -> None: diff --git a/tests/models/test_glm52_exact_moe_construction.py b/tests/models/test_glm52_exact_moe_construction.py index 69b4ac9e..32ef7a72 100644 --- a/tests/models/test_glm52_exact_moe_construction.py +++ b/tests/models/test_glm52_exact_moe_construction.py @@ -7,6 +7,7 @@ from torch.distributed._tensor import Replicate, Shard from tests.models.test_glm52_qlora import _meta_model, _official_config +from xorl.models.exact_contract import set_glm52_exact_active_lora from xorl.models.transformers.glm5.exact_lm_head_qlora import ( Glm52ExactTP16LmHeadLoraLinear, Glm52ExactTP16LmHeadSelectedLogprob, @@ -189,6 +190,72 @@ def test_glm52_exact_moe_construction_preserves_complete_global_inventory_and_so } +def test_glm52_exact_qlora_router_training_admits_only_the_75_bf16_gate_weights( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_exact_world16_rank7(monkeypatch) + config = _exact_moe_config() + config._glm52_exact_active_lora_lm_head_component = True + config.train_router = True + model = _meta_model(config) + + inventory = prepare_glm52_block_fp8_qlora(model, config, adapter_rank=1, adapter_alpha=1) + + router_names = {f"model.layers.{layer_idx}.mlp.gate.weight" for layer_idx in range(3, 78)} + trainable = {name: parameter for name, parameter in model.named_parameters() if parameter.requires_grad} + assert set(trainable) == inventory.factor_names | router_names + assert all(trainable[name].dtype is torch.bfloat16 for name in router_names) + assert all(trainable[name].dtype is torch.float32 for name in inventory.factor_names) + + +def test_glm52_exact_moe_construction_admits_deepep_sparse_leaf_transport( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_ep16_rank7(monkeypatch) + config = _exact_moe_config() + config._ep_dispatch = "deepep" + config._deepep_buffer_size_gb = 1.25 + config._deepep_num_sms = 24 + config._deepep_async_combine = True + model = _meta_model(config) + + prepare_glm52_block_fp8_qlora(model, config, adapter_rank=1, adapter_alpha=1) + + routed = model.get_submodule("model.layers.3.mlp.experts") + assert type(routed) is Glm52ExactEP16BlockFP8QLoRARoutedExperts + assert routed.ep_dispatch == "deepep" + assert routed.deepep_buffer_size_gb == pytest.approx(1.25) + assert routed.deepep_num_sms == 24 + assert routed.deepep_async_combine is True + assert routed.expert_adapter_gradient_contract.factor_layout == "gkn_gate_up_down" + + +def test_glm52_exact_moe_construction_admits_native_deepep_split_contract( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_exact_world16_rank7(monkeypatch) + config = _exact_moe_config() + config._ep_dispatch = "deepep" + config._deepep_native_exact = True + config._deepep_buffer_size_gb = 1.5 + config._deepep_num_sms = 24 + config._deepep_async_combine = False + config.train_router = False + set_glm52_exact_active_lora(config, enabled=True) + model = _meta_model(config) + + prepare_glm52_block_fp8_qlora(model, config, adapter_rank=1, adapter_alpha=1) + + block = model.get_submodule("model.layers.3.mlp") + routed = block.experts + assert block.deepep_native_exact is True + assert type(routed) is Glm52ExactEP16BlockFP8QLoRARoutedExperts + assert routed.ep_dispatch == "deepep" + assert routed.deepep_buffer_size_gb == pytest.approx(1.5) + assert routed.deepep_num_sms == 24 + assert routed.deepep_async_combine is False + + def test_glm52_exact_moe_post_ep_layout_preserves_factor_fqns_and_owner_logical_shapes( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -297,7 +364,6 @@ def test_glm52_complete_exact_construction_attaches_only_selected_logprob_lm_hea "requires the exact active-LoRA routed-expert component", ), ({"_sparse_mla_enabled": False}, "requires sparse_mla_enabled=true"), - ({"_ep_dispatch": "deepep"}, "requires ep_dispatch='alltoall'"), ), ) def test_glm52_exact_moe_construction_rejects_incomplete_dependency_flags_before_mutation( diff --git a/tests/models/test_glm52_exact_routed_experts_qlora.py b/tests/models/test_glm52_exact_routed_experts_qlora.py index a4bb552b..0d7b8ee2 100644 --- a/tests/models/test_glm52_exact_routed_experts_qlora.py +++ b/tests/models/test_glm52_exact_routed_experts_qlora.py @@ -195,13 +195,14 @@ def test_routed_bank_topology_remap_and_physical_buffer_policy() -> None: ) with pytest.raises(ValueError, match="MoE-TP1"): Glm52ExactEP16BlockFP8QLoRARoutedExperts(_HIDDEN, _INTERMEDIATE, ep_rank=0, moe_tp_size=2) - with pytest.raises(ValueError, match="DeepEP is not admitted"): - Glm52ExactEP16BlockFP8QLoRARoutedExperts( - _HIDDEN, - _INTERMEDIATE, - ep_rank=0, - ep_dispatch="deepep", - ) + deepep = Glm52ExactEP16BlockFP8QLoRARoutedExperts( + _HIDDEN, + _INTERMEDIATE, + ep_rank=0, + ep_dispatch="deepep", + device="meta", + ) + assert deepep.expert_adapter_gradient_contract.factor_layout == "gkn_gate_up_down" with pytest.raises(ValueError, match=r"in \[0, 15\]"): Glm52ExactEP16BlockFP8QLoRARoutedExperts(_HIDDEN, _INTERMEDIATE, ep_rank=16) module.set_runtime_lora_config(1, 1) @@ -296,6 +297,73 @@ def _assert_post_ep_owner_local_factor_banks_produce_same_views() -> None: assert torch.equal(global_buffers[name], local_buffers[name]), name +def test_sampler_value_requests_fused_combined_leaf(monkeypatch) -> None: + import sys + from types import ModuleType, SimpleNamespace + + from xorl.models.transformers.glm5 import exact_routed_experts_qlora as qlora_module + + module = _module(0) + hidden = torch.ones((2, _HIDDEN), dtype=torch.bfloat16) + routing = torch.tensor([[0.25, 0.5], [0.75, 0.125]], dtype=torch.float32) + local_ids = torch.tensor([[0, -1], [1, 2]], dtype=torch.int32) + fused_leaf = torch.tensor([[0.5] * _HIDDEN, [4.0] * _HIDDEN], dtype=torch.bfloat16) + observed = {} + + monkeypatch.setattr(qlora_module.MoEExperts, "_ensure_sglang_server_args", lambda: None) + + def fake_prepare(*_args, **_kwargs): + return ( + {}, + None, + False, + torch.empty(0, dtype=torch.int32), + torch.empty(0, dtype=torch.int32), + torch.tensor(0, dtype=torch.int32), + ) + + def fake_build_lora_hooks(*_args, **kwargs): + observed["mul_routed_weight"] = kwargs.get("mul_routed_weight", True) + return SimpleNamespace(after_gate_up=None, after_down=None) + + def fake_kernel_sequence(*_args, **kwargs): + observed["no_combine"] = kwargs["no_combine"] + observed["routed_scaling_factor"] = kwargs["routed_scaling_factor"] + return fused_leaf + + fake_fused_moe = ModuleType("sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe") + fake_fused_moe._prepare_fused_moe_run = fake_prepare + fake_fused_moe._fused_moe_kernel_sequence = fake_kernel_sequence + monkeypatch.setitem( + sys.modules, + "sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe", + fake_fused_moe, + ) + + fake_lora_module = ModuleType("sglang.srt.lora.lora_moe_runners") + fake_lora_module.LoRAInfo = lambda **kwargs: SimpleNamespace(**kwargs) + fake_lora_module.build_lora_hooks = fake_build_lora_hooks + monkeypatch.setitem(sys.modules, "sglang.srt.lora.lora_moe_runners", fake_lora_module) + factors = tuple( + getattr(module, name).detach().to(torch.bfloat16).contiguous() for name in module.logical_factor_names + ) + actual, trace = module._sampler_value( + hidden, + routing, + local_ids, + *factors, + routed_scaling_factor=2.0, + capture_trace=False, + ) + assert trace is None + assert observed == { + "mul_routed_weight": True, + "no_combine": False, + "routed_scaling_factor": 2.0, + } + assert actual is fused_leaf + + @pytest.mark.gpu @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") def test_routed_experts_literal_sampler_and_gradient_policy() -> None: diff --git a/tests/models/test_glm52_exact_shared_expert_qlora.py b/tests/models/test_glm52_exact_shared_expert_qlora.py index d6ba82f2..072f458d 100644 --- a/tests/models/test_glm52_exact_shared_expert_qlora.py +++ b/tests/models/test_glm52_exact_shared_expert_qlora.py @@ -11,6 +11,7 @@ from xorl.models.transformers.glm5.exact_shared_expert_qlora import ( GLM52_EXACT_TP16_SHARED_EXPERT_QLORA_CONTRACT_VERSION, Glm52ExactTP16SharedExpertBlockFP8QLoRA, + _Glm52ExactTP16SharedExpertAllContributorsFunction, ) from xorl.ops.fused_silu_and_mul import exact_fp32_silu_and_mul @@ -284,6 +285,87 @@ def _assert_shared_expert_native_base_views_use_output_rows_and_input_columns() assert torch.equal(actual.down_scales[:, 0], down_scales[:, ordinal]) +def test_shared_expert_runtime_contract_fails_before_sglang_kernel_import() -> None: + module = Glm52ExactTP16SharedExpertBlockFP8QLoRA(device="cpu") + with pytest.raises(TypeError, match="contributor_ordinal must be an integer"): + module(torch.zeros(1, 6144, dtype=torch.bfloat16), contributor_ordinal=True) + with pytest.raises(ValueError, match=r"must be in \[0, 16\)"): + module(torch.zeros(1, 6144, dtype=torch.bfloat16), contributor_ordinal=16) + with pytest.raises(TypeError, match="requires BF16 activations"): + module(torch.zeros(1, 6144), contributor_ordinal=0) + with pytest.raises(ValueError, match="input width"): + module(torch.zeros(1, 128, dtype=torch.bfloat16), contributor_ordinal=0) + with pytest.raises(ValueError, match="contiguous sampler-layout"): + module(torch.zeros(6144, 2, dtype=torch.bfloat16).transpose(0, 1), contributor_ordinal=0) + with pytest.raises(RuntimeError, match="requires CUDA"): + module(torch.zeros(1, 6144, dtype=torch.bfloat16), contributor_ordinal=0) + with pytest.raises(RuntimeError, match="cannot run independently"): + module.gate_proj(torch.zeros(1, 6144, dtype=torch.bfloat16)) + with pytest.raises(RuntimeError, match="cannot bypass active LoRA"): + module.gate_proj.forward_partition( + torch.zeros(1, 6144, dtype=torch.bfloat16), + output_range=(0, 128), + ) + + module.tp_size = 8 + with pytest.raises(RuntimeError, match="runtime contract was mutated"): + module(torch.zeros(1, 6144, dtype=torch.bfloat16), contributor_ordinal=0) + module.tp_size = 16 + module.gate_proj.lora_A = nn.Parameter(module.gate_proj.lora_A.to(torch.bfloat16)) + with pytest.raises(TypeError, match="gate_proj.lora_A must remain FP32"): + module(torch.zeros(1, 6144, dtype=torch.bfloat16), contributor_ordinal=0) + + +def test_all_contributor_custom_function_stacks_values_and_accumulates_vjps() -> None: + class FakeSharedRoot: + tp_size = 4 + hidden_size = 3 + + def _exact_forward_value(self, input, *effective, contributor_ordinal): + assert all(factor.dtype is torch.bfloat16 for factor in effective) + ordinal = contributor_ordinal + 1 + return SimpleNamespace( + output=input + ordinal, + gate_up=input + 2 * ordinal, + activated=input + 3 * ordinal, + ) + + def _surrogate_vjp( + self, + input, + *args, + contributor_ordinal, + needs_input_grad, + ): + effective = args[:6] + grad_output = args[-1] + scale = contributor_ordinal + 1 + gradients = [grad_output * scale] + gradients.extend( + torch.full_like(factor, scale, dtype=torch.float32) if needed else None + for factor, needed in zip(effective, needs_input_grad[1:], strict=True) + ) + return tuple(gradients) + + input = torch.zeros((2, 3), dtype=torch.bfloat16, requires_grad=True) + factors = [torch.ones((1, 1), dtype=torch.float32, requires_grad=True) for _ in range(6)] + output = _Glm52ExactTP16SharedExpertAllContributorsFunction.apply( + input, + *factors, + FakeSharedRoot(), + ) + + assert output.shape == (4, 2, 3) + for ordinal in range(4): + assert torch.equal(output[ordinal], torch.full_like(input, ordinal + 1)) + output.backward(torch.ones_like(output)) + assert torch.equal(input.grad, torch.full_like(input, 10)) + for factor in factors: + assert factor.grad is not None + assert factor.grad.dtype is torch.float32 + assert torch.equal(factor.grad, torch.full_like(factor, 10)) + + def _manual_local_vjp( module: Glm52ExactTP16SharedExpertBlockFP8QLoRA, input: torch.Tensor, @@ -498,6 +580,11 @@ def test_official_shared_expert_actual_operands_fold_and_surrogate_vjp() -> None [module(fold_input, contributor_ordinal=rank) for rank in range(16)], dim=0, ) + all_contributor_partials = module(fold_input, all_contributors=True) + assert torch.equal( + all_contributor_partials.view(torch.uint8), + partials.view(torch.uint8), + ) metadata = CanonicalMoEGraphMetadata.build( torch.tensor([0], dtype=torch.int64, device=device), torch.tensor([0], dtype=torch.int64, device=device), diff --git a/tests/models/test_glm52_fullparam_reduced_backward_gate.py b/tests/models/test_glm52_fullparam_reduced_backward_gate.py index a97f921d..1eaf09b0 100644 --- a/tests/models/test_glm52_fullparam_reduced_backward_gate.py +++ b/tests/models/test_glm52_fullparam_reduced_backward_gate.py @@ -77,10 +77,11 @@ def _single_contributor_experts_with_shared( routing_weights: torch.Tensor, selected_experts: torch.Tensor, absolute_positions: torch.Tensor | None = None, + backward_layer_dependency: torch.Tensor | None = None, ): """EP1 projection of the canonical dispatch through the REAL partial seams.""" - del absolute_positions + del absolute_positions, backward_layer_dependency batch_size, seq_len, hidden_dim = hidden_states.shape flat = hidden_states.reshape(-1, hidden_dim) rows = flat.shape[0] diff --git a/tests/models/test_glm52_native_fp8.py b/tests/models/test_glm52_native_fp8.py index 031e90b3..6d1c4ca9 100644 --- a/tests/models/test_glm52_native_fp8.py +++ b/tests/models/test_glm52_native_fp8.py @@ -13,6 +13,7 @@ NativeBlockFP8ExpertPairBuffer, NativeBlockFP8PairBuffer, native_fp8_dense_source_map, + reduce_glm52_no_combine_routes, validate_glm52_native_fp8_config, ) from xorl.ops.block_fp8_native import NativeBlockFP8Linear, unpack_float32_as_fp8 @@ -291,6 +292,45 @@ def test_glm_expert_state_is_frozen_exact_and_scoring_only(): module(hidden.detach(), routing) +def test_glm_no_combine_reducer_keeps_bf16_wire_and_fp32_local_sum(): + routed_values = torch.tensor( + [ + [[1.0, 2.0], [4.0, 8.0], [float("nan"), float("nan")]], + [[0.5, -1.0], [2.0, 3.0], [8.0, -4.0]], + ], + dtype=torch.bfloat16, + ) + routing = torch.tensor([[0.25, 0.5, 1000.0], [0.125, 0.25, 0.5]], dtype=torch.float32) + local_ids = torch.tensor([[0, 1, -1], [2, -1, 3]], dtype=torch.int32) + + actual = reduce_glm52_no_combine_routes( + routed_values, + routing, + local_ids, + routed_scaling_factor=2.0, + ) + safe = torch.where((local_ids >= 0).unsqueeze(-1), routed_values, torch.zeros_like(routed_values)) + weights = torch.where(local_ids >= 0, routing, torch.zeros_like(routing)) + expected = (safe.float() * weights.unsqueeze(-1)).sum(dim=1).mul(2.0).to(torch.bfloat16) + + assert actual.dtype is torch.bfloat16 + assert torch.equal(actual, expected) + assert bool(torch.all(torch.isfinite(actual))) + + +def test_glm_no_combine_reducer_rejects_wider_wire_values(): + routes = torch.zeros((1, 2, 4), dtype=torch.float32) + routing = torch.ones((1, 2), dtype=torch.float32) + local_ids = torch.zeros((1, 2), dtype=torch.int32) + with pytest.raises(TypeError, match="must be BF16"): + reduce_glm52_no_combine_routes( + routes, + routing, + local_ids, + routed_scaling_factor=1.0, + ) + + def test_expert_pair_buffer_fuses_local_gate_up_down_bytes_and_scales(): model = _TinyExpertModel() buffer = NativeBlockFP8ExpertPairBuffer(model, ep_rank=1, ep_size=2, num_experts=4) diff --git a/tests/models/test_glm52_qlora.py b/tests/models/test_glm52_qlora.py index 78d8896b..009d7569 100644 --- a/tests/models/test_glm52_qlora.py +++ b/tests/models/test_glm52_qlora.py @@ -19,7 +19,10 @@ GLM52_QLORA_ROUTED_BANK_COUNT, prepare_glm52_block_fp8_qlora, ) -from xorl.models.transformers.glm5.support import validate_glm5_training_mode +from xorl.models.transformers.glm5.support import ( + validate_glm5_training_mode, + validate_glm52_local_router_inventory, +) from xorl.ops.block_fp8_native import NativeBlockFP8Linear from xorl.qlora.modules.block_fp8_linear import BlockFP8QLoRALinear from xorl.qlora.modules.moe_experts import BlockFP8QLoRAMoeExperts @@ -230,7 +233,7 @@ def test_glm52_qlora_rejects_missing_official_indexer_exclusion_before_adapteriz ("override", "message"), [ ({"_moe_implementation": "eager"}, "moe_implementation='triton'"), - ({"_ep_dispatch": "alltoall"}, "ep_dispatch='deepep'"), + ({"_ep_dispatch": "alltoall"}, r"ep_dispatch in \['deepep'\]"), ({"_glm52_exact_contract": True}, "cannot use the scoring-only exact contract"), ({"_glm52_block_fp8_qlora": False}, "block_fp8_qlora_training=true"), ], @@ -292,22 +295,34 @@ def test_glm5_training_mode_uses_alltoall_only_for_complete_exact_active_lora() moe_hybrid_shared_lora=True, ) - with pytest.raises(ValueError, match="ep_dispatch='deepep'.*requires 'alltoall'"): - validate_glm5_training_mode( - config, - enable_qlora=True, - freeze_router=True, - merge_qkv=True, - block_fp8_qlora_training=True, - quant_format="block_fp8", - quant_group_size=128, - moe_implementation="triton", - ep_dispatch="deepep", - moe_hybrid_shared_lora=True, - ) + validate_glm5_training_mode( + config, + enable_qlora=True, + freeze_router=True, + merge_qkv=True, + block_fp8_qlora_training=True, + quant_format="block_fp8", + quant_group_size=128, + moe_implementation="triton", + ep_dispatch="deepep", + moe_hybrid_shared_lora=True, + ) def test_block_fp8_qlora_scale_storage_covers_partial_edge_tiles() -> None: module = BlockFP8QLoRALinear(6144, 576, r=4, lora_alpha=4, device=torch.device("meta")) assert module.weight_block_scales.shape == (5, 192) + + +def test_router_inventory_allows_dense_only_pipeline_stage_and_checks_sparse_stage() -> None: + class Glm5MoEBlock(nn.Module): + pass + + dense_stage = nn.Sequential(nn.Linear(2, 2)) + sparse_stage = nn.Sequential(Glm5MoEBlock(), Glm5MoEBlock()) + + assert validate_glm52_local_router_inventory(dense_stage, retained_router_count=0) == 0 + assert validate_glm52_local_router_inventory(sparse_stage, retained_router_count=2) == 2 + with pytest.raises(RuntimeError, match="retained=0, expected=2"): + validate_glm52_local_router_inventory(sparse_stage, retained_router_count=0) diff --git a/tests/models/test_lora_merged_forward.py b/tests/models/test_lora_merged_forward.py index 97262d45..0a36fdce 100644 --- a/tests/models/test_lora_merged_forward.py +++ b/tests/models/test_lora_merged_forward.py @@ -7,6 +7,8 @@ sglang postfold serving) live in experiments/k3_tests/lora_path_xengine.py. """ +import sys +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -352,6 +354,7 @@ def _module(self, hybrid=True): mod.down_proj.normal_(std=0.02) for proj in ("gate_proj", "up_proj", "down_proj"): getattr(mod, f"{proj}_lora_B").normal_(std=0.02) + mod.lora_serving_mode = "merged" return mod def _assert_merged_weight_and_cache_policy(self): @@ -420,13 +423,55 @@ def masked_partial(got_hidden, got_routing, got_ids): self._assert_native_ep_no_grad_uses_canonical_fold_and_filter(monkeypatch) + def test_native_deepep_widens_only_bf16_routing_metadata(self, monkeypatch): + import xorl.distributed.moe.deepep_native_exact as native_exact + + mod = self._module() + mod.exact_merged_forward = True + mod.deepep_native_exact = True + mod.ep_dispatch = "deepep" + mod.deepep_async_combine = False + mod.deepep_buffer_size_gb = 2.0 + mod.deepep_num_sms = 20 + hidden = torch.randn(3, H).to(torch.bfloat16) + routing_bf16 = torch.rand(3, 2).to(torch.bfloat16) + selected = torch.randint(0, E, (3, 2)) + expected = torch.randn_like(hidden) + captured = {} + + def fake_program(got_hidden, got_routing, got_selected, **kwargs): + captured.update( + hidden=got_hidden, + routing=got_routing, + selected=got_selected, + kwargs=kwargs, + ) + return expected + + monkeypatch.setattr(native_exact, "native_dispatch_runner_combine", fake_program) + got = mod._ep_forward( + hidden, + routing_bf16, + selected, + SimpleNamespace(ep_group=object()), + ) + + assert got is expected + assert captured["hidden"] is hidden + assert captured["selected"] is selected + assert captured["routing"].dtype is torch.float32 + assert torch.equal( + captured["routing"], + routing_bf16.to(torch.float32), + ) + def _assert_native_ep_no_grad_uses_canonical_fold_and_filter(self, monkeypatch): mod = self._module() mod.exact_merged_forward = True hidden = torch.randn(3, H).to(torch.bfloat16) - routing = torch.rand(3, 2).to(torch.bfloat16) + routing = torch.rand(3, 2).to(torch.float32) local_ids = torch.tensor([[0, -1], [1, -1], [-1, 2]], dtype=torch.int32) - expected = torch.randn_like(hidden) + expected = torch.randn(3, H, dtype=torch.bfloat16) gate_up_f, down_f = mod._merged_weights() captured = {} @@ -462,7 +507,8 @@ def fake_kernel( ): got = mod.sglang_ep_native_routed_partial(hidden, routing, local_ids) - assert got is expected + assert got.dtype is torch.bfloat16 + assert torch.equal(got, expected) assert captured == { "hidden": hidden, "gate_up": gate_up_f, @@ -473,6 +519,166 @@ def fake_kernel( "filter_expert": True, } + def test_native_ep_trainable_path_requests_fused_local_combine(self): + mod = self._module() + mod.exact_merged_forward = True + hidden = torch.randn(3, H).to(torch.bfloat16) + routing = torch.rand(3, 2).to(torch.float32) + local_ids = torch.tensor([[0, -1], [1, -1], [-1, 2]], dtype=torch.int32) + expected = torch.randn_like(hidden) + + with ( + patch("xorl.models.layers.moe.experts.MoEExperts._load_sglang_fused_experts_impl", return_value=object()), + patch( + "xorl.models.layers.moe.experts._SglangFusedExpertsTrainFunction.apply", + return_value=expected, + ) as apply, + ): + got = mod.sglang_ep_native_routed_partial(hidden, routing, local_ids) + + assert got is expected + assert apply.call_args.args[-1] is True + + def test_native_ep_separate_mode_is_distinct_from_merged(self, monkeypatch): + mod = self._module() + mod.exact_merged_forward = True + mod.lora_serving_mode = "separate" + hidden = torch.randn(3, H).to(torch.bfloat16) + routing = torch.rand(3, 2).to(torch.float32) + local_ids = torch.tensor([[0, -1], [1, -1], [-1, 2]], dtype=torch.int32) + expected = torch.randn_like(hidden) + captured = {} + + def hook_value(got_hidden, got_routing, got_ids, *factors): + captured.update( + hidden=got_hidden, + routing=got_routing, + ids=got_ids, + factors=factors, + ) + return expected + + monkeypatch.setattr(mod, "_sglang_native_lora_hook_value", hook_value) + monkeypatch.setattr( + mod, + "_merged_weights", + lambda: pytest.fail("separate mode must not fold LoRA into base weights"), + ) + with torch.no_grad(): + got = mod.sglang_ep_native_routed_partial(hidden, routing, local_ids) + + assert got is expected + assert captured["hidden"] is hidden + assert captured["routing"] is routing + assert captured["ids"] is local_ids + assert len(captured["factors"]) == 6 + assert all(factor.dtype is torch.bfloat16 for factor in captured["factors"]) + + def test_separate_zero_b_still_builds_active_shared_outer_info(self, monkeypatch): + mod = self._module() + mod.lora_serving_mode = "separate" + monkeypatch.setitem( + sys.modules, + "sglang.srt.lora.lora_moe_runners", + SimpleNamespace(LoRAInfo=SimpleNamespace), + ) + with torch.no_grad(): + for projection in ("gate_proj", "up_proj", "down_proj"): + getattr(mod, f"{projection}_lora_B").zero_() + factors = tuple( + value.to(torch.bfloat16) + for projection in ("gate_proj", "up_proj", "down_proj") + for value in mod._active_lora_views(projection) + ) + physical = mod._sglang_native_lora_physical_buffers(*factors) + info = mod._sglang_native_lora_info(3, physical) + + assert physical["gate_up_lora_a_weights"].shape == (1, 1, 2 * R, H) + assert physical["gate_up_lora_b_weights"].shape == (1, E, 2 * I, R) + assert physical["down_lora_a_weights"].shape == (1, E, R, I) + assert physical["down_lora_b_weights"].shape == (1, 1, H, R) + assert torch.count_nonzero(physical["gate_up_lora_b_weights"]) == 0 + assert torch.count_nonzero(physical["down_lora_b_weights"]) == 0 + assert info.has_active_lora is True + assert info.experts_shared_outer_loras is True + assert info.adapter_enabled.tolist() == [1] + + def test_native_ep_separate_trainable_path_uses_value_surrogate_boundary(self): + mod = self._module() + mod.exact_merged_forward = True + mod.lora_serving_mode = "separate" + hidden = torch.randn(3, H).to(torch.bfloat16) + routing = torch.rand(3, 2).to(torch.float32) + local_ids = torch.tensor([[0, -1], [1, -1], [-1, 2]], dtype=torch.int32) + expected = torch.randn_like(hidden) + + with patch( + "xorl.models.layers.moe.lora._SglangNativeLoRAHooksTrainFunction.apply", + return_value=expected, + ) as apply: + got = mod.sglang_ep_native_routed_partial(hidden, routing, local_ids) + + assert got is expected + assert apply.call_args.args[:3] == (hidden, routing, local_ids) + assert apply.call_args.args[-1] is mod + + def test_separate_surrogate_uses_post_shard_physical_expert_count(self): + mod = self._module() + # Qwen constructs a global expert bank and ParallelPlan/FSDP shards + # its tensors afterwards. Preserve the stale construction attribute + # while presenting the two-expert physical tensors seen by each rank. + local_experts = E // 2 + with torch.no_grad(): + mod.gate_up_proj = torch.nn.Parameter( + mod.gate_up_proj[:local_experts].detach().clone(), requires_grad=False + ) + mod.down_proj = torch.nn.Parameter(mod.down_proj[:local_experts].detach().clone(), requires_grad=False) + for name in ("gate_proj_lora_B", "up_proj_lora_B", "down_proj_lora_A"): + value = getattr(mod, name) + setattr( + mod, + name, + torch.nn.Parameter(value[:local_experts].detach().clone()), + ) + assert mod.num_local_experts == E + assert mod.gate_up_proj.shape[0] == local_experts + + hidden = torch.randn(3, H, dtype=torch.bfloat16, requires_grad=True) + routing = torch.rand(3, 2, dtype=torch.float32, requires_grad=True) + local_ids = torch.tensor([[0, -1], [1, -1], [-1, 0]], dtype=torch.int32) + factors = tuple( + value + for projection in ("gate_proj", "up_proj", "down_proj") + for value in mod._active_lora_views(projection) + ) + needs_input_grad = (True, True, False, True, True, True, True, True, True, False) + + def fake_grad(_output, inputs, **_kwargs): + return tuple(torch.zeros_like(value) for value in inputs) + + with ( + patch( + "xorl.models.layers.moe.experts.MoEExperts._load_sglang_fused_experts_impl", + return_value=object(), + ), + patch( + "xorl.models.layers.moe.experts._SglangFusedExpertsTrainFunction.apply", + return_value=torch.zeros_like(hidden, requires_grad=True), + ) as apply, + patch("torch.autograd.grad", side_effect=fake_grad), + ): + mod._sglang_native_lora_hook_surrogate_vjp( + hidden, + routing, + local_ids, + *factors, + grad_output=torch.ones_like(hidden), + needs_input_grad=needs_input_grad, + ) + + assert apply.call_args.args[9] == local_experts + assert apply.call_args.args[11] is True + class TestTrunkWrapComposition: def _model(self): diff --git a/tests/models/test_moe_ep_native_combine.py b/tests/models/test_moe_ep_native_combine.py index e5c423d4..357e485b 100644 --- a/tests/models/test_moe_ep_native_combine.py +++ b/tests/models/test_moe_ep_native_combine.py @@ -122,7 +122,7 @@ def unexpected_transport(*_args): assert not called -def _qwen_block(*, exact: bool = False): +def _qwen_block(*, exact: bool = False, native_deepep: bool = False): from transformers import PretrainedConfig # noqa: PLC0415 from xorl.models.transformers.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeSparseMoeBlock # noqa: PLC0415 @@ -137,6 +137,8 @@ def _qwen_block(*, exact: bool = False): shared_expert_intermediate_size=24, train_router=False, _qwen35_exact_contract=exact, + _deepep_native_exact=native_deepep, + _ep_dispatch="deepep" if native_deepep else "alltoall", ) return Qwen3_5MoeSparseMoeBlock(cfg, moe_implementation="eager", layer_idx=0).to(torch.bfloat16) @@ -157,6 +159,16 @@ def test_exact_native_combine_is_structural(): assert blk.router._exact_batch_invariant +def test_qwen35_native_deepep_selects_shared_transport_and_thin_shared_join(): + blk = _qwen_block(exact=True, native_deepep=True) + + assert blk.deepep_native_exact + assert blk.experts.deepep_native_exact + assert blk.experts.ep_dispatch == "deepep" + assert not blk._native_ep_combine + assert not blk.supports_routing_replay() + + def test_native_routed_partial_enters_through_module_call(monkeypatch): """The EP serving-kernel lane must run inside FSDP's pre-forward hooks.""" blk = _qwen_block() @@ -211,6 +223,106 @@ def fake_reduce_scatter(out, grad, op=None, group=None): assert torch.equal(x.grad, torch.full_like(x, 2.0)) +def test_token_gather_orders_backward_before_dependency_producer(monkeypatch): + """The shared c10d branch must queue before its routed DeepEP sibling.""" + import xorl.models.layers.moe.ep_native_combine as combine # noqa: PLC0415 + + events = [] + + class RoutedBoundary(torch.autograd.Function): + @staticmethod + def forward(ctx, value): + return value.clone() + + @staticmethod + def backward(ctx, grad_output): + events.append("routed") + return grad_output + + monkeypatch.setattr(combine.dist, "get_world_size", lambda _group: 1) + monkeypatch.setattr( + combine.dist, + "all_gather_into_tensor", + lambda output, local, group=None: output.copy_(local), + ) + + def fake_reduce_scatter(output, grad, op=None, group=None): + del op, group + events.append("shared") + output.copy_(grad) + + monkeypatch.setattr(combine.dist, "reduce_scatter_tensor", fake_reduce_scatter) + + x = torch.tensor([[1.0, 2.0]], requires_grad=True) + routed = RoutedBoundary.apply(x) + gathered = gather_tokens_for_ep_combine( + x, + group=None, + padded_rows=1, + backward_dependency=routed, + ) + (gathered.sum() + routed.sum()).backward() + + assert events == ["shared", "routed"] + assert torch.equal(x.grad, torch.full_like(x, 2.0)) + + +def test_shared_then_routed_dependency_holds_transformer_residual(monkeypatch): + """Do not release an earlier layer through the residual bypass.""" + import xorl.models.layers.moe.ep_native_combine as combine # noqa: PLC0415 + + events = [] + + class ResidualBoundary(torch.autograd.Function): + @staticmethod + def forward(ctx, value): + return value.clone() + + @staticmethod + def backward(ctx, grad_output): + events.append("residual") + return grad_output + + class RoutedBoundary(torch.autograd.Function): + @staticmethod + def forward(ctx, value, backward_layer_dependency): + del backward_layer_dependency + return value.clone() + + @staticmethod + def backward(ctx, grad_output): + events.append("routed") + return grad_output, None + + monkeypatch.setattr(combine.dist, "get_world_size", lambda _group: 1) + monkeypatch.setattr( + combine.dist, + "all_gather_into_tensor", + lambda output, local, group=None: output.copy_(local), + ) + + def fake_reduce_scatter(output, grad, op=None, group=None): + del op, group + events.append("shared") + output.copy_(grad) + + monkeypatch.setattr(combine.dist, "reduce_scatter_tensor", fake_reduce_scatter) + + x = torch.tensor([[1.0, 2.0]], requires_grad=True) + residual = ResidualBoundary.apply(x) + routed = RoutedBoundary.apply(x, residual) + gathered = gather_tokens_for_ep_combine( + x, + group=None, + padded_rows=1, + backward_dependency=routed, + ) + (gathered.sum() + routed.sum() + residual.sum()).backward() + + assert events == ["shared", "routed", "residual"] + assert torch.equal(x.grad, torch.full_like(x, 3.0)) + + def test_variable_row_id_gather_uses_invalid_padding(monkeypatch): import xorl.models.layers.moe.ep_native_combine as combine # noqa: PLC0415 diff --git a/tests/models/test_qwen3_5_trunk_wrap.py b/tests/models/test_qwen3_5_trunk_wrap.py index bf101e8c..5d215414 100644 --- a/tests/models/test_qwen3_5_trunk_wrap.py +++ b/tests/models/test_qwen3_5_trunk_wrap.py @@ -11,6 +11,8 @@ router gate (contracted separately by the exact model program) and lm_head/embed. """ +from types import SimpleNamespace + import pytest import torch @@ -82,6 +84,39 @@ def test_exact_qwen_hook_enables_merged_lora_before_trunk_wrap(): set_trunk_linear_contract(False) +def test_exact_qwen_hook_preserves_native_deepep_transport_ownership(monkeypatch): + class NativeBlock(torch.nn.Module): + def __init__(self): + super().__init__() + self._native_ep_combine = True + self.deepep_native_exact = True + + class ResolvedNorm(torch.nn.Module): + rmsnorm_family = "v2" + + model = torch.nn.Module() + model.config = SimpleNamespace( + model_type="xorl_qwen3_5_moe", + _qwen35_rmsnorm_family="v2", + _rmsnorm_mode="sglang_fused", + _deepep_native_exact=True, + ) + model.block = NativeBlock() + model.norm = ResolvedNorm() + model.q_proj = torch.nn.Linear(8, 8, bias=False, dtype=torch.bfloat16) + + monkeypatch.setattr( + "xorl.distributed.parallel_state.get_parallel_state", + lambda: SimpleNamespace(ep_enabled=True, ep_size=8), + ) + try: + _apply_qwen35_gdn_exact(model) + assert model.block.deepep_native_exact + assert not model.block._native_ep_combine + finally: + set_trunk_linear_contract(False) + + def test_qwen3_5_hybrid_trunk_wrap_selection(): model = _build() try: @@ -111,11 +146,25 @@ def _is_wrapped(module): return getattr(module, "_xorl_bi_trunk_wrapped", False) # Full-attention projections wrap. - assert all(_is_wrapped(m) for m in (full_attn.q_proj, full_attn.k_proj, full_attn.v_proj, full_attn.o_proj)) + assert all( + _is_wrapped(m) + for m in ( + full_attn.q_proj, + full_attn.k_proj, + full_attn.v_proj, + full_attn.o_proj, + ) + ) # Linear-attention q/k/v/o_proj ALSO wrap (same leaf names) — audit fact, # not a contract guarantee: the GDN kernel chain is uncontracted. assert all( - _is_wrapped(m) for m in (linear_attn.q_proj, linear_attn.k_proj, linear_attn.v_proj, linear_attn.o_proj) + _is_wrapped(m) + for m in ( + linear_attn.q_proj, + linear_attn.k_proj, + linear_attn.v_proj, + linear_attn.o_proj, + ) ) # GDN a/b/g projections and short convolutions are silently skipped. assert not _is_wrapped(linear_attn.a_proj) @@ -141,11 +190,15 @@ def _is_wrapped(module): @requires_cuda @pytest.mark.gpu -def test_qwen3_5_full_attn_forward_runs_under_trunk_wrap(): +def test_qwen3_5_full_attn_forward_runs_under_trunk_wrap(monkeypatch): """Wrapped full-attention + shared-expert + dense projections must run the bf16 contract GEMM end-to-end (the runtime guard raises on any non-bf16 operand).""" torch.manual_seed(1) + # This test isolates the trunk-linear contract. The local fused-expert + # kernel has its own parity tests and can leave asynchronous CUDA failures + # indistinguishable from a trunk projection failure here. + monkeypatch.setenv("XORL_MOE_SGLANG_FUSED_EXPERTS", "0") config = _hybrid_config(layer_types=["full_attention", "full_attention"], _moe_implementation="eager") model = Qwen3_5MoeForCausalLM(config).to(device="cuda", dtype=torch.bfloat16).eval() try: diff --git a/tests/models/test_rmsnorm_family_contract.py b/tests/models/test_rmsnorm_family_contract.py index e4a1e2a4..175686ef 100644 --- a/tests/models/test_rmsnorm_family_contract.py +++ b/tests/models/test_rmsnorm_family_contract.py @@ -220,6 +220,48 @@ def forward(self, hidden_states, *args, **kwargs): assert final_norm.family_values == [None] +def test_shared_attention_exposes_cold_path_component_diagnostics(): + cfg = Qwen3Config( + hidden_size=64, + intermediate_size=128, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + num_hidden_layers=1, + _attn_implementation="eager", + ) + attn = MultiHeadAttention(cfg, layer_idx=0) + captured = {} + attn._diagnostic_capture_component = lambda name, value: captured.setdefault(name, value.detach().clone()) + + hidden = torch.randn(1, 3, 64) + cos = torch.ones(1, 3, 16) + sin = torch.zeros(1, 3, 16) + q, k, v = attn._project_qkv(hidden, (cos, sin)) + output = attn._project_output(torch.randn(1, 3, 4, 16)) + + assert { + "attention_input", + "qkv", + "q_pre_qk_norm", + "k_pre_qk_norm", + "q_post_qk_norm", + "k_post_qk_norm", + "v", + "rope_cos", + "rope_sin", + "q", + "k", + "attn_output", + "o_proj_output", + } == set(captured) + torch.testing.assert_close(captured["attention_input"], hidden) + torch.testing.assert_close(captured["q"], q) + torch.testing.assert_close(captured["k"], k) + torch.testing.assert_close(captured["v"], v) + torch.testing.assert_close(captured["o_proj_output"], output) + + # --------------------------------------------------------------------------- # # Loud tripwire for undeclared parity-lane calls. # --------------------------------------------------------------------------- # diff --git a/tests/models/test_topk_router.py b/tests/models/test_topk_router.py index d536893e..34e3c33f 100644 --- a/tests/models/test_topk_router.py +++ b/tests/models/test_topk_router.py @@ -21,6 +21,25 @@ from xorl.models.layers.moe.router import TopKRouter +def test_native_exact_router_retains_fp32_metadata_after_fixed_renorm(): + router = TopKRouter( + num_experts=4, + top_k=2, + norm_topk_prob=True, + exact_batch_invariant=True, + exact_weights_fp32=True, + ) + logits = torch.tensor([[1.0, 3.0, 2.0, -1.0]], dtype=torch.float32) + + weights, ids = router(logits, torch.bfloat16) + + assert weights.dtype is torch.float32 + assert ids.tolist() == [[1, 2]] + expected_scores = torch.softmax(logits, dim=1).gather(1, ids) + expected_denom = expected_scores[:, 0] + expected_scores[:, 1] + assert torch.equal(weights, expected_scores / expected_denom.unsqueeze(-1)) + + pytestmark = pytest.mark.cpu diff --git a/tests/ops/dsv4/test_exact_attention.py b/tests/ops/dsv4/test_exact_attention.py index 7e043b72..0eb51280 100644 --- a/tests/ops/dsv4/test_exact_attention.py +++ b/tests/ops/dsv4/test_exact_attention.py @@ -18,7 +18,12 @@ def _exact_attention_module(): def _cache_state(num_tokens: int, ratio: int): exact_attention = _exact_attention_module() state = exact_attention.Dsv4DecodeCarryState() - state.kvcache = exact_attention._ensure_paged_kvcache(None, num_tokens, 128, torch.device("cpu")) + state.kvcache = exact_attention._ensure_paged_kvcache( + None, + num_tokens, + exact_attention._SERVING_SWA_PAGE_SIZE, + torch.device("cpu"), + ) state.num_tokens = num_tokens if ratio: state.num_compressed = num_tokens // ratio @@ -31,16 +36,52 @@ def _cache_state(num_tokens: int, ratio: int): return state -def _install_fake_flash_mla(monkeypatch, implementation): +def _install_fake_flash_mla(monkeypatch, implementation, scheduler=None): package = types.ModuleType("sgl_kernel") module = types.ModuleType("sgl_kernel.flash_mla") module.flash_mla_with_kvcache = implementation - module.get_mla_metadata = lambda: (object(), None) + module.get_mla_metadata = lambda: ( + object() if scheduler is None else scheduler, + None, + ) package.flash_mla = module monkeypatch.setitem(sys.modules, "sgl_kernel", package) monkeypatch.setitem(sys.modules, "sgl_kernel.flash_mla", module) +@pytest.mark.cpu +@pytest.mark.parametrize( + ("page_size", "expected_pages"), + [(256, 4), (64, 33), (2, 33)], +) +def test_serving_cache_extent_preserves_qualified_sampler_pool_shape(page_size, expected_pages): + exact_attention = _exact_attention_module() + cache = exact_attention._ensure_paged_kvcache(None, 1, page_size, torch.device("cpu")) + assert cache.shape == (expected_pages, exact_attention._flashmla_page_bytes(page_size)) + + +@pytest.mark.cpu +def test_operator_capture_extracts_split_plane_flashmla_rows(): + exact_attention = _exact_attention_module() + page_size = 64 + storage = exact_attention._ensure_paged_kvcache(None, 1, page_size, torch.device("cpu")) + reference = page_size + 5 + storage[1, 5 * 576 : 6 * 576] = 17 + storage[1, page_size * 576 + 5 * 8 : page_size * 576 + 6 * 8] = 29 + kernel_view = exact_attention._paged_cache_kernel_view(storage, page_size) + + references, rows = exact_attention._referenced_cache_rows( + kernel_view, + torch.tensor([[[reference]]], dtype=torch.int32), + torch.tensor([1], dtype=torch.int32), + ) + + assert references.tolist() == [reference] + assert rows.shape == (1, 584) + assert (rows[0, :576] == 17).all() + assert (rows[0, 576:] == 29).all() + + @pytest.mark.cpu @pytest.mark.parametrize( ("ratio", "num_tokens", "positions", "expected_blocks"), @@ -89,7 +130,13 @@ def fake_flash_mla_with_kvcache(**kwargs): assert call["q"].shape == (1, 1, 64, 512) assert torch.equal(call["q"], q[:, row : row + 1]) swa_length = min(position + 1, 128) - expected_swa = list(range(position, position - swa_length, -1)) + expected_swa = list( + range( + exact_attention._SERVING_SWA_PAGE_SIZE + position, + exact_attention._SERVING_SWA_PAGE_SIZE + position - swa_length, + -1, + ) + ) actual_swa = call["indices"][0, 0] assert actual_swa[:swa_length].tolist() == expected_swa assert (actual_swa[swa_length:] == -1).all() @@ -98,13 +145,76 @@ def fake_flash_mla_with_kvcache(**kwargs): blocks = expected_blocks[row] actual_extra = call["extra_indices_in_kvcache"][0, 0] assert call["extra_indices_in_kvcache"].shape == (1, 1, 512 if ratio == 4 else 64) - assert actual_extra[:blocks].tolist() == list(range(blocks)) + extra_page_size = 256 // ratio + assert actual_extra[:blocks].tolist() == list(range(extra_page_size, extra_page_size + blocks)) assert (actual_extra[blocks:] == -1).all() assert call["extra_topk_length"].tolist() == [max(blocks, 1)] else: assert "extra_indices_in_kvcache" not in call +@pytest.mark.cpu +def test_exact_trainer_rejects_non_fixed_k2_flashmla_scheduler(monkeypatch): + exact_attention = _exact_attention_module() + + def fake_flash_mla_with_kvcache(**kwargs): + return kwargs["q"].clone(), torch.empty(0) + + scheduler = types.SimpleNamespace(num_splits=torch.tensor([0, 3])) + _install_fake_flash_mla( + monkeypatch, + fake_flash_mla_with_kvcache, + scheduler=scheduler, + ) + state = _cache_state(65, 4) + + with pytest.raises(RuntimeError, match="fixed-K2"): + exact_attention._serving_decode_attention( + torch.zeros((1, 1, 64, 512), dtype=torch.bfloat16), + state, + 64, + 4, + torch.zeros(64), + 512**-0.5, + ) + + +@pytest.mark.cpu +def test_operator_capture_records_semantic_cache_rows(tmp_path, monkeypatch): + exact_attention = _exact_attention_module() + + def fake_flash_mla_with_kvcache(**kwargs): + return kwargs["q"].clone(), torch.empty(0) + + _install_fake_flash_mla(monkeypatch, fake_flash_mla_with_kvcache) + monkeypatch.setenv("XORL_DSV4_TRAINER_OPERATOR_CAPTURE_DIR", str(tmp_path)) + monkeypatch.setenv("XORL_DSV4_OPERATOR_CAPTURE_LAYER", "2") + monkeypatch.setenv("XORL_DSV4_OPERATOR_CAPTURE_POSITIONS", "63,64") + exact_attention._CAPTURED_OPERATOR_INPUTS.clear() + state = _cache_state(65, 4) + q = torch.zeros((1, 1, 64, 512), dtype=torch.bfloat16) + + with exact_attention.exact_attention_layer(2): + exact_attention._serving_decode_attention( + q, + state, + 64, + 4, + torch.zeros(64), + 512**-0.5, + ) + + [capture_path] = list(tmp_path.glob("*.pt")) + payload = torch.load(capture_path, map_location="cpu", weights_only=True) + assert payload["schema"] == "xorl.dsv4_trainer_flashmla_operator_capture.v2" + assert payload["layer"] == 2 + assert payload["position"] == 64 + assert payload["raw_topk_length"].tolist() == [65] + assert payload["raw_referenced_cache_rows"].shape[0] == 65 + assert payload["extra_topk_length"].tolist() == [16] + assert payload["extra_referenced_cache_rows"].shape[0] == 16 + + @pytest.mark.cpu @pytest.mark.parametrize( ("ratio", "num_tokens", "position"), @@ -143,21 +253,27 @@ def fake_flash_mla_with_kvcache(**kwargs): _install_fake_flash_mla(monkeypatch, fake_flash_mla_with_kvcache) state = _cache_state(num_tokens, ratio) - raw = exact_attention._paged_cache_kernel_view(state.kvcache, 128) + raw_page_size = exact_attention._SERVING_SWA_PAGE_SIZE + assert raw_page_size == 256 + raw = exact_attention._paged_cache_kernel_view(state.kvcache, raw_page_size) + assert raw.dtype == torch.float8_e4m3fn + raw_storage = raw.view(torch.uint8) raw_values = torch.arange(1, num_tokens + 1, dtype=torch.int64).remainder(251).to(torch.uint8) - raw[ - torch.arange(num_tokens) // 128, - torch.arange(num_tokens) % 128, + raw_storage[ + (raw_page_size + torch.arange(num_tokens)) // raw_page_size, + (raw_page_size + torch.arange(num_tokens)) % raw_page_size, 0, 0, ] = raw_values if ratio: extra_page_size = 256 // ratio extra = exact_attention._paged_cache_kernel_view(state.compressed_kvcache, extra_page_size) + assert extra.dtype == torch.float8_e4m3fn + extra_storage = extra.view(torch.uint8) extra_locs = torch.arange(state.num_compressed) - extra[ - extra_locs // extra_page_size, - extra_locs % extra_page_size, + extra_storage[ + (extra_page_size + extra_locs) // extra_page_size, + (extra_page_size + extra_locs) % extra_page_size, 0, 0, ] = torch.arange(1, state.num_compressed + 1, dtype=torch.uint8) @@ -165,11 +281,13 @@ def fake_flash_mla_with_kvcache(**kwargs): q = torch.zeros((1, 1, 64, 512), dtype=torch.bfloat16) before = exact_attention._serving_decode_attention(q, state, position, ratio, torch.zeros(64), 512**-0.5) future_raw = torch.arange(position + 1, num_tokens) - raw[future_raw // 128, future_raw % 128, 0, 0] = 251 + future_raw = raw_page_size + future_raw + raw_storage[future_raw // raw_page_size, future_raw % raw_page_size, 0, 0] = 251 if ratio: first_future_block = (position + 1) // ratio future_extra = torch.arange(first_future_block, state.num_compressed) - extra[ + future_extra = extra_page_size + future_extra + extra_storage[ future_extra // extra_page_size, future_extra % extra_page_size, 0, @@ -185,8 +303,8 @@ def fake_flash_mla_with_kvcache(**kwargs): ) assert torch.equal(before, after_future_mutation) - page, slot = divmod(position, 128) - raw[page, slot, 0, 0] = (int(raw[page, slot, 0, 0]) + 17) % 251 + page, slot = divmod(raw_page_size + position, raw_page_size) + raw_storage[page, slot, 0, 0] = (int(raw_storage[page, slot, 0, 0]) + 17) % 251 after_visible_mutation = exact_attention._serving_decode_attention( q, state, @@ -299,7 +417,7 @@ def fake_attention(q, state, positions, ratio, _sink, _scale): @pytest.mark.gpu @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("ratio", [0, 4, 128]) -def test_full_sequence_surrogate_backward_is_finite_and_reaches_prior_kv(monkeypatch, ratio): +def test_full_sequence_custom_vjp_is_finite_and_reaches_prior_kv(monkeypatch, ratio): exact_attention = _exact_attention_module() def fake_raw_store(kv, _weight, _freqs, _eps, state, _offset, *, dequantize): diff --git a/tests/optim/test_ep_optimizer_groups.py b/tests/optim/test_ep_optimizer_groups.py new file mode 100644 index 00000000..57b39275 --- /dev/null +++ b/tests/optim/test_ep_optimizer_groups.py @@ -0,0 +1,72 @@ +"""Contracts for EP optimizer parameter-group construction.""" + +import pytest +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed.tensor import DeviceMesh, DTensor, Shard + +import xorl.optim.optimizer as optimizer_module + + +pytestmark = pytest.mark.cpu + + +def test_mixed_dtensor_parameters_are_split_for_fused_optimizer(monkeypatch): + class FakeDTensor: + pass + + dtensor = FakeDTensor() + local = object() + monkeypatch.setattr(optimizer_module, "DTensor", FakeDTensor) + + groups = optimizer_module._split_dtensor_parameter_groups([local, dtensor]) + + assert groups == [[dtensor], [local]] + + +def test_homogeneous_parameters_preserve_one_group(monkeypatch): + class FakeDTensor: + pass + + first = object() + second = object() + monkeypatch.setattr(optimizer_module, "DTensor", FakeDTensor) + + assert optimizer_module._split_dtensor_parameter_groups([first, second]) == [[first, second]] + + +def test_split_groups_execute_real_fused_adamw_with_dtensor(tmp_path): + initialized_here = False + if not dist.is_initialized(): + dist.init_process_group( + "gloo", + init_method=f"file://{tmp_path / 'single_rank_pg'}", + rank=0, + world_size=1, + ) + initialized_here = True + elif dist.get_world_size() != 1: + pytest.skip("single-rank DTensor optimizer smoke test requires a single-rank process group") + + try: + mesh = DeviceMesh("cpu", [0], mesh_dim_names=("dp",)) + local = nn.Parameter(torch.ones(2, 3)) + distributed = nn.Parameter(DTensor.from_local(torch.ones(2, 3), mesh, [Shard(0)], run_check=False)) + local.grad = torch.full_like(local, 0.5) + distributed.grad = DTensor.from_local(torch.full((2, 3), 0.5), mesh, [Shard(0)], run_check=False) + + split = optimizer_module._split_dtensor_parameter_groups([local, distributed]) + optimizer = torch.optim.AdamW( + [{"params": parameters} for parameters in split], + lr=0.1, + fused=True, + foreach=False, + ) + optimizer.step() + + assert torch.allclose(local, torch.full_like(local, 0.899)) + assert torch.allclose(distributed.to_local(), torch.full((2, 3), 0.899)) + finally: + if initialized_here: + dist.destroy_process_group() diff --git a/tests/server/api_server/test_checkpoint_paths.py b/tests/server/api_server/test_checkpoint_paths.py index d3ac3af6..6ed8fdb2 100644 --- a/tests/server/api_server/test_checkpoint_paths.py +++ b/tests/server/api_server/test_checkpoint_paths.py @@ -362,17 +362,16 @@ def _assert_sampler_listing_deletion_and_adapter_tracking(self): assert self.server.loaded_sampling_loras == {} self._assert_sampling_adapter_reconciliation_and_session_tracking_policy() - self._assert_save_weights_for_sampler_uses_normalized_lora_session_spec() + self._assert_save_weights_for_sampler_respects_lora_serving_mode() - def _assert_save_weights_for_sampler_uses_normalized_lora_session_spec(self): - """Normalized session specs should still export adapter-only sampler weights.""" + def _assert_save_weights_for_sampler_respects_lora_serving_mode(self): + """Sampler publication must use the explicitly selected LoRA contract.""" self.server._running = True self.server.orchestrator_client = MagicMock() self.server.orchestrator_client.send_request = AsyncMock(return_value=AsyncMock()) - self.server.model_configs["adapter-run"] = { + common_config = { "base_model": "Qwen/Qwen3-8B", "is_lora": True, - "lora_config": {"lora_rank": 8, "lora_alpha": 16}, "optimizer_config": { "type": "signsgd", "learning_rate": 2e-4, @@ -383,6 +382,14 @@ def _assert_save_weights_for_sampler_uses_normalized_lora_session_spec(self): "optimizer_kwargs": {}, }, } + self.server.model_configs["adapter-run"] = { + **common_config, + "lora_config": { + "lora_rank": 8, + "lora_alpha": 16, + "lora_serving_mode": "separate", + }, + } mock_output = MagicMock() mock_output.outputs = [{"lora_path": os.path.join(self.temp_dir, "sampler_weights", "adapter-export")}] @@ -399,6 +406,24 @@ def _assert_save_weights_for_sampler_uses_normalized_lora_session_spec(self): assert request.payload.lora_path.endswith("sampler_weights/adapter-export") assert response.path == "xorl://adapter-run/sampler_weights/adapter-export" + self.server.orchestrator_client.send_request.reset_mock() + self.server.model_configs["merged-run"] = { + **common_config, + "lora_config": { + "lora_rank": 8, + "lora_alpha": 16, + "lora_serving_mode": "merged", + }, + } + with pytest.raises(HTTPException, match="does not materialize a folded") as exc_info: + asyncio.run( + self.server.save_weights_for_sampler( + SaveWeightsForSamplerRequest(model_id="merged-run", name="merged-export") + ) + ) + assert exc_info.value.status_code == 500 + self.server.orchestrator_client.send_request.assert_not_awaited() + def _assert_sampling_adapter_reconciliation_and_session_tracking_policy(self): """Stale tracked adapters should not force a bogus unload before loading a fresh adapter.""" fresh_path = os.path.join(self.temp_dir, "sampler_weights", "fresh-001") diff --git a/tests/server/api_server/test_training_ops.py b/tests/server/api_server/test_training_ops.py index 97dc0aab..8867a87a 100644 --- a/tests/server/api_server/test_training_ops.py +++ b/tests/server/api_server/test_training_ops.py @@ -124,6 +124,8 @@ async def _wait_for_response(self, response_future, request_id, timeout, timeout "executor_total_s": 0.92, "forward_compute_time": 0.35, "backward_compute_time": 0.45, + "dsv4_grad_norm": 3.25, + "dsv4_router_grad_present_count": 0, } ] ) @@ -151,3 +153,5 @@ async def _wait_for_response(self, response_future, request_id, timeout, timeout assert response.metrics["executor_total_s"] == pytest.approx(0.92) assert response.metrics["backward_compute_time"] == pytest.approx(0.45) assert response.metrics["forward_compute_time"] == pytest.approx(0.35) + assert response.metrics["dsv4_grad_norm:mean"] == pytest.approx(3.25) + assert response.metrics["dsv4_router_grad_present_count:mean"] == 0 diff --git a/tests/server/orchestrator/test_packing.py b/tests/server/orchestrator/test_packing.py index ae315eab..13981931 100644 --- a/tests/server/orchestrator/test_packing.py +++ b/tests/server/orchestrator/test_packing.py @@ -288,6 +288,113 @@ def test_packing_disabled_derives_sampler_boundaries_before_advantage_masking(): assert batches[0]["target_tokens"] == [[IGNORE_INDEX] * 4] +@pytest.mark.parametrize("enable_packing", [False, True]) +@pytest.mark.parametrize("mask_name", ["weights", "advantages"]) +def test_sampler_boundary_uses_action_mask_when_targets_are_ordinary_ids(enable_packing, mask_name): + packer = SequentialPacker( + enable_packing=enable_packing, + log_stats=False, + pad_to_multiple_of=1, + ) + datum = { + "input_ids": [11, 12, 13, 21], + "target_tokens": [12, 13, 21, 22], + mask_name: [0.0, 0.0, 1.0, 1.0], + } + + batch = packer.pack([datum], max_seq_len=16, request_id=f"ordinary-targets-{mask_name}")[0] + + assert batch["sampler_prefill_lengths"] == [3] + assert batch["labels"] == [[IGNORE_INDEX, IGNORE_INDEX, 21, 22]] + + +@pytest.mark.parametrize("enable_packing", [False, True]) +def test_sampler_boundary_rejects_ambiguous_all_zero_advantages(enable_packing): + packer = SequentialPacker( + enable_packing=enable_packing, + log_stats=False, + pad_to_multiple_of=1, + ) + datum = { + "input_ids": [11, 12, 13, 21], + "target_tokens": [12, 13, 21, 22], + "advantages": [0.0, 0.0, 0.0, 0.0], + } + + with pytest.raises(ValueError, match="cannot infer the sampler prefill boundary"): + packer.pack([datum], max_seq_len=16, request_id="ambiguous-boundary") + + +def test_packing_disabled_emits_full_boundary_for_prompt_only_shifted_request(): + packer = SequentialPacker(enable_packing=False, log_stats=False, pad_to_multiple_of=1) + + batch = packer.pack( + [{"input_ids": [11, 12, 13], "target_tokens": [IGNORE_INDEX] * 3}], + max_seq_len=16, + request_id="prompt-only-boundary", + )[0] + + assert batch["sampler_prefill_lengths"] == [3] + + +def test_packing_enabled_preserves_one_sampler_boundary_per_request(): + packer = SequentialPacker(enable_packing=True, log_stats=False, pad_to_multiple_of=8) + data = [ + { + "model_input": {"input_ids": [11, 12, 13, 21]}, + "loss_fn_inputs": { + "target_tokens": [IGNORE_INDEX, IGNORE_INDEX, 21, 22], + "advantages": [0.0, 0.0, 0.0, 0.0], + }, + }, + { + "model_input": {"input_ids": [31, 41]}, + "loss_fn_inputs": { + "target_tokens": [41, 42], + "advantages": [1.0, 1.0], + }, + }, + ] + + batches = packer.pack(data, max_seq_len=16, request_id="test-packed-sampler-boundaries") + + assert len(batches) == 1 + batch = batches[0] + assert batch["_r3_sample_lengths"] == [4, 2] + assert batch["sampler_prefill_lengths"] == [3, 1] + assert batch["labels"][0][:4] == [IGNORE_INDEX] * 4 + assert len(batch["input_ids"][0]) == 8 + + +def test_packing_enabled_emits_full_boundary_for_prompt_only_shifted_request(): + packer = SequentialPacker(enable_packing=True, log_stats=False, pad_to_multiple_of=1) + data = [ + {"input_ids": [11, 12, 13], "target_tokens": [IGNORE_INDEX] * 3}, + {"input_ids": [21, 22], "target_tokens": [IGNORE_INDEX, 23]}, + ] + + batch = packer.pack(data, max_seq_len=16, request_id="prompt-only-boundary")[0] + + assert batch["_r3_sample_lengths"] == [3, 2] + assert batch["sampler_prefill_lengths"] == [3, 2] + + +def test_packing_enabled_splits_mixed_sampler_boundary_contracts(): + packer = SequentialPacker(enable_packing=True, log_stats=False, pad_to_multiple_of=1) + data = [ + {"input_ids": [11, 12, 13], "labels": [12, 13, 14]}, + {"input_ids": [21, 22], "target_tokens": [IGNORE_INDEX, 23]}, + ] + + batches = packer.pack(data, max_seq_len=16, request_id="mixed-boundaries") + + assert len(batches) == 2 + assert batches[0]["_r3_sample_lengths"] == [2] + assert "sampler_prefill_lengths" not in batches[0] + assert batches[1]["_r3_sample_lengths"] == [2] + assert batches[1]["sampler_prefill_lengths"] == [2] + + def test_packing_disabled_warns_on_hf_shift(monkeypatch): """HF labels should warn when shifted in the non-packed path.""" packer = SequentialPacker(enable_packing=False, log_stats=False, pad_to_multiple_of=1) diff --git a/tests/server/orchestrator/test_request_processor.py b/tests/server/orchestrator/test_request_processor.py index 0e7c6a8b..a88e074b 100644 --- a/tests/server/orchestrator/test_request_processor.py +++ b/tests/server/orchestrator/test_request_processor.py @@ -159,6 +159,11 @@ async def test_forward_backward_preserves_runner_result_fields(processor): result = { "backward_compute_time": 0.75, "forward_compute_time": 0.5, + "router_grad_norm": 1.5, + "router_grad_tensor_count": 1200, + "dsv4_grad_norm": 3.25, + "dsv4_grad_staged_factor_count_min": 948, + "dsv4_router_grad_present_count": 0, "total_loss": 0.5, "global_valid_tokens": 2, "forward_backward_time": 1.25, @@ -178,6 +183,36 @@ async def test_forward_backward_preserves_runner_result_fields(processor): assert output.outputs[0]["backward_compute_time"] == pytest.approx(0.75) assert output.outputs[0]["forward_compute_time"] == pytest.approx(0.5) assert output.outputs[0]["forward_backward_time"] == pytest.approx(1.25) + assert output.outputs[0]["router_grad_norm"] == pytest.approx(1.5) + assert output.outputs[0]["router_grad_tensor_count"] == 1200 + assert output.outputs[0]["dsv4_grad_norm"] == pytest.approx(3.25) + assert output.outputs[0]["dsv4_grad_staged_factor_count_min"] == 948 + assert output.outputs[0]["dsv4_router_grad_present_count"] == 0 + + +@pytest.mark.asyncio +async def test_optim_step_preserves_router_update_receipt(processor): + processor.backend.optim_step = AsyncMock( + return_value={ + "grad_norm": 2.0, + "step": 1, + "router_update_tensor_count": 1200, + "router_update_changed_tensor_count": 1200, + "router_update_changed_element_count": 4096, + } + ) + request = OrchestratorRequest( + request_id="req-router-update", + request_type=RequestType.ADD, + operation="optim_step", + payload=OptimStepData(model_id="default", lr=1e-4), + ) + + output = await processor.execute_optim_step(request) + + assert output.outputs[0]["router_update_tensor_count"] == 1200 + assert output.outputs[0]["router_update_changed_tensor_count"] == 1200 + assert output.outputs[0]["router_update_changed_element_count"] == 4096 def test_teacher_sort_key_reads_nested_loss_inputs(): @@ -186,6 +221,21 @@ def test_teacher_sort_key_reads_nested_loss_inputs(): assert RequestProcessor._teacher_sort_key({"teacher_id": 1, "loss_fn_inputs": {"teacher_id": 4}}) == 1 +def test_restore_datum_order_inverts_balanced_dp_packer_order(): + outputs = [{"case": name} for name in ("long", "medium", "short")] + + restored = RequestProcessor._restore_datum_order(outputs, [2, 0, 1]) + + assert [output["case"] for output in restored] == ["medium", "short", "long"] + + +def test_restore_datum_order_rejects_cardinality_and_duplicate_indices(): + with pytest.raises(RuntimeError, match="output count"): + RequestProcessor._restore_datum_order([{"case": "only"}], [1, 0]) + with pytest.raises(RuntimeError, match="duplicate"): + RequestProcessor._restore_datum_order([{"case": "a"}, {"case": "b"}], [0, 0]) + + @pytest.mark.asyncio async def test_nccl_sync_uses_request_scoped_group_name(): class CapturingBackend(DummyBackend): @@ -567,6 +617,60 @@ def _pack_samples(*args, **kwargs): await exec.stop() +@pytest.mark.asyncio +async def test_model_pass_outputs_restore_packer_datum_order(monkeypatch): + backend = DummyBackend() + exec = RequestProcessor(backend=backend, sample_packing_sequence_len=100) + + def _pack_samples(*args, **kwargs): + del args, kwargs + return ( + [ + { + "input_ids": [[1, 2, 3, 4]], + "labels": [[2, 3, 4, 5]], + "position_ids": [[0, 1, 0, 1]], + "request_id": "req-output-order", + "batch_id": 0, + "num_samples": 2, + } + ], + [1, 0], + ) + + monkeypatch.setattr(request_processor_module, "pack_samples", _pack_samples) + await exec.start() + try: + exec.backend.forward = AsyncMock( + return_value={ + "total_loss": 0.0, + "global_valid_tokens": 4, + "packed_logprobs": [[10.0, 11.0, 20.0, 21.0]], + "packed_position_ids": [[0, 1, 0, 1]], + } + ) + request = OrchestratorRequest( + request_id="req-output-order", + request_type=RequestType.ADD, + operation="forward", + payload=ModelPassData( + data=[ + {"input_ids": [1, 2], "labels": [2, 3]}, + {"input_ids": [3, 4], "labels": [4, 5]}, + ] + ), + ) + + output = await exec.execute_forward(request) + finally: + await exec.stop() + + assert output.outputs[0]["per_sample_outputs"] == [ + {"logprobs": [20.0, 21.0]}, + {"logprobs": [10.0, 11.0]}, + ] + + @pytest.mark.asyncio async def test_model_pass_cleans_externalized_routing_payloads_by_default(tmp_path): backend = DummyBackend() @@ -1085,6 +1189,7 @@ async def test_packed_row_batching_groups_single_row_packed_batches(): assert batch["input_ids"][0] == [1, 2, 3, 4, 5, 6, 7, 8, 9] assert batch["position_ids"][0] == [0, 1, 2, 0, 1, 2, 0, 1, 2] assert batch["cu_seq_lens_q"] == [0, 3, 6, 9] + assert batch["sampler_prefill_lengths"] == [1, 1, 1] assert batch["num_samples"] == 3 assert output.outputs[0]["executor_original_batches"] == 2 assert output.outputs[0]["executor_batches"] == 1 diff --git a/tests/server/runner/test_adapter_pp.py b/tests/server/runner/test_adapter_pp.py index f3cd9f41..68d71421 100644 --- a/tests/server/runner/test_adapter_pp.py +++ b/tests/server/runner/test_adapter_pp.py @@ -104,6 +104,22 @@ def test_pp1_publication_keeps_direct_path(monkeypatch) -> None: assert manager.materialize_logical_state_dict("policy", destination_rank=0) is state +def test_live_model_publication_does_not_restore_stale_adapter_slots(monkeypatch) -> None: + manager = _bare_manager(pp_size=1, stage_group=object()) + live_state = {"lm_head.lora_B": torch.tensor([[7.0]], dtype=torch.float32)} + manager.prepare_forward = lambda _model_id: pytest.fail( + "detached publication must not overwrite shared-optimizer updates" + ) + monkeypatch.setattr("xorl.lora.utils.get_lora_state_dict", lambda _model: live_state) + monkeypatch.setattr(manager_impl, "_optimizer_shard_rank_world", lambda: (0, 1)) + + publisher = manager.make_live_model_lora_publisher() + published = publisher.materialize_live_model_logical_state_dict(destination_rank=0) + + assert published is live_state + assert not hasattr(publisher, "adapters") + + def test_pp_load_filters_a_combined_checkpoint_to_the_local_stage() -> None: local_name = "model.layers.7.self_attn.q_proj.lora_A" other_stage_name = "model.layers.1.self_attn.q_proj.lora_A" diff --git a/tests/server/runner/test_checkpoint_manager_save_failures.py b/tests/server/runner/test_checkpoint_manager_save_failures.py index e6e6fc7b..03b3fcc2 100644 --- a/tests/server/runner/test_checkpoint_manager_save_failures.py +++ b/tests/server/runner/test_checkpoint_manager_save_failures.py @@ -8,6 +8,7 @@ import torch import torch.nn as nn +from xorl.models.exact_contract import set_glm52_exact_active_lora from xorl.server.runner.adapters.manager import LoRAAdapterManager from xorl.server.session_spec import normalize_session_spec @@ -123,6 +124,21 @@ def _build_checkpoint_manager() -> CheckpointManager: return manager +def test_detached_single_tenant_publication_uses_live_pp_publisher() -> None: + manager = object.__new__(CheckpointManager) + manager._adapter_manager = None + expected = {"model.layers.7.self_attn.q_proj.lora_A": torch.ones(2, 3)} + + class _Publisher: + def materialize_live_model_logical_state_dict(self, *, destination_rank): + assert destination_rank == 0 + return expected + + manager._detached_adapter_publisher = _Publisher() + + assert manager._gather_adapter_lora_params("default") is expected + + def _build_fast_save_manager(tmp_path: Path) -> CheckpointManager: model = _DummyLoRAModel(max_rank=4) adapter_manager = LoRAAdapterManager( @@ -409,6 +425,68 @@ def _capture_save_lora_checkpoint(**kwargs): assert captured["r"] == 4 +def test_glm52_active_lora_save_publishes_frozen_router_bundle(monkeypatch, tmp_path): + manager = _build_fast_save_manager(tmp_path) + manager.global_step = 0 + manager._adapter_manager.get_adapter_state("policy-a").global_step = 7 + manager.model.exact_component = _ExactActiveLoRAComponent() + manager.model.config = type( + "Config", + (), + { + "train_router": False, + "first_k_dense_replace": 3, + "num_hidden_layers": 5, + }, + )() + + router_state = { + "layer.3.weight": torch.ones(2, 2, dtype=torch.bfloat16), + "layer.4.weight": torch.full((2, 2), 2, dtype=torch.bfloat16), + } + calls = {} + + def _gather_router(model, *, destination_rank): + calls["gather"] = (model, destination_rank) + return router_state + + def _save_router(directory, state, *, weight_step, expected_layer_ids): + calls["save"] = (directory, state, weight_step, expected_layer_ids) + return { + "schema": "xorl.glm52_router_bundle.v1", + "tensor_file": "xorl_router/xorl_glm52_router.safetensors", + "sha256": "a" * 64, + "router_count": 2, + "layer_ids": [3, 4], + "weight_step": weight_step, + } + + monkeypatch.setattr(_MODULE, "gather_glm52_router_weights_across_ranks", _gather_router) + monkeypatch.setattr(_MODULE, "save_glm52_router_bundle", _save_router) + monkeypatch.setattr( + _MODULE, + "mark_adapter_config_with_glm52_router_bundle", + lambda directory, manifest: calls.setdefault("mark", (directory, manifest)), + ) + + export_dir = tmp_path / "glm52-export" + manager._save_lora_weights(str(export_dir), "policy-a") + + assert calls["gather"] == (manager.model, 0) + assert calls["save"] == (str(export_dir), router_state, 7, [3, 4]) + assert calls["mark"][0] == str(export_dir) + assert calls["mark"][1]["router_count"] == 2 + + +def test_glm52_active_lora_publication_uses_replicated_config_stamp_on_dense_stage(): + manager = object.__new__(CheckpointManager) + manager.model = nn.Sequential(nn.Linear(2, 2)) + manager.model.config = type("Config", (), {})() + set_glm52_exact_active_lora(manager.model.config, enabled=True) + + assert manager._has_glm52_exact_active_lora() + + def test_moe_lora_save_uses_resolved_target_modules_for_detection(monkeypatch, tmp_path): manager = _build_checkpoint_manager() @@ -449,6 +527,7 @@ def test_lora_save_forwards_export_format(monkeypatch, tmp_path): manager = object.__new__(CheckpointManager) manager.rank = 0 manager.local_rank = 0 + manager.global_step = 0 manager.model = nn.Module() manager.model_config = {"model_path": "Qwen/Qwen3-8B"} manager._adapter_manager = None diff --git a/tests/server/runner/test_glm52_exact_gate_up_ownership.py b/tests/server/runner/test_glm52_exact_gate_up_ownership.py index cadfb0b5..11d26794 100644 --- a/tests/server/runner/test_glm52_exact_gate_up_ownership.py +++ b/tests/server/runner/test_glm52_exact_gate_up_ownership.py @@ -121,10 +121,7 @@ def __init__(self) -> None: def _assert_exact_routed_ownership_guard_rejects_invalid_runtime_ownership(tmp_path, monkeypatch) -> None: runner, _manager, model = _exact_routed_runner(tmp_path, monkeypatch) - with pytest.raises(AdapterGradientOwnershipError, match="requires managed FSDP ownership"): - runner._compile_registered_adapter_gradient_ownership("policy") - - model.experts.ep_dispatch = "deepep" - - with pytest.raises(AdapterGradientOwnershipError, match="exact EP16 alltoall routed lane"): - runner._compile_registered_adapter_gradient_ownership("policy") + for ep_dispatch in ("alltoall", "deepep"): + model.experts.ep_dispatch = ep_dispatch + with pytest.raises(AdapterGradientOwnershipError, match="requires managed FSDP ownership"): + runner._compile_registered_adapter_gradient_ownership("policy") diff --git a/tests/server/runner/test_model_runner_drgrpo.py b/tests/server/runner/test_model_runner_drgrpo.py index 63bbd044..5a37df67 100644 --- a/tests/server/runner/test_model_runner_drgrpo.py +++ b/tests/server/runner/test_model_runner_drgrpo.py @@ -200,6 +200,57 @@ def test_compute_micro_batch_loss_dispatches_drgrpo_and_filters_loss_inputs(): assert metric_ops is None +def test_exact_dsv4_drgrpo_uses_only_independently_computed_model_values(monkeypatch): + runner = object.__new__(ModelRunner) + runner.model = _TinyModel() + runner.model.config = SimpleNamespace(_dsv4_flash_exact_mode=True) + runner.ce_mode = "eager" + runner.lm_head_fp32 = False + + def invalid_replay(**_kwargs): + raise AssertionError("DR-GRPO must differentiate the independent trainer forward") + + monkeypatch.setattr(runner, "_compute_decode_cache_micro_batch_loss", invalid_replay) + micro_batch = { + "input_ids": torch.tensor([[1, 2, 3]]), + "target_tokens": torch.tensor([[2, 3, 4]]), + "old_logprobs": torch.zeros((1, 3)), + "advantages": torch.ones((1, 3)), + } + + loss, per_token_outputs, metrics, _metric_ops, _outputs = runner._compute_micro_batch_loss( + micro_batch, + "drgrpo", + {"beta": 0.0, "return_per_token": True}, + ) + + assert not torch.equal(per_token_outputs["logprobs"], micro_batch["old_logprobs"]) + assert not any("surrogate" in name or "external_value" in name for name in metrics) + loss.backward() + assert runner.model.embed.weight.grad is not None + assert runner.model.lm_head.weight.grad is not None + + +def test_drgrpo_rejects_external_exact_value_injection(): + runner = object.__new__(ModelRunner) + runner.model = _TinyModel() + runner.model.config = SimpleNamespace(_dsv4_flash_exact_mode=True) + runner.ce_mode = "eager" + runner.lm_head_fp32 = False + with pytest.raises(ValueError, match="is not a valid training input"): + runner._compute_micro_batch_loss( + { + "input_ids": torch.tensor([[1, 2, 3]]), + "target_tokens": torch.tensor([[2, 3, 4]]), + "old_logprobs": torch.tensor([[-1.0, -1.5, -2.0]]), + "exact_value_logprobs": torch.tensor([[-1.0, -1.5, -2.0]]), + "advantages": torch.ones((1, 3)), + }, + "drgrpo", + {"beta": 0.0, "return_per_token": True}, + ) + + def test_compute_micro_batch_loss_drgrpo_accepts_legacy_logprobs_key(): runner = object.__new__(ModelRunner) runner.model = _TinyModel() @@ -481,3 +532,47 @@ def fake_forward_loop(micro_batches, loss_fn, loss_fn_params, **kwargs): assert result["model_id"] == "policy-a" assert runner.global_forward_backward_step == 8 assert runner._routing_handler.calls == [(micro_batches, None, None)] + + +def test_dsv4_gradient_metrics_measure_staged_optimizer_numerators_and_frozen_routers() -> None: + class _Dsv4TinyModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(_dsv4_flash_exact_active_lora=True) + self.lora_A = torch.nn.Parameter(torch.ones(2)) + self.lora_B = torch.nn.Parameter(torch.ones(2)) + self.mlp = torch.nn.Module() + self.mlp.gate = torch.nn.Linear(2, 2, bias=False) + self.mlp.gate.weight.requires_grad_(False) + + planned = [SimpleNamespace(fqn=f"layer.{index}.lora_A") for index in range(948)] + scratch = SimpleNamespace( + capture_open=True, + capture_staged=True, + staged_parameter_fqns=("layer.0.lora_A", "layer.1.lora_A"), + staged_numerators={ + "layer.0.lora_A": torch.tensor([3.0, 4.0]), + "layer.1.lora_A": torch.tensor([0.0, 0.0]), + }, + ) + state = SimpleNamespace( + gradient_ownership_plan=SimpleNamespace(parameters=planned), + gradient_scratch=scratch, + ) + runner = object.__new__(ModelRunner) + runner.model = _Dsv4TinyModel() + runner.model_parts = None + runner.pp_enabled = False + runner._adapter_manager = SimpleNamespace(get_adapter_state=lambda _model_id: state) + + metrics = runner._collect_dsv4_adapter_gradient_metrics("default") + + assert metrics["dsv4_grad_norm"] == pytest.approx(5.0) + assert metrics["dsv4_grad_nonzero_count"] == 2 + assert metrics["dsv4_grad_nonfinite_count"] == 0 + assert metrics["dsv4_grad_planned_factor_count_min"] == 948 + assert metrics["dsv4_grad_staged_factor_count_min"] == 2 + assert metrics["dsv4_grad_trainable_non_lora_count"] == 0 + assert metrics["dsv4_router_tensor_count_min"] == 1 + assert metrics["dsv4_router_requires_grad_count"] == 0 + assert metrics["dsv4_router_grad_present_count"] == 0 diff --git a/tests/server/runner/test_model_runner_pp_objectives.py b/tests/server/runner/test_model_runner_pp_objectives.py index 29c5ea4a..c64cfbb2 100644 --- a/tests/server/runner/test_model_runner_pp_objectives.py +++ b/tests/server/runner/test_model_runner_pp_objectives.py @@ -128,6 +128,27 @@ def test_physical_pp_dispatcher_fails_when_schedule_skips_or_mixes_ids(): assert not dispatcher.active +def test_physical_pp_terminal_objective_rejects_storage_row_mismatch(monkeypatch): + monkeypatch.setattr( + model_runner_module, + "get_parallel_state", + lambda: SimpleNamespace(tp_enabled=False, cp_enabled=False, lm_head_tp_group=None), + ) + runner = _runner() + micro_batch = _micro_batch() + micro_batch["target_tokens"] = micro_batch["target_tokens"][:, :2] + metadata = model_runner_module._PhysicalPPMicrobatchObjective( + microbatch_id=0, + micro_batch=micro_batch, + loss_fn="drgrpo", + loss_fn_params={}, + model_id="policy-a", + ) + + with pytest.raises(ValueError, match=r"target_tokens.*\(1, 2\).*terminal hidden.*\(1, 4\)"): + runner._compute_pp_terminal_objective(torch.randn(1, 4, 8), metadata) + + def test_physical_pp_rejects_only_intermediate_capture_objectives(): ModelRunner._validate_physical_pp_objective("opd_loss", {}) with pytest.raises(NotImplementedError, match="OPRD/intermediate-layer"): diff --git a/tests/server/runner/test_model_runner_session_registry.py b/tests/server/runner/test_model_runner_session_registry.py index a8b2bd78..d8209629 100644 --- a/tests/server/runner/test_model_runner_session_registry.py +++ b/tests/server/runner/test_model_runner_session_registry.py @@ -87,6 +87,12 @@ def __init__(self) -> None: self.param = torch.nn.Parameter(torch.tensor([2.0])) +class Glm5TopkRouter(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor([[1.0, 2.0]])) + + class _FakeOptimizer: def __init__(self) -> None: self.param_groups = [{"lr": 0.1}] @@ -113,6 +119,77 @@ def _build_runner() -> ModelRunner: return runner +def test_exact_glm52_router_session_is_single_tenant(): + runner = _build_runner() + runner._adapter_manager = None + runner._glm52_exact_router_single_tenant = True + runner._active_session_id = "default" + + result = ModelRunner.register_session(runner, "default", _session_spec(0.05), materialize=True) + assert result["materialized"] is True + + with pytest.raises(ValueError, match="single-tenant"): + ModelRunner.register_session(runner, "policy-b", _session_spec(0.05), materialize=True) + + +def test_exact_glm52_shared_optimizer_is_rebuilt_after_adapter_preparation(monkeypatch): + runner = _build_runner() + calls = [] + + class _Manager: + publisher = object() + + def has_adapter(self, model_id): + return model_id == "default" + + def prepare_forward(self, model_id): + calls.append(("prepare", model_id)) + + def make_live_model_lora_publisher(self): + return self.publisher + + runner._adapter_manager = _Manager() + runner._checkpoint_mgr = SimpleNamespace(optimizer="stale", _adapter_manager=runner._adapter_manager) + runner.model = _TinyModule() + runner.optimizer = "stale" + + monkeypatch.setattr(model_runner_module, "get_parallel_state", lambda: SimpleNamespace(ep_enabled=True)) + monkeypatch.setattr(model_runner_module, "refresh_ep_param_groups", lambda model: calls.append(("refresh", model))) + + def _rebuild(): + calls.append(("rebuild", None)) + runner.optimizer = "current" + + runner._initialize_optimizer = _rebuild + ModelRunner._promote_exact_glm52_default_adapter_to_shared_optimizer(runner) + + assert calls == [("prepare", "default"), ("refresh", runner.model), ("rebuild", None)] + assert runner.optimizer == "current" + assert runner._checkpoint_mgr.optimizer == "current" + assert runner._adapter_manager is None + assert runner._checkpoint_mgr._adapter_manager is None + assert runner._checkpoint_mgr._detached_adapter_publisher is _Manager.publisher + assert runner._active_session_id == "default" + + +def test_exact_glm52_router_step_receipt_proves_parameter_movement(): + runner = _build_runner() + runner.model = torch.nn.Module() + runner.model.router = Glm5TopkRouter() + runner.optimizer = torch.optim.SGD(runner.model.parameters(), lr=0.25) + runner.model.router.weight.grad = torch.ones_like(runner.model.router.weight) + + snapshots = ModelRunner._snapshot_glm52_router_weights_for_step(runner) + runner.optimizer.step() + metrics = ModelRunner._collect_glm52_router_update_metrics(runner, snapshots) + + assert metrics == { + "router_update_tensor_count": 1, + "router_update_changed_tensor_count": 1, + "router_update_changed_element_count": 2, + } + + def test_lora_session_registry_syncs_after_optimizer_checkpoint_load_and_kill(monkeypatch, tmp_path): runner = _build_runner() runner._adapter_manager = _FakeAdapterManager(lr=0.05) diff --git a/tests/server/runner/test_model_runner_token_diagnostics.py b/tests/server/runner/test_model_runner_token_diagnostics.py index db587514..dba66c1b 100644 --- a/tests/server/runner/test_model_runner_token_diagnostics.py +++ b/tests/server/runner/test_model_runner_token_diagnostics.py @@ -221,6 +221,41 @@ def test_compute_token_diagnostics_reports_hidden_state_summaries(): assert summary["layers"][1]["mean"] == pytest.approx(7.0) +def test_compute_token_diagnostics_summarizes_dsv4_hyperconnection_rows(): + hidden_states = torch.tensor([[[1.0, 2.0], [3.0, 4.0]]]) + hyperconnection_states = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]]) + labels = torch.tensor([[IGNORE_INDEX, 1]]) + + out = ModelRunner._compute_token_diagnostics( + hidden_states, + torch.eye(2), + labels, + topk=1, + all_hidden_states=(hyperconnection_states,), + hidden_sample_count=2, + ) + + layer = out["hidden_state_summaries"][0]["layers"][0] + assert layer["mean"] == pytest.approx(6.5) + assert layer["sample_values"] == pytest.approx([5.0, 6.0]) + + +def test_compute_hidden_state_diagnostics_does_not_require_lm_head(): + labels = torch.tensor([[IGNORE_INDEX, 7]]) + residual = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]]) + + out = ModelRunner._compute_hidden_state_diagnostics( + all_hidden_states=(residual,), + labels=labels, + hidden_sample_count=2, + ) + + assert out["valid_positions"] == [1] + layer = out["hidden_state_summaries"][0]["layers"][0] + assert layer["mean"] == pytest.approx(6.5) + assert layer["sample_values"] == pytest.approx([5.0, 6.0]) + + def test_compute_token_diagnostics_uses_explicit_hidden_sample_indices(): hidden_states = torch.tensor([[[1.0, 2.0, 3.0], [4.0, 6.0, 8.0]]]) weight = torch.eye(3) @@ -502,7 +537,14 @@ class DummyMlp(torch.nn.Module): def forward(self, hidden_states): capture = self._diagnostic_capture_component capture("moe_native_gathered_input", hidden_states) + capture("moe_native_recv_hidden", hidden_states) + capture("moe_native_recv_weights", hidden_states[..., :1]) + capture("moe_native_recv_local_ids", torch.zeros_like(hidden_states[..., :1], dtype=torch.long)) + capture("moe_native_expert_start", torch.tensor([0], device=hidden_states.device)) + capture("moe_native_gate_up_packed_local_0", hidden_states + 0.25) + capture("moe_native_recv_leaf", hidden_states + 0.75) capture("moe_native_local_partial", hidden_states + 1.0) + capture("moe_native_shared_folded", hidden_states + 1.5) capture("moe_native_combined", hidden_states + 2.0) return hidden_states + 2.0 @@ -537,7 +579,14 @@ def __init__(self, layer): by_name = {capture["name"]: capture for capture in captures} assert by_name["moe_native_gathered_input"]["order"] == 80 + assert by_name["moe_native_recv_hidden"]["order"] == 80 + assert by_name["moe_native_recv_weights"]["order"] == 81 + assert by_name["moe_native_recv_local_ids"]["order"] == 82 + assert by_name["moe_native_expert_start"]["order"] == 83 + assert by_name["moe_native_gate_up_packed_local_0"]["order"] == 83 + assert by_name["moe_native_recv_leaf"]["order"] == 84 assert by_name["moe_native_local_partial"]["order"] == 89 + assert by_name["moe_native_shared_folded"]["order"] == 89 assert by_name["moe_native_combined"]["order"] == 90 torch.testing.assert_close(by_name["moe_native_local_partial"]["tensor"], hidden_states + 1.0) diff --git a/tests/server/runner/test_runner_dispatcher.py b/tests/server/runner/test_runner_dispatcher.py index 2a1cc84b..a4302008 100644 --- a/tests/server/runner/test_runner_dispatcher.py +++ b/tests/server/runner/test_runner_dispatcher.py @@ -775,10 +775,27 @@ def _assert_rank_local_row_batching_and_provenance_policy(monkeypatch): assert my_batches[0]["packed_row_source_num_samples"] == [1, 1] assert my_batches[0]["packed_row_source_token_spans"] == [[0, 2], [2, 4]] assert my_batches[0]["packed_row_source_group_size"] == 2 + assert "sampler_prefill_lengths" not in my_batches[0] _assert_unmerged_rows_record_source_provenance() +def test_packed_row_batching_keeps_only_complete_sampler_boundaries(): + first = _batch(10) + second = _batch(20) + first.update(_r3_sample_lengths=[2], sampler_prefill_lengths=[1]) + second.update(_r3_sample_lengths=[2], sampler_prefill_lengths=[2]) + + complete = batch_packed_rows([first, second], row_batch_size=2)[0] + assert complete["sampler_prefill_lengths"] == [1, 2] + + second.pop("sampler_prefill_lengths") + mixed_contracts = batch_packed_rows([first, second], row_batch_size=2) + assert len(mixed_contracts) == 2 + assert mixed_contracts[0]["sampler_prefill_lengths"] == [1] + assert "sampler_prefill_lengths" not in mixed_contracts[1] + + def _assert_unmerged_rows_record_source_provenance(): first = _batch(10) second = _batch(20) diff --git a/tests/server/test_server_arguments.py b/tests/server/test_server_arguments.py index 80e762bb..b61dfde3 100644 --- a/tests/server/test_server_arguments.py +++ b/tests/server/test_server_arguments.py @@ -182,6 +182,8 @@ def test_removed_field_inventory_allows_unrelated_unknown_fields(): ("fp8_cfg", {"enabled": True}), ("externalize_r3_payloads", True), ("keep_r3_payloads", True), + ("deepep_native_lora_mode", "separate"), + ("deepep_native_combine_mode", "deterministic"), ), ) def test_removed_training_aliases_are_absent_from_server_arguments_schema(field_name, value): @@ -195,6 +197,29 @@ def test_server_arguments_config_does_not_emit_removed_training_aliases(): assert {"fp8_cfg", "externalize_r3_payloads", "keep_r3_payloads"}.isdisjoint(config["train"]) +def test_server_arguments_propagates_native_deepep_modes(): + config = ServerArguments( + model_path="Qwen/Qwen3.5-35B-A3B", + deepep_native_exact=True, + expert_parallel_size=2, + enable_lora=True, + lora_serving_mode="separate", + ).to_config_dict() + + assert config["model"]["deepep_native_exact"] is True + assert config["lora"]["lora_serving_mode"] == "separate" + assert config["train"]["gradient_checkpointing_method"] == "recompute_before_dispatch" + + +def test_server_arguments_reject_native_deepep_at_ep1(): + with pytest.raises(ValueError, match="expert_parallel_size > 1"): + ServerArguments( + model_path="Qwen/Qwen3.5-35B-A3B", + deepep_native_exact=True, + expert_parallel_size=1, + ) + + _SHIPPED_EXACT_QWEN35_MOE_LORA_CONFIG = "examples/server/configs/lora/qwen3_5_35b_a3b_lora.yaml" _SHIPPED_MOE_LORA_CONFIGS = ("examples/server/configs/lora/qwen3_coder_30b_a3b_lora.yaml",) _SHIPPED_QWEN35_LORA_CONFIGS = ( @@ -800,6 +825,19 @@ def test_load_server_arguments_admits_any_positive_exact_glm52_rank(tmp_path, ra assert args.get_total_gpus() == 16 +def test_load_server_arguments_admits_exact_glm52_router_training(tmp_path): + payload = _exact_glm52_rank1_server_config(tmp_path) + payload["model"]["train_router"] = True + payload["train"]["freeze_router"] = False + config_path = tmp_path / "server_config.yaml" + config_path.write_text(yaml.safe_dump(payload), encoding="utf-8") + + args = load_server_arguments(str(config_path)) + + assert args.train_router is True + assert args.freeze_router is False + + def test_load_server_arguments_admits_dp_owned_exact_glm52_row(tmp_path): payload = _exact_glm52_rank1_server_config(tmp_path) payload["train"].update(ulysses_parallel_size=1, data_parallel_shard_size=16) diff --git a/tests/server/weight_sync/test_glm52_router_bundle.py b/tests/server/weight_sync/test_glm52_router_bundle.py new file mode 100644 index 00000000..5a442682 --- /dev/null +++ b/tests/server/weight_sync/test_glm52_router_bundle.py @@ -0,0 +1,85 @@ +import hashlib +import json + +import torch + +from xorl.server.weight_sync.glm52_router_bundle import ( + GLM52_ROUTER_BUNDLE_SCHEMA, + GLM52_ROUTER_MANIFEST, + GLM52_ROUTER_TENSORS, + _merge_glm52_router_states, + gather_glm52_router_weights, + mark_adapter_config_with_glm52_router_bundle, + save_glm52_router_bundle, +) + + +class Glm5TopkRouter(torch.nn.Module): + def __init__(self, value: float): + super().__init__() + self.weight = torch.nn.Parameter(torch.full((3, 4), value, dtype=torch.float32)) + + +class _MLP(torch.nn.Module): + def __init__(self, value: float): + super().__init__() + self.gate = Glm5TopkRouter(value) + + +class _Layer(torch.nn.Module): + def __init__(self, value: float): + super().__init__() + self.mlp = _MLP(value) + + +class _Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.layers = torch.nn.ModuleList([_Layer(1.0), _Layer(2.0)]) + + +def test_router_bundle_is_complete_checksummed_and_bound_to_adapter(tmp_path): + state = gather_glm52_router_weights(_Model()) + assert list(state) == ["layer.0.weight", "layer.1.weight"] + assert all(tensor.dtype is torch.bfloat16 for tensor in state.values()) + + manifest = save_glm52_router_bundle(tmp_path, state, weight_step=7, expected_layer_ids=[0, 1]) + (tmp_path / "adapter_config.json").write_text("{}\n") + mark_adapter_config_with_glm52_router_bundle(tmp_path, manifest) + + tensor_bytes = (tmp_path / GLM52_ROUTER_TENSORS).read_bytes() + persisted = json.loads((tmp_path / GLM52_ROUTER_MANIFEST).read_text()) + marker = json.loads((tmp_path / "adapter_config.json").read_text())["_xorl_glm52_router_bundle"] + assert persisted == manifest + assert persisted["schema"] == GLM52_ROUTER_BUNDLE_SCHEMA + assert persisted["sha256"] == hashlib.sha256(tensor_bytes).hexdigest() + assert marker == {key: manifest[key] for key in marker} + assert marker["layer_ids"] == [0, 1] + assert marker["weight_step"] == 7 + assert not list(tmp_path.glob("xorl_glm52_router*.safetensors")) + assert (tmp_path / GLM52_ROUTER_TENSORS).parent.name == "xorl_router" + + +def test_router_bundle_rejects_incomplete_inventory(tmp_path): + state = {"layer.0.weight": torch.zeros((3, 4), dtype=torch.bfloat16)} + try: + save_glm52_router_bundle(tmp_path, state, weight_step=1, expected_layer_ids=[0, 1]) + except RuntimeError as error: + assert "Incomplete GLM-5.2 router sidecar" in str(error) + else: + raise AssertionError("incomplete router publication was accepted") + + +def test_pipeline_router_states_merge_disjoint_stages_and_identical_replicas(): + layer_0 = torch.ones((3, 4), dtype=torch.bfloat16) + layer_1 = torch.full((3, 4), 2, dtype=torch.bfloat16) + + merged = _merge_glm52_router_states( + [ + {"layer.0.weight": layer_0}, + {"layer.1.weight": layer_1}, + {"layer.1.weight": layer_1.clone()}, + ] + ) + + assert list(merged) == ["layer.0.weight", "layer.1.weight"] diff --git a/tests/server/weight_sync/test_handler_config.py b/tests/server/weight_sync/test_handler_config.py index e14948ed..d745b327 100644 --- a/tests/server/weight_sync/test_handler_config.py +++ b/tests/server/weight_sync/test_handler_config.py @@ -351,6 +351,10 @@ def register_lora_adapter(self, model_id, lr): assert handler._prepare_lora_adapter_for_sync(None) == "current-adapter" assert trainer.adapter_manager.synced == ["policy-b", "current-adapter"] + trainer.lora_config = {"lora_serving_mode": "separate"} + with pytest.raises(RuntimeError, match="publishes A/B factors"): + handler._prepare_lora_adapter_for_sync("policy-a") + _assert_extract_params_for_sync_policy() _assert_unfuse_for_inference_layout_policy() with monkeypatch.context() as case_patch: @@ -996,6 +1000,109 @@ def __init__(self): ) +def test_collect_ep_moe_data_separate_mode_omits_frozen_expert_base(): + from xorl.models.layers.moe.lora import MoEExpertsLoRA + + class Wrapper(torch.nn.Module): + def __init__(self): + super().__init__() + base = MoEExperts( + num_experts=2, + hidden_dim=3, + intermediate_size=5, + hidden_act="silu", + moe_implementation="eager", + gated=True, + ) + self.experts = MoEExpertsLoRA.from_module( + base, + r=2, + lora_alpha=2, + target_modules=["gate_proj", "up_proj", "down_proj"], + hybrid_shared=True, + ) + + wrapper = Wrapper() + experts = wrapper.experts + torch.nn.init.normal_(experts.gate_up_proj) + torch.nn.init.normal_(experts.down_proj) + with torch.no_grad(): + for name, parameter in experts.named_parameters(): + if "lora_B" in name: + parameter.fill_(0.25) + experts.exact_merged_forward = True + experts.lora_serving_mode = "separate" + + handler = WeightSyncHandler(rank=0, world_size=1, trainer=None) + contexts = handler._collect_ep_moe_data(wrapper, "(root)", None) + + # Separate mode publishes routed-expert LoRA factors separately. The base + # checkpoint is immutable, so no per-expert HF keys may reach SGLang's + # online loader (whose active-LoRA wrapper nests the base FusedMoE). The + # skip-only context must suppress the ordinary dense extraction path while + # producing no EP transfer context. + assert contexts == [ + { + "type": "frozen_active_lora_base", + "prefix": "experts", + "local_experts": None, + } + ] + prefixes, transferable = handler._split_ep_moe_contexts_for_sync(contexts, "(root)") + assert prefixes == {"experts"} + assert transferable == [] + extracted = handler._extract_params_for_sync( + wrapper, + "(root)", + object, + skip_moe_prefixes=prefixes, + ) + assert extracted == [] + + +def test_collect_ep_moe_data_separate_mode_omits_metadata_only_context(): + from xorl.models.layers.moe.lora import MoEExpertsLoRA + + base = MoEExperts( + num_experts=2, + hidden_dim=3, + intermediate_size=5, + hidden_act="silu", + moe_implementation="eager", + gated=True, + ) + wrapper = torch.nn.Module() + wrapper.experts = MoEExpertsLoRA.from_module( + base, + r=2, + lora_alpha=2, + target_modules=["gate_proj", "up_proj", "down_proj"], + hybrid_shared=True, + ) + wrapper.experts.lora_serving_mode = "separate" + + handler = WeightSyncHandler(rank=1, world_size=8, trainer=None) + contexts = handler._collect_ep_moe_data( + wrapper, + "model.layers.0.mlp", + None, + collect_tensors=False, + ) + assert contexts == [ + { + "type": "frozen_active_lora_base", + "prefix": "model.layers.0.mlp.experts", + "local_experts": None, + } + ] + prefixes, transferable = handler._split_ep_moe_contexts_for_sync( + contexts, + "model.layers.0.mlp", + ) + assert prefixes == {"experts"} + assert transferable == [] + + def _assert_compile_wrapper_name_normalization_policy(): config = SimpleNamespace( hidden_size=8, diff --git a/tests/test_arguments.py b/tests/test_arguments.py index 14ab9f55..83c2d4ee 100644 --- a/tests/test_arguments.py +++ b/tests/test_arguments.py @@ -70,6 +70,8 @@ def test_parse_args_optimizer_packing_and_numeric_policy(tmp_path, monkeypatch): assert args.model.activation_native is True assert args.model.rope_native is True assert args.model.attention_cast_bf16 is True + assert args.model.deepep_native_exact is False + assert args.lora.lora_serving_mode is None muon_root = tmp_path / "muon" muon_root.mkdir() @@ -170,6 +172,33 @@ def _assert_parse_args_checkpoint_policy(tmp_path, monkeypatch): assert args.train.gradient_checkpointing_method == "recompute_before_dispatch" assert args.train.moe_recomputed is False + native_config_path = tmp_path / "native-exact-default.yaml" + native_config_path.write_text( + yaml.safe_dump( + { + "model": { + "model_path": "Qwen/Qwen3-8B", + "deepep_native_exact": True, + }, + "data": { + "datasets": [{"path": "dummy", "type": "tokenized"}], + }, + "train": { + "init_device": "meta", + "output_dir": str(tmp_path / "native-outputs"), + "expert_parallel_size": 2, + "use_wandb": False, + }, + } + ), + encoding="utf-8", + ) + monkeypatch.setattr(sys, "argv", ["train.py", str(native_config_path)]) + native_args = parse_args(Arguments) + + assert native_args.train.gradient_checkpointing_method == "recompute_before_dispatch" + assert native_args.train.moe_recomputed is False + auto_root = tmp_path / "auto-checkpoint" auto_root.mkdir() _assert_parse_args_resolves_auto_checkpoint_before_validation(auto_root, monkeypatch) @@ -178,6 +207,36 @@ def _assert_parse_args_checkpoint_policy(tmp_path, monkeypatch): _assert_parse_args_load_optimizer_flag(optimizer_root, monkeypatch) +def test_native_exact_training_rejects_ep1(tmp_path, monkeypatch): + config_path = tmp_path / "native-exact-ep1.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model": { + "model_path": "Qwen/Qwen3-8B", + "deepep_native_exact": True, + }, + "data": {"datasets": [{"path": "dummy", "type": "tokenized"}]}, + "train": { + "init_device": "meta", + "output_dir": str(tmp_path / "outputs"), + "expert_parallel_size": 1, + "use_wandb": False, + }, + } + ), + encoding="utf-8", + ) + monkeypatch.setenv("WORLD_SIZE", "1") + monkeypatch.setenv("LOCAL_WORLD_SIZE", "1") + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("LOCAL_RANK", "0") + monkeypatch.setattr(sys, "argv", ["train.py", str(config_path)]) + + with pytest.raises(ValueError, match="expert_parallel_size > 1"): + parse_args(Arguments) + + def _assert_parse_args_low_precision_configuration_policy(tmp_path, monkeypatch): config_path = tmp_path / "config.yaml" config_path.write_text( diff --git a/tests/trainers/test_bi_trunk_linear_server_mode.py b/tests/trainers/test_bi_trunk_linear_server_mode.py index cc4b527e..5a65471c 100644 --- a/tests/trainers/test_bi_trunk_linear_server_mode.py +++ b/tests/trainers/test_bi_trunk_linear_server_mode.py @@ -24,6 +24,8 @@ def __init__(self): nn.ModuleDict( { "q_proj": nn.Linear(16, 16, bias=False, dtype=torch.bfloat16), + "k_proj": nn.Linear(16, 16, bias=False, dtype=torch.bfloat16), + "v_proj": nn.Linear(16, 16, bias=False, dtype=torch.bfloat16), "o_proj": nn.Linear(16, 16, bias=False, dtype=torch.bfloat16), "gate_proj": nn.Linear(16, 32, bias=False, dtype=torch.bfloat16), "up_proj": nn.Linear(16, 32, bias=False, dtype=torch.bfloat16), @@ -42,6 +44,15 @@ def _apply_qwen35_gdn_exact(self): return wrap_trunk_linears_batch_invariant(self) +class NativeExactTinyTrunkModel(TinyTrunkModel): + def _apply_deepep_native_exact(self): + from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import ( + _apply_qwen3_deepep_native_exact, + ) + + return _apply_qwen3_deepep_native_exact(self) + + @pytest.fixture(autouse=True) def _reset_contract_state(): yield @@ -57,6 +68,7 @@ def _build( server_training=False, freeze_router=False, enable_lora=False, + lora_serving_mode=None, ): def fake_parallelize(model, **_kwargs): captured["wrapped_at_parallelize"] = sum( @@ -82,11 +94,10 @@ def fake_build_foundation_model(**kwargs): server_training=server_training, freeze_router=freeze_router, enable_lora=enable_lora, + lora_serving_mode=lora_serving_mode, lora_rank=2, lora_alpha=4, - # TinyTrunkModel has no k_proj/v_proj, and injection rejects a target that - # matches nothing, so request only the projections the fixture defines. - lora_target_modules=["q_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + lora_target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], ) @@ -95,7 +106,7 @@ def test_exact_server_model_wraps_trunk_linears_before_parallelize(monkeypatch): result = _build(monkeypatch, captured, model=ExactTinyTrunkModel(), server_training=True) assert captured["server_training_at_build"] is True - assert captured["wrapped_at_parallelize"] == 10, "wrap must land before parallelization (pre-FSDP2)" + assert captured["wrapped_at_parallelize"] == 14, "wrap must land before parallelization (pre-FSDP2)" assert not getattr(result.model.lm_head, "_xorl_bi_trunk_wrapped", False) assert is_trunk_linear_contract_enabled() @@ -146,16 +157,53 @@ def test_dense_qwen_exact_program_is_reinstalled_after_lora_replacement(monkeypa wrap_trunk_linears_batch_invariant(model) result = _build(monkeypatch, captured, model=model, server_training=True, enable_lora=True) - assert captured["wrapped_at_parallelize"] == 10 + assert captured["wrapped_at_parallelize"] == 14 assert bi_families_v2.families_v2_enabled() is True for layer in result.model.layers: - for name in ("q_proj", "o_proj", "gate_proj", "up_proj", "down_proj"): + for name in ("q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"): module = layer[name] assert type(module) is LoraLinear assert module.exact_merged_forward is True assert module._xorl_bi_trunk_wrapped is True +def test_qwen3_moe_native_deepep_installs_exact_program_before_parallelize(monkeypatch): + captured = {} + model = NativeExactTinyTrunkModel() + model.config = SimpleNamespace(model_type="qwen3_moe", _deepep_native_exact=True) + monkeypatch.setenv("XORL_FAMILIES_V2", "0") + + _build(monkeypatch, captured, model=model, server_training=True) + + assert captured["wrapped_at_parallelize"] == 14 + assert bi_families_v2.families_v2_enabled() is True + + +@pytest.mark.parametrize("lora_serving_mode", ["merged", "separate"]) +def test_qwen3_moe_native_deepep_installs_selected_lora_program(monkeypatch, lora_serving_mode): + captured = {} + model = NativeExactTinyTrunkModel() + model.config = SimpleNamespace( + model_type="qwen3_moe", + _deepep_native_exact=True, + _lora_serving_mode=lora_serving_mode, + ) + + result = _build( + monkeypatch, + captured, + model=model, + server_training=True, + enable_lora=True, + lora_serving_mode=lora_serving_mode, + ) + + for layer in result.model.layers: + for module in layer.values(): + assert module.lora_serving_mode == lora_serving_mode + assert module.exact_merged_forward is True + + def test_ordinary_model_restores_nonexact_family_selection(monkeypatch): captured = {} bi_families_v2._select_glm52_families_v2() diff --git a/tests/trainers/test_deepseek_v3_training_guards.py b/tests/trainers/test_deepseek_v3_training_guards.py index 654ecd39..8f6f242c 100644 --- a/tests/trainers/test_deepseek_v3_training_guards.py +++ b/tests/trainers/test_deepseek_v3_training_guards.py @@ -53,6 +53,15 @@ def _assert_build_foundation_model_rejects_train_router_for_deepseek(): build_foundation_model(_tiny_config(), train_router=True) +def test_build_foundation_model_rejects_unwired_exact_deepep(monkeypatch): + with pytest.raises(ValueError, match="model-declared BF16 native DeepEP"): + build_foundation_model( + _tiny_config(), + deepep_native_exact=True, + ep_dispatch="deepep", + ) + + def test_deepseek_training_and_parallelization_admission_policy(monkeypatch): _assert_build_foundation_model_rejects_train_router_for_deepseek() diff --git a/tests/trainers/test_fp8_model_builder.py b/tests/trainers/test_fp8_model_builder.py index 61624908..78146e21 100644 --- a/tests/trainers/test_fp8_model_builder.py +++ b/tests/trainers/test_fp8_model_builder.py @@ -16,6 +16,16 @@ pytestmark = [pytest.mark.cpu] +def test_build_training_model_rejects_native_exact_without_ep(): + with pytest.raises(ValueError, match="expert_parallel_size > 1"): + build_training_model( + config_path="unused", + weights_path="unused", + deepep_native_exact=True, + expert_parallel_size=1, + ) + + class TinyDenseMoEModel(nn.Module): _no_split_modules = [] @@ -136,6 +146,126 @@ def fake_parallelize(model, **kwargs): assert all(param.requires_grad for param in result.model.parameters()) +def test_build_training_model_threads_sharded_lm_head_loss_to_parallelize(monkeypatch): + captured = {} + + def fake_build_foundation_model(**_kwargs): + return TinyDenseOnlyModel() + + def fake_parallelize(model, **kwargs): + captured.update(kwargs) + return model + + monkeypatch.setattr("xorl.trainers.model_builder.build_foundation_model", fake_build_foundation_model) + monkeypatch.setattr("xorl.trainers.model_builder._parallelize", fake_parallelize) + monkeypatch.setattr("xorl.trainers.model_builder.helper.print_device_mem_info", lambda *args, **kwargs: None) + + build_training_model( + config_path="unused", + weights_path="unused", + fsdp_sharded_lm_head_loss=True, + enable_mixed_precision=False, + enable_gradient_checkpointing=False, + ) + + assert captured["fsdp_sharded_lm_head_loss"] is True + + +def test_build_training_model_threads_glm52_block_fp8_qlora_mode(monkeypatch): + captured = {} + inventory = object() + + def fake_build_foundation_model(**kwargs): + captured["foundation_flag"] = kwargs["block_fp8_qlora_training"] + captured["foundation_rank"] = kwargs["lora_rank"] + captured["foundation_alpha"] = kwargs["lora_alpha"] + return TinyDenseOnlyModel() + + def fake_inject_qlora(model, **kwargs): + captured["inject_flag"] = kwargs["block_fp8_qlora_training"] + captured["quant_format"] = kwargs["quant_format"] + captured["quant_group_size"] = kwargs["quant_group_size"] + model._glm52_adapter_inventory = inventory + return True, "block_fp8", set() + + monkeypatch.setattr("xorl.trainers.model_builder.build_foundation_model", fake_build_foundation_model) + monkeypatch.setattr("xorl.trainers.model_builder._inject_qlora", fake_inject_qlora) + monkeypatch.setattr("xorl.trainers.model_builder._deferred_qlora_quantize", lambda *args, **kwargs: None) + monkeypatch.setattr("xorl.trainers.model_builder._parallelize", lambda model, **_kwargs: model) + monkeypatch.setattr("xorl.trainers.model_builder.helper.print_device_mem_info", lambda *args, **kwargs: None) + + result = build_training_model( + config_path="unused", + weights_path="unused", + moe_implementation="triton", + ep_dispatch="deepep", + enable_lora=True, + enable_qlora=True, + block_fp8_qlora_training=True, + quant_format="block_fp8", + quant_group_size=128, + moe_hybrid_shared_lora=True, + freeze_router=True, + enable_mixed_precision=False, + enable_gradient_checkpointing=False, + ) + + assert captured == { + "foundation_flag": True, + "foundation_rank": 32, + "foundation_alpha": 16, + "inject_flag": True, + "quant_format": "block_fp8", + "quant_group_size": 128, + } + assert result.glm52_adapter_inventory is inventory + + +def test_build_training_model_retains_exact_glm52_router_after_post_fsdp_qlora_freeze(monkeypatch): + inventory = object() + + def fake_build_foundation_model(**_kwargs): + model = TinyDenseOnlyModel() + model.config.train_router = True + # Match the production local-router inventory seam, which counts + # sparse blocks rather than arbitrary modules that happen to own a gate. + model.mlp = type("Glm5MoEBlock", (nn.Module,), {})() + model.mlp.gate = nn.Linear(16, 2, bias=False) + return model + + def fake_inject_qlora(model, **_kwargs): + model.proj.register_parameter("lora_A", nn.Parameter(model.proj.weight.new_zeros((1, 16)))) + model._glm52_adapter_inventory = inventory + return True, "block_fp8", set() + + monkeypatch.setattr("xorl.trainers.model_builder.build_foundation_model", fake_build_foundation_model) + monkeypatch.setattr("xorl.trainers.model_builder._inject_qlora", fake_inject_qlora) + monkeypatch.setattr("xorl.trainers.model_builder._deferred_qlora_quantize", lambda *args, **kwargs: None) + monkeypatch.setattr("xorl.trainers.model_builder._parallelize", lambda model, **_kwargs: model) + monkeypatch.setattr("xorl.trainers.model_builder.glm52_exact_active_lora_enabled", lambda _config: True) + monkeypatch.setattr("xorl.trainers.model_builder.helper.print_device_mem_info", lambda *args, **kwargs: None) + + result = build_training_model( + config_path="unused", + weights_path="unused", + moe_implementation="triton", + ep_dispatch="deepep", + enable_lora=True, + enable_qlora=True, + block_fp8_qlora_training=True, + quant_format="block_fp8", + quant_group_size=128, + moe_hybrid_shared_lora=True, + freeze_router=False, + enable_mixed_precision=False, + enable_gradient_checkpointing=False, + ) + + assert result.model.mlp.gate.weight.requires_grad + assert result.model.proj.lora_A.requires_grad + assert not result.model.proj.weight.requires_grad + + def _assert_build_training_model_rejects_glm52_block_fp8_mode_without_qlora(): with pytest.raises(ValueError, match="requires enable_lora=True and enable_qlora=True"): build_training_model( diff --git a/tests/trainers/test_rope_class_b_config.py b/tests/trainers/test_rope_class_b_config.py index 51f811e8..80719c68 100644 --- a/tests/trainers/test_rope_class_b_config.py +++ b/tests/trainers/test_rope_class_b_config.py @@ -402,6 +402,18 @@ def test_exact_qwen35_moe_admits_structural_defaults(): ) +def test_exact_qwen35_moe_admits_shared_native_deepep_program(): + config = _exact_qwen35_moe_config() + config._qwen35_exact_contract = True + _validate_exact_qwen35_moe_program( + config, + moe_implementation="triton", + ep_dispatch="deepep", + deepep_async_combine=False, + deepep_native_exact=True, + ) + + def test_exact_glm52_model_scope_accepts_only_official_geometry(): _validate_canonical_glm52_model_scope(_exact_glm52_config()) config = _exact_glm52_config() diff --git a/tests/trainers/test_trainer_model_alignment_flags.py b/tests/trainers/test_trainer_model_alignment_flags.py index ed7fca9c..56fe24c0 100644 --- a/tests/trainers/test_trainer_model_alignment_flags.py +++ b/tests/trainers/test_trainer_model_alignment_flags.py @@ -82,6 +82,7 @@ def _trainer_args(): enable_lora=False, enable_qlora=False, unfuse_for_lora=False, + lora_serving_mode=None, lora_rank=16, lora_alpha=16, ), diff --git a/tests/trainers/test_training_utils.py b/tests/trainers/test_training_utils.py index 577e678d..15a1ae6a 100644 --- a/tests/trainers/test_training_utils.py +++ b/tests/trainers/test_training_utils.py @@ -47,7 +47,11 @@ def remote_owner_has_twelve_storage_rows(negotiation, op=None, group=None): { "input_ids": torch.arange(8).view(1, -1), "labels": torch.arange(8).view(1, -1), + "target_tokens": torch.arange(8).view(1, -1), + "old_logprobs": torch.arange(8, dtype=torch.float32).view(1, -1), + "advantages": torch.ones((1, 8), dtype=torch.float32), "position_ids": torch.arange(16).view(1, -1), + "attention_mask": torch.ones((1, 16), dtype=torch.long), "cu_seq_lens_q": torch.tensor([0, 16], dtype=torch.int32), "cu_seq_lens_k": torch.tensor([0, 16], dtype=torch.int32), "max_length_q": 16, @@ -63,10 +67,15 @@ def remote_owner_has_twelve_storage_rows(negotiation, op=None, group=None): assert align_dsv4_pp_storage_rows(micro_batches, cp_size=2) == 12 batch = micro_batches[0] assert batch["input_ids"].shape == (1, 12) + assert batch["target_tokens"].shape == (1, 12) + assert batch["target_tokens"][0, 8:].tolist() == [IGNORE_INDEX] * 4 + torch.testing.assert_close(batch["old_logprobs"][0, 8:], torch.zeros(4)) + torch.testing.assert_close(batch["advantages"][0, 8:], torch.zeros(4)) assert batch["_cp_live_mask"].shape == (1, 12) assert torch.count_nonzero(batch["_cp_live_mask"]) == 6 assert batch["_r3_sample_lengths"] is sample_lengths assert batch["position_ids"].shape == (1, 24) + assert batch["attention_mask"].shape == (1, 24) assert int(batch["cu_seq_lens_q"][-1]) == 24 assert int(batch["cu_seq_lens_k"][-1]) == 24 @@ -263,6 +272,12 @@ def test_pp_padding_uses_exact_sampling_transform_identities(): { "input_ids": torch.tensor([[7, 8]], dtype=torch.int64), "labels": torch.tensor([[8, 9]], dtype=torch.int64), + "target_tokens": torch.tensor([[8, 9]], dtype=torch.int64), + "logprobs": torch.tensor([[-1.0, -2.0]], dtype=torch.float32), + "old_logprobs": torch.tensor([[-1.0, -2.0]], dtype=torch.float32), + "ref_logprobs": torch.tensor([[-1.5, -2.5]], dtype=torch.float32), + "rollout_logprobs": torch.tensor([[-1.0, -2.0]], dtype=torch.float32), + "advantages": torch.tensor([[0.25, -0.5]], dtype=torch.float32), "logprob_temperatures": torch.tensor([[0.7, 1.3]], dtype=torch.float32), "logprob_top_ks": torch.tensor([[4, 9]], dtype=torch.int64), "logprob_top_ps": torch.tensor([[0.8, 0.9]], dtype=torch.float32), @@ -275,12 +290,27 @@ def test_pp_padding_uses_exact_sampling_transform_identities(): batch = micro_batches[0] assert batch["input_ids"].tolist() == [[7, 8, 0, 0]] assert batch["labels"].tolist() == [[8, 9, IGNORE_INDEX, IGNORE_INDEX]] + assert batch["target_tokens"].tolist() == [[8, 9, IGNORE_INDEX, IGNORE_INDEX]] + for key in ("logprobs", "old_logprobs", "ref_logprobs", "rollout_logprobs", "advantages"): + torch.testing.assert_close(batch[key][:, 2:], torch.zeros((1, 2))) torch.testing.assert_close(batch["logprob_temperatures"], torch.tensor([[0.7, 1.3, 1.0, 1.0]])) assert batch["logprob_top_ks"].tolist() == [[4, 9, 1 << 30, 1 << 30]] torch.testing.assert_close(batch["logprob_top_ps"], torch.tensor([[0.8, 0.9, 1.0, 1.0]])) torch.testing.assert_close(batch["logprob_min_ps"], torch.tensor([[0.1, 0.2, 0.0, 0.0]])) +def test_pp_padding_rejects_misaligned_objective_rows(): + micro_batches = [ + { + "input_ids": torch.tensor([[7, 8]], dtype=torch.int64), + "target_tokens": torch.tensor([[8]], dtype=torch.int64), + } + ] + + with pytest.raises(ValueError, match="target_tokens.*1 rows.*input_ids.*2"): + pad_micro_batches_for_pp(micro_batches, sample_packing_sequence_len=4) + + def test_pp_chunked_ce_matches_eager_loss_and_grad(monkeypatch): monkeypatch.setenv("XORL_PP_CE_CHUNK_TOKENS", "2") labels = torch.tensor([[1, 2, IGNORE_INDEX], [3, 4, 0]]) diff --git a/tests/trainers/test_unfuse_for_lora.py b/tests/trainers/test_unfuse_for_lora.py index 74640592..2f2af19c 100644 --- a/tests/trainers/test_unfuse_for_lora.py +++ b/tests/trainers/test_unfuse_for_lora.py @@ -248,6 +248,7 @@ def _trainer(monkeypatch, calls, model): enable_lora=True, enable_qlora=False, unfuse_for_lora=True, + lora_serving_mode=None, lora_rank=2, lora_alpha=2, ), diff --git a/vendor/deepep-release.lock.json b/vendor/deepep-release.lock.json new file mode 100644 index 00000000..e80134a5 --- /dev/null +++ b/vendor/deepep-release.lock.json @@ -0,0 +1,32 @@ +{ + "schema": "xorl.deepep_release_lock.v2", + "release_repository": "https://github.com/togethercomputer/xorl-wheels", + "release_tag": "deepep_sglang_kernel_torch212_cu132_sm90_723b8b3", + "source": { + "xorl_wheels_commit": "723b8b394519cc3d39274480a82933322b2bc037", + "deepep_commit": "65538abaaec51d5a06b92a96f874b3f9274ebdc8", + "deepep_tree": "54927caff5e877b9b390729258086e3d679f760f", + "url": "https://github.com/togethercomputer/xorl-wheels/releases/download/deepep_sglang_kernel_torch212_cu132_sm90_723b8b3/deep_ep-1.2.1%2B65538ab.xorl.c85744ca7250.tar.gz", + "sha256": "1c191d78d0132049ac6caa988266d2fa4502d78a26735d65bd46b5fa18ee5f77" + }, + "wheel": { + "package": "deep-ep==1.2.1+65538ab.xorl.c85744ca7250", + "url": "https://github.com/togethercomputer/xorl-wheels/releases/download/deepep_sglang_kernel_torch212_cu132_sm90_723b8b3/deep_ep-1.2.1%2B65538ab.xorl.c85744ca7250-cp312-cp312-linux_x86_64.whl", + "sha256": "f30485d585a4cd935f44ffadd276d7ea8603a919cff60302bbb40ded344cdccb", + "extension_sha256": "1233ae5405e2a10439f331b07cb221697a451cd3c8a2c28bccddf7d209551e59", + "python_abi": "cp312-cp312-linux_x86_64", + "torch": "2.12.1+cu132", + "torch_cuda_runtime": "13.2", + "build_cuda_toolkit": "13.2", + "gpu_arch": "sm90" + }, + "contract": { + "deterministic_mode": "explicit_opt_in", + "protocol": "deepep_deterministic_hierarchical_bf16_v2", + "supported_ep": [2, 4, 8, 16], + "model_route_wire_dtype": "bfloat16", + "receiver_weight_modes": ["unit", "fp32"], + "fold": "tree8_fp64_bf16_node_leaf_then_ascending_node_fp64_fold", + "cross_mode": "normal_equals_low_latency_byte_for_byte" + } +} diff --git a/vendor/sglang-kernel-release.lock.json b/vendor/sglang-kernel-release.lock.json new file mode 100644 index 00000000..7cb72618 --- /dev/null +++ b/vendor/sglang-kernel-release.lock.json @@ -0,0 +1,23 @@ +{ + "schema": "xorl.sglang_kernel_release_lock.v1", + "release_repository": "https://github.com/togethercomputer/xorl-wheels", + "release_tag": "deepep_sglang_kernel_torch212_cu132_sm90_723b8b3", + "source": { + "xorl_wheels_commit": "723b8b394519cc3d39274480a82933322b2bc037", + "sglang_commit": "65d5a0ec2596d314fbcd88bf179d2ea912aaaa02", + "sglang_aot_tree": "5dea732d55c9f413ea698054342a56e6df4356d1", + "url": "https://github.com/togethercomputer/xorl-wheels/releases/download/deepep_sglang_kernel_torch212_cu132_sm90_723b8b3/sglang_kernel-0.4.5%2Bxorl.torch212.cu132.sm90.tar.gz", + "sha256": "d0934b3871805c1c9233f0e6437b85cc04ecd459cbfb3452cefd20d5ab23ca88" + }, + "wheel": { + "package": "sglang-kernel==0.4.5+xorl.torch212.cu132.sm90", + "url": "https://github.com/togethercomputer/xorl-wheels/releases/download/deepep_sglang_kernel_torch212_cu132_sm90_723b8b3/sglang_kernel-0.4.5%2Bxorl.torch212.cu132.sm90-cp312-cp312-linux_x86_64.whl", + "sha256": "f02e35414c18fd311ce29b7d27c9a07678f3932a9be3cabfec8d322f7b738f21", + "python_abi": "cp312-cp312-linux_x86_64", + "torch": "2.12.1+cu132", + "torch_cuda_runtime": "13.2", + "build_cuda_toolkit": "13.2", + "gpu_arch": "sm90", + "full_fa3": true + } +}