[DataFlow runtime] Domino disaggregated training (example + launcher + strategy params + capture/tracker fixes)#660
Merged
Conversation
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
jiapingW
requested review from
FlamingoPg,
FrankLeeeee,
shuaills and
sleepcoo
as code owners
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
Collaborator
Author
Curve here, also checked the accept length, 5.9 main and 5.75 refactored (acceptable diff?)
|
maocheng23
force-pushed
the
dataflow-up-16-zerocopy
branch
from
dc96359 to
d3e4c4f
Compare
jiapingW
force-pushed
the
domino-disagg
branch
2 times, most recently
from
13067d2 to
ef75987
Compare
… 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>
…nly (move-only)
Move the domain planes out of specforge/runtime/ into their S-homes:
- specforge/runtime/training/* -> specforge/training/{controller,backend,strategies/{base,registry}}.py
- specforge/runtime/inference/* -> specforge/inference/{rollout_worker,capture,adapters/{eagle3,dflash}}.py
- specforge/modeling/target/** -> specforge/inference/target_engine/** (incl. sglang_backend/)
- specforge/runtime/launch.py -> specforge/launch.py
Every old path keeps a thin import shim re-exporting from the new home;
tests/scripts/examples switch to the new import paths. No behavior change.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
DRAFT_REGISTRY + @register_draft in modeling/draft/registry.py; the auto loaders resolve architecture -> (class, config_class) from the registry first (AutoDraftModelConfig.from_file, AutoEagle3DraftModel.from_config), keeping the hardcoded mappings only as fallback. LlamaForCausalLMEagle3 and DFlashDraftModel register themselves; DFlashDraftModel config JSONs become loadable through AutoDraftModelConfig for the first time. New tests/test_modeling/test_draft_registry.py. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
config/schema.py: validated Config (model/data/training) mapping 1:1 onto the DataFlow launch builders; YAML/JSON loader + dotted overrides re-validate through the schema. New specforge console script drives eagle3 offline/online via the same builder wiring, threading dtype/device, the capture contract, and the resume/rotation/control-plane knobs; build_online_runtime gains log_interval. Tests: schema mechanics + a config-built == programmatic gate. Co-Authored-By: Claude Fable 5 <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>
2. del some test code to increase speed 3. refactor code without influence training
…--num-epochs
- examples/disagg/run_qwen3_8b_domino_disagg_1srv_dp7.sh: single inference
server (1 GPU) + DP=7 trainer (7 GPUs) variant of the multiserver launcher.
The domino disagg workload is trainer-bound and the capture server is
heavily overprovisioned in the 2-server layout (~6% util); this rebalances
to 1 server + 7 trainers, which raises both throughput and server util.
- examples/disagg/run_disagg_domino.py: the producer now replicates the prompt
pool by --num-epochs before streaming. The online streaming channel is
consume-once, so previously the pool streamed once regardless of
--num-epochs; multi-epoch training now works as expected.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
perf(domino-disagg): env-tunable data-plane knobs + findings (28.8->50.1 samples/s)
Config-only tuning that ~doubles single-node online domino-disagg throughput,
with util/SM/mem healthy on all 8 GPUs. All knobs are off by default.
Launchers (1srv+dp7 and multiserver):
- DISAGG_SERVER_URLS now overridable; repeating one URL N times spawns N
concurrent producer workers with disjoint leases against one server, which
breaks the single-producer supply ceiling (the blocking HTTP prefill releases
the GIL so workers genuinely overlap).
- SERVER_MEM_FRACTION knob (default 0.85): 0.5 right-sizes the capture-only
server from ~126 GB to ~78 GB with no perf cost.
- NUM_ANCHORS / BLOCK_SIZE knobs on the 1srv launcher (quality<->throughput).
- multiserver launcher now honors BATCH_SIZE / ACCUM (were hardcoded to 2 / none).
FeatureDataLoader:
- LOADER_PREFETCH=N: background-thread batch prefetch so the training step never
pays mooncake get/collate latency inline (ack still after the trainer consumes).
- CLONE_ON_FETCH=0: skip the redundant defensive clone on the mooncake zero-copy
path (get() already allocates a fresh tensor).
See examples/disagg/PERF_FINDINGS.md for the full analysis: the pipeline is a
supply/demand seesaw (1:7 beats 2:6), bigger batch does not help (compute-latency
bound, MFU ~14%), and ~78% of per-sample trainer compute is the num_anchors=256
draft expansion.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
perf(domino-disagg): env-gated profiling instrumentation for the data path
Default-off timers that produced the PERF_FINDINGS analysis. Each is gated on an
env var and adds nothing when unset.
- PROFILE_PRODUCER=N -> [prod] (backpressure-park / run_once / publish + in_flight)
in launch.py, and [prod-http] (build / HTTP RTT / parse)
in the server-capture adapter.
- PROFILE_LOADER=N -> [loader rK] queue-wait / get / clone / release / gc.
- PROFILE_STORE=N -> [store rK] get/put GB/s, rm_ok/rm_fail(-706), pending depth.
- PROFILE_DISTRIB=sec -> [dist] commit / count / inbox-publish + per-rank lag.
- PROFILE_STEPS=N -> [profile] data-wait vs compute, [profile2] fwd/bwd/opt + padded len.
- PROFILE_TORCH=N -> rank-0 torch.profiler top-ops + chrome trace.
- FSDP_SHARDING -> override the FSDP sharding strategy (e.g. NO_SHARD).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
revert(domino-disagg): drop NUM_ANCHORS/BLOCK_SIZE knob; num_anchors is not a throughput lever
Reducing num_anchors raises sequences/s but LOWERS training-signal/s
(anchor-updates/s): each anchor is a supervised target, and a fixed ~53ms/
microstep per-sequence overhead means fewer anchors amortize that cost over less
signal (256 -> ~2130 vs 64 -> ~1280 anchor-updates/s). It also forces ~4x more
captured sequences for equal signal, loading the supply bottleneck. So it is a
quality/data-efficiency setting, not a speed knob — revert the launcher override
back to the fixed --num-anchors 256 / --block-size 16 and correct PERF_FINDINGS.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
docs(domino-disagg): measured MFU/FLOPs + correct the inference-ceiling analysis
Adds bench_domino_mfu.py (isolated single-GPU, FLOP-counted, CUDA-event-timed
trainer benchmark) and folds its results into PERF_FINDINGS.md:
- Trainer MFU ≈ 44% (measured, b2/b4/b8; ~45 TFLOP/sample fwd+bwd; per-sample
time flat across batch → compute-bound, not stalled). Supersedes the earlier
"MFU ~14% / occupied-but-stalled" note, which came from a cuda.synchronize-
distorted timing and was wrong.
- Inference (target prefill) ≈ 43% MFU, estimated from ~27k tok/s at ~550-tok
prompts (2·params·tokens for the 8B target).
- Corrects the "supply ceiling = capture sink" framing: at >=8 producer workers
the server is PREFILL-COMPUTE-BOUND (GPU prefills ~100% duty at ~27k tok/s);
the sink thread keeps up (~55/s) and is only the *next* wall. So an async/
parallel sink is a second-order lever, not the primary fix — the real supply
levers are higher prefill MFU (CUDA graphs / bigger prefill batches / lighter
aux-hidden capture copy), fewer aux layers (training change), or more inference
GPUs (the seesaw). Trainer-MFU work does not raise system throughput while
supply-bound.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
fix lint
fix lint
Match stock train_dflash.py's logging so the disagg (online/DP) curves overlay the in-process trainer: - controller._result: all_reduce(loss)/world and all_reduce(acc)/world before logging (was: single rank's local-batch value -> ~sqrt(world) noisier, spiky). - _dflash_family_disagg.py logger: log accuracy under train/accuracy (was train/acc) and add train/lr from the optimizer. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…eps) Adds the DFlash counterpart of the Domino 1srv+DP7 launcher so DFlash disagg training is reproducible from the repo: - run_disagg_dflash.py entry, qwen3-8b-dflash.json, drops domino-only --lambda-base-* args, NUM_ANCHORS/WANDB_NAME overridable. - auto-derives DISAGG_TOTAL_STEPS = ceil(NUM_EPOCHS * prompts / (TRAIN_DP * BATCH_SIZE)) from the dataset when unset, so only --num-epochs is needed. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add a reusable two-node Qwen3-8B DFlash launcher with rank-specific inference and training roles, cross-node Mooncake coordination, allocated-GPU UUID handling, and symmetric failure cleanup. Document the RCLI launch flow and apply the formatter changes required by PR 660's lint job.
Apply suggestion from @gemini-code-assist[bot] Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Apply suggestion from @gemini-code-assist[bot] Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> fix lint and update filename
feat(data): add validated Qwen ShareGPT regeneration pipelines Apply suggestion from @gemini-code-assist[bot] Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Apply suggestion from @gemini-code-assist[bot] Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> fix lint and update filename fix lint and reuse code in conversation_valiation.py
Apply suggestion from @gemini-code-assist[bot] Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> fix lint and upgrade comment and dir fix ckpt_sort
jiapingW
approved these changes
Jul 13, 2026
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
3 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
Adds end-to-end Domino disaggregated training on top of the DataFlow runtime integration branch (
dataflow-up-16-zerocopy), plus a few robustness fixes surfaced along the way. Cleanly stacked: 2 commits, +718 / −71 across 12 files, no rebase needed.Commits
6cd68297579d8eWhat's in it
New Domino disagg entrypoints
examples/disagg/run_disagg_domino.py(+298) — end-to-end domino disaggregated training exampleexamples/.../run_qwen3_8b_domino_disagg_multiserver.sh(+216) — multiserver launcher; plus edits to the qwen3.6-27b dflash disagg launcherconfigs/qwen3-8b-domino.json— addsemb_dim: 256Parameterized domino strategy
training/strategies/registry.py— domino strategy now acceptslambda_start/decay_ratiovia a new_make_domino_strategybuildertraining/trainer.py+launch.py— thread a newstrategy_kwargsdict throughTrainer,_assemble_trainer, andbuild_disagg_online_consumerso strategy construction can be parameterizedregistry.py— online collate fix: bothdflashanddominonow use_dflash_offline_collatefor the online path instead ofconcat_collate(behavioral change — worth a close look)Robustness fixes
inference/adapters/server_capture.py(+51) — handle batch/list-wrapped capture results from the spec-capture server; select this task's row bysample_id, with explicit errors on ambiguous or non-object results (_flatten_list_wrappers,_capture_result_for_task)tracker.py— treat a localwandb/output dir mis-imported as a namespace package as "wandb unavailable"; raise a clear error if--report-to wandbis used without the real clientruntime/data_plane/mooncake_store.py— include the underlying exception + CUDA-runtime hint in the mooncake load-failure messageTests
tests/test_runtime/test_server_capture.py(+19) — batch-result selectiontests/test_runtime/test_strategy_registry.py(+22) — domino strategy params / online collateReviewer notes
fix bugs,[bug]) — worth squashing/rewording before merge.make_online_collatechange fromconcat_collate→_dflash_offline_collateaffects the online path for both dflash and domino, not just domino — confirm this is intended for dflash too.tests/test_runtimeshould validate.Base
Targets
dataflow-up-16-zerocopy(the DataFlow runtime integration branch), consistent with how the rest of the stack lands. The merge-base is exactly that branch's current tip (#657 o3/o4 rollup).🤖 Generated with Claude Code