Merge DataFlow runtime branch into main#648
Open
maocheng23 wants to merge 93 commits into
Open
Conversation
Replace the torch.save/torch.load pickle round-trip with Mooncake's native
raw-buffer DMA. One hard-pinned object per *tensor*, keyed
{store_id}/{sid}/g{gen}/{name}: put() writes each tensor straight from its
storage via put_from(ptr); get() reads each straight into a tensor allocated
from the ref's FeatureSpec via get_into(ptr). Shape/dtype travel on the ref, so
there is no serialized header. The generation lives in the key (like
SharedDirFeatureStore's filename gen), so a re-put supersedes the old key set
and a stale ref's keys are gone -> get() raises (B5). Hard-pin is preserved
(put_from carries the ReplicateConfig). Falls back to the pickle blob path when
the backend lacks put_from/get_into (zero_copy=False or an older mooncake).
Source + receive buffers are registered with the transfer engine around the DMA:
RDMA rejects an unregistered address (AddressNotRegistered, -800); TCP ignores
the registration. Validated cross-node on 2x H200 over BOTH tcp and rdma --
producer put_from on node 0, consumer get_into + FSDP train on node 1.
Tests: the fake now simulates put_from/get_into via ctypes, so the full contract
runs on the zero-copy path (28 tests, real gated test green on a live master).
Adds zero-copy specifics (per-tensor keys, no-pickle-on-the-wire, re-put
supersede, pickle fallback) + cross-process contract tests (producer/consumer as
separate instances over one backend) + abort-under-lease-defer and
re-put-remove-failure edge cases.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…or_keys _store_get_tensor only checked get_into's return for a negative error code, so a short read (0 <= rc < nb) into the freshly torch.empty'd receive buffer was accepted, handing the trainer a tensor with an uninitialized garbage tail. Unlike the pickle path (torch.load reconstructs whole tensors), the raw-buffer path cannot otherwise detect under-fill. get_into returns the bytes read (a full read == nb), so require rc == nb and raise KeyError on a short read. Adds a zero-copy regression test that truncates the stored blob. Also remove the dead _tensor_keys() helper: it has no call sites, while _try_physical_free and _sample_exists each inline the same per-tensor-key comprehension. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ocess online ref stream The offline disagg path hands the consumer a STATIC ref manifest written once. Online disaggregation needs a continuous stream: the rollout producer commits SampleRefs while the trainer consumes them, on another node. StreamingRefChannel is that control-plane channel: * tensor-free append-only JSONL (asserts no-tensor on publish); feature tensors go through the FeatureStore (Mooncake), so no shared *data* mount is needed. * poll() tail-reads complete lines from the last offset, buffering a partial trailing line so a half-written record is never parsed. * mark_consumed()/consumed_remote() give the producer a cross-process backpressure signal (in_flight_remote) with no shared in-process state. * close() drops an EOF sentinel so stream() terminates once drained; idle_timeout_s guards against a dead producer. Filesystem-backed (any shared control mount); a networked control plane slots in behind the same publish/poll API later. 7 CPU tests. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…consumer} + StreamingRefQueue Wires online disaggregated training: a rollout producer pool streams features to a trainer pool on a different node, tensors over Mooncake, refs over the StreamingRefChannel. * StreamingRefQueue: adapts the channel to the SampleRefQueue protocol (get/ack/fail) the FeatureDataLoader consumes. get() blocks until refs are available or the channel is closed-and-drained; ack() advances the channel's consumed counter (the producer's backpressure signal). * build_disagg_online_producer: RolloutWorker(s) (HF/SGLang target via SGLangAdapter) put() consume-once features into a Mooncake store and publish refs to the channel. drive_producer() runs until the prompt pool drains, pausing while in_flight_remote() exceeds a high-watermark so a lagging trainer can't overrun the segment, then closes the channel. * build_disagg_online_consumer: the online trainer assembly (target_head=None) reading refs from a StreamingRefQueue + tensors from a consume-once Mooncake store. The loader frees each sample on read (get -> release -> remote remove). The cross-pool consume-once free works through the shared Mooncake remove (proven in the mooncake cross-process tests); from SampleRef down the trainer path is identical to colocated online. Tests: test_disagg_online (CPU integration -- stream -> loader -> consume-once free, backpressure, blocks-until-close) + test_disagg_online_launch (GPU -- producer streams, consumer trains through FSDP end to end). StreamingRefQueue covered via the integration test. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
[DataFlow runtime · online] Online disaggregated training (StreamingRefChannel + build_disagg_online_*)
Wire a shared, durable metadata store into the online disaggregated
producer/consumer so commit/dedup/ack are cross-process, and make the
streaming consumer restart-safe.
- build_disagg_online_{producer,consumer} accept metadata_store /
metadata_db_path; both processes share one SQLiteMetadataStore instead of
a private InMemoryMetadataStore each. New _resolve_metadata_store helper.
- consumer gains resume=: reconcile_on_restart derives the already-trained
set and hands it to StreamingRefQueue as skip_ids, so a restarted consumer
drops durably-trained refs on the append-only channel re-read (no duplicate
train); the committed-but-unacked tail re-streams and re-trains.
- StreamingRefQueue gains skip_ids: drops matching refs on read and counts
them consumed so the producer's in_flight_remote backpressure stays exact.
CPU tests (test_disagg_online_shared_plane.py): cross-process commit/dedup,
durable ack visible across processes, reconcile releases-acked/requeues-
unacked (and requeues-all when the optimizer step wasn't durable), restart
skip. Implements stage O1.1 of the online roadmap (#618).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…loop Replace the drain-then-fit shape (generate the whole prompt pool, then train) with a live loop in which the producer and trainer run concurrently. - run_disagg_online_interleaved: the producer streams refs on a daemon thread while trainer.fit consumes on the main thread (StreamingRefQueue blocks until closed-and-drained, so the trainer tracks the producer). Symmetric, hang-free shutdown: trainer-finishes-first sets a cooperative should_stop so the producer doesn't block on the in-flight watermark after the consumer stops draining; producer-finishes-first closes the channel; producer-raise closes the channel (so the consumer can't hang) and re-raises on the main thread. - build_disagg_online_eagle3_runtime: the named single-process builder from the roadmap; composes producer + consumer over one shared metadata store, one consume-once feature store, and a producer/consumer StreamingRefChannel pair, and returns (trainer, loader, run). Uses the in-process generate stub (no live SGLang server, no Ray -- those are O1.3 / O2). - drive_producer gains a cooperative should_stop predicate + finally-close. CPU tests (test_disagg_online_interleave.py): full-drain-then-terminate, trainer-stops-first cooperative wind-down, producer-exception propagation + consumer unblock. GPU test for the named builder's interleaved run() added to test_disagg_online_launch.py. Stacked on O1.1; implements stage O1.2 (#618). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…erized builders Adding a draft model is now a StrategySpec entry, not a new build_*_runtime family. The topology stays a named builder; the model becomes a `strategy=` parameter resolved through a registry. launch.py no longer grows as (topologies x models). - registry.py (new): StrategySpec + register_strategy/resolve_strategy/ available_strategies + concat_collate. eagle3 spec fully wired (reader/transform/collate/online-collate/adapter). - launch.py: extract _assemble_trainer + _assemble_rollout_workers shared by every topology (offline / disagg-offline / online / disagg-online producer+consumer + one-process + interleaved); each builder takes `strategy=` and resolves a spec. eagle3-named builders kept as back-compat aliases; eagle3 behavior is byte-identical. - scripts/train_eagle3_dataflow.py, examples/disagg/run_disagg_eagle3.py: use the strategy-neutral builders. - tests/test_runtime/test_strategy_registry.py (new, CPU): registry / alias / unwired-strategy-guard contract. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…e + online)
DFlash now trains through the runtime via a StrategySpec + a DFlashAdapter, with
ZERO launch.py changes (the spec seam from the previous commit carries it).
- registry.py: dflash spec — offline reader (OfflineManifestReader with dflash
feature_keys, no aux/target swap), per-sample transform, padding collate; online
via DFlashAdapter; supports_online=True.
- inference/dflash_adapter.py (new): wraps generate_dflash_data, emits
{input_ids, hidden_states, loss_mask}; verify_capture self-skips the eagle3
aux/target checks (different feature names + __aux_layer_ids__=None).
- tests/_fixtures.py: write_offline_files_dflash + build_dflash (tiny Qwen3 target
-> DFlash draft + TargetEmbeddingsAndHead -> OnlineDFlashModel).
- tests/test_dflash_launch.py + test_dflash_online_launch.py (new, GPU): offline
and online dflash train end-to-end through FSDP.
- tests/test_strategy_registry.py: dflash-fully-wired assertions.
DFlash is online-only in production (no offline dumper exists yet — prepare_
hidden_states.py is eagle3-only), so the offline path is exercised with synthetic
fixtures while online is its real workflow.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…dent loss
Domino is the third algorithm on the composable launch — a StrategySpec plus the
ONE genuine shared-contract extension the analysis predicted.
- strategy.py: StepContext{global_step, total_steps} threaded into forward_loss
(optional; eagle3/dflash ignore it). DominoTrainStrategy: reuses the DFlash
feature schema + adapter; its forward_loss reads ctx to compute the decaying
lambda_base that blends Domino's base loss (mirrors train_domino.get_lambda_base).
- trainer.py: TrainerCore.train_step / eval_step accept a StepContext; fit passes
StepContext(global_step, total_steps=max_steps). Backward-compatible.
- contracts.py: DraftStrategyName += "domino".
- registry.py: domino spec — reuses DFlash transform/collate/adapter, domino reader
(strategy tag) + DominoTrainStrategy. No new builder, no launch.py change.
- tests/_fixtures.py: build_domino (DFlash draft w/ projector_type="domino" head ->
OnlineDominoModel).
- tests/test_domino_launch.py (new): CPU lambda-schedule test + offline/online GPU
end-to-end.
Adding domino touched ZERO launch.py and reused the dflash data path — exactly the
"new algorithm = a spec + its loss" goal.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… boundary
Extract a backend-agnostic `TargetEngine` ABC (modeling/target/base.py) with a
generic `capture(...)` entry point + a real `backend` tag, replacing the two
EAGLE3-/DFlash-named ABCs as the shared base:
TargetEngine
├── Eagle3TargetEngine (was Eagle3TargetModel) {HF,SGLang,Custom}
└── DFlashTargetEngine (was DFlashTargetModel) {SGLang,HF}
- `capture()` / `set_capture_layers()` are thin dispatchers onto the unchanged
`generate_eagle3_data` / `generate_dflash_data` / `set_aux_hidden_states_layers`,
so extraction is byte-identical — this PR is pure structure/naming, no logic.
- Real `backend` class attr on every leaf ("sglang"/"hf"/"custom"); the inference
adapters' health() stops reading getattr(..., "unknown").
- Generic `get_target_engine(strategy=, backend=)` factory (modeling/target/factory.py),
loaders imported lazily so `import specforge` still works without the pinned sglang
(dflash_target_model imports sglang unconditionally — kept off the eager path).
- All pre-Phase-B names (`Eagle3TargetModel`, `get_eagle3_target_model`, ...) kept as
aliases; scripts/tests untouched. `sglang_server` factory branch lands in B2.
- Add tests/test_runtime/test_target_engine_abc.py (hierarchy, backend tags, capture
dispatch, alias identity, factory dispatch).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ang version
Extract EVERY sglang internal + the duplicated extend/capture forward into one
version-pinned boundary, `sglang_backend/capture.py::SGLangCaptureBackend`, and
have the algorithm engines COMPOSE it instead of embedding it:
SGLangCaptureBackend (the only place that imports sglang.srt.* for capture)
· build() ServerArgs / ModelConfig / SGLangRunner wiring (unified)
· _forward_extend() the single ScheduleBatch/ForwardBatch capture forward
· _maybe_prepare_mlp_sync_batch() ONE (0.5.9) prepare_mlp_sync signature
· extend / extend_vlm / extend_dflash / get_rope_index / set_eagle3_capture_layers
SGLangEagle3TargetEngine / SGLangDFlashTargetEngine now hold a backend and do
only torch-side output shaping — they import ZERO sglang internals (verified by
tests/test_runtime/test_sglang_capture_backend.py, a pure-AST invariant).
Why: before this, both sglang engines imported ~20 sglang symbols and each carried
its own near-duplicate `_extend`; the two copies had drifted to DIFFERENT sglang
API versions (eagle3 = module-level prepare_mlp_sync_batch_raw(attn_cp_size=);
dflash = the removed Scheduler.prepare_mlp_sync_batch_raw(spec_algorithm=)). A
sglang bump touched every subclass and the copies could silently diverge. Now a
bump touches one file; "put the pieces together" (capture backend + shaping +
adapter) instead of tangling the version into each algorithm.
Behavior:
- Byte-identical on the test configs (TP=1/2, dp=1): require_mlp_sync is False so
the unified mlp-sync branch is skipped identically; construction, req building,
the forward, splitting/shard logic, and pool-clear ordering are transplanted
verbatim (`import specforge` stays sglang-optional via lazy import in
from_pretrained; the engine forward is still under @torch.no_grad).
- Two deliberate, flagged changes: (1) DFlash's mlp-sync now uses the same 0.5.9
signature as eagle3 — its old Scheduler.* call was latent-broken for dp>1;
(2) dropped a stray debug print() in DFlash set_capture_layers.
Also adds the `sglang_server` backend (SGLangServerEagle3TargetEngine): selectable
via get_eagle3_target_model(backend="sglang_server"), construction raises an
actionable NotImplementedError until the live-capture depth is set by the O1.3
spike (docs/roadmap/online-disaggregation.md §O1.3).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Introduce the domain training layer `specforge/training/` with a caller-facing
`Trainer` that composes the whole training spine behind one object + `.fit()`:
FeatureDataLoader + FSDPTrainingBackend.prepare_model (FSDP wrap)
+ spec.make_strategy -> TrainerCore -> TrainerController
`Trainer` is the canonical assembler now; `launch._assemble_trainer` delegates to
it and returns the same `(TrainerController, FeatureDataLoader)` tuple, so every
`build_*_runtime` path is byte-for-byte unchanged (no fork — one wiring path).
The runtime seam (TrainerController / TrainerCore / DraftTrainStrategy /
FSDPTrainingBackend) is untouched; this is the domain facade over it. Topology
(offline/online/disagg) stays invisible to Trainer — absorbed by the (ref source
+ FeatureStore) it is handed. No HiddenStateStream: the loader is the stream.
- specforge/training/{__init__.py (PEP 562 lazy Trainer export), trainer.py}
- launch.py: _assemble_trainer delegates to Trainer; drops the now-unused
FeatureDataLoader / FSDPTrainingBackend / ParallelConfig / TrainerCore /
TrainerController imports.
- tests/test_runtime/test_domain_trainer.py: fakes the runtime pieces and asserts
the composition (refs enqueued, loader/backend/core/controller args, ack_fn
wired to the DataFlowController, .fit() delegates over the loader).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…docs + gate) Makes the Phase B abstractions load-bearing instead of forward-only scaffolding. 1. adapter cutover: SGLangAdapter / DFlashAdapter now call the engine's generic `capture(...)` instead of `generate_eagle3_data` / `generate_dflash_data` (which stay as byte-identical back-compat aliases). This is the dataflow runtime's real capture path, so `TargetEngine.capture()` is now exercised. 2. factory cutover: the dataflow test fixtures + online launch tests build targets via `get_target_engine(strategy=, backend=)`; the old `get_eagle3_target_model` / `get_dflash_target_model` stay as shims (legacy scripts untouched). 3. docs: runtime/inference/DESIGN.md, runtime/ARCHITECTURE.md, and the adapter docstrings now describe the boundary as `TargetEngine.capture` (sglang glue in SGLangCaptureBackend) instead of `Eagle3TargetModel.generate_eagle3_data`. 4. explicit B gate: tests/test_runtime/test_phase_b_gate.py — `capture() ≡ legacy generate_*_data()` **bytewise** (eagle3 + dflash), a stable captured-feature digest, and per-step-loss reproducibility through the full online capture→train path (which now routes through `capture()`), instead of relying only on the full-suite's indirect coverage. No behavior change: `capture()` dispatches to the legacy method, so byte-identical. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Colocated runs now pay nothing for the disagg control plane on the SAME code path (plan.md §C / roadmap domain-refactor.md §C). DeploymentMode becomes load-bearing rather than decorative. - NoOpMetadataStore: a MetadataStore that retains nothing (no dedup index, no durable marker) for local_colocated runs. - resolve_control_plane(mode, run_id) -> (controller, durable_ack): local_colocated gets the no-op store + durable_ack=False (loader releases features as it consumes them); dataflow_colocated/disaggregated keep the durable store + optimizer-boundary ack. Disagg builders are untouched. - build_offline_runtime / build_online_runtime route through the selector and thread durable_ack; the domain Trainer skips the ack transaction and the offline enqueue when durable_ack is False. Gates: test_noop_metadata_store (contract + selection, CPU) and test_colocated_vs_disagg_equiv (per-step loss bit-identical across the no-op colocated and full disagg control planes — only the control plane differs). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…resolver, stronger gate Post-review fixes for #636 (adversarial review vs the #630 roadmap): - build_offline_runtime / build_online_runtime take deployment_mode (default local_colocated) + metadata_db_path, so the mode parameter actually varies at the builder surface instead of being a hardcoded literal, and tests/tools can drive any mode through the ONE builder. - resolve_control_plane: drop the never-passed sample_queue/backpressure kwargs; the non-colocated arm now delegates the default store to DataFlowController's own InMemory default instead of re-encoding the selection policy (SQLite only when metadata_db_path is given). - Domain Trainer: the colocated offline path (durable_ack=False) keeps the control plane's assert_no_tensors contract on every ref even though the durable enqueue is skipped. - test_colocated_vs_disagg_equiv: both legs now run through build_offline_runtime (no hand-copied assembly); the disagg leg uses a real SQLiteMetadataStore and asserts the durable ack marker was written; the deterministic-algorithms flag is restored after the test instead of leaking process-wide. - test_noop_metadata_store: cover metadata_db_path -> SQLite selection (and local_colocated ignoring it). - Docs: DESIGN.md + NoOpMetadataStore state the scope of the no-op axis explicitly (leasing/queue bookkeeping stays shared; backpressure is already opt-in) and the consequences of retaining nothing (status() zeros, no dedup, no TrainLease ref reconstruction). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Self-review of the fix commit confirmed the new deployment_mode / metadata_db_path parameters were a hazard on the ONLINE builder: the online loader consumes a rank-private in-process SampleRefQueue that is fed only by commit_sample()==True, so a shared durable store dedups each sample onto exactly one rank's queue (divergent per-rank batch counts -> mismatched FSDP collectives -> hang), and a reused db file starves a restarted run to zero steps. The offline builder keeps the parameter (its refs-mode loader iterates the full ref list regardless of commit dedup, which is what the equivalence gate exercises); online durable-store runs remain the disagg online builders' job. Documented at the pin site. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…-control-plane [DataFlow runtime · online] O1.1 — shared cross-process control plane
…loop [DataFlow runtime · online] O1.2 — named builder + interleaved async loop
…unch [DataFlow runtime] Composable launch: StrategySpec registry + parameterized builders
[DataFlow runtime] DFlash end-to-end on the composable launch (offline + online)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
[DataFlow runtime] Domino end-to-end + StepContext for schedule-dependent loss
[DataFlow runtime] Phase B1 — TargetEngine ABC + de-EAGLE3 the target boundary
…e-backend [DataFlow runtime] Phase B2 — decouple the target engine from the sglang version
[DataFlow runtime] Phase B3 — domain Trainer wrapping the runtime seam
…-cutover [DataFlow runtime] Phase B4 — adopt the de-EAGLE3 surface (cutover + docs + gate)
…htweight [DataFlow runtime] Phase C — colocated lightweight control plane
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Ports the sglang 0.5.14 upgrade onto up-16's new-layout target-capture backend. On `main` the 0.5.14 bump landed via #605 (eagle3 / custom backends) plus the follow-up `patch/fix_sgl0.5.14_dflash` (dflash), but both target the OLD `specforge/modeling/target/dflash_target_model.py` layout. up-16 consolidated eagle3 + dflash extend into a single `SGLangCaptureBackend` (`inference/target_engine/sglang_backend/capture.py`), so the same API adaptations are applied there once, covering both paths: - import `prepare_mlp_sync_batch_raw` from `managers.scheduler_components.dp_attn` (moved from `managers.scheduler_dp_attn_mixin` in 0.5.14). - `build()`: replicate the post-load setup 0.5.14 split out of `ModelRunner.initialize()` — `alloc_memory_pool()` / `init_attention_backends()` / `init_cuda_graphs()` — since we drive the ModelRunner directly; and set `chunked_prefill_size=-1` (we prefill the whole batch in one `_forward_extend`, no scheduler chunking). - `_forward_extend()`: drop the removed `get_model_worker_batch()` step — `ForwardBatch.init_new` now consumes the `ScheduleBatch` directly and reads `capture_hidden_mode` off it; and materialize `prefill_input_ids_cpu` -> device (the scheduler normally does this in `resolve_forward_inputs`). - Req construction (dflash + eagle3 + eagle3-vlm): `fill_ids` was removed in favor of `full_untruncated_fill_ids` (array) + integer `fill_len`. - capture `input_lens` BEFORE the forward — `origin_input_ids` is released during prepare_for_extend/forward on 0.5.13+. - bump pyproject pin to `sglang==0.5.14`. Benefits eagle3 + dflash (+ dspark, which builds on the dflash path). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Two more 0.5.14 API drifts in the vendored parallel_state setup (surfaced building the target on 0.5.14, mirrors main's post-#605 patch.py): - init_model_parallel_group() dropped the pynccl_use_current_stream kwarg (removed in 0.5.13) — remove it from the tp and pdmux_prefill_tp calls. - compute_dp_attention_world_info() now returns a 4-tuple (attn_tp_rank, attn_tp_size, attn_dp_rank, attn_dp_size); the attn-tp rank/size are no longer module globals, so keep only _ATTN_DP_RANK. Validated: sglang DFlash capture (Qwen2.5-0.5B) runs on sglang 0.5.14 and returns correctly-shaped concatenated hidden states. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The 0.5.14 port added multi-line rationale comments at each sglang-API adaptation; several were 5-6 lines and the Req.fill_ids explanation was duplicated verbatim across the three req-construction sites. Collapse each to a concise "<version>: what changed / what we do" line, keeping the version markers (the useful part in a version-coupling boundary file) and dropping the prose. Comments only; no code changes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…iter replaced_logits_processor_forward_for_eagle3 read self.multi_item_delimiter, which sglang removed from the LogitsProcessor in 0.5.13 (now carried per-request on the ForwardBatch as multi_item_delimiter_indices). On 0.5.14 that attribute access raises AttributeError on the eagle3 capture path (wrap_eagle3_logits=True); dflash sets it False, so the dflash smoke test did not exercise this. Extract the indices off the ForwardBatch before the LogitsMetadata conversion, mirroring sglang's own logits_processor. Matches the fix in #605. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
[SpecForge] examples + curves for DFlash/EAGLE3 (Qwen3.6-27B & Qwen2.5-0.5B); retire docs/reports/e2e
fix(sglang): port the DataFlow capture backend to sglang 0.5.14 (up-16)
…py Mooncake transport The engine-side transport PR #641's reforward anticipated: a patched live SGLang server captures aux/last hidden states during prefill and writes them straight into the Mooncake store in MooncakeFeatureStore's zero-copy layout; the rollout side commits SampleRefs built from response metadata alone — feature tensors never cross the HTTP/worker path, one forward per sample. sglang side (patches/sglang/v0.5.14/spec-capture.patch, 8 hooks + 1 module): - server args: --enable-spec-capture, --spec-capture-aux-layer-ids (answers the stock model_runner TODO: aux capture without a speculative draft worker) - per-request spec_capture dict threaded GenerateReqInput -> Req; capture rides the existing return_hidden_states/CaptureHiddenMode.FULL path - logits processor additionally exposes post-norm last_hidden_states under FULL+aux (the eagle3 target_repr="hidden_state" representation) - scheduler sink (sglang/srt/spec_capture_sink.py, self-contained): raw-bytes put_from at {store_id}/{sample_id}/g{gen}/{name}, hard-pinned; results ride the existing customized_info channel into meta_info (no new IPC fields) SpecForge side: - adapters/server_capture.py: SGLangServerCaptureAdapter (RefSource) + per-strategy ServerCaptureSchema registry (eagle3 / dflash / domino) — the server knows only generic artifacts; naming/shapes are client-decided - rollout_worker: RefSource path (produce_refs -> verify-from-specs -> commit), per-task failure markers keep lease bookkeeping exact - capture.verify_feature_contract_specs: the same loud extraction-boundary checks, run from FeatureSpec shapes (no tensor fetch) - MooncakeFeatureStore.adopt(): lifecycle registration for server-written samples so producer-side abort/gc frees them - tests: CPU contract suite over a stub server + fake mooncake backend; opt-in GPU gate (SPECFORGE_RUN_SERVER_CAPTURE_TESTS=1) boots a patched server + mooncake master, checks HF-reference parity, trains one step
…ention target_repr='hidden_state' routes through TargetHead.preprocess, which adds the trailing mask dim itself; storing (B, L, 1) broadcast position_mask to (B, T, T) in the loss. Caught by the GPU gate's real train step.
- run_disagg_dflash.py + run_qwen3.6_27b_dflash_disagg.sh: the disaggregated DFlash example now uses the server-capture transport — a patched live SGLang server captures to Mooncake, a thin CPU producer drives prompts, the trainer consumes zero-copy. Replaces the in-process staged producer. - launch.py: build_disagg_online_producer / _assemble_rollout_workers accept an optional feature_source (RefSource) so the server-capture adapter drives the producer without an in-process target_model. - trim over-verbose comments/docstrings across the sglang patch + new modules; regenerate patches/sglang/v0.5.14/spec-capture.patch. - exclude *.patch from whitespace/eof hooks (stripping context whitespace breaks git apply); add CPU test for the feature_source producer wiring.
[Online track] O2 — server-side spec capture: sglang 0.5.14 patch + zero-copy Mooncake transport
…tion metric means) so the DFlash consumer is monitorable without W&B
…e-method; fixes dflash server path (VL targets) + example uses it
…stomized_info misalignment under batched traffic); client reads result dict + error marker
…erver segment Hard-pinned features live in mounted client segments: objects placed in the producer's segment died with its exit, and 10GB total capacity failed puts at ~140 samples (master log: memory_ratio>0.95, eviction frees nothing under hard pin). Producer/consumer now contribute no storage; the server sink's segment is the only store and must exceed the in-flight watermark.
Mooncake defers a physical free while the consumer's get() read-lease is live (remove -> -706); release() parks it in _release_pending and nothing retried: every consumed sample leaked until the segment hit the eviction watermark (hard-pinned objects never evict), failing further puts. The loader now pumps gc() at gc_interval_s (default 15s > lease TTL, spacing retries so the store's max_release_attempts isn't burned while the lease is still live).
…ase) Probed on a live master: is_exist -> immediate remove fails -706; a bare remove after the lease TTL succeeds. gc()'s exists pre-check therefore re-leased every pending object right before its retry, exhausting max_release_attempts and leaking 100% of consumed samples (both 64GB and 200GB segments filled at exactly cumulative-produced bytes). Removes now run first; the exist probe only classifies an already-gone key after a failed remove.
…ger ack One book-keeping authority instead of N: trainer rank 0 hosts a RefDistributor that alone reads the producer's ref channel, commits into the run's ONE durable metadata store, and round-robin dispatches aligned windows (dp_size refs, one per rank) to per-rank inboxes — equal per-rank counts keep FSDP ranks in lockstep. Ranks hold no channel offsets, no partition math, no ledger. - control_plane/dp_ack.py: DPAckController makes ack_train_refs a DP collective — every rank gathers its shard's ids at the optimizer boundary (the existing ack_fn seam), only rank 0 records the union + advances the producer's backpressure counter. - data_plane/ref_distributor.py: commit -> aligned dispatch -> inbox; skips released refs on restart re-read (marked consumed); end-of-stream remainder stays committed-unacked and replays on restart. - streaming_ref_channel.seed_consumed(): a restarted consumer adopts the sidecar count instead of rewinding it (fixes producer in-flight inflation). - launch.build_disagg_online_consumer: dp_rank/dp_size/inbox_dir (defaults from the torchrun world); single-rank path unchanged. - example: CPU-only producer (no torchrun/GPU/process group; prompts built without DistributedSampler), private producer dedup — the ledger is trainer-side only; TRAIN_DP knob in the qwen3.6 disagg script. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…CUDA probe) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…rt guards, loud failure - Producer counter now mirrors per-rank inbox acks (micro-batch granularity) instead of optimizer-boundary acks: kills the watermark deadlock when dp*batch*accum > in_flight_high_watermark, and the restart double-count of released refs (seed covers the prior run's counts; skips no longer marked). - Distributor death poisons every inbox with a .failed sentinel; InboxChannel readers raise immediately — no silent multi-rank hang, distinct from the clean .closed end-of-stream. DP path defaults idle_timeout_s to 1800s. - Rank 0 refuses a stale ledger without resume=True (commit-dedup would silently drop the whole re-streamed channel) and NoOp stores; db-arg and world-size preconditions raise symmetrically on ALL ranks (no barrier stranding). - End-of-stream partial window releases its queue leases (fail, not leak); honest log (re-queued only on a resume=True restart). - inbox_dir now opts a single-rank run into the distributor path (example dp1 keeps a consistent, self-contained ledger); DISAGG_IDLE_TIMEOUT env knob. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…captures Found by the first DP=2 run: one prefill batch hit a 64-token cached prefix, the server captured hidden states 64 tokens short of the prompt, and the trainer crashed on the rope shape mismatch 1570 steps in (the distributor's new .failed poisoning turned it into a loud failure instead of a silent hang). - patch: scheduler startup now also requires --disable-radix-cache (same guard style as chunked-prefill; capture must cover the whole sequence). - adapter: reject any capture whose seq len != prompt len (defense in depth, frees the corrupt keys and fails just that sample). - example server args: --disable-radix-cache. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
PR654's fix #1 added the flag (the VL target only wires aux capture via the dflash path; the default eagle3 hook crashes at capture time on Qwen3.6-27B), but the example script never passed it — the validated runs launched the server by hand. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…vers build_disagg_online_producer(feature_source=[...]) now maps one adapter to one RolloutWorker (1 server : 1 adapter : 1 worker); drive_producer runs one thread per worker so N patched servers prefill concurrently over disjoint prompt leases into the one Mooncake namespace + ref channel. Each ref records its capture server in metadata (post-hoc provenance audit). Failure semantics replace the truncating `if not refs: break`: - drained = no prompt pending OR leased (an all-failed round is not EOF — the PR654 known rough edge); - a dead server's worker fails its leases retryable (survivors re-lease), is dropped after max_worker_failures consecutive errors; all workers dead with prompts remaining raises instead of silently truncating; - max_prompt_attempts bounds per-task retryable failures (poisoned prompt goes terminal, not infinite spin) via a new DataFlowController knob. Example: run_qwen3.6_27b_dflash_disagg_multiserver.sh — 2x TP=2 servers + DP=2 trainer on one 8-GPU node over DISAGG_SERVER_URLS in the dflash launcher (comma-separated; single-URL path unchanged). CPU tests: tests/test_runtime/test_disagg_multiserver.py (disjoint+complete 2-server production incl. per-server ref provenance, dead-server failover, all-dead raise, poisoned-prompt termination, source/worker-count conflict). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
[Online track] O2 follow-up — server-capture fixes from the first real 5k-step disaggregated run
[Online track] O4: multi-server producer — fan the disagg producer out over N capture servers
[Online track] Roll up O3/O4 into up-16
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
dataflow-up-16-zerocopyintomain.Notes
Testing