tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 12:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

llama: correct 8b init (#15397)

Browse Source
This commit is contained in:
wozeparrot
2026-03-24 13:41:41 -07:00
committed by GitHub
parent 652bab8aad
commit da2031266a
4 changed files with 63 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/mlperf/model_train.py
+8 -10
View File
@@ -13,8 +13,6 @@ from extra.bench_log import BenchEvent, WallTimeEvent
# TODO: fix benchmark logging and use tinygrad tqdm
from tqdm import tqdm
from tinygrad.uop.ops import UOp
def train_resnet():
from extra.models import resnet
from examples.mlperf.dataloader import batch_load_resnet
@@ -1284,7 +1282,7 @@ def train_bert():
previous_step = i
def train_llama3():
from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad
from examples.mlperf.models.flat_llama import FlatTransformer
from examples.llama3 import MODEL_PARAMS
from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
from examples.mlperf.optim import GradAccClipAdamW
@@ -1348,8 +1346,6 @@ def train_llama3():
model = FlatTransformer(**model_params, max_context=SEQLEN)
params = get_parameters(model)
# weights are all bfloat16 for now
assert params and all(p.dtype == dtypes.bfloat16 for p in params)
if getenv("FAKEDATA"):
for v in get_parameters(model):
@@ -1373,7 +1369,9 @@ def train_llama3():
eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
# init grads
grads = [Tensor.zeros_like(p).contiguous() for p in optim.params]
for p in optim.params:
p.grad = Tensor.zeros_like(p).contiguous()
grads = [p.grad for p in optim.params]
scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
@@ -1394,8 +1392,8 @@ def train_llama3():
logits:Tensor = model(tokens[:, :-1])
loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
for i,(t,g) in enumerate(zip(optim.params, loss.gradient(*optim.params))):
grads[i].replace(Tensor(grads[i].uop.after(UOp.group(*apply_grad(grads[i].uop, g.uop))), device=t.device))
loss.backward()
assert all(p.grad is g for p,g in zip(optim.params, grads))
loss_cpu = loss.flatten().float().to("CPU")
return loss_cpu.realize(*grads)
@@ -1491,7 +1489,7 @@ def train_llama3():
mem_gb = GlobalCounters.mem_used / 1e9
gflops = GlobalCounters.global_ops / 1e9 / dev_time
mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * max(getenv("DP", 1), getenv("MP", 1)) * 2.3e15)) * 100
mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 2.3e15)) * 100
tqdm.write(
f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
@@ -1546,7 +1544,7 @@ def train_llama3():
if BENCHMARK and (j+1) == min(BENCHMARK, EVAL_SAMPLES//EVAL_BS):
return
log_perplexity = Tensor(eval_losses).mean().float().item()
log_perplexity = sum(eval_losses) / len(eval_losses)
tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")
examples/mlperf/models/flat_llama.py
+44 -23
View File
@@ -18,6 +18,7 @@ from tinygrad.uop.ops import Ops, UOp
from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
FP8 = getenv("FP8", 0)
WQKV = getenv("WQKV", 0)
FP8_DTYPE = dtypes.fp8e4m3
FP8_MAX = 448.0
@@ -49,13 +50,20 @@ class FlatTransformer:
self.head_dim = dim // n_heads
self.n_rep = self.n_heads // self.n_kv_heads
scaled_std = 0.02 / math.sqrt(2 * n_layers)
# Attention
self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
self.wo = self.lin_per_layer(self.n_heads * self.head_dim, dim)
if WQKV:
self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
else:
self.wq = self.lin_per_layer(dim, self.n_heads * self.head_dim)
self.wk = self.lin_per_layer(dim, self.n_kv_heads * self.head_dim)
self.wv = self.lin_per_layer(dim, self.n_kv_heads * self.head_dim)
self.wo = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)
# FeedForward
self.w1 = self.lin_per_layer(dim, hidden_dim)
self.w2 = self.lin_per_layer(hidden_dim, dim)
self.w2 = self.lin_per_layer(hidden_dim, dim, std=scaled_std)
self.w3 = self.lin_per_layer(dim, hidden_dim)
self.norm_eps = norm_eps
@@ -65,25 +73,31 @@ class FlatTransformer:
# output
self.norm = nn.RMSNorm(dim, norm_eps)
self.tok_embeddings = nn.Embedding(vocab_size, dim)
self.output = nn.Linear(dim, vocab_size, bias=False)
self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=0.02)
self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02)
self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
def lin_per_layer(self, in_features:int, out_features:int):
bound = 1 / math.sqrt(in_features)
def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02):
dt = FP8_DTYPE if FP8 else None
if getenv("ZEROS"): return Tensor.zeros(self.n_layers, out_features, in_features, dtype=dt)
return Tensor.uniform(self.n_layers, out_features, in_features, low=-bound, high=bound, dtype=dt)
return Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std, dtype=dt)
def attention(self, x:Tensor, freqs_cis:Tensor, attention_norm:Tensor, wqkv:Tensor, wo:Tensor):
def attention(self, x:Tensor, freqs_cis:Tensor, attention_norm:Tensor, wo:Tensor, wqkv:Tensor|None=None,
wq:Tensor|None=None, wk:Tensor|None=None, wv:Tensor|None=None):
x = rmsnorm(x, self.norm_eps) * attention_norm
xqkv = matmul(x, wqkv)
bsz, seqlen, _ = x.shape
bsz, seqlen, _ = xqkv.shape
# interleaved layout: each kv group has [n_rep q heads, 1 k head, 1 v head] for clean MP sharding
xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
if wqkv is not None:
xqkv = matmul(x, wqkv)
xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
else:
assert wq is not None and wk is not None and wv is not None
xq = matmul(x, wq).reshape(bsz, seqlen, self.n_heads, self.head_dim)
xk = matmul(x, wk).reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
xv = matmul(x, wv).reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
@@ -99,9 +113,10 @@ class FlatTransformer:
@function(precompile=True, precompile_backward=True)
def run_layer(self, x:Tensor, freqs_cis:Tensor,
attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor):
h = x + self.attention(x, freqs_cis, attention_norm, wqkv, wo)
attention_norm:Tensor, wo:Tensor,
ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor,
wqkv:Tensor|None=None, wq:Tensor|None=None, wk:Tensor|None=None, wv:Tensor|None=None):
h = x + self.attention(x, freqs_cis, attention_norm, wo, wqkv=wqkv, wq=wq, wk=wk, wv=wv)
return h + self.feed_forward(h, ffn_norm, w1, w2, w3)
def shard(self, device:tuple[str, ...], mp:bool=False):
@@ -110,7 +125,12 @@ class FlatTransformer:
for v in get_parameters(self): v.shard_(device, axis=None)
else:
# flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer
self.wqkv.shard_(device, axis=1).realize() # (n_layers, out, dim) shard out
if WQKV:
self.wqkv.shard_(device, axis=1).realize() # (n_layers, out, dim) shard out
else:
self.wq.shard_(device, axis=1).realize() # (n_layers, n_heads*head_dim, dim) shard out
self.wk.shard_(device, axis=1).realize() # (n_layers, n_kv_heads*head_dim, dim) shard out
self.wv.shard_(device, axis=1).realize() # (n_layers, n_kv_heads*head_dim, dim) shard out
self.wo.shard_(device, axis=2).realize() # (n_layers, dim, in) shard in
self.w1.shard_(device, axis=1).realize() # (n_layers, hidden, dim) shard out
self.w2.shard_(device, axis=2).realize() # (n_layers, dim, hidden) shard in
@@ -119,17 +139,18 @@ class FlatTransformer:
self.ffn_norm.shard_(device, axis=None).realize()
self.norm.weight.shard_(device, axis=None).realize()
self.tok_embeddings.weight.shard_(device, axis=0).realize()
self.output.weight.shard_(device, axis=0).realize()
self.output.shard_(device, axis=1).realize()
self.freqs_cis.shard_(device, axis=None).realize()
def __call__(self, tokens:Tensor):
h = self.tok_embeddings(tokens)
freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
for i in range(self.n_layers):
attn_kwargs = {"wqkv": self.wqkv[i]} if WQKV else {"wq": self.wq[i], "wk": self.wk[i], "wv": self.wv[i]}
h = self.run_layer(h, freqs_cis,
self.attention_norm[i], self.wqkv[i], self.wo[i],
self.ffn_norm[i], self.w1[i], self.w2[i], self.w3[i])
logits = self.output(self.norm(h))
self.attention_norm[i], self.wo[i],
self.ffn_norm[i], self.w1[i], self.w2[i], self.w3[i], **attn_kwargs)
logits = self.norm(h) @ self.output[0].T
return logits
# TODO: this shouldn't be needed, but it prevents a copy of the grads. CAT can help
examples/mlperf/optim.py
+10 -5
View File
@@ -11,6 +11,8 @@ class GradAccClipAdamW(Optimizer):
self.m = self._new_optim_param()
self.v = self._new_optim_param()
self.grad_acc, self.clip_norm = grad_acc, clip_norm
# fp32 master weights for mixed precision training
self.master_params:list[Tensor]|None = [p.float().contiguous() for p in self.params] if self.params[0].dtype != dtypes.float32 else None
def fstep(self, grads:list[Tensor]):
if self.fused:
@@ -18,8 +20,8 @@ class GradAccClipAdamW(Optimizer):
updates = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
else:
updates, extra = self._step([], grads)
for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i]))
to_realize = extra+self.params+self.buffers
for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i], self.master_params[i] if self.master_params else None))
to_realize = extra+self.params+self.buffers+(self.master_params or [])
Tensor.realize(*to_realize)
return extra[-1]
@@ -53,7 +55,10 @@ class GradAccClipAdamW(Optimizer):
ret.append((self.lr * up).cast(g.dtype))
return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
def _apply_update(self, t:Tensor, up:Tensor) -> Tensor:
def _apply_update(self, t:Tensor, up:Tensor, master:Tensor|None=None) -> Tensor:
w = master if master is not None else t
wd = self.wd if t.ndim >= 3 else 0.0
up = up.shard_like(t) + self.lr.to(t.device) * wd * t.detach()
return t.detach() - up.cast(t.dtype)
up = up.float().shard_like(w) + self.lr.to(w.device) * wd * w.detach()
new_w = w.detach() - up
if master is not None: master.assign(new_w)
return new_w.cast(t.dtype)
examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
+1 -1
View File
@@ -14,7 +14,7 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
export ASM_GEMM=${ASM_GEMM:-1}
export WQKV=${WQKV:-0}
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32"
export DP=${DP:-8} MP=${MP:-1} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-4}
export GBS=$((BS * GRADIENT_ACC_STEPS))