Skip to content

Add experimental loader-matched AF4G64 projection training - #551

Open
linson007 wants to merge 14 commits into
incoai:mainfrom
linson007:experiment/dflash-affine4-training
Open

linson007 wants to merge 14 commits into
incoai:mainfrom
linson007:experiment/dflash-affine4-training

Conversation

@linson007

@linson007 linson007 commented Oct 22, 2026 •

Copy link
Copy Markdown
Contributor

Adds opt-in AF4G64 projection training with native-loader-compatible BF16 source/scale/bias rounding, group-64 packing and straight-through gradients. Context/selector and full-projection scopes preserve FP32 masters and normal BF16 export names; norms, base convolutions and codebooks retain their existing precision.

Merge scope and dependencies

Depends on #537; merge the training chain #499 → #537 → #551 → #552 → #553 → #554 in that order. 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.

Implements the projection factor of #543, not a quantized-teacher adapter or deployed-feature corpus. Reference matmul is not a Metal accumulation parity claim. 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.

Validation at d8b15fbe5415

  • 57 local training/reference tests passed with no skips. Python 3.12.14, torch 2.13.0 and Transformers 5.15.1 on Apple M5 Pro/macOS 26.6.
  • The pinned upstream numerical reference (07ebd93db9f472af339b644bb70221ad8428328a) ran. Actual MPS was required and CPU fallback disabled. Regression coverage reproduces and fixes an MPS evaluation bug where switching quantization scope could orphan FP32 optimizer masters. Packed quantization bytes match the independent native preparation oracle; the native preparation golden harness passed.
  • Ruff lint/format, architecture boundaries and diff whitespace checks passed. CI checks the current head separately for CPU/reference training, Python 3.12–3.14, native production/CPU builds and sanitizers.
DFLASH_REFERENCE=/path/to/dflash/dflash/model.py \
DFLASH_REQUIRE_MPS=1 PYTORCH_ENABLE_MPS_FALLBACK=0 \
  python -m unittest discover -s dev/tests/training -t . -v

CPU/MPS fixture and reference checks above are implementation evidence. Additional real-checkpoint CUDA qualification is recorded below.

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 QAT recipe completed 20 selector steps with AF4G64 selector/context projections. Its native acceptance lengths equaled the matched plain selector continuation in all four short prompts; the smoke does not demonstrate QAT effectiveness.

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. Future capture controls should make thinking mode explicit and include it in cache/run provenance.

@linson007

Copy link
Copy Markdown
Contributor Author

Real-checkpoint QAT smoke results

The RunPod source was the combined training stack at d066629ff370d94e645bbdd6d32f73a4c8194838 in #554, which includes this PR and its prerequisites. This is integration evidence, not an isolated CUDA run of this PR's own head. One A100-SXM4-80GB ran the actual pinned BF16 Qwen/Qwen3.8-27B teacher and incoai/Qwen3.8-27B-DFlash2 draft (Python 3.12.3, torch 2.13.0+cu130, CUDA 13.0, Transformers 5.15.1).

  • affine4_scope: selector_context, selector training, 20 steps; finite gradients, checkpoint selection and finite BF16 export passed.
  • Near-2K optimizer preflight passed: 12.0 GiB peak allocated during training; teacher capture 51.8 GiB.
  • Export hashes matched the cloud selection manifest exactly. Native loading and 108 capture/replay cycles passed.
  • Validation matching proposals: 3.875 → 3.875. Native acceptance lengths equaled the equal-step plain selector continuation on all four short prompts. This smoke did not demonstrate a QAT effectiveness gain.

The exact PR head d8b15fbe541566c664c68bc6350856384d72ddc7 separately has 57 passing local training/reference tests with actual MPS required and green scheduled CI. Its independent native packing/reference and optimizer-master restoration tests remain separate from the CUDA integration run.

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/qat/metrics.jsonl

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