Add exact deterministic DeepEP training

.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

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

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
 

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

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

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,20 +2396,43 @@ 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),
                 "lora.quant_format": (self.lora.quant_format, "block_fp8"),
                 "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:

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(

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",
 ]