Skip to content

Add experimental DFlash2 training for dense and MoE Qwen models - #499

Open
linson007 wants to merge 11 commits into
incoai:mainfrom
linson007:feat/dflash2-training
Open

linson007 wants to merge 11 commits into
incoai:mainfrom
linson007:feat/dflash2-training

Conversation

@linson007

@linson007 linson007 commented Oct 19, 2026 •

Copy link
Copy Markdown
Contributor

Adds an opt-in, source-checkout workflow for users to adapt released Qwen3.8-27B (dense) and Qwen3.6-35B-A3B (MoE) DFlash2 drafts with their own text dataset, compatible local BF16 checkpoints and sufficient training hardware.

Both families share preparation, tokenized-data checking, isolated teacher capture, selector-only/full-draft optimization, exact optimizer/RNG resume, finite BF16 export and held-out comparison. The teacher stays frozen; MoE routers and experts are not fine-tuned. draft_training/families.py registers explicit target loaders and released draft geometries. Validation rejects mismatched target/draft pairs, incompatible routing geometry, missing expert weights and quantized training inputs. Packed and individual expert storage are tested.

prepare-data normalizes messages, prompt/response, prompt/completion or prompt-only JSONL into reproducible train/validation/test splits grouped by source and normalized prompt identity. check-data renders the examples before model weights are loaded; new training runs perform this check automatically.

Own-data training quickstart covers environment setup, pinned checkpoints for both families, data schemas, both recipes per family, preflight, training, resume, export and held-out evaluation. Improvement remains data- and deployment-dependent; reports distinguish improved, tied and regressed validation outcomes.

Real MoE CUDA qualification: the isolated #499 trainer at 1aae5d5dcb45805aad63915de54e845dd6be7c37 completed the full-size BF16 Qwen3.6-35B-A3B workflow on one A100-SXM4-80GB. Both selector-only and full-draft recipes completed 20 optimizer updates, finite BF16 export and held-out baseline comparison. The full run recovered from SIGINT after a committed step-10 checkpoint; resuming the completed run also passed. Standalone full export matched the automatic export tensor for tensor, and all downloaded export-file hashes match the frozen selections. No remote-only source changes were required.

MoE recipes, prepared splits, raw metrics, provenance and reproduction are published in documentation-only follow-up a4380587. Separate long-context preflight peaks were 65.9 GiB allocated during teacher capture, 5.1 GiB during selector optimization and 11.8 GiB during full-draft optimization. These measurements are not total hardware requirements.

On this authored 8/4/4 smoke with one seed and eight replay blocks per evaluation, selector validation tied at 3.5 → 3.5, and full-draft validation rose 3.5 → 3.75. Both held-out results tied at 4.5 → 4.5; no held-out matching-proposal gain was demonstrated. The recipes were fixed before final testing, with thinking explicitly disabled. Trained MoE native loading/acceptance/throughput qualification and representative effectiveness testing remain pending. Arbitrary target architectures and fresh draft initialization are not included.

Merge scope and dependencies

First PR in the training chain: #499 → #537 → #551 → #552 → #553 → #554. All PRs target main per CONTRIBUTING; review this component against its immediate prerequisite. Parent changes remain visible against main until those parents merge. Merge commits preserve the tested ancestry; after squash/rebase merges, refresh dependent heads against main and rerun their checks.

Optional implementation infrastructure; real-checkpoint CUDA training, native export qualification and paired serving performance are separate checkpoint/deployment gates. Ready-for-review status applies to these optional tools and their implementation tests. No checkpoint acceptance or speedup improvement is claimed, and no newly trained/experimental architecture becomes a serving default.

