nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

update median_step_time in model_train.py (#15649)

Browse Source 
BENCHMARK=5 used to pick the 4th largest, not the middle one
This commit is contained in:
chenyu
2026-04-08 09:53:59 -04:00
committed by GitHub
parent dae9dea903
commit 839d37b7bc
1 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/mlperf/model_train.py
+5 -5
View File
@@ -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")