Text Generation
Transformers
Safetensors
English
llama
gpt-u
tiny-lm
pretrained-from-scratch
text-generation-inference
GPT-U-20M / scripts /common.py
DedeProGames's picture
GPT-U-20M: final weights, tokenizer, model card, scripts and logs
13caab9 verified
Raw History Blame Contribute Delete
1.69 kB
"""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