Repository navigation
Honour every_n_layers and drop redundant offload in gradient_checkpointing_enable - #14403
danielhanchen wants to merge 8 commits into
Conversation
…nting_enable every_n_layers and offload are gradient_checkpointing_enable arguments (transformers 5.17+), not torch.utils.checkpoint ones. Passed inside gradient_checkpointing_kwargs (TrainingArguments on older transformers, direct calls) the Unsloth checkpoint shims bound them into every decoder layer call, so both were silently ignored. every_n_layers is now applied on any transformers version through the skip-aware layer flag, so for_training / for_inference no longer reset it. offload=True is dropped while Unsloth's offloading checkpoint is active, where it only re-offloaded the inputs Unsloth keeps on the GPU.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2f48d4866e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| "Unsloth: gradient checkpointing already offloads activations to host memory, " | ||
| "so transformers' offload=True is ignored.", | ||
| ) | ||
| offload = False |
There was a problem hiding this comment.
When the active checkpoint is unsloth_checkpoint, this unconditionally disables a direct Transformers 5.17 offload=True request, but the smart checkpointer only offloads inputs above its MINIMUM_SIZE and deliberately keeps the final checkpoint input on-device. The added test even sets MINIMUM_SIZE extremely high and then verifies that save_on_cpu is never entered, demonstrating that small activations receive no offloading at all. This violates the caller's explicit memory-saving request and can leave a marginal run OOM; native offload should be suppressed only for checkpoints the shim will actually offload.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Not changing this. The tensors Unsloth's checkpointer keeps on the GPU are by design: inputs under its 2 MiB floor (unsloth_zoo gradient_checkpointing.py MINIMUM_SIZE), whose host round trip costs more than it saves, and the last checkpoint input, which backward consumes first. Running save_on_cpu on top copies those same small or immediately needed tensors to the host and back, which is the double offload this removes; every input large enough to matter is still offloaded by the shim. The test raises MINIMUM_SIZE only to prove the native hook is not entered, not to model real sizes.
| layers = [m for m in model.modules() if isinstance(m, GradientCheckpointingLayer)] | ||
| skipped = [] | ||
| for i, layer in enumerate(layers): | ||
| if every_n_layers > 1 and i % every_n_layers != 0 and hasattr(layer, "gradient_checkpointing"): |
There was a problem hiding this comment.
When every_n_layers > 1 is combined with offload_layers, this loop can mark block-swapped decoder layers as skipped because it ignores model._unsloth_block_swap.indices. For example, spread placement on four layers with two swapped layers selects indices 1 and 3, and every_n_layers=2 disables checkpointing on exactly those layers. install_block_swap requires per-layer checkpointing for swapped blocks; otherwise their weights remain resident through backward and the slot pool grows, defeating the memory plan and potentially causing the OOM that block swapping was intended to avoid. Swapped layer objects must remain enabled regardless of the interval.
Useful? React with 👍 / 👎.
for more information, see https://pre-commit.ci
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
for more information, see https://pre-commit.ci
…y_n_layers to remote layers vision.py's _nonre_checkpoint already offloads through _offload_activation_hooks, so it now carries an _unsloth_offloads marker that the offload=True suppression reads. Remote-code decoder layers wrapped by install_remote_gradient_checkpointing read a per-layer every_n_layers mark in their wrapper, so every_n_layers > 1 skips them instead of only warning. [skip ci]
for more information, see https://pre-commit.ci
[skip ci]
This branch has not been deployed
1 participant
Summary
transformers 5.17 added
every_n_layersandoffloadtogradient_checkpointing_enable(). Under Unsloth:gradient_checkpointing_kwargs(TrainingArguments on transformers < 5.17, direct calls, PEFT'sprepare_model_for_kbit_training), both keys were bound into every decoder layer call by the Unsloth checkpoint shims. They were silently ignored, and every layer stayed checkpointed.for_training/for_inference(every GRPO generation step) rewrote every layer's flag, so a nativeevery_n_layerswas reset after the first generation.offload=Truestackedsave_on_cpuon top of Unsloth's own offload, copying the input Unsloth keeps on the GPU.This wraps
PreTrainedModel.gradient_checkpointing_enableonce:every_n_layersuses transformers' rule (GradientCheckpointingLayer 0, N, 2N, ...) on any version, held by the existing skip-aware layer flag sofor_trainingkeeps it.offload=Trueis dropped with a one-time message while Unsloth's offloading checkpoint is active.Tests
tests/test_gc_enable_kwargs.py(CPU, tiny Qwen3, fp32 grads vs a no-checkpointing reference): main 3 failed / 2 passed, this PR 5 passed on torch 2.10, 2.13, 2.14 with transformers 5.17; transformers 4.57.6 main 2 failed, this PR 3 passed (the 2 native-offload cases need 5.17).every_n_layers=2checkpoints 14 layers per step instead of 28;offload=Truemakes 0 extra host copies instead of 1.Not covered
GRPO end to end with
every_n_layers; transformers < 4.52 (no GradientCheckpointingLayer: every layer stays checkpointed).