Non-reentrant offloading, fix race condition, fix activation offloading - #1649
Open
dxqb wants to merge 1 commit into
Open
Non-reentrant offloading, fix race condition, fix activation offloading#1649dxqb wants to merge 1 commit into
dxqb wants to merge 1 commit into
Conversation
Makes the LoadBoundary/EvictBoundary path the only offloading path and removes the reentrant one. Weight movement is driven from autograd Functions instead of the `use_reentrant=True` recompute, which is the fix for Nerogar#1306 — the train-stream event is now recorded after the block's backward kernels rather than after its forward, so an offload transfer can no longer overlap live compute. The block is compiled together with its checkpoint rather than bare, so AOTAutograd's min-cut partitioner prunes the recompute instead of the whole block re-running. Measured **-40.5 ms/step** (1061.6 → 1021.0 ms of compute-stream work, 4009 → 3959 kernels) on a 32-block transformer with clocks locked; `CompiledFunction` drops 64 → 32. `OT_BOUNDARY_OFFLOAD` is gone: `BoundaryOffloadCheckpointLayer` serves compiled and uncompiled parts alike, so `OffloadCheckpointLayer`, the conductor's `before_layer`/`after_layer`, the reentrant activation machinery and the `use_reentrant=True` dummy-grad helpers all go with it. Offloading still requires gradient checkpointing, and now raises for compiled parts too. Autograd's `SavedVariable` keeps a shallow copy of a saved weight, so repointing `param.data` at a reloaded buffer never redirects it, and the conductor recycles those buffers between layers — an uncheckpointed eager backward reads another layer's weights and produces plausible but wrong gradients. Checkpointing cures it by re-reading the weights during the recompute. A compiled backward escapes it (AOT reads parameters at call time, verified against a deliberately clobbered buffer), but that is an implementation detail rather than a guarantee, so the requirement applies there too. Activation-offload selection is fixed. `LoadBoundary` returns fresh output tensors aliasing its inputs and autograd saves an alias rather than the arg object, so matching declared args by `id()` missed almost everything — 1 of 8 blocks offloaded where 8 were expected. Matching is now on `(data_ptr, shape, dtype)`. The declared-arg list stays: the partitioner sees one block at a time and saves tensors shared across all blocks (rotary embeddings, masks) once per block, so offloading its whole saved set would cost bandwidth and free nothing. Known and out of scope: activation reloads run ahead of the GPU early in the backward (~28 of 40 reloads issue in one burst), because prefetch is bounded in host time rather than GPU time. Confirmed present on `master` too, so it is not this branch's to fix. 🤖 Drafted by Claude
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.
Summary
This PR replaces the offloading triggers using reentrant checkpointing with autograd functions and non-reentrant checkpointing.
This does a few things:
This PR also fixes activation offloading: Activation offloading technically worked, but it wasn't of any practical use.
Activations were offloaded, and moved back to GPU when the backward started. But: this wasn't limited. CPU can run ahead of GPU for many layers. This caused vram allocations on GPU too early and you ran out of vram anyway.
With this fix, high batch sizes are actually possible. Tested on Flux2, 4070 16 GB, batch size 20: 7.8 s/it
Test plan
pre-commit run --all-filespassesAI assistance