From 4a8edd5776c42cc8c05e7ce8d0bc0f19792cc537 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 2 Feb 2026 23:57:09 -0800 Subject: [PATCH] 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 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/_utils.py | 36 ++++++++++++++++++++++++++++++++++++ unsloth/models/llama.py | 12 +++++++----- unsloth/models/loader.py | 15 +++++++++------ 3 files changed, 52 insertions(+), 11 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 76952b00a5..301cdf7b80 100644 --- a/unsloth/models/_utils.py +++ b/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() diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 39f2ba1460..fb3b8133a5 100644 --- a/unsloth/models/llama.py +++ b/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.") diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index eb3b21e206..ef7e67b468 100644 --- a/unsloth/models/loader.py +++ b/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(