geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

late loss.to("CPU") in llama (#17476)

Browse Source 
* late loss.to("CPU") in llama

* acc = 0
This commit is contained in:
qazal
2026-08-10 17:31:50 +09:00
committed by GitHub
parent 44f1f45cd5
commit 2821bd646f
1 changed files with 11 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/mlperf/model_train.py
+11 -10
View File
@@ -1458,7 +1458,8 @@ def train_llama3():
# realize everything here
if optim.master_params: Tensor.realize(*optim.master_params)
Tensor.realize(*optim.params, *fp8_inv_scales, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
loss_acc = Tensor.zeros(1, dtype=dtypes.float32, device=device)
Tensor.realize(loss_acc, *optim.params, *fp8_inv_scales, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
@TinyJit
def minibatch(tokens:Tensor):
@@ -1476,8 +1477,8 @@ def train_llama3():
for g, new_g in zip(grads, loss.gradient(*optim.params)):
apply_grad(g, new_g.uop)
loss_cpu = loss.flatten().float().to("CPU")
return loss_cpu.realize(*grads, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
loss_acc.assign(loss_acc + loss.flatten().float())
return loss_acc.realize(*grads, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
@TinyJit
def optim_step():
@@ -1490,9 +1491,10 @@ def train_llama3():
lr_cpu = optim.lr.float().to("CPU")
grad_norm_cpu = grad_norm.float().to("CPU")
Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax)
loss_cpu = loss_acc.to("CPU")
Tensor.realize(lr_cpu, grad_norm_cpu, loss_cpu, loss_acc.assign(0), *grads, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax)
return lr_cpu, grad_norm_cpu
return lr_cpu, grad_norm_cpu, loss_cpu
@TinyJit
@Context(TRAINING=0)
@@ -1547,8 +1549,8 @@ def train_llama3():
st = time.perf_counter()
stopped = False
losses, data_time, dev_time = [], 0, 0
for _ in range(grad_acc if i >= 2 else 1):
data_time, dev_time = 0, 0
for _ in range(accum_steps:=grad_acc if i >= 2 else 1):
ist = time.perf_counter()
try: tokens = next(train_iter)
except StopIteration:
@@ -1556,16 +1558,15 @@ def train_llama3():
break
mst = time.perf_counter()
data_time += mst - ist
losses.append(minibatch(tokens).item())
minibatch(tokens)
dev_time += time.perf_counter() - mst
if stopped: break
gt = time.perf_counter()
ret = optim_step()
lr, grad_norm = ret[0].item(), ret[1].item()
lr, grad_norm, loss = ret[0].item(), ret[1].item(), ret[2].item() / accum_steps
et = time.perf_counter()
loss = sum(losses) / len(losses)
optim_time = et - gt
dev_time += optim_time
step_time = et - st