Use standard gradient checkpointing for small sequence lengths (#3867)

* Use standard gradient checkpointing for small sequence lengths

When max_seq_length < 512, the overhead of gradient offloading in
gc="unsloth" mode is not worth it. Benchmarks on B200 show:

| seq_len | gc=unsloth | gc=True  | Difference |
|---------|------------|----------|------------|
| 256     | 6,803 t/s  | 6,993 t/s| +2.8%      |
| 384     | 9,889 t/s  | 9,963 t/s| +0.7%      |
| 512     | 13,151 t/s | 13,092 t/s| -0.4%     |
| 1024    | 26,662 t/s | 25,094 t/s| -5.9%     |

The crossover point is around seq_len 384-512. For sequences shorter
than 512, we now automatically use standard gradient checkpointing
instead of the custom offloading implementation.

Additionally, when user explicitly sets use_gradient_checkpointing to
True or False in get_peft_model, it now correctly overrides any
previous "unsloth" patching from from_pretrained. This ensures
consistent behavior regardless of the order of function calls.

Updated in three locations:
- FastLlamaModel.get_peft_model (llama.py)
- FastLanguageModel.from_pretrained (loader.py)
- FastModel.from_pretrained (loader.py)

* Refactor: extract gradient checkpointing heuristic into utility function

Addresses code review feedback to reduce duplication. The gradient
checkpointing heuristic logic was duplicated in 3 places:
- FastLlamaModel.get_peft_model (llama.py)
- FastLanguageModel.from_pretrained (loader.py)
- FastModel.from_pretrained (loader.py)

Created apply_unsloth_gradient_checkpointing() utility function in
_utils.py that handles:
- Heuristic: seq < 512 falls back to standard gc
- Explicit True/False overrides unpatch previous patching
- Returns the effective use_gradient_checkpointing value

Net reduction of ~6 lines while improving maintainability.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: danielhanchen <danielhanchen@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

unsloth/models/_utils.py

@@ -59,6 +59,7 @@ __all__ = [
     "unsloth_fused_ce_loss",
     "patch_unsloth_smart_gradient_checkpointing",
     "unpatch_unsloth_smart_gradient_checkpointing",
+    "apply_unsloth_gradient_checkpointing",
     "patch_compiled_autograd",
     "process_vision_info",
     "unsloth_compile_transformers",
@@ -148,6 +149,41 @@ from unsloth_zoo.temporary_patches import (
     TEMPORARY_PATCHES,
 )
 
+
+def apply_unsloth_gradient_checkpointing(
+    use_gradient_checkpointing, max_seq_length, dtype
+):
+    """
+    Apply gradient checkpointing with smart heuristics.
+
+    For seq < 512, the overhead of gradient offloading in gc="unsloth" mode
+    is not worth it. Benchmarks show standard gc is faster for small sequences.
+
+    Args:
+        use_gradient_checkpointing: "unsloth", True, False, or None
+        max_seq_length: The maximum sequence length
+        dtype: The model dtype for patching
+
+    Returns:
+        The effective use_gradient_checkpointing value (may change from "unsloth" to True)
+    """
+    if use_gradient_checkpointing == "unsloth":
+        # Gradient offloading overhead is not worth it for small sequences.
+        # Benchmarks show crossover point is around seq_len 384-512.
+        # For seq < 512, standard gradient checkpointing is faster.
+        if max_seq_length < 512:
+            unpatch_unsloth_smart_gradient_checkpointing()
+            return True
+        else:
+            patch_unsloth_smart_gradient_checkpointing(dtype = dtype)
+            return "unsloth"
+    elif use_gradient_checkpointing in (True, False):
+        # User explicitly set True or False - unpatch any previous "unsloth" patching
+        unpatch_unsloth_smart_gradient_checkpointing()
+        return use_gradient_checkpointing
+    return use_gradient_checkpointing
+
+
 for temporary_patch in TEMPORARY_PATCHES:
     temporary_patch()
 

unsloth/models/llama.py

@@ -19,7 +19,7 @@ import functools
 from typing import Optional, Tuple, List, Union
 
 from ._utils import *
-from ._utils import patch_unsloth_smart_gradient_checkpointing
+from ._utils import apply_unsloth_gradient_checkpointing
 from ._utils import __version__, importlib_version
 from ._utils import move_to_device
 from ._utils import (
@@ -2693,10 +2693,12 @@ class FastLlamaModel:
             return model
         transformers_set_seed(random_state)
 
-        if use_gradient_checkpointing == "unsloth":
-            patch_unsloth_smart_gradient_checkpointing(
-                dtype = model.get_input_embeddings().weight.dtype
-            )
+        # Apply gradient checkpointing with smart heuristics
+        max_seq = getattr(model, "max_seq_length", 512)
+        dtype = model.get_input_embeddings().weight.dtype
+        use_gradient_checkpointing = apply_unsloth_gradient_checkpointing(
+            use_gradient_checkpointing, max_seq, dtype
+        )
 
         if type(r) is not int:
             raise TypeError(f"Unsloth: Rank of {str(r)} must be an integer.")

unsloth/models/loader.py

@@ -88,7 +88,7 @@ from ._utils import (
     patch_compiling_bitsandbytes,
     patch_model_and_tokenizer,
     prepare_model_for_kbit_training,
-    patch_unsloth_smart_gradient_checkpointing,
+    apply_unsloth_gradient_checkpointing,
     patch_compiled_autograd,
     process_vision_info,
     unsloth_compile_transformers,
@@ -559,8 +559,10 @@ class FastLanguageModel(FastLlamaModel):
                 **kwargs,
             )
 
-        if use_gradient_checkpointing == "unsloth":
-            patch_unsloth_smart_gradient_checkpointing(dtype = dtype)
+        # Apply gradient checkpointing with smart heuristics
+        use_gradient_checkpointing = apply_unsloth_gradient_checkpointing(
+            use_gradient_checkpointing, max_seq_length, dtype
+        )
 
         # Check if this is local model since the tokenizer gets overwritten
         if (
@@ -1188,9 +1190,10 @@ class FastModel(FastBaseModel):
                 os.environ["UNSLOTH_FORCE_FLOAT32"] = "1"
                 dtype = torch.bfloat16  # Change to bfloat16 loading
                 break
-        # Patch gradient checkpointing
-        if use_gradient_checkpointing == "unsloth":
-            patch_unsloth_smart_gradient_checkpointing(dtype = dtype)
+        # Apply gradient checkpointing with smart heuristics
+        use_gradient_checkpointing = apply_unsloth_gradient_checkpointing(
+            use_gradient_checkpointing, max_seq_length, dtype
+        )
         with redirector:
             patch_loss_functions(torch_compile = False)
             model_types, supports_sdpa = unsloth_compile_transformers(