Why gradient checkpointing breaks your fine-tuning run, and how to fix it
UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly
Also appears as
- RuntimeError: Trying to backward through the graph a second time
- RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
Short answer
Gradient checkpointing errors during fine-tuning almost always come from three sources: the use_reentrant parameter left unset (it now must be explicit and False is usually correct for transformer models), an attention implementation that isn't fully compatible with checkpointing's re-computation approach, or leaving use_cache=True enabled while checkpointing is on, which conflicts because checkpointing recomputes the forward pass and a live KV cache assumes it won't be recomputed. Set use_reentrant=False and use_cache=False together.
Affects: HuggingFace transformers with gradient_checkpointing=True, any model, especially with FlashAttention or custom attention kernels
Fastest path to working gradient checkpointing
- 1Explicitly set gradient_checkpointing_kwargs={"use_reentrant": False} when calling model.gradient_checkpointing_enable(), since the default has changed across PyTorch versions and leaving it implicit causes warnings or errors.
- 2Set model.config.use_cache = False before training whenever gradient_checkpointing is enabled; the two are fundamentally incompatible since checkpointing recomputes forward passes and caching assumes it won't.
- 3If using FlashAttention 2, confirm your transformers and flash-attn versions are compatible with reentrant=False checkpointing; some older combinations only support the reentrant=True path.
- 4If the error mentions backward through the graph a second time, check for any code that calls loss.backward() more than once on the same graph, which conflicts with checkpointing's single recomputation assumption.
- 5For PEFT/LoRA training, call model.enable_input_require_grads() before enabling gradient checkpointing, since checkpointing needs the input embeddings to require gradients even though the base model itself is frozen.
How to confirm this is your problem
- A UserWarning appears at startup about use_reentrant needing to be passed explicitly
- Training crashes with a graph already freed or backward through the graph a second time error
- Error mentions a tensor that does not require grad and does not have a grad_fn, often the input embeddings
- Training works without gradient_checkpointing enabled but fails immediately when it's turned on
- Memory usage does not drop as expected even though gradient_checkpointing is enabled
Root causes and fixes
use_reentrant parameter left unset
PyTorch changed torch.utils.checkpoint to require this parameter explicitly because the old reentrant-based implementation has known limitations (it doesn't support certain autograd features and can silently produce incorrect gradients in edge cases). Leaving it unset triggers a warning in newer PyTorch versions and can cause outright failures in newer transformers versions that assume it will be set.
Fix: Always pass gradient_checkpointing_kwargs={"use_reentrant": False} when calling gradient_checkpointing_enable(), which uses the newer, more robust non-reentrant implementation that most current transformer architectures support correctly.
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False})use_cache left True while gradient checkpointing is enabled
use_cache=True tells the model to store key/value states from previous forward passes for reuse, which is meant for autoregressive generation. Gradient checkpointing works by discarding intermediate activations and recomputing them during the backward pass, an approach that is fundamentally incompatible with also trying to cache and reuse those same values, producing a graph inconsistency error.
Fix: Set model.config.use_cache = False explicitly before starting training whenever gradient checkpointing is active; most trainers do this automatically but custom training loops often miss it.
model.config.use_cache = False
Input embeddings do not require gradients under PEFT/LoRA
When only LoRA adapter weights are trainable and the base model (including the embedding layer) is frozen, the very first operation in the forward pass produces a tensor with requires_grad=False. Gradient checkpointing needs at least the checkpointed segment's input to require gradients so it can correctly recompute and backpropagate through that segment; without this, the backward pass has no gradient function to call.
Fix: Call model.enable_input_require_grads() (a PEFT/transformers helper) before enabling gradient checkpointing, which forces the embedding output to require gradients without making the embedding weights themselves trainable.
model.enable_input_require_grads()
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False})Incompatibility with a specific attention implementation
Some custom or fused attention kernels (certain FlashAttention integration paths, some SDPA backend combinations) were written assuming a particular checkpointing mode and can produce shape errors or silent incorrect gradients when combined with the other mode. This is more common right after upgrading either the attention library or transformers without re-testing the combination.
Fix: Check the release notes for your specific attention implementation and transformers version for known gradient checkpointing compatibility notes, and fall back to eager attention implementation temporarily to confirm whether the attention kernel is the source of the error.
model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation='eager')
Calling backward() twice on the same checkpointed graph
Non-reentrant checkpointing recomputes the forward pass exactly once during backward and then discards the recomputation graph. Custom training loops that accidentally call loss.backward() a second time (for example, once for logging gradient norms and again for the optimizer step) violate this single-use assumption and raise a graph already freed style error.
Fix: Ensure backward() is called exactly once per forward pass in custom training loops, and use retain_graph=True only if you have a specific, deliberate reason to reuse the graph, understanding it increases memory usage.
Diagnostic commands
Check current gradient checkpointing and cache settings
python -c "print(model.is_gradient_checkpointing, model.config.use_cache)"
is_gradient_checkpointing should be True and use_cache should be False during training. If use_cache is still True, that mismatch is likely causing the error.
Confirm input embeddings require gradients under LoRA
python -c "print(model.get_input_embeddings().weight.requires_grad)"
This will correctly show False for a frozen base model under LoRA; the relevant check is whether enable_input_require_grads() was called, which affects the embedding output tensor, not the weight itself.
Test with eager attention to isolate the attention kernel
python -c "from transformers import AutoModelForCausalLM; m=AutoModelForCausalLM.from_pretrained('MODEL', attn_implementation='eager')"If the error disappears with eager attention but reappears with flash_attention_2 or sdpa, the specific attention kernel's interaction with checkpointing is the root cause, not your training script logic.
Stopping it from happening again
- Always pass use_reentrant explicitly (False in almost all current transformer training setups) rather than relying on framework defaults.
- Standardize a training script template that always pairs gradient_checkpointing_enable with use_cache=False and enable_input_require_grads for PEFT.
- Pin transformers, PEFT, and attention library versions together and re-test the combination after any upgrade.
- Run a short smoke test (a few dozen steps) with gradient checkpointing enabled before committing to a multi-hour training run.
- Keep eager attention as a documented fallback path for quickly isolating whether a new error is attention-kernel related.
When this becomes an architecture problem
If gradient checkpointing errors persist across attention implementations and PyTorch versions on a model architecture you need in production, that points to either an upstream bug worth reporting or a need to redesign the training pipeline around a different memory-saving strategy (like sequence packing or lower precision) instead of checkpointing. If you're scaling training across multiple GPUs or nodes and checkpointing interacts unpredictably with your parallelism strategy, that combination is worth validating with someone experienced in multi-GPU training infrastructure.
Frequently asked questions
What should use_reentrant be set to for transformer fine-tuning?
False, in almost all current transformer training setups. The non-reentrant implementation is more robust and supports more autograd features correctly, and current transformers and PEFT code paths are built and tested against use_reentrant=False as the expected configuration.
Why do I need use_cache=False with gradient checkpointing?
Gradient checkpointing recomputes forward-pass activations during the backward pass instead of storing them, while use_cache stores key/value states specifically so they don't need to be recomputed. These two mechanisms make contradictory assumptions about recomputation, so use_cache must be disabled during training whenever checkpointing is active.
Why does gradient checkpointing fail specifically with LoRA but not full fine-tuning?
With LoRA, the base model's embedding layer is frozen and its output does not require gradients by default. Gradient checkpointing needs a gradient-requiring tensor entering each checkpointed segment to correctly recompute and backpropagate, so PEFT-based training needs the extra enable_input_require_grads() call that full fine-tuning doesn't.
Does gradient checkpointing actually save the memory it claims to?
Yes, typically 30-60% activation memory savings depending on model depth, at the cost of roughly 20-30% more compute time from the recomputation during backward. If memory usage doesn't drop after enabling it, check that is_gradient_checkpointing is actually True and use_cache is False, since a misconfiguration can silently make checkpointing a no-op.
Size it properly next time
Free calculators that prevent this class of failure before you provision hardware.
Fine-Tuning GPU-Hours Estimator
Estimate GPU-hours for LoRA, QLoRA, and full fine-tuning on the same model size and dataset, so you can compare method tradeoffs before choosing.
Free ToolLoRA Fine-Tuning Cost Calculator
Turn model size, dataset tokens, epochs, and rank into a GPU-hour and dollar estimate for a LoRA fine-tuning run on rented or owned hardware.
Free ToolGPU Sizing Calculator for LLM Inference
Work out how many GPUs you need to serve a given open-weight model to your user base, based on memory footprint and token throughput.
Free ToolKV Cache Memory Calculator
Calculate KV cache memory per sequence and per batch from model architecture and context length, then see how many concurrent sequences your GPU can hold.
Related problems
Loss becomes NaN during fine-tuning
NaN loss during training is most often caused by fp16 numeric overflow in gradients or activations, which bf16 avoids because of its wider exponent range. Other common causes are a learning rate spike (especially right after warmup), a small number of corrupt or malformed training samples, and unsafe division or log operations in a custom loss function. Switch to bf16 first if your hardware supports it, then check for corrupt samples and unstable LR.
CUDA out of memory even though nvidia-smi shows free VRAM
This almost always means memory fragmentation: the allocator has enough total free memory but no single contiguous block large enough for the requested allocation, because the address space is broken into many small free-and-used segments from prior allocations of different sizes. The fix is enabling PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, reducing allocation size variability, or restarting the process to reset the address space.
Training loss not decreasing during fine-tuning
Training loss that stays flat is most often caused by a LoRA adapter that targets the wrong modules (missing q_proj/k_proj/v_proj/o_proj), a learning rate that is too low for LoRA or too high and bouncing, or labels that were never masked so the model is trying to learn the prompt tokens as if they were random noise. Check target_modules first, then the label mask, then the learning rate.
Out of memory when merging a LoRA adapter into the base model
Merging a LoRA adapter mathematically requires the base weights in a real-valued (not 4-bit quantized) format so the adapter delta can be added in, which means a QLoRA workflow that trained happily in 4-bit suddenly needs a full fp16/bf16 copy of the base model just to merge, often doubling or more the memory footprint versus either training or inference alone. The fix is to merge on CPU, merge in a lower-footprint dtype, or skip merging entirely by serving the adapter unmerged.
GuideFine-Tuning Failure Modes: What Actually Goes Wrong
Fine-tuning failure modes that actually derail enterprise projects: catastrophic forgetting, eval overfitting, data leakage, and how to catch each one.
GuideLoRA vs QLoRA: Choosing the Right Fine-Tuning Method
LoRA vs QLoRA for enterprise fine-tuning: rank and alpha choices, real VRAM math by model size, and when each method actually wins.
GuideMulti-Node LLM Training Infrastructure: Networking and Storage
Multi-node LLM training infrastructure explained: InfiniBand vs RoCE tradeoffs, storage throughput needs, and cluster topology for enterprise fine-tuning.
GuideKV Cache Optimization: Prefix Caching and Chunked Prefill
KV cache optimization techniques for production LLM serving: prefix caching, chunked prefill, PagedAttention, and sizing memory for concurrent users.
Still stuck, or tired of fighting your own infrastructure?
Netray deploys and operates on-prem AI for regulated manufacturers and defense suppliers. We have debugged this stack in production, on air-gapped networks, at scale.