Skip to content

Add experimental Apple Silicon MPS draft training - #537

Draft
linson007 wants to merge 7 commits into
incoai:mainfrom
linson007:feat/dflash2-mps-training
Draft

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

Conversation

@linson007

@linson007 linson007 commented Oct 21, 2026 •

Copy link
Copy Markdown
Contributor

Apple Silicon users can select --device mps for optional DFlash2 capture, preflight, training, resume, and final evaluation.

Depends on #499. Please merge the parent first. This upstream PR targets main, so it currently includes the parent training changes. The Mac-only follow-up is commit 2fe5fe39. After #499 merges, rebase that follow-up onto current main to isolate the final diff. This remains draft pending real-checkpoint qualification.

The Mac follow-up adds:

  • MPS availability checks, teacher-cache release, synchronized timing, and preflight memory snapshots. Snapshots are explicitly distinguished from peak memory.
  • MPS RNG checkpoint/restore and execution provenance covering Apple GPU, macOS, and explicit CPU fallback.
  • Shared CPU BF16 conversion for export and MPS validation, with original FP32 optimizer parameter objects restored after successful or failed validation. This addresses a validation/export mismatch reproduced locally on MPS.
  • Nine Mac-specific checks covering CLI selection, unavailable/required backend behavior, actual teacher generation and capture, selector and backbone optimization, exact resumed-versus-uninterrupted checkpoints, BF16 export, final evaluation, and failure recovery.
  • A contributor testing guide and a manual workflow for an existing trusted Apple Silicon runner. The workflow requires registration on the default branch before dispatch; local smoke commands work directly from this PR.

Validation on an Apple M5 Pro with 48 GB unified memory, macOS 26.6, Python 3.12.14, PyTorch 2.13.0, and Transformers 5.15.1:

  • All 50 optional training tests passed, including all nine Mac checks and pinned Z Lab reference parity, with explicit CPU fallback disabled.
  • All 35 ordinary launcher tests passed.
  • Without PyTorch, optional discovery passed 17 metadata/launcher checks with 33 expected skips.
  • Ruff, formatting, compileall, architecture, workflow YAML/trigger checks, and whitespace checks passed.

The full ordinary/native suite was not rerun for the Mac follow-up. No real Qwen3.8-27B MPS training, production CUDA run, or trained-checkpoint native qualification is claimed.

Testing help wanted: contributors with sufficient unified memory for the full BF16 teacher can help measure real-checkpoint preflight, training memory/step time, resume, export, and native acceptance/throughput. Teacher weights alone are approximately 54 GB before working memory. The current command runs capture and training on the selected backend and does not provide cross-device cached-feature training.

The smoke test needs no model downloads after installing training/requirements.txt in a separate environment:

DFLASH_REQUIRE_MPS=1 PYTORCH_ENABLE_MPS_FALLBACK=0 python -m unittest dev.tests.training.test_mps -v

Please follow the guide's reporting checklist for hardware, versions, input revisions, recipe/split fingerprints, memory, timings, and errors.


View with [code]smith Autofix with [code]smith
Need help on this PR? Tag @codesmith-bot with what you need.

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