From 656a0875936c84c81dacb602e60564bfe5020d97 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 10 Jun 2026 05:45:00 -0700 Subject: [PATCH] Restore config use_cache in for_inference after gradient checkpointing prep (#6137) * Restore config use_cache in for_inference after gradient checkpointing prep unsloth_zoo's prepare_model_for_training now sets use_cache=False on the model config and nested sub-configs when gradient checkpointing is on (unsloth-zoo PR 715) and records the original values. Wire the counterpart into both for_inference implementations so the original values come back for inference, and re-disable in for_training when a record exists so resumed training keeps the config consistent. Both calls import lazily and tolerate older unsloth_zoo without the helpers, so version skew in either direction is a no-op. The MLX shims in __init__.py are deliberately untouched. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/llama.py | 20 ++++++++++++++++++++ unsloth/models/vision.py | 20 ++++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 6018a4ccb4..d5883e4346 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -3467,6 +3467,14 @@ class FastLlamaModel: embeddings = model.get_output_embeddings() if hasattr(embeddings, "training"): embeddings.training = False + + # Restore use_cache values that prepare_model_for_training disabled + # for gradient checkpointing (older unsloth_zoo has no restore helper) + try: + from unsloth_zoo.training_utils import restore_use_cache + restore_use_cache(model) + except ImportError: + pass return model @staticmethod @@ -3514,6 +3522,18 @@ class FastLlamaModel: embeddings = model.get_output_embeddings() if hasattr(embeddings, "training"): embeddings.training = True + + # Re-disable use_cache if prepare_model_for_training had disabled it + # and for_inference restored it (record only exists after a disable) + if ( + use_gradient_checkpointing + and getattr(model, "_unsloth_use_cache_originals", None) is not None + ): + try: + from unsloth_zoo.training_utils import disable_use_cache + disable_use_cache(model) + except ImportError: + pass return model diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index a75d23e710..34336aa80f 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -1706,6 +1706,14 @@ class FastBaseModel: embeddings = model.get_output_embeddings() if hasattr(embeddings, "training"): embeddings.training = False + # Restore use_cache values that prepare_model_for_training disabled + # for gradient checkpointing (older unsloth_zoo has no restore helper) + try: + from unsloth_zoo.training_utils import restore_use_cache + restore_use_cache(model) + except ImportError: + pass + # Must disable returning hidden states in the case for GRPO os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "0" # Must enable returning logits @@ -1764,6 +1772,18 @@ class FastBaseModel: embeddings = model.get_output_embeddings() if hasattr(embeddings, "training"): embeddings.training = True + # Re-disable use_cache if prepare_model_for_training had disabled it + # and for_inference restored it (record only exists after a disable) + if ( + use_gradient_checkpointing + and getattr(model, "_unsloth_use_cache_originals", None) is not None + ): + try: + from unsloth_zoo.training_utils import disable_use_cache + disable_use_cache(model) + except ImportError: + pass + # Can re-enable not returning logits os.environ["UNSLOTH_RETURN_LOGITS"] = "0" # Turn off skip guards and set stance to default