Add Domino draft model training

[SeaTuKeMa]

Replaces #118, which was auto-closed for inactivity.

Summary

• add a Domino draft model with a DFlash backbone and causal low-rank correction head
• add the base-anchored curriculum training wrapper and trainer dispatch
• register Domino configuration/model loading and provide a reference draft config
• add CPU tests for configuration, forward behavior, curriculum semantics, gradient flow, and learning

Why

This adds training support for the Domino speculative decoding method while reusing TorchSpec's existing DFlash anchor sampling, masking, metrics, and distributed trainer path.

Addresses #114.

Usage

# Prepare data
python scripts/tools/prepare_perfectblend.py \
  --output data/perfectblend_50k.jsonl --sample-size 50000

# Train (8x H100: 4 inference + 4 training FSDP)
# Use a persistent volume such as /workspace/checkpoints on RunPod.
RUN_NAME=qwen3-8b-domino-8h100-500 \
TRAIN_DATA_PATH=data/perfectblend_50k.jsonl \
OUTPUT_ROOT=/workspace/checkpoints \
NUM_TRAIN_STEPS=500 \
SAVE_INTERVAL=250 \
MAX_CHECKPOINTS=1 \
SAVE_OPTIMIZER=false \
./examples/qwen3-8b-domino-8h100/run.sh

# Verify checkpoint output
cat /workspace/checkpoints/qwen3-8b-domino-8h100-500/checkpoints/latest_checkpointed_iteration.txt
ls /workspace/checkpoints/qwen3-8b-domino-8h100-500/checkpoints/iter_NNNNNNN

The run script accepts normal config overrides after the optional config path, for example:

./examples/qwen3-8b-domino-8h100/run.sh \
  configs/sglang_qwen3_8b_domino_2gpu.yaml \
  training.num_train_steps=20

Validation

• pytest tests/test_domino.py tests/test_dflash.py -q (90 passed)
• ruff check on all changed Python files
• git diff --check
• JSON and Python syntax validation

GPU Phase 1

• Environment: Qwen3-8B Domino training with SGLang target inference on 2x A100-SXM4-80GB.
• Stability:
  • 100-step calibration completed; checkpoints saved at steps 50 and 100.
  • 500-step run completed; checkpoint rotation saved through final iter_0000501.
  • Final checkpoint state was verified: model, optimizer, LR scheduler, RNG, metadata, and latest-step marker.
• Notes:
  • SGLang max-context truncation warnings appeared for a few 3070-3072 token requests but did not stop training.
  • A Ray queue actor warning appeared during shutdown after the final checkpoint had already been saved.

Run	Data	Steps	Time	Inference	Training	Checkpoint result
Phase 1A	2K PerfectBlend subset	100/100	705.5s	2.4 entries/s	2.3 entries/s	step 50, step 100
Phase 1B	10K PerfectBlend subset	500/500	2541.2s	3.2 entries/s	3.1 entries/s	final iter_0000501

Phase 1B 20-step window	Avg loss	Avg acceptance	Avg accepted length
Steps 1-20	13.504	0.019	0.020
Steps 81-100	6.569	0.083	0.108
Steps 181-200	5.484	0.175	0.303
Steps 281-300	4.917	0.229	0.436
Steps 381-400	4.666	0.251	0.499
Steps 481-500	4.557	0.253	0.512

GPU Phase 2

• Environment: Qwen3-8B Domino training with SGLang target inference on 8x H100-SXM.
• Configuration: 4 SGLang inference engines plus 4 FSDP training ranks via examples/qwen3-8b-domino-8h100/run.sh.
• Data: PerfectBlend 50K request produced 47,484 valid normalized samples and 47,265 tokenized train samples.
• Result:
  • 500/500 steps completed in 614.6s.
  • Checkpoints were written directly to a RunPod Network Volume under /workspace/checkpoints.
  • Step 250 checkpoint saved successfully at iter_0000251.
  • Final checkpoint saved successfully at iter_0000501; latest_checkpointed_iteration.txt reported 501.
  • Run completed with average inference throughput 13.1 entries/s and training throughput 13.0 entries/s.
• Profiling:
  • TorchSpec perf metrics were enabled for all steps.
  • 1-second GPU sampling was collected during the run for later bottleneck review.
  • Per-step progress reported steady-state training throughput around 20-22 entries/s while the async sample pool stayed full.

Phase 2 window	Avg loss	Avg acceptance	Avg accepted length
Steps 1-20	13.987	0.017	0.017
Steps 231-250	5.147	0.207	0.360
Steps 451-500	4.447	0.271	0.541
Steps 481-500	4.426	0.271	0.546

loss:       13.987 -> 5.147 -> 4.426
acceptance:  0.017 -> 0.207 -> 0.271
acc_len:     0.017 -> 0.360 -> 0.546

Scope

This PR implements training support. Serving/export integration for the fused Domino head remains follow-up work.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 6852d5d928

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 364cd3778b

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: afe66884e7

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Preserve explicit zero curriculum steps

When a run sets training.domino_curriculum_steps=0 to disable the curriculum, this or treats the explicit zero as missing and falls back to lr_total_steps, so _compute_curriculum_lambda() starts at 1.0 and anneals instead of returning 0.0 as its total <= 0 branch intends. This silently changes no-curriculum ablation experiments; use an is None check so only omitted values default to the LR schedule length.

Useful? React with 👍 / 👎.

[Dogacel]

Hi @SeaTuKeMa , thanks for the PR. Can you help me understand how this is different from DSpark? I would like to see if their performance numbers are similar or whether one of them is a clear winner.

The implementation looks mostly clean to me. I recommend cleaning the run script.

[SeaTuKeMa]

Hi @SeaTuKeMa , thanks for the PR. Can you help me understand how this is different from DSpark? I would like to see if their performance numbers are similar or whether one of them is a clear winner.

The implementation looks mostly clean to me. I recommend cleaning the run script.

Thanks @Dogacel . I’m trying to resolve this issue (#114), which was created before DSpark was released.

My understanding is that Domino and DSpark are still architecturally different, but I don’t have an apples-to-apples performance comparison yet. I’d like to involve the Domino author @jianuo-huang, since he may already have run relevant experiments and may have results that help answer this question.

Yes, I can clean up the run script in the next iteration.