Non-reentrant offloading, fix race condition, fix activation offloading - #1649
Open
dxqb wants to merge 5 commits into
Open
Non-reentrant offloading, fix race condition, fix activation offloading#1649dxqb wants to merge 5 commits into
dxqb wants to merge 5 commits 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
The TorchMemoryRecorder and TorchProfiler around the training step were hardcoded to enabled=False, so profiling a run meant editing GenericTrainer and reverting afterwards. Setting OT_DEBUG_PROFILES now writes a CUDA memory snapshot for steps 0 and 1 and a profiler trace for steps 10, 11, 40 and 41; with the variable unset the step tuples are empty and nothing is recorded, exactly as before. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Resolves the conductor's to(device) against the materialize()/evict() split that landed upstream, matching the resolution already reviewed on the source branch.
OT_DEBUG_PROFILES is under review separately (Nerogar#1682); this branch carries it so the profiling switch is present alongside the offload work.
BaseModel._move_part now allocates each part's component and its LoRA into a per-stem torch.cuda.MemPool, so the two land contiguously and are released together when the part is evicted. Moving a part back to a device without MemPool support drops the stem's pool, letting the collection in evict() release its segments. supports_mem_pool(), create_mem_pool() and mem_pool_context() in torch_util keep the CUDA-only parts out of the caller. The load and evict boundaries take a `dummy` grad-requiring leaf. A torch.autograd.Function whose inputs all lack grad creates no node, so a block with nothing trainable in front of it would never reach its boundary and would run against weights still on the temp device. When no input carries grad, _apply_boundary routes the first floating-point tensor through against `dummy` so the node exists and the backward duty fires. pin_tensor_ now explains a failed registration of 2 GiB or more instead of reporting a bare CUDA error code: on linux 6.11 and 6.12 the kernel refuses such a request, and the refused call leaves the pages pinned with no handle to release them. The message names the running kernel and the memory that is lost. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Collaborator
Author
|
contains github.com//pull/1620 |
|
|
||
| # OT_DEBUG_PROFILES dumps a CUDA memory snapshot for the first two steps, where the allocator is still | ||
| # growing, and a profiler trace at steps 10 and 40, past compilation and warmup. | ||
| _DEBUG_PROFILES = bool(os.environ.get("OT_DEBUG_PROFILES")) |
Contributor
There was a problem hiding this comment.
Suggested change
| _DEBUG_PROFILES = bool(os.environ.get("OT_DEBUG_PROFILES")) | |
| _DEBUG_PROFILES = bool(int(os.environ.get("OT_DEBUG_PROFILES"))) |
Since os.environ.get(...) is always a string, and even "0" ends up being mapped to True as len("0") != 0.
Collaborator
Author
There was a problem hiding this comment.
this is #1682 and already merged to main (in a better form)
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
2 participants
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