`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.
execute_tool() relied on signal.SIGALRM, which does not exist on Windows.
The AttributeError was swallowed by a bare except, so every tool call
silently returned None and the model always observed
{"error": "tool not found"} (issue #846).
calculate_math also ran eval() on model-generated text. Two of the three
call sites (scripts/eval_toolcall.py, scripts/web_demo.py) passed no
globals at all, leaving full builtins reachable from sampled output.
safe_math_eval() in trainer_utils.py parses the expression with ast and
walks it, allowing only arithmetic operators and a math whitelist, with
caps on expression length and exponent size. It cannot loop forever, so
the signal-based timeout is no longer needed and is removed. Both sqrt(x)
and math.sqrt(x) spellings keep working.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Corey52HZ
Lingxu
tkgcc
jingyaogong
jingyaogong
gongjy