Feature: add opt-in bfloat16 mixed precision to PPO and Distillation#219
Open
lzyang2000 wants to merge 1 commit into
Open
Feature: add opt-in bfloat16 mixed precision to PPO and Distillation#219lzyang2000 wants to merge 1 commit into
lzyang2000 wants to merge 1 commit into
Conversation
Add a `use_mixed_precision` option (True / False / "auto", default False) that runs the forward pass and loss computation under torch.amp.autocast in bfloat16, while keeping backward, gradient clipping, multi-GPU gradient reduction, and the optimizer step in fp32. No GradScaler is required since bfloat16 shares fp32's exponent range. "auto" enables bf16 only on bfloat16-capable CUDA devices and falls back to fp32 otherwise. Advantage normalization and the adaptive-LR/KL step stay in fp32 (outside the autocast region). The default is off to preserve existing numerics and reproducibility; "auto" is the recommended value when opting in.
lzyang2000
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add opt-in bfloat16 mixed precision to PPO and Distillation
Adds an opt-in
use_mixed_precisionoption to PPO and Distillation. When enabled, the forward pass and loss run undertorch.amp.autocastin bfloat16; backward, gradient clipping, multi-GPU all-reduce, and the optimizer step stay in fp32. Off by default — existing runs are unchanged. (Hopefully I did not step on any toes)API
use_mixed_precision(defaultFalse):FalseTrue"auto"What changed
rsl_rl/utils/utils.py: newresolve_mixed_precision(setting, device_type)helper (interprets"auto"), exported fromutils/__init__.py.ppo.py/distillation.py: newuse_mixed_precision: bool | str = Falsearg;update()wraps only the forward + loss intorch.amp.autocast(..., dtype=torch.bfloat16). Backward/clip/all-reduce/step stay outside (gradients stay fp32, so the distributed path is unchanged).save()/load()untouched.Validation
Full test suite passes (181). Isolated
update()benchmark (RTX 4090, MLP[1024,1024,512]):update()End-to-end (mjlab Go1 velocity, 4096 envs × 300 iters):
auto)update())bf16 is stable (no non-finite values, reward matches fp32). The end-to-end gain is small only because this task is collection-bound — env sim is ~87% of wall time; the update itself is 1.3× faster.
Default is
False, so the change is purely additive and backward compatible.