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