`device_type` currently only distinguishes "cuda" from "cpu", so running
with `--device mps` on Apple Silicon falls into the "cpu" branch, disables
autocast entirely and forces fp32 training.
Changes (9 files, mechanical):
- `device_type` now resolves to cuda / mps / cpu
- `torch.cuda.amp.autocast(...)` -> `torch.amp.autocast(device_type, ...)`
- `torch.cuda.amp.GradScaler(...)` -> `torch.amp.GradScaler(device_type, ...)`,
enabled only for CUDA + fp16 (MPS has no gradient scaler)
- `web_demo.py` device selection now prefers mps when available
This also silences the `torch.cuda.amp.*` deprecation warnings on torch 2.x.
CUDA behaviour is unchanged.
Measured on Apple M5 Pro (24GB), hidden_size=512 / 8 layers, bs=16, seq=512:
- before (fp32, no autocast): 11,784 tok/s
- after (bf16 autocast): 15,492 tok/s (+31%)
Verified end-to-end on MPS: pretrain, full_sft, lora, dpo, ppo, grpo,
distillation and agent RL all run to completion, and the resulting weights
export correctly to transformers / GGUF.
Note for macOS users: `--num_workers 0` is currently required, as DataLoader
worker processes deadlock on macOS + Python 3.12 due to fork safety.
Lingxu
jingyaogong
jingyaogong