From 0e64336ac9a09f72f3b2b1525e85f13ded5bfccc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Feb 2026 08:18:45 -0800 Subject: [PATCH] Fix transformers v5 RoPE inv_freq corruption and generate() BatchEncoding compat (#4112) * Fix transformers v5 RoPE inv_freq corruption during model loading Transformers v5 initializes models on the meta device, then _move_missing_keys_from_meta_to_device() replaces all non-persistent buffers with torch.empty_like() (uninitialized memory). Vanilla transformers restores inv_freq via _init_weights() checking for original_inv_freq, but Unsloth's LlamaRotaryEmbedding subclasses lack this attribute, so inv_freq stays corrupted with garbage values. This caused 5-11x higher training loss on transformers v5 for all models using Unsloth's rope (Llama 3.x, Qwen3, Mistral, TinyLlama, Granite). Models using native transformers rope (Gemma, Phi-4, Falcon-H1) were unaffected. The fix recomputes inv_freq from the stored base/dim after model loading, applies model-specific scaling via _apply_inv_freq_scaling(), and rebuilds cos/sin caches. Also handles LongRopeRotaryEmbedding (Phi-3.5 style short/long inv_freq). Guarded by transformers >= 5.0.0 so it is a no-op on v4. Tested on: Llama 3.1 8B, Llama 3.2 3B, Qwen3 14B, Qwen3 4B, Phi-4, TinyLlama, Mistral 7B, Gemma2 2B, Falcon-H1 -- all v5 losses now match v4 baselines to < 0.004 absolute difference. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Unpack BatchEncoding in generate() for v4/v5 backwards compatibility Old notebooks pass the full tokenizer output as input_ids: inputs = tokenizer(..., return_tensors="pt").to("cuda") model.generate(input_ids=inputs, ...) This worked on transformers v4 because generate() internally extracted the tensor. Transformers v5 calls .shape on input_ids directly, which crashes since BatchEncoding has no .shape attribute. Fix: in unsloth_fast_generate(), detect when input_ids is a dict-like object (BatchEncoding) and unpack its contents into separate kwargs before forwarding to the underlying generate(). This makes both old and new notebook patterns work on both v4 and v5. * Remove redundant seen_ids dedup in _fix_rope_inv_freq named_modules() already deduplicates with remove_duplicate=True (default). Also clarify that native v5 rotary classes (Gemma3 etc.) have original_inv_freq which transformers v5's _init_weights() uses to restore inv_freq, so they do not need this fix. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/llama.py | 18 ++++++-- unsloth/models/loader.py | 99 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 114 insertions(+), 3 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index f80a55fdd2..6fe60cf940 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2062,6 +2062,21 @@ def unsloth_fast_generate( FastLlamaModel.for_inference(self) + # Unpack BatchEncoding passed as input_ids for backwards compatibility. + # Old notebooks do model.generate(input_ids=tokenizer(...)) where the tokenizer + # output is a BatchEncoding (dict-like). Transformers v5 generate() calls + # .shape on it directly and crashes. Unpack into separate kwargs so both + # v4 and v5 work transparently. + _maybe_encoding = kwargs.get("input_ids", None) + if ( + _maybe_encoding is not None + and not isinstance(_maybe_encoding, torch.Tensor) + and hasattr(_maybe_encoding, "items") + ): + batch_data = kwargs.pop("input_ids") + for key, val in batch_data.items(): + kwargs.setdefault(key, val) + dtype = _get_dtype(dtype_from_config(self.config)) if hasattr(self, "config") and hasattr(self.config, "max_position_embeddings"): @@ -2071,9 +2086,6 @@ def unsloth_fast_generate( and "max_new_tokens" in kwargs ): _ids = kwargs["input_ids"] - # Handle BatchEncoding from transformers 5.0+ (no .shape attribute) - if hasattr(_ids, "input_ids"): - _ids = _ids["input_ids"] if hasattr(_ids, "shape") and ( _ids.shape[-1] + kwargs["max_new_tokens"] > self.config.max_position_embeddings diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 72a6d7d10f..711476b759 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -78,6 +78,9 @@ SUPPORTS_QWEN3_MOE = transformers_version >= Version("4.50.3") SUPPORTS_FALCON_H1 = transformers_version >= Version("4.53.0") SUPPORTS_GEMMA3N = transformers_version >= Version("4.53.0") SUPPORTS_GPTOSS = transformers_version >= Version("4.55.0") +# Transformers v5 meta-device loading corrupts non-persistent buffers (inv_freq). +# See _fix_rope_inv_freq() below for details. +_NEEDS_ROPE_FIX = transformers_version >= Version("5.0.0") if SUPPORTS_GEMMA: from .gemma import FastGemmaModel if SUPPORTS_GEMMA2: @@ -121,6 +124,100 @@ DISABLE_SDPA_MODEL_NAMES = [ ] +def _fix_rope_inv_freq(model): + """Fix inv_freq corruption caused by transformers v5 meta-device loading. + + Transformers v5 initializes models on the meta device, then + _move_missing_keys_from_meta_to_device() (modeling_utils.py) replaces ALL + non-persistent buffers with torch.empty_like() -- uninitialized memory. + + Vanilla transformers restores inv_freq via _init_weights() which checks for + hasattr(module, "original_inv_freq"). Unsloth's LlamaRotaryEmbedding and + subclasses do not have this attribute, so inv_freq stays corrupted. This + produces wrong positional encodings and causes 5-11x higher training loss. + + This function recomputes inv_freq from the stored base and dim, applies + any model-specific scaling, and rebuilds the cos/sin caches. + + Only runs on transformers >= 5.0.0. No-op on v4. + """ + if not _NEEDS_ROPE_FIX: + return model + + for name, module in model.named_modules(): + # Unsloth's LlamaRotaryEmbedding and subclasses (Extended, LinearScaling, + # Granite). Native v5 rotary classes (Gemma3, etc.) have original_inv_freq + # which v5's _init_weights() uses to restore inv_freq, so they are fine. + if ( + hasattr(module, "inv_freq") + and hasattr(module, "base") + and hasattr(module, "dim") + and hasattr(module, "_apply_inv_freq_scaling") + and hasattr(module, "multi_gpu_cos_cached") + ): + inv_freq = 1.0 / ( + module.base + ** ( + torch.arange( + 0, module.dim, 2, dtype = torch.int64, device = "cpu" + ).float() + / module.dim + ) + ) + inv_freq = module._apply_inv_freq_scaling(inv_freq) + module.inv_freq = inv_freq + for device_idx in range(len(module.multi_gpu_cos_cached)): + if module.multi_gpu_cos_cached[device_idx] is not None: + module._set_cos_sin_cache( + seq_len = module.current_rope_size, + device = torch.device(device_idx), + dtype = torch.get_default_dtype(), + ) + + # LongRopeRotaryEmbedding (Phi-3.5 style with short_inv_freq + long_inv_freq) + elif ( + hasattr(module, "short_inv_freq") + and hasattr(module, "long_inv_freq") + and hasattr(module, "base") + and hasattr(module, "dim") + ): + config = getattr(model, "config", None) + rope_scaling = getattr(config, "rope_scaling", None) if config else None + if rope_scaling is not None: + short_factor = rope_scaling.get("short_factor", None) + long_factor = rope_scaling.get("long_factor", None) + if short_factor is not None and long_factor is not None: + inv_freq_shape = ( + torch.arange( + 0, module.dim, 2, dtype = torch.int64, device = "cpu" + ).float() + / module.dim + ) + sf = torch.tensor(short_factor, device = "cpu", dtype = torch.float32) + lf = torch.tensor(long_factor, device = "cpu", dtype = torch.float32) + module.short_inv_freq = 1.0 / (sf * module.base**inv_freq_shape) + module.long_inv_freq = 1.0 / (lf * module.base**inv_freq_shape) + + dtype = torch.bfloat16 if is_bfloat16_supported() else torch.float16 + t = torch.arange( + module.original_max_position_embeddings, + device = module.short_inv_freq.device, + dtype = torch.int64, + ).float() + freqs = torch.outer(t, module.short_inv_freq) + emb = torch.cat((freqs, freqs), dim = -1) + for device_idx in range(len(module.multi_gpu_short_cos_cached)): + if module.multi_gpu_short_cos_cached[device_idx] is not None: + device_obj = torch.device(device_idx) + module.multi_gpu_short_cos_cached[device_idx] = ( + emb.cos() * module.scaling_factor + ).to(dtype = dtype, device = device_obj, non_blocking = True) + module.multi_gpu_short_sin_cached[device_idx] = ( + emb.sin() * module.scaling_factor + ).to(dtype = dtype, device = device_obj, non_blocking = True) + return model + + class FastLanguageModel(FastLlamaModel): @staticmethod def from_pretrained( @@ -685,6 +782,7 @@ class FastLanguageModel(FastLlamaModel): if patch_tiled_mlp_choice != "0" or unsloth_tiled_mlp: patch_tiled_mlp(model, patch_options_str = patch_tiled_mlp_choice) + model = _fix_rope_inv_freq(model) return model, tokenizer @@ -1408,6 +1506,7 @@ class FastModel(FastBaseModel): if patch_tiled_mlp_choice != "0" or unsloth_tiled_mlp: patch_tiled_mlp(model, patch_options_str = patch_tiled_mlp_choice) + model = _fix_rope_inv_freq(model) return model, tokenizer