186 lines
7.1 KiB
Python
186 lines
7.1 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import math
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.distributed.checkpoint as dcp
|
|
import torch.nn.functional as F
|
|
from tokenizers import Tokenizer
|
|
|
|
|
|
CKPT = Path(
|
|
"runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/"
|
|
"checkpoints/iter_0107500"
|
|
)
|
|
TOKENIZER = Path("tokenizer/glm5.2/tokenizer.json")
|
|
|
|
HIDDEN = 1536
|
|
HEADS = 24
|
|
KV_HEADS = 4
|
|
HEAD_DIM = 64
|
|
DENSE_FFN = 4608
|
|
MOE_FFN = 6144
|
|
MOE_LAYERS = {2, 4, 6, 8, 10}
|
|
TOP_K = 4
|
|
N_LAYERS = 12
|
|
ROPE_BASE = 10000.0
|
|
|
|
|
|
def init_dist() -> None:
|
|
if not dist.is_initialized():
|
|
dist.init_process_group("gloo", init_method="tcp://127.0.0.1:29641", rank=0, world_size=1)
|
|
|
|
|
|
def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
|
|
y = x.float() * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + eps)
|
|
return (y * weight.float()).to(torch.bfloat16)
|
|
|
|
|
|
def rotate_half(x: torch.Tensor) -> torch.Tensor:
|
|
x1 = x[..., ::2]
|
|
x2 = x[..., 1::2]
|
|
out = torch.stack((-x2, x1), dim=-1)
|
|
return out.flatten(-2)
|
|
|
|
|
|
def apply_rope(q: torch.Tensor, k: torch.Tensor, seq_len: int) -> tuple[torch.Tensor, torch.Tensor]:
|
|
inv_freq = 1.0 / (ROPE_BASE ** (torch.arange(0, HEAD_DIM, 2, dtype=torch.float32) / HEAD_DIM))
|
|
pos = torch.arange(seq_len, dtype=torch.float32)
|
|
freqs = torch.einsum("t,d->td", pos, inv_freq)
|
|
emb = torch.repeat_interleave(freqs, 2, dim=-1)
|
|
cos = emb.cos()[None, :, None, :]
|
|
sin = emb.sin()[None, :, None, :]
|
|
qf = q.float()
|
|
kf = k.float()
|
|
return (qf * cos + rotate_half(qf) * sin).to(torch.bfloat16), (
|
|
kf * cos + rotate_half(kf) * sin
|
|
).to(torch.bfloat16)
|
|
|
|
|
|
def swiglu(x: torch.Tensor) -> torch.Tensor:
|
|
a, b = x.chunk(2, dim=-1)
|
|
return F.silu(a.float()) * b.float()
|
|
|
|
|
|
def load_model_tensors(ckpt: Path) -> dict[str, torch.Tensor]:
|
|
init_dist()
|
|
metadata = dcp.FileSystemReader(str(ckpt)).read_metadata()
|
|
state: dict[str, torch.Tensor] = {}
|
|
for key, meta in metadata.state_dict_metadata.items():
|
|
if key.startswith("optimizer.") or not hasattr(meta, "size"):
|
|
continue
|
|
state[key] = torch.empty(tuple(meta.size), dtype=meta.properties.dtype)
|
|
print(f"loading {len(state)} model tensors from {ckpt}", flush=True)
|
|
dcp.load(state, checkpoint_id=str(ckpt))
|
|
print("loaded model tensors", flush=True)
|
|
return state
|
|
|
|
|
|
def dense_mlp(x: torch.Tensor, state: dict[str, torch.Tensor], layer: int) -> torch.Tensor:
|
|
prefix = f"decoder.layers.{layer}.mlp"
|
|
h = rms_norm(x, state[f"{prefix}.linear_fc1.layer_norm_weight"])
|
|
fc1 = F.linear(h.float(), state[f"{prefix}.linear_fc1.weight"].float())
|
|
hidden = swiglu(fc1)
|
|
return F.linear(hidden, state[f"{prefix}.linear_fc2.weight"].float()).to(torch.bfloat16)
|
|
|
|
|
|
def moe_mlp(x: torch.Tensor, state: dict[str, torch.Tensor], layer: int) -> torch.Tensor:
|
|
prefix = f"decoder.layers.{layer}.mlp"
|
|
h = rms_norm(x, state[f"decoder.layers.{layer}.pre_mlp_layernorm.weight"])
|
|
router = F.linear(h.float(), state[f"{prefix}.router.weight"].float())
|
|
topv, topi = torch.topk(router, TOP_K, dim=-1)
|
|
probs = torch.softmax(topv, dim=-1)
|
|
|
|
fc1_w = state[f"{prefix}.experts.experts.linear_fc1.weight"].float()
|
|
fc2_w = state[f"{prefix}.experts.experts.linear_fc2.weight"].float()
|
|
out = torch.zeros_like(h.float())
|
|
flat_h = h.float().reshape(-1, HIDDEN)
|
|
flat_topi = topi.reshape(-1, TOP_K)
|
|
flat_probs = probs.reshape(-1, TOP_K)
|
|
flat_out = out.reshape(-1, HIDDEN)
|
|
for expert_id in range(fc1_w.shape[0]):
|
|
mask = flat_topi == expert_id
|
|
if not mask.any():
|
|
continue
|
|
token_idx, choice_idx = mask.nonzero(as_tuple=True)
|
|
expert_in = flat_h[token_idx]
|
|
expert_hidden = swiglu(F.linear(expert_in, fc1_w[expert_id]))
|
|
expert_out = F.linear(expert_hidden, fc2_w[expert_id])
|
|
flat_out[token_idx] += expert_out * flat_probs[token_idx, choice_idx].unsqueeze(-1)
|
|
return out.to(torch.bfloat16)
|
|
|
|
|
|
def layer_forward(x: torch.Tensor, state: dict[str, torch.Tensor], layer: int) -> torch.Tensor:
|
|
prefix = f"decoder.layers.{layer}"
|
|
attn_in = rms_norm(x, state[f"{prefix}.self_attention.linear_qkv.layer_norm_weight"])
|
|
qkv = F.linear(attn_in.float(), state[f"{prefix}.self_attention.linear_qkv.weight"].float())
|
|
q, k, v = torch.split(qkv, [HEADS * HEAD_DIM, KV_HEADS * HEAD_DIM, KV_HEADS * HEAD_DIM], dim=-1)
|
|
bsz, seq_len, _ = q.shape
|
|
q = q.view(bsz, seq_len, HEADS, HEAD_DIM)
|
|
k = k.view(bsz, seq_len, KV_HEADS, HEAD_DIM)
|
|
v = v.view(bsz, seq_len, KV_HEADS, HEAD_DIM)
|
|
q, k = apply_rope(q, k, seq_len)
|
|
repeat = HEADS // KV_HEADS
|
|
k = k.repeat_interleave(repeat, dim=2)
|
|
v = v.repeat_interleave(repeat, dim=2)
|
|
q = q.transpose(1, 2).float()
|
|
k = k.transpose(1, 2).float()
|
|
v = v.transpose(1, 2).float()
|
|
scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(HEAD_DIM)
|
|
mask = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool), diagonal=1)
|
|
scores = scores.masked_fill(mask[None, None, :, :], float("-inf"))
|
|
ctx = torch.matmul(torch.softmax(scores, dim=-1), v)
|
|
ctx = ctx.transpose(1, 2).contiguous().view(bsz, seq_len, HIDDEN)
|
|
attn_out = F.linear(ctx, state[f"{prefix}.self_attention.linear_proj.weight"].float()).to(torch.bfloat16)
|
|
x = x + attn_out
|
|
x = x + (moe_mlp(x, state, layer) if layer in MOE_LAYERS else dense_mlp(x, state, layer))
|
|
return x
|
|
|
|
|
|
@torch.inference_mode()
|
|
def logits_for_ids(ids: list[int], state: dict[str, torch.Tensor]) -> torch.Tensor:
|
|
input_ids = torch.tensor(ids, dtype=torch.long).unsqueeze(0)
|
|
x = F.embedding(input_ids, state["embedding.word_embeddings.weight"])
|
|
for layer in range(N_LAYERS):
|
|
x = layer_forward(x, state, layer)
|
|
x = rms_norm(x, state["decoder.final_layernorm.weight"])
|
|
return torch.matmul(x[:, -1, :].float(), state["embedding.word_embeddings.weight"].float().T)
|
|
|
|
|
|
def generate(prompt: str, state: dict[str, torch.Tensor], tokenizer: Tokenizer, max_new_tokens: int) -> str:
|
|
ids = tokenizer.encode(prompt).ids
|
|
print(f"prompt={prompt!r} input_tokens={len(ids)}", flush=True)
|
|
for step in range(max_new_tokens):
|
|
logits = logits_for_ids(ids, state)
|
|
next_id = int(torch.argmax(logits, dim=-1)[0])
|
|
ids.append(next_id)
|
|
piece = tokenizer.decode([next_id], skip_special_tokens=False)
|
|
print(f"step={step} id={next_id} piece={piece!r}", flush=True)
|
|
return tokenizer.decode(ids, skip_special_tokens=False)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--max-new-tokens", type=int, default=8)
|
|
parser.add_argument("--prompt", action="append", default=None)
|
|
args = parser.parse_args()
|
|
|
|
tokenizer = Tokenizer.from_file(str(TOKENIZER))
|
|
prompts = args.prompt or [
|
|
"The capital of France is",
|
|
"中国的首都是",
|
|
]
|
|
state = load_model_tensors(CKPT)
|
|
for prompt in prompts:
|
|
text = generate(prompt, state, tokenizer, args.max_new_tokens)
|
|
print("=== GENERATED ===")
|
|
print(text)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|