12345678910111213141516171819202122232425262728293031323334353637383940414243444546 |
- # Copyright (c) Microsoft Corporation.
- # SPDX-License-Identifier: Apache-2.0
- # DeepSpeed Team
- from deepspeed.runtime.config_utils import get_scalar_param, DeepSpeedConfigObject
- from deepspeed.profiling.constants import *
- class DeepSpeedFlopsProfilerConfig(DeepSpeedConfigObject):
- def __init__(self, param_dict):
- super(DeepSpeedFlopsProfilerConfig, self).__init__()
- self.enabled = None
- self.recompute_fwd_factor = None
- self.profile_step = None
- self.module_depth = None
- self.top_modules = None
- if FLOPS_PROFILER in param_dict.keys():
- flops_profiler_dict = param_dict[FLOPS_PROFILER]
- else:
- flops_profiler_dict = {}
- self._initialize(flops_profiler_dict)
- def _initialize(self, flops_profiler_dict):
- self.enabled = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_ENABLED, FLOPS_PROFILER_ENABLED_DEFAULT)
- self.recompute_fwd_factor = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_RECOMPUTE_FWD_FACTOR,
- FLOPS_PROFILER_RECOMPUTE_FWD_FACTOR_DEFAULT)
- self.profile_step = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_PROFILE_STEP,
- FLOPS_PROFILER_PROFILE_STEP_DEFAULT)
- self.module_depth = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_MODULE_DEPTH,
- FLOPS_PROFILER_MODULE_DEPTH_DEFAULT)
- self.top_modules = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_TOP_MODULES,
- FLOPS_PROFILER_TOP_MODULES_DEFAULT)
- self.detailed = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_DETAILED, FLOPS_PROFILER_DETAILED_DEFAULT)
- self.output_file = get_scalar_param(flops_profiler_dict, FLOPS_PROFILER_OUTPUT_FILE,
- FLOPS_PROFILER_OUTPUT_FILE_DEFAULT)
|