Skip to content

feat(training): enable opt-in Liger kernels for Domino and DSpark - #963

Merged
jiapingW merged 4 commits into
mainfrom
codex/pr942-07-domino-optimizations
Oct 26, 2026
Merged

feat(training): enable opt-in Liger kernels for Domino and DSpark#963
jiapingW merged 4 commits into
mainfrom
codex/pr942-07-domino-optimizations

Conversation

@jiapingW

@jiapingW jiapingW commented Oct 14, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Domino and DSpark can opt into the existing Liger RMSNorm/SwiGLU kernels with model.use_liger_kernel: true. Their default null and explicit false retain native kernels. Providers and constructors reuse the DFlash kernel bundle and its platform, activation, and dependency validation; preflight follows the same opt-in behavior.

Parameter layouts remain compatible, and the implementation does not globally patch Transformers.

Stack: main → #960 → #961 → #963 → #965. One Liger-only commit on #961; projection caching, flattened projections, and selective checkpoint contexts are excluded.

Full-stack validation on H20 (Python 3.12, PyTorch 2.14.1+cu130, Liger 0.8.5): 285 tests passed, 3 skipped (FlashAttention 2 unavailable and two Blackwell-only tests). Black 24.10.0, isort 5.13.2, and git diff --check passed. Coverage includes BF16 loss/gradient parity, opt-in/native defaults, configuration, and model loading.

This was referenced Oct 14, 2026
Closed
Open
Closed
@jiapingW
jiapingW force-pushed the codex/pr942-07-domino-optimizations branch from d19a88d to f8b7482 Compare October 14, 2026 09:00
@jiapingW jiapingW mentioned this pull request Oct 14, 2026
Merged
@jiapingW jiapingW changed the title [7/11] perf(domino): add Liger kernels and frozen projection caching [7/10] perf(training): share Domino and DSpark vocabulary projection controls Oct 14, 2026
yushengsu-thu added a commit to yushengsu-thu/SpecForge that referenced this pull request Oct 15, 2026
@yushengsu-thu @claude
dense-1/8: use the sgl-project#963 projection key names after the sta…
e44fee6 
…ck merge

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
yushengsu-thu added a commit to yushengsu-thu/SpecForge that referenced this pull request Oct 15, 2026
@yushengsu-thu @claude
dense-same: sync with spark-Qwen3-8B c3ff230 and rename exp_readme.md…
c884a97 
… to exp_read.me

Same base as the dense-1/8 branch (PR sgl-project#942 stack sgl-project#957-sgl-project#963 merged there, DDP-wrapped draft lookup for MoE
metrics, managed-local NCCL port pin); the experiment log is exp_read.me like on dense-1/8.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
yushengsu-thu added a commit to yushengsu-thu/SpecForge that referenced this pull request Oct 15, 2026
@yushengsu-thu @claude
Merge upstream PR stack sgl-project#957-sgl-project#963 (jiapingW, la…
5a6b088 
…yers 1-7 of the sgl-project#942 split), matching the dense-1/8 branch

The combined PR sgl-project#942 merge was reverted in the previous commit; this takes the split stack instead
(sgl-project#957 streaming-read fixes, sgl-project#958 DDP sync once per accumulation window, sgl-project#959 no implicit qwen3-instruct
system prompt, sgl-project#960 is_causal, sgl-project#961 fused DFlash head, sgl-project#962 selective checkpoint contexts, sgl-project#963 shared
flatten_vocab_projection / cache_vocab_projection), so the MoE arm carries exactly the same upstream code
as the dense-1/8 arm. Conflicts in specforge/training/backend.py and controller.py resolved in the stack's
favour.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
yushengsu-thu added a commit to yushengsu-thu/SpecForge that referenced this pull request Oct 15, 2026
@yushengsu-thu @claude
exp: align backend, controller, launch_plan and the recipe with the d…
b2a84fc 
…ense-1/8 branch

