diff --git a/xfuser/model_executor/models/runner_models/base_model.py b/xfuser/model_executor/models/runner_models/base_model.py index de00e8c4..c8570564 100644 --- a/xfuser/model_executor/models/runner_models/base_model.py +++ b/xfuser/model_executor/models/runner_models/base_model.py @@ -841,18 +841,25 @@ def prepare_run(self, input_args: dict) -> None: self._vae_manager.prepare_run(self._decoding_vaes(), input_args) def _run_timed_pipe(self, input_args: dict) -> Tuple[DiffusionOutput, float]: - """ Run a a full pipeline with timing information """ + """ Run the pipeline and time its latency from the synchronized across all ranks beginning + of the model execution, till the moment the current rank finishes. + + Later, we typically discard timings of all ranks except the last one, which is assumed to + be the rank providing model's output. + """ self.prepare_run(input_args) start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + get_world_group().barrier() # aligns all ranks as closely as possible start.record() out = self._run_pipe(input_args) end.record() + end.synchronize() # we don't care about other streams if there are any - torch.cuda.synchronize() elapsed_time = start.elapsed_time(end) / 1000 # Convert to seconds return out, elapsed_time