Autocast (scope-based prototype) — mixed precision - #2702
Draft
CarloLucibello wants to merge 1 commit into
Draft
Conversation
…ispatch barrier Draft/comparison implementation of Flux.autocast. Layers consult an ambient ScopedValue at forward time through a per-layer dispatch barrier; the machinery is compiled out until the first `autocast` call (a one-time world-age flip of `autocast_active()`), so code that never uses autocast keeps exact inference and zero overhead. Once enabled, forward passes of the affected layers infer as the small union of the Float32/Float16/BFloat16 paths. Also splits the precision casts: f16/bf16 are full casts (like PyTorch model.half()); f16mix/bf16mix keep BatchNorm/InstanceNorm/GroupNorm statistics and affine parameters in Float32. Known: after the flip, active-scope forward passes are a small union rather than a single concrete type; bf16 + AutoEnzyme is blocked upstream (EnzymeAD/Enzyme.jl#3430). This branch is kept for comparison against a wrapper-based implementation. Co-Authored-By: Claude Fable 5 <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.
Draft — kept for comparison, not intended to merge.
Scope-based prototype of
Flux.autocast(PyTorch-torch.autocast-style mixed precision). A companion wrapper-based implementation is being developed on a separate branch; this PR preserves the scope-based approach so the two can be compared (ergonomics, inference, performance).Approach
Affected layers consult an ambient
ScopedValueat forward time via a per-layer dispatch barrier. The whole mechanism is compiled out until the firstautocastcall: a one-time world-age flip ofautocast_active()fromfalsetotrue. So:Float32/Float16/BFloat16paths (union-split by the compiler).Parameters stay
Float32("master weights"); gradients come backFloat32, so the usual optimiser setup is unchanged. Normalization layers and losses compute inFloat32.Also
f16/bf16are full casts (PyTorchmodel.half()/.bfloat16()); newf16mix/bf16mixkeepBatchNorm/InstanceNorm/GroupNormstatistics + affine parameters inFloat32.autocast=Tkeyword added togradient/withgradient/train!.Known limitations
BFloat16autocast underAutoEnzymeis blocked upstream: EnzymeAD/Enzyme.jl#3430.🤖 Generated with Claude Code