diff --git a/.github/workflows/publish_pypi.yaml b/.github/workflows/publish_pypi.yaml index af118e434..d0540b5a7 100644 --- a/.github/workflows/publish_pypi.yaml +++ b/.github/workflows/publish_pypi.yaml @@ -3,31 +3,56 @@ name: Publish to PyPI on: workflow_dispatch: +permissions: + contents: read + +concurrency: + group: publish-pypi + cancel-in-progress: false + jobs: - build-n-publish: - if: github.event_name == 'workflow_dispatch' - name: Build and publish Python distributions to PyPI + build: + name: Build Python distributions + runs-on: ubuntu-latest + timeout-minutes: 20 + steps: + - uses: actions/checkout@v7 + with: + persist-credentials: false + + - uses: actions/setup-python@v7 + with: + python-version: "3.11" + + - name: Build and validate distributions + run: | + python -m pip install --upgrade build twine + python -m build + python -m twine check --strict dist/* + + - name: Upload distributions + uses: actions/upload-artifact@v7 + with: + name: python-distributions + path: dist/ + if-no-files-found: error + + publish: + name: Publish distributions to PyPI + needs: build runs-on: ubuntu-latest timeout-minutes: 20 environment: name: pypi - url: https://pypi.org/p/specforgeee + url: https://pypi.org/p/specforge permissions: id-token: write steps: - - uses: actions/checkout@v2 - - - uses: actions/setup-python@v2 - with: - python-version: '3.11' - - - run: pip install build && python -m build --sdist - - # publish to PyPI if executed on the main branch - - name: Publish package to PyPI - id: publish - uses: pypa/gh-action-pypi-publish@release/v1 - with: - user: __token__ - password: ${{ secrets.PYPI_TOKEN }} - verbose: true + - name: Download distributions + uses: actions/download-artifact@v8 + with: + name: python-distributions + path: dist/ + + - name: Publish distributions to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 520b60555..8ada955c1 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -13,7 +13,24 @@ permissions: contents: read jobs: + cleanup-runner: + if: (github.repository == 'sgl-project/SpecForge' || github.event_name == 'pull_request') && + github.event.pull_request.draft == false + runs-on: [self-hosted] + concurrency: + group: specforge-gpu-ci + cancel-in-progress: false + queue: max + timeout-minutes: 5 + steps: + # Job containers need a fresh Docker network before any container step can + # run, so stale networks must be pruned in a host-side prerequisite job. + - name: Prune stale Docker networks + shell: bash + run: docker network prune --force --filter "until=1h" + unit-test: + needs: cleanup-runner if: (github.repository == 'sgl-project/SpecForge' || github.event_name == 'pull_request') && github.event.pull_request.draft == false runs-on: [self-hosted] diff --git a/README.md b/README.md index d34a68055..8c39db113 100644 --- a/README.md +++ b/README.md @@ -16,11 +16,12 @@ SpecForge is an ecosystem project developed by the SGLang team. It is a framework for training speculative decoding models so that you can smoothly port them over to the SGLang serving framework to speed up your inference. We have seen many open-source projects for speculative decoding, but most of them are not well-maintained or not directly compatible with SGLang. We prepared this project because we wish that the open-source community can enjoy a speculative decoding framework that is + - regularly maintained by the SpecForge team: the code is runnable out-of-the-box -- directly compatible with SGLang: there is no additional efforts for porting to SGLang -- provides local offline and server-only online-disaggregated training through - one runtime, including the supported data, tensor, and sequence parallel - topologies +- directly compatible with SGLang: no additional porting effort is required +- able to run online disaggregated training and both colocated and + disaggregated offline training through one runtime, including the supported + data, tensor, and sequence parallel topologies Check out [**our documentation**](https://docs.sglang.ai/SpecForge/) to get started. @@ -31,23 +32,25 @@ Check out [**our documentation**](https://docs.sglang.ai/SpecForge/) to get star Every method uses the same typed training entry point: ```bash -specforge train --config examples/configs/qwen3-8b-eagle3-disaggregated.yaml +specforge train --config examples/configs/online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml ``` -The typed `deployment.trainer` topology self-launches trainer DP and EAGLE3 -offline USP process groups. A single-node disaggregated config also supervises -its SpecForge producer and consumer; Mooncake and SGLang remain externally -managed services, and online target parallelism belongs to SGLang. There are no -method-specific Python training entry points. +The path under `examples/configs` identifies feature mode, topology, and online +service ownership. The command above uses an `external` recipe: SpecForge +supervises the producer and consumer on one trainer node, while the user or +scheduler owns Mooncake and SGLang. Recipes under `managed-local` also start +those services on the local host. Online target parallelism belongs to SGLang; +`deployment.trainer` owns trainer DP and offline EAGLE3 USP process groups. +There are no method-specific Python training entry points. | Method | Description | Example config | Optimization | | --- | --- | --- | --- | -| **[EAGLE3](https://arxiv.org/abs/2503.01840)** | Feature-based autoregressive drafting | [Online](./examples/configs/qwen3-8b-eagle3-disaggregated.yaml) / [Offline](./examples/configs/qwen3-8b-eagle3-offline.yaml) / [Disaggregated offline](./examples/configs/qwen3-8b-eagle3-offline-disaggregated.yaml) | [LK loss](https://arxiv.org/pdf/2602.23881) | -| **[P-EAGLE](https://arxiv.org/abs/2602.01469)** | Parallel EAGLE | [Online](./examples/configs/qwen3-8b-peagle-disaggregated.yaml) | — | -| **EAGLE3.1** | Feature-based autoregressive drafting with attention drift | [Online](./examples/configs/qwen3-30b-a3b-eagle3.1-online.yaml) | — | -| **[DFlash](https://arxiv.org/abs/2602.06036)** | Block-parallel drafting | [Online](./examples/configs/qwen3-8b-dflash-online.yaml) / [Disaggregated](./examples/configs/qwen3-8b-dflash-disaggregated.yaml) | [D-PACE](https://arxiv.org/abs/2605.18810) | -| **[Domino](https://arxiv.org/html/2605.29707v1)** | DFlash with GRU logit correction | [Online](./examples/configs/qwen3-8b-domino-online.yaml) / [Disaggregated](./examples/configs/qwen3-8b-domino-disaggregated.yaml) | — | -| **[DSpark](https://arxiv.org/abs/2607.05147)** | Confidence-Scheduled Semi-Autoregressive Generation | [Disaggregated](./examples/configs/qwen3-4b-dspark-disaggregated.yaml) | — | +| **[EAGLE3](https://arxiv.org/abs/2503.01840)** | Feature-based autoregressive drafting | [Online external](./examples/configs/online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml) / [Offline colocated](./examples/configs/offline/colocated/qwen3-8b-eagle3-offline.yaml) / [Offline disaggregated](./examples/configs/offline/disaggregated/qwen3-8b-eagle3-offline-disaggregated.yaml) | [LK loss](https://arxiv.org/pdf/2602.23881) | +| **[P-EAGLE](https://arxiv.org/abs/2602.01469)** | Parallel EAGLE | [Online external](./examples/configs/online/disaggregated/external/qwen3-8b-peagle-disaggregated.yaml) | — | +| **EAGLE3.1** | Feature-based autoregressive drafting with attention drift | [Online external](./examples/configs/online/disaggregated/external/qwen3-30b-a3b-eagle3.1-online.yaml) | — | +| **[DFlash](https://arxiv.org/abs/2602.06036)** | Block-parallel drafting | [Online external](./examples/configs/online/disaggregated/external/qwen3-8b-dflash-online.yaml) / [Offline colocated](./examples/configs/offline/colocated/qwen3-8b-dflash-offline.yaml) / [Online managed-local](./examples/configs/online/disaggregated/managed-local/qwen3-8b-dflash-1server-dp7-disaggregated.yaml) | [D-PACE](https://arxiv.org/abs/2605.18810) | +| **[Domino](https://arxiv.org/html/2605.29707v1)** | DFlash with GRU logit correction | [Online external](./examples/configs/online/disaggregated/external/qwen3-8b-domino-online.yaml) / [Offline colocated](./examples/configs/offline/colocated/qwen3-8b-domino-offline.yaml) / [Online managed-local](./examples/configs/online/disaggregated/managed-local/qwen3-8b-domino-multiserver-disaggregated.yaml) | — | +| **[DSpark](https://arxiv.org/abs/2607.05147)** | Confidence-Scheduled Semi-Autoregressive Generation | [Online external](./examples/configs/online/disaggregated/external/qwen3-4b-dspark-disaggregated.yaml) / [Offline colocated](./examples/configs/offline/colocated/qwen3-4b-dspark-offline.yaml) | — | See the [training guide](./docs/basic_usage/training.md) for the supported method/topology matrix and the @@ -69,11 +72,14 @@ SpecBundle is a collection of production-grade speculative decoding models that ## 🎉 News - +- [2026-08] 🎉 Released SpecBundle (phase 2) and SpecForge v0.3.0. Check out our blog at [LMSYS.org](https://www.lmsys.org/blog/2026-08-04-specforge-v0-3) +- [2026-07] 🚀 Day0 supported two flagship dspark draft model, [Inkling](https://huggingface.co/RadixArk/Inkling-DSpark-Preview) and [Kimi-K3](https://huggingface.co/RadixArk/Kimi-K3-DSpark). +- [2026-07] 🔥 Supported full disaggregation of training and inference in online training. +- [2026-07] 🔥 Added DSpark online training for DFlash draft models. - [2026-06] 🔥 Added D-PACE as an optional loss for DFlash training. - [2026-06] 🔥 Added Domino online training for DFlash draft models. - [2026-01] 🔥 Added DFlash block-parallel online training with SGLang serving support. -- [2025-12] 🎉 Released SpecBundle (phase 1) and SpecForge v0.2. Check out our blog at [LMSYS.org](https://lmsys.org/blog/2025-12-23-spec-bundle-phase-1/) +- [2025-12] 🎉 Released SpecBundle (phase 1) and SpecForge v0.2.0. Check out our blog at [LMSYS.org](https://lmsys.org/blog/2025-12-23-spec-bundle-phase-1/) - [2025-08] 🔔 SpecForge is listed as a [flagship project](https://lmsys.org/about/) in LMSYS. Congratulations to the SpecForge team! - [2025-08] 🔥 SpecForge powered the Eagle3 draft model for GPT-OSS. Check out the blog at [LMSYS.org](https://lmsys.org/blog/2025-08-27-gpt-oss/) - [2025-07] 🔥 SpecForge is released together with Llama4-Eagle3 checkpoints. Check out our blog at [LMSYS.org](https://lmsys.org/blog/2025-07-25-spec-forge/) diff --git a/configs/kimi-k3-dspark.json b/configs/kimi-k3-dspark.json new file mode 100644 index 000000000..f59c98d61 --- /dev/null +++ b/configs/kimi-k3-dspark.json @@ -0,0 +1,53 @@ +{ + "architectures": ["DSparkDraftModel"], + "attention_bias": false, + "attention_dropout": 0.0, + "auto_map": {"AutoModel": "dspark.DSparkDraftModel"}, + "block_size": 7, + "bos_token_id": 163584, + "dflash_config": { + "attention_mode": "gqa", + "confidence_head_alpha": 1.0, + "confidence_head_with_markov": true, + "enable_confidence_head": true, + "markov_head_type": "vanilla", + "markov_rank": 256, + "mask_token_id": 163824, + "projector_type": "dspark", + "target_layer_ids": [7, 23, 51, 67, 83] + }, + "dtype": "bfloat16", + "eos_token_id": 163586, + "head_dim": 64, + "hidden_act": "silu", + "hidden_size": 7168, + "initializer_range": 0.02, + "intermediate_size": 14336, + "layer_types": [ + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention" + ], + "max_position_embeddings": 65536, + "max_window_layers": 5, + "model_type": "qwen3", + "num_attention_heads": 64, + "num_hidden_layers": 5, + "num_key_value_heads": 16, + "num_target_layers": 93, + "pad_token_id": 163839, + "rms_norm_eps": 1e-05, + "rope_parameters": { + "factor": 16.0, + "original_max_position_embeddings": 65536, + "rope_theta": 10000.0, + "rope_type": "yarn" + }, + "sliding_window": null, + "tie_word_embeddings": false, + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 163840 +} diff --git a/configs/qwen3.5-4b-vl-dflash.json b/configs/qwen3.5-4b-mtp.json similarity index 50% rename from configs/qwen3.5-4b-vl-dflash.json rename to configs/qwen3.5-4b-mtp.json index 366df304b..29a3858db 100644 --- a/configs/qwen3.5-4b-vl-dflash.json +++ b/configs/qwen3.5-4b-mtp.json @@ -1,50 +1,37 @@ { "architectures": [ - "DFlashDraftModel" + "Qwen3_5MTPDraftModel" ], "attention_bias": false, "attention_dropout": 0.0, "auto_map": { - "AutoModel": "dflash.DFlashDraftModel" + "AutoModel": "mtp.Qwen3_5MTPDraftModel" }, - "block_size": 16, + "attn_output_gate": true, "bos_token_id": 248043, - "dflash_config": { - "mask_token_id": 248070, - "target_layer_ids": [1, 8, 15, 22, 29] - }, - "dtype": "bfloat16", "eos_token_id": 248044, "head_dim": 256, "hidden_act": "silu", "hidden_size": 2560, "initializer_range": 0.02, "intermediate_size": 9216, - "layer_types": [ - "full_attention", - "full_attention", - "full_attention", - "full_attention", - "full_attention" - ], "max_position_embeddings": 262144, - "model_type": "qwen3_vl_text", + "model_type": "qwen3", + "mtp_config": { + "share_lm_head": true + }, "num_attention_heads": 16, - "num_hidden_layers": 5, + "num_hidden_layers": 1, "num_key_value_heads": 4, - "num_target_layers": 32, "pad_token_id": 248044, "partial_rotary_factor": 0.25, "rms_norm_eps": 1e-06, - "rope_scaling": { - "mrope_interleaved": true, - "mrope_section": [11, 11, 10], - "rope_type": "default", - "partial_rotary_factor": 0.25 - }, + "rope_scaling": null, "rope_theta": 10000000, "sliding_window": null, "tie_word_embeddings": true, + "torch_dtype": "bfloat16", "use_cache": true, + "use_sliding_window": false, "vocab_size": 248320 } diff --git a/docs/advanced_features/customization.md b/docs/advanced_features/customization.md index dd4776680..fa71be886 100644 --- a/docs/advanced_features/customization.md +++ b/docs/advanced_features/customization.md @@ -13,7 +13,7 @@ For one-off changes, use dotted overrides rather than adding another launcher: ```bash specforge train \ - --config examples/configs/qwen3-8b-eagle3-disaggregated.yaml \ + --config examples/configs/online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml \ model.target_model_path=/models/my-target \ data.train_data_path=/datasets/my-training-data.jsonl \ training.learning_rate=5e-5 diff --git a/docs/advanced_features/vlm_dflash.md b/docs/advanced_features/vlm_dflash.md index e6ef712b3..23aa530e9 100644 --- a/docs/advanced_features/vlm_dflash.md +++ b/docs/advanced_features/vlm_dflash.md @@ -1,25 +1,27 @@ -# VLM DFlash Support — Status (add_vl_support) +# VLM DFlash Support - Status (add_vl_support) This branch ports sgl-project/SpecForge PR #585 (commit `9323a510`, author zyk42) onto the server-only unified runtime **and implements end-to-end multimodal (image+text) DFlash training on top of it**. This document records what landed, what was deliberately not ported, and the validation status. +> **Rope decision**: multimodal DFlash drafts use the plain 1D rope convention - +> the same as the text-only path - on purpose. Visual information reaches the +> draft exclusively through the captured target hidden states, never through the +> draft's own positional embedding, so the draft has no use for the target's +> (3, L) mRoPE positions. Staying on one position convention keeps training and +> serving byte-identical to the text pipeline and reuses the stock +> `configs/qwen3.5-4b-dflash.json` draft geometry. + ## What this branch contains ### Foundation (ported from PR #585) -- **Draft model (VLM-capable)** — `specforge/modeling/draft/dflash.py`: +- **Draft model (VLM-capable)** - `specforge/modeling/draft/dflash.py`: partial rotation in `apply_rotary_pos_emb` (`rotary_dim < head_dim`, for - Qwen3.5/Qwen3.6 `partial_rotary_factor=0.25`) and - `Qwen3InterleavedMultiRotaryEmbedding` selected by - `rope_scaling.mrope_interleaved`. -- **Draft config** — `configs/qwen3.5-4b-vl-dflash.json` (new, mirrors the - Qwen3.5-4B target geometry: head_dim 256, mrope_section [11,11,10], - partial_rotary_factor 0.25). The wider #585 config set (Qwen3-VL-8B/30B-A3B, - Qwen3.5-9B/35B-A3B) is intentionally out of scope for this branch; it can be - added verbatim in a follow-up once more targets are validated. -- **Weight-key resolution** — `resolve_target_weight_keys()` in + Qwen3.5/Qwen3.6 `partial_rotary_factor=0.25`). The draft always uses the + stock `Qwen3RotaryEmbedding`. +- **Weight-key resolution** - `resolve_target_weight_keys()` in `specforge/modeling/target/target_utils.py` auto-selects `model.language_model.embed_tokens.weight` for VLM targets; `populate_dflash_generated_config` reads language-model depth via @@ -27,46 +29,49 @@ what was deliberately not ported, and the validation status. ### Multimodal capture (new in this branch) -End-to-end data flow: **JSONL (+ image) → expanded ids/loss mask → capture -request (single-placeholder ids + base64 image) → patched SGLang server -expands, runs the ViT, captures aux hidden states + mRoPE positions → -Mooncake → collator → training forward with 3D position ids**. +End-to-end data flow: **JSONL (+ image) -> expanded ids/loss mask -> capture +request (single-placeholder ids + base64 image) -> patched SGLang server +expands, runs the ViT, captures aux hidden states -> Mooncake -> collator -> +training forward on plain 1D positions**. - `model.input_modality: multimodal` (DFlash only): a `FeatureContract` - (`{input_ids, loss_mask, hidden_states, position_ids}`) and a + (`{input_ids, loss_mask, hidden_states}`) and a `ServerStreamingProvider` with a VLM `ServerInputAdapter` - (`specforge/algorithms/common/vlm_input.py`). + (`specforge/algorithms/common/vlm_input.py`). Multimodal capture stores the + same three tensors as text capture - no `position_ids` artifact is requested + or consumed. - `specforge/data/vlm_preprocessing.py`: ShareGPT-style records with an optional `image` field (path or base64); the target's own chat template and HF processor produce the expanded `input_ids`/`loss_mask` (image region expanded in id space, mask zeros). One image per sample max (v1); text-only samples work in the same run. -- `ServerCaptureLayout.position_ids_feature` → the capture request's - `features["position_ids"]`; the patched server writes the request's mRoPE - positions `(1, L, 3) int64` into Mooncake (`_spec_capture_position_ids` in - the scheduler sink; text requests get the arange broadcast fallback). -- `patches/sglang/v0.5.14/spec-capture.patch`: regenerated with the - `position_ids` artifact (`SpecCaptureSink.put_sample(position_ids=...)`). - Multimodal capture requests ride the stock `input_ids` + `image_data` - `/generate` path with `SGLANG_MM_AVOID_RETOKENIZE=1` (set by the managed - launcher for `input_modality=multimodal`), so the server re-expands - placeholders in id space with zero retokenization drift — and the - passthrough/seq-len checks fail loudly if client and server expansions ever - disagree. -- Training: `OnlineDFlashModel.forward(..., position_ids=None)` gathers 3D - mRoPE positions for context + anchor-offset draft slots - (`(3, B, S + N·bs)`); `DFlashTrainStrategy` passes the collated - `position_ids` tensor through. Text runs are byte-identical to before. -- Recipe: `examples/configs/qwen3.5-4b-vl-dflash-disaggregated.yaml` - (single-node Ascend NPU managed stack). +- `patches/sglang/v0.5.14/spec-capture.patch`: tracks upstream's rewritten + async streaming sink. Multimodal capture requests ride the stock + `input_ids` + `image_data` `/generate` path with + `SGLANG_MM_AVOID_RETOKENIZE=1` (set by the managed launcher for + `input_modality=multimodal`), so the server re-expands placeholders in id + space with zero retokenization drift - and the passthrough/seq-len checks + fail loudly if client and server expansions ever disagree. +- Training: `OnlineDFlashModel._forward_draft_blocks` builds positions with + the unconditional text-path 1D `arange` convention; multimodal batches flow + through the identical forward as text batches. Text runs are + byte-identical to before. +- Recipe: `examples/configs/online/disaggregated/external/qwen3.5-4b-vl-dflash-disaggregated.yaml` + (single-node Ascend NPU managed stack; draft config + `configs/qwen3.5-4b-dflash.json`). ## Not ported (by design) - HF-backend VLM capture (`dflash_target_model.py`, `_build_vlm_reqs`, `mm_token_type_ids`) and the `train_dflash.py --is-vlm` plumbing from the - pre-#678 script stack — superseded by server capture. -- `QwenVLOnlineDFlashModel` wiring — PR #585 referenced this class but never + pre-#678 script stack - superseded by server capture. +- `QwenVLOnlineDFlashModel` wiring - PR #585 referenced this class but never defined it; the unified runtime needs no separate VLM wrapper class. +- mRoPE draft support (`Qwen3InterleavedMultiRotaryEmbedding`, the + `rope_scaling.mrope_interleaved` switch, and the server `position_ids` + capture artifact): the draft consumes visual information only through the + captured target hidden states, so the (3, L) target positions carry no + signal for it. Retired in favor of the single plain-rope convention above. - Two accidental reverts in the original #585 diff (domino projector code, D-PACE CLI args) were dropped during the cherry-pick. - Offline (precomputed hidden states) multimodal capture: the offline path @@ -77,14 +82,13 @@ Mooncake → collator → training forward with 3D position ids**. - **Verified (CPU, this repo)**: registration parity and provider gates, request/payload construction, expansion math, collator, golden - topology/recipe tests — `tests/test_algorithms/test_dflash_multimodal.py` - plus the updated `test_config` suites (80 passed locally, 2 torch tests - skipped pending GPU). + topology/recipe tests - `tests/test_algorithms/test_dflash_multimodal.py` + plus the updated `test_config` suites. - **Verified statically**: the regenerated patch applies cleanly both ways to pristine sglang v0.5.14 (`git apply --check` / `--reverse --check`). - **Not yet verified (needs GPU/NPU)**: a live multimodal capture run - (Qwen3-VL / Qwen3.5 target + ViT + mRoPE positions) and an end-to-end - training run. This is the next step; see the recipe above. + (Qwen3-VL / Qwen3.5 target + ViT) and an end-to-end training run. This is + the next step; see the recipe above. ## Reference results from PR #585 (HF stack, author-validated) diff --git a/docs/basic_usage/AMD/amd_rocm.md b/docs/basic_usage/AMD/amd_rocm.md new file mode 100644 index 000000000..5a75e5527 --- /dev/null +++ b/docs/basic_usage/AMD/amd_rocm.md @@ -0,0 +1,499 @@ +# 🚀 AMD ROCm Tutorial + +This is an end-to-end tutorial for running SpecForge on AMD Instinct GPUs +(ROCm). It walks through **installation → data preparation → offline colocated +training → online disaggregated training**. Within the online workflow, it +covers external services under one supervisor, the managed-local shortcut, and +external producer/consumer roles split across process pools or nodes. + +All commands assume a ROCm host with the AMD driver stack and Docker already +installed. Validated on MI300X (gfx942) and MI355X (gfx950). + +--- + +## 1. Installation + +On ROCm, install SpecForge into an environment that already provides a ROCm +PyTorch and a ROCm SGLang, and install the package **without dependencies** so +pip does not pull CUDA wheels over the working ROCm stack. + +The recommended base is an official SGLang ROCm release container. These ship a +ROCm PyTorch and an editable ROCm SGLang build, so SpecForge only needs to be +cloned and installed on top. + +### Step 1: Pull the image for your accelerator + +The accelerator is baked into the tag, so use the image that matches your +hardware: + +```bash +# AMD Instinct MI300X (gfx942) +docker pull lmsysorg/sglang:v0.5.14-rocm720-mi30x + +# AMD Instinct MI355X (gfx950) +docker pull lmsysorg/sglang:v0.5.14-rocm700-mi35x +``` + +### Step 2: Start the container + +Expose the ROCm device nodes (swap in the tag for your accelerator). Use `--name` +and omit `--rm` so the checkout survives across sessions: + +```bash +docker run -it --name specforge \ + --device=/dev/kfd --device=/dev/dri \ + --group-add video --cap-add SYS_PTRACE --security-opt seccomp=unconfined \ + --ipc=host --shm-size=16g \ + lmsysorg/sglang:v0.5.14-rocm720-mi30x \ + bash +``` + +`--device=/dev/kfd --device=/dev/dri --group-add video` are required for ROCm +GPU access; `--ipc=host --shm-size=16g` gives Mooncake and PyTorch enough shared +memory. Re-enter the running container later with `docker exec -it specforge bash`. + +### Step 3: Clone and install SpecForge + +Inside the container, clone SpecForge into `/workspace/SpecForge` and register it +in editable mode without touching the image's torch/sglang: + +```bash +git clone https://github.com/sgl-project/SpecForge.git /workspace/SpecForge +cd /workspace/SpecForge +python -m pip install -e . --no-deps +``` + +`--no-deps` is mandatory: a full resolve pulls the CUDA SGLang stack and +clobbers the image's ROCm torch/sglang. If a later step reports a missing +lightweight dependency (for example `accelerate`), install just that package, +also with `--no-deps`. + +### Step 4: Apply the capture patch (online runs only) + +These images pin SGLang to exactly `0.5.14` (editable at `/sgl-workspace/sglang`), +so the online capture patch applies with a plain `git apply`. Skip this step for +offline training, which reads features from disk and needs no capture service: + +```bash +cd /sgl-workspace/sglang +git apply /workspace/SpecForge/patches/sglang/v0.5.14/spec-capture.patch +cd /workspace/SpecForge +``` + +The patch adds the `--enable-spec-capture`, `--spec-capture-method`, and +`--spec-capture-aux-layer-ids` server flags plus the `sglang.srt.spec_capture_sink` +module used by online capture. + +### Step 5: Attention backends on ROCm + +Use the `sdpa` or `flex_attention` attention backends for the **trainer** on +ROCm. The `fa` (flash-attn) and `usp` backends, and `yunchang`-based +Ulysses/Ring sequence parallel (`sp_ulysses_size` / `sp_ring_size` > 1), depend +on a CUDA flash-attn build; the single-GPU / data-parallel path never loads +`yunchang`, and selecting those backends raises a clear error. The checked-in +[`qwen3.5-4b-dflash-offline-amd.yaml`](../../../examples/configs/offline/colocated/qwen3.5-4b-dflash-offline-amd.yaml) +recipe already uses `flex_attention`, so it runs on ROCm unchanged as a +single-GPU offline DFlash example. + +The **capture side** (the SGLang target that materializes hidden states, both +offline and online) has an extra ROCm requirement for Qwen3.5-4B, a hybrid +linear-attention/Mamba target: run it under **AITER** and disable the radix +cache. Section 3 covers this in detail. + +--- + +## 2. Data preparation + +Data preparation is platform independent — the same scripts run on ROCm. Write a +ShareGPT training set into `cache/dataset` from the repository root: + +```bash +python scripts/prepare_data.py --dataset sharegpt +``` + +This produces `./cache/dataset/sharegpt_train.jsonl` in the stable +`id` + `conversations` contract used by every checked-in recipe. For the full +preset list, custom datasets, preformatted text, and target-model regeneration, +see the [Data Preparation](../data_preparation.md) guide. + +--- + +## 3. Offline colocated training + +Offline training reads target features from disk, so the trainer only has to fit +the draft model. It uses more storage but keeps target inference out of the +training loop, and needs no capture patch or Mooncake. This section trains a +**Qwen3.5-4B DFlash** draft (`configs/qwen3.5-4b-dflash.json`). + +### Step 1: Capture hidden states + +Feature preparation is a data-processing step, not a second training entry point. +Qwen3.5-4B is a hybrid linear-attention/Mamba target, so run the capture under +**AITER** and pass `--sglang-disable-radix-cache`. Without it, SGLang's Mamba +radix cache selects the `extra_buffer` strategy, which asserts CUDA/MUSA/NPU +(FLA) at server init and fails on ROCm: + +```bash +SGLANG_USE_AITER=1 SGLANG_USE_AITER_UNIFIED_ATTN=1 AITER_FLYDSL_FORCE=1 \ +torchrun --standalone --nproc_per_node 1 \ + scripts/prepare_hidden_states.py \ + --target-model-path Qwen/Qwen3.5-4B \ + --strategy dflash \ + --draft-model-config configs/qwen3.5-4b-dflash.json \ + --trust-remote-code \ + --data-path ./cache/dataset/sharegpt_train.jsonl \ + --output-path ./cache/hidden_states/qwen3.5-4b-dflash-sharegpt \ + --chat-template qwen3.5 \ + --max-length 2048 \ + --tp-size 1 \ + --batch-size 8 \ + --sglang-attention-backend aiter \ + --sglang-disable-radix-cache \ + --sglang-mem-fraction-static 0.8 \ + --sglang-context-length 2560 +``` + +The output path matches `data.hidden_states_path` in the checked-in offline +recipe. See [Data Preparation](../data_preparation.md#option-2-pre-formatted-text-format) +for preformatted inputs and other options. + +> **Data note:** `prepare_hidden_states.py` truncates each rendered conversation +> at `max_length`. A long prompt can push the assistant reply past the cutoff, +> leaving an empty loss region (fewer than two anchorable tokens), which trips +> DFlash's anchor sampler (`ValueError: should preprocess the data.`). Drop the +> captured samples with `< 2` loss-mask tokens before the last `block_size` +> positions before training. The online path (Section 4) never hits this — its +> producer regenerates full-length responses. + +### Step 2: Train + +The checked-in offline recipe already uses `flex_attention` for the trainer, so +it runs on ROCm unchanged: + +```bash +specforge train --config examples/configs/offline/colocated/qwen3.5-4b-dflash-offline-amd.yaml +``` + +Override any field inline without copying the YAML, e.g. a quick smoke run: + +```bash +specforge train --config examples/configs/offline/colocated/qwen3.5-4b-dflash-offline-amd.yaml \ + training.max_steps=20 output_dir=./outputs/dflash-offline-smoke +``` + +See the [Training](../training.md) guide for the full run schema, +checkpoint/resume rules, and evaluation. + +### Optional: Offline disaggregated deployment + +The walkthrough above is offline colocated: the trainer reads prepared feature +files directly. If a producer must ingest those files for a separate trainer +pool, choose a recipe under +[`examples/configs/offline/disaggregated/`](../../../examples/configs/offline/disaggregated/). +The feature source remains offline; only the deployment topology changes. This +path does not start SGLang. A `shared_dir` backend requires storage visible to +the producer and consumers, while a Mooncake backend requires an existing +Mooncake deployment. See +[Offline shared-directory and Mooncake stores](../disaggregated_training.md#offline-shared-directory-and-mooncake-stores) +for the complete contract. + +--- + +## 4. Online disaggregated training + +Online training captures target features live from a patched SGLang server and +streams them through Mooncake to the trainer. Every online run is +**disaggregated**: a producer drives prompts through the capture server and a +consumer trains the draft model. The choices below change service ownership and +process placement; they do not create additional training modes. + +### External services with a single-node supervisor + +With `deployment.trainer.nnodes: 1` and no `--role`, one `specforge train` +command supervises the producer and consumer. Mooncake and SGLang remain +external services started by the user. + +This section uses the external +[`qwen3.5-4b-dflash-online-amd.yaml`](../../../examples/configs/online/disaggregated/external/qwen3.5-4b-dflash-online-amd.yaml) +recipe as a single-node smoke test. `external` means that the user starts +Mooncake and SGLang; the services still run locally in this example. Complete +Step 4 of the installation first. + +#### Step 1: One-time run inputs + +DFlash needs **no shared vocabulary mapping** (that is an EAGLE3-only +requirement, where a reduced draft vocabulary must be derived once and shared by +producer and consumer). DFlash keeps the full vocabulary and derives its target +signal from the draft's `target_layer_ids`, so there is nothing to precompute. + +Qwen3.5-4B is also a large sharded checkpoint that already ships a +`*.index.json` weight map and a resolvable head, so it needs no local target +directory or index workaround. Just make sure `cache/dataset/sharegpt_train.jsonl` +exists (Section 2). + +#### Step 2: Start Mooncake and the capture server + +Start the Mooncake master. **Set `--default_kv_lease_ttl=500`**: the consumer's +teardown drain now allows about 19.5s for leases to settle, while the shorter +managed TTL keeps a normal shutdown from waiting several seconds for an expired +read lease. + +```bash +mooncake_master --enable_http_metadata_server=true \ + --rpc_port=35551 --http_metadata_server_port=35880 \ + --metrics_port=35903 --enable_metric_reporting=false \ + --default_kv_lease_ttl=500 & +``` + +Start the patched capture server on GPU 0. **The `--spec-capture-aux-layer-ids` +must match the draft's `target_layer_ids`** — for DFlash these are read straight +from `configs/qwen3.5-4b-dflash.json`: `1 8 15 22 29` (this is not the EAGLE3 +`[1, num_layers//2 - 1, num_layers - 4]` formula). A mismatch produces zero +features with no error. Because Qwen3.5-4B is a hybrid Mamba target, the server +must run under **AITER** with `--attention-backend aiter` and +`--disable-radix-cache` (see Section 3 for why): + +```bash +SGLANG_USE_AITER=1 SGLANG_USE_AITER_UNIFIED_ATTN=1 AITER_FLYDSL_FORCE=1 \ +HIP_VISIBLE_DEVICES=0 CUDA_VISIBLE_DEVICES=0 \ +MOONCAKE_LOCAL_HOSTNAME=127.0.0.1 \ +MOONCAKE_METADATA_SERVER=http://127.0.0.1:35880/metadata \ +MOONCAKE_MASTER_SERVER_ADDR=127.0.0.1:35551 \ +MOONCAKE_PROTOCOL=tcp \ +MOONCAKE_GLOBAL_SEGMENT_SIZE=$((32<<30)) \ +python -m sglang.launch_server \ + --model-path Qwen/Qwen3.5-4B \ + --trust-remote-code \ + --skip-tokenizer-init \ + --tp-size 1 \ + --context-length 4096 \ + --mem-fraction-static 0.8 \ + --attention-backend aiter \ + --chunked-prefill-size -1 \ + --disable-radix-cache \ + --enable-spec-capture --spec-capture-method dflash \ + --spec-capture-aux-layer-ids 1 8 15 22 29 \ + --host 127.0.0.1 --port 30000 & +``` + +Wait for `curl --fail http://127.0.0.1:30000/health` to return 200 (the first +health check can take a few minutes while AITER kernels compile). +`--context-length` must exceed `data.max_length` (2048) or `/generate` returns +`400 input longer than context length`. + +#### Step 3: Launch training + +One command supervises producer and consumer on GPU 1: + +```bash +CUDA_VISIBLE_DEVICES=1 HIP_VISIBLE_DEVICES=1 \ +MOONCAKE_LOCAL_HOSTNAME=127.0.0.1 \ +MOONCAKE_METADATA_SERVER=http://127.0.0.1:35880/metadata \ +MOONCAKE_MASTER_SERVER_ADDR=127.0.0.1:35551 \ +MOONCAKE_PROTOCOL=tcp \ +MOONCAKE_GLOBAL_SEGMENT_SIZE=$((32<<30)) \ +specforge train -c examples/configs/online/disaggregated/external/qwen3.5-4b-dflash-online-amd.yaml \ + training.max_steps=20 training.num_epochs=1 \ + training.save_interval=20 training.log_interval=5 +``` + +The trainer needs no AITER env — it runs `flex_attention` on ROCm; only the +capture server (Step 2) drives the Mamba target. Before rerunning, clear stale +control state: `rm -rf outputs/qwen3.5-4b-dflash-online`. + +#### Success criteria + +- Producer log: `drive_producer returning produced= prompts_failed=0`. +- Consumer log: `step N: {...loss..., acc...}` lines, and **no** + `could not drain` error or traceback at teardown. +- Checkpoint: `outputs/qwen3.5-4b-dflash-online/qwen3.5-4b-dflash-online-step20/` + contains `training_state.pt` and `training_state_rank0.pt`. + +> If `produced=0`, the capture aux-layer ids do not match the producer contract +> (see Step 2). If training succeeds but teardown reports `could not drain`, the +> Mooncake lease TTL is above the drain window (see Step 2). + +### Managed-local shortcut + +Instead of starting Mooncake and the capture server by hand, a +`deployment.disaggregated.managed_local` block lets one `specforge train` +command own those local processes and derive their endpoints. It defaults +`default_kv_lease_ttl_ms` to 500, so the lease-TTL fix is applied automatically. +Managed-local owns one local process tree and cannot be launched with +`--role producer` or `--role consumer` or split across nodes. +See [Multi-server capture](../disaggregated_training.md#multi-server-capture) +for the managed-local profile. + +### Split an external run across process pools or nodes + +This is the split-pool form of the external workflow above, not another +training mode. Both roles must use the same resolved run contract, but each is +launched explicitly: + +```bash +# Inference / capture pool +specforge train -c examples/configs/online/disaggregated/external/qwen3.5-4b-dflash-online-amd.yaml --role producer + +# Trainer pool +specforge train -c examples/configs/online/disaggregated/external/qwen3.5-4b-dflash-online-amd.yaml --role consumer +``` + +The checked-in AMD recipe is a single-node example and deliberately uses +`127.0.0.1`. For separate hosts, copy it to `run.yaml` and replace the loopback +Mooncake and SGLang endpoints with addresses reachable by the relevant roles. +Keep these values consistent across the deployment: + +| Shared run contract | Node-local values | +| --- | --- | +| `run_id`, data/training settings, `store_id`, capture contract, routable `server_urls`, and Mooncake metadata/master endpoints | `MOONCAKE_LOCAL_HOSTNAME` and GPU visibility | +| `control_dir`, visible to the producer and consumers unless an inbox relay is configured | `consumer_state_dir`, on reliable local storage for consumer rank 0 | +| `output_dir`, visible to every consumer rank, and the complete `deployment.trainer` topology | `--node-rank` on each consumer host | + +For multiple consumer nodes, record `deployment.trainer.nnodes`, +`nproc_per_node`, `master_addr`, and `master_port` once in the config, then pass +only the node-local identity on each trainer host: + +```bash +specforge train -c run.yaml --role consumer --node-rank 0 # trainer-0 +specforge train -c run.yaml --role consumer --node-rank 1 # trainer-1 +``` + +A fresh attempt requires fresh control and consumer-state directories, and every +capture server must use the same target model, revision, capture method, and +auxiliary layer ids — for this recipe `--spec-capture-method dflash` with aux ids +`1 8 15 22 29`, and on ROCm each must run under AITER with `--disable-radix-cache` +(see External services, Step 2). For external-service prerequisites, freshness +rules, multi-server capture, inbox relays, and resume, see the +[Disaggregated training](../disaggregated_training.md) guide. + +--- + +## Reference results on MI355X + +The offline and online paths were run end-to-end on a single AMD Instinct +**MI355X** (gfx950) inside the `lmsysorg/sglang:v0.5.14-rocm720-mi35x` container, +training a **Qwen3.5-4B DFlash** draft on ShareGPT. Qwen3.5-4B is a hybrid +linear-attention/Mamba target (`Qwen3_5ForConditionalGeneration`); its draft is a +5-layer DFlash head (`hidden_size=2560`, `block_size=16`, +`target_layer_ids=[1, 8, 15, 22, 29]`). Both runs used `max_length=2048`, +`chat_template=qwen3.5`, `batch_size=2`, `accumulation_steps=4`, +`learning_rate=6e-4`, `num_anchors=512`, `loss_decay_gamma=7`, a `flex_attention` +trainer, and ~10 epochs (~680 optimizer steps). + +The SGLang side (offline capture and online capture server) runs the hybrid +Mamba target under **AITER** on ROCm — export +`SGLANG_USE_AITER=1 SGLANG_USE_AITER_UNIFIED_ATTN=1 AITER_FLYDSL_FORCE=1` and use +`--attention-backend aiter`. A ROCm-specific requirement: the target needs +the radix cache disabled (offline capture: `--sglang-disable-radix-cache`; +external online server: `--disable-radix-cache`; managed-local config: +`model.sglang_disable_radix_cache: true`). SGLang's Mamba radix cache +auto-selects the `extra_buffer` strategy, which asserts CUDA/MUSA/NPU (FLA) at +server init and fails on ROCm; disabling the radix cache bypasses that path. +Offline consumes hidden states captured to disk by +`prepare_hidden_states.py`; online consumes the same features streamed live +from the AITER capture server through Mooncake. Both paths converge together — +the online capture path reproduces offline quality on ROCm. + +### Training loss + +![Qwen3.5-4B DFlash training loss on MI355X](imgs/mi355x_qwen35_4b_dflash_loss.png) + +Draft loss falls from ~9 to ~5.6 over ~680 steps. Faint lines are raw per-step +values; bold lines are an exponential moving average. + +### Draft accuracy + +![Qwen3.5-4B DFlash draft accuracy on MI355X](imgs/mi355x_qwen35_4b_dflash_acc.png) + +Top-1 draft-token accuracy (`acc`) — the training-time proxy for serving-time +acceptance — rises from ~0.03 to ~0.12–0.13 and the two paths track each other +closely. + +### Summary + +| Metric (final) | Offline | Online | +| --- | --- | --- | +| Draft loss (start → end) | 8.3 → 5.6 | 9.1 → 5.7 | +| Top-1 draft accuracy (`acc`) | ~0.12 (peak ~0.22) | ~0.13 (peak ~0.17) | +| Epochs / steps | 10 / 687 | 10 / 670 | + +**Throughput** (single MI355X, `batch_size=2`, `max_length=2048`): + +- **Offline** capture: the AITER server generated hidden states for 572 prompts + (286 batches) in ~48 s (~8 batches/s). The GPU-local trainer then ran at + ~0.3 steps/s — sequences up to 2,048 tokens on a 4B target are much heavier + than a small draft at short context. +- **Online** trainer: ~1.3 steps/s end-to-end (670 steps in ~520 s) with the + capture server on GPU 0 and the trainer on GPU 1. A single AITER capture server + produced 5,410 prompts across 10 epochs with **0 failures** (~10 prompts/s); + the single-command `managed_local` stack (Mooncake master + capture server + + trainer) came up and tore down cleanly (`default_kv_lease_ttl_ms=500`). + +> **Data note (offline only):** `prepare_hidden_states.py` truncates each rendered +> conversation at `max_length`. Long-prompt samples whose assistant reply is +> pushed past the cutoff end up with an empty loss region, i.e. fewer than two +> anchorable tokens, which trips DFlash's anchor sampler +> (`ValueError: should preprocess the data.`). Drop those captured samples (any +> with `< 2` loss-mask tokens before the last `block_size` positions) before +> training. The online path never hits this — its producer regenerates full-length +> responses, so every streamed sample has a non-empty loss region. + +These numbers are a functional reference for a 4B DFlash draft on ROCm, not a +tuned performance benchmark — longer sequences and multi-GPU trainers scale +differently. + +--- + +## Reference results on MI300X + +The same **Qwen3.5-4B DFlash** recipe was reproduced end-to-end on a single AMD +Instinct **MI300X** (gfx942) inside the `lmsysorg/sglang:v0.5.14-rocm720-mi30x` +container, using identical hyperparameters (offline and online, `max_length=2048`, +`chat_template=qwen3.5`, `batch_size=2`, `accumulation_steps=4`, +`learning_rate=6e-4`, `num_anchors=512`, `loss_decay_gamma=7`, `flex_attention` +trainer, ~10 epochs). The ROCm requirements are the same as on MI355X: run the +hybrid Mamba target under **AITER** +(`SGLANG_USE_AITER=1 SGLANG_USE_AITER_UNIFIED_ATTN=1 AITER_FLYDSL_FORCE=1`, +`--attention-backend aiter`) and **`--disable-radix-cache`** to bypass the +`extra_buffer` Mamba radix-cache FLA assertion. + +### Training loss + +![Qwen3.5-4B DFlash training loss on MI300X](imgs/mi300x_qwen35_4b_dflash_loss.png) + +Draft loss falls from ~8.4 to ~5.4 over ~680 steps; faint lines are raw per-step +values, bold lines an exponential moving average. + +### Draft accuracy + +![Qwen3.5-4B DFlash draft accuracy on MI300X](imgs/mi300x_qwen35_4b_dflash_acc.png) + +Top-1 draft-token accuracy (`acc`) climbs from ~0.02 to ~0.14, and the offline and +online paths converge to the same quality — matching the MI355X result. + +### Summary + +| Metric (final) | Offline | Online | +| --- | --- | --- | +| Draft loss (start → end) | 8.3 → 5.4 | 8.5 → 5.5 | +| Top-1 draft accuracy (`acc`) | ~0.14 (peak ~0.21) | ~0.14 (peak ~0.16) | +| Epochs / steps | 10 / 687 | 10 / 666 | + +**Throughput** (single MI300X, `batch_size=2`, `max_length=2048`): + +- **Offline** capture: the AITER server captured hidden states for all 572 + prompts; 21 truncated samples with an empty loss region were dropped (see the + data note below), leaving 551 for training. +- **Online** trainer: 666 steps in ~858 s (~0.78 steps/s) with the capture server + on GPU 0 and the trainer on GPU 1. A single AITER capture server produced 5,330 + prompts across 10 epochs with **0 failures** (~6 prompts/s) and streamed 15,990 + feature objects through Mooncake; the single-command `managed_local` stack came + up and tore down cleanly (`default_kv_lease_ttl_ms=500`). + +> **Data note (offline only):** identical to the MI355X run — the offline capture +> produced the same 21 empty-loss-region samples (mostly `max_length`-truncated +> conversations), which must be dropped before training or DFlash's anchor sampler +> raises `ValueError: should preprocess the data.`. The online path never hits this. + +Results on MI300X track MI355X closely, confirming the ROCm DFlash flow (AITER + +`--disable-radix-cache`) is portable across gfx942 and gfx950. diff --git a/docs/basic_usage/AMD/imgs/mi300x_qwen35_4b_dflash_acc.png b/docs/basic_usage/AMD/imgs/mi300x_qwen35_4b_dflash_acc.png new file mode 100644 index 000000000..5937cae29 Binary files /dev/null and b/docs/basic_usage/AMD/imgs/mi300x_qwen35_4b_dflash_acc.png differ diff --git a/docs/basic_usage/AMD/imgs/mi300x_qwen35_4b_dflash_loss.png b/docs/basic_usage/AMD/imgs/mi300x_qwen35_4b_dflash_loss.png new file mode 100644 index 000000000..d93c01a90 Binary files /dev/null and b/docs/basic_usage/AMD/imgs/mi300x_qwen35_4b_dflash_loss.png differ diff --git a/docs/basic_usage/AMD/imgs/mi355x_qwen35_4b_dflash_acc.png b/docs/basic_usage/AMD/imgs/mi355x_qwen35_4b_dflash_acc.png new file mode 100644 index 000000000..cd5a2007d Binary files /dev/null and b/docs/basic_usage/AMD/imgs/mi355x_qwen35_4b_dflash_acc.png differ diff --git a/docs/basic_usage/AMD/imgs/mi355x_qwen35_4b_dflash_loss.png b/docs/basic_usage/AMD/imgs/mi355x_qwen35_4b_dflash_loss.png new file mode 100644 index 000000000..e6be59c79 Binary files /dev/null and b/docs/basic_usage/AMD/imgs/mi355x_qwen35_4b_dflash_loss.png differ diff --git a/docs/basic_usage/Ascend/ascend_npu.md b/docs/basic_usage/Ascend/ascend_npu.md new file mode 100644 index 000000000..4c73522af --- /dev/null +++ b/docs/basic_usage/Ascend/ascend_npu.md @@ -0,0 +1,262 @@ +# Ascend NPU Tutorial + +This is an end-to-end tutorial for running SpecForge on Ascend NPU hosts. It +walks through **installation → data preparation → online disaggregated training +with external services → the managed-local full stack → split multi-node +roles**, using Qwen3.5-4B DFlash as the running example. Validated on a 16-card +A3 (64GB) host. + +--- + +## 1. Installation + +You need an Ascend host with the driver, CANN, and a `torch_npu`-enabled +PyTorch already installed, plus SGLang `0.5.14` with NPU support. Then install +SpecForge without touching that stack: + +```bash +git clone https://github.com/sgl-project/SpecForge.git +cd SpecForge +python -m pip install -e . --no-deps +``` + +`--no-deps` keeps pip from pulling CUDA wheels over the working NPU +torch/sglang. If a later step reports a missing lightweight dependency, install +just that package, also with `--no-deps`. + +### Apply the SGLang capture patches (online runs only) + +Online capture needs two patches on top of the installed SGLang, applied **in +this order**: + +```bash +# Base capture patch: --enable-spec-capture server flags + the capture sink +bash scripts/apply_sglang_spec_capture_patch.sh + +# Ascend companion patch: skip the wildcard segment mount that Ascend +# Mooncake rejects, and mount the feature segment with location="cpu" +SGLANG_DIR=$(python -c "import sglang, os; print(os.path.dirname(os.path.dirname(sglang.__file__)))") +cd "$SGLANG_DIR" && git apply /path/to/SpecForge/patches/sglang/v0.5.14/spec-capture-ascend-mount.patch +``` + +Skip both for offline training, which reads features from disk. The companion +patch is a no-op on non-Ascend hosts (it keys off `ASCEND_RT_VISIBLE_DEVICES`), +so a shared installation stays CUDA-safe. + +### Device visibility on Ascend + +Ascend selects devices through `ASCEND_RT_VISIBLE_DEVICES`, and the driver +rejects an *empty* value — hiding devices from a process means **unsetting** +the variable, not setting it to `""`. SpecForge handles this internally: device +ordinals from the config are injected through `CUDA_VISIBLE_DEVICES` on CUDA +hosts and `ASCEND_RT_VISIBLE_DEVICES` on Ascend hosts, and a "hide all devices" +role unsets the variable. You only need to export the visible set once for the +supervisor: + +```bash +export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 +``` + +Two more environment settings are recommended for long online runs: + +```bash +export HCCL_CONNECT_TIMEOUT=7200 HCCL_EXEC_TIMEOUT=7200 +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +``` + +### Attention backends on Ascend + +Use `sdpa` for the **trainer**. Sequence parallelism (`sp_ulysses_size` / +`sp_ring_size` > 1, the `usp` backend) requires `yunchang`, whose import probes +the CUDA device and crashes NPU-only torch builds, so USP is CUDA-only for now. +SpecForge imports `yunchang` lazily only when SP sizes exceed 1, so the default +SP=1 path never touches it. + +For the **capture server**, the recipes set `model.sglang_attention_backend: +ascend`. If a config leaves it at the flashinfer default, the launcher falls +back to `ascend` automatically on Ascend hosts (flashinfer does not exist +there). + +--- + +## 2. Data preparation + +Data preparation is platform independent. From the repository root: + +```bash +python scripts/prepare_data.py --dataset sharegpt +``` + +This writes `./cache/dataset/sharegpt_train.jsonl`. For custom datasets and +target-model regeneration, see the [Data Preparation](../data_preparation.md) +guide. + +--- + +## 3. Online training (external capture server) + +Online training is always **disaggregated**: a producer drives prompts through +a patched SGLang capture server, features stream through Mooncake, and a +consumer trains the draft. The checked-in +[`qwen3.5-4b-dflash-online-npu.yaml`](../../../examples/configs/online/disaggregated/external/qwen3.5-4b-dflash-online-npu.yaml) +recipe targets an externally started capture server. + +### Step 1: Start Mooncake and the capture server + +```bash +mooncake_master --enable_http_metadata_server=true \ + --rpc_port=35551 --http_metadata_server_port=35880 \ + --metrics_port=35903 --enable_metric_reporting=false & +``` + +Start the patched capture server on NPU 0. The `--spec-capture-aux-layer-ids` +must match the draft's `target_layer_ids` — for DFlash these come from +`configs/qwen3.5-4b-dflash.json`: `1 8 15 22 29`. A mismatch produces zero +features with no error: + +```bash +ASCEND_RT_VISIBLE_DEVICES=0 \ +MOONCAKE_LOCAL_HOSTNAME=127.0.0.1 \ +MOONCAKE_METADATA_SERVER=http://127.0.0.1:35880/metadata \ +MOONCAKE_MASTER_SERVER_ADDR=127.0.0.1:35551 \ +MOONCAKE_PROTOCOL=tcp \ +MOONCAKE_GLOBAL_SEGMENT_SIZE=$((32<<30)) \ +python -m sglang.launch_server \ + --model-path Qwen/Qwen3.5-4B \ + --trust-remote-code \ + --skip-tokenizer-init \ + --tp-size 1 \ + --mem-fraction-static 0.8 \ + --attention-backend ascend \ + --enable-spec-capture --spec-capture-method dflash \ + --spec-capture-aux-layer-ids 1 8 15 22 29 \ + --host 127.0.0.1 --port 30000 & +``` + +Wait for `curl --fail http://127.0.0.1:30000/health` to return 200. + +### Step 2: Launch training + +On the remaining NPUs: + +```bash +ASCEND_RT_VISIBLE_DEVICES=1,2,3,4,5,6,7,8 \ +specforge train -c examples/configs/online/disaggregated/external/qwen3.5-4b-dflash-online-npu.yaml +``` + +Before rerunning, clear stale control state: +`rm -rf outputs/qwen3.5-4b-dflash-npu-online`. + +### Success criteria + +- Producer log ends with `prompts_failed=0`. +- Consumer prints `step N: {...loss..., acc...}` lines and no + `could not drain` error at teardown. +- If the capture producer dies with `ACL_ERROR_RT_CONTEXT_NULL` (107002), the + installed SpecForge predates the NPU transport bind — upgrade past #722. + +--- + +## 4. Managed-local full stack (one command) + +Instead of starting Mooncake and capture servers by hand, the +[`qwen3.5-4b-dflash-disaggregated-npu.yaml`](../../../examples/configs/online/disaggregated/managed-local/qwen3.5-4b-dflash-disaggregated-npu.yaml) +recipe lets a single `specforge train` command own the whole single-node +stack — Mooncake, capture server(s), and the trainer — and derives their +endpoints and device assignments: + +```bash +export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 +export HCCL_CONNECT_TIMEOUT=7200 HCCL_EXEC_TIMEOUT=7200 +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True + +specforge train -c examples/configs/online/disaggregated/managed-local/qwen3.5-4b-dflash-disaggregated-npu.yaml +``` + +The checked-in layout parks the capture server on device 0 and runs a 14-rank +trainer on devices 2-15 (8-card hosts: trainer on devices 1-7 with +`deployment.trainer.nproc_per_node=7`). + +### Scaling capture throughput + +One capture server is enough for correctness but bounds the pipeline: the +trainer's data wait then dominates step time. To give capture more cards, +override the layout inline — e.g. a 6:10 split with six TP=1 capture servers: + +```bash +specforge train -c examples/configs/online/disaggregated/managed-local/qwen3.5-4b-dflash-disaggregated-npu.yaml \ + 'deployment.trainer.nproc_per_node=10' \ + 'deployment.disaggregated.managed_local.trainer_cuda_visible_devices=["6","7","8","9","10","11","12","13","14","15"]' \ + 'deployment.disaggregated.managed_local.capture_servers=[{port: 40000, cuda_visible_devices: ["0"], tp_size: 1}, {port: 40001, cuda_visible_devices: ["1"], tp_size: 1}, {port: 40002, cuda_visible_devices: ["2"], tp_size: 1}, {port: 40003, cuda_visible_devices: ["3"], tp_size: 1}, {port: 40004, cuda_visible_devices: ["4"], tp_size: 1}, {port: 40005, cuda_visible_devices: ["5"], tp_size: 1}]' +``` + +In the validated 16-card run this cut step time from ~20s to ~4.5s. + +### Cleanup between runs + +Managed children exit with the supervisor, but after an interrupted run check +for strays before relaunching: + +```bash +pkill -9 -f specforge; pkill -9 -f mooncake_master; pkill -9 -f torch.distributed.run +rm -rf outputs/qwen3.5-4b-dflash-npu-managed +``` + +--- + +## 5. Split producer and consumer roles across nodes + +The same configs split across nodes with an explicit `--role`: + +```bash +specforge train -c run.yaml --role producer # inference / capture pool +specforge train -c run.yaml --role consumer --node-rank 0 # trainer-0 +specforge train -c run.yaml --role consumer --node-rank 1 # trainer-1 +``` + +Every capture server must use the same target model, capture method, and aux +layer ids. For external-service prerequisites, freshness rules, and resume, see +the [Disaggregated training](../disaggregated_training.md) guide. + +--- + +## 6. NPU notes and troubleshooting + +- **Empty `ASCEND_RT_VISIBLE_DEVICES` is invalid** — the Ascend driver rejects + it. SpecForge unsets the variable instead of emptying it; do not export + `ASCEND_RT_VISIBLE_DEVICES=` by hand. +- **USP is CUDA-only for now** — `yunchang` probes CUDA at import. Keep + `sp_ulysses_size` / `sp_ring_size` at 1 and use the `sdpa` trainer backend. +- **Ascend Mooncake rejects wildcard buffer registration** — the trainer side + forces `local_buffer_size=0` automatically (SpecForge roles are pure + zero-copy clients), and the capture side needs the companion patch from + Section 1. +- **`ACL_ERROR_RT_CONTEXT_NULL` in the capture producer** — the Mooncake + transfer engine needs a bound device context; SpecForge binds the local NPU + before `setup()`. Seeing this error means the installation predates #722. +- **Teardown drain** — the lifecycle drain window (~20s) covers Mooncake's + read-lease TTL, and managed-local masters start with + `default_kv_lease_ttl_ms=500`. A `could not drain pending removals` error + after an otherwise successful run means one of these is missing — upgrade + SpecForge. + +--- + +## Reference results on 16x A3 (64GB) + +End-to-end online run on the managed-local stack, training a **Qwen3.5-4B +DFlash** draft (`configs/qwen3.5-4b-dflash.json`) with 6 TP=1 capture servers +(devices 0-5) and 10 trainer ranks (devices 6-15): ~9.5k prompts x 10 epochs, +global batch 80/step, ~1.18k optimizer steps, ~1.5h total. + +- Loss falls from ~6.0 to ~2.0 over the run with cosine LR (6e-4 -> 0); no + NaN, no stall. +- ~4.5s/optimizer step with six capture servers (~20s/step with one). +- Memory: ~29GB per trainer rank, ~31-33GB per capture server; no OOM with + `num_anchors: 512` (lower it on smaller cards). +- Clean terminal drain on shutdown. + +![Qwen3.5-4B DFlash online training on 16x A3](https://github.com/user-attachments/assets/1b94b455-cb61-4374-a548-9a219a37ed64) + +These numbers are a functional reference for the NPU managed-local stack, not +a tuned performance benchmark. diff --git a/docs/basic_usage/data_preparation.md b/docs/basic_usage/data_preparation.md index 3537309d1..d156f6e19 100644 --- a/docs/basic_usage/data_preparation.md +++ b/docs/basic_usage/data_preparation.md @@ -247,7 +247,7 @@ to create the text. SpecForge uses it to identify assistant spans and build the loss mask. ```bash -# After copying a disaggregated online example YAML, set these data fields: +# After copying an online disaggregated example YAML, set these data fields: # data: # train_data_path: ./your_preformatted_dataset.jsonl # is_preformatted: true @@ -277,7 +277,7 @@ torchrun --nproc_per_node=8 \ --batch-size 32 ``` -Use these strategy/model pairs for the checked-in local offline recipes: +Use these strategy/model pairs for the checked-in offline colocated recipes: | Strategy | Target model | Draft config | Output path used by the recipe | | --- | --- | --- | --- | @@ -319,7 +319,7 @@ torchrun --nproc_per_node=8 \ scripts/prepare_hidden_states.py \ --strategy eagle3 \ --target-model-path meta-llama/Llama-3.1-8B-Instruct \ - --draft-model-config configs/llama3.1-8b-eagle3.json \ + --draft-model-config configs/llama3-8B-eagle3.json \ --data-path ./your_preformatted_dataset.jsonl \ --output-path ./cache/hidden_states/llama3.1-8b-eagle3 \ --chat-template llama3 \ @@ -334,7 +334,7 @@ Launch the matching recipe after its `data.hidden_states_path` points at the generated directory: ```bash -specforge train --config examples/configs/qwen3-8b-dflash-offline.yaml +specforge train --config examples/configs/offline/colocated/qwen3-8b-dflash-offline.yaml ``` See the [Training](training.md) guide for the complete run schema and supported diff --git a/docs/basic_usage/disaggregated_training.md b/docs/basic_usage/disaggregated_training.md index b4d5f9cb3..0b84373c3 100644 --- a/docs/basic_usage/disaggregated_training.md +++ b/docs/basic_usage/disaggregated_training.md @@ -1,31 +1,39 @@ # Disaggregated training -Disaggregation is a launch topology of the canonical training command: +Disaggregation is a producer/consumer launch topology of the canonical training +command: ```bash specforge train -c run.yaml ``` The producer captures or ingests features and the consumer runs the canonical -trainer. Online training always uses this producer/consumer topology; there is -no colocated target-inference path and no separate Python training entry. - -The checked-in recipes are: - -| Workflow | Config | -| --- | --- | -| Online EAGLE3 | `examples/configs/qwen3-8b-eagle3-disaggregated.yaml` | -| Online P-EAGLE | `examples/configs/qwen3-8b-peagle-disaggregated.yaml` | -| Online DFlash | `examples/configs/qwen3-8b-dflash-disaggregated.yaml` | -| Online one-server + DP7 Qwen3-8B DFlash | `examples/configs/qwen3-8b-dflash-1server-dp7-disaggregated.yaml` | -| Online Domino | `examples/configs/qwen3-8b-domino-disaggregated.yaml` | -| Online two-server Qwen3-8B Domino | `examples/configs/qwen3-8b-domino-multiserver-disaggregated.yaml` | -| Online DSpark | `examples/configs/qwen3-4b-dspark-disaggregated.yaml` | -| Online Qwen3.6 DFlash | `examples/configs/qwen3.6-27b-dflash-disaggregated.yaml` | -| Online one-server + DP2 Qwen3.6 DFlash | `examples/configs/qwen3.6-27b-dflash-1server-dp2-disaggregated.yaml` | -| Online two-server Qwen3.6 DFlash | `examples/configs/qwen3.6-27b-dflash-multiserver-disaggregated.yaml` | -| Offline EAGLE3 | `examples/configs/qwen3-8b-eagle3-offline-disaggregated.yaml` | -| Offline Qwen2.5-7B EAGLE3 | `examples/configs/qwen2.5-7b-eagle3-offline-disaggregated.yaml` | +trainer. + +| Category | Producer responsibility | Feature store | Service ownership | +| --- | --- | --- | --- | +| Offline disaggregated | Ingest existing feature files and publish a static manifest | `shared_dir` or Mooncake | No SGLang lifecycle; configure the storage backend directly | +| Online disaggregated, external | Send prompts to existing capture endpoints and publish streaming refs | Mooncake | User or scheduler owns Mooncake and SGLang | +| Online disaggregated, managed-local | Send prompts to endpoints started from the same run config | Mooncake | SpecForge owns local Mooncake and SGLang | + +Online training always uses the producer/consumer topology; there is currently +no colocated online target-inference path. `external` describes ownership, not +distance—a loopback service started by the user is still external. Conversely, +`managed-local` still uses distinct producer and consumer roles even though one +supervisor owns the local process tree. + +Choose checked-in recipes by directory: + +| Workflow | Catalog | Representative config | +| --- | --- | --- | +| Offline disaggregated | [`offline/disaggregated/`](../../examples/configs/offline/disaggregated/) | [`qwen3-8b-eagle3-offline-disaggregated.yaml`](../../examples/configs/offline/disaggregated/qwen3-8b-eagle3-offline-disaggregated.yaml) | +| Online, external services | [`online/disaggregated/external/`](../../examples/configs/online/disaggregated/external/) | [`qwen3-8b-dflash-disaggregated.yaml`](../../examples/configs/online/disaggregated/external/qwen3-8b-dflash-disaggregated.yaml) | +| Online, managed-local stack | [`online/disaggregated/managed-local/`](../../examples/configs/online/disaggregated/managed-local/) | [`qwen3-8b-dflash-1server-dp7-disaggregated.yaml`](../../examples/configs/online/disaggregated/managed-local/qwen3-8b-dflash-1server-dp7-disaggregated.yaml) | + +The full model and strategy index is in the +[recipe catalog](../../examples/configs/README.md). Filename suffixes are +historical identifiers; the directory and typed YAML fields define the runtime +semantics. ## One config owns the launch topology @@ -58,9 +66,12 @@ The control directory is attempt-scoped. The launcher deterministically derives the online reference channel and lifecycle markers, or the offline manifest, beneath that root. Online consumers put the rank-0 SQLite/WAL under `consumer_state_dir`; that path should be node-local and is required for -multi-node trainers. Rank inboxes stay under the shared `control_dir` when -`deployment.trainer.nnodes > 1`, so remote ranks never need to access the -SQLite filesystem. A fresh attempt requires fresh control and consumer-state +multi-node trainers. By default rank inboxes stay under the shared +`control_dir`. When a cluster cannot mount that path on every trainer node, set +`inbox_server_url` to a private rank-0 HTTP origin: rank 0 keeps the inbox files +locally and relays only tensor-free references and consumed counters to remote +ranks. The producer and consumer rank 0 must still resolve `control_dir` to the +same path. A fresh attempt requires fresh control and consumer-state directories. `store_id` defaults to `run_id` and can be set explicitly when the Mooncake deployment requires another namespace. @@ -81,14 +92,14 @@ With `deployment.trainer.nnodes: 1`, omitting `--role` starts a supervisor for both SpecForge roles: ```bash -specforge train -c examples/configs/qwen3-8b-dflash-disaggregated.yaml +specforge train -c examples/configs/online/disaggregated/external/qwen3-8b-dflash-disaggregated.yaml ``` -The checked-in online recipes use the local demo endpoints shown above, so the -command and `--plan` need no hidden topology variables. Point those typed fields -at the real services for a remote deployment. Environment values override the -typed Mooncake endpoint fields when the same recipe runs on another node; -`MOONCAKE_LOCAL_HOSTNAME` remains node-local. +The checked-in external online recipes use the local demo endpoints shown +above, so the command and `--plan` need no hidden topology variables. Point +those typed fields at the actual services for another deployment. Environment +values override the typed Mooncake endpoint fields when the same recipe runs on +another node; `MOONCAKE_LOCAL_HOSTNAME` remains node-local. The producer is a direct single process. The consumer is automatically launched with the configured local process count. If either role fails, the supervisor @@ -99,7 +110,7 @@ canceling it. Inspect the resolved plan without starting processes: ```bash -specforge train -c examples/configs/qwen3-8b-dflash-disaggregated.yaml --plan +specforge train -c examples/configs/online/disaggregated/external/qwen3-8b-dflash-disaggregated.yaml --plan ``` Plan output redacts secret-shaped overrides and credentials embedded in URLs. @@ -141,7 +152,7 @@ four producer workers against one server while keeping the checked-in YAML as the source of every other setting: ```bash -specforge train -c examples/configs/qwen3-8b-domino-disaggregated.yaml \ +specforge train -c examples/configs/online/disaggregated/external/qwen3-8b-domino-disaggregated.yaml \ 'deployment.disaggregated.server_urls=["http://127.0.0.1:30000","http://127.0.0.1:30000","http://127.0.0.1:30000","http://127.0.0.1:30000"]' ``` @@ -161,10 +172,10 @@ TP=1 capture server on GPU 0 and a DP=7 trainer on GPUs 1–7: ```bash specforge train -c \ - examples/configs/qwen3-8b-dflash-1server-dp7-disaggregated.yaml + examples/configs/online/disaggregated/managed-local/qwen3-8b-dflash-1server-dp7-disaggregated.yaml specforge train -c \ - examples/configs/qwen3-8b-domino-1server-dp7-disaggregated.yaml + examples/configs/online/disaggregated/managed-local/qwen3-8b-domino-1server-dp7-disaggregated.yaml ``` The checked-in multi-server Qwen3-8B Domino recipe records two TP=1 capture @@ -174,7 +185,7 @@ the unified training entry: ```bash specforge train -c \ - examples/configs/qwen3-8b-domino-multiserver-disaggregated.yaml + examples/configs/online/disaggregated/managed-local/qwen3-8b-domino-multiserver-disaggregated.yaml ``` The Qwen3.6 DFlash one-server recipe owns a TP=1 server on GPU 0 and a DP=2 @@ -183,10 +194,10 @@ GPUs 0–3 and a DP=2 trainer on GPUs 4–5: ```bash specforge train -c \ - examples/configs/qwen3.6-27b-dflash-1server-dp2-disaggregated.yaml + examples/configs/online/disaggregated/managed-local/qwen3.6-27b-dflash-1server-dp2-disaggregated.yaml specforge train -c \ - examples/configs/qwen3.6-27b-dflash-multiserver-disaggregated.yaml + examples/configs/online/disaggregated/managed-local/qwen3.6-27b-dflash-multiserver-disaggregated.yaml ``` The launcher starts Mooncake first, waits for its metadata and RPC endpoints, @@ -217,7 +228,7 @@ the launcher-provided rank to those same two commands: ```bash rcli exec --per-node \ - 'CONFIG=examples/configs/qwen2.5-7b-eagle3-offline-disaggregated.yaml bash examples/disagg/run_offline_2node.sh' + 'CONFIG=examples/configs/offline/disaggregated/qwen2.5-7b-eagle3-offline-disaggregated.yaml bash examples/disagg/run_offline_2node.sh' ``` `RCLI_NODE_RANK=0` selects the producer and rank 1 selects the consumer. The @@ -236,7 +247,11 @@ deployment: master_addr: trainer-0.example master_port: 29500 disaggregated: - control_dir: /shared/control/attempt-001 + # With inbox_server_url, only the producer and consumer rank 0 need this + # path; remote ranks receive their tensor-free reference stream over HTTP. + control_dir: /local/rank0/control/attempt-001 + consumer_state_dir: /local/rank0/consumer-state/attempt-001 + inbox_server_url: http://trainer-0.example:35900 backend: mooncake server_urls: [http://capture-server:30000] ``` @@ -257,6 +272,13 @@ detected and used as the worker environment rather than nesting another torchrun. A producer is rejected inside a multi-rank torchrun to prevent duplicate capture/ingestion roles. +`inbox_server_url` is optional; omit it when every trainer rank shares +`control_dir`. When set, rank 0 binds the configured port on all interfaces and +remote ranks pull only their private metadata stream. The endpoint has no +built-in authentication or TLS, so expose it only on a trusted cluster network. +The producer must run on the rank-0 host (or otherwise share rank 0's +`control_dir`); payload tensors never traverse this HTTP service. + The separate-inference-node Qwen3-8B example keeps that scheduler boundary but restores full development-stack orchestration: @@ -283,10 +305,12 @@ path. This split-state form currently supports one trainer node only. ## External and managed-local services Without `deployment.disaggregated.managed_local`, the unified supervisor owns -only SpecForge producer and consumer processes. Mooncake and patched SGLang are -external, usually long-lived services managed by Kubernetes, Slurm, systemd, or -the development environment. In this default mode, the training CLI does not -start, stop, or assign GPUs to them. +only the requested SpecForge roles. With the default single-node launch it +starts producer and consumer; with an explicit `--role` it starts only that +role. Mooncake and patched SGLang remain external, usually long-lived services +managed by Kubernetes, Slurm, systemd, or the development environment. They may +also run on the same host: external means that the training CLI does not start, +stop, or assign GPUs to them. Install a Mooncake client compatible with the store's `put_from`/`get_into` wire contract on the producer and consumer images: @@ -315,18 +339,20 @@ node-local deployment values. The online producer sends prompts to the URLs in `deployment.disaggregated.server_urls`. Start a patched SGLang server separately with the model, capture method, and auxiliary layer ids matching the draft -config. DFlash, Domino, and DSpark use the DFlash capture contract; EAGLE3 and -P-EAGLE use the EAGLE3 capture contract. Capture rejects chunked prefill and +config. DFlash and Domino use the DFlash capture contract, DSpark uses its +dedicated K3 capture contract, and EAGLE3 and P-EAGLE use the EAGLE3 capture +contract. Capture rejects chunked prefill and gives every request attempt a unique radix-cache namespace so cached prefixes cannot truncate the captured sequence. Online capture is text-only: VLM training, including Qwen2.5-VL, is not supported. Online evaluation is also not supported. -The repository's strict e2e gate remains a full local test-stack orchestrator: +The `external`/`managed-local` subdivision applies only to online +disaggregation. Offline disaggregation has no capture-server lifecycle and is +instead distinguished by its `shared_dir` or Mooncake feature-store backend. -```bash -bash scripts/gates/run_disaggregated_overfit_gate.sh -``` +For the strict e2e one-sample overfit and serving validation procedure, follow +[`scripts/gates/README.md`](../../scripts/gates/README.md). It starts and health-checks Mooncake and SGLang, invokes the canonical training entry, verifies overfit and serving behavior, and cleans up every process it diff --git a/docs/basic_usage/training.md b/docs/basic_usage/training.md index 8107cf2a9..88fd05045 100644 --- a/docs/basic_usage/training.md +++ b/docs/basic_usage/training.md @@ -4,65 +4,62 @@ SpecForge has one public training entry point for every strategy and runtime topology: ```bash -specforge train --config examples/configs/qwen3-8b-eagle3-disaggregated.yaml +specforge train --config examples/configs/online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml ``` The YAML file is the run contract. It selects the draft strategy, target model, -data source, optimizer settings, and deployment mode. Method-specific Python -trainers are not part of the public interface. +data source, optimizer settings, and deployment topology. Method-specific +Python trainers are not part of the public interface. This is an intentional hard cutover. The old `scripts/train_*.py` commands and temporary move-only Python import paths were removed rather than deprecated. Downstream launchers should migrate to a typed run config and `specforge train`; there is no compatibility dispatch to the previous trainers. -### Defaults when migrating removed trainers +## Choose a recipe -The typed schema defaults existed before the old trainers were removed, but -they are not identical to defaults embedded in every deleted script. If an old -launch omitted these flags, write the legacy value explicitly in its new YAML -when reproducing that run: +The config catalog separates three concepts: -| Removed CLI default | Typed run field and default | Legacy value to preserve | +| Question | Config source | Catalog level | | --- | --- | --- | -| DFlash/Domino `--num-epochs=6` | `training.num_epochs: 1` | `6` | -| DFlash/Domino `--learning-rate=6e-4` | `training.learning_rate: 1e-4` | `6e-4` | -| DFlash/Domino `--warmup-ratio=0.04` | `training.warmup_ratio: 0.015` | `0.04` | -| DFlash/Domino `--max-grad-norm=1.0` | `training.max_grad_norm: 0.5` | `1.0` | -| DFlash/Domino `--max-length=3072` | `data.max_length: 2048` | `3072` | -| DFlash/Domino `--chat-template=qwen` | `data.chat_template: llama3` | `qwen` | -| DFlash/Domino `--save-interval=1000` | `training.save_interval: 0` | `1000` | -| DFlash/Domino `--dist-timeout=30` | `training.dist_timeout: 10` | `30` | -| EAGLE3 `--kl-decay=3.0` | `training.kl_decay: 1.0` | `3.0` | - -The old DFlash/Domino `--eval-interval=1000` did not identify an evaluation -source by itself. In the unified runtime, evaluation is deliberately off by -default and must be paired with `data.eval_hidden_states_path`. - -Two numerical lifecycle details are also deliberate. All unified FSDP methods -keep buffers in float32; the removed EAGLE3 and DFlash scripts used bfloat16 -buffers, while the removed Domino script already used float32. Consequently, -bit-for-bit comparisons to old EAGLE3/DFlash baselines must account for that -dtype change. Also, `global_step`, LR/loss horizons, logging, saving, and Domino -lambda decay are all expressed in completed optimizer updates. Fixed datasets -are validated before backend/optimizer assembly to contain complete accumulation windows; -finite online plans train only complete global optimizer quanta. The old -scripts mixed micro-batch counters with a ceil-derived optimizer horizon, so -accumulation greater than one did not have the same boundary semantics. +| Are target features captured during training or prepared earlier? | `data.train_data_path` / `data.prompts_path` versus `data.hidden_states_path` | `online/` versus `offline/` | +| Does the trainer read offline hidden states files directly or consume refs from a producer? | `deployment.mode: local_colocated` versus `deployment.mode: disaggregated` | `colocated/` versus `disaggregated/` | +| Who starts Mooncake and SGLang for online disaggregation? | Presence of `deployment.disaggregated.managed_local` | `external/` versus `managed-local/` | + +The supported catalog layout is: + +```text +examples/configs/ +├── offline/ +│ ├── colocated/ +│ └── disaggregated/ +└── online/ + ├── colocated/ # reserved; currently unsupported + └── disaggregated/ + ├── external/ + └── managed-local/ +``` + +`external` is an ownership boundary, not a statement that the services are on +another machine. It is `external` when the user starts +the server. `managed-local` owns those services on one host, while producer and +consumer remain separate SpecForge roles. The runtime reads the YAML fields, +not filename suffixes, to determine these semantics. See the complete +[recipe catalog](../../examples/configs/README.md) for representative configs. ## Launch a run Use the command directly for every checked-in topology: ```bash -specforge train --config examples/configs/qwen3-8b-eagle3-disaggregated.yaml +specforge train --config examples/configs/online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml ``` `deployment.trainer.nproc_per_node` records the audited local process count. When it is greater than one, the CLI starts torch distributed itself: ```bash -specforge train -c examples/configs/qwen3-30b-a3b-eagle3-online.yaml +specforge train -c examples/configs/online/disaggregated/external/qwen3-30b-a3b-eagle3-online.yaml ``` Online target inference never runs in the trainer. A patched SGLang server owns @@ -79,7 +76,7 @@ validated `section.field=value` syntax: ```bash specforge train \ - --config examples/configs/qwen3-8b-eagle3-disaggregated.yaml \ + --config examples/configs/online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml \ training.learning_rate=5e-5 \ training.max_steps=100 \ output_dir=./outputs/eagle3-smoke @@ -154,6 +151,31 @@ model: draft_block_size: 8 # DFlash only ``` +For DFlash, configure the attention layout in the referenced draft JSON. Each +entry corresponds to one draft layer; sliding layers share one positive window: + +```json +{ + "num_hidden_layers": 5, + "layer_types": [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention" + ], + "use_sliding_window": true, + "sliding_window": 2048 +} +``` + +Use `"full_attention"` for every entry, `"use_sliding_window": false`, and +`"sliding_window": null` for a full-only draft. The layout length must equal +`num_hidden_layers`. A layer-count override may resize a uniform layout, but a +mixed layout must be edited explicitly in the draft JSON. + +The `eager`, `sdpa`, and `flex_attention` backends support both layouts. + Domino and DSpark need their projector/head metadata, so they require an explicit draft config (or a pretrained warm-start source that contains `config.json`). The old Domino parser exposed an optional config flag, but its @@ -193,22 +215,22 @@ Set exactly one data source: The checked-in examples are the canonical starting points: -| Strategy and mode | Config | -| --- | --- | -| EAGLE3 online | [`qwen3-8b-eagle3-disaggregated.yaml`](../../examples/configs/qwen3-8b-eagle3-disaggregated.yaml) | -| EAGLE3 offline | [`qwen3-8b-eagle3-offline.yaml`](../../examples/configs/qwen3-8b-eagle3-offline.yaml) | -| DFlash online | [`qwen3-8b-dflash-online.yaml`](../../examples/configs/qwen3-8b-dflash-online.yaml) | -| DFlash offline | [`qwen3-8b-dflash-offline.yaml`](../../examples/configs/qwen3-8b-dflash-offline.yaml) | -| Domino online | [`qwen3-8b-domino-online.yaml`](../../examples/configs/qwen3-8b-domino-online.yaml) | -| Domino offline | [`qwen3-8b-domino-offline.yaml`](../../examples/configs/qwen3-8b-domino-offline.yaml) | -| P-EAGLE online | [`qwen3-8b-peagle-disaggregated.yaml`](../../examples/configs/qwen3-8b-peagle-disaggregated.yaml) | -| DFlash disaggregated | [`qwen3-8b-dflash-disaggregated.yaml`](../../examples/configs/qwen3-8b-dflash-disaggregated.yaml) | -| Domino disaggregated | [`qwen3-8b-domino-disaggregated.yaml`](../../examples/configs/qwen3-8b-domino-disaggregated.yaml) | -| DSpark disaggregated | [`qwen3-4b-dspark-disaggregated.yaml`](../../examples/configs/qwen3-4b-dspark-disaggregated.yaml) | -| DSpark offline | [`qwen3-4b-dspark-offline.yaml`](../../examples/configs/qwen3-4b-dspark-offline.yaml) | -| EAGLE3 offline disaggregated | [`qwen3-8b-eagle3-offline-disaggregated.yaml`](../../examples/configs/qwen3-8b-eagle3-offline-disaggregated.yaml) | -| Ascend NPU DFlash online | [`qwen3.5-4b-dflash-online-npu.yaml`](../../examples/configs/qwen3.5-4b-dflash-online-npu.yaml) | -| Ascend NPU Domino online | [`qwen3.5-4b-domino-online-npu.yaml`](../../examples/configs/qwen3.5-4b-domino-online-npu.yaml) | +| Strategy | Category | Config | +| --- | --- | --- | +| EAGLE3 | Online disaggregated, external | [`qwen3-8b-eagle3-disaggregated.yaml`](../../examples/configs/online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml) | +| EAGLE3 | Offline colocated | [`qwen3-8b-eagle3-offline.yaml`](../../examples/configs/offline/colocated/qwen3-8b-eagle3-offline.yaml) | +| EAGLE3 | Offline disaggregated | [`qwen3-8b-eagle3-offline-disaggregated.yaml`](../../examples/configs/offline/disaggregated/qwen3-8b-eagle3-offline-disaggregated.yaml) | +| P-EAGLE | Online disaggregated, external | [`qwen3-8b-peagle-disaggregated.yaml`](../../examples/configs/online/disaggregated/external/qwen3-8b-peagle-disaggregated.yaml) | +| DFlash | Online disaggregated, external | [`qwen3-8b-dflash-online.yaml`](../../examples/configs/online/disaggregated/external/qwen3-8b-dflash-online.yaml) | +| DFlash | Online disaggregated, managed-local | [`qwen3-8b-dflash-1server-dp7-disaggregated.yaml`](../../examples/configs/online/disaggregated/managed-local/qwen3-8b-dflash-1server-dp7-disaggregated.yaml) | +| DFlash | Offline colocated | [`qwen3-8b-dflash-offline.yaml`](../../examples/configs/offline/colocated/qwen3-8b-dflash-offline.yaml) | +| Domino | Online disaggregated, external | [`qwen3-8b-domino-online.yaml`](../../examples/configs/online/disaggregated/external/qwen3-8b-domino-online.yaml) | +| Domino | Online disaggregated, managed-local | [`qwen3-8b-domino-multiserver-disaggregated.yaml`](../../examples/configs/online/disaggregated/managed-local/qwen3-8b-domino-multiserver-disaggregated.yaml) | +| Domino | Offline colocated | [`qwen3-8b-domino-offline.yaml`](../../examples/configs/offline/colocated/qwen3-8b-domino-offline.yaml) | +| DSpark | Online disaggregated, external | [`qwen3-4b-dspark-disaggregated.yaml`](../../examples/configs/online/disaggregated/external/qwen3-4b-dspark-disaggregated.yaml) | +| DSpark | Offline colocated | [`qwen3-4b-dspark-offline.yaml`](../../examples/configs/offline/colocated/qwen3-4b-dspark-offline.yaml) | +| DFlash (Ascend) | Online disaggregated, external | [`qwen3.5-4b-dflash-online-npu.yaml`](../../examples/configs/online/disaggregated/external/qwen3.5-4b-dflash-online-npu.yaml) | +| MTP (Ascend) | Online disaggregated, managed-local | [`qwen3.5-4b-mtp-disaggregated-npu.yaml`](../../examples/configs/online/disaggregated/managed-local/qwen3.5-4b-mtp-disaggregated-npu.yaml) | ## Online and offline data @@ -220,7 +242,7 @@ more storage. | Mode | Target during training | Disk use | Data config | | --- | --- | --- | --- | -| Online | External/managed SGLang capture server | Low | `train_data_path` or `prompts_path` | +| Online | External or managed-local SGLang capture server| Low | `train_data_path` or `prompts_path` | | Offline | Not loaded by the trainer | High | `hidden_states_path` | Prepare raw datasets and offline features as described in [Data @@ -231,12 +253,13 @@ launching it. The unified runtime supports text training in these combinations: -| Strategy | SGLang server online | Local/dataflow offline | Disaggregated offline | +| Strategy | Online disaggregated | Offline colocated | Offline disaggregated | | --- | --- | --- | --- | | EAGLE3 | Yes, consumer DP | Yes, DP + USP | Yes, consumer DP | | DFlash | Yes, consumer DP | Yes, DP | Yes, consumer DP | | Domino | Yes, consumer DP | Yes, DP | Yes, consumer DP | | DSpark | Yes, consumer DP | Yes, DP | Yes, consumer DP | +| MTP | Yes, consumer DP | Yes, DP | Yes, consumer DP | | P-EAGLE | Yes, consumer DP, batch size 1 | No | No | Unsupported combinations fail explicitly during config validation or run @@ -249,24 +272,25 @@ assembly. In particular: - attention backends are strategy-specific: EAGLE3 accepts `sdpa`, `flex_attention`, `fa`, or offline `usp`; P-EAGLE requires `flex_attention`; DFlash, Domino, and DSpark accept `eager`, `sdpa`, or - `flex_attention`; + `flex_attention`; MTP accepts `eager` or `sdpa`; - P-EAGLE requires `training.batch_size=1` and reuses EAGLE3's server capture schema; -- offline feature training supports EAGLE3, DFlash, Domino, and DSpark; +- offline feature training supports EAGLE3, DFlash, Domino, DSpark, and MTP; - every online run is disaggregated and uses `model.target_backend=sglang`; finite runs may omit both step fields so the producer can publish the exact optimizer horizon derived from the prepared prompt plan; -- EAGLE3 local offline runs derive and cache a deterministic vocabulary mapping +- EAGLE3 offline colocated runs derive and cache a deterministic vocabulary mapping from the feature corpus when `model.vocab_mapping_path` is empty. EAGLE3 disaggregated runs require an explicit shared mapping so producer and consumer cannot derive different artifacts. -There is no fallback to a removed training script. - -Step limits are global optimizer updates. `training.max_steps` is a stop cap and, -when set without `training.total_steps`, the fallback optimizer/loss schedule -horizon. `training.total_steps` can describe a longer schedule, but does not by -itself stop an online stream. When a finite online run omits both, the producer +Step limits, LR/loss horizons, logging, saving, and Domino lambda decay are +expressed in completed optimizer updates. Fixed datasets are validated to +contain complete accumulation windows, and finite online plans do not train an +incomplete final quantum. `training.max_steps` is a stop cap and, when set +without `training.total_steps`, the fallback optimizer/loss schedule horizon. +`training.total_steps` can describe a longer schedule, but does not by itself +stop an online stream. When a finite online run omits both, the producer publishes the exact schedule horizon and the consumer trains to EOF. ## Parallel topologies @@ -389,7 +413,7 @@ export HCCL_CONNECT_TIMEOUT=7200 export HCCL_EXEC_TIMEOUT=7200 export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True -specforge train -c examples/configs/qwen3.5-4b-dflash-online-npu.yaml +specforge train -c examples/configs/online/disaggregated/external/qwen3.5-4b-dflash-online-npu.yaml ``` The unified launcher supplies rank, world-size, and rendezvous variables. The @@ -398,13 +422,15 @@ is active. ## Disaggregated roles -A single-node disaggregated config supervises producer and consumer with one -command. Split deployments use the same config with `--role producer` or -`--role consumer`; multi-node consumers add only `--node-rank` on each host. -The optional `examples/disagg/run_online.sh` and `run_offline.sh` files are thin -delegates, not topology wrappers. See the -[disaggregated training guide](disaggregated_training.md) for external -Mooncake/SGLang prerequisites, freshness rules, and both launch forms. +A single-node disaggregated config supervises the SpecForge producer and +consumer with one command. This does not make the run colocated: the roles and +their data-plane contract remain separate. For an external recipe, the command +also does not start Mooncake or SGLang. Split deployments use the same config +with `--role producer` or `--role consumer`; multi-node consumers add only +`--node-rank` on each host. The optional `examples/disagg/run_online.sh` and +`run_offline.sh` files are thin delegates, not topology wrappers. See the +[disaggregated training guide](disaggregated_training.md) for external-service +prerequisites, managed-local ownership, freshness rules, and both launch forms. ## Checkpoints and resume @@ -414,14 +440,14 @@ Mooncake/SGLang prerequisites, freshness rules, and both launch forms. even when `save_interval` is zero or the final step is not an interval boundary. The `-latest` symlink resolves to the newest complete checkpoint. -Local offline runs restore draft weights, optimizer/scheduler, epoch/step/data +Offline colocated runs restore draft weights, optimizer/scheduler, epoch/step/data position, and per-rank RNG. Offline disaggregated consumers have the same -checkpoint contract. For a local offline run, override +checkpoint contract. For an offline colocated run, override `training.resume_from`: ```bash specforge train \ - --config examples/configs/qwen3-8b-eagle3-offline.yaml \ + --config examples/configs/offline/colocated/qwen3-8b-eagle3-offline.yaml \ training.resume_from=./outputs/qwen3-8b-eagle3-offline/qwen3-8b-eagle3-offline-latest ``` @@ -484,6 +510,23 @@ specforge export --to hf \ Pass `--vocab-mapping /path/to/mapping.pt` when the checkpoint predates the mapping buffers or when you intentionally need to refresh them. +MTP is deployed by merging its trained native head back into the target model. +The merge command accepts the same runtime checkpoint shapes as the generic +exporter (`training_state.pt`, a step/latest directory, or the run output +directory): + +```bash +python scripts/merge_mtp_to_base.py \ + --base-model-path Qwen/Qwen3.5-4B \ + --mtp-checkpoint-path ./outputs/qwen3.5-4b-mtp/qwen3.5-4b-mtp-latest \ + --draft-config configs/qwen3.5-4b-mtp.json \ + --output-path ./exports/Qwen3.5-4B-MTP +``` + +An already-exported HF MTP draft directory can also be passed as +`--mtp-checkpoint-path`; in that case its own `config.json` is used and +`--draft-config` may be omitted. + ## Troubleshooting ### Late OOM or non-finite hidden states on online runs diff --git a/docs/benchmarks/domino-disaggregated-performance.md b/docs/benchmarks/domino-disaggregated-performance.md index 7fb6e111c..118b53dae 100644 --- a/docs/benchmarks/domino-disaggregated-performance.md +++ b/docs/benchmarks/domino-disaggregated-performance.md @@ -32,7 +32,7 @@ export MOONCAKE_PROTOCOL=rdma export MOONCAKE_RDMA_DEVICES=mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_6,mlx5_7 # Start the external capture server with --mem-fraction-static 0.5, then run: -specforge train -c examples/configs/qwen3-8b-domino-disaggregated.yaml \ +specforge train -c examples/configs/online/disaggregated/external/qwen3-8b-domino-disaggregated.yaml \ 'deployment.disaggregated.server_urls=["http://127.0.0.1:30000","http://127.0.0.1:30000","http://127.0.0.1:30000","http://127.0.0.1:30000","http://127.0.0.1:30000","http://127.0.0.1:30000","http://127.0.0.1:30000","http://127.0.0.1:30000"]' \ deployment.trainer.nproc_per_node=7 \ training.batch_size=2 \ @@ -52,7 +52,7 @@ seven trainer ranks on GPUs 1–7: ```bash specforge train -c \ - examples/configs/qwen3-8b-domino-1server-dp7-disaggregated.yaml + examples/configs/online/disaggregated/managed-local/qwen3-8b-domino-1server-dp7-disaggregated.yaml ``` The recipe records batch size 2, accumulation 8, and capture-server memory diff --git a/docs/benchmarks/eagle3-disaggregated-parity.md b/docs/benchmarks/eagle3-disaggregated-parity.md index a146a99f5..64a5f1048 100644 --- a/docs/benchmarks/eagle3-disaggregated-parity.md +++ b/docs/benchmarks/eagle3-disaggregated-parity.md @@ -29,18 +29,18 @@ measurement as serving-time accepted length. ## Unified entry for a new run The current equivalent recipe is -`examples/configs/qwen2.5-7b-eagle3-offline-disaggregated.yaml`. To place the +`examples/configs/offline/disaggregated/qwen2.5-7b-eagle3-offline-disaggregated.yaml`. To place the offline producer and consumer on different nodes, run one checked-in command through the cluster launcher: ```bash rcli exec --per-node \ - 'CONFIG=examples/configs/qwen2.5-7b-eagle3-offline-disaggregated.yaml bash examples/disagg/run_offline_2node.sh' + 'CONFIG=examples/configs/offline/disaggregated/qwen2.5-7b-eagle3-offline-disaggregated.yaml bash examples/disagg/run_offline_2node.sh' ``` Rank 0 invokes `specforge train --role producer`; rank 1 invokes `specforge train --role consumer`. Both nodes must resolve the config's `control_dir`, `store_root`, hidden-state input, and vocabulary mapping to the same data. Use a fresh attempt directory, then compare against the colocated -`examples/configs/qwen2.5-7b-eagle3-offline.yaml` recipe with the same inputs, +`examples/configs/offline/colocated/qwen2.5-7b-eagle3-offline.yaml` recipe with the same inputs, seed, and training overrides. diff --git a/docs/examples/llama3-eagle3-offline.md b/docs/examples/llama3-eagle3-offline.md index 69a52c2a2..da89b57ed 100644 --- a/docs/examples/llama3-eagle3-offline.md +++ b/docs/examples/llama3-eagle3-offline.md @@ -20,7 +20,7 @@ torchrun --standalone --nproc_per_node 8 \ scripts/prepare_hidden_states.py \ --strategy eagle3 \ --target-model-path meta-llama/Llama-3.1-8B-Instruct \ - --draft-model-config configs/llama3.1-8b-eagle3.json \ + --draft-model-config configs/llama3-8B-eagle3.json \ --data-path ./cache/dataset/sharegpt_train.jsonl \ --output-path ./cache/hidden_states/sharegpt_train_Llama-3.1-8B-Instruct \ --chat-template llama3 \ @@ -36,13 +36,13 @@ and `vocab_mapping/vocab_mapping.pt`, derived from the same processed corpus. ## 3. Use the checked-in run config The canonical recipe is -[`examples/configs/llama3.1-8b-eagle3-offline.yaml`](../../examples/configs/llama3.1-8b-eagle3-offline.yaml). +[`examples/configs/offline/colocated/llama3.1-8b-eagle3-offline.yaml`](../../examples/configs/offline/colocated/llama3.1-8b-eagle3-offline.yaml). It records the same target, feature directory, draft architecture, and trainer settings used by this walkthrough; edit the checked-in recipe or use dotted overrides instead of copying its YAML into another document. The recipe carries a conventional `model.vocab_mapping_path` for deployments -that prepare and share a mapping artifact. For a local offline run, leaving +that prepare and share a mapping artifact. For an offline colocated run, leaving that field empty makes SpecForge count effective tokens in the exact feature corpus, derive `t2d`/`d2t` deterministically, and cache the reusable mapping under `data.cache_dir/vocab_mapping`. Equal target and draft vocabularies need @@ -52,7 +52,7 @@ no mapping. ```bash specforge train \ - --config examples/configs/llama3.1-8b-eagle3-offline.yaml \ + --config examples/configs/offline/colocated/llama3.1-8b-eagle3-offline.yaml \ model.vocab_mapping_path=./cache/hidden_states/sharegpt_train_Llama-3.1-8B-Instruct/vocab_mapping/vocab_mapping.pt ``` diff --git a/docs/examples/llama3-eagle3-online.md b/docs/examples/llama3-eagle3-online.md index aef57b22c..8e0b586c9 100644 --- a/docs/examples/llama3-eagle3-online.md +++ b/docs/examples/llama3-eagle3-online.md @@ -15,7 +15,7 @@ python ./scripts/prepare_data.py --dataset sharegpt ## 2. Use the checked-in run config The canonical recipe is -[`examples/configs/llama3.1-8b-eagle3-online.yaml`](../../examples/configs/llama3.1-8b-eagle3-online.yaml). +[`examples/configs/online/disaggregated/external/llama3.1-8b-eagle3-online.yaml`](../../examples/configs/online/disaggregated/external/llama3.1-8b-eagle3-online.yaml). It already points at the ShareGPT output from step 1 and records the target, draft architecture, SGLang backend, optimizer settings, and output directory in the same typed contract used by every other training method. Edit that file or @@ -25,7 +25,7 @@ create a second copy of the recipe in documentation. ## 3. Train ```bash -specforge train --config examples/configs/llama3.1-8b-eagle3-online.yaml +specforge train --config examples/configs/online/disaggregated/external/llama3.1-8b-eagle3-online.yaml ``` The recipe points at an external SGLang capture server and starts the diff --git a/docs/get_started/about.md b/docs/get_started/about.md index 4702b2aac..76ad3af5d 100644 --- a/docs/get_started/about.md +++ b/docs/get_started/about.md @@ -5,10 +5,10 @@ Speculative decoding is an important and powerful technique for speeding up inference without losing performance. Industries have used it extensively in production to better serve their users with lower latency and higher throughput. We have seen some open-source projects for training speculative decoding models, but most of them are not well-maintained or not directly compatible with SGLang. We prepared this project because we wish that the open-source community can enjoy a speculative decoding framework that is - regularly maintained by the SGLang team: the code is runnable out-of-the-box -- directly compatible with SGLang: there is no additional efforts for porting to SGLang -- provide SGLang-server online training and local/disaggregated offline training - through one runtime, including consumer DP, offline USP, evaluation, - checkpoint selection, and CUDA/ROCm/Ascend portability +- directly compatible with SGLang: no additional porting effort is required +- able to provide online disaggregated training and both colocated and disaggregated + offline training through one runtime, including consumer DP, offline USP, + evaluation, checkpoint selection, and CUDA/ROCm/Ascend portability ## ✅ SGLang-ready diff --git a/docs/get_started/installation.md b/docs/get_started/installation.md index f920b1a00..daf8a18d9 100644 --- a/docs/get_started/installation.md +++ b/docs/get_started/installation.md @@ -12,11 +12,11 @@ git clone https://github.com/sgl-project/SpecForge.git cd SpecForge # create a new virtual environment -uv venv -p 3.11 +uv venv -p 3.11 --seed source .venv/bin/activate # install specforge -uv pip install -v . --prerelease=allow +uv pip install -e . ``` - **Install from PyPI** @@ -35,27 +35,31 @@ the same `specforge train` entry. ### AMD ROCm -For the pinned ROCm environment, install the checked-in requirements before the -package: +On ROCm, install SpecForge into an environment that already provides a ROCm +PyTorch and a ROCm SGLang (an official SGLang ROCm release container is the +recommended base), and install the package **without dependencies** so pip does +not pull CUDA wheels over the working ROCm stack: ```bash -python -m pip install -r requirements-rocm.txt -python -m pip install -e . +# Inside the ROCm SGLang container +git clone https://github.com/sgl-project/SpecForge.git /workspace/SpecForge +cd /workspace/SpecForge +python -m pip install -e . --no-deps ``` -The file pins a ROCm 7.2 PyTorch stack. Use a wheel index and driver combination -compatible with the host if your ROCm version differs. Online runs require a -ROCm-compatible SGLang capture service; offline feature consumers can start -without target inference. PyTorch exposes ROCm accelerators through its -`torch.cuda` API and uses NCCL for distributed runs. +For the complete container setup and an end-to-end walkthrough covering +installation, data preparation, offline colocated training, online +disaggregated training, and its single-supervisor and split external launch +forms on AMD Instinct GPUs, follow the +[AMD ROCm Tutorial](../basic_usage/AMD/amd_rocm.md). ### Ascend NPU Install the vendor-matched PyTorch and `torch_npu` packages first, then install SpecForge. The checked-in -[`qwen3.5-4b-dflash-online-npu.yaml`](../../examples/configs/qwen3.5-4b-dflash-online-npu.yaml) +[`qwen3.5-4b-dflash-online-npu.yaml`](../../examples/configs/online/disaggregated/external/qwen3.5-4b-dflash-online-npu.yaml) and -[`qwen3.5-4b-domino-online-npu.yaml`](../../examples/configs/qwen3.5-4b-domino-online-npu.yaml) +[`qwen3.5-4b-domino-online-npu.yaml`](../../examples/configs/online/disaggregated/external/qwen3.5-4b-domino-online-npu.yaml) recipes use external SGLang server capture with SDPA consumers. Install a compatible SGLang/Mooncake service first. The unified launcher detects the NPU device, self-launches the process count recorded in YAML, and selects HCCL; see diff --git a/docs/index.rst b/docs/index.rst index 189e78361..c03c6d493 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -26,6 +26,8 @@ SpecForge is an ecosystem project developed by the SGLang team. It is a framewor basic_usage/data_preparation.md basic_usage/training.md basic_usage/disaggregated_training.md + basic_usage/AMD/amd_rocm.md + basic_usage/Ascend/ascend_npu.md .. toctree:: :maxdepth: 1 @@ -54,3 +56,9 @@ SpecForge is an ecosystem project developed by the SGLang team. It is a framewor examples/llama3-eagle3-online.md examples/llama3-eagle3-offline.md + +.. toctree:: + :maxdepth: 1 + :caption: Recipes + + recipes/kimi-k3-dspark-disaggregated.md diff --git a/docs/recipes/kimi-k3-dspark-disaggregated.md b/docs/recipes/kimi-k3-dspark-disaggregated.md new file mode 100644 index 000000000..2a051aada --- /dev/null +++ b/docs/recipes/kimi-k3-dspark-disaggregated.md @@ -0,0 +1,151 @@ +# Kimi K3 DSpark disaggregated reproduction + +This recipe migrates the prior four-node colocated Kimi K3 continual run +to one TP8 capture node and one four-rank trainer node. It preserves the draft +architecture, weights-only warm start, regenerated agentic prompt order, +effective global batch, constant learning rate, and DSpark loss weights. + +## Required source revisions + +- SpecForge with configurable LR scheduling, an independent online prompt + seed, and dedicated DSpark capture support. +- Kimi K3 SGLang revision `ee560a2b2df5dafe18fd835d2e546eff019ca5ba`, + the recommended deployment revision and reproducibility baseline validated + by this recipe. +- The K3 SGLang tree patched with: + + ```bash + scripts/apply_sglang_spec_capture_patch.sh --target kimi-k3-ee560a2 + ``` + +The patch is generated against `ee560a2`. `f8493a4` and `9acd9cb` are +compatibility-validation targets only; their accepted `--target` names remain +aliases for existing automation. The patch's historical directory name is also +retained. It makes `--spec-capture-method dspark` call K3's +`set_dspark_layers_to_capture` hook. The generic DFlash capture method is not +equivalent for K3. The same versioned patch carries the three required 64K +correctness guards: 64-bit Triton token offsets, scale-stable residual scoring, +and the Marlin grid.y fallback above 65,535 tokens. + +## Artifacts + +The checked-in recipe uses paths already provisioned on the deployment. Other +deployments should override them without editing the recipe: + +- target revision `cdd2e49a2c1cf8d4713b513955e415ed75405a72`; +- the weights-only `epoch_0_step_0` draft checkpoint; +- the 462-row regenerated dataset whose SHA-256 is + `6d50e6bb9ee59095eed91bfba035081efef9fea43bece9ad5dd01c6648a8ef24`. + +Do not put Hugging Face or W&B credentials in YAML. Supply `HF_TOKEN` and +`WANDB_API_KEY` through protected node-local files or the process environment. + +## Capture node + +Start Mooncake with at least a 1 TiB global segment, then start the patched K3 +server. Replace `CAPTURE_IP` with the routable address used by both nodes. + +```bash +export MOONCAKE_LOCAL_HOSTNAME="$CAPTURE_IP" +export MC_TCP_BIND_ADDRESS="$CAPTURE_IP" +export MC_TRANSFER_TIMEOUT=300 +export MOONCAKE_GLOBAL_SEGMENT_SIZE=1099511627776 +export MOONCAKE_LOCAL_BUFFER_SIZE=1073741824 +mooncake_master \ + --enable_http_metadata_server=true \ + --http_metadata_server_host=0.0.0.0 \ + --rpc_port=35551 \ + --http_metadata_server_port=35880 \ + --metrics_port=35903 \ + --default_kv_lease_ttl=5m +``` + +In a second process: + +```bash +export MOONCAKE_MASTER_SERVER_ADDR="$CAPTURE_IP:35551" +export MOONCAKE_METADATA_SERVER="http://$CAPTURE_IP:35880/metadata" +export MOONCAKE_LOCAL_HOSTNAME="$CAPTURE_IP" +export MC_TCP_BIND_ADDRESS="$CAPTURE_IP" +export MC_TRANSFER_TIMEOUT=300 +export MOONCAKE_PROTOCOL=tcp +export MOONCAKE_GLOBAL_SEGMENT_SIZE=1099511627776 +export MOONCAKE_LOCAL_BUFFER_SIZE=1073741824 +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m sglang.launch_server \ + --host 0.0.0.0 \ + --port 30000 \ + --model-path /workspace/models/Kimi-K3-cdd2e49a \ + --trust-remote-code \ + --skip-tokenizer-init \ + --tp-size 8 \ + --mem-fraction-static 0.76 \ + --context-length 66048 \ + --max-running-requests 1 \ + --max-total-tokens 66048 \ + --attention-backend trtllm_mla \ + --moe-runner-backend marlin \ + --mamba-radix-cache-strategy extra_buffer \ + --max-mamba-cache-size 5 \ + --disable-cuda-graph \ + --chunked-prefill-size -1 \ + --enable-spec-capture \ + --spec-capture-method dspark \ + --spec-capture-aux-layer-ids 7 23 51 67 83 +``` + +## Trainer node + +Resolve `capture-node` through DNS or override the three endpoint fields with +the capture node's IP. The default CLI role starts one CPU producer and a +four-rank FSDP consumer on the same trainer host. + +```bash +export MOONCAKE_LOCAL_HOSTNAME="$TRAINER_IP" +export MC_TCP_BIND_ADDRESS="$TRAINER_IP" +export MC_TRANSFER_TIMEOUT=300 +export WANDB_API_KEY="$(< /protected/path/wandb-api-key)" +export WANDB_ENTITY=your-entity +unset RANK LOCAL_RANK WORLD_SIZE MASTER_ADDR MASTER_PORT NODE_RANK +CUDA_VISIBLE_DEVICES=0,1,2,3 specforge train \ + -c examples/configs/online/disaggregated/external/kimi-k3-dspark-disaggregated.yaml \ + --role both \ + "deployment.disaggregated.server_urls=[\"http://$CAPTURE_IP:30000\"]" \ + "deployment.disaggregated.mooncake_metadata_server=http://$CAPTURE_IP:35880/metadata" \ + "deployment.disaggregated.mooncake_master_server_addr=$CAPTURE_IP:35551" +``` + +`MC_TCP_BIND_ADDRESS` is required on multi-interface or containerized hosts. +Without it Mooncake may publish a Docker bridge address even when +`MOONCAKE_LOCAL_HOSTNAME` names the routable inter-node address, causing remote +`get_into` operations to fail. The five-minute master lease and matching +transfer timeout cover a 5.25 GiB 64K feature object over a shared TCP link; +Mooncake's short default lease can expire while that object is still in flight. +The rendezvous variables are cleared because cluster base images sometimes +inject a partial multi-node environment; SpecForge intentionally rejects that +instead of guessing which world the four trainer ranks should join. + +For a one-update smoke run, additionally override the pre-tokenized four-row +fixture and shrink the optimizer quantum: + +```bash +specforge train \ + -c examples/configs/online/disaggregated/external/kimi-k3-dspark-disaggregated.yaml \ + --role both \ + data.train_data_path= \ + data.prompts_path=/workspace/k3_dspark/data/longest-smoke-pretokenized-4rows-65536.jsonl \ + training.num_epochs=1 \ + training.max_steps=1 \ + training.accumulation_steps=1 \ + tracking.report_to=none \ + runtime.in_flight_high_watermark=4 \ + runtime.in_flight_low_watermark=4 +``` + +The smoke low watermark must stay at least as large as the four-rank global +optimizer-step quantum. A lower value is rejected before capture starts so a +producer cannot pause while consumers are waiting for an incomplete step. + +Validate in order: config plan, patch dry-run, server health, one captured +sample's tensor shapes/dtypes, one finite optimizer update and checkpoint, then +the full run. A smoke pass does not establish numerical parity; compare the +full run's loss/accuracy/tau trajectory and final checkpoint hashes separately. diff --git a/examples/README.md b/examples/README.md index bf614ee76..a02026603 100644 --- a/examples/README.md +++ b/examples/README.md @@ -5,49 +5,37 @@ its model and data paths, and launch it directly; multi-process topology is already recorded in the YAML: ```bash -specforge train --config examples/configs/qwen3-8b-eagle3-disaggregated.yaml +specforge train --config examples/configs/online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml ``` -The representative configs are below. The complete recipe catalog, including -NPU, offline, and managed/external-service variants, is in -[`examples/configs/README.md`](./configs/README.md). +The directory path identifies feature mode, topology, and—in the online +disaggregated case—service ownership. Start with one of these categories, then +choose the model and strategy inside it: -| Config | Mode | Strategy | +| Category | Use it when | Representative config | | --- | --- | --- | -| `examples/configs/qwen3-8b-eagle3-disaggregated.yaml` | Disaggregated SGLang server capture | EAGLE3 | -| `examples/configs/qwen3-8b-eagle3-offline.yaml` | Precomputed features | EAGLE3 | -| `examples/configs/qwen3-8b-eagle3-offline-disaggregated.yaml` | Disaggregated precomputed features | EAGLE3 | -| `examples/configs/qwen2.5-7b-eagle3-offline-disaggregated.yaml` | Disaggregated precomputed features | EAGLE3 | -| `examples/configs/qwen3-30b-a3b-eagle3.1-online.yaml` | Disaggregated SGLang server capture with normalized EAGLE3 inputs | EAGLE3.1 | -| `examples/configs/qwen3-8b-dflash-online.yaml` | Disaggregated SGLang server capture | DFlash | -| `examples/configs/qwen3-8b-dflash-offline.yaml` | Precomputed features | DFlash | -| `examples/configs/qwen3-8b-dpace-online.yaml` | Online D-PACE objective | DFlash | -| `examples/configs/qwen3-8b-dflash-disaggregated.yaml` | Disaggregated server capture | DFlash | -| `examples/configs/qwen3-8b-dflash-1server-dp7-disaggregated.yaml` | Managed local one capture server + DP7 | DFlash | -| `examples/configs/qwen3-8b-domino-online.yaml` | Disaggregated SGLang server capture | Domino | -| `examples/configs/qwen3-8b-domino-offline.yaml` | Precomputed features | Domino | -| `examples/configs/qwen3-8b-domino-disaggregated.yaml` | Disaggregated server capture | Domino | -| `examples/configs/qwen3-8b-domino-1server-dp7-disaggregated.yaml` | Managed local one capture server + DP7 | Domino | -| `examples/configs/qwen3-8b-domino-multiserver-disaggregated.yaml` | Managed local Mooncake + two capture servers | Domino | -| `examples/configs/qwen3-8b-peagle-disaggregated.yaml` | Disaggregated SGLang server capture | P-EAGLE | -| `examples/configs/qwen3-4b-dspark-disaggregated.yaml` | Disaggregated server capture | DSpark | -| `examples/configs/qwen3-4b-dspark-offline.yaml` | Precomputed features | DSpark | -| `examples/configs/qwen3.6-27b-dflash-multiserver-disaggregated.yaml` | Managed local Mooncake + two capture servers | DFlash | -| `examples/configs/qwen3.6-27b-dflash-1server-dp2-disaggregated.yaml` | Managed local one capture server + DP2 | DFlash | -| `examples/configs/qwen3.5-4b-dflash-online-npu.yaml` | Disaggregated NPU SGLang capture | DFlash | -| `examples/configs/qwen3.5-4b-domino-online-npu.yaml` | Disaggregated NPU SGLang capture | Domino | +| Offline colocated | Feature checkpoints already exist and the trainer can read them directly | [`offline/colocated/qwen3-8b-eagle3-offline.yaml`](./configs/offline/colocated/qwen3-8b-eagle3-offline.yaml) | +| Offline disaggregated | A producer must ingest existing features for a separate trainer pool | [`offline/disaggregated/qwen3-8b-eagle3-offline-disaggregated.yaml`](./configs/offline/disaggregated/qwen3-8b-eagle3-offline-disaggregated.yaml) | +| Online disaggregated, external | Mooncake and patched SGLang are started by the user or scheduler | [`online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml`](./configs/online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml) | +| Online disaggregated, managed-local | One command should own a single-node Mooncake/SGLang/trainer stack | [`online/disaggregated/managed-local/qwen3-8b-dflash-1server-dp7-disaggregated.yaml`](./configs/online/disaggregated/managed-local/qwen3-8b-dflash-1server-dp7-disaggregated.yaml) | + +The complete model, strategy, CUDA/ROCm/Ascend, and resource-layout catalog is +documented in [`examples/configs/README.md`](./configs/README.md). The +`online/colocated` directory is reserved for a topology that the current +runtime does not implement. Online configs point `data.train_data_path` at raw conversation data. Offline configs expect strategy-specific feature checkpoints in `data.hidden_states_path`. -Local offline EAGLE3 derives and caches its vocabulary map when no path is set; +Offline colocated EAGLE3 derives and caches its vocabulary map when no path is set; disaggregated EAGLE3 requires one explicit shared `model.vocab_mapping_path`. -Online training always uses an external or managed SGLang capture server and -the disaggregated producer/consumer data plane. Colocated online target loading -and the HF/custom online backends are intentionally unsupported. Online capture -is text-only: VLM training, including Qwen2.5-VL, is not supported. Online -evaluation is also not supported. +Online training always uses the disaggregated producer/consumer data plane and +an external or managed-local SGLang capture server. Here `external` means that +SpecForge does not own the service lifecycle; the server may still be on the +same host. Colocated online target loading and the HF/custom online backends are +intentionally unsupported. Online capture is text-only: VLM training, +including Qwen2.5-VL, is not supported. Online evaluation is also unsupported. The same CLI owns offline DP, EAGLE3 offline USP, and managed capture-server topology. Trainer `tp_size` remains 1; target TP belongs to SGLang capture @@ -68,7 +56,7 @@ projection for offline text EAGLE3, and W&B, TensorBoard, SwanLab, or MLflow tracking. See the [training guide](../docs/basic_usage/training.md) for the full capability matrix, ROCm installation, and Ascend NPU/HCCL launch example. -Local offline and disaggregated offline resume are supported. +Offline colocated and offline disaggregated resume are supported. Disaggregated online resume is consumer-only and requires the retained SQLite ledger, channel/inboxes, Mooncake data, and an exactly matching checkpoint; the producer is never resumed. diff --git a/examples/configs/README.md b/examples/configs/README.md index eca197da1..08969e0a3 100644 --- a/examples/configs/README.md +++ b/examples/configs/README.md @@ -1,10 +1,47 @@ # Unified training recipe catalog Every draft model JSON under `configs/` has at least one typed YAML recipe in -this directory. Run any recipe through the one public training entry: +this catalog. Choose a recipe by answering three separate questions: + +1. **When are target features produced?** `online` captures them during the + run; `offline` reads feature files prepared earlier. +2. **How are feature supply and training deployed?** `colocated` maps to + `deployment.mode: local_colocated` and lets the trainer read offline files + directly; `disaggregated` uses separate producer and consumer roles. +3. **Who owns the online infrastructure?** `external` means the user or + scheduler owns Mooncake and SGLang; `managed-local` means SpecForge starts + and stops those services on the local host. + +Those decisions map directly to the directory tree: + +```text +examples/configs/ +├── offline/ +│ ├── colocated/ +│ └── disaggregated/ +└── online/ + ├── colocated/ + └── disaggregated/ + ├── external/ + └── managed-local/ +``` + +| Directory | Feature source | SpecForge roles | Infrastructure ownership | +| --- | --- | --- | --- | +| `offline/colocated` | Precomputed `.ckpt` files | Trainer only | Not applicable | +| `offline/disaggregated` | Precomputed files ingested into `shared_dir` or Mooncake | Producer + consumer | Storage is configured separately; there is no `external`/`managed-local` subdivision | +| `online/colocated` | — | — | Reserved; the current runtime does not support online colocated training | +| `online/disaggregated/external` | Live SGLang capture | Producer + consumer | User or scheduler starts Mooncake and SGLang | +| `online/disaggregated/managed-local` | Live SGLang capture | Producer + consumer | SpecForge starts local Mooncake and SGLang from `managed_local` | + +The directory is the source of truth for mode, topology, and service ownership. +Some filenames retain historical `-online`, `-offline`, or `-disaggregated` +labels so existing recipe identities and run names remain recognizable. + +Run any recipe through the one public training entry: ```bash -specforge train --config examples/configs/qwen3-8b-eagle3-disaggregated.yaml +specforge train --config examples/configs/online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml ``` `model.draft_model_config` may name a local JSON file, a local model directory, @@ -18,13 +55,13 @@ Every recipe records its audited process count under `deployment.trainer`. Multi-process configs self-launch through torch distributed: ```bash -specforge train -c examples/configs/qwen3-30b-a3b-eagle3-online.yaml +specforge train -c examples/configs/online/disaggregated/external/qwen3-30b-a3b-eagle3-online.yaml ``` The Qwen3-30B-A3B EAGLE3.1 variant uses the same unified entry point: ```bash -specforge train -c examples/configs/qwen3-30b-a3b-eagle3.1-online.yaml +specforge train -c examples/configs/online/disaggregated/external/qwen3-30b-a3b-eagle3.1-online.yaml ``` Its draft config enables per-layer RMS normalization before the three captured @@ -32,26 +69,23 @@ target hidden states are concatenated and projected. It remains registered as the `eagle3` strategy; EAGLE3.1 is a draft-model configuration variant, not a second runtime or launch path. -The filename is the index: `*-online.yaml` performs SGLang server capture while -training, `*-offline.yaml` consumes precomputed features, and -`*-disaggregated.yaml` highlights a producer/consumer topology. Every online -recipe is disaggregated even when its historical filename only says `online`. Multimodal (image+text) training is supported for `training.strategy=dflash` via `model.input_modality: multimodal` (see -`qwen3.5-4b-vl-dflash-disaggregated.yaml`); the remaining catalog is -text-only. - -The `qwen3-8b-dflash-1server-dp7-disaggregated.yaml`, -`qwen3-8b-domino-1server-dp7-disaggregated.yaml`, -`qwen3-8b-domino-multiserver-disaggregated.yaml`, -`qwen3.6-27b-dflash-1server-dp2-disaggregated.yaml`, and -`qwen3.6-27b-dflash-multiserver-disaggregated.yaml`, and -`qwen3.6-27b-dspark-disaggregated.yaml` recipes are opt-in local +`online/disaggregated/external/qwen3.5-4b-vl-dflash-disaggregated.yaml`); the +remaining catalog is text-only. + +Recipes under `online/disaggregated/managed-local` are opt-in, single-node full-stack examples. Their typed `managed_local` blocks own Mooncake, one or -two patched SGLang capture servers, and the trainer GPU allocation; the same -`specforge train -c ...` command starts and cleans up each complete stack. -Disaggregated recipes without `managed_local` keep Mooncake and SGLang external -for scheduler- or service-managed deployments. +more patched SGLang capture servers, and the trainer GPU allocation; the same +`specforge train -c ...` command starts and cleans up the complete stack. +Recipes under `online/disaggregated/external` omit `managed_local`, so Mooncake +and SGLang remain owned by the user, scheduler, or service platform. + +The `kimi-k3-dspark-disaggregated.yaml` recipe is the external +two-node migration of the 64K Kimi K3 continual run. Its dedicated +[runbook](../../docs/recipes/kimi-k3-dspark-disaggregated.md) pins the K3 +SGLang revision and patch target, preserves the old effective global batch and +prompt order, and documents the TP8 capture plus four-rank trainer topology. Before running a recipe, update model/data paths and create any referenced offline feature or vocabulary-mapping artifacts. Managed-local recipes @@ -70,7 +104,8 @@ errors; YAML files and dotted CLI overrides go through the same validation. New checked-in recipes should explicitly set `training.strategy`, `deployment.mode`, `deployment.trainer.nnodes`, `deployment.trainer.nproc_per_node`, `run_id`, and `output_dir`, even when the -schema has the same default. A minimal server-only online recipe looks like: +schema has the same default. A minimal online recipe with externally managed +services looks like: ```yaml model: @@ -115,16 +150,18 @@ assume the command runs from the repository root. | Workflow | Canonical starting point | | --- | --- | -| EAGLE3 colocated offline | `qwen3-8b-eagle3-offline.yaml` | -| DFlash colocated offline | `qwen3-8b-dflash-offline.yaml` | -| Domino colocated offline | `qwen3-8b-domino-offline.yaml` | -| DSpark colocated offline | `qwen3-4b-dspark-offline.yaml` | -| External-service online | `qwen3-8b-eagle3-disaggregated.yaml` | -| Managed-local disaggregated online | `qwen3-8b-domino-multiserver-disaggregated.yaml` | -| Disaggregated offline | `qwen3-8b-eagle3-offline-disaggregated.yaml` | - -The online/offline mode is derived from the selected `data` source, not from -the filename. The filename is a discoverability convention. +| EAGLE3 offline, colocated | [`offline/colocated/qwen3-8b-eagle3-offline.yaml`](offline/colocated/qwen3-8b-eagle3-offline.yaml) | +| DFlash offline, colocated | [`offline/colocated/qwen3-8b-dflash-offline.yaml`](offline/colocated/qwen3-8b-dflash-offline.yaml) | +| Domino offline, colocated | [`offline/colocated/qwen3-8b-domino-offline.yaml`](offline/colocated/qwen3-8b-domino-offline.yaml) | +| DSpark offline, colocated | [`offline/colocated/qwen3-4b-dspark-offline.yaml`](offline/colocated/qwen3-4b-dspark-offline.yaml) | +| EAGLE3 offline, disaggregated | [`offline/disaggregated/qwen3-8b-eagle3-offline-disaggregated.yaml`](offline/disaggregated/qwen3-8b-eagle3-offline-disaggregated.yaml) | +| EAGLE3 online, external services | [`online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml`](online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml) | +| DFlash online, managed-local stack | [`online/disaggregated/managed-local/qwen3-8b-dflash-1server-dp7-disaggregated.yaml`](online/disaggregated/managed-local/qwen3-8b-dflash-1server-dp7-disaggregated.yaml) | +| Domino online, managed-local stack | [`online/disaggregated/managed-local/qwen3-8b-domino-multiserver-disaggregated.yaml`](online/disaggregated/managed-local/qwen3-8b-domino-multiserver-disaggregated.yaml) | + +The runtime derives online/offline mode from the selected `data` source and +reads topology from `deployment.mode`; it does not parse the filename or its +historical suffixes. ### Top-level fields @@ -148,7 +185,7 @@ should make their training strategy and topology explicit. | `model.draft_block_size` | `null` | Positive DFlash block-size override; generated DFlash configs default to 16. | | `model.target_backend` | `sglang` | `sglang` is the only accepted value; retired `hf`/`custom` names fail at config load. Offline feature consumers do not instantiate a target inference backend. | | `model.input_modality` | `text` | The provider modality. Built-ins support `text`; DFlash additionally supports `multimodal` (image+text server capture). Other identifiers are rejected at application resolution. | -| `model.shard_target_output` | `false` | Retained for config migration; leave it false on the server-only online path. | +| `model.shard_target_output` | `false` | Retained for config migration; leave it false on the online disaggregated path. | | `model.trust_remote_code` | `false` | Enable only for model repositories that require custom loading code. | | `model.use_liger_kernel` | `false` | Enable Liger Qwen3 RMSNorm/SwiGLU kernels for DFlash training. Requires the `specforge[liger]` extra. | | `model.embedding_key` | `model.embed_tokens.weight` | Target checkpoint key copied into or used by the draft embedding. | @@ -163,6 +200,7 @@ should make their training strategy and topology explicit. | `model.sglang_attention_backend` | `flashinfer` | SGLang attention implementation for an in-process or managed capture server. | | `model.sglang_mm_attention_backend` | `null` | Vision-encoder attention backend for capture servers. On Ascend NPU with a non-text `model.input_modality` it defaults to `ascend_attn` (fused); `sdpa` materializes N² vision scores and can OOM on large images. | | `model.sglang_mem_fraction_static` | `0.4` | SGLang static-memory fraction in `(0, 1]`; inherited by managed capture servers unless they override it. | +| `model.sglang_disable_radix_cache` | `true` | Preserve the historical managed-capture behavior. Set `false` for hybrid targets such as Inkling that require the radix tree. Unique per-attempt cache namespaces still force complete capture prefills. | | `model.sglang_context_length` | `null` | Positive explicit context limit. Managed capture requires at least `data.max_length + 7`; omitting it derives that value. | | `model.sglang_enable_nccl_nvls` | `false` | Pass the matching SGLang NCCL NVLS optimization flag. | | `model.sglang_enable_symm_mem` | `false` | Pass the matching SGLang symmetric-memory flag. | @@ -222,6 +260,7 @@ Common fields: | `training.accumulation_steps` | `1` | Positive microbatches per optimizer update. | | `training.fsdp_sharding` | `SHARD_GRAD_OP` | Trainer FSDP mode: `SHARD_GRAD_OP`, `FULL_SHARD`, or `NO_SHARD`. | | `training.learning_rate` | `1e-4` | Positive peak learning rate. | +| `training.lr_scheduler` | `cosine` | Learning-rate schedule after warmup: `cosine` or `constant`. | | `training.warmup_ratio` | `0.015` | Fraction in `[0, 1]` used for scheduler warmup. | | `training.max_grad_norm` | `0.5` | Positive gradient-clipping norm. | | `training.optimizer_cpu_offload` | `false` | Keep the optimizer's FP32 master parameters and Adam state on CPU. | @@ -237,8 +276,10 @@ Common fields: | `training.resume_from` | `null` | Full-run checkpoint/run root: draft, optimizer/scheduler, counters, data position, and RNG. Mutually exclusive with `model.draft_checkpoint_path`. | | `training.compact_teacher` | `false` | Exact lower-peak-memory teacher projection for offline text EAGLE3. | | `training.compact_teacher_chunk_size` | `null` | Positive vocabulary chunk size; requires `compact_teacher: true`. | -| `training.role` | `all` | Use `all` for local offline training; disaggregated entrypoints select `auto`, `producer`, or `consumer`. | +| `training.trim_loss_positions` | `false` | EAGLE3 only. Compute the teacher target_p, draft logits, and loss only at supervised positions (batch size 1, plain KL loss); mathematically equivalent to the full-length path. | +| `training.role` | `all` | Use `all` for offline colocated training; disaggregated entrypoints select `auto`, `producer`, or `consumer`. | | `training.seed` | `42` | Run and per-rank RNG seed. | +| `training.prompt_seed` | `null` | Optional online prompt-shuffle seed. `null` preserves the historical behavior of using `training.seed`. | Strategy-specific fields should be written only when tuning that objective: @@ -274,9 +315,10 @@ For `deployment.mode: disaggregated`, also write: | Field | Default | What to write | | --- | --- | --- | -| `deployment.disaggregated.control_dir` | required | Fresh attempt-scoped shared directory for refs/manifest and lifecycle markers. | +| `deployment.disaggregated.control_dir` | required | Fresh attempt-scoped directory for refs/manifest and lifecycle markers. Shared by default; with `inbox_server_url`, only producer and consumer rank 0 must share it. | | `deployment.disaggregated.backend` | required | `mooncake` or `shared_dir`. Online disaggregated runs require Mooncake. | | `deployment.disaggregated.consumer_state_dir` | `null` | Node-local rank-0 SQLite/WAL root. Required for multi-node online consumers; their rank inboxes remain under shared `control_dir`. | +| `deployment.disaggregated.inbox_server_url` | `null` | Optional private `http://host:port` rank-0 relay for tensor-free inbox refs when remote trainer ranks cannot share `control_dir`. Online multi-node only; no credentials, path, query, TLS, or built-in authentication. | | `deployment.disaggregated.store_root` | `null` | Shared feature directory; required when `backend: shared_dir`. | | `deployment.disaggregated.store_id` | `null` | Feature-store namespace; defaults to `run_id`. | | `deployment.disaggregated.server_urls` | `[]` | External patched SGLang capture endpoints. One rollout worker is created per entry. Do not set with `managed_local`. | @@ -345,6 +387,7 @@ Managed-local fields: | `deployment.disaggregated.managed_local.mooncake.global_segment_size_bytes` | `34359738368` | Owned global segment size. | | `deployment.disaggregated.managed_local.mooncake.local_buffer_size_bytes` | `1073741824` | Owned local client buffer. | | `deployment.disaggregated.managed_local.mooncake.startup_timeout_s` | `60` | Positive Mooncake readiness timeout. | +| `deployment.disaggregated.managed_local.mooncake.default_kv_lease_ttl_ms` | `500` | Master key-lease TTL (ms) forwarded to `mooncake_master --default_kv_lease_ttl`. Kept below the consumer's teardown drain window so managed_local shuts down cleanly; set `null` to inherit Mooncake's stock default. | | `deployment.disaggregated.managed_local.capture_servers[].port` | required | Unique capture HTTP port. | | `deployment.disaggregated.managed_local.capture_servers[].cuda_visible_devices` | required | Device tokens for this server. Their count must equal its `tp_size`. | | `deployment.disaggregated.managed_local.capture_servers[].tp_size` | `1` | Target-model tensor parallelism for this server. | @@ -418,24 +461,25 @@ unless tuning throughput or memory pressure. requires batch size 1. - `training.compact_teacher` is offline text EAGLE3 only. - Multimodal (image+text) training requires `training.strategy=dflash` with - `model.input_modality: multimodal`, a VLM draft config - (`configs/*-dflash-vlm-*.json`), and the v0.5.14 spec-capture patch with the - `position_ids` artifact. Vendor modalities such as `qwen2_5_vl` remain + `model.input_modality: multimodal`, a plain-rope DFlash draft config (e.g. + `configs/qwen3.5-4b-dflash.json`; the draft needs no mRoPE - visual + information arrives through the captured target hidden states), and the + v0.5.14 spec-capture patch. Vendor modalities such as `qwen2_5_vl` remain unsupported. - Online evaluation is not supported. Offline `data.eval_hidden_states_path` and `training.eval_interval` must be configured together. -Validate the complete schema and inspect the resolved processes without -starting a run: +After copying a checked-in recipe to `./my-run.yaml` and editing it, validate +the complete schema and inspect the resolved processes without starting a run: ```bash -specforge train -c examples/configs/my-run.yaml --plan +specforge train -c ./my-run.yaml --plan ``` Use validated dotted overrides for temporary changes: ```bash -specforge train -c examples/configs/my-run.yaml \ +specforge train -c ./my-run.yaml \ training.learning_rate=5e-5 \ 'deployment.disaggregated.server_urls=["http://capture-0:30000"]' ``` @@ -446,7 +490,7 @@ For deeper lifecycle and recovery semantics, see the ## Capability matrix -| Strategy | SGLang server online | Local offline | Disaggregated offline | +| Strategy | Online disaggregated | Offline colocated | Offline disaggregated | | --- | --- | --- | --- | | EAGLE3 | consumer DP | DP + USP | consumer DP | | DFlash | consumer DP | DP | consumer DP | @@ -476,14 +520,14 @@ export HCCL_CONNECT_TIMEOUT=7200 export HCCL_EXEC_TIMEOUT=7200 export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True -specforge train -c examples/configs/qwen3.5-4b-dflash-online-npu.yaml +specforge train -c examples/configs/online/disaggregated/external/qwen3.5-4b-dflash-online-npu.yaml ``` The unified launcher provides rank/world/rendezvous variables and the runtime selects HCCL when `torch_npu` is active. For AMD GPUs, install `requirements-rocm.txt`; HF + SDPA is the portable ROCm starting point. -Local offline and disaggregated offline resume are supported. +Offline colocated and offline disaggregated resume are supported. Disaggregated online recovery resumes only the consumer against retained control/data-plane state; capture producers always start a fresh attempt. diff --git a/examples/configs/deepseek-v3-671b-eagle3-offline.yaml b/examples/configs/offline/colocated/deepseek-v3-671b-eagle3-offline.yaml similarity index 100% rename from examples/configs/deepseek-v3-671b-eagle3-offline.yaml rename to examples/configs/offline/colocated/deepseek-v3-671b-eagle3-offline.yaml diff --git a/examples/configs/ling-flash-2.0-eagle3-offline.yaml b/examples/configs/offline/colocated/ling-flash-2.0-eagle3-offline.yaml similarity index 100% rename from examples/configs/ling-flash-2.0-eagle3-offline.yaml rename to examples/configs/offline/colocated/ling-flash-2.0-eagle3-offline.yaml diff --git a/examples/configs/llama3.1-8b-eagle3-offline.yaml b/examples/configs/offline/colocated/llama3.1-8b-eagle3-offline.yaml similarity index 100% rename from examples/configs/llama3.1-8b-eagle3-offline.yaml rename to examples/configs/offline/colocated/llama3.1-8b-eagle3-offline.yaml diff --git a/examples/configs/qwen2.5-7b-eagle3-offline.yaml b/examples/configs/offline/colocated/qwen2.5-7b-eagle3-offline.yaml similarity index 100% rename from examples/configs/qwen2.5-7b-eagle3-offline.yaml rename to examples/configs/offline/colocated/qwen2.5-7b-eagle3-offline.yaml diff --git a/examples/configs/qwen3-4b-dspark-offline.yaml b/examples/configs/offline/colocated/qwen3-4b-dspark-offline.yaml similarity index 100% rename from examples/configs/qwen3-4b-dspark-offline.yaml rename to examples/configs/offline/colocated/qwen3-4b-dspark-offline.yaml diff --git a/examples/configs/qwen3-8b-dflash-offline.yaml b/examples/configs/offline/colocated/qwen3-8b-dflash-offline.yaml similarity index 100% rename from examples/configs/qwen3-8b-dflash-offline.yaml rename to examples/configs/offline/colocated/qwen3-8b-dflash-offline.yaml diff --git a/examples/configs/qwen3-8b-domino-offline.yaml b/examples/configs/offline/colocated/qwen3-8b-domino-offline.yaml similarity index 100% rename from examples/configs/qwen3-8b-domino-offline.yaml rename to examples/configs/offline/colocated/qwen3-8b-domino-offline.yaml diff --git a/examples/configs/qwen3-8b-dspark-offline.yaml b/examples/configs/offline/colocated/qwen3-8b-dspark-offline.yaml similarity index 100% rename from examples/configs/qwen3-8b-dspark-offline.yaml rename to examples/configs/offline/colocated/qwen3-8b-dspark-offline.yaml diff --git a/examples/configs/qwen3-8b-eagle3-offline.yaml b/examples/configs/offline/colocated/qwen3-8b-eagle3-offline.yaml similarity index 100% rename from examples/configs/qwen3-8b-eagle3-offline.yaml rename to examples/configs/offline/colocated/qwen3-8b-eagle3-offline.yaml diff --git a/examples/configs/qwen3-coder-480b-a35b-eagle3-offline.yaml b/examples/configs/offline/colocated/qwen3-coder-480b-a35b-eagle3-offline.yaml similarity index 100% rename from examples/configs/qwen3-coder-480b-a35b-eagle3-offline.yaml rename to examples/configs/offline/colocated/qwen3-coder-480b-a35b-eagle3-offline.yaml diff --git a/examples/configs/qwen3.5-35b-a3b-eagle3-offline.yaml b/examples/configs/offline/colocated/qwen3.5-35b-a3b-eagle3-offline.yaml similarity index 100% rename from examples/configs/qwen3.5-35b-a3b-eagle3-offline.yaml rename to examples/configs/offline/colocated/qwen3.5-35b-a3b-eagle3-offline.yaml diff --git a/examples/configs/offline/colocated/qwen3.5-4b-dflash-offline-amd.yaml b/examples/configs/offline/colocated/qwen3.5-4b-dflash-offline-amd.yaml new file mode 100644 index 000000000..70bedc43b --- /dev/null +++ b/examples/configs/offline/colocated/qwen3.5-4b-dflash-offline-amd.yaml @@ -0,0 +1,45 @@ +# Qwen3.5-4B DFlash draft, offline (features captured to disk) — AMD ROCm. +# Trainer runs on ROCm unchanged (flex_attention). Capture the hidden states with +# scripts/prepare_hidden_states.py using the AITER + --sglang-disable-radix-cache +# flags (see docs/basic_usage/AMD/amd_rocm.md, Section 3). +model: + target_model_path: "Qwen/Qwen3.5-4B" + draft_model_config: "configs/qwen3.5-4b-dflash.json" + target_backend: "sglang" + trust_remote_code: true + embedding_key: "model.language_model.embed_tokens.weight" + torch_dtype: "bfloat16" + # Liger fused RMSNorm/SwiGLU: validated on MI355X (+10% tokens/s, -15% peak + # mem, loss parity). Requires the flex-backend fix (torch<2.11) to run. + use_liger_kernel: true +data: + hidden_states_path: "./cache/hidden_states/qwen3.5-4b-dflash-sharegpt" + max_length: 2048 + chat_template: "qwen3.5" + cache_dir: "./cache" +training: + strategy: "dflash" + num_epochs: 10 + max_steps: 10000 + batch_size: 2 + accumulation_steps: 4 + learning_rate: 6.0e-4 + warmup_ratio: 0.04 + max_grad_norm: 1 + attention_backend: "flex_attention" + num_anchors: 512 + loss_decay_gamma: 7 + save_interval: 1000 + log_interval: 50 + dist_timeout: 30 + seed: 42 +tracking: + report_to: "tensorboard" +run_id: "qwen3.5-4b-dflash-offline" +output_dir: "./outputs/qwen3.5-4b-dflash-offline" + +deployment: + mode: local_colocated + trainer: + nnodes: 1 + nproc_per_node: 1 diff --git a/examples/configs/qwen2.5-7b-eagle3-offline-disaggregated.yaml b/examples/configs/offline/disaggregated/qwen2.5-7b-eagle3-offline-disaggregated.yaml similarity index 100% rename from examples/configs/qwen2.5-7b-eagle3-offline-disaggregated.yaml rename to examples/configs/offline/disaggregated/qwen2.5-7b-eagle3-offline-disaggregated.yaml diff --git a/examples/configs/qwen3-8b-eagle3-offline-disaggregated.yaml b/examples/configs/offline/disaggregated/qwen3-8b-eagle3-offline-disaggregated.yaml similarity index 100% rename from examples/configs/qwen3-8b-eagle3-offline-disaggregated.yaml rename to examples/configs/offline/disaggregated/qwen3-8b-eagle3-offline-disaggregated.yaml diff --git a/examples/configs/online/colocated/README.md b/examples/configs/online/colocated/README.md new file mode 100644 index 000000000..506c54208 --- /dev/null +++ b/examples/configs/online/colocated/README.md @@ -0,0 +1,6 @@ +# Online colocated recipes + +This directory is reserved so the catalog keeps the same mode/topology axes for +online and offline recipes. All supported online recipes +use separate producer and consumer roles under `../disaggregated/`, regardless +of whether one supervisor starts both roles. diff --git a/examples/configs/deepseek-v2-lite-eagle3-online.yaml b/examples/configs/online/disaggregated/external/deepseek-v2-lite-eagle3-online.yaml similarity index 100% rename from examples/configs/deepseek-v2-lite-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/deepseek-v2-lite-eagle3-online.yaml diff --git a/examples/configs/deepseek-v3-671b-eagle3-online.yaml b/examples/configs/online/disaggregated/external/deepseek-v3-671b-eagle3-online.yaml similarity index 100% rename from examples/configs/deepseek-v3-671b-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/deepseek-v3-671b-eagle3-online.yaml diff --git a/examples/configs/gemma3-1b-eagle3-online.yaml b/examples/configs/online/disaggregated/external/gemma3-1b-eagle3-online.yaml similarity index 100% rename from examples/configs/gemma3-1b-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/gemma3-1b-eagle3-online.yaml diff --git a/examples/configs/glm-5.2-dspark-disaggregated.yaml b/examples/configs/online/disaggregated/external/glm-5.2-dspark-disaggregated.yaml similarity index 100% rename from examples/configs/glm-5.2-dspark-disaggregated.yaml rename to examples/configs/online/disaggregated/external/glm-5.2-dspark-disaggregated.yaml diff --git a/examples/configs/gpt-oss-120b-eagle3-online.yaml b/examples/configs/online/disaggregated/external/gpt-oss-120b-eagle3-online.yaml similarity index 100% rename from examples/configs/gpt-oss-120b-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/gpt-oss-120b-eagle3-online.yaml diff --git a/examples/configs/gpt-oss-20b-eagle3-online.yaml b/examples/configs/online/disaggregated/external/gpt-oss-20b-eagle3-online.yaml similarity index 100% rename from examples/configs/gpt-oss-20b-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/gpt-oss-20b-eagle3-online.yaml diff --git a/examples/configs/inkling-dspark-disaggregated.yaml b/examples/configs/online/disaggregated/external/inkling-dspark-disaggregated.yaml similarity index 69% rename from examples/configs/inkling-dspark-disaggregated.yaml rename to examples/configs/online/disaggregated/external/inkling-dspark-disaggregated.yaml index afadeac74..46e18e8dd 100644 --- a/examples/configs/inkling-dspark-disaggregated.yaml +++ b/examples/configs/online/disaggregated/external/inkling-dspark-disaggregated.yaml @@ -2,9 +2,22 @@ model: target_model_path: thinkingmachines/Inkling draft_model_config: configs/inkling-dspark.json target_backend: sglang + trust_remote_code: true tokenizer_pad_token_id: 200006 embedding_key: model.llm.embed.weight lm_head_key: model.llm.unembed.weight + # Inkling's hybrid cache path requires the unified radix tree. These values + # match the SGLang #31847 configuration validated by the two-node launcher. + sglang_attention_backend: fa4 + sglang_mem_fraction_static: 0.85 + sglang_disable_radix_cache: false + sglang_context_length: 4103 + sglang_moe_runner_backend: flashinfer_trtllm_routed + sglang_page_size: 128 + sglang_quantization: modelopt_fp4 + sglang_mamba_radix_cache_strategy: extra_buffer + sglang_max_mamba_cache_size: 64 + sglang_swa_full_tokens_ratio: 0.2 data: train_data_path: ./cache/dataset/inkling_dspark_train.jsonl max_length: 4096 @@ -20,6 +33,7 @@ training: learning_rate: 0.0006 warmup_ratio: 0.04 max_grad_norm: 1.0 + attention_backend: eager num_anchors: 512 loss_decay_gamma: 4.0 objective_chunk_blocks: 128 diff --git a/examples/configs/online/disaggregated/external/kimi-k3-dspark-disaggregated.yaml b/examples/configs/online/disaggregated/external/kimi-k3-dspark-disaggregated.yaml new file mode 100644 index 000000000..5f453c1f2 --- /dev/null +++ b/examples/configs/online/disaggregated/external/kimi-k3-dspark-disaggregated.yaml @@ -0,0 +1,91 @@ +model: + target_model_path: moonshotai/Kimi-K3 + draft_model_config: configs/kimi-k3-dspark.json + target_backend: sglang + trust_remote_code: true + embedding_key: language_model.model.embed_tokens.weight + lm_head_key: language_model.lm_head.weight + mask_token_id: 163824 + torch_dtype: bfloat16 + sglang_attention_backend: trtllm_mla + sglang_mem_fraction_static: 0.76 + sglang_context_length: 66048 + sglang_max_running_requests: 1 + sglang_max_total_tokens: 66048 + sglang_moe_runner_backend: marlin + sglang_mamba_radix_cache_strategy: extra_buffer + sglang_max_mamba_cache_size: 5 + +data: + train_data_path: ./cache/dataset/kimi_k3_dspark_train.jsonl + max_length: 65536 + chat_template: kimi-k3-thinking + cache_dir: cache + build_dataset_num_proc: 64 + dataloader_num_workers: 0 + +training: + strategy: dspark + num_epochs: 10 + # Four consumer ranks x one sequence x 32 microbatches preserves the old + # effective global batch of 128 sequences. + batch_size: 1 + accumulation_steps: 32 + learning_rate: 0.0006 + lr_scheduler: constant + warmup_ratio: 0 + max_grad_norm: 1 + attention_backend: flex_attention + num_anchors: 512 + loss_decay_gamma: 4.0 + objective_chunk_blocks: 128 + dspark_ce_loss_alpha: 0.1 + dspark_l1_loss_alpha: 0.9 + dspark_confidence_head_alpha: 1.0 + save_interval: 8 + log_interval: 10 + dist_timeout: 30 + seed: 42 + prompt_seed: 1 + +tracking: + report_to: wandb + wandb_project: specforge-dspark + wandb_name: kimi-k3-dspark-specforge-disaggregated + wandb_dir: outputs/kimi-k3-dspark-disaggregated/wandb + +runtime: + producer_lease: 1 + producer_concurrency: 1 + # One 64K K3 sample is roughly 5.25 GiB of BF16 DSpark features. A complete + # optimizer quantum is 128 samples, so the server segment must exceed 672 GiB. + in_flight_high_watermark: 128 + # Consumers dispatch complete optimizer windows. Keep the resume threshold + # at least one global quantum to avoid a producer/consumer backpressure + # deadlock while a window is still incomplete. + in_flight_low_watermark: 128 + resident_high_watermark_bytes: 858993459200 + resident_low_watermark_bytes: 697932185600 + feature_store_max_resident_bytes: 966367641600 + +run_id: kimi-k3-dspark-disaggregated +output_dir: outputs/kimi-k3-dspark-disaggregated + +deployment: + mode: disaggregated + trainer: + nnodes: 1 + nproc_per_node: 4 + disaggregated: + control_dir: outputs/kimi-k3-dspark-disaggregated/control + consumer_state_dir: outputs/kimi-k3-dspark-disaggregated/consumer-state + backend: mooncake + store_id: kimi-k3-dspark-specforge-disaggregated + server_urls: + - http://capture-node:30000 + mooncake_metadata_server: http://capture-node:35880/metadata + mooncake_master_server_addr: capture-node:35551 + mooncake_protocol: tcp + client_buffer_size: 1073741824 + idle_timeout_s: 7200 + peer_wait_timeout_s: 7200 diff --git a/examples/configs/lfm2.5-1.2b-instruct-dflash-online.yaml b/examples/configs/online/disaggregated/external/lfm2.5-1.2b-instruct-dflash-online.yaml similarity index 100% rename from examples/configs/lfm2.5-1.2b-instruct-dflash-online.yaml rename to examples/configs/online/disaggregated/external/lfm2.5-1.2b-instruct-dflash-online.yaml diff --git a/examples/configs/ling-flash-2.0-eagle3-online.yaml b/examples/configs/online/disaggregated/external/ling-flash-2.0-eagle3-online.yaml similarity index 100% rename from examples/configs/ling-flash-2.0-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/ling-flash-2.0-eagle3-online.yaml diff --git a/examples/configs/llama3.1-8b-eagle3-online.yaml b/examples/configs/online/disaggregated/external/llama3.1-8b-eagle3-online.yaml similarity index 100% rename from examples/configs/llama3.1-8b-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/llama3.1-8b-eagle3-online.yaml diff --git a/examples/configs/llama3.3-70b-eagle3-online.yaml b/examples/configs/online/disaggregated/external/llama3.3-70b-eagle3-online.yaml similarity index 100% rename from examples/configs/llama3.3-70b-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/llama3.3-70b-eagle3-online.yaml diff --git a/examples/configs/llama4-scout-17b-16e-eagle3-online.yaml b/examples/configs/online/disaggregated/external/llama4-scout-17b-16e-eagle3-online.yaml similarity index 100% rename from examples/configs/llama4-scout-17b-16e-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/llama4-scout-17b-16e-eagle3-online.yaml diff --git a/examples/configs/longcat-flash-dflash-online.yaml b/examples/configs/online/disaggregated/external/longcat-flash-dflash-online.yaml similarity index 100% rename from examples/configs/longcat-flash-dflash-online.yaml rename to examples/configs/online/disaggregated/external/longcat-flash-dflash-online.yaml diff --git a/examples/configs/longcat-flash-eagle3-online.yaml b/examples/configs/online/disaggregated/external/longcat-flash-eagle3-online.yaml similarity index 100% rename from examples/configs/longcat-flash-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/longcat-flash-eagle3-online.yaml diff --git a/examples/configs/phi4-eagle3-online.yaml b/examples/configs/online/disaggregated/external/phi4-eagle3-online.yaml similarity index 100% rename from examples/configs/phi4-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/phi4-eagle3-online.yaml diff --git a/examples/configs/qwen2.5-0.5b-dflash-online.yaml b/examples/configs/online/disaggregated/external/qwen2.5-0.5b-dflash-online.yaml similarity index 100% rename from examples/configs/qwen2.5-0.5b-dflash-online.yaml rename to examples/configs/online/disaggregated/external/qwen2.5-0.5b-dflash-online.yaml diff --git a/examples/configs/qwen2.5-0.5b-eagle3-online.yaml b/examples/configs/online/disaggregated/external/qwen2.5-0.5b-eagle3-online.yaml similarity index 100% rename from examples/configs/qwen2.5-0.5b-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/qwen2.5-0.5b-eagle3-online.yaml diff --git a/examples/configs/qwen3-235b-a22b-eagle3-online.yaml b/examples/configs/online/disaggregated/external/qwen3-235b-a22b-eagle3-online.yaml similarity index 100% rename from examples/configs/qwen3-235b-a22b-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/qwen3-235b-a22b-eagle3-online.yaml diff --git a/examples/configs/qwen3-30b-a3b-eagle3-online.yaml b/examples/configs/online/disaggregated/external/qwen3-30b-a3b-eagle3-online.yaml similarity index 100% rename from examples/configs/qwen3-30b-a3b-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/qwen3-30b-a3b-eagle3-online.yaml diff --git a/examples/configs/qwen3-30b-a3b-eagle3.1-online.yaml b/examples/configs/online/disaggregated/external/qwen3-30b-a3b-eagle3.1-online.yaml similarity index 100% rename from examples/configs/qwen3-30b-a3b-eagle3.1-online.yaml rename to examples/configs/online/disaggregated/external/qwen3-30b-a3b-eagle3.1-online.yaml diff --git a/examples/configs/qwen3-32b-eagle3-online.yaml b/examples/configs/online/disaggregated/external/qwen3-32b-eagle3-online.yaml similarity index 100% rename from examples/configs/qwen3-32b-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/qwen3-32b-eagle3-online.yaml diff --git a/examples/configs/qwen3-4b-dflash-online.yaml b/examples/configs/online/disaggregated/external/qwen3-4b-dflash-online.yaml similarity index 100% rename from examples/configs/qwen3-4b-dflash-online.yaml rename to examples/configs/online/disaggregated/external/qwen3-4b-dflash-online.yaml diff --git a/examples/configs/qwen3-4b-dspark-disaggregated.yaml b/examples/configs/online/disaggregated/external/qwen3-4b-dspark-disaggregated.yaml similarity index 100% rename from examples/configs/qwen3-4b-dspark-disaggregated.yaml rename to examples/configs/online/disaggregated/external/qwen3-4b-dspark-disaggregated.yaml diff --git a/examples/configs/qwen3-4b-eagle3-online.yaml b/examples/configs/online/disaggregated/external/qwen3-4b-eagle3-online.yaml similarity index 100% rename from examples/configs/qwen3-4b-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/qwen3-4b-eagle3-online.yaml diff --git a/examples/configs/qwen3-8b-dflash-disaggregated.yaml b/examples/configs/online/disaggregated/external/qwen3-8b-dflash-disaggregated.yaml similarity index 100% rename from examples/configs/qwen3-8b-dflash-disaggregated.yaml rename to examples/configs/online/disaggregated/external/qwen3-8b-dflash-disaggregated.yaml diff --git a/examples/configs/qwen3-8b-dflash-online.yaml b/examples/configs/online/disaggregated/external/qwen3-8b-dflash-online.yaml similarity index 100% rename from examples/configs/qwen3-8b-dflash-online.yaml rename to examples/configs/online/disaggregated/external/qwen3-8b-dflash-online.yaml diff --git a/examples/configs/qwen3-8b-domino-disaggregated.yaml b/examples/configs/online/disaggregated/external/qwen3-8b-domino-disaggregated.yaml similarity index 100% rename from examples/configs/qwen3-8b-domino-disaggregated.yaml rename to examples/configs/online/disaggregated/external/qwen3-8b-domino-disaggregated.yaml diff --git a/examples/configs/qwen3-8b-domino-online.yaml b/examples/configs/online/disaggregated/external/qwen3-8b-domino-online.yaml similarity index 100% rename from examples/configs/qwen3-8b-domino-online.yaml rename to examples/configs/online/disaggregated/external/qwen3-8b-domino-online.yaml diff --git a/examples/configs/qwen3-8b-dpace-online.yaml b/examples/configs/online/disaggregated/external/qwen3-8b-dpace-online.yaml similarity index 100% rename from examples/configs/qwen3-8b-dpace-online.yaml rename to examples/configs/online/disaggregated/external/qwen3-8b-dpace-online.yaml diff --git a/examples/configs/qwen3-8b-dspark-disaggregated.yaml b/examples/configs/online/disaggregated/external/qwen3-8b-dspark-disaggregated.yaml similarity index 100% rename from examples/configs/qwen3-8b-dspark-disaggregated.yaml rename to examples/configs/online/disaggregated/external/qwen3-8b-dspark-disaggregated.yaml diff --git a/examples/configs/qwen3-8b-eagle3-disaggregated.yaml b/examples/configs/online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml similarity index 100% rename from examples/configs/qwen3-8b-eagle3-disaggregated.yaml rename to examples/configs/online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml diff --git a/examples/configs/qwen3-8b-peagle-disaggregated.yaml b/examples/configs/online/disaggregated/external/qwen3-8b-peagle-disaggregated.yaml similarity index 100% rename from examples/configs/qwen3-8b-peagle-disaggregated.yaml rename to examples/configs/online/disaggregated/external/qwen3-8b-peagle-disaggregated.yaml diff --git a/examples/configs/qwen3-coder-30b-a3b-eagle3-online.yaml b/examples/configs/online/disaggregated/external/qwen3-coder-30b-a3b-eagle3-online.yaml similarity index 100% rename from examples/configs/qwen3-coder-30b-a3b-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/qwen3-coder-30b-a3b-eagle3-online.yaml diff --git a/examples/configs/qwen3-coder-480b-a35b-eagle3-online.yaml b/examples/configs/online/disaggregated/external/qwen3-coder-480b-a35b-eagle3-online.yaml similarity index 100% rename from examples/configs/qwen3-coder-480b-a35b-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/qwen3-coder-480b-a35b-eagle3-online.yaml diff --git a/examples/configs/qwen3-coder-next-eagle3-online.yaml b/examples/configs/online/disaggregated/external/qwen3-coder-next-eagle3-online.yaml similarity index 100% rename from examples/configs/qwen3-coder-next-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/qwen3-coder-next-eagle3-online.yaml diff --git a/examples/configs/qwen3-next-80b-a3b-eagle3-online.yaml b/examples/configs/online/disaggregated/external/qwen3-next-80b-a3b-eagle3-online.yaml similarity index 100% rename from examples/configs/qwen3-next-80b-a3b-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/qwen3-next-80b-a3b-eagle3-online.yaml diff --git a/examples/configs/qwen3.5-35b-a3b-dflash-online.yaml b/examples/configs/online/disaggregated/external/qwen3.5-35b-a3b-dflash-online.yaml similarity index 100% rename from examples/configs/qwen3.5-35b-a3b-dflash-online.yaml rename to examples/configs/online/disaggregated/external/qwen3.5-35b-a3b-dflash-online.yaml diff --git a/examples/configs/qwen3.5-35b-a3b-eagle3-online.yaml b/examples/configs/online/disaggregated/external/qwen3.5-35b-a3b-eagle3-online.yaml similarity index 100% rename from examples/configs/qwen3.5-35b-a3b-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/qwen3.5-35b-a3b-eagle3-online.yaml diff --git a/examples/configs/online/disaggregated/external/qwen3.5-4b-dflash-online-amd.yaml b/examples/configs/online/disaggregated/external/qwen3.5-4b-dflash-online-amd.yaml new file mode 100644 index 000000000..a37cd4d61 --- /dev/null +++ b/examples/configs/online/disaggregated/external/qwen3.5-4b-dflash-online-amd.yaml @@ -0,0 +1,56 @@ +# Qwen3.5-4B DFlash draft, online disaggregated — AMD ROCm single node. +# The producer/consumer stream target features live through Mooncake. Launch the +# patched capture server with AITER + --disable-radix-cache (Qwen3.5-4B is a hybrid +# Mamba target); the trainer runs on ROCm unchanged (flex_attention). +# See docs/basic_usage/AMD/amd_rocm.md, Section 4. +model: + target_model_path: "Qwen/Qwen3.5-4B" + draft_model_config: "configs/qwen3.5-4b-dflash.json" + target_backend: "sglang" + trust_remote_code: true + embedding_key: "model.language_model.embed_tokens.weight" + torch_dtype: "bfloat16" + # Liger fused RMSNorm/SwiGLU: validated on MI355X (+10% tokens/s, -15% peak + # mem, loss parity). Requires the flex-backend fix (torch<2.11) to run. + use_liger_kernel: true +data: + train_data_path: "./cache/dataset/sharegpt_train.jsonl" + max_length: 2048 + chat_template: "qwen3.5" + build_dataset_num_proc: 32 + cache_dir: "./cache" +training: + strategy: "dflash" + num_epochs: 10 + max_steps: 10000 + batch_size: 2 + accumulation_steps: 4 + learning_rate: 6.0e-4 + warmup_ratio: 0.04 + max_grad_norm: 1 + attention_backend: "flex_attention" + num_anchors: 512 + loss_decay_gamma: 7 + save_interval: 1000 + log_interval: 50 + dist_timeout: 30 + seed: 42 +tracking: + report_to: "tensorboard" +run_id: "qwen3.5-4b-dflash-online" +output_dir: "./outputs/qwen3.5-4b-dflash-online" + +deployment: + mode: disaggregated + trainer: + nnodes: 1 + nproc_per_node: 1 + disaggregated: + control_dir: outputs/qwen3.5-4b-dflash-online/control + consumer_state_dir: outputs/qwen3.5-4b-dflash-online/consumer-state + backend: mooncake + server_urls: + - http://127.0.0.1:30000 + mooncake_metadata_server: http://127.0.0.1:35880/metadata + mooncake_master_server_addr: 127.0.0.1:35551 + mooncake_protocol: tcp diff --git a/examples/configs/qwen3.5-4b-dflash-online-npu.yaml b/examples/configs/online/disaggregated/external/qwen3.5-4b-dflash-online-npu.yaml similarity index 100% rename from examples/configs/qwen3.5-4b-dflash-online-npu.yaml rename to examples/configs/online/disaggregated/external/qwen3.5-4b-dflash-online-npu.yaml diff --git a/examples/configs/qwen3.5-4b-domino-online-npu.yaml b/examples/configs/online/disaggregated/external/qwen3.5-4b-domino-online-npu.yaml similarity index 100% rename from examples/configs/qwen3.5-4b-domino-online-npu.yaml rename to examples/configs/online/disaggregated/external/qwen3.5-4b-domino-online-npu.yaml diff --git a/examples/configs/qwen3.5-4b-vl-dflash-disaggregated.yaml b/examples/configs/online/disaggregated/external/qwen3.5-4b-vl-dflash-disaggregated.yaml similarity index 91% rename from examples/configs/qwen3.5-4b-vl-dflash-disaggregated.yaml rename to examples/configs/online/disaggregated/external/qwen3.5-4b-vl-dflash-disaggregated.yaml index 6b846bfba..099ac219a 100644 --- a/examples/configs/qwen3.5-4b-vl-dflash-disaggregated.yaml +++ b/examples/configs/online/disaggregated/external/qwen3.5-4b-vl-dflash-disaggregated.yaml @@ -1,13 +1,13 @@ # Multimodal (image+text) DFlash on Qwen3.5-4B: disaggregated online training # against an external patched SGLang capture server. Requires sglang v0.5.14 -# with patches/sglang/v0.5.14/spec-capture.patch (position_ids artifact) and +# with patches/sglang/v0.5.14/spec-capture.patch and # SGLANG_MM_AVOID_RETOKENIZE=1 on the capture server. # # Data: ShareGPT-style JSONL with an optional top-level "image" field (file # path or base64) per record; text-only records are supported in the same run. model: target_model_path: Qwen/Qwen3.5-4B - draft_model_config: configs/qwen3.5-4b-vl-dflash.json + draft_model_config: configs/qwen3.5-4b-dflash.json target_backend: sglang trust_remote_code: true input_modality: multimodal diff --git a/examples/configs/qwen3.6-27b-dflash-disaggregated.yaml b/examples/configs/online/disaggregated/external/qwen3.6-27b-dflash-disaggregated.yaml similarity index 100% rename from examples/configs/qwen3.6-27b-dflash-disaggregated.yaml rename to examples/configs/online/disaggregated/external/qwen3.6-27b-dflash-disaggregated.yaml diff --git a/examples/configs/qwen3.6-27b-dflash-online.yaml b/examples/configs/online/disaggregated/external/qwen3.6-27b-dflash-online.yaml similarity index 100% rename from examples/configs/qwen3.6-27b-dflash-online.yaml rename to examples/configs/online/disaggregated/external/qwen3.6-27b-dflash-online.yaml diff --git a/examples/configs/qwen3.6-27b-domino-online.yaml b/examples/configs/online/disaggregated/external/qwen3.6-27b-domino-online.yaml similarity index 81% rename from examples/configs/qwen3.6-27b-domino-online.yaml rename to examples/configs/online/disaggregated/external/qwen3.6-27b-domino-online.yaml index 17bc8642f..6bdc70e0b 100644 --- a/examples/configs/qwen3.6-27b-domino-online.yaml +++ b/examples/configs/online/disaggregated/external/qwen3.6-27b-domino-online.yaml @@ -16,6 +16,7 @@ training: strategy: "domino" num_epochs: 6 batch_size: 1 + accumulation_steps: 8 learning_rate: 0.0006 warmup_ratio: 0.04 max_grad_norm: 1 @@ -31,11 +32,24 @@ training: run_id: "qwen3.6-27b-domino-online" output_dir: "./outputs/qwen3.6-27b-domino-online" +profiling: + enabled: false + start_step: 0 + num_steps: 20 + record_shapes: false + +runtime: + producer_lease: 64 + producer_concurrency: 32 + in_flight_high_watermark: 1024 + in_flight_low_watermark: 512 + + deployment: mode: disaggregated trainer: nnodes: 1 - nproc_per_node: 8 + nproc_per_node: 1 disaggregated: control_dir: outputs/qwen3.6-27b-domino-online/control consumer_state_dir: outputs/qwen3.6-27b-domino-online/consumer-state @@ -44,4 +58,5 @@ deployment: - http://127.0.0.1:30000 mooncake_metadata_server: http://127.0.0.1:35880/metadata mooncake_master_server_addr: 127.0.0.1:35551 + mooncake_local_hostname: 127.0.0.1 mooncake_protocol: tcp diff --git a/examples/configs/online/disaggregated/external/qwen3.6-27b-dspark-online.yaml b/examples/configs/online/disaggregated/external/qwen3.6-27b-dspark-online.yaml new file mode 100644 index 000000000..98045ec91 --- /dev/null +++ b/examples/configs/online/disaggregated/external/qwen3.6-27b-dspark-online.yaml @@ -0,0 +1,50 @@ +model: + target_model_path: Qwen/Qwen3.6-27B + draft_model_config: configs/qwen3.6-27b-dspark.json + target_backend: sglang + trust_remote_code: true + embedding_key: "model.language_model.embed_tokens.weight" + torch_dtype: "bfloat16" + mask_token_id: 248070 + +data: + train_data_path: "./cache/dataset/nemotron_v2_train.jsonl" + max_length: 4096 + chat_template: qwen3.5 + cache_dir: cache + build_dataset_num_proc: 64 + +training: + strategy: dspark + num_epochs: 10 + # One trainer rank with 512 microbatches preserves global batch size 512. + batch_size: 1 + accumulation_steps: 512 + learning_rate: 0.0006 + warmup_ratio: 0.04 + max_grad_norm: 1.0 + num_anchors: 512 + loss_decay_gamma: 4.0 + objective_chunk_blocks: 128 + save_interval: 125 + dist_timeout: 30 + seed: 42 + +run_id: qwen3.6-27b-dspark-online +output_dir: outputs/qwen3.6-27b-dspark-online + +deployment: + mode: disaggregated + trainer: + nnodes: 1 + nproc_per_node: 1 + disaggregated: + control_dir: outputs/qwen3.6-27b-dspark-online/control + consumer_state_dir: outputs/qwen3.6-27b-dspark-online/consumer-state + backend: mooncake + server_urls: + - http://127.0.0.1:30000 + mooncake_metadata_server: http://127.0.0.1:35880/metadata + mooncake_master_server_addr: 127.0.0.1:35551 + mooncake_local_hostname: 127.0.0.1 + mooncake_protocol: tcp diff --git a/examples/configs/qwq-32b-eagle3-online.yaml b/examples/configs/online/disaggregated/external/qwq-32b-eagle3-online.yaml similarity index 100% rename from examples/configs/qwq-32b-eagle3-online.yaml rename to examples/configs/online/disaggregated/external/qwq-32b-eagle3-online.yaml diff --git a/examples/configs/qwen3-8b-dflash-1server-dp7-disaggregated.yaml b/examples/configs/online/disaggregated/managed-local/qwen3-8b-dflash-1server-dp7-disaggregated.yaml similarity index 100% rename from examples/configs/qwen3-8b-dflash-1server-dp7-disaggregated.yaml rename to examples/configs/online/disaggregated/managed-local/qwen3-8b-dflash-1server-dp7-disaggregated.yaml diff --git a/examples/configs/qwen3-8b-domino-1server-dp7-disaggregated.yaml b/examples/configs/online/disaggregated/managed-local/qwen3-8b-domino-1server-dp7-disaggregated.yaml similarity index 100% rename from examples/configs/qwen3-8b-domino-1server-dp7-disaggregated.yaml rename to examples/configs/online/disaggregated/managed-local/qwen3-8b-domino-1server-dp7-disaggregated.yaml diff --git a/examples/configs/qwen3-8b-domino-multiserver-disaggregated.yaml b/examples/configs/online/disaggregated/managed-local/qwen3-8b-domino-multiserver-disaggregated.yaml similarity index 100% rename from examples/configs/qwen3-8b-domino-multiserver-disaggregated.yaml rename to examples/configs/online/disaggregated/managed-local/qwen3-8b-domino-multiserver-disaggregated.yaml diff --git a/examples/configs/online/disaggregated/managed-local/qwen3.5-4b-dflash-disaggregated-npu.yaml b/examples/configs/online/disaggregated/managed-local/qwen3.5-4b-dflash-disaggregated-npu.yaml new file mode 100644 index 000000000..6617a8da7 --- /dev/null +++ b/examples/configs/online/disaggregated/managed-local/qwen3.5-4b-dflash-disaggregated-npu.yaml @@ -0,0 +1,61 @@ +# Single-node managed full stack for Ascend NPU, launched with one +# `specforge train` command. Device ordinals are injected via +# ASCEND_RT_VISIBLE_DEVICES. +model: + target_model_path: "Qwen/Qwen3.5-4B" + draft_model_config: "configs/qwen3.5-4b-dflash.json" + target_backend: sglang + trust_remote_code: true + embedding_key: "model.language_model.embed_tokens.weight" + torch_dtype: "bfloat16" + sglang_attention_backend: ascend +data: + train_data_path: "./cache/dataset/train_regen.jsonl" + max_length: 3072 + chat_template: "qwen3.5" + build_dataset_num_proc: 32 + cache_dir: "./cache" +training: + strategy: "dflash" + num_epochs: 10 + max_steps: 10000 + batch_size: 2 + accumulation_steps: 4 + learning_rate: 0.0006 + warmup_ratio: 0.04 + max_grad_norm: 1 + attention_backend: "sdpa" + # If OOM occurs on low-memory NPU devices, retry with a smaller num_anchors value. + num_anchors: 512 + loss_decay_gamma: 7 + save_interval: 10000 + log_interval: 50 + dist_timeout: 30 + seed: 42 +tracking: + report_to: "tensorboard" +run_id: "qwen3.5-4b-dflash-npu-managed" +output_dir: "./outputs/qwen3.5-4b-dflash-npu-managed" + +deployment: + mode: disaggregated + trainer: + nnodes: 1 + nproc_per_node: 14 + disaggregated: + control_dir: outputs/qwen3.5-4b-dflash-npu-managed/control + consumer_state_dir: outputs/qwen3.5-4b-dflash-npu-managed/consumer-state + backend: mooncake + managed_local: + # 16-card A3 layout: capture server owns device 0, trainer owns 2-15. + # On 8-card hosts use devices 1-7 with trainer.nproc_per_node: 7. + trainer_cuda_visible_devices: + ["2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13", "14", "15"] + mooncake: + # Host segment mounted by the capture-server sink for feature storage. + # 16 GiB fits a 64GB-card host next to the target model and KV cache. + global_segment_size_bytes: 17179869184 + capture_servers: + - port: 30000 + cuda_visible_devices: ["0"] + tp_size: 1 diff --git a/examples/configs/online/disaggregated/managed-local/qwen3.5-4b-mtp-disaggregated-npu.yaml b/examples/configs/online/disaggregated/managed-local/qwen3.5-4b-mtp-disaggregated-npu.yaml new file mode 100644 index 000000000..8b22ab424 --- /dev/null +++ b/examples/configs/online/disaggregated/managed-local/qwen3.5-4b-mtp-disaggregated-npu.yaml @@ -0,0 +1,77 @@ +# Single-node managed full stack for Ascend NPU, launched with one +# `specforge train` command. Device ordinals are injected via +# ASCEND_RT_VISIBLE_DEVICES. +model: + target_model_path: "Qwen/Qwen3.5-4B" + draft_model_config: "configs/qwen3.5-4b-mtp.json" + target_backend: sglang + trust_remote_code: true + # Qwen3.5-4B nests its text decoder under model.language_model and ties + # lm_head to the embedding, so both keys point at the same tensor. + embedding_key: "model.language_model.embed_tokens.weight" + lm_head_key: "model.language_model.embed_tokens.weight" + torch_dtype: "bfloat16" + sglang_attention_backend: ascend +data: + train_data_path: "./cache/dataset/train_regen.jsonl" + max_length: 32768 + chat_template: "qwen3.5" + build_dataset_num_proc: 32 + cache_dir: "./cache" +training: + strategy: "mtp" + num_epochs: 10 + max_steps: 10000 + batch_size: 2 + accumulation_steps: 4 + learning_rate: 0.0006 + warmup_ratio: 0.04 + max_grad_norm: 1 + attention_backend: "sdpa" + save_interval: 10000 + log_interval: 50 + dist_timeout: 30 + seed: 42 +tracking: + report_to: "tensorboard" +run_id: "qwen3.5-4b-mtp-npu-managed" +output_dir: "./outputs/qwen3.5-4b-mtp-npu-managed" + +deployment: + mode: disaggregated + trainer: + nnodes: 1 + nproc_per_node: 10 + disaggregated: + control_dir: outputs/qwen3.5-4b-mtp-npu-managed/control + consumer_state_dir: outputs/qwen3.5-4b-mtp-npu-managed/consumer-state + backend: mooncake + managed_local: + # Validated 16-card A3 layout: 6 TP=1 capture servers own devices 0-5, + # the 10-rank trainer owns devices 6-15 (global batch = 2*4*10 = 80). + # The pipeline is capture-bound, so scale capture servers first. + trainer_cuda_visible_devices: + ["6", "7", "8", "9", "10", "11", "12", "13", "14", "15"] + mooncake: + # Host segment mounted by the capture-server sink for feature storage. + # 16 GiB fits a 64GB-card host next to the target model and KV cache. + global_segment_size_bytes: 17179869184 + capture_servers: + - port: 30000 + cuda_visible_devices: ["0"] + tp_size: 1 + - port: 30001 + cuda_visible_devices: ["1"] + tp_size: 1 + - port: 30002 + cuda_visible_devices: ["2"] + tp_size: 1 + - port: 30003 + cuda_visible_devices: ["3"] + tp_size: 1 + - port: 30004 + cuda_visible_devices: ["4"] + tp_size: 1 + - port: 30005 + cuda_visible_devices: ["5"] + tp_size: 1 diff --git a/examples/configs/qwen3.6-27b-dflash-1server-dp2-disaggregated.yaml b/examples/configs/online/disaggregated/managed-local/qwen3.6-27b-dflash-1server-dp2-disaggregated.yaml similarity index 100% rename from examples/configs/qwen3.6-27b-dflash-1server-dp2-disaggregated.yaml rename to examples/configs/online/disaggregated/managed-local/qwen3.6-27b-dflash-1server-dp2-disaggregated.yaml diff --git a/examples/configs/qwen3.6-27b-dflash-multiserver-disaggregated.yaml b/examples/configs/online/disaggregated/managed-local/qwen3.6-27b-dflash-multiserver-disaggregated.yaml similarity index 100% rename from examples/configs/qwen3.6-27b-dflash-multiserver-disaggregated.yaml rename to examples/configs/online/disaggregated/managed-local/qwen3.6-27b-dflash-multiserver-disaggregated.yaml diff --git a/examples/configs/qwen3.6-27b-dspark-disaggregated.yaml b/examples/configs/online/disaggregated/managed-local/qwen3.6-27b-dspark-disaggregated.yaml similarity index 100% rename from examples/configs/qwen3.6-27b-dspark-disaggregated.yaml rename to examples/configs/online/disaggregated/managed-local/qwen3.6-27b-dspark-disaggregated.yaml diff --git a/examples/disagg/README.md b/examples/disagg/README.md index 6be6e0841..cda1305ac 100644 --- a/examples/disagg/README.md +++ b/examples/disagg/README.md @@ -1,10 +1,11 @@ # Disaggregated examples -Disaggregated training uses the same typed config and public command as every -other run: +This directory contains optional launch helpers for producer/consumer runs; the +YAML recipes themselves live under `examples/configs`. Disaggregated training +uses the same typed config and public command as every other run: ```bash -specforge train -c examples/configs/qwen3-8b-dflash-disaggregated.yaml +specforge train -c examples/configs/online/disaggregated/external/qwen3-8b-dflash-disaggregated.yaml ``` For a single trainer node, that command supervises the SpecForge producer and @@ -12,15 +13,28 @@ consumer together. The producer is one direct process; the consumer topology comes from `deployment.trainer` and is launched through torch distributed when `nproc_per_node > 1`. +Choose the YAML from the category that matches the data flow: + +| Category | Recipe directory | What the launcher owns | +| --- | --- | --- | +| Offline disaggregated | `examples/configs/offline/disaggregated/` | Producer and consumer; feature storage is `shared_dir` or Mooncake | +| Online disaggregated, external | `examples/configs/online/disaggregated/external/` | Producer and consumer only; Mooncake and SGLang already exist | +| Online disaggregated, managed-local | `examples/configs/online/disaggregated/managed-local/` | Producer, consumer, local Mooncake, and local SGLang servers | + +`external` is a lifecycle boundary, not a network-location label. Services on +`127.0.0.1` are `external` when the user started them. Similarly, supervising +producer and consumer with one command does not make the topology colocated; +they remain separate roles connected through the disaggregated data plane. + The scripts in this directory are optional thin examples. The single-node wrappers only add the config path and forward arguments to `specforge train`; they contain no second trainer, torchrun construction, or transport validation: ```bash -CONFIG=examples/configs/qwen3-8b-dflash-disaggregated.yaml \ +CONFIG=examples/configs/online/disaggregated/external/qwen3-8b-dflash-disaggregated.yaml \ examples/disagg/run_online.sh -CONFIG=examples/configs/qwen3-8b-eagle3-offline-disaggregated.yaml \ +CONFIG=examples/configs/offline/disaggregated/qwen3-8b-eagle3-offline-disaggregated.yaml \ examples/disagg/run_offline.sh ``` @@ -30,7 +44,7 @@ rank dispatcher on both nodes. The cluster launcher supplies ```bash rcli exec --per-node \ - 'CONFIG=examples/configs/qwen2.5-7b-eagle3-offline-disaggregated.yaml bash examples/disagg/run_offline_2node.sh' + 'CONFIG=examples/configs/offline/disaggregated/qwen2.5-7b-eagle3-offline-disaggregated.yaml bash examples/disagg/run_offline_2node.sh' ``` This wrapper only maps rank to `specforge train --role`; the YAML still owns @@ -44,10 +58,10 @@ not expose the corresponding `RCLI_*` variables. Use the same YAML when producer and consumer belong to different pools: ```bash -specforge train -c examples/configs/qwen3-8b-dflash-disaggregated.yaml \ +specforge train -c examples/configs/online/disaggregated/external/qwen3-8b-dflash-disaggregated.yaml \ --role producer -specforge train -c examples/configs/qwen3-8b-dflash-disaggregated.yaml \ +specforge train -c examples/configs/online/disaggregated/external/qwen3-8b-dflash-disaggregated.yaml \ --role consumer ``` @@ -88,18 +102,76 @@ directory, while the wrapper's lifecycle markers stay at the shared run root. The consumer's SQLite/WAL and rank inboxes default to the trainer-node-local `/tmp/specforge/$DISAGG_STORE_ID/consumer-state`; override `DISAGG_CONSUMER_STATE_DIR` or `LOCAL_SCRATCH` when `/tmp` is unsuitable. -Node-local consumer state currently supports one trainer node only. +For a multi-node trainer, set `deployment.disaggregated.inbox_server_url` to a +private HTTP origin on trainer node 0. Rank 0 owns the SQLite/WAL and relays +only tensor-free `SampleRef` metadata to the other trainer nodes; feature +tensors continue to move directly through the selected feature store. + +## Checkpoints without shared storage + +Each distributed rank writes its own `training_state_rankN.pt`. If every +trainer node resolves `output_dir` to the same shared filesystem, no extra +step is required. If `output_dir` is node-local, run the checked-in relay on +both trainer nodes before training. It exchanges rank-local archives over the +private trainer network, verifies SHA-256, and atomically assembles a complete +checkpoint directory on each node: + +```bash +# Trainer node 0 (ranks 0-7) +python examples/disagg/sync_distributed_checkpoints.py \ + --run-root /workspace/runs/$RUN_ID \ + --run-id "$RUN_ID" \ + --local-ranks 0-7 --peer-ranks 8-15 \ + --serve-host 10.0.0.3 --serve-port 35914 \ + --peer-url http://10.0.0.4:35915 \ + --max-archives 3 + +# Trainer node 1 (ranks 8-15) +python examples/disagg/sync_distributed_checkpoints.py \ + --run-root /workspace/runs/$RUN_ID \ + --run-id "$RUN_ID" \ + --local-ranks 8-15 --peer-ranks 0-7 \ + --serve-host 10.0.0.4 --serve-port 35915 \ + --peer-url http://10.0.0.3:35914 \ + --max-archives 3 +``` + +The relay has no authentication or TLS; bind it only to a trusted private +interface. Set `--max-archives` to the same retention window as +`training.max_checkpoints` so relay archives cannot grow without bound. + +The Inkling DSpark variant uses the same two-node lifecycle with the target +settings validated against SGLang +[#31847](https://github.com/sgl-project/sglang/pull/31847): + +```bash +export DISAGG_STORE_ID=inkling-two-node-attempt-001 +export DISAGG_RUN_ROOT=/shared/specforge/$DISAGG_STORE_ID + +rcli exec --per-node \ + 'bash examples/disagg/run_inkling_dspark_disagg_2node.sh' +``` + +Rank 0 uses four GPUs for TP4 ModelOpt-FP4 capture; rank 1 defaults to four +FSDP trainer ranks. Override `TARGET_MODEL_PATH`, `SERVER_GPUS`, +`TRAINER_GPUS`, or `TRAINER_NPROC` for another allocation. The launcher keeps +the unified radix tree enabled and does not pass `--disable-radix-cache`. +Until #31847 is available in a supported SGLang release, install that PR's +checkout into both nodes' environment. The wrapper applies SpecForge's +checked-in capture patch before starting the server; the patch is dry-run +validated against both v0.5.14 and #31847 commit `b7252cc`. ## External and managed-local services -By default, online capture requires an already-running Mooncake deployment and -patched SGLang capture server. Those long-lived services are not started or -stopped by `specforge train`. Put stable, non-secret topology in the typed -`deployment.disaggregated` section; inject authentication tokens, each node's -Mooncake hostname, and device visibility through the deployment environment. -The checked-in external-service recipes point at the standard local demo ports; -replace those endpoint fields, or override them with environment values, for a -remote deployment. +Recipes under `online/disaggregated/external` require an already-running +Mooncake deployment and patched SGLang capture server. Those services are not +started or stopped by `specforge train`. Put stable, non-secret topology in the +typed `deployment.disaggregated` section; inject authentication tokens, each +node's Mooncake hostname, and device visibility through the deployment +environment. +The checked-in external recipes point at standard local demo ports; replace +those endpoint fields, or override them with environment values, for another +deployment. For a self-contained single-node development run, the managed-local recipes record Mooncake, one or more capture servers, their GPU placement, and the DP @@ -107,10 +179,10 @@ trainer in one YAML: ```bash specforge train -c \ - examples/configs/qwen3-8b-dflash-1server-dp7-disaggregated.yaml + examples/configs/online/disaggregated/managed-local/qwen3-8b-dflash-1server-dp7-disaggregated.yaml specforge train -c \ - examples/configs/qwen3-8b-domino-1server-dp7-disaggregated.yaml + examples/configs/online/disaggregated/managed-local/qwen3-8b-domino-1server-dp7-disaggregated.yaml ``` These recipes preserve the old DFlash and Domino one-server + DP7 @@ -118,17 +190,17 @@ self-contained topologies. The genuine two-server Domino recipe is: ```bash specforge train -c \ - examples/configs/qwen3-8b-domino-multiserver-disaggregated.yaml + examples/configs/online/disaggregated/managed-local/qwen3-8b-domino-multiserver-disaggregated.yaml ``` Qwen3.6 DFlash has both one-server and larger two-TP2-server managed recipes: ```bash specforge train -c \ - examples/configs/qwen3.6-27b-dflash-1server-dp2-disaggregated.yaml + examples/configs/online/disaggregated/managed-local/qwen3.6-27b-dflash-1server-dp2-disaggregated.yaml specforge train -c \ - examples/configs/qwen3.6-27b-dflash-multiserver-disaggregated.yaml + examples/configs/online/disaggregated/managed-local/qwen3.6-27b-dflash-multiserver-disaggregated.yaml ``` The first command preserves the historical one-server + DP2 self-contained @@ -138,12 +210,6 @@ That opt-in profile starts, health-checks, and cleans up the owned local services. It does not change the default external-service boundary or attempt to schedule services on remote hosts. -The strict e2e gate at -`scripts/gates/run_disaggregated_overfit_gate.sh` retains full local test-stack -automation: it starts and health-checks Mooncake and SGLang, runs the unified -producer/consumer entry, verifies training and serving, and cleans up owned -processes. That test harness is not the production service supervisor. - Online configs use Mooncake. Offline configs may use either a typed `shared_dir` store or Mooncake. `deployment.disaggregated.control_dir` is the one attempt root from which the launcher derives the reference channel or diff --git a/examples/disagg/run_inkling_dspark_disagg_2node.sh b/examples/disagg/run_inkling_dspark_disagg_2node.sh new file mode 100755 index 000000000..fee19df0c --- /dev/null +++ b/examples/disagg/run_inkling_dspark_disagg_2node.sh @@ -0,0 +1,43 @@ +#!/usr/bin/env bash +# Two-node Inkling DSpark recipe: +# rank 0: Mooncake + SGLang #31847 TP4 capture + CPU producer +# rank 1: four-rank FSDP consumer/trainer +# +# Launch this command on both nodes. The cluster launcher supplies +# RCLI_NODE_RANK, RCLI_NUM_NODES, and RCLI_HEAD_IP; both nodes must share the +# fresh DISAGG_RUN_ROOT. Install SGLang #31847 in the active environment; the +# shared launcher applies the checked-in SpecForge capture patch before start. +set -Eeuo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +ROOT_DIR="$(dirname "$(dirname "$SCRIPT_DIR")")" + +export CONFIG="${CONFIG:-$ROOT_DIR/examples/configs/online/disaggregated/external/inkling-dspark-disaggregated.yaml}" +export RUN_LABEL="${RUN_LABEL:-inkling-dspark-2node}" +export TARGET_MODEL_PATH="${TARGET_MODEL_PATH:-thinkingmachines/Inkling}" + +export SERVER_GPUS="${SERVER_GPUS:-0,1,2,3}" +export SERVER_TP="${SERVER_TP:-4}" +export SERVER_MEM_FRACTION="${SERVER_MEM_FRACTION:-0.85}" +export CAPTURE_LAYER_IDS="${CAPTURE_LAYER_IDS:-5 17 35 47 59}" + +export TRAINER_GPUS="${TRAINER_GPUS:-0,1,2,3}" +export TRAINER_NPROC="${TRAINER_NPROC:-4}" +TRAINER_ACCUMULATION_STEPS="${TRAINER_ACCUMULATION_STEPS:-128}" + +export APPLY_SGLANG_CAPTURE_PATCH="${APPLY_SGLANG_CAPTURE_PATCH:-1}" +export SGLANG_ENABLE_UNIFIED_RADIX_TREE="${SGLANG_ENABLE_UNIFIED_RADIX_TREE:-1}" +export SGLANG_OPT_USE_INKLING_CUSTOM_AR="${SGLANG_OPT_USE_INKLING_CUSTOM_AR:-1}" + +DEFAULT_SERVER_EXTRA_ARGS="--dtype bfloat16 --attention-backend fa4" +DEFAULT_SERVER_EXTRA_ARGS+=" --context-length 4103 --quantization modelopt_fp4" +DEFAULT_SERVER_EXTRA_ARGS+=" --moe-runner-backend flashinfer_trtllm_routed" +DEFAULT_SERVER_EXTRA_ARGS+=" --page-size 128" +DEFAULT_SERVER_EXTRA_ARGS+=" --mamba-radix-cache-strategy extra_buffer" +DEFAULT_SERVER_EXTRA_ARGS+=" --max-mamba-cache-size 64" +DEFAULT_SERVER_EXTRA_ARGS+=" --swa-full-tokens-ratio 0.2" +export SERVER_EXTRA_ARGS="${SERVER_EXTRA_ARGS:-$DEFAULT_SERVER_EXTRA_ARGS}" + +exec "$SCRIPT_DIR/run_qwen3_8b_dflash_disagg_2node.sh" \ + "training.accumulation_steps=$TRAINER_ACCUMULATION_STEPS" \ + "$@" diff --git a/examples/disagg/run_qwen3.6-27b-domino.sh b/examples/disagg/run_qwen3.6-27b-domino.sh new file mode 100644 index 000000000..bf6f55bd1 --- /dev/null +++ b/examples/disagg/run_qwen3.6-27b-domino.sh @@ -0,0 +1,61 @@ +# producer +mooncake_master \ + --enable_http_metadata_server=true \ + --http_metadata_server_host=127.0.0.1 \ + --rpc_port=35551 \ + --http_metadata_server_port=35880 \ + --metrics_port=35903 + + +export FLASHINFER_DISABLE_VERSION_CHECK=1 +export MOONCAKE_METADATA_SERVER=http://127.0.0.1:35880/metadata +export MOONCAKE_MASTER_SERVER_ADDR=127.0.0.1:35551 +export MOONCAKE_LOCAL_HOSTNAME=127.0.0.1 +export MOONCAKE_PROTOCOL=tcp +export MOONCAKE_GLOBAL_SEGMENT_SIZE=68719476736 +export MOONCAKE_LOCAL_BUFFER_SIZE=1073741824 + +export LD_LIBRARY_PATH="/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH:-}" +CUDA_VISIBLE_DEVICES=2 \ + python -m sglang.launch_server \ + --model-path Qwen/Qwen3.6-27B \ + --dtype bfloat16 \ + --trust-remote-code \ + --tp-size 1 \ + --context-length 16384 \ + --attention-backend flashinfer \ + --chunked-prefill-size -1 \ + --disable-radix-cache \ + --disable-cuda-graph \ + --enable-spec-capture \ + --spec-capture-method dflash \ + --spec-capture-aux-layer-ids 1 16 31 46 61 \ + --host 127.0.0.1 + +# Check sglang capture whether run success. +curl -f http://127.0.0.1:30000/health + + +RUN_ID=qwen3.6-27b-domino-split-disaggregated-08082328 +OUTPUT_DIR=outputs/${RUN_ID} +CONTROL_DIR=${OUTPUT_DIR}/control +CONSUMER_STATE_DIR=${OUTPUT_DIR}/consumer-state + +# producer +CUDA_VISIBLE_DEVICES=2 specforge train \ + -c examples/configs/online/disaggregated/external/qwen3.6-27b-domino-online.yaml \ + --role producer \ + "run_id=${RUN_ID}" \ + "output_dir=${OUTPUT_DIR}" \ + "deployment.disaggregated.control_dir=${CONTROL_DIR}" \ + "deployment.disaggregated.consumer_state_dir=${CONSUMER_STATE_DIR}" + + +# consumer +CUDA_VISIBLE_DEVICES=3 specforge train \ + -c examples/configs/online/disaggregated/external/qwen3.6-27b-domino-online.yaml \ + --role consumer \ + "run_id=${RUN_ID}" \ + "output_dir=${OUTPUT_DIR}" \ + "deployment.disaggregated.control_dir=${CONTROL_DIR}" \ + "deployment.disaggregated.consumer_state_dir=${CONSUMER_STATE_DIR}" diff --git a/examples/disagg/run_qwen3_8b_dflash_disagg_2node.sh b/examples/disagg/run_qwen3_8b_dflash_disagg_2node.sh index 60531035c..5e110e1b8 100755 --- a/examples/disagg/run_qwen3_8b_dflash_disagg_2node.sh +++ b/examples/disagg/run_qwen3_8b_dflash_disagg_2node.sh @@ -17,7 +17,8 @@ HEAD_IP="${HEAD_IP:-${RCLI_HEAD_IP:-}}" RUN_ID="${DISAGG_STORE_ID:-}" RUN_ROOT="${DISAGG_RUN_ROOT:-}" CONSUMER_STATE_DIR="${DISAGG_CONSUMER_STATE_DIR:-${LOCAL_SCRATCH:-/tmp}/specforge/$RUN_ID/consumer-state}" -CONFIG="${CONFIG:-$ROOT_DIR/examples/configs/qwen3-8b-dflash-disaggregated.yaml}" +CONFIG="${CONFIG:-$ROOT_DIR/examples/configs/online/disaggregated/external/qwen3-8b-dflash-disaggregated.yaml}" +RUN_LABEL="${RUN_LABEL:-qwen3-8b-dflash-2node}" SERVER_GPUS="${SERVER_GPUS:-0}" SERVER_TP="${SERVER_TP:-1}" @@ -27,6 +28,15 @@ CAPTURE_LAYER_IDS="${CAPTURE_LAYER_IDS:-1 9 17 25 33}" TRAINER_GPUS="${TRAINER_GPUS:-0,1,2,3}" TRAINER_NPROC="${TRAINER_NPROC:-4}" TARGET_MODEL_PATH="${TARGET_MODEL_PATH:-Qwen/Qwen3-8B}" +# Whitespace-separated SGLang CLI tokens for model-specific server settings. +# Each flag and value must be one shell token; embedded whitespace is +# unsupported. +SERVER_EXTRA_ARGS="${SERVER_EXTRA_ARGS:-}" +APPLY_SGLANG_CAPTURE_PATCH="${APPLY_SGLANG_CAPTURE_PATCH:-1}" +SERVER_EXTRA_ARGV=() +if [[ -n "$SERVER_EXTRA_ARGS" ]]; then + read -r -a SERVER_EXTRA_ARGV <<< "$SERVER_EXTRA_ARGS" +fi MOONCAKE_RPC_PORT="${MOONCAKE_RPC_PORT:-35551}" MOONCAKE_HTTP_PORT="${MOONCAKE_HTTP_PORT:-35880}" @@ -36,7 +46,7 @@ START_TIMEOUT_S="${START_TIMEOUT_S:-1800}" PEER_TIMEOUT_S="${PEER_TIMEOUT_S:-1800}" log() { - printf '[qwen3-8b-dflash-2node][rank=%s] %s\n' "${NODE_RANK:-?}" "$*" + printf '[%s][rank=%s] %s\n' "$RUN_LABEL" "${NODE_RANK:-?}" "$*" } fail() { @@ -120,6 +130,9 @@ validate_identity() { [[ "$SERVER_TP" =~ ^[1-9][0-9]*$ ]] || fail "SERVER_TP must be positive" [[ "$TRAINER_NPROC" =~ ^[1-9][0-9]*$ ]] || \ fail "TRAINER_NPROC must be positive" + [[ "$APPLY_SGLANG_CAPTURE_PATCH" == "0" || \ + "$APPLY_SGLANG_CAPTURE_PATCH" == "1" ]] || \ + fail "APPLY_SGLANG_CAPTURE_PATCH must be 0 or 1" [[ "$(count_devices "$SERVER_GPUS")" == "$SERVER_TP" ]] || \ fail "SERVER_GPUS must contain exactly SERVER_TP=$SERVER_TP devices" [[ "$(count_devices "$TRAINER_GPUS")" == "$TRAINER_NPROC" ]] || \ @@ -161,17 +174,28 @@ run_inference_node() { local producer_result=1 if [[ "${DRY_RUN:-0}" == "1" ]]; then + local -a dry_run_server_command=( + python -m sglang.launch_server --host 0.0.0.0 + --model-path "$TARGET_MODEL_PATH" + --trust-remote-code + --skip-tokenizer-init + --tp-size "$SERVER_TP" + --mem-fraction-static "$SERVER_MEM_FRACTION" + --chunked-prefill-size -1 + --enable-spec-capture --spec-capture-method dflash + --spec-capture-aux-layer-ids $CAPTURE_LAYER_IDS + --port "$SERVER_PORT" + ) + if [[ -n "$SERVER_EXTRA_ARGS" ]]; then + dry_run_server_command+=("${SERVER_EXTRA_ARGV[@]}") + fi print_command mooncake_master --enable_http_metadata_server=true \ --http_metadata_server_host=0.0.0.0 \ --rpc_port="$MOONCAKE_RPC_PORT" \ --http_metadata_server_port="$MOONCAKE_HTTP_PORT" \ --metrics_port="$MOONCAKE_METRICS_PORT" print_command env "CUDA_VISIBLE_DEVICES=$SERVER_GPUS" \ - python -m sglang.launch_server --host 0.0.0.0 \ - --model-path "$TARGET_MODEL_PATH" --tp-size "$SERVER_TP" \ - --enable-spec-capture --spec-capture-method dflash \ - --spec-capture-aux-layer-ids $CAPTURE_LAYER_IDS \ - --port "$SERVER_PORT" + "${dry_run_server_command[@]}" print_command env CUDA_VISIBLE_DEVICES= specforge train -c "$CONFIG" \ --role producer "${COMMON_OVERRIDES[@]}" "$@" result=0 @@ -195,7 +219,9 @@ run_inference_node() { command -v mooncake_master >/dev/null || fail "mooncake_master is not on PATH" command -v curl >/dev/null || fail "curl is not on PATH" - "$ROOT_DIR/scripts/apply_sglang_spec_capture_patch.sh" + if [[ "$APPLY_SGLANG_CAPTURE_PATCH" == "1" ]]; then + "$ROOT_DIR/scripts/apply_sglang_spec_capture_patch.sh" + fi export MOONCAKE_LOCAL_HOSTNAME="${INFERENCE_NODE_IP:-$HEAD_IP}" export MOONCAKE_GLOBAL_SEGMENT_SIZE="${MOONCAKE_GLOBAL_SEGMENT_SIZE:-$((32 << 30))}" export MOONCAKE_LOCAL_BUFFER_SIZE="${MOONCAKE_LOCAL_BUFFER_SIZE:-$((1 << 30))}" @@ -227,19 +253,25 @@ run_inference_node() { done read -r -a capture_layers <<< "$CAPTURE_LAYER_IDS" + local -a server_command=( + python -m sglang.launch_server + --host 0.0.0.0 + --model-path "$TARGET_MODEL_PATH" + --trust-remote-code + --skip-tokenizer-init + --tp-size "$SERVER_TP" + --mem-fraction-static "$SERVER_MEM_FRACTION" + --chunked-prefill-size -1 + --enable-spec-capture + --spec-capture-method dflash + --spec-capture-aux-layer-ids "${capture_layers[@]}" + --port "$SERVER_PORT" + ) + if [[ -n "$SERVER_EXTRA_ARGS" ]]; then + server_command+=("${SERVER_EXTRA_ARGV[@]}") + fi setsid env CUDA_VISIBLE_DEVICES="$SERVER_GPUS" \ - python -m sglang.launch_server \ - --host 0.0.0.0 \ - --model-path "$TARGET_MODEL_PATH" \ - --trust-remote-code \ - --skip-tokenizer-init \ - --tp-size "$SERVER_TP" \ - --mem-fraction-static "$SERVER_MEM_FRACTION" \ - --chunked-prefill-size -1 \ - --enable-spec-capture \ - --spec-capture-method dflash \ - --spec-capture-aux-layer-ids "${capture_layers[@]}" \ - --port "$SERVER_PORT" \ + "${server_command[@]}" \ > "$RUN_ROOT/sglang-server.log" 2>&1 & server_pid="$!" diff --git a/examples/disagg/sync_distributed_checkpoints.py b/examples/disagg/sync_distributed_checkpoints.py new file mode 100755 index 000000000..88fbc0373 --- /dev/null +++ b/examples/disagg/sync_distributed_checkpoints.py @@ -0,0 +1,361 @@ +#!/usr/bin/env python3 +"""Mirror rank-local SpecForge checkpoints between two trainer nodes. + +SpecForge writes ``training_state.pt`` on rank 0 and one +``training_state_rankN.pt`` file per rank. On platforms without a shared output +filesystem, each trainer node therefore owns only part of a multi-node +checkpoint. Run one relay on each node to package the files written locally, +serve them over the private trainer network, fetch the peer package, and +atomically assemble a complete checkpoint directory on both nodes. + +The HTTP server intentionally has no authentication or TLS. Bind it only to a +trusted private network interface and do not expose its port publicly. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import re +import signal +import tarfile +import threading +from functools import partial +from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from urllib.error import HTTPError, URLError +from urllib.request import urlopen + +CHUNK_BYTES = 8 * 1024 * 1024 +STATE_FILE = "training_state.pt" + + +class _PrivateHTTPServer(ThreadingHTTPServer): + request_queue_size = 64 + daemon_threads = True + + +class _QuietHandler(SimpleHTTPRequestHandler): + def log_message(self, _format, *_args) -> None: + return + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(CHUNK_BYTES): + digest.update(chunk) + return digest.hexdigest() + + +def _atomic_json(path: Path, payload: object) -> None: + tmp = path.with_name(path.name + ".tmp") + with tmp.open("w", encoding="utf-8") as stream: + json.dump(payload, stream, sort_keys=True) + stream.flush() + os.fsync(stream.fileno()) + os.replace(tmp, path) + + +def _atomic_text(path: Path, value: str) -> None: + tmp = path.with_name(path.name + ".tmp") + with tmp.open("w", encoding="utf-8") as stream: + stream.write(value) + stream.flush() + os.fsync(stream.fileno()) + os.replace(tmp, path) + + +def _rank_range(value: str) -> tuple[int, ...]: + match = re.fullmatch(r"(\d+)-(\d+)", value) + if match is None: + raise argparse.ArgumentTypeError("rank range must look like 0-7") + first, last = (int(item) for item in match.groups()) + if first < 0 or last < first: + raise argparse.ArgumentTypeError("rank range must be increasing") + return tuple(range(first, last + 1)) + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed < 1: + raise argparse.ArgumentTypeError("value must be at least 1") + return parsed + + +class CheckpointRelay: + def __init__(self, args: argparse.Namespace) -> None: + self.run_root = Path(args.run_root).resolve() + self.output_dir = self.run_root / "output" + self.relay_dir = self.run_root / "checkpoint-relay" + self.run_id = args.run_id + self.local_ranks = tuple(args.local_ranks) + self.peer_ranks = tuple(args.peer_ranks) + self.peer_url = args.peer_url.rstrip("/") + self.poll_s = float(args.poll_s) + self.max_archives = int(args.max_archives) + self.relay_dir.mkdir(parents=True, exist_ok=True) + self._stop = threading.Event() + handler = partial(_QuietHandler, directory=str(self.relay_dir)) + self._httpd = _PrivateHTTPServer( + (args.serve_host, int(args.serve_port)), handler + ) + self._server_thread = threading.Thread( + target=self._httpd.serve_forever, + name="checkpoint-relay-http", + daemon=True, + ) + self._step_pattern = re.compile(rf"^{re.escape(self.run_id)}-step(\d+)$") + self._local_archive_pattern = self._archive_pattern(self.local_ranks) + self._peer_archive_pattern = self._archive_pattern(self.peer_ranks) + + def _archive_pattern(self, ranks: tuple[int, ...]) -> re.Pattern[str]: + return re.compile( + rf"^{re.escape(self.run_id)}-step(\d+)-" + rf"ranks{ranks[0]}-{ranks[-1]}\.tar$" + ) + + def _local_names(self) -> tuple[str, ...]: + names = [f"training_state_rank{rank}.pt" for rank in self.local_ranks] + if 0 in self.local_ranks: + names.insert(0, STATE_FILE) + return tuple(names) + + def _peer_names(self) -> tuple[str, ...]: + names = [f"training_state_rank{rank}.pt" for rank in self.peer_ranks] + if 0 in self.peer_ranks: + names.insert(0, STATE_FILE) + return tuple(names) + + def _checkpoint_dirs(self) -> list[tuple[int, Path]]: + found = [] + try: + children = list(self.output_dir.iterdir()) + except FileNotFoundError: + return [] + for child in children: + match = self._step_pattern.fullmatch(child.name) + if match is not None and child.is_dir(): + found.append((int(match.group(1)), child)) + return sorted(found) + + def _archive_name(self, step: int) -> str: + return ( + f"{self.run_id}-step{step}-" + f"ranks{self.local_ranks[0]}-{self.local_ranks[-1]}.tar" + ) + + def _local_archives(self) -> list[tuple[int, Path]]: + found = [] + for path in self.relay_dir.iterdir(): + match = self._local_archive_pattern.fullmatch(path.name) + if match is not None and path.is_file(): + found.append((int(match.group(1)), path)) + return sorted(found) + + def _prune_local_archives(self, keep: set[str]) -> None: + for path in self.relay_dir.iterdir(): + if self._local_archive_pattern.fullmatch(path.name) is None: + continue + if path.name in keep: + continue + path.unlink(missing_ok=True) + path.with_name(path.name + ".sha256").unlink(missing_ok=True) + + def _prune_peer_archives(self, keep: set[str]) -> None: + for path in self.relay_dir.iterdir(): + if not path.name.startswith("peer-"): + continue + archive_name = path.name.removeprefix("peer-") + if archive_name.endswith(".partial"): + archive_name = archive_name.removesuffix(".partial") + if self._peer_archive_pattern.fullmatch(archive_name) is None: + continue + if archive_name not in keep: + path.unlink(missing_ok=True) + + def _publish_local(self) -> None: + local_names = self._local_names() + checkpoint_dirs = self._checkpoint_dirs()[-self.max_archives :] + for step, checkpoint_dir in checkpoint_dirs: + sources = [checkpoint_dir / name for name in local_names] + if not all(path.is_file() and path.stat().st_size > 0 for path in sources): + continue + archive = self.relay_dir / self._archive_name(step) + sha_path = archive.with_name(archive.name + ".sha256") + if not archive.is_file(): + tmp = archive.with_name(archive.name + ".tmp") + with tarfile.open(tmp, mode="w") as bundle: + for source in sources: + bundle.add(source, arcname=source.name, recursive=False) + os.replace(tmp, archive) + _atomic_text(sha_path, _sha256(archive)) + print( + f"PACKED step={step} bytes={archive.stat().st_size} " + f"archive={archive.name}", + flush=True, + ) + archives = self._local_archives()[-self.max_archives :] + self._prune_local_archives({archive.name for _, archive in archives}) + entries = [] + for step, archive in archives: + sha_path = archive.with_name(archive.name + ".sha256") + try: + archive_sha = sha_path.read_text(encoding="utf-8").strip() + except FileNotFoundError: + archive_sha = _sha256(archive) + _atomic_text(sha_path, archive_sha) + entries.append( + { + "step": step, + "archive": archive.name, + "sha256": archive_sha, + "files": list(local_names), + } + ) + _atomic_json( + self.relay_dir / "manifest.json", + {"run_id": self.run_id, "entries": entries}, + ) + + def _peer_manifest(self) -> dict | None: + try: + with urlopen(f"{self.peer_url}/manifest.json", timeout=5.0) as response: + payload = json.load(response) + except (HTTPError, URLError, OSError, TimeoutError, ValueError): + return None + if payload.get("run_id") != self.run_id: + return None + return payload + + def _download(self, name: str, expected_sha: str) -> Path: + match = self._peer_archive_pattern.fullmatch(name) + if match is None or Path(name).name != name: + raise ValueError(f"unexpected peer archive name {name!r}") + destination = self.relay_dir / f"peer-{name}" + if destination.is_file() and _sha256(destination) == expected_sha: + return destination + partial_path = destination.with_name(destination.name + ".partial") + digest = hashlib.sha256() + with ( + urlopen(f"{self.peer_url}/{name}", timeout=300.0) as response, + partial_path.open("wb") as stream, + ): + while chunk := response.read(CHUNK_BYTES): + stream.write(chunk) + digest.update(chunk) + if digest.hexdigest() != expected_sha: + partial_path.unlink(missing_ok=True) + raise ValueError(f"SHA-256 mismatch for peer archive {name}") + os.replace(partial_path, destination) + return destination + + def _install_peer_archive(self, entry: dict) -> None: + step = int(entry["step"]) + archive_name = str(entry["archive"]) + archive_match = self._peer_archive_pattern.fullmatch(archive_name) + if archive_match is None or int(archive_match.group(1)) != step: + raise ValueError( + f"unexpected peer archive for step {step}: {archive_name!r}" + ) + expected_names = set(self._peer_names()) + if set(entry.get("files", ())) != expected_names: + raise ValueError(f"unexpected peer file set at step {step}") + checkpoint_dir = self.output_dir / f"{self.run_id}-step{step}" + if not checkpoint_dir.is_dir(): + return + marker = checkpoint_dir / ( + f".checkpoint-relay-ranks{self.peer_ranks[0]}-" + f"{self.peer_ranks[-1]}.json" + ) + expected_sha = str(entry["sha256"]) + try: + with marker.open(encoding="utf-8") as stream: + if json.load(stream).get("sha256") == expected_sha: + return + except (FileNotFoundError, OSError, ValueError): + pass + archive = self._download(archive_name, expected_sha) + with tarfile.open(archive, mode="r") as bundle: + members = bundle.getmembers() + names = {member.name for member in members} + if names != expected_names or any( + not member.isfile() for member in members + ): + raise ValueError(f"unsafe or incomplete peer archive {archive.name}") + for member in members: + source = bundle.extractfile(member) + if source is None: + raise ValueError(f"cannot read {member.name} from {archive.name}") + destination = checkpoint_dir / member.name + tmp = destination.with_name(destination.name + ".relay-tmp") + with source, tmp.open("wb") as stream: + while chunk := source.read(CHUNK_BYTES): + stream.write(chunk) + os.replace(tmp, destination) + _atomic_json( + marker, + {"sha256": expected_sha, "archive": entry["archive"], "step": step}, + ) + print( + f"INSTALLED step={step} peer_files={len(expected_names)} " + f"sha256={expected_sha}", + flush=True, + ) + + def _pull_peer(self) -> None: + manifest = self._peer_manifest() + if manifest is None: + return + entries = sorted( + manifest.get("entries", ()), key=lambda item: int(item["step"]) + )[-self.max_archives :] + self._prune_peer_archives({str(entry["archive"]) for entry in entries}) + for entry in entries: + self._install_peer_archive(entry) + + def stop(self, *_args) -> None: + self._stop.set() + + def run(self) -> None: + signal.signal(signal.SIGTERM, self.stop) + signal.signal(signal.SIGINT, self.stop) + self._server_thread.start() + print( + f"STARTED run={self.run_id} local_ranks={self.local_ranks[0]}-" + f"{self.local_ranks[-1]} peer={self.peer_url} " + f"max_archives={self.max_archives}", + flush=True, + ) + try: + while not self._stop.is_set(): + try: + self._publish_local() + self._pull_peer() + except Exception as exc: # noqa: BLE001 - retry loop boundary + print(f"RETRY {type(exc).__name__}: {exc}", flush=True) + self._stop.wait(self.poll_s) + finally: + self._httpd.shutdown() + self._httpd.server_close() + self._server_thread.join(timeout=5.0) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--run-root", required=True) + parser.add_argument("--run-id", required=True) + parser.add_argument("--local-ranks", required=True, type=_rank_range) + parser.add_argument("--peer-ranks", required=True, type=_rank_range) + parser.add_argument("--serve-host", required=True) + parser.add_argument("--serve-port", required=True, type=int) + parser.add_argument("--peer-url", required=True) + parser.add_argument("--poll-s", type=float, default=15.0) + parser.add_argument("--max-archives", type=_positive_int, default=3) + return parser.parse_args() + + +if __name__ == "__main__": + CheckpointRelay(_parse_args()).run() diff --git a/patches/sglang/kimi-k3-f8493a4/spec-capture.patch b/patches/sglang/kimi-k3-f8493a4/spec-capture.patch new file mode 100644 index 000000000..dcb3c4870 --- /dev/null +++ b/patches/sglang/kimi-k3-f8493a4/spec-capture.patch @@ -0,0 +1,1153 @@ +diff --git a/python/sglang/srt/layers/attn_residual.py b/python/sglang/srt/layers/attn_residual.py +index 266d2f2cd..2a75af790 100644 +--- a/python/sglang/srt/layers/attn_residual.py ++++ b/python/sglang/srt/layers/attn_residual.py +@@ -115,10 +115,28 @@ def _score_kernel( + BLOCK_H: tl.constexpr, + ): + """One CTA per (token, row): scan H, output one scalar score.""" +- pid_t = tl.program_id(0) ++ # bank stride is 8 * 7168; 64K token offsets exceed signed int32. ++ pid_t = tl.program_id(0).to(tl.int64) + j = tl.program_id(1) + if j > NVB: + return ++ ++ # The raw K3 residual stream can remain finite near the bf16 limit. A ++ # direct fp32 sum(v * v) then overflows even though RMSNorm(v) is well ++ # defined, making this DSpark-only capture path emit NaN scores. Scale ++ # each row before the norm/projection reduction to keep intermediates ++ # representable without changing the normalized score. ++ max_abs = 0.0 ++ for h0 in tl.static_range(0, H, BLOCK_H): ++ offs_h = h0 + tl.arange(0, BLOCK_H) ++ if j < NVB: ++ v = tl.load(bank_ptr + pid_t * stride_bm + j * stride_bb + offs_h).to( ++ tl.float32 ++ ) ++ else: ++ v = tl.load(prefix_ptr + pid_t * stride_pm + offs_h).to(tl.float32) ++ max_abs = tl.maximum(max_abs, tl.max(tl.abs(v), axis=0)) ++ scale = tl.maximum(max_abs, 1.0) + sumsq = 0.0 + dotv = 0.0 + for h0 in tl.static_range(0, H, BLOCK_H): +@@ -130,9 +148,10 @@ def _score_kernel( + else: + v = tl.load(prefix_ptr + pid_t * stride_pm + offs_h).to(tl.float32) + cw = tl.load(cw_ptr + offs_h) +- sumsq += tl.sum(v * v) +- dotv += tl.sum(v * cw) +- rrms = 1.0 / tl.sqrt(sumsq / H + eps) ++ v_scaled = v / scale ++ sumsq += tl.sum(v_scaled * v_scaled) ++ dotv += tl.sum(v_scaled * cw) ++ rrms = 1.0 / tl.sqrt(sumsq / H + eps / scale / scale) + tl.store(scores_ptr + pid_t * stride_sm + j, dotv * rrms) + + +@@ -159,7 +178,8 @@ def _combine_kernel( + Softmax is redundantly computed by each H-chunk CTA (≤16 elements, trivial). + This gives full H-parallelism: 7 CTAs for H=7168/1024. + """ +- pid_t = tl.program_id(0) ++ # Keep bank/scores/output pointer arithmetic in the same 64-bit domain. ++ pid_t = tl.program_id(0).to(tl.int64) + pid_h = tl.program_id(1) + h0 = pid_h * BLOCK_H + +diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py +index 480ed81a5..f84d489a1 100644 +--- a/python/sglang/srt/layers/logits_processor.py ++++ b/python/sglang/srt/layers/logits_processor.py +@@ -163,6 +163,9 @@ class LogitsProcessorOutput: + # Used by speculative decoding (EAGLE) + # The last hidden layers + hidden_states: Optional[torch.Tensor] = None ++ # Spec-training capture: under FULL+aux, `hidden_states` is the aux ++ # concatenation, so the post-norm last hidden is exposed here separately. ++ last_hidden_states: Optional[torch.Tensor] = None + + ## Part 2: This part will be assigned in python/sglang/srt/layers/sampler.py::Sampler + # he log probs of output tokens, if SGLANG_RETURN_ORIGINAL_LOGPROB = True, will get the log probs before applying temperature. If False, will get the log probs before applying temperature. +@@ -447,6 +450,16 @@ class LogitsProcessor(nn.Module): + sample_indices, + logits_metadata, + ) ++ # Spec-training capture: keep the post-norm last hidden (the training ++ # target) alongside the aux concatenation in hidden_states_to_store. ++ last_hidden_states_to_store = ( ++ hidden_states ++ if ( ++ logits_metadata.capture_hidden_mode.is_full() ++ and aux_hidden_states is not None ++ ) ++ else None ++ ) + del hidden_states + + if not logits_metadata.extend_return_logprob: +@@ -460,6 +473,7 @@ class LogitsProcessor(nn.Module): + return LogitsProcessorOutput( + next_token_logits=sampled_logits, + hidden_states=hidden_states_to_store, ++ last_hidden_states=last_hidden_states_to_store, + # FIXME: These fields are not logits-related but are passed through here as a + # workaround since ForwardBatch is local to forward_batch_generation(). + # They should be moved to GenerationBatchResult to keep this class clean. +diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py +index 1bf6aa7d2..e7a303074 100644 +--- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py ++++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py +@@ -402,6 +402,9 @@ def fused_marlin_moe( + and intermediate_cache3.is_contiguous() + and output.is_contiguous() + and intermediate_cache3.shape[-1] % 8 == 0 ++ # The JIT helper maps M to CUDA grid.y, whose architectural limit ++ # is 65,535. Long-context prefills use the generic reduction. ++ and intermediate_cache3.shape[0] <= 65_535 + ): + from sglang.kernels.ops.moe.moe_topk_sum import moe_topk_sum + +diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py +index 19a7f4938..b3c5220f3 100644 +--- a/python/sglang/srt/managers/detokenizer_manager.py ++++ b/python/sglang/srt/managers/detokenizer_manager.py +@@ -475,6 +475,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin): + output_token_sampling_mask=recv_obj.output_token_sampling_mask, + output_token_sampling_logprobs=recv_obj.output_token_sampling_logprobs, + output_hidden_states=recv_obj.output_hidden_states, ++ spec_capture=recv_obj.spec_capture, + routed_experts=routed_experts, + indexer_topk=indexer_topk, + customized_info=recv_obj.customized_info, +diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py +index 06282f331..4675de4ef 100644 +--- a/python/sglang/srt/managers/io_struct.py ++++ b/python/sglang/srt/managers/io_struct.py +@@ -309,6 +309,10 @@ class GenerateReqInput: + # For Unlimited-OCR + images_config: Optional[dict] = None + ++ # Spec-training capture sink instructions (see spec_capture_sink.py). ++ # Batch-level: List[Optional[dict]]; per-request after __getitem__. ++ spec_capture: Optional[Union[List[Optional[Dict]], Dict]] = None ++ + # Pre-computed delimiter indices for multi-item scoring. + # Batch-level: List[List[int]] (one per request). After __getitem__: List[int]. + multi_item_delimiter_indices: Optional[Union[List[List[int]], List[int]]] = None +@@ -804,6 +808,11 @@ class GenerateReqInput: + if self.multi_item_delimiter_indices is not None + else None + ), ++ spec_capture=( ++ self.spec_capture[i] ++ if isinstance(self.spec_capture, list) ++ else self.spec_capture ++ ), + ) + cache[i] = sub + return sub +@@ -902,6 +911,9 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True): + # Internal IPC only. + encoder_urls: Optional[List[str]] = None + ++ # Spec-training capture sink instructions (see GenerateReqInput.spec_capture) ++ spec_capture: Optional[Dict] = None ++ + # Pre-computed delimiter indices for multi-item scoring + multi_item_delimiter_indices: Optional[List[int]] = None + +@@ -1328,6 +1340,9 @@ class BatchTokenIDOutput(BaseBatchReq, kw_only=True): + # Number of times each request was retracted. + retraction_counts: Optional[List[int]] = None + ++ # Spec-training capture: one result dict per request (see spec_capture_sink). ++ spec_capture: Optional[List[Any]] = None ++ + # The trainer step id. Used to know which step's weights are used for sampling. + token_steps: Optional[List[List[int]]] = None + +@@ -1419,6 +1434,9 @@ class BatchStrOutput(BaseBatchReq, kw_only=True): + # Number of times each request was retracted. + retraction_counts: Optional[List[int]] = None + ++ # Spec-training capture: one result dict per request (see spec_capture_sink). ++ spec_capture: Optional[List[Any]] = None ++ + # The trainer step id. Used to know which step's weights are used for sampling. + token_steps: Optional[List[List[int]]] = None + +diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py +index 0ef94dd95..d3cada1fa 100755 +--- a/python/sglang/srt/managers/schedule_batch.py ++++ b/python/sglang/srt/managers/schedule_batch.py +@@ -768,6 +768,7 @@ class Req(ReqDllmMixin): + Union[APIServerReqTimeStats, DPControllerReqTimeStats] + ] = None, + return_pooled_hidden_states: bool = False, ++ spec_capture: Optional[Dict[str, Any]] = None, + multi_item_delimiter_indices: Optional[List[int]] = None, + session_id: Optional[str] = None, + ): +@@ -832,7 +833,13 @@ class Req(ReqDllmMixin): + } + self.sampling_params = sampling_params + self.custom_logit_processor = custom_logit_processor +- self.return_hidden_states = return_hidden_states ++ # Spec-training capture: piggyback the return_hidden_states path (fires ++ # CaptureHiddenMode.FULL); the sink consumes the slices, not the response. ++ self.spec_capture = spec_capture ++ self.return_hidden_states = return_hidden_states or spec_capture is not None ++ self.spec_capture_aux: List[torch.Tensor] = [] ++ self.spec_capture_last_hidden: List[torch.Tensor] = [] ++ self.spec_capture_result = None # per-request sink result -> output field + + # extra key for classifying the request (e.g. cache_salt) + if lora_id is not None: +diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py +index 6589c6a78..d31be1b92 100644 +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -637,6 +637,19 @@ class Scheduler( + + self.init_batch_result_processor() + ++ if server_args.enable_spec_capture: ++ # Capture needs single-pass prefill; chunking would drop all but the ++ # final chunk's hidden rows. ++ if server_args.chunked_prefill_size != -1: ++ raise ValueError( ++ "--enable-spec-capture requires --chunked-prefill-size -1 " ++ "(single-pass prefill) so captured hidden states cover the " ++ "whole sequence" ++ ) ++ from sglang.srt import spec_capture_sink ++ ++ spec_capture_sink.maybe_init_sink(server_args) ++ + self.is_initializing = False + + def init_zbal_on_npu(self): +@@ -2283,6 +2296,7 @@ class Scheduler( + dllm_config=self.dllm_config, + time_stats=recv_req.time_stats, + multi_item_delimiter_indices=recv_req.multi_item_delimiter_indices, ++ spec_capture=recv_req.spec_capture, + ) + req.tokenizer = self.tokenizer + +@@ -3385,6 +3399,13 @@ class Scheduler( + # GenerationBatchResult.extra_keep_alive_refs after forward returns. + self.batch_record_buf[self.batch_record_ct] = [batch, attr_snapshot] + ++ def _should_copy_hidden_states_to_cpu(self, batch: ScheduleBatch) -> bool: ++ """Avoid redundant multi-GiB capture D2H on non-writer TP ranks.""" ++ return batch.return_hidden_states and ( ++ self.ps.attn_tp_rank == 0 ++ or any(req.spec_capture is None for req in batch.reqs) ++ ) ++ + @contextmanager + def _forward_isolation(self, batch: ScheduleBatch, *, overlap: bool): + """Make SB transactional across one forward (overlap and non-overlap). +@@ -3531,7 +3552,9 @@ class Scheduler( + # overlaps. + batch_result.copy_to_cpu( + return_logprob=batch.return_logprob, +- return_hidden_states=batch.return_hidden_states, ++ return_hidden_states=self._should_copy_hidden_states_to_cpu( ++ batch ++ ), + ) + else: + # Result D2H on copy_stream overlaps the next forward +@@ -3541,7 +3564,9 @@ class Scheduler( + with self.copy_stream_ctx: + batch_result.copy_to_cpu( + return_logprob=batch.return_logprob, +- return_hidden_states=batch.return_hidden_states, ++ return_hidden_states=self._should_copy_hidden_states_to_cpu( ++ batch ++ ), + ) + else: + batch_result.future_indices = future_indices +@@ -3580,7 +3605,9 @@ class Scheduler( + batch_result.copy_done = self.device_module.Event() + batch_result.copy_to_cpu( + return_logprob=batch.return_logprob, +- return_hidden_states=batch.return_hidden_states, ++ return_hidden_states=self._should_copy_hidden_states_to_cpu( ++ batch ++ ), + ) + else: + kwargs = ( +@@ -3699,7 +3726,9 @@ class Scheduler( + self._relay_forward_payload(batch_result.future_indices, batch_result) + batch_result.copy_to_cpu( + return_logprob=cur_batch.return_logprob, +- return_hidden_states=cur_batch.return_hidden_states, ++ return_hidden_states=self._should_copy_hidden_states_to_cpu( ++ cur_batch ++ ), + ) + + # Release the closure and large GPU tensors that are no longer needed. +@@ -3718,6 +3747,10 @@ class Scheduler( + batch: ScheduleBatch, + result: Union[GenerationBatchResult, EmbeddingBatchResult], + ): ++ # Complete capture transfers on the scheduler thread before processing ++ # any kind of next result (prefill, decode, or idle). This preserves ++ # the output socket's single-thread ownership under mixed traffic. ++ self.batch_result_processor.drain_spec_captures() + self.publish_load_snapshot(force=batch.forward_mode.is_extend()) + + if batch.forward_mode.is_decode(): +@@ -3826,6 +3859,15 @@ class Scheduler( + + def on_idle(self): + """Idle housekeeping: guard, check, metrics, reset, sleep.""" ++ # A capture response is intentionally delayed until its background ++ # Mooncake batch write completes. Poll it on the scheduler thread so ++ # the ZeroMQ output socket remains single-threaded. Do not enter the ++ # idle sleeper while a future is outstanding or the producer waiting ++ # for that response could deadlock. ++ self.batch_result_processor.drain_spec_captures() ++ if self.batch_result_processor.has_pending_spec_captures(): ++ time.sleep(0.001) ++ return + if not self.is_fully_idle(): + return + +@@ -3879,6 +3921,7 @@ class Scheduler( + and not self.dllm_manager.any_staging_reqs() + and (self.last_batch is None or self.last_batch.is_empty()) + and (not self.enable_overlap or len(self.result_queue) == 0) ++ and not self.batch_result_processor.has_pending_spec_captures() + and self._pp_microbatches_drained() + ) + +diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +index f75eab004..5af4c1d6b 100644 +--- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py ++++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +@@ -1,7 +1,9 @@ + from __future__ import annotations + + import logging +-from dataclasses import dataclass ++import os ++import time ++from dataclasses import dataclass, field + from typing import ( + TYPE_CHECKING, + Callable, +@@ -87,6 +89,12 @@ class SchedulerBatchResultProcessor: + logprob_result_processor: SchedulerLogprobResultProcessor + output_streamer: SchedulerOutputStreamer + abort_request: Callable ++ _spec_capture_batches: List = field( ++ default_factory=list, ++ init=False, ++ repr=False, ++ compare=False, ++ ) + + def process_batch_result_prebuilt(self, batch: ScheduleBatch): + assert self.disaggregation_mode == DisaggregationMode.DECODE +@@ -190,6 +198,7 @@ class SchedulerBatchResultProcessor: + result: Union[GenerationBatchResult, EmbeddingBatchResult], + ): + skip_stream_req = None ++ pending_spec_captures = [] + + if self.is_generation: + if result.copy_done is not None: +@@ -269,6 +278,20 @@ class SchedulerBatchResultProcessor: + self.add_sampling_mask_return_values(i, req, logits_output) + + if ( ++ req.spec_capture is not None ++ and logits_output.hidden_states is not None ++ ): ++ # Spec-training capture: keep tensor slices, sink on finish. ++ hidden_state_offset = self._append_spec_capture_states( ++ req=req, ++ logits_output=logits_output, ++ hidden_state_offset=hidden_state_offset, ++ ) ++ if req.finished(): ++ pending = self._sink_spec_capture(req) ++ if pending is not None: ++ pending_spec_captures.append(pending) ++ elif ( + req.return_hidden_states + and logits_output.hidden_states is not None + ): +@@ -343,9 +366,25 @@ class SchedulerBatchResultProcessor: + req.inflight_middle_chunks -= 1 + req.time_stats.set_last_chunked_prefill_finish_time() + +- self.output_streamer.stream_output( +- batch.reqs, batch.return_logprob, skip_stream_req +- ) ++ if pending_spec_captures: ++ self._queue_spec_captures( ++ pending_spec_captures, ++ return_logprob=batch.return_logprob, ++ ) ++ capture_req_ids = {id(item[0]) for item in pending_spec_captures} ++ ready_reqs = [ ++ req ++ for req in batch.reqs ++ if id(req) not in capture_req_ids and req is not skip_stream_req ++ ] ++ if ready_reqs: ++ self.output_streamer.stream_output( ++ ready_reqs, batch.return_logprob ++ ) ++ else: ++ self.output_streamer.stream_output( ++ batch.reqs, batch.return_logprob, skip_stream_req ++ ) + + can_run_cuda_graph = result.can_run_cuda_graph + self.metrics_reporter.report_prefill_stats( +@@ -476,6 +515,181 @@ class SchedulerBatchResultProcessor: + f"Placeholder zeros would be appended to output_ids." + ) + ++ def _append_spec_capture_states( ++ self, ++ *, ++ req: Req, ++ logits_output: LogitsProcessorOutput, ++ hidden_state_offset: int, ++ ) -> int: ++ """Accumulate captured rows as CPU tensors for the Mooncake sink. ++ ++ Same offset arithmetic as ``_append_prefill_hidden_states`` but keeps ++ tensor slices (aux concat in ``hidden_states``, post-norm last in ++ ``last_hidden_states``) rather than the JSON-able response payload. ++ """ ++ start = hidden_state_offset ++ end = start + len(req.origin_input_ids) ++ # Only attention-TP rank 0 owns the Mooncake sink. Copying these very ++ # large tensors to host on every TP rank creates eight identical D2H ++ # transfers and seven immediately-discarded CPU copies on TP8. ++ if self.output_streamer.ps.attn_tp_rank != 0: ++ return end ++ # Materialize each scheduler result on host once, rather than issuing ++ # one synchronous D2H transfer per request. In overlap mode the writer ++ # rank has already copied both tensors asynchronously on copy_stream; ++ # ``.cpu()`` is then a no-op. The fallback keeps non-overlap correct. ++ features = dict(req.spec_capture.get("features") or {}) ++ aux_cpu = getattr(logits_output, "_spec_capture_aux_cpu", None) ++ if "aux" in features and aux_cpu is None: ++ aux_cpu = logits_output.hidden_states.cpu() ++ logits_output._spec_capture_aux_cpu = aux_cpu ++ last_hidden_cpu = getattr( ++ logits_output, "_spec_capture_last_hidden_cpu", None ++ ) ++ if ( ++ "last_hidden" in features ++ and logits_output.last_hidden_states is not None ++ and last_hidden_cpu is None ++ ): ++ last_hidden_cpu = logits_output.last_hidden_states.cpu() ++ logits_output._spec_capture_last_hidden_cpu = last_hidden_cpu ++ if "aux" in features and aux_cpu is not None: ++ req.spec_capture_aux.append(aux_cpu[start:end]) ++ if "last_hidden" in features and last_hidden_cpu is not None: ++ req.spec_capture_last_hidden.append(last_hidden_cpu[start:end]) ++ return end ++ ++ def _sink_spec_capture(self, req: Req): ++ """Write a finished capture request's tensors to the Mooncake sink. ++ ++ Runs on the attention-TP rank that streams output; the per-request result ++ (or an ``{"error": ...}`` marker) is set on ``req.spec_capture_result``, ++ returned to the client via the dedicated ``spec_capture`` output field ++ (a per-request channel, unlike per-token ``customized_info``). ++ """ ++ from sglang.srt import spec_capture_sink ++ ++ sink = spec_capture_sink.get_sink() ++ if sink is None or self.output_streamer.ps.attn_tp_rank != 0: ++ return None ++ timing_enabled = os.environ.get("SGLANG_SPEC_CAPTURE_TIMING", "0") == "1" ++ cat_start = time.perf_counter() ++ # With chunked prefill disabled (the K3 training configuration), each ++ # request owns exactly one contiguous view into the batch-level pinned ++ # D2H buffer. torch.cat([view]) needlessly copied the whole capture a ++ # second time on CPU -- about 18 GiB per 128 reference samples. Keep ++ # that view zero-copy; concatenate only the genuinely chunked case. ++ aux = self._coalesce_spec_capture_chunks(req.spec_capture_aux) ++ last_hidden = self._coalesce_spec_capture_chunks( ++ req.spec_capture_last_hidden ++ ) ++ cat_ms = (time.perf_counter() - cat_start) * 1000.0 ++ req.spec_capture_aux = [] ++ req.spec_capture_last_hidden = [] ++ return req, req.spec_capture, aux, last_hidden, timing_enabled, cat_ms ++ ++ def _queue_spec_captures(self, pending, *, return_logprob: bool) -> None: ++ """Queue one batch transfer while the scheduler runs the next prefill.""" ++ if not pending: ++ return ++ from sglang.srt import spec_capture_sink ++ ++ sink = spec_capture_sink.get_sink() ++ samples = [ ++ (spec, aux, last_hidden) ++ for _, spec, aux, last_hidden, _, _ in pending ++ ] ++ self._spec_capture_batches.append( ++ ( ++ pending, ++ sink.submit_samples(samples), ++ return_logprob, ++ time.perf_counter(), ++ ) ++ ) ++ # Bound retained D2H buffers. With producer concurrency=2, reaching ++ # this point means target prefill N+1 already overlapped host transfer ++ # N; wait only if N is still finishing before admitting N+2. ++ max_pending = int( ++ os.environ.get("SGLANG_SPEC_CAPTURE_MAX_PENDING_BATCHES", "2") ++ ) ++ if max_pending < 1: ++ raise ValueError("SGLANG_SPEC_CAPTURE_MAX_PENDING_BATCHES must be >= 1") ++ if len(self._spec_capture_batches) >= max_pending: ++ self.drain_spec_captures(block=True, max_batches=1) ++ ++ def has_pending_spec_captures(self) -> bool: ++ return bool(self._spec_capture_batches) ++ ++ def drain_spec_captures( ++ self, *, block: bool = False, max_batches: Optional[int] = None ++ ) -> int: ++ """Finish ready transfers and stream their responses on this thread.""" ++ completed = 0 ++ while self._spec_capture_batches: ++ if max_batches is not None and completed >= max_batches: ++ break ++ pending, future, return_logprob, queued_at = self._spec_capture_batches[0] ++ if not block and not future.done(): ++ break ++ self._spec_capture_batches.pop(0) ++ self._complete_spec_capture_batch( ++ pending, ++ future, ++ return_logprob=return_logprob, ++ queued_at=queued_at, ++ ) ++ completed += 1 ++ return completed ++ ++ def _complete_spec_capture_batch( ++ self, pending, future, *, return_logprob: bool, queued_at: float ++ ) -> None: ++ try: ++ results = future.result() ++ if len(results) != len(pending): ++ raise RuntimeError( ++ f"spec-capture sink returned {len(results)} results for " ++ f"{len(pending)} samples" ++ ) ++ except Exception as e: ++ logger.error( ++ "spec-capture batch sink failed for %d requests: %s", len(pending), e ++ ) ++ for req, spec, _, _, _, _ in pending: ++ req.spec_capture_result = { ++ "sample_id": spec.get("sample_id"), ++ "error": str(e), ++ } ++ else: ++ for pending_item, result in zip(pending, results): ++ req, _, _, _, _, _ = pending_item ++ req.spec_capture_result = result ++ ++ timing_enabled = any(item[4] for item in pending) ++ if timing_enabled: ++ logger.info( ++ "[spec-capture-timing] async_batch_complete samples=%d " ++ "queue_to_stream_ms=%.3f cat_ms=%.3f", ++ len(pending), ++ (time.perf_counter() - queued_at) * 1000.0, ++ sum(item[5] for item in pending), ++ ) ++ self.output_streamer.stream_output( ++ [item[0] for item in pending], return_logprob ++ ) ++ ++ @staticmethod ++ def _coalesce_spec_capture_chunks( ++ chunks: List[torch.Tensor], ++ ) -> Optional[torch.Tensor]: ++ if not chunks: ++ return None ++ if len(chunks) == 1: ++ return chunks[0] ++ return torch.cat(chunks, dim=0) ++ + def _append_prefill_hidden_states( + self, + *, +diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py +index 9f5b2329a..cce5b4afe 100644 +--- a/python/sglang/srt/managers/scheduler_components/output_streamer.py ++++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py +@@ -162,6 +162,14 @@ class SchedulerOutputStreamer: + for req in reqs: + if req is skip_req: + continue ++ if self.server_args.enable_spec_capture and req.spec_capture is not None: ++ # Capture requests are not complete until the background sink ++ # has durably published every feature object. Several ++ # scheduler paths can ask the common streamer to emit a ++ # finished request; centralize the completion barrier here so ++ # none of them can win the race and send a metadata-less 200. ++ if req.finished() and req.spec_capture_result is None: ++ continue + if req.finished() and req.finished_output: + # With the overlap schedule, a request will try to output twice and hit this line twice + # because of the one additional delayed token. This "continue" prevented the dummy output. +@@ -302,6 +310,7 @@ class _GenerationStreamAccumulator: + spec_cap_lens_histogram: list = field(default_factory=list) + retraction_counts: list = field(default_factory=list) + output_hidden_states: Optional[list] = None ++ spec_capture: list = field(default_factory=list) + routed_experts: Optional[list] = None + indexer_topk: Optional[list] = None + customized_info: dict = field(default_factory=dict) +@@ -571,6 +580,8 @@ class _GenerationStreamAccumulator: + self.output_hidden_states.append(hs) + else: + self.output_hidden_states.append(None) ++ # Per-request spec-capture result (aligned with rids), like the field above. ++ self.spec_capture.append(getattr(req, "spec_capture_result", None)) + if self.return_routed_experts: + self.routed_experts.append( + req.routed_experts if req.return_routed_experts else None +@@ -663,6 +674,7 @@ class _GenerationStreamAccumulator: + output_token_sampling_mask=self.output_token_sampling_mask, + output_token_sampling_logprobs=self.output_token_sampling_logprobs, + output_hidden_states=self.output_hidden_states, ++ spec_capture=self.spec_capture or None, + routed_experts=self.routed_experts, + indexer_topk=self.indexer_topk, + customized_info=( +diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py +index 41cc4c246..5d429a91f 100644 +--- a/python/sglang/srt/managers/tokenizer_manager.py ++++ b/python/sglang/srt/managers/tokenizer_manager.py +@@ -1337,6 +1337,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): + multi_item_delimiter_indices=obj.multi_item_delimiter_indices, + mm_data_mooncake=obj.mm_data_mooncake, + encoder_urls=obj.encoder_urls, ++ spec_capture=obj.spec_capture, + ) + elif isinstance(obj, EmbeddingReqInput): + # Resolve unresolved embed overrides now that input_ids are available +@@ -2124,6 +2125,10 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): + hidden_states = recv_obj.output_hidden_states[i] + if hidden_states is not None: + meta_info["hidden_states"] = hidden_states ++ if getattr(recv_obj, "spec_capture", None): ++ sc = recv_obj.spec_capture[i] ++ if sc is not None: ++ meta_info["spec_capture"] = sc + if getattr(recv_obj, "routed_experts", None): + val = recv_obj.routed_experts[i] + if val is not None: +diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py +index fe883c264..2b7206a13 100644 +--- a/python/sglang/srt/managers/utils.py ++++ b/python/sglang/srt/managers/utils.py +@@ -148,6 +148,10 @@ class GenerationBatchResult: + self.logits_output.hidden_states = _async_d2h( + self.logits_output.hidden_states + ) ++ if self.logits_output.last_hidden_states is not None: ++ self.logits_output.last_hidden_states = _async_d2h( ++ self.logits_output.last_hidden_states ++ ) + self.next_token_ids = _async_d2h(self.next_token_ids) + + if self.accept_lens is not None: +diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py +index 661bd5ad1..e1fce6365 100644 +--- a/python/sglang/srt/model_executor/model_runner.py ++++ b/python/sglang/srt/model_executor/model_runner.py +@@ -489,6 +489,26 @@ class ModelRunner: + is_draft_worker=self.is_draft_worker, + ) + ) ++ if self.server_args.enable_spec_capture and not self.is_draft_worker: ++ # Spec-training capture runs without a speculative draft worker, so ++ # populate the same aux-state configuration that online decoding ++ # would normally derive from the draft model. ++ capture_method = self.server_args.spec_capture_method ++ if capture_method in ("dflash", "dspark"): ++ self.spec_aux_config.dflash_use_aux_hidden_state = True ++ self.spec_aux_config.dflash_target_layer_ids = ( ++ self.server_args.spec_capture_aux_layer_ids ++ ) ++ elif capture_method == "eagle3": ++ self.spec_aux_config.eagle_use_aux_hidden_state = True ++ self.spec_aux_config.eagle_aux_hidden_state_layer_ids = ( ++ self.server_args.spec_capture_aux_layer_ids ++ ) ++ else: ++ raise ValueError( ++ "--spec-capture-method must be one of: eagle3, dflash, dspark; " ++ f"got {capture_method!r}" ++ ) + + def init_weight_exporter(self): + self.weight_exporter = WeightExporter( +@@ -866,7 +886,13 @@ class ModelRunner: + eagle_aux_hidden_state_layer_ids=self.spec_aux_config.eagle_aux_hidden_state_layer_ids, + dflash_use_aux_hidden_state=self.spec_aux_config.dflash_use_aux_hidden_state, + dflash_target_layer_ids=self.spec_aux_config.dflash_target_layer_ids, +- is_dspark=self.spec_algorithm.is_dspark(), ++ is_dspark=( ++ self.spec_algorithm.is_dspark() ++ or ( ++ self.server_args.enable_spec_capture ++ and self.server_args.spec_capture_method == "dspark" ++ ) ++ ), + ) + backends = build_attention_backends(model_runner=self) + self.attn_backend = backends.attn_backend +diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py +index e9f84cb17..dd64be9f1 100644 +--- a/python/sglang/srt/server_args.py ++++ b/python/sglang/srt/server_args.py +@@ -3349,6 +3349,26 @@ class ServerArgs: + enable_return_hidden_states: A[ + bool, "Enable returning hidden states with responses.", NS("exec.features") + ] = False ++ enable_spec_capture: A[ ++ bool, ++ "Enable server-side speculative-training capture: per-request aux/last " ++ "hidden states are written to a Mooncake store (SpecForge DataFlow " ++ "layout) instead of the response payload. Enables aux-hidden-state " ++ "capture on the target model without a speculative draft worker.", ++ ] = False ++ spec_capture_aux_layer_ids: A[ ++ Optional[List[int]], ++ "Target layer ids whose hidden states are captured (concatenated) for " ++ "spec-capture requests. Defaults to the model's EAGLE3 default layers " ++ "(low/mid/high) when unset.", ++ ] = None ++ spec_capture_method: A[ ++ str, ++ "Capture method for --enable-spec-capture: 'eagle3', 'dflash', or " ++ "'dspark'. Must " ++ "match the draft strategy being trained; they wire capture onto " ++ "different target-model submodules.", ++ ] = "eagle3" + enable_return_routed_experts: A[ + bool, + "Enable returning routed experts of each layer with responses.", +diff --git a/python/sglang/srt/spec_capture_sink.py b/python/sglang/srt/spec_capture_sink.py +new file mode 100644 +index 000000000..8ab06e25d +--- /dev/null ++++ b/python/sglang/srt/spec_capture_sink.py +@@ -0,0 +1,388 @@ ++# Copyright 2024 SGLang Team ++# Licensed under the Apache License, Version 2.0 (the "License"); ++# you may not use this file except in compliance with the License. ++# You may obtain a copy of the License at ++# ++# http://www.apache.org/licenses/LICENSE-2.0 ++"""Server-side spec-training capture sink (SpecForge DataFlow transport). ++ ++Under ``--enable-spec-capture``, a request's ``spec_capture`` dict tells this ++sink to write the prefill's captured tensors straight into a Mooncake store ++(one hard-pinned object per tensor at ``{store_id}/{sample_id}/g{gen}/{name}``, ++raw bytes — shape/dtype travel on the returned spec). Feature tensors never ++touch the response path; ``meta_info["spec_capture"]`` returns only keys + ++shapes/dtypes. Strategy naming is the client's (the ``features`` mapping); the ++server knows only generic artifacts. Self-contained: the scheduler hooks are ++one-liners, every capture decision lives here. ++ ++Request schema:: ++ ++ {"store_id", "sample_id", "gen", "replace", # key namespace / retry policy ++ "features": {"aux": , "last_hidden": }, # artifact -> feature ++ "passthrough": [{"name", "data", "shape", "dtype"}]} # client tensors verbatim ++ ++Response (``meta_info["spec_capture"]``): ``{"sample_id", "store_id", "gen", ++"aux_layer_ids", "features": {name: {"shape", "dtype"}}}``. ++ ++Mooncake connection uses the standard ``MOONCAKE_*`` env vars (see ++``MooncakeFeatureStore``). ++""" ++ ++from __future__ import annotations ++ ++import logging ++import os ++import threading ++import time ++from concurrent.futures import Future, ThreadPoolExecutor ++from typing import Any, Dict, List, Optional, Tuple ++ ++import torch ++ ++logger = logging.getLogger(__name__) ++ ++# torch dtype -> the FeatureSpec dtype string SpecForge's zero-copy get() maps ++# back to a torch dtype. Keep in sync with MooncakeFeatureStore._TORCH_DTYPES. ++_DTYPE_STR = { ++ torch.float32: "float32", ++ torch.float64: "float64", ++ torch.float16: "float16", ++ torch.bfloat16: "bfloat16", ++ torch.int64: "int64", ++ torch.int32: "int32", ++ torch.int16: "int16", ++ torch.int8: "int8", ++ torch.uint8: "uint8", ++ torch.bool: "bool", ++} ++_STR_DTYPE = {v: k for k, v in _DTYPE_STR.items()} ++ ++_ARTIFACT_AUX = "aux" ++_ARTIFACT_LAST_HIDDEN = "last_hidden" ++ ++ ++class SpecCaptureSink: ++ """Writes captured per-request tensors into Mooncake in SpecForge layout.""" ++ ++ def __init__(self, aux_layer_ids: Optional[List[int]] = None) -> None: ++ self.aux_layer_ids = list(aux_layer_ids) if aux_layer_ids else None ++ self._store = None ++ self._put_config = None ++ self._lock = threading.Lock() ++ # One store writer is sufficient: Mooncake already stripes a batched ++ # transfer internally. The executor decouples that host transfer from ++ # the scheduler so the next target prefill can run concurrently. ++ self._executor = ThreadPoolExecutor( ++ max_workers=1, ++ thread_name_prefix="spec-capture-batch-put", ++ ) ++ ++ # -- connection --------------------------------------------------------- ++ def _connect(self): ++ if self._store is not None: ++ return self._store ++ with self._lock: ++ if self._store is not None: ++ return self._store ++ from mooncake.store import MooncakeDistributedStore, ReplicateConfig ++ ++ store = MooncakeDistributedStore() ++ rc = store.setup( ++ local_hostname=os.environ.get("MOONCAKE_LOCAL_HOSTNAME", "localhost"), ++ metadata_server=os.environ.get( ++ "MOONCAKE_METADATA_SERVER", "http://localhost:8080/metadata" ++ ), ++ global_segment_size=int( ++ os.environ.get("MOONCAKE_GLOBAL_SEGMENT_SIZE", 1 << 30) ++ ), ++ local_buffer_size=int( ++ os.environ.get("MOONCAKE_LOCAL_BUFFER_SIZE", 1 << 30) ++ ), ++ protocol=os.environ.get("MOONCAKE_PROTOCOL", "tcp"), ++ rdma_devices=os.environ.get("MOONCAKE_RDMA_DEVICES", ""), ++ master_server_addr=os.environ.get( ++ "MOONCAKE_MASTER_SERVER_ADDR", "localhost:50051" ++ ), ++ ) ++ if rc is not None and int(rc) != 0: ++ raise RuntimeError(f"spec-capture mooncake setup failed (status {rc})") ++ # Hard-pin every object: SpecForge (not Mooncake's LRU) is the ++ # lifetime authority — a committed feature must never be evicted ++ # before the trainer consumes it. ++ cfg = ReplicateConfig() ++ cfg.replica_num = 1 ++ cfg.with_hard_pin = True ++ self._put_config = cfg ++ self._store = store ++ logger.info("spec-capture mooncake sink connected") ++ return store ++ ++ # -- key/put primitives (pinned to MooncakeFeatureStore's layout) -------- ++ @staticmethod ++ def _tkey(store_id: str, sample_id: str, gen: int, name: str) -> str: ++ return f"{store_id}/{sample_id}/g{gen}/{name}" ++ ++ def _remove_quiet(self, key: str) -> None: ++ try: ++ self._connect().remove(key) ++ except Exception: ++ pass ++ ++ def _remove_many_quiet(self, keys: List[str]) -> None: ++ if not keys: ++ return ++ store = self._connect() ++ batch_remove = getattr(store, "batch_remove", None) ++ if batch_remove is not None: ++ try: ++ batch_remove(keys) ++ return ++ except Exception: ++ pass ++ for key in keys: ++ self._remove_quiet(key) ++ ++ # -- the batch entry point ------------------------------------------------ ++ def submit_samples( ++ self, ++ samples: List[ ++ Tuple[ ++ Dict[str, Any], Optional[torch.Tensor], Optional[torch.Tensor] ++ ] ++ ], ++ ) -> Future[List[Dict[str, Any]]]: ++ """Queue one scheduler batch without blocking the scheduler thread.""" ++ return self._executor.submit(self.put_samples, samples) ++ ++ def put_samples( ++ self, ++ samples: List[ ++ Tuple[ ++ Dict[str, Any], Optional[torch.Tensor], Optional[torch.Tensor] ++ ] ++ ], ++ ) -> List[Dict[str, Any]]: ++ """Publish a scheduler batch with one native Mooncake batch RPC. ++ ++ A K3 prefill normally finishes 16 samples together. Calling ++ ``put_from`` four times per sample paid 64 metadata/transport round ++ trips on the writer rank and serialized the response behind them. ++ ``batch_put_from`` preserves the existing per-feature keys and raw ++ tensor layout while amortizing that fixed cost across the whole ++ scheduler batch. The response is still emitted only after every ++ status succeeds, so refs can never point at incomplete samples. ++ """ ++ if not samples: ++ return [] ++ ++ store = self._connect() ++ timing_enabled = os.environ.get("SGLANG_SPEC_CAPTURE_TIMING", "0") == "1" ++ started = time.perf_counter() ++ keys: List[str] = [] ++ tensors: List[torch.Tensor] = [] ++ sizes: List[int] = [] ++ replace_keys: List[str] = [] ++ results: List[Dict[str, Any]] = [] ++ ++ def _stage( ++ result_feats: Dict[str, Dict[str, Any]], ++ *, ++ store_id: str, ++ sample_id: str, ++ gen: int, ++ replace: bool, ++ name: str, ++ tensor: torch.Tensor, ++ ) -> None: ++ tensor = tensor.detach().to("cpu").contiguous() ++ key = self._tkey(store_id, sample_id, gen, name) ++ keys.append(key) ++ tensors.append(tensor) ++ sizes.append(tensor.element_size() * tensor.numel()) ++ if replace: ++ replace_keys.append(key) ++ result_feats[name] = { ++ "shape": list(tensor.shape), ++ "dtype": _DTYPE_STR.get( ++ tensor.dtype, str(tensor.dtype).replace("torch.", "") ++ ), ++ } ++ ++ for spec, aux, last_hidden in samples: ++ store_id = str(spec["store_id"]) ++ sample_id = str(spec["sample_id"]) ++ gen = int(spec.get("gen", 1)) ++ replace = bool(spec.get("replace", False)) ++ features: Dict[str, str] = dict(spec.get("features") or {}) ++ result_feats: Dict[str, Dict[str, Any]] = {} ++ ++ aux_name = features.get(_ARTIFACT_AUX) ++ if aux_name is not None: ++ if aux is None: ++ raise RuntimeError( ++ "spec_capture requested 'aux' but no aux hidden states were " ++ "captured -- launch the server with --enable-spec-capture " ++ "(and optionally --spec-capture-aux-layer-ids)" ++ ) ++ _stage( ++ result_feats, ++ store_id=store_id, ++ sample_id=sample_id, ++ gen=gen, ++ replace=replace, ++ name=aux_name, ++ tensor=aux.unsqueeze(0), ++ ) ++ lh_name = features.get(_ARTIFACT_LAST_HIDDEN) ++ if lh_name is not None: ++ if last_hidden is None: ++ raise RuntimeError( ++ "spec_capture requested 'last_hidden' but the logits " ++ "processor did not return it (is aux capture enabled?)" ++ ) ++ _stage( ++ result_feats, ++ store_id=store_id, ++ sample_id=sample_id, ++ gen=gen, ++ replace=replace, ++ name=lh_name, ++ tensor=last_hidden.unsqueeze(0), ++ ) ++ for item in spec.get("passthrough") or []: ++ dtype = _STR_DTYPE.get(str(item.get("dtype", "int64"))) ++ if dtype is None: ++ raise RuntimeError( ++ f"spec_capture passthrough {item.get('name')!r}: " ++ f"unsupported dtype {item.get('dtype')!r}" ++ ) ++ tensor = torch.tensor(item["data"], dtype=dtype).reshape( ++ [int(d) for d in item["shape"]] ++ ) ++ _stage( ++ result_feats, ++ store_id=store_id, ++ sample_id=sample_id, ++ gen=gen, ++ replace=replace, ++ name=str(item["name"]), ++ tensor=tensor, ++ ) ++ results.append( ++ { ++ "sample_id": sample_id, ++ "store_id": store_id, ++ "gen": gen, ++ "aux_layer_ids": self.aux_layer_ids, ++ "features": result_feats, ++ } ++ ) ++ ++ materialize_ms = (time.perf_counter() - started) * 1000.0 ++ self._remove_many_quiet(replace_keys) ++ registered: List[torch.Tensor] = [] ++ register_started = time.perf_counter() ++ try: ++ for tensor, nbytes in zip(tensors, sizes): ++ try: ++ store.register_buffer(tensor.data_ptr(), nbytes) ++ registered.append(tensor) ++ except Exception: ++ pass # TCP and some Mooncake builds auto-register ++ register_ms = (time.perf_counter() - register_started) * 1000.0 ++ put_started = time.perf_counter() ++ batch_put = getattr(store, "batch_put_from", None) ++ if batch_put is None: ++ statuses = [ ++ store.put_from(key, tensor.data_ptr(), nbytes, self._put_config) ++ for key, tensor, nbytes in zip(keys, tensors, sizes) ++ ] ++ else: ++ statuses = batch_put( ++ keys, ++ [tensor.data_ptr() for tensor in tensors], ++ sizes, ++ self._put_config, ++ ) ++ put_ms = (time.perf_counter() - put_started) * 1000.0 ++ except Exception: ++ self._remove_many_quiet(keys) ++ raise ++ finally: ++ for tensor in registered: ++ try: ++ store.unregister_buffer(tensor.data_ptr()) ++ except Exception: ++ pass ++ ++ if statuses is None: ++ statuses = [0] * len(keys) ++ if len(statuses) != len(keys): ++ self._remove_many_quiet(keys) ++ raise RuntimeError( ++ "spec-capture batch_put_from returned " ++ f"{len(statuses)} statuses for {len(keys)} keys" ++ ) ++ failed = [ ++ (key, status) ++ for key, status in zip(keys, statuses) ++ if status is not None and int(status) < 0 ++ ] ++ if failed: ++ self._remove_many_quiet(keys) ++ raise RuntimeError( ++ "spec-capture batch_put_from failed for " ++ f"{len(failed)}/{len(keys)} keys; first={failed[0]}" ++ ) ++ ++ if timing_enabled: ++ logger.info( ++ "[spec-capture-timing] batch_sink samples=%d objects=%d " ++ "bytes=%d materialize_ms=%.3f register_ms=%.3f put_ms=%.3f " ++ "total_ms=%.3f", ++ len(samples), ++ len(keys), ++ sum(sizes), ++ materialize_ms, ++ register_ms, ++ put_ms, ++ (time.perf_counter() - started) * 1000.0, ++ ) ++ return results ++ ++ def put_sample( ++ self, ++ spec: Dict[str, Any], ++ *, ++ aux: Optional[torch.Tensor], ++ last_hidden: Optional[torch.Tensor], ++ ) -> Dict[str, Any]: ++ """Write one sample's artifacts; return the meta_info result dict. ++ ++ ``aux``/``last_hidden`` are the per-request (L, W) captured tensors, ++ stored with a leading batch dim of 1. On any failure the keys already ++ written are best-effort removed (no partial sample is consumable). ++ """ ++ # Keep the single-sample entry point for compatibility with tests and ++ # callers outside the scheduler; production uses put_samples(). ++ return self.put_samples([(spec, aux, last_hidden)])[0] ++ ++ ++_SINK: Optional[SpecCaptureSink] = None ++ ++ ++def maybe_init_sink(server_args) -> None: ++ """Called from Scheduler init on the writer rank when spec capture is on. ++ ++ Connection to Mooncake is lazy (first put), so a capture-enabled server ++ without a reachable Mooncake master still boots and serves normal traffic. ++ """ ++ global _SINK ++ if getattr(server_args, "enable_spec_capture", False) and _SINK is None: ++ _SINK = SpecCaptureSink( ++ aux_layer_ids=getattr(server_args, "spec_capture_aux_layer_ids", None) ++ ) ++ ++ ++def get_sink() -> Optional[SpecCaptureSink]: ++ return _SINK diff --git a/patches/sglang/v0.5.14/spec-capture-ascend-mount.patch b/patches/sglang/v0.5.14/spec-capture-ascend-mount.patch new file mode 100644 index 000000000..40333a3e6 --- /dev/null +++ b/patches/sglang/v0.5.14/spec-capture-ascend-mount.patch @@ -0,0 +1,76 @@ +diff --git a/python/sglang/srt/spec_capture_sink.py b/python/sglang/srt/spec_capture_sink.py +index 4955e11..902f0b1 100644 +--- a/python/sglang/srt/spec_capture_sink.py ++++ b/python/sglang/srt/spec_capture_sink.py +@@ -81,18 +81,31 @@ class SpecCaptureSink: + from mooncake.store import MooncakeDistributedStore, ReplicateConfig + + store = MooncakeDistributedStore() ++ global_segment_size = int( ++ os.environ.get("MOONCAKE_GLOBAL_SEGMENT_SIZE", 1 << 30) ++ ) ++ local_buffer_size = int( ++ os.environ.get("MOONCAKE_LOCAL_BUFFER_SIZE", 1 << 30) ++ ) ++ protocol = os.environ.get("MOONCAKE_PROTOCOL", "tcp") ++ # Ascend Mooncake rejects the wildcard location ("location:* is ++ # not supported"); skip it in setup() and mount with location="cpu". ++ ascend_host = bool(os.environ.get("ASCEND_RT_VISIBLE_DEVICES")) ++ segment_to_mount = global_segment_size if ascend_host else 0 ++ if ascend_host: ++ global_segment_size = 0 ++ local_buffer_size = 0 + rc = store.setup( + local_hostname=os.environ.get("MOONCAKE_LOCAL_HOSTNAME", "localhost"), + metadata_server=os.environ.get( + "MOONCAKE_METADATA_SERVER", "http://localhost:8080/metadata" + ), +- global_segment_size=int( +- os.environ.get("MOONCAKE_GLOBAL_SEGMENT_SIZE", 1 << 30) +- ), +- local_buffer_size=int( +- os.environ.get("MOONCAKE_LOCAL_BUFFER_SIZE", 1 << 30) +- ), +- protocol=os.environ.get("MOONCAKE_PROTOCOL", "tcp"), ++ global_segment_size=global_segment_size, ++ local_buffer_size=local_buffer_size, ++ protocol=protocol, + rdma_devices=os.environ.get("MOONCAKE_RDMA_DEVICES", ""), + master_server_addr=os.environ.get( + "MOONCAKE_MASTER_SERVER_ADDR", "localhost:50051" +@@ -100,12 +113,34 @@ class SpecCaptureSink: + ) + if rc is not None and int(rc) != 0: + raise RuntimeError(f"spec-capture mooncake setup failed (status {rc})") ++ if segment_to_mount: ++ mount = getattr(store, "allocate_and_mount_segment", None) ++ if mount is None: ++ raise RuntimeError( ++ "Mooncake build on this Ascend host cannot register a " ++ "wildcard segment and has no allocate_and_mount_segment; " ++ "upgrade mooncake-transfer-engine" ++ ) ++ result = mount(segment_to_mount, protocol, "cpu") ++ mrc = result.get("ret", -1) if isinstance(result, dict) else result ++ if mrc is not None and int(mrc) != 0: ++ raise RuntimeError( ++ f"spec-capture mooncake mount segment failed (status {mrc})" ++ ) ++ logger.info( ++ "spec-capture mooncake segment mounted with location=cpu " ++ "(%d bytes)", ++ segment_to_mount, ++ ) + # Hard-pin every object: SpecForge (not Mooncake's LRU) is the + # lifetime authority — a committed feature must never be evicted + # before the trainer consumes it. + cfg = ReplicateConfig() + cfg.replica_num = 1 +- cfg.with_hard_pin = True ++ # Older Mooncake builds expose no with_hard_pin field; objects ++ # then follow the store's default pin behavior until remove(). ++ if hasattr(cfg, "with_hard_pin"): ++ cfg.with_hard_pin = True + self._put_config = cfg + self._store = store + logger.info("spec-capture mooncake sink connected") diff --git a/patches/sglang/v0.5.14/spec-capture.patch b/patches/sglang/v0.5.14/spec-capture.patch index efe9f4a89..5fd29d1e1 100644 --- a/patches/sglang/v0.5.14/spec-capture.patch +++ b/patches/sglang/v0.5.14/spec-capture.patch @@ -1,5 +1,5 @@ diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py -index a99d252677..f5d0b35e79 100644 +index a99d252..bde8a84 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -94,6 +94,9 @@ class LogitsProcessorOutput: @@ -12,7 +12,7 @@ index a99d252677..f5d0b35e79 100644 ## Part 2: This part will be assigned in python/sglang/srt/layers/sampler.py::Sampler # he log probs of output tokens, if SGLANG_RETURN_ORIGINAL_LOGPROB = True, will get the log probs before applying temperature. If False, will get the log probs before applying temperature. -@@ -361,6 +364,16 @@ class LogitsProcessor(nn.Module): +@@ -361,6 +364,30 @@ class LogitsProcessor(nn.Module): sample_indices, logits_metadata, ) @@ -26,10 +26,24 @@ index a99d252677..f5d0b35e79 100644 + ) + else None + ) ++ # muP targets pass LM-head-scaled hidden states into LogitsProcessor, ++ # while SpecForge folds the same multiplier into its frozen target ++ # head. Restore the pre-head-scale representation before capture so ++ # the multiplier is applied exactly once when training recomputes logits. ++ logits_mup_width_multiplier = getattr( ++ self.config, "logits_mup_width_multiplier", None ++ ) ++ if ( ++ last_hidden_states_to_store is not None ++ and logits_mup_width_multiplier ++ ): ++ last_hidden_states_to_store = ( ++ last_hidden_states_to_store * float(logits_mup_width_multiplier) ++ ) del hidden_states if not logits_metadata.extend_return_logprob: -@@ -374,6 +387,7 @@ class LogitsProcessor(nn.Module): +@@ -374,6 +401,7 @@ class LogitsProcessor(nn.Module): return LogitsProcessorOutput( next_token_logits=sampled_logits, hidden_states=hidden_states_to_store, @@ -38,7 +52,7 @@ index a99d252677..f5d0b35e79 100644 # workaround since ForwardBatch is local to forward_batch_generation(). # They should be moved to GenerationBatchResult to keep this class clean. diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py -index b05334deaa..2853c82315 100644 +index b05334d..2853c82 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -441,6 +441,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin): @@ -50,20 +64,20 @@ index b05334deaa..2853c82315 100644 indexer_topk=indexer_topk, customized_info=recv_obj.customized_info, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py -index 951f35495e..4de78534a4 100644 +index 951f354..cd6ebf2 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py -@@ -280,6 +280,10 @@ class GenerateReqInput(BaseReq): - image_max_dynamic_patch: Optional[int] = None - video_max_dynamic_patch: Optional[int] = None - +@@ -188,6 +188,10 @@ class GenerateReqInput(BaseReq): + log_metrics: bool = True + # Whether to return hidden states + return_hidden_states: Union[List[bool], bool] = False + # Spec-training capture sink instructions (see spec_capture_sink.py). + # Batch-level: List[Optional[dict]]; per-request after __getitem__. + spec_capture: Optional[Union[List[Optional[Dict]], Dict]] = None + - # Pre-computed delimiter indices for multi-item scoring. - # Batch-level: List[List[int]] (one per request). After __getitem__: List[int]. - multi_item_delimiter_indices: Optional[Union[List[List[int]], List[int]]] = None + # Whether to return captured routed experts + return_routed_experts: bool = False + return_indexer_topk: bool = False @@ -743,6 +747,11 @@ class GenerateReqInput(BaseReq): if self.multi_item_delimiter_indices is not None else None @@ -107,7 +121,7 @@ index 951f35495e..4de78534a4 100644 token_steps: List[List[int]] = None diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py -index f1dc81d174..acefaace89 100755 +index f1dc81d..acefaac 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -707,6 +707,7 @@ class Req(ReqDllmMixin): @@ -134,7 +148,7 @@ index f1dc81d174..acefaace89 100755 # extra key for classifying the request (e.g. cache_salt) if lora_id is not None: diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py -index abba37441d..3018b82473 100644 +index abba374..241da3f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -576,6 +576,19 @@ class Scheduler( @@ -165,11 +179,103 @@ index abba37441d..3018b82473 100644 ) req.tokenizer = self.tokenizer +@@ -3097,6 +3111,13 @@ class Scheduler( + # GenerationBatchResult.extra_keep_alive_refs after forward returns. + self.batch_record_buf[self.batch_record_ct] = [batch, attr_snapshot] + ++ def _should_copy_hidden_states_to_cpu(self, batch: ScheduleBatch) -> bool: ++ """Avoid redundant multi-GiB capture D2H on non-writer TP ranks.""" ++ return batch.return_hidden_states and ( ++ self.ps.attn_tp_rank == 0 ++ or any(req.spec_capture is None for req in batch.reqs) ++ ) ++ + @contextmanager + def _forward_isolation(self, batch: ScheduleBatch, *, overlap: bool): + """Make SB transactional across one forward (overlap and non-overlap). +@@ -3258,7 +3279,9 @@ class Scheduler( + batch_result.copy_done = self.device_module.Event() + batch_result.copy_to_cpu( + return_logprob=batch.return_logprob, +- return_hidden_states=batch.return_hidden_states, ++ return_hidden_states=self._should_copy_hidden_states_to_cpu( ++ batch ++ ), + ) + else: + kwargs = ( +@@ -3369,6 +3392,10 @@ class Scheduler( + batch: ScheduleBatch, + result: Union[GenerationBatchResult, EmbeddingBatchResult], + ): ++ # Complete capture transfers on the scheduler thread before processing ++ # any kind of next result (prefill, decode, or idle). This preserves ++ # the output socket's single-thread ownership under mixed traffic. ++ self.batch_result_processor.drain_spec_captures() + self.publish_load_snapshot(force=batch.forward_mode.is_extend()) + + if batch.forward_mode.is_decode(): +@@ -3448,6 +3475,15 @@ class Scheduler( + + def on_idle(self): + """Idle housekeeping: guard, check, metrics, reset, sleep.""" ++ # A capture response is intentionally delayed until its background ++ # sink transfer completes; finish it on this thread so ownership of ++ # the ZeroMQ output socket remains single-threaded. Do not enter the ++ # idle sleeper while a future is outstanding or the producer waiting ++ # for that response could deadlock. ++ self.batch_result_processor.drain_spec_captures() ++ if self.batch_result_processor.has_pending_spec_captures(): ++ time.sleep(0.001) ++ return + if not self.is_fully_idle(): + return + +@@ -3496,6 +3532,7 @@ class Scheduler( + and (self.last_batch is None or self.last_batch.is_empty()) + and (self.cur_batch is None or self.cur_batch.is_empty()) + and (not self.enable_overlap or len(self.result_queue) == 0) ++ and not self.batch_result_processor.has_pending_spec_captures() + and self._pp_microbatches_drained() + ) + diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py -index a9d5f0c28b..c9f8fdf7d3 100644 +index a9d5f0c..f134354 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py -@@ -257,6 +257,18 @@ class SchedulerBatchResultProcessor: +@@ -1,7 +1,9 @@ + from __future__ import annotations + + import logging +-from dataclasses import dataclass ++import os ++import time ++from dataclasses import dataclass, field + from typing import ( + TYPE_CHECKING, + Callable, +@@ -78,6 +80,12 @@ class SchedulerBatchResultProcessor: + logprob_result_processor: SchedulerLogprobResultProcessor + output_streamer: SchedulerOutputStreamer + abort_request: Callable ++ _spec_capture_batches: List = field( ++ default_factory=list, ++ init=False, ++ repr=False, ++ compare=False, ++ ) + + def process_batch_result_prebuilt(self, batch: ScheduleBatch): + assert self.disaggregation_mode == DisaggregationMode.DECODE +@@ -181,6 +189,7 @@ class SchedulerBatchResultProcessor: + result: Union[GenerationBatchResult, EmbeddingBatchResult], + ): + skip_stream_req = None ++ pending_spec_captures: List = [] + + if self.is_generation: + if result.copy_done is not None: +@@ -257,6 +266,20 @@ class SchedulerBatchResultProcessor: ) if ( @@ -183,12 +289,43 @@ index a9d5f0c28b..c9f8fdf7d3 100644 + hidden_state_offset=hidden_state_offset, + ) + if req.finished(): -+ self._sink_spec_capture(req) ++ pending = self._sink_spec_capture(req) ++ if pending is not None: ++ pending_spec_captures.append(pending) + elif ( req.return_hidden_states and logits_output.hidden_states is not None ): -@@ -462,6 +474,95 @@ class SchedulerBatchResultProcessor: +@@ -329,9 +352,25 @@ class SchedulerBatchResultProcessor: + req.inflight_middle_chunks -= 1 + req.time_stats.set_last_chunked_prefill_finish_time() + +- self.output_streamer.stream_output( +- batch.reqs, batch.return_logprob, skip_stream_req +- ) ++ if pending_spec_captures: ++ self._queue_spec_captures( ++ pending_spec_captures, ++ return_logprob=batch.return_logprob, ++ ) ++ capture_req_ids = {id(item[0]) for item in pending_spec_captures} ++ ready_reqs = [ ++ req ++ for req in batch.reqs ++ if id(req) not in capture_req_ids and req is not skip_stream_req ++ ] ++ if ready_reqs: ++ self.output_streamer.stream_output( ++ ready_reqs, batch.return_logprob ++ ) ++ else: ++ self.output_streamer.stream_output( ++ batch.reqs, batch.return_logprob, skip_stream_req ++ ) + + can_run_cuda_graph = result.can_run_cuda_graph + self.metrics_reporter.report_prefill_stats( +@@ -462,6 +501,183 @@ class SchedulerBatchResultProcessor: f"Placeholder zeros would be appended to output_ids." ) @@ -207,20 +344,43 @@ index a9d5f0c28b..c9f8fdf7d3 100644 + """ + start = hidden_state_offset + end = start + len(req.origin_input_ids) -+ req.spec_capture_aux.append( -+ logits_output.hidden_states[start:end].cpu().clone() ++ # Only attention-TP rank 0 owns the Mooncake sink. Copying these very ++ # large tensors to host on every TP rank creates identical D2H ++ # transfers whose results are immediately discarded on TP > 1. ++ if self.output_streamer.ps.attn_tp_rank != 0: ++ return end ++ # Materialize each scheduler result on host once, rather than issuing ++ # one synchronous D2H transfer per request. In overlap mode the writer ++ # rank has already copied both tensors asynchronously on copy_stream; ++ # ``.cpu()`` is then a no-op. The fallback keeps non-overlap correct. ++ features = dict(req.spec_capture.get("features") or {}) ++ aux_cpu = getattr(logits_output, "_spec_capture_aux_cpu", None) ++ if "aux" in features and aux_cpu is None: ++ aux_cpu = logits_output.hidden_states.cpu() ++ logits_output._spec_capture_aux_cpu = aux_cpu ++ last_hidden_cpu = getattr( ++ logits_output, "_spec_capture_last_hidden_cpu", None + ) -+ if logits_output.last_hidden_states is not None: -+ req.spec_capture_last_hidden.append( -+ logits_output.last_hidden_states[start:end].cpu().clone() -+ ) ++ if ( ++ "last_hidden" in features ++ and logits_output.last_hidden_states is not None ++ and last_hidden_cpu is None ++ ): ++ last_hidden_cpu = logits_output.last_hidden_states.cpu() ++ logits_output._spec_capture_last_hidden_cpu = last_hidden_cpu ++ if "aux" in features and aux_cpu is not None: ++ req.spec_capture_aux.append(aux_cpu[start:end]) ++ if "last_hidden" in features and last_hidden_cpu is not None: ++ req.spec_capture_last_hidden.append(last_hidden_cpu[start:end]) + return end + -+ def _sink_spec_capture(self, req: Req) -> None: -+ """Write a finished capture request's tensors to the Mooncake sink. ++ def _sink_spec_capture(self, req: Req): ++ """Stage a finished capture request for the background Mooncake sink. + -+ Runs on the attention-TP rank that streams output; the per-request result -+ (or an ``{"error": ...}`` marker) is set on ``req.spec_capture_result``, ++ Runs on the attention-TP rank that streams output; returns the pending ++ tuple queued by ``_queue_spec_captures`` (or ``None`` off the writer ++ rank). The per-request result (or an ``{"error": ...}`` marker) is set ++ on ``req.spec_capture_result`` when the batch transfer completes, and + returned to the client via the dedicated ``spec_capture`` output field + (a per-request channel, unlike per-token ``customized_info``). + """ @@ -228,67 +388,147 @@ index a9d5f0c28b..c9f8fdf7d3 100644 + + sink = spec_capture_sink.get_sink() + if sink is None or self.output_streamer.ps.attn_tp_rank != 0: ++ return None ++ timing_enabled = os.environ.get("SGLANG_SPEC_CAPTURE_TIMING", "0") == "1" ++ cat_start = time.perf_counter() ++ # With chunked prefill disabled (the recommended capture-server ++ # configuration), each request owns exactly one contiguous view into ++ # the batch-level D2H buffer. torch.cat([view]) needlessly copied the ++ # whole capture a second time on CPU. Keep that view zero-copy; ++ # concatenate only the genuinely chunked case. ++ aux = self._coalesce_spec_capture_chunks(req.spec_capture_aux) ++ last_hidden = self._coalesce_spec_capture_chunks( ++ req.spec_capture_last_hidden ++ ) ++ cat_ms = (time.perf_counter() - cat_start) * 1000.0 ++ req.spec_capture_aux = [] ++ req.spec_capture_last_hidden = [] ++ return req, req.spec_capture, aux, last_hidden, timing_enabled, cat_ms ++ ++ def _queue_spec_captures(self, pending, *, return_logprob: bool) -> None: ++ """Queue one batch transfer while the scheduler runs the next prefill.""" ++ if not pending: + return -+ aux = torch.cat(req.spec_capture_aux, dim=0) if req.spec_capture_aux else None -+ last_hidden = ( -+ torch.cat(req.spec_capture_last_hidden, dim=0) -+ if req.spec_capture_last_hidden -+ else None ++ from sglang.srt import spec_capture_sink ++ ++ sink = spec_capture_sink.get_sink() ++ samples = [ ++ (spec, aux, last_hidden) ++ for _, spec, aux, last_hidden, _, _ in pending ++ ] ++ self._spec_capture_batches.append( ++ ( ++ pending, ++ sink.submit_samples(samples), ++ return_logprob, ++ time.perf_counter(), ++ ) + ) -+ position_ids = self._spec_capture_position_ids(req) -+ try: -+ req.spec_capture_result = sink.put_sample( -+ req.spec_capture, -+ aux=aux, -+ last_hidden=last_hidden, -+ position_ids=position_ids, ++ # Bound retained D2H buffers. Reaching this point means the target ++ # prefill for batch N+1 already overlapped host transfer N; wait only ++ # if N is still finishing before admitting N+2. ++ max_pending = int( ++ os.environ.get("SGLANG_SPEC_CAPTURE_MAX_PENDING_BATCHES", "2") ++ ) ++ if max_pending < 1: ++ raise ValueError("SGLANG_SPEC_CAPTURE_MAX_PENDING_BATCHES must be >= 1") ++ if len(self._spec_capture_batches) >= max_pending: ++ self.drain_spec_captures(block=True, max_batches=1) ++ ++ def has_pending_spec_captures(self) -> bool: ++ return bool(self._spec_capture_batches) ++ ++ def drain_spec_captures( ++ self, *, block: bool = False, max_batches: Optional[int] = None ++ ) -> int: ++ """Finish ready transfers and stream their responses on this thread.""" ++ completed = 0 ++ while self._spec_capture_batches: ++ if max_batches is not None and completed >= max_batches: ++ break ++ pending, future, return_logprob, queued_at = self._spec_capture_batches[0] ++ if not block and not future.done(): ++ break ++ self._spec_capture_batches.pop(0) ++ self._complete_spec_capture_batch( ++ pending, ++ future, ++ return_logprob=return_logprob, ++ queued_at=queued_at, + ) ++ completed += 1 ++ return completed ++ ++ def _complete_spec_capture_batch( ++ self, pending, future, *, return_logprob: bool, queued_at: float ++ ) -> None: ++ try: ++ results = future.result() ++ if len(results) != len(pending): ++ raise RuntimeError( ++ f"spec-capture sink returned {len(results)} results for " ++ f"{len(pending)} samples" ++ ) + except Exception as e: -+ logger.error("spec-capture sink failed for %s: %s", req.rid, e) -+ req.spec_capture_result = { -+ "sample_id": req.spec_capture.get("sample_id"), -+ "error": str(e), -+ } -+ req.spec_capture_aux = [] -+ req.spec_capture_last_hidden = [] ++ logger.error( ++ "spec-capture batch sink failed for %d requests: %s", len(pending), e ++ ) ++ for req, spec, _, _, _, _ in pending: ++ req.spec_capture_result = { ++ "sample_id": spec.get("sample_id"), ++ "error": str(e), ++ } ++ else: ++ for pending_item, result in zip(pending, results): ++ req, _, _, _, _, _ = pending_item ++ req.spec_capture_result = result ++ ++ timing_enabled = any(item[4] for item in pending) ++ if timing_enabled: ++ logger.info( ++ "[spec-capture-timing] async_batch_complete samples=%d " ++ "queue_to_stream_ms=%.3f cat_ms=%.3f", ++ len(pending), ++ (time.perf_counter() - queued_at) * 1000.0, ++ sum(item[5] for item in pending), ++ ) ++ self.output_streamer.stream_output( ++ [item[0] for item in pending], return_logprob ++ ) + + @staticmethod -+ def _spec_capture_position_ids(req: Req) -> Optional[torch.Tensor]: -+ """Prompt position ids for the sink's ``position_ids`` artifact. -+ -+ Multimodal (mRoPE) requests carry the processor-computed positions on -+ ``req.multimodal_inputs.mrope_positions`` ([3, L]); text requests fall -+ back to the plain arange broadcast the model uses for non-mm prompts. -+ Returns an (L, 3) CPU int64 tensor, or None when the client did not -+ request the artifact. -+ """ -+ features = (req.spec_capture or {}).get("features") or {} -+ if features.get("position_ids") is None: ++ def _coalesce_spec_capture_chunks( ++ chunks: List[torch.Tensor], ++ ) -> Optional[torch.Tensor]: ++ if not chunks: + return None -+ seq_len = len(req.origin_input_ids) -+ mrope_positions = ( -+ req.multimodal_inputs.mrope_positions -+ if req.multimodal_inputs is not None -+ else None -+ ) -+ if mrope_positions is not None: -+ return mrope_positions[:, :seq_len].t().contiguous().cpu() -+ return ( -+ torch.arange(seq_len, dtype=torch.int64) -+ .unsqueeze(0) -+ .expand(3, -1) -+ .t() -+ .contiguous() -+ ) ++ if len(chunks) == 1: ++ return chunks[0] ++ return torch.cat(chunks, dim=0) + def _append_prefill_hidden_states( self, *, diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py -index f95c59f6f2..9e64326a13 100644 +index f95c59f..6c51abc 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py -@@ -273,6 +273,7 @@ class _GenerationStreamAccumulator: +@@ -145,6 +145,14 @@ class SchedulerOutputStreamer: + for req in reqs: + if req is skip_req: + continue ++ if self.server_args.enable_spec_capture and req.spec_capture is not None: ++ # Capture requests are not complete until the background sink ++ # has durably published every feature object. Several ++ # scheduler paths can ask the common streamer to emit a ++ # finished request; centralize the completion barrier here so ++ # none of them can win the race and send a metadata-less 200. ++ if req.finished() and req.spec_capture_result is None: ++ continue + if req.finished() and req.finished_output: + # With the overlap schedule, a request will try to output twice and hit this line twice + # because of the one additional delayed token. This "continue" prevented the dummy output. +@@ -273,6 +281,7 @@ class _GenerationStreamAccumulator: spec_correct_drafts_histogram: list = field(default_factory=list) retraction_counts: list = field(default_factory=list) output_hidden_states: Optional[list] = None @@ -296,7 +536,7 @@ index f95c59f6f2..9e64326a13 100644 routed_experts: Optional[list] = None indexer_topk: Optional[list] = None customized_info: dict = field(default_factory=dict) -@@ -482,6 +483,8 @@ class _GenerationStreamAccumulator: +@@ -482,6 +491,8 @@ class _GenerationStreamAccumulator: self.output_hidden_states.append(hs) else: self.output_hidden_states.append(None) @@ -305,7 +545,7 @@ index f95c59f6f2..9e64326a13 100644 if self.return_routed_experts: self.routed_experts.append( req.routed_experts if req.return_routed_experts else None -@@ -540,6 +543,7 @@ class _GenerationStreamAccumulator: +@@ -540,6 +551,7 @@ class _GenerationStreamAccumulator: output_token_ids_logprobs_idx=self.output_token_ids_logprobs_idx, output_token_entropy_val=None, output_hidden_states=self.output_hidden_states, @@ -314,7 +554,7 @@ index f95c59f6f2..9e64326a13 100644 indexer_topk=self.indexer_topk, customized_info=self.customized_info, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py -index bf932611a5..f0e821ec47 100644 +index bf93261..f0e821e 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1172,6 +1172,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): @@ -336,11 +576,32 @@ index bf932611a5..f0e821ec47 100644 if getattr(recv_obj, "routed_experts", None): val = recv_obj.routed_experts[i] if val is not None: +diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py +index e9ede57..1d624fd 100644 +--- a/python/sglang/srt/managers/utils.py ++++ b/python/sglang/srt/managers/utils.py +@@ -77,6 +77,16 @@ class GenerationBatchResult: + Only the tensors which are needed for processing results are copied, + e.g., next_token_ids, logits outputs + """ ++ if ( ++ return_hidden_states ++ and self.logits_output.last_hidden_states is not None ++ ): ++ # Spec-capture's post-norm last hidden rides the same async D2H ++ # as the aux hidden states (the field exists only when the ++ # spec-capture patch is applied). ++ self.logits_output.last_hidden_states = ( ++ self.logits_output.last_hidden_states.to("cpu", non_blocking=True) ++ ) + if return_logprob: + if self.logits_output.next_token_logprobs is not None: + self.logits_output.next_token_logprobs = ( diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py -index 1cff5c9839..b842aeccb9 100644 +index 1cff5c9..5309965 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py -@@ -515,6 +515,22 @@ class ModelRunner(ModelRunnerKVCacheMixin): +@@ -515,6 +515,42 @@ class ModelRunner(ModelRunnerKVCacheMixin): draft_num_layers=int(draft_num_layers), ) @@ -348,26 +609,46 @@ index 1cff5c9839..b842aeccb9 100644 + # Aux capture without a draft worker, routed to the strategy's own + # capture method (they wire different submodules — e.g. VL models + # only populate layers_to_capture via the dflash path). -+ if getattr(server_args, "spec_capture_method", "eagle3") == "dflash": ++ capture_method = getattr(server_args, "spec_capture_method", "eagle3") ++ if capture_method in ("dflash", "dspark"): ++ # DSpark rides the DFlash aux plumbing: this patch's supported ++ # targets expose dflash capture hooks (targets with a native ++ # dspark hook are served by the kimi-k3 patch variant). + self.dflash_use_aux_hidden_state = True + self.dflash_target_layer_ids = server_args.spec_capture_aux_layer_ids ++ if hasattr(self, "spec_aux_config"): ++ self.spec_aux_config.dflash_use_aux_hidden_state = True ++ self.spec_aux_config.dflash_target_layer_ids = ( ++ server_args.spec_capture_aux_layer_ids ++ ) + self.dflash_family_use_aux_hidden_state = True + self.dflash_family_target_layer_ids = ( + server_args.spec_capture_aux_layer_ids + ) -+ else: ++ elif capture_method == "eagle3": + self.eagle_use_aux_hidden_state = True + self.eagle_aux_hidden_state_layer_ids = ( + server_args.spec_capture_aux_layer_ids + ) ++ if hasattr(self, "spec_aux_config"): ++ self.spec_aux_config.eagle_use_aux_hidden_state = True ++ self.spec_aux_config.eagle_aux_hidden_state_layer_ids = ( ++ server_args.spec_capture_aux_layer_ids ++ ) ++ else: ++ raise ValueError( ++ "--spec-capture-method must be one of: eagle3, dflash, dspark; " ++ f"got {capture_method!r}" ++ ) ++ # Apply the rank zero filter to logger if server_args.show_time_cost: enable_show_time_cost() diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py -index c7162c16d5..d515b3c699 100644 +index c7162c1..cfbce7a 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py -@@ -2127,6 +2127,25 @@ class ServerArgs: +@@ -2127,6 +2127,26 @@ class ServerArgs: bool, "Enable returning hidden states with responses.", ] = False @@ -386,19 +667,20 @@ index c7162c16d5..d515b3c699 100644 + ] = None + spec_capture_method: A[ + str, -+ "Capture method for --enable-spec-capture: 'eagle3' or 'dflash'. Must " -+ "match the draft strategy being trained; they wire capture onto " -+ "different submodules (VL targets only populate the dflash path).", ++ "Capture method for --enable-spec-capture: 'eagle3', 'dflash', or " ++ "'dspark'. Must match the draft strategy being trained; they wire " ++ "capture onto different submodules (VL targets only populate the " ++ "dflash path; dspark rides the dflash aux plumbing).", + ] = "eagle3" enable_return_routed_experts: A[ bool, "Enable returning routed experts of each layer with responses.", diff --git a/python/sglang/srt/spec_capture_sink.py b/python/sglang/srt/spec_capture_sink.py new file mode 100644 -index 0000000000..a678bd85e7 +index 0000000..d54b920 --- /dev/null +++ b/python/sglang/srt/spec_capture_sink.py -@@ -0,0 +1,256 @@ +@@ -0,0 +1,395 @@ +# Copyright 2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. @@ -419,8 +701,7 @@ index 0000000000..a678bd85e7 +Request schema:: + + {"store_id", "sample_id", "gen", "replace", # key namespace / retry policy -+ "features": {"aux": , "last_hidden": , # artifact -> feature -+ "position_ids": }, # (L, 3) prompt positions ++ "features": {"aux": , "last_hidden": }, # artifact -> feature + "passthrough": [{"name", "data", "shape", "dtype"}]} # client tensors verbatim + +Response (``meta_info["spec_capture"]``): ``{"sample_id", "store_id", "gen", @@ -435,7 +716,9 @@ index 0000000000..a678bd85e7 +import logging +import os +import threading -+from typing import Any, Dict, List, Optional ++import time ++from concurrent.futures import Future, ThreadPoolExecutor ++from typing import Any, Dict, List, Optional, Tuple + +import torch + @@ -459,7 +742,6 @@ index 0000000000..a678bd85e7 + +_ARTIFACT_AUX = "aux" +_ARTIFACT_LAST_HIDDEN = "last_hidden" -+_ARTIFACT_POSITION_IDS = "position_ids" + + +class SpecCaptureSink: @@ -470,9 +752,13 @@ index 0000000000..a678bd85e7 + self._store = None + self._put_config = None + self._lock = threading.Lock() -+ # Retried HTTP requests reuse deterministic keys. Striped locks keep -+ # replacement atomic per key without retaining one lock per sample. -+ self._write_locks = [threading.Lock() for _ in range(256)] ++ # One store writer is sufficient: Mooncake already stripes a batched ++ # transfer internally. The executor decouples that host transfer from ++ # the scheduler so the next target prefill can run concurrently. ++ self._executor = ThreadPoolExecutor( ++ max_workers=1, ++ thread_name_prefix="spec-capture-batch-put", ++ ) + + # -- connection --------------------------------------------------------- + def _connect(self): @@ -508,7 +794,13 @@ index 0000000000..a678bd85e7 + # before the trainer consumes it. + cfg = ReplicateConfig() + cfg.replica_num = 1 -+ cfg.with_hard_pin = True ++ # `with_hard_pin` exists only on newer Mooncake builds; older ROCm ++ # sglang images expose only `with_soft_pin`. Map the hard-pin intent ++ # onto whatever the installed build supports. ++ if hasattr(cfg, "with_hard_pin"): ++ cfg.with_hard_pin = True ++ elif hasattr(cfg, "with_soft_pin"): ++ cfg.with_soft_pin = True + self._put_config = cfg + self._store = store + logger.info("spec-capture mooncake sink connected") @@ -519,83 +811,118 @@ index 0000000000..a678bd85e7 + def _tkey(store_id: str, sample_id: str, gen: int, name: str) -> str: + return f"{store_id}/{sample_id}/g{gen}/{name}" + -+ def _put_tensor( -+ self, key: str, t: torch.Tensor, *, replace: bool = False -+ ) -> None: -+ store = self._connect() -+ t = t.detach().to("cpu").contiguous() -+ nbytes = t.element_size() * t.numel() -+ lock = self._write_locks[hash(key) % len(self._write_locks)] -+ with lock: -+ if replace: -+ # Do not probe with is_exist(): Mooncake existence checks can -+ # acquire a read lease that prevents the following removal. -+ self._remove_quiet(key) -+ try: -+ store.register_buffer(t.data_ptr(), nbytes) -+ except Exception: -+ pass # some builds auto-register -+ try: -+ rc = store.put_from(key, t.data_ptr(), nbytes, self._put_config) -+ finally: -+ try: -+ store.unregister_buffer(t.data_ptr()) -+ except Exception: -+ pass -+ if rc is not None and int(rc) < 0: -+ raise RuntimeError(f"spec-capture put_from failed (status {rc}) for {key}") -+ + def _remove_quiet(self, key: str) -> None: + try: + self._connect().remove(key) + except Exception: + pass + -+ # -- the one entry point -------------------------------------------------- -+ def put_sample( ++ def _remove_many_quiet(self, keys: List[str]) -> None: ++ if not keys: ++ return ++ store = self._connect() ++ batch_remove = getattr(store, "batch_remove", None) ++ if batch_remove is not None: ++ try: ++ batch_remove(keys) ++ return ++ except Exception: ++ pass ++ for key in keys: ++ self._remove_quiet(key) ++ ++ # -- the batch entry point ------------------------------------------------ ++ def submit_samples( + self, -+ spec: Dict[str, Any], -+ *, -+ aux: Optional[torch.Tensor], -+ last_hidden: Optional[torch.Tensor], -+ position_ids: Optional[torch.Tensor] = None, -+ ) -> Dict[str, Any]: -+ """Write one sample's artifacts; return the meta_info result dict. ++ samples: List[ ++ Tuple[ ++ Dict[str, Any], Optional[torch.Tensor], Optional[torch.Tensor] ++ ] ++ ], ++ ) -> Future[List[Dict[str, Any]]]: ++ """Queue one scheduler batch without blocking the scheduler thread.""" ++ return self._executor.submit(self.put_samples, samples) + -+ ``aux``/``last_hidden`` are the per-request (L, W) captured tensors, -+ ``position_ids`` the per-request (L, 3) prompt positions (mRoPE rows -+ for multimodal requests, arange broadcast for text), all stored with a -+ leading batch dim of 1. On any failure the keys already written are -+ best-effort removed (no partial sample is consumable). ++ def put_samples( ++ self, ++ samples: List[ ++ Tuple[ ++ Dict[str, Any], Optional[torch.Tensor], Optional[torch.Tensor] ++ ] ++ ], ++ ) -> List[Dict[str, Any]]: ++ """Publish a scheduler batch with one native Mooncake batch RPC. ++ ++ A capture prefill batch normally finishes many samples together. ++ Calling ``put_from`` once per feature object paid dozens of metadata ++ and transport round trips per batch on the writer rank and serialized ++ the response behind them. ++ ``batch_put_from`` preserves the existing per-feature keys and raw ++ tensor layout while amortizing that fixed cost across the whole ++ scheduler batch. The response is still emitted only after every ++ status succeeds, so refs can never point at incomplete samples. + """ -+ store_id = str(spec["store_id"]) -+ sample_id = str(spec["sample_id"]) -+ gen = int(spec.get("gen", 1)) -+ replace = bool(spec.get("replace", False)) -+ features: Dict[str, str] = dict(spec.get("features") or {}) ++ if not samples: ++ return [] + -+ written: List[str] = [] -+ result_feats: Dict[str, Dict[str, Any]] = {} ++ store = self._connect() ++ timing_enabled = os.environ.get("SGLANG_SPEC_CAPTURE_TIMING", "0") == "1" ++ started = time.perf_counter() ++ keys: List[str] = [] ++ tensors: List[torch.Tensor] = [] ++ sizes: List[int] = [] ++ replace_keys: List[str] = [] ++ results: List[Dict[str, Any]] = [] + -+ def _write(name: str, t: torch.Tensor) -> None: ++ def _stage( ++ result_feats: Dict[str, Dict[str, Any]], ++ *, ++ store_id: str, ++ sample_id: str, ++ gen: int, ++ replace: bool, ++ name: str, ++ tensor: torch.Tensor, ++ ) -> None: ++ tensor = tensor.detach().to("cpu").contiguous() + key = self._tkey(store_id, sample_id, gen, name) -+ self._put_tensor(key, t, replace=replace) -+ written.append(key) ++ keys.append(key) ++ tensors.append(tensor) ++ sizes.append(tensor.element_size() * tensor.numel()) ++ if replace: ++ replace_keys.append(key) + result_feats[name] = { -+ "shape": list(t.shape), -+ "dtype": _DTYPE_STR.get(t.dtype, str(t.dtype).replace("torch.", "")), ++ "shape": list(tensor.shape), ++ "dtype": _DTYPE_STR.get( ++ tensor.dtype, str(tensor.dtype).replace("torch.", "") ++ ), + } + -+ try: ++ for spec, aux, last_hidden in samples: ++ store_id = str(spec["store_id"]) ++ sample_id = str(spec["sample_id"]) ++ gen = int(spec.get("gen", 1)) ++ replace = bool(spec.get("replace", False)) ++ features: Dict[str, str] = dict(spec.get("features") or {}) ++ result_feats: Dict[str, Dict[str, Any]] = {} ++ + aux_name = features.get(_ARTIFACT_AUX) + if aux_name is not None: + if aux is None: + raise RuntimeError( + "spec_capture requested 'aux' but no aux hidden states were " -+ "captured — launch the server with --enable-spec-capture " ++ "captured -- launch the server with --enable-spec-capture " + "(and optionally --spec-capture-aux-layer-ids)" + ) -+ _write(aux_name, aux.unsqueeze(0)) ++ _stage( ++ result_feats, ++ store_id=store_id, ++ sample_id=sample_id, ++ gen=gen, ++ replace=replace, ++ name=aux_name, ++ tensor=aux.unsqueeze(0), ++ ) + lh_name = features.get(_ARTIFACT_LAST_HIDDEN) + if lh_name is not None: + if last_hidden is None: @@ -603,15 +930,15 @@ index 0000000000..a678bd85e7 + "spec_capture requested 'last_hidden' but the logits " + "processor did not return it (is aux capture enabled?)" + ) -+ _write(lh_name, last_hidden.unsqueeze(0)) -+ pos_name = features.get(_ARTIFACT_POSITION_IDS) -+ if pos_name is not None: -+ if position_ids is None: -+ raise RuntimeError( -+ "spec_capture requested 'position_ids' but no prompt " -+ "positions were collected on the scheduler" -+ ) -+ _write(pos_name, position_ids.unsqueeze(0)) ++ _stage( ++ result_feats, ++ store_id=store_id, ++ sample_id=sample_id, ++ gen=gen, ++ replace=replace, ++ name=lh_name, ++ tensor=last_hidden.unsqueeze(0), ++ ) + for item in spec.get("passthrough") or []: + dtype = _STR_DTYPE.get(str(item.get("dtype", "int64"))) + if dtype is None: @@ -619,22 +946,116 @@ index 0000000000..a678bd85e7 + f"spec_capture passthrough {item.get('name')!r}: " + f"unsupported dtype {item.get('dtype')!r}" + ) -+ t = torch.tensor(item["data"], dtype=dtype).reshape( ++ tensor = torch.tensor(item["data"], dtype=dtype).reshape( + [int(d) for d in item["shape"]] + ) -+ _write(str(item["name"]), t) ++ _stage( ++ result_feats, ++ store_id=store_id, ++ sample_id=sample_id, ++ gen=gen, ++ replace=replace, ++ name=str(item["name"]), ++ tensor=tensor, ++ ) ++ results.append( ++ { ++ "sample_id": sample_id, ++ "store_id": store_id, ++ "gen": gen, ++ "aux_layer_ids": self.aux_layer_ids, ++ "features": result_feats, ++ } ++ ) ++ ++ materialize_ms = (time.perf_counter() - started) * 1000.0 ++ self._remove_many_quiet(replace_keys) ++ registered: List[torch.Tensor] = [] ++ register_started = time.perf_counter() ++ try: ++ for tensor, nbytes in zip(tensors, sizes): ++ try: ++ store.register_buffer(tensor.data_ptr(), nbytes) ++ registered.append(tensor) ++ except Exception: ++ pass # TCP and some Mooncake builds auto-register ++ register_ms = (time.perf_counter() - register_started) * 1000.0 ++ put_started = time.perf_counter() ++ batch_put = getattr(store, "batch_put_from", None) ++ if batch_put is None: ++ statuses = [ ++ store.put_from(key, tensor.data_ptr(), nbytes, self._put_config) ++ for key, tensor, nbytes in zip(keys, tensors, sizes) ++ ] ++ else: ++ statuses = batch_put( ++ keys, ++ [tensor.data_ptr() for tensor in tensors], ++ sizes, ++ self._put_config, ++ ) ++ put_ms = (time.perf_counter() - put_started) * 1000.0 + except Exception: -+ for key in written: -+ self._remove_quiet(key) ++ self._remove_many_quiet(keys) + raise ++ finally: ++ for tensor in registered: ++ try: ++ store.unregister_buffer(tensor.data_ptr()) ++ except Exception: ++ pass + -+ return { -+ "sample_id": sample_id, -+ "store_id": store_id, -+ "gen": gen, -+ "aux_layer_ids": self.aux_layer_ids, -+ "features": result_feats, -+ } ++ if statuses is None: ++ statuses = [0] * len(keys) ++ if len(statuses) != len(keys): ++ self._remove_many_quiet(keys) ++ raise RuntimeError( ++ "spec-capture batch_put_from returned " ++ f"{len(statuses)} statuses for {len(keys)} keys" ++ ) ++ failed = [ ++ (key, status) ++ for key, status in zip(keys, statuses) ++ if status is not None and int(status) < 0 ++ ] ++ if failed: ++ self._remove_many_quiet(keys) ++ raise RuntimeError( ++ "spec-capture batch_put_from failed for " ++ f"{len(failed)}/{len(keys)} keys; first={failed[0]}" ++ ) ++ ++ if timing_enabled: ++ logger.info( ++ "[spec-capture-timing] batch_sink samples=%d objects=%d " ++ "bytes=%d materialize_ms=%.3f register_ms=%.3f put_ms=%.3f " ++ "total_ms=%.3f", ++ len(samples), ++ len(keys), ++ sum(sizes), ++ materialize_ms, ++ register_ms, ++ put_ms, ++ (time.perf_counter() - started) * 1000.0, ++ ) ++ return results ++ ++ def put_sample( ++ self, ++ spec: Dict[str, Any], ++ *, ++ aux: Optional[torch.Tensor], ++ last_hidden: Optional[torch.Tensor], ++ ) -> Dict[str, Any]: ++ """Write one sample's artifacts; return the meta_info result dict. ++ ++ ``aux``/``last_hidden`` are the per-request (L, W) captured tensors, ++ stored with a leading batch dim of 1. On any failure the keys already ++ written are best-effort removed (no partial sample is consumable). ++ """ ++ # Keep the single-sample entry point for compatibility with tests and ++ # callers outside the scheduler; production uses put_samples(). ++ return self.put_samples([(spec, aux, last_hidden)])[0] + + +_SINK: Optional[SpecCaptureSink] = None diff --git a/pyproject.toml b/pyproject.toml index 9d8097322..ff68a1949 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta" name = "specforge" dynamic = ["version"] readme = "README.md" -requires-python = ">=3.11" +requires-python = ">=3.10" description = "SpecForge: Speculative Decoding Training Framework" authors = [{name = "SGLang Team"}] urls = {Homepage = "https://github.com/sgl-project/SpecForge"} @@ -35,7 +35,7 @@ dependencies = [ specforge = "specforge.cli:main" [tool.setuptools.packages.find] -exclude = ["configs*", "scripts*", "tests*"] +include = ["specforge*"] [project.optional-dependencies] data = ["openai"] diff --git a/scripts/apply_sglang_spec_capture_patch.sh b/scripts/apply_sglang_spec_capture_patch.sh index cb59afced..1f2594d6b 100755 --- a/scripts/apply_sglang_spec_capture_patch.sh +++ b/scripts/apply_sglang_spec_capture_patch.sh @@ -12,31 +12,78 @@ # when a reverse dry-run proves it matches the current patch byte-for-byte; # anything else fails loudly rather than testing against unknown server code. # -# Usage: scripts/apply_sglang_spec_capture_patch.sh [--reverse] +# Usage: scripts/apply_sglang_spec_capture_patch.sh +# [--target v0.5.14|kimi-k3-ee560a2|kimi-k3-9acd9cb|kimi-k3-f8493a4] +# [--reverse] set -euo pipefail HERE="$(cd "$(dirname "$0")/.." && pwd)" -PATCH="$HERE/patches/sglang/v0.5.14/spec-capture.patch" +TARGET="v0.5.14" +PATCH_TARGET="" +REVERSE=0 +while [[ $# -gt 0 ]]; do + case "$1" in + --target) + if [[ $# -lt 2 ]]; then + echo "ERROR: --target requires a value" >&2 + exit 2 + fi + TARGET="$2" + shift 2 + ;; + --reverse) + REVERSE=1 + shift + ;; + *) + echo "ERROR: unknown argument: $1" >&2 + exit 2 + ;; + esac +done + +case "$TARGET" in + v0.5.14) + EXPECTED_VERSION_PREFIX="0.5.14" + PATCH_TARGET="$TARGET" + ;; + kimi-k3-ee560a2|kimi-k3-9acd9cb|kimi-k3-f8493a4) + # Kimi K3's SGLang fork currently reports a base-package version that + # does not uniquely identify this source revision, so patch --check is + # the authoritative compatibility gate below. + EXPECTED_VERSION_PREFIX="" + # One patch is generated against ee560a2 and compatibility-checked + # against the original f8493a4 integration point and the 9acd9cb tip. + # Keep the historical directory name so existing automation remains + # source-compatible. + PATCH_TARGET="kimi-k3-f8493a4" + ;; + *) + echo "ERROR: unsupported SGLang patch target: $TARGET" >&2 + exit 2 + ;; +esac +PATCH="$HERE/patches/sglang/$PATCH_TARGET/spec-capture.patch" SGL_PARENT="$(python -c 'import sglang, os; print(os.path.dirname(os.path.dirname(sglang.__file__)))')" SGL_VERSION="$(python -c 'import sglang; print(sglang.__version__)')" APPLIED_COPY="$SGL_PARENT/sglang/.spec_capture_patch.applied" SINK="$SGL_PARENT/sglang/srt/spec_capture_sink.py" -if [[ "$SGL_VERSION" != 0.5.14* ]]; then - echo "WARNING: installed sglang is $SGL_VERSION; the patch targets v0.5.14" >&2 +if [[ -n "$EXPECTED_VERSION_PREFIX" && "$SGL_VERSION" != "$EXPECTED_VERSION_PREFIX"* ]]; then + echo "WARNING: installed sglang is $SGL_VERSION; the patch targets $TARGET" >&2 fi -if [[ "${1:-}" == "--reverse" ]]; then +if [[ "$REVERSE" == 1 ]]; then patch --reverse -p2 --batch -N -d "$SGL_PARENT" < "$PATCH" rm -f "$APPLIED_COPY" - echo "spec-capture patch --reverse at $SGL_PARENT/sglang (sglang $SGL_VERSION)" + echo "spec-capture patch $TARGET --reverse at $SGL_PARENT/sglang (sglang $SGL_VERSION)" exit 0 fi if [[ -f "$APPLIED_COPY" ]]; then if cmp -s "$APPLIED_COPY" "$PATCH"; then - echo "spec-capture patch already applied at $SGL_PARENT/sglang" + echo "spec-capture patch $TARGET already applied at $SGL_PARENT/sglang" exit 0 fi echo "spec-capture patch changed; reversing the recorded version first" @@ -52,7 +99,7 @@ elif [[ -f "$SINK" ]]; then fi if matches; then cp "$PATCH" "$APPLIED_COPY" - echo "spec-capture patch already applied at $SGL_PARENT/sglang (adopted)" + echo "spec-capture patch $TARGET already applied at $SGL_PARENT/sglang (adopted)" exit 0 fi echo "ERROR: $SGL_PARENT/sglang carries an unknown spec-capture patch state" >&2 @@ -62,4 +109,4 @@ fi patch -p2 --batch -N -d "$SGL_PARENT" < "$PATCH" cp "$PATCH" "$APPLIED_COPY" -echo "spec-capture patch applied at $SGL_PARENT/sglang (sglang $SGL_VERSION)" +echo "spec-capture patch $TARGET applied at $SGL_PARENT/sglang (sglang $SGL_VERSION)" diff --git a/scripts/gates/README.md b/scripts/gates/README.md index a5d4875e6..7a1465589 100644 --- a/scripts/gates/README.md +++ b/scripts/gates/README.md @@ -1,10 +1,10 @@ # One-sample overfit validation -This guide contains only two stages: +This guide contains three stages: 1. regenerate a small subset of datasets by target model; -1. overfit one sample on Qwen3.6-27B Dspark training with `specforge train`; -2. export the checkpoint, serve it with SGLang, and verify that one complete +2. overfit one sample on Qwen3.6-27B Dspark training with `specforge train`; +3. export the checkpoint, serve it with SGLang, and verify that one complete 16-token draft block is accepted. The commands below use GPU 0 for target capture and GPU 1 for training and @@ -28,7 +28,7 @@ gate_report_tcp_ports python 127.0.0.1 \ The check always returns to the shell. If one or more ports are occupied, it prints every conflict without stopping or killing the owning processes. -s + ## Prepare one sample ### Stage 1: Regen the datasets Launch the sglang server: @@ -66,7 +66,7 @@ export MODEL_NAME=Qwen3.6-27B export SPEC_METHOD=Dspark export DRAFT_MODEL_CONFIG=configs/qwen3.6-27b-dspark.json export MODEL=Qwen/Qwen3.6-27B -export TRAINING_CONFIG=examples/configs/qwen3.6-27b-dspark-disaggregated.yaml +export TRAINING_CONFIG=examples/configs/online/disaggregated/managed-local/qwen3.6-27b-dspark-disaggregated.yaml specforge train \ --config ${TRAINING_CONFIG} \ @@ -117,7 +117,7 @@ The `control` and `consumer-state` directories must be fresh. To repeat the experiment, use a new suffix consistently in `run_id`, `output_dir`, `control_dir`, `consumer_state_dir`, and the sample paths. -### Stage 2: serve with SGLang and check accept length +### Stage 3: serve with SGLang and check accept length First export the trained draft: @@ -198,5 +198,5 @@ The validation passes only when the result contains: ``` `spec_accept_length >= 16` and `target_prefix_match_tokens >= 16` mean that the -complete DFlash block was accepted and agrees with the target continuation. +complete DSpark block was accepted and agrees with the target continuation. Stop the SGLang server with `Ctrl-C` after validation. diff --git a/scripts/merge_mtp_to_base.py b/scripts/merge_mtp_to_base.py new file mode 100755 index 000000000..ca695b9c2 --- /dev/null +++ b/scripts/merge_mtp_to_base.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 +# coding=utf-8 +"""Merge a trained MTP draft checkpoint back into the base target model. + +Thin CLI wrapper around ``specforge.export.mtp.merge_mtp_into_base``; the merge +logic and the per-family key mapping live in the package. + +Example: + python scripts/merge_mtp_to_base.py \ + --base-model-path PATH/TO/Qwen3.5-4B \ + --mtp-checkpoint-path PATH/TO/outputs/qwen3.5-4b-mtp/RUN-latest \ + --draft-config configs/qwen3.5-4b-mtp.json \ + --output-path PATH/TO/Qwen3.5-4B-MTP \ + --key-format sglang +""" + +import argparse + +from specforge.export.mtp import merge_mtp_into_base + + +def main(): + parser = argparse.ArgumentParser( + description="Merge trained MTP weights back into the base Qwen3.5 model." + ) + parser.add_argument( + "--base-model-path", + type=str, + required=True, + help="Path to the original Qwen3.5 base model checkpoint.", + ) + parser.add_argument( + "--mtp-checkpoint-path", + type=str, + required=True, + help=( + "SpecForge runtime checkpoint/output path, or an already-exported " + "HF MTP draft directory." + ), + ) + parser.add_argument( + "--draft-config", + type=str, + default=None, + help=( + "Draft config JSON (required for a SpecForge runtime checkpoint; " + "the exported HF directory already contains config.json)." + ), + ) + parser.add_argument( + "--output-path", + type=str, + required=True, + help="Directory to write the merged checkpoint.", + ) + parser.add_argument( + "--key-format", + type=str, + default="sglang", + choices=["sglang", "hf"], + help=( + "MTP key layout. Both 'sglang' and 'hf' produce the flat native " + "layout (mtp.layers.0.* / mtp.norm.weight) that SGLang's flat " + "Qwen3_5ForCausalLMMTP and HF/vLLM MTP modules expect; the argument " + "is kept for backward compatibility." + ), + ) + args = parser.parse_args() + + merge_mtp_into_base( + args.base_model_path, + args.mtp_checkpoint_path, + args.output_path, + args.key_format, + draft_config_path=args.draft_config, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/prepare_hidden_states.py b/scripts/prepare_hidden_states.py index 47dbb06f2..109c0142f 100644 --- a/scripts/prepare_hidden_states.py +++ b/scripts/prepare_hidden_states.py @@ -44,13 +44,12 @@ import gc import gzip import hashlib -import json import os import uuid from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from pathlib import Path -from typing import Dict, List, Mapping, Optional +from typing import Callable, Dict, List, Mapping, Optional import torch import torch.distributed as dist @@ -92,6 +91,7 @@ class OfflineCapturePlan: capture_method: str capture_layers: tuple[int, ...] layout: OfflineCaptureLayout + loss_mask_filter: Optional[Callable[[object], bool]] def parse_args(): @@ -192,34 +192,30 @@ def parse_args(): sglang_group.add_argument("--sglang-enable-dp-attention", action="store_true") sglang_group.add_argument("--sglang-enable-dp-lm-head", action="store_true") sglang_group.add_argument("--sglang-ep-size", type=int, default=1) + sglang_group.add_argument( + "--sglang-disable-radix-cache", + action="store_true", + help=( + "Disable the SGLang radix (prefix) cache (default: enabled). " + "Required for hybrid linear-attention/Mamba targets on ROCm, whose " + "mamba radix-cache extra_buffer strategy asserts CUDA/MUSA/NPU " + "(FLA) at server init." + ), + ) return parser.parse_args() -def _resolve_draft_vocab_size(source: str) -> int: - """Load ``draft_vocab_size`` from one existing local JSON file.""" +def _resolve_draft_vocab_size(draft_config: object) -> int: + """Read the vocabulary size from the canonical resolved draft config.""" - expanded = Path(source).expanduser() - if not expanded.is_file(): - raise FileNotFoundError( - "--draft-model-config must point to an existing local JSON file: " - f"{source}" - ) - if expanded.suffix.lower() != ".json": - raise ValueError( - f"--draft-model-config must point to a local .json file: {source}" - ) - try: - with expanded.open(encoding="utf-8") as stream: - payload = json.load(stream) - except json.JSONDecodeError as exc: - raise ValueError(f"invalid draft config JSON {expanded}: {exc}") from exc - if not isinstance(payload, dict): - raise ValueError(f"draft config JSON {expanded} must contain an object") - value = payload.get("draft_vocab_size", payload.get("vocab_size")) + missing = object() + value = getattr(draft_config, "draft_vocab_size", missing) + if value is missing: + value = getattr(draft_config, "vocab_size", None) if isinstance(value, bool) or not isinstance(value, int) or value <= 0: raise ValueError( - f"draft model config {source!r} must define a positive " + "resolved draft model config must define a positive " f"draft_vocab_size (or vocab_size fallback), got {value!r}" ) return value @@ -308,6 +304,7 @@ def _sglang_kwargs(args: argparse.Namespace) -> Dict[str, object]: "enable_dp_attention": args.sglang_enable_dp_attention, "enable_dp_lm_head": args.sglang_enable_dp_lm_head, "ep_size": args.sglang_ep_size, + "disable_radix_cache": getattr(args, "sglang_disable_radix_cache", False), "max_running_requests": args.batch_size, "max_total_tokens": args.batch_size * args.max_length, } @@ -341,6 +338,7 @@ def resolve_offline_capture_plan( capture_method=resolved.capture_method, capture_layers=resolved.capture_layers, layout=resolved.layout, + loss_mask_filter=resolved.loss_mask_filter, ) @@ -773,7 +771,7 @@ def main(): trust_remote_code=args.trust_remote_code, ) capture_plan = resolve_offline_capture_plan(args, target_model_config) - draft_vocab_size = _resolve_draft_vocab_size(args.draft_model_config) + draft_vocab_size = _resolve_draft_vocab_size(capture_plan.draft_config) # Initialize distributed environment (TP + DP) init_distributed(timeout=args.dist_timeout, tp_size=args.tp_size) @@ -826,7 +824,7 @@ def main(): ), num_proc=min(args.build_dataset_num_proc, 32), ) - if args.num_samples is not None: + if args.num_samples is not None and capture_plan.loss_mask_filter is None: dataset = dataset.select(range(args.num_samples)) # Tokenizer and cache key tokenizer = load_tokenizer( @@ -847,6 +845,15 @@ def main(): cache_key=cache_key, is_preformatted=args.is_preformatted, num_proc=args.build_dataset_num_proc, + loss_mask_filter=capture_plan.loss_mask_filter, + ) + if capture_plan.loss_mask_filter is not None and args.num_samples is not None: + eagle3_dataset = eagle3_dataset.select( + range(min(args.num_samples, len(eagle3_dataset))) + ) + if not len(eagle3_dataset): + raise ValueError( + f"no samples satisfy {capture_plan.strategy} training eligibility" ) print_with_rank(f"Dataset prepared with {len(eagle3_dataset)} samples.") diff --git a/scripts/regenerate_train_data.py b/scripts/regenerate_train_data.py index f2f737fd2..252c989b2 100644 --- a/scripts/regenerate_train_data.py +++ b/scripts/regenerate_train_data.py @@ -244,6 +244,29 @@ def build_query_kwargs(args, messages, max_tokens=None): return query_kwargs +def _extract_record_images(data: Dict[str, Any]) -> List[str]: + """Resolve image references of a record (``image``/``image_path`` string or + ``images`` list), in insertion order.""" + refs: List[str] = [] + single = data.get("image") or data.get("image_path") + if isinstance(single, str): + refs.append(single) + images = data.get("images") + if isinstance(images, list): + refs.extend(r for r in images if isinstance(r, str)) + return refs + + +def _image_url_part(path: str) -> Dict[str, Any]: + import base64 + import mimetypes + + mime = mimetypes.guess_type(path)[0] or "image/jpeg" + with open(path, "rb") as f: + b64 = base64.b64encode(f.read()).decode("ascii") + return {"type": "image_url", "image_url": {"url": f"data:{mime};base64,{b64}"}} + + def call_sglang( args, server_address: str, @@ -260,6 +283,8 @@ def call_sglang( messages = data["conversations"] regenerated_messages = [] + record_images = _extract_record_images(data) + image_attached = False # ignore data which starts with an assistant message if messages[0]["role"] == "assistant": @@ -273,6 +298,26 @@ def call_sglang( elif message["role"] == "assistant": continue elif message["role"] == "user": + # Multimodal records: attach the record's images to the first user + # turn that carries the placeholder (OpenAI content parts). + content = message.get("content") + if ( + record_images + and not image_attached + and isinstance(content, str) + and "" in content + ): + try: + parts = [_image_url_part(p) for p in record_images] + except OSError as exc: + data["status"] = "error" + data["error"] = f"unreadable image file: {exc}" + return data + text = content.replace("\n", "").replace("", "") + parts.append({"type": "text", "text": text}) + message = dict(message) + message["content"] = parts + image_attached = True regenerated_messages.append(message) query_kwargs = build_query_kwargs(args, regenerated_messages, max_tokens) diff --git a/specforge/algorithms/builtin.py b/specforge/algorithms/builtin.py index c2d8527bd..8e4ea7e9a 100644 --- a/specforge/algorithms/builtin.py +++ b/specforge/algorithms/builtin.py @@ -6,6 +6,7 @@ from specforge.algorithms.domino.providers import create_registration as domino from specforge.algorithms.dspark.providers import create_registration as dspark from specforge.algorithms.eagle3.providers import create_registration as eagle3 +from specforge.algorithms.mtp.providers import create_registration as mtp from specforge.algorithms.peagle.providers import create_registration as peagle from specforge.algorithms.registry import AlgorithmRegistry @@ -13,7 +14,7 @@ def builtin_algorithm_registry() -> AlgorithmRegistry: """Return a fresh immutable catalog without module-level mutation.""" - return AlgorithmRegistry((eagle3(), peagle(), dflash(), domino(), dspark())) + return AlgorithmRegistry((eagle3(), peagle(), dflash(), domino(), dspark(), mtp())) __all__ = ["builtin_algorithm_registry"] diff --git a/specforge/algorithms/common/dflash_family_model.py b/specforge/algorithms/common/dflash_family_model.py index 9c0e09fb7..12319633f 100644 --- a/specforge/algorithms/common/dflash_family_model.py +++ b/specforge/algorithms/common/dflash_family_model.py @@ -1,6 +1,7 @@ # coding=utf-8 """DFlash-family training models and shared masking helpers.""" +import os from typing import Dict, Optional, Tuple import torch @@ -9,6 +10,7 @@ from specforge.core.chunking import checkpointed_chunk_reduce from specforge.modeling.draft.dflash import DFlashDraftModel +from specforge.modeling.draft.flex_attention_backend import flex_attention_backend try: from torch.nn.attention.flex_attention import BlockMask, create_block_mask @@ -44,7 +46,18 @@ def compute_accept_len( return accept_prefix.sum(dim=2).float() -def create_dflash_sdpa_mask(anchor_positions, block_keep_mask, S, block_size, device): +def create_dflash_sdpa_mask( + anchor_positions, + block_keep_mask, + S, + block_size, + device, + sliding_window: Optional[int] = None, +): + """Construct a full or sliding dense boolean DFlash mask.""" + + if sliding_window is not None and sliding_window <= 0: + raise ValueError("sliding_window must be > 0") B, N = anchor_positions.shape Q_LEN = N * block_size KV_LEN = S + N * block_size @@ -55,16 +68,24 @@ def create_dflash_sdpa_mask(anchor_positions, block_keep_mask, S, block_size, de ) # (1, 1, 1, KV_LEN) q_block_ids = q_indices // block_size + q_block_offsets = q_indices % block_size anchor_expanded = anchor_positions.view(B, 1, N, 1).repeat_interleave( block_size, dim=2 ) mask_context = (kv_indices < S) & (kv_indices < anchor_expanded) + if sliding_window is not None: + # The current draft token occupies one slot in the window. + context_lower_bound = anchor_expanded + q_block_offsets - (sliding_window - 1) + mask_context = mask_context & (kv_indices >= context_lower_bound) is_draft = kv_indices >= S kv_block_ids = (kv_indices - S) // block_size mask_draft = is_draft & (q_block_ids == kv_block_ids) + if sliding_window is not None: + kv_block_offsets = (kv_indices - S) % block_size + mask_draft = mask_draft & (kv_block_offsets <= q_block_offsets) valid_block = block_keep_mask.view(B, 1, N, 1).repeat_interleave(block_size, dim=2) @@ -78,21 +99,17 @@ def create_dflash_block_mask( S: int, block_size: int, device: torch.device, + flex_block_size=None, + sliding_window: Optional[int] = None, ): - """Construct Flex Attention BlockMask for DFlash training. + """Construct a full or sliding Flex Attention mask for DFlash training.""" - KV: [Context (S tokens) | Block_0 | Block_1 | ... | Block_{n-1}] - Q: [Block_0 | Block_1 | ... | Block_{n-1}] - - Rules: - 1. Each block sees context strictly before its anchor (kv_idx < anchor_pos). - 2. Intra-block attention is bidirectional. - 3. Different blocks are invisible to each other. - 4. Invalid blocks (block_keep_mask=False) see nothing. - """ + if sliding_window is not None and sliding_window <= 0: + raise ValueError("sliding_window must be > 0") def dflash_mask_mod(b, h, q_idx, kv_idx): q_block_id = q_idx // block_size + q_block_offset = q_idx % block_size safe_q_block_id = q_block_id.clamp(max=N - 1) anchor_pos = anchor_positions[b, safe_q_block_id] @@ -100,10 +117,17 @@ def dflash_mask_mod(b, h, q_idx, kv_idx): # Strictly less than: matches inference where target_hidden[anchor_pos] # is not available as context. mask_context = is_context & (kv_idx < anchor_pos) + if sliding_window is not None: + # The current draft token occupies one slot in the window. + context_lower_bound = anchor_pos + q_block_offset - (sliding_window - 1) + mask_context = mask_context & (kv_idx >= context_lower_bound) is_draft = kv_idx >= S kv_block_id = (kv_idx - S) // block_size mask_draft = is_draft & (q_block_id == kv_block_id) + if sliding_window is not None: + kv_block_offset = (kv_idx - S) % block_size + mask_draft = mask_draft & (kv_block_offset <= q_block_offset) is_valid_block = block_keep_mask[b, safe_q_block_id] in_bounds = q_block_id < N @@ -113,8 +137,17 @@ def dflash_mask_mod(b, h, q_idx, kv_idx): Q_LEN = N * block_size KV_LEN = S + N * block_size + kwargs = {} + if flex_block_size is not None: + kwargs["BLOCK_SIZE"] = flex_block_size return create_block_mask( - dflash_mask_mod, B=B, H=None, Q_LEN=Q_LEN, KV_LEN=KV_LEN, device=device + dflash_mask_mod, + B=B, + H=None, + Q_LEN=Q_LEN, + KV_LEN=KV_LEN, + device=device, + **kwargs, ) @@ -162,42 +195,46 @@ def __init__( self._cached_bsz: Optional[int] = None def _sample_anchor_positions( - self, seq_len: int, loss_mask: torch.Tensor, device: torch.device + self, + seq_len: int, + loss_mask: torch.Tensor, + device: torch.device, + max_valid_anchors: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - """Randomly sample anchor positions per sample; returns (anchors, keep_mask).""" - bs = self.block_size - bsz = loss_mask.shape[0] - max_anchor = max(seq_len - bs, 0) + """Sample anchors whose clean token and first target are supervised.""" - valid = loss_mask[:, : max_anchor + 1] > 0.5 - valid_counts = valid.sum(dim=1) - max_n = min(self.num_anchors, int(valid_counts.max().item()) - 1) - - if max_n <= 0: - raise ValueError("should preprocess the data.") - - indices = ( - torch.arange(max_anchor + 1, device=device).unsqueeze(0).expand(bsz, -1) - ) - masked_indices = torch.where( - valid, indices, torch.tensor(seq_len + 1, device=device) + num_candidates = max(seq_len - 1, 0) + valid = (loss_mask[:, :num_candidates] > 0.5) & ( + loss_mask[:, 1 : num_candidates + 1] > 0.5 ) + valid_counts = valid.sum(dim=1) + if max_valid_anchors is None: + # Direct model callers may supply an already-device-resident mask. + # Training strategies pass the CPU-computed value and avoid this + # synchronizing fallback on CUDA. + max_valid_anchors = int(valid_counts.max().item()) + width = min(self.num_anchors, max(0, int(max_valid_anchors))) + if width == 0: + raise ValueError( + "DFlash-family training requires two consecutive supervised tokens" + ) - random_vals = torch.rand(bsz, max_anchor + 1, device=device) - random_vals = torch.where(valid, random_vals, torch.tensor(2.0, device=device)) - - _, sorted_idx = random_vals.sort(dim=1) - gathered = torch.gather(masked_indices, 1, sorted_idx) - anchors = gathered[:, :max_n].sort(dim=1).values - - keep_mask = torch.arange(max_n, device=device).unsqueeze( + random_values = torch.rand(valid.shape, device=device) + random_values.masked_fill_(~valid, 2.0) + candidates = random_values.argsort(dim=1)[:, :width] + keep_mask = torch.arange(width, device=device).unsqueeze( 0 - ) < valid_counts.unsqueeze(1).clamp(max=max_n) + ) < valid_counts.clamp(max=width).unsqueeze(1) + + sentinel = valid.shape[1] anchors = torch.where( - keep_mask, anchors, torch.tensor(0, dtype=torch.long, device=device) + keep_mask, + candidates, + torch.full_like(candidates, sentinel), ) - - return anchors, keep_mask + anchors = anchors.sort(dim=1).values + keep_mask = anchors < sentinel + return torch.where(keep_mask, anchors, 0), keep_mask def _create_position_ids(self, anchor_positions: torch.Tensor) -> torch.Tensor: """Create absolute position IDs for parallel draft blocks.""" @@ -269,68 +306,74 @@ def _forward_draft_blocks( input_ids: torch.Tensor, hidden_states: torch.Tensor, loss_mask: torch.Tensor, - position_ids: Optional[torch.Tensor] = None, + max_valid_anchors: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: bsz, seq_len = input_ids.shape device = input_ids.device anchor_positions, block_keep_mask = self._sample_anchor_positions( - seq_len, loss_mask, device + seq_len, + loss_mask, + device, + max_valid_anchors=max_valid_anchors, ) noise_embedding = self._create_noise_embed( input_ids, anchor_positions, block_keep_mask ) - if position_ids is None: - context_position_ids = ( - torch.arange(seq_len, device=device).unsqueeze(0).expand(bsz, -1) - ) - draft_position_ids = self._create_position_ids(anchor_positions) - full_position_ids = torch.cat( - [context_position_ids, draft_position_ids], dim=1 - ) - else: - if not getattr(self.draft_model, "use_interleaved_mrope", False): - raise ValueError( - "multimodal capture carries mRoPE position_ids, but the " - "draft config does not enable rope_scaling.mrope_interleaved; " - "use a VLM draft config (e.g. configs/*-dflash-vlm-*.json)" - ) - # Server-produced mRoPE positions, (B, S, 3) -> (3, B, S + N*bs). - offsets = torch.arange(self.block_size, device=device).view(1, 1, -1) - draft_indices = (anchor_positions.unsqueeze(-1) + offsets).view(bsz, -1) - draft_position_ids = torch.gather( - position_ids, - 1, - draft_indices.unsqueeze(-1).expand(-1, -1, 3), - ) - full_position_ids = torch.cat( - [position_ids, draft_position_ids], dim=1 - ).permute(2, 0, 1) + context_position_ids = ( + torch.arange(seq_len, device=device).unsqueeze(0).expand(bsz, -1) + ) + draft_position_ids = self._create_position_ids(anchor_positions) + full_position_ids = torch.cat( + [context_position_ids, draft_position_ids], dim=1 + ) + mask_builder = ( + create_dflash_block_mask + if self.attention_backend == "flex_attention" + else create_dflash_sdpa_mask + ) + mask_args = { + "anchor_positions": anchor_positions, + "block_keep_mask": block_keep_mask, + "S": seq_len, + "block_size": self.block_size, + "device": device, + } + if ( + self.attention_backend == "flex_attention" + and flex_attention_backend() == "FLASH" + ): + # FLASH requires a minimum of this block size. + mask_args["flex_block_size"] = (256, 128) + full_attn_mask = mask_builder(**mask_args) + sliding_window = self.draft_model.sliding_window + dflash_attn_mask = full_attn_mask + if sliding_window is not None: + dflash_attn_mask = { + "full_attention": full_attn_mask, + "sliding_attention": mask_builder( + **mask_args, + sliding_window=sliding_window, + ), + } + + draft_kwargs = {} if self.attention_backend == "flex_attention": - dflash_attn_mask = create_dflash_block_mask( - anchor_positions=anchor_positions, - block_keep_mask=block_keep_mask, - S=seq_len, - block_size=self.block_size, - device=device, - ) - else: - dflash_attn_mask = create_dflash_sdpa_mask( - anchor_positions=anchor_positions, - block_keep_mask=block_keep_mask, - S=seq_len, - block_size=self.block_size, - device=device, - ) - + # DFlash's dynamic short-query batches are training/prefill shaped, + # not autoregressive decoding. AUTO may route q_len < 128 to the + # more restrictive flex-decoding kernel, whose config set can be + # empty for DFlash's sparse BlockMask. Keep the general Triton + # Flex Attention kernel for every DFlash-family batch. + draft_kwargs["kernel_options"] = {"BACKEND": "TRITON"} output_hidden = self.draft_model( position_ids=full_position_ids, noise_embedding=noise_embedding, target_hidden=hidden_states, attention_mask=dflash_attn_mask, + **draft_kwargs, ) return anchor_positions, block_keep_mask, output_hidden @@ -392,14 +435,10 @@ def forward( input_ids: torch.Tensor, hidden_states: torch.Tensor, loss_mask: torch.Tensor, - position_ids: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor]]: + max_valid_anchors: Optional[int] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, object]]: """Parallel block-wise training forward pass; returns - (loss, accuracy, metrics) — same shape as Domino's forward. - - ``position_ids`` is the optional server-captured mRoPE position tensor - ``(B, S, 3)`` for multimodal runs; text runs leave it None and use - internally synthesized flat positions. + (loss, accuracy, metrics) - same shape as Domino's forward. """ if self.attention_backend == "flex_attention" and not FLEX_ATTENTION_AVAILABLE: raise ValueError( @@ -412,7 +451,7 @@ def forward( input_ids=input_ids, hidden_states=hidden_states, loss_mask=loss_mask, - position_ids=position_ids, + max_valid_anchors=max_valid_anchors, ) # --- Labels: same-position prediction (position k predicts token anchor+k) --- @@ -457,13 +496,20 @@ def forward( chunk_size=self.objective_chunk_blocks, dim=1, ) - if self.loss_type == "dflash": - loss = loss_num / (loss_den + 1e-6) - else: - loss = loss_num / float(bsz) - accuracy = correct_num / (accuracy_denom + 1e-6) - - return loss, accuracy, {"accuracy_denom": accuracy_denom.detach()} + ratio_metrics = { + "acc": (correct_num.detach(), accuracy_denom.detach()), + } + metrics: Dict[str, object] = { + "accuracy_denom": accuracy_denom.detach(), + "ratio_metrics": ratio_metrics, + } + loss_denominator = ( + loss_den if self.loss_type == "dflash" else loss_num.new_tensor(float(bsz)) + ) + loss = loss_num / loss_denominator + metrics["loss_terms"] = (loss_num, loss_denominator.detach()) + accuracy = correct_num / accuracy_denom + return loss, accuracy, metrics class OnlineDominoModel(OnlineDFlashModel): @@ -495,42 +541,10 @@ def __init__( loss_type="dflash", ) self.shift_label = shift_label - - def _sample_anchor_positions( - self, seq_len: int, loss_mask: torch.Tensor, device: torch.device - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Randomly sample anchor positions per sample; returns (anchors, keep_mask).""" - bs = self.block_size - bsz = loss_mask.shape[0] - max_anchor = max(seq_len - bs, 0) - - valid = loss_mask[:, : max_anchor + 1] > 0.5 - valid_counts = valid.sum(dim=1) - max_n = max(1, min(self.num_anchors, int(valid_counts.max().item()) - 1)) - - indices = ( - torch.arange(max_anchor + 1, device=device).unsqueeze(0).expand(bsz, -1) - ) - masked_indices = torch.where( - valid, indices, torch.tensor(seq_len + 1, device=device) + self._use_fused_domino_ce = ( + os.environ.get("SPECFORGE_DOMINO_TRITON_CE", "1") == "1" ) - random_vals = torch.rand(bsz, max_anchor + 1, device=device) - random_vals = torch.where(valid, random_vals, torch.tensor(2.0, device=device)) - - _, sorted_idx = random_vals.sort(dim=1) - gathered = torch.gather(masked_indices, 1, sorted_idx) - anchors = gathered[:, :max_n].sort(dim=1).values - - keep_mask = torch.arange(max_n, device=device).unsqueeze( - 0 - ) < valid_counts.unsqueeze(1).clamp(max=max_n) - anchors = torch.where( - keep_mask, anchors, torch.tensor(0, dtype=torch.long, device=device) - ) - - return anchors, keep_mask - def _build_domino_head_inputs( self, input_ids: torch.Tensor, @@ -557,21 +571,6 @@ def _build_domino_head_inputs( return hidden4d, prev_ids - def _apply_domino_head( - self, - base_logits4d: torch.Tensor, - hidden4d: torch.Tensor, - prev_ids: torch.Tensor, - target_ids: torch.Tensor, - ) -> torch.Tensor: - head_token_ids = prev_ids if self.shift_label else target_ids - head_token_embeddings = self.embed_tokens(head_token_ids) - return self.draft_model.apply_logits_head( - base_logits4d, - hidden_states=hidden4d, - prev_token_embeddings=head_token_embeddings, - ) - def _domino_objective_chunk_terms( self, hidden: torch.Tensor, @@ -581,34 +580,34 @@ def _domino_objective_chunk_terms( eval_weight_mask: torch.Tensor, ) -> Tuple[torch.Tensor, ...]: """Return additive Domino loss and telemetry terms for one block slice.""" + from specforge.core.domino_loss import domino_weighted_cross_entropy batch_size, num_blocks, block_size, hidden_size = hidden.shape base_logits = self.lm_head( hidden.reshape(batch_size, num_blocks * block_size, hidden_size) ).reshape(batch_size, num_blocks, block_size, -1) - final_logits = self._apply_domino_head( - base_logits4d=base_logits, - hidden4d=hidden, - prev_ids=prev_ids, - target_ids=target_ids, + head_token_ids = prev_ids if self.shift_label else target_ids + head_token_embeddings = self.embed_tokens(head_token_ids) + correction_logits = self.draft_model.compute_correction_logits( + hidden_states=hidden, + prev_token_embeddings=head_token_embeddings, + ) + final_num, base_num, predicted_ids, base_predicted_ids = ( + domino_weighted_cross_entropy( + base_logits.reshape(-1, base_logits.shape[-1]), + correction_logits.reshape(-1, correction_logits.shape[-1]), + target_ids.reshape(-1), + weight_mask.reshape(-1), + block_size=block_size, + suffix_start=self.draft_model.suffix_start, + use_fused=self._use_fused_domino_ce and base_logits.is_cuda, + ) ) - final_ce = F.cross_entropy( - final_logits.reshape(-1, final_logits.shape[-1]), - target_ids.reshape(-1), - reduction="none", - ).reshape_as(target_ids) - base_ce = F.cross_entropy( - base_logits.reshape(-1, base_logits.shape[-1]), - target_ids.reshape(-1), - reduction="none", - ).reshape_as(target_ids) - final_num = (final_ce * weight_mask).sum() - base_num = (base_ce * weight_mask).sum() loss_den = weight_mask.sum() with torch.no_grad(): - predicted_ids = final_logits.argmax(dim=-1) - base_predicted_ids = base_logits.argmax(dim=-1) + predicted_ids = predicted_ids.reshape_as(target_ids) + base_predicted_ids = base_predicted_ids.reshape_as(target_ids) binary_accuracy_mask = eval_weight_mask > 0.5 correct_num = ( ((predicted_ids == target_ids) & binary_accuracy_mask).sum().float() @@ -650,6 +649,7 @@ def forward( hidden_states: torch.Tensor, loss_mask: torch.Tensor, lambda_base: float = 0.0, + max_valid_anchors: Optional[int] = None, ): """Parallel Domino training forward pass.""" if self.attention_backend == "flex_attention" and not FLEX_ATTENTION_AVAILABLE: @@ -663,6 +663,7 @@ def forward( input_ids=input_ids, hidden_states=hidden_states, loss_mask=loss_mask, + max_valid_anchors=max_valid_anchors, ) label_start = 1 if self.shift_label else 0 @@ -744,7 +745,10 @@ def forward( "base_accuracy": (base_correct_num / (accuracy_denom + 1e-6)).detach(), "accept_len": (accept_num / (accept_den + 1e-6)).detach(), "base_accept_len": (base_accept_num / (accept_den + 1e-6)).detach(), - "lambda_base": torch.tensor(lambda_base, device=loss.device), + # Telemetry does not participate in the objective. Keeping this as + # a host scalar avoids a tiny H2D copy that otherwise synchronizes + # the whole forward stream once per micro-step. + "lambda_base": float(lambda_base), "accuracy_denom": accuracy_denom.detach(), } @@ -792,55 +796,6 @@ def __init__( self.dspark_l1_loss_alpha = float(dspark_l1_loss_alpha) self.dspark_confidence_head_alpha = float(dspark_confidence_head_alpha) - def _build_anchor_candidate_mask( - self, - seq_len: int, - loss_mask: torch.Tensor, - ) -> torch.Tensor: - num_candidates = max(seq_len - 1, 0) - if num_candidates == 0: - return loss_mask[:, :0].bool() - anchor_valid = loss_mask[:, :num_candidates] > 0.5 - first_target_valid = loss_mask[:, 1 : num_candidates + 1] > 0.5 - return anchor_valid & first_target_valid - - def _sample_anchor_positions( - self, seq_len: int, loss_mask: torch.Tensor, device: torch.device - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Sample only anchors with a valid first target, without dummy width.""" - - valid = self._build_anchor_candidate_mask(seq_len, loss_mask) - if valid.shape[1] == 0: - raise ValueError("DSpark needs sequences with at least two tokens") - valid_counts = valid.sum(dim=1) - width = min(self.num_anchors, int(valid_counts.max().item())) - if width <= 0: - raise ValueError( - "DSpark found no valid anchor with two consecutive loss tokens" - ) - indices = torch.arange(valid.shape[1], device=device).expand( - loss_mask.shape[0], -1 - ) - random_values = torch.rand(valid.shape, device=device) - random_values.masked_fill_(~valid, 2.0) - order = random_values.argsort(dim=1) - candidates = torch.gather(indices, 1, order)[:, :width] - keep_mask = torch.arange(width, device=device).unsqueeze( - 0 - ) < valid_counts.clamp(max=width).unsqueeze(1) - anchors = ( - torch.where( - keep_mask, - candidates, - torch.full_like(candidates, valid.shape[1]), - ) - .sort(dim=1) - .values - ) - keep_mask = anchors < valid.shape[1] - anchors = torch.where(keep_mask, anchors, torch.zeros_like(anchors)) - return anchors, keep_mask - def _build_dspark_labels_and_mask( self, input_ids: torch.Tensor, @@ -1157,6 +1112,7 @@ def forward( hidden_states: torch.Tensor, loss_mask: torch.Tensor, target_last_hidden_states: Optional[torch.Tensor] = None, + max_valid_anchors: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, object]]: """Parallel DSpark training forward pass.""" if self.attention_backend == "flex_attention" and not FLEX_ATTENTION_AVAILABLE: @@ -1167,6 +1123,7 @@ def forward( input_ids=input_ids, hidden_states=hidden_states, loss_mask=loss_mask, + max_valid_anchors=max_valid_anchors, ) ( diff --git a/specforge/algorithms/common/dflash_family_data.py b/specforge/algorithms/common/hidden_states_data.py similarity index 68% rename from specforge/algorithms/common/dflash_family_data.py rename to specforge/algorithms/common/hidden_states_data.py index 6d774907d..770ed710c 100644 --- a/specforge/algorithms/common/dflash_family_data.py +++ b/specforge/algorithms/common/hidden_states_data.py @@ -1,13 +1,18 @@ -"""Shared DFlash-family normalization and padding adapters.""" +"""Shared hidden-states normalization and padding adapters. + +Used by the DFlash-family (DFlash/Domino/DSpark) and MTP algorithms. +""" from __future__ import annotations from functools import partial from specforge.algorithms.common.collation import pad_and_concatenate_features +from specforge.data.loss_mask import has_consecutive_supervised_tokens NORMALIZER_ID = "dflash_family_offline_v1" DSPARK_NORMALIZER_ID = "dspark_offline_v1" +MTP_NORMALIZER_ID = "mtp_offline_v1" def _normalize_hidden_states( @@ -56,6 +61,10 @@ def normalize_offline_sample(raw, max_len: int): f"loss_mask={loss_mask.shape[1]}, " f"hidden_states={hidden_states.shape[1]}" ) + if not has_consecutive_supervised_tokens(loss_mask[0]): + raise ValueError( + "offline DFlash-family samples require two consecutive supervised tokens" + ) return { "input_ids": input_ids, "loss_mask": loss_mask, @@ -159,14 +168,7 @@ def collate(features): return collate -def build_vlm_collator(): - """Collator for multimodal capture: text tensors + mRoPE position ids. - - ``position_ids`` arrives as ``(1, L, 3)`` int64 per sample (temporal / - height / width mRoPE rows produced by the capture server) and is padded - along the sequence axis like every other per-token feature. - """ - +def build_dspark_collator(): def collate(features): return pad_and_concatenate_features( features, @@ -174,28 +176,91 @@ def collate(features): "input_ids": 1, "loss_mask": 1, "hidden_states": 1, - "position_ids": 1, + "target_last_hidden_states": 1, }, - required_keys=("input_ids", "loss_mask", "hidden_states", "position_ids"), + required_keys=( + "input_ids", + "loss_mask", + "hidden_states", + "target_last_hidden_states", + ), ) return collate -def build_dspark_collator(): +def normalize_mtp_offline_sample(raw, max_len: int): + """Normalize MTP capture tensors (no aux-layer concat, final hidden only).""" + + input_ids = raw["input_ids"][:max_len].unsqueeze(0) + loss_mask = raw["loss_mask"][:max_len].unsqueeze(0) + target_last_hidden_states = _normalize_hidden_states( + raw, + "target_last_hidden_states", + max_len, + description="MTP target_last_hidden_states", + ) + lengths = { + input_ids.shape[1], + loss_mask.shape[1], + target_last_hidden_states.shape[1], + } + if len(lengths) != 1: + raise ValueError( + "offline MTP features have mismatched sequence lengths after " + f"truncation: input_ids={input_ids.shape[1]}, " + f"loss_mask={loss_mask.shape[1]}, " + f"target_last_hidden_states={target_last_hidden_states.shape[1]}" + ) + return { + "input_ids": input_ids, + "loss_mask": loss_mask, + "target_last_hidden_states": target_last_hidden_states, + } + + +def build_mtp_offline_reader( + strategy, + hidden_states_path, + *, + run_id, + ttt_length, + max_len, +): + # Transitional runtime import; the composition root will inject this port. + from specforge.runtime.data_plane.offline_reader import OfflineManifestReader + + return OfflineManifestReader( + hidden_states_path, + run_id=run_id, + strategy=strategy, + feature_keys=( + "input_ids", + "loss_mask", + "target_last_hidden_states", + ), + target_repr="hidden_state", + ttt_length=ttt_length, + max_len=max_len, + ) + + +def build_mtp_offline_normalizer(max_len, **_topology): + return partial(normalize_mtp_offline_sample, max_len=max_len) + + +def build_mtp_collator(): def collate(features): return pad_and_concatenate_features( features, sequence_axes={ "input_ids": 1, "loss_mask": 1, - "hidden_states": 1, "target_last_hidden_states": 1, }, required_keys=( "input_ids", "loss_mask", - "hidden_states", "target_last_hidden_states", ), ) @@ -205,14 +270,18 @@ def collate(features): __all__ = [ "DSPARK_NORMALIZER_ID", + "MTP_NORMALIZER_ID", "NORMALIZER_ID", "build_collator", "build_dspark_collator", "build_dspark_offline_normalizer", "build_dspark_offline_reader", + "build_mtp_collator", + "build_mtp_offline_normalizer", + "build_mtp_offline_reader", "build_offline_normalizer", "build_offline_reader", - "build_vlm_collator", "normalize_dspark_offline_sample", + "normalize_mtp_offline_sample", "normalize_offline_sample", ] diff --git a/specforge/algorithms/common/providers.py b/specforge/algorithms/common/providers.py index e2f28bc93..b3ae2facf 100644 --- a/specforge/algorithms/common/providers.py +++ b/specforge/algorithms/common/providers.py @@ -166,9 +166,24 @@ class DraftConfigProvider: target_defaults: TargetDerivedDraftDefaults | None = None expected_auto_map_model: str | None = None apply_overrides: Factory | None = None + compatible_architectures: FrozenSet[str] | None = None def __post_init__(self) -> None: _non_empty(self.architecture, field_name="architecture") + compatible = self.compatible_architectures + if compatible is None: + compatible = frozenset({self.architecture}) + else: + compatible = frozenset(compatible) + if not compatible: + raise ValueError("compatible_architectures must not be empty") + for item in compatible: + _non_empty(item, field_name="compatible_architectures item") + if self.architecture not in compatible: + raise ValueError( + "architecture must be included in compatible_architectures" + ) + object.__setattr__(self, "compatible_architectures", compatible) if self.expected_auto_map_model is not None: _non_empty( self.expected_auto_map_model, @@ -361,6 +376,7 @@ class ModelProvider: needs_input_tools: Factory default_dataloader_num_workers: int allow_missing_warm_start_embedding: bool = False + loss_mask_filter: Factory | None = None def __post_init__(self) -> None: if not isinstance(self.draft_config, DraftConfigProvider): @@ -384,6 +400,8 @@ def __post_init__(self) -> None: ) if not isinstance(self.allow_missing_warm_start_embedding, bool): raise TypeError("allow_missing_warm_start_embedding must be a bool") + if self.loss_mask_filter is not None and not callable(self.loss_mask_filter): + raise TypeError("loss_mask_filter must be callable or None") @dataclass(frozen=True) @@ -493,25 +511,18 @@ def __post_init__(self) -> None: @dataclass(frozen=True) class ServerCaptureLayout: - """Maps generic server artifacts onto algorithm-ready feature names. - - ``position_ids_feature`` names an optional server-produced position-id - artifact (mRoPE positions for multimodal targets), stored ``(1, L, 3)`` - int64. Text providers leave it unset. - """ + """Maps generic server artifacts onto algorithm-ready feature names.""" aux_feature: str | None last_hidden_feature: str | None passthrough: Tuple[Tuple[str, str, Tuple[int, ...]], ...] attention_mask_feature: str | None = None - position_ids_feature: str | None = None def __post_init__(self) -> None: for field_name in ( "aux_feature", "last_hidden_feature", "attention_mask_feature", - "position_ids_feature", ): value = getattr(self, field_name) if value is not None: @@ -743,7 +754,6 @@ def make_registration( layout.aux_feature, layout.last_hidden_feature, layout.attention_mask_feature, - layout.position_ids_feature, ) if feature is not None ), diff --git a/specforge/algorithms/contracts.py b/specforge/algorithms/contracts.py index ae65537b3..1d97da89f 100644 --- a/specforge/algorithms/contracts.py +++ b/specforge/algorithms/contracts.py @@ -238,6 +238,7 @@ class AlgorithmCapabilities: attention_backends: FrozenSet[str] required_batch_size: int | None = None supports_compact_teacher: bool = False + supports_trim_loss_positions: bool = False supports_vocab_mapping: bool = False allows_aux_layer_override: bool = False @@ -254,6 +255,7 @@ def __post_init__(self) -> None: raise ValueError("required_batch_size must be a positive integer or None") for field_name in ( "supports_compact_teacher", + "supports_trim_loss_positions", "supports_vocab_mapping", "allows_aux_layer_override", ): diff --git a/specforge/algorithms/dflash/providers.py b/specforge/algorithms/dflash/providers.py index 84a8b38ce..98832d5c2 100644 --- a/specforge/algorithms/dflash/providers.py +++ b/specforge/algorithms/dflash/providers.py @@ -8,12 +8,11 @@ empty_options, no_missing_checkpoint_keys, ) -from specforge.algorithms.common.dflash_family_data import ( +from specforge.algorithms.common.hidden_states_data import ( NORMALIZER_ID, build_collator, build_offline_normalizer, build_offline_reader, - build_vlm_collator, ) from specforge.algorithms.common.providers import ( AlgorithmProviders, @@ -36,6 +35,7 @@ FeatureMode, OfflineStorageContract, ) +from specforge.data.loss_mask import has_consecutive_supervised_tokens ALGORITHM_NAME = "dflash" DRAFT_ARCHITECTURE = "DFlashDraftModel" @@ -134,7 +134,6 @@ def needs_input_tools(config, draft_model): def algorithm_spec() -> AlgorithmSpec: ready = {"input_ids", "loss_mask", "hidden_states"} - vlm_ready = {"input_ids", "loss_mask", "hidden_states", "position_ids"} return AlgorithmSpec( name=ALGORITHM_NAME, draft=DraftRequirement( @@ -161,7 +160,7 @@ def algorithm_spec() -> AlgorithmSpec: FeatureContract( mode=FeatureMode.STREAMING, modality="multimodal", - required_tensors=vlm_ready, + required_tensors=ready, ), ), capabilities=AlgorithmCapabilities( @@ -198,6 +197,7 @@ def algorithm_providers() -> AlgorithmProviders: minimum_loss_tokens=minimum_loss_tokens, needs_input_tools=needs_input_tools, default_dataloader_num_workers=8, + loss_mask_filter=has_consecutive_supervised_tokens, ), offline=( OfflineDataProvider( @@ -243,9 +243,8 @@ def algorithm_providers() -> AlgorithmProviders: ("input_ids", "input_ids", ()), ("loss_mask", "loss_mask", ()), ), - position_ids_feature="position_ids", ), - build_collator=build_vlm_collator, + build_collator=collator, build_input_adapter=build_vlm_input_adapter, ), ), diff --git a/specforge/algorithms/domino/providers.py b/specforge/algorithms/domino/providers.py index d1b3a97a8..16dfcca1b 100644 --- a/specforge/algorithms/domino/providers.py +++ b/specforge/algorithms/domino/providers.py @@ -5,7 +5,7 @@ from functools import partial from specforge.algorithms.common.defaults import no_missing_checkpoint_keys -from specforge.algorithms.common.dflash_family_data import ( +from specforge.algorithms.common.hidden_states_data import ( NORMALIZER_ID, build_collator, build_offline_normalizer, @@ -30,6 +30,7 @@ FeatureMode, OfflineStorageContract, ) +from specforge.data.loss_mask import has_consecutive_supervised_tokens ALGORITHM_NAME = "domino" DRAFT_ARCHITECTURE = "DominoDraftModel" @@ -169,6 +170,7 @@ def algorithm_providers() -> AlgorithmProviders: minimum_loss_tokens=minimum_loss_tokens, needs_input_tools=needs_input_tools, default_dataloader_num_workers=8, + loss_mask_filter=has_consecutive_supervised_tokens, ), offline=( OfflineDataProvider( diff --git a/specforge/algorithms/dspark/providers.py b/specforge/algorithms/dspark/providers.py index 153575d5d..c48951cc0 100644 --- a/specforge/algorithms/dspark/providers.py +++ b/specforge/algorithms/dspark/providers.py @@ -8,7 +8,7 @@ empty_options, no_missing_checkpoint_keys, ) -from specforge.algorithms.common.dflash_family_data import ( +from specforge.algorithms.common.hidden_states_data import ( DSPARK_NORMALIZER_ID, build_dspark_collator, build_dspark_offline_normalizer, @@ -33,9 +33,11 @@ FeatureMode, OfflineStorageContract, ) +from specforge.data.loss_mask import has_consecutive_supervised_tokens ALGORITHM_NAME = "dspark" DRAFT_ARCHITECTURE = "DSparkDraftModel" +COMPATIBLE_DRAFT_ARCHITECTURES = frozenset({DRAFT_ARCHITECTURE}) def build_step(wrapped_model, *, target_head=None, **_options): @@ -112,7 +114,7 @@ def algorithm_spec() -> AlgorithmSpec: return AlgorithmSpec( name=ALGORITHM_NAME, draft=DraftRequirement( - compatible_architectures={DRAFT_ARCHITECTURE}, + compatible_architectures=COMPATIBLE_DRAFT_ARCHITECTURES, default_architecture=DRAFT_ARCHITECTURE, ), feature_contracts=( @@ -155,6 +157,7 @@ def algorithm_providers() -> AlgorithmProviders: model=ModelProvider( draft_config=DraftConfigProvider( architecture=DRAFT_ARCHITECTURE, + compatible_architectures=COMPATIBLE_DRAFT_ARCHITECTURES, expected_auto_map_model="dspark.DSparkDraftModel", ), build_draft=build_draft, @@ -163,13 +166,14 @@ def algorithm_providers() -> AlgorithmProviders: minimum_loss_tokens=minimum_loss_tokens, needs_input_tools=needs_input_tools, default_dataloader_num_workers=8, + loss_mask_filter=has_consecutive_supervised_tokens, ), offline=( OfflineDataProvider( modality="text", normalizer_id=DSPARK_NORMALIZER_ID, capture_layout=OfflineCaptureLayout( - capture_method="dflash", + capture_method="dspark", aux_feature="hidden_states", last_hidden_feature="target_last_hidden_states", passthrough=( @@ -185,7 +189,7 @@ def algorithm_providers() -> AlgorithmProviders: server_streaming=( ServerStreamingProvider( modality="text", - capture_method="dflash", + capture_method="dspark", target_representation="hidden_state", layout=ServerCaptureLayout( aux_feature="hidden_states", diff --git a/specforge/algorithms/eagle3/model.py b/specforge/algorithms/eagle3/model.py index c16022401..d89cf5d0c 100644 --- a/specforge/algorithms/eagle3/model.py +++ b/specforge/algorithms/eagle3/model.py @@ -149,6 +149,8 @@ def _acc_and_loss( position_mask: torch.Tensor, loss_mask: torch.Tensor, adapter: BackendAdapter, + loss_scale: float = 1.0, + full_positions: Optional[int] = None, ) -> Tuple[ torch.Tensor, torch.Tensor, @@ -183,8 +185,14 @@ def _acc_and_loss( reduce_metrics_fn=adapter.reduce_metrics, reduce_loss_fn=adapter.reduce_loss, ) + if loss_scale != 1.0: + # The trimmed loss kernel averages over n_sup supervised positions, but + # the full-length semantics average over L; rescale to recover it. Only + # valid when lk_loss_type is None (a plain KL loss is linearly scalable). + loss = loss * loss_scale loss_denom = torch.tensor( - logits.shape[0] * logits.shape[1], + logits.shape[0] + * (full_positions if full_positions is not None else logits.shape[1]), device=logits.device, dtype=torch.float32, ) @@ -253,6 +261,7 @@ def forward( target_hidden_for_compact: Optional[torch.Tensor] = None, target_head_weight: Optional[torch.Tensor] = None, compact_teacher_chunk_size: int = DEFAULT_VOCAB_CHUNK_SIZE, + trim_loss_positions: bool = False, ) -> Tuple[ List[torch.Tensor], List[torch.Tensor], @@ -274,7 +283,10 @@ def forward( target_hidden_for_compact, target_head_weight, compact_teacher_chunk_size: when the first two are given, the padded teacher is built from hidden states in draft-vocab space and ``target`` is ignored. + trim_loss_positions: compute the teacher, draft logits and loss only at + supervised positions when the batch/objective supports it. """ + adapter = self._make_adapter() # Step 1: handle vocab size if target_hidden_for_compact is not None: ( @@ -291,18 +303,59 @@ def forward( chunk_size=compact_teacher_chunk_size, ) del target_hidden_for_compact + trim_pack = None else: - ( - target_p_padded, - target_p_on_draft_padded, - target_token_ids_padded, - position_mask, - ) = _compute_target_p_padded( - target=target, - t2d=self.draft_model.t2d, - loss_mask=loss_mask, - length=self.length, + # A-level trim: with batch==1 and no lk_loss, compute the teacher only at + # supervised positions; fall back to the full path otherwise. Under USP + # the backbone keeps running on this rank's own chunk (usp_chunk_size = + # local_len - ttt_length); the local buffer's ttt_length overlap tail may + # only act as teacher positions for own-chunk rows, never emit loss rows + # itself (those rows belong to the next rank), so the per-step row sets + # are bounded by chunk_len. + # chunk_len must come from the SAME source the full path uses for its + # slicing/normalization: the hidden-state sequence length (the loss + # kernel means over backbone rows). loss_mask can carry an extra + # zero-padded slot in the offline pipeline, so deriving from it would + # be off by one (wrong rows, wrong denominator, and under USP a + # backbone length that disagrees with full-path ranks). + trim_chunk_len = adapter.backbone_row_count( + seq_length=hidden_states.shape[1], ttt_length=self.length ) + _trim_ok = ( + trim_loss_positions + and self.lk_loss_type is None + and loss_mask.shape[0] == 1 + and trim_chunk_len > 0 + ) + trim_pack = None + if _trim_ok: + # Returns None when no supervised position can reach any row + # (e.g. supervision only beyond the reachable window); then we + # fall through to the full path below. + trim_pack = _build_trim_pack( + target, + self.draft_model.t2d, + loss_mask, + self.length, + chunk_len=trim_chunk_len, + ) + if trim_pack is not None: + target_p_padded = None + target_p_on_draft_padded = None + target_token_ids_padded = None + position_mask = trim_pack["position_mask_sup"] + else: + ( + target_p_padded, + target_p_on_draft_padded, + target_token_ids_padded, + position_mask, + ) = _compute_target_p_padded( + target=target, + t2d=self.draft_model.t2d, + loss_mask=loss_mask, + length=self.length, + ) del target torch.cuda.empty_cache() @@ -349,7 +402,6 @@ def forward( metric_denoms = [] metric_losses = [] metric_loss_denoms = [] - adapter = self._make_adapter() # for sequence paralle, position mask and input ids will split by sequence dim, need to keep origin for ttt shift global_input_ids = input_ids if self.attention_backend in ["sdpa", "fa", "usp"]: @@ -362,33 +414,54 @@ def forward( raise ValueError(f"Unknown attention backend: {self.attention_backend}") for idx in range(self.length): - state = adapter.step_view( - idx=idx, - ttt_length=self.length, - global_input_ids=global_input_ids, - attention_mask=attention_mask, - loss_mask=loss_mask, - position_ids=position_ids, - hidden_states=hidden_states, - target_p_padded=target_p_padded, - target_p_on_draft_padded=target_p_on_draft_padded, - target_token_ids_padded=target_token_ids_padded, - position_mask=position_mask, - seq_length=seq_length, - ) + if trim_pack is not None: + # A-level: the teacher tables are already compacted to supervised + # positions; the backbone runs exactly the same inputs as the full + # path (per-rank chunk under USP, full length otherwise) and only + # supervised rows go through logits/loss below. + backbone = adapter.backbone_view( + row_count=trim_pack["full_len"], + global_input_ids=global_input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + hidden_states=hidden_states, + ) + step_input_ids = backbone.input_ids + step_hidden = backbone.hidden_states + step_attn = backbone.attention_mask + step_pos = backbone.position_ids + else: + state = adapter.step_view( + idx=idx, + ttt_length=self.length, + global_input_ids=global_input_ids, + attention_mask=attention_mask, + loss_mask=loss_mask, + position_ids=position_ids, + hidden_states=hidden_states, + target_p_padded=target_p_padded, + target_p_on_draft_padded=target_p_on_draft_padded, + target_token_ids_padded=target_token_ids_padded, + position_mask=position_mask, + seq_length=seq_length, + ) + step_input_ids = state.input_ids + step_hidden = state.hidden_states + step_attn = state.attention_mask + step_pos = state.position_ids is_last = idx == self.length - 1 # Step 5.1: embed the input ids - inputs_embeds = self.draft_model.embed_input_ids(state.input_ids) + inputs_embeds = self.draft_model.embed_input_ids(step_input_ids) inputs_embeds = inputs_embeds.to(hidden_states.dtype) # Step 5.2: run the draft model backbone hidden_states_out = self.draft_model.backbone( input_embeds=inputs_embeds, - hidden_states=state.hidden_states, + hidden_states=step_hidden, cache_hidden=cache_hidden, - attention_mask=state.attention_mask, - position_ids=state.position_ids, + attention_mask=step_attn, + position_ids=step_pos, past_key_values=past_key_values, use_cache=True, ) @@ -396,27 +469,65 @@ def forward( # update hidden states for next step hidden_states = hidden_states_out - # Step 5.4: get logits - logits = self.draft_model.compute_logits(hidden_states) - - # Step 5.5 + 5.6: metric and loss - ( - acc, - acceptance_rate, - loss, - correct, - denom, - metric_loss, - loss_denom, - ) = self._acc_and_loss( - logits=logits, - target_p=state.target_p, - target_p_on_draft=state.target_p_on_draft, - target_token_ids=state.target_token_ids, - position_mask=state.position_mask, - loss_mask=state.loss_mask, - adapter=adapter, - ) + # Step 5.4 + 5.5 + 5.6: logits, metric and loss + if trim_pack is not None: + # A-level: only the rows that can carry loss at this step go through + # norm + lm_head. Rows shift down by one per step (rows = sup - idx) + # while the teacher/mask stay pinned at the supervised positions. + rows_j = trim_pack["rows_steps"][idx] + keep_j = trim_pack["keep_steps"][idx] + nrows_j = trim_pack["nrows_steps"][idx] + logits = self.draft_model.compute_logits( + hidden_states.index_select(1, rows_j) + ) + pm_j = trim_pack["position_mask_sup"].index_select(1, keep_j) + lm_j = trim_pack["loss_mask_sup"].index_select(1, keep_j) + if nrows_j == 0: + # Dead step: no own-chunk row carries loss here (e.g. all local + # supervised positions sit in the USP overlap tail at this + # depth). The pack padded one dummy entry; zeroing its masks + # makes the contribution exactly zero while every rank still + # runs the same kernels and collective calls. + pm_j = torch.zeros_like(pm_j) + lm_j = torch.zeros_like(lm_j) + ( + acc, + acceptance_rate, + loss, + correct, + denom, + metric_loss, + loss_denom, + ) = self._acc_and_loss( + logits=logits, + target_p=trim_pack["target_p_c"].index_select(1, keep_j), + target_p_on_draft=trim_pack["on_draft_c"].index_select(1, keep_j), + target_token_ids=trim_pack["token_ids_c"].index_select(1, keep_j), + position_mask=pm_j, + loss_mask=lm_j, + adapter=adapter, + loss_scale=nrows_j / trim_pack["full_len"], + full_positions=trim_pack["full_len"], + ) + else: + logits = self.draft_model.compute_logits(hidden_states) + ( + acc, + acceptance_rate, + loss, + correct, + denom, + metric_loss, + loss_denom, + ) = self._acc_and_loss( + logits=logits, + target_p=state.target_p, + target_p_on_draft=state.target_p_on_draft, + target_token_ids=state.target_token_ids, + position_mask=state.position_mask, + loss_mask=state.loss_mask, + adapter=adapter, + ) acces.append(acc) acceptance_rates.append(acceptance_rate) plosses.append(loss) @@ -516,3 +627,127 @@ def _compute_metric_counts(logits, target_token_ids, loss_mask, d2t): ).sum() denom = loss_mask.sum().clamp_min(1e-6) return correct, denom + + +def _compute_target_p_eager(target, t2d, loss_mask, row_chunk=256): + """Uncompiled variant of the teacher target_p computation. + + Kept uncompiled because the supervised-row count varies per batch, which would + trigger repeated torch.compile recompilation. Mathematically identical to the + compiled path; chunks over rows to bound the transient full-vocab fp32 + activation (which can reach several GB when the row count is large). + """ + tps, tpds, toks, pms = [], [], [], [] + n = target.shape[1] + for s in range(0, n, row_chunk): + t = target[:, s : s + row_chunk].float() + ids = t.argmax(-1) + tm = t2d[ids][..., None].int() + pms.append(tm * loss_mask[:, s : s + row_chunk]) + dth = t[..., t2d] + tps.append(F.softmax(dth, dim=2).detach()) + lse = torch.logsumexp(t, dim=-1, keepdim=True) + tpds.append(torch.exp(dth - lse).detach()) + toks.append(ids.detach()) + return ( + torch.cat(tps, 1), + torch.cat(tpds, 1), + torch.cat(toks, 1), + torch.cat(pms, 1), + ) + + +def _build_trim_pack(target, t2d, loss_mask, length, chunk_len=None): + """A-level trim (--trim-loss-positions): keep only the rows that can carry loss. + + Derivation of the per-step row set. On the full-length path the loop applies + ``padding(..., left=False)`` to ``position_mask`` / ``loss_mask`` once per TTT + step, so at step j the mask seen at row p is ``mask[p + j]``; meanwhile + ``step_view`` slices the padded teacher so row p is supervised by the teacher at + absolute position ``p + j``. A row therefore contributes at step j iff + ``p + j`` is supervised, i.e. ``p = s - j`` for some supervised position s. + + So the rows shift *down* by one per step while the teacher/mask positions stay + pinned at the supervised set: + + step j: rows = {s - j : s in sup, s >= j} teacher/mask at those s + + That makes the teacher cheap: it only ever has to be evaluated at ``sup`` + (no sliding window), and each step just drops the entries whose row would fall + off the front of the sequence. + + ``chunk_len`` is the number of rows the backbone actually produces per step. + On single-rank backends it equals the full local length L (default). Under USP + the backbone runs on this rank's own chunk (``usp_chunk_size = L - + ttt_length``) while the local buffer keeps a ``ttt_length`` overlap tail: + tail positions may act as teachers for own-chunk rows at deeper steps, but + tail rows belong to the next rank and must never emit loss here — hence the + additional ``s - j < chunk_len`` bound on the row set, and ``full_len`` (the + loss denominator) becomes ``chunk_len`` to match the full path's + mean-over-chunk semantics. + + A step whose row set comes out empty (all supervised positions out of reach + at that depth) is padded with one dummy entry and reported with + ``nrows_steps[j] == 0``; the caller zeroes its masks so the step contributes + exactly zero loss while kernel launches and collective calls stay aligned + across ranks. + + batch == 1 only (online training uses batch == 1 per rank); the caller falls + back to the full-length path otherwise. + + Returns a dict with, per step j: ``rows_steps[j]`` (row indices into the draft + hidden states), ``keep_steps[j]`` (which supervised entries survive), + ``nrows_steps[j]`` (real row count, 0 for dead steps), plus the teacher + tables evaluated once at ``sup`` and the mask values at ``sup``. + """ + with torch.no_grad(): + B, L = loss_mask.shape[0], loss_mask.shape[1] + assert B == 1, "trim path requires batch==1" + if chunk_len is None: + chunk_len = L + sup = loss_mask.view(-1).nonzero(as_tuple=False).squeeze(-1) # [n_sup] + # Positions beyond chunk_len + length - 2 can never supervise any row at + # any step (would need j >= length); dropping them keeps every later + # index within the hidden/target sequence range even when loss_mask is + # longer than the hidden states (offline pipelines pad it by one). + sup = sup[sup < chunk_len + length - 1] + if sup.numel() == 0: + # Nothing reachable at any step; tell the caller to use the full path. + return None + # Teacher is only ever needed at the supervised positions themselves. + target_sel = target[:, sup] # [1, n_sup, V_target] + lm_sel = loss_mask[:, sup] + target_p_c, on_draft_c, token_ids_c, pm_sup = _compute_target_p_eager( + target_sel, t2d, lm_sel + ) + lm_sup = loss_mask.view(-1)[sup].view(1, -1, 1) + + rows_steps, keep_steps, nrows_steps = [], [], [] + pad_idx = torch.zeros(1, dtype=sup.dtype, device=sup.device) + for j in range(length): + keep = ( + ((sup >= j) & (sup - j < chunk_len)).nonzero(as_tuple=False).squeeze(-1) + ) + n = int(keep.numel()) + if n == 0: + # Dead step: pad with one dummy entry; the caller zeroes its + # masks so it contributes nothing. + keep = pad_idx + rows = pad_idx + else: + rows = sup[keep] - j + rows_steps.append(rows) + keep_steps.append(keep) + nrows_steps.append(n) + return dict( + sup=sup, + rows_steps=rows_steps, + keep_steps=keep_steps, + nrows_steps=nrows_steps, + target_p_c=target_p_c, + on_draft_c=on_draft_c, + token_ids_c=token_ids_c, + position_mask_sup=pm_sup, + loss_mask_sup=lm_sup, + full_len=chunk_len, + ) diff --git a/specforge/algorithms/eagle3/providers.py b/specforge/algorithms/eagle3/providers.py index 4e2d25cf7..835be387c 100644 --- a/specforge/algorithms/eagle3/providers.py +++ b/specforge/algorithms/eagle3/providers.py @@ -68,6 +68,7 @@ def resume_contract(config, draft_model, training_model): "eagle3_lk_loss_type": training_model.lk_loss_type, "eagle3_kl_scale": float(training_model.kl_scale), "eagle3_kl_decay": float(training_model.kl_decay), + "eagle3_trim_loss_positions": bool(config.training.trim_loss_positions), "eagle3_compact_teacher": bool(config.training.compact_teacher), "eagle3_compact_teacher_chunk_size": ( config.training.compact_teacher_chunk_size @@ -160,6 +161,7 @@ def algorithm_spec() -> AlgorithmSpec: capabilities=AlgorithmCapabilities( attention_backends={"sdpa", "flex_attention", "fa", "usp"}, supports_compact_teacher=True, + supports_trim_loss_positions=True, supports_vocab_mapping=True, allows_aux_layer_override=True, ), diff --git a/specforge/algorithms/model_providers.py b/specforge/algorithms/model_providers.py index a485c8af5..e09cfe05a 100644 --- a/specforge/algorithms/model_providers.py +++ b/specforge/algorithms/model_providers.py @@ -211,6 +211,22 @@ def resolve_eagle_capture_layers( return layers +def _validate_dflash_block_size(draft_config: Any) -> None: + block_size = ( + draft_config.get("block_size") + if isinstance(draft_config, dict) + else getattr(draft_config, "block_size", None) + ) + if ( + not isinstance(block_size, int) + or isinstance(block_size, bool) + or block_size < 2 + ): + raise ValueError( + "DFlash-family draft config must define an integer block_size >= 2" + ) + + def resolve_dflash_capture_layers( _cfg: Config, draft_config: Any, _target_config: Any ) -> List[int]: @@ -318,6 +334,7 @@ def _build_dflash_family_model( ) -> AlgorithmModelParts: from specforge.modeling.target.target_utils import TargetEmbeddingsAndHead + _validate_dflash_block_size(draft_model) mask_token_id = _resolve_mask_token_id(cfg, draft_model, tokenizer) draft_model.mask_token_id = mask_token_id method_config = getattr(draft_model.config, "dflash_config", None) @@ -419,6 +436,7 @@ def build_dspark_model( def eagle3_strategy_kwargs(cfg: Config) -> Dict[str, Any]: return { + "trim_loss_positions": cfg.training.trim_loss_positions, "compact_teacher": cfg.training.compact_teacher, "compact_teacher_chunk_size": cfg.training.compact_teacher_chunk_size, } @@ -432,16 +450,8 @@ def domino_strategy_kwargs(cfg: Config) -> Dict[str, Any]: def dflash_min_loss_tokens(_cfg: Config, draft_config: Any) -> int: - block_size = getattr(draft_config, "block_size", None) - if ( - not isinstance(block_size, int) - or isinstance(block_size, bool) - or block_size < 1 - ): - raise ValueError( - "DFlash-family draft config must define a positive integer block_size" - ) - return 2 * block_size + _validate_dflash_block_size(draft_config) + return 2 def populate_dflash_generated_config( @@ -461,22 +471,41 @@ def populate_dflash_generated_config( ) payload["num_target_layers"] = target_layers payload["block_size"] = 16 + payload["layer_types"] = ["full_attention"] * int(payload["num_hidden_layers"]) + payload["sliding_window"] = None + payload["use_sliding_window"] = False payload["dflash_config"] = { "target_layer_ids": build_target_layer_ids(target_layers, 1) } def apply_dflash_overrides(cfg: Config, draft_config: Any) -> None: - if cfg.model.draft_num_hidden_layers is None: - return - from specforge.modeling.draft.dflash import build_target_layer_ids - - target_layers = int(draft_config.num_target_layers) - method_config = dict(getattr(draft_config, "dflash_config", None) or {}) - method_config["target_layer_ids"] = build_target_layer_ids( - target_layers, cfg.model.draft_num_hidden_layers + from specforge.modeling.draft.dflash import ( + build_target_layer_ids, + resolve_dflash_attention_layout, ) - draft_config.dflash_config = method_config + + requested_layers = cfg.model.draft_num_hidden_layers + if requested_layers is not None: + layer_types = draft_config.layer_types + if len(layer_types) != requested_layers: + if len(set(layer_types)) > 1: + raise ValueError( + "model.draft_num_hidden_layers cannot resize a mixed " + "DFlash layer_types layout; provide a draft config with " + "exactly the requested per-layer layout" + ) + draft_config.layer_types = [layer_types[0]] * requested_layers + + draft_config.dflash_config = { + **dict(getattr(draft_config, "dflash_config", None) or {}), + "target_layer_ids": build_target_layer_ids( + int(draft_config.num_target_layers), + requested_layers, + ), + } + + resolve_dflash_attention_layout(draft_config) __all__ = [ diff --git a/specforge/algorithms/mtp/__init__.py b/specforge/algorithms/mtp/__init__.py new file mode 100644 index 000000000..2b409daf0 --- /dev/null +++ b/specforge/algorithms/mtp/__init__.py @@ -0,0 +1,5 @@ +"""MTP algorithm registration.""" + +from specforge.algorithms.mtp.providers import create_registration + +__all__ = ["create_registration"] diff --git a/specforge/algorithms/mtp/providers.py b/specforge/algorithms/mtp/providers.py new file mode 100644 index 000000000..bde2aa7b3 --- /dev/null +++ b/specforge/algorithms/mtp/providers.py @@ -0,0 +1,318 @@ +"""Built-in MTP (Multi-Token Prediction) registration and executable providers. + +MTP fine-tunes the single-layer draft head shipped natively with Qwen3.5-style +checkpoints. The training signal is the target model's *final* (post-norm) +hidden state — unlike EAGLE3 it needs no aux-layer concat, so the offline +capture persists only ``input_ids`` / ``loss_mask`` / ``target_last_hidden_states`` +(the streaming layout still carries the aux tensor the server patch always +produces; the strategy ignores it). + +Draft construction initializes from the *native* ``mtp.*`` weights inside the +target checkpoint (fine-tuning) and shares + freezes the target embedding and +lm_head, mirroring the serving layout consumed by SGLang's +``Qwen3_5ForCausalLMMTP``. +""" + +from __future__ import annotations + +from functools import partial + +from specforge.algorithms.common.defaults import ( + empty_options, + no_missing_checkpoint_keys, + one_loss_token, + online_needs_input_tools, +) +from specforge.algorithms.common.hidden_states_data import ( + MTP_NORMALIZER_ID, + build_mtp_collator, + build_mtp_offline_normalizer, + build_mtp_offline_reader, +) +from specforge.algorithms.common.providers import ( + AlgorithmProviders, + DraftConfigProvider, + ModelProvider, + OfflineCaptureLayout, + OfflineDataProvider, + ServerCaptureLayout, + ServerStreamingProvider, + StepProvider, + make_registration, +) +from specforge.algorithms.contracts import ( + AlgorithmCapabilities, + AlgorithmSpec, + DraftRequirement, + FeatureContract, + FeatureMode, + OfflineStorageContract, +) + +ALGORITHM_NAME = "mtp" +DRAFT_ARCHITECTURE = "Qwen3_5MTPDraftModel" + +# MTP persists no aux-layer tensor, but the capture plan requires a non-empty +# layer list; a single layer keeps the (discarded) aux capture cheap. +_CAPTURE_LAYER_IDS = [1] + + +def build_step(wrapped_model, *, target_head=None, **_options): + del target_head + from specforge.training.strategies.base import MTPTrainStrategy + + return MTPTrainStrategy(wrapped_model) + + +def resume_contract(_config, draft_model, training_model): + """Persist resolved MTP model and objective semantics.""" + + mtp_config = getattr(draft_model.config, "mtp_config", None) or {} + return { + "mtp_draft_num_hidden_layers": int( + getattr(draft_model.config, "num_hidden_layers", 1) + ), + "mtp_draft_vocab_size": int(getattr(draft_model.config, "vocab_size", 0)), + "mtp_share_lm_head": bool(mtp_config.get("share_lm_head", True)), + "mtp_attention_backend": str( + getattr(draft_model.config, "_attn_implementation", "") + ), + } + + +def _init_from_native_mtp(cfg, draft_model) -> None: + """Initialize the draft's ``mtp.*`` weights from the target checkpoint. + + Qwen3.5-style target checkpoints ship a native MTP head whose keys match + the draft's flat ``mtp.*`` layout. The target lm_head may be shared rather + than duplicated under that prefix. Loading the required keys turns training + into fine-tuning of the native head — the only training mode for this + algorithm. Initialization is strict by default: missing or partial native + state fails rather than silently leaving trainable tensors randomized. + """ + + if cfg.model.draft_checkpoint_path: + print( + "[mtp] native target initialization skipped; weights come from the " + "warm-start draft checkpoint." + ) + return + + from specforge.modeling.target.checkpoint import ( + load_selected_tensors, + resolve_checkpoint_dir, + ) + + target_path = cfg.model.target_model_path + prefix = draft_model.NATIVE_KEY_PREFIX + try: + checkpoint_dir = resolve_checkpoint_dir( + target_path, cache_dir=cfg.model.cache_dir + ) + native_mtp = load_selected_tensors( + checkpoint_dir, lambda key: key.startswith(prefix) + ) + scan_error = None + except Exception as exc: # pragma: no cover - depends on target checkpoint + native_mtp = {} + scan_error = exc + + if native_mtp: + model_keys = set(draft_model.native_state_dict()) + required_keys = set(draft_model.required_native_state_keys()) + extra_keys = set(draft_model.allowed_extra_native_state_keys()) + loaded_keys = set(native_mtp) + missing = sorted(required_keys - loaded_keys) + unexpected = sorted(loaded_keys - model_keys - extra_keys) + if missing or unexpected: + details = [] + if missing: + details.append(f"missing required native keys: {missing}") + if unexpected: + details.append(f"unexpected native keys: {unexpected}") + raise RuntimeError( + f"[mtp] incompatible native {prefix}* state in {target_path}: " + + "; ".join(details) + ) + draft_model.load_state_dict(native_mtp, strict=False) + print( + f"[mtp] initialized {len(native_mtp)} native {prefix}* weights from " + f"{target_path} (native-MTP fine-tune)." + ) + return + + detail = f" (scan failed: {scan_error})" if scan_error is not None else "" + raise RuntimeError( + f"[mtp] no native {prefix}* weights found in {target_path}{detail}. MTP " + "training fine-tunes the native MTP head shipped with the target " + "checkpoint and does not start from random initialization by default. " + "Point model.target_model_path at a checkpoint that ships native MTP " + "weights (e.g. Qwen3.5), or set model.draft_checkpoint_path to resume " + "from a trained MTP draft." + ) + + +def _share_target_embeddings(cfg, draft_model, torch_dtype) -> None: + """Share (and freeze) the target checkpoint's embed_tokens and lm_head.""" + + from specforge.modeling.target.target_utils import TargetEmbeddingsAndHead + + target_components = TargetEmbeddingsAndHead.from_pretrained( + cfg.model.target_model_path, + embed_key=cfg.model.embedding_key, + lm_head_key=cfg.model.lm_head_key, + cache_dir=cfg.model.cache_dir, + device="cpu", + dtype=torch_dtype, + trust_remote_code=cfg.model.trust_remote_code, + ) + draft_model.share_target_embeddings( + target_components.embed_tokens.weight, + lm_head_weight=target_components.lm_head.weight, + ) + print("[mtp] shared target embed_tokens/lm_head with the draft (frozen).") + + +def build_draft(cfg, draft_config): + import torch + + from specforge.modeling.auto import AutoDraftModel + from specforge.training.model_loading import warm_start_draft_model + from specforge.utils import get_local_device + + torch_dtype = getattr(torch, cfg.model.torch_dtype) + draft_config._attn_implementation = cfg.training.attention_backend + draft_model = AutoDraftModel.from_config(draft_config, torch_dtype=torch_dtype) + + _init_from_native_mtp(cfg, draft_model) + _share_target_embeddings(cfg, draft_model, torch_dtype) + + if cfg.model.draft_checkpoint_path: + warm_start_draft_model( + draft_model, + cfg.model.draft_checkpoint_path, + draft_config=draft_config, + strategy=cfg.training.strategy, + cache_dir=cfg.model.cache_dir, + trust_remote_code=cfg.model.trust_remote_code, + ) + return draft_model.to(device=get_local_device(), dtype=torch_dtype) + + +def build_training_model(config, draft_model, draft_config, target_config, tokenizer): + from specforge.algorithms.model_providers import AlgorithmModelParts + from specforge.core.mtp import OnlineMTPModel + + return AlgorithmModelParts( + model=OnlineMTPModel(draft_model=draft_model), + capture_layers=None, + ) + + +def resolve_capture_layers(config, draft_config, target_config): + return list(_CAPTURE_LAYER_IDS) + + +def create_registration(): + return make_registration(algorithm_spec(), algorithm_providers()) + + +def algorithm_spec() -> AlgorithmSpec: + ready = { + "input_ids", + "loss_mask", + "target_last_hidden_states", + } + return AlgorithmSpec( + name=ALGORITHM_NAME, + draft=DraftRequirement( + compatible_architectures={DRAFT_ARCHITECTURE}, + default_architecture=DRAFT_ARCHITECTURE, + ), + feature_contracts=( + FeatureContract( + mode=FeatureMode.OFFLINE, + modality="text", + required_tensors=ready, + allowed_target_representations={"hidden_state"}, + default_target_representation="hidden_state", + storage=OfflineStorageContract( + format="specforge_hidden_states_v1", + required_tensors=ready, + normalizer=MTP_NORMALIZER_ID, + ), + ), + FeatureContract( + mode=FeatureMode.STREAMING, + modality="text", + required_tensors=ready, + allowed_target_representations={"hidden_state"}, + default_target_representation="hidden_state", + ), + ), + capabilities=AlgorithmCapabilities( + attention_backends={"eager", "sdpa"}, + ), + ) + + +def algorithm_providers() -> AlgorithmProviders: + return AlgorithmProviders( + algorithm_name=ALGORITHM_NAME, + step=StepProvider( + build=build_step, + options=empty_options, + resume_contract=resume_contract, + allowed_missing_checkpoint_keys=no_missing_checkpoint_keys, + uses_external_target_head=False, + ), + model=ModelProvider( + draft_config=DraftConfigProvider( + architecture=DRAFT_ARCHITECTURE, + expected_auto_map_model="mtp.Qwen3_5MTPDraftModel", + ), + build_draft=build_draft, + build_training_model=build_training_model, + resolve_capture_layers=resolve_capture_layers, + minimum_loss_tokens=one_loss_token, + needs_input_tools=online_needs_input_tools, + default_dataloader_num_workers=8, + ), + offline=( + OfflineDataProvider( + modality="text", + normalizer_id=MTP_NORMALIZER_ID, + capture_layout=OfflineCaptureLayout( + capture_method="dflash", + aux_feature=None, + last_hidden_feature="target_last_hidden_states", + passthrough=( + ("input_ids", "input_ids"), + ("loss_mask", "loss_mask"), + ), + ), + build_reader=partial(build_mtp_offline_reader, ALGORITHM_NAME), + build_normalizer=build_mtp_offline_normalizer, + build_collator=build_mtp_collator, + ), + ), + server_streaming=( + ServerStreamingProvider( + modality="text", + capture_method="dflash", + target_representation="hidden_state", + layout=ServerCaptureLayout( + aux_feature="hidden_states", + last_hidden_feature="target_last_hidden_states", + passthrough=( + ("input_ids", "input_ids", ()), + ("loss_mask", "loss_mask", ()), + ), + ), + build_collator=build_mtp_collator, + ), + ), + ) + + +__all__ = ["algorithm_providers", "algorithm_spec", "create_registration"] diff --git a/specforge/application/composition.py b/specforge/application/composition.py index 735084f16..0eb271c17 100644 --- a/specforge/application/composition.py +++ b/specforge/application/composition.py @@ -3,6 +3,7 @@ from __future__ import annotations from dataclasses import dataclass +from typing import Callable from specforge.algorithms.common.providers import OfflineCaptureLayout from specforge.algorithms.registry import AlgorithmRegistration, AlgorithmRegistry @@ -26,6 +27,7 @@ class ResolvedOfflineCapture: capture_method: str capture_layers: tuple[int, ...] layout: OfflineCaptureLayout + loss_mask_filter: Callable[[object], bool] | None def bind_run(cfg: Config, algorithm: AlgorithmRegistration) -> ResolvedRun: @@ -127,6 +129,7 @@ def resolve_offline_capture( capture_method=offline.capture_layout.capture_method, capture_layers=layers, layout=offline.capture_layout, + loss_mask_filter=model_provider.loss_mask_filter, ) diff --git a/specforge/application/planning.py b/specforge/application/planning.py index b59a5b82a..44830a927 100644 --- a/specforge/application/planning.py +++ b/specforge/application/planning.py @@ -120,6 +120,12 @@ def _validate_algorithm_capabilities( f"mode={mode.value!r}, modality={cfg.model.input_modality!r}" ) + if training.trim_loss_positions and not capabilities.supports_trim_loss_positions: + raise ValueError( + f"algorithm {algorithm.name!r} does not support " + "training.trim_loss_positions" + ) + def _validate_training_topology( cfg: Config, diff --git a/specforge/cli.py b/specforge/cli.py index bd057dd53..92fed155a 100644 --- a/specforge/cli.py +++ b/specforge/cli.py @@ -256,7 +256,12 @@ def main(argv: Optional[List[str]] = None) -> int: print(plan.render()) return 0 if plan.kind == "worker": - os.environ.update(plan.worker_env) + for key, value in plan.worker_env.items(): + if value is None: + # CommandSpec.env contract: None unsets the variable. + os.environ.pop(key, None) + else: + os.environ[key] = value role_config = _config_for_role(resolved.config, plan.role) try: with _worker_signal_unwind(): diff --git a/specforge/config/schema.py b/specforge/config/schema.py index ea6c1c532..eb955e7c5 100644 --- a/specforge/config/schema.py +++ b/specforge/config/schema.py @@ -20,6 +20,7 @@ import json import os from typing import List, Literal, Optional +from urllib.parse import urlparse from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -85,6 +86,12 @@ class ModelConfig(StrictConfigModel): #: large images. sglang_mm_attention_backend: Optional[str] = None sglang_mem_fraction_static: float = Field(default=0.4, gt=0.0, le=1.0) + #: Keep the historical managed-local behavior by default. Hybrid targets + #: such as Inkling require the radix tree and can opt back in explicitly. + #: Disabling it is also required for hybrid linear-attention/Mamba targets + #: on ROCm, whose mamba radix-cache ``extra_buffer`` strategy asserts + #: CUDA/MUSA/NPU (FLA) at server init. + sglang_disable_radix_cache: bool = True sglang_context_length: Optional[int] = Field(default=None, gt=0) sglang_enable_nccl_nvls: bool = False sglang_enable_symm_mem: bool = False @@ -288,6 +295,12 @@ class ManagedLocalMooncakeConfig(StrictConfigModel): global_segment_size_bytes: int = Field(default=32 << 30, gt=0) local_buffer_size_bytes: int = Field(default=1 << 30, gt=0) startup_timeout_s: float = Field(default=60.0, gt=0) + #: Master key-lease TTL (ms) forwarded to ``mooncake_master + #: --default_kv_lease_ttl``. The consumer's teardown drain allows about + #: 19.5s for leases to settle. Keep the managed-local default at 500ms so + #: normal shutdown does not spend several seconds waiting for an expired + #: read lease; set null to inherit Mooncake's server default. + default_kv_lease_ttl_ms: Optional[int] = Field(default=500, gt=0) @model_validator(mode="after") def _validate_endpoint(self): @@ -306,7 +319,11 @@ def _validate_endpoint(self): class ManagedLocalCaptureServerConfig(StrictConfigModel): - """One patched SGLang capture server owned by the local supervisor.""" + """One patched SGLang capture server owned by the local supervisor. + + ``cuda_visible_devices`` holds plain device ordinals; the launcher maps + them to the accelerator's visibility env var. + """ port: int = Field(gt=0, le=65535) cuda_visible_devices: List[str] = Field(min_length=1) @@ -332,7 +349,11 @@ def _validate_devices(self): class ManagedLocalStackConfig(StrictConfigModel): - """Opt-in ownership of a complete single-node online capture stack.""" + """Opt-in ownership of a complete single-node online capture stack. + + ``trainer_cuda_visible_devices`` holds plain device ordinals; the + launcher maps them to the accelerator's visibility env var. + """ trainer_cuda_visible_devices: List[str] = Field(min_length=1) mooncake: ManagedLocalMooncakeConfig = Field( @@ -388,9 +409,13 @@ class DisaggregatedDeploymentConfig(StrictConfigModel): #: Attempt-scoped shared directory. The launcher derives refs, manifest, #: and lifecycle markers beneath it. control_dir: str - #: Optional node-local root for online consumer SQLite/WAL and rank inboxes. - #: When omitted, the historical control_dir-derived paths remain in use. + #: Optional node-local root for the online consumer SQLite/WAL. When + #: omitted, the historical control_dir-derived path remains in use. consumer_state_dir: Optional[str] = None + #: Optional rank-0 HTTP relay for per-rank online inboxes. This removes the + #: shared-filesystem requirement between trainer nodes while keeping the + #: authority-owned source channel and SQLite/WAL on trainer node 0. + inbox_server_url: Optional[str] = None backend: Literal["shared_dir", "mooncake"] store_root: Optional[str] = None store_id: Optional[str] = None @@ -428,6 +453,24 @@ def _validate_store(self): "deployment.disaggregated.consumer_state_dir must be non-empty " "and must not contain surrounding whitespace" ) + if self.inbox_server_url is not None: + parsed = urlparse(self.inbox_server_url) + if ( + parsed.scheme != "http" + or not parsed.hostname + or parsed.port is None + or parsed.username is not None + or parsed.password is not None + or parsed.path not in ("", "/") + or parsed.params + or parsed.query + or parsed.fragment + ): + raise ValueError( + "deployment.disaggregated.inbox_server_url must be an " + "http://host:port origin without credentials, path, query, " + "or fragment" + ) if self.backend == "shared_dir" and not self.store_root: raise ValueError( "deployment.disaggregated.store_root is required for shared_dir" @@ -452,8 +495,7 @@ def _validate_store(self): explicit = [name for name, value in configured_endpoints.items() if value] if explicit: raise ValueError( - "managed_local derives Mooncake endpoints; do not set " - f"{explicit}" + f"managed_local derives Mooncake endpoints; do not set {explicit}" ) if self.producer_segment_size is not None: raise ValueError( @@ -493,6 +535,7 @@ class TrainingConfig(StrictConfigModel): accumulation_steps: int = Field(default=1, gt=0) fsdp_sharding: Literal["SHARD_GRAD_OP", "FULL_SHARD", "NO_SHARD"] = "SHARD_GRAD_OP" learning_rate: float = Field(default=1e-4, gt=0.0) + lr_scheduler: Literal["cosine", "constant"] = "cosine" warmup_ratio: float = Field(default=0.015, ge=0.0, le=1.0) max_grad_norm: float = Field(default=0.5, gt=0.0) #: Keep FP32 Adam masters and moments on CPU while the trainable draft @@ -512,6 +555,12 @@ class TrainingConfig(StrictConfigModel): lk_loss_type: Optional[Literal["lambda", "alpha"]] = None kl_scale: float = 1.0 kl_decay: float = 1.0 + #: Compute the teacher target_p, draft logits and loss only at supervised + #: (loss-masked) positions instead of over the full sequence. Mathematically + #: equivalent (the mean denominator is rescaled) and saves memory/compute on + #: prompt-heavy data. Falls back to the full-length path for batch > 1 or when + #: an lk_loss objective is used. + trim_loss_positions: bool = False #: DFlash-family objective/model knobs. num_anchors: int = Field(default=512, gt=0) loss_decay_gamma: Optional[float] = None @@ -551,6 +600,9 @@ class TrainingConfig(StrictConfigModel): #: and different roles. role: Literal["auto", "all", "producer", "consumer"] = "all" seed: int = 42 + #: Deterministic online prompt ordering. ``None`` preserves the historical + #: behavior of using the run RNG seed for both model and prompt sampling. + prompt_seed: Optional[int] = None @model_validator(mode="after") def _validate_training_shape(self): @@ -727,8 +779,7 @@ def _validate_run_structure(self): ) if self.data.eval_hidden_states_path and mode != "offline": raise ValueError( - "data.eval_hidden_states_path requires an offline training data " - "source" + "data.eval_hidden_states_path requires an offline training data source" ) if ( not self.training.compact_teacher @@ -758,12 +809,28 @@ def _validate_run_structure(self): if self.deployment.disaggregated is not None else None ) + inbox_server_url = ( + self.deployment.disaggregated.inbox_server_url + if self.deployment.disaggregated is not None + else None + ) if consumer_state_dir is not None: if mode != "online" or deployment != "disaggregated": raise ValueError( "deployment.disaggregated.consumer_state_dir is valid only " "for online disaggregated training" ) + if inbox_server_url is not None: + if mode != "online" or deployment != "disaggregated": + raise ValueError( + "deployment.disaggregated.inbox_server_url is valid only " + "for online disaggregated training" + ) + if self.deployment.trainer.nnodes < 2: + raise ValueError( + "deployment.disaggregated.inbox_server_url requires a " + "multi-node trainer" + ) if ( mode == "online" and deployment == "disaggregated" diff --git a/specforge/core/domino_loss.py b/specforge/core/domino_loss.py new file mode 100644 index 000000000..7c43bff22 --- /dev/null +++ b/specforge/core/domino_loss.py @@ -0,0 +1,144 @@ +"""Domino cross-entropy with a portable reference and optional Triton path.""" + +import torch +import torch.nn.functional as F + + +def domino_weighted_cross_entropy( + base_logits, + correction_logits, + targets, + weights, + block_size, + suffix_start, + *, + use_fused: bool, +): + """Return weighted final/base loss sums and predictions for Domino logits. + + ``use_fused=True`` requires CUDA and imports the Triton implementation + lazily. Otherwise, the reference path materializes final logits without + requiring Triton. Only the two logit tensors are differentiable; targets, + weights, and layout arguments are treated as fixed inputs. + + Tensor layout: + ``num_rows = num_blocks * block_size`` + ``base_logits``: ``[num_rows, vocab_size]`` + ``correction_logits``: ``[num_blocks * suffix_size, vocab_size]`` + ``targets``: ``[num_rows]`` integer indices in ``[0, vocab_size)`` + ``weights``: ``[num_rows]`` + + Returns: + ``final_loss_sum``: weighted corrected-logit losses summed over rows. + ``base_loss_sum``: weighted base-logit losses summed over rows. + ``final_pred``: ``[num_rows]`` predictions from corrected logits. + ``base_pred``: ``[num_rows]`` predictions from base logits. + """ + if base_logits.dim() != 2 or correction_logits.dim() != 2: + raise ValueError("Domino cross entropy expects 2D logits") + if block_size < 1 or not 0 <= suffix_start < block_size: + raise ValueError("suffix_start must select a non-empty Domino suffix") + num_rows, vocab_size = base_logits.shape + if vocab_size < 1: + raise ValueError("Domino cross entropy requires a non-empty vocabulary") + if num_rows % block_size: + raise ValueError("base-logit rows must be divisible by block_size") + num_blocks = num_rows // block_size + suffix_size = block_size - suffix_start + expected_correction_rows = num_blocks * suffix_size + if correction_logits.shape != (expected_correction_rows, vocab_size): + raise ValueError("correction logits do not match the configured suffix width") + if targets.shape != (num_rows,) or weights.shape != (num_rows,): + raise ValueError("targets and weights must have one value per logit row") + if targets.dtype != torch.long: + raise ValueError("targets must contain torch.long class indices") + if base_logits.device != correction_logits.device or any( + tensor.device != base_logits.device for tensor in (targets, weights) + ): + raise ValueError("Domino cross-entropy inputs must be on the same device") + if base_logits.dtype != correction_logits.dtype: + raise ValueError("base and correction logits must have the same dtype") + if not base_logits.is_floating_point(): + raise ValueError("base and correction logits must be floating point") + + if not use_fused: + return _domino_weighted_cross_entropy_reference( + base_logits, + correction_logits, + targets, + weights, + block_size, + suffix_start, + ) + if not base_logits.is_cuda: + raise ValueError("Fused Domino cross entropy requires CUDA logits") + if base_logits.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError( + "Fused Domino cross entropy requires FP16, BF16, or FP32 logits" + ) + + try: + from specforge.core.domino_loss_triton import ( + domino_weighted_cross_entropy_fused, + ) + except ModuleNotFoundError as exc: + if exc.name == "triton" or exc.name.startswith("triton."): + raise ImportError( + "Fused Domino cross entropy requires Triton; install Triton or " + "call with use_fused=False." + ) from exc + raise + + return domino_weighted_cross_entropy_fused( + base_logits, + correction_logits, + targets, + weights, + block_size, + suffix_start, + ) + + +def _domino_weighted_cross_entropy_reference( + base_logits, + correction_logits, + targets, + weights, + block_size, + suffix_start, +): + """Reference implementation that materializes the corrected logits.""" + num_blocks = base_logits.shape[0] // block_size + suffix_size = block_size - suffix_start + base_logits_3d = base_logits.reshape(num_blocks, block_size, -1) + correction_logits_3d = correction_logits.reshape(num_blocks, suffix_size, -1) + final_logits = torch.cat( + [ + base_logits_3d[:, :suffix_start], + base_logits_3d[:, suffix_start:] + correction_logits_3d, + ], + dim=1, + ).reshape_as(base_logits) + loss_logits = ( + final_logits.float() + if final_logits.dtype in (torch.float16, torch.bfloat16) + else final_logits + ) + base_loss_logits = ( + base_logits.float() + if base_logits.dtype in (torch.float16, torch.bfloat16) + else base_logits + ) + final_losses = F.cross_entropy(loss_logits, targets, reduction="none") + base_losses = F.cross_entropy(base_loss_logits, targets, reduction="none") + final_loss_sum = (final_losses * weights).sum() + base_loss_sum = (base_losses * weights).sum() + return ( + final_loss_sum, + base_loss_sum, + final_logits.argmax(-1), + base_logits.argmax(-1), + ) + + +__all__ = ["domino_weighted_cross_entropy"] diff --git a/specforge/core/domino_loss_triton.py b/specforge/core/domino_loss_triton.py new file mode 100644 index 000000000..17a28c9e3 --- /dev/null +++ b/specforge/core/domino_loss_triton.py @@ -0,0 +1,464 @@ +"""Triton implementation of Domino cross-entropy. + +The CUDA path processes base and compact-correction logits in vocabulary tiles. +It never writes full corrected logits, padded corrections, base/final softmax or +log-softmax tensors, or a separate corrected-logit gradient. Backward rebuilds +each probability tile from the saved row maximum and shifted exponential sum, +then writes only the required base and compact-correction gradients. +""" + +import torch +import triton +import triton.language as tl + +__all__ = ["domino_weighted_cross_entropy_fused"] + + +def domino_weighted_cross_entropy_fused( + base_logits, + correction_logits, + targets, + weights, + block_size, + suffix_start, +): + """Launch the fused autograd implementation after public validation.""" + return _DominoCrossEntropy.apply( + base_logits, + correction_logits, + targets, + weights, + block_size, + suffix_start, + ) + + +class _DominoCrossEntropy(torch.autograd.Function): + """Return weighted loss sums and predictions from the fused kernels. + + Saves each row's maximum logit and shifted exponential sum so backward can + reconstruct both softmaxes without saving full probability tensors. + """ + + @staticmethod + def forward( + ctx, + base_logits, + correction_logits, + targets, + weights, + block_size, + suffix_start, + ): + num_rows, vocab_size = base_logits.shape + base_logits = base_logits.contiguous() + correction_logits = correction_logits.contiguous() + targets = targets.contiguous() + weights = weights.contiguous() + + device = base_logits.device + + # kernel outputs + final_losses = torch.empty(num_rows, device=device, dtype=torch.float32) + base_losses = torch.empty_like(final_losses) + final_max = torch.empty_like(final_losses) + final_exp_sum = torch.empty_like(final_losses) + base_max = torch.empty_like(final_losses) + base_exp_sum = torch.empty_like(final_losses) + final_pred = torch.empty(num_rows, device=device, dtype=torch.long) + base_pred = torch.empty_like(final_pred) + + vocab_block_size, num_warps = _calculate_domino_ce_settings(vocab_size) + _domino_cross_entropy_forward_kernel[(num_rows,)]( + base_logits, + base_logits.stride(0), + correction_logits, + correction_logits.stride(0), + targets, + weights, + final_losses, + base_losses, + final_max, + final_exp_sum, + base_max, + base_exp_sum, + final_pred, + base_pred, + vocab_size, + BLOCK_SIZE=block_size, + SUFFIX_START=suffix_start, + VOCAB_BLOCK_SIZE=vocab_block_size, + num_warps=num_warps, + ) + + ctx.block_size = block_size + ctx.suffix_start = suffix_start + ctx.save_for_backward( + base_logits, + correction_logits, + targets, + weights, + final_max, + final_exp_sum, + base_max, + base_exp_sum, + ) + ctx.mark_non_differentiable(final_pred, base_pred) + return ( + final_losses.sum(), + base_losses.sum(), + final_pred, + base_pred, + ) + + @staticmethod + def backward( + ctx, + grad_final_loss, + grad_base_loss, + _grad_final_pred, + _grad_base_pred, + ): + ( + base_logits, + correction_logits, + targets, + weights, + final_max, + final_exp_sum, + base_max, + base_exp_sum, + ) = ctx.saved_tensors + num_rows, vocab_size = base_logits.shape + grad_base_logits = torch.empty_like(base_logits) + grad_correction_logits = torch.empty_like(correction_logits) + vocab_block_size, num_warps = _calculate_domino_ce_settings(vocab_size) + _domino_cross_entropy_backward_kernel[(num_rows,)]( + base_logits, + base_logits.stride(0), + correction_logits, + correction_logits.stride(0), + targets, + weights, + grad_final_loss, + grad_base_loss, + grad_base_logits, + grad_base_logits.stride(0), + grad_correction_logits, + grad_correction_logits.stride(0), + final_max, + final_exp_sum, + base_max, + base_exp_sum, + vocab_size, + BLOCK_SIZE=ctx.block_size, + SUFFIX_START=ctx.suffix_start, + VOCAB_BLOCK_SIZE=vocab_block_size, + num_warps=num_warps, + ) + return ( + grad_base_logits, + grad_correction_logits, + None, + None, + None, + None, + ) + + +def _calculate_domino_ce_settings(vocab_size): + """Choose the vocabulary tile and warp count for the active GPU backend.""" + vocab_block_size = min(triton.next_power_of_2(vocab_size), 2048) + num_warps = 8 if vocab_block_size >= 2048 else 4 + + # Preserve the NVIDIA thread count on AMD targets with 64-lane wavefronts. + # Note: this isn't tested, someone should tune this + if hasattr(torch.version, "hip") and torch.version.hip is not None: + warp_size = triton.runtime.driver.active.get_current_target().warp_size + num_warps = num_warps * 32 // warp_size + + return vocab_block_size, max(num_warps, 1) + + +@triton.jit +def _domino_cross_entropy_forward_kernel( + base_logits_ptr, + base_logits_stride, + correction_logits_ptr, + correction_logits_stride, + targets_ptr, + weights_ptr, + final_losses_ptr, + base_losses_ptr, + final_max_ptr, + final_exp_sum_ptr, + base_max_ptr, + base_exp_sum_ptr, + final_pred_ptr, + base_pred_ptr, + vocab_size, + BLOCK_SIZE: tl.constexpr, + SUFFIX_START: tl.constexpr, + VOCAB_BLOCK_SIZE: tl.constexpr, +): + """Write one set of outputs per flattened logit row. + + For each row, the kernel writes: + + - ``final_losses`` / ``base_losses``: weighted cross-entropy per row. + ``_DominoCrossEntropy.forward`` sums them for its scalar loss outputs. + - ``final_max`` / ``base_max``: largest logit in the row. + - ``final_exp_sum`` / ``base_exp_sum``: sum of + ``exp(logit - largest_logit)`` over the row. + - ``final_pred`` / ``base_pred``: index of the largest logit. + + For corrected logits ``[2, 5, 4]`` with target index 2 and weight 1: + + - ``final_max[row] = 5`` + - ``final_exp_sum[row] = exp(2 - 5) + exp(5 - 5) + exp(4 - 5)`` + - ``final_losses[row] = 5 + log(final_exp_sum[row]) - 4`` + - ``final_pred[row] = 1`` + + Backward reconstructs each probability as + ``exp(logit - final_max) / final_exp_sum``. The base outputs use the same + calculation with uncorrected base logits. + """ + row = tl.program_id(0).to(tl.int64) + has_correction, correction_row = _domino_correction_mapping( + row, + BLOCK_SIZE, + SUFFIX_START, + ) + + base_logits_ptr += row * base_logits_stride + correction_logits_ptr += correction_row * correction_logits_stride + target = tl.load(targets_ptr + row) + target_in_bounds = (target >= 0) & (target < vocab_size) + tl.device_assert(target_in_bounds, "target index is outside the vocabulary") + weight = tl.load(weights_ptr + row).to(tl.float32) + base_target_raw = tl.load( + base_logits_ptr + target, + mask=target_in_bounds, + other=float("nan"), + ) + correction_target_raw = tl.load( + correction_logits_ptr + target, + mask=has_correction & target_in_bounds, + other=0.0, + ) + base_target = base_target_raw.to(tl.float32) + final_target = (base_target_raw + correction_target_raw).to(tl.float32) + + final_max = float("-inf") + final_exp_sum = 0.0 + final_argmax = 0 + base_max = float("-inf") + base_exp_sum = 0.0 + base_argmax = 0 + + # Iterates over the one row of base logits. + # Each iteration: + # - loads one chunk of base_logits + # - loads the corresponding correction_logits chunk (or zeros if base-only row). + # - updates the {base/final}{max, exp_sum, argmax} variables + for i in range(0, vocab_size, VOCAB_BLOCK_SIZE): + offsets = i + tl.arange(0, VOCAB_BLOCK_SIZE) + mask = offsets < vocab_size + base_block_raw = tl.load( + base_logits_ptr + offsets, + mask=mask, + other=float("-inf"), + ) + correction_block_raw = tl.load( + correction_logits_ptr + offsets, + mask=mask & has_correction, + other=0.0, + ) + base_block = base_block_raw.to(tl.float32) + final_block = (base_block_raw + correction_block_raw).to(tl.float32) + + final_max, final_exp_sum, final_argmax = _update_online_logsumexp_and_argmax( + final_block, + mask, + i, + final_max, + final_exp_sum, + final_argmax, + ) + base_max, base_exp_sum, base_argmax = _update_online_logsumexp_and_argmax( + base_block, + mask, + i, + base_max, + base_exp_sum, + base_argmax, + ) + + final_loss = weight * (tl.log(final_exp_sum) + final_max - final_target) + base_loss = weight * (tl.log(base_exp_sum) + base_max - base_target) + tl.store(final_losses_ptr + row, final_loss) + tl.store(base_losses_ptr + row, base_loss) + tl.store(final_max_ptr + row, final_max.to(tl.float32)) + tl.store(final_exp_sum_ptr + row, final_exp_sum.to(tl.float32)) + tl.store(base_max_ptr + row, base_max.to(tl.float32)) + tl.store(base_exp_sum_ptr + row, base_exp_sum.to(tl.float32)) + tl.store(final_pred_ptr + row, final_argmax) + tl.store(base_pred_ptr + row, base_argmax) + + +@triton.jit +def _domino_cross_entropy_backward_kernel( + base_logits_ptr, + base_logits_stride, + correction_logits_ptr, + correction_logits_stride, + targets_ptr, + weights_ptr, + grad_final_loss_ptr, + grad_base_loss_ptr, + grad_base_logits_ptr, + grad_base_logits_stride, + grad_correction_logits_ptr, + grad_correction_logits_stride, + final_max_ptr, + final_exp_sum_ptr, + base_max_ptr, + base_exp_sum_ptr, + vocab_size, + BLOCK_SIZE: tl.constexpr, + SUFFIX_START: tl.constexpr, + VOCAB_BLOCK_SIZE: tl.constexpr, +): + """Reconstruct both softmaxes and write one row of input gradients. + + For one ``x[vocab_size]`` logit row and fixed scalar target index ``y``:: + + CE(x, y) = logsumexp(x) - x[y] + d CE(x, y) / d x[j] = softmax(x)[j] - 1[j == y] + + ``logsumexp(x)`` reduces the vocabulary row to one scalar, as does selecting + ``x[y]``. Their difference is the row's scalar cross-entropy; its gradient + with respect to ``x`` has shape ``[vocab_size]``. + + The kernel evaluates the softmax term stably using the forward statistics: + ``softmax(x)[j] = exp(x[j] - row_max) / row_exp_sum``. + + Including the row weight and upstream loss gradients gives:: + + grad_from_final_loss = grad_final_loss * weight + * (softmax(final_logits) - one_hot(y)) + grad_base_logits = grad_from_final_loss + + grad_base_loss * weight + * (softmax(base_logits) - one_hot(y)) + grad_correction_logits = grad_from_final_loss + + Any normalization applied to the returned loss sums is already included in + the incoming ``grad_final_loss`` and ``grad_base_loss`` values. + + Base logits feed both losses. Correction logits feed only the final loss + and exist only for suffix rows; prefix rows have no correction to update. + """ + row = tl.program_id(0).to(tl.int64) + has_correction, correction_row = _domino_correction_mapping( + row, + BLOCK_SIZE, + SUFFIX_START, + ) + + base_logits_ptr += row * base_logits_stride + grad_base_logits_ptr += row * grad_base_logits_stride + correction_logits_ptr += correction_row * correction_logits_stride + grad_correction_logits_ptr += correction_row * grad_correction_logits_stride + + target = tl.load(targets_ptr + row) + weight = tl.load(weights_ptr + row).to(tl.float32) + final_scale = tl.load(grad_final_loss_ptr).to(tl.float32) * weight + base_scale = tl.load(grad_base_loss_ptr).to(tl.float32) * weight + final_max = tl.load(final_max_ptr + row).to(tl.float32) + final_exp_sum = tl.load(final_exp_sum_ptr + row).to(tl.float32) + base_max = tl.load(base_max_ptr + row).to(tl.float32) + base_exp_sum = tl.load(base_exp_sum_ptr + row).to(tl.float32) + + # Iterates over one row of base logits. + # Each iteration: + # - loads base_logits and correction_logits chunks + # - reconstructs the base and final softmax probabilities + # - writes grad_base_logits with contributions from both losses + # - writes grad_correction_logits from the final loss for suffix rows + for i in range(0, vocab_size, VOCAB_BLOCK_SIZE): + offsets = i + tl.arange(0, VOCAB_BLOCK_SIZE) + mask = offsets < vocab_size + base_block_raw = tl.load( + base_logits_ptr + offsets, + mask=mask, + other=0.0, + ) + correction_block_raw = tl.load( + correction_logits_ptr + offsets, + mask=mask & has_correction, + other=0.0, + ) + base_block = base_block_raw.to(tl.float32) + final_block = (base_block_raw + correction_block_raw).to(tl.float32) + target_grad = tl.where(offsets == target, 1.0, 0.0) + final_grad = final_scale * ( + tl.exp(final_block - final_max) / final_exp_sum - target_grad + ) + base_grad = final_grad + base_scale * ( + tl.exp(base_block - base_max) / base_exp_sum - target_grad + ) + tl.store(grad_base_logits_ptr + offsets, base_grad, mask=mask) + tl.store( + grad_correction_logits_ptr + offsets, + final_grad, + mask=mask & has_correction, + ) + + +@triton.jit +def _domino_correction_mapping( + row, + BLOCK_SIZE: tl.constexpr, + SUFFIX_START: tl.constexpr, +): + """Map a flattened base row to Domino's compact correction layout. + + For ``BLOCK_SIZE=4`` and ``SUFFIX_START=1``:: + + base row: 0 1 2 3 | 4 5 6 7 + correction row: - 0 1 2 | - 3 4 5 + + ``-`` marks a base-only row with no correction. + """ + row_in_block = row % BLOCK_SIZE + suffix_size: tl.constexpr = BLOCK_SIZE - SUFFIX_START + has_correction = row_in_block >= SUFFIX_START + correction_row = (row // BLOCK_SIZE) * suffix_size + (row_in_block - SUFFIX_START) + correction_row = tl.where(has_correction, correction_row, 0) + return has_correction, correction_row + + +@triton.jit +def _update_online_logsumexp_and_argmax( + logits, + mask, + block_offset, + previous_max, + previous_exp_sum, + previous_argmax, +): + """Update the row maximum, shifted exponential sum, and leftmost argmax.""" + block_max, block_argmax = tl.max( + tl.where(mask, logits, float("-inf")), + axis=0, + return_indices=True, + return_indices_tie_break_left=True, + ) + block_argmax += block_offset + take_block = block_max > previous_max + argmax = tl.where(take_block, block_argmax, previous_argmax) + max_logit = tl.maximum(previous_max, block_max) + exp_sum = previous_exp_sum * tl.exp(previous_max - max_logit) + tl.sum( + tl.where(mask, tl.exp(logits - max_logit), 0.0) + ) + return max_logit, exp_sum, argmax diff --git a/specforge/core/eagle3_adapters.py b/specforge/core/eagle3_adapters.py index b03db8ee4..dfcac43ef 100644 --- a/specforge/core/eagle3_adapters.py +++ b/specforge/core/eagle3_adapters.py @@ -11,11 +11,15 @@ @dataclass -class StepState: +class BackboneStepState: input_ids: torch.Tensor hidden_states: torch.Tensor position_ids: torch.Tensor attention_mask: torch.Tensor + + +@dataclass +class StepState(BackboneStepState): target_p: torch.Tensor target_p_on_draft: torch.Tensor target_token_ids: torch.Tensor @@ -27,6 +31,30 @@ class BackendAdapter: def __init__(self, model: "OnlineEagle3Model"): self.m = model + def backbone_row_count(self, *, seq_length: int, ttt_length: int) -> int: + return seq_length + + def backbone_view( + self, + *, + row_count: int, + global_input_ids: torch.Tensor, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + ) -> BackboneStepState: + if row_count != hidden_states.shape[1]: + raise ValueError( + f"backbone row_count ({row_count}) must equal local hidden-state " + f"length ({hidden_states.shape[1]})" + ) + return BackboneStepState( + input_ids=global_input_ids, + hidden_states=hidden_states, + position_ids=position_ids, + attention_mask=attention_mask, + ) + def step_view( self, *, @@ -71,6 +99,13 @@ def step_view( position_mask: torch.Tensor, seq_length: int, ) -> StepState: + backbone = self.backbone_view( + row_count=seq_length, + global_input_ids=global_input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + hidden_states=hidden_states, + ) if target_p_on_draft_padded is None: target_p_on_draft_padded = target_p_padded if target_token_ids_padded is None: @@ -83,10 +118,10 @@ def step_view( :, idx : idx + seq_length ].contiguous() return StepState( - input_ids=global_input_ids, - hidden_states=hidden_states, - position_ids=position_ids, - attention_mask=attention_mask, + input_ids=backbone.input_ids, + hidden_states=backbone.hidden_states, + position_ids=backbone.position_ids, + attention_mask=backbone.attention_mask, target_p=target_p, target_p_on_draft=target_p_on_draft, target_token_ids=target_token_ids, @@ -103,6 +138,31 @@ def __init__(self, model: "OnlineEagle3Model"): self.ulysses_pg = get_sp_ulysses_group() self.sp_ulysses_degree = dist.get_world_size(self.ulysses_pg) + def backbone_row_count(self, *, seq_length: int, ttt_length: int) -> int: + row_count = seq_length - ttt_length + if row_count <= 0: + raise ValueError( + f"USP local seq_length ({seq_length}) must be larger than " + f"ttt_length ({ttt_length})" + ) + return row_count + + def backbone_view( + self, + *, + row_count: int, + global_input_ids: torch.Tensor, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + ) -> BackboneStepState: + return BackboneStepState( + input_ids=global_input_ids[:, :row_count], + hidden_states=hidden_states[:, :row_count, :], + position_ids=position_ids[:, : row_count * self.sp_ulysses_degree], + attention_mask=attention_mask[:, :row_count], + ) + def step_view( self, *, @@ -123,20 +183,24 @@ def step_view( target_p_on_draft_padded = target_p_padded if target_token_ids_padded is None: target_token_ids_padded = target_p_padded.argmax(dim=-1) - usp_chunk_size = seq_length - ttt_length - if usp_chunk_size <= 0: - raise ValueError( - f"USP local seq_length ({seq_length}) must be larger than " - f"ttt_length ({ttt_length})" - ) + usp_chunk_size = self.backbone_row_count( + seq_length=seq_length, ttt_length=ttt_length + ) + backbone = self.backbone_view( + row_count=usp_chunk_size, + global_input_ids=global_input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + hidden_states=hidden_states, + ) target_p = target_p_padded[:, idx : idx + usp_chunk_size, :] target_p_on_draft = target_p_on_draft_padded[:, idx : idx + usp_chunk_size, :] target_token_ids = target_token_ids_padded[:, idx : idx + usp_chunk_size] return StepState( - input_ids=global_input_ids[:, :usp_chunk_size], - hidden_states=hidden_states[:, :usp_chunk_size, :], - position_ids=position_ids[:, : usp_chunk_size * self.sp_ulysses_degree], - attention_mask=attention_mask[:, :usp_chunk_size], + input_ids=backbone.input_ids, + hidden_states=backbone.hidden_states, + position_ids=backbone.position_ids, + attention_mask=backbone.attention_mask, target_p=target_p, target_p_on_draft=target_p_on_draft, target_token_ids=target_token_ids, diff --git a/specforge/core/mtp.py b/specforge/core/mtp.py new file mode 100644 index 000000000..230a66edf --- /dev/null +++ b/specforge/core/mtp.py @@ -0,0 +1,144 @@ +# coding=utf-8 +"""Online training wrapper for single-layer MTP (architecture-independent). + +MTP predicts the next token from the current token's embedding plus the target +model's last hidden state. Shift is performed inside this wrapper; the target +backend is expected to return *raw* input_ids and last_hidden_states (DFlash +style), not the pre-shifted output of generate_eagle3_data. +""" + +from typing import List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class OnlineMTPModel(nn.Module): + """ + Online MTP training wrapper. + + Architecture-agnostic: any registered MTP draft module exposing + ``forward(input_ids, hidden_states, attention_mask, position_ids)`` and a + ``config`` with ``pad_token_id`` can be plugged in (see + ``specforge/modeling/draft/mtp/``). + + Args: + draft_model: The MTP draft model (e.g. ``modeling/draft/mtp/qwen3_5.py``). + ploss_decay: Per-layer loss decay. For a single MTP layer this is + unused, but kept for multi-layer extension. + """ + + def __init__( + self, + draft_model: nn.Module, + ploss_decay: float = 1.0, + ) -> None: + super().__init__() + self.draft_model = draft_model + self.ploss_decay = ploss_decay + + def _shift_for_next_token( + self, + logits: torch.Tensor, + input_ids: torch.Tensor, + loss_mask: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Shift logits/labels/mask to match vLLM speculative decoding. + + In serving, the draft model's input_ids are the target input_ids shifted + right by one (the draft fuses token x_{t+1} with target hidden state + h_t) and it predicts the token after that (x_{t+2}). Training therefore + uses: + - draft input: input_ids[:, 1:] (x_1..x_T, padded) + - label: x_2..x_T followed by a pad (length matches logits) + """ + shift_logits = logits[:, :-1, :].contiguous() + # x_2..x_T has length seq_len-2; pad one position so its length equals + # seq_len-1 (same as shift_logits). The padded position is ignored. + shift_labels = F.pad(input_ids[:, 2:], (0, 1), value=-100).contiguous() + shift_mask = F.pad(loss_mask[:, 2:], (0, 1), value=0).contiguous() + return shift_logits, shift_labels, shift_mask + + def forward( + self, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + loss_mask: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]: + """ + Args: + input_ids: raw token ids, [batch, seq_len]. + hidden_states: target model last hidden states, [batch, seq_len, hidden]. + loss_mask: [batch, seq_len]. + attention_mask: optional padding mask, [batch, seq_len]. + position_ids: optional position ids, [batch, seq_len]. + + Returns: + loss: scalar weighted loss. + acc_corrects: per-layer per-position correct tensors. + acc_denoms: per-layer per-position denominator tensors. + """ + # Draft input is the target sequence shifted right by one. The last + # position is padded because there is no x_{T+1}; the corresponding + # logit is dropped in _shift_for_next_token. + pad_token_id = getattr(self.draft_model.config, "pad_token_id", 0) + shifted_input_ids = F.pad(input_ids[:, 1:], (0, 1), value=pad_token_id) + + # The padding mask must follow the same shift so the synthetic pad token + # at the last position is not attended to. + if attention_mask is not None: + shifted_attention_mask = F.pad(attention_mask[:, 1:], (0, 1), value=0).to( + attention_mask.dtype + ) + else: + shifted_attention_mask = None + + # Serving evaluates the shifted draft token x[t+1] at its own position + # p[t+1], even though it is fused with the target hidden state h[t]. + # Preserve caller-supplied offsets (for packed/non-zero-based sequences) + # and give the synthetic final token the next monotonic position. + batch_size, seq_len = input_ids.shape + if position_ids is None: + position_ids = ( + torch.arange(seq_len, dtype=torch.long, device=input_ids.device) + .unsqueeze(0) + .expand(batch_size, -1) + ) + elif position_ids.shape != input_ids.shape: + raise ValueError( + "position_ids must have the same [batch, seq_len] shape as " + f"input_ids; got {tuple(position_ids.shape)} and " + f"{tuple(input_ids.shape)}" + ) + shifted_position_ids = torch.cat( + (position_ids[:, 1:], position_ids[:, -1:] + 1), dim=1 + ) + + outputs = self.draft_model( + input_ids=shifted_input_ids, + hidden_states=hidden_states, + attention_mask=shifted_attention_mask, + position_ids=shifted_position_ids, + ) + logits = outputs.logits + + shift_logits, shift_labels, shift_mask = self._shift_for_next_token( + logits, input_ids, loss_mask + ) + + flat_logits = shift_logits.view(-1, shift_logits.size(-1)) + flat_labels = shift_labels.view(-1) + losses = F.cross_entropy(flat_logits, flat_labels, reduction="none") + losses = losses * shift_mask.view(-1).float() + loss = losses.sum() / shift_mask.sum().clamp_min(1) + + with torch.no_grad(): + preds = shift_logits.argmax(dim=-1) + corrects = (preds == shift_labels).float() * shift_mask.float() + denoms = shift_mask.float() + + # Single-layer MTP: wrap in length-1 lists for E1 evaluator compatibility. + return loss, [corrects], [denoms] diff --git a/specforge/data/loss_mask.py b/specforge/data/loss_mask.py new file mode 100644 index 000000000..d52b30240 --- /dev/null +++ b/specforge/data/loss_mask.py @@ -0,0 +1,21 @@ +"""Loss-mask predicates shared by online and offline data paths.""" + +from collections.abc import Sequence +from typing import Any + + +def has_consecutive_supervised_tokens(loss_mask: Any) -> bool: + """Return whether one sample contains adjacent supervised tokens.""" + + values = loss_mask.tolist() if hasattr(loss_mask, "tolist") else list(loss_mask) + if values and isinstance(values[0], Sequence): + if len(values) != 1: + raise ValueError("expected one loss-mask row") + values = list(values[0]) + return any( + bool(current) and bool(following) + for current, following in zip(values, values[1:]) + ) + + +__all__ = ["has_consecutive_supervised_tokens"] diff --git a/specforge/data/preprocessing.py b/specforge/data/preprocessing.py index 5d27dea77..a2f268101 100644 --- a/specforge/data/preprocessing.py +++ b/specforge/data/preprocessing.py @@ -27,7 +27,7 @@ import re import warnings from collections import Counter -from typing import Dict, List, Optional, Tuple, Union +from typing import Callable, Dict, List, Optional, Tuple, Union import torch import torch.nn.functional as F @@ -36,6 +36,7 @@ from transformers import PreTrainedTokenizer from ..distributed import get_draft_sp_group, get_sp_ring_group +from .loss_mask import has_consecutive_supervised_tokens from .parse import GeneralParser, GLMParser, HarmonyParser, ThinkingParser from .template import TEMPLATE_REGISTRY, ChatTemplate @@ -180,6 +181,7 @@ def build_eagle3_dataset( is_preformatted: Optional[bool] = False, train_only_last_turn: Optional[bool] = False, minimum_valid_tokens: Optional[int] = None, + loss_mask_filter: Optional[Callable[[object], bool]] = None, ) -> HFDataset: """ build eagle3 dataset @@ -206,12 +208,16 @@ def build_eagle3_dataset( train_only_last_turn: If True, only the last assistant turn contributes to the loss. Useful for thinking models where history may not contain thoughts. minimum_valid_tokens: If set, drops samples with fewer trainable tokens. + loss_mask_filter: Optional algorithm-owned predicate applied after + tokenization and truncation. Returns: The processed HF dataset. """ if minimum_valid_tokens is not None and minimum_valid_tokens < 0: raise ValueError("minimum_valid_tokens must be >= 0") + if loss_mask_filter is not None and not callable(loss_mask_filter): + raise TypeError("loss_mask_filter must be callable or None") # Validate chat_template requirement if chat_template is None: @@ -344,6 +350,21 @@ def has_minimum_valid_tokens(example): f"Filtered dataset by trainable tokens: {before_filter} -> {len(dataset)}" ) + if loss_mask_filter is not None: + before_filter = len(dataset) + + def has_eligible_loss_mask(example): + return loss_mask_filter(example["loss_mask"]) + + dataset = dataset.filter( + has_eligible_loss_mask, + num_proc=num_proc, + desc="Filtering samples by algorithm loss-mask eligibility", + ) + print( + f"Filtered dataset by loss-mask eligibility: {before_filter} -> {len(dataset)}" + ) + dataset.set_format(type="torch") return dataset @@ -685,6 +706,10 @@ def process_offline_dflash_sample( f"loss_mask={loss_mask.shape[1]}, " f"hidden_states={hidden_states.shape[1]}" ) + if not has_consecutive_supervised_tokens(loss_mask[0]): + raise ValueError( + "offline DFlash samples require two consecutive supervised tokens" + ) return { "input_ids": input_ids, "loss_mask": loss_mask, diff --git a/specforge/data/prompt_builder.py b/specforge/data/prompt_builder.py index 74df210cb..a3461d6c2 100644 --- a/specforge/data/prompt_builder.py +++ b/specforge/data/prompt_builder.py @@ -12,7 +12,7 @@ import os from collections.abc import Iterable, Iterator, Mapping, Sequence from numbers import Integral -from typing import Any +from typing import Any, Callable PromptTaskDict = dict[str, Any] @@ -30,14 +30,17 @@ def prepare_prompt_tasks( num_proc: int | None, min_loss_tokens: int = 1, max_prompts: int | None = None, -) -> list[PromptTaskDict]: + loss_mask_filter: Callable[[Sequence[int]], bool] | None = None, +) -> Sequence[PromptTaskDict]: """Prepare runtime prompt dictionaries from a JSONL file. Each returned item has the control-plane shape ``{"payload": {"input_ids": [...], "loss_mask": [...]}}`` and contains no tensors. Files whose first record contains ``input_ids`` and ``loss_mask`` are treated as pre-tokenized. Other files are treated as raw conversation - data and processed through :func:`build_eagle3_dataset`. + data and processed through :func:`build_eagle3_dataset`, then exposed as a + lazy random-access sequence so large Arrow datasets are not expanded into + Python token lists before rollout starts. ``max_prompts`` caps accepted prompts; ``None`` and ``0`` mean no cap. """ @@ -48,6 +51,7 @@ def prepare_prompt_tasks( max_prompts=max_prompts, cache_dir=cache_dir, cache_key=cache_key, + loss_mask_filter=loss_mask_filter, ) path_string = os.fspath(path) first_record = next(_iter_records(path_string), None) @@ -75,6 +79,7 @@ def prepare_prompt_tasks( max_length=max_length, min_loss_tokens=min_loss_tokens, limit=limit, + loss_mask_filter=loss_mask_filter, ) return _prepare_raw_prompts( @@ -89,6 +94,7 @@ def prepare_prompt_tasks( num_proc=num_proc, min_loss_tokens=min_loss_tokens, limit=limit, + loss_mask_filter=None, ) @@ -105,7 +111,8 @@ def _prepare_raw_prompts( num_proc: int | None, min_loss_tokens: int, limit: int | None, -) -> list[PromptTaskDict]: + loss_mask_filter: Callable[[Sequence[int]], bool] | None, +) -> Sequence[PromptTaskDict]: try: from datasets import load_dataset except ImportError as exc: # pragma: no cover - package dependency in production @@ -116,7 +123,7 @@ def _prepare_raw_prompts( from .preprocessing import build_eagle3_dataset dataset = load_dataset("json", data_files=path, split="train") - if limit is not None and limit < len(dataset): + if loss_mask_filter is None and limit is not None and limit < len(dataset): dataset = dataset.select(range(limit)) processed_dataset = build_eagle3_dataset( @@ -130,63 +137,121 @@ def _prepare_raw_prompts( is_preformatted=is_preformatted, train_only_last_turn=train_only_last_turn, minimum_valid_tokens=min_loss_tokens, + loss_mask_filter=loss_mask_filter, ) - rows = ( - (record, f"processed dataset row {index}") - for index, record in enumerate(processed_dataset) - ) - return _materialize_prompt_tasks( - rows, + return _ProcessedPromptSequence( + processed_dataset, max_length=max_length, min_loss_tokens=min_loss_tokens, - limit=limit, + loss_mask_filter=loss_mask_filter, ) +class _ProcessedPromptSequence(Sequence[PromptTaskDict]): + """Normalize memory-mapped processed rows only when the producer ingests them.""" + + def __init__( + self, + dataset, + *, + max_length: int, + min_loss_tokens: int, + loss_mask_filter: Callable[[Sequence[int]], bool] | None, + ) -> None: + self._dataset = dataset + self._max_length = max_length + self._min_loss_tokens = min_loss_tokens + self._loss_mask_filter = loss_mask_filter + + def __len__(self) -> int: + return len(self._dataset) + + def __iter__(self) -> Iterator[PromptTaskDict]: + for index in range(len(self)): + yield self[index] + + def __getitem__(self, index): + if isinstance(index, slice): + return [self[item] for item in range(*index.indices(len(self)))] + record = self._dataset[index] + prompt = _prompt_from_record( + record, + source=f"processed dataset row {index}", + max_length=self._max_length, + min_loss_tokens=self._min_loss_tokens, + ) + if prompt is None: + raise ValueError( + f"processed dataset row {index} violates the preprocessing " + f"minimum of {self._min_loss_tokens} trainable tokens" + ) + return prompt + + def _materialize_prompt_tasks( rows: Iterable[tuple[Mapping[str, Any], str]], *, max_length: int, min_loss_tokens: int, limit: int | None, + loss_mask_filter: Callable[[Sequence[int]], bool] | None, ) -> list[PromptTaskDict]: prompts: list[PromptTaskDict] = [] for record, source in rows: - if "input_ids" not in record or "loss_mask" not in record: - raise ValueError(f"{source} must contain both input_ids and loss_mask") - - input_ids = _normalize_integer_sequence( - record["input_ids"], field="input_ids", source=source, binary=False - ) - loss_mask = _normalize_integer_sequence( - record["loss_mask"], field="loss_mask", source=source, binary=True + prompt = _prompt_from_record( + record, + source=source, + max_length=max_length, + min_loss_tokens=min_loss_tokens, ) - if len(input_ids) != len(loss_mask): - raise ValueError( - f"{source} has mismatched input_ids/loss_mask lengths: " - f"{len(input_ids)} != {len(loss_mask)}" - ) - if not input_ids: - raise ValueError(f"{source} contains an empty token sequence") - - input_ids = input_ids[:max_length] - loss_mask = loss_mask[:max_length] - if sum(loss_mask) < min_loss_tokens: + if prompt is None: continue - - prompts.append( - { - "payload": { - "input_ids": input_ids, - "loss_mask": loss_mask, - } - } - ) + if loss_mask_filter is not None and not loss_mask_filter( + prompt["payload"]["loss_mask"] + ): + continue + prompts.append(prompt) if limit is not None and len(prompts) >= limit: break return prompts +def _prompt_from_record( + record: Mapping[str, Any], + *, + source: str, + max_length: int, + min_loss_tokens: int, +) -> PromptTaskDict | None: + if "input_ids" not in record or "loss_mask" not in record: + raise ValueError(f"{source} must contain both input_ids and loss_mask") + + input_ids = _normalize_integer_sequence( + record["input_ids"], field="input_ids", source=source, binary=False + ) + loss_mask = _normalize_integer_sequence( + record["loss_mask"], field="loss_mask", source=source, binary=True + ) + if len(input_ids) != len(loss_mask): + raise ValueError( + f"{source} has mismatched input_ids/loss_mask lengths: " + f"{len(input_ids)} != {len(loss_mask)}" + ) + if not input_ids: + raise ValueError(f"{source} contains an empty token sequence") + + input_ids = input_ids[:max_length] + loss_mask = loss_mask[:max_length] + if sum(loss_mask) < min_loss_tokens: + return None + return { + "payload": { + "input_ids": input_ids, + "loss_mask": loss_mask, + } + } + + def _normalize_integer_sequence( value: Any, *, @@ -275,6 +340,7 @@ def _validate_options( max_prompts: int | None, cache_dir: str | None, cache_key: str | None, + loss_mask_filter: Callable[[Sequence[int]], bool] | None, ) -> None: if ( not isinstance(max_length, int) @@ -300,6 +366,8 @@ def _validate_options( ) if (cache_dir is None) != (cache_key is None): raise ValueError("cache_dir and cache_key must be provided together") + if loss_mask_filter is not None and not callable(loss_mask_filter): + raise TypeError("loss_mask_filter must be callable or None") __all__ = ["PromptTaskDict", "prepare_prompt_tasks"] diff --git a/specforge/data/template.py b/specforge/data/template.py index cfb409829..26d1e8218 100644 --- a/specforge/data/template.py +++ b/specforge/data/template.py @@ -249,6 +249,25 @@ def get_all_template_names(self) -> List[str]: ), ) +# Kimi K3 uses the checkpoint tokenizer's XTML renderer rather than a Jinja +# template. The rendered assistant turn opens the thinking segment before the +# stored reasoning content, so supervision starts after this exact scaffold and +# excludes the stop-trimmed end-of-message token. +TEMPLATE_REGISTRY.register( + name="kimi-k3-thinking", + template=ChatTemplate( + assistant_header=( + '<|open|>message role="assistant"<|sep|><|open|>think<|sep|>' + ), + user_header='<|open|>message role="user"<|sep|>', + system_prompt=None, + end_of_turn_token="<|end_of_msg|>", + parser_type="thinking", + enable_thinking=False, + ignore_token=["<|end_of_msg|>"], + ), +) + TEMPLATE_REGISTRY.register( name="deepseek-v3", template=ChatTemplate( diff --git a/specforge/distributed.py b/specforge/distributed.py index 4016c0f4a..d8e7d00ca 100644 --- a/specforge/distributed.py +++ b/specforge/distributed.py @@ -145,10 +145,15 @@ def init_distributed( # initialization; doing the same for NCCL also removes ambiguous rank/device # inference on heterogeneous hosts. local_rank = _bind_local_device(device_type) - # Yunchang probes the active CUDA device while importing. Keep it behind - # the trainer-only, device-bound initialization boundary so config loading - # and prompt preprocessing remain safe in CPU-only producer processes. - process_group, set_seq_parallel_pg = _load_yunchang_globals() + # Yunchang probes the active CUDA device at import, which crashes + # NPU-only torch builds. Only USP needs it, so load it lazily when SP + # sizes exceed 1. For SP=1, the getters are populated from the singleton + # DeviceMesh group below; publishing None would make torch.distributed + # operations silently use the default WORLD group instead. + if sp_ulysses_size * sp_ring_size > 1: + process_group, set_seq_parallel_pg = _load_yunchang_globals() + else: + process_group, set_seq_parallel_pg = None, None dist.init_process_group(backend=backend, timeout=timedelta(minutes=timeout)) print_with_rank(f"bind to {device_type} device {local_rank}") @@ -173,14 +178,23 @@ def init_distributed( (draft_dp_size, sp_ulysses_size * sp_ring_size), mesh_dim_names=("draft_dp", "sp"), ) - set_seq_parallel_pg(sp_ulysses_size, sp_ring_size, dist.get_rank(), world_size) + if set_seq_parallel_pg is not None: + set_seq_parallel_pg(sp_ulysses_size, sp_ring_size, dist.get_rank(), world_size) print_with_rank(f"device mesh: {device_mesh}") tp_group = device_mesh.get_group("tp") dp_group = device_mesh.get_group("dp") + draft_dp_group = draft_device_mesh.get_group("draft_dp") + draft_sp_group = draft_device_mesh.get_group("sp") - sp_ulysses_group = process_group.ULYSSES_PG - sp_ring_group = process_group.RING_PG + if process_group is not None: + sp_ulysses_group = process_group.ULYSSES_PG + sp_ring_group = process_group.RING_PG + else: + # Both SP dimensions are 1 here. Reuse the per-rank singleton SP group + # so callers always receive a real group whose world size is 1. + sp_ulysses_group = draft_sp_group + sp_ring_group = draft_sp_group # we need to create a 1D submesh tp_device_mesh = dist.DeviceMesh.from_group(tp_group, device_type=device_type) @@ -191,8 +205,8 @@ def init_distributed( _SP_ULYSSES_GROUP = sp_ulysses_group _SP_RING_GROUP = sp_ring_group _DP_GROUP = dp_group - _DRAFT_DP_GROUP = draft_device_mesh.get_group("draft_dp") - _DRAFT_SP_GROUP = draft_device_mesh.get_group("sp") + _DRAFT_DP_GROUP = draft_dp_group + _DRAFT_SP_GROUP = draft_sp_group _DP_DEVICE_MESH = dist.DeviceMesh.from_group(dp_group, device_type=device_type) diff --git a/specforge/eval/evaluator.py b/specforge/eval/evaluator.py index 9b10a7400..1f560f300 100644 --- a/specforge/eval/evaluator.py +++ b/specforge/eval/evaluator.py @@ -39,10 +39,10 @@ def run( ) -> Dict[str, Any]: """Run the pass; returns ``{}`` if zero batches were processed globally. - Scalar accuracy is weighted by ``metrics['accuracy_denom']`` when present, - else by the loss-token count — only approximately batch-size invariant - when the accuracy counts a different token set than the loss. In a mixed - pass, scalar batches feed avg_loss only; their accuracy is not merged. + Additive loss and accuracy terms are used directly when the strategy + provides them. Scalar fallbacks use ``metrics['accuracy_denom']`` when + present, else the loss-token count. In a mixed pass, scalar batches feed + avg_loss only; their accuracy is not merged. """ # pp rows: [correct, denom, acceptance_rate*w, ploss*w] per TTT # position, float64 so counts stay exact past 2**24. @@ -63,8 +63,13 @@ def run( if sums is None: sums = torch.zeros(7, dtype=torch.float64, device=loss.device) tokens = self._token_count(batch, m, device=sums.device) - sums[0] += loss.to(sums.device) * tokens - sums[1] += tokens + if out.loss_terms is None: + sums[0] += loss.to(sums.device) * tokens + sums[1] += tokens + else: + loss_numerator, loss_denominator = out.loss_terms + sums[0] += self._sum64(loss_numerator, sums.device) + sums[1] += self._sum64(loss_denominator, sums.device) sums[4] += 1.0 if "acc_corrects" in m and "acc_denoms" in m: @@ -86,6 +91,10 @@ def run( if "plosses" in m: pp[3] += self._stack(m["plosses"]) * w sums[6] += tokens + elif "acc" in out.ratio_metrics: + accuracy_numerator, accuracy_denominator = out.ratio_metrics["acc"] + sums[2] += self._sum64(accuracy_numerator, sums.device) + sums[3] += self._sum64(accuracy_denominator, sums.device) elif "accuracy" in m: acc = m["accuracy"] acc = ( @@ -96,11 +105,7 @@ def run( ) ) denom = m.get("accuracy_denom") - w = ( - torch.as_tensor(denom).detach().double().sum().to(sums.device) - if denom is not None - else tokens - ) + w = self._sum64(denom, sums.device) if denom is not None else tokens sums[2] += acc * w sums[3] += w @@ -162,6 +167,10 @@ def run( def _stack(values: Iterable[Any]) -> torch.Tensor: return torch.stack([torch.as_tensor(v).detach().double() for v in values]) + @staticmethod + def _sum64(value: Any, device: torch.device) -> torch.Tensor: + return torch.as_tensor(value).detach().double().sum().to(device) + @staticmethod def _comm_device() -> torch.device: """Return the bound device required by the active collective backend.""" diff --git a/specforge/export/checkpoint_io.py b/specforge/export/checkpoint_io.py index 2c294c0f9..87aad6929 100644 --- a/specforge/export/checkpoint_io.py +++ b/specforge/export/checkpoint_io.py @@ -11,6 +11,7 @@ from __future__ import annotations import glob +import json import os import re from typing import Any, Dict, Optional @@ -18,6 +19,75 @@ import torch STATE_FILE = "training_state.pt" +_DISABLE_LEGACY_ROPE_SCALING_ENV = "SPECFORGE_DISABLE_LEGACY_ROPE_SCALING" + + +def _env_flag_enabled(name: str) -> bool: + value = os.environ.get(name) + if value is None: + return False + return value.strip().lower() in {"1", "true", "yes", "on"} + + +def apply_legacy_rope_scaling(output_dir: str) -> bool: + """Keep modern and legacy RoPE scaling fields compatible in an export. + + Transformers 5 writes ``rope_parameters`` while older serving stacks read + only ``rope_scaling`` and the top-level ``rope_theta``. Mirror non-default + scaling representations and preserve the modern ``rope_theta`` for legacy + readers. On by default; set + ``SPECFORGE_DISABLE_LEGACY_ROPE_SCALING=1`` to skip. Returns whether the + config was rewritten. + """ + if _env_flag_enabled(_DISABLE_LEGACY_ROPE_SCALING_ENV): + return False + + config_path = os.path.join(output_dir, "config.json") + with open(config_path, encoding="utf-8") as handle: + config = json.load(handle) + + rope_parameters = config.get("rope_parameters") + rope_scaling = config.get("rope_scaling") + + def rope_kind(payload): + return (payload or {}).get("rope_type") or (payload or {}).get("type") + + changed = False + if rope_parameters and "rope_theta" in rope_parameters: + rope_theta = rope_parameters["rope_theta"] + if config.get("rope_theta") != rope_theta: + config["rope_theta"] = rope_theta + changed = True + + if ( + rope_parameters + and not rope_scaling + and rope_kind(rope_parameters) not in (None, "default") + ): + config["rope_scaling"] = { + key: value for key, value in rope_parameters.items() if key != "rope_theta" + } + changed = True + elif ( + rope_scaling + and not rope_parameters + and rope_kind(rope_scaling) not in (None, "default") + ): + mirrored = dict(rope_scaling) + if "rope_theta" in config: + mirrored.setdefault("rope_theta", config["rope_theta"]) + config["rope_parameters"] = mirrored + changed = True + + if not changed: + return False + + temporary = f"{config_path}.{os.getpid()}.tmp" + with open(temporary, "w", encoding="utf-8") as handle: + json.dump(config, handle, indent=2, sort_keys=True) + handle.write("\n") + os.replace(temporary, config_path) + return True def resolve_training_state(checkpoint_path: str) -> Dict[str, Any]: @@ -111,4 +181,9 @@ def materialize_draft( return model -__all__ = ["resolve_training_state", "materialize_draft", "STATE_FILE"] +__all__ = [ + "apply_legacy_rope_scaling", + "resolve_training_state", + "materialize_draft", + "STATE_FILE", +] diff --git a/specforge/export/mtp.py b/specforge/export/mtp.py new file mode 100644 index 000000000..49a665066 --- /dev/null +++ b/specforge/export/mtp.py @@ -0,0 +1,381 @@ +# coding=utf-8 +"""Merge a trained MTP draft checkpoint back into the base target checkpoint. + +This module owns the MTP-specific merge policy (native key prefix handling, +shared-embedding backfill, config patching). The model-independent merge +machinery — copying non-weight files, replacing keys by prefix, shard/index +writing — lives in ``specforge/modeling/target/checkpoint.py``. +Architecture-specific knowledge (target-side embed/lm_head key candidates, +native key prefix) comes from the registered MTP draft class — see +``specforge/modeling/draft/mtp/``. +""" + +from __future__ import annotations + +import glob +import json +import os +import shutil +from typing import Dict, List, Optional, Tuple + +import torch + +from specforge.modeling.target.checkpoint import ( + load_selected_tensors, + load_tensors_by_keys, + merge_state_into_checkpoint, +) + + +def _default_key_candidates() -> Tuple[List[str], List[str], str]: + """Base-class defaults, imported lazily to keep this module import-light.""" + + from specforge.modeling.draft.mtp.base import MTPDraftModel + + return ( + list(MTPDraftModel.TARGET_EMBED_KEY_CANDIDATES), + list(MTPDraftModel.TARGET_HEAD_KEY_CANDIDATES), + MTPDraftModel.NATIVE_KEY_PREFIX, + ) + + +def _resolve_key_candidates( + draft_config_source: str, +) -> Tuple[List[str], List[str], str]: + """Return (embed candidates, head candidates, native prefix) for the draft. + + Reads the draft ``config.json`` and resolves its + ``architectures[0]`` through the draft registry, so each MTP family can + override its target-side key candidates on the draft class. + """ + + embed, head, prefix = _default_key_candidates() + config_path = ( + draft_config_source + if os.path.isfile(draft_config_source) + else os.path.join(draft_config_source, "config.json") + ) + if os.path.exists(config_path): + try: + with open(config_path, "r") as f: + architectures = json.load(f).get("architectures") or [] + if architectures: + from specforge.modeling.draft.registry import DRAFT_REGISTRY + + draft_cls = DRAFT_REGISTRY.get(architectures[0]) + if draft_cls is not None: + embed = list( + getattr(draft_cls, "TARGET_EMBED_KEY_CANDIDATES", embed) + ) + head = list(getattr(draft_cls, "TARGET_HEAD_KEY_CANDIDATES", head)) + prefix = getattr(draft_cls, "NATIVE_KEY_PREFIX", prefix) + except Exception as exc: # pragma: no cover - defensive + print(f" warning: could not resolve draft key candidates: {exc}") + return embed, head, prefix + + +def convert_mtp_keys( + state_dict: Dict[str, torch.Tensor], fmt: str, prefix: str = "mtp." +) -> Dict[str, torch.Tensor]: + """Convert MTP weight keys to the requested output format. + + Training already saves the flat native layout that both SGLang and + HF/vLLM MTP modules expect, so ``sglang`` and ``hf`` both return it + unchanged; ``fmt`` is kept for backward compatibility. A legacy nested + layout (``mtp.model.layers.0.*``) is normalized to flat. + """ + + converted = {} + for k, v in state_dict.items(): + # Normalize legacy nested keys (mtp.model.layers.* -> mtp.layers.*). + if k.startswith(f"{prefix}model.layers."): + new_k = k.replace(f"{prefix}model.layers.", f"{prefix}layers.", 1) + elif k == f"{prefix}model.norm.weight": + new_k = f"{prefix}norm.weight" + # Promote bare embed_tokens / lm_head saved by the training script to the + # native namespace expected by vLLM/SGLang. + elif k == "embed_tokens.weight": + new_k = f"{prefix}embed_tokens.weight" + elif k == "lm_head.weight": + new_k = f"{prefix}lm_head.weight" + else: + new_k = k + converted[new_k] = v + return _unshare_storage(converted) + + +def _unshare_storage(state: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """Clone tensors whose storage is aliased under another key. + + A tied target shares one Parameter between ``embed_tokens.weight`` and + ``mtp.lm_head.weight``; after promotion both keys keep aliasing the same + storage, which safetensors' ``save_file`` rejects ("tensors share + memory"). Cloning the later alias preserves the values while giving every + key its own storage. + """ + + seen: set[int] = set() + out: Dict[str, torch.Tensor] = {} + for key, value in state.items(): + ptr = value.untyped_storage().data_ptr() + if ptr in seen: + value = value.clone() + else: + seen.add(ptr) + out[key] = value + return out + + +def _find_base_key(state_dict: Dict[str, torch.Tensor], *candidates: str) -> str | None: + """Return the first candidate key that exists in ``state_dict``.""" + + for key in candidates: + if key in state_dict: + return key + return None + + +def _copy_shared_embeddings( + base_state: Dict[str, torch.Tensor], + mtp_state: Dict[str, torch.Tensor], + tie_word_embeddings: bool, + embed_key_candidates: List[str], + head_key_candidates: List[str], + prefix: str, +) -> Dict[str, torch.Tensor]: + """Copy base embed_tokens/lm_head into the MTP state if they are missing. + + During training the draft model typically shares ``embed_tokens`` and + ``lm_head`` with the target model, so the saved MTP checkpoint does not + contain those tensors. vLLM/SGLang, however, instantiate their own + ``mtp.embed_tokens`` (and a separate ``lm_head`` when weights are not tied), + and expect them in the checkpoint. Copying them from the base model keeps + the merged checkpoint self-contained and avoids random-initialization of the + MTP input/output embeddings at serving time. + """ + + embed_target = f"{prefix}embed_tokens.weight" + head_target = f"{prefix}lm_head.weight" + + if embed_target not in mtp_state: + embed_key = _find_base_key(base_state, *embed_key_candidates) + if embed_key: + mtp_state[embed_target] = base_state[embed_key] + print(f" copied {embed_key} -> {embed_target}") + else: + print( + " warning: base embed_tokens.weight not found; " + f"{embed_target} will be randomly initialized" + ) + + if not tie_word_embeddings and head_target not in mtp_state: + lm_head_key = _find_base_key(base_state, *head_key_candidates) + if lm_head_key: + mtp_state[head_target] = base_state[lm_head_key] + print(f" copied {lm_head_key} -> {head_target}") + else: + print( + " warning: base lm_head.weight not found; " + f"{head_target} will be randomly initialized" + ) + + return mtp_state + + +def _patch_text_config(base_config: dict, draft_config: dict) -> dict: + """Ensure base text_config contains MTP-critical dims from the draft config. + + Some Qwen3.5 base checkpoints omit ``head_dim`` in ``text_config``; vLLM's + ``Qwen3_5TextConfig`` then falls back to its default (``head_dim=256``), + which mismatches the trained MTP weights (e.g. q_norm/k_norm shape 128). + Only the structural dims that must agree between base and draft are synced. + """ + + keys_to_sync = [ + "head_dim", + "hidden_size", + "intermediate_size", + "num_attention_heads", + "num_key_value_heads", + ] + + target = base_config + if "text_config" in base_config: + target = base_config["text_config"] + + source = draft_config + if "text_config" in draft_config: + source = draft_config["text_config"] + + for key in keys_to_sync: + if key not in source: + continue + old = target.get(key) + new = source[key] + if old != new: + target[key] = new + print(f" overriding text_config.{key}: {old} -> {new}") + + return base_config + + +def _load_first_checkpoint(checkpoint_dir: str) -> Dict[str, torch.Tensor]: + """Load every tensor of a single-file checkpoint directory.""" + + safetensors = glob.glob(os.path.join(checkpoint_dir, "*.safetensors")) + bins = glob.glob(os.path.join(checkpoint_dir, "*.bin")) + if safetensors: + return load_selected_tensors(checkpoint_dir, lambda _key: True) + if bins: + return torch.load(bins[0], map_location="cpu", weights_only=True) + raise FileNotFoundError(f"No safetensors/bin weights found in {checkpoint_dir}") + + +def _has_model_weights(path: str) -> bool: + """Return whether ``path`` is already an exported model directory.""" + + if not os.path.isdir(path): + return False + patterns = ("model*.safetensors", "pytorch_model*.bin") + return any(glob.glob(os.path.join(path, pattern)) for pattern in patterns) + + +def _load_mtp_source( + checkpoint_path: str, + draft_config_path: Optional[str], +) -> Tuple[Dict[str, torch.Tensor], str, Optional[str]]: + """Load MTP weights from either runtime state or an exported draft. + + Returns ``(state_dict, config_source, model_source_dir)``. The last item is + set only for an exported model directory, where companion modeling files + may also need to be copied. + """ + + path = checkpoint_path + if path.startswith("file://"): + path = path[len("file://") :] + if _has_model_weights(path): + return _load_first_checkpoint(path), path, path + + from specforge.export.checkpoint_io import resolve_training_state + + state = resolve_training_state(checkpoint_path) + if state.get("strategy") != "mtp": + raise ValueError( + "MTP merge requires a training checkpoint written by strategy='mtp'; " + f"got strategy={state.get('strategy')!r}" + ) + draft_state = state.get("draft_state_dict") + if not isinstance(draft_state, dict): + raise ValueError("MTP training checkpoint has no draft_state_dict") + if not draft_config_path: + raise ValueError( + "draft_config_path is required when merging a runtime training " + "checkpoint" + ) + if not os.path.isfile(draft_config_path): + raise FileNotFoundError(f"draft config not found: {draft_config_path}") + return dict(draft_state), draft_config_path, None + + +def merge_mtp_into_base( + base_model_path: str, + mtp_checkpoint_path: str, + output_path: str, + key_format: str = "sglang", + *, + draft_config_path: Optional[str] = None, +) -> None: + """Merge trained MTP weights into a copy of the base checkpoint. + + The output directory is a self-contained HF checkpoint loadable directly by + SGLang's native MTP modules (no separate draft-model path). Runtime + checkpoints require ``draft_config_path``; an exported HF draft supplies its + own ``config.json``. + """ + + mtp_state, config_source, model_source_dir = _load_mtp_source( + mtp_checkpoint_path, draft_config_path + ) + embed_key_candidates, head_key_candidates, prefix = _resolve_key_candidates( + config_source + ) + os.makedirs(output_path, exist_ok=True) + + mtp_state = convert_mtp_keys(mtp_state, key_format, prefix) + + # Determine whether word embeddings are tied to decide whether a separate + # lm_head must be materialized for the MTP module. + tie_word_embeddings = True + base_config_path = os.path.join(base_model_path, "config.json") + if os.path.exists(base_config_path): + with open(base_config_path, "r") as f: + base_cfg = json.load(f) + # VLM checkpoints nest text config under "text_config". + text_cfg = base_cfg.get("text_config", base_cfg) + tie_word_embeddings = text_cfg.get("tie_word_embeddings", True) + + # If the trained checkpoint did not save shared embeddings, copy them from + # the base checkpoint so vLLM/SGLang can initialise the MTP embed_tokens/ + # lm_head from the merged checkpoint. + embed_target = f"{prefix}embed_tokens.weight" + head_target = f"{prefix}lm_head.weight" + if embed_target not in mtp_state or ( + not tie_word_embeddings and head_target not in mtp_state + ): + base_state = load_tensors_by_keys( + base_model_path, embed_key_candidates + head_key_candidates + ) + mtp_state = _copy_shared_embeddings( + base_state, + mtp_state, + tie_word_embeddings, + embed_key_candidates, + head_key_candidates, + prefix, + ) + + # The generic merge machinery (copy, prefix-key replacement, shard/index + # writing) lives in modeling/target/checkpoint.py. + merge_state_into_checkpoint( + base_model_path, + mtp_state, + output_path, + shard_name="mtp-merged.safetensors", + drop_prefixes=(prefix,), + ) + + # Ensure the merged config exposes the MTP structural dims. vLLM/SGLang + # use these values to build the MTP module; if the base config omits + # ``head_dim`` (common for some Qwen3.5 checkpoints), the loader will use + # its default and fail with a shape mismatch. + resolved_draft_config_path = ( + config_source + if os.path.isfile(config_source) + else os.path.join(config_source, "config.json") + ) + output_config_path = os.path.join(output_path, "config.json") + if os.path.exists(resolved_draft_config_path) and os.path.exists( + output_config_path + ): + with open(resolved_draft_config_path, "r") as f: + draft_config = json.load(f) + with open(output_config_path, "r") as f: + base_config = json.load(f) + patched_config = _patch_text_config(base_config, draft_config) + with open(output_config_path, "w") as f: + json.dump(patched_config, f, indent=2) + + # Copy over the MTP modeling file if present; some loaders need it for + # trust_remote_code / auto_map resolution. + if model_source_dir is not None: + mtp_py_src = os.path.join(model_source_dir, "mtp.py") + if os.path.exists(mtp_py_src): + shutil.copy2(mtp_py_src, os.path.join(output_path, "mtp.py")) + + print(f"Merged checkpoint saved to {output_path}") + print(f" key format: {key_format}") + print(f" MTP tensors merged: {len(mtp_state)}") + + +__all__ = ["convert_mtp_keys", "merge_mtp_into_base"] diff --git a/specforge/export/to_hf.py b/specforge/export/to_hf.py index 4bea4067f..f97e98c60 100644 --- a/specforge/export/to_hf.py +++ b/specforge/export/to_hf.py @@ -26,7 +26,11 @@ from huggingface_hub import snapshot_download from safetensors import safe_open -from specforge.export.checkpoint_io import materialize_draft, resolve_training_state +from specforge.export.checkpoint_io import ( + apply_legacy_rope_scaling, + materialize_draft, + resolve_training_state, +) def _load_embedding_tensor(source: str, key: str) -> torch.Tensor: @@ -110,6 +114,7 @@ def export_to_hf( ) full_state.update(state["draft_state_dict"]) # trained keys win model.save_pretrained(output_dir, state_dict=full_state) + apply_legacy_rope_scaling(output_dir) return output_dir diff --git a/specforge/export/to_sglang.py b/specforge/export/to_sglang.py index 1d281f1c8..b08b856ec 100644 --- a/specforge/export/to_sglang.py +++ b/specforge/export/to_sglang.py @@ -23,7 +23,11 @@ import argparse from typing import Dict, Optional -from specforge.export.checkpoint_io import materialize_draft, resolve_training_state +from specforge.export.checkpoint_io import ( + apply_legacy_rope_scaling, + materialize_draft, + resolve_training_state, +) #: per-architecture trainer-key -> serving-key renames ({} = identity). WEIGHT_MAPS: Dict[str, Dict[str, str]] = { @@ -80,6 +84,7 @@ def export_to_sglang( # embeddings exactly as the trainer-side checkpoint filter does. full = {k: v for k, v in model.state_dict().items() if "embed" not in k.lower()} model.save_pretrained(output_dir, state_dict=_serving_state(full, weight_map)) + apply_legacy_rope_scaling(output_dir) return output_dir diff --git a/specforge/inference/adapters/server_capture.py b/specforge/inference/adapters/server_capture.py index f9b730272..6647c627f 100644 --- a/specforge/inference/adapters/server_capture.py +++ b/specforge/inference/adapters/server_capture.py @@ -48,16 +48,12 @@ class ServerCaptureSchema: ``(feature_name, payload_key, trailing_shape)`` for client tensors stored verbatim (``trailing_shape`` is appended after ``(1, L)``). ``attention_mask_feature`` is synthesized all-ones (PromptTasks are unpadded). - ``position_ids_feature`` names the server-produced position-id artifact - (mRoPE positions for multimodal targets, stored ``(1, L, 3)`` int64); - None = not requested. """ aux_feature: Optional[str] last_hidden_feature: Optional[str] passthrough: Tuple[Tuple[str, str, Tuple[int, ...]], ...] attention_mask_feature: Optional[str] = None - position_ids_feature: Optional[str] = None @dataclass(frozen=True) @@ -212,8 +208,6 @@ def _spec_capture_payload(self, task: PromptTask) -> Dict[str, Any]: features["aux"] = self.schema.aux_feature if self.schema.last_hidden_feature is not None: features["last_hidden"] = self.schema.last_hidden_feature - if self.schema.position_ids_feature is not None: - features["position_ids"] = self.schema.position_ids_feature passthrough: List[Dict[str, Any]] = [] for feature_name, payload_key, trailing in self.schema.passthrough: if payload_key == "input_ids": @@ -329,7 +323,7 @@ def produce_refs( # targets. Retries need a new key because the prior attempt may have # populated its namespace before the response was lost. body["extra_key"] = [uuid.uuid4().hex for _ in tasks] - body["sampling_params"] = {"temperature": 0.0, "max_new_tokens": 1} + body["sampling_params"] = {"temperature": 0.0, "max_new_tokens": 0} capture_payloads = [self._spec_capture_payload(t) for t in tasks] body["spec_capture"] = capture_payloads for payload in capture_payloads: diff --git a/specforge/inference/sglang_patch_inventory.md b/specforge/inference/sglang_patch_inventory.md index 5ecd2e08f..3dc7266b5 100644 --- a/specforge/inference/sglang_patch_inventory.md +++ b/specforge/inference/sglang_patch_inventory.md @@ -1,20 +1,41 @@ # SGLang patch inventory and supported version -SpecForge pins `sglang==0.5.14`. The online patch is also kept compatible with -SGLang's public `inkling-support` layout. There are two deliberately separate -SGLang integration surfaces. +SpecForge pins `sglang==0.5.14` by default. The online patch is also kept +compatible with SGLang's public `inkling-support` layout, and a separately +versioned patch supports the Kimi K3 SGLang fork at the current validated +`kimi-k3` branch tip `9acd9cb` (and its original `f8493a4` integration point). +There are two deliberately separate SGLang integration surfaces. ## Online: external spec-capture server -Online training uses -[`patches/sglang/v0.5.14/spec-capture.patch`](../../patches/sglang/v0.5.14/spec-capture.patch). +Online training uses one of these source-specific patches: + +| Target | Patch | Capture methods | +|---|---|---| +| SGLang v0.5.14 / `inkling-support` | [`patches/sglang/v0.5.14/spec-capture.patch`](../../patches/sglang/v0.5.14/spec-capture.patch) | EAGLE3, DFlash, DSpark | +| Kimi K3 SGLang `9acd9cb` (`f8493a4` compatible) | [`patches/sglang/kimi-k3-f8493a4/spec-capture.patch`](../../patches/sglang/kimi-k3-f8493a4/spec-capture.patch) | EAGLE3, DFlash, DSpark | + The patch adds `--enable-spec-capture` and a server-side sink that: 1. captures requested auxiliary and final hidden states during prefill; -2. writes tensors directly into Mooncake using - `MooncakeFeatureStore`'s key layout; and +2. writes tensors into Mooncake using `MooncakeFeatureStore`'s key layout + from a background writer thread (one `batch_put_from` RPC per scheduler + batch), off the scheduler's critical path; and 3. returns only key, shape, and dtype metadata in - `meta_info["spec_capture"]`. + `meta_info["spec_capture"]`, and only after every feature object of the + request has been durably published — a response therefore guarantees the + refs it names are readable. + +Capture transfers overlap the next target prefill instead of blocking it: +the aux/last-hidden D2H rides the overlap scheduler's `copy_to_cpu` copy +stream, per-request tensors stay zero-copy views into the batch-level host +buffer, and the scheduler retains at most +`SGLANG_SPEC_CAPTURE_MAX_PENDING_BATCHES` (default 2) in-flight batches +before it blocks on the oldest — bounding pinned-host memory while keeping +backpressure. The idle loop never enters the sleeper while a transfer is +outstanding, so a lone in-flight response cannot deadlock a waiting +producer. Set `SGLANG_SPEC_CAPTURE_TIMING=1` to log per-stage +materialize/register/put timings and queue-to-stream latency. The client boundary is [`adapters/server_capture.py`](adapters/server_capture.py). Algorithm-owned @@ -22,14 +43,31 @@ providers map generic server artifacts (`aux`, `last_hidden`, passthrough inputs) to training feature names. No trainer or producer process imports SGLang model-runner internals or loads a target model. -The same patch is dry-run validated against the v0.5.14 tag and the public -`inkling-support` branch. Capture requests carry a unique `extra_key`, so every +The same patch is dry-run validated against the v0.5.14 tag and SGLang #31847 +commit `b7252cc`. Capture requests carry a unique `extra_key`, so every training sample executes a full prefill even when radix cache support is -present. Capture launch configs therefore leave radix cache enabled, including -for hybrid targets that require the unified radix tree. +present. Managed-local launch preserves the historical disabled-cache default; +hybrid targets that require the unified radix tree set +`model.sglang_disable_radix_cache: false`. + +For targets that declare `logits_mup_width_multiplier`, the SGLang model passes +an LM-head-scaled hidden state into the logits processor. The capture patch +restores the pre-head-scale post-norm representation because SpecForge folds +the same multiplier into the frozen target head used during training. -Apply the patch with `scripts/apply_sglang_spec_capture_patch.sh`. The -server-capture unit and GPU gates must pass before updating the SGLang pin. +Apply the default patch with `scripts/apply_sglang_spec_capture_patch.sh`, or +the K3 patch with +`scripts/apply_sglang_spec_capture_patch.sh --target kimi-k3-9acd9cb`. +On the default patch, `--spec-capture-method dspark` rides the DFlash aux +plumbing (`set_dflash_layers_to_capture`), which both stock v0.5.14 targets +and `inkling-support`'s Inkling model implement; DSpark and DFlash capture +the same aux/last-hidden artifacts, so managed-local DSpark launches work +unchanged. The K3 patch instead routes `--spec-capture-method dspark` to the +model's dedicated `set_dspark_layers_to_capture` hook. It also keeps 64K capture correct by using +64-bit Triton pointer arithmetic, scale-stable residual scoring, and a generic +Marlin reduction fallback when the token dimension exceeds CUDA grid.y's +65,535 limit. The server-capture unit and GPU gates must pass before updating +either supported source revision. ## Offline: dedicated local capture @@ -40,13 +78,13 @@ version-pinned APIs required for offline EAGLE3 preprocessing: | Dependency | Upgrade risk | |---|---| | `CaptureHiddenMode.FULL` and logits-processor replacement | hidden-state output fields or pruning behavior may change | -| `set_eagle3_layers_to_capture` / `set_dflash_layers_to_capture` | strategy-specific layer-selection APIs may move | +| `set_eagle3_layers_to_capture` / `set_dflash_layers_to_capture` / `set_dspark_layers_to_capture` | strategy-specific layer-selection APIs may move | | `ScheduleBatch`, `ForwardBatch`, and `ModelRunner` construction | constructor and memory-pool setup may change | | splitting captured states by request input length | token packing conventions may change | | DP-attention/model-parallel initialization patches | distributed group signatures may change | -This package computes no logits and supports text EAGLE3 and DFlash-family -state capture needed by the preprocessing script. It does not provide +This package computes no logits and supports text EAGLE3, DFlash, Domino, and +K3 DSpark state capture needed by the preprocessing script. It does not provide HF/custom backends, VLM capture, online rollout, or a general target-engine factory. diff --git a/specforge/launch.py b/specforge/launch.py index c2d73da7c..001c56696 100644 --- a/specforge/launch.py +++ b/specforge/launch.py @@ -11,6 +11,7 @@ from __future__ import annotations import logging +import os from typing import Any, Callable, List, Mapping, Optional, Tuple from specforge.algorithms.registry import AlgorithmRegistration @@ -409,28 +410,55 @@ def _epoch_online_prompts( seed: int = 0, ): """Build one deterministic, epoch-specific online prompt plan.""" + return [ + _epoch_online_prompt(prompts[index], index, epoch, prompt_epochs) + for index in _epoch_prompt_indices(prompts, epoch, seed=seed) + ] + + +def _epoch_prompt_indices(prompts, epoch: int, *, seed: int = 0): + """Return the legacy deterministic shuffle without materializing payloads.""" import random - indexed_prompts = list(enumerate(prompts)) - random.Random(int(seed) + int(epoch)).shuffle(indexed_prompts) + indices = list(range(len(prompts))) + random.Random(int(seed) + int(epoch)).shuffle(indices) + return indices + + +def _epoch_online_prompt(prompt, index: int, epoch: int, prompt_epochs: int): + """Apply epoch identity while preserving the single-epoch prompt shape.""" if prompt_epochs == 1: - return [prompt for _idx, prompt in indexed_prompts] - - out = [] - for idx, prompt in indexed_prompts: - item = dict(prompt) - metadata = dict(prompt.get("metadata") or {}) - if "task_id" in prompt: - metadata.setdefault("base_task_id", str(prompt["task_id"])) - metadata["prompt_index"] = idx - metadata["epoch"] = epoch - metadata["prompt_epochs"] = prompt_epochs - item["metadata"] = metadata - # The online feature store is consume-once and commit dedups by - # sample_id, so every epoch pass must mint distinct task/sample ids. - item["task_id"] = f"epoch{epoch:04d}-prompt{idx:012d}" - out.append(item) - return out + return prompt + + item = dict(prompt) + metadata = dict(prompt.get("metadata") or {}) + if "task_id" in prompt: + metadata.setdefault("base_task_id", str(prompt["task_id"])) + metadata["prompt_index"] = index + metadata["epoch"] = epoch + metadata["prompt_epochs"] = prompt_epochs + item["metadata"] = metadata + # The online feature store is consume-once and commit dedups by + # sample_id, so every epoch pass must mint distinct task/sample ids. + item["task_id"] = f"epoch{epoch:04d}-prompt{index:012d}" + return item + + +def _iter_epoch_online_prompt_batches( + prompts, + epoch: int, + prompt_epochs: int, + *, + seed: int = 0, + batch_size: int = 4096, +): + """Yield a shuffled epoch while bounding expanded token-list residency.""" + indices = _epoch_prompt_indices(prompts, epoch, seed=seed) + for start in range(0, len(indices), batch_size): + yield [ + _epoch_online_prompt(prompts[index], index, epoch, prompt_epochs) + for index in indices[start : start + batch_size] + ] def _assemble_server_rollout_workers( @@ -791,6 +819,7 @@ def build_disagg_online_producer( sleep=None, prompt_epochs: int = 1, prompt_seed: int = 0, + prompt_ingest_batch_size: int = 4096, ): """Producer side of an ONLINE disaggregated run (rollout pool). @@ -814,7 +843,9 @@ def build_disagg_online_producer( ``prompt_epochs`` repeats the prompt stream on the producer side by minting epoch-tagged task/sample ids. Each pass uses the deterministic ``prompt_seed + epoch`` order, matching sampler-style epoch semantics while - keeping a reconstructed plan stable across restarts. + keeping a reconstructed plan stable across restarts. Prompt payloads are + normalized and ingested in ``prompt_ingest_batch_size`` chunks so a large + memory-mapped dataset does not expand every token list before rollout. Failure semantics: a worker whose source raises (dead/unreachable server) has already failed its leases retryable — the surviving workers re-lease @@ -870,6 +901,9 @@ def elapsed(start: float) -> str: producer_concurrency = int(producer_concurrency) if producer_concurrency < 1: raise ValueError("producer_concurrency must be >= 1") + prompt_ingest_batch_size = int(prompt_ingest_batch_size) + if prompt_ingest_batch_size < 1: + raise ValueError("prompt_ingest_batch_size must be >= 1") flow_control = ProducerFlowControl( FlowControlLimits( high_watermark_refs=in_flight_high_watermark, @@ -890,20 +924,20 @@ def elapsed(start: float) -> str: and feature_store_max_resident_bytes < resident_high_watermark_bytes ): raise ValueError( - "feature_store_max_resident_bytes must be >= " - "resident_high_watermark_bytes" + "feature_store_max_resident_bytes must be >= resident_high_watermark_bytes" ) worker_lease = flow_control.prompt_lease(lease) build_start = time.perf_counter() prompt_epochs = _normalize_prompt_epochs(prompt_epochs) - if prompt_epochs > 1: + if not hasattr(prompts, "__len__") or not hasattr(prompts, "__getitem__"): prompts = list(prompts) - base_prompt_count = len(prompts) if hasattr(prompts, "__len__") else "unknown" + base_prompt_count = len(prompts) producer_timing( "build_disagg_online_producer enter " f"algorithm={algorithm.name} modality={modality} " f"base_prompts={base_prompt_count} " f"prompt_epochs={prompt_epochs} " + f"prompt_ingest_batch_size={prompt_ingest_batch_size} " f"lease={worker_lease} workers={num_rollout_workers} " f"concurrency={producer_concurrency} " f"watermarks={in_flight_high_watermark}/" @@ -1150,10 +1184,17 @@ def run_worker(w) -> None: in_flight = channel.in_flight_remote() current_resident_bytes = resident_bytes() - paused = flow_control.should_pause( + # The consumer cannot acknowledge anything until one + # complete optimizer window is available. Byte + # backpressure below that ref quantum would deadlock + # both roles; allow the required window to fill while + # the explicit hard byte cap remains enforced by + # publish_refs(). + policy_paused = flow_control.should_pause( in_flight_refs=in_flight, resident_bytes=current_resident_bytes, ) + paused = in_flight >= consumer_quantum and policy_paused if paused: if backpressure_started is None: backpressure_started = time.monotonic() @@ -1234,8 +1275,7 @@ def run_worker(w) -> None: # retryable. Other active calls keep draining. failures += 1 logger.warning( - "rollout worker %s capture call failed " - "(%d/%d): %s", + "rollout worker %s capture call failed (%d/%d): %s", w.worker_id, failures, max_worker_failures, @@ -1258,26 +1298,19 @@ def run_worker(w) -> None: abort_unpublished(futures) raise - def ingest_epoch(epoch: int) -> None: - epoch_prompts = _epoch_online_prompts( - prompts, - epoch, - prompt_epochs, - seed=prompt_seed, - ) - epoch_count = ( - len(epoch_prompts) if hasattr(epoch_prompts, "__len__") else "unknown" - ) + def ingest_prompt_batch(epoch: int, batch_index: int, epoch_prompts) -> None: phase = time.perf_counter() producer_timing( "controller.ingest_prompts start " - f"epoch={epoch + 1}/{prompt_epochs} prompts={epoch_count}" + f"epoch={epoch + 1}/{prompt_epochs} batch={batch_index + 1} " + f"prompts={len(epoch_prompts)}" ) task_ids = controller.ingest_prompts(epoch_prompts) status = controller.status() producer_timing( "controller.ingest_prompts done " - f"epoch={epoch + 1}/{prompt_epochs} tasks={len(task_ids)} " + f"epoch={epoch + 1}/{prompt_epochs} batch={batch_index + 1} " + f"tasks={len(task_ids)} " f"pending={status['prompts_pending']} elapsed={elapsed(phase)}" ) @@ -1314,21 +1347,35 @@ def run_worker_guarded(w) -> None: for epoch in range(prompt_epochs): if should_stop is not None and should_stop(): break - ingest_epoch(epoch) - if not live_workers: - raise RuntimeError( - f"all rollout workers were already dropped before " - f"epoch {epoch + 1}/{prompt_epochs} could run — " - f"dead workers: {dead}" - ) - run_epoch_workers(live_workers) - stopped = should_stop is not None and should_stop() - live_workers = [w for w in live_workers if w.worker_id not in dead] - if dead and not stopped and not pool_drained(): - raise RuntimeError( - f"all rollout workers exited with {len(dead)} dropped as " - f"dead and prompts remaining — dead workers: {dead}" - ) + epoch_batches = _iter_epoch_online_prompt_batches( + prompts, + epoch, + prompt_epochs, + seed=prompt_seed, + batch_size=prompt_ingest_batch_size, + ) + stopped = False + for batch_index, prompt_batch in enumerate(epoch_batches): + if should_stop is not None and should_stop(): + stopped = True + break + ingest_prompt_batch(epoch, batch_index, prompt_batch) + if not live_workers: + raise RuntimeError( + f"all rollout workers were already dropped before " + f"epoch {epoch + 1}/{prompt_epochs} batch " + f"{batch_index + 1} could run — dead workers: {dead}" + ) + run_epoch_workers(live_workers) + stopped = should_stop is not None and should_stop() + live_workers = [w for w in live_workers if w.worker_id not in dead] + if dead and not stopped and not pool_drained(): + raise RuntimeError( + f"all rollout workers exited with {len(dead)} dropped " + f"as dead and prompts remaining — dead workers: {dead}" + ) + if stopped: + break if stopped: break st = controller.status() @@ -1492,6 +1539,7 @@ def build_disagg_online_consumer( inbox_dir = channel.path + ".inboxes" distributor = None + inbox_server = None store = None setup_exc = None if dp_rank == 0: @@ -1563,6 +1611,18 @@ def build_disagg_online_consumer( requeued_ids=requeued_ids, idle_timeout_s=idle_timeout_s, ) + inbox_server_url = os.environ.get("DISAGG_INBOX_SERVER_URL") + if inbox_server_url: + from specforge.runtime.data_plane.http_inbox import InboxHTTPServer + + inbox_server = InboxHTTPServer( + inbox_dir, + dp_size, + inbox_server_url, + bind_host=os.environ.get( + "DISAGG_INBOX_SERVER_BIND_HOST", "0.0.0.0" + ), + ).start() channel.publish_consumer_quantum( dp_size * batch_size * accumulation_steps, allow_existing=resume_from is not None, @@ -1576,6 +1636,8 @@ def build_disagg_online_consumer( dist.broadcast_object_list(payload, src=0) setup_error = payload[0] if setup_error is not None: + if dp_rank == 0 and inbox_server is not None: + inbox_server.stop() if dp_rank == 0 and store is not None and hasattr(store, "close"): store.close() if not distributed or world == 1: @@ -1592,7 +1654,16 @@ def build_disagg_online_consumer( # The successful rank-0 setup broadcast guarantees inbox recreation and the # optimizer-window sidecar are visible before any rank opens its reader. - inbox = InboxChannel(RefDistributor.inbox_path(inbox_dir, dp_rank)) + # An optional rank-0 HTTP relay removes the shared-mount requirement for + # non-authority ranks; rank 0 remains local so the durable authority never + # depends on its own network service. + inbox_server_url = os.environ.get("DISAGG_INBOX_SERVER_URL") + if inbox_server_url and dp_rank != 0: + from specforge.runtime.data_plane.http_inbox import RemoteInboxChannel + + inbox = RemoteInboxChannel(inbox_server_url, dp_rank) + else: + inbox = InboxChannel(RefDistributor.inbox_path(inbox_dir, dp_rank)) queue = StreamingRefQueue(inbox, idle_timeout_s=idle_timeout_s) drain_state = {"attempted": False} @@ -1641,6 +1712,8 @@ def mark_consumer_failed(exc: BaseException) -> None: def stop_distributor_and_drain() -> None: if distributor is not None: distributor.stop() + if inbox_server is not None: + inbox_server.stop() # The success hook already drained before publishing consumer_done. On # an exception, finalization still makes one bounded local attempt and # reports a cleanup failure loudly without replacing the primary fit @@ -1663,7 +1736,7 @@ def stop_distributor_and_drain() -> None: ) except Exception as signal_exc: print( - "failed to publish consumer cleanup failure: " f"{signal_exc}", + f"failed to publish consumer cleanup failure: {signal_exc}", flush=True, ) logging.getLogger(__name__).error("%s", combined) diff --git a/specforge/launch_plan.py b/specforge/launch_plan.py index 4597d768d..74460483f 100644 --- a/specforge/launch_plan.py +++ b/specforge/launch_plan.py @@ -78,12 +78,12 @@ def _redacted(value: str) -> str: return f"{name}={raw}" if name is not None else raw -def _redacted_env(values: Mapping[str, str]) -> dict[str, str]: +def _redacted_env(values: Mapping[str, Optional[str]]) -> dict[str, Optional[str]]: return { name: ( "" if any(fragment in name.lower() for fragment in _SECRET_NAMES) - else _redacted(value) + else (_redacted(value) if value is not None else None) ) for name, value in sorted(values.items()) } @@ -93,7 +93,9 @@ def _redacted_env(values: Mapping[str, str]) -> dict[str, str]: class CommandSpec: label: str argv: tuple[str, ...] - env: Mapping[str, str] = field(default_factory=dict) + #: Child environment overrides; a None value unsets the inherited + #: variable (Ascend rejects an empty ASCEND_RT_VISIBLE_DEVICES). + env: Mapping[str, Optional[str]] = field(default_factory=dict) def as_dict(self) -> dict: return { @@ -142,7 +144,9 @@ class LaunchPlan: kind: PlanKind role: Literal["all", "producer", "consumer", "both"] commands: tuple[CommandSpec, ...] = () - worker_env: Mapping[str, str] = field(default_factory=dict) + #: In-process environment overrides for kind="worker". Shares the + #: CommandSpec.env contract: a None value unsets the variable. + worker_env: Mapping[str, Optional[str]] = field(default_factory=dict) services: tuple[ServiceSpec, ...] = () managed_root: Optional[str] = None managed_ports: tuple[int, ...] = () @@ -203,7 +207,7 @@ def _resolve_role( raise ValueError("--role producer/consumer/both requires disaggregated mode") if requested == "both" and cfg.training.resume_from: raise ValueError( - "--role both cannot resume a disaggregated producer; use " "--role consumer" + "--role both cannot resume a disaggregated producer; use --role consumer" ) if distributed and requested == "both": raise ValueError( @@ -222,7 +226,7 @@ def _resolved_node_rank( node_rank = int(env["NODE_RANK"]) if node_rank is not None and not 0 <= node_rank < cfg.deployment.trainer.nnodes: raise ValueError( - f"node_rank={node_rank} must be in [0, " f"{cfg.deployment.trainer.nnodes})" + f"node_rank={node_rank} must be in [0, {cfg.deployment.trainer.nnodes})" ) return node_rank @@ -265,6 +269,8 @@ def _disaggregated_env( "DISAGG_BACKEND": deployment.backend, "DISAGG_STORE_ID": deployment.store_id or cfg.run_id, } + if deployment.inbox_server_url: + values["DISAGG_INBOX_SERVER_URL"] = deployment.inbox_server_url if cfg.mode == "online": if deployment.backend != "mooncake": raise ValueError("online disaggregated training requires Mooncake") @@ -275,9 +281,9 @@ def _disaggregated_env( { "DISAGG_REF_CHANNEL": str(control_dir / "refs.jsonl"), "DISAGG_DB": str(consumer_state_dir / "consumer.sqlite"), - # SQLite/WAL stays on rank 0's local filesystem. Inboxes are - # ordinary append-only channels and must remain visible to - # ranks on every trainer node. + # SQLite/WAL stays on rank 0's local filesystem. Inboxes are + # shared normally or served by rank 0 when the HTTP relay is + # configured. "DISAGG_INBOX_DIR": str( ( control_dir @@ -337,30 +343,6 @@ def _disaggregated_env( return values -def _device_visibility_env_var() -> str: - """Name of the visible-devices env var for the active accelerator. - - CUDA hosts use ``CUDA_VISIBLE_DEVICES``; Ascend NPU hosts use - ``ASCEND_RT_VISIBLE_DEVICES``. Kept torch-free so launch planning works - in supervisor processes without torch: ``torch_npu`` is detected by - module availability (no import) after explicit env markers. - """ - forced = os.environ.get("SPECFORGE_DEVICE") - if forced == "npu": - return "ASCEND_RT_VISIBLE_DEVICES" - if forced == "cuda": - return "CUDA_VISIBLE_DEVICES" - if os.environ.get("ASCEND_RT_VISIBLE_DEVICES") or os.environ.get( - "ASCEND_VISIBLE_DEVICES" - ): - return "ASCEND_RT_VISIBLE_DEVICES" - if os.environ.get("CUDA_VISIBLE_DEVICES"): - return "CUDA_VISIBLE_DEVICES" - if importlib.util.find_spec("torch_npu") is not None: - return "ASCEND_RT_VISIBLE_DEVICES" - return "CUDA_VISIBLE_DEVICES" - - def _managed_local_environment(cfg: Config) -> dict[str, str]: deployment = cfg.deployment.disaggregated assert deployment is not None and deployment.managed_local is not None @@ -383,6 +365,37 @@ def _managed_local_environment(cfg: Config) -> dict[str, str]: return values +def _device_visibility_env_var() -> str: + """Visible-devices env var for the active accelerator. + + Kept torch-free for supervisor processes: ``torch_npu`` is detected by + module availability, after explicit env markers. + """ + forced = os.environ.get("SPECFORGE_DEVICE") + if forced == "npu": + return "ASCEND_RT_VISIBLE_DEVICES" + if forced == "cuda": + return "CUDA_VISIBLE_DEVICES" + if os.environ.get("ASCEND_RT_VISIBLE_DEVICES") or os.environ.get( + "ASCEND_VISIBLE_DEVICES" + ): + return "ASCEND_RT_VISIBLE_DEVICES" + if os.environ.get("CUDA_VISIBLE_DEVICES"): + return "CUDA_VISIBLE_DEVICES" + if importlib.util.find_spec("torch_npu") is not None: + return "ASCEND_RT_VISIBLE_DEVICES" + return "CUDA_VISIBLE_DEVICES" + + +def _hidden_devices_env_value(device_visibility_env: str) -> Optional[str]: + """Env value that hides accelerators from a managed child process. + + The Ascend driver rejects an empty ``ASCEND_RT_VISIBLE_DEVICES``, so + there the variable must be unset (``None``) instead of emptied. + """ + return "" if device_visibility_env == "CUDA_VISIBLE_DEVICES" else None + + def _sglang_argv( model: ModelConfig, *, @@ -424,7 +437,7 @@ def _managed_local_services( control_dir = Path(deployment.control_dir) log_dir = control_dir / "logs" shared_env = _managed_local_environment(cfg) - visibility_env = _device_visibility_env_var() + device_visibility_env = _device_visibility_env_var() capture_context_length = cfg.model.sglang_context_length or ( cfg.data.max_length + SGLANG_CAPTURE_CONTEXT_HEADROOM ) @@ -439,11 +452,13 @@ def _managed_local_services( f"--rpc_port={mooncake.rpc_port}", f"--http_metadata_server_port={mooncake.metadata_port}", f"--metrics_port={mooncake.metrics_port}", + *( + (f"--default_kv_lease_ttl={mooncake.default_kv_lease_ttl_ms}",) + if mooncake.default_kv_lease_ttl_ms is not None + else () + ), ), - # Hide accelerators from the mooncake process. The Ascend driver - # rejects an empty ASCEND_RT_VISIBLE_DEVICES, so only CUDA hosts - # get the (empty) variable. - ({visibility_env: ""} if visibility_env == "CUDA_VISIBLE_DEVICES" else {}), + {device_visibility_env: _hidden_devices_env_value(device_visibility_env)}, ), readiness=ReadinessSpec( "mooncake", @@ -479,7 +494,6 @@ def _managed_local_services( str(server.tp_size), "--chunked-prefill-size", "-1", - "--disable-radix-cache", "--enable-spec-capture", "--spec-capture-method", contract.method, @@ -491,16 +505,21 @@ def _managed_local_services( str(server.port), ] ) - visibility_env = ( - "ASCEND_RT_VISIBLE_DEVICES" - if "ASCEND_RT_VISIBLE_DEVICES" in os.environ - else "CUDA_VISIBLE_DEVICES" + attention_backend = ( + server.attention_backend or cfg.model.sglang_attention_backend ) + if ( + server.attention_backend is None + and attention_backend == "flashinfer" + and device_visibility_env == "ASCEND_RT_VISIBLE_DEVICES" + ): + # flashinfer does not exist on Ascend; default to the NPU backend. + attention_backend = "ascend" mm_attention_backend = cfg.model.sglang_mm_attention_backend if ( mm_attention_backend is None and cfg.model.input_modality != "text" - and visibility_env == "ASCEND_RT_VISIBLE_DEVICES" + and device_visibility_env == "ASCEND_RT_VISIBLE_DEVICES" ): # The sdpa vision backend materializes [heads, N, N] attention # scores and OOMs on large images; ascend_attn is the fused @@ -516,16 +535,14 @@ def _managed_local_services( if server.mem_fraction_static is not None else cfg.model.sglang_mem_fraction_static ), - "sglang_attention_backend": ( - server.attention_backend or cfg.model.sglang_attention_backend - ), + "sglang_attention_backend": attention_backend, "sglang_mm_attention_backend": mm_attention_backend, }, ) ) service_env = { **shared_env, - visibility_env: ",".join(server.cuda_visible_devices), + device_visibility_env: ",".join(server.cuda_visible_devices), "FLASHINFER_DISABLE_VERSION_CHECK": "1", "MOONCAKE_GLOBAL_SEGMENT_SIZE": str(mooncake.global_segment_size_bytes), "MOONCAKE_LOCAL_BUFFER_SIZE": str(mooncake.local_buffer_size_bytes), @@ -641,8 +658,7 @@ def _validate_consumer_database( if cfg.training.resume_from: if state_owner and not os.path.exists(database): raise ValueError( - "consumer resume requires the retained metadata database: " - f"{database}" + f"consumer resume requires the retained metadata database: {database}" ) return stale = [ @@ -765,11 +781,14 @@ def build_launch_plan( if role in ("consumer", "both"): consumer_env = _disaggregated_env(cfg, role_base_env, role="consumer") if managed_local is not None: + device_visibility_env = _device_visibility_env_var() producer_env.update(managed_environment) - producer_env["CUDA_VISIBLE_DEVICES"] = "" + producer_env[device_visibility_env] = _hidden_devices_env_value( + device_visibility_env + ) producer_env[_MANAGED_CHILD_ENV] = "1" consumer_env.update(managed_environment) - consumer_env[_device_visibility_env_var()] = ",".join( + consumer_env[device_visibility_env] = ",".join( managed_local.trainer_cuda_visible_devices ) consumer_env[_MANAGED_CHILD_ENV] = "1" @@ -1040,7 +1059,11 @@ def _spawn_command( stderr=None, ) -> subprocess.Popen: child_env = os.environ.copy() - child_env.update(command.env) + for key, value in command.env.items(): + if value is None: + child_env.pop(key, None) + else: + child_env[key] = value kwargs = {"env": child_env, "start_new_session": True} if stdout is not None: kwargs["stdout"] = stdout diff --git a/specforge/lr_scheduler.py b/specforge/lr_scheduler.py index caf6b4cec..c375842bd 100644 --- a/specforge/lr_scheduler.py +++ b/specforge/lr_scheduler.py @@ -119,4 +119,32 @@ def __init__( super().__init__(optimizer, warmup_steps, base_scheduler, last_epoch=last_epoch) -__all__ = ["CosineAnnealingWarmupLR"] +class _FlatLR(_LRScheduler): + """Keep every parameter group at its configured base learning rate.""" + + def get_lr(self): + return self.base_lrs + + +class ConstantWarmupLR(_WarmupScheduler): + """Linear warmup followed by a constant learning rate.""" + + def __init__( + self, + optimizer, + total_steps: int, + warmup_steps: int = 0, + last_epoch: int = -1, + ): + if total_steps <= 0: + raise ValueError(f"total_steps must be positive, got {total_steps}") + if not 0 <= warmup_steps < total_steps: + raise ValueError( + "warmup_steps must be in [0, total_steps), got " + f"{warmup_steps} for total_steps={total_steps}" + ) + base_scheduler = _FlatLR(optimizer, last_epoch=last_epoch) + super().__init__(optimizer, warmup_steps, base_scheduler, last_epoch=last_epoch) + + +__all__ = ["ConstantWarmupLR", "CosineAnnealingWarmupLR"] diff --git a/specforge/modeling/draft/__init__.py b/specforge/modeling/draft/__init__.py index 839869dd6..ce7c43f22 100644 --- a/specforge/modeling/draft/__init__.py +++ b/specforge/modeling/draft/__init__.py @@ -8,6 +8,7 @@ from .domino import DominoDraftModel from .dspark import DSparkDraftModel from .llama3_eagle import LlamaForCausalLMEagle3 +from .mtp import Qwen3_5MTPDraftModel from .peagle import PEagleDraftModel from .registry import DRAFT_REGISTRY, available_drafts, register_draft, resolve_draft @@ -18,6 +19,7 @@ "DSparkDraftModel", "LlamaForCausalLMEagle3", "PEagleDraftModel", + "Qwen3_5MTPDraftModel", "build_target_layer_ids", "extract_context_feature", "sample", diff --git a/specforge/modeling/draft/dflash.py b/specforge/modeling/draft/dflash.py index c671828ae..39f86568e 100644 --- a/specforge/modeling/draft/dflash.py +++ b/specforge/modeling/draft/dflash.py @@ -4,6 +4,7 @@ from torch import nn from transformers import DynamicCache from transformers.cache_utils import Cache +from transformers.integrations.flex_attention import compile_friendly_flex_attention from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.models.qwen3.modeling_qwen3 import ( ALL_ATTENTION_FUNCTIONS, @@ -18,8 +19,13 @@ from typing_extensions import Tuple, Unpack from .dflash_kernels import DEFAULT_DFLASH_KERNELS, DFlashKernels +from .flex_attention_backend import flex_attention_backend from .registry import register_draft +FULL_ATTENTION = "full_attention" +SLIDING_ATTENTION = "sliding_attention" +_VALID_DFLASH_LAYER_TYPES = {FULL_ATTENTION, SLIDING_ATTENTION} + def sample(logits: torch.Tensor, temperature: float = 0.0) -> torch.Tensor: if temperature < 1e-5: @@ -31,6 +37,39 @@ def sample(logits: torch.Tensor, temperature: float = 0.0) -> torch.Tensor: return torch.multinomial(probs, num_samples=1).view(bsz, seq_len) +def resolve_dflash_attention_layout( + config: Qwen3Config, +) -> tuple[tuple[str, ...], Optional[int]]: + """Validate and return the configured per-layer DFlash attention layout.""" + + num_hidden_layers = config.num_hidden_layers + layer_types = tuple(config.layer_types) + + if len(layer_types) != num_hidden_layers: + raise ValueError( + "DFlash config.layer_types must contain exactly " + f"num_hidden_layers={num_hidden_layers} entries, got " + f"{len(layer_types)}" + ) + invalid = set(layer_types) - _VALID_DFLASH_LAYER_TYPES + if invalid: + raise ValueError( + "DFlash config.layer_types supports only full_attention and " + f"sliding_attention, got {sorted(invalid)}" + ) + + if SLIDING_ATTENTION not in layer_types: + return layer_types, None + + sliding_window = config.sliding_window + if sliding_window is None or sliding_window <= 0: + raise ValueError( + "DFlash sliding_attention layers require use_sliding_window=true " + "and a positive config.sliding_window" + ) + return layer_types, sliding_window + + def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) @@ -52,62 +91,21 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): return q_embed, k_embed -def get_rope_scaling_value(config: Qwen3Config, key: str, default=None): - rope_scaling = getattr(config, "rope_scaling", None) - if rope_scaling is None: - return default - if isinstance(rope_scaling, dict): - return rope_scaling.get(key, default) - return getattr(rope_scaling, key, default) - - -class Qwen3InterleavedMultiRotaryEmbedding(Qwen3RotaryEmbedding): - """Interleaved mRoPE for Qwen3-VL style multimodal position ids.""" - - def __init__(self, config: Qwen3Config): - super().__init__(config) - self.mrope_section = get_rope_scaling_value( - config, "mrope_section", [24, 20, 20] - ) - - def _apply_interleaved_mrope(self, freqs: torch.Tensor) -> torch.Tensor: - freqs_t = freqs[0] - for dim_idx, offset in enumerate((1, 2), start=1): - length = self.mrope_section[dim_idx] * 3 - idx_slice = slice(offset, length, 3) - freqs_t[..., idx_slice] = freqs[dim_idx, ..., idx_slice] - return freqs_t - - @torch.no_grad() - def forward( - self, x: torch.Tensor, position_ids: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - if position_ids.ndim == 2: - position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) - - inv_freq_expanded = ( - self.inv_freq[None, None, :, None] - .float() - .expand(3, position_ids.shape[1], -1, 1) - ) - position_ids_expanded = position_ids[:, :, None, :].float() +def _prepare_dflash_eager_mask( + attention_mask: Optional[torch.Tensor], + dtype: torch.dtype, +) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """Convert a boolean allow-mask to eager's additive representation.""" - device_type = ( - x.device.type - if isinstance(x.device.type, str) and x.device.type != "mps" - else "cpu" - ) - with torch.autocast(device_type=device_type, enabled=False): - freqs = ( - inv_freq_expanded.float() @ position_ids_expanded.float() - ).transpose(2, 3) - interleaved_freqs = self._apply_interleaved_mrope(freqs) - emb = torch.cat((interleaved_freqs, interleaved_freqs), dim=-1) - scaling = getattr(self, "attention_scaling", 1.0) - cos = emb.cos() * scaling - sin = emb.sin() * scaling + if attention_mask is None or attention_mask.dtype != torch.bool: + return attention_mask, None - return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + valid_queries = attention_mask.any(dim=-1, keepdim=True) + # A finite minimum keeps eager softmax stable. Fully masked query rows are + # explicitly zeroed after attention so they cannot average forbidden values. + additive_mask = torch.zeros_like(attention_mask, dtype=dtype) + additive_mask.masked_fill_(~attention_mask, torch.finfo(dtype).min) + return additive_mask, valid_queries class Qwen3DFlashAttention(nn.Module): @@ -130,6 +128,10 @@ def __init__( ) self.scaling = self.head_dim**-0.5 self.attention_dropout = config.attention_dropout + if config._attn_implementation == "flex_attention": + assert ( + config.attention_dropout == 0.0 + ), "DFlash FlexAttention requires attention_dropout=0.0" self.is_causal = False self.q_proj = nn.Linear( config.hidden_size, @@ -155,7 +157,7 @@ def __init__( self.k_norm = kernels.make_rms_norm(self.head_dim, config.rms_norm_eps) self.sliding_window = ( config.sliding_window - if config.layer_types[layer_idx] == "sliding_attention" + if config.layer_types[layer_idx] == SLIDING_ATTENTION else None ) @@ -191,22 +193,52 @@ def forward( if past_key_values is not None: cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs) - attn_fn: Callable = eager_attention_forward - if self.config._attn_implementation != "eager": - attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] - attn_output, attn_weights = attn_fn( - self, - q, - k, - v, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - sliding_window=self.sliding_window, - **kwargs, - ) + valid_queries = None + if self.config._attn_implementation == "flex_attention": + kernel_options = dict(kwargs.pop("kernel_options", None) or {}) + backend = flex_attention_backend() + if backend is not None: + kernel_options["BACKEND"] = backend + + attn_output = compile_friendly_flex_attention( + q, + k, + v, + block_mask=attention_mask, + enable_gqa=True, + scale=self.scaling, + kernel_options=kernel_options or None, + ).transpose(1, 2) + attn_weights = None + else: + attn_fn: Callable = eager_attention_forward + if self.config._attn_implementation == "eager": + attention_mask, valid_queries = _prepare_dflash_eager_mask( + attention_mask, + q.dtype, + ) + else: + attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + attn_output, attn_weights = attn_fn( + self, + q, + k, + v, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, + **kwargs, + ) + if valid_queries is not None and attn_weights is not None: + attn_weights = attn_weights.masked_fill(~valid_queries, 0) attn_output = attn_output.reshape(bsz, q_len, -1) attn_output = self.o_proj(attn_output) + if valid_queries is not None: + attn_output = attn_output.masked_fill( + ~valid_queries.any(dim=1), + 0, + ) return attn_output, attn_weights @@ -347,6 +379,7 @@ def __init__( ) -> None: super().__init__(config) self.config = config + self.layer_types, self.sliding_window = resolve_dflash_attention_layout(config) kernels = dflash_kernels or DEFAULT_DFLASH_KERNELS self.layers = nn.ModuleList( [ @@ -360,13 +393,7 @@ def __init__( build_target_layer_ids(config.num_target_layers, config.num_hidden_layers), ) self.norm = kernels.make_rms_norm(config.hidden_size, config.rms_norm_eps) - self.use_interleaved_mrope = bool( - get_rope_scaling_value(config, "mrope_interleaved", False) - ) - if self.use_interleaved_mrope: - self.rotary_emb = Qwen3InterleavedMultiRotaryEmbedding(config) - else: - self.rotary_emb = Qwen3RotaryEmbedding(config) + self.rotary_emb = Qwen3RotaryEmbedding(config) self.fc = nn.Linear( len(self.target_layer_ids) * config.hidden_size, config.hidden_size, @@ -439,7 +466,7 @@ def _sample_draft_tokens( def forward( self, position_ids: torch.LongTensor, - attention_mask: Optional[torch.Tensor] = None, + attention_mask: Optional[object] = None, noise_embedding: Optional[torch.Tensor] = None, target_hidden: Optional[torch.Tensor] = None, past_key_values: Optional[Cache] = None, @@ -449,11 +476,16 @@ def forward( hidden_states = noise_embedding target_hidden = self.hidden_norm(self.fc(target_hidden)) position_embeddings = self.rotary_emb(hidden_states, position_ids) - for layer in self.layers: + for layer_type, layer in zip(self.layer_types, self.layers): + layer_attention_mask = ( + attention_mask[layer_type] + if isinstance(attention_mask, dict) + else attention_mask + ) hidden_states = layer( hidden_states=hidden_states, target_hidden=target_hidden, - attention_mask=attention_mask, + attention_mask=layer_attention_mask, position_ids=position_ids, past_key_value=past_key_values, use_cache=use_cache, diff --git a/specforge/modeling/draft/domino.py b/specforge/modeling/draft/domino.py index ca1bde824..8c152e969 100644 --- a/specforge/modeling/draft/domino.py +++ b/specforge/modeling/draft/domino.py @@ -113,10 +113,25 @@ def apply_logits_head( hidden_states: torch.Tensor, ) -> torch.Tensor: del prev_token_ids + correction_logits = self.compute_correction_logits( + prev_token_embeddings=prev_token_embeddings, + hidden_states=hidden_states, + ) + prefix_logits = base_logits[:, :, : self.suffix_start, :] + suffix_logits = base_logits[:, :, self.suffix_start :, :] + correction_logits + return torch.cat([prefix_logits, suffix_logits], dim=2) + + def compute_correction_logits( + self, + *, + prev_token_embeddings: Optional[torch.Tensor], + hidden_states: torch.Tensor, + ) -> torch.Tensor: + """Return suffix-only Domino logits without materializing final logits.""" if prev_token_embeddings is None: raise ValueError("DominoDraftModel requires prev_token_embeddings") - bsz, n_blocks, block_size = base_logits.shape[:3] + bsz, n_blocks, block_size = hidden_states.shape[:3] if self.shift_label: gru_inputs = prev_token_embeddings.reshape(bsz * n_blocks, block_size, -1) gru_out = self._run_gru(gru_inputs) @@ -132,11 +147,7 @@ def apply_logits_head( z_n = hidden_states[:, :, self.suffix_start :, :] concat_features = torch.cat([z_n, prefix_states], dim=-1) - logits_e = self.embed_proj(concat_features) - - prefix_logits = base_logits[:, :, : self.suffix_start, :] - suffix_logits = base_logits[:, :, self.suffix_start :, :] + logits_e - return torch.cat([prefix_logits, suffix_logits], dim=2) + return self.embed_proj(concat_features) __all__ = ["DominoDraftModel"] diff --git a/specforge/modeling/draft/flex_attention_backend.py b/specforge/modeling/draft/flex_attention_backend.py new file mode 100644 index 000000000..80c127425 --- /dev/null +++ b/specforge/modeling/draft/flex_attention_backend.py @@ -0,0 +1,28 @@ +"""FlexAttention backend selection shared by DFlash training components.""" + +from __future__ import annotations + +import os +from typing import Optional + +from specforge.torch_compat import patch_inductor_cutedsl_lowerings + +_VALID_BACKENDS = {"AUTO", "TRITON", "FLASH", "TRITON_DECODE"} +_BACKEND_ENV = "SPECFORGE_FLEX_ATTENTION_BACKEND" + + +def flex_attention_backend() -> Optional[str]: + backend = os.environ.get(_BACKEND_ENV, "").upper() + if not backend: + return None + if backend not in _VALID_BACKENDS: + raise ValueError( + f"{_BACKEND_ENV} must be one of {sorted(_VALID_BACKENDS)}, " + f"got {backend!r}" + ) + if backend == "FLASH": + patch_inductor_cutedsl_lowerings() + return backend + + +__all__ = ["flex_attention_backend"] diff --git a/specforge/modeling/draft/mtp/__init__.py b/specforge/modeling/draft/mtp/__init__.py new file mode 100644 index 000000000..77ee2881c --- /dev/null +++ b/specforge/modeling/draft/mtp/__init__.py @@ -0,0 +1,7 @@ +# coding=utf-8 +"""MTP draft architectures: one registered module per model family.""" + +from specforge.modeling.draft.mtp.base import MTPDraftModel +from specforge.modeling.draft.mtp.qwen3_5 import Qwen3_5MTPDraftModel + +__all__ = ["MTPDraftModel", "Qwen3_5MTPDraftModel"] diff --git a/specforge/modeling/draft/mtp/base.py b/specforge/modeling/draft/mtp/base.py new file mode 100644 index 000000000..e79b63606 --- /dev/null +++ b/specforge/modeling/draft/mtp/base.py @@ -0,0 +1,120 @@ +# coding=utf-8 +"""Architecture-independent MTP draft contract. + +One registered subclass per model family (e.g. ``qwen3_5.Qwen3_5MTPDraftModel``) +owns the actual trainable network and the native checkpoint key layout for +that family. This base owns the shared contract the MTP algorithm code relies +on: + +- ``embed_tokens`` plus a trainable ``mtp`` module with an ``lm_head`` +- ``forward(input_ids, hidden_states, ...)`` -> object exposing ``logits`` +- the native checkpoint key prefix (``mtp.*``) used for native-head init and + export round-trips +- sharing/freezing the target checkpoint's embedding (and optional lm_head) +""" + +from __future__ import annotations + +from typing import Optional + +import torch +from torch import nn + + +class MTPDraftModel(nn.Module): + """Contract shared by all MTP draft architectures.""" + + #: Native MTP weights live under this checkpoint key prefix. + NATIVE_KEY_PREFIX = "mtp." + + #: Target-checkpoint key candidates consulted when merging trained MTP + #: weights back into a base checkpoint (see ``specforge/export/mtp.py``). + #: Families override these when the target nests its text decoder + #: differently (VLM ``model.language_model.*`` vs plain ``model.*``). + TARGET_EMBED_KEY_CANDIDATES = [ + "model.language_model.embed_tokens.weight", + "model.embed_tokens.weight", + "embed_tokens.weight", + ] + TARGET_HEAD_KEY_CANDIDATES = [ + "model.language_model.lm_head.weight", + "model.lm_head.weight", + "lm_head.weight", + ] + + embed_tokens: nn.Embedding + mtp: nn.Module + + def forward( + self, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + ): + """Run the draft on shifted tokens plus target last hidden states. + + Returns an object exposing ``logits`` of shape [batch, seq, vocab]. + """ + raise NotImplementedError + + def share_target_embeddings( + self, + embed_weight: torch.Tensor, + lm_head_weight: Optional[torch.Tensor] = None, + ) -> None: + """Share and freeze the target checkpoint's embedding and lm_head. + + lm_head sharing follows the family config's ``mtp_config.share_lm_head`` + flag (default True) and then requires ``lm_head_weight``. + """ + self.embed_tokens.weight = embed_weight + self.embed_tokens.requires_grad_(False) + mtp_config = getattr(self.config, "mtp_config", None) or {} + if mtp_config.get("share_lm_head", True): + if lm_head_weight is None: + raise ValueError( + "share_lm_head is enabled but no target lm_head weight given" + ) + self.mtp.lm_head.weight = lm_head_weight + self.mtp.lm_head.requires_grad_(False) + + def native_state_dict(self) -> dict[str, torch.Tensor]: + """Return the draft weights in the native ``mtp.*`` serving layout.""" + return { + key: value + for key, value in self.state_dict().items() + if key.startswith(self.NATIVE_KEY_PREFIX) + } + + def required_native_state_keys(self) -> set[str]: + """Return native keys that must exist for safe target initialization. + + A shared lm_head is deliberately reconstructed from the target model, + and native Qwen3.5 checkpoints do not need to duplicate it under the + MTP prefix. Every other native draft tensor must be present; accepting a + partial prefix match would leave part of the trainable head randomized. + """ + required = set(self.native_state_dict()) + mtp_config = getattr(self.config, "mtp_config", None) or {} + if mtp_config.get("share_lm_head", True): + required.discard(f"{self.NATIVE_KEY_PREFIX}lm_head.weight") + return required + + def allowed_extra_native_state_keys(self) -> set[str]: + """Native keys a merged checkpoint may carry that the draft never owns. + + ``export/mtp.merge_mtp_into_base`` backfills shared embeddings into the + native namespace so serving can instantiate ``mtp.embed_tokens`` (and a + separate ``mtp.lm_head`` for untied targets). Re-finetuning a merged + checkpoint must tolerate those keys instead of rejecting the + checkpoint as incompatible. + """ + extra = {f"{self.NATIVE_KEY_PREFIX}embed_tokens.weight"} + mtp_config = getattr(self.config, "mtp_config", None) or {} + if mtp_config.get("share_lm_head", True): + extra.add(f"{self.NATIVE_KEY_PREFIX}lm_head.weight") + return extra + + +__all__ = ["MTPDraftModel"] diff --git a/specforge/modeling/draft/mtp/qwen3_5.py b/specforge/modeling/draft/mtp/qwen3_5.py new file mode 100644 index 000000000..d1e023cfb --- /dev/null +++ b/specforge/modeling/draft/mtp/qwen3_5.py @@ -0,0 +1,543 @@ +# coding=utf-8 +"""Multi-Token Prediction (MTP) draft model for Qwen3.5. + +Architecture follows the Qwen3.5 MTP design: + 1. Normalize input embeddings and target last hidden states separately. + 2. Concatenate and project via fc( [norm(emb); norm(hidden)] ). + 3. Run a 1-layer Qwen3 transformer. + 4. Compute logits with a (shared) lm_head. + +Weight key layout matches SGLang's Qwen3_5ForCausalLMMTP: + mtp.pre_fc_norm_embedding.weight + mtp.pre_fc_norm_hidden.weight + mtp.fc.weight + mtp.layers.0.self_attn.q_proj.weight + mtp.layers.0.mlp.gate_proj.weight + mtp.lm_head.weight +""" + +import copy +from typing import Optional, Tuple + +import torch +from torch import nn +from transformers.cache_utils import Cache +from transformers.modeling_outputs import CausalLMOutputWithPast +from transformers.models.qwen3.modeling_qwen3 import ( + ALL_ATTENTION_FUNCTIONS, + FlashAttentionKwargs, + GradientCheckpointingLayer, + Qwen3Config, + Qwen3MLP, + Qwen3PreTrainedModel, + eager_attention_forward, + rotate_half, +) +from typing_extensions import Unpack + +from specforge.modeling._mask_utils import _expand_mask, _make_causal_mask +from specforge.modeling.draft.mtp.base import MTPDraftModel +from specforge.modeling.draft.registry import register_draft + + +class Qwen3_5RMSNorm(nn.Module): + """Gemma-style RMSNorm used by Qwen3.5 in vLLM. + + The official Qwen3.5 checkpoints and vLLM use ``x * (1 + weight)`` instead of + the standard HuggingFace ``x * weight``. The parameter is initialized to + zeros so that ``1 + weight`` starts at one. This must match the inference + implementation, otherwise every RMSNorm weight trained in SpecForge would be + off by +1.0 when loaded into vLLM. + """ + + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.zeros(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + # Gemma-style: multiply by (1 + weight). + hidden_states = hidden_states * (1.0 + self.weight.float()) + return hidden_states.to(input_dtype) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Apply rotary positional embeddings with partial rotary support. + + When partial_rotary_factor < 1.0, only the first ``rotary_dim`` dimensions + of q/k are rotated; the rest pass through unchanged. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + rotary_dim = cos.shape[-1] + q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin) + k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin) + q_embed = torch.cat([q_embed, q_pass], dim=-1) + k_embed = torch.cat([k_embed, k_pass], dim=-1) + return q_embed, k_embed + + +class PartialRotaryEmbedding(nn.Module): + """Rotary embedding that computes inv_freq for only a fraction of head_dim. + + Matches the official Qwen3.5 ``partial_rotary_factor`` behaviour where + ``dim = int(head_dim * partial_rotary_factor)``. + """ + + inv_freq: torch.Tensor + + def __init__(self, config, head_dim): + super().__init__() + rope_theta = getattr(config, "rope_theta", 10000.0) + partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0) + dim = int(head_dim * partial_rotary_factor) + inv_freq = 1.0 / ( + rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + @torch.no_grad() + def forward(self, x, position_ids): + inv_freq_expanded = ( + self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + ) + position_ids_expanded = position_ids[:, None, :].float() + freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class Qwen3MTPAttention(nn.Module): + """Causal self-attention with optional output gate (Qwen3.5 style). + + When ``attn_output_gate`` is True (the default for Qwen3.5), ``q_proj`` + outputs ``num_heads * head_dim * 2`` and the extra half is used as a + sigmoid gate applied to the attention output before ``o_proj``. + """ + + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + self.num_key_value_groups = ( + config.num_attention_heads // config.num_key_value_heads + ) + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.attn_output_gate = getattr(config, "attn_output_gate", False) + + if self.attn_output_gate: + self.q_proj = nn.Linear( + config.hidden_size, + config.num_attention_heads * self.head_dim * 2, + bias=config.attention_bias, + ) + else: + self.q_proj = nn.Linear( + config.hidden_size, + config.num_attention_heads * self.head_dim, + bias=config.attention_bias, + ) + self.k_proj = nn.Linear( + config.hidden_size, + config.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.v_proj = nn.Linear( + config.hidden_size, + config.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, + config.hidden_size, + bias=config.attention_bias, + ) + self.q_norm = Qwen3_5RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = Qwen3_5RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + past_key_value: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + bsz, q_len, _ = hidden_states.size() + + if self.attn_output_gate: + query_states, gate = torch.chunk( + self.q_proj(hidden_states).view(bsz, q_len, -1, self.head_dim * 2), + 2, + dim=-1, + ) + gate = gate.reshape(bsz, q_len, -1) + else: + query_states = self.q_proj(hidden_states).view( + bsz, q_len, -1, self.head_dim + ) + + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view( + bsz, q_len, self.config.num_attention_heads, self.head_dim + ).transpose(1, 2) + key_states = key_states.view( + bsz, q_len, self.config.num_key_value_heads, self.head_dim + ).transpose(1, 2) + value_states = value_states.view( + bsz, q_len, self.config.num_key_value_heads, self.head_dim + ).transpose(1, 2) + + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb( + query_states, key_states, cos, sin + ) + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + + attn_fn = eager_attention_forward + if self.config._attn_implementation != "eager": + attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attn_fn( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(bsz, q_len, -1) + + if self.attn_output_gate: + attn_output = attn_output * torch.sigmoid(gate) + + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class Qwen3MTPDecoderLayer(GradientCheckpointingLayer): + """A single Qwen3-style decoder layer for the MTP draft model.""" + + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = Qwen3MTPAttention(config=config, layer_idx=layer_idx) + self.mlp = Qwen3MLP(config) + self.input_layernorm = Qwen3_5RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = Qwen3_5RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[ + torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] + ]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + past_key_value=past_key_value, + cache_position=cache_position, + **kwargs, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + if output_attentions: + outputs += (self_attn_weights,) + return outputs + + +class Qwen3_5MTPModel(nn.Module): + """The core MTP module wrapped under the ``mtp.`` prefix. + + SGLang's ``Qwen3_5ForCausalLMMTP`` wraps a *flat* ``Qwen3_5ForCausalLM`` + (``self.layers`` directly on the ForCausalLM, not nested under + ``self.model``), so after the ``mtp.`` -> ``model.`` remap the flat keys + ``mtp.layers.0.*`` / ``mtp.norm.weight`` become ``model.layers.0.*`` / + ``model.norm.weight``, matching ``self.model.layers`` / ``self.model.norm``. + """ + + def __init__(self, config: Qwen3Config): + super().__init__() + self.config = config + + # Fusion projection: fc( concat( norm(input_embeds), norm(target_hidden) ) ) + self.pre_fc_norm_embedding = Qwen3_5RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.pre_fc_norm_hidden = Qwen3_5RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.fc = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) + + # Single-layer Qwen3 transformer, flat under `mtp.layers.*` / `mtp.norm` + # to match the native Qwen3.5 checkpoint layout consumed by SGLang. + mtp_config = copy.deepcopy(config) + mtp_config.num_hidden_layers = 1 + self.layers = nn.ModuleList([Qwen3MTPDecoderLayer(mtp_config, layer_idx=0)]) + self.norm = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = PartialRotaryEmbedding( + mtp_config, + getattr( + mtp_config, + "head_dim", + mtp_config.hidden_size // mtp_config.num_attention_heads, + ), + ) + + # LM head (shared with target model during training) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + def forward( + self, + inputs_embeds: torch.Tensor, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + ) -> torch.Tensor: + # Fusion + normed_emb = self.pre_fc_norm_embedding(inputs_embeds) + normed_hidden = self.pre_fc_norm_hidden(hidden_states) + hidden_states = self.fc(torch.cat([normed_emb, normed_hidden], dim=-1)) + + bsz, seq_len, _ = hidden_states.size() + if position_ids is None: + device = hidden_states.device + position_ids = ( + torch.arange(seq_len, dtype=torch.long, device=device) + .unsqueeze(0) + .expand(bsz, -1) + ) + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # Causal mask + if attention_mask is not None and attention_mask.dim() == 2: + # [bsz, seq_len] -> [bsz, 1, seq_len, seq_len] + combined_mask = _make_causal_mask( + (bsz, seq_len), hidden_states.dtype, device=hidden_states.device + ) + expanded_mask = _expand_mask( + attention_mask, hidden_states.dtype, tgt_len=seq_len + ).to(hidden_states.device) + attention_mask = expanded_mask + combined_mask + elif attention_mask is None and seq_len > 1: + attention_mask = _make_causal_mask( + (bsz, seq_len), hidden_states.dtype, device=hidden_states.device + ) + + for layer in self.layers: + hidden_states = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + position_embeddings=position_embeddings, + )[0] + + hidden_states = self.norm(hidden_states) + return hidden_states + + +@register_draft +class Qwen3_5MTPDraftModel(MTPDraftModel, Qwen3PreTrainedModel): + """ + Qwen3.5 MTP draft model for SpecForge training. + + The embed_tokens table is loaded from the target model and frozen by default; + the lm_head is optionally shared with the target model. All trainable MTP + parameters live under the `mtp.*` prefix so that checkpoints can be loaded + directly by SGLang's Qwen3_5ForCausalLMMTP. + """ + + config_class = Qwen3Config + _no_split_modules = ["Qwen3MTPDecoderLayer"] + + def __init__(self, config: Qwen3Config) -> None: + super().__init__(config) + self.config = config + + # Shared embedding with the target model (loaded externally, frozen) + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id + ) + + # All trainable MTP weights under `mtp.*` + self.mtp = Qwen3_5MTPModel(config) + + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def get_output_embeddings(self): + return self.mtp.lm_head + + def set_output_embeddings(self, value): + self.mtp.lm_head = value + + def forward( + self, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + **kwargs, + ) -> CausalLMOutputWithPast: + inputs_embeds = self.embed_tokens(input_ids) + hidden_states = self.mtp( + inputs_embeds=inputs_embeds, + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + ) + logits = self.mtp.lm_head(hidden_states) + return CausalLMOutputWithPast(logits=logits) + + @torch.inference_mode() + def spec_generate( + self, + target: nn.Module, + input_ids: torch.LongTensor, + max_new_tokens: int, + stop_token_ids: Optional[list[int]] = None, + temperature: float = 0.0, + ) -> torch.LongTensor: + """Sequential MTP speculative generation (single MTP layer).""" + self.eval() + device = input_ids.device + num_input_tokens = input_ids.shape[1] + max_length = num_input_tokens + max_new_tokens + output_ids = input_ids.clone() + + from transformers.cache_utils import DynamicCache + + past_key_values_target = DynamicCache() + + # Prefill target once to get initial last hidden state + target_out = target( + input_ids, + past_key_values=past_key_values_target, + use_cache=True, + output_hidden_states=True, + ) + next_token_logits = target_out.logits[:, -1, :] + if temperature < 1e-5: + next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True) + else: + next_token = torch.multinomial( + torch.softmax(next_token_logits / temperature, dim=-1), num_samples=1 + ) + output_ids = torch.cat([output_ids, next_token], dim=1) + target_hidden = target_out.hidden_states[-1][:, -1:, :] + + # Loop invariant: the target KV cache covers output_ids[:-1], and + # target_hidden is the post-norm hidden state at position + # len(output_ids)-2 (driving the draft's view of the newest token). + while output_ids.shape[1] < max_length: + committed = output_ids.shape[1] + draft_input_ids = output_ids[:, -1:] + draft_position_ids = torch.tensor( + [[committed - 1]], dtype=torch.long, device=device + ) + draft_embeds = self.embed_tokens(draft_input_ids) + draft_hidden = self.mtp( + inputs_embeds=draft_embeds, + hidden_states=target_hidden, + position_ids=draft_position_ids, + ) + draft_logits = self.mtp.lm_head(draft_hidden) + if temperature < 1e-5: + draft_token = torch.argmax(draft_logits[:, -1, :], dim=-1, keepdim=True) + else: + draft_token = torch.multinomial( + torch.softmax(draft_logits[:, -1, :] / temperature, dim=-1), + num_samples=1, + ) + + # Verify against the target. logits[:, 0] is the target's prediction + # for the position the draft token occupies; logits[:, -1] is already + # conditioned on the (possibly rejected) draft token. + verify_input_ids = torch.cat([draft_input_ids, draft_token], dim=1) + verify_position_ids = torch.arange( + committed - 1, + committed + 1, + dtype=torch.long, + device=device, + ).unsqueeze(0) + target_out = target( + verify_input_ids, + position_ids=verify_position_ids, + past_key_values=past_key_values_target, + use_cache=True, + output_hidden_states=True, + ) + target_token = torch.argmax( + target_out.logits[:, 0:1, :], dim=-1, keepdim=True + ) + # h at the newest committed token drives the next draft round. + target_hidden = target_out.hidden_states[-1][:, 0:1, :] + # Trim the scored draft token's KV again: wrong on rejection, and + # re-appended by next round's verify on acceptance. + past_key_values_target.crop(committed) + + if torch.equal(draft_token, target_token): + output_ids = torch.cat([output_ids, draft_token], dim=1) + else: + output_ids = torch.cat([output_ids, target_token], dim=1) + + if ( + stop_token_ids is not None + and output_ids[0, -1].item() in stop_token_ids + ): + break + + return output_ids diff --git a/specforge/modeling/target/checkpoint.py b/specforge/modeling/target/checkpoint.py new file mode 100644 index 000000000..e27dfec6b --- /dev/null +++ b/specforge/modeling/target/checkpoint.py @@ -0,0 +1,209 @@ +# coding=utf-8 +"""Model-agnostic selective loading from local or Hugging Face checkpoints. + +These helpers know nothing about any model family or key naming convention; +callers provide the keys or the predicate. Both sharded checkpoints +(``*.safetensors.index.json``) and single-file checkpoints are supported. +""" + +from __future__ import annotations + +import glob +import json +import os +from typing import Callable, Dict, Iterable, List, Optional + +import torch +from safetensors import safe_open + + +def resolve_checkpoint_dir( + path_or_repo: str, + cache_dir: Optional[str] = None, + allow_patterns: Optional[List[str]] = None, +) -> str: + """Return a local checkpoint directory, downloading from the Hub if needed.""" + + if os.path.exists(path_or_repo): + return path_or_repo + from huggingface_hub import snapshot_download + + return snapshot_download( + repo_id=path_or_repo, + cache_dir=cache_dir, + allow_patterns=allow_patterns or ["*.json", "*.safetensors", "*.bin"], + ) + + +def read_weight_map(checkpoint_dir: str) -> Dict[str, str]: + """Return the ``weight_map`` of a sharded checkpoint, or {} if unsharded.""" + + index_files = glob.glob(os.path.join(checkpoint_dir, "*.index.json")) + if not index_files: + return {} + with open(index_files[0], "r") as f: + index = json.load(f) + return index.get("weight_map", {}) + + +def list_checkpoint_keys(checkpoint_dir: str) -> List[str]: + """List all tensor keys without loading tensor payloads.""" + + weight_map = read_weight_map(checkpoint_dir) + if weight_map: + return sorted(weight_map.keys()) + for pattern in ("*.safetensors", "*.bin"): + files = sorted(glob.glob(os.path.join(checkpoint_dir, pattern))) + if files: + target = files[0] + if target.endswith(".safetensors"): + with safe_open(target, framework="pt") as f: + return sorted(f.keys()) + state = torch.load(target, map_location="cpu", weights_only=True) + return sorted(state.keys()) + raise FileNotFoundError(f"No checkpoint found in {checkpoint_dir}") + + +def load_selected_tensors( + checkpoint_dir: str, + predicate: Callable[[str], bool], +) -> Dict[str, torch.Tensor]: + """Load only the tensors whose key matches ``predicate``. + + Sharded checkpoints open just the shards that hold selected keys. + """ + + weight_map = read_weight_map(checkpoint_dir) + selected: Dict[str, torch.Tensor] = {} + if weight_map: + shards = sorted({weight_map[k] for k in weight_map if predicate(k)}) + for shard in shards: + shard_path = os.path.join(checkpoint_dir, shard) + if not os.path.exists(shard_path): + continue + with safe_open(shard_path, framework="pt") as f: + for key in f.keys(): + if predicate(key): + selected[key] = f.get_tensor(key) + return selected + + for pattern in ("*.safetensors", "*.bin"): + files = sorted(glob.glob(os.path.join(checkpoint_dir, pattern))) + if files: + target = files[0] + if target.endswith(".safetensors"): + with safe_open(target, framework="pt") as f: + for key in f.keys(): + if predicate(key): + selected[key] = f.get_tensor(key) + else: + state = torch.load(target, map_location="cpu", weights_only=True) + for key, value in state.items(): + if predicate(key): + selected[key] = value + return selected + raise FileNotFoundError(f"No checkpoint found in {checkpoint_dir}") + + +def load_tensors_by_keys( + checkpoint_dir: str, keys: Iterable[str] +) -> Dict[str, torch.Tensor]: + """Load exactly ``keys`` (missing keys are simply absent from the result).""" + + wanted = set(keys) + return load_selected_tensors(checkpoint_dir, lambda key: key in wanted) + + +def merge_state_into_checkpoint( + base_checkpoint_dir: str, + state: Dict[str, torch.Tensor], + output_dir: str, + *, + shard_name: str, + drop_prefixes: Iterable[str] = (), +) -> None: + """Merge a state dict into a copy of a base checkpoint (model-agnostic). + + Copies non-weight files, drops base weight entries under ``drop_prefixes``, + and merges ``state``. Sharded bases get ``state`` written to a new + ``shard_name`` shard with the index ``weight_map`` updated in place (the + large base shards are never rewritten); single-file bases are rewritten + whole under their original file name. + """ + + import shutil + + from safetensors.torch import save_file + + os.makedirs(output_dir, exist_ok=True) + prefixes = tuple(drop_prefixes) + + # Copy non-weight files so the output directory is self-contained. + for fname in os.listdir(base_checkpoint_dir): + src = os.path.join(base_checkpoint_dir, fname) + if os.path.isfile(src): + shutil.copy2(src, os.path.join(output_dir, fname)) + + index_files = glob.glob(os.path.join(base_checkpoint_dir, "*.index.json")) + if index_files: + with open(index_files[0], "r") as f: + index = json.load(f) + weight_map = index.get("weight_map", {}) + + old_keys = [k for k in weight_map if k.startswith(prefixes)] + for key in old_keys: + del weight_map[key] + if old_keys: + print( + f"Replaced {len(old_keys)} weight entries under {prefixes} " + "from base model." + ) + + # Write the incoming tensors to a dedicated shard; base shards untouched. + save_file(state, os.path.join(output_dir, shard_name)) + for key in state.keys(): + weight_map[key] = shard_name + + index["weight_map"] = weight_map + with open(os.path.join(output_dir, os.path.basename(index_files[0])), "w") as f: + json.dump(index, f, indent=2) + return + + # Single-file base: load, drop, merge, rewrite under the original name. + base_safetensors = glob.glob(os.path.join(base_checkpoint_dir, "*.safetensors")) + base_bins = glob.glob(os.path.join(base_checkpoint_dir, "*.bin")) + if not base_safetensors and not base_bins: + raise FileNotFoundError(f"No checkpoint found in {base_checkpoint_dir}") + base_state = ( + load_selected_tensors(base_checkpoint_dir, lambda _key: True) + if base_safetensors + else torch.load(base_bins[0], map_location="cpu", weights_only=True) + ) + out_name = os.path.basename( + base_safetensors[0] if base_safetensors else base_bins[0] + ) + + old_keys = [k for k in base_state if k.startswith(prefixes)] + for key in old_keys: + del base_state[key] + if old_keys: + print( + f"Replaced {len(old_keys)} weight entries under {prefixes} " + "from base model." + ) + + merged = {**base_state, **state} + if out_name.endswith(".safetensors"): + save_file(merged, os.path.join(output_dir, out_name)) + else: + torch.save(merged, os.path.join(output_dir, out_name)) + + +__all__ = [ + "list_checkpoint_keys", + "load_selected_tensors", + "load_tensors_by_keys", + "merge_state_into_checkpoint", + "read_weight_map", + "resolve_checkpoint_dir", +] diff --git a/specforge/offline_capture/sglang_backend/capture.py b/specforge/offline_capture/sglang_backend/capture.py index 5721ee576..306bab3fe 100644 --- a/specforge/offline_capture/sglang_backend/capture.py +++ b/specforge/offline_capture/sglang_backend/capture.py @@ -97,10 +97,12 @@ def set_capture_layers( setter_name = { "eagle3": "set_eagle3_layers_to_capture", "dflash": "set_dflash_layers_to_capture", + "dspark": "set_dspark_layers_to_capture", }.get(capture_method) if setter_name is None: raise ValueError( - "offline SGLang capture method must be 'eagle3' or 'dflash', " + "offline SGLang capture method must be 'eagle3', 'dflash', or " + "'dspark', " f"got {capture_method!r}" ) setter = getattr(self.model_runner.model, setter_name, None) diff --git a/specforge/optimizer.py b/specforge/optimizer.py index 3d5e1aabe..10c3fb137 100644 --- a/specforge/optimizer.py +++ b/specforge/optimizer.py @@ -3,7 +3,7 @@ import torch import torch.distributed as dist -from specforge.lr_scheduler import CosineAnnealingWarmupLR +from specforge.lr_scheduler import ConstantWarmupLR, CosineAnnealingWarmupLR from specforge.utils import print_on_rank0 logger = logging.getLogger(__name__) @@ -11,7 +11,7 @@ class BF16Optimizer: """AdamW over fp32 master copies of the bf16 trainable params, with grad - clipping and cosine warmup scheduling.""" + clipping and configurable warmup scheduling.""" def __init__( self, @@ -21,6 +21,7 @@ def __init__( max_grad_norm=0.5, total_steps=800_000, warmup_ratio=0.015, + lr_scheduler="cosine", offload_master=False, ): # defaults copied from EAGLE traineagle3 ds_config.json @@ -44,7 +45,17 @@ def __init__( self.last_grad_norm = None self._grad_norm_process_group = None self._reduce_grad_norm_across_ranks = True - self.scheduler = CosineAnnealingWarmupLR( + scheduler_types = { + "constant": ConstantWarmupLR, + "cosine": CosineAnnealingWarmupLR, + } + if lr_scheduler not in scheduler_types: + raise ValueError( + f"unsupported lr_scheduler={lr_scheduler!r}; " + f"expected one of {sorted(scheduler_types)}" + ) + self.lr_scheduler_type = lr_scheduler + self.scheduler = scheduler_types[lr_scheduler]( self.optimizer, total_steps=total_steps, warmup_steps=int(warmup_ratio * total_steps), @@ -160,6 +171,13 @@ def load_state_dict(self, state_dict): """Restore optimizer/scheduler state and, when present, the rank-local fp32 master params; without them the masters are re-cloned from the bf16 weights and the resume is not numerically faithful.""" + saved_scheduler_type = state_dict.get("lr_scheduler_type", "cosine") + if saved_scheduler_type != self.lr_scheduler_type: + raise ValueError( + "checkpoint optimizer used lr_scheduler=" + f"{saved_scheduler_type!r} but this run has " + f"lr_scheduler={self.lr_scheduler_type!r}" + ) saved_max_grad_norm = state_dict.get("max_grad_norm") if saved_max_grad_norm is not None and float(saved_max_grad_norm) != float( self.max_grad_norm @@ -204,6 +222,7 @@ def state_dict(self): return { "optimizer_state_dict": self.optimizer.state_dict(), "scheduler_state_dict": self.scheduler.state_dict(), + "lr_scheduler_type": self.lr_scheduler_type, "max_grad_norm": self.max_grad_norm, # rank-local fp32 masters; without them a resume re-quantizes from bf16 "fp32_params": [t.detach().cpu() for t in self.fp32_params], diff --git a/specforge/runtime/contracts.py b/specforge/runtime/contracts.py index 7b6958b06..6d30421eb 100644 --- a/specforge/runtime/contracts.py +++ b/specforge/runtime/contracts.py @@ -33,7 +33,7 @@ SCHEMA_VERSION = 1 RunMode = Literal["online", "offline"] -DraftStrategyName = Literal["eagle3", "dflash", "domino", "dspark", "peagle"] +DraftStrategyName = Literal["eagle3", "dflash", "domino", "dspark", "peagle", "mtp"] # Tagged union for the EAGLE3 target feature. The *strategy* owns the # projection so the trainer core stays branch-free: # - pruned_logits: rollout applied the t2d vocab map; stored (seq, draft_vocab) diff --git a/specforge/runtime/control_plane/dp_ack.py b/specforge/runtime/control_plane/dp_ack.py index 94a6b95eb..ae829b521 100644 --- a/specforge/runtime/control_plane/dp_ack.py +++ b/specforge/runtime/control_plane/dp_ack.py @@ -25,6 +25,7 @@ from __future__ import annotations +from collections import OrderedDict from typing import Callable, List, Optional from specforge.runtime.control_plane.controller import DataFlowController @@ -66,6 +67,8 @@ def _broadcast_authority_error(error: Optional[str]) -> Optional[str]: return error if not (dist.is_available() and dist.is_initialized()): return error + if dist.get_world_size() == 1: + return error payload = [error] dist.broadcast_object_list(payload, src=0) return payload[0] @@ -115,6 +118,9 @@ class DPAckController(DataFlowController): optimizer-boundary counts onto the source counter. """ + _CLEANUP_LAG_BOUNDARIES = 1 + _CLEANUP_ESCALATION_BOUNDARIES = 4 + def __init__( self, run_id: str, @@ -136,6 +142,8 @@ def __init__( self._sync_error = sync_error self._sync_cleanup_error = sync_cleanup_error self.feature_store = feature_store + self._cleanup_boundary = 0 + self._cleanup_pending: OrderedDict[str, int] = OrderedDict() def ack_train_refs( self, @@ -167,6 +175,8 @@ def ack_train_refs( cleanup_error = None if optimizer_durable and self.feature_store is not None: + self._cleanup_boundary += 1 + boundary = self._cleanup_boundary failures = [] for sample_id in local_ids: try: @@ -175,6 +185,54 @@ def ack_train_refs( ) except BaseException as exc: failures.append(f"{sample_id}: {type(exc).__name__}: {exc}") + self._cleanup_pending.setdefault(sample_id, boundary) + + eligible_ids = [ + sample_id + for sample_id, first_boundary in self._cleanup_pending.items() + if boundary - first_boundary >= self._CLEANUP_LAG_BOUNDARIES + ] + try: + from specforge.runtime.data_plane.feature_store import ( + drain_feature_store_sample_removals, + retry_feature_store_sample_removals, + ) + + # Give short Mooncake read leases one optimizer window to + # expire, then make one selective, no-sleep batched attempt. + # A store-wide retry could delete prefetched refs that still + # need crash replay. + report = retry_feature_store_sample_removals( + self.feature_store, eligible_ids + ) + remaining_ids = set(report.get("remaining_ids", ())) + for sample_id in eligible_ids: + if sample_id not in remaining_ids: + self._cleanup_pending.pop(sample_id, None) + except BaseException as exc: + failures.append( + "optimizer-boundary selective retry: " + f"{type(exc).__name__}: {exc}" + ) + + overdue_ids = [ + sample_id + for sample_id, first_boundary in self._cleanup_pending.items() + if boundary - first_boundary >= self._CLEANUP_ESCALATION_BOUNDARIES + ] + if overdue_ids: + try: + # Sustained failure is exceptional. Use the existing + # bounded strong drain only for the old durable ids, never + # for this boundary's fresh or merely-prefetched samples. + drain_feature_store_sample_removals(self.feature_store, overdue_ids) + for sample_id in overdue_ids: + self._cleanup_pending.pop(sample_id, None) + except BaseException as exc: + failures.append( + "optimizer-boundary selective drain: " + f"{type(exc).__name__}: {exc}" + ) if failures: cleanup_error = ", ".join(failures) cleanup_error = self._sync_cleanup_error(cleanup_error) diff --git a/specforge/runtime/data_plane/DESIGN.md b/specforge/runtime/data_plane/DESIGN.md index 1caf60275..3f91e504a 100644 --- a/specforge/runtime/data_plane/DESIGN.md +++ b/specforge/runtime/data_plane/DESIGN.md @@ -161,7 +161,9 @@ requeues the unacknowledged tail. For an acked remote-rank ref, the fresh authority first adopts its durable key metadata before deleting it. A second pass over consumed refs remains unsupported. Producer/consumer failures are propagated explicitly, and both roles run a bounded pending-remove drain that -fails loudly instead of hiding a hard-pinned Mooncake leak. +fails loudly instead of hiding a Mooncake object leak. Hard pinning is used +when the installed Mooncake client exposes `with_hard_pin`; older clients fall +back to their default pin behavior and emit a warning. Offline manifests and feature objects are intentionally stable instead. They remain available for repeated epochs and checkpoint resume. diff --git a/specforge/runtime/data_plane/__init__.py b/specforge/runtime/data_plane/__init__.py index b0592f482..90eaf9b5f 100644 --- a/specforge/runtime/data_plane/__init__.py +++ b/specforge/runtime/data_plane/__init__.py @@ -11,6 +11,7 @@ "FeatureStore", "LocalFeatureStore", "drain_feature_store_removals", + "drain_feature_store_sample_removals", "load_feature_file", "spec_from_tensor", "SampleRefQueue", @@ -27,6 +28,7 @@ "FeatureStore": "feature_store", "LocalFeatureStore": "feature_store", "drain_feature_store_removals": "feature_store", + "drain_feature_store_sample_removals": "feature_store", "load_feature_file": "feature_store", "spec_from_tensor": "feature_store", "SampleRefQueue": "sample_ref_queue", diff --git a/specforge/runtime/data_plane/feature_dataloader.py b/specforge/runtime/data_plane/feature_dataloader.py index 2112b9063..9c4263cfc 100644 --- a/specforge/runtime/data_plane/feature_dataloader.py +++ b/specforge/runtime/data_plane/feature_dataloader.py @@ -106,6 +106,7 @@ def __init__( ack: bool = True, gc_interval_s: Optional[float] = 15.0, num_workers: int = 0, + pin_memory: bool = False, ) -> None: if (queue is None) == (refs is None): raise ValueError( @@ -129,6 +130,7 @@ def __init__( if num_workers < 0: raise ValueError("num_workers must be >= 0") self.num_workers = int(num_workers) + self.pin_memory = bool(pin_memory) self._seek_batches = 0 # Remote stores defer a physical free while the get() read-lease is live # (Mooncake remove -> -706): release() parks it and gc() must retry. @@ -213,6 +215,15 @@ def _make_batch(self, refs: List[SampleRef]) -> TrainBatch: ] if non_tensors: raise TypeError(f"collate_fn returned non-tensors for {non_tensors}") + if self.pin_memory: + batch_tensors = { + name: ( + value.pin_memory() + if value.device.type == "cpu" and not value.is_pinned() + else value + ) + for name, value in batch_tensors.items() + } return TrainBatch( sample_ids=[r.sample_id for r in refs], strategy=self.strategy, diff --git a/specforge/runtime/data_plane/feature_store.py b/specforge/runtime/data_plane/feature_store.py index eab5cda57..df40a1332 100644 --- a/specforge/runtime/data_plane/feature_store.py +++ b/specforge/runtime/data_plane/feature_store.py @@ -68,6 +68,11 @@ logger = logging.getLogger(__name__) +DEFAULT_PENDING_DRAIN_MAX_ATTEMPTS = 40 +DEFAULT_PENDING_DRAIN_RETRY_INTERVAL_S = 0.25 +DEFAULT_SAMPLE_DRAIN_MAX_ATTEMPTS = 40 +DEFAULT_SAMPLE_DRAIN_RETRY_INTERVAL_S = 0.25 + _DTYPE_BYTES = { # best-effort; falls back to element_size() for real tensors "float32": 4, "float16": 2, @@ -159,8 +164,8 @@ def gc(self, *, now: Optional[float] = None) -> Dict[str, int]: def drain_feature_store_removals( store: FeatureStore, *, - max_attempts: int = 8, - retry_interval_s: float = 0.25, + max_attempts: int = DEFAULT_PENDING_DRAIN_MAX_ATTEMPTS, + retry_interval_s: float = DEFAULT_PENDING_DRAIN_RETRY_INTERVAL_S, sleep: Callable[[float], None] = time.sleep, ) -> Dict[str, int]: """Bound lifecycle shutdown until deferred physical removes settle. @@ -169,10 +174,11 @@ def drain_feature_store_removals( may expose ``drain_pending_removals`` to retry fallible RPCs. Keeping this small adapter at the FeatureStore boundary lets online producer/consumer finalization enforce the same loud contract without depending on Mooncake's - concrete class. The default is bounded to eight attempts and 1.75 seconds of + concrete class. The default is bounded to 40 attempts and 9.75 seconds of inter-attempt waiting; Mooncake's implementation avoids existence probes between attempts so those waits let an existing read lease expire rather - than renewing it. + than renewing it. The window must exceed Mooncake's read-lease TTL, or + removals with live leases fail at shutdown (remove -706). """ if max_attempts < 1: raise ValueError("max_attempts must be >= 1") @@ -198,6 +204,72 @@ def drain_feature_store_removals( return {"removed": 0, "removed_bytes": 0, "release_pending": 0} +def drain_feature_store_sample_removals( + store: FeatureStore, + sample_ids: List[str], + *, + max_attempts: int = DEFAULT_SAMPLE_DRAIN_MAX_ATTEMPTS, + retry_interval_s: float = DEFAULT_SAMPLE_DRAIN_RETRY_INTERVAL_S, + sleep: Callable[[float], None] = time.sleep, +) -> Dict[str, int]: + """Physically reclaim only optimizer-durable samples from a remote store. + + A streaming loader can have removal-pending objects from prefetched batches + that have not reached an optimizer boundary yet. Draining the store-wide + pending set at an acknowledgement boundary would delete their crash-replay + source too early. Backends that need lease-authority removal therefore + expose a selective hook; synchronously-freeing stores need no extra work. + """ + if max_attempts < 1: + raise ValueError("max_attempts must be >= 1") + if retry_interval_s < 0: + raise ValueError("retry_interval_s must be >= 0") + ids = list(dict.fromkeys(sample_ids)) + if not ids: + return {"removed": 0, "removed_bytes": 0, "release_pending": 0} + drain = getattr(store, "drain_sample_removals", None) + if not callable(drain): + return {"removed": 0, "removed_bytes": 0, "release_pending": 0} + return drain( + ids, + max_attempts=max_attempts, + retry_interval_s=retry_interval_s, + sleep=sleep, + ) + + +def retry_feature_store_sample_removals( + store: FeatureStore, + sample_ids: List[str], +) -> Dict[str, Any]: + """Make one non-blocking removal attempt for selected durable samples. + + Remote stores may need a later optimizer boundary to outlive a short read + lease. Unlike the lifecycle drain, this steady-state hook never sleeps and + reports the ids that remain pending so the control plane can batch a later + retry. Synchronously-freeing stores need no extra work after ``abort``. + """ + ids = list(dict.fromkeys(sample_ids)) + if not ids: + return { + "removed": 0, + "removed_bytes": 0, + "release_pending": 0, + "remaining_ids": [], + "attempts": 0, + } + retry = getattr(store, "retry_sample_removals", None) + if not callable(retry): + return { + "removed": 0, + "removed_bytes": 0, + "release_pending": 0, + "remaining_ids": [], + "attempts": 0, + } + return retry(ids) + + def load_feature_file(path: str) -> Dict[str, torch.Tensor]: """Load one prepared SpecForge offline feature file.""" if path.endswith(".gz"): @@ -596,6 +668,8 @@ def health(self) -> Dict[str, Any]: "FeatureStore", "LocalFeatureStore", "drain_feature_store_removals", + "drain_feature_store_sample_removals", + "retry_feature_store_sample_removals", "load_feature_file", "spec_from_tensor", ] diff --git a/specforge/runtime/data_plane/http_inbox.py b/specforge/runtime/data_plane/http_inbox.py new file mode 100644 index 000000000..b24ff6eee --- /dev/null +++ b/specforge/runtime/data_plane/http_inbox.py @@ -0,0 +1,324 @@ +# Copyright 2024 The SpecForge team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Private-network HTTP transport for rank-local online reference inboxes. + +Online feature tensors remain in Mooncake. This module relays only the small, +tensor-free ``SampleRef`` JSONL records that :class:`RefDistributor` already +writes. It lets multi-node consumers run on container platforms where trainer +nodes cannot mount a shared control filesystem. + +Rank 0 owns the server and its local inbox files. Every other rank tail-reads +its private stream by byte offset and posts only its durable consumed-count +target. Requests are idempotent with respect to a client retry: a read offset +is advanced only after the response arrives, and an absolute consumed target is +applied under a per-rank server lock. +""" + +from __future__ import annotations + +import base64 +import json +import os +import threading +import time +from http import HTTPStatus +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from urllib.error import HTTPError, URLError +from urllib.parse import parse_qs, urlparse +from urllib.request import Request, urlopen + +from specforge.runtime.contracts import SampleRef +from specforge.runtime.data_plane.ref_distributor import RefDistributor +from specforge.runtime.data_plane.ref_serialization import ref_from_dict +from specforge.runtime.data_plane.streaming_ref_channel import StreamingRefChannel + +_CLOSED_SUFFIX = ".closed" +_FAILED_SUFFIX = ".failed" + + +class _InboxThreadingHTTPServer(ThreadingHTTPServer): + # socketserver.TCPServer defaults to a backlog of five. A DP consumer can + # have dozens of ranks polling in lockstep, so that default turns a healthy + # rank-0 relay into intermittent connection resets at every step boundary. + request_queue_size = 256 + daemon_threads = True + + +def _validated_origin(origin: str): + parsed = urlparse(origin) + if ( + parsed.scheme != "http" + or not parsed.hostname + or parsed.port is None + or parsed.username is not None + or parsed.password is not None + or parsed.path not in ("", "/") + or parsed.params + or parsed.query + or parsed.fragment + ): + raise ValueError( + "inbox HTTP origin must be http://host:port without credentials, " + "path, query, or fragment" + ) + return parsed + + +class InboxHTTPServer: + """Serve rank-0 inbox files and consumed counters on a private origin.""" + + def __init__( + self, + inbox_dir: str, + dp_size: int, + origin: str, + *, + bind_host: str = "0.0.0.0", + ) -> None: + parsed = _validated_origin(origin) + if dp_size < 2: + raise ValueError("inbox HTTP server requires dp_size >= 2") + self.inbox_dir = os.path.abspath(inbox_dir) + self.dp_size = dp_size + self.origin = origin.rstrip("/") + self._ack_channels = [ + StreamingRefChannel(RefDistributor.inbox_path(self.inbox_dir, rank)) + for rank in range(dp_size) + ] + # A target ack may be retried after an ambiguous connection reset. The + # per-rank lock makes the read/advance/respond sequence atomic across + # the original request and its retry. + self._ack_locks = [threading.Lock() for _ in range(dp_size)] + self._httpd = _InboxThreadingHTTPServer( + (bind_host, parsed.port), self._handler_type() + ) + self._thread: threading.Thread | None = None + + def _handler_type(self): + owner = self + + class Handler(BaseHTTPRequestHandler): + server_version = "SpecForgeInbox/1" + + def log_message(self, _format, *_args): + return + + def _rank(self, *, consumed: bool = False) -> int | None: + suffix = "/consumed" if consumed else "" + path = urlparse(self.path).path + prefix = "/v1/inboxes/" + if not path.startswith(prefix) or not path.endswith(suffix): + return None + token = path[len(prefix) :] + if suffix: + token = token[: -len(suffix)] + try: + rank = int(token) + except ValueError: + return None + return rank if 0 <= rank < owner.dp_size else None + + def _json(self, status: HTTPStatus, payload) -> None: + body = json.dumps(payload, separators=(",", ":")).encode("utf-8") + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): + rank = self._rank() + if rank is None: + self._json(HTTPStatus.NOT_FOUND, {"error": "unknown inbox"}) + return + query = parse_qs(urlparse(self.path).query) + try: + offset = int(query.get("offset", ["0"])[0]) + except ValueError: + offset = -1 + if offset < 0: + self._json(HTTPStatus.BAD_REQUEST, {"error": "invalid offset"}) + return + path = RefDistributor.inbox_path(owner.inbox_dir, rank) + data = b"" + next_offset = offset + try: + with open(path, "rb") as stream: + stream.seek(offset) + data = stream.read() + next_offset = stream.tell() + except FileNotFoundError: + pass + failure = None + try: + with open(path + _FAILED_SUFFIX, encoding="utf-8") as stream: + failure = stream.read() + except FileNotFoundError: + pass + self._json( + HTTPStatus.OK, + { + "data": base64.b64encode(data).decode("ascii"), + "next_offset": next_offset, + "closed": os.path.exists(path + _CLOSED_SUFFIX), + "failure": failure, + }, + ) + + def do_POST(self): + rank = self._rank(consumed=True) + if rank is None: + self._json(HTTPStatus.NOT_FOUND, {"error": "unknown inbox"}) + return + try: + length = int(self.headers.get("Content-Length", "0")) + if length < 1 or length > 1024: + raise ValueError("invalid body size") + payload = json.loads(self.rfile.read(length)) + target = int(payload["target"]) + if target < 1: + raise ValueError("target must be positive") + except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: + self._json(HTTPStatus.BAD_REQUEST, {"error": str(exc)}) + return + with owner._ack_locks[rank]: + channel = owner._ack_channels[rank] + current = channel.consumed_remote() + if target < current: + self._json( + HTTPStatus.CONFLICT, + { + "error": "consumed target moved backwards", + "consumed": current, + }, + ) + return + if target > current: + channel.mark_consumed(target - current) + self._json(HTTPStatus.OK, {"consumed": target}) + + return Handler + + def start(self) -> InboxHTTPServer: + if self._thread is not None: + return self + self._thread = threading.Thread( + target=self._httpd.serve_forever, + name="specforge-inbox-http", + daemon=True, + ) + self._thread.start() + return self + + def stop(self) -> None: + if self._thread is None: + return + self._httpd.shutdown() + self._httpd.server_close() + self._thread.join(timeout=5.0) + self._thread = None + + +class RemoteInboxChannel: + """Read one rank inbox and acknowledge it through ``InboxHTTPServer``.""" + + def __init__(self, origin: str, dp_rank: int, *, timeout_s: float = 10.0) -> None: + _validated_origin(origin) + if dp_rank < 0: + raise ValueError("dp_rank must be non-negative") + self.path = f"{origin.rstrip('/')}/v1/inboxes/{dp_rank}" + self.timeout_s = timeout_s + self._read_offset = 0 + self._buf = "" + self._closed = False + self._failure: str | None = None + self._pending: list[SampleRef] = [] + self._pull_lock = threading.Lock() + self._consumed_target = 0 + + def _pull(self) -> None: + with self._pull_lock: + request = Request(f"{self.path}?offset={self._read_offset}") + try: + with urlopen(request, timeout=self.timeout_s) as response: + payload = json.load(response) + except (URLError, OSError, TimeoutError): + # Let StreamingRefQueue's configured idle timeout distinguish a + # transient network outage from a dead rank-0 service. + return + data = base64.b64decode(payload["data"]) + next_offset = int(payload["next_offset"]) + if ( + next_offset < self._read_offset + or next_offset - self._read_offset != len(data) + ): + raise RuntimeError("inbox HTTP server returned an invalid byte range") + self._read_offset = next_offset + self._closed = bool(payload["closed"]) + self._failure = payload.get("failure") + self._buf += data.decode("utf-8") + lines = self._buf.split("\n") + self._buf = lines.pop() + self._pending.extend( + ref_from_dict(json.loads(line)) for line in lines if line + ) + + def poll(self) -> list[SampleRef]: + self._pull() + refs, self._pending = self._pending, [] + return refs + + def is_closed(self) -> bool: + self._pull() + return self._closed + + def failure(self) -> str | None: + self._pull() + if self._failure is None: + return None + return f"ref-distributor died:\n{self._failure}" + + def mark_consumed(self, n: int) -> None: + if n < 1: + return + target = self._consumed_target + n + body = json.dumps({"target": target}, separators=(",", ":")).encode("utf-8") + transient_error = None + for attempt in range(4): + request = Request( + f"{self.path}/consumed", + data=body, + headers={"Content-Type": "application/json"}, + method="POST", + ) + try: + with urlopen(request, timeout=self.timeout_s) as response: + payload = json.load(response) + except HTTPError: + raise + except (URLError, OSError, TimeoutError) as exc: + transient_error = exc + if attempt < 3: + time.sleep(0.05 * (attempt + 1)) + continue + if int(payload.get("consumed", -1)) != target: + raise RuntimeError("inbox HTTP server did not confirm consumed target") + self._consumed_target = target + return + raise RuntimeError( + "inbox HTTP consumed acknowledgement failed after retries" + ) from transient_error + + +__all__ = ["InboxHTTPServer", "RemoteInboxChannel"] diff --git a/specforge/runtime/data_plane/mooncake_store.py b/specforge/runtime/data_plane/mooncake_store.py index 93e5a708e..a5a7e6ddb 100644 --- a/specforge/runtime/data_plane/mooncake_store.py +++ b/specforge/runtime/data_plane/mooncake_store.py @@ -33,14 +33,18 @@ Lifetime: Mooncake's default eviction is approximate-LRU for a KV *cache*, which would silently drop a committed-but-unacked feature when the trainer lags hours (turning ``get()`` into a ``KeyError`` and violating the controller's -no-data-loss guarantee). We therefore **hard-pin** every object on ``put`` and -free it only by explicit ``remove()`` on consume/abort — SpecForge is the sole -lifetime authority, not Mooncake's LRU. Because ``remove()`` is a real (fallible) -RPC, ``release()`` parks a failed free in ``_release_pending`` and ``gc()`` -retries up to ``max_release_attempts`` during steady state. Lifecycle shutdown -calls :meth:`drain_pending_removals`, a separate bounded retry that raises if -physical removal never succeeds; failed hard-pinned objects are never silently -dropped from bookkeeping. +no-data-loss guarantee). When Mooncake exposes ``with_hard_pin``, SpecForge +therefore **hard-pins** every object on ``put`` and frees it only by explicit +``remove()`` on consume/abort — SpecForge is the sole lifetime authority, not +Mooncake's LRU. Older and Ascend Mooncake clients may not expose that field; the +store logs a warning and inherits their default pin behavior, so deployments +that require the strict no-eviction guarantee must use a hard-pin-capable +client. Because ``remove()`` is a real (fallible) RPC, ``release()`` parks a +failed free in ``_release_pending`` and ``gc()`` retries up to +``max_release_attempts`` during steady state. Lifecycle shutdown calls +:meth:`drain_pending_removals`, a separate bounded retry that raises if physical +removal never succeeds; failed removals are never silently dropped from +bookkeeping. Concurrency: ``release``/``abort``/``gc`` hold ``self._lock`` across the ``remove()`` RPC. The lock is what makes consume-once free race-free against a @@ -64,7 +68,14 @@ from specforge.runtime.contracts import SCHEMA_VERSION, FeatureHandle, SampleRef from specforge.runtime.data_plane.disaggregated import AuthPolicy -from specforge.runtime.data_plane.feature_store import FeatureStore, spec_from_tensor +from specforge.runtime.data_plane.feature_store import ( + DEFAULT_PENDING_DRAIN_MAX_ATTEMPTS, + DEFAULT_PENDING_DRAIN_RETRY_INTERVAL_S, + DEFAULT_SAMPLE_DRAIN_MAX_ATTEMPTS, + DEFAULT_SAMPLE_DRAIN_RETRY_INTERVAL_S, + FeatureStore, + spec_from_tensor, +) logger = logging.getLogger(__name__) @@ -91,21 +102,17 @@ def __init__( self.with_soft_pin = with_soft_pin -# Ascend selects visible NPUs through these env vars, mirroring -# ``CUDA_VISIBLE_DEVICES``. Their presence marks a host as Ascend even before -# ``torch_npu`` has been imported (e.g. in a capture producer that never runs -# ``init_distributed``). +# Ascend's CUDA_VISIBLE_DEVICES equivalent; its presence marks an Ascend +# host even before torch_npu is imported. _ASCEND_VISIBLE_DEVICE_ENVS = ("ASCEND_RT_VISIBLE_DEVICES", "ASCEND_VISIBLE_DEVICES") def _ascend_runtime_available() -> bool: - """Report whether a usable Ascend NPU runtime is present. + """Whether a usable Ascend NPU runtime is present. - ``torch.npu`` only exists once ``torch_npu`` is imported, which the canonical - trainer does lazily inside ``init_distributed``. A disaggregated capture - producer skips that path, so activate ``torch_npu`` here (only on a host that - actually selects Ascend devices) to detect the runtime without forcing the - import on CUDA/CPU hosts. + ``torch.npu`` exists only after ``torch_npu`` is imported; do that here, + but only on hosts already selecting Ascend devices via env var, so the + import never fires on CUDA/CPU hosts. """ if getattr(torch, "npu", None) is None: if not any(os.environ.get(name) for name in _ASCEND_VISIBLE_DEVICE_ENVS): @@ -124,19 +131,10 @@ def _ascend_runtime_available() -> bool: def _bind_transport_device() -> None: """Bind this process's local NPU before Mooncake installs its transport. - On Ascend, ``MooncakeDistributedStore.setup()`` installs the - ``AscendDirectTransport``, which calls ``aclrtGetDevice`` and fails with - ``ACL_ERROR_RT_CONTEXT_NULL`` (107002) when no device context exists for the - calling process. The store is constructed at run-assembly time; a trainer - (consumer) has already bound its NPU in ``init_distributed``, but a capture - *producer* deliberately never initializes an accelerator, so the transport - would have no device to allocate its local segment against. - - Bind only for NPU: CUDA's transport defaults to device 0 without an explicit - ``set_device`` and the producer intentionally avoids initializing CUDA. - ``_bind_local_device`` reads ``LOCAL_RANK``/``RANK`` (set by ``torchrun``), so - every rank pins the transport to its own NPU, and it is idempotent with the - trainer's earlier binding. + Ascend's transport calls ``aclrtGetDevice`` in ``setup()`` and fails with + ``ACL_ERROR_RT_CONTEXT_NULL`` when no device context exists — the case in + a capture producer, which never runs ``init_distributed``. No-op on CUDA, + where the transport defaults to device 0. """ from specforge.utils import get_device_type @@ -163,17 +161,12 @@ def _connect_store(setup_kwargs: Dict[str, Any]) -> Tuple[Any, Any]: "official wheel (`mooncake-transfer-engine` for CUDA < 13, or " "`mooncake-transfer-engine-cuda13` for CUDA >= 13)." ) from e - # Ascend's transport is installed inside setup() and needs a bound device - # context; bind the local accelerator first so the transfer engine can - # allocate its local segment (see _bind_transport_device). + # Ascend's transport needs a bound device context (see _bind_transport_device). _bind_transport_device() setup_kwargs = dict(setup_kwargs) if _ascend_runtime_available(): - # AscendDirectTransport rejects the store client's wildcard-location - # registration of the CPU staging buffer ("location:* is not - # supported" -> INVALID_PARAMS). SpecForge roles are pure zero-copy - # clients (put_from/get_into), so force the staging buffer to 0; the - # store client skips registration entirely when local_buffer_size == 0. + # Ascend rejects the wildcard-location staging-buffer registration + # ("location:* is not supported"); zero-copy clients can drop it. setup_kwargs["local_buffer_size"] = 0 store = MooncakeDistributedStore() rc = store.setup(**setup_kwargs) @@ -232,7 +225,7 @@ def _nbytes(t: torch.Tensor) -> int: class MooncakeFeatureStore(FeatureStore): """A disaggregated :class:`FeatureStore` backed by the Mooncake store. - **Zero-copy transport.** One hard-pinned Mooncake object per + **Zero-copy transport.** One Mooncake object per *tensor*, keyed ``{store_id}/{sample_id}/g{gen}/{name}``. ``put()`` writes each tensor straight from its storage with ``put_from(ptr)``; ``get()`` reads each straight into a tensor allocated from the ref's ``FeatureSpec`` with @@ -280,15 +273,23 @@ def __init__( _require_store_api(store) self._store = store put_config.replica_num = replica_num - # Older/Ascend Mooncake builds expose no with_hard_pin on - # ReplicateConfig; objects then follow the store's default pin - # behavior until remove(). Set it only when the field exists. + # Prefer true hard pinning when the installed Mooncake supports it. + # Older ROCm builds expose only `with_soft_pin`; that is a best-effort + # fallback rather than the same no-eviction guarantee. Some older + # Ascend builds expose neither field and must use the store default. if hasattr(put_config, "with_hard_pin"): put_config.with_hard_pin = hard_pin + elif hasattr(put_config, "with_soft_pin"): + put_config.with_soft_pin = hard_pin + if hard_pin: + logger.warning( + "Mooncake ReplicateConfig has no with_hard_pin field; " + "falling back to with_soft_pin" + ) elif hard_pin: logger.warning( - "Mooncake ReplicateConfig has no with_hard_pin field; " - "objects use the store's default pin behavior" + "Mooncake ReplicateConfig exposes neither with_hard_pin nor " + "with_soft_pin; objects use the store's default pin behavior" ) self._put_config = put_config self.max_resident_bytes = max_resident_bytes @@ -310,7 +311,7 @@ def __init__( # Server capture registers deterministic keys before issuing HTTP. If # the response is lost, no SampleRef exists to adopt/abort them. Keep a # shared (multi-adapter) provisional index so terminal producer cleanup - # can reclaim those hard-pinned objects; a successful adopt clears it. + # can reclaim those provisional objects; a successful adopt clears it. self._external_provisional: Dict[Tuple[str, int], List[str]] = {} self._active_leases: Dict[str, FeatureHandle] = {} # Samples whose remote remove() failed. gc() performs bounded @@ -341,7 +342,7 @@ def _store_exists(self, key: str) -> bool: return int(self._store.is_exist(key)) == 1 def _store_put_tensor(self, key: str, t: torch.Tensor) -> None: - """Zero-copy publish: DMA straight from the tensor's storage, hard-pinned. + """Zero-copy publish, requesting a hard pin when the client supports it. ``t`` must be contiguous + CPU (caller stages it). The bytes are the raw tensor buffer; shape/dtype travel on the ref's FeatureSpec, so get() @@ -394,10 +395,23 @@ def _store_get_tensor(self, key: str, out: torch.Tensor) -> None: f"mooncake get_into short read for {key}: got {rc} of {nb} bytes" ) - def _store_remove(self, key: str) -> bool: - """Best-effort physical free. Returns True on confirmed removal.""" + def _store_remove(self, key: str, *, force: bool = False) -> bool: + """Best-effort physical free. Returns True on confirmed removal. + + Recent Mooncake bindings expose ``remove(key, force=True)`` so a + lifecycle authority can reclaim an object after all application-level + leases have closed without waiting for Mooncake's (potentially + minutes-long) KV lease TTL. Older bindings only accept ``key``; keep + those usable and let their normal bounded retry behavior apply. + """ try: - rc = self._store.remove(key) + if force: + try: + rc = self._store.remove(key, force=True) + except TypeError: + rc = self._store.remove(key) + else: + rc = self._store.remove(key) except Exception: # pragma: no cover - transient RPC failure return False return rc is None or int(rc) == 0 @@ -429,7 +443,8 @@ def put( gen = self._gen_counter prior_gen = self._generation.get(sample_id) prior_names = self._sample_names.get(sample_id, []) - # One hard-pinned object per tensor, DMA'd straight from its storage. + # One object per tensor, DMA'd straight from its storage. The shared put + # config requests hard pinning when the Mooncake client supports it. # staged keeps the source tensors alive across the synchronous puts. for name, t in staged.items(): self._store_put_tensor(self._tkey(sample_id, gen, name), t) @@ -444,7 +459,7 @@ def put( if leaked: logger.warning( "MooncakeFeatureStore re-put of %s gen %s: removing prior " - "generation %s tensors %s failed; hard-pinned objects may be " + "generation %s tensors %s failed; remote objects may be " "orphaned (and the stale ref stays readable until reclaimed)", sample_id, prior_gen, @@ -645,6 +660,7 @@ def _try_physical_free( sample_id: str, *, confirm_absent_on_failure: bool = True, + force: bool = False, ) -> bool: """Remove all tensor objects. False on a retryable RPC failure. @@ -661,7 +677,7 @@ def _try_physical_free( ok = True for name in self._sample_names.get(sample_id, []): key = self._tkey(sample_id, gen, name) - if self._store_remove(key): + if self._store_remove(key, force=force): continue if confirm_absent_on_failure and not self._store_exists(key): continue # already gone (freed remotely) counts as freed @@ -716,11 +732,77 @@ def abort(self, sample_id: str, *, reason: str = "aborted") -> None: else: self._release_pending.setdefault(sample_id, 0) + def drain_sample_removals( + self, + sample_ids: List[str], + *, + max_attempts: int = DEFAULT_SAMPLE_DRAIN_MAX_ATTEMPTS, + retry_interval_s: float = DEFAULT_SAMPLE_DRAIN_RETRY_INTERVAL_S, + sleep: Callable[[float], None] = time.sleep, + ) -> Dict[str, int]: + """Force-remove only the named optimizer-durable samples. + + Other pending samples may belong to prefetched, not-yet-durable + batches and must remain available for crash replay. + """ + return self._drain_removals( + sample_ids=sample_ids, + max_attempts=max_attempts, + retry_interval_s=retry_interval_s, + sleep=sleep, + ) + + def retry_sample_removals(self, sample_ids: List[str]) -> Dict[str, Any]: + """Try selected removals once, without probing or sleeping. + + The optimizer path calls this only for samples made durable at an + earlier boundary. A failed remove remains in ``_release_pending`` for + the next batched attempt; avoiding ``is_exist`` here is important + because that probe renews Mooncake's read lease. + """ + target_ids = set(sample_ids) + removed = removed_bytes = 0 + with self._lock: + pending = [ + sample_id + for sample_id in self._release_pending + if sample_id in target_ids + ] + for sample_id in pending: + physically_removed = self._try_physical_free( + sample_id, + force=True, + confirm_absent_on_failure=False, + ) + if physically_removed: + sample_bytes = self._free_bookkeeping_locked(sample_id) + removed += 1 + removed_bytes += sample_bytes + self._stats["force_freed"] += 1 + self._stats["force_freed_bytes"] += sample_bytes + else: + self._release_pending[sample_id] = min( + self.max_release_attempts, + self._release_pending.get(sample_id, 0) + 1, + ) + remaining = [ + sample_id + for sample_id in self._release_pending + if sample_id in target_ids + ] + return { + "removed": removed, + "removed_bytes": removed_bytes, + "release_pending": len(remaining), + "remaining_ids": remaining, + "attempts": 1 if pending else 0, + } + def drain_pending_removals( self, *, - max_attempts: int = 40, - retry_interval_s: float = 0.5, + max_attempts: int = DEFAULT_PENDING_DRAIN_MAX_ATTEMPTS, + retry_interval_s: float = DEFAULT_PENDING_DRAIN_RETRY_INTERVAL_S, sleep: Callable[[float], None] = time.sleep, ) -> Dict[str, int]: """Retry deferred removes at lifecycle shutdown or fail loudly. @@ -732,17 +814,37 @@ def drain_pending_removals( ``sleep`` is injectable so protocol tests can advance a fake lease clock without wall-clock delays. """ + return self._drain_removals( + sample_ids=None, + max_attempts=max_attempts, + retry_interval_s=retry_interval_s, + sleep=sleep, + ) + + def _drain_removals( + self, + *, + sample_ids: Optional[List[str]], + max_attempts: int, + retry_interval_s: float, + sleep: Callable[[float], None], + ) -> Dict[str, int]: if max_attempts < 1: raise ValueError("max_attempts must be >= 1") if retry_interval_s < 0: raise ValueError("retry_interval_s must be >= 0") + target_ids = None if sample_ids is None else set(sample_ids) removed = removed_bytes = 0 last_errors: Dict[str, str] = {} attempts_run = 0 for attempt in range(max_attempts): attempts_run = attempt + 1 with self._lock: - pending = list(self._release_pending) + pending = [ + sample_id + for sample_id in self._release_pending + if target_ids is None or sample_id in target_ids + ] if not pending: return { "removed": removed, @@ -755,6 +857,12 @@ def drain_pending_removals( try: physically_removed = self._try_physical_free( sample_id, + # The application lease has already been released + # before a sample enters _release_pending. Use the + # lifecycle-authority path in current Mooncake so + # its default multi-minute KV lease does not turn a + # clean trainer shutdown into a false failure. + force=True, # Intermediate retries must not renew Mooncake's # read lease. The final probe only classifies an # already-absent key and has no following retry to @@ -776,7 +884,11 @@ def drain_pending_removals( self.max_release_attempts, self._release_pending.get(sample_id, 0) + 1, ) - remaining = list(self._release_pending) + remaining = [ + sample_id + for sample_id in self._release_pending + if target_ids is None or sample_id in target_ids + ] if not remaining: return { "removed": removed, @@ -788,7 +900,11 @@ def drain_pending_removals( sleep(retry_interval_s) with self._lock: - remaining = list(self._release_pending) + remaining = [ + sample_id + for sample_id in self._release_pending + if target_ids is None or sample_id in target_ids + ] preview = remaining[:16] detail = f"; last errors={last_errors}" if last_errors else "" raise RuntimeError( @@ -822,7 +938,7 @@ def gc(self, *, now: Optional[float] = None) -> Dict[str, int]: # Keep the physical key metadata and surface the pending # sample. Lifecycle drain owns the final bounded retry and # loud failure; silently dropping this bookkeeping would - # make a hard-pinned remote leak invisible. + # make a remote object leak invisible. continue attempts = self._release_pending[sid] + 1 if self._try_physical_free(sid, confirm_absent_on_failure=False): diff --git a/specforge/runtime/data_plane/streaming_ref_channel.py b/specforge/runtime/data_plane/streaming_ref_channel.py index f1008b165..927962cd7 100644 --- a/specforge/runtime/data_plane/streaming_ref_channel.py +++ b/specforge/runtime/data_plane/streaming_ref_channel.py @@ -41,6 +41,7 @@ import os import threading import time +from collections import deque from dataclasses import dataclass from typing import Iterator, List, Optional, Sequence @@ -114,7 +115,8 @@ def __init__(self, path: str) -> None: self._published = 0 # consumer-side self._read_offset = 0 - self._buf = "" + self._partial_line = "" + self._complete_lines = deque[str]() self._consumed = 0 # ``mark_consumed`` is called from both the trainer thread (batch acks) # and the prefetch worker (failure settlement), so the increment and @@ -327,15 +329,16 @@ def poll(self, max_n: Optional[int] = None) -> List[SampleRef]: except FileNotFoundError: chunk = "" if chunk: - self._buf += chunk - # parse the buffer even when no new bytes arrived -- a previous max_n call - # may have left complete lines buffered. + lines = (self._partial_line + chunk).split("\n") + self._partial_line = lines.pop() + self._complete_lines.extend(lines) + # Parse queued lines even when no new bytes arrived -- a previous max_n + # call may have left complete lines buffered. out: List[SampleRef] = [] - while "\n" in self._buf: + while self._complete_lines: if max_n is not None and len(out) >= max_n: break - line, self._buf = self._buf.split("\n", 1) - line = line.strip() + line = self._complete_lines.popleft().strip() if line: out.append(ref_from_dict(json.loads(line))) return out diff --git a/specforge/torch_compat.py b/specforge/torch_compat.py new file mode 100644 index 000000000..2bf868853 --- /dev/null +++ b/specforge/torch_compat.py @@ -0,0 +1,62 @@ +"""PyTorch compatibility shims used by optional fast paths.""" + +from __future__ import annotations + +import importlib + +import sympy +import torch +from packaging.version import InvalidVersion, Version + + +def patch_inductor_cutedsl_lowerings() -> bool: + """Backfill CuteDSL lowering needed by Torch 2.11 FLASH FlexAttention.""" + try: + torch_version = Version(torch.__version__.split("+", 1)[0]) + except InvalidVersion: + return False + if torch_version.major != 2 or torch_version.minor != 11: + return False + + try: + module = importlib.import_module( + "torch._inductor.codegen.cutedsl.cutedsl_op_overrides" + ) + except ImportError: + return False + from torch._inductor.utils import get_bounds_index_expr + from torch._inductor.virtualized import V + + overrides = module.CuteDSLOpOverrides + if getattr(overrides, "_specforge_cutedsl_patch", False): + return True + + def _minimum(a, b): + return overrides.where(overrides.lt(a, b), a, b) + + def _maximum(a, b): + return overrides.where(overrides.gt(a, b), a, b) + + def _index_expr(expr: sympy.Expr, dtype: torch.dtype): + if isinstance(expr, (int, sympy.Integer)): + return overrides.constant(int(expr), dtype) + + idx_str = V.kernel.kexpr(V.kernel.rename_indexing(expr)) + result = V.kernel.cse.generate( + V.kernel.body, + idx_str, + bounds=get_bounds_index_expr(expr), + dtype=dtype, + ) + result.is_scalar_expr = True + result.index_expr = V.graph.sizevars.simplify(expr) + return result + + overrides.minimum = staticmethod(_minimum) + overrides.maximum = staticmethod(_maximum) + overrides.index_expr = staticmethod(_index_expr) + overrides._specforge_cutedsl_patch = True + return True + + +__all__ = ["patch_inductor_cutedsl_lowerings"] diff --git a/specforge/tracker.py b/specforge/tracker.py index ef886028d..1189ed6cf 100644 --- a/specforge/tracker.py +++ b/specforge/tracker.py @@ -3,6 +3,7 @@ import abc import netrc import os +from collections.abc import Mapping from typing import Any, Dict, Optional import torch.distributed as dist @@ -40,15 +41,35 @@ # --- End Lazy Imports --- +def _is_secret_field(name: str) -> bool: + lowered = name.lower() + return lowered in { + "key", + "token", + "password", + "secret", + "wandb_key", + "swanlab_key", + "hf_key", + } or lowered.endswith(("_api_key", "_auth_token", "_token", "_password", "_secret")) + + +def _redact_config(value: Any, *, field: str | None = None) -> Any: + if field is not None and _is_secret_field(field): + return None if value is None else "" + if isinstance(value, Mapping): + return { + str(name): _redact_config(item, field=str(name)) + for name, item in value.items() + } + if isinstance(value, (list, tuple)): + return [_redact_config(item) for item in value] + return value + + def _public_config(args) -> Dict[str, Any]: - """Return tracker metadata without copying credentials into run logs.""" - config = dict(vars(args)) - for name in list(config): - lowered = name.lower() - if any(secret in lowered for secret in ("key", "token", "password")): - if config[name] is not None: - config[name] = "" - return config + """Return recursively redacted tracker metadata safe for run logs.""" + return _redact_config(vars(args)) class Tracker(abc.ABC): @@ -158,6 +179,7 @@ def validate_args(cls, parser, args): def __init__(self, args, output_dir: str): super().__init__(args, output_dir) + self._run = None if wandb is None: raise RuntimeError( "To use --report-to wandb, install the W&B client: " @@ -178,16 +200,25 @@ def __init__(self, args, output_dir: str): } if args.wandb_offline: init_kwargs["mode"] = "offline" - wandb.init(**init_kwargs) - self.is_initialized = True + # Keep the run handle owned by this tracker. The module-level + # ``wandb.run`` singleton is mutable process-global state; relying + # on it after a multiprocessing-heavy setup can silently detach + # explicit training history from the initialized run. + self._run = wandb.init(**init_kwargs) + self.is_initialized = self._run is not None def log(self, log_dict: Dict[str, Any], step: Optional[int] = None): - if self.rank == 0 and self.is_initialized: - wandb.log(log_dict, step=step) + if self.rank == 0 and self.is_initialized and self._run is not None: + # W&B defaults ``commit`` to False whenever an explicit step is + # supplied. Finalize each trainer record so live runs publish the + # point immediately instead of keeping the newest step buffered. + self._run.log(log_dict, step=step, commit=True) def close(self): - if self.rank == 0 and self.is_initialized and wandb.run: - wandb.finish() + if self.rank == 0 and self.is_initialized: + if self._run is not None: + self._run.finish() + self._run = None self.is_initialized = False diff --git a/specforge/training/assembly.py b/specforge/training/assembly.py index b16209f20..5d773ccb9 100644 --- a/specforge/training/assembly.py +++ b/specforge/training/assembly.py @@ -30,7 +30,7 @@ import os from collections import Counter from dataclasses import dataclass, field -from typing import Any, Callable, Dict, List, Mapping, Optional +from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence from specforge.algorithms.contracts import FeatureMode from specforge.algorithms.registry import AlgorithmRegistration @@ -268,6 +268,7 @@ def __call__(self, draft_module): lr=t.learning_rate, max_grad_norm=t.max_grad_norm, warmup_ratio=t.warmup_ratio, + lr_scheduler=t.lr_scheduler, total_steps=self.total_steps, offload_master=t.optimizer_cpu_offload, ) @@ -306,6 +307,11 @@ def _configured_logger(cfg: Config): options["swanlab_name"] = options["swanlab_name"] or cfg.run_id options["mlflow_experiment_name"] = options["mlflow_experiment_name"] or "specforge" options["mlflow_run_name"] = options["mlflow_run_name"] or cfg.run_id + if cfg.tracking.report_to == "wandb": + # W&B is the canonical reproduction record for the disaggregated K3 + # runs. Keep the complete resolved config next to the metric stream; + # tracker._public_config recursively redacts credentials before init. + options["specforge_config"] = cfg.model_dump(mode="json") return create_tracker_logger( SimpleNamespace(**options), cfg.output_dir, console_logger=_logger ) @@ -366,7 +372,7 @@ def _prepare_prompts( draft_config, path: Optional[str] = None, cache_key: Optional[str] = None, -) -> List[dict]: +) -> Sequence[dict]: """Prepare one prompt source with an optional path/cache namespace override. Training keeps the configured cache key. Evaluation supplies its own path @@ -403,6 +409,7 @@ def _prepare_prompts( num_proc=cfg.data.build_dataset_num_proc, min_loss_tokens=min_loss_tokens, max_prompts=cfg.data.max_prompts, + loss_mask_filter=algorithm.providers.model.loss_mask_filter, ) diff --git a/specforge/training/backend.py b/specforge/training/backend.py index 853c08730..ce9475b5c 100644 --- a/specforge/training/backend.py +++ b/specforge/training/backend.py @@ -132,6 +132,12 @@ def prepare_model(self, model: nn.Module) -> nn.Module: ... @abc.abstractmethod def backward(self, loss: torch.Tensor, *, is_boundary: bool = True) -> None: ... + def scale_gradients(self, factor: torch.Tensor) -> None: + """Scale synchronized gradients before clipping and stepping.""" + raise NotImplementedError( + f"{type(self).__name__} does not implement gradient scaling" + ) + @abc.abstractmethod def step(self) -> Optional[torch.Tensor]: ... @@ -313,6 +319,14 @@ def backward(self, loss: torch.Tensor, *, is_boundary: bool = True) -> None: with self.module.no_sync(): loss.backward() + def scale_gradients(self, factor: torch.Tensor) -> None: + if self.module is None: + raise RuntimeError("scale_gradients called before prepare_model") + with torch.no_grad(): + for parameter in self.module.parameters(): + if parameter.grad is not None: + parameter.grad.mul_(factor) + def step(self) -> Optional[torch.Tensor]: """Run the optimizer step, which clips and returns the global grad norm.""" if self.optimizer is None: diff --git a/specforge/training/controller.py b/specforge/training/controller.py index de6fcb29e..c910855a8 100644 --- a/specforge/training/controller.py +++ b/specforge/training/controller.py @@ -20,6 +20,7 @@ import logging import os import sys +import time from dataclasses import dataclass, field from typing import Any, Callable, Dict, Iterable, List, Optional @@ -48,28 +49,80 @@ class Checkpoint: metadata: Dict[str, Any] = field(default_factory=dict) -@dataclass(frozen=True) class StepResult: """Result of one TrainerCore step; ``optimizer_stepped`` is the authoritative - grad-accumulation boundary signal.""" + grad-accumulation boundary signal. + + Metric tensors stay on the training device until a consumer actually asks + for host values. This keeps ``train_step`` asynchronous on non-logging + steps while preserving the existing float-valued public properties. + """ + + def __init__( + self, + *, + optimizer_stepped: bool, + metric_values: Dict[str, Any], + has_grad_norm: bool, + ) -> None: + self.optimizer_stepped = optimizer_stepped + self._metric_values = metric_values + self._has_grad_norm = has_grad_norm + self._materialized_metrics: Optional[Dict[str, float]] = None + + def materialize_metrics(self) -> Dict[str, float]: + """Copy all device metrics to the host in one synchronization.""" + if self._materialized_metrics is None: + self._materialized_metrics = _materialize_metrics(self._metric_values) + return self._materialized_metrics - optimizer_stepped: bool - loss: float - grad_norm: Optional[float] - metrics: Dict[str, Any] = field(default_factory=dict) + @property + def metrics(self) -> Dict[str, float]: + return self.materialize_metrics() + @property + def loss(self) -> float: + return self.materialize_metrics()["loss"] -def _scalar(x: Any) -> float: - if isinstance(x, torch.Tensor): - return float(x.detach().float().mean().item()) - if isinstance(x, (list, tuple)) and x: - return float(torch.stack([t.detach().float() for t in x]).mean().item()) - return float(x) + @property + def grad_norm(self) -> Optional[float]: + if not self._has_grad_norm: + return None + return self.materialize_metrics()["grad_norm"] + + +def _materialize_metrics(values: Dict[str, Any]) -> Dict[str, float]: + """Materialize scalar metrics with at most one device-to-host transfer.""" + host_values: Dict[str, float] = {} + tensor_names = [] + tensors = [] + for name, value in values.items(): + if isinstance(value, torch.Tensor): + if value.numel() != 1: + raise ValueError( + f"metric {name!r} must be scalar before materialization" + ) + tensor_names.append(name) + tensors.append(value.detach().float().reshape(())) + else: + host_values[name] = float(value) + + if tensors: + device = tensors[0].device + packed = torch.stack([tensor.to(device) for tensor in tensors]) + materialized = packed.cpu().tolist() + host_values.update( + {name: float(value) for name, value in zip(tensor_names, materialized)} + ) + return host_values def _dp_mean_scalars( - values: Dict[str, torch.Tensor], *, process_group: Any = None -) -> Dict[str, torch.Tensor]: + values: Dict[str, Any], + *, + device: torch.device, + process_group: Any = None, +) -> Dict[str, Any]: """Average scalar metrics across DP ranks with one collective. Uses the established DFlash metric convention (DP mean): @@ -82,17 +135,34 @@ def _dp_mean_scalars( """ import torch.distributed as dist - if not values or not (dist.is_available() and dist.is_initialized()): - return values + normalized = { + name: ( + value.detach().float().reshape(()) + if isinstance(value, torch.Tensor) + else float(value) + ) + for name, value in values.items() + } + if not normalized or not (dist.is_available() and dist.is_initialized()): + return normalized world = ( dist.get_world_size() if process_group is None else dist.get_world_size(group=process_group) ) if world <= 1: - return values - names = list(values) - packed = torch.stack([values[name].detach().float().reshape(()) for name in names]) + return normalized + names = list(normalized) + packed = torch.stack( + [ + ( + value.to(device) + if isinstance(value, torch.Tensor) + else torch.tensor(value, dtype=torch.float32, device=device) + ) + for value in normalized.values() + ] + ) if process_group is None: dist.all_reduce(packed) else: @@ -107,7 +177,7 @@ def _reduce_ratio_metrics( device: torch.device, process_group: Any, reduce: bool, -) -> Dict[str, float]: +) -> Dict[str, torch.Tensor]: """Form telemetry ratios only after summing their numerators and counts.""" if not values: @@ -139,12 +209,18 @@ def _reduce_ratio_metrics( import torch.distributed as dist if dist.is_available() and dist.is_initialized(): - if process_group is None: - dist.all_reduce(packed) - else: - dist.all_reduce(packed, group=process_group) + world = ( + dist.get_world_size() + if process_group is None + else dist.get_world_size(group=process_group) + ) + if world > 1: + if process_group is None: + dist.all_reduce(packed) + else: + dist.all_reduce(packed, group=process_group) - output: Dict[str, float] = {} + output: Dict[str, torch.Tensor] = {} cursor = 0 for name, numerator, _denominator in normalized: width = numerator.numel() @@ -154,11 +230,9 @@ def _reduce_ratio_metrics( cursor += width ratios = summed_numerator / summed_denominator.clamp_min(1e-12) if width == 1: - output[name] = float(ratios.item()) + output[name] = ratios.reshape(()) else: - output.update( - {f"{name}_{index}": float(value) for index, value in enumerate(ratios)} - ) + output.update({f"{name}_{index}": ratios[index] for index in range(width)}) return output @@ -196,7 +270,8 @@ def _reduce_eagle3_metrics( device: torch.device, process_group: Any, ploss_decay: float, -) -> Optional[Dict[str, float]]: + reduce: bool, +) -> Optional[Dict[str, torch.Tensor]]: """Reduce EAGLE3's per-position training telemetry as numerators/counts. Accuracy and p-loss are ratios, so averaging rank-local ratios biases the @@ -268,32 +343,31 @@ def _reduce_eagle3_metrics( import torch.distributed as dist - if dist.is_available() and dist.is_initialized(): + if reduce and dist.is_available() and dist.is_initialized(): world = dist.get_world_size(process_group) if world > 1: dist.all_reduce(packed, op=dist.ReduceOp.SUM, group=process_group) reduced_acc = packed[0] / packed[1].clamp_min(1e-6) reduced_ploss = packed[2] / packed[3].clamp_min(1e-6) - result: Dict[str, float] = {} - for index, value in enumerate(reduced_acc.tolist()): - result[f"acc_{index}"] = float(value) - for index, value in enumerate(reduced_ploss.tolist()): - result[f"ploss_{index}"] = float(value) + result: Dict[str, torch.Tensor] = {} + for index in range(length): + result[f"acc_{index}"] = reduced_acc[index] + result[f"ploss_{index}"] = reduced_ploss[index] - result["acc"] = float(packed[0].sum().div(packed[1].sum().clamp_min(1e-6)).item()) + result["acc"] = packed[0].sum().div(packed[1].sum().clamp_min(1e-6)) weights = torch.tensor( [ploss_decay**index for index in range(length)], dtype=reduced_ploss.dtype, device=reduced_ploss.device, ) - result["loss"] = float((reduced_ploss * weights).sum().item()) + result["loss"] = (reduced_ploss * weights).sum() if acceptance_rates is not None: reduced_acceptance = packed[4] / packed[5].clamp_min(1e-6) - for index, value in enumerate(reduced_acceptance.tolist()): - result[f"acceptance_rate_{index}"] = float(value) - result["acceptance_rate"] = float(reduced_acceptance.mean().item()) + for index in range(length): + result[f"acceptance_rate_{index}"] = reduced_acceptance[index] + result["acceptance_rate"] = reduced_acceptance.mean() return result @@ -311,6 +385,7 @@ def __init__( self.backend = backend self.accumulation_steps = max(1, accumulation_steps) self._micro = 0 + self._ratio_totals: Dict[str, tuple[torch.Tensor, torch.Tensor]] = {} @property def accumulation_remainder(self) -> int: @@ -321,16 +396,79 @@ def train_step( self, batch: TrainBatch, ctx: Optional[StepContext] = None ) -> StepResult: out: StepOutput = self.strategy.forward_loss(batch, ctx) - loss = out.loss / self.accumulation_steps + loss = out.loss + ratio_metrics = dict(out.ratio_metrics) + if out.loss_terms is not None: + numerator, denominator = out.loss_terms + if numerator.numel() != 1 or denominator.numel() != 1: + raise ValueError("loss_terms must contain scalar tensors") + loss = numerator.reshape(()) + denominator = denominator.detach().reshape(()) + ratio_metrics["loss"] = ( + numerator.detach().reshape(()), + denominator, + ) + self._accumulate_ratio_metrics(ratio_metrics) + loss = loss / self.accumulation_steps self._micro += 1 # The boundary is known before backward so the backend can defer the FSDP # gradient reduction (no_sync) on non-boundary micro-steps. stepped = self._micro % self.accumulation_steps == 0 self.backend.backward(loss, is_boundary=stepped) + if stepped and out.loss_terms is not None: + self._normalize_gradients(self._ratio_totals["loss"][1]) grad_norm = self.backend.step() if stepped else None - return self._result(out, grad_norm, stepped) + result_ratio_metrics = self._ratio_totals if stepped else ratio_metrics + result = self._result( + out, + grad_norm, + stepped, + ratio_metrics=result_ratio_metrics, + ) + if stepped: + self._ratio_totals = {} + return result + + def _accumulate_ratio_metrics(self, values: Dict[str, Any]) -> None: + for name, (raw_numerator, raw_denominator) in values.items(): + numerator = torch.as_tensor(raw_numerator).detach() + denominator = torch.as_tensor(raw_denominator).detach() + previous = self._ratio_totals.get(name) + if previous is not None: + numerator = previous[0] + numerator + denominator = previous[1] + denominator + self._ratio_totals[name] = (numerator, denominator) + + def _normalize_gradients(self, local_denominator: torch.Tensor) -> None: + import torch.distributed as dist - def _result(self, out: StepOutput, grad_norm, stepped: bool) -> StepResult: + denominator = local_denominator.clone() + parallel_config = getattr(self.backend, "parallel_config", None) + process_group = getattr(parallel_config, "fsdp_process_group", None) + world_size = 1 + if dist.is_available() and dist.is_initialized(): + world_size = dist.get_world_size(group=process_group) + if world_size > 1: + dist.all_reduce( + denominator, + op=dist.ReduceOp.SUM, + group=process_group, + ) + if denominator.item() <= 0: + raise ValueError("global loss denominator must be positive") + scale = ( + denominator.new_tensor(world_size * self.accumulation_steps) / denominator + ) + self.backend.scale_gradients(scale) + + def _result( + self, + out: StepOutput, + grad_norm, + stepped: bool, + *, + ratio_metrics: Optional[Dict[str, Any]] = None, + ) -> StepResult: # EAGLE3 carries per-TTT numerators and denominators. Preserve those # positions and reduce counts before ratios; scalarizing its lists here # would both collapse the TTT structure and log one rank's local data. @@ -346,6 +484,7 @@ def _result(self, out: StepOutput, grad_norm, stepped: bool) -> StepResult: device=metric_device, process_group=process_group, ploss_decay=float(getattr(self.strategy, "ploss_decay", 1.0)), + reduce=stepped, ) # Structured EAGLE3 metrics are already globally reduced. Remaining # scalar diagnostics are DP-averaged in a single collective at optimizer @@ -353,54 +492,59 @@ def _result(self, out: StepOutput, grad_norm, stepped: bool) -> StepResult: metrics: Dict[str, Any] = dict(structured or {}) metrics.update( _reduce_ratio_metrics( - out.ratio_metrics, + out.ratio_metrics if ratio_metrics is None else ratio_metrics, device=metric_device, process_group=process_group, reduce=stepped, ) ) - scalar_metrics: Dict[str, torch.Tensor] = {} + scalar_metrics: Dict[str, Any] = {} if "loss" not in metrics: - scalar_metrics["loss"] = out.loss + scalar_metrics["loss"] = out.loss.detach() if "accuracy" in out.metrics and "acc" not in metrics: accuracy = out.metrics["accuracy"] if isinstance(accuracy, torch.Tensor): - scalar_metrics["acc"] = ( - accuracy.detach().float().mean().to(out.loss.device) - ) + scalar_metrics["acc"] = accuracy.detach().float().mean() elif isinstance(accuracy, (int, float)) and not isinstance(accuracy, bool): - scalar_metrics["acc"] = out.loss.detach().new_tensor(float(accuracy)) + scalar_metrics["acc"] = float(accuracy) # Strategies may expose additional scalar diagnostics without teaching - # the generic trainer their algorithm-specific names. Move CPU schedule - # scalars (for example Domino's lambda_base) onto the loss device before - # the DP reduction so NCCL-backed runs do not all-reduce a CPU tensor. - reserved_metric_keys = _EAGLE3_STRUCTURED_METRIC_KEYS | {"accuracy", "loss"} + # the generic trainer their algorithm-specific names. Keep host schedule + # scalars (for example Domino's lambda_base) on the host unless a real + # multi-rank DP reduction requires moving them to the loss device. + reserved_metric_keys = _EAGLE3_STRUCTURED_METRIC_KEYS | { + "accuracy", + "accuracy_denom", + "loss", + } for key, value in out.metrics.items(): if key in reserved_metric_keys: continue if isinstance(value, torch.Tensor): if value.numel() != 1: continue - scalar = value.detach().reshape(()).to(out.loss.device) + scalar = value.detach().reshape(()) elif isinstance(value, (int, float)) and not isinstance(value, bool): - scalar = out.loss.detach().new_tensor(float(value)) + scalar = float(value) else: continue scalar_metrics[key] = scalar if stepped: scalar_metrics = _dp_mean_scalars( scalar_metrics, + device=metric_device, process_group=process_group, ) - metrics.update({key: _scalar(value) for key, value in scalar_metrics.items()}) - gn = _scalar(grad_norm) if grad_norm is not None else None - if gn is not None: - metrics["grad_norm"] = gn + metrics.update(scalar_metrics) + if grad_norm is not None: + if isinstance(grad_norm, torch.Tensor): + grad_norm_metric = grad_norm.detach().float().mean() + else: + grad_norm_metric = float(grad_norm) + metrics["grad_norm"] = grad_norm_metric return StepResult( optimizer_stepped=stepped, - loss=metrics["loss"], - grad_norm=gn, - metrics=metrics, + metric_values=metrics, + has_grad_norm=grad_norm is not None, ) @@ -484,7 +628,8 @@ def __init__( # with its trained prefix already removed. Suppress the generic iterable # seek exactly once; consuming ``start_batch`` again would skip fresh data. self._data_prepositioned = bool(data_prepositioned) - self.last_metrics: Dict[str, Any] = {} + self._last_result: Optional[StepResult] = None + self._last_eval_metrics: Dict[str, Any] = {} self.last_checkpoint_step: Optional[int] = None from specforge.training.profiling import ProfilingOptions, StepProfiler @@ -494,6 +639,16 @@ def __init__( output_dir=output_dir, ) + @property + def last_metrics(self) -> Dict[str, Any]: + metrics = ( + dict(self._last_result.materialize_metrics()) + if self._last_result is not None + else {} + ) + metrics.update(self._last_eval_metrics) + return metrics + def _make_progress_bar(self): """Build a rank-0 optimizer-step bar for interactive terminals only.""" if not sys.stderr.isatty() or ( @@ -540,6 +695,12 @@ def _fit(self, data: Iterable[TrainBatch], progress: Optional[Any]) -> int: self.eval_interval > 0 and self.eval_data_factory is not None ) pending_ack: List[str] = [] + perf_window_started = time.perf_counter() + perf_window_steps = 0 + perf_window_samples = 0 + perf_data_wait_s = 0.0 + perf_train_compute_s = 0.0 + perf_durable_ack_s = 0.0 for epoch in range(self.epoch, self.num_epochs): self.epoch = epoch if hasattr(data, "set_epoch"): @@ -563,40 +724,88 @@ def _fit(self, data: Iterable[TrainBatch], progress: Optional[Any]) -> int: stream = it _it = iter(stream) while True: + data_wait_started = time.perf_counter() try: batch = next(_it) except StopIteration: break + perf_data_wait_s += time.perf_counter() - data_wait_started + perf_window_samples += len(batch.sample_ids) self._epoch_batch += 1 self._epoch_samples += len(batch.sample_ids) self.micro_step += 1 if self.ack_fn is not None: pending_ack.extend(batch.sample_ids) self._step_profiler.before_micro_step(self.global_step) + train_compute_started = time.perf_counter() result = self.core.train_step( batch, ctx=StepContext( global_step=self.global_step, total_steps=self.total_steps ), ) - self.last_metrics = result.metrics + perf_train_compute_s += time.perf_counter() - train_compute_started # grad accumulated but optimizer has not stepped yet; everything # keyed on optimizer steps fires only at the boundary. if not result.optimizer_stepped: continue self.global_step += 1 + self._last_result = result + self._last_eval_metrics = {} + perf_window_steps += 1 self._step_profiler.after_optimizer_step(self.global_step) if self.ack_fn is not None: # durable ack transaction at the optimizer-step boundary + durable_ack_started = time.perf_counter() self.ack_fn(pending_ack, self.global_step) + perf_durable_ack_s += time.perf_counter() - durable_ack_started pending_ack = [] if self.logger and self.global_step % max(1, self.log_interval) == 0: - log_metrics = dict(result.metrics) + log_metrics = dict(result.materialize_metrics()) optimizer = getattr(self.core.backend, "optimizer", None) get_learning_rate = getattr(optimizer, "get_learning_rate", None) if callable(get_learning_rate): log_metrics["lr"] = float(get_learning_rate()) + perf_elapsed_s = max( + time.perf_counter() - perf_window_started, + 1e-12, + ) + parallel = getattr(self.core.backend, "parallel_config", None) + world_size = int(getattr(parallel, "world_size", 1)) + tp_size = int(getattr(parallel, "tp_size", 1)) + sp_size = int(getattr(parallel, "sp_size", 1)) + data_parallel_size = max(1, world_size // (tp_size * sp_size)) + log_metrics.update( + { + "perf/optimizer_steps_per_hour": ( + perf_window_steps * 3600.0 / perf_elapsed_s + ), + "perf/optimizer_step_time_s": ( + perf_elapsed_s / max(1, perf_window_steps) + ), + "perf/data_wait_time_s": ( + perf_data_wait_s / max(1, perf_window_steps) + ), + "perf/train_compute_time_s": ( + perf_train_compute_s / max(1, perf_window_steps) + ), + "perf/durable_ack_time_s": ( + perf_durable_ack_s / max(1, perf_window_steps) + ), + "perf/global_samples_per_second": ( + perf_window_samples + * data_parallel_size + / perf_elapsed_s + ), + } + ) self.logger(log_metrics, self.global_step) + perf_window_started = time.perf_counter() + perf_window_steps = 0 + perf_window_samples = 0 + perf_data_wait_s = 0.0 + perf_train_compute_s = 0.0 + perf_durable_ack_s = 0.0 eval_metrics: Optional[Dict[str, Any]] = None if eval_enabled and self.global_step % self.eval_interval == 0: eval_metrics = self.evaluate_configured() @@ -604,7 +813,7 @@ def _fit(self, data: Iterable[TrainBatch], progress: Optional[Any]) -> int: if eval_metrics: if self.logger: self.logger(eval_metrics, self.global_step) - self.last_metrics = {**self.last_metrics, **eval_metrics} + self._last_eval_metrics = dict(eval_metrics) # ``is_better`` is collective (rank0 verdict broadcast inside # the manager); its guard is rank-identical because eval metrics # are DP-reduced. Empty eval metrics skip best tracking. diff --git a/specforge/training/disaggregated.py b/specforge/training/disaggregated.py index d638f0ff2..47d175e80 100644 --- a/specforge/training/disaggregated.py +++ b/specforge/training/disaggregated.py @@ -201,6 +201,62 @@ def _consumer_database_path(cfg: Config) -> Optional[str]: return os.path.join(state_dir, "consumer.sqlite") +def _online_prompt_seed(cfg: Config) -> int: + """Resolve prompt ordering independently while preserving old configs.""" + configured = getattr(cfg.training, "prompt_seed", None) + return cfg.training.seed if configured is None else configured + + +def _online_flow_window(cfg: Config) -> tuple[int, Optional[int]]: + """Resolve and validate producer ref watermarks before prompt preparation. + + The consumer dispatches one complete global optimizer window at a time. + Validating this contract only after tokenizing/materializing prompts can + waste tens of minutes and many GiB for a large online dataset. + """ + from specforge.runtime.control_plane.flow_control import FlowControlLimits + + high_override = os.environ.get("DISAGG_IN_FLIGHT_HIGH_WATERMARK") + high = int(high_override or cfg.runtime.in_flight_high_watermark) + low_override = os.environ.get("DISAGG_IN_FLIGHT_LOW_WATERMARK") + # Preserve the legacy one-watermark environment override: when only the + # old high value is supplied, resume at that same threshold. + low = ( + int(low_override) + if low_override is not None + else ( + None if high_override is not None else cfg.runtime.in_flight_low_watermark + ) + ) + limits = FlowControlLimits( + high_watermark_refs=high, + low_watermark_refs=low, + max_prompt_lease_per_worker=cfg.runtime.producer_lease, + ) + trainer = cfg.deployment.trainer + consumer_quantum = ( + trainer.nnodes + * trainer.nproc_per_node + * cfg.training.batch_size + * cfg.training.accumulation_steps + ) + if high < consumer_quantum: + raise ValueError( + "producer in-flight high watermark " + f"{high} is smaller than the consumer's global optimizer-step " + f"quantum {consumer_quantum}; set " + "DISAGG_IN_FLIGHT_HIGH_WATERMARK to at least that value" + ) + if limits.resolved_low_watermark_refs < consumer_quantum: + raise ValueError( + "producer in-flight low watermark " + f"{limits.resolved_low_watermark_refs} is smaller than the " + f"consumer's global optimizer-step quantum {consumer_quantum}; " + "set DISAGG_IN_FLIGHT_LOW_WATERMARK to at least that value" + ) + return high, low + + def _online_schedule_payload(cfg: Config, *, num_prompts: int) -> dict: """Describe the exact finite online schedule prepared by the producer.""" from specforge.training.schedule import resolve_online_total_steps @@ -219,7 +275,7 @@ def _online_schedule_payload(cfg: Config, *, num_prompts: int) -> dict: "total_steps": total_steps, "num_prompts": num_prompts, "prompt_epochs": cfg.training.num_epochs, - "prompt_seed": cfg.training.seed, + "prompt_seed": _online_prompt_seed(cfg), "dp_size": dp_size, "batch_size": cfg.training.batch_size, "accumulation_steps": cfg.training.accumulation_steps, @@ -246,7 +302,7 @@ def _read_online_total_steps(cfg: Config, channel_path: str) -> int: expected = { "version": 1, "prompt_epochs": cfg.training.num_epochs, - "prompt_seed": cfg.training.seed, + "prompt_seed": _online_prompt_seed(cfg), "dp_size": trainer.nnodes * trainer.nproc_per_node, "batch_size": cfg.training.batch_size, "accumulation_steps": cfg.training.accumulation_steps, @@ -539,6 +595,9 @@ def _build_online( from specforge.launch import build_disagg_online_producer from specforge.training.model_loading import resolve_draft_config + # This check is independent of dataset size. Keep it before tokenizer + # loading and prompt preparation so an invalid window fails cheaply. + in_flight_high_watermark, in_flight_low_watermark = _online_flow_window(cfg) input_adapter = streaming.create_input_adapter(cfg) input_tools = _load_input_tools( cfg, @@ -561,6 +620,10 @@ def _build_online( input_tools, draft_config=draft_config, ) + if not prompts: + raise ValueError( + f"no prompts satisfy {algorithm.name} training eligibility" + ) if cfg.training.total_steps is None and cfg.training.max_steps is None: schedule = _online_schedule_payload(cfg, num_prompts=len(prompts)) _write_control( @@ -576,7 +639,6 @@ def _build_online( last_hidden_feature=layout.last_hidden_feature, passthrough=layout.passthrough, attention_mask_feature=layout.attention_mask_feature, - position_ids_feature=layout.position_ids_feature, ) adapters = [ SGLangServerCaptureAdapter( @@ -592,22 +654,6 @@ def _build_online( ] target_repr = streaming.target_representation peer_wait_timeout_s = _optional_timeout_s("DISAGG_PEER_WAIT_TIMEOUT") - high_watermark_override = os.environ.get("DISAGG_IN_FLIGHT_HIGH_WATERMARK") - in_flight_high_watermark = int( - high_watermark_override or cfg.runtime.in_flight_high_watermark - ) - low_watermark_override = os.environ.get("DISAGG_IN_FLIGHT_LOW_WATERMARK") - # Preserve the legacy one-watermark environment override: when only the - # old high value is supplied, resume at that same threshold. - in_flight_low_watermark = ( - int(low_watermark_override) - if low_watermark_override is not None - else ( - None - if high_watermark_override is not None - else cfg.runtime.in_flight_low_watermark - ) - ) _workers, drive = build_disagg_online_producer( algorithm=algorithm, modality=modality, @@ -624,7 +670,7 @@ def _build_online( target_repr=target_repr, aux_hidden_state_layer_ids=layers, prompt_epochs=cfg.training.num_epochs, - prompt_seed=cfg.training.seed, + prompt_seed=_online_prompt_seed(cfg), lease=cfg.runtime.producer_lease, in_flight_high_watermark=in_flight_high_watermark, in_flight_low_watermark=in_flight_low_watermark, diff --git a/specforge/training/model_loading.py b/specforge/training/model_loading.py index 75fb32f82..a4f8d8616 100644 --- a/specforge/training/model_loading.py +++ b/specforge/training/model_loading.py @@ -114,7 +114,9 @@ def _draft_config_from_dict(payload: Dict[str, Any]) -> "PretrainedConfig": payload["architectures"] = [architecture] payload["tie_word_embeddings"] = False payload.setdefault("use_cache", True) - if payload.get("draft_vocab_size") is None: + if "draft_vocab_size" in payload and payload["draft_vocab_size"] is None: + raise ValueError("draft config draft_vocab_size cannot be null") + if "draft_vocab_size" not in payload: payload["draft_vocab_size"] = payload.get("vocab_size") return DRAFT_REGISTRY[architecture].config_class.from_dict(payload) @@ -288,11 +290,12 @@ def resolve_draft_config( draft_config = _generate_draft_config(cfg, provider) expected = provider.architecture + compatible = provider.compatible_architectures or frozenset({expected}) architectures = list(getattr(draft_config, "architectures", None) or []) - if architectures != [expected]: + if len(architectures) != 1 or architectures[0] not in compatible: raise ValueError( f"training.strategy={cfg.training.strategy!r} requires draft " - f"architecture {expected}, got {architectures!r}" + f"architecture in {sorted(compatible)!r}, got {architectures!r}" ) _apply_draft_overrides(cfg, draft_config, provider) return draft_config diff --git a/specforge/training/strategies/base.py b/specforge/training/strategies/base.py index 2ce45943d..69495a295 100644 --- a/specforge/training/strategies/base.py +++ b/specforge/training/strategies/base.py @@ -29,11 +29,17 @@ @dataclass(frozen=True) class StepOutput: """Per-step result: loss + strategy-specific metrics, kept generic so - per-position (TTT) and single-scalar strategies share one trainer loop.""" + per-position (TTT) and single-scalar strategies share one trainer loop. + + ``loss_terms`` carries an additive objective numerator and denominator when + gradients and reported loss must be normalized across the full optimizer + window. + """ loss: torch.Tensor metrics: Dict[str, Any] ratio_metrics: Dict[str, Tuple[Any, Any]] = field(default_factory=dict) + loss_terms: Optional[Tuple[torch.Tensor, torch.Tensor]] = None @dataclass(frozen=True) @@ -60,6 +66,25 @@ def linear_lambda_base( return max(0.0, min(1.0, lambda_start * (1.0 - progress))) +def _cpu_max_valid_anchors(loss_mask: torch.Tensor) -> Optional[int]: + """Count the widest valid anchor row without synchronizing the GPU. + + Online/offline loaders hand strategies CPU tensors. Computing this one + scalar before the asynchronous H2D copies lets DFlash-family models size + their anchor tensors without a CUDA ``item()`` in the forward critical + path. Direct GPU callers keep the model's existing fallback. + """ + if loss_mask.device.type != "cpu": + return None + num_candidates = max(loss_mask.shape[1] - 1, 0) + valid = (loss_mask[:, :num_candidates] > 0.5) & ( + loss_mask[:, 1 : num_candidates + 1] > 0.5 + ) + if valid.shape[0] == 0: + return 0 + return int(valid.sum(dim=1).max().item()) + + class DraftTrainStrategy(abc.ABC): name: str required_features: set @@ -110,9 +135,17 @@ def _prepare_eagle_target( input_ids, target, loss_mask = target_head.preprocess( input_ids, target, loss_mask ) - target = target_head(target.to(device)) - return input_ids.to(device), target, loss_mask.to(device) - return input_ids.to(device), target.to(device), loss_mask.to(device) + target = target_head(target.to(device, non_blocking=True)) + return ( + input_ids.to(device, non_blocking=True), + target, + loss_mask.to(device, non_blocking=True), + ) + return ( + input_ids.to(device, non_blocking=True), + target.to(device, non_blocking=True), + loss_mask.to(device, non_blocking=True), + ) class Eagle3TrainStrategy(DraftTrainStrategy): @@ -139,12 +172,14 @@ def __init__( *, target_head: Optional[nn.Module] = None, ploss_decay: float = 0.8, + trim_loss_positions: bool = False, compact_teacher: bool = False, compact_teacher_chunk_size: Optional[int] = None, ) -> None: self.eagle3_model = eagle3_model self.target_head = target_head self.ploss_decay = ploss_decay + self.trim_loss_positions = trim_loss_positions self.compact_teacher = compact_teacher self.compact_teacher_chunk_size = compact_teacher_chunk_size if compact_teacher: @@ -249,9 +284,9 @@ def forward_loss( input_ids, target_hidden, loss_mask = self.target_head.preprocess( t["input_ids"], t["target"], t["loss_mask"] ) - input_ids = input_ids.to(device) - target_hidden = target_hidden.to(device) - loss_mask = loss_mask.to(device) + input_ids = input_ids.to(device, non_blocking=True) + target_hidden = target_hidden.to(device, non_blocking=True) + loss_mask = loss_mask.to(device, non_blocking=True) from specforge.core.compact_teacher import build_offline_teacher_inputs target, compact_kwargs = build_offline_teacher_inputs( @@ -275,11 +310,16 @@ def forward_loss( metric_loss_denoms, ) = self.eagle3_model( input_ids=input_ids, - attention_mask=t["attention_mask"].to(device), + attention_mask=t["attention_mask"].to(device, non_blocking=True), loss_mask=loss_mask, target=target, - hidden_states=t["hidden_state"].to(device), - position_ids=position_ids.to(device) if position_ids is not None else None, + hidden_states=t["hidden_state"].to(device, non_blocking=True), + position_ids=( + position_ids.to(device, non_blocking=True) + if position_ids is not None + else None + ), + trim_loss_positions=self.trim_loss_positions, **compact_kwargs, ) weights = [self.ploss_decay**i for i in range(len(plosses))] @@ -369,11 +409,11 @@ def forward_loss( loss, model_metrics = self.peagle_model( input_ids=input_ids, - attention_mask=tensors["attention_mask"].to(device), + attention_mask=tensors["attention_mask"].to(device, non_blocking=True), loss_mask=loss_mask, target=target, - hidden_states=tensors["hidden_state"].to(device), - lengths=lengths.to(device), + hidden_states=tensors["hidden_state"].to(device, non_blocking=True), + lengths=lengths.to(device, non_blocking=True), ) if not isinstance(loss, torch.Tensor) or loss.numel() != 1: raise ValueError( @@ -433,23 +473,22 @@ def forward_loss( self.validate_batch(batch) t = batch.tensors device = self._device() - # Multimodal capture additionally stores server-produced mRoPE - # position ids (B, S, 3); text runs do not carry the tensor at all, - # and wrapped models without mRoPE support must not see the kwarg. - forward_kwargs: Dict[str, Any] = {} - position_ids = t.get("position_ids") - if position_ids is not None: - forward_kwargs["position_ids"] = position_ids.to(device) + max_valid_anchors = _cpu_max_valid_anchors(t["loss_mask"]) loss, accuracy, model_metrics = self.dflash_model( - input_ids=t["input_ids"].to(device), - hidden_states=t["hidden_states"].to(device), - loss_mask=t["loss_mask"].to(device), - **forward_kwargs, + input_ids=t["input_ids"].to(device, non_blocking=True), + hidden_states=t["hidden_states"].to(device, non_blocking=True), + loss_mask=t["loss_mask"].to(device, non_blocking=True), + max_valid_anchors=max_valid_anchors, ) metrics = {"accuracy": accuracy.detach()} if "accuracy_denom" in model_metrics: metrics["accuracy_denom"] = model_metrics["accuracy_denom"] - return StepOutput(loss=loss, metrics=metrics) + return StepOutput( + loss=loss, + metrics=metrics, + ratio_metrics=model_metrics.get("ratio_metrics", {}), + loss_terms=model_metrics.get("loss_terms"), + ) def checkpoint_state_filter(self, state_dict: Dict[str, Any]) -> Dict[str, Any]: # Everything trainable lives under draft_model.; the target @@ -487,11 +526,15 @@ def forward_loss( self.validate_batch(batch) t = batch.tensors device = self._device() + max_valid_anchors = _cpu_max_valid_anchors(t["loss_mask"]) loss, accuracy, model_metrics = self.dspark_model( - input_ids=t["input_ids"].to(device), - hidden_states=t["hidden_states"].to(device), - loss_mask=t["loss_mask"].to(device), - target_last_hidden_states=t["target_last_hidden_states"].to(device), + input_ids=t["input_ids"].to(device, non_blocking=True), + hidden_states=t["hidden_states"].to(device, non_blocking=True), + loss_mask=t["loss_mask"].to(device, non_blocking=True), + target_last_hidden_states=t["target_last_hidden_states"].to( + device, non_blocking=True + ), + max_valid_anchors=max_valid_anchors, ) metrics = { "accuracy": accuracy.detach(), @@ -520,6 +563,63 @@ def checkpoint_state_filter(self, state_dict: Dict[str, Any]) -> Dict[str, Any]: } +class MTPTrainStrategy(DraftTrainStrategy): + """MTP strategy over ``OnlineMTPModel`` with final-hidden supervision.""" + + name = "mtp" + required_features = { + "input_ids", + "loss_mask", + "target_last_hidden_states", + } + + def __init__(self, mtp_model: nn.Module) -> None: + self.mtp_model = mtp_model + + def trainable_module(self) -> nn.Module: + return self.mtp_model + + def _device(self) -> torch.device: + return next(self.mtp_model.parameters()).device + + def forward_loss( + self, batch: TrainBatch, ctx: Optional[StepContext] = None + ) -> StepOutput: + self.validate_batch(batch) + t = batch.tensors + device = self._device() + # OnlineMTPModel performs the next-token shift internally and returns + # per-position correct/denominator tensors (single-layer: length-1 lists). + loss, corrects, denoms = self.mtp_model( + input_ids=t["input_ids"].to(device), + hidden_states=t["target_last_hidden_states"].to(device), + loss_mask=t["loss_mask"].to(device), + ) + correct_sum = corrects[0].sum() + denom_sum = denoms[0].sum() + metrics = { + "accuracy": (correct_sum / denom_sum.clamp_min(1)).detach(), + "accuracy_denom": denom_sum.detach(), + } + return StepOutput( + loss=loss, + metrics=metrics, + ratio_metrics={"accuracy": (correct_sum, denom_sum)}, + # TrainerCore backpropagates additive numerators and divides by the + # global token denominator across accumulation steps / DP ranks. + loss_terms=(loss * denom_sum, denom_sum), + ) + + def checkpoint_state_filter(self, state_dict: Dict[str, Any]) -> Dict[str, Any]: + # Everything trainable lives under draft_model.; persisting the stripped + # keys (embed_tokens.* + mtp.*) matches the native serving layout. + return { + k.replace("draft_model.", ""): v + for k, v in state_dict.items() + if "draft_model." in k + } + + class DominoTrainStrategy(DraftTrainStrategy): """Domino block-parallel strategy wrapping ``OnlineDominoModel``. @@ -563,17 +663,19 @@ def forward_loss( t = batch.tensors device = self._device() lambda_base = self._lambda_base(ctx) + max_valid_anchors = _cpu_max_valid_anchors(t["loss_mask"]) loss, accuracy, model_metrics = self.domino_model( - input_ids=t["input_ids"].to(device), - hidden_states=t["hidden_states"].to(device), - loss_mask=t["loss_mask"].to(device), + input_ids=t["input_ids"].to(device, non_blocking=True), + hidden_states=t["hidden_states"].to(device, non_blocking=True), + loss_mask=t["loss_mask"].to(device, non_blocking=True), lambda_base=lambda_base, + max_valid_anchors=max_valid_anchors, ) metrics = dict(model_metrics) metrics["accuracy"] = accuracy.detach() metrics.setdefault( "lambda_base", - torch.tensor(float(lambda_base), device=loss.device), + float(lambda_base), ) return StepOutput(loss=loss, metrics=metrics) diff --git a/specforge/training/trainer.py b/specforge/training/trainer.py index b1279a708..d465b947b 100644 --- a/specforge/training/trainer.py +++ b/specforge/training/trainer.py @@ -22,6 +22,8 @@ from contextlib import nullcontext from typing import Callable, Mapping, Optional +import torch + from specforge.algorithms.common.providers import ( MODEL_PROVENANCE_CONTRACT_KEY, OMITTED_STATE_FINGERPRINT_CONTRACT_KEY, @@ -152,6 +154,9 @@ def __init__( strategy=algorithm_name, ack=not defer_queue_ack, num_workers=dataloader_num_workers, + # Pin in the existing loader workers so Domino's non-blocking H2D + # copies do not add pinning work to the training thread. + pin_memory=dataloader_num_workers > 0 and torch.cuda.is_available(), ) if refs_for_epoch is not None: expected_refs = len(ref_source["refs"]) @@ -539,7 +544,7 @@ def close_loader() -> None: self._on_fit_failure(exc) raise finally: - primary_exception = sys.exception() + primary_exception = sys.exc_info()[1] cleanup_errors: list[tuple[str, BaseException]] = [] def capture_cleanup(label: str, action) -> None: diff --git a/tests/test_algorithms/test_builtin_parity.py b/tests/test_algorithms/test_builtin_parity.py index a4d62d356..50f366034 100644 --- a/tests/test_algorithms/test_builtin_parity.py +++ b/tests/test_algorithms/test_builtin_parity.py @@ -144,6 +144,11 @@ def test_peagle_reuses_the_eagle_server_contract_explicitly(self): self.assertEqual("hidden_state", peagle_stream.target_representation) self.assertTrue(peagle.step.uses_external_target_head) + def test_dspark_uses_the_dedicated_server_capture_method(self): + stream = self.registry.resolve("dspark").providers.server_streaming_for("text") + + self.assertEqual("dspark", stream.capture_method) + def test_step_factories_preserve_concrete_strategy_types(self): expected = { "eagle3": "Eagle3TrainStrategy", diff --git a/tests/test_algorithms/test_builtin_providers.py b/tests/test_algorithms/test_builtin_providers.py index 1a6234909..59b8750a2 100644 --- a/tests/test_algorithms/test_builtin_providers.py +++ b/tests/test_algorithms/test_builtin_providers.py @@ -22,7 +22,7 @@ from specforge.algorithms.contracts import AlgorithmSpec, FeatureMode REPO_ROOT = Path(__file__).resolve().parents[2] -BUILTINS = ("dflash", "domino", "dspark", "eagle3", "peagle") +BUILTINS = ("dflash", "domino", "dspark", "eagle3", "mtp", "peagle") class BuiltinProviderContractTest(unittest.TestCase): @@ -55,6 +55,19 @@ def test_every_registration_pairs_contract_and_providers(self): } self.assertEqual(contract_keys, provider_keys) + def test_dflash_family_requires_a_trainable_block_size(self): + for algorithm in ("dflash", "domino", "dspark"): + minimum_loss_tokens = self.registry.resolve( + algorithm + ).providers.model.minimum_loss_tokens + with self.subTest(algorithm=algorithm): + self.assertEqual( + minimum_loss_tokens(None, SimpleNamespace(block_size=2)), + 2, + ) + with self.assertRaisesRegex(ValueError, "block_size >= 2"): + minimum_loss_tokens(None, SimpleNamespace(block_size=1)) + def test_algorithm_metadata_has_no_factories_or_topology_flags(self): field_names = {field.name for field in fields(AlgorithmSpec)} self.assertEqual( @@ -101,6 +114,7 @@ def test_draft_resolution_stays_outside_algorithm_spec(self): self.assertEqual( { "architecture", + "compatible_architectures", "target_defaults", "expected_auto_map_model", "apply_overrides", @@ -135,7 +149,7 @@ def test_target_derived_defaults_and_overrides_are_provider_owned(self): self.assertEqual(vocab_size, defaults.draft_vocab_size) self.assertEqual(has_override, policy.apply_overrides is not None) - for name in ("domino", "dspark"): + for name in ("domino", "dspark", "mtp"): with self.subTest(algorithm=name): policy = self.registry.resolve(name).providers.model.draft_config self.assertIsNone(policy.target_defaults) @@ -158,6 +172,7 @@ def test_builtin_modalities_are_pinned_per_algorithm(self): def test_builtin_resume_contracts_cover_resolved_objective_semantics(self): training = SimpleNamespace( attention_backend="flex_attention", + trim_loss_positions=True, compact_teacher=True, compact_teacher_chunk_size=1024, lambda_base_start=0.75, @@ -201,6 +216,7 @@ def test_builtin_resume_contracts_cover_resolved_objective_semantics(self): "dflash": dflash_family, "domino": dflash_family, "dspark": dflash_family, + "mtp": SimpleNamespace(), } expected_keys = { "eagle3": { @@ -208,6 +224,7 @@ def test_builtin_resume_contracts_cover_resolved_objective_semantics(self): "eagle3_lk_loss_type", "eagle3_kl_scale", "eagle3_kl_decay", + "eagle3_trim_loss_positions", "eagle3_compact_teacher", }, "peagle": { @@ -235,6 +252,12 @@ def test_builtin_resume_contracts_cover_resolved_objective_semantics(self): "dspark_l1_loss_alpha", "dspark_confidence_head_alpha", }, + "mtp": { + "mtp_draft_num_hidden_layers", + "mtp_draft_vocab_size", + "mtp_share_lm_head", + "mtp_attention_backend", + }, } for name in BUILTINS: @@ -254,6 +277,7 @@ def test_step_runtime_config_binds_options_contract_and_missing_key_policy(self) config = SimpleNamespace( training=SimpleNamespace( attention_backend="flex_attention", + trim_loss_positions=False, compact_teacher=False, compact_teacher_chunk_size=None, ) @@ -289,6 +313,7 @@ def test_step_runtime_config_binds_options_contract_and_missing_key_policy(self) ( ("compact_teacher", False), ("compact_teacher_chunk_size", None), + ("trim_loss_positions", False), ), ) self.assertEqual( @@ -419,7 +444,7 @@ def test_building_catalog_does_not_import_training_or_torch(self): code = ( "import sys; " "from specforge.algorithms.builtin import builtin_algorithm_registry; " - "r=builtin_algorithm_registry(); assert len(r)==5; " + "r=builtin_algorithm_registry(); assert len(r)==6; " "assert 'torch' not in sys.modules; " "assert 'specforge.training.strategies.registry' not in sys.modules" ) diff --git a/tests/test_algorithms/test_dflash_multimodal.py b/tests/test_algorithms/test_dflash_multimodal.py index af10a7e79..7bfc5d8c2 100644 --- a/tests/test_algorithms/test_dflash_multimodal.py +++ b/tests/test_algorithms/test_dflash_multimodal.py @@ -28,11 +28,10 @@ def test_dflash_registers_the_multimodal_streaming_contract(self): ) self.assertEqual( set(contract.required_tensors), - {"input_ids", "loss_mask", "hidden_states", "position_ids"}, + {"input_ids", "loss_mask", "hidden_states"}, ) provider = registration.providers.server_streaming_for("multimodal") self.assertEqual(provider.capture_method, "dflash") - self.assertEqual(provider.layout.position_ids_feature, "position_ids") self.assertEqual(provider.layout.aux_feature, "hidden_states") def test_other_builtins_have_no_multimodal_contract(self): @@ -171,111 +170,57 @@ def test_build_request_inputs_uses_collapsed_ids_and_image_data(self): self.assertEqual(request["image_data"], ["aGVsbG8=", None]) -class ServerCapturePositionIdsTest(unittest.TestCase): - def _adapter(self, position_ids_feature): - from specforge.inference.adapters.server_capture import ( - ServerCaptureSchema, - SGLangServerCaptureAdapter, - ) - - class _FakeStore: - store_id = "store" - - def adopt(self, ref): - pass - - def discard_external_attempts(self, *args, **kwargs): - pass +@unittest.skipUnless(TORCH_AVAILABLE, "requires torch") +class DraftPositionsTest(unittest.TestCase): + """Drafts always train on the plain 1D position convention, identical to + the text-only path; multimodal capture stores no position ids.""" - def track_external_attempt(self, *args, **kwargs): - pass + def _build_model(self): + import torch as t + from torch import nn - schema = ServerCaptureSchema( - aux_feature="hidden_states", - last_hidden_feature=None, - passthrough=( - ("input_ids", "input_ids", ()), - ("loss_mask", "loss_mask", ()), - ), - position_ids_feature=position_ids_feature, - ) - return SGLangServerCaptureAdapter( - "http://localhost:1", - _FakeStore(), - run_id="run", - algorithm="dflash", - schema=schema, - ) + from specforge.algorithms.common.dflash_family_model import OnlineDFlashModel - def _task(self): - return SimpleNamespace( - task_id="t0", - attempt=0, - payload={"input_ids": [5, 6, 7], "loss_mask": [0, 1, 1]}, - metadata={}, + class _StubDraftModel(nn.Module): + def __init__(self): + super().__init__() + self.recorded = {} + + def forward( + self, + position_ids=None, + noise_embedding=None, + target_hidden=None, + attention_mask=None, + ): + self.recorded["position_ids"] = position_ids + return t.zeros(1) + + return OnlineDFlashModel( + draft_model=_StubDraftModel(), + target_lm_head=nn.Linear(8, 32, bias=False), + target_embed_tokens=nn.Embedding(32, 8), + mask_token_id=31, + block_size=2, + attention_backend="sdpa", + num_anchors=4, ) - def test_payload_requests_position_ids_artifact_when_configured(self): - adapter = self._adapter("position_ids") - payload = adapter._spec_capture_payload(self._task()) - self.assertEqual( - payload["features"], - {"aux": "hidden_states", "position_ids": "position_ids"}, - ) - - def test_payload_omits_position_ids_artifact_when_unset(self): - adapter = self._adapter(None) - payload = adapter._spec_capture_payload(self._task()) - self.assertEqual(payload["features"], {"aux": "hidden_states"}) - - -@unittest.skipUnless(TORCH_AVAILABLE, "requires torch") -class VlmCollatorTest(unittest.TestCase): - def test_collator_pads_position_ids_like_other_features(self): - from specforge.algorithms.common.dflash_family_data import build_vlm_collator - - collate = build_vlm_collator() - features = [ - { - "input_ids": torch.tensor([[1, 2, 3]]), - "loss_mask": torch.tensor([[0, 1, 1]]), - "hidden_states": torch.zeros(1, 3, 8), - "position_ids": torch.arange(9).reshape(1, 3, 3), - }, - { - "input_ids": torch.tensor([[4]]), - "loss_mask": torch.tensor([[1]]), - "hidden_states": torch.zeros(1, 1, 8), - "position_ids": torch.arange(3).reshape(1, 1, 3), - }, - ] - batch = collate(features) - self.assertEqual(tuple(batch["input_ids"].shape), (2, 3)) - self.assertEqual(tuple(batch["position_ids"].shape), (2, 3, 3)) - # Padding is zeros on the sequence axis. - self.assertTrue((batch["position_ids"][1, 1:] == 0).all()) - self.assertEqual(batch["position_ids"][0, 2].tolist(), [6, 7, 8]) - - -@unittest.skipUnless(TORCH_AVAILABLE, "requires torch") -class MropeDraftPositionsTest(unittest.TestCase): - def test_gathered_draft_positions_follow_anchor_offsets(self): + def test_positions_follow_the_1d_convention(self): import torch as t - from specforge.algorithms.common.dflash_family_model import OnlineDFlashModel - - model = OnlineDFlashModel.__new__(OnlineDFlashModel) - model.block_size = 2 - anchors = t.tensor([[1, 3]]) - stored = t.arange(5 * 3).reshape(1, 5, 3) - offsets = t.arange(model.block_size).view(1, 1, -1) - draft_indices = (anchors.unsqueeze(-1) + offsets).view(1, -1) - gathered = t.gather(stored, 1, draft_indices.unsqueeze(-1).expand(-1, -1, 3)) - full = t.cat([stored, gathered], dim=1).permute(2, 0, 1) - self.assertEqual(tuple(full.shape), (3, 1, 5 + 4)) - # Draft slot for anchor=1: positions of indices 1 and 2. - self.assertEqual(full[:, 0, 5].tolist(), [3, 4, 5]) - self.assertEqual(full[:, 0, 6].tolist(), [6, 7, 8]) + model = self._build_model() + b, s = 2, 8 + input_ids = t.randint(0, 31, (b, s)) + hidden_states = t.randn(b, s, 16) + loss_mask = t.ones(b, s) + t.manual_seed(0) + model._forward_draft_blocks(input_ids, hidden_states, loss_mask) + got = model.draft_model.recorded["position_ids"] + self.assertEqual(got.ndim, 2) + self.assertTrue( + t.equal(got[:, :s], t.arange(s).unsqueeze(0).expand(b, -1)) + ) if __name__ == "__main__": diff --git a/tests/test_algorithms/test_offline_capture_layout.py b/tests/test_algorithms/test_offline_capture_layout.py index 3acc1087d..a7407f6ae 100644 --- a/tests/test_algorithms/test_offline_capture_layout.py +++ b/tests/test_algorithms/test_offline_capture_layout.py @@ -40,6 +40,12 @@ def test_builtin_offline_layouts_materialize_exact_storage_schemas(self): "target_last_hidden_states": "last_hidden_states", }, } + expected_capture_methods = { + "eagle3": "eagle3", + "dflash": "dflash", + "domino": "dflash", + "dspark": "dspark", + } sources = { "input_ids": torch.tensor([1, 2, 3]), "loss_mask": torch.tensor([1, 1, 0]), @@ -54,7 +60,7 @@ def test_builtin_offline_layouts_materialize_exact_storage_schemas(self): record = provider.capture_layout.materialize(sources) self.assertEqual( - "eagle3" if strategy == "eagle3" else "dflash", + expected_capture_methods[strategy], provider.capture_layout.capture_method, ) @@ -98,6 +104,23 @@ def test_materialize_preserves_arbitrary_auxiliary_layer_counts(self): record["hidden_states"].shape[-1], ) + def test_dflash_family_normalizers_require_adjacent_supervision(self): + raw = { + "input_ids": torch.tensor([1, 2, 3]), + "loss_mask": torch.tensor([1, 0, 1]), + "hidden_states": torch.randn(1, 3, 8), + "target_last_hidden_states": torch.randn(1, 3, 8), + } + for strategy in ("dflash", "domino", "dspark"): + with self.subTest(strategy=strategy): + normalizer = ( + self.registry.resolve(strategy) + .providers.offline_for("text") + .build_normalizer(3) + ) + with self.assertRaisesRegex(ValueError, "two consecutive"): + normalizer(raw) + def test_duplicate_output_names_are_rejected(self): with self.assertRaisesRegex(ValueError, "duplicate.*hidden_states"): OfflineCaptureLayout( @@ -134,16 +157,19 @@ def test_local_capture_forwards_the_algorithm_capture_method(self): backend = mock.Mock() capture = OfflineSGLangCapture(backend) - capture.set_capture_layers( - [1, 9, 17, 25, 33], - capture_method="dflash", - ) + for capture_method in ("dflash", "dspark"): + with self.subTest(capture_method=capture_method): + backend.reset_mock() + capture.set_capture_layers( + [1, 9, 17, 25, 33], + capture_method=capture_method, + ) - self.assertEqual("dflash", capture.capture_method) - backend.set_capture_layers.assert_called_once_with( - [1, 9, 17, 25, 33], - capture_method="dflash", - ) + self.assertEqual(capture_method, capture.capture_method) + backend.set_capture_layers.assert_called_once_with( + [1, 9, 17, 25, 33], + capture_method=capture_method, + ) if __name__ == "__main__": diff --git a/tests/test_config/test_example_draft_config_wiring.py b/tests/test_config/test_example_draft_config_wiring.py index 74e610f9f..e8fd059af 100644 --- a/tests/test_config/test_example_draft_config_wiring.py +++ b/tests/test_config/test_example_draft_config_wiring.py @@ -29,7 +29,7 @@ def _yaml_scalar(path: Path, key: str) -> Optional[str]: def _local_draft_configs(): - for recipe in sorted(EXAMPLE_CONFIG_DIR.glob("*.yaml")): + for recipe in sorted(EXAMPLE_CONFIG_DIR.rglob("*.yaml")): source = _yaml_scalar(recipe, "draft_model_config") if source is None or not source.endswith(".json"): continue @@ -49,16 +49,29 @@ def test_local_draft_architecture_matches_recipe_strategy(self): draft_provider = algorithm.providers.model.draft_config self.assertTrue(draft_config.is_file(), draft_config) payload = json.loads(draft_config.read_text()) - self.assertEqual( - payload.get("architectures"), - [draft_provider.architecture], + architectures = payload.get("architectures") + self.assertIsInstance(architectures, list) + self.assertEqual(len(architectures), 1) + architecture = architectures[0] + self.assertIn( + architecture, + draft_provider.compatible_architectures, ) expected_auto_model = draft_provider.expected_auto_map_model - if expected_auto_model is not None: + actual_auto_model = payload.get("auto_map", {}).get("AutoModel") + if ( + expected_auto_model is not None + and architecture == draft_provider.architecture + ): self.assertEqual( - payload.get("auto_map", {}).get("AutoModel"), + actual_auto_model, expected_auto_model, ) + elif actual_auto_model is not None: + self.assertEqual( + actual_auto_model.rsplit(".", 1)[-1], + architecture, + ) def test_only_future_vlm_draft_configs_lack_a_unified_recipe(self): referenced = { diff --git a/tests/test_config/test_launch_topology.py b/tests/test_config/test_launch_topology.py index d92e1fe73..3c431c21d 100644 --- a/tests/test_config/test_launch_topology.py +++ b/tests/test_config/test_launch_topology.py @@ -10,419 +10,61 @@ REPO_ROOT = Path(__file__).resolve().parents[2] EXAMPLE_CONFIG_DIR = REPO_ROOT / "examples" / "configs" -EXPECTED_NPROC_PER_NODE = { - "deepseek-v2-lite-eagle3-online.yaml": 8, - "deepseek-v3-671b-eagle3-offline.yaml": 8, - "deepseek-v3-671b-eagle3-online.yaml": 8, - "gemma3-1b-eagle3-online.yaml": 1, - "glm-5.2-dspark-disaggregated.yaml": 1, - "gpt-oss-120b-eagle3-online.yaml": 8, - "gpt-oss-20b-eagle3-online.yaml": 8, - "lfm2.5-1.2b-instruct-dflash-online.yaml": 8, - "inkling-dspark-disaggregated.yaml": 1, - "ling-flash-2.0-eagle3-offline.yaml": 8, - "ling-flash-2.0-eagle3-online.yaml": 8, - "llama3.1-8b-eagle3-offline.yaml": 1, - "llama3.1-8b-eagle3-online.yaml": 1, - "llama3.3-70b-eagle3-online.yaml": 8, - "llama4-scout-17b-16e-eagle3-online.yaml": 8, - "longcat-flash-dflash-online.yaml": 1, - "longcat-flash-eagle3-online.yaml": 1, - "phi4-eagle3-online.yaml": 1, - "qwen2.5-0.5b-dflash-online.yaml": 1, - "qwen2.5-0.5b-eagle3-online.yaml": 1, - "qwen2.5-7b-eagle3-offline.yaml": 1, - "qwen2.5-7b-eagle3-offline-disaggregated.yaml": 1, - "qwen3-235b-a22b-eagle3-online.yaml": 8, - "qwen3-30b-a3b-eagle3.1-online.yaml": 4, - "qwen3-30b-a3b-eagle3-online.yaml": 4, - "qwen3-32b-eagle3-online.yaml": 4, - "qwen3-4b-dflash-online.yaml": 8, - "qwen3-4b-dspark-offline.yaml": 1, - "qwen3-4b-dspark-disaggregated.yaml": 1, - "qwen3-4b-eagle3-online.yaml": 1, - "qwen3-8b-dspark-offline.yaml": 1, - "qwen3-8b-dflash-disaggregated.yaml": 4, - "qwen3-8b-dflash-1server-dp7-disaggregated.yaml": 7, - "qwen3-8b-dflash-offline.yaml": 1, - "qwen3-8b-dflash-online.yaml": 8, - "qwen3-8b-domino-1server-dp7-disaggregated.yaml": 7, - "qwen3-8b-domino-disaggregated.yaml": 4, - "qwen3-8b-domino-multiserver-disaggregated.yaml": 2, - "qwen3-8b-domino-offline.yaml": 1, - "qwen3-8b-domino-online.yaml": 8, - "qwen3-8b-dpace-online.yaml": 8, - "qwen3-8b-dspark-disaggregated.yaml": 1, - "qwen3-8b-eagle3-offline-disaggregated.yaml": 1, - "qwen3-8b-eagle3-offline.yaml": 1, - "qwen3-8b-eagle3-disaggregated.yaml": 1, - "qwen3-8b-peagle-disaggregated.yaml": 1, - "qwen3-coder-30b-a3b-eagle3-online.yaml": 4, - "qwen3-coder-480b-a35b-eagle3-offline.yaml": 8, - "qwen3-coder-480b-a35b-eagle3-online.yaml": 8, - "qwen3-coder-next-eagle3-online.yaml": 8, - "qwen3-next-80b-a3b-eagle3-online.yaml": 8, - "qwen3.5-35b-a3b-dflash-online.yaml": 4, - "qwen3.5-35b-a3b-eagle3-offline.yaml": 4, - "qwen3.5-35b-a3b-eagle3-online.yaml": 2, - "qwen3.5-4b-dflash-online-npu.yaml": 8, - "qwen3.5-4b-domino-online-npu.yaml": 8, - "qwen3.5-4b-vl-dflash-disaggregated.yaml": 4, - "qwen3.6-27b-dflash-disaggregated.yaml": 2, - "qwen3.6-27b-dflash-1server-dp2-disaggregated.yaml": 2, - "qwen3.6-27b-dflash-multiserver-disaggregated.yaml": 2, - "qwen3.6-27b-dflash-online.yaml": 8, - "qwen3.6-27b-domino-online.yaml": 8, - "qwen3.6-27b-dspark-disaggregated.yaml": 1, - "qwq-32b-eagle3-online.yaml": 4, -} - -LOCAL_MOONCAKE_ENDPOINTS = { - "mooncake_metadata_server": "http://127.0.0.1:35880/metadata", - "mooncake_master_server_addr": "127.0.0.1:35551", - "mooncake_protocol": "tcp", -} - -EXPECTED_DISAGGREGATED = { - "glm-5.2-dspark-disaggregated.yaml": { - "control_dir": "outputs/glm-5.2-dspark-disaggregated/control", - "backend": "mooncake", - "server_urls": ["http://127.0.0.1:30000"], - **LOCAL_MOONCAKE_ENDPOINTS, - }, - "inkling-dspark-disaggregated.yaml": { - "control_dir": "outputs/inkling-dspark-disaggregated/control", - "backend": "mooncake", - "server_urls": ["http://127.0.0.1:30000"], - **LOCAL_MOONCAKE_ENDPOINTS, - }, - "qwen2.5-7b-eagle3-offline-disaggregated.yaml": { - "control_dir": ("outputs/qwen2.5-7b-eagle3-offline-disaggregated/control"), - "backend": "shared_dir", - "store_root": ("outputs/qwen2.5-7b-eagle3-offline-disaggregated/features"), - }, - "qwen3-4b-dspark-disaggregated.yaml": { - "control_dir": "outputs/qwen3-4b-dspark-disaggregated/control", - "backend": "mooncake", - "server_urls": ["http://127.0.0.1:30000"], - **LOCAL_MOONCAKE_ENDPOINTS, - }, - "qwen3-8b-dflash-disaggregated.yaml": { - "control_dir": "outputs/qwen3-8b-dflash-disaggregated/control", - "consumer_state_dir": ("outputs/qwen3-8b-dflash-disaggregated/consumer-state"), - "backend": "mooncake", - "server_urls": ["http://127.0.0.1:30000"], - **LOCAL_MOONCAKE_ENDPOINTS, - }, - "qwen3-8b-dspark-disaggregated.yaml": { - "control_dir": "outputs/qwen3-8b-dspark-disaggregated/control", - "backend": "mooncake", - "server_urls": ["http://127.0.0.1:30000"], - **LOCAL_MOONCAKE_ENDPOINTS, - }, - "qwen3-8b-dflash-1server-dp7-disaggregated.yaml": { - "control_dir": ("outputs/qwen3-8b-dflash-1server-dp7-disaggregated/control"), - "backend": "mooncake", - "managed_local": { - "trainer_cuda_visible_devices": [ - "1", - "2", - "3", - "4", - "5", - "6", - "7", - ], - "mooncake": { - "protocol": "tcp", - "global_segment_size_bytes": 34359738368, - "local_buffer_size_bytes": 1073741824, - }, - "capture_servers": [ - { - "port": 30000, - "cuda_visible_devices": ["0"], - "tp_size": 1, - "mem_fraction_static": 0.5, - } - ], - }, - }, - "qwen3-8b-domino-1server-dp7-disaggregated.yaml": { - "control_dir": ("outputs/qwen3-8b-domino-1server-dp7-disaggregated/control"), - "backend": "mooncake", - "managed_local": { - "trainer_cuda_visible_devices": [ - "1", - "2", - "3", - "4", - "5", - "6", - "7", - ], - "mooncake": { - "protocol": "tcp", - "global_segment_size_bytes": 34359738368, - "local_buffer_size_bytes": 1073741824, - }, - "capture_servers": [ - { - "port": 30000, - "cuda_visible_devices": ["0"], - "tp_size": 1, - "mem_fraction_static": 0.5, - } - ], - }, - }, - "qwen3-8b-domino-disaggregated.yaml": { - "control_dir": "outputs/qwen3-8b-domino-disaggregated/control", - "backend": "mooncake", - "server_urls": ["http://127.0.0.1:30000"], - **LOCAL_MOONCAKE_ENDPOINTS, - }, - "qwen3-8b-domino-multiserver-disaggregated.yaml": { - "control_dir": ("outputs/qwen3-8b-domino-multiserver-disaggregated/control"), - "backend": "mooncake", - "managed_local": { - "trainer_cuda_visible_devices": ["2", "3"], - "mooncake": { - "protocol": "tcp", - "global_segment_size_bytes": 34359738368, - "local_buffer_size_bytes": 1073741824, - }, - "capture_servers": [ - { - "port": 30000, - "cuda_visible_devices": ["0"], - "tp_size": 1, - "mem_fraction_static": 0.85, - }, - { - "port": 30001, - "cuda_visible_devices": ["1"], - "tp_size": 1, - "mem_fraction_static": 0.85, - }, - ], - }, - }, - "qwen3-8b-eagle3-offline-disaggregated.yaml": { - "control_dir": ("outputs/qwen3-8b-eagle3-offline-disaggregated/control"), - "backend": "shared_dir", - "store_root": ("outputs/qwen3-8b-eagle3-offline-disaggregated/features"), - }, - "qwen3-8b-eagle3-disaggregated.yaml": { - "control_dir": "outputs/qwen3-8b-eagle3-disaggregated/control", - "consumer_state_dir": ("outputs/qwen3-8b-eagle3-disaggregated/consumer-state"), - "backend": "mooncake", - "server_urls": ["http://127.0.0.1:30000"], - **LOCAL_MOONCAKE_ENDPOINTS, - }, - "qwen3-8b-peagle-disaggregated.yaml": { - "control_dir": "outputs/qwen3-8b-peagle-disaggregated/control", - "consumer_state_dir": ("outputs/qwen3-8b-peagle-disaggregated/consumer-state"), - "backend": "mooncake", - "server_urls": ["http://127.0.0.1:30000"], - **LOCAL_MOONCAKE_ENDPOINTS, - }, - "qwen3.6-27b-dflash-disaggregated.yaml": { - "control_dir": "outputs/qwen3.6-27b-dflash-disaggregated/control", - "backend": "mooncake", - "server_urls": ["http://127.0.0.1:30000"], - **LOCAL_MOONCAKE_ENDPOINTS, - }, - "qwen3.5-4b-vl-dflash-disaggregated.yaml": { - "control_dir": "outputs/qwen3.5-4b-vl-dflash/control", - "consumer_state_dir": "outputs/qwen3.5-4b-vl-dflash/consumer-state", - "backend": "mooncake", - "server_urls": ["http://127.0.0.1:30000"], - **LOCAL_MOONCAKE_ENDPOINTS, - }, - "qwen3.6-27b-dflash-1server-dp2-disaggregated.yaml": { - "control_dir": ("outputs/qwen3.6-27b-dflash-1server-dp2-disaggregated/control"), - "backend": "mooncake", - "managed_local": { - "trainer_cuda_visible_devices": ["1", "2"], - "mooncake": { - "protocol": "tcp", - "global_segment_size_bytes": 68719476736, - "local_buffer_size_bytes": 1073741824, - }, - "capture_servers": [ - { - "port": 30000, - "cuda_visible_devices": ["0"], - "tp_size": 1, - "mem_fraction_static": 0.85, - } - ], - }, - }, - "qwen3.6-27b-dflash-multiserver-disaggregated.yaml": { - "control_dir": ("outputs/qwen3.6-27b-dflash-multiserver-disaggregated/control"), - "backend": "mooncake", - "managed_local": { - "trainer_cuda_visible_devices": ["4", "5"], - "mooncake": { - "protocol": "tcp", - "global_segment_size_bytes": 51539607552, - "local_buffer_size_bytes": 1073741824, - }, - "capture_servers": [ - { - "port": 30000, - "cuda_visible_devices": ["0", "1"], - "tp_size": 2, - "mem_fraction_static": 0.85, - }, - { - "port": 30001, - "cuda_visible_devices": ["2", "3"], - "tp_size": 2, - "mem_fraction_static": 0.85, - }, - ], - }, - }, - "qwen3.6-27b-dspark-disaggregated.yaml": { - "control_dir": "outputs/qwen3.6-27b-dspark-disaggregated/control", - "backend": "mooncake", - "managed_local": { - "trainer_cuda_visible_devices": ["1"], - "mooncake": { - "protocol": "tcp", - "global_segment_size_bytes": 68719476736, - "local_buffer_size_bytes": 1073741824, - }, - "capture_servers": [ - { - "port": 30000, - "cuda_visible_devices": ["0"], - "tp_size": 1, - "mem_fraction_static": 0.7, - } - ], - }, - }, -} - def _recipes() -> dict[str, Path]: - return { - path.name: path - for path in sorted(EXAMPLE_CONFIG_DIR.glob("*.yaml")) + paths = [ + path + for path in sorted(EXAMPLE_CONFIG_DIR.rglob("*.yaml")) if not path.name.startswith(".") - } + ] + recipes = {path.name: path for path in paths} + if len(recipes) != len(paths): + raise AssertionError("example recipe filenames must be globally unique") + return recipes class ExampleLaunchTopologyTest(unittest.TestCase): - def test_every_recipe_has_the_explicit_golden_topology(self): - recipes = _recipes() - self.assertEqual(len(EXPECTED_NPROC_PER_NODE), 64) - self.assertEqual(set(recipes), set(EXPECTED_NPROC_PER_NODE)) - - for filename, nproc_per_node in EXPECTED_NPROC_PER_NODE.items(): + def test_every_recipe_matches_its_directory_topology(self): + self.assertTrue( + (EXAMPLE_CONFIG_DIR / "online" / "colocated" / "README.md").is_file() + ) + for filename, path in _recipes().items(): with self.subTest(config=filename): - payload = yaml.safe_load(recipes[filename].read_text()) - training = payload["training"] - for legacy_field in ( - "deployment_mode", - "server_urls", - "metadata_db_path", - ): - self.assertNotIn(legacy_field, training) - self.assertNotIn("role", training) - - deployment = payload["deployment"] + payload = yaml.safe_load(path.read_text()) data = payload["data"] - online = bool(data.get("train_data_path") or data.get("prompts_path")) - expected_mode = ( - "disaggregated" - if online or filename in EXPECTED_DISAGGREGATED - else "local_colocated" + mode = ( + "online" + if data.get("train_data_path") or data.get("prompts_path") + else "offline" ) - self.assertEqual(deployment["mode"], expected_mode) - self.assertEqual( - deployment["trainer"], - {"nnodes": 1, "nproc_per_node": nproc_per_node}, + deployment = payload["deployment"] + topology = ( + "colocated" + if deployment["mode"] == "local_colocated" + else "disaggregated" ) - - expected_keys = {"mode", "trainer"} - if expected_mode == "disaggregated": - expected_keys.add("disaggregated") - if filename in EXPECTED_DISAGGREGATED: - self.assertEqual( - deployment["disaggregated"], - EXPECTED_DISAGGREGATED[filename], + expected_parent = Path(mode) / topology + if mode == "online" and topology == "disaggregated": + ownership = ( + "managed-local" + if "managed_local" in deployment["disaggregated"] + else "external" ) - elif online: - services = deployment["disaggregated"] - self.assertEqual(services["backend"], "mooncake") - self.assertTrue(services.get("server_urls")) - self.assertTrue(services.get("consumer_state_dir")) - self.assertEqual(payload["model"]["target_backend"], "sglang") - self.assertGreater(payload["training"]["num_epochs"], 0) - self.assertEqual(set(deployment), expected_keys) - - def test_golden_topologies_validate_for_their_declared_world_size(self): - for filename, path in _recipes().items(): - with self.subTest(config=filename): - config = Config.from_file(str(path)) - topology = config.deployment.trainer - expected_nproc = EXPECTED_NPROC_PER_NODE[filename] - self.assertEqual(topology.nnodes, 1) - self.assertEqual(topology.nproc_per_node, expected_nproc) - config.validate_world_size(topology.nnodes * expected_nproc) + expected_parent /= ownership + self.assertEqual( + path.relative_to(EXAMPLE_CONFIG_DIR).parent, + expected_parent, + ) - def test_every_recipe_keeps_trainer_tp_at_one(self): + def test_every_recipe_validates_for_its_declared_world_size(self): recipes = _recipes() + self.assertTrue(recipes) + for filename, path in recipes.items(): with self.subTest(config=filename): config = Config.from_file(str(path)) - self.assertEqual(config.training.tp_size, 1) - if config.deployment.mode != "disaggregated": - continue - self.assertEqual(config.training.role, "auto") - self.assertEqual(config.training.sp_ulysses_size, 1) - self.assertEqual(config.training.sp_ring_size, 1) - - def test_migrated_dspark_recipes_match_source_training_contract(self): - expected_save_intervals = { - "qwen3-8b-dspark-disaggregated.yaml": 125, - "glm-5.2-dspark-disaggregated.yaml": 16, - "inkling-dspark-disaggregated.yaml": 8, - } - for filename, save_interval in expected_save_intervals.items(): - with self.subTest(config=filename): - config = Config.from_file(str(EXAMPLE_CONFIG_DIR / filename)) topology = config.deployment.trainer - global_batch = ( - topology.nnodes - * topology.nproc_per_node - * config.training.batch_size - * config.training.accumulation_steps - ) - self.assertEqual(global_batch, 512) - self.assertEqual(config.training.num_epochs, 10) - self.assertIsNone(config.training.max_steps) - self.assertAlmostEqual(config.training.learning_rate, 6e-4) - self.assertEqual(config.training.num_anchors, 512) - self.assertEqual(config.training.loss_decay_gamma, 4.0) - self.assertEqual(config.training.objective_chunk_blocks, 128) - self.assertEqual(config.training.save_interval, save_interval) - - qwen = Config.from_file( - str(EXAMPLE_CONFIG_DIR / "qwen3-8b-dspark-disaggregated.yaml") - ) - self.assertEqual(qwen.data.chat_template, "qwen") - - qwen4b = Config.from_file( - str(EXAMPLE_CONFIG_DIR / "qwen3-4b-dspark-disaggregated.yaml") - ) - self.assertEqual(qwen4b.training.loss_decay_gamma, 4.0) - self.assertEqual(qwen4b.training.objective_chunk_blocks, 128) + config.validate_world_size(topology.nnodes * topology.nproc_per_node) if __name__ == "__main__": diff --git a/tests/test_config/test_schema.py b/tests/test_config/test_schema.py index c78e739a9..690835fba 100644 --- a/tests/test_config/test_schema.py +++ b/tests/test_config/test_schema.py @@ -626,9 +626,16 @@ def test_offline_dp_and_usp_topologies_are_validated(self): def test_overrides_coerce_and_revalidate(self): cfg = Config.model_validate(MINIMAL) out = apply_overrides( - cfg, ["training.learning_rate=1e-3", "training.max_steps=7", "run_id=r2"] + cfg, + [ + "training.learning_rate=1e-3", + "training.lr_scheduler=constant", + "training.max_steps=7", + "run_id=r2", + ], ) self.assertEqual(out.training.learning_rate, 1e-3) + self.assertEqual(out.training.lr_scheduler, "constant") self.assertEqual(out.training.max_steps, 7) self.assertEqual(out.run_id, "r2") # original untouched diff --git a/tests/test_config/test_server_only_online.py b/tests/test_config/test_server_only_online.py index 392b65a3b..a5bc1443f 100644 --- a/tests/test_config/test_server_only_online.py +++ b/tests/test_config/test_server_only_online.py @@ -73,7 +73,7 @@ def test_multimodal_modality_resolves_for_dflash(self): resolved = resolve_run(Config.model_validate(payload)) self.assertEqual(resolved.algorithm.name, "dflash") streaming = resolved.algorithm.providers.server_streaming_for("multimodal") - self.assertEqual(streaming.layout.position_ids_feature, "position_ids") + self.assertEqual(streaming.layout.aux_feature, "hidden_states") self.assertIsNotNone(streaming.create_input_adapter(resolved.config)) def test_multimodal_modality_is_rejected_for_text_only_algorithms(self): @@ -110,7 +110,7 @@ def _offline_payload(backend): def test_every_online_recipe_uses_the_server_data_plane(self): online_recipes = [] - for path in sorted(EXAMPLE_CONFIG_DIR.glob("*.yaml")): + for path in sorted(EXAMPLE_CONFIG_DIR.rglob("*.yaml")): payload = yaml.safe_load(path.read_text()) data = payload["data"] if not (data.get("train_data_path") or data.get("prompts_path")): @@ -124,10 +124,10 @@ def test_every_online_recipe_uses_the_server_data_plane(self): ) self.assertTrue(online_recipes) self.assertFalse( - (EXAMPLE_CONFIG_DIR / "qwen2.5-vl-7b-eagle3-online.yaml").exists() + any(EXAMPLE_CONFIG_DIR.rglob("qwen2.5-vl-7b-eagle3-online.yaml")) ) self.assertFalse( - (EXAMPLE_CONFIG_DIR / "qwen2.5-vl-32b-eagle3-online.yaml").exists() + any(EXAMPLE_CONFIG_DIR.rglob("qwen2.5-vl-32b-eagle3-online.yaml")) ) def test_application_resolution_accepts_the_server_only_contract(self): diff --git a/tests/test_config/test_unified_feature_reachability.py b/tests/test_config/test_unified_feature_reachability.py index f08bd01e2..680069324 100644 --- a/tests/test_config/test_unified_feature_reachability.py +++ b/tests/test_config/test_unified_feature_reachability.py @@ -28,185 +28,26 @@ "data": {"hidden_states_path": "features"}, } -EXPECTED_TRACKING = { - "lfm2.5-1.2b-instruct-dflash-online.yaml": { - "report_to": "wandb", - "wandb_project": "specforge-lfm2.5-1.2b-instruct-dflash", - "wandb_name": "lfm2.5-1.2b-instruct-dflash-perfectblend-8layers", - }, - "longcat-flash-dflash-online.yaml": { - "report_to": "wandb", - "wandb_project": "specforge-longcat-flash-dflash", - "wandb_name": "longcat-flash-dflash-sharegpt", - }, - "qwen3-4b-dflash-online.yaml": { - "report_to": "wandb", - "wandb_project": "specforge-qwen3-4b-dflash", - "wandb_name": "qwen3-4b-dflash-perfectblend", - }, - "qwen3-8b-dflash-disaggregated.yaml": { - "report_to": "none", - "wandb_project": "qwen3-8b-dflash-disagg", - "wandb_name": "qwen3-8b-dflash-disagg-dp4", - }, - "qwen3-8b-dflash-1server-dp7-disaggregated.yaml": { - "report_to": "none", - "wandb_project": "qwen3-8b-dflash-disagg", - "wandb_name": "qwen3-8b-dflash-1srv-dp7", - }, - "qwen3-8b-dflash-online.yaml": { - "report_to": "wandb", - "wandb_project": "specforge-qwen3-8b-dflash", - "wandb_name": "qwen3-8b-dflash-perfectblend", - }, - "qwen3-8b-domino-disaggregated.yaml": { - "report_to": "none", - "wandb_project": "qwen3-8b-domino-disagg", - "wandb_name": "qwen3-8b-domino-disagg-dp4", - }, - "qwen3-8b-domino-1server-dp7-disaggregated.yaml": { - "report_to": "none", - "wandb_project": "qwen3-8b-domino-disagg", - "wandb_name": "qwen3-8b-domino-1srv-dp7", - }, - "qwen3-8b-domino-multiserver-disaggregated.yaml": { - "report_to": "none", - "wandb_project": "qwen3-8b-domino-disagg", - "wandb_name": "qwen3-8b-domino-2srv-dp2", - }, - "qwen3-8b-domino-online.yaml": { - "report_to": "wandb", - "wandb_project": "specforge-qwen3-8b-domino", - "wandb_name": "qwen3-8b-domino_sharegpt", - }, - "qwen3-8b-dpace-online.yaml": { - "report_to": "wandb", - "wandb_project": "dpace-qwen3-8b", - "wandb_name": "qwen3-8b-dpace", - }, - "qwen3-8b-eagle3-disaggregated.yaml": {"report_to": "tensorboard"}, - "qwen3-8b-peagle-disaggregated.yaml": {"report_to": "wandb"}, - "qwen3-coder-30b-a3b-eagle3-online.yaml": { - "report_to": "wandb", - "wandb_project": "specforge-qwen3-coder", - "wandb_name": "qwen3-coder-30b-eagle3-tp4-opc-regen", - }, - "qwen3-coder-480b-a35b-eagle3-online.yaml": { - "report_to": "wandb", - "wandb_project": "specforge-qwen3-480-coder-fp8", - "wandb_name": "qwen3-coder-480b-a35b-eagle3-tp8-ep2-opc-regen", - }, - "qwen3.5-35b-a3b-dflash-online.yaml": {"report_to": "tensorboard"}, - "qwen3.5-35b-a3b-eagle3-online.yaml": {"report_to": "tensorboard"}, - "qwen3.5-4b-dflash-online-npu.yaml": {"report_to": "tensorboard"}, - "qwen3.5-4b-domino-online-npu.yaml": {"report_to": "tensorboard"}, - "qwen3.6-27b-dflash-disaggregated.yaml": { - "report_to": "wandb", - "wandb_project": "qwen36-dflash-disagg", - "wandb_name": "qwen36-27b-dflash-server-capture-dp2", - }, - "qwen3.6-27b-dflash-1server-dp2-disaggregated.yaml": { - "report_to": "none", - "wandb_project": "qwen36-dflash-disagg", - "wandb_name": "qwen36-27b-dflash-1srv-dp2", - }, - "qwen3.6-27b-dflash-multiserver-disaggregated.yaml": { - "report_to": "none", - "wandb_project": "qwen36-dflash-disagg", - "wandb_name": "qwen36-27b-dflash-2srv-dp2", - }, - "qwen3.6-27b-dflash-online.yaml": { - "report_to": "wandb", - "wandb_project": "qwen36-dflash-pr645", - "wandb_name": "qwen36-27b-dflash-nemotron-6ep", - }, -} - -LEGACY_DIST_TIMEOUT_OVERRIDES = { - "deepseek-v3-671b-eagle3-online.yaml": 60, - "gpt-oss-120b-eagle3-online.yaml": 60, - "gpt-oss-20b-eagle3-online.yaml": 60, - "ling-flash-2.0-eagle3-online.yaml": 60, - "qwen3-8b-peagle-disaggregated.yaml": 120, - "qwen3-coder-30b-a3b-eagle3-online.yaml": 60, -} - -LEGACY_EPOCH_OVERRIDES = { - "qwen2.5-0.5b-dflash-online.yaml": 10, - "qwen3.5-35b-a3b-dflash-online.yaml": 10, - "qwen3.5-4b-dflash-online-npu.yaml": 10, - "qwen3.5-4b-domino-online-npu.yaml": 10, - "qwen3.5-4b-vl-dflash-disaggregated.yaml": 10, -} - class UnifiedFeatureReachabilityTest(unittest.TestCase): def test_all_example_configs_validate_through_the_typed_entry(self): paths = sorted( path - for path in EXAMPLE_CONFIG_DIR.glob("*.yaml") + for path in EXAMPLE_CONFIG_DIR.rglob("*.yaml") if not path.name.startswith(".") ) - self.assertEqual(len(paths), 64) + self.assertTrue(paths) - resolved_runs = { - path.name: resolve_run(Config.from_file(str(path))) for path in paths - } - configs = { - filename: resolved.config for filename, resolved in resolved_runs.items() - } - - dpace = configs["qwen3-8b-dpace-online.yaml"] - self.assertEqual(dpace.training.strategy, "dflash") - self.assertEqual(dpace.training.loss_type, "dpace") - self.assertEqual(dpace.tracking.report_to, "wandb") - self.assertEqual(dpace.tracking.wandb_project, "dpace-qwen3-8b") - - for filename, expected in EXPECTED_TRACKING.items(): - with self.subTest(config=filename): - tracking = configs[filename].tracking - for field, value in expected.items(): - self.assertEqual(getattr(tracking, field), value) - - for filename, config in configs.items(): - with self.subTest(config=filename, contract="legacy runtime defaults"): - is_eagle = config.training.strategy in ("eagle3", "peagle") - expected_timeout = LEGACY_DIST_TIMEOUT_OVERRIDES.get( - filename, 20 if is_eagle else 30 - ) - self.assertEqual(config.training.dist_timeout, expected_timeout) - self.assertEqual(config.training.seed, 0 if is_eagle else 42) - - for filename, epochs in LEGACY_EPOCH_OVERRIDES.items(): - with self.subTest(config=filename, contract="legacy epochs"): - self.assertEqual(configs[filename].training.num_epochs, epochs) - - self.assertEqual( - configs[ - "qwen3-8b-eagle3-disaggregated.yaml" - ].model.sglang_mem_fraction_static, - 0.3, - ) - self.assertTrue( - configs["longcat-flash-dflash-online.yaml"].tracking.wandb_offline - ) - self.assertEqual( - configs["qwen3-next-80b-a3b-eagle3-online.yaml"].training.batch_size, - 2, - ) - for filename, config in configs.items(): - if config.mode != "online": - continue - with self.subTest(config=filename, contract="server-only online"): - self.assertEqual(config.deployment.mode, "disaggregated") - self.assertEqual(config.model.target_backend, "sglang") - self.assertIn(config.model.input_modality, {"text", "multimodal"}) + for path in paths: + with self.subTest(config=path.name): + resolve_run(Config.from_file(str(path))) def test_compact_teacher_reaches_the_eagle3_step_provider(self): cfg = Config.model_validate( { **OFFLINE_EAGLE3, "training": { + "trim_loss_positions": True, "compact_teacher": True, "compact_teacher_chunk_size": 2048, }, @@ -217,11 +58,29 @@ def test_compact_teacher_reaches_the_eagle3_step_provider(self): self.assertEqual( resolved.algorithm.providers.step.options(cfg), { + "trim_loss_positions": True, "compact_teacher": True, "compact_teacher_chunk_size": 2048, }, ) + def test_trim_loss_positions_rejects_non_eagle3_strategy(self): + cfg = Config.model_validate( + { + **OFFLINE_EAGLE3, + "training": { + "strategy": "dflash", + "trim_loss_positions": True, + }, + } + ) + + with self.assertRaisesRegex( + ValueError, + "algorithm 'dflash' does not support training.trim_loss_positions", + ): + resolve_run(cfg) + def test_loader_and_profiler_options_reach_the_canonical_trainer(self): eagle = resolve_run(Config.model_validate(OFFLINE_EAGLE3)) dflash = resolve_run( @@ -338,6 +197,11 @@ def test_tracking_config_reaches_the_existing_tracker_adapter(self): self.assertEqual(args.wandb_name, "experiment") self.assertTrue(args.wandb_offline) self.assertEqual(args.wandb_dir, "/tmp/wandb") + self.assertEqual(args.specforge_config["run_id"], "run") + self.assertEqual( + args.specforge_config["training"]["strategy"], + cfg.training.strategy, + ) self.assertEqual(output_dir, "/tmp/output") self.assertIs(create.call_args.kwargs["console_logger"], _logger) diff --git a/tests/test_data/test_kimi_k3_template.py b/tests/test_data/test_kimi_k3_template.py new file mode 100644 index 000000000..50c6c5d4b --- /dev/null +++ b/tests/test_data/test_kimi_k3_template.py @@ -0,0 +1,25 @@ +"""Kimi-K3 template registration and draft-config contract.""" + +import unittest + +from specforge.data.template import TEMPLATE_REGISTRY + + +class TestKimiK3Template(unittest.TestCase): + def test_registered_with_thinking_contract(self): + t = TEMPLATE_REGISTRY.get("kimi-k3-thinking") + self.assertIsNotNone(t) + self.assertEqual(t.parser_type, "thinking") + # Reasoning is stored inline in assistant content (deepspec-style + # regenerations), so the split-field thinking path stays off. + self.assertFalse(t.enable_thinking) + # The assistant header must end inside the think block: the chat + # template emits the opening think tag as part of the generation + # prompt, so it is never model output. + self.assertTrue(t.assistant_header.endswith("<|open|>think<|sep|>")) + self.assertEqual(t.end_of_turn_token, "<|end_of_msg|>") + self.assertIn("<|end_of_msg|>", t.ignore_token) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_data/test_prompt_builder.py b/tests/test_data/test_prompt_builder.py index 5182618bc..16a2274d3 100644 --- a/tests/test_data/test_prompt_builder.py +++ b/tests/test_data/test_prompt_builder.py @@ -6,12 +6,14 @@ import unittest from unittest.mock import patch +from specforge.data.loss_mask import has_consecutive_supervised_tokens from specforge.data.prompt_builder import prepare_prompt_tasks class _FakeDataset: def __init__(self, rows): self.rows = list(rows) + self.getitem_calls = 0 def __iter__(self): return iter(self.rows) @@ -19,6 +21,10 @@ def __iter__(self): def __len__(self): return len(self.rows) + def __getitem__(self, index): + self.getitem_calls += 1 + return self.rows[index] + def select(self, indices): return _FakeDataset(self.rows[index] for index in indices) @@ -86,6 +92,35 @@ def test_pre_tokenized_path_truncates_filters_and_caps(self): ], ) + def test_algorithm_loss_mask_filter_applies_after_truncation(self): + records = [ + {"input_ids": [1, 2, 3, 4, 5], "loss_mask": [1, 0, 0, 1, 1]}, + {"input_ids": [4, 5, 6, 7], "loss_mask": [0, 0, 1, 1]}, + ] + with tempfile.TemporaryDirectory() as tmp_dir: + path = os.path.join(tmp_dir, "prompts.jsonl") + _write_jsonl(path, records) + + prompts = prepare_prompt_tasks( + path, + tokenizer=None, + chat_template=None, + max_length=4, + is_preformatted=False, + train_only_last_turn=False, + cache_dir=None, + cache_key=None, + num_proc=1, + min_loss_tokens=2, + max_prompts=1, + loss_mask_filter=has_consecutive_supervised_tokens, + ) + + self.assertEqual( + prompts, + [{"payload": {"input_ids": [4, 5, 6, 7], "loss_mask": [0, 0, 1, 1]}}], + ) + def test_raw_conversations_use_lazy_dataset_preprocessing(self): raw_rows = [ {"conversations": [{"role": "user", "content": "one"}]}, @@ -141,8 +176,11 @@ def fake_build_eagle3_dataset(**kwargs): max_prompts=1, ) + self.assertEqual(processed_dataset.getitem_calls, 0) + materialized_prompts = list(prompts) + self.assertEqual(processed_dataset.getitem_calls, 1) self.assertEqual( - prompts, + materialized_prompts, [ { "payload": { diff --git a/tests/test_data/test_template_registry.py b/tests/test_data/test_template_registry.py index 7609494ee..6355f5ffc 100644 --- a/tests/test_data/test_template_registry.py +++ b/tests/test_data/test_template_registry.py @@ -21,6 +21,21 @@ def test_deepseek_v2_uses_its_plain_text_tokenizer_headers(self): template.assistant_header, ) + def test_kimi_k3_template_matches_target_xtml_contract(self): + template = TEMPLATE_REGISTRY.get("kimi-k3-thinking") + self.assertEqual( + template.assistant_header, + '<|open|>message role="assistant"<|sep|><|open|>think<|sep|>', + ) + self.assertEqual( + template.user_header, + '<|open|>message role="user"<|sep|>', + ) + self.assertEqual(template.end_of_turn_token, "<|end_of_msg|>") + self.assertEqual(template.parser_type, "thinking") + self.assertFalse(template.enable_thinking) + self.assertEqual(template.ignore_token, ["<|end_of_msg|>"]) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_modeling/test_dflash_eager_attention.py b/tests/test_modeling/test_dflash_eager_attention.py new file mode 100644 index 000000000..ee1e8d60d --- /dev/null +++ b/tests/test_modeling/test_dflash_eager_attention.py @@ -0,0 +1,120 @@ +import unittest + +import torch +from torch import nn +from torch.testing import assert_close +from transformers import Qwen3Config + +from specforge.algorithms.common.dflash_family_model import create_dflash_sdpa_mask +from specforge.modeling.draft.dflash import Qwen3DFlashAttention +from specforge.modeling.draft.dflash_kernels import DFlashKernels + + +def _make_attention(layer_type, implementation, sliding_window): + config = Qwen3Config( + hidden_size=8, + intermediate_size=16, + num_attention_heads=2, + num_key_value_heads=1, + num_hidden_layers=1, + head_dim=4, + max_position_embeddings=64, + vocab_size=32, + layer_types=[layer_type], + sliding_window=sliding_window, + use_sliding_window=sliding_window is not None, + attention_bias=False, + attention_dropout=0.0, + ) + config._attn_implementation = implementation + kernels = DFlashKernels( + make_rms_norm=lambda *_: nn.Identity(), + make_mlp=lambda *_: nn.Identity(), + ) + return Qwen3DFlashAttention(config, layer_idx=0, kernels=kernels).eval() + + +def _forward(attention, hidden_states, target_hidden, attention_mask): + total_length = target_hidden.shape[1] + hidden_states.shape[1] + position_embeddings = ( + hidden_states.new_ones(1, total_length, attention.head_dim), + hidden_states.new_zeros(1, total_length, attention.head_dim), + ) + return attention( + hidden_states=hidden_states, + target_hidden=target_hidden, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + ) + + +class TestDFlashEagerAttentionMasking(unittest.TestCase): + def test_eager_matches_sdpa_for_full_and_sliding_masks(self): + for layer_type, sliding_window in ( + ("full_attention", None), + ("sliding_attention", 2), + ): + with self.subTest(layer_type=layer_type): + torch.manual_seed(17) + eager = _make_attention(layer_type, "eager", sliding_window) + sdpa = _make_attention(layer_type, "sdpa", sliding_window) + sdpa.load_state_dict(eager.state_dict()) + + mask = create_dflash_sdpa_mask( + anchor_positions=torch.tensor([[2, 4]]), + block_keep_mask=torch.tensor([[True, False]]), + S=4, + block_size=2, + device=torch.device("cpu"), + sliding_window=sliding_window, + ) + eager_hidden = torch.randn(1, 4, 8, requires_grad=True) + eager_target = torch.randn(1, 4, 8, requires_grad=True) + sdpa_hidden = eager_hidden.detach().clone().requires_grad_(True) + sdpa_target = eager_target.detach().clone().requires_grad_(True) + + eager_output, eager_weights = _forward( + eager, + eager_hidden, + eager_target, + mask, + ) + sdpa_output, _ = _forward( + sdpa, + sdpa_hidden, + sdpa_target, + mask, + ) + assert_close(eager_output, sdpa_output, rtol=1e-5, atol=1e-6) + + allowed = mask.expand_as(eager_weights) + forbidden_weights = eager_weights.masked_select(~allowed) + assert_close( + forbidden_weights, + torch.zeros_like(forbidden_weights), + rtol=0, + atol=0, + ) + invalid_rows = ~mask.any(dim=-1).squeeze(1) + assert_close( + eager_output[invalid_rows], + torch.zeros_like(eager_output[invalid_rows]), + rtol=0, + atol=0, + ) + + output_grad = torch.randn_like(eager_output) + eager_grads = torch.autograd.grad( + (eager_output * output_grad).sum(), + (eager_hidden, eager_target), + ) + sdpa_grads = torch.autograd.grad( + (sdpa_output * output_grad).sum(), + (sdpa_hidden, sdpa_target), + ) + for eager_grad, sdpa_grad in zip(eager_grads, sdpa_grads): + assert_close(eager_grad, sdpa_grad, rtol=1e-5, atol=1e-6) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_modeling/test_dflash_sliding.py b/tests/test_modeling/test_dflash_sliding.py new file mode 100644 index 000000000..34cbc1487 --- /dev/null +++ b/tests/test_modeling/test_dflash_sliding.py @@ -0,0 +1,234 @@ +import unittest +from pathlib import Path +from unittest import mock + +import torch +from torch import nn +from transformers import Qwen3Config + +from specforge.algorithms.common.dflash_family_model import OnlineDFlashModel +from specforge.modeling.draft.dflash import ( + DFlashDraftModel, + resolve_dflash_attention_layout, +) + + +def _draft_config(layer_types, sliding_window=None): + config = Qwen3Config( + architectures=["DFlashDraftModel"], + block_size=2, + hidden_size=8, + intermediate_size=16, + num_attention_heads=2, + num_key_value_heads=1, + num_hidden_layers=len(layer_types), + num_target_layers=6, + head_dim=4, + max_position_embeddings=64, + vocab_size=32, + layer_types=list(layer_types), + sliding_window=sliding_window, + use_sliding_window=sliding_window is not None, + ) + config._attn_implementation = "sdpa" + return config + + +class _CaptureLayer(nn.Module): + def __init__(self): + super().__init__() + self.attention_mask = None + self.kernel_options = None + + def forward(self, *, hidden_states, attention_mask, kernel_options=None, **_): + self.attention_mask = attention_mask + self.kernel_options = kernel_options + return hidden_states + + +class _RotaryStub(nn.Module): + def forward(self, *_): + return (torch.empty(0), torch.empty(0)) + + +def _capture_model(layer_types, sliding_window=None): + model = DFlashDraftModel(_draft_config(layer_types, sliding_window)) + capture_layers = [_CaptureLayer() for _ in layer_types] + model.layers = nn.ModuleList(capture_layers) + model.fc = nn.Identity() + model.hidden_norm = nn.Identity() + model.norm = nn.Identity() + model.rotary_emb = _RotaryStub() + return model, capture_layers + + +def _forward(model, attention_mask): + noise_embedding = torch.randn(1, 2, model.config.hidden_size) + target_hidden = torch.randn(1, 4, model.config.hidden_size) + position_ids = torch.arange(6).unsqueeze(0) + return model( + position_ids=position_ids, + noise_embedding=noise_embedding, + target_hidden=target_hidden, + attention_mask=attention_mask, + ) + + +class TestDFlashSlidingDispatch(unittest.TestCase): + def test_full_only_model_keeps_single_mask_compatibility(self): + model, layers = _capture_model(["full_attention", "full_attention"]) + full_mask = torch.tensor([1]) + + _forward(model, full_mask) + + self.assertIs(layers[0].attention_mask, full_mask) + self.assertIs(layers[1].attention_mask, full_mask) + + def test_online_wrapper_builds_both_masks_for_hybrid_model(self): + model, layers = _capture_model( + ["sliding_attention", "full_attention"], + sliding_window=4, + ) + wrapper = OnlineDFlashModel( + draft_model=model, + target_lm_head=nn.Identity(), + target_embed_tokens=nn.Embedding(32, model.config.hidden_size), + mask_token_id=31, + block_size=2, + attention_backend="sdpa", + num_anchors=1, + ) + anchors = torch.tensor([[2]]) + keep = torch.tensor([[True]]) + noise_embedding = torch.randn(1, 2, model.config.hidden_size) + full_mask = torch.tensor([1]) + sliding_mask = torch.tensor([2]) + + with ( + mock.patch.object( + wrapper, + "_sample_anchor_positions", + return_value=(anchors, keep), + ), + mock.patch.object( + wrapper, + "_create_noise_embed", + return_value=noise_embedding, + ), + mock.patch( + "specforge.algorithms.common.dflash_family_model." + "create_dflash_sdpa_mask", + side_effect=(full_mask, sliding_mask), + ) as create_mask, + ): + wrapper._forward_draft_blocks( + input_ids=torch.ones(1, 4, dtype=torch.long), + hidden_states=torch.randn(1, 4, model.config.hidden_size), + loss_mask=torch.ones(1, 4), + ) + + self.assertEqual(create_mask.call_count, 2) + self.assertIs(layers[0].attention_mask, sliding_mask) + self.assertIs(layers[1].attention_mask, full_mask) + self.assertTrue(all(layer.kernel_options is None for layer in layers)) + + def test_online_wrapper_forces_standard_triton_flex_backend(self): + model, layers = _capture_model(["full_attention"]) + wrapper = OnlineDFlashModel( + draft_model=model, + target_lm_head=nn.Identity(), + target_embed_tokens=nn.Embedding(32, model.config.hidden_size), + mask_token_id=31, + block_size=2, + attention_backend="flex_attention", + num_anchors=1, + ) + anchors = torch.tensor([[2]]) + keep = torch.tensor([[True]]) + + with ( + mock.patch.object( + wrapper, + "_sample_anchor_positions", + return_value=(anchors, keep), + ), + mock.patch.object( + wrapper, + "_create_noise_embed", + return_value=torch.randn(1, 2, model.config.hidden_size), + ), + mock.patch( + "specforge.algorithms.common.dflash_family_model." + "create_dflash_block_mask", + return_value=torch.tensor([1]), + ), + ): + wrapper._forward_draft_blocks( + input_ids=torch.ones(1, 4, dtype=torch.long), + hidden_states=torch.randn(1, 4, model.config.hidden_size), + loss_mask=torch.ones(1, 4), + ) + + self.assertEqual(layers[0].kernel_options, {"BACKEND": "TRITON"}) + + +class TestDFlashSlidingConfig(unittest.TestCase): + def test_checked_in_qwen36_config_preserves_hybrid_layout(self): + config_path = ( + Path(__file__).resolve().parents[2] / "configs" / "qwen3.6-27b-dflash.json" + ) + config = Qwen3Config.from_json_file(str(config_path)) + + layer_types, sliding_window = resolve_dflash_attention_layout(config) + + self.assertEqual( + list(layer_types), + [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ], + ) + self.assertEqual(sliding_window, 2048) + + def test_configures_attention_modules_from_layer_types(self): + model = DFlashDraftModel( + _draft_config( + ["sliding_attention", "full_attention", "sliding_attention"], + sliding_window=7, + ) + ) + + self.assertEqual( + list(model.layer_types), + ["sliding_attention", "full_attention", "sliding_attention"], + ) + self.assertEqual(model.sliding_window, 7) + self.assertEqual(model.layers[0].self_attn.sliding_window, 7) + self.assertIsNone(model.layers[1].self_attn.sliding_window) + self.assertEqual(model.layers[2].self_attn.sliding_window, 7) + + def test_rejects_invalid_attention_layouts(self): + cases = ( + (["full_attention"], None), + (["full_attention", "unknown"], None), + (["sliding_attention", "full_attention"], None), + (["sliding_attention", "full_attention"], 0), + (["sliding_attention", "full_attention"], -1), + ) + for layer_types, sliding_window in cases: + with self.subTest( + layer_types=layer_types, + sliding_window=sliding_window, + ): + config = _draft_config(["full_attention", "full_attention"]) + config.layer_types = layer_types + config.sliding_window = sliding_window + with self.assertRaises(ValueError): + resolve_dflash_attention_layout(config) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_modeling/test_domino_model.py b/tests/test_modeling/test_domino_model.py index c80d917e5..4426e9e5f 100644 --- a/tests/test_modeling/test_domino_model.py +++ b/tests/test_modeling/test_domino_model.py @@ -78,8 +78,14 @@ def test_training_loss_updates_gru_and_projection(self): ) fixed_hidden = torch.randn(1, block_size, hidden_size) - def fixed_draft_blocks(self, input_ids, hidden_states, loss_mask): - del hidden_states, loss_mask + def fixed_draft_blocks( + self, + input_ids, + hidden_states, + loss_mask, + max_valid_anchors=None, + ): + del hidden_states, loss_mask, max_valid_anchors return ( torch.zeros(1, 1, dtype=torch.long, device=input_ids.device), torch.ones(1, 1, dtype=torch.bool, device=input_ids.device), @@ -87,12 +93,13 @@ def fixed_draft_blocks(self, input_ids, hidden_states, loss_mask): ) model._forward_draft_blocks = MethodType(fixed_draft_blocks, model) - loss, _accuracy, _metrics = model( + loss, _accuracy, metrics = model( input_ids=torch.tensor([[1, 2, 3, 4]]), hidden_states=torch.zeros(1, block_size, hidden_size), loss_mask=torch.ones(1, block_size), lambda_base=0.0, ) + self.assertIsInstance(metrics["lambda_base"], float) loss.backward() for module in (draft.prefix_gru, draft.embed_proj): @@ -147,8 +154,20 @@ def test_chunked_objective_matches_full_loss_metrics_and_gradients(self): chunked_hidden = full_hidden.detach().clone().requires_grad_() def fixed_blocks(output_hidden): - def _forward(self, input_ids, hidden_states, loss_mask): - del self, input_ids, hidden_states, loss_mask + def _forward( + self, + input_ids, + hidden_states, + loss_mask, + max_valid_anchors=None, + ): + del ( + self, + input_ids, + hidden_states, + loss_mask, + max_valid_anchors, + ) return anchors, keep_mask, output_hidden return _forward @@ -220,6 +239,86 @@ def _forward(self, input_ids, hidden_states, loss_mask): atol=1e-7, ) + @unittest.skipUnless(torch.cuda.is_available(), "Triton loss requires CUDA") + def test_fused_loss_matches_standard_model_path(self): + from specforge.algorithms.common.dflash_family_model import OnlineDominoModel + + torch.manual_seed(19) + hidden_size, vocab_size, block_size = 4, 7, 4 + draft = _bare_domino( + hidden_size=hidden_size, + gru_hidden_size=3, + embedding_size=2, + vocab_size=vocab_size, + block_size=block_size, + ) + model = OnlineDominoModel( + draft_model=draft, + target_lm_head=nn.Linear(hidden_size, vocab_size, bias=False), + target_embed_tokens=nn.Embedding(vocab_size, hidden_size), + mask_token_id=0, + block_size=block_size, + attention_backend="sdpa", + num_anchors=2, + objective_chunk_blocks=1, + shift_label=False, + ).cuda() + fixed_hidden = torch.randn( + 1, 2 * block_size, hidden_size, device="cuda", requires_grad=True + ) + + def fixed_draft_blocks( + self, + input_ids, + hidden_states, + loss_mask, + max_valid_anchors=None, + ): + del hidden_states, loss_mask, max_valid_anchors + return ( + torch.tensor([[0, 4]], device=input_ids.device), + torch.ones(1, 2, dtype=torch.bool, device=input_ids.device), + fixed_hidden, + ) + + model._forward_draft_blocks = MethodType(fixed_draft_blocks, model) + inputs = { + "input_ids": torch.tensor([[1, 2, 3, 4, 5, 6, 0, 1]], device="cuda"), + "hidden_states": torch.zeros(1, 2 * block_size, hidden_size, device="cuda"), + "loss_mask": torch.ones(1, 2 * block_size, device="cuda"), + "lambda_base": 0.25, + } + model._use_fused_domino_ce = False + standard = model(**inputs) + standard[0].backward() + standard_hidden_grad = fixed_hidden.grad.clone() + standard_parameter_grads = { + name: None if parameter.grad is None else parameter.grad.clone() + for name, parameter in model.named_parameters() + } + fixed_hidden.grad = None + model.zero_grad(set_to_none=True) + + model._use_fused_domino_ce = True + fused = model(**inputs) + fused[0].backward() + + torch.testing.assert_close(fused[0], standard[0]) + torch.testing.assert_close(fused[1], standard[1]) + for key in standard[2]: + torch.testing.assert_close(fused[2][key], standard[2][key]) + + torch.testing.assert_close( + fixed_hidden.grad, standard_hidden_grad, rtol=1e-4, atol=1e-4 + ) + for name, parameter in model.named_parameters(): + standard_grad = standard_parameter_grads[name] + self.assertEqual(parameter.grad is None, standard_grad is None, name) + if standard_grad is not None: + torch.testing.assert_close( + parameter.grad, standard_grad, rtol=1e-4, atol=1e-4 + ) + def test_npu_bf16_gru_gradients_reach_registered_weights(self): torch.manual_seed(7) model = _bare_domino().to(dtype=torch.bfloat16) diff --git a/tests/test_modeling/test_flex_attention_backend.py b/tests/test_modeling/test_flex_attention_backend.py new file mode 100644 index 000000000..e2cdff95b --- /dev/null +++ b/tests/test_modeling/test_flex_attention_backend.py @@ -0,0 +1,316 @@ +import os +import unittest +from unittest import mock + +import torch +from torch.nn.attention.flex_attention import flex_attention +from transformers import Qwen3Config + +from specforge.algorithms.common.dflash_family_model import ( + create_dflash_block_mask, + create_dflash_sdpa_mask, +) +from specforge.modeling.draft.dflash import DFlashDraftModel, Qwen3DFlashAttention +from specforge.modeling.draft.dflash_kernels import DEFAULT_DFLASH_KERNELS +from specforge.modeling.draft.flex_attention_backend import flex_attention_backend + + +class FlexAttentionBackendTest(unittest.TestCase): + @unittest.skipUnless(torch.cuda.is_available(), "FlexAttention requires CUDA") + def test_sliding_block_mask_matches_sdpa(self): + torch.manual_seed(0) + device = torch.device("cuda") + dtype = torch.bfloat16 + context_len, draft_block_size = 16, 4 + anchors = torch.tensor([[8, 12]], device=device) + keep_blocks = torch.ones(1, 2, dtype=torch.bool, device=device) + query_len = anchors.shape[1] * draft_block_size + kv_len = context_len + query_len + query = torch.randn(1, 2, query_len, 64, device=device, dtype=dtype) + key = torch.randn(1, 2, kv_len, 64, device=device, dtype=dtype) + value = torch.randn(1, 2, kv_len, 64, device=device, dtype=dtype) + + block_mask = create_dflash_block_mask( + anchor_positions=anchors, + block_keep_mask=keep_blocks, + S=context_len, + block_size=draft_block_size, + device=device, + sliding_window=8, + ) + dense_mask = create_dflash_sdpa_mask( + anchor_positions=anchors, + block_keep_mask=keep_blocks, + S=context_len, + block_size=draft_block_size, + device=device, + sliding_window=8, + ) + compiled_attention = torch.compile( + lambda q, k, v, mask: flex_attention( + q, + k, + v, + block_mask=mask, + ), + fullgraph=True, + ) + + flex_output = compiled_attention(query, key, value, block_mask) + sdpa_output = torch.nn.functional.scaled_dot_product_attention( + query, + key, + value, + attn_mask=dense_mask, + ) + torch.testing.assert_close( + flex_output, + sdpa_output, + atol=3e-3, + rtol=2e-2, + ) + + def test_mixed_layer_types_select_their_own_masks(self): + config = Qwen3Config( + hidden_size=16, + intermediate_size=32, + num_attention_heads=2, + num_key_value_heads=1, + num_hidden_layers=2, + num_target_layers=4, + head_dim=8, + layer_types=["sliding_attention", "full_attention"], + use_sliding_window=True, + sliding_window=8, + attention_dropout=0.0, + block_size=2, + dflash_config={"target_layer_ids": [1, 2]}, + ) + config._attn_implementation = "eager" + model = DFlashDraftModel(config) + sliding_mask = torch.ones(1, 1, 2, 5, dtype=torch.bool) + full_mask = torch.ones(1, 1, 2, 5, dtype=torch.bool) + + for layer in model.layers: + layer.self_attn.forward = mock.Mock( + side_effect=lambda hidden_states, **kwargs: (hidden_states, None) + ) + + model( + position_ids=torch.arange(5).unsqueeze(0), + noise_embedding=torch.randn(1, 2, config.hidden_size), + target_hidden=torch.randn(1, 3, 2 * config.hidden_size), + attention_mask={ + "sliding_attention": sliding_mask, + "full_attention": full_mask, + }, + ) + + self.assertIs( + model.layers[0].self_attn.forward.call_args.kwargs["attention_mask"], + sliding_mask, + ) + self.assertIs( + model.layers[1].self_attn.forward.call_args.kwargs["attention_mask"], + full_mask, + ) + + def test_eager_converts_boolean_mask_to_additive_mask(self): + config = Qwen3Config( + hidden_size=8, + intermediate_size=16, + num_attention_heads=1, + num_key_value_heads=1, + num_hidden_layers=1, + head_dim=8, + layer_types=["sliding_attention"], + sliding_window=8, + attention_dropout=0.0, + ) + config._attn_implementation = "eager" + attention = Qwen3DFlashAttention( + config, + layer_idx=0, + kernels=DEFAULT_DFLASH_KERNELS, + ) + boolean_mask = torch.tensor([[[[True, False]]]]) + cos = torch.ones(1, 2, config.head_dim) + sin = torch.zeros_like(cos) + + with mock.patch( + "specforge.modeling.draft.dflash.eager_attention_forward", + return_value=(torch.zeros(1, 1, 1, config.head_dim), None), + ) as eager: + attention( + hidden_states=torch.randn(1, 1, config.hidden_size), + target_hidden=torch.randn(1, 1, config.hidden_size), + position_embeddings=(cos, sin), + attention_mask=boolean_mask, + ) + + additive_mask = eager.call_args.args[4] + self.assertEqual(additive_mask[0, 0, 0, 0].item(), 0.0) + self.assertEqual( + additive_mask[0, 0, 0, 1].item(), + torch.finfo(additive_mask.dtype).min, + ) + + # This correctness regression test can be deleted when we require + # torch>=2.13; it tests the Torch 2.11 Inductor monkeypatch for CuteDSL + # operations in patch_inductor_cutedsl_lowerings(). + @unittest.skipUnless( + torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10, + "FLASH FlexAttention correctness requires a Blackwell CUDA device", + ) + def test_flash_matches_triton_forward_and_backward(self): + torch.manual_seed(0) + device = torch.device("cuda") + dtype = torch.bfloat16 + batch_size, num_query_heads, num_key_value_heads = 1, 3, 1 + context_len, head_dim = 256, 64 + num_blocks, draft_block_size = 4, 64 + query_len = num_blocks * draft_block_size + kv_len = context_len + query_len + anchors = torch.tensor([[64, 128, 192, 224]], device=device) + keep_blocks = torch.ones( + (batch_size, num_blocks), dtype=torch.bool, device=device + ) + + inputs = ( + torch.randn( + batch_size, + num_query_heads, + query_len, + head_dim, + device=device, + dtype=dtype, + ), + torch.randn( + batch_size, + num_key_value_heads, + kv_len, + head_dim, + device=device, + dtype=dtype, + ), + torch.randn( + batch_size, + num_key_value_heads, + kv_len, + head_dim, + device=device, + dtype=dtype, + ), + ) + + def run_backend(backend, flex_block_size=None): + block_mask = create_dflash_block_mask( + anchor_positions=anchors, + block_keep_mask=keep_blocks, + S=context_len, + block_size=draft_block_size, + device=device, + flex_block_size=flex_block_size, + sliding_window=128, + ) + + compiled_attention = torch.compile( + lambda query, key, value, mask: flex_attention( + query, + key, + value, + block_mask=mask, + enable_gqa=True, + kernel_options={"BACKEND": backend}, + ), + fullgraph=True, + ) + + query, key, value = [ + tensor.detach().clone().requires_grad_(True) for tensor in inputs + ] + output = compiled_attention(query, key, value, block_mask) + output.float().square().mean().backward() + torch.cuda.synchronize() + return output.detach(), tuple( + tensor.grad.detach() for tensor in (query, key, value) + ) + + triton_output, triton_grads = run_backend("TRITON") + with mock.patch.dict(os.environ, {"SPECFORGE_FLEX_ATTENTION_BACKEND": "FLASH"}): + self.assertEqual(flex_attention_backend(), "FLASH") + flash_output, flash_grads = run_backend("FLASH", (256, 128)) + + self.assertTrue(torch.isfinite(flash_output).all()) + torch.testing.assert_close(flash_output, triton_output, atol=3e-3, rtol=2e-2) + for flash_grad, triton_grad in zip(flash_grads, triton_grads): + self.assertTrue(torch.isfinite(flash_grad).all()) + torch.testing.assert_close(flash_grad, triton_grad, atol=5e-6, rtol=2e-2) + + @unittest.skipUnless( + torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10, + "FLASH FlexAttention correctness requires a Blackwell CUDA device", + ) + def test_dflash_flash_attention_forward_backward_smoke(self): + config = Qwen3Config( + hidden_size=256, + intermediate_size=512, + num_attention_heads=4, + num_key_value_heads=2, + num_hidden_layers=1, + head_dim=64, + layer_types=["full_attention"], + attention_dropout=0.0, + ) + config._attn_implementation = "flex_attention" + attention = Qwen3DFlashAttention( + config, + layer_idx=0, + kernels=DEFAULT_DFLASH_KERNELS, + ).to(device="cuda", dtype=torch.bfloat16) + hidden_states = torch.randn( + 1, + 256, + config.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + target_hidden = torch.randn( + 1, + 256, + config.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + block_mask = create_dflash_block_mask( + anchor_positions=torch.tensor([[64, 128, 192, 224]], device="cuda"), + block_keep_mask=torch.ones(1, 4, dtype=torch.bool, device="cuda"), + S=256, + block_size=64, + device=torch.device("cuda"), + flex_block_size=(256, 128), + ) + cos = torch.ones(1, 512, config.head_dim, device="cuda", dtype=torch.bfloat16) + sin = torch.zeros_like(cos) + + with mock.patch.dict(os.environ, {"SPECFORGE_FLEX_ATTENTION_BACKEND": "FLASH"}): + output, weights = attention( + hidden_states=hidden_states, + target_hidden=target_hidden, + position_embeddings=(cos, sin), + attention_mask=block_mask, + ) + output.float().square().mean().backward() + torch.cuda.synchronize() + + self.assertIsNone(weights) + self.assertIsNotNone(hidden_states.grad) + self.assertIsNotNone(target_hidden.grad) + self.assertTrue(torch.isfinite(hidden_states.grad).all()) + self.assertTrue(torch.isfinite(target_hidden.grad).all()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_optimizer/test_lr_scheduler.py b/tests/test_optimizer/test_lr_scheduler.py new file mode 100644 index 000000000..7f30dd146 --- /dev/null +++ b/tests/test_optimizer/test_lr_scheduler.py @@ -0,0 +1,59 @@ +import unittest + +import torch + +from specforge.optimizer import BF16Optimizer + + +def _optimizer(*, scheduler="cosine", total_steps=4, warmup_ratio=0.0): + model = torch.nn.Linear(2, 2, bias=False) + optimizer = BF16Optimizer( + model, + lr=1e-3, + max_grad_norm=1.0, + total_steps=total_steps, + warmup_ratio=warmup_ratio, + lr_scheduler=scheduler, + ) + return model, optimizer + + +class TestLearningRateScheduler(unittest.TestCase): + def test_constant_scheduler_keeps_base_lr_without_warmup(self): + model, optimizer = _optimizer(scheduler="constant") + observed = [optimizer.get_learning_rate()] + for _ in range(4): + model.weight.grad = torch.ones_like(model.weight) + optimizer.step() + observed.append(optimizer.get_learning_rate()) + self.assertEqual(observed, [1e-3] * 5) + + def test_constant_scheduler_supports_linear_warmup(self): + _model, optimizer = _optimizer( + scheduler="constant", total_steps=4, warmup_ratio=0.5 + ) + self.assertAlmostEqual(optimizer.get_learning_rate(), 5e-4) + optimizer.scheduler.step() + self.assertAlmostEqual(optimizer.get_learning_rate(), 1e-3) + optimizer.scheduler.step() + self.assertAlmostEqual(optimizer.get_learning_rate(), 1e-3) + + def test_unknown_scheduler_is_rejected(self): + with self.assertRaisesRegex(ValueError, "unsupported lr_scheduler"): + _optimizer(scheduler="linear") + + def test_resume_rejects_scheduler_change(self): + _model, cosine = _optimizer(scheduler="cosine") + _other_model, constant = _optimizer(scheduler="constant") + with self.assertRaisesRegex(ValueError, "checkpoint optimizer used"): + constant.load_state_dict(cosine.state_dict()) + + def test_legacy_checkpoint_defaults_to_cosine(self): + _model, cosine = _optimizer(scheduler="cosine") + state = cosine.state_dict() + state.pop("lr_scheduler_type") + self.assertEqual(state.get("lr_scheduler_type", "cosine"), "cosine") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_runtime/_fixtures.py b/tests/test_runtime/_fixtures.py index ad850771f..6024685a8 100644 --- a/tests/test_runtime/_fixtures.py +++ b/tests/test_runtime/_fixtures.py @@ -308,6 +308,7 @@ def build_dflash( draft_config = AutoConfig.from_pretrained(target_dir) draft_config.num_hidden_layers = draft_layers + draft_config.layer_types = ["full_attention"] * draft_layers draft_config.block_size = block_size draft_config.num_target_layers = target_layers draft_config.dflash_config = {"mask_token_id": mask_token_id} @@ -375,6 +376,7 @@ def build_domino( draft_config = AutoConfig.from_pretrained(target_dir) draft_config.num_hidden_layers = draft_layers + draft_config.layer_types = ["full_attention"] * draft_layers draft_config.block_size = block_size draft_config.num_target_layers = target_layers draft_config.dflash_config = { @@ -445,6 +447,7 @@ def build_dspark( draft_config = AutoConfig.from_pretrained(target_dir) draft_config.num_hidden_layers = draft_layers + draft_config.layer_types = ["full_attention"] * draft_layers draft_config.block_size = block_size draft_config.num_target_layers = target_layers draft_config.dflash_config = { diff --git a/tests/test_runtime/test_compact_teacher_strategy.py b/tests/test_runtime/test_compact_teacher_strategy.py index fa07a2f85..acd6f8f0c 100644 --- a/tests/test_runtime/test_compact_teacher_strategy.py +++ b/tests/test_runtime/test_compact_teacher_strategy.py @@ -315,6 +315,7 @@ def test_default_path_remains_full_vocab_projection(self): self.assertEqual(head.forward_calls, 1) self.assertEqual(model.kwargs["target"].shape, (1, 3, 8)) self.assertNotIn("target_hidden_for_compact", model.kwargs) + self.assertFalse(model.kwargs["trim_loss_positions"]) def test_compact_path_rejects_online_target_repr(self): strategy = Eagle3TrainStrategy( @@ -323,13 +324,18 @@ def test_compact_path_rejects_online_target_repr(self): with self.assertRaisesRegex(ValueError, "offline-only"): strategy.forward_loss(_batch(target_repr="logits")) - def test_step_provider_forwards_compact_strategy_kwargs(self): + def test_step_provider_forwards_eagle3_strategy_kwargs(self): + model = _Eagle3() strategy = EAGLE3.providers.step.build( - _Eagle3(), + model, target_head=_TargetHead(), + trim_loss_positions=True, compact_teacher=True, compact_teacher_chunk_size=4, ) + strategy.forward_loss(_batch()) + self.assertTrue(strategy.trim_loss_positions) + self.assertTrue(model.kwargs["trim_loss_positions"]) self.assertTrue(strategy.compact_teacher) self.assertEqual(strategy.compact_teacher_chunk_size, 4) diff --git a/tests/test_runtime/test_disagg_multiserver.py b/tests/test_runtime/test_disagg_multiserver.py index da7e96d0e..5110d6118 100644 --- a/tests/test_runtime/test_disagg_multiserver.py +++ b/tests/test_runtime/test_disagg_multiserver.py @@ -338,6 +338,53 @@ def run_producer(): self.assertGreaterEqual(snapshot["pause_transitions"], 1) self.assertGreaterEqual(snapshot["resume_transitions"], 1) + def test_byte_watermark_cannot_block_the_first_optimizer_window(self): + backend = _FakeMooncakeStore() + stub = _StubCaptureServer(backend) + store = MooncakeFeatureStore(store=backend, store_id="run0") + channel = StreamingRefChannel(os.path.join(self._workdir(), "refs.jsonl")) + _workers, drive = _build( + [_adapter(store, stub)], + _prompts(3), + store, + channel, + consumer_quantum=3, + lease=1, + resident_high_watermark_bytes=1, + resident_low_watermark_bytes=0, + ) + + outcome = {} + + def run_producer(): + try: + outcome["produced"] = drive() + except BaseException as exc: # expose a thread failure to the test + outcome["error"] = exc + + thread = threading.Thread(target=run_producer, daemon=True) + thread.start() + deadline = time.monotonic() + 2 + while channel.published < 3 and time.monotonic() < deadline: + time.sleep(0.001) + published_before_ack = channel.published + + # Keep cleanup bounded if this invariant regresses and the producer + # pauses before publishing a complete window. + reader = StreamingRefChannel(channel.path) + cleanup_deadline = time.monotonic() + 2 + while thread.is_alive() and time.monotonic() < cleanup_deadline: + refs = reader.poll() + if refs: + reader.mark_consumed(len(refs)) + time.sleep(0.001) + thread.join(2) + + self.assertEqual(published_before_ack, 3) + self.assertFalse(thread.is_alive()) + self.assertNotIn("error", outcome) + self.assertEqual(outcome.get("produced"), 3) + def test_hard_byte_cap_aborts_unpublished_capture_and_fails_channel(self): backend = _FakeMooncakeStore() stub = _StubCaptureServer(backend) @@ -549,6 +596,35 @@ def test_prompt_epochs_republish_with_unique_sample_ids(self): ) self.assertTrue(channel.is_closed()) + def test_prompt_ingest_chunks_preserve_epoch_ids_and_release_payloads(self): + backend = _FakeMooncakeStore() + stub = _StubCaptureServer(backend) + store = MooncakeFeatureStore(store=backend, store_id="run0") + N, E = 5, 2 + channel = StreamingRefChannel(os.path.join(self._workdir(), "refs.jsonl")) + workers, drive = _build( + [_adapter(store, stub)], + _prompts(N), + store, + channel, + lease=2, + prompt_epochs=E, + prompt_ingest_batch_size=2, + ) + + produced = drive() + + self.assertEqual(produced, N * E) + self.assertEqual(workers[0].controller.status()["prompts"], 0) + self.assertEqual( + set(_published_sample_ids(channel.path)), + { + f"run0:epoch{epoch:04d}-prompt{idx:012d}" + for epoch in range(E) + for idx in range(N) + }, + ) + def test_prompt_epoch_order_is_seeded_and_reconstruction_stable(self): prompts = _prompts(12) diff --git a/tests/test_runtime/test_disaggregated_model_loading.py b/tests/test_runtime/test_disaggregated_model_loading.py index 8b98c893f..2593fcb02 100644 --- a/tests/test_runtime/test_disaggregated_model_loading.py +++ b/tests/test_runtime/test_disaggregated_model_loading.py @@ -135,6 +135,37 @@ def test_domino_uses_the_dflash_server_capture_method(self): self.assertEqual(contract.method, "dflash") self.assertEqual(contract.aux_layer_ids, (3, 7)) + def test_dspark_uses_its_dedicated_server_capture_method(self): + resolved = resolve_run( + _config(strategy="dspark", draft_model_config="dspark-draft") + ) + draft_payload = { + "architectures": ["DSparkDraftModel"], + "vocab_size": 128, + "num_target_layers": 2, + "dflash_config": { + "projector_type": "dspark", + "target_layer_ids": [3, 7], + }, + } + with ( + mock.patch( + "transformers.AutoConfig.from_pretrained", + return_value=SimpleNamespace(hidden_size=64, vocab_size=128), + ), + mock.patch( + "specforge.training.model_loading.draft_config_dict", + return_value=draft_payload, + ), + ): + contract = resolve_server_capture_contract( + resolved.config, + algorithm=resolved.algorithm, + ) + + self.assertEqual(contract.method, "dspark") + self.assertEqual(contract.aux_layer_ids, (3, 7)) + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/tests/test_runtime/test_domino_launch.py b/tests/test_runtime/test_domino_launch.py index 7c141d9f3..52ff16650 100644 --- a/tests/test_runtime/test_domino_launch.py +++ b/tests/test_runtime/test_domino_launch.py @@ -63,7 +63,7 @@ def __init__(self): } def forward(self, **kwargs): - del kwargs + self.last_kwargs = kwargs return self.weight.square(), torch.tensor(0.75), self.model_metrics model = DiagnosticModel() @@ -90,6 +90,7 @@ def forward(self, **kwargs): for name, value in model.model_metrics.items(): self.assertIs(output.metrics[name], value) self.assertEqual(float(output.metrics["accuracy"]), 0.75) + self.assertEqual(model.last_kwargs["max_valid_anchors"], 1) @unittest.skipUnless(CUDA, "Domino offline launcher path requires CUDA") diff --git a/tests/test_runtime/test_equiv_trim_loss_positions.py b/tests/test_runtime/test_equiv_trim_loss_positions.py new file mode 100644 index 000000000..e45b434e8 --- /dev/null +++ b/tests/test_runtime/test_equiv_trim_loss_positions.py @@ -0,0 +1,89 @@ +# coding=utf-8 +"""Equivalence: trim_loss_positions must not change the training loss. + +A-level position trimming computes the teacher target_p, the draft logits and the +loss only at supervised (loss-masked) positions instead of over the full sequence. +It is mathematically equivalent to the full-length path (the mean denominator is +rescaled from n_sup back to the full length). This test runs the identical forward +with trimming off and on and asserts the per-step losses match within bf16 +tolerance. + +GPU-only, matching the other EAGLE3 equivalence tests in this directory. +""" + +import os +import shutil +import tempfile +import unittest + +import torch + +CUDA = torch.cuda.is_available() + + +@unittest.skipUnless(CUDA, "trim_loss_positions equivalence requires CUDA") +class TestEquivTrimLossPositions(unittest.TestCase): + def test_trim_loss_positions_matches_full(self): + torch.manual_seed(0) + from tests.test_runtime import _fixtures as fx + + fx.build_single_rank_distributed(port="29567") + + workdir = tempfile.mkdtemp(prefix="equiv_trim_") + self.addCleanup(shutil.rmtree, workdir, ignore_errors=True) + + model, target_head = fx.build_eagle3(workdir, ttt=3) + model.eval() + + # One offline sample gives us (input_ids, target, loss_mask, hidden_state). + feature_dir = os.path.join(workdir, "features") + os.makedirs(feature_dir, exist_ok=True) + fx.write_offline_files(feature_dir, n=1, seq=16) + batch = torch.load( + os.path.join(feature_dir, sorted(os.listdir(feature_dir))[0]), + map_location="cpu", + ) + + # `hidden_state` is the target-side capture that the target head turns into + # the teacher distribution; `aux_hidden_state` is the draft backbone input. + input_ids, target, loss_mask = target_head.preprocess( + batch["input_ids"].unsqueeze(0), + batch["hidden_state"], + batch["loss_mask"].unsqueeze(0), + ) + target = target_head(target.cuda()) + hidden_states = batch["aux_hidden_state"].cuda() + input_ids = input_ids.cuda() + + # Prompt-heavy mask so trimming is non-trivial: first half unsupervised. + loss_mask = loss_mask.cuda().clone() + loss_mask[:, : loss_mask.shape[1] // 2] = 0 + attention_mask = torch.ones_like(input_ids) + + @torch.no_grad() + def step_losses(trim: bool): + plosses, *_ = model( + input_ids=input_ids, + attention_mask=attention_mask, + loss_mask=loss_mask, + target=target, + hidden_states=hidden_states, + trim_loss_positions=trim, + ) + return [float(p.item()) for p in plosses] + + full = step_losses(False) + trimmed = step_losses(True) + + self.assertEqual(len(full), len(trimmed)) + for i, (a, b) in enumerate(zip(full, trimmed)): + tol = 5e-3 * max(abs(a), abs(b)) + 1e-4 + self.assertLessEqual( + abs(a - b), + tol, + msg=f"step {i}: full={a} trimmed={b} (tol={tol})", + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_runtime/test_equiv_trim_usp.py b/tests/test_runtime/test_equiv_trim_usp.py new file mode 100644 index 000000000..c8a831afc --- /dev/null +++ b/tests/test_runtime/test_equiv_trim_usp.py @@ -0,0 +1,376 @@ +# coding=utf-8 +"""trim_loss_positions under USP sequence parallelism. + +Three layers, so CI keeps guarding the USP row-selection math even on hosts +without four GPUs: + +1. ``TestTrimPackGolden`` (CPU) -- hand-derived literal tables for the per-step + row sets: the overlap-tail bound (``s - j < chunk_len``), dead-step padding, + the unreachable-supervision fallback, and non-USP back-compat. +2. ``TestTrimLossAnalytic`` (one GPU) -- with zero logits and normalized + teachers every masked row's loss is exactly ``ln(draft_vocab)``, so the + trim-scaled and full-shaped losses must both hit a closed-form constant. +3. ``TestEquivTrimUspFourRank`` (four GPUs + flash-attn, like + ``test_equiv_4rank``) -- per-step loss parity, trim ON vs OFF, on a real + ring-4 offline pipeline across three adversarial masks: all-supervised + (trivial-equality boundary: any row/denominator error is exposed without + tolerance cover), supervision straddling every rank boundary + (overlap-tail-as-teacher), and supervision only inside one rank's overlap + tail (dead steps plus ranks with no supervision at all, which fall back to + the full path -- the mixed-path collective-alignment hazard). +""" + +import json +import math +import os +import shutil +import tempfile +import unittest +from unittest import mock + +import torch + +CUDA = torch.cuda.is_available() +NGPU = torch.cuda.device_count() if CUDA else 0 +WORLD_SIZE = 4 +SEQ = 48 +TTT = 3 + + +def _has_standard_flash_attention() -> bool: + try: + from flash_attn import flash_attn_varlen_func # noqa: F401 + from flash_attn.bert_padding import pad_input, unpad_input # noqa: F401 + from flash_attn.flash_attn_interface import ( # noqa: F401 + _flash_attn_varlen_backward, + ) + except Exception: + return False + return True + + +class TestTrimPackGolden(unittest.TestCase): + """Hand-derived expected outputs for _build_trim_pack (CPU only).""" + + def _mk(self, mask_list, seed=1, vocab=32, draft=8): + g = torch.Generator().manual_seed(seed) + ids = torch.randperm(vocab, generator=g)[:draft].sort().values + t2d = torch.zeros(vocab, dtype=torch.bool) + t2d[ids] = True + L = len(mask_list) + lm = torch.tensor(mask_list, dtype=torch.long).view(1, L, 1) + tgt = torch.randn(1, L, vocab, generator=g) + return tgt, t2d, lm + + def test_usp_tail_bound(self): + # C=6, ttt=2, local_len=8, supervised {4,5,6}; 6 is an overlap-tail + # position: a legal teacher, never a loss row. + # step0: {s : s-0 < 6} = {4,5} -> rows [4,5], n=2 + # step1: all of {4,5,6} -> rows [3,4,5], n=3 + from specforge.algorithms.eagle3.model import _build_trim_pack + + tgt, t2d, lm = self._mk([0, 0, 0, 0, 1, 1, 1, 0]) + p = _build_trim_pack(tgt, t2d, lm, length=2, chunk_len=6) + self.assertEqual(p["full_len"], 6) + self.assertEqual(p["sup"].tolist(), [4, 5, 6]) + self.assertEqual(p["rows_steps"][0].tolist(), [4, 5]) + self.assertEqual(p["keep_steps"][0].tolist(), [0, 1]) + self.assertEqual(p["nrows_steps"][0], 2) + self.assertEqual(p["rows_steps"][1].tolist(), [3, 4, 5]) + self.assertEqual(p["nrows_steps"][1], 3) + + def test_usp_dead_step(self): + # Supervision only at {6,7} (pure tail). Position 7 can never reach a + # row (needs j >= 2) and is filtered; 6 is unreachable at step 0. + from specforge.algorithms.eagle3.model import _build_trim_pack + + tgt, t2d, lm = self._mk([0, 0, 0, 0, 0, 0, 1, 1]) + p = _build_trim_pack(tgt, t2d, lm, length=2, chunk_len=6) + self.assertEqual(p["sup"].tolist(), [6]) + self.assertEqual(p["nrows_steps"][0], 0) # dead step + self.assertEqual(p["rows_steps"][0].tolist(), [0]) # padded dummy + self.assertEqual(p["rows_steps"][1].tolist(), [5]) + self.assertEqual(p["nrows_steps"][1], 1) + + def test_unreachable_supervision_falls_back(self): + from specforge.algorithms.eagle3.model import _build_trim_pack + + tgt, t2d, lm = self._mk([0, 0, 0, 0, 0, 0, 0, 1]) + self.assertIsNone(_build_trim_pack(tgt, t2d, lm, length=2, chunk_len=6)) + + def test_empty_supervision_falls_back(self): + from specforge.algorithms.eagle3.model import _build_trim_pack + + tgt, t2d, lm = self._mk([0] * 8) + self.assertIsNone(_build_trim_pack(tgt, t2d, lm, length=2, chunk_len=6)) + + def test_reuses_position_mask_from_teacher_computation(self): + from specforge.algorithms.eagle3 import model as eagle_model + + tgt, t2d, lm = self._mk([0, 1, 0, 1]) + nrows = 2 + draft_vocab = int(t2d.sum()) + teacher = ( + torch.zeros(1, nrows, draft_vocab), + torch.zeros(1, nrows, draft_vocab), + torch.zeros(1, nrows, dtype=torch.long), + torch.full((1, nrows, 1), 7), + ) + with mock.patch.object( + eagle_model, "_compute_target_p_eager", return_value=teacher + ): + pack = eagle_model._build_trim_pack(tgt, t2d, lm, length=2) + + self.assertIs(pack["position_mask_sup"], teacher[3]) + + def test_non_usp_backcompat(self): + # chunk_len=None -> C=L: the pre-USP semantics, plus full_len now comes + # from the row count rather than the (possibly padded) mask length. + from specforge.algorithms.eagle3.model import _build_trim_pack + + tgt, t2d, lm = self._mk([1, 0, 0, 0, 1, 0, 0, 1]) + p = _build_trim_pack(tgt, t2d, lm, length=2) + self.assertEqual(p["full_len"], 8) + self.assertEqual(p["rows_steps"][0].tolist(), [0, 4, 7]) + self.assertEqual(p["rows_steps"][1].tolist(), [3, 6]) + self.assertEqual(p["keep_steps"][1].tolist(), [1, 2]) + self.assertEqual(p["nrows_steps"][1], 2) + + +class TestTrimAdapterViews(unittest.TestCase): + def _inputs(self): + return { + "global_input_ids": torch.arange(8).view(1, 8), + "hidden_states": torch.arange(24).view(1, 8, 3), + "attention_mask": torch.ones(1, 8), + "position_ids": torch.arange(16).view(1, 16), + } + + def test_default_adapter_keeps_full_backbone_view(self): + from specforge.core.eagle3_adapters import BackendAdapter + + adapter = BackendAdapter(model=None) + inputs = self._inputs() + self.assertEqual(adapter.backbone_row_count(seq_length=8, ttt_length=2), 8) + + view = adapter.backbone_view(row_count=8, **inputs) + + self.assertIs(view.input_ids, inputs["global_input_ids"]) + self.assertIs(view.hidden_states, inputs["hidden_states"]) + self.assertIs(view.attention_mask, inputs["attention_mask"]) + self.assertIs(view.position_ids, inputs["position_ids"]) + + def test_usp_adapter_owns_chunk_and_position_slicing(self): + from specforge.core import eagle3_adapters + + world_sizes = {"sp": 4, "ulysses": 2} + with ( + mock.patch.object(eagle3_adapters, "get_draft_sp_group", return_value="sp"), + mock.patch.object( + eagle3_adapters, "get_sp_ulysses_group", return_value="ulysses" + ), + mock.patch.object( + eagle3_adapters.dist, + "get_world_size", + side_effect=lambda group: world_sizes[group], + ), + ): + adapter = eagle3_adapters.UspAdapter(model=None) + + inputs = self._inputs() + row_count = adapter.backbone_row_count(seq_length=8, ttt_length=2) + view = adapter.backbone_view(row_count=row_count, **inputs) + + self.assertEqual(row_count, 6) + self.assertEqual(view.input_ids.shape, (1, 6)) + self.assertEqual(view.hidden_states.shape, (1, 6, 3)) + self.assertEqual(view.attention_mask.shape, (1, 6)) + self.assertEqual(view.position_ids.shape, (1, 12)) + + +@unittest.skipUnless(CUDA, "loss kernel is a Triton kernel") +class TestTrimLossAnalytic(unittest.TestCase): + """Zero logits + normalized teachers => masked row loss == ln(D) exactly.""" + + def test_trim_and_full_hit_closed_form(self): + from specforge.core.loss import LogSoftmaxLoss + + D = 64 + g = torch.Generator().manual_seed(2) + + def one_hot_rows(n): + t = torch.zeros(1, n, D, device="cuda") + t[0, torch.arange(n), torch.randint(0, D, (n,), generator=g)] = 1.0 + return t + + # Trim-shaped: 3 selected rows, all masked-in; kernel mean == ln(D); + # rescaled by nrows/C = 3/6 -> 3*ln(64)/6. + kernel = LogSoftmaxLoss.apply( + torch.zeros(1, 3, D, device="cuda"), + one_hot_rows(3), + torch.ones(1, 3, 1, device="cuda"), + ) + self.assertAlmostEqual(kernel.item(), math.log(D), places=5) + self.assertAlmostEqual((kernel * (3 / 6)).item(), 2.0794415417, places=5) + + # Full-shaped: 6 rows, 3 masked-in -> same constant with no rescale. + pm = torch.tensor([0, 1, 0, 1, 1, 0], device="cuda").view(1, 6, 1) + full = LogSoftmaxLoss.apply( + torch.zeros(1, 6, D, device="cuda"), one_hot_rows(6), pm + ) + self.assertAlmostEqual(full.item(), 2.0794415417, places=5) + + +def _write_workdir(workdir): + from tests.test_runtime import _fixtures as fx + + fx.write_draft_config(os.path.join(workdir, "draft.json")) + fx.write_target_head_dir(os.path.join(workdir, "target")) + fx.write_vocab_mapping(os.path.join(workdir, "vocab_mapping.pt")) + masks = {} + m1 = torch.ones(SEQ, dtype=torch.long) + m1[-1] = 0 + masks["allones"] = m1 + m3 = torch.zeros(SEQ, dtype=torch.long) + m3[[10, 11, 12, 13, 22, 23, 24, 25, 34, 35, 36, 37]] = 1 + masks["boundary"] = m3 + m4 = torch.zeros(SEQ, dtype=torch.long) + m4[[12, 13]] = 1 + masks["tailonly"] = m4 + g = torch.Generator().manual_seed(11) + base_input = torch.randint(0, fx.V, (SEQ,), generator=g) + base_hid = torch.randn(1, SEQ, fx.H, generator=g).to(torch.bfloat16) + base_aux = torch.randn(1, SEQ, 3 * fx.H, generator=g).to(torch.bfloat16) + for name, lm in masks.items(): + d = os.path.join(workdir, f"features_{name}") + os.makedirs(d, exist_ok=True) + torch.save( + { + "input_ids": base_input.clone(), + "loss_mask": lm.clone(), + "hidden_state": base_hid.clone(), + "aux_hidden_state": base_aux.clone(), + }, + os.path.join(d, "0000.ckpt"), + ) + return list(masks) + + +def _worker(rank, world_size, port, workdir): + from tests.test_runtime import _fixtures as fx + + fx.init_rank_distributed( + rank, world_size, tp_size=1, sp_ulysses_size=1, sp_ring_size=4, port=str(port) + ) + try: + import torch.distributed as dist + + from specforge.algorithms.builtin import builtin_algorithm_registry + from specforge.algorithms.eagle3.model import OnlineEagle3Model + from specforge.modeling.auto import AutoDraftModel, AutoDraftModelConfig + from specforge.modeling.target.target_head import TargetHead + from specforge.runtime.data_plane import FeatureDataLoader, LocalFeatureStore + + torch.manual_seed(0) + torch.cuda.manual_seed_all(0) + torch.use_deterministic_algorithms(True, warn_only=True) + cfg = AutoDraftModelConfig.from_file(os.path.join(workdir, "draft.json")) + dm = AutoDraftModel.from_config( + cfg, attention_backend="usp", torch_dtype=torch.bfloat16 + ).cuda() + dm.load_vocab_mapping(os.path.join(workdir, "vocab_mapping.pt")) + dm.freeze_embedding() + model = OnlineEagle3Model( + draft_model=dm, length=TTT, attention_backend="usp" + ).cuda() + model.train() + target_head = TargetHead.from_pretrained( + os.path.join(workdir, "target"), lm_head_key="lm_head.weight" + ) + algorithm = builtin_algorithm_registry().resolve("eagle3") + provider = algorithm.providers.offline_for("text") + + results = {} + for case in ("allones", "boundary", "tailonly"): + refs = provider.build_reader( + os.path.join(workdir, f"features_{case}"), + run_id=f"trimusp-{case}", + ttt_length=TTT, + max_len=SEQ, + ).read() + loader = FeatureDataLoader( + LocalFeatureStore(f"trimusp-{case}-{rank}"), + refs=refs, + batch_size=1, + collate_fn=provider.build_collator(), + per_sample_transform=provider.build_normalizer( + SEQ, ttt_length=TTT, use_usp_preprocess=True + ), + strategy=algorithm.name, + ) + batch = next(iter(loader)) + + def step_losses(trim): + strat = algorithm.providers.step.build( + model, + target_head=target_head, + trim_loss_positions=trim, + ) + with torch.no_grad(): + out = strat.forward_loss(batch) + return [float(p.item()) for p in out.metrics["plosses"]] + + results[case] = {"full": step_losses(False), "trim": step_losses(True)} + + gathered = [None] * world_size + dist.all_gather_object(gathered, results) + if rank == 0: + with open(os.path.join(workdir, "results.json"), "w") as fh: + json.dump(gathered, fh) + dist.barrier() + finally: + from specforge.distributed import destroy_distributed + + destroy_distributed() + + +@unittest.skipUnless( + CUDA and NGPU >= WORLD_SIZE and _has_standard_flash_attention(), + "requires four CUDA devices and the standard flash-attn USP interfaces", +) +class TestEquivTrimUspFourRank(unittest.TestCase): + def test_trim_matches_full_per_step_on_ring4(self): + import torch.multiprocessing as mp + + workdir = tempfile.mkdtemp(prefix="trim_usp_") + self.addCleanup(shutil.rmtree, workdir, ignore_errors=True) + _write_workdir(workdir) + mp.spawn( + _worker, + args=(WORLD_SIZE, 29871, workdir), + nprocs=WORLD_SIZE, + join=True, + ) + with open(os.path.join(workdir, "results.json")) as fh: + gathered = json.load(fh) + for case in ("allones", "boundary", "tailonly"): + for rank, res in enumerate(gathered): + full, trim = res[case]["full"], res[case]["trim"] + self.assertEqual(len(full), TTT) + for j, (a, b) in enumerate(zip(full, trim)): + if case == "allones": + # expected near-bit-equal; 1e-6 is ~4 orders below the + # smallest possible discrete error (one row's worth, + # ~loss/C) while allowing 1-2 ulp of fp32 noise + tol = 1e-6 + else: + tol = max(1e-3 * abs(a), 1e-4) + self.assertLessEqual( + abs(a - b), + tol, + msg=f"{case} rank{rank} step{j}: full={a} trim={b}", + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_runtime/test_evaluator_aggregation.py b/tests/test_runtime/test_evaluator_aggregation.py index 8055b3130..b7436b357 100644 --- a/tests/test_runtime/test_evaluator_aggregation.py +++ b/tests/test_runtime/test_evaluator_aggregation.py @@ -54,6 +54,19 @@ def _scalar_out(loss, acc, tokens, denom=None): return StepOutput(loss=torch.tensor(float(loss)), metrics=metrics) +def _additive_scalar_out(loss_num, loss_den, accuracy_num, accuracy_den): + loss_num = torch.tensor(float(loss_num)) + loss_den = torch.tensor(float(loss_den)) + accuracy_num = torch.tensor(float(accuracy_num)) + accuracy_den = torch.tensor(float(accuracy_den)) + return StepOutput( + loss=loss_num / loss_den, + metrics={"accuracy": accuracy_num / accuracy_den}, + ratio_metrics={"acc": (accuracy_num, accuracy_den)}, + loss_terms=(loss_num, loss_den), + ) + + class TestEvaluatorAggregation(unittest.TestCase): def _run(self, outputs): from specforge.eval import Evaluator @@ -171,6 +184,19 @@ def test_scalar_accuracy_weighted_by_accuracy_denom(self): # loss-token weighting would skew to (0.75*10 + 0.5*50)/60 ~ 0.542 self.assertNotAlmostEqual(m["eval/avg_acc"], 32.5 / 60, places=2) + def test_additive_scalar_terms_are_partition_invariant(self): + split = self._run( + [ + _additive_scalar_out(2, 1, 1, 1), + _additive_scalar_out(30, 3, 1, 3), + ] + ) + combined = self._run([_additive_scalar_out(32, 4, 2, 4)]) + + self.assertEqual(split, combined) + self.assertEqual(split["eval/avg_loss"], 8.0) + self.assertEqual(split["eval/avg_acc"], 0.5) + def test_reports_per_position_acceptance(self): m = self._run([_step_output(1.0, corrects=[3, 2], denoms=[4, 4])]) self.assertAlmostEqual(m["eval/per_position_acc"][0], 0.75, places=6) diff --git a/tests/test_runtime/test_export.py b/tests/test_runtime/test_export.py index c5d7a5dbe..9a2dc0993 100644 --- a/tests/test_runtime/test_export.py +++ b/tests/test_runtime/test_export.py @@ -10,6 +10,7 @@ round trips require GPU and can be run on the H200 box via rcli. """ +import json import os import tempfile import unittest @@ -19,6 +20,94 @@ CUDA = torch.cuda.is_available() +class TestRoPEConfigCompatibility(unittest.TestCase): + def _write_config(self, directory, payload): + path = os.path.join(directory, "config.json") + with open(path, "w", encoding="utf-8") as handle: + json.dump(payload, handle) + return path + + def test_modern_rope_parameters_are_mirrored_for_legacy_readers(self): + from specforge.export.checkpoint_io import apply_legacy_rope_scaling + + with tempfile.TemporaryDirectory() as directory: + path = self._write_config( + directory, + { + "rope_parameters": { + "rope_type": "yarn", + "factor": 128.0, + "rope_theta": 8_000_000, + } + }, + ) + self.assertTrue(apply_legacy_rope_scaling(directory)) + with open(path, encoding="utf-8") as handle: + config = json.load(handle) + + self.assertEqual( + config["rope_scaling"], + {"rope_type": "yarn", "factor": 128.0}, + ) + self.assertEqual(config["rope_theta"], 8_000_000) + + def test_legacy_rope_scaling_is_mirrored_for_modern_readers(self): + from specforge.export.checkpoint_io import apply_legacy_rope_scaling + + with tempfile.TemporaryDirectory() as directory: + path = self._write_config( + directory, + { + "rope_theta": 8_000_000, + "rope_scaling": {"type": "yarn", "factor": 128.0}, + }, + ) + self.assertTrue(apply_legacy_rope_scaling(directory)) + with open(path, encoding="utf-8") as handle: + config = json.load(handle) + + self.assertEqual( + config["rope_parameters"], + {"type": "yarn", "factor": 128.0, "rope_theta": 8_000_000}, + ) + + def test_default_rope_config_is_not_rewritten(self): + from specforge.export.checkpoint_io import apply_legacy_rope_scaling + + with tempfile.TemporaryDirectory() as directory: + path = self._write_config( + directory, + {"rope_parameters": {"rope_type": "default"}}, + ) + with open(path, "rb") as handle: + before = handle.read() + self.assertFalse(apply_legacy_rope_scaling(directory)) + with open(path, "rb") as handle: + after = handle.read() + + self.assertEqual(after, before) + + def test_default_rope_theta_is_mirrored_for_legacy_readers(self): + from specforge.export.checkpoint_io import apply_legacy_rope_scaling + + with tempfile.TemporaryDirectory() as directory: + path = self._write_config( + directory, + { + "rope_parameters": { + "rope_type": "default", + "rope_theta": 1_000_000, + } + }, + ) + self.assertTrue(apply_legacy_rope_scaling(directory)) + with open(path, encoding="utf-8") as handle: + config = json.load(handle) + + self.assertEqual(config["rope_theta"], 1_000_000) + self.assertNotIn("rope_scaling", config) + + class TestLegacyVocabMappingCompatibility(unittest.TestCase): def setUp(self): from specforge.modeling.auto import AutoDraftModel, AutoDraftModelConfig diff --git a/tests/test_runtime/test_feature_dataloader.py b/tests/test_runtime/test_feature_dataloader.py index 7e5151d09..b1afcd4b9 100644 --- a/tests/test_runtime/test_feature_dataloader.py +++ b/tests/test_runtime/test_feature_dataloader.py @@ -174,6 +174,25 @@ def test_offline_loader_emits_trainbatch(self): self.assertEqual(q.in_flight(), 0) self.assertEqual(q.depth(), 0) + @unittest.skipUnless(torch.cuda.is_available(), "pinned memory requires CUDA") + def test_pin_memory_pins_collated_cpu_tensors(self): + store = LocalFeatureStore("st") + ref = store.put( + {"x": torch.arange(8)}, + sample_id="sample-0", + metadata={"run_id": "run", "target_repr": "hidden_state"}, + ) + loader = FeatureDataLoader( + store, + refs=[ref], + drop_last=False, + pin_memory=True, + ) + + batch = next(iter(loader)) + + self.assertTrue(batch.tensors["x"].is_pinned()) + def test_drop_last(self): with tempfile.TemporaryDirectory() as d: self._write_offline_files(d, n=3) diff --git a/tests/test_runtime/test_http_inbox.py b/tests/test_runtime/test_http_inbox.py new file mode 100644 index 000000000..c119fd2a8 --- /dev/null +++ b/tests/test_runtime/test_http_inbox.py @@ -0,0 +1,115 @@ +"""Private-network inbox relay contracts.""" + +from __future__ import annotations + +import json +import socket +import tempfile +import threading +import unittest +from unittest import mock +from urllib.request import Request, urlopen + +from specforge.runtime.contracts import FeatureSpec, SampleRef +from specforge.runtime.data_plane.http_inbox import InboxHTTPServer, RemoteInboxChannel +from specforge.runtime.data_plane.ref_distributor import RefDistributor +from specforge.runtime.data_plane.streaming_ref_channel import StreamingRefChannel + + +def _ref(sample_id: str) -> SampleRef: + return SampleRef( + sample_id=sample_id, + run_id="run0", + source_task_id=f"task-{sample_id}", + feature_store_uri=f"mooncake://run0/{sample_id}", + feature_keys={"hidden_state": f"{sample_id}/hidden_state"}, + feature_specs={ + "hidden_state": FeatureSpec( + name="hidden_state", shape=(2, 4), dtype="float32" + ) + }, + strategy="dspark", + metadata={"target_repr": "hidden_state"}, + ) + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +class TestHTTPInbox(unittest.TestCase): + def setUp(self): + self.work = tempfile.mkdtemp(prefix="http-inbox-") + self.path = RefDistributor.inbox_path(self.work, 1) + self.local = StreamingRefChannel(self.path) + self.origin = f"http://127.0.0.1:{_free_port()}" + self.server = InboxHTTPServer( + self.work, 2, self.origin, bind_host="127.0.0.1" + ).start() + self.remote = RemoteInboxChannel(self.origin, 1) + + def tearDown(self): + self.server.stop() + + def test_tail_read_close_and_consumed_counter(self): + self.local.publish_batch([_ref("s0"), _ref("s1")]) + self.assertEqual([ref.sample_id for ref in self.remote.poll()], ["s0", "s1"]) + self.assertEqual(self.remote.poll(), []) + + self.remote.mark_consumed(2) + self.assertEqual(self.local.consumed_remote(), 2) + + self.remote.mark_consumed(1) + self.assertEqual(self.local.consumed_remote(), 3) + + self.local.close() + self.assertTrue(self.remote.is_closed()) + + def test_status_probe_does_not_discard_unpolled_refs(self): + self.local.publish(_ref("s0")) + self.assertFalse(self.remote.is_closed()) + self.assertIsNone(self.remote.failure()) + self.assertEqual([ref.sample_id for ref in self.remote.poll()], ["s0"]) + + def test_failure_is_forwarded(self): + failure = self.path + ".failed" + with open(failure, "w", encoding="utf-8") as stream: + stream.write("capture failed") + self.assertIn("capture failed", self.remote.failure()) + + def test_consumed_target_is_idempotent_under_concurrent_retries(self): + body = json.dumps({"target": 5}).encode("utf-8") + + def post_target(): + request = Request( + f"{self.origin}/v1/inboxes/1/consumed", + data=body, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urlopen(request, timeout=2.0) as response: + self.assertEqual(json.load(response)["consumed"], 5) + + threads = [threading.Thread(target=post_target) for _ in range(12)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + self.assertEqual(self.local.consumed_remote(), 5) + + def test_pull_treats_connection_reset_as_transient(self): + with mock.patch( + "specforge.runtime.data_plane.http_inbox.urlopen", + side_effect=ConnectionResetError("peer reset"), + ): + self.assertEqual(self.remote.poll(), []) + + def test_invalid_origin_is_rejected(self): + with self.assertRaisesRegex(ValueError, "http://host:port"): + RemoteInboxChannel("https://trainer.example:35900/path", 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_runtime/test_launch_plan.py b/tests/test_runtime/test_launch_plan.py index 226410221..feef5a80c 100644 --- a/tests/test_runtime/test_launch_plan.py +++ b/tests/test_runtime/test_launch_plan.py @@ -30,6 +30,7 @@ from specforge.launch_plan import build_launch_plan as _build_launch_plan from specforge.launch_plan import run_commands from specforge.training.capture_contract import ServerCaptureContract +from tests.utils import wait_for_processes_to_stop ALGORITHM = builtin_algorithm_registry().resolve("dflash") @@ -391,8 +392,12 @@ def test_multi_node_requires_node_local_consumer_state(self): Config.model_validate(raw) def test_multi_node_keeps_wal_local_and_inboxes_shared(self): + raw = _config(mode="disaggregated", nproc=2, nnodes=2).model_dump() + raw["deployment"]["disaggregated"][ + "inbox_server_url" + ] = "http://trainer-0:35900" plan = build_launch_plan( - _config(mode="disaggregated", nproc=2, nnodes=2), + Config.model_validate(raw), config_path="run.yaml", requested_role="consumer", node_rank=0, @@ -404,6 +409,30 @@ def test_multi_node_keeps_wal_local_and_inboxes_shared(self): command = plan.commands[0] self.assertEqual("/local/attempt-1/consumer.sqlite", command.env["DISAGG_DB"]) self.assertEqual("/shared/attempt-1/inboxes", command.env["DISAGG_INBOX_DIR"]) + self.assertEqual( + "http://trainer-0:35900", command.env["DISAGG_INBOX_SERVER_URL"] + ) + + def test_inbox_server_url_is_typed_and_online_multinode_only(self): + cfg = _config(mode="disaggregated", nproc=2, nnodes=2) + invalid = { + "https": "https://trainer-0:35900", + "missing port": "http://trainer-0", + "path": "http://trainer-0:35900/inboxes", + } + for name, value in invalid.items(): + with self.subTest(case=name): + raw = cfg.model_dump() + raw["deployment"]["disaggregated"]["inbox_server_url"] = value + with self.assertRaisesRegex(ValidationError, "http://host:port"): + Config.model_validate(raw) + + raw = _config(mode="disaggregated", nproc=2, nnodes=1).model_dump() + raw["deployment"]["disaggregated"][ + "inbox_server_url" + ] = "http://trainer-0:35900" + with self.assertRaisesRegex(ValidationError, "multi-node trainer"): + Config.model_validate(raw) def test_disaggregated_roles_are_independently_selectable(self): producer = build_launch_plan( @@ -539,7 +568,7 @@ def test_managed_local_rejects_external_and_nonlocal_modes(self): with self.assertRaisesRegex(ValidationError, message): Config.model_validate(raw) - def test_managed_local_accepts_minimum_context_and_disables_radix_cache(self): + def test_managed_local_accepts_minimum_context_and_configures_radix_cache(self): with tempfile.TemporaryDirectory() as root: cfg = _managed_config(os.path.join(root, "attempt")) raw = cfg.model_dump() @@ -556,6 +585,19 @@ def test_managed_local_accepts_minimum_context_and_disables_radix_cache(self): self.assertEqual(argv[argv.index("--context-length") + 1], "135") self.assertIn("--disable-radix-cache", argv) + with tempfile.TemporaryDirectory() as root: + cfg = _managed_config(os.path.join(root, "attempt")) + raw = cfg.model_dump() + raw["model"]["sglang_disable_radix_cache"] = False + validated = Config.model_validate(raw) + with mock.patch( + "specforge.training.capture_contract.resolve_server_capture_contract", + return_value=CAPTURE_CONTRACT, + ): + plan = build_launch_plan(validated, config_path="run.yaml", env={}) + + self.assertNotIn("--disable-radix-cache", plan.services[1].command.argv) + def test_managed_local_plan_owns_mooncake_and_multiple_capture_servers(self): servers = [ { @@ -626,6 +668,7 @@ def test_managed_local_plan_owns_mooncake_and_multiple_capture_servers(self): "--rpc_port=35551", "--http_metadata_server_port=35880", "--metrics_port=35903", + "--default_kv_lease_ttl=500", ), ) self.assertEqual(mooncake.readiness.kind, "mooncake") @@ -685,7 +728,7 @@ def test_managed_local_plan_owns_mooncake_and_multiple_capture_servers(self): def test_multiserver_example_yaml_builds_the_managed_plan(self): path = ( Path(__file__).resolve().parents[2] - / "examples/configs/qwen3.6-27b-dflash-multiserver-disaggregated.yaml" + / "examples/configs/online/disaggregated/managed-local/qwen3.6-27b-dflash-multiserver-disaggregated.yaml" ) cfg = Config.from_file(str(path)) with ( @@ -1453,6 +1496,7 @@ def test_sigterm_cleans_the_real_child_process_group(self): (sys.executable, "-c", {child_code!r}, {marker!r}), ), ), + shutdown_grace_s=1.0, ) raise SystemExit(run_commands(plan)) """ @@ -1464,6 +1508,7 @@ def test_sigterm_cleans_the_real_child_process_group(self): text=True, ) child_pid = None + grandchild_pid = None try: deadline = time.monotonic() + 10 while time.monotonic() < deadline: @@ -1471,12 +1516,15 @@ def test_sigterm_cleans_the_real_child_process_group(self): with open(marker, encoding="utf-8") as stream: raw = stream.read().strip() if len(raw.split()) == 2: - child_pid = int(raw.split()[0]) + child_pid, grandchild_pid = map(int, raw.split()) break if runner.poll() is not None: break time.sleep(0.02) self.assertIsNotNone(child_pid, "managed child never became ready") + self.assertIsNotNone( + grandchild_pid, "managed grandchild never became ready" + ) self.assertEqual(os.getpgid(child_pid), child_pid) os.kill(runner.pid, signal.SIGTERM) @@ -1487,15 +1535,13 @@ def test_sigterm_cleans_the_real_child_process_group(self): f"stdout={stdout!r} stderr={stderr!r}", ) - deadline = time.monotonic() + 5 - while time.monotonic() < deadline: - try: - os.killpg(child_pid, 0) - except ProcessLookupError: - break - time.sleep(0.02) - else: - self.fail(f"managed process group {child_pid} survived SIGTERM") + survivors = wait_for_processes_to_stop( + (child_pid, grandchild_pid), timeout_s=5 + ) + self.assertFalse( + survivors, + f"managed processes survived SIGTERM: {survivors}", + ) finally: if runner.poll() is None: runner.kill() diff --git a/tests/test_runtime/test_model_loading.py b/tests/test_runtime/test_model_loading.py index 87f053957..419b28ffe 100644 --- a/tests/test_runtime/test_model_loading.py +++ b/tests/test_runtime/test_model_loading.py @@ -106,6 +106,17 @@ def _draft_payload(architecture: str, *, layers: int = 1, block_size=None): class DraftConfigResolutionTest(unittest.TestCase): + def test_explicit_null_draft_vocab_size_does_not_fallback(self): + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "draft.json") + payload = _draft_payload("LlamaForCausalLMEagle3") + payload["draft_vocab_size"] = None + with open(path, "w", encoding="utf-8") as stream: + json.dump(payload, stream) + + with self.assertRaisesRegex(ValueError, "draft_vocab_size cannot be null"): + load_draft_config_source(path) + def test_config_resolution_does_not_initialize_cuda_model_dependencies(self): with tempfile.TemporaryDirectory() as directory: path = os.path.join(directory, "draft.json") @@ -157,6 +168,9 @@ def test_target_derived_defaults_match_legacy_trainers(self): self.assertEqual(resolved.block_size, block_size) self.assertEqual(resolved.num_target_layers, 12) self.assertEqual(len(resolved.dflash_config["target_layer_ids"]), 1) + self.assertEqual(resolved.layer_types, ["full_attention"]) + self.assertIsNone(resolved.sliding_window) + self.assertFalse(resolved.use_sliding_window) else: self.assertEqual(resolved.draft_vocab_size, 32000) @@ -180,6 +194,36 @@ def test_dflash_typed_overrides_rebuild_capture_layers(self): self.assertEqual(resolved.num_hidden_layers, 2) self.assertEqual(resolved.block_size, 8) self.assertEqual(len(resolved.dflash_config["target_layer_ids"]), 2) + self.assertEqual( + resolved.layer_types, + ["full_attention", "full_attention"], + ) + + def test_dflash_layer_override_rejects_ambiguous_hybrid_resize(self): + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "draft.json") + payload = _draft_payload("DFlashDraftModel", layers=3, block_size=16) + payload.update( + layer_types=[ + "sliding_attention", + "sliding_attention", + "full_attention", + ], + sliding_window=128, + use_sliding_window=True, + ) + with open(path, "w", encoding="utf-8") as stream: + json.dump(payload, stream) + cfg = _run_config( + "dflash", + draft_model_config=path, + draft_num_hidden_layers=2, + ) + with self.assertRaisesRegex(ValueError, "mixed DFlash layer_types"): + resolve_draft_config( + cfg, + provider=_draft_config_provider("dflash"), + ) def test_local_json_and_directory_are_equivalent_sources(self): with tempfile.TemporaryDirectory() as directory: diff --git a/tests/test_runtime/test_mooncake_store.py b/tests/test_runtime/test_mooncake_store.py index 3c1590b9f..cf49bc864 100644 --- a/tests/test_runtime/test_mooncake_store.py +++ b/tests/test_runtime/test_mooncake_store.py @@ -11,15 +11,23 @@ import ctypes import importlib.util import unittest +from inspect import signature +from unittest import mock import torch from specforge.runtime.control_plane.controller import DataFlowController +from specforge.runtime.control_plane.dp_ack import DPAckController from specforge.runtime.control_plane.metadata_store import InMemoryMetadataStore from specforge.runtime.data_plane.disaggregated import AuthPolicy from specforge.runtime.data_plane.feature_store import ( + DEFAULT_PENDING_DRAIN_MAX_ATTEMPTS, + DEFAULT_PENDING_DRAIN_RETRY_INTERVAL_S, + DEFAULT_SAMPLE_DRAIN_MAX_ATTEMPTS, + DEFAULT_SAMPLE_DRAIN_RETRY_INTERVAL_S, LocalFeatureStore, drain_feature_store_removals, + drain_feature_store_sample_removals, ) from specforge.runtime.data_plane.mooncake_store import MooncakeFeatureStore @@ -108,6 +116,35 @@ def _store(**kw): class TestMooncakeFeatureStore(unittest.TestCase): + def test_drain_interfaces_share_retry_defaults(self): + groups = ( + ( + ( + drain_feature_store_removals, + MooncakeFeatureStore.drain_pending_removals, + ), + DEFAULT_PENDING_DRAIN_MAX_ATTEMPTS, + DEFAULT_PENDING_DRAIN_RETRY_INTERVAL_S, + ), + ( + ( + drain_feature_store_sample_removals, + MooncakeFeatureStore.drain_sample_removals, + ), + DEFAULT_SAMPLE_DRAIN_MAX_ATTEMPTS, + DEFAULT_SAMPLE_DRAIN_RETRY_INTERVAL_S, + ), + ) + for drains, attempts, interval in groups: + for drain in drains: + with self.subTest(drain=drain.__qualname__): + parameters = signature(drain).parameters + self.assertEqual(parameters["max_attempts"].default, attempts) + self.assertEqual( + parameters["retry_interval_s"].default, + interval, + ) + def test_put_get_roundtrip_bit_exact(self): fs = _store() src = _tensors() @@ -133,6 +170,49 @@ def test_hard_pin_config_on_put(self): self.assertTrue(getattr(fake.last_config, "with_hard_pin", False)) self.assertTrue(fs.health()["hard_pin"]) + def test_constructor_falls_back_to_soft_pin(self): + import specforge.runtime.data_plane.mooncake_store as mooncake_store + + class _SoftPinOnlyConfig: + def __init__(self): + self.replica_num = 1 + self.with_soft_pin = False + + fake = _FakeMooncakeStore() + with ( + mock.patch.object( + mooncake_store, + "_connect_store", + return_value=(fake, _SoftPinOnlyConfig), + ), + self.assertLogs(mooncake_store.logger, level="WARNING") as logs, + ): + fs = MooncakeFeatureStore(store_id="run0", setup_kwargs={}) + + self.assertTrue(fs._put_config.with_soft_pin) + self.assertIn("falling back to with_soft_pin", "\n".join(logs.output)) + + def test_constructor_tolerates_config_without_pin_fields(self): + import specforge.runtime.data_plane.mooncake_store as mooncake_store + + class _ConfigWithoutPinFields: + def __init__(self): + self.replica_num = 1 + + fake = _FakeMooncakeStore() + with ( + mock.patch.object( + mooncake_store, + "_connect_store", + return_value=(fake, _ConfigWithoutPinFields), + ), + self.assertLogs(mooncake_store.logger, level="WARNING") as logs, + ): + fs = MooncakeFeatureStore(store_id="run0", setup_kwargs={}) + + self.assertEqual(fs._put_config.replica_num, 1) + self.assertIn("neither with_hard_pin nor with_soft_pin", "\n".join(logs.output)) + def test_get_after_release_raises(self): fs = _store() ref = fs.put(_tensors(), sample_id="s0", metadata=_meta()) @@ -267,6 +347,94 @@ def release_remote_lease(interval): self.assertEqual(fs.health()["force_freed_total"], 1) self.assertFalse(_phys_resident(fake)) + def test_lifecycle_drain_forces_removal_after_application_lease_closes(self): + class ForceAwareFake(_FakeMooncakeStore): + def __init__(self): + super().__init__() + self.force_values = [] + + def remove(self, key, force=False): + self.remove_calls += 1 + self.force_values.append(force) + if not force: + return -706 + self._d.pop(key, None) + return 0 + + fake = ForceAwareFake() + fs = MooncakeFeatureStore(store=fake, store_id="run0") + ref = fs.put(_tensors(), sample_id="s0", metadata=_meta()) + _, handle = fs.get(ref) + + fs.release(handle) + self.assertEqual(fs.health()["release_pending"], 1) + self.assertTrue(_phys_resident(fake)) + + report = drain_feature_store_removals(fs) + + self.assertEqual(report["attempts"], 1) + self.assertEqual(fs.health()["release_pending"], 0) + self.assertFalse(_phys_resident(fake)) + num_features = len(ref.feature_keys) + self.assertEqual(fake.force_values[:num_features], [False] * num_features) + self.assertEqual(fake.force_values[num_features:], [True] * num_features) + + def test_optimizer_ack_forces_only_durable_samples(self): + class ForceAwareFake(_FakeMooncakeStore): + def __init__(self): + super().__init__() + self.force_values = [] + + def remove(self, key, force=False): + self.remove_calls += 1 + self.force_values.append((key, force)) + if not force: + return -706 + self._d.pop(key, None) + return 0 + + fake = ForceAwareFake() + fs = MooncakeFeatureStore(store=fake, store_id="run0") + durable = fs.put(_tensors(), sample_id="durable", metadata=_meta()) + prefetched = fs.put(_tensors(), sample_id="prefetched", metadata=_meta()) + for ref in (durable, prefetched): + _, handle = fs.get(ref) + fs.release(handle) + self.assertEqual(fs.health()["release_pending"], 2) + + controller = DPAckController( + "run0", + feature_store=fs, + metadata_store=InMemoryMetadataStore(), + ) + controller.commit_samples("distributor", [durable, prefetched]) + controller.ack_train_refs( + "trainer", + [durable.sample_id], + global_step=1, + optimizer_durable=True, + ) + + # The current optimizer window is only tombstoned. Its short remote + # read lease gets one full window to expire, so ack itself never sleeps. + self.assertTrue(_phys_resident(fake, sid="durable")) + self.assertTrue(_phys_resident(fake, sid="prefetched")) + self.assertEqual(fs.health()["release_pending"], 2) + + controller.ack_train_refs( + "trainer", + [], + global_step=2, + optimizer_durable=True, + ) + + self.assertFalse(_phys_resident(fake, sid="durable")) + self.assertTrue(_phys_resident(fake, sid="prefetched")) + self.assertEqual(fs.health()["release_pending"], 1) + marker = controller.store.durable_marker() + self.assertEqual(marker["global_step"], 2) + self.assertEqual(marker["acked"], {"durable"}) + def test_lifecycle_drain_does_not_renew_read_lease_between_retries(self): clock = _FakeClock() lease_ttl = 1.0 diff --git a/tests/test_runtime/test_npu_portability.py b/tests/test_runtime/test_npu_portability.py index 95b4d632a..906005785 100644 --- a/tests/test_runtime/test_npu_portability.py +++ b/tests/test_runtime/test_npu_portability.py @@ -101,7 +101,7 @@ def test_npu_init_uses_hccl_binding_and_npu_meshes(self): sf_dist, "_load_yunchang_globals", return_value=(process_group, set_seq_parallel_pg), - ), + ) as load_yunchang, ): sf_dist.init_distributed( timeout=7, tp_size=2, sp_ulysses_size=1, sp_ring_size=1 @@ -116,7 +116,15 @@ def test_npu_init_uses_hccl_binding_and_npu_meshes(self): [call.kwargs["device_type"] for call in from_group.call_args_list], ["npu", "npu"], ) - set_seq_parallel_pg.assert_called_once_with(1, 1, 3, 8) + # SP sizes of 1 must not touch yunchang: its import probes CUDA and + # crashes NPU-only torch builds. The public SP getters must still + # return the singleton draft-SP group because PyTorch interprets a + # None group as the full WORLD group. + load_yunchang.assert_not_called() + set_seq_parallel_pg.assert_not_called() + self.assertEqual(sf_dist.get_draft_sp_group(), "draft:sp") + self.assertEqual(sf_dist.get_sp_ulysses_group(), "draft:sp") + self.assertEqual(sf_dist.get_sp_ring_group(), "draft:sp") class DistributedTeardownTest(unittest.TestCase): diff --git a/tests/test_runtime/test_package_architecture.py b/tests/test_runtime/test_package_architecture.py index 82e4e3b71..c775c1760 100644 --- a/tests/test_runtime/test_package_architecture.py +++ b/tests/test_runtime/test_package_architecture.py @@ -114,26 +114,26 @@ "docs/benchmarks/domino-disaggregated-performance.md", ), "examples/disagg/run_qwen2.5_7b_eagle3_disagg.sh": ( - "examples/configs/qwen2.5-7b-eagle3-offline-disaggregated.yaml", + "examples/configs/offline/disaggregated/qwen2.5-7b-eagle3-offline-disaggregated.yaml", "examples/disagg/run_offline_2node.sh", "docs/benchmarks/eagle3-disaggregated-parity.md", ), "examples/disagg/run_qwen3.6_27b_dflash_disagg.sh": ( - "examples/configs/qwen3.6-27b-dflash-1server-dp2-disaggregated.yaml", - "examples/configs/qwen3.6-27b-dflash-disaggregated.yaml", + "examples/configs/online/disaggregated/managed-local/qwen3.6-27b-dflash-1server-dp2-disaggregated.yaml", + "examples/configs/online/disaggregated/external/qwen3.6-27b-dflash-disaggregated.yaml", ), "examples/disagg/run_qwen3.6_27b_dflash_disagg_multiserver.sh": ( - "examples/configs/qwen3.6-27b-dflash-multiserver-disaggregated.yaml", + "examples/configs/online/disaggregated/managed-local/qwen3.6-27b-dflash-multiserver-disaggregated.yaml", ), "examples/disagg/run_qwen3_8b_dflash_disagg_1srv_dp7.sh": ( - "examples/configs/qwen3-8b-dflash-1server-dp7-disaggregated.yaml", + "examples/configs/online/disaggregated/managed-local/qwen3-8b-dflash-1server-dp7-disaggregated.yaml", ), "examples/disagg/run_qwen3_8b_domino_disagg_1srv_dp7.sh": ( - "examples/configs/qwen3-8b-domino-1server-dp7-disaggregated.yaml", + "examples/configs/online/disaggregated/managed-local/qwen3-8b-domino-1server-dp7-disaggregated.yaml", "docs/benchmarks/domino-disaggregated-performance.md", ), "examples/disagg/run_qwen3_8b_domino_disagg_multiserver.sh": ( - "examples/configs/qwen3-8b-domino-multiserver-disaggregated.yaml", + "examples/configs/online/disaggregated/managed-local/qwen3-8b-domino-multiserver-disaggregated.yaml", ), "examples/disagg/run_domino_dflash_serving_gate.sh": ( "scripts/gates/README.md", @@ -713,6 +713,7 @@ def test_dspark_configs_are_qwen3_gqa_only(self): { "glm-5.2-dspark.json", "inkling-dspark.json", + "kimi-k3-dspark.json", "qwen3-4b-dspark.json", "qwen3-8b-dspark.json", "qwen3.6-27b-dspark.json", @@ -751,6 +752,7 @@ def test_examples_and_scripts_do_not_bypass_the_cli(self): Path("examples/disagg/run_offline.sh"), Path("examples/disagg/run_offline_2node.sh"), Path("examples/disagg/run_qwen3_8b_dflash_disagg_2node.sh"), + Path("examples/disagg/run_inkling_dspark_disagg_2node.sh"), } bypasses = [] train_command = re.compile(r"\btrain\s+(?:--config|-c)\b") diff --git a/tests/test_runtime/test_ref_distributor.py b/tests/test_runtime/test_ref_distributor.py index a55510369..1039fda04 100644 --- a/tests/test_runtime/test_ref_distributor.py +++ b/tests/test_runtime/test_ref_distributor.py @@ -732,6 +732,67 @@ def abort(self, sample_id, *, reason): self.assertEqual(reason, "optimizer-boundary-durable-ack") controller.store.close() + def test_selective_cleanup_waits_one_boundary_without_strong_drain(self): + retries = [] + + class FeatureStore: + def abort(self, sample_id, *, reason): + pass + + def retry_sample_removals(self, sample_ids): + retries.append(list(sample_ids)) + return {"remaining_ids": []} + + def drain_sample_removals(self, sample_ids, **kwargs): + raise AssertionError("normal cleanup must not enter the sleeping drain") + + store = FeatureStore() + controller = DPAckController( + "run0", + is_authority=True, + feature_store=store, + metadata_store=SQLiteMetadataStore(os.path.join(self.dir, "lag.db")), + ) + controller.commit_samples("w0", [_ref("s0")]) + controller.ack_train_refs("t0", ["s0"], global_step=1, optimizer_durable=True) + self.assertEqual(retries, []) + + controller.ack_train_refs("t0", [], global_step=2, optimizer_durable=True) + self.assertEqual(retries, [["s0"]]) + controller.store.close() + + def test_selective_drain_failure_is_reported_after_bounded_deferral(self): + class FeatureStore: + def abort(self, sample_id, *, reason): + pass + + def retry_sample_removals(self, sample_ids): + return {"remaining_ids": list(sample_ids)} + + def drain_sample_removals(self, sample_ids, **kwargs): + raise OSError(f"remove stayed pinned for {sample_ids}") + + controller = DPAckController( + "run0", + is_authority=True, + feature_store=FeatureStore(), + metadata_store=SQLiteMetadataStore(os.path.join(self.dir, "drain.db")), + ) + controller.commit_samples("w0", [_ref("s0")]) + controller.ack_train_refs("t0", ["s0"], global_step=1, optimizer_durable=True) + for step in range(2, 5): + controller.ack_train_refs( + "t0", [], global_step=step, optimizer_durable=True + ) + with self.assertRaisesRegex( + RuntimeError, "optimizer-boundary selective drain.*remove stayed pinned" + ): + controller.ack_train_refs("t0", [], global_step=5, optimizer_durable=True) + marker = controller.store.durable_marker() + self.assertEqual(marker["global_step"], 5) + self.assertTrue(marker["optimizer_durable"]) + controller.store.close() + def test_non_authority_participates_but_records_nothing(self): calls = [] diff --git a/tests/test_runtime/test_schedule.py b/tests/test_runtime/test_schedule.py index 4a28c3adb..cd0e826d5 100644 --- a/tests/test_runtime/test_schedule.py +++ b/tests/test_runtime/test_schedule.py @@ -1,10 +1,13 @@ import json +import os import tempfile import unittest from types import SimpleNamespace +from unittest import mock from specforge.training.disaggregated import ( _ONLINE_SCHEDULE_SUFFIX, + _online_flow_window, _online_schedule_payload, _read_online_total_steps, _write_control, @@ -17,6 +20,44 @@ class TestResolveTotalSteps(unittest.TestCase): + @staticmethod + def _online_flow_config(*, high=1152, low=1024): + return SimpleNamespace( + runtime=SimpleNamespace( + in_flight_high_watermark=high, + in_flight_low_watermark=low, + producer_lease=8, + ), + training=SimpleNamespace(batch_size=8, accumulation_steps=32), + deployment=SimpleNamespace( + trainer=SimpleNamespace(nnodes=1, nproc_per_node=4) + ), + ) + + def test_online_flow_window_accepts_one_global_optimizer_window(self): + self.assertEqual( + _online_flow_window(self._online_flow_config()), + (1152, 1024), + ) + + def test_online_flow_window_rejects_small_watermarks_before_data_build(self): + cfg = self._online_flow_config(high=64, low=32) + with self.assertRaisesRegex(ValueError, "high watermark 64.*quantum 1024"): + _online_flow_window(cfg) + + cfg = self._online_flow_config(high=1152, low=32) + with self.assertRaisesRegex(ValueError, "low watermark 32.*quantum 1024"): + _online_flow_window(cfg) + + def test_online_flow_window_preserves_high_only_environment_override(self): + cfg = self._online_flow_config(high=64, low=32) + environment = { + "DISAGG_IN_FLIGHT_HIGH_WATERMARK": "1024", + } + with mock.patch.dict(os.environ, environment, clear=False): + os.environ.pop("DISAGG_IN_FLIGHT_LOW_WATERMARK", None) + self.assertEqual(_online_flow_window(cfg), (1024, None)) + def test_finite_data_horizon_counts_optimizer_steps(self): self.assertEqual( resolve_total_steps( @@ -94,6 +135,7 @@ def test_online_schedule_sidecar_round_trips_the_producer_horizon(self): training=SimpleNamespace( num_epochs=3, seed=17, + prompt_seed=None, batch_size=2, accumulation_steps=4, ), @@ -117,6 +159,35 @@ def test_online_schedule_sidecar_round_trips_the_producer_horizon(self): with self.assertRaisesRegex(ValueError, "does not match"): _read_online_total_steps(cfg, channel_path) + def test_online_prompt_seed_is_independent_from_model_seed(self): + cfg = SimpleNamespace( + training=SimpleNamespace( + num_epochs=3, + seed=17, + prompt_seed=5, + batch_size=2, + accumulation_steps=4, + ), + deployment=SimpleNamespace( + trainer=SimpleNamespace(nnodes=2, nproc_per_node=2) + ), + ) + payload = _online_schedule_payload(cfg, num_prompts=100) + self.assertEqual(payload["prompt_seed"], 5) + + with tempfile.TemporaryDirectory() as directory: + channel_path = f"{directory}/refs.jsonl" + _write_control( + channel_path + _ONLINE_SCHEDULE_SUFFIX, + json.dumps(payload), + ) + cfg.training.seed = 18 + self.assertEqual(_read_online_total_steps(cfg, channel_path), 9) + + cfg.training.prompt_seed = 6 + with self.assertRaisesRegex(ValueError, "does not match"): + _read_online_total_steps(cfg, channel_path) + def test_fixed_plan_rejects_partial_accumulation_before_training(self): with self.assertRaisesRegex( ValueError, "ends with incomplete gradient accumulation" diff --git a/tests/test_runtime/test_seam_fixes.py b/tests/test_runtime/test_seam_fixes.py index bcecc7469..242bffef1 100644 --- a/tests/test_runtime/test_seam_fixes.py +++ b/tests/test_runtime/test_seam_fixes.py @@ -106,6 +106,16 @@ def __init__(self): self.assertEqual(ignored, (model.lm_head, model.embed_tokens)) + def test_backend_scales_gradients_before_optimizer_step(self): + model = nn.Linear(2, 1, bias=False) + backend = FSDPTrainingBackend(ParallelConfig()) + backend.prepare_model(model, wrap=False) + model.weight.grad = torch.tensor([[4.0, 8.0]]) + + backend.scale_gradients(torch.tensor(0.25)) + + torch.testing.assert_close(model.weight.grad, torch.tensor([[1.0, 2.0]])) + class _FakeBackend(TrainingBackend): name = "fake" @@ -140,9 +150,11 @@ class _FakeDFlashModel(nn.Module): def __init__(self): super().__init__() self.w = nn.Parameter(torch.ones(1)) + self.max_valid_anchors = None - def forward(self, input_ids, hidden_states, loss_mask): + def forward(self, input_ids, hidden_states, loss_mask, max_valid_anchors=None): # mirrors OnlineDFlashModel's (loss, accuracy, metrics) contract + self.max_valid_anchors = max_valid_anchors loss = (self.w * hidden_states.float().sum()).abs() acc = torch.tensor(0.5) return loss, acc, {"accuracy_denom": loss_mask.sum()} @@ -167,6 +179,7 @@ def test_dflash_strategy_plugs_into_trainer_core(self): rep = core.train_step(batch) self.assertTrue(rep.optimizer_stepped) # optimizer stepped via the shared core self.assertEqual(backend.steps, 1) + self.assertEqual(model.max_valid_anchors, 3) self.assertAlmostEqual(rep.metrics["acc"], 0.5) out = strat.forward_loss(batch) self.assertAlmostEqual(float(out.metrics["accuracy"]), 0.5) diff --git a/tests/test_runtime/test_server_capture.py b/tests/test_runtime/test_server_capture.py index 3b6a41635..412f5a927 100644 --- a/tests/test_runtime/test_server_capture.py +++ b/tests/test_runtime/test_server_capture.py @@ -322,7 +322,7 @@ def recording_server(url, json_body, timeout): request["multi_modal_data"], ) self.assertEqual( - {"temperature": 0.0, "max_new_tokens": 1}, + {"temperature": 0.0, "max_new_tokens": 0}, request["sampling_params"], ) self.assertEqual(2, len(request["spec_capture"])) diff --git a/tests/test_runtime/test_tracking_logger.py b/tests/test_runtime/test_tracking_logger.py index 9e5da434c..97d09ffd5 100644 --- a/tests/test_runtime/test_tracking_logger.py +++ b/tests/test_runtime/test_tracking_logger.py @@ -2,10 +2,11 @@ """Backend-neutral experiment tracking at the Trainer logger seam.""" import unittest +from tempfile import TemporaryDirectory from types import SimpleNamespace from unittest import mock -from specforge.tracker import _public_config +from specforge.tracker import WandbTracker, _public_config from specforge.training.tracking import ( TrackerLogger, create_tracker_logger, @@ -73,11 +74,31 @@ def test_tracker_metadata_redacts_credentials(self): wandb_key="secret", auth_token="also-secret", wandb_project="specforge", + specforge_config={ + "tracking": {"wandb_key": "nested-secret"}, + "model": { + "embedding_key": "language_model.embed_tokens.weight", + "lm_head_key": "language_model.lm_head.weight", + }, + "data": {"cache_key": "reproduction-cache"}, + }, ) ) self.assertEqual(config["wandb_key"], "") self.assertEqual(config["auth_token"], "") self.assertEqual(config["wandb_project"], "specforge") + self.assertEqual( + config["specforge_config"]["tracking"]["wandb_key"], + "", + ) + self.assertEqual( + config["specforge_config"]["model"]["embedding_key"], + "language_model.embed_tokens.weight", + ) + self.assertEqual( + config["specforge_config"]["data"]["cache_key"], + "reproduction-cache", + ) def test_normalizes_scalars_and_expands_vectors(self): self.assertEqual( @@ -136,6 +157,38 @@ def test_factory_adapts_existing_tracker_registry(self): logger({"loss": 2.0}, 3) self.assertEqual(tracker.logged, [({"train/loss": 2.0}, 3)]) + def test_wandb_tracker_logs_through_its_owned_run_handle(self): + run = mock.Mock() + wandb = mock.Mock() + wandb.init.return_value = run + args = SimpleNamespace( + wandb_dir=None, + wandb_offline=True, + wandb_key=None, + wandb_project="specforge", + wandb_name="disaggregated-trainer", + ) + + with ( + TemporaryDirectory() as output_dir, + mock.patch("specforge.tracker.wandb", wandb), + mock.patch("specforge.tracker.dist.is_available", return_value=True), + mock.patch("specforge.tracker.dist.is_initialized", return_value=True), + mock.patch("specforge.tracker.dist.get_rank", return_value=0), + ): + tracker = WandbTracker(args, output_dir) + # A later multiprocessing lifecycle may clear W&B's module-global + # current run. The tracker-owned handle must remain authoritative. + wandb.run = None + tracker.log({"train/loss": 1.25}, step=10) + tracker.close() + tracker.close() + + run.log.assert_called_once_with({"train/loss": 1.25}, step=10, commit=True) + run.finish.assert_called_once_with() + wandb.log.assert_not_called() + wandb.finish.assert_not_called() + def test_noop_tracker_does_not_require_initialized_distributed(self): logger = create_tracker_logger(SimpleNamespace(report_to="none"), "/tmp/output") logger({"loss": 1.0}, 1) diff --git a/tests/test_runtime/test_trainer.py b/tests/test_runtime/test_trainer.py index efc8ebcfc..e8f3a64a4 100644 --- a/tests/test_runtime/test_trainer.py +++ b/tests/test_runtime/test_trainer.py @@ -22,6 +22,7 @@ Checkpoint, TrainerController, TrainerCore, + _materialize_metrics, _reduce_ratio_metrics, ) from specforge.training.strategies.base import DraftTrainStrategy, StepOutput @@ -61,6 +62,23 @@ def forward_loss(self, batch, ctx=None): return super().forward_loss(batch, ctx) +class WeightedStrategy(FakeStrategy): + def forward_loss(self, batch, ctx=None): + self.validate_batch(batch) + coefficient = batch.tensors["x"].reshape(()) + denominator = batch.tensors["denominator"].reshape(()) + correct = batch.tensors["correct"].reshape(()) + numerator = self.model.w.reshape(()) * coefficient + return StepOutput( + loss=numerator / denominator, + metrics={"accuracy": correct / denominator}, + ratio_metrics={ + "acc": (correct, denominator), + }, + loss_terms=(numerator, denominator), + ) + + class FakeBackend(TrainingBackend): name = "fake" @@ -78,6 +96,11 @@ def backward(self, loss, *, is_boundary=True): self.boundaries.append(is_boundary) loss.backward() + def scale_gradients(self, factor): + for parameter in self.model.parameters(): + if parameter.grad is not None: + parameter.grad.mul_(factor) + def step(self): self.steps += 1 return torch.tensor(1.0) @@ -103,6 +126,19 @@ def _batch(): ) +def _weighted_batch(coefficient, denominator, correct): + return TrainBatch( + sample_ids=["s"], + strategy="fake", + tensors={ + "x": torch.tensor(float(coefficient)), + "denominator": torch.tensor(float(denominator)), + "correct": torch.tensor(float(correct)), + }, + metadata={}, + ) + + class TestTrainerCore(unittest.TestCase): def test_accumulation_boundary(self): strat = FakeStrategy() @@ -148,15 +184,70 @@ def test_metrics_carry_no_mode(self): rep = core.train_step(_batch()) self.assertNotIn("mode", rep.metrics) + def test_global_loss_normalization_matches_combined_batch(self): + split_strategy = WeightedStrategy() + split_core = TrainerCore( + split_strategy, + FakeBackend(split_strategy.model), + accumulation_steps=2, + ) + split_core.train_step(_weighted_batch(2, 1, 1)) + split_result = split_core.train_step(_weighted_batch(30, 3, 1)) + + combined_strategy = WeightedStrategy() + combined_core = TrainerCore( + combined_strategy, + FakeBackend(combined_strategy.model), + ) + combined_result = combined_core.train_step(_weighted_batch(32, 4, 2)) + + torch.testing.assert_close( + split_strategy.model.w.grad, + combined_strategy.model.w.grad, + ) + self.assertEqual(split_strategy.model.w.grad.item(), 8.0) + self.assertEqual(split_result.loss, combined_result.loss) + self.assertEqual(split_result.metrics["acc"], 0.5) + self.assertEqual(combined_result.metrics["acc"], 0.5) + + def test_global_loss_normalization_compensates_rank_averaging(self): + strategy = FakeStrategy() + backend = FakeBackend(strategy.model) + backend.parallel_config = mock.Mock(fsdp_process_group="dp") + core = TrainerCore(strategy, backend) + strategy.model.w.grad = torch.tensor([9.0]) + + def add_remote_denominator(denominator, *, op, group): + self.assertEqual(op, torch.distributed.ReduceOp.SUM) + self.assertEqual(group, "dp") + denominator.add_(7.0) + + with ( + mock.patch("torch.distributed.is_available", return_value=True), + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.distributed.get_world_size", return_value=2), + mock.patch( + "torch.distributed.all_reduce", + side_effect=add_remote_denominator, + ), + ): + core._normalize_gradients(torch.tensor(3.0)) + + torch.testing.assert_close( + strategy.model.w.grad, + torch.tensor([1.8]), + ) + def test_strategy_scalar_metrics_are_preserved(self): strat = FakeStrategy() core = TrainerCore(strat, FakeBackend(strat.model), accumulation_steps=1) result = core._result( StepOutput( - loss=torch.tensor(2.0), + loss=torch.tensor(2.0, requires_grad=True), metrics={ "loss": torch.tensor(99.0), "accuracy": torch.tensor(0.5), + "accuracy_denom": torch.tensor(4.0), "ce_loss": torch.tensor(1.25), "lambda_base": 0.75, "non_scalar_debug": torch.tensor([1.0, 2.0]), @@ -170,8 +261,28 @@ def test_strategy_scalar_metrics_are_preserved(self): self.assertEqual(result.metrics["acc"], 0.5) self.assertEqual(result.metrics["ce_loss"], 1.25) self.assertEqual(result.metrics["lambda_base"], 0.75) + self.assertIsInstance(result._metric_values["lambda_base"], float) + self.assertFalse(result._metric_values["loss"].requires_grad) + self.assertNotIn("accuracy_denom", result.metrics) self.assertNotIn("non_scalar_debug", result.metrics) + def test_metrics_are_materialized_lazily_and_cached(self): + strat = FakeStrategy() + core = TrainerCore(strat, FakeBackend(strat.model), accumulation_steps=1) + + with mock.patch( + "specforge.training.controller._materialize_metrics", + wraps=_materialize_metrics, + ) as materialize: + result = core.train_step(_batch()) + materialize.assert_not_called() + + self.assertEqual(result.loss, 2.0) + materialize.assert_called_once() + self.assertEqual(result.metrics["acc"], 0.5) + self.assertEqual(result.grad_norm, 1.0) + materialize.assert_called_once() + def test_ratio_metrics_override_mean_of_means_accuracy(self): strat = FakeStrategy() core = TrainerCore(strat, FakeBackend(strat.model), accumulation_steps=1) @@ -195,35 +306,59 @@ def test_ratio_metrics_override_mean_of_means_accuracy(self): self.assertEqual(result.metrics["ce_position_0"], 0.5) self.assertEqual(result.metrics["ce_position_1"], 0.5) + _RATIO_INPUTS = { + "acc": (torch.tensor(2.0), torch.tensor(4.0)), + "ce_position": ( + torch.tensor([1.0, 3.0]), + torch.tensor([2.0, 6.0]), + ), + } + def test_ratio_metrics_sum_numerators_and_denominators_before_dividing(self): remote = torch.tensor([8.0, 6.0, 3.0, 1.0, 2.0, 2.0]) + reduced_groups = [] def all_reduce(packed, *, group): - self.assertEqual(group, "dp") + reduced_groups.append(group) packed.add_(remote) with ( mock.patch("torch.distributed.is_available", return_value=True), mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.distributed.get_world_size", return_value=2), mock.patch("torch.distributed.all_reduce", side_effect=all_reduce), ): metrics = _reduce_ratio_metrics( - { - "acc": (torch.tensor(2.0), torch.tensor(4.0)), - "ce_position": ( - torch.tensor([1.0, 3.0]), - torch.tensor([2.0, 6.0]), - ), - }, + self._RATIO_INPUTS, device=torch.device("cpu"), process_group="dp", reduce=True, ) + self.assertEqual(reduced_groups, ["dp"]) self.assertEqual(metrics["acc"], 1.0) self.assertEqual(metrics["ce_position_0"], 1.0) self.assertEqual(metrics["ce_position_1"], 0.5) + def test_ratio_metrics_skip_all_reduce_when_world_size_is_one(self): + with ( + mock.patch("torch.distributed.is_available", return_value=True), + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.distributed.get_world_size", return_value=1), + mock.patch("torch.distributed.all_reduce") as all_reduce_mock, + ): + metrics = _reduce_ratio_metrics( + self._RATIO_INPUTS, + device=torch.device("cpu"), + process_group="dp", + reduce=True, + ) + + all_reduce_mock.assert_not_called() + self.assertEqual(metrics["acc"], 0.5) + self.assertEqual(metrics["ce_position_0"], 0.5) + self.assertEqual(metrics["ce_position_1"], 0.5) + @staticmethod def _eagle_output( *, @@ -305,7 +440,7 @@ def all_reduce(value, *, op, group): "torch.distributed.all_reduce", side_effect=all_reduce ) as reduce, ): - result = core._result(self._eagle_output(), grad_norm=None, stepped=False) + result = core._result(self._eagle_output(), grad_norm=None, stepped=True) reduce.assert_called_once() self.assertAlmostEqual(result.metrics["acc_0"], 9 / 12, places=6) @@ -327,6 +462,47 @@ def test_validate_batch_missing_feature(self): class TestTrainerController(unittest.TestCase): + def test_training_log_reports_pipeline_throughput_breakdown(self): + strat = FakeStrategy() + backend = FakeBackend(strat.model) + core = TrainerCore(strat, backend, accumulation_steps=1) + logged = [] + with ( + tempfile.TemporaryDirectory() as d, + mock.patch( + "specforge.training.controller._materialize_metrics", + wraps=_materialize_metrics, + ) as materialize, + ): + ctrl = TrainerController( + core, + run_id="r", + output_dir=d, + max_steps=2, + num_epochs=1, + log_interval=2, + logger=lambda metrics, step: logged.append((dict(metrics), step)), + ) + self.assertEqual(ctrl.fit([_batch(), _batch()]), 2) + + materialize.assert_called_once() + + self.assertEqual(len(logged), 1) + metrics, step = logged[0] + self.assertEqual(step, 2) + for name in ( + "perf/optimizer_steps_per_hour", + "perf/optimizer_step_time_s", + "perf/data_wait_time_s", + "perf/train_compute_time_s", + "perf/durable_ack_time_s", + "perf/global_samples_per_second", + ): + self.assertIn(name, metrics) + self.assertGreaterEqual(metrics[name], 0.0) + self.assertGreater(metrics["perf/optimizer_steps_per_hour"], 0.0) + self.assertGreater(metrics["perf/global_samples_per_second"], 0.0) + def test_progress_bar_tracks_optimizer_steps_on_rank_zero(self): strat = FakeStrategy() backend = FakeBackend(strat.model) diff --git a/tests/test_scripts/test_disagg_launchers.py b/tests/test_scripts/test_disagg_launchers.py index fa06f9be1..bdaaec6dd 100644 --- a/tests/test_scripts/test_disagg_launchers.py +++ b/tests/test_scripts/test_disagg_launchers.py @@ -13,6 +13,10 @@ OFFLINE = ROOT / "examples" / "disagg" / "run_offline.sh" OFFLINE_TWO_NODE = ROOT / "examples" / "disagg" / "run_offline_2node.sh" TWO_NODE = ROOT / "examples" / "disagg" / "run_qwen3_8b_dflash_disagg_2node.sh" +INKLING_TWO_NODE = ROOT / "examples" / "disagg" / "run_inkling_dspark_disagg_2node.sh" +KIMI_K3_CAPTURE_PATCH = ( + ROOT / "patches" / "sglang" / "kimi-k3-f8493a4" / "spec-capture.patch" +) class DisaggregatedWrapperTest(unittest.TestCase): @@ -123,6 +127,23 @@ def test_help_describes_auto_and_explicit_roles(self): self.assertIn("--role", result.stdout) self.assertIn("producer and consumer", result.stdout) + def test_kimi_k3_capture_patch_only_copies_features_on_the_writer_rank(self): + source = KIMI_K3_CAPTURE_PATCH.read_text(encoding="utf-8") + self.assertIn("self.output_streamer.ps.attn_tp_rank != 0", source) + self.assertNotIn(".cpu().clone()", source) + self.assertIn('getattr(logits_output, "_spec_capture_aux_cpu", None)', source) + self.assertIn("logits_output.hidden_states.cpu()", source) + self.assertIn('"aux" in features', source) + self.assertIn('"last_hidden" in features', source) + self.assertIn("_should_copy_hidden_states_to_cpu", source) + self.assertIn("self.ps.attn_tp_rank == 0", source) + self.assertIn("self.logits_output.last_hidden_states = _async_d2h(", source) + self.assertIn("len(chunks) == 1", source) + self.assertIn("ThreadPoolExecutor(", source) + self.assertIn('getattr(store, "batch_put_from", None)', source) + self.assertIn("SGLANG_SPEC_CAPTURE_MAX_PENDING_BATCHES", source) + self.assertIn("req.finished() and req.spec_capture_result is None", source) + def test_two_node_wrapper_keeps_training_on_the_unified_cli(self): self.assertTrue(os.access(TWO_NODE, os.X_OK)) syntax = subprocess.run( @@ -179,6 +200,54 @@ def test_two_node_wrapper_keeps_training_on_the_unified_cli(self): self.assertNotIn("torchrun", "".join(outputs.values())) self.assertFalse(shared_root.exists()) + def test_inkling_two_node_wrapper_pins_the_validated_server_contract(self): + self.assertTrue(os.access(INKLING_TWO_NODE, os.X_OK)) + syntax = subprocess.run( + ["bash", "-n", str(INKLING_TWO_NODE)], + text=True, + capture_output=True, + check=False, + ) + self.assertEqual(syntax.returncode, 0, syntax.stderr) + + env = self._env() + env.update( + { + "NODE_RANK": "0", + "NUM_NODES": "2", + "HEAD_IP": "10.0.0.1", + "DISAGG_STORE_ID": "inkling-two-node-test", + "DISAGG_RUN_ROOT": str(self.root / "inkling-shared-attempt"), + "DRY_RUN": "1", + } + ) + result = subprocess.run( + [str(INKLING_TWO_NODE), "training.max_steps=1"], + cwd=ROOT, + env=env, + text=True, + capture_output=True, + check=False, + ) + self.assertEqual(result.returncode, 0, result.stderr) + output = result.stdout + for expected in ( + "thinkingmachines/Inkling", + "--tp-size 4", + "--spec-capture-aux-layer-ids 5 17 35 47 59", + "--attention-backend fa4", + "--quantization modelopt_fp4", + "--mamba-radix-cache-strategy extra_buffer", + "training.accumulation_steps=128", + ): + with self.subTest(expected=expected): + self.assertIn(expected, output) + self.assertNotIn("--disable-radix-cache", output) + + source = INKLING_TWO_NODE.read_text(encoding="utf-8") + self.assertIn("SGLANG_ENABLE_UNIFIED_RADIX_TREE", source) + self.assertIn("SGLANG_OPT_USE_INKLING_CUSTOM_AR", source) + def test_offline_two_node_wrapper_dispatches_roles_to_the_unified_cli(self): self.assertTrue(os.access(OFFLINE_TWO_NODE, os.X_OK)) syntax = subprocess.run( diff --git a/tests/test_scripts/test_expand_reasoning_conversations.py b/tests/test_scripts/test_expand_reasoning_conversations.py index 5a687f326..d5665658a 100644 --- a/tests/test_scripts/test_expand_reasoning_conversations.py +++ b/tests/test_scripts/test_expand_reasoning_conversations.py @@ -34,7 +34,7 @@ def _load_preprocessing_stack(): distributed_module.get_sp_ring_group = lambda: None sys.modules[f"{package_name}.distributed"] = distributed_module - for module_name in ("template", "parse", "preprocessing"): + for module_name in ("template", "parse", "loss_mask", "preprocessing"): full_name = f"{data_package_name}.{module_name}" module_path = repo_root / "specforge" / "data" / f"{module_name}.py" spec = importlib.util.spec_from_file_location(full_name, module_path) diff --git a/tests/test_scripts/test_prepare_hidden_states.py b/tests/test_scripts/test_prepare_hidden_states.py index d30738054..13257f77a 100644 --- a/tests/test_scripts/test_prepare_hidden_states.py +++ b/tests/test_scripts/test_prepare_hidden_states.py @@ -1,5 +1,4 @@ import gzip -import json import tempfile import unittest from pathlib import Path @@ -35,10 +34,25 @@ def test_cli_defaults_to_legacy_eagle3_capture(self): self.assertEqual("eagle3", args.strategy) self.assertIsNone(args.draft_model_config) + self.assertFalse(args.sglang_disable_radix_cache) self.assertFalse(hasattr(args, "draft_num_hidden_layers")) self.assertFalse(hasattr(args, "draft_block_size")) self.assertFalse(hasattr(args, "capture_layers")) + def test_cli_can_explicitly_disable_radix_cache(self): + argv = [ + "prepare_hidden_states.py", + "--target-model-path", + "target", + "--data-path", + "data.jsonl", + "--sglang-disable-radix-cache", + ] + with mock.patch("sys.argv", argv): + args = parse_args() + + self.assertTrue(args.sglang_disable_radix_cache) + def test_cli_accepts_dflash_family_config(self): argv = [ "prepare_hidden_states.py", @@ -109,12 +123,18 @@ def test_strategy_capture_plans_use_draft_owned_layers_and_schemas(self): with self.subTest(strategy=strategy): plan = resolve_offline_capture_plan(args, target_config) self.assertEqual(strategy, plan.strategy) - self.assertEqual( - "eagle3" if strategy == "eagle3" else "dflash", - plan.capture_method, - ) + expected_capture_method = { + "eagle3": "eagle3", + "dspark": "dspark", + }.get(strategy, "dflash") + self.assertEqual(expected_capture_method, plan.capture_method) self.assertEqual(layers, plan.capture_layers) self.assertEqual(feature_names, set(plan.layout.output_names)) + if strategy == "eagle3": + self.assertIsNone(plan.loss_mask_filter) + else: + self.assertTrue(plan.loss_mask_filter([0, 1, 1])) + self.assertFalse(plan.loss_mask_filter([1, 0, 1])) def test_build_uses_dedicated_offline_loader(self): config = SimpleNamespace(num_hidden_layers=32, dtype=None) @@ -130,6 +150,7 @@ def test_build_uses_dedicated_offline_loader(self): sglang_enable_dp_attention=False, sglang_enable_dp_lm_head=False, sglang_ep_size=1, + sglang_disable_radix_cache=False, batch_size=4, max_length=128, ) @@ -147,6 +168,7 @@ def test_build_uses_dedicated_offline_loader(self): self.assertEqual(load.call_args.args, ("target",)) self.assertNotIn("device", load.call_args.kwargs) self.assertNotIn("cache_dir", load.call_args.kwargs) + self.assertFalse(load.call_args.kwargs["disable_radix_cache"]) target.set_capture_layers.assert_called_once_with( [2, 7, 19], capture_method="eagle3", @@ -166,6 +188,7 @@ def test_build_accepts_pre_resolved_arbitrary_capture_layers(self): sglang_enable_dp_attention=False, sglang_enable_dp_lm_head=False, sglang_ep_size=1, + sglang_disable_radix_cache=True, batch_size=4, max_length=128, ) @@ -174,7 +197,7 @@ def test_build_accepts_pre_resolved_arbitrary_capture_layers(self): with mock.patch( "scripts.prepare_hidden_states.load_offline_capture", return_value=target, - ): + ) as load: self.assertIs( build_target_model( args, @@ -185,6 +208,7 @@ def test_build_accepts_pre_resolved_arbitrary_capture_layers(self): target, ) + self.assertTrue(load.call_args.kwargs["disable_radix_cache"]) target.set_capture_layers.assert_called_once_with( capture_layers, capture_method="dflash", @@ -269,29 +293,24 @@ def test_nan_in_any_strategy_mapped_tensor_skips_the_record(self): class PrepareHiddenStatesVocabMappingTest(unittest.TestCase): - def test_resolves_draft_vocab_size_from_local_json_file(self): - with tempfile.TemporaryDirectory() as directory: - config_path = Path(directory) / "config.json" - config_path.write_text( - json.dumps({"vocab_size": 16, "draft_vocab_size": 8}), - encoding="utf-8", - ) - - self.assertEqual(_resolve_draft_vocab_size(str(config_path)), 8) - - def test_rejects_directory_and_hugging_face_repo_id(self): - with tempfile.TemporaryDirectory() as directory: - with self.assertRaisesRegex(FileNotFoundError, "local JSON file"): - _resolve_draft_vocab_size(directory) - with self.assertRaisesRegex(FileNotFoundError, "local JSON file"): - _resolve_draft_vocab_size("org/draft-model") + def test_prefers_resolved_draft_vocab_size(self): + config = SimpleNamespace(vocab_size=16, draft_vocab_size=8) + self.assertEqual(_resolve_draft_vocab_size(config), 8) def test_resolves_vocab_size_fallback(self): - with tempfile.TemporaryDirectory() as directory: - config_path = Path(directory) / "draft.json" - config_path.write_text(json.dumps({"vocab_size": 16}), encoding="utf-8") - - self.assertEqual(_resolve_draft_vocab_size(str(config_path)), 16) + self.assertEqual(_resolve_draft_vocab_size(SimpleNamespace(vocab_size=16)), 16) + + def test_does_not_fallback_when_draft_vocab_size_is_explicitly_none(self): + config = SimpleNamespace(vocab_size=16, draft_vocab_size=None) + with self.assertRaisesRegex(ValueError, "positive"): + _resolve_draft_vocab_size(config) + + def test_rejects_invalid_resolved_vocab_size(self): + for value in (None, True, 0, -1, "16"): + with self.subTest(value=value): + config = SimpleNamespace(draft_vocab_size=value) + with self.assertRaisesRegex(ValueError, "positive"): + _resolve_draft_vocab_size(config) def test_global_rank_zero_generates_fixed_mapping_path(self): with tempfile.TemporaryDirectory() as directory: diff --git a/tests/test_scripts/test_sync_distributed_checkpoints.py b/tests/test_scripts/test_sync_distributed_checkpoints.py new file mode 100644 index 000000000..b6a5b067f --- /dev/null +++ b/tests/test_scripts/test_sync_distributed_checkpoints.py @@ -0,0 +1,175 @@ +"""Dependency-light tests for the non-shared-filesystem checkpoint relay.""" + +from __future__ import annotations + +import importlib.util +import json +import os +import shutil +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace + +ROOT = Path(__file__).resolve().parents[2] +SCRIPT = ROOT / "examples" / "disagg" / "sync_distributed_checkpoints.py" +SPEC = importlib.util.spec_from_file_location("checkpoint_relay_example", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +RELAY_MODULE = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(RELAY_MODULE) + + +class DistributedCheckpointRelayTest(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory(prefix="checkpoint_relay_") + self.root = Path(self._tmp.name) + self.run_id = "relay-test" + self.relays = [] + + def tearDown(self): + for relay in self.relays: + relay._httpd.server_close() + self._tmp.cleanup() + + def _checkpoint(self, root: Path, step: int, ranks: range) -> Path: + checkpoint = root / "output" / f"{self.run_id}-step{step}" + checkpoint.mkdir(parents=True) + if 0 in ranks: + (checkpoint / "training_state.pt").write_bytes( + f"shared-step-{step}".encode() + ) + for rank in ranks: + (checkpoint / f"training_state_rank{rank}.pt").write_bytes( + f"rank-{rank}-step-{step}".encode() + ) + return checkpoint + + def _relay( + self, + root: Path, + *, + local_ranks: range, + peer_ranks: range, + max_archives: int = 2, + ): + relay = RELAY_MODULE.CheckpointRelay( + SimpleNamespace( + run_root=str(root), + run_id=self.run_id, + local_ranks=tuple(local_ranks), + peer_ranks=tuple(peer_ranks), + peer_url="file:///not-configured", + poll_s=0.01, + max_archives=max_archives, + serve_host="127.0.0.1", + serve_port=0, + ) + ) + self.relays.append(relay) + return relay + + def test_two_nodes_assemble_complete_checkpoints_and_bound_archives(self): + node0 = self.root / "node0" + node1 = self.root / "node1" + for step in (1, 2, 3): + self._checkpoint(node0, step, range(0, 2)) + self._checkpoint(node1, step, range(2, 4)) + + relay0 = self._relay(node0, local_ranks=range(0, 2), peer_ranks=range(2, 4)) + relay1 = self._relay(node1, local_ranks=range(2, 4), peer_ranks=range(0, 2)) + relay0.peer_url = relay1.relay_dir.as_uri() + relay1.peer_url = relay0.relay_dir.as_uri() + + relay0._publish_local() + relay1._publish_local() + relay0._pull_peer() + relay1._pull_peer() + + for root in (node0, node1): + for step in (2, 3): + checkpoint = root / "output" / f"{self.run_id}-step{step}" + expected = {"training_state.pt"} + expected.update(f"training_state_rank{rank}.pt" for rank in range(4)) + self.assertTrue( + expected.issubset(path.name for path in checkpoint.iterdir()) + ) + + for relay in (relay0, relay1): + manifest = json.loads( + (relay.relay_dir / "manifest.json").read_text(encoding="utf-8") + ) + self.assertEqual([entry["step"] for entry in manifest["entries"]], [2, 3]) + local_archives = [ + path + for path in relay.relay_dir.glob("*.tar") + if not path.name.startswith("peer-") + ] + peer_archives = list(relay.relay_dir.glob("peer-*.tar")) + self.assertEqual(len(local_archives), 2) + self.assertEqual(len(peer_archives), 2) + + self._checkpoint(node0, 4, range(0, 2)) + self._checkpoint(node1, 4, range(2, 4)) + relay0._publish_local() + relay1._publish_local() + relay0._pull_peer() + relay1._pull_peer() + + for relay in (relay0, relay1): + names = {path.name for path in relay.relay_dir.glob("*.tar")} + self.assertFalse(any("step2-" in name for name in names)) + self.assertTrue(any("step3-" in name for name in names)) + self.assertTrue(any("step4-" in name for name in names)) + + # A transiently absent or already-pruned output tree must not erase the + # relay's bounded recovery copies. + shutil.rmtree(node0 / "output") + shutil.rmtree(node1 / "output") + relay0._publish_local() + relay1._publish_local() + for relay in (relay0, relay1): + manifest = json.loads( + (relay.relay_dir / "manifest.json").read_text(encoding="utf-8") + ) + self.assertEqual([entry["step"] for entry in manifest["entries"]], [3, 4]) + + def test_rank_ranges_and_archive_retention_are_validated(self): + self.assertTrue(os.access(SCRIPT, os.X_OK)) + self.assertEqual(RELAY_MODULE._rank_range("8-15"), tuple(range(8, 16))) + with self.assertRaisesRegex(Exception, "rank range"): + RELAY_MODULE._rank_range("15-8") + self.assertEqual(RELAY_MODULE._positive_int("3"), 3) + with self.assertRaisesRegex(Exception, "at least 1"): + RELAY_MODULE._positive_int("0") + + def test_peer_archive_name_cannot_escape_or_disagree_with_step(self): + relay = self._relay( + self.root / "node0", + local_ranks=range(0, 2), + peer_ranks=range(2, 4), + ) + base_entry = { + "step": 3, + "sha256": "0" * 64, + "files": [ + "training_state_rank2.pt", + "training_state_rank3.pt", + ], + } + with self.assertRaisesRegex(ValueError, "unexpected peer archive"): + relay._install_peer_archive( + {**base_entry, "archive": "../outside-step3-ranks2-3.tar"} + ) + with self.assertRaisesRegex(ValueError, "unexpected peer archive"): + relay._install_peer_archive( + { + **base_entry, + "archive": f"{self.run_id}-step4-ranks2-3.tar", + } + ) + with self.assertRaisesRegex(ValueError, "unexpected peer archive name"): + relay._download("../outside.tar", "0" * 64) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_utils/test_dflash_losses.py b/tests/test_utils/test_dflash_losses.py index 3a0457044..6001ce336 100644 --- a/tests/test_utils/test_dflash_losses.py +++ b/tests/test_utils/test_dflash_losses.py @@ -64,10 +64,15 @@ class _DFlashDraftStub(nn.Module): OnlineDSparkModel = _dflash_module.OnlineDSparkModel +def _anchor_sampler_subject(num_anchors: int = 8): + return types.SimpleNamespace(num_anchors=num_anchors) + + class _FixedDraft(nn.Module): def __init__(self, hidden_size: int): super().__init__() self.hidden_size = hidden_size + self.sliding_window = None def forward(self, position_ids, noise_embedding, target_hidden, attention_mask): bsz, draft_len = noise_embedding.shape[:2] @@ -166,7 +171,7 @@ def _fixed_noise_embed(self, input_ids, anchor_positions, block_keep_mask): def _fixed_anchor_sampler(anchors, keep_mask): - def _sample(self, seq_len, loss_mask, device): + def _sample(self, seq_len, loss_mask, device, max_valid_anchors=None): return anchors.to(device), keep_mask.to(device) return _sample @@ -317,7 +322,7 @@ def _naive_dflash_loss(neg_log_q, binary_mask, gamma): positions = torch.arange(block_size, dtype=neg_log_q.dtype).view(1, 1, -1) decay = torch.exp(-(positions - 1).clamp(min=0) / gamma) weight = weight * decay - return (neg_log_q * weight).sum() / (weight.sum() + 1e-6) + return (neg_log_q * weight).sum() / weight.sum() class TestDFlashLosses(unittest.TestCase): @@ -362,6 +367,49 @@ def test_dflash_decay_gamma_is_preserved(self): want = _naive_dflash_loss(self.neg_log_q, self.binary_mask, gamma=gamma) torch.testing.assert_close(got, want, rtol=0, atol=1e-8) + def test_dflash_exposes_additive_loss_and_accuracy_terms(self): + head = nn.Linear(4, self.logits.shape[-1], bias=False).double() + model = _make_model( + self.logits, + self.anchors, + self.keep_mask, + draft_model=_LearnableDSparkDraft(4).double(), + lm_head=head, + ) + + loss, accuracy, metrics = model( + input_ids=self.input_ids, + hidden_states=self.hidden_states, + loss_mask=self.loss_mask, + ) + + loss_num, loss_den = metrics["loss_terms"] + self.assertTrue(loss_num.requires_grad) + torch.testing.assert_close(loss, loss_num / loss_den) + accuracy_num, accuracy_den = metrics["ratio_metrics"]["acc"] + torch.testing.assert_close(accuracy, accuracy_num / accuracy_den) + + def test_dflash_partial_tail_has_one_finite_target(self): + vocab_size = 7 + logits = torch.randn(1, 1, 5, vocab_size, dtype=torch.double) + input_ids = torch.tensor([[1, 2, 3, 4]]) + loss_mask = torch.ones_like(input_ids, dtype=torch.double) + model = _make_model( + logits, + anchors=torch.tensor([[2]]), + keep_mask=torch.tensor([[True]]), + ) + + loss, _accuracy, metrics = model( + input_ids=input_ids, + hidden_states=torch.zeros(1, 4, 4, dtype=torch.double), + loss_mask=loss_mask, + ) + + expected = F.cross_entropy(logits[0, 0, 1].unsqueeze(0), input_ids[:, 3]) + torch.testing.assert_close(loss, expected) + torch.testing.assert_close(metrics["loss_terms"][1], loss.new_tensor(1.0)) + def test_dpace_full_matches_naive_reference(self): alpha = 0.5 got = self._forward_loss(loss_type="dpace", dpace_alpha=alpha) @@ -405,12 +453,28 @@ def test_continuation_value_ablation_matches_naive_reference(self): def test_dpace_loss_reduces_by_batch_size(self): alpha = 0.5 - got = self._forward_loss(loss_type="dpace", dpace_alpha=alpha) + model = _make_model( + self.logits, + self.anchors, + self.keep_mask, + loss_type="dpace", + dpace_alpha=alpha, + ) + got, _accuracy, metrics = model( + input_ids=self.input_ids, + hidden_states=self.hidden_states, + loss_mask=self.loss_mask, + ) weight = _naive_dpace_weight(self.q, self.binary_mask, alpha, "dpace") weighted_sum = (self.neg_log_q * weight * self.binary_mask).sum() token_count_loss = weighted_sum / ((weight * self.binary_mask).sum() + 1e-6) batch_loss = weighted_sum / float(self.input_ids.shape[0]) torch.testing.assert_close(got, batch_loss, rtol=0, atol=1e-10) + torch.testing.assert_close(metrics["loss_terms"][0], weighted_sum) + torch.testing.assert_close( + metrics["loss_terms"][1], + got.new_tensor(float(self.input_ids.shape[0])), + ) self.assertFalse(torch.allclose(got, token_count_loss)) def test_alpha_changes_dpace_loss(self): @@ -773,6 +837,59 @@ def test_dspark_sampler_keeps_sparse_high_index_anchor(self): self.assertEqual(anchors[0, 0].item(), 4) self.assertEqual(keep[0].tolist(), [True, False]) + def test_shared_sampler_uses_adjacent_targets_and_partial_tails(self): + sampler = OnlineDFlashModel._sample_anchor_positions + model = _anchor_sampler_subject() + loss_mask = torch.tensor([[1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0]]) + + anchors, keep = sampler( + model, + seq_len=loss_mask.shape[1], + loss_mask=loss_mask, + device=loss_mask.device, + ) + + self.assertEqual(anchors[keep].tolist(), [0, 5]) + + def test_shared_sampler_is_batch_padding_invariant(self): + sampler = OnlineDFlashModel._sample_anchor_positions + model = _anchor_sampler_subject() + short_mask = torch.tensor([[0.0, 0.0, 1.0, 1.0]]) + short_anchors, short_keep = sampler( + model, + seq_len=4, + loss_mask=short_mask, + device=short_mask.device, + ) + padded_batch = torch.tensor( + [ + [0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0], + [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + [1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0], + ] + ) + batch_anchors, batch_keep = sampler( + model, + seq_len=7, + loss_mask=padded_batch, + device=padded_batch.device, + ) + + self.assertEqual(short_anchors[short_keep].tolist(), [2]) + self.assertEqual(batch_anchors[0][batch_keep[0]].tolist(), [2]) + self.assertFalse(batch_keep[2].any()) + + def test_shared_sampler_rejects_a_batch_without_adjacent_targets(self): + model = _anchor_sampler_subject() + loss_mask = torch.tensor([[1.0, 0.0, 1.0]]) + with self.assertRaisesRegex(ValueError, "two consecutive"): + OnlineDFlashModel._sample_anchor_positions( + model, + seq_len=loss_mask.shape[1], + loss_mask=loss_mask, + device=loss_mask.device, + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_utils/test_dflash_mask.py b/tests/test_utils/test_dflash_mask.py index 8db5c732c..ab46def04 100644 --- a/tests/test_utils/test_dflash_mask.py +++ b/tests/test_utils/test_dflash_mask.py @@ -1,15 +1,28 @@ import unittest +from types import SimpleNamespace +from unittest import mock import torch +from torch import nn from specforge.algorithms.common.dflash_family_model import ( + OnlineDFlashModel, + OnlineDominoModel, + OnlineDSparkModel, create_dflash_block_mask, create_dflash_sdpa_mask, ) -def _reference_dflash_mask(anchor_positions, block_keep_mask, S, block_size, device): - """Element-level reference mask mirroring the mask_mod inside create_dflash_block_mask. +def _reference_dflash_mask( + anchor_positions, + block_keep_mask, + S, + block_size, + device, + sliding_window=None, +): + """Element-level reference for full and sliding DFlash attention. This uses plain Python loops so correctness is obvious by inspection. """ @@ -21,30 +34,61 @@ def _reference_dflash_mask(anchor_positions, block_keep_mask, S, block_size, dev for b in range(B): for q_idx in range(Q_LEN): q_block_id = q_idx // block_size + q_offset = q_idx % block_size anchor_pos = anchor_positions[b, q_block_id].item() is_valid = block_keep_mask[b, q_block_id].item() if not is_valid: continue for kv_idx in range(KV_LEN): is_context = kv_idx < S - ctx_visible = is_context and (kv_idx < anchor_pos) + ctx_visible = is_context and kv_idx < anchor_pos is_draft = kv_idx >= S kv_block_id = (kv_idx - S) // block_size draft_visible = is_draft and (q_block_id == kv_block_id) + if sliding_window is not None: + q_offset = q_idx % block_size + kv_offset = (kv_idx - S) % block_size + ctx_visible = ctx_visible and ( + kv_idx >= anchor_pos + q_offset - (sliding_window - 1) + ) + draft_visible = draft_visible and kv_offset <= q_offset + if ctx_visible or draft_visible: mask[b, 0, q_idx, kv_idx] = True return mask +class _RecordingDraftModel(nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace( + layer_types=["sliding_attention", "full_attention"], + sliding_window=8, + ) + self.sliding_window = 8 + self.attention_mask = None + + def forward(self, noise_embedding, attention_mask, **kwargs): + self.attention_mask = attention_mask + return noise_embedding + + class TestDFlashMask(unittest.TestCase): def setUp(self): torch.manual_seed(42) self.device = torch.device("cuda") - def _compare_masks(self, anchor_positions, block_keep_mask, S, block_size): + def _compare_masks( + self, + anchor_positions, + block_keep_mask, + S, + block_size, + sliding_window=None, + ): """Compare create_dflash_sdpa_mask against element-level reference (ground truth).""" anchor_positions = anchor_positions.to(self.device) block_keep_mask = block_keep_mask.to(self.device) @@ -55,6 +99,7 @@ def _compare_masks(self, anchor_positions, block_keep_mask, S, block_size): S=S, block_size=block_size, device=self.device, + sliding_window=sliding_window, ) ref_mask = _reference_dflash_mask( @@ -63,6 +108,7 @@ def _compare_masks(self, anchor_positions, block_keep_mask, S, block_size): S=S, block_size=block_size, device=self.device, + sliding_window=sliding_window, ) self.assertEqual( @@ -73,12 +119,18 @@ def _compare_masks(self, anchor_positions, block_keep_mask, S, block_size): self.assertTrue( torch.equal(sdpa_mask, ref_mask), f"Mask mismatch with S={S}, block_size={block_size}, " - f"anchors={anchor_positions.tolist()}, keep={block_keep_mask.tolist()}\n" + f"sliding_window={sliding_window}, anchors={anchor_positions.tolist()}, " + f"keep={block_keep_mask.tolist()}\n" f"Diff positions: {(sdpa_mask != ref_mask).nonzero(as_tuple=False).tolist()}", ) def _compare_block_mask_consistency( - self, anchor_positions, block_keep_mask, S, block_size + self, + anchor_positions, + block_keep_mask, + S, + block_size, + sliding_window=None, ): """Verify create_dflash_block_mask block-level mask is consistent with reference.""" anchor_positions = anchor_positions.to(self.device) @@ -90,6 +142,7 @@ def _compare_block_mask_consistency( S=S, block_size=block_size, device=self.device, + sliding_window=sliding_window, ) ref_mask = _reference_dflash_mask( @@ -98,6 +151,7 @@ def _compare_block_mask_consistency( S=S, block_size=block_size, device=self.device, + sliding_window=sliding_window, ) dense_blocks = block_mask.to_dense() # (B, H, Q_blocks, KV_blocks) @@ -118,12 +172,11 @@ def _compare_block_mask_consistency( k_end = min(k_start + BM_BLOCK, KV_LEN) has_nonzero = ref_int[b, q_start:q_end, k_start:k_end].any().item() block_val = dense_blocks[b, 0, qi, ki].item() - if has_nonzero: - self.assertEqual( - block_val, - 1, - f"Block ({qi},{ki}) for batch {b} should be 1 but got 0", - ) + self.assertEqual( + block_val, + int(has_nonzero), + f"Block ({qi},{ki}) for batch {b} has incorrect occupancy", + ) def test_basic_single_batch_single_block(self): """Single batch, single draft block.""" @@ -179,6 +232,97 @@ def test_block_size_1(self): block_keep_mask = torch.tensor([[True, True, True]]) self._compare_masks(anchor_positions, block_keep_mask, S=64, block_size=1) + def test_sliding_window_moves_with_query_offset(self): + """The context window advances while the draft block stays causal.""" + anchor_positions = torch.tensor([[6]]) + block_keep_mask = torch.tensor([[True]]) + mask = create_dflash_sdpa_mask( + anchor_positions=anchor_positions.to(self.device), + block_keep_mask=block_keep_mask.to(self.device), + S=12, + block_size=4, + device=self.device, + sliding_window=4, + ) + expected_visible_keys = ( + [3, 4, 5, 12], + [4, 5, 12, 13], + [5, 12, 13, 14], + [12, 13, 14, 15], + ) + for query_offset, expected in enumerate(expected_visible_keys): + with self.subTest(query_offset=query_offset): + actual = mask[0, 0, query_offset].nonzero().flatten().tolist() + self.assertEqual(actual, expected) + + def test_sliding_window_one_has_no_context_and_causal_draft(self): + """A one-token window removes context but keeps causal own-block keys.""" + anchor_positions = torch.tensor([[4, 9]]) + block_keep_mask = torch.tensor([[True, True]]) + self._compare_masks( + anchor_positions, + block_keep_mask, + S=12, + block_size=3, + sliding_window=1, + ) + + def test_sliding_window_block_mask_consistency(self): + anchor_positions = torch.tensor([[12, 24]]) + block_keep_mask = torch.tensor([[True, True]]) + self._compare_block_mask_consistency( + anchor_positions, + block_keep_mask, + S=32, + block_size=4, + sliding_window=8, + ) + + def test_invalid_sliding_window(self): + anchor_positions = torch.tensor([[12]], device=self.device) + block_keep_mask = torch.tensor([[True]], device=self.device) + for factory in (create_dflash_sdpa_mask, create_dflash_block_mask): + with self.subTest(factory=factory.__name__): + with self.assertRaisesRegex(ValueError, "sliding_window must be > 0"): + factory( + anchor_positions=anchor_positions, + block_keep_mask=block_keep_mask, + S=16, + block_size=4, + device=self.device, + sliding_window=0, + ) + + def test_all_dflash_families_build_mixed_layer_masks(self): + anchors = torch.tensor([[12]], device=self.device) + keep = torch.tensor([[True]], device=self.device) + for model_class in (OnlineDFlashModel, OnlineDominoModel, OnlineDSparkModel): + with self.subTest(model_class=model_class.__name__): + draft_model = _RecordingDraftModel().to(self.device) + model = model_class( + draft_model=draft_model, + target_lm_head=nn.Identity(), + target_embed_tokens=nn.Embedding(32, 8).to(self.device), + mask_token_id=31, + block_size=4, + attention_backend="sdpa", + ) + with mock.patch.object( + model, + "_sample_anchor_positions", + return_value=(anchors, keep), + ): + model._forward_draft_blocks( + input_ids=torch.arange(16, device=self.device).unsqueeze(0), + hidden_states=torch.randn(1, 16, 8, device=self.device), + loss_mask=torch.ones(1, 16, device=self.device), + ) + + masks = draft_model.attention_mask + self.assertEqual(set(masks), {"full_attention", "sliding_attention"}) + self.assertTrue(masks["full_attention"][0, 0, 0, 0].item()) + self.assertFalse(masks["sliding_attention"][0, 0, 0, 0].item()) + def test_mixed_validity_multi_batch(self): """Multi-batch with mixed block validity patterns.""" anchor_positions = torch.tensor([[10, 40, 70, 100], [20, 50, 80, 110]]) @@ -263,6 +407,25 @@ def test_block_mask_consistency_mixed(self): anchor_positions, block_keep_mask, S=128, block_size=8 ) + def test_sliding_block_mask_matches_element_reference(self): + """Flex and dense masks implement the same sliding-layer rule.""" + anchors = torch.tensor([[6, 12]], device=self.device) + keep = torch.tensor([[True, False]], device=self.device) + mask_args = { + "anchor_positions": anchors, + "block_keep_mask": keep, + "S": 16, + "block_size": 4, + "device": self.device, + "sliding_window": 5, + } + dense_mask = create_dflash_sdpa_mask(**mask_args) + block_mask = create_dflash_block_mask(**mask_args) + q_idx = torch.arange(8, device=self.device).unsqueeze(1) + kv_idx = torch.arange(24, device=self.device).unsqueeze(0) + flex_mask = block_mask.mask_mod(0, 0, q_idx, kv_idx) + self.assertTrue(torch.equal(flex_mask, dense_mask[0, 0])) + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/tests/test_utils/test_domino_cross_entropy.py b/tests/test_utils/test_domino_cross_entropy.py new file mode 100644 index 000000000..079ae1769 --- /dev/null +++ b/tests/test_utils/test_domino_cross_entropy.py @@ -0,0 +1,238 @@ +import subprocess +import sys +import unittest + +import torch +import torch.nn.functional as F + +from specforge.core.domino_loss import domino_weighted_cross_entropy + + +def _reference(base, correction, targets, weights, block_size, suffix_start): + """Materialize corrected logits and accumulate cross-entropy in FP32.""" + num_blocks = base.shape[0] // block_size + base_3d = base.reshape(num_blocks, block_size, -1) + suffix_size = block_size - suffix_start + final = torch.cat( + ( + base_3d[:, :suffix_start], + base_3d[:, suffix_start:] + correction.reshape(num_blocks, suffix_size, -1), + ), + dim=1, + ).reshape_as(base) + final_loss_sum = ( + F.cross_entropy(final.float(), targets, reduction="none") * weights + ).sum() + base_loss_sum = ( + F.cross_entropy(base.float(), targets, reduction="none") * weights + ).sum() + return ( + final_loss_sum, + base_loss_sum, + final.argmax(dim=-1), + base.argmax(dim=-1), + ) + + +class DominoCrossEntropyTest(unittest.TestCase): + def _compare( + self, + device, + *, + block_size=4, + suffix_start=1, + num_blocks=3, + vocab_size=19, + dtype=torch.float32, + ): + torch.manual_seed(7) + rows = num_blocks * block_size + correction_rows = num_blocks * (block_size - suffix_start) + targets = torch.randint(vocab_size, (rows,), device=device) + weights = torch.rand(rows, device=device) + weights[::5] = 0 + + actual_base = torch.randn( + rows, + vocab_size, + device=device, + dtype=dtype, + requires_grad=True, + ) + actual_correction = torch.randn( + correction_rows, + vocab_size, + device=device, + dtype=dtype, + requires_grad=True, + ) + expected_base = actual_base.detach().clone().requires_grad_(True) + expected_correction = actual_correction.detach().clone().requires_grad_(True) + + actual = domino_weighted_cross_entropy( + actual_base, + actual_correction, + targets, + weights, + block_size, + suffix_start, + use_fused=device.type == "cuda", + ) + expected = _reference( + expected_base, + expected_correction, + targets, + weights, + block_size, + suffix_start, + ) + rtol, atol = (1e-2, 5e-3) if dtype == torch.bfloat16 else (1e-4, 1e-4) + torch.testing.assert_close(actual[0], expected[0], rtol=rtol, atol=atol) + torch.testing.assert_close(actual[1], expected[1], rtol=rtol, atol=atol) + torch.testing.assert_close(actual[2], expected[2]) + torch.testing.assert_close(actual[3], expected[3]) + + (0.7 * actual[0] + 0.3 * actual[1]).backward() + (0.7 * expected[0] + 0.3 * expected[1]).backward() + torch.testing.assert_close( + actual_base.grad, expected_base.grad, rtol=rtol, atol=atol + ) + torch.testing.assert_close( + actual_correction.grad, + expected_correction.grad, + rtol=rtol, + atol=atol, + ) + + def test_cpu_fallback_matches_pytorch(self): + self._compare(torch.device("cpu")) + + def test_cpu_fallback_does_not_import_triton(self): + script = """ +import builtins +import torch + +real_import = builtins.__import__ + +def guarded_import(name, *args, **kwargs): + if name == "triton" or name.startswith("triton."): + raise AssertionError("CPU fallback imported Triton") + return real_import(name, *args, **kwargs) + +builtins.__import__ = guarded_import +from specforge.core.domino_loss import domino_weighted_cross_entropy + +domino_weighted_cross_entropy( + torch.randn(4, 8), + torch.randn(3, 8), + torch.zeros(4, dtype=torch.long), + torch.ones(4), + block_size=4, + suffix_start=1, + use_fused=False, +) +""" + subprocess.run([sys.executable, "-c", script], check=True) + + def test_invalid_correction_shape_is_rejected(self): + with self.assertRaisesRegex(ValueError, "correction logits"): + domino_weighted_cross_entropy( + torch.randn(4, 8), + torch.randn(4, 8), + torch.zeros(4, dtype=torch.long), + torch.ones(4), + block_size=4, + suffix_start=1, + use_fused=False, + ) + + def test_non_integer_targets_are_rejected(self): + with self.assertRaisesRegex(ValueError, "torch.long"): + domino_weighted_cross_entropy( + torch.randn(4, 8), + torch.randn(3, 8), + torch.zeros(4), + torch.ones(4), + block_size=4, + suffix_start=1, + use_fused=False, + ) + + def test_fused_path_requires_cuda(self): + with self.assertRaisesRegex(ValueError, "requires CUDA"): + domino_weighted_cross_entropy( + torch.randn(4, 8), + torch.randn(3, 8), + torch.zeros(4, dtype=torch.long), + torch.ones(4), + block_size=4, + suffix_start=1, + use_fused=True, + ) + + @unittest.skipUnless(torch.cuda.is_available(), "Triton loss requires CUDA") + def test_fused_path_rejects_unsupported_logit_dtype(self): + with self.assertRaisesRegex(ValueError, "FP16, BF16, or FP32"): + domino_weighted_cross_entropy( + torch.randn(4, 8, device="cuda", dtype=torch.float64), + torch.randn(3, 8, device="cuda", dtype=torch.float64), + torch.zeros(4, device="cuda", dtype=torch.long), + torch.ones(4, device="cuda"), + block_size=4, + suffix_start=1, + use_fused=True, + ) + + @unittest.skipUnless(torch.cuda.is_available(), "Triton loss requires CUDA") + def test_triton_matches_pytorch(self): + device = torch.device("cuda") + self._compare(device) + self._compare( + device, + block_size=4, + suffix_start=0, + num_blocks=2, + vocab_size=2053, + ) + self._compare( + device, + block_size=4, + suffix_start=3, + num_blocks=2, + vocab_size=2053, + ) + + # Both reductions must preserve PyTorch's leftmost tie break. + base = torch.zeros(8, 2053, device=device) + correction = torch.zeros(6, 2053, device=device) + targets = torch.arange(8, device=device) + weights = torch.ones(8, device=device) + actual = domino_weighted_cross_entropy( + base, + correction, + targets, + weights, + block_size=4, + suffix_start=1, + use_fused=True, + ) + torch.testing.assert_close(actual[2], torch.zeros_like(actual[2])) + torch.testing.assert_close(actual[3], torch.zeros_like(actual[3])) + + @unittest.skipUnless( + torch.cuda.is_available() and torch.cuda.is_bf16_supported(), + "Triton BF16 loss requires a BF16 CUDA device", + ) + def test_triton_bfloat16_matches_float32_reference(self): + self._compare( + torch.device("cuda"), + dtype=torch.bfloat16, + block_size=16, + suffix_start=2, + num_blocks=2, + vocab_size=5003, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_utils/test_mtp.py b/tests/test_utils/test_mtp.py new file mode 100644 index 000000000..2fcc7cd33 --- /dev/null +++ b/tests/test_utils/test_mtp.py @@ -0,0 +1,619 @@ +# coding=utf-8 +"""CPU unit tests for the MTP (native-head fine-tune) algorithm. + +Covers the pieces that do not need a GPU or a real target checkpoint: + - OnlineMTPModel forward/shift/loss/accuracy plumbing + - MTPTrainStrategy batch -> StepOutput adaptation + - strict native ``mtp.*`` weight initialization from a target checkpoint + - selective checkpoint loading (modeling/target/checkpoint.py) + - merge-back round trip (export/mtp.py) + - draft-architecture and built-in algorithm registration +""" + +from __future__ import annotations + +import json +import os +import tempfile +import unittest +from types import SimpleNamespace + +import torch +from transformers.models.qwen3.modeling_qwen3 import Qwen3Config + +from specforge.algorithms.builtin import builtin_algorithm_registry +from specforge.algorithms.mtp.providers import _init_from_native_mtp +from specforge.core.mtp import OnlineMTPModel +from specforge.modeling.draft import available_drafts, resolve_draft +from specforge.modeling.draft.mtp import Qwen3_5MTPDraftModel +from specforge.training.strategies.base import MTPTrainStrategy + + +def _tiny_config(**overrides) -> Qwen3Config: + payload = dict( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + vocab_size=128, + max_position_embeddings=512, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + rms_norm_eps=1e-6, + attn_output_gate=True, + partial_rotary_factor=0.25, + mtp_config={"share_lm_head": True}, + tie_word_embeddings=False, + ) + payload.update(overrides) + config = Qwen3Config(**payload) + config._attn_implementation = "eager" + return config + + +def _tiny_batch(config: Qwen3Config, seq_len: int = 16, batch: int = 2): + input_ids = torch.randint(0, config.vocab_size, (batch, seq_len)) + hidden_states = torch.randn(batch, seq_len, config.hidden_size) + loss_mask = torch.ones(batch, seq_len) + return input_ids, hidden_states, loss_mask + + +def _cfg(target_model_path: str, draft_checkpoint_path: str = ""): + return SimpleNamespace( + model=SimpleNamespace( + target_model_path=target_model_path, + draft_checkpoint_path=draft_checkpoint_path, + cache_dir=None, + ) + ) + + +class OnlineMTPModelTest(unittest.TestCase): + def test_forward_returns_finite_loss_and_accuracy_lists(self): + config = _tiny_config() + model = OnlineMTPModel(Qwen3_5MTPDraftModel(config)) + input_ids, hidden_states, loss_mask = _tiny_batch(config) + + loss, corrects, denoms = model( + input_ids=input_ids, hidden_states=hidden_states, loss_mask=loss_mask + ) + + self.assertEqual(0, loss.dim()) + self.assertTrue(torch.isfinite(loss)) + self.assertEqual(1, len(corrects)) + self.assertEqual(1, len(denoms)) + # next-token shift drops one position + self.assertEqual((2, 15), corrects[0].shape) + self.assertEqual((2, 15), denoms[0].shape) + + def test_loss_backward_populates_mtp_grads(self): + config = _tiny_config() + draft = Qwen3_5MTPDraftModel(config) + model = OnlineMTPModel(draft) + input_ids, hidden_states, loss_mask = _tiny_batch(config) + + loss, _, _ = model( + input_ids=input_ids, hidden_states=hidden_states, loss_mask=loss_mask + ) + loss.backward() + + self.assertIsNotNone(draft.mtp.fc.weight.grad) + self.assertTrue(torch.isfinite(draft.mtp.fc.weight.grad).all()) + + def test_shift_for_next_token_matches_serving_alignment(self): + model = OnlineMTPModel(Qwen3_5MTPDraftModel(_tiny_config())) + logits = torch.zeros(1, 5, 7) + input_ids = torch.tensor([[10, 11, 12, 13, 14]]) + loss_mask = torch.ones(1, 5) + + shift_logits, shift_labels, shift_mask = model._shift_for_next_token( + logits, input_ids, loss_mask + ) + + self.assertEqual((1, 4, 7), shift_logits.shape) + # labels are x_2..x_T padded with one ignore index + self.assertEqual([12, 13, 14, -100], shift_labels[0].tolist()) + # the padded position is masked out + self.assertEqual([1, 1, 1, 0], shift_mask[0].tolist()) + + def test_forward_shifts_position_ids_with_draft_tokens(self): + class _RecordingDraft(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(pad_token_id=0) + self.position_ids = None + + def forward( + self, + input_ids, + hidden_states, + attention_mask=None, + position_ids=None, + ): + self.position_ids = position_ids.detach().clone() + logits = torch.zeros( + input_ids.shape[0], input_ids.shape[1], 32, requires_grad=True + ) + return SimpleNamespace(logits=logits) + + draft = _RecordingDraft() + model = OnlineMTPModel(draft) + model( + input_ids=torch.tensor([[10, 11, 12, 13]]), + hidden_states=torch.zeros(1, 4, 8), + loss_mask=torch.ones(1, 4), + position_ids=torch.tensor([[4, 5, 6, 7]]), + ) + + # x[t+1] is fused with h[t], but RoPE must use x[t+1]'s serving + # position. The synthetic final token is assigned the next position. + self.assertEqual([[5, 6, 7, 8]], draft.position_ids.tolist()) + + +class MTPTrainStrategyTest(unittest.TestCase): + def test_forward_loss_adapts_model_outputs(self): + loss = torch.tensor(1.5, requires_grad=True) + corrects = [torch.tensor([[1.0, 0.0, 1.0]])] + denoms = [torch.tensor([[1.0, 1.0, 1.0]])] + + class _Stub(torch.nn.Module): + def __init__(self): + super().__init__() + self.proj = torch.nn.Linear(1, 1) + + def forward(self, input_ids, hidden_states, loss_mask): + return loss, corrects, denoms + + strategy = MTPTrainStrategy(_Stub()) + batch = SimpleNamespace( + tensors={ + "input_ids": torch.zeros(1, 3, dtype=torch.long), + "loss_mask": torch.ones(1, 3), + "target_last_hidden_states": torch.zeros(1, 3, 4), + } + ) + + out = strategy.forward_loss(batch) + + self.assertIs(out.loss, loss) + self.assertAlmostEqual(2.0 / 3.0, out.metrics["accuracy"].item()) + self.assertEqual(3.0, out.metrics["accuracy_denom"].item()) + num, den = out.ratio_metrics["accuracy"] + self.assertEqual(2.0, num.item()) + self.assertEqual(3.0, den.item()) + loss_num, loss_den = out.loss_terms + self.assertIsNotNone(loss_num.grad_fn) + self.assertEqual(4.5, loss_num.item()) + self.assertEqual(3.0, loss_den.item()) + + def test_forward_loss_rejects_missing_features(self): + strategy = MTPTrainStrategy(torch.nn.Linear(1, 1)) + batch = SimpleNamespace(tensors={"input_ids": torch.zeros(1, 3)}) + with self.assertRaisesRegex(ValueError, "missing required features"): + strategy.forward_loss(batch) + + def test_checkpoint_state_filter_strips_draft_prefix(self): + strategy = MTPTrainStrategy(torch.nn.Linear(1, 1)) + state = { + "draft_model.mtp.fc.weight": torch.zeros(2, 2), + "draft_model.embed_tokens.weight": torch.zeros(4, 2), + "other.weight": torch.zeros(1), + } + filtered = strategy.checkpoint_state_filter(state) + self.assertEqual({"mtp.fc.weight", "embed_tokens.weight"}, set(filtered)) + + +class NativeMTPInitTest(unittest.TestCase): + def test_loads_native_mtp_weights_from_target_checkpoint(self): + from safetensors.torch import save_file + + config = _tiny_config() + draft = Qwen3_5MTPDraftModel(config) + replacement = torch.ones_like(draft.mtp.fc.weight) + native_state = { + key: torch.ones_like(value) + for key, value in draft.native_state_dict().items() + if key in draft.required_native_state_keys() + } + + with tempfile.TemporaryDirectory(prefix="mtp-native-init-") as tmpdir: + save_file(native_state, f"{tmpdir}/model.safetensors") + _init_from_native_mtp(_cfg(tmpdir), draft) + + self.assertTrue(torch.equal(draft.mtp.fc.weight, replacement)) + + def test_partial_native_weights_raise(self): + from safetensors.torch import save_file + + draft = Qwen3_5MTPDraftModel(_tiny_config()) + with tempfile.TemporaryDirectory(prefix="mtp-native-init-") as tmpdir: + save_file( + {"mtp.fc.weight": torch.ones_like(draft.mtp.fc.weight)}, + f"{tmpdir}/model.safetensors", + ) + with self.assertRaisesRegex(RuntimeError, "missing required native"): + _init_from_native_mtp(_cfg(tmpdir), draft) + + def test_missing_native_weights_raise_by_default(self): + draft = Qwen3_5MTPDraftModel(_tiny_config()) + with tempfile.TemporaryDirectory(prefix="mtp-native-init-") as tmpdir: + with self.assertRaisesRegex(RuntimeError, "no native mtp"): + _init_from_native_mtp(_cfg(tmpdir), draft) + + def test_missing_native_weights_allowed_for_warm_start(self): + draft = Qwen3_5MTPDraftModel(_tiny_config()) + before = draft.mtp.fc.weight.detach().clone() + with tempfile.TemporaryDirectory(prefix="mtp-native-init-") as tmpdir: + # warm start from a trained draft checkpoint skips the strict check + _init_from_native_mtp( + _cfg(tmpdir, draft_checkpoint_path="some/ckpt"), draft + ) + self.assertTrue(torch.equal(draft.mtp.fc.weight, before)) + + def test_native_init_covers_all_mtp_parameters(self): + """A native checkpoint shipping the full mtp.* key set must overwrite + every draft native parameter — none may keep its random init.""" + from safetensors.torch import save_file + + config = _tiny_config() + draft = Qwen3_5MTPDraftModel(config) + native_keys = draft.required_native_state_keys() + replacement = { + key: torch.ones_like(value) + for key, value in draft.native_state_dict().items() + if key in native_keys + } + + with tempfile.TemporaryDirectory(prefix="mtp-native-init-") as tmpdir: + save_file(replacement, f"{tmpdir}/model.safetensors") + _init_from_native_mtp(_cfg(tmpdir), draft) + + after = draft.native_state_dict() + self.assertTrue(native_keys.issubset(after)) + for key in native_keys: + value = after[key] + self.assertTrue( + torch.all(value == 1), f"{key} was not loaded from native weights" + ) + + def test_native_init_tolerates_merged_checkpoint_keys(self): + """A previously merged checkpoint carries backfilled shared embeddings + (mtp.embed_tokens.weight / mtp.lm_head.weight); re-finetuning it must + not be rejected as an incompatible native state.""" + from safetensors.torch import save_file + + config = _tiny_config() + draft = Qwen3_5MTPDraftModel(config) + state = { + key: torch.ones_like(value) + for key, value in draft.native_state_dict().items() + if key in draft.required_native_state_keys() + } + state["mtp.embed_tokens.weight"] = torch.randn( + config.vocab_size, config.hidden_size + ) + state["mtp.lm_head.weight"] = torch.randn(config.vocab_size, config.hidden_size) + + with tempfile.TemporaryDirectory(prefix="mtp-native-init-") as tmpdir: + save_file(state, f"{tmpdir}/model.safetensors") + _init_from_native_mtp(_cfg(tmpdir), draft) # must not raise + + self.assertTrue(torch.all(draft.mtp.fc.weight == 1)) + + +class DraftBaseContractTest(unittest.TestCase): + def test_share_target_embeddings_freezes_and_shares(self): + config = _tiny_config() + draft = Qwen3_5MTPDraftModel(config) + embed_w = torch.nn.Parameter(torch.randn(config.vocab_size, config.hidden_size)) + head_w = torch.nn.Parameter(torch.randn(config.vocab_size, config.hidden_size)) + + draft.share_target_embeddings(embed_w, lm_head_weight=head_w) + + self.assertIs(draft.embed_tokens.weight, embed_w) + self.assertIs(draft.mtp.lm_head.weight, head_w) + self.assertFalse(draft.embed_tokens.weight.requires_grad) + self.assertFalse(draft.mtp.lm_head.weight.requires_grad) + + def test_share_lm_head_disabled_keeps_own_head(self): + config = _tiny_config(mtp_config={"share_lm_head": False}) + draft = Qwen3_5MTPDraftModel(config) + own_head = draft.mtp.lm_head.weight + embed_w = torch.nn.Parameter(torch.randn(config.vocab_size, config.hidden_size)) + + draft.share_target_embeddings(embed_w) + + self.assertIs(draft.embed_tokens.weight, embed_w) + self.assertIs(draft.mtp.lm_head.weight, own_head) + + def test_native_state_dict_uses_native_prefix(self): + draft = Qwen3_5MTPDraftModel(_tiny_config()) + native = draft.native_state_dict() + self.assertTrue(native) + self.assertTrue(all(key.startswith(draft.NATIVE_KEY_PREFIX) for key in native)) + self.assertIn("mtp.fc.weight", native) + + def test_required_native_state_respects_shared_lm_head(self): + shared = Qwen3_5MTPDraftModel(_tiny_config()) + own = Qwen3_5MTPDraftModel(_tiny_config(mtp_config={"share_lm_head": False})) + + self.assertNotIn("mtp.lm_head.weight", shared.required_native_state_keys()) + self.assertIn("mtp.lm_head.weight", own.required_native_state_keys()) + + +class SelectiveCheckpointLoadingTest(unittest.TestCase): + def test_sharded_selective_loading(self): + from safetensors.torch import save_file + + from specforge.modeling.target.checkpoint import ( + list_checkpoint_keys, + load_selected_tensors, + read_weight_map, + ) + + with tempfile.TemporaryDirectory(prefix="mtp-ckpt-") as tmpdir: + save_file( + { + "mtp.fc.weight": torch.zeros(4, 4), + "model.embed_tokens.weight": torch.zeros(2, 2), + }, + os.path.join(tmpdir, "model-00001-of-00002.safetensors"), + ) + save_file( + { + "mtp.norm.weight": torch.ones(4), + "lm_head.weight": torch.zeros(2, 2), + }, + os.path.join(tmpdir, "model-00002-of-00002.safetensors"), + ) + weight_map = { + "mtp.fc.weight": "model-00001-of-00002.safetensors", + "model.embed_tokens.weight": "model-00001-of-00002.safetensors", + "mtp.norm.weight": "model-00002-of-00002.safetensors", + "lm_head.weight": "model-00002-of-00002.safetensors", + } + with open(os.path.join(tmpdir, "model.safetensors.index.json"), "w") as f: + json.dump({"weight_map": weight_map}, f) + + self.assertEqual(weight_map, read_weight_map(tmpdir)) + self.assertEqual(4, len(list_checkpoint_keys(tmpdir))) + selected = load_selected_tensors(tmpdir, lambda key: key.startswith("mtp.")) + + self.assertEqual({"mtp.fc.weight", "mtp.norm.weight"}, set(selected)) + self.assertTrue(torch.equal(selected["mtp.norm.weight"], torch.ones(4))) + + def test_single_file_selective_loading(self): + from safetensors.torch import save_file + + from specforge.modeling.target.checkpoint import load_selected_tensors + + with tempfile.TemporaryDirectory(prefix="mtp-ckpt-") as tmpdir: + save_file( + {"mtp.fc.weight": torch.zeros(4, 4), "other.weight": torch.zeros(1)}, + os.path.join(tmpdir, "model.safetensors"), + ) + selected = load_selected_tensors(tmpdir, lambda key: key.startswith("mtp.")) + self.assertEqual({"mtp.fc.weight"}, set(selected)) + + +class ExportRoundTripTest(unittest.TestCase): + """merge_mtp_into_base round trip on synthetic single-file checkpoints.""" + + def test_merge_replaces_native_and_copies_embeddings(self): + from safetensors.torch import save_file + + from specforge.export.mtp import merge_mtp_into_base + from specforge.modeling.target.checkpoint import load_selected_tensors + + with tempfile.TemporaryDirectory() as tmpdir: + base = os.path.join(tmpdir, "base") + draft = os.path.join(tmpdir, "draft") + out = os.path.join(tmpdir, "out") + os.makedirs(base) + os.makedirs(draft) + + base_embed = torch.randn(128, 64) + stale_native = torch.zeros(64, 128) + save_file( + { + "model.embed_tokens.weight": base_embed, + "mtp.fc.weight": stale_native, + }, + os.path.join(base, "model.safetensors"), + ) + with open(os.path.join(base, "config.json"), "w") as f: + json.dump({"hidden_size": 64, "tie_word_embeddings": True}, f) + + trained = torch.ones(64, 128) + save_file( + {"mtp.fc.weight": trained}, + os.path.join(draft, "model.safetensors"), + ) + with open(os.path.join(draft, "config.json"), "w") as f: + json.dump( + { + "architectures": ["Qwen3_5MTPDraftModel"], + "hidden_size": 64, + "head_dim": 16, + }, + f, + ) + + merge_mtp_into_base(base, draft, out) + + merged = load_selected_tensors(out, lambda _key: True) + # trained weights replace the stale native ones + self.assertTrue(torch.equal(merged["mtp.fc.weight"], trained)) + # shared embedding copied into the native namespace + self.assertTrue(torch.equal(merged["mtp.embed_tokens.weight"], base_embed)) + # base weights untouched + self.assertTrue( + torch.equal(merged["model.embed_tokens.weight"], base_embed) + ) + # config patched with the draft's structural dims + with open(os.path.join(out, "config.json")) as f: + merged_config = json.load(f) + self.assertEqual(16, merged_config["head_dim"]) + + def test_merge_accepts_runtime_training_checkpoint(self): + from safetensors.torch import save_file + + from specforge.export.mtp import merge_mtp_into_base + from specforge.modeling.target.checkpoint import load_selected_tensors + + with tempfile.TemporaryDirectory() as tmpdir: + base = os.path.join(tmpdir, "base") + runtime = os.path.join(tmpdir, "run-step1") + out = os.path.join(tmpdir, "out") + draft_config = os.path.join(tmpdir, "draft-config.json") + os.makedirs(base) + os.makedirs(runtime) + + base_embed = torch.randn(128, 64) + save_file( + { + "model.embed_tokens.weight": base_embed, + "mtp.fc.weight": torch.zeros(64, 128), + }, + os.path.join(base, "model.safetensors"), + ) + with open(os.path.join(base, "config.json"), "w") as f: + json.dump({"hidden_size": 64, "tie_word_embeddings": True}, f) + with open(draft_config, "w") as f: + json.dump( + { + "architectures": ["Qwen3_5MTPDraftModel"], + "hidden_size": 64, + "head_dim": 16, + }, + f, + ) + + trained = torch.ones(64, 128) + torch.save( + { + "strategy": "mtp", + "draft_state_dict": {"mtp.fc.weight": trained}, + }, + os.path.join(runtime, "training_state.pt"), + ) + + merge_mtp_into_base( + base, + runtime, + out, + draft_config_path=draft_config, + ) + + merged = load_selected_tensors(out, lambda _key: True) + self.assertTrue(torch.equal(merged["mtp.fc.weight"], trained)) + self.assertTrue(torch.equal(merged["mtp.embed_tokens.weight"], base_embed)) + with open(os.path.join(out, "config.json")) as f: + merged_config = json.load(f) + self.assertEqual(16, merged_config["head_dim"]) + + def test_runtime_checkpoint_requires_draft_config(self): + from specforge.export.mtp import merge_mtp_into_base + + with tempfile.TemporaryDirectory() as tmpdir: + runtime = os.path.join(tmpdir, "run-step1") + os.makedirs(runtime) + torch.save( + { + "strategy": "mtp", + "draft_state_dict": {"mtp.fc.weight": torch.ones(1)}, + }, + os.path.join(runtime, "training_state.pt"), + ) + + with self.assertRaisesRegex(ValueError, "draft_config_path is required"): + merge_mtp_into_base("unused", runtime, os.path.join(tmpdir, "out")) + + def test_merge_runtime_checkpoint_with_shared_tied_weights(self): + """Regression: a tied target shares one storage between the draft's + embed_tokens.weight and mtp.lm_head.weight; safetensors must not choke + on the aliased pair when writing the merged checkpoint.""" + from safetensors.torch import save_file + + from specforge.export.mtp import merge_mtp_into_base + from specforge.modeling.target.checkpoint import load_selected_tensors + + with tempfile.TemporaryDirectory() as tmpdir: + base = os.path.join(tmpdir, "base") + runtime = os.path.join(tmpdir, "run-step1") + out = os.path.join(tmpdir, "out") + draft_config = os.path.join(tmpdir, "draft-config.json") + os.makedirs(base) + os.makedirs(runtime) + + save_file( + {"model.embed_tokens.weight": torch.randn(128, 64)}, + os.path.join(base, "model.safetensors"), + ) + with open(os.path.join(base, "config.json"), "w") as f: + json.dump({"hidden_size": 64, "tie_word_embeddings": True}, f) + with open(draft_config, "w") as f: + json.dump( + { + "architectures": ["Qwen3_5MTPDraftModel"], + "hidden_size": 64, + "head_dim": 16, + }, + f, + ) + + shared = torch.randn(128, 64) + trained = torch.ones(64, 128) + torch.save( + { + "strategy": "mtp", + "draft_state_dict": { + "embed_tokens.weight": shared, + "mtp.lm_head.weight": shared, + "mtp.fc.weight": trained, + }, + }, + os.path.join(runtime, "training_state.pt"), + ) + + merge_mtp_into_base(base, runtime, out, draft_config_path=draft_config) + + merged = load_selected_tensors(out, lambda _key: True) + self.assertTrue(torch.equal(merged["mtp.embed_tokens.weight"], shared)) + self.assertTrue(torch.equal(merged["mtp.lm_head.weight"], shared)) + self.assertTrue(torch.equal(merged["mtp.fc.weight"], trained)) + + +class MTPRegistrationTest(unittest.TestCase): + def test_draft_architecture_is_registered(self): + self.assertIn("Qwen3_5MTPDraftModel", available_drafts()) + self.assertIs(resolve_draft("Qwen3_5MTPDraftModel"), Qwen3_5MTPDraftModel) + + def test_builtin_registry_resolves_mtp(self): + registration = builtin_algorithm_registry().resolve("mtp") + self.assertEqual("mtp", registration.spec.name) + self.assertEqual( + "Qwen3_5MTPDraftModel", + registration.providers.model.draft_config.architecture, + ) + + def test_offline_layout_persists_only_final_hidden(self): + providers = builtin_algorithm_registry().resolve("mtp").providers + layout = providers.offline_for("text").capture_layout + self.assertEqual( + ("input_ids", "loss_mask", "target_last_hidden_states"), + layout.output_names, + ) + self.assertIsNone(layout.aux_feature) + + def test_streaming_layout_exposes_final_hidden(self): + providers = builtin_algorithm_registry().resolve("mtp").providers + layout = providers.server_streaming_for("text").layout + self.assertEqual("target_last_hidden_states", layout.last_hidden_feature) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_utils/test_process_cleanup.py b/tests/test_utils/test_process_cleanup.py index d423eada4..bc5bef7a6 100644 --- a/tests/test_utils/test_process_cleanup.py +++ b/tests/test_utils/test_process_cleanup.py @@ -5,7 +5,11 @@ import time import unittest -from tests.utils import execute_shell_command, terminate_process_trees +from tests.utils import ( + execute_shell_command, + terminate_process_trees, + wait_for_processes_to_stop, +) @unittest.skipUnless( @@ -54,15 +58,13 @@ def test_terminates_descendants_after_the_group_leader_exits(self): process.wait(timeout=5) terminate_process_trees(process, grace_s=1) - deadline = time.monotonic() + 5 - while time.monotonic() < deadline: - try: - os.killpg(leader_pid, 0) - except ProcessLookupError: - break - time.sleep(0.02) - else: - self.fail(f"process group {leader_pid} survived cleanup") + survivors = wait_for_processes_to_stop( + (leader_pid, grandchild_pid), timeout_s=5 + ) + self.assertFalse( + survivors, + f"processes survived cleanup: {survivors}", + ) finally: try: os.killpg(process.pid, signal.SIGKILL) diff --git a/tests/utils.py b/tests/utils.py index ccdc78d79..b72f14921 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -82,6 +82,48 @@ def terminate_process_trees(*processes: subprocess.Popen, grace_s: float = 30) - ) +def process_is_running(pid: int) -> bool: + """Return whether *pid* is live, treating Linux zombies as terminated.""" + stat_path = f"/proc/{pid}/stat" + if os.path.isdir("/proc"): + try: + with open(stat_path, encoding="utf-8") as stream: + stat = stream.read() + except FileNotFoundError: + return False + except OSError: + pass + else: + # The command name is parenthesized and may contain spaces or ')', + # so split only after its final closing parenthesis. The next field + # is the process state. Zombies still answer kill(pid, 0), even + # though no code can run and no signal can make them exit again. + command_end = stat.rfind(")") + fields = stat[command_end + 1 :].split() + if fields and fields[0] in {"Z", "X", "x"}: + return False + + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +def wait_for_processes_to_stop( + pids: tuple[int, ...], *, timeout_s: float +) -> tuple[int, ...]: + """Wait for PIDs to stop running and return any live survivors.""" + deadline = time.monotonic() + timeout_s + while True: + survivors = tuple(pid for pid in pids if process_is_running(pid)) + if not survivors or time.monotonic() >= deadline: + return survivors + time.sleep(0.02) + + def wait_for_server( base_url: str, timeout: int | None = None,