From 1a9fa0f2c985d5fd4c6c44b6b9ef6c4f94627c9f Mon Sep 17 00:00:00 2001 From: Jae-Won Chung Date: Fri, 13 Oct 2023 11:08:16 -0400 Subject: [PATCH] Add cuda sync --- zeus/optimizer/perseus/optimizer.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/zeus/optimizer/perseus/optimizer.py b/zeus/optimizer/perseus/optimizer.py index 6a102105..ce4d87c1 100644 --- a/zeus/optimizer/perseus/optimizer.py +++ b/zeus/optimizer/perseus/optimizer.py @@ -38,6 +38,7 @@ FrequencySchedule, ) from zeus.util.env import resolve_gpu_indices +from zeus.util.framework import cuda_sync class PerseusOptimizer(Callback): @@ -213,6 +214,8 @@ def on_instruction_begin(self, name: str) -> None: expected instruction matches the name of the instruction, and set the frequency accordingly. """ + cuda_sync(self.cuda_device_id) + # Retrieve the next frequency from the schedule. item = next(self.freq_schedule_iter, None) if item is None: @@ -228,3 +231,6 @@ def on_instruction_begin(self, name: str) -> None: ) self.frequency_controller.set_frequency(frequency) + + def on_instruction_end(self, _: str) -> None: + """Mark the end of an instruction, like forward and backward."""