Add offload to gradient checkpointing

[qgallouedec]

Gradient checkpointing still keeps one activation per checkpointed layer on the device. At long sequence lengths that is the dominant term: layers x sequence x hidden bytes, so Qwen3-8B at 128k tokens is 36 GB of it. This adds an option to hold those in pinned host memory instead.

args = TrainingArguments(
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"offload": True},
)

It is a wrap of gradient_checkpointing_func around torch's save_on_cpu(pin_memory=True), so it follows whichever path the model already uses and needs no per-model support.

It frees exactly what the arithmetic says

Memory resident after the forward, where every layer's saved input is still live. Qwen3 config, bf16, one H100:

layers x hidden x seq	resident, plain	resident, offload	freed	layers x seq x hidden x 2
8 x 1024 x 8k	194.1 MB	66.1 MB	128.0 MB	128.0 MB
16 x 1024 x 16k	644.2 MB	132.2 MB	512.0 MB	512.0 MB
28 x 2048 x 32k	4112.4 MB	528.4 MB	3584.0 MB	3584.0 MB

(How much of that shows up in peak memory depends on where your peak is. On the 28-layer run above the peak only moved 11.13 -> 10.16 GB, because it sits at the end of the backward where gradients dominate and the saved inputs are already freed. The option pays off when activations dominate the peak, which is the long-sequence case it is meant for.)

Cost

Both copies run on the compute stream, so the step gets slower. There is no separate copy stream and no backward prefetch; that is pytorch/pytorch#158657, and this option does not wait for it.

Setups

Qwen3-1.7B base, 32k tokens, peak memory. The DeepSpeed row is at 8k: at 32k that engine takes a materializing attention path and asks for a 64 GiB allocation, unrelated to this option.

	plain	offload
single GPU	11.13 GB	10.16 GB
FSDP2, 2 GPUs	9.73 GB	8.06 GB
DeepSpeed ZeRO-3, 2 GPUs (8k tokens)	34.12 GB	32.93 GB

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 33466313732:1
Result: success | Jobs: 16 | Tests: 175,490 | Failures: 0 | Duration: 14h 44m

[ArthurZucker]

Perfect PR!

Maybe a todo is to auto enable this in case of 90% memory taken, or we detect long sequence length etc
Would recompile potentially, but its a great default no?

[qgallouedec]

I'll have to measure this!