From f993228b73bd86d761e425fa7938eaaeb79c14e2 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Tue, 4 Aug 2026 17:41:57 +0800 Subject: [PATCH] llama: accurate mxfp4 mfu (#17388) * llama: accurate mxfp4 mfu * train_llama3 import --- examples/mlperf/model_train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index fec53400be..e9d8086aab 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -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")