Skip to content

Commit 4544b7d

Browse files
cli99cli99
authored
Improve flops profiler functionality (deepspeedai#1065)
* use the original function's name as the key to old_functions dict * update profile output format * print at global rank 0 * add flops calculation in bwd pass using time from ds timers * improve aggregated profiling out to show all depth * print samples/second * update readme and examples * update docs * fix typo and reorder printing * fix format
1 parent 1f82ab7 commit 4544b7d

11 files changed

Lines changed: 838 additions & 745 deletions

File tree

‎DeepSpeedExamples‎

‎deepspeed/profiling/config.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,3 +43,7 @@ def _initialize(self, flops_profiler_dict):
4343
self.detailed = get_scalar_param(flops_profiler_dict,
4444
FLOPS_PROFILER_DETAILED,
4545
FLOPS_PROFILER_DETAILED_DEFAULT)
46+
47+
self.output_file = get_scalar_param(flops_profiler_dict,
48+
FLOPS_PROFILER_OUTPUT_FILE,
49+
FLOPS_PROFILER_OUTPUT_FILE_DEFAULT)

‎deepspeed/profiling/constants.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
"module_depth": -1,
1818
"top_modules": 3,
1919
"detailed": true,
20+
"output_file": null
2021
}
2122
}
2223
'''
@@ -33,7 +34,10 @@
3334
FLOPS_PROFILER_MODULE_DEPTH_DEFAULT = -1
3435

3536
FLOPS_PROFILER_TOP_MODULES = "top_modules"
36-
FLOPS_PROFILER_TOP_MODULES_DEFAULT = 3
37+
FLOPS_PROFILER_TOP_MODULES_DEFAULT = 1
3738

3839
FLOPS_PROFILER_DETAILED = "detailed"
3940
FLOPS_PROFILER_DETAILED_DEFAULT = True
41+
42+
FLOPS_PROFILER_OUTPUT_FILE = "output_file"
43+
FLOPS_PROFILER_OUTPUT_FILE_DEFAULT = None

‎deepspeed/profiling/flops_profiler/README.md‎

Lines changed: 243 additions & 269 deletions
Large diffs are not rendered by default.

‎deepspeed/profiling/flops_profiler/profiler.py‎

Lines changed: 209 additions & 99 deletions
Large diffs are not rendered by default.

‎deepspeed/runtime/activation_checkpointing/checkpointing.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -405,7 +405,7 @@ def save_args_for_backward(*all_args):
405405
f"----contiguous Memory Checkpointing {CONTIGUOUS_CHECKPOINTING} with {num_layers} total layers"
406406
)
407407
logger.info(f"----Synchronization {SYNCHRONIZE}")
408-
logger.info(f"----Profiling {PROFILE_TIME}")
408+
logger.info(f"----Profiling time in checkpointing {PROFILE_TIME}")
409409

410410
cuda_device = torch.cuda.current_device()
411411
transport_stream = torch.cuda.Stream(device=cuda_device)

‎deepspeed/runtime/config.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -646,8 +646,9 @@ def _initialize_params(self, param_dict):
646646
self.scheduler_name = get_scheduler_name(param_dict)
647647
self.scheduler_params = get_scheduler_params(param_dict)
648648

649-
self.wall_clock_breakdown = get_wall_clock_breakdown(param_dict)
650649
self.flops_profiler_config = DeepSpeedFlopsProfilerConfig(param_dict)
650+
self.wall_clock_breakdown = get_wall_clock_breakdown(
651+
param_dict) | self.flops_profiler_config.enabled
651652
self.memory_breakdown = get_memory_breakdown(param_dict)
652653
self.tensorboard_enabled = get_tensorboard_enabled(param_dict)
653654
self.tensorboard_output_path = get_tensorboard_output_path(param_dict)

‎deepspeed/runtime/engine.py‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,13 @@ def __init__(self,
160160
steps_per_output=self.steps_per_print(),
161161
monitor_memory=False)
162162

163+
if dist.get_rank() == 0:
164+
logger.info(
165+
f"DeepSpeed Flops Profiler Enabled: {self.flops_profiler_enabled()}")
166+
167+
if self.flops_profiler_enabled():
168+
self.flops_profiler = FlopsProfiler(self.module, self)
169+
163170
if training_data:
164171
self.training_dataloader = self.deepspeed_io(training_data)
165172
else:
@@ -287,6 +294,9 @@ def flops_profiler_top_modules(self):
287294
def flops_profiler_detailed(self):
288295
return self._config.flops_profiler_config.detailed
289296

297+
def flops_profiler_output_file(self):
298+
return self._config.flops_profiler_config.output_file
299+
290300
def memory_breakdown(self):
291301
return self._config.memory_breakdown
292302

@@ -924,7 +934,6 @@ def forward(self, *inputs, **kwargs):
924934
if self.flops_profiler_enabled(
925935
) and self.global_steps == self.flops_profiler_profile_step(
926936
) and self.global_rank == 0:
927-
self.flops_profiler = FlopsProfiler(self.module)
928937
self.flops_profiler.start_profile(ignore_list=None)
929938

930939
if self.module.training and self.progressive_layer_drop:
@@ -943,6 +952,7 @@ def forward(self, *inputs, **kwargs):
943952

944953
if self.training_dataloader is None:
945954
self.tput_timer.start()
955+
946956
loss = self.module(*inputs, **kwargs)
947957

948958
if self.zero_optimization_partition_weights():
@@ -961,11 +971,13 @@ def forward(self, *inputs, **kwargs):
961971
if self.flops_profiler_enabled(
962972
) and self.global_steps == self.flops_profiler_profile_step(
963973
) and self.global_rank == 0:
974+
self.flops_profiler.stop_profile()
964975
self.flops_profiler.print_model_profile(
965976
profile_step=self.global_steps,
966977
module_depth=self.flops_profiler_module_depth(),
967978
top_modules=self.flops_profiler_top_modules(),
968-
detailed=self.flops_profiler_detailed())
979+
detailed=self.flops_profiler_detailed(),
980+
output_file=self.flops_profiler_output_file())
969981
self.flops_profiler.end_profile()
970982

971983
return loss
@@ -1177,7 +1189,9 @@ def step(self, lr_kwargs=None):
11771189
'backward_allreduce_microstep',
11781190
'step_microstep'
11791191
]
1180-
self.timers.log(names=timer_names, memory_breakdown=self.memory_breakdown())
1192+
self.timers.log(names=timer_names,
1193+
reset=False,
1194+
memory_breakdown=self.memory_breakdown())
11811195

11821196
# Log timing
11831197
if self.is_gradient_accumulation_boundary():
@@ -1212,7 +1226,8 @@ def step(self, lr_kwargs=None):
12121226
'backward_inner',
12131227
'backward_allreduce',
12141228
'step'
1215-
])
1229+
],
1230+
reset=False)
12161231

12171232
self.micro_steps += 1
12181233

‎deepspeed/utils/timer.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,11 +33,14 @@ def start(self):
3333
self.start_time = time.time()
3434
self.started_ = True
3535

36-
def stop(self):
36+
def stop(self, reset=True):
3737
"""Stop the timer."""
3838
assert self.started_, 'timer is not started'
3939
torch.cuda.synchronize()
40-
self.elapsed_ += (time.time() - self.start_time)
40+
if reset:
41+
self.elapsed_ = (time.time() - self.start_time)
42+
else:
43+
self.elapsed_ += (time.time() - self.start_time)
4144
self.started_ = False
4245

4346
def reset(self):

0 commit comments

Comments
 (0)