File size: 10,395 Bytes
a4019dd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 | from __future__ import annotations
import math
import os
import importlib.util
from dataclasses import dataclass
import torch
import torch.nn.functional as F
from huggingface_hub import login
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
from .checkpoints import load_token_checkpoint
from .config import PipelineConfig
def _load_causal_lm(model_name: str, **kwargs):
"""Load a causal LM, falling back to architecture-specific classes when
AutoModelForCausalLM doesn't recognize the model type."""
try:
return AutoModelForCausalLM.from_pretrained(model_name, **kwargs)
except (ValueError, ModuleNotFoundError):
config = AutoConfig.from_pretrained(model_name)
model_type = getattr(config, "model_type", "")
if model_type == "gemma3":
from transformers import Gemma3ForConditionalGeneration
return Gemma3ForConditionalGeneration.from_pretrained(model_name, **kwargs)
if model_type == "gemma4":
from transformers import Gemma4ForConditionalGeneration
return Gemma4ForConditionalGeneration.from_pretrained(model_name, **kwargs)
raise
def resolve_hf_token() -> str | None:
return os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
def maybe_login_hf() -> None:
token = resolve_hf_token()
if token:
login(token=token, add_to_git_credential=False)
def resolve_dtype() -> torch.dtype:
return torch.bfloat16 if torch.cuda.is_available() else torch.float32
def has_accelerate() -> bool:
return importlib.util.find_spec("accelerate") is not None
@dataclass(slots=True)
class ModelBundle:
config: PipelineConfig
tokenizer: AutoTokenizer
model: AutoModelForCausalLM
initial_tokenizer_len: int
def initialize_model_bundle(config: PipelineConfig) -> ModelBundle:
maybe_login_hf()
hf_token = resolve_hf_token()
dtype = resolve_dtype()
tokenizer = AutoTokenizer.from_pretrained(config.model.model_name, use_fast=True, token=hf_token)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
initial_len = len(tokenizer)
tokenizer.add_special_tokens(
{"additional_special_tokens": [config.model.ai_token, config.model.human_token]}
)
load_kwargs = {
"token": hf_token,
"dtype": dtype,
}
if torch.cuda.is_available() and has_accelerate():
load_kwargs["device_map"] = "auto"
model = _load_causal_lm(config.model.model_name, **load_kwargs)
if torch.cuda.is_available() and not has_accelerate():
model = model.to("cuda")
elif not torch.cuda.is_available():
model = model.to("cpu")
try:
model.resize_token_embeddings(len(tokenizer), mean_resizing=False)
except TypeError:
model.resize_token_embeddings(len(tokenizer))
# Mean-fill Gemma 4's auxiliary per-layer token table. Some Transformers
# versions resize it internally, while others leave it at the old size.
for name, module in model.named_modules():
if isinstance(module, torch.nn.Embedding) and module is not model.get_input_embeddings():
if module.weight.shape[0] == initial_len:
mean_row = module.weight.data[:initial_len].mean(
dim=0, dtype=torch.float32
).to(dtype=module.weight.dtype)
new_emb = torch.nn.Embedding(
len(tokenizer), module.weight.shape[1],
device=module.weight.device, dtype=module.weight.dtype,
)
new_emb.weight.data[:initial_len] = module.weight.data
new_emb.weight.data[initial_len:] = mean_row
if not torch.equal(
new_emb.weight.data[initial_len], mean_row
):
raise RuntimeError(
f"Failed to mean-fill resized embedding {name}"
)
parent_name, attr_name = name.rsplit(".", 1)
parent = dict(model.named_modules())[parent_name]
setattr(parent, attr_name, new_emb)
print(
f"Mean-filled secondary embedding {name}: "
f"{initial_len} -> {len(tokenizer)}"
)
elif module.weight.shape[0] == len(tokenizer):
mean_row = module.weight.data[:initial_len].mean(
dim=0, dtype=torch.float32
).to(dtype=module.weight.dtype)
module.weight.data[initial_len:] = mean_row
if not torch.equal(module.weight.data[initial_len], mean_row):
raise RuntimeError(
f"Failed to mean-fill expanded embedding {name}"
)
print(
f"Mean-filled expanded secondary embedding {name}: "
f"rows {initial_len}:{len(tokenizer)}"
)
model.config.use_cache = False
if hasattr(model, "gradient_checkpointing_enable"):
model.gradient_checkpointing_enable()
model.eval()
input_emb = model.get_input_embeddings()
with torch.no_grad():
mean_in = input_emb.weight[:initial_len].mean(dim=0)
for token in (config.model.ai_token, config.model.human_token):
token_id = tokenizer.convert_tokens_to_ids(token)
input_emb.weight[token_id].copy_(mean_in + torch.randn_like(mean_in) * 1e-5)
bundle = ModelBundle(
config=config,
tokenizer=tokenizer,
model=model,
initial_tokenizer_len=initial_len,
)
apply_initial_checkpoints(bundle)
return bundle
def apply_initial_checkpoints(bundle: ModelBundle) -> None:
input_emb = bundle.model.get_input_embeddings()
token_dir = bundle.config.output.model_tokens_dir(bundle.config.model.model_name)
ai_path = bundle.config.init_checkpoints.ai_token_path or token_dir / "ai_token.pt"
human_path = bundle.config.init_checkpoints.human_token_path or token_dir / "human_token.pt"
secondary_embs = [
module
for module in bundle.model.modules()
if (
isinstance(module, torch.nn.Embedding)
and module is not input_emb
and module.weight.shape[0] == len(bundle.tokenizer)
)
]
def install(path, token):
if not path.exists():
return
checkpoint = load_token_checkpoint(path)
token_id = bundle.tokenizer.convert_tokens_to_ids(token)
saved_secondary = checkpoint.secondary_embeddings or []
if len(saved_secondary) != len(secondary_embs):
raise ValueError(
f"{path} has {len(saved_secondary)} secondary rows, but "
f"{bundle.config.model.model_name} exposes "
f"{len(secondary_embs)} secondary token embeddings."
)
input_emb.weight[token_id].copy_(
checkpoint.embedding.to(
input_emb.weight.device, dtype=input_emb.weight.dtype
)
)
for embedding, row in zip(secondary_embs, saved_secondary):
embedding.weight[token_id].copy_(
row.to(embedding.weight.device, dtype=embedding.weight.dtype)
)
with torch.no_grad():
install(ai_path, bundle.config.model.ai_token)
install(human_path, bundle.config.model.human_token)
def build_prompt(bundle: ModelBundle, token: str) -> str:
content = bundle.config.model.prompt_template.format(token=token)
messages = [{"role": "user", "content": content}]
return bundle.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
def encode_response(bundle: ModelBundle, prompt_text: str, response_text: str) -> tuple[torch.Tensor, int]:
full_text = prompt_text + response_text
prompt_ids = bundle.tokenizer(prompt_text, return_tensors="pt", add_special_tokens=False)["input_ids"][0]
full_ids = bundle.tokenizer(
full_text,
return_tensors="pt",
add_special_tokens=False,
truncation=True,
max_length=bundle.config.model.max_length,
)["input_ids"][0]
return full_ids, int(len(prompt_ids))
def compute_average_logprob(bundle: ModelBundle, input_ids: torch.Tensor, prompt_len: int) -> torch.Tensor:
token_logps = compute_token_logprobs(bundle, input_ids, prompt_len)
return token_logps.mean()
def _model_forward(bundle: ModelBundle, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None):
"""Run a forward pass, injecting token_type_ids for models that require it."""
fwd_kwargs: dict = {"input_ids": input_ids}
if attention_mask is not None:
fwd_kwargs["attention_mask"] = attention_mask
model_type = getattr(bundle.model.config, "model_type", "")
if model_type in ("gemma3", "gemma4"):
fwd_kwargs["token_type_ids"] = torch.zeros_like(input_ids)
return bundle.model(**fwd_kwargs)
def compute_token_logprobs(bundle: ModelBundle, input_ids: torch.Tensor, prompt_len: int) -> torch.Tensor:
input_ids = input_ids.unsqueeze(0).to(bundle.model.device)
logits = _model_forward(bundle, input_ids).logits[0]
shift_logits = logits[prompt_len - 1 : -1]
shift_labels = input_ids[0, prompt_len:]
log_probs = F.log_softmax(shift_logits, dim=-1)
return log_probs[torch.arange(len(shift_labels), device=bundle.model.device), shift_labels]
def compute_sequence_logprob(bundle: ModelBundle, input_ids: torch.Tensor, prompt_len: int) -> torch.Tensor:
input_ids = input_ids.unsqueeze(0).to(bundle.model.device)
logits = _model_forward(bundle, input_ids).logits[0]
shift_logits = logits[prompt_len - 1 : -1]
shift_labels = input_ids[0, prompt_len:]
log_probs = F.log_softmax(shift_logits, dim=-1)
token_logps = log_probs[torch.arange(len(shift_labels), device=bundle.model.device), shift_labels]
return token_logps.sum()
def cosine_with_floor(
step: int,
total_steps: int,
base_lr: float,
*,
min_lr: float,
warmup_steps: int,
) -> float:
if step < warmup_steps:
return base_lr * float(step + 1) / float(max(1, warmup_steps))
progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps))
cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
return min_lr + (base_lr - min_lr) * cosine
|