Load policy checkpoints directly in bf16 to halve peak load memory#725
Open
maxwbuckley wants to merge 2 commits into
Open
Load policy checkpoints directly in bf16 to halve peak load memory#725maxwbuckley wants to merge 2 commits into
maxwbuckley wants to merge 2 commits into
Conversation
Gr00tPolicy previously called AutoModel.from_pretrained with no torch_dtype, so the bf16 checkpoint was materialized in fp32 (~2x the steady-state size) before being cast back down by the .to(device, dtype) call. Passing torch_dtype=torch.bfloat16 skips the transient fp32 copy; the final state is bit-identical because bf16 -> fp32 -> bf16 is a lossless round-trip, and the existing .to() is kept to normalize any stragglers exactly as before. Measured on a 100M-param synthetic bf16 checkpoint (fresh process per mode, ru_maxrss): peak RSS 1281 MiB -> 708 MiB; the ~578 MiB delta matches the removed fp32 materialization. For the ~3B production checkpoint this avoids a ~10 GB transient, which is most of the gap between the documented 16 GB inference floor and the model's ~6 GB steady-state footprint. Tests pin both the wiring (from_pretrained receives torch_dtype, .to() and eval() still applied) and the guarantee (state dict and produced actions are bitwise identical to the legacy fp32-load-then-cast path). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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.
What
One-line change in
Gr00tPolicy: passtorch_dtype=torch.bfloat16toAutoModel.from_pretrainedso the checkpoint (stored in bf16) is materialized directly in bf16 instead of being upcast to fp32 and cast back down. The existingmodel.to(device, dtype=torch.bfloat16)is kept unchanged, so any buffer normalization it performed still happens exactly as before.Why
Without
torch_dtype, transformers materializes the bf16 checkpoint in fp32 — a transient ~2x memory spike (~10 GB for the ~3B production checkpoint) that exists only to be cast away one line later. Removing it lowers the real memory floor for inference on edge and shared-memory devices (Orin, Thor, Spark) and roughly halves model load time.The end state is bit-identical: bf16 → fp32 → bf16 is a lossless round-trip, and steady-state inference already ran fully in bf16.
Verification
torch.equalbetween the direct-bf16 load and the legacy fp32-load-then-cast path, through bothAutoModeldirectly and the realGr00tPolicyload path.from_pretrainedreceivestorch_dtype=torch.bfloat16;.to(device, dtype)and.eval()still applied.ru_maxrss): peak RSS 1281 MiB → 708 MiB; the ~578 MiB delta matches the removed fp32 materialization (100.7M params x 4 B + bf16 copy). Both modes end with identical bf16 parameters.tests/gr00t/policy/CPU suite: 43 passed.tests/getting_started/CPU suite: 7 passed.pre-commitclean.Notes
low_cpu_mem_usagewas deliberately left untouched to keep this minimal — the nested backbonefrom_pretrainedinsideGr00tN1d7.__init__makes meta-device init a separate, riskier change.🤖 Generated with Claude Code