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. |
|
A correction to my earlier CI summary for unchanged head I checked |
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.
Current code HEAD:
edcdf4829c66fcf1ebab9f9c28d296f8da543e04.Retained local validation
The results below were generated during the original A100 validation; they were not rerun for the September 16 description/CI clarification.
On VIME base
ce92eff12ecdc81396bf41a2f94e62dd5b0aca32, the new regression suite fails all 18 cases before the production fix on an NVIDIA A100 80GB PCIe.After the fix, the previously reported results are:
These selections/probes are reported separately, not as an additive total or as upstream CI results.
Environment for the retained 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 commit1dcf0dafa884ad52ffb243625717a3471643e087; NumPy 1.26.4 was injected through an isolated user path to match VIME's Docker pin.Reproduction in a compatible CUDA environment
The first selection requires one CUDA GPU; BF16 cases also require BF16 support. The existing TP=2 control in the second selection requires two GPUs. A CPU-only invocation that skips the CUDA cases is not a passing GPU validation.
Upstream CI and review — checked September 16, 2026
Buildkite #1199 maps to GitHub
success, but its detailed status ispassed and blocked. This is not evidence that the GPU suite completed.The existing pipeline separates always-on CPU jobs from a manual GPU gate. On this HEAD,
.buildkite/gpu_suites.pyregisterstest_ppo_entropy_autocast_gpu.pyin themegatronselection withNUM_GPUS=1. Its presence in the selection is routing evidence only, not execution evidence.The remaining CI item is confirmation/execution of that regression through the project's authorized GPU workflow. No GPU gate, permissions, labels, or test requirements have been bypassed or weakened. Maintainer review remains pending; the earlier automated review is not human approval.
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. No full-training convergence, model-quality improvement, or full upstream GPU-suite pass is claimed.
This description update changes no production code, tests, pipeline files, commit history, or existing sign-offs.
AI assistance
This contribution was developed and reviewed with ChatGPT assistance. ChatGPT also assisted with the September 16 source/CI clarification. The human submitter remains responsible for the code and validation.