ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix transformers v5 RoPE inv_freq corruption and generate() BatchEncoding compat (#4112)

Browse source 
* 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>
This commit is contained in:
Daniel Han 2026-02-25 08:18:45 -08:00 • committed by GitHub
commit 0e64336ac9
2 changed files with 114 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

18
unsloth/models/llama.py
View file

@ -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

99
unsloth/models/loader.py
View file

@ -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