train_grpo.py, train_ppo.py and train_agent.py wrap the policy in
DistributedDataParallel but run the loss-bearing forward through
.module, bypassing the wrapper. DDP arms its gradient all-reduce hooks
inside DistributedDataParallel.forward(), so calling .module directly
means backward() performs no all-reduce: every rank steps on purely
local gradients and the replicas silently diverge. Only rank 0 is
saved, so multi-GPU RL training quietly degrades to single-GPU
training on 1/N of the data with no error.
The supervised trainers are unaffected; they already call the wrapped
module. no_grad uses of .module (rollout generation, checkpoint
saving) are left untouched, where unwrapping is correct.
Single-GPU behaviour is unchanged: with no process group, model is the
bare module and these calls are identical to before.
`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.
Corey52HZ
basil-k-aji-dev
jingyaogong
Lingxu
jingyaogong