On macOS there's no CUDA, so these trainer scripts fall back to CPU. Detect MPS in the --device default and switch mixed precision to the device-agnostic torch.autocast / torch.amp.GradScaler so it works there.
CUDA and CPU paths are unchanged.
Move log_softmax for mb_resp_logp inside autocast_ctx so it is upcast
to fp32, matching the rollout path (compute_per_token_logps runs under
no_grad+autocast). Previously log_softmax ran outside autocast on
fp16/bf16 logits, introducing a ~1e-2 systematic bias between
mb_resp_logp and old_resp_logp, so ratio != 1 and KL != 0 even before
any weight update. Under sparse rule-based rewards (e.g. math PPO)
this bias destabilises training.
Also add:
- --rollout_eval_mode (default 1): switch actor to eval during rollout
to disable dropout and keep old/new logp consistent.
- --debug_log_ratio: optional diagnostic printing |mb-old| on the first
minibatch to verify ratio ~= 1.
Jiyang Liu
jingyaogong
FloatFrank
jingyaogong
readlnh