diff --git a/thunder/benchmarks/benchmark_litgpt.py b/thunder/benchmarks/benchmark_litgpt.py index f68dd1f0fa..5bf61b4d6c 100644 --- a/thunder/benchmarks/benchmark_litgpt.py +++ b/thunder/benchmarks/benchmark_litgpt.py @@ -286,7 +286,7 @@ def train(self): iter_times = [] if i == self.warmup_iter: # warmup t0 = iter_t0 - iter_times = [] #reset the iter time list + iter_times = [] # reset the iter time list for step_idx in range(self.gradient_accumulation_steps): input_ids, targets = next(self.train_data_iter) @@ -328,9 +328,7 @@ def train(self): t1 = time.perf_counter() if global_rank in [0, None]: iter_time = (t1 - iter_t0) * 1000 - print( - f"iter {i}: loss {loss_item:.4f}, iter time: {iter_time:.2f}ms, t: {input_ids.size(1)}" - ) + print(f"iter {i}: loss {loss_item:.4f}, iter time: {iter_time:.2f}ms, t: {input_ids.size(1)}") iter_times.append(iter_time) # if global_rank in [0, None] and i >=warmup_iter: @@ -348,7 +346,7 @@ def train(self): if global_rank in [0, None]: self.perf_metrics["average_iter_time"] = ((t1 - t0) * 1000) / (self.max_iters - self.warmup_iter) - self.perf_metrics["median_iter_time"] = np.median(iter_times) #To avoid outliers + self.perf_metrics["median_iter_time"] = np.median(iter_times) # To avoid outliers def add_perf_metrics(self): # tokens_per_sec = total number of benchmarked iterations x global BS x block_size / total elapsed time (s)