Conversation
added 3 commits
July 8, 2026 14:47
Port the multi-stream design from branch cpu_gpu2 onto main, gated by new flag grad.async_cpu_transfer (default False; requires pinned layer-input and KV cache checkpoint memory). When enabled: - Layer input checkpoints are written GPU -> CPU on a separate device_to_host_stream during the forward pass. - In GradientAccumulator.run, inputs for the next cell are prefetched CPU -> GPU on host_to_device_stream and bottom head gradients of the previous cell are written GPU -> CPU on device_to_host_stream, both in parallel with the backward computation of the current cell (double buffering). Fixes for races present on cpu_gpu2 (NaN weights, annotation errors): - The d2h checkpoint copy now waits for a main-stream event recorded when the copy source has been produced (the input to layer L+1 is created by torch.cat at the end of layer L's iteration, after the old synchronization point; layer 0 raced with wte the same way). - One torch.cuda.synchronize() at backward start makes the forward's d2h checkpoint writes host-visible before backward reads them. - Explicit cross-stream ordering events at the start of GradientAccumulator.run, instead of relying on the per-row synchronize which only runs when do_checkpointing. All stream ordering inside the cell loop uses GPU-side events; the host thread is never blocked there.
Runs LongContextGradientModel forward+backward over two batches, once with async_cpu_transfer=False and once with True, and compares losses and gradients. Parametrized over cache policies (lastrec, h2o) and KV cache checkpoint quantization (default, torch-quantized8). Also checks that async_cpu_transfer=True requires pinned checkpoint memory.
GradientAccumulator.run contains a torch.cuda.synchronize() after the KV cache checkpoint computation, blocking the host thread for about a second per row of cells (36 times per batch for a 36-layer model). With multiple streams, the same ordering is already expressed by GPU-side events at the start of the cell loop: the transfer streams wait for an event recorded on the main stream, which is ordered after the checkpoint computation. So the host-blocking synchronize is only needed in the single-stream case, where reads of the CPU checkpoint buffers are issued by the host without stream ordering.
Contributor
|
Thanks! But you said the impact is minor? I still wonder why that is. Sure, maybe the CPU/GPU transfer is anyway cheap in our case (but I think we could use more of it). But why is it we see all these events in profiling which seem to block GPU computations? |
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.
Adds opt-in async CPU↔GPU transfer during gradient computation (#62), behind
--grad.async_cpu_transfer(defaultFalse; requires pinned checkpointmemory). Layer-input checkpoints, KV-cache checkpoints, and head gradients are
moved on dedicated
host_to_device/device_to_hostCUDA streams thatoverlap main-stream compute, via double buffering in
GradientAccumulator.run.Ordering in the cell loop uses GPU-side events, so the host thread isn't blocked.
Re-implements the
cpu_gpu2design on currentmainand fixes the races thatmade that branch produce NaN weights.
Race fixes vs
cpu_gpu2source is produced (on
cpu_gpu2, layer L+1's input could be copied beforethe
torch.catthat computes it — the main NaN source).torch.cuda.synchronize()at backward start makes forward checkpointwrites host-visible before backward reads them.
run, replacing reliance on theper-row synchronize.
Also skips that per-row
torch.cuda.synchronize()(host-blocking ~1 s/row) inthe async path, where events already order things.
async_cpu_transfer=Falseleaves all paths unchanged.
Testing
test/kvcache/test_gradient_async.py: compares losses/gradients with theflag off vs on over two batches, across cache policy × checkpoint
quantization. 5/5 pass; existing gradient suites green.
cpu_gpu2): 3 clean steps,no NaNs, losses matching the async-off control; ~4–8% faster/iter.
Closes #<62>.
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.