Modelpublic

alunxu/llama-3.1-8b-tuned-lens

sourceHugging Facemitupdated 5mo agoView on Hugging Face

0likes5downloads

Model Card

Llama-3.1-8B Tuned Lens

Linear tuned lens for meta-llama/Llama-3.1-8B. Used in the Relocation Hypothesis project (CS-552 / EMNLP working paper) for depth-axis trajectory analysis under chain-of-thought conditions.

Usage

python

from huggingface_hub import snapshot_download
import json, torch
from tuned_lens import TunedLens
from tuned_lens.nn.lenses import TunedLensConfig
from tuned_lens.nn.unembed import Unembed
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B",
                                              torch_dtype=torch.float16,
                                              device_map="cuda")
lens_dir = snapshot_download("alunxu/llama-3.1-8b-tuned-lens")
with open(f"{lens_dir}/config.json") as f:
    cfg = TunedLensConfig.from_dict(json.load(f))
lens = TunedLens(Unembed(model), cfg)
state = torch.load(f"{lens_dir}/params.pt", map_location="cuda")
lens.layer_translators.load_state_dict(state)
lens = lens.to("cuda").eval()

fp16 inference works for Llama-3.1 (vocab 128k, no overflow).