Backward CPU offloading: Asynchronous transfer

[mseeger]

Is your feature request related to a problem? Please describe.
CPU offloading in backward works by iterating this over layers:

• Load weights from CPU
• Run backward from head gradients
• Store gradients to CPU

These steps are currently run sequentially.

Describe the solution you'd like
Can we run these steps in parallel?

• Switch between two shards on GPU
• Run backward on one while loading weights for the other, and storing gradients for the previous

Needs clear understanding how async CPU <-> GPU transfer works! We know how transfer between GPUs works.

[mseeger]

This may be relevant:
https://docs.pytorch.org/tutorials/intermediate/pinmem_nonblock.html

According to this, memcpy is blocking on the device (GPU) in general, but there is a (complex) way to get around this.

[mseeger]

This would help in several places (not just backward CPU offloading):

• Forward pass: Write activation checkpoints -> CPU in parallel to comp
• Backward pass: Load activation checkpoints and head gradients from CPU in parallel to comp

If this could somehow be done generically, this would help a lot!

[mseeger]

Got really helpful feedback here: https://discuss.pytorch.org/t/transfer-data-gpu-cpu-and-compute-on-gpu-in-parallel/224695/4

It’s not entirely clear to me what you mean by “blocking on the device”. On the device, operations are typically stream-ordered, so if your device supports enough physical work queues, and you place operations in separate streams, and there is no contention between resources, nothing will be blocking on the device.
For your specific use case, you need at least 3 streams: one for the actual compute (Forward or backward), one for offloading activations of the previous layer and one for onloading activations for the next layer.

There are several utilities to help you do that: if you are using Nvidia transformer engine (or are able to switch to it), it supports everything you need through get_cpu_offload_context: PyTorch — Transformer Engine 2.12.0 documentation

There seems to be a feature request for this for standard PyTorch FSDP: Feature Request: Activation offloading with async prefetch in FSDP · Issue #174960 · pytorch/pytorch · GitHub

And if you want to roll your own implementation, I used this for a GTC 2024 presentation: llama-recipes/src/llama_recipes/interleaved_offload.py at 60307ec06d7a4edf6d02f402d29c0eb093aac9fa · MatthiasKohl/llama-recipes · GitHub

These implementations might be much more generic and complicated than what you actually need, but they will always be based on some core principle: place computation in one stream, place copies in one direction in a second stream and place copies in the other direction in a third stream. Also note that this kind of offloading doesn’t always make sense at all : it might make more sense to distribute across multiple devices or simply use a device with larger memory if it’s possible.

[mseeger]

Check this out: https://docs.nvidia.com/deeplearning/transformer-engine/. Also: #68

[mseeger]

Apart from the technical aspects (need 3 streams, see above), there is another difficulty for the application during the backward pass. For the computation on a layer/shard and cell:

• Inputs come from CPU: Head gradient, layer input
• Outputs go to CPU: Input gradient, gradients for weights

A safe way to decouple this is to reserve extra buffers on GPU where the streams write to (from CPU) or read from (to CPU). Then, the following can run independent in parallel:

• Transfer input and head gradient for cell k - 1 from CPU
• Do computation for cell k, based on input/head grad for k (transferred in prev step)
• Write input and weights gradients to CPU for cell k + 1 (computed in prev step)

The drawback of this is that GPU memory requirements for these buffers are doubled, since we need to flip between two copies of each: one copy taking part in the current computation, the other in the current transfers.

I don't really see a way around this. To make this work, cells need to be smaller than they could be.

[mseeger]

The issue title is misleading: CPU-GPU transfer plays a key role in both forward and backward. Here are all possible cases:

• forward: Write layer input checkpoints to CPU
• forward, KV cache offloading: CPU-GPU bidirectional streaming
• backward: Load layer input checkpoints from CPU
• backward, head gradients: CPU-GPU bidirectional streaming
• backward, weights/gradients offloading: CPU-GPU bidirectional streaming

There are two steps, do one after the other:

1. Decouple CPU-GPU transfer from GPU compute, using 3 streams: load (CPU->GPU), compute, store (GPU->CPU)
2. Ensure (by cycling buffers) that 3 streams do not interfere (load target is read in compute; store source is written to in compute)

Step 2 is harder in backward, where we have less free GPU memory.

[mseeger]

Here is a plan. For the naming:

• layer is clear
• chunk: Sequence batch split into chunks. Unit passed through the model, unit of KV cache updates
• cell: List of neighboring chunks, until total size about cache length. Forward passes loop over cells (outermost), then layers, then chunks of cell (innermost). Activations for one cell fit in GPU memory, but not activations for complete sequence.

Implements 3 streams solution: load (CPU->GPU), compute (GPU), store (GPU->CPU)

Follow code examples posted above. Don't use Transformer Engine dependence.

Implement tests that this works if there is no interference:

• load target is read by compute
• store source is written by compute
• load target overlaps with store source

The following should be OK:

• store source is read by compute (both are read accesses)

forward: Write layer input checkpoints to CPU

Cell by cell. store and compute work on the same layer input. Synchronization at the end of each cell. No interference.

backward: Load layer input checkpoints from CPU

• Use ring of 2 arrays
• In parallel: load for next cell needed; compute for current cell

forward, KV cache offloading: CPU-GPU bidirectional streaming

Not so critical for gradient computation, because KV cache offloading not really needed then.

TODO!