Conversation
…ocast Keep the CUDA entropy dot-product outside autocast so mixed-precision training cannot downcast an already-fp32 reduction. Add single-GPU regression coverage and CI registration. AI-assisted: reviewed with ChatGPT; submitter must add DCO sign-off before publication. Signed-off-by: luozijian <luozijian0924@gmail.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request disables CUDA autocast during the torch.einsum calculation in sum_softmax_logits to prevent downcasting of FP32 inputs, which could corrupt entropy values and backward statistics. It also adds a new GPU test suite (test_ppo_entropy_autocast_gpu.py) to verify shift invariance and correctness of autocast entropy against FP32 references. I have no feedback to provide.
|
Could a maintainer review the scoped FP32 entropy-reduction fix at edcdf48? The current DCO, Buildkite and documentation checks are successful; the previously posted A100 red/green evidence is unchanged. The change only excludes the CUDA reduction from autocast and does not alter PPO math, optimization or rollout behavior. To keep related review routing in one place: #420 contains the signed FP32 rotary-frequency follow-up to its earlier feedback, and #416 fixes the cross-epoch sample cursor. Both also have successful visible checks, but they are independent PRs, not prerequisites for this one. Please point me to the appropriate owner for any of these scopes. Thanks for taking a look. |
1 participant
Summary
Fix CUDA autocast from downcasting the FP32 reduction used by VIME's fused PPO entropy path.
Under CUDA autocast,
torch.einsumcan run at a reduced precision even when both operands were explicitly promoted to FP32. The result is saved for entropy backward, so the precision loss affects both the reported entropy value and its gradient.This change keeps only that CUDA reduction outside autocast and adds a one-GPU regression suite covering FP16/BF16 autocast, masked/unmasked log-prob paths, metric-only/gradient entropy, and chunked/non-chunked execution.
Reproduction
On VIME main
ce92eff12ecdc81396bf41a2f94e62dd5b0aca32, the new regression suite fails all 18 cases before the production fix on an NVIDIA A100 80GB PCIe.After the fix:
Environment for GPU validation: NVIDIA A100 80GB PCIe, PyTorch 2.3.0a0+ebedce2, CUDA runtime 12.3, Python 3.10.12. Megatron-LM was checked out at VIME's pinned Docker commit
1dcf0dafa884ad52ffb243625717a3471643e087; NumPy 1.26.4 was injected through an isolated user path to match VIME's Docker pin.Scope
The production change is deliberately narrow: disable autocast only around the CUDA entropy dot-product. No API, optimizer, rollout, or distributed topology behavior changes.
AI assistance
This contribution was developed and reviewed with ChatGPT assistance. The human submitter remains responsible for the code and validation.