GaudiTrainer: Take accuracy measurement outside of timer (#2331)

ugolowic · web-flow · commit 00a103d955e7 · 2025-10-31T16:39:18.000+01:00
Signed-off-by: Urszula &lt;urszula.golowicz@intel.com&gt;
diff --git a/optimum/habana/diffusers/pipelines/controlnet/pipeline_controlnet.py b/optimum/habana/diffusers/pipelines/controlnet/pipeline_controlnet.py
@@ -634,11 +634,13 @@ def __call__(
                     self.htcore.mark_step()
 
             hb_profiler.stop()
+            end_time = time.time()
 
             speed_metrics_prefix = "generation"
             speed_measures = speed_metrics(
                 split=speed_metrics_prefix,
                 start_time=t0,
+                end_time=end_time,
                 num_samples=num_batches * batch_size
                 if t1 == t0 or use_warmup_inference_steps
                 else (num_batches - throughput_warmup_steps) * batch_size,
diff --git a/optimum/habana/diffusers/pipelines/controlnet/pipeline_stable_video_diffusion_controlnet.py b/optimum/habana/diffusers/pipelines/controlnet/pipeline_stable_video_diffusion_controlnet.py
@@ -534,10 +534,13 @@ def __call__(
 
                 outputs["frames"].append(frames)
 
+            end_time = time.time()
+
             speed_metrics_prefix = "generation"
             speed_measures = speed_metrics(
                 split=speed_metrics_prefix,
                 start_time=t0,
+                end_time=end_time,
                 num_samples=num_batches * batch_size
                 if t1 == t0
                 else (num_batches - throughput_warmup_steps) * batch_size,
diff --git a/optimum/habana/diffusers/pipelines/ddpm/pipeline_ddpm.py b/optimum/habana/diffusers/pipelines/ddpm/pipeline_ddpm.py
@@ -192,6 +192,8 @@ def __call__(
             if not self.use_hpu_graphs:  # For checking output resutls
                 self.htcore.mark_step()
 
+        end_time = time.time()
+
         if self.gaudi_config.use_torch_autocast:
             image = image.float()
 
@@ -204,6 +206,7 @@ def __call__(
         speed_measures = speed_metrics(
             split=speed_metrics_prefix,
             start_time=start_time,
+            end_time=end_time,
             num_samples=batch_size,
             num_steps=batch_size * len(num_inference_steps),
             start_time_after_warmup=time_after_warmup,
diff --git a/optimum/habana/diffusers/pipelines/flux/pipeline_flux.py b/optimum/habana/diffusers/pipelines/flux/pipeline_flux.py
@@ -583,12 +583,15 @@ def __call__(
             finalize_calibration(self.transformer)
 
         ht.hpu.synchronize()
+        end_time = time.time()
+
         speed_metrics_prefix = "generation"
         if use_warmup_inference_steps:
             t1 = warmup_inference_steps_time_adjustment(t1, t1, num_inference_steps, throughput_warmup_steps)
         speed_measures = speed_metrics(
             split=speed_metrics_prefix,
             start_time=t0,
+            end_time=end_time,
             num_samples=batch_size
             if t1 == t0 or use_warmup_inference_steps
             else (num_batches - throughput_warmup_steps) * batch_size,
diff --git a/optimum/habana/diffusers/pipelines/flux/pipeline_flux_img2img.py b/optimum/habana/diffusers/pipelines/flux/pipeline_flux_img2img.py
@@ -615,12 +615,15 @@ def __call__(
             finalize_calibration(self.transformer)
 
         ht.hpu.synchronize()
+        end_time = time.time()
+
         speed_metrics_prefix = "generation"
         if use_warmup_inference_steps:
             t1 = warmup_inference_steps_time_adjustment(t1, t1, num_inference_steps, throughput_warmup_steps)
         speed_measures = speed_metrics(
             split=speed_metrics_prefix,
             start_time=t0,
+            end_time=end_time,
             num_samples=batch_size
             if t1 == t0 or use_warmup_inference_steps
             else (num_batches - throughput_warmup_steps) * batch_size,
diff --git a/optimum/habana/diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py b/optimum/habana/diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py
@@ -624,11 +624,14 @@ def __call__(
                 if not self.use_hpu_graphs:
                     self.htcore.mark_step()
 
+            end_time = time.time()
             hb_profiler.stop()
+
             speed_metrics_prefix = "generation"
             speed_measures = speed_metrics(
                 split=speed_metrics_prefix,
                 start_time=t0,
+                end_time=end_time,
                 num_samples=num_batches * batch_size
                 if t1 == t0 or use_warmup_inference_steps
                 else (num_batches - throughput_warmup_steps) * batch_size,
diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
@@ -596,6 +596,7 @@ def __call__(
                     self.htcore.mark_step()
 
             hb_profiler.stop()
+            end_time = time.time()
 
             speed_metrics_prefix = "generation"
             if t1 == t0 or use_warmup_inference_steps:
@@ -608,6 +609,7 @@ def __call__(
             speed_measures = speed_metrics(
                 split=speed_metrics_prefix,
                 start_time=t0,
+                end_time=end_time,
                 num_samples=num_samples,
                 num_steps=num_steps,
                 start_time_after_warmup=t1,
diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py
@@ -372,11 +372,14 @@ def __call__(
                 if not self.use_hpu_graphs:
                     self.htcore.mark_step()
 
+            end_time = time.time()
             hb_profiler.stop()
+
             speed_metrics_prefix = "generation"
             speed_measures = speed_metrics(
                 split=speed_metrics_prefix,
                 start_time=t0,
+                end_time=end_time,
                 num_samples=num_batches * batch_size
                 if t1 == t0 or use_warmup_inference_steps
                 else (num_batches - throughput_warmup_steps) * batch_size,
diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py
@@ -607,12 +607,14 @@ def __call__(
 
                 outputs["images"].append(image)
 
+            end_time = time.time()
             hb_profiler.stop()
 
             speed_metrics_prefix = "generation"
             speed_measures = speed_metrics(
                 split=speed_metrics_prefix,
                 start_time=t0,
+                end_time=end_time,
                 num_samples=num_batches * batch_size
                 if t1 == t0 or use_warmup_inference_steps
                 else (num_batches - throughput_warmup_steps) * batch_size,
diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py
@@ -708,6 +708,8 @@ def __call__(
                 if not self.use_hpu_graphs:
                     self.htcore.mark_step()
 
+            end_time = time.time()
+
             # Remove dummy generations if needed
             if num_dummy_samples > 0:
                 outputs["images"][-1] = outputs["images"][-1][:-num_dummy_samples]
@@ -716,6 +718,7 @@ def __call__(
             speed_measures = speed_metrics(
                 split=speed_metrics_prefix,
                 start_time=t0,
+                end_time=end_time,
                 num_samples=num_batches * batch_size
                 if t1 == t0 or use_warmup_inference_steps
                 else (num_batches - throughput_warmup_steps) * batch_size,
diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py
@@ -491,11 +491,14 @@ def __call__(
                 if not self.use_hpu_graphs:
                     self.htcore.mark_step()
 
+            end_time = time.time()
             hb_profiler.stop()
+
             speed_metrics_prefix = "generation"
             speed_measures = speed_metrics(
                 split=speed_metrics_prefix,
                 start_time=t0,
+                end_time=end_time,
                 num_samples=num_batches * batch_size
                 if t1 == t0 or use_warmup_inference_steps
                 else (num_batches - throughput_warmup_steps) * batch_size,
diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_ldm3d.py b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_ldm3d.py
@@ -416,10 +416,13 @@ def __call__(
                 if not self.use_hpu_graphs:
                     self.htcore.mark_step()
 
+            end_time = time.time()
+
             speed_metrics_prefix = "generation"
             speed_measures = speed_metrics(
                 split=speed_metrics_prefix,
                 start_time=t0,
+                end_time=end_time,
                 num_samples=num_batches * batch_size
                 if t1 == t0 or use_warmup_inference_steps
                 else (num_batches - throughput_warmup_steps) * batch_size,
diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py b/optimum/habana/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py
@@ -543,10 +543,13 @@ def __call__(
                 if not self.use_hpu_graphs:
                     self.htcore.mark_step()
 
+            end_time = time.time()
+
             speed_metrics_prefix = "generation"
             speed_measures = speed_metrics(
                 split=speed_metrics_prefix,
                 start_time=t0,
+                end_time=end_time,
                 num_samples=num_batches * batch_size
                 if t1 == t0 or use_warmup_inference_steps
                 else (num_batches - throughput_warmup_steps) * batch_size,
diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py b/optimum/habana/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py
@@ -698,6 +698,7 @@ def __call__(
 
             # End of Denoising loop
 
+            end_time = time.time()
             hb_profiler.stop()
 
             ht.hpu.synchronize()
@@ -707,6 +708,7 @@ def __call__(
             speed_measures = speed_metrics(
                 split=speed_metrics_prefix,
                 start_time=t0,
+                end_time=end_time,
                 num_samples=batch_size
                 if t1 == t0 or use_warmup_inference_steps
                 else (num_batches - throughput_warmup_steps) * batch_size,
diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py b/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py
@@ -817,6 +817,7 @@ def __call__(
                 if not self.use_hpu_graphs:
                     self.htcore.mark_step()
 
+            end_time = time.time()
             hb_profiler.stop()
 
             speed_metrics_prefix = "generation"
@@ -829,6 +830,7 @@ def __call__(
             speed_measures = speed_metrics(
                 split=speed_metrics_prefix,
                 start_time=t0,
+                end_time=end_time,
                 num_samples=num_samples,
                 num_steps=num_steps,
                 start_time_after_warmup=t1,
diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py b/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py
@@ -667,12 +667,14 @@ def denoising_value_valid(dnv):
                 if not self.use_hpu_graphs:
                     self.htcore.mark_step()
 
+            end_time = time.time()
             hb_profiler.stop()
 
             speed_metrics_prefix = "generation"
             speed_measures = speed_metrics(
                 split=speed_metrics_prefix,
                 start_time=t0,
+                end_time=end_time,
                 num_samples=num_batches * batch_size
                 if t1 == t0 or use_warmup_inference_steps
                 else (num_batches - throughput_warmup_steps) * batch_size,
diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py b/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py
@@ -913,6 +913,8 @@ def denoising_value_valid(dnv):
                 if not self.use_hpu_graphs:
                     self.htcore.mark_step()
 
+            end_time = time.time()
+
             # Remove dummy generations if needed
             if num_dummy_samples > 0:
                 outputs["images"][-1] = outputs["images"][-1][:-num_dummy_samples]
@@ -921,6 +923,7 @@ def denoising_value_valid(dnv):
             speed_measures = speed_metrics(
                 split=speed_metrics_prefix,
                 start_time=t0,
+                end_time=end_time,
                 num_samples=num_batches * batch_size
                 if t1 == t0 or use_warmup_inference_steps
                 else (num_batches - throughput_warmup_steps) * batch_size,
diff --git a/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_mlperf.py b/optimum/habana/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_mlperf.py
@@ -933,6 +933,7 @@ def __call__(
                 t_vae_e = time.time()
                 t1 = t1 + t_vae_e - t_vae_b
 
+        end_time = time.time()
         hb_profiler.stop()
 
         speed_metrics_prefix = "generation"
@@ -946,6 +947,7 @@ def __call__(
         speed_measures = speed_metrics(
             split=speed_metrics_prefix,
             start_time=t0,
+            end_time=end_time,
             num_samples=num_samples,
             num_steps=num_steps,
             start_time_after_warmup=t1,
diff --git a/optimum/habana/diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py b/optimum/habana/diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py
@@ -581,11 +581,14 @@ def __call__(
                 if not self.use_hpu_graphs:
                     self.htcore.mark_step()
 
+            end_time = time.time()
             hb_profiler.stop()
+
             speed_metrics_prefix = "generation"
             speed_measures = speed_metrics(
                 split=speed_metrics_prefix,
                 start_time=t0,
+                end_time=end_time,
                 num_samples=num_batches * batch_size
                 if t1 == t0 or use_warmup_inference_steps
                 else (num_batches - throughput_warmup_steps) * batch_size,
diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py
@@ -1166,6 +1166,7 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio
             if self.control.should_training_stop:
                 break
 
+        end_time = time.time()
         hb_profiler.stop()
 
         if args.past_index and hasattr(self, "_past"):
@@ -1191,6 +1192,7 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio
         metrics = speed_metrics(
             "train",
             start_time,
+            end_time,
             num_samples=num_samples_for_speed_metrics,
             num_steps=num_steps_for_speed_metrics,
             num_tokens=num_train_tokens,
@@ -1624,7 +1626,9 @@ def _load_optimizer_and_scheduler(self, checkpoint):
             if self.args.use_habana:
                 to_device_dtype(self.optimizer.state.values(), target_device=torch.device("hpu"))
 
-    def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None:
+    def log(
+        self, logs: dict[str, float], start_time: Optional[float] = None, end_time: Optional[float] = None
+    ) -> None:
         """
         Log `logs` on the various objects watching training.
 
@@ -1640,8 +1644,8 @@ def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> Non
             logs["epoch"] = self.state.epoch
         if self.args.include_num_input_tokens_seen:
             logs["num_input_tokens_seen"] = self.state.num_input_tokens_seen
-            if start_time is not None:
-                logs.update(speed_metrics("train", start_time, num_tokens=self.state.num_input_tokens_seen))
+            if start_time is not None and end_time is not None:
+                logs.update(speed_metrics("train", start_time, end_time, num_tokens=self.state.num_input_tokens_seen))
 
         mem_stats = get_hpu_memory_stats(self.args.device)
         logs.update(mem_stats)
@@ -1930,6 +1934,8 @@ def evaluate(
             start_time += output.metrics[f"{metric_key_prefix}_model_preparation_time"]
         num_samples = output.num_samples - self.args.throughput_warmup_steps * total_batch_size
         num_steps = math.ceil(output.num_samples / total_batch_size) - self.args.throughput_warmup_steps
+        if f"{metric_key_prefix}_end_time" in output.metrics:
+            end_time = output.metrics[f"{metric_key_prefix}_end_time"]
 
         eval_steps = math.ceil(output.num_samples / total_batch_size)
         if eval_steps <= self.args.throughput_warmup_steps:
@@ -1943,6 +1949,7 @@ def evaluate(
             speed_metrics(
                 metric_key_prefix,
                 start_time,
+                end_time,
                 num_samples=num_samples,
                 num_steps=num_steps,
                 start_time_after_warmup=self.start_time_after_warmup,
@@ -1975,6 +1982,7 @@ def predict(
         output = eval_loop(
             test_dataloader, description="Prediction", ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix
         )
+        end_time = time.time()
         total_batch_size = self.args.eval_batch_size * self.args.world_size
         if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
             start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
@@ -1990,6 +1998,7 @@ def predict(
             speed_metrics(
                 metric_key_prefix,
                 start_time,
+                end_time,
                 num_samples=num_samples,
                 num_steps=num_steps,
                 start_time_after_warmup=self.start_time_after_warmup,
@@ -2201,6 +2210,7 @@ def evaluation_loop(
 
             hb_profiler.step()
 
+        end_time = time.time()
         hb_profiler.stop()
 
         # After all calls to `.gather_function`, reset to `gather_for_metrics`:
@@ -2260,6 +2270,7 @@ def evaluation_loop(
             metrics[f"{metric_key_prefix}_model_preparation_time"] = self.model_preparation_time
         if hasattr(self, "compilation_time"):
             metrics[f"{metric_key_prefix}_graph_compliation_duration"] = self.compilation_time
+        metrics[f"{metric_key_prefix}_end_time"] = end_time
 
         # Prefix all keys with metric_key_prefix + '_'
         for key in list(metrics.keys()):
diff --git a/optimum/habana/utils/misc.py b/optimum/habana/utils/misc.py
@@ -66,6 +66,7 @@ def to_device_dtype(my_input: Any, target_device: torch.device = None, target_dt
 def speed_metrics(
     split: str,
     start_time: float,
+    end_time: float,
     num_samples: int = None,
     num_steps: int = None,
     num_tokens: int = None,
@@ -91,7 +92,7 @@ def speed_metrics(
         Dict[str, float]: dictionary with performance metrics.
     """
 
-    runtime = time.time() - start_time
+    runtime = end_time - start_time
     result = {f"{split}_runtime": round(runtime, 4)}
     if runtime == 0:
         return result