Taking the stack's side of the merge conflicts had dropped this branch's own changes to
specforge/training/backend.py and controller.py: the MoE checkpoint state-dict naming conversion
(to_checkpoint_state_dict / from_checkpoint_state_dict), the fsdp_replicate_when_frozen candidates, the
SPECFORGE_FSDP_SHARDED_GRAD_ACCUM path and the perf/mem_* allocator metrics. The files are now the
dense-1/8 versions (base + PR stack sgl-project#957-sgl-project#963), plus dense-1/8's launch_plan.py fix pinning each
capture server's --nccl-port, so specforge/ is identical across the two arms.

Recipe: flatten_vocab_projection / cache_vocab_projection (the sgl-project#963 names), save_interval 1196 with 12
checkpoints kept (9 per epoch, the 9th on the epoch boundary), in-flight window 264/132 and no
consumer_dispatch override, all as on dense-1/8.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
yushengsu-thu added a commit to yushengsu-thu/SpecForge that referenced this pull request Oct 15, 2026
@yushengsu-thu @claude
Merge the latest spark-Qwen3-8B (48c6038: PR stack sgl-project#957-sg…
da2fafd 
…l-project#963 + sgl-project#966) into the MoE experiment branch

The same upstream content was already merged here independently; this records the fork branch's
current head as an ancestor. The two MoE config files that spark-Qwen3-8B also carries are kept in
this branch's version (the one the running experiment uses).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
@yushengsu-thu yushengsu-thu mentioned this pull request Oct 21, 2026
Open
@maocheng23

maocheng23 commented Oct 21, 2026

Copy link
Copy Markdown
Collaborator

Two suggestions here:

  1. For Liger, can we split into a new PR instead of mixing together?
  2. flatten_vocab_projection, are we sure about the improvement here? I am not fully convinced here if we need a new param

@jiapingW
jiapingW force-pushed the codex/pr942-06-projection-checkpointing branch from f85e6ae to f5b38c6 Compare October 22, 2026 00:58
@jiapingW
jiapingW requested a review from zyksir as a code owner October 22, 2026 00:58
@jiapingW
jiapingW force-pushed the codex/pr942-07-domino-optimizations branch from f8b7482 to 554a0f1 Compare October 22, 2026 01:25
@jiapingW jiapingW changed the title [7/10] perf(training): share Domino and DSpark vocabulary projection controls perf(training): share Domino and DSpark vocabulary projection controls Oct 22, 2026
@jiapingW
jiapingW changed the base branch from codex/pr942-06-projection-checkpointing to main October 22, 2026 01:26
@jiapingW

jiapingW commented Oct 22, 2026

Copy link
Copy Markdown
Collaborator Author
Checkpoint Projection cache Flatten DSpark time (ms) DSpark peak memory (GiB) DSpark throughput change Domino time (ms) Domino peak memory (GiB) Domino throughput change
Off Off Off 378.08 17.09 Baseline 676.11 24.61 Baseline
Off Off On 378.10 17.09 -0.01% 676.36 24.61 -0.04%
On Off Off 456.09 14.52 -17.10% 760.52 22.43 -11.10%
On Off On 455.86 14.52 -17.06% 760.54 22.43 -11.10%
On On Off 392.34 16.04 -3.64% 687.92 23.43 -1.72%
On On On 392.40 16.04 -3.65% 688.45 23.43 -1.79%

Setup: H20, Qwen3-8B , a 5-layer draft model, BF16, FlexAttention, frozen LM head, Liger disabled, and objective chunk size fixed at 128. DSpark: batch=1, seq_len=3072, anchors=512. Domino: batch=2, seq_len=3072, anchors=256.

Timing covers draft forward, loss, backward, and optimizer step, excluding teacher inference, data loading, and distributed communication. Throughput changes are relative to each model's baseline with checkpointing, projection cache, and flatten all disabled. All configurations produced identical losses and trainable parameter gradients to the baseline. The flatten has no gain and the projection cache is a medium status between using checkpoint and not using checkpoint. This is not a necessary optimization. I will update this PR to include only the Liger kernel.

@jiapingW jiapingW changed the title perf(training): share Domino and DSpark vocabulary projection controls feat(training): enable opt-in Liger kernels for Domino and DSpark Oct 22, 2026
@jiapingW
jiapingW requested a review from maocheng23 October 22, 2026 04:40
@jiapingW
jiapingW force-pushed the codex/pr942-07-domino-optimizations branch from 45b7c51 to 8c0a361 Compare October 22, 2026 05:46
@jiapingW
jiapingW changed the base branch from main to codex/pr942-05-dflash-fused-head October 22, 2026 05:46
@jiapingW
jiapingW force-pushed the codex/pr942-05-dflash-fused-head branch from 6d7b8c9 to 158cfdd Compare October 22, 2026 06:03
@jiapingW
jiapingW force-pushed the codex/pr942-07-domino-optimizations branch from 8c0a361 to c0fef19 Compare October 22, 2026 06:03
This was referenced Oct 22, 2026
Merged
Merged
maocheng23
maocheng23 approved these changes Oct 23, 2026
View reviewed changes

@maocheng23 maocheng23 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, merge after 961 please

yushengsu-thu added a commit to yushengsu-thu/SpecForge that referenced this pull request Oct 23, 2026
@yushengsu-thu
configs: drop use_liger_kernel from the qwen-8b-moe recipe
6f87117 
planning.py rejects model.use_liger_kernel=true for dspark: on main only the
dflash algorithm wires Liger kernels into its draft. The exp branch could set
it because it merged sgl-project#963. The first online smoke of this recipe failed on it.
Base automatically changed from codex/pr942-05-dflash-fused-head to main October 26, 2026 06:13
@maocheng23 @claude
test(training): fold Domino/DSpark Liger tests into the CI-run Liger …
119e6da 
…suite

Replace tests/test_training/test_domino_dspark_liger.py with two CPU tests
in test_liger_kernel.py: Domino and DSpark constructors forward the
injected kernel factories, and their providers resolve Liger only for an
explicit `model.use_liger_kernel: true`. tests/test_training is not a
package, so `unittest discover` never collected the old module, while CI
runs test_liger_kernel explicitly. The backbone's Liger numerics are
already covered by the DFlash/DFlash2 parity tests, and the Domino/DSpark
heads contain no Liger modules.

Drop the Domino/DSpark Liger section from training/DESIGN.md;
examples/configs/README.md documents the flag. This also resolves the
conflict with the consumer-dispatch section on main.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@jiapingW
jiapingW merged commit 348585f into main Oct 26, 2026
6 checks passed
@jiapingW
jiapingW deleted the codex/pr942-07-domino-optimizations branch October 26, 2026 08:06
@yushengsu-thu yushengsu-thu mentioned this pull request Oct 28, 2026
Open
yushengsu-thu added a commit to yushengsu-thu/SpecForge that referenced this pull request Oct 28, 2026
@yushengsu-thu
configs: turn on Liger kernels in the MoE DSpark recipe
e9e23a9 
DSpark accepts use_liger_kernel since sgl-project#963. Liger replaces the draft's
RMSNorms; the MoE FFN keeps its grouped_mm experts, which Liger does not
cover yet. A 10-step offline run of a Liger-built MoE DSpark draft trained
with finite loss, bias and counts.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@maocheng23 maocheng23 maocheng23 approved these changes
@shuaills shuaills Awaiting requested review from shuaills shuaills is a code owner
@sleepcoo sleepcoo Awaiting requested review from sleepcoo sleepcoo is a code owner
@FlamingoPg FlamingoPg Awaiting requested review from FlamingoPg FlamingoPg is a code owner
@FrankLeeeee FrankLeeeee Awaiting requested review from FrankLeeeee FrankLeeeee is a code owner
@zyksir zyksir Awaiting requested review from zyksir

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@jiapingW @maocheng23