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

llama: accurate mxfp4 mfu (#17388)

Browse Source 
* llama: accurate mxfp4 mfu

* train_llama3 import
This commit is contained in:
qazal
2026-08-04 18:41:57 +09:00
committed by GitHub
parent c9cd44b751
commit f993228b73
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/mlperf/model_train.py
+2 -2
View File
@@ -1282,7 +1282,7 @@ def train_bert():
previous_step = i
def train_llama3():
from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8_DTYPE, MXFP8
from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8_DTYPE, MXFP8, MXFP4
from examples.llama3 import MODEL_PARAMS
from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
from examples.mlperf.optim import GradAccClipAdamW, clip_grads
@@ -1577,7 +1577,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 * device_count * 4.6e15)) * 100
mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * (9.2e15 if MXFP4 else 4.6e15))) * 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")