Fix LLama RoPE

[mntss]

Description

The current LLama3.x implementations use incorrect scaling factors for RoPE

https://github.com/huggingface/transformers/blob/3165eb7c2808832d0de86c8f508d9da6b2124044/src/transformers/modeling_rope_utils.py#L407-L410

Type of change

Please delete options that are not relevant.

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Screenshots

I used this script for testing:

# %%
import os
import torch
from torch.utils.data import DataLoader
from transformer_lens import HookedTransformer
from transformers import AutoModelForCausalLM
from datasets import load_dataset

torch.set_grad_enabled(False)

# %%
# Load models
dtype = torch.float32
model_tl = HookedTransformer.from_pretrained_no_processing("meta-llama/Llama-3.2-1B", dtype=dtype)
model_hf = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", torch_dtype=dtype)

# %%
# Load dataset
dataset = load_dataset("EleutherAI/fineweb-edu-dedup-10b", split="train")

# %%
# Create dataloader helper function
def create_dataloader(dataset, model, batch_size=64, max_length=256):
    """Create a dataloader with the dataset."""
    
    def tokenize_collate_fn(batch):
        texts = [item["text"] for item in batch]
        tokenized = model.tokenizer(
            texts,
            padding="max_length",
            truncation=True,
            max_length=max_length,
            return_tensors="pt",
        )
        return tokenized.to(model.cfg.device, non_blocking=True)
    
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=False,
        collate_fn=tokenize_collate_fn,
    )

# %%
# Create dataloader and get a batch
dataloader = create_dataloader(
    dataset.take(1024), model_tl, batch_size=16
)
batch = next(iter(dataloader))

# %%
# Move models and data to CUDA
model_tl = model_tl.to('cuda').eval()
model_hf = model_hf.to('cuda').eval()
batch = batch.to('cuda')

# %%
# Run both models
with torch.no_grad():
    logits_hf = model_hf(batch.input_ids).logits
    logits_tl = model_tl(batch.input_ids)

# %%
# Compare logits (optional - to verify models produce similar results)
logits_diff = logits_hf - logits_tl
print(f"Logits difference max: {logits_diff.abs().max()}")
print(f"Logits difference mean: {logits_diff.abs().mean()}")

# %%
# Compare top predictions
top_k_hf = logits_hf.argsort(dim=-1, descending=True)[:, :, :5]
top_k_tl = logits_tl.argsort(dim=-1, descending=True)[:, :, :5]

# Check if top predictions match
top_matches = (top_k_hf == top_k_tl).float().mean()
print(f"Top-5 prediction match rate: {top_matches}")

Before:

Logits difference max: 2.952045440673828
Logits difference mean: 0.05180485174059868
Top-5 prediction match rate: 0.899121105670929

After:

Logits difference max: 0.00018334388732910156
Logits difference mean: 8.470005013805348e-06
Top-5 prediction match rate: 1.0

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility