# Why gradient checkpointing breaks your fine-tuning run, and how to fix it > Source: https://netray.co/troubleshoot/gradient-checkpointing-errors > Publisher: Netray (netray.co), AI and integration engineering for ERP and manufacturing systems. > Section: On-Prem AI Troubleshooting ## 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. ## Error Title Gradient checkpointing errors during fine-tuning ## Error Signature UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly ## Alternate Signatures - 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 ## Category fine-tuning ## Stack - pytorch - huggingface - transformers - accelerate ## Affects HuggingFace transformers with gradient_checkpointing=True, any model, especially with FlashAttention or custom attention kernels ## Quick Fix ### Fastest path to working gradient checkpointing **Steps:** Explicitly 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. Set 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. If 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. If 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. For 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. ## Symptoms - 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 ## Causes ### use_reentrant parameter left unset **Likelihood:** most-common **Explanation:** 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. **Commands:** - model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False}) ### use_cache left True while gradient checkpointing is enabled **Likelihood:** common **Explanation:** 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. **Commands:** - model.config.use_cache = False ### Input embeddings do not require gradients under PEFT/LoRA **Likelihood:** common **Explanation:** 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. **Commands:** - model.enable_input_require_grads() - model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False}) ### Incompatibility with a specific attention implementation **Likelihood:** occasional **Explanation:** 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. **Commands:** - model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation='eager') ### Calling backward() twice on the same checkpointed graph **Likelihood:** rare **Explanation:** 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. ## Diagnostics ### Check current gradient checkpointing and cache settings **Command:** python -c "print(model.is_gradient_checkpointing, model.config.use_cache)" **Interpretation:** 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 **Command:** python -c "print(model.get_input_embeddings().weight.requires_grad)" **Interpretation:** 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 **Command:** python -c "from transformers import AutoModelForCausalLM; m=AutoModelForCausalLM.from_pretrained('MODEL', attn_implementation='eager')" **Interpretation:** 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. ## Prevention - 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 To Escalate 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. ## Faq ### What should use_reentrant be set to for transformer fine-tuning? **Answer:** 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? **Answer:** 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? **Answer:** 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? **Answer:** 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. ## Search terms gradient checkpointing error, use_reentrant parameter, backward through graph a second time, use_cache gradient checkpointing conflict, enable_input_require_grads, gradient checkpointing LoRA error, does not require grad no grad_fn, flash attention gradient checkpointing --- Talk to an engineer: https://netray.co/contact?intent=erp_ai | Machine catalog: https://netray.co/agent/catalog.json