Implementation validation at 1aae5d5dcb45

  • 77 component/reference tests passed with no skips. Python 3.12.14, torch 2.13.0, Transformers 5.15.1 and safetensors 0.8.0 on Apple M5 Pro. The upstream numerical reference exactly matches revision 07ebd93db9f472af339b644bb70221ad8428328a.
  • The capture/training/export/final-evaluation suite runs with both a real small dense teacher and a real small MoE teacher. MoE checks cover frozen routers/experts, exact packed/individual expert loading, full-draft optimization, split preparation, invalid-data and incompatible-resume rejection, and frozen-evaluation contracts. Deterministic optimizer/RNG resume is independently tested on cached-feature training.
  • A broader ordinary Python run passed 1,394 tests with 77 environment-dependent skips. After the final expert-format regression was added, the current family module was checked again in that environment: 21 tests, 16 tensor-dependent skips. All optional tensor tests ran separately in the 77-test component suite.
  • A temporary merge with Add selective recall training and reached-prefix coverage metrics #554 at 310e75d7f22e7a86f2141c5cf93e82aba62a9342 passed 112 combined-stack tests with no skips, including the existing actual-MPS paths with DFLASH_REQUIRE_MPS=1 and PYTORCH_ENABLE_MPS_FALLBACK=0. MoE pipeline tests in that matrix run on CPU; this does not qualify a full-size MoE accelerator run.
  • For Qwen/Qwen3.6-35B-A3B at 995ad96eacd98c81ed38be0c5b274b04031597b0 and its released draft at 51ef7b6923ad6c14cb1bb41c37a9041446496ab6, the real config/safetensors-header contracts passed. A full-size teacher constructed on the meta device had no missing required parameters or shape mismatches against the released headers; additional MTP tensors are outside this teacher graph. No full weight files were downloaded or optimized in these checks.
  • Both MoE recipes passed tokenizer-only checks with thinking enabled and disabled. Existing 27B tokenizer checks and both family geometry fixtures remain covered.
  • Ruff lint/format, architecture boundaries and diff whitespace checks passed. Read-only merge rehearsals are clean against Add experimental Apple Silicon MPS draft training #537, Add experimental loader-matched AF4G64 projection training #551, Add experimental prefix-survival loss and position weighting #552, Add training-only pair-context and nonlinear selector prototypes #553 and Add selective recall training and reached-prefix coverage metrics #554.
  • Tested-trainer CI independently checks CPU/reference training, Python 3.12–3.14, native builds and sanitizers.
DFLASH_REFERENCE=/path/to/dflash/dflash/model.py \
  python -m unittest discover -s dev/tests/training -t . -v

The historical CUDA/native results below were produced by the earlier Qwen3.8-27B combined stack. They have a separate source/model scope from the current isolated-head MoE report above.

Real-checkpoint CUDA and native qualification

Recipes, model revisions, raw metrics and replayable native captures are published in #554. The CUDA source is the combined training stack at d066629ff370; this is additional integration evidence, not a claim that every isolated parent head ran on CUDA. One A100-SXM4-80GB ran the actual pinned BF16 Qwen3.8-27B teacher and released DFlash2 draft. Python 3.12.3, torch 2.13.0+cu130, CUDA 13.0, Transformers 5.15.1.

The baseline completed 20 selector steps. A separate near-2K full-draft backward/optimizer preflight also passed (40.7 GiB peak allocated during training, 51.8 GiB during teacher capture). This was one full-draft preflight update, not a full-draft training study.

All four 20-step recipes exported finite BF16 weights. Capture tensor payloads matched exactly across recipes; frozen export tensors were checked by payload hash. Validation mean matching proposals remained 3.875 → 3.875 in every run. All four exports passed native loading/capture/replay against a fixed low-bit target, with exact cloud export hashes verified. The released draft plus four exports produced 540 native lane-cycles. No consistent acceptance improvement or serving speedup was demonstrated. Tiny authored splits, one seed and correlated cycles limit effectiveness conclusions. The BF16 teacher and quantized native target remain distinct artifacts. No new checkpoint becomes a default; research/deployment issues remain open.

The pinned teacher template defaults to thinking enabled with xhigh reasoning effort; native probes explicitly disabled thinking. This mode mismatch, together with the BF16/low-bit target difference and short captures, prevents a matched-distribution effectiveness conclusion. This PR now makes thinking mode explicit and freezes it in cache/run provenance; the historical dense measurements retain their original mode mismatch.

@linson007

Copy link
Copy Markdown
Contributor Author

Test-run follow-up — 2026-10-22 UTC

The base training workflow has passing implementation tests and successful real-checkpoint integration evidence. The exact scope matters:

The plain selector continuation completed 20 optimizer steps, validation, checkpoint selection and finite BF16 export. Its near-2K preflight passed at 10.6 GiB peak allocated during optimization; teacher capture peaked at 51.8 GiB. A separate single full-draft backward/optimizer preflight update passed at 40.7 GiB; this was not a 20-step full-draft run. The probe used 1,895 prompt tokens and at most 128 generated tokens, with one complete draft block's backward pass.

The plain trained export loaded in native Splash on Apple M5 Pro and passed 108 capture/replay cycles across code, math, prose and JSON. Together with QAT, prefix-loss, recall and released-draft comparisons, the matrix contains 540 native cycles. Validation mean matching proposals stayed 3.875 → 3.875 in every 20-step recipe; the native probes showed no consistent acceptance improvement.

