def _get_torch_gpu_allocator_function_addresses(self): if self._use_external_gpu_allocator and torch.cuda.is_available(): # CPP extension to get torch GPU allocator's alloc and free function addresses from onnxruntime.training.ortmodule.torch_cpp_extensions import torch_gpu_allocator self._torch_alloc = torch_gpu_allocator.gpu_caching_allocator_raw_alloc_address() self._torch_free = torch_gpu_allocator.gpu_caching_allocator_raw_delete_address() self._torch_empty_cache = torch_gpu_allocator.gpu_caching_allocator_empty_cache_address()
def __init__(self, module, debug_options: DebugOptions): """Manages building and execution of onnx graphs This class is an abstract class and should not directly be instantiated. Please use one of the concrete implementations of GraphExecutionManager. Interacts with OrtModuleGraphBuilder to build and optimize the onnx graph, and ExecutionAgent to run the onnx graph. """ super(GraphExecutionManager, self).__init__(module._original_module) # Original and flattened (tranformed) output module self._flattened_module = module # onnx models self._onnx_models = _onnx_models.ONNXModels() # Model after inference optimization or gradient building. self._optimized_onnx_model = None self._graph_builder = None self._graph_info = None self._graph_initializer_names = None self._graph_initializer_names_to_train = None self._graph_initializers = None # TrainingAgent or InferenceAgent self._execution_agent = None # indicators of some logic have been executed previously thus could be skipped for faster training self._skip_check = _SkipCheck.SKIP_CHECK_DISABLED # Debug flags self._debug_options = debug_options # Graph transformer config # Specify cast propagation strategy. Currently three strategies are available, NONE, INSERT-AND-REDUCE and FLOOD-FILL # The default is NONE, which implies the transformer does no cast-propagation transformation. self._propagate_cast_ops_strategy = C.PropagateCastOpsStrategy.NONE # Optimize by moving Cast operations if propagate_cast_ops_level is non-negative. # - If the _propagate_cast_ops_level is set to zero, then the transformation considers only the opcodes specified by _propagate_cast_ops_allow # as "FP16 safe", in order to insert/(re)move cast operations before/after to perform such operations in reduced (16-bit) precision. # - If propagate_cast_ops_level is positive, 1 or 2, then in addition to opcode codes specified by propagate_cast_ops_allow use onnxruntime # predetermined list of opcodes considered safe to move before/after cast operation. # - Onnxruntime Level 1 predetermind "FP16 safe" opcodes include only opcode that do not perform any computation such as Transpose, Split, Reshape, etc. # whereas Level 2 perdetermined "FP16 safe" opcodes include opcodes that perform computation using contrib ops, GeLU, Dropout, LayerNormalization, etc. self._propagate_cast_ops_level = 1 # List of opcodes to be considered safe to move before/after cast operation if propagate_cast_ops_level is zero. self._propagate_cast_ops_allow = [] # Whether allow fusion of layer norm subgraph if doing so will cause modified precision. self._allow_layer_norm_mod_precision = False # Value can be either torch.onnx.TrainingMode.TRAINING or torch.onnx.TrainingMode.EVAL # To be instantiated in the concrete implementation of GraphExecutionManager self._export_mode = None # Related to training graph shape inference self._current_input_shape = None # default execution order is priority-based for both dynamic/static shape input for now # if we observe benefit of static shape, we can expose this flag to user self._use_static_shape = False # flag to enable symbolic shape inference for dynamic shape inputs to improve performance self._run_symbolic_shape_infer = True # A flag saying if custom autograd.Function should be allowed. True means yes and otherwise False. self._enable_custom_autograd_function = False self._input_info = None self._module_output_schema = None # TODO: Single device support for now self._device = _utils.get_device_from_module(module) self._module_parameters = inspect.signature( self._original_module.forward).parameters.values() # TODO: remove after PyTorch ONNX exporter supports VAR_KEYWORD parameters. for input_parameter in self._module_parameters: if input_parameter.kind == inspect.Parameter.VAR_KEYWORD: if self._debug_options.logging.log_level <= LogLevel.WARNING: warnings.warn( "The model's forward method has **kwargs parameter which has EXPERIMENTAL support!", UserWarning) self.is_rocm_pytorch = (True if ((torch.version.hip is not None) and (ROCM_HOME is not None)) else False) self._use_external_gpu_allocator = True if self._use_external_gpu_allocator and torch.cuda.is_available(): # CPP extension to get torch GPU allocator's alloc and free function addresses from onnxruntime.training.ortmodule.torch_cpp_extensions import torch_gpu_allocator self._torch_alloc = torch_gpu_allocator.gpu_caching_allocator_raw_alloc_address( ) self._torch_free = torch_gpu_allocator.gpu_caching_allocator_raw_delete_address( ) # WIP feature to enable caching in Gradient accumulation scenario. self._enable_grad_acc_optimization = False