From 839d37b7bc33154379db38097cb8443d2aaa6665 Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 8 Apr 2026 09:53:59 -0400 Subject: [PATCH] update median_step_time in model_train.py (#15649) BENCHMARK=5 used to pick the 4th largest, not the middle one --- examples/mlperf/model_train.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index 8ac1dbd4fa..29c0f73b1e 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -246,7 +246,7 @@ def train_resnet(): if i == BENCHMARK: assert not math.isnan(loss) - median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2] # in seconds + median_step_time = sorted(step_times)[BENCHMARK // 2] # in seconds estimated_total_minutes = int(median_step_time * steps_in_train_epoch * epochs / 60) print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m") print(f"epoch global_ops: {steps_in_train_epoch * GlobalCounters.global_ops:_}, " @@ -593,7 +593,7 @@ def train_retinanet(): if i == BENCHMARK: assert not math.isnan(loss) - median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2] # in seconds + median_step_time = sorted(step_times)[BENCHMARK // 2] # in seconds estimated_total_minutes = int(median_step_time * steps_in_train_epoch * EPOCHS / 60) print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m") print(f"epoch global_ops: {steps_in_train_epoch * GlobalCounters.global_ops:_}, " @@ -868,7 +868,7 @@ def train_unet3d(): i += 1 if i == BENCHMARK: - median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2] # in seconds + median_step_time = sorted(step_times)[BENCHMARK // 2] # in seconds estimated_total_minutes = int(median_step_time * SAMPLES_PER_EPOCH * NUM_EPOCHS / 60) print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m") if (TRAIN_BEAM or EVAL_BEAM) and epoch == start_epoch: break @@ -1167,7 +1167,7 @@ def train_bert(): i += 1 if i == BENCHMARK: - median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2] # in seconds + median_step_time = sorted(step_times)[BENCHMARK // 2] # in seconds estimated_total_minutes = int(median_step_time * train_steps / 60) print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m") print(f"epoch global_ops: {train_steps * GlobalCounters.global_ops:_}, " @@ -1577,7 +1577,7 @@ def train_llama3(): safe_save(get_state_dict(scheduler), fn) if i == BENCHMARK: - median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2] + median_step_time = sorted(step_times)[BENCHMARK // 2] estimated_steps = 200_000 // GBS if getenv("LLAMA3_SIZE", "8B") == "8B" else MAX_STEPS estimated_total_minutes = int(median_step_time * estimated_steps / 60) print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")