Changes and publication audit: all 12 trainer-module hashes in the RunPod selection manifest exactly match the already-published d066629ff370d94e645bbdd6d32f73a4c8194838 source. No remote-only or unpushed training/runtime hotfix was needed. Later features are intentionally separate: MPS and export/RNG handling (#537), QAT and its MPS optimizer-restoration fix (#551), prefix/position objectives (#552), research selectors (#553), and recall scopes/metrics (#554). They are not all present in #499 alone. The two commits after the tested source add only documentation and measured artifacts, published in #554 at 310e75d7f22e7a86f2141c5cf93e82aba62a9342.

This was an 8-train/4-validation/4-native-test authored smoke, one seed, with eight validation blocks and short continuations. The teacher used template-default thinking/xhigh; native probes disabled thinking and used a distinct 4-bit target. These differences and the tiny sample prevent a matched-distribution effectiveness conclusion. No serving throughput or sampled-decoding benchmark was run, and no newly trained weights become a serving default.

Full report and reproduction · cuda/baseline/selection.json · summary.json

Merge order remains #499 → #537 → #551 → #552 → #553 → #554. This supports reviewing the optional implementation; research/deployment gates remain open.

@linson007 linson007 changed the title Add experimental DFlash2 draft fine-tuning command Add experimental Qwen3.8-27B DFlash2 training workflow Oct 23, 2026
@linson007 linson007 changed the title Add experimental Qwen3.8-27B DFlash2 training workflow Add experimental DFlash2 training for dense and MoE Qwen models Oct 23, 2026
@linson007

Copy link
Copy Markdown
Contributor Author

Full-size MoE RunPod qualification — 2026-10-23 UTC

The own-data workflow now has real-checkpoint CUDA evidence for Qwen3.6-35B-A3B, separately from the earlier dense Qwen3.8 combined-stack smoke.

Tested source: isolated #499 head 1aae5d5dcb45805aad63915de54e845dd6be7c37. No downstream training PRs or pod-only source fixes were used. Documentation-only follow-up a4380587bf11de1725a31fcf1ea1f10f269b714b publishes the results.

Environment: one A100-SXM4-80GB; Python 3.12.3, torch 2.13.0+cu130, CUDA 13.0, Transformers 5.15.1 and safetensors 0.8.0.

Pinned checkpoints:

  • BF16 teacher: Qwen/Qwen3.6-35B-A3B at 995ad96eacd98c81ed38be0c5b274b04031597b0.
  • Released draft: incoai/Qwen3.6-35B-A3B-DFlash2 at 51ef7b6923ad6c14cb1bb41c37a9041446496ab6.

The MoE teacher, routers and experts remain frozen; training adapts the dense DFlash2 draft.

Recipe Optimizer updates Validation matching proposals Held-out matching proposals
Selector only 20 3.5 → 3.5 4.5 → 4.5
Full draft 20 3.5 → 3.75 4.5 → 4.5

Both workflows passed data preparation/checking, teacher capture, optimization, validation selection, finite BF16 export and held-out baseline comparison.

Additional checks:

  • Full training was interrupted with SIGINT after its step-10 checkpoint committed, then resumed to step 20: 20 updates, cursor 80, metric steps 1–20 exactly once. Resuming the completed run also verified the existing export/selection successfully. This was not a comparison against a separate uninterrupted CUDA run.
  • Selector export changed only its three selector tensors; the other 93 tensor payloads stayed identical to the released draft. Full-draft export changed 70 of 96 BF16 tensor payloads. Standalone full export matched the automatic export tensor for tensor.
  • All downloaded export-file hashes matched the frozen selection manifests. All 13 trainer-module and five native-contract file hashes matched the published tested source.

Separate long-context preflights passed with 1,773 prompt tokens and at most 128 generated tokens. Peak allocated CUDA memory was 65.9 GiB during capture, 5.1 GiB during selector optimization, and 11.8 GiB during full-draft optimization. These are separate allocator peaks from a single-block backward/optimizer probe, not total hardware requirements.

Interpretation and remaining work: this used authored 8/4/4 splits, one seed, explicit non-thinking capture and only eight replay blocks per evaluation. Recipes were fixed before final testing. The small full-draft validation gain did not produce a held-out matching-proposal gain. Native loading, acceptance, latency and throughput qualification of these MoE exports, plus representative effectiveness experiments, remain pending. No trained checkpoint becomes a serving default; arbitrary model families and fresh draft initialization are outside this workflow.

Full report, reproduction and raw artifacts · Own-data quickstart

The tested trainer's CI passed; the documentation follow-up's CI is running at posting time. The RunPod was stopped after artifact collection.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant