github-actions[bot]
de77c89cf6
version: sunnypilot v2026.002.000 (staging) date: 2026-07-21T07:44:02 master commit: 2c334ede443d7391d27575af8c854a095ba702a8
36 lines
1.1 KiB
Python
36 lines
1.1 KiB
Python
from tinygrad import nn, Tensor, dtypes
|
|
from tinygrad.helpers import DEV, Timing
|
|
|
|
from extra.models.llama import Transformer
|
|
from examples.llama3 import MODEL_PARAMS
|
|
|
|
if __name__ == "__main__":
|
|
DEV.value = "NULL"
|
|
Tensor.training = True
|
|
#model_size = "8B"
|
|
model_size = "405B"
|
|
|
|
with Timing("total "):
|
|
with Timing("***** create model in "):
|
|
model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=nn.Linear, embedding=nn.Embedding,
|
|
max_context=1024, jit=True, disable_kv_cache=True)
|
|
|
|
with Timing("***** fake state in "):
|
|
Tensor.realize(*[p.assign(Tensor.empty(*p.shape, device=p.device, dtype=p.dtype)) for p in nn.state.get_parameters(model)])
|
|
|
|
with Timing("***** create optim in "):
|
|
opt = nn.optim.AdamW(nn.state.get_parameters(model))
|
|
|
|
with Timing("***** run model in "):
|
|
toks = Tensor.empty(1, 1024, dtype=dtypes.int)
|
|
out = model(toks, 0, temperature=float('nan'))
|
|
|
|
with Timing("***** backward in "):
|
|
out.mean().backward()
|
|
|
|
with Timing("***** realize in "):
|
|
out.realize()
|
|
|
|
with Timing("***** step in "):
|
|
opt.step()
|