File size: 1,691 Bytes
13caab9 | 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 | """Shared helpers for GPT-U-20M: repo paths, model config and model construction."""
import json
import math
from pathlib import Path
import torch
from transformers import LlamaConfig, LlamaForCausalLM
ROOT = Path(__file__).resolve().parents[1]
CONFIG_PATH = ROOT / "configs" / "model.json"
TOKENIZER_DIR = ROOT / "tokenizer"
DATA_DIR = ROOT / "data" / "tokens"
OUT_DIR = ROOT / "out"
CKPT_DIR = OUT_DIR / "ckpt"
LOG_DIR = ROOT / "logs"
EXPECTED_PARAMS = 20_453_760
INIT_STD = 0.02
def load_config() -> LlamaConfig:
return LlamaConfig.from_dict(json.loads(CONFIG_PATH.read_text()), attn_implementation="sdpa")
def init_weights(model: LlamaForCausalLM) -> None:
"""N(0, 0.02) everywhere; residual output projections (o_proj, down_proj) scaled by 1/sqrt(2 * n_layers)."""
resid_std = INIT_STD / math.sqrt(2 * model.config.num_hidden_layers)
with torch.no_grad():
for name, module in model.named_modules():
if name == "lm_head" and model.config.tie_word_embeddings:
continue # shares its weight with embed_tokens
if isinstance(module, torch.nn.Linear):
module.weight.normal_(0.0, resid_std if name.endswith(("o_proj", "down_proj")) else INIT_STD)
elif isinstance(module, torch.nn.Embedding):
module.weight.normal_(0.0, INIT_STD)
elif type(module).__name__.endswith("RMSNorm"):
module.weight.fill_(1.0)
def build_model(init: bool = True) -> LlamaForCausalLM:
# Master weights stay in fp32; bf16 only enters through autocast.
model = LlamaForCausalLM(load_config()).to(torch.float32)
if init:
init_weights(model)
return model
|