From 43a5147e015e105547aa0e6862462a352fa43c5f Mon Sep 17 00:00:00 2001 From: pengwa Date: Thu, 23 Nov 2023 11:39:00 +0800 Subject: [PATCH] Memory optimization refactor and refinement (#17481) ### Memory optimization refactor and refinement Currently memory optimizer runs graph transformations and print recompute opportunities in INFO level, while ORT backend has many many INFO level logs making users hard to find those information. So we are looking for a Python binding API to retrieve the memory optimization opportunities instead of depending on the MemoryOptimizer's default logging. Then we can print ORTModule feature statistics using this information. Also, with such an API, we can create an ORT session created, where allocation plan is done, the analysis will consider buffer reuse as well. This can void giving some recomputation subgraphs that are reusing other subgraphs' output buffers. Check https://github.com/microsoft/onnxruntime/blob/pengwa/add_devinfo_level/docs/Memory_Optimizer.md for the new flow using `MemoryOptimizer`. This pull requests made following refactoring: 1. Print the log in ORTModule Python script, along with ORTModule feature enabling stats. This is implemented by exposing an API `get_serialized_ortmodule_memory_stat` to retrieve the memory optimization opportunities. 2. We are analyzing memory optimization opportunities considering ORT memory planning. This is done by firstly creating the execution graph without enabling MemoryOptimizer, then we call `execution_agent.get_serialized_ortmodule_memory_stat` which internally will consider the session memory allocation planner when analyzing memory optimization opportunity. As a direct result, the memory optimization opportunities can show those stashed activations that are reusing other buffers. 3. Move recompute analysis logic from memory_optimizer.h/cc to recompute_analysis.h/cc. 4. Abstract optimization strategies for their own implementation. This will make introducing new strategies (for example compression and decompression ) easier. New logging matrix (INFO Level), in WARNING level, the details will NOT show. ``` 2023-09-13 13:25:09,249 orttraining.rank-0 [WARNING] - ***** ONNX Runtime Training (ORTModule) is accelerating your model ***** ORTModule is enabled with following features ON/OFF for [training] mode: ATen Executor : ON : Dispatch ATen operators to ORT's ATen executor Cast Propagation : ON : Level 1 enabled Custom Function : ON : Support custom torch.autograd.Function export and execution Memory Optimizer : ON : RecomputeConfig: Reshape+Where+BiasSoftmax+:1:-1,Cast+:1:-1, ProbeLevel: 1, available configs: Config Freq Saving(B) Saving Symbolic(Bytes) - Plan 1 : ON : Reshape+Where+BiasSoftmax+:1:-1 5 671,088,640 640.0*inputs_input_ids_dim0*inputs_input_ids_dim1**2 - Plan 2 : ON : Cast+:1:-1 6 402,587,648 inputs_input_ids_dim0*inputs_input_ids_dim1*(384.0*inputs_input_ids_dim1 - 64.0) - Plan 3 : OFF : Reshape+Where+:1:-1 1 134,217,728 128.0*inputs_input_ids_dim0*inputs_input_ids_dim1**2 - Plan 4 : OFF : BiasSoftmax+:1:-1 1 134,086,656 128.0*inputs_input_ids_dim0*inputs_input_ids_dim1*(inputs_input_ids_dim1 - 1) - Plan 5 : OFF : BiasGelu+:1:-1 6 125,808,640 inputs_input_ids_dim0*(122880.0*inputs_input_ids_dim1 - 20480.0) - Plan 6 : OFF : FusedMatMul+:1:-1 6 125,808,640 inputs_input_ids_dim0*(122880.0*inputs_input_ids_dim1 - 20480.0) - Plan 7 : OFF : FusedMatMul+Add+FusedMatMul+Add+Add+Add+:1:-1 5 26,214,400 25600.0*inputs_input_ids_dim0*inputs_input_ids_dim1 - Plan 8 : OFF : Add+:1:-1 1 5,237,760 5120.0*inputs_input_ids_dim0*(inputs_input_ids_dim1 - 1) - Plan 9 : OFF : Reshape+Unsqueeze+Unsqueeze+Cast+Sub+Mul+Cast+:1:-1 1 4,096 4.0*inputs_input_ids_dim0*inputs_input_ids_dim1 - Plan 10 : OFF : Cast+:2:-1 1 2,048 2.0*inputs_input_ids_dim0*inputs_input_ids_dim1 Compute Optimizer : ON : Enable/Disable with env ORTMODULE_ENABLE_COMPUTE_OPTIMIZER=1/0 - FLOPReduction : ON : Reduce FLOPs by upstreaming shrinking-sized ops Auto Fallback : ON : Fallback to PyTorch when encountering unsupported ops TritonOp Enabled : OFF : ORT will switch to Triton for executing some ops to further accelerate training. ZeRO Stage3 Support : OFF : Enable/Disable with env ORTMODULE_ENABLE_ZERO_STAGE3=1/0 Total ORT initialization overhead is 10.73s where export takes 8.39s. Other overhead details: graph builder init takes 0.06s, runtime detection takes 0.01s, graph building takes 0.31s, session creation takes 1.96s Versions: ONNX Runtime - 1.16.0+cu118, ONNX - 1.11.0 Note 1: use comma to enable multiple plans at the same time. export ORTMODULE_MEMORY_OPT_CONFIG=,,... Note 2: saving is calculated based on the 1st batch symbolic dim values: inputs_input_ids_dim0=1, inputs_input_ids_dim1=1024, inputs_attention_mask_dim0=1, inputs_attention_mask_dim1=1024, inputs_labels_dim0=1, inputs_labels_dim1=1024, ************************************************************************ ``` If DEVINFO level is enabled, then more details about the memory optimizations are printed. ``` MemoryInsight Summary - User config: BiasGelu+:1:-1,Cast+:2:-1 ========================================================================================================================================== |Freq | Memory Optimization Opportunities (Clustered by node-level activation patterns) | |_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ | |3 |For each row options are mutually exclusive, only one of them can be enabled. | | | | | |>>Option 1 : Recompute subgraph FusedMatMul+Add+Reshape+ | | | Status : Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=FusedMatMul+Add+Reshape+:1:-1 | | | Stashed Activations: | | | - ReuseFreq : Output 0(3), | | | - Output 0 : [inputs_input_ids_dim0 x inputs_input_ids_dim1 x 32 x 240 x ], byte/elem: 2, 100% saved | |_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ | |2 |For each row options are mutually exclusive, only one of them can be enabled. | | | | | |>>Option 1 : Recompute subgraph Reshape+ | | | Status : Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=Reshape+:1:-1 | | | Stashed Activations: | | | - ReuseFreq : Output 0(2), | | | - Output 0 : [ x 2560 x ], byte/elem: 2, 100% saved | |_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ | |2 |For each row options are mutually exclusive, only one of them can be enabled. | | | | | |>>Option 1 : Recompute subgraph FusedMatMul+ | | | Status : Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=FusedMatMul+:1:-1 | | | Stashed Activations: | | | - Output 0 : [inputs_input_ids_dim0 x inputs_input_ids_dim1 x 10240 x ], byte/elem: 2, 100% saved | |_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ | |2 |For each row options are mutually exclusive, only one of them can be enabled. | | | | | |>>Option 1 : Recompute subgraph Cast+ | | | Status : Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=Cast+:1:-1 | | | Stashed Activations: | | | - Output 0 : [inputs_input_ids_dim0 x 32 x inputs_input_ids_dim1 x inputs_input_ids_dim1 x ], byte/elem: 2, 100% saved | |_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ | |2 |For each row options are mutually exclusive, only one of them can be enabled. | | | | | |>>Option 1 : Recompute subgraph Reshape+Where+BiasSoftmax+ | | | Status : Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=Reshape+Where+BiasSoftmax+:1:-1 | | | Stashed Activations: | | | - Output 0 : [inputs_input_ids_dim0 x 32 x inputs_input_ids_dim1 x inputs_input_ids_dim1 x ], byte/elem: 4, 100% saved | |_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ | |2 |For each row options are mutually exclusive, only one of them can be enabled. | | | | | |>>Option 1 : Recompute subgraph BiasGelu+ | | | Status : Enabled, requested count=-1, actual applied count=2 | | | Stashed Activations: | | | - Output 0 : [inputs_input_ids_dim0 x inputs_input_ids_dim1 x 10240 x ], byte/elem: 2, 100% saved | |_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ | |2 |For each row options are mutually exclusive, only one of them can be enabled. | | | | | |>>Option 1 : Recompute subgraph FusedMatMul+Add+FusedMatMul+Add+Add+Add+ | | | Status : Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=FusedMatMul+Add+FusedMatMul+Add+Add+Add+:1:-1 | | | Stashed Activations: | | | - Output 0 : [inputs_input_ids_dim0 x inputs_input_ids_dim1 x 2560 x ], byte/elem: 2, 100% saved | |_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ | |1 |For each row options are mutually exclusive, only one of them can be enabled. | | | | | |>>Option 1 : Recompute subgraph Reshape+Where+ | | | Status : Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=Reshape+Where+:1:-1 | | | Stashed Activations: | | | - Output 0 : [inputs_input_ids_dim0 x 32 x inputs_input_ids_dim1 x inputs_input_ids_dim1 x ], byte/elem: 4, 100% saved | |_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ | |1 |For each row options are mutually exclusive, only one of them can be enabled. | | | | | |>>Option 1 : Recompute subgraph FusedMatMul+ | | | Status : Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=FusedMatMul+:1:-1 | | | Stashed Activations: | | | - Output 0 : [inputs_input_ids_dim0*(inputs_input_ids_dim1 - 1) x 10240 x ], byte/elem: 2, 100% saved | |_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ | |1 |For each row options are mutually exclusive, only one of them can be enabled. | | | | | |>>Option 1 : Recompute subgraph Cast+ | | | Status : Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=Cast+:1:-1 | | | Stashed Activations: | | | - Output 0 : [inputs_input_ids_dim0 x 32 x inputs_input_ids_dim1 - 1 x inputs_input_ids_dim1 x ], byte/elem: 2, 100% saved | |_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ | |1 |For each row options are mutually exclusive, only one of them can be enabled. | | | | | |>>Option 1 : Recompute subgraph Reshape+Unsqueeze+Unsqueeze+Cast+Sub+Mul+Cast+ | | | Status : Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=Reshape+Unsqueeze+Unsqueeze+Cast+Sub+Mul+Cast+:1:-1 | | | Stashed Activations: | | | - Output 0 : [inputs_input_ids_dim0 x 1 x 1 x inputs_input_ids_dim1 x ], byte/elem: 4, 100% saved | | | | | |>>Option 2 : RecomputeWithCompromise subgraph Cast+ | | | Status : Enabled, requested count=-1, actual applied count=1 | | | Stashed Activations: | | | - Output 0 : [inputs_input_ids_dim0 x 1 x 1 x inputs_input_ids_dim1 x ], byte/elem: 4, 50% saved | |_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ | |1 |For each row options are mutually exclusive, only one of them can be enabled. | | | | | |>>Option 1 : Recompute subgraph BiasSoftmax+ | | | Status : Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=BiasSoftmax+:1:-1 | | | Stashed Activations: | | | - Output 0 : [inputs_input_ids_dim0 x 32 x inputs_input_ids_dim1 - 1 x inputs_input_ids_dim1 x ], byte/elem: 4, 100% saved | |_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ | |1 |For each row options are mutually exclusive, only one of them can be enabled. | | | | | |>>Option 1 : Recompute subgraph BiasGelu+ | | | Status : Enabled, requested count=-1, actual applied count=1 | | | Stashed Activations: | | | - Output 0 : [inputs_input_ids_dim0*(inputs_input_ids_dim1 - 1) x 10240 x ], byte/elem: 2, 100% saved | |_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ | |1 |For each row options are mutually exclusive, only one of them can be enabled. | | | | | |>>Option 1 : Recompute subgraph Add+ | | | Status : Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=Add+:1:-1 | | | Stashed Activations: | | | - Output 0 : [inputs_input_ids_dim0*(inputs_input_ids_dim1 - 1) x 2560 x ], byte/elem: 2, 100% saved | |_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ | ========================================================================================================================================== Note: use comma as a separator for enabling more than one subgraphs. ************************************************************************ ``` ### Motivation and Context --- cmake/onnxruntime_optimizer.cmake | 2 + docs/Memory_Optimizer.md | 151 ++-- .../onnxruntime_session_options_config_keys.h | 4 +- onnxruntime/core/common/string_utils.h | 28 + .../core/optimizer/graph_transformer_utils.cc | 13 - onnxruntime/core/session/inference_session.cc | 15 + .../orttraining/core/agent/training_agent.cc | 30 +- .../orttraining/core/agent/training_agent.h | 9 + .../compute_optimizer/padding_elimination.cc | 3 +- .../core/optimizer/memory_optimizer.cc | 671 ++------------- .../core/optimizer/memory_optimizer.h | 287 +------ .../core/optimizer/memory_optimizer/common.cc | 149 ++++ .../core/optimizer/memory_optimizer/common.h | 76 ++ .../memory_optimizer/memory_insight.cc | 763 ++++++++++++++++++ .../memory_optimizer/memory_insight.h | 129 +++ .../memory_optimizer/optimization_planner.cc | 140 ++++ .../memory_optimizer/optimization_planner.h | 133 +++ .../memory_optimizer/recompute_analysis.cc | 405 ++++++++++ .../memory_optimizer/recompute_analysis.h | 104 +++ .../core/optimizer/scaled_sum_fusion.cc | 4 +- .../python/orttraining_pybind_state.cc | 15 +- .../training/ortmodule/_execution_agent.py | 12 + .../ortmodule/_graph_execution_manager.py | 158 ++-- .../python/training/ortmodule/_io.py | 7 + .../python/training/ortmodule/_logger.py | 2 +- .../training/ortmodule/_runtime_inspector.py | 247 ++++-- .../training/ortmodule/_training_manager.py | 44 +- .../python/training/ortmodule/options.py | 6 +- .../python/training/utils/__init__.py | 2 + .../utils/hooks/_zero_offload_subscriber.py | 2 +- .../python/training/utils/ptable.py | 64 ++ 31 files changed, 2637 insertions(+), 1038 deletions(-) create mode 100644 orttraining/orttraining/core/optimizer/memory_optimizer/common.cc create mode 100644 orttraining/orttraining/core/optimizer/memory_optimizer/common.h create mode 100644 orttraining/orttraining/core/optimizer/memory_optimizer/memory_insight.cc create mode 100644 orttraining/orttraining/core/optimizer/memory_optimizer/memory_insight.h create mode 100644 orttraining/orttraining/core/optimizer/memory_optimizer/optimization_planner.cc create mode 100644 orttraining/orttraining/core/optimizer/memory_optimizer/optimization_planner.h create mode 100644 orttraining/orttraining/core/optimizer/memory_optimizer/recompute_analysis.cc create mode 100644 orttraining/orttraining/core/optimizer/memory_optimizer/recompute_analysis.h create mode 100644 orttraining/orttraining/python/training/utils/ptable.py diff --git a/cmake/onnxruntime_optimizer.cmake b/cmake/onnxruntime_optimizer.cmake index baea52e84a..6f09583199 100644 --- a/cmake/onnxruntime_optimizer.cmake +++ b/cmake/onnxruntime_optimizer.cmake @@ -86,6 +86,8 @@ if (onnxruntime_ENABLE_TRAINING) "${ORTTRAINING_SOURCE_DIR}/core/optimizer/*.cc" "${ORTTRAINING_SOURCE_DIR}/core/optimizer/compute_optimizer/*.h" "${ORTTRAINING_SOURCE_DIR}/core/optimizer/compute_optimizer/*.cc" + "${ORTTRAINING_SOURCE_DIR}/core/optimizer/memory_optimizer/*.h" + "${ORTTRAINING_SOURCE_DIR}/core/optimizer/memory_optimizer/*.cc" ) endif() diff --git a/docs/Memory_Optimizer.md b/docs/Memory_Optimizer.md index e9ceae00a6..0147a937db 100644 --- a/docs/Memory_Optimizer.md +++ b/docs/Memory_Optimizer.md @@ -20,70 +20,115 @@ Not all models and recipes need this optimizer technique. Imagine if your traini ## Quick trial 1. Make sure ONNX Runtime training wheel is installed and correctly configured. -2. Integrate models using `ORTModule`, be noted log_level should be equal to or lower than DEVINFO. - > ort_model = ORTModule(pt_model, DebugOptions(log_level=LogLevel.DEVINFO)) -3. Run the training as usual and redirect all outputs into the log file; then stop it after training a few steps. -4. Check the logging file, and search "Summary", you could find something like this: +2. Integrate models using `ORTModule`, be noted log_level should be equal or lower than INFO. + > ort_model = ORTModule(pt_model, DebugOptions(log_level=LogLevel.INFO)) +3. Run the training as usual; then stop it after training few steps. +4. Check the logs, you could find something like this: ``` - MemoryOptimizer Summary: - User config: + Memory Optimizer : OFF : Enable with env ORTMODULE_MEMORY_OPT_CONFIG=, available configs: + Config Freq Max Saving(B) Saving Symbolic(Bytes) + - Plan 1 : OFF : Reshape+Where+BiasSoftmax+:1:-1 5 671,088,640 640.0*inputs_input_ids_dim0*inputs_input_ids_dim1**2 + - Plan 2 : OFF : Cast+:1:-1 6 402,587,648 inputs_input_ids_dim0*inputs_input_ids_dim1*(384.0*inputs_input_ids_dim1 - 64.0) + - Plan 3 : OFF : Reshape+Where+:1:-1 1 134,217,728 128.0*inputs_input_ids_dim0*inputs_input_ids_dim1**2 + - Plan 4 : OFF : BiasSoftmax+:1:-1 1 134,086,656 128.0*inputs_input_ids_dim0*inputs_input_ids_dim1*(inputs_input_ids_dim1 - 1) + - Plan 5 : OFF : BiasGelu+:1:-1 6 125,808,640 inputs_input_ids_dim0*(122880.0*inputs_input_ids_dim1 - 20480.0) + - Plan 6 : OFF : FusedMatMul+:1:-1 6 125,808,640 inputs_input_ids_dim0*(122880.0*inputs_input_ids_dim1 - 20480.0) + - Plan 7 : OFF : FusedMatMul+Add+FusedMatMul+Add+Add+Add+:1:-1 5 26,214,400 25600.0*inputs_input_ids_dim0*inputs_input_ids_dim1 + - Plan 8 : OFF : Add+:1:-1 1 5,237,760 5120.0*inputs_input_ids_dim0*(inputs_input_ids_dim1 - 1) + - Plan 9 : OFF : Reshape+Unsqueeze+Unsqueeze+Cast+Sub+Mul+Cast+:1:-1 1 4,096 4.0*inputs_input_ids_dim0*inputs_input_ids_dim1 + - Plan 10 : OFF : Cast+:2:-1 1 2,048 2.0*inputs_input_ids_dim0*inputs_input_ids_dim1 - ================================= - ########Recompute######## - Subgraph: CumSum+Sub+Mul+Unsqueeze+Cast+Mul+Cast+Reshape+Mul+FusedMatMul+Add+Reshape+Cast+Where+Softmax+ - OptimizationType: Disabled - Patterns: - PatternShape:input_ids_dim0 x 16 x input_ids_dim1 x input_ids_dim1 x Frequency:23 - -------------------------------- - Subgraph: FastGelu+ - OptimizationType: Disabled - Patterns: - PatternShape:input_ids_dim0 x input_ids_dim1 x 4096 x Frequency:24 - ================================= - ########RecomputeWithCompromise######## - Subgraph: Cast+Where+Softmax+ - OptimizationType: Disabled - Patterns: - PatternShape:input_ids_dim0 x 16 x input_ids_dim1 x input_ids_dim1 x Frequency:24 - -------------------------------- - ================================= + + Note 1: use comma as delimiter to enable multiple memory optimization plans at the same time: + export ORTMODULE_MEMORY_OPT_CONFIG=,,... + Note 2: memory saving is calculated based on the 1st batch symbolic dim values: + inputs_input_ids_dim0=1, inputs_input_ids_dim1=1024, inputs_attention_mask_dim0=1, inputs_attention_mask_dim1=1024, inputs_labels_dim0=1, inputs_labels_dim1=1024, ``` -5. As shown above, 'Subgraph' shows 1) a string representative for a re-computable subgraph; and 2) current status of memory optimization. All are disabled for recompute in this case. -6. Set environment variable `ORTMODULE_MEMORY_OPT_CONFIG` to enable some of the subgraph to do recompute. In below example, 12 FastGelu related subgraphs are allowed to recompute. -`FastGelu+` is the subgraph string representative; `1` in the middle indicates 'Recompute' is enabled (0, on the contrary indicates it's disabled); `12` means the initial 12 subgraph occurrences will be recomputed, all others are left as it is, filling `-1` will make all occurrences be recomputed. +5. As shown above, `Config` is a string representative for a re-computable subgraph. All are disabled for recompute in this case. +6. Set environment variable `ORTMODULE_MEMORY_OPT_CONFIG` to enable some of the subgraph to do recompute. In below example, `6` `BiasGelu+` related subgraphs are allowed to recompute. +`BiasGelu+` is the subgraph string representative; `1` in the middle indicates 'Recompute' is enabled (0, on the contrary indicates it's disabled); `6` means the initial 6 subgraph occurrences will be recomputed, all others are left as it is, filling `-1` will make all occurrences be recomputed. ``` - export ORTMODULE_MEMORY_OPT_CONFIG="FastGelu+:1:12" + export ORTMODULE_MEMORY_OPT_CONFIG="BiasGelu+:1:6" # Use comma as separator for enabling more than one subgraphs. ``` -7. Then run the training again, you will see logs like this: +7. Then run the training again, and you will see logs like this: ``` - MemoryOptimizer Summary: - User config: - **FastGelu+:1:12** - ================================= - ########Recompute######## - Subgraph: CumSum+Sub+Mul+Unsqueeze+Cast+Mul+Cast+Reshape+Mul+FusedMatMul+Add+Reshape+Cast+Where+Softmax+ - OptimizationType: Disabled - Patterns: - PatternShape:input_ids_dim0 x 16 x input_ids_dim1 x input_ids_dim1 x Frequency:23 - -------------------------------- - Subgraph: FastGelu+ - OptimizationType: **Recompute (requested_count=12, actual applied_count=12)** - Patterns: - PatternShape:input_ids_dim0 x input_ids_dim1 x 4096 x Frequency:24 - ================================= - ########RecomputeWithCompromise######## - Subgraph: Cast+Where+Softmax+ - OptimizationType: Disabled - Patterns: - PatternShape:input_ids_dim0 x 16 x input_ids_dim1 x input_ids_dim1 x Frequency:24 - -------------------------------- - ================================= + Memory Optimizer : ON : User config: Reshape+Where+BiasSoftmax+:1:-1, probe level: 1, available configs: + Config Freq Max Saving(B) Saving Symbolic(Bytes) + - Plan 1 : OFF : Reshape+Where+BiasSoftmax+:1:-1 5 671,088,640 640.0*inputs_input_ids_dim0*inputs_input_ids_dim1**2 + - Plan 2 : OFF : Cast+:1:-1 6 402,587,648 inputs_input_ids_dim0*inputs_input_ids_dim1*(384.0*inputs_input_ids_dim1 - 64.0) + - Plan 3 : OFF : Reshape+Where+:1:-1 1 134,217,728 128.0*inputs_input_ids_dim0*inputs_input_ids_dim1**2 + - Plan 4 : OFF : BiasSoftmax+:1:-1 1 134,086,656 128.0*inputs_input_ids_dim0*inputs_input_ids_dim1*(inputs_input_ids_dim1 - 1) + - Plan 5 : ON : BiasGelu+:1:-1 6 125,808,640 inputs_input_ids_dim0*(122880.0*inputs_input_ids_dim1 - 20480.0) + - Plan 6 : OFF : FusedMatMul+:1:-1 6 125,808,640 inputs_input_ids_dim0*(122880.0*inputs_input_ids_dim1 - 20480.0) + - Plan 7 : OFF : FusedMatMul+Add+FusedMatMul+Add+Add+Add+:1:-1 5 26,214,400 25600.0*inputs_input_ids_dim0*inputs_input_ids_dim1 + - Plan 8 : OFF : Add+:1:-1 1 5,237,760 5120.0*inputs_input_ids_dim0*(inputs_input_ids_dim1 - 1) + - Plan 9 : OFF : Reshape+Unsqueeze+Unsqueeze+Cast+Sub+Mul+Cast+:1:-1 1 4,096 4.0*inputs_input_ids_dim0*inputs_input_ids_dim1 + - Plan 10 : OFF : Cast+:2:-1 1 2,048 2.0*inputs_input_ids_dim0*inputs_input_ids_dim1 ``` 8. You may need iterate few times on step 6 and 7 until you find a good config for this model to run a bigger batch size. Or you may fail to find if memory optimization does not apply to the model well. +## Optimization Configuration + +The basic optimization unit is represented with a unique `cluster id`, for example `BiasGelu+` is one `cluster id`. +Following `cluster id` is the `optimization strategy`: 0 - none, 1 - recompute, 2 - recompute with compromised memory saving. +Following `optimization strategy` is the `request count` to apply the given optimization. Using `-1` to apply all. This would give user a bit more flexibility to avoid unnecessary memory saving. + ## Compromised Recompute -If you check the above logs, there is a separate section called "RecomputeWithCompromise". Recompute the subgraphs under it usually will save part of the activation (for example half of them), not all of them. Follow the same way to enable it. +If you check the above logs, there is a config `Cast+:2:-1`, `2` indicates it's a recomputation than can save part of the stashed activation size, not all. Recompute the subgraphs under it usually will save part of the activation (for example half of them), not all of them. Follow the same way to enable it. + +## Memory Optimization Debug Infos + +Using following log level +> ort_model = ORTModule(pt_model, DebugOptions(log_level=LogLevel.DEVINFO)) + +Besides the logs shown in `LogLevel.INFO`, you can also see different node patterns that can apply different optimization options. + +The way we get the table: +- For a specific node, it might has different optimization options, we [generates](../orttraining/orttraining/core/optimizer/memory_optimizer/common.h#L124C26-L124C26) a hash (called `Node Cluster ID`) for the node according to all available optimization options. +- Map all nodes having same `Node Cluster ID` in buckets, each bucket is displayed as one row. + +``` +MemoryInsight Summary - User config: not provided +=========================================================================================================================================== +|Freq | Memory Optimization Opportunities (Clustered by node-level activation patterns) | +|_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _| +|6 |For each row options are mutually exclusive, only one of them can be enabled. | +| | | +| |>>Option 1 : Recompute subgraph FusedMatMul+Add+Reshape+ | +| | Status : Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=FusedMatMul+Add+Reshape+:1:-1 | +| | Stashed Activations: | +| | - ReuseFreq : Output 0(6), | +| | - Output 0 : [((inputs_input_ids_dim0)*(inputs_input_ids_dim1)*(32)*(240))], byte/elem: 2, 100% saved | +|_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _| +|5 |For each row options are mutually exclusive, only one of them can be enabled. | +| | | +| |>>Option 1 : Recompute subgraph FusedMatMul+ | +| | Status : Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=FusedMatMul+:1:-1 | +| | Stashed Activations: | +| | - Output 0 : [((inputs_input_ids_dim0)*(inputs_input_ids_dim1)*(10240))], byte/elem: 2, 100% saved | +|_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _| +|5 |For each row options are mutually exclusive, only one of them can be enabled. | +| | | +| |>>Option 1 : Recompute subgraph Cast+ | +| | Status : Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=Cast+:1:-1 | +| | Stashed Activations: | +| | - Output 0 : [((inputs_input_ids_dim0)*(32)*(inputs_input_ids_dim1)*(inputs_input_ids_dim1))], byte/elem: 2, 100% saved | +|_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _| +|1 |For each row options are mutually exclusive, only one of them can be enabled. | +| | | +| |>>Option 1 : Recompute subgraph Reshape+Unsqueeze+Unsqueeze+Cast+Sub+Mul+Cast+ | +| | Status : Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=Reshape+Unsqueeze+Unsqueeze+Cast+Sub+Mul+Cast+:1:-1 | +| | Stashed Activations: | +| | - Output 0 : [((inputs_input_ids_dim0)*(1)*(1)*(inputs_input_ids_dim1))], byte/elem: 4, 100% saved | +| | | +| |>>Option 2 : RecomputeWithCompromise subgraph Cast+ | +| | Status : Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=Cast+:2:-1 | +| | Stashed Activations: | +| | - Output 0 : [((inputs_input_ids_dim0)*(1)*(1)*(inputs_input_ids_dim1))], byte/elem: 4, 50% saved | +|_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _| + +``` ## Notes diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h index 831def24e4..4628afbb5a 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -80,13 +80,13 @@ static const char* const kOrtSessionOptionsDisableAheadOfTimeFunctionInlining = #ifdef ENABLE_TRAINING // Specifies a list of op types for memory footprint reduction. // The value should be a ","-delimited list of pair of -// . +// . // For example, "Gelu+Cast+:1:0,Dropout+:1:1". // A valid "subgraph string" should be one subgraph representation output by ORT graph transformations. // "optimization strategy" currently has valid values: 0 - disabled, 1 - recompute. // "number of subgraph to apply" is used to control how many subgraphs to apply optimization, to avoid "oversaving" // the memory. -static const char* const kOrtSessionOptionsMemoryOptimizerEnabler = "optimization.enable_memory_optimizer"; +static const char* const kOrtSessionOptionsMemoryOptimizerEnabler = "optimization.memory_optimizer_config"; // Specifies the level for detecting subgraphs for memory footprint reduction. // The value should be an integer. The default value is 0. diff --git a/onnxruntime/core/common/string_utils.h b/onnxruntime/core/common/string_utils.h index 6e0eb460d2..eca1221e84 100644 --- a/onnxruntime/core/common/string_utils.h +++ b/onnxruntime/core/common/string_utils.h @@ -3,6 +3,7 @@ #pragma once +#include #include #include @@ -37,5 +38,32 @@ inline InlinedVector SplitString(std::string_view string_to_sp return result; } +/** + * Trim a string from start inplace. + * @param s The string to trim. + */ +inline void TrimStringFromLeft(std::string& s) { + s.erase(s.begin(), std::find_if(s.begin(), s.end(), [](unsigned char ch) { return !std::isspace(ch); })); +} + +/** + * Trim a string from end inplace. + * @param s The string to trim. + */ +inline void TrimStringFromRight(std::string& s) { + s.erase(std::find_if(s.rbegin(), s.rend(), [](unsigned char ch) { return !std::isspace(ch); }).base(), s.end()); +} + +/** + * Trim a string from both ends. + * @param s The string to trim. + * @return The trimmed string. + */ +inline std::string TrimString(std::string s) { + TrimStringFromRight(s); + TrimStringFromLeft(s); + return s; +} + } // namespace utils } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index c1397e92d9..3d6251a694 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -77,7 +77,6 @@ #include "orttraining/core/optimizer/bias_softmax_dropout_fusion.h" #include "orttraining/core/optimizer/bitmask_dropout_replacement.h" #include "orttraining/core/optimizer/sce_loss_grad_bias_fusion.h" -#include "orttraining/core/optimizer/memory_optimizer.h" #endif #ifdef ENABLE_TRITON #include "orttraining/core/optimizer/triton_fusion.h" @@ -354,18 +353,6 @@ InlinedVector> GenerateTransformers( // fusions might be prevented if this one removes a Q/DQ node too early. transformers.emplace_back(std::make_unique(enable_quant_qdq_cleanup)); -#ifdef ENABLE_TRAINING - // Put memory optimization transformer at last (which is done after most of fusions are done) by intention. - // Known issue: after memory optimization is completed, if some fusion happens, it is possible that the - // node priority got changed. This may disorder the execution order of nodes to recompute. - // TODO(pengwa): need to fix this issue. - const std::string enable_memory_optimizer = - session_options.config_options.GetConfigOrDefault(kOrtSessionOptionsMemoryOptimizerEnabler, ""); - const std::string probe_level = - session_options.config_options.GetConfigOrDefault(kOrtSessionOptionsMemoryOptimizerProbeLevel, "0"); - transformers.emplace_back(std::make_unique(enable_memory_optimizer, probe_level)); -#endif - } break; case TransformerLevel::Level3: { diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index f02d180ab1..75be72658f 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -74,6 +74,7 @@ #ifdef ENABLE_TRAINING #include "core/framework/partial_graph_execution_state.h" #include "core/framework/stream_execution_context.h" +#include "orttraining/core/optimizer/memory_optimizer.h" #endif using namespace ONNX_NAMESPACE; @@ -1149,6 +1150,20 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, bool ORT_RETURN_IF_ERROR_SESSIONID_(apply_transformer_once(copy_transformer, *session_logger_, graph)); } +#ifdef ENABLE_TRAINING + // Enable memory optimizations (mainly insert recomputation nodes with priority). + // Only applicable for training scenarios. + { + const std::string memory_optimizer_config = + session_options_.config_options.GetConfigOrDefault(kOrtSessionOptionsMemoryOptimizerEnabler, ""); + const std::string probe_level = + session_options_.config_options.GetConfigOrDefault(kOrtSessionOptionsMemoryOptimizerProbeLevel, "0"); + + MemoryOptimizer mem_transformer{memory_optimizer_config, probe_level}; + ORT_RETURN_IF_ERROR_SESSIONID_(apply_transformer_once(mem_transformer, *session_logger_, graph)); + } +#endif + return Status::OK(); } #endif // !defined(ORT_MINIMAL_BUILD) diff --git a/orttraining/orttraining/core/agent/training_agent.cc b/orttraining/orttraining/core/agent/training_agent.cc index 3b701fa8bf..0b38a79cc2 100644 --- a/orttraining/orttraining/core/agent/training_agent.cc +++ b/orttraining/orttraining/core/agent/training_agent.cc @@ -1,11 +1,17 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include +#include +#include +#include + #include "orttraining/core/agent/training_agent.h" #include "core/framework/utils.h" #include "core/framework/feeds_fetches_manager.h" #include "core/framework/partial_graph_execution_state.h" #include "core/framework/stream_execution_context.h" +#include "orttraining/core/optimizer/memory_optimizer/memory_insight.h" namespace onnxruntime { namespace training { @@ -25,7 +31,8 @@ TrainingAgent::TrainingAgent(InferenceSession& session, std::vector bw_feed_names; size_t break_point = 0; - auto& training_node_execution_order = session_state.GetGraphViewer().GetNodesInTopologicalOrder(session.GetSessionOptions().execution_order); + auto& training_node_execution_order = session_state.GetGraphViewer().GetNodesInTopologicalOrder( + session.GetSessionOptions().execution_order); for (auto node_index : training_node_execution_order) { if (session_state.GetKernel(node_index)->KernelDef().OpName() == "YieldOp") { auto& node = *(session_state.GetGraphViewer().GetGraph().GetNode(node_index)); @@ -89,7 +96,8 @@ void TrainingAgent::CreateAndInitializeFeedsFetchesManager(const SessionState& s const std::vector& feed_names, const std::vector& fetches_names, const std::vector& outputs_device_info, - std::unique_ptr& feeds_fetches_manager) { + std::unique_ptr& + feeds_fetches_manager) { ORT_THROW_IF_ERROR(FeedsFetchesManager::Create(feed_names, fetches_names, session_state.GetOrtValueNameIdxMap(), feeds_fetches_manager)); auto& fetch_info = feeds_fetches_manager->GetMutableFetchesDeviceCopyInfo(); @@ -100,5 +108,23 @@ void TrainingAgent::CreateAndInitializeFeedsFetchesManager(const SessionState& s ORT_ENFORCE(utils::InitializeFeedFetchCopyInfo(session_state, *feeds_fetches_manager) == Status::OK()); } +std::string TrainingAgent::GetSerializedORTModuleMemoryStat(std::string_view memory_optimization_config, + std::string_view recompute_probe_level, + std::map>& + cluster_id_combinations_to_saved_symbolic_byte_map) + const { + auto& session_state = inference_session_.GetSessionState(); + const OrtValueNameIdxMap& ortvalue_name_to_idx_map = session_state.GetOrtValueNameIdxMap(); + const SequentialExecutionPlan& p_seq_exec_plan = *session_state.GetExecutionPlan(); + return optimizer::memory_optimizer::GetSerializedORTModuleMemoryStat( + session_state.GetGraphViewer(), + memory_optimization_config, + recompute_probe_level, + *inference_session_.GetLogger(), + cluster_id_combinations_to_saved_symbolic_byte_map, + &ortvalue_name_to_idx_map, + &p_seq_exec_plan); +} + } // namespace training } // namespace onnxruntime diff --git a/orttraining/orttraining/core/agent/training_agent.h b/orttraining/orttraining/core/agent/training_agent.h index b12f5e6d75..37e5272f66 100644 --- a/orttraining/orttraining/core/agent/training_agent.h +++ b/orttraining/orttraining/core/agent/training_agent.h @@ -5,11 +5,15 @@ #include #include +#include +#include +#include #include "core/common/common.h" #include "core/common/logging/logging.h" #include "core/framework/framework_common.h" #include "core/session/inference_session.h" +#include "orttraining/core/optimizer/memory_optimizer/memory_insight.h" namespace onnxruntime { struct PartialGraphExecutionState; @@ -45,6 +49,11 @@ class TrainingAgent { const std::vector& outputs_device_info, std::unique_ptr& feeds_fetches_manager); + std::string GetSerializedORTModuleMemoryStat(std::string_view memory_optimization_config, + std::string_view recompute_probe_level, + std::map>& + cluster_id_combinations_to_saved_symbolic_byte_map) const; + private: // TrainingAgent runs on a InferenceSession under the hood InferenceSession& inference_session_; diff --git a/orttraining/orttraining/core/optimizer/compute_optimizer/padding_elimination.cc b/orttraining/orttraining/core/optimizer/compute_optimizer/padding_elimination.cc index 73638e8ba6..2d75a02004 100644 --- a/orttraining/orttraining/core/optimizer/compute_optimizer/padding_elimination.cc +++ b/orttraining/orttraining/core/optimizer/compute_optimizer/padding_elimination.cc @@ -470,7 +470,8 @@ Status PaddingElimination::ApplyImpl(Graph& graph, bool& modified, int graph_lev // Get the first two dims value of input_ids which is [batch_size, seq_len] NodeArg* first_two_dims_arg = GetDimsValue(graph, input_ids_arg, - CreateInitializerFromVector(graph, {2}, {0, 1}, graph.GenerateNodeArgName("first_two_indices")), + CreateInitializerFromVector(graph, {2}, {0, 1}, + graph.GenerateNodeArgName("first_two_indices")), *embedding_node); // Add flatten pattern to each input node of the subgraph diff --git a/orttraining/orttraining/core/optimizer/memory_optimizer.cc b/orttraining/orttraining/core/optimizer/memory_optimizer.cc index 88c786d693..834e5ebb5f 100644 --- a/orttraining/orttraining/core/optimizer/memory_optimizer.cc +++ b/orttraining/orttraining/core/optimizer/memory_optimizer.cc @@ -1,233 +1,84 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include +#include +#include +#include +#include +#include + #include "core/framework/random_seed.h" #include "core/framework/tensorprotoutils.h" #include "core/graph/graph_utils.h" #include "core/optimizer/utils.h" #include "orttraining/core/graph/recompute_graph_utils.h" #include "orttraining/core/optimizer/memory_optimizer.h" +#include "orttraining/core/optimizer/memory_optimizer/common.h" +#include "orttraining/core/optimizer/memory_optimizer/optimization_planner.h" +#include "orttraining/core/optimizer/memory_optimizer/recompute_analysis.h" +#include "orttraining/core/optimizer/memory_optimizer/memory_insight.h" namespace onnxruntime { namespace { -constexpr int32_t MAXIMUM_RECOMPUTE_NODE_COUNT = 15; - -std::string TensorShapeProtoToString(const ONNX_NAMESPACE::TensorShapeProto* shape) { - std::ostringstream shape_oss; - if (shape != nullptr) { - for (int dim_index = 0; dim_index < shape->dim_size(); dim_index++) { - auto dim = shape->dim(dim_index); - if (utils::HasDimValue(dim)) { - shape_oss << dim.dim_value() << " x "; - } else { - shape_oss << dim.dim_param() << " x "; - } - } - } else { - shape_oss << "unknown"; - } - - return shape_oss.str(); -} - -int ParseIntValueFromString(std::string_view str) { - int int_value = 0; - auto result = std::from_chars(str.data(), str.data() + str.size(), int_value); - ORT_ENFORCE(result.ec != std::errc::invalid_argument, "Fail to convert to int from string: ", str); - return int_value; -} - -constexpr bool IsForwardPassOperator(ptrdiff_t op_order_in_topological_sort, ptrdiff_t boundary_op_order_in_topological_sort) { +constexpr bool IsForwardPassOperator(ptrdiff_t op_order_in_topological_sort, + ptrdiff_t boundary_op_order_in_topological_sort) { return op_order_in_topological_sort <= boundary_op_order_in_topological_sort; } -static size_t GetElementSize(const ONNX_NAMESPACE::DataType& tensor_type) { - const ONNX_NAMESPACE::TypeProto& type_proto = ONNX_NAMESPACE::Utils::DataTypeUtils::ToTypeProto(tensor_type); - MLDataType ml_data_type = DataTypeImpl::TypeFromProto(type_proto); - const TensorTypeBase* tensor_type_base = ml_data_type->AsTensorType(); - ORT_ENFORCE(nullptr != tensor_type_base); - MLDataType elt_type = tensor_type_base->GetElementType(); - return elt_type->Size(); -} - -// TODO(pengwa): extend this function to be more general. -float InputOutputSizeRatio(const Node* node) { - if (node->OpType().compare("Cast") == 0) { - const NodeArg* input = node->InputDefs()[0]; - const NodeArg* output = node->OutputDefs()[0]; - if (input->TypeAsProto()->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_STRING || - output->TypeAsProto()->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_STRING) { - return 1.0f; - } - const auto& ptype1 = input->Type(); - const auto& ptype2 = output->Type(); - float ratio = float(GetElementSize(ptype1)) / (float)GetElementSize(ptype2); - return ratio; - } - - return 1.0f; -} - } // namespace -Status MemoryOptimizer::ParseConfigFromString(const std::string& enable_memory_optimizer, +Status MemoryOptimizer::ParseConfigFromString(const std::string& memory_optimizer_config, const std::string& level) { - optimizer_config_ = enable_memory_optimizer; - if (!enable_memory_optimizer.empty()) { - const auto user_config_strs = utils::SplitString(enable_memory_optimizer, ","); - for (const auto& user_config_str : user_config_strs) { - const auto user_config = utils::SplitString(user_config_str, ":"); - ORT_RETURN_IF_NOT(user_config.size() == 3, - "User config should be in format of SubgraphStr:OptimizationType:RequestApplyCount."); + optimizer_config_ = memory_optimizer_config; - const std::string subgraph_string_representation(user_config[0]); - int optimization_type_int = ParseIntValueFromString(user_config[1]); - int requested_apply_count = ParseIntValueFromString(user_config[2]); - ORT_RETURN_IF_NOT(optimization_type_int < static_cast(OptimizationType::TypeMax) && - optimization_type_int >= 0, - "Invalid optimization type specified for subgraph: ", - subgraph_string_representation); + ORT_RETURN_IF_ERROR(optimizer::memory_optimizer::ParseConfigFromString( + memory_optimizer_config, + pattern_subgraph_to_user_optimizer_config_map_)); - ORT_RETURN_IF_NOT(requested_apply_count == -1 || requested_apply_count >= 0, - "Invalid requested_apply_count specified for subgraph: ", requested_apply_count); - - // At this point, subgraph_string_representation is a pattern graph string representation. - pattern_subgraph_to_user_optimizer_config_map_[subgraph_string_representation] = - UserConfig{static_cast(optimization_type_int), requested_apply_count}; - } - } - - int probe_level = ParseIntValueFromString(level); - ORT_RETURN_IF_NOT(probe_level < static_cast(ProbeLevel::LevelMax) && probe_level >= 0, + int probe_level = optimizer::memory_optimizer::ParseIntValueFromString(level); + ORT_RETURN_IF_NOT(probe_level < static_cast(optimizer::memory_optimizer::ProbeLevel::LevelMax) && + probe_level >= 0, "Invalid probe level specified: ", level); - recompute_probe_level_ = static_cast(probe_level); - - return Status::OK(); -} - -int64_t MemoryOptimizer::PrepareForTransformation(const Graph& graph, - ActivationUsedMap& fw_op_output_arg_used_map, - InlinedHashMap& - node_index_to_its_order_in_topological_sort_map) const { - fw_op_output_arg_used_map.clear(); - - GraphViewer graph_viewer(graph); - const auto& node_ids = graph_viewer.GetNodesInTopologicalOrder(); - - // Find boundary ops between forward and backward pass, currently, it's limited to YieldOp. - ptrdiff_t yield_op_order_in_topological_sort = -1; - for (size_t i = 0; i < node_ids.size(); ++i) { - const Node* p_node = graph.GetNode(node_ids[i]); - if (p_node == nullptr) { /* skip removed nodes*/ - continue; - } - - if (p_node->OpType() == "YieldOp") { - yield_op_order_in_topological_sort = static_cast(i); - } - - node_index_to_its_order_in_topological_sort_map[p_node->Index()] = i; - } - - // If boundary op found, create forward op output arg used map. - if (yield_op_order_in_topological_sort >= 0) { - for (size_t i = 0; i < node_ids.size(); ++i) { - const Node* p_node = graph.GetNode(node_ids[i]); - if (p_node == nullptr /* skip removed nodes*/) { - continue; - } - - const Node& node = *p_node; - bool is_forward_op = IsForwardPassOperator(static_cast(i), yield_op_order_in_topological_sort); - if (!is_forward_op) { - continue; - } - - for (auto& output_arg : node.OutputDefs()) { - bool used_in_fw = false; - bool used_in_bw = false; - for (auto& consumer_node : graph.GetConsumerNodes(output_arg->Name())) { - size_t consumer_node_index_in_topological_order = - node_index_to_its_order_in_topological_sort_map.at(consumer_node->Index()); - if (IsForwardPassOperator(static_cast(consumer_node_index_in_topological_order), - yield_op_order_in_topological_sort)) { - used_in_fw = true; - } else { - used_in_bw = true; - } - } - fw_op_output_arg_used_map.insert({{output_arg->Name(), std::make_pair(used_in_fw, used_in_bw)}}); - } - } - } - - // Return whether boundary op is found or not. - return yield_op_order_in_topological_sort; -} - -Status MemoryOptimizer::GetStashedActivationCandidates(const Graph& graph, - const InlinedHashMap>& - fw_op_output_arg_used_map, - InlinedHashMap>& - candidate_output_args_map, - const logging::Logger& logger) const { - for (auto& kv : fw_op_output_arg_used_map) { - // used by fw and bw, then it is a candidates. - if (kv.second.first && kv.second.second) { - const Node* n = graph.GetProducerNode(kv.first); - ORT_ENFORCE(n, "Activation should have a producer node"); - size_t k = 0; - for (k = 0; k < n->OutputDefs().size(); ++k) { - if (n->OutputDefs()[k]->Name().compare(kv.first) == 0) { - break; - } - } - - candidate_output_args_map[n].push_back(k); - LOGS(logger, VERBOSE) << "Find candidate output named [" << kv.first << "] of Node " << n->Name() << "(" - << n->OpType() << ")"; - } - } + recompute_probe_level_ = static_cast(probe_level); return Status::OK(); } bool MemoryOptimizer::ModifyGraph(Graph& graph, - const InlinedHashMap& + const InlinedHashMap& node_index_to_its_order_in_topological_sort_map, const InlinedHashMap>& candidate_output_args_map, const logging::Logger& logger, - int64_t boundary_op_order_in_topological_sort, - SubGraphStores& subgraph_stores, - Node* node) const { + ptrdiff_t boundary_op_order_in_topological_sort, + Node* node, + std::shared_ptr& node_plan, + std::shared_ptr& apply_context) + const { bool graph_is_modified = false; - if (subgraph_stores.SubGraphDescCount() == 0) { - return graph_is_modified; - } - - SubGraphStores::GraphInstanceInfo& sub_graph_instance_info = - subgraph_stores.GetSubGraphInstance(node); - - SubGraphDesc& subgraph_desc = subgraph_stores.GetSubGraphDesc(sub_graph_instance_info.second); - UserConfig user_config = subgraph_desc.user_optimizer_config; - int skip_count = (user_config.requested_count == -1) + int skip_count = (apply_context->requested_count == -1) ? 0 - : std::max(0, subgraph_desc.total_frequency - user_config.requested_count); + : std::max(0, apply_context->total_frequency - apply_context->requested_count); - subgraph_desc.skip_count += 1; + apply_context->skip_count += 1; - if (user_config.type != OptimizationType::None && subgraph_desc.skip_count > skip_count) { - subgraph_desc.applied_count += 1; + if (apply_context->skip_count > skip_count) { + apply_context->applied_count += 1; Node* replacement_node_ptr = nullptr; - LOGS(logger, WARNING) << "[Modify Graph] Node " << node->Name() << "(" << node->OpType() << ") is " - << UserConfigToString(user_config); - if (user_config.type == OptimizationType::Recompute) { - ORT_ENFORCE(CreateRecomputeGraph(graph, sub_graph_instance_info.first, replacement_node_ptr).IsOK()); + LOGS(logger, INFO) << "Node " << node->Name() << "(" << node->OpType() << ") is applying following optimization:" + << "type [" << optimizer::memory_optimizer::OptimizationTypeToString(apply_context->type) + << "], request count [" << apply_context->requested_count << "]"; + if (apply_context->type == optimizer::memory_optimizer::OptimizationType::Recompute || + apply_context->type == optimizer::memory_optimizer::OptimizationType::RecomputeWithCompromise) { + optimizer::memory_optimizer::NodeRecomputePlan* recompute_plan = + dynamic_cast(node_plan.get()); + ORT_ENFORCE(recompute_plan != nullptr); + ORT_ENFORCE(CreateRecomputeGraph(graph, recompute_plan->GetNodesInTopoOrder(), replacement_node_ptr).IsOK()); } else { - ORT_THROW("unsupported optimization type found: " + UserConfigToString(user_config)); + ORT_THROW("unsupported optimization type found."); } ORT_ENFORCE(replacement_node_ptr); @@ -278,60 +129,44 @@ Status MemoryOptimizer::ApplyImpl(Graph& graph, bool& modified, int /*graph_leve LOGS(logger, VERBOSE) << "Memory optimization config: " << optimizer_config_ << ", probe level: " << static_cast(recompute_probe_level_); - InlinedHashMap> fw_op_output_arg_used_map; - InlinedHashMap node_index_to_its_order_in_topological_sort_map; - int64_t boundary_op_order_in_topological_sort = - PrepareForTransformation(graph, fw_op_output_arg_used_map, - node_index_to_its_order_in_topological_sort_map); - if (boundary_op_order_in_topological_sort < 0) { - LOGS(logger, VERBOSE) << "No boundary op found. Skip memory optimization."; + if (pattern_subgraph_to_user_optimizer_config_map_.empty()) { + LOGS(logger, VERBOSE) << "No optimization pattern is specified, skip memory optimization."; return Status::OK(); } + ptrdiff_t yield_op_order_in_topological_sort; InlinedHashMap> candidate_output_args_map; - ORT_RETURN_IF_ERROR(GetStashedActivationCandidates(graph, fw_op_output_arg_used_map, candidate_output_args_map, - logger)); - - SubGraphStores recompute_subgraph_stores; - SubGraphStores recompute_with_compromise_subgraph_stores; - GraphViewer graph_viewer(graph); - const auto& node_ids = graph_viewer.GetNodesInTopologicalOrder(); + InlinedHashMap node_index_to_its_order_in_topological_sort_map; // The first pass - find the candidate subgraphs. - for (int i = static_cast(node_ids.size()) - 1; i >= 0; --i) { - Node* p_node = graph.GetNode(node_ids[i]); - if (p_node == nullptr) { - continue; - } + GraphViewer graph_viewer(graph); + optimizer::memory_optimizer::MemoryOptimizationPlanner memory_opt_planner; + ORT_ENFORCE(optimizer::memory_optimizer::FindORTModuleMemoryOpportunity( + graph_viewer, + recompute_probe_level_, + logger, + node_index_to_its_order_in_topological_sort_map, + yield_op_order_in_topological_sort, + candidate_output_args_map, + memory_opt_planner) + .IsOK()); - if (candidate_output_args_map.find(p_node) == candidate_output_args_map.end()) { - continue; - } - - bool can_compromise_stashed_activation = false; - CheckNodeForRecompute(*p_node, fw_op_output_arg_used_map, - node_index_to_its_order_in_topological_sort_map, - candidate_output_args_map, - recompute_subgraph_stores, logger, false, - can_compromise_stashed_activation); - - if (can_compromise_stashed_activation) { - LOGS(logger, VERBOSE) << "Searching Node " << p_node->Name() << "(" << p_node->OpType() - << ") for compromised recompute"; - // If the subgraph recompute can save memory by comprising the assumption - recompute graphs' input must exist - // during backward pass, then we can try to compromise the assumption. - CheckNodeForRecompute(*p_node, fw_op_output_arg_used_map, node_index_to_its_order_in_topological_sort_map, - candidate_output_args_map, - recompute_with_compromise_subgraph_stores, logger, true, - can_compromise_stashed_activation); - } - } + // Finalize the plan according to user config, + // then create a ClusterApplyContext for each unique cluster (having the same node pattern) + InlinedHashMap> + node_to_opt_plan_map; + optimizer::memory_optimizer::NodeToClusterApplyContextMap node_to_apply_context_map; + ORT_ENFORCE(memory_opt_planner.FinalizeNodePlansFromUserConfig(pattern_subgraph_to_user_optimizer_config_map_, + node_to_opt_plan_map, + node_to_apply_context_map) + .IsOK()); // The second pass - apply the transformation. // Iterate through the nodes in reversed topological order and find the subgraph that can be alleviated. // The reason we do reversed topological order is that we want the later layers' recompute nodes can be appended // earlier than the earlier layers, in this way, the execution order of later layers will be in front of the earlier // layers. + const auto& node_ids = graph_viewer.GetNodesInTopologicalOrder(); for (int i = static_cast(node_ids.size()) - 1; i >= 0; --i) { Node* p_node = graph.GetNode(node_ids[i]); if (p_node == nullptr) { @@ -339,374 +174,40 @@ Status MemoryOptimizer::ApplyImpl(Graph& graph, bool& modified, int /*graph_leve } bool has_been_modified = false; - if (recompute_subgraph_stores.ContainsSubGraphInstance(p_node)) { + if (node_to_opt_plan_map.find(p_node) != node_to_opt_plan_map.end()) { has_been_modified = ModifyGraph(graph, node_index_to_its_order_in_topological_sort_map, candidate_output_args_map, logger, - boundary_op_order_in_topological_sort, - recompute_subgraph_stores, p_node); - } - - // If there are other recompute plan for this node, we skip them because the graph is already modified. - if (!has_been_modified && recompute_with_compromise_subgraph_stores.ContainsSubGraphInstance(p_node)) { - has_been_modified = ModifyGraph(graph, node_index_to_its_order_in_topological_sort_map, - candidate_output_args_map, logger, - boundary_op_order_in_topological_sort, - recompute_with_compromise_subgraph_stores, p_node); + yield_op_order_in_topological_sort, + p_node, + node_to_opt_plan_map[p_node], + node_to_apply_context_map[p_node]); } modified = modified || has_been_modified; } - PrintSummary(recompute_subgraph_stores, recompute_with_compromise_subgraph_stores, logger); + PrintSummary(memory_opt_planner, node_to_apply_context_map, logger); return Status::OK(); } -void MemoryOptimizer::NodesInTopoOrderToString(const InlinedVector& nodes_in_topological_order, - std::string& subgraph_string_representation, - std::string& log_info) const { - std::ostringstream oss; - std::ostringstream subgraph_string_representation_oss; - size_t node_count = nodes_in_topological_order.size(); - for (size_t i = 0; i < node_count; ++i) { - if (i < node_count - 1) { // Ignore the last node. - oss << "(name:" << nodes_in_topological_order[i]->Name() << ", type:" << nodes_in_topological_order[i]->OpType() - << "),"; - } - - subgraph_string_representation_oss << nodes_in_topological_order[i]->OpType() << "+"; - } - - subgraph_string_representation = subgraph_string_representation_oss.str(); - log_info = oss.str(); - if (log_info.size() > 0) { - log_info = " with its precedent nodes: " + log_info; - } -} - -std::string MemoryOptimizer::UserConfigToString(const UserConfig& config) const { - std::string type_str; - switch (config.type) { - case OptimizationType::None: { - type_str = "Disabled"; - } break; - case OptimizationType::Recompute: { - type_str = "Recomputed"; - } break; - default: { - type_str = "Unknown"; - } break; - } - return type_str; -} - -void MemoryOptimizer::PrintSummary(const SubGraphStores& recompute_stores, - const SubGraphStores& recompute_with_compromise_stores, +void MemoryOptimizer::PrintSummary(const optimizer::memory_optimizer::MemoryOptimizationPlanner& memory_opt_planner, + const InlinedHashMap< + const Node*, + std::shared_ptr>& + node_to_apply_contexts_map, const logging::Logger& logger) const { - if (recompute_stores.SubGraphDescCount() == 0 && recompute_with_compromise_stores.SubGraphDescCount() == 0) { - return; - } - - std::ostringstream summary; - summary << "\nMemoryOptimizer Summary:\n"; - summary << "\tUser config:\n\t" << optimizer_config_ << "\n"; - summary << "\t=================================\n"; - - auto print_info_from_stores = [&summary, this](std::string store_name, const SubGraphStores& stores) { - summary << "\t########" << store_name << "########\n"; - for (auto subgraph_it = stores.subgraph_descs.begin(); subgraph_it != stores.subgraph_descs.end(); - ++subgraph_it) { - std::string freq_info; - if (subgraph_it->second.user_optimizer_config.type != OptimizationType::None) - freq_info = " (requested_count=" + std::to_string(subgraph_it->second.user_optimizer_config.requested_count) + - ", actual applied_count=" + - std::to_string(subgraph_it->second.applied_count) + ")"; - summary << "\tSubgraph: " << subgraph_it->first << "\n" - << "\t\tOptimizationType: " - << UserConfigToString(subgraph_it->second.user_optimizer_config) << freq_info << "\n" - << "\t\tPatterns: \n"; - for (auto shape_stat_it = subgraph_it->second.shape_str_frequency.begin(); - shape_stat_it != subgraph_it->second.shape_str_frequency.end(); - ++shape_stat_it) { - summary << "\t\t\tPatternShape:" << shape_stat_it->first << "\tFrequency:" << shape_stat_it->second << "\n"; - } - summary << "\t--------------------------------\n"; - } - summary << "\t=================================\n"; - }; - - print_info_from_stores("Recompute", recompute_stores); - print_info_from_stores("RecomputeWithCompromise", recompute_with_compromise_stores); - - LOGS(logger, INFO) << summary.str() << "\n"; + std::vector> records_grouped_by_node_cluster_id; + optimizer::memory_optimizer::GetMemoryRecordsGroupedByNodeClusterId(memory_opt_planner, + node_to_apply_contexts_map, + records_grouped_by_node_cluster_id); + LOGS(logger, INFO) << SerializeMemoryRecords(records_grouped_by_node_cluster_id, optimizer_config_) << "\n"; } /****************************************************** ** Recompute related function implementation starts ** ******************************************************/ -void MemoryOptimizer::RegisterAllowedRecomputeOps() { - if (static_cast(recompute_probe_level_) >= static_cast(ProbeLevel::Basic)) { - recomputable_op_type_to_input_arg_index_map_.insert({ - // Binary elementwise - {"Add", AllowedRecomputeNodeConfig{{0, 1}}}, - {"BiasGelu", AllowedRecomputeNodeConfig{{0, 1}}}, - {"Div", AllowedRecomputeNodeConfig{{0, 1}}}, - {"Mul", AllowedRecomputeNodeConfig{{0, 1}}}, - {"Sub", AllowedRecomputeNodeConfig{{0, 1}}}, - - // Data layout - /// The shape input is trivial whether it exists or not in backward. - {"Reshape", AllowedRecomputeNodeConfig{{0}}}, - {"Squeeze", AllowedRecomputeNodeConfig{{0}}}, - {"Unsqueeze", AllowedRecomputeNodeConfig{{0}}}, - - // Unary elementwise - /// The ratio and mode input are trivial whether they exist or not in backward - {"BitmaskDropout", AllowedRecomputeNodeConfig{{0}}}, - /// The axis input is trivial whether it exists or not in backward - {"CumSum", AllowedRecomputeNodeConfig{{0}}}, - {"Dropout", AllowedRecomputeNodeConfig{{0}}}, - {"Gelu", AllowedRecomputeNodeConfig{{0}}}, - {"FastGelu", AllowedRecomputeNodeConfig{{0}}}, - - // Ternary elementwise - {"Where", AllowedRecomputeNodeConfig{{0, 1, 2}}}, - - // Data copy - {"Tile", AllowedRecomputeNodeConfig{{0}}}, - {"Cast", AllowedRecomputeNodeConfig{{0}}}, - }); - } - - if (static_cast(recompute_probe_level_) >= static_cast(ProbeLevel::Advanced)) { - recomputable_op_type_to_input_arg_index_map_.insert({ - {"MatMul", AllowedRecomputeNodeConfig{{0, 1}}}, - {"FusedMatMul", AllowedRecomputeNodeConfig{{0, 1}}}, - {"Softmax", AllowedRecomputeNodeConfig{{0}}}, - {"BiasSoftmax", AllowedRecomputeNodeConfig{{0, 1}}}, - {"BiasSoftmaxDropout", AllowedRecomputeNodeConfig{{0, 1}}}, - }); - } -} - -Status MemoryOptimizer::SelectRecomputeSubgraph(const Node& entry_node, - const InlinedVector& node_output_index_candidates, - const ActivationUsedMap& fw_op_output_arg_used_map, - const InlinedHashMap& - node_index_to_its_order_in_topological_sort_map, - InlinedVector& nodes, - const logging::Logger& logger, - bool compromise_stashed_activation, - bool& can_compromise_stashed_activation) const { - can_compromise_stashed_activation = false; - - LOGS(logger, VERBOSE) << "Enter SelectRecomputeSubgraph for Node " << entry_node.Name() << "(" << entry_node.OpType() << ")"; - nodes.clear(); - - std::deque q; - for (auto output_index : node_output_index_candidates) { - q.push_back(NodeOutputPort(&entry_node, static_cast(output_index))); - } - - bool early_stop = false; - std::set visited_output_arg_set; - std::set visited_node_set; - - // For the initial activations in queue, they are stashed ones, so we do differently when scan the queue for them. - bool is_first_queue_scan = true; - while (nodes.size() < MAXIMUM_RECOMPUTE_NODE_COUNT && !q.empty() && !early_stop) { - // Loop all candidate NodeOutputPort, and find the next layer of input nodes. - size_t current_queue_size = q.size(); - for (size_t i = 0; i < current_queue_size; ++i) { - NodeOutputPort p = q.front(); - q.pop_front(); - const Node* curr_node = p.first; - - // Skip if the node output is already visited. - if (std::find(visited_output_arg_set.begin(), visited_output_arg_set.end(), p) != - visited_output_arg_set.end()) { - continue; - } - - visited_output_arg_set.insert({p}); - - // If the node already visited by from it's other output index, skip it. - if (visited_node_set.find(curr_node) != visited_node_set.end()) { - continue; - } - - visited_node_set.insert(curr_node); - - // Bottom-up search rules. - // If current op is entry output node (that generates stashed activations): - // 1. If the op is not in recomputable_op_type_to_input_arg_index_map_, skip it. - // Otherwise: - // If current op is in allowed list, check its input args, and append the producers' NodeOutputPorts to next_q. - // If current op is NOT in allowed list: - // 1). the output does not exist in backward, we cannot find a good solution for so, search terminates. - // 2). the output is used in backward, we don't need trace back further, continue searching. - auto op_recompute_config_it = recomputable_op_type_to_input_arg_index_map_.find(curr_node->OpType()); - auto cur_output_arg_name = curr_node->OutputDefs()[p.second]->Name(); - if (is_first_queue_scan) { - // We handle the entry node outputs differently because, we don't want this case falls into and succeed one of - // the checks in the other branch - // 1. "op is not in recompute op list, but its output is used in backward" - // 2. "op is in recompute op list, but its output is used in backward" - // (either of the above checks is true for entry node outputs) - if (op_recompute_config_it == recomputable_op_type_to_input_arg_index_map_.end()) { - early_stop = true; - LOGS(logger, VERBOSE) << "Entry Node " << curr_node->Name() << "(" << curr_node->OpType() << ") is **NOT** " - << "in recompute op list, search terminates."; - break; - } - } else { - if (op_recompute_config_it == recomputable_op_type_to_input_arg_index_map_.end()) { - if (fw_op_output_arg_used_map.at(cur_output_arg_name).second) { - LOGS(logger, VERBOSE) << "Node " << curr_node->Name() << "(" << curr_node->OpType() << ") is **NOT** in " - << "recompute op list, but its output [" << cur_output_arg_name << "] is used in " - << "backward, we don't need trace bottom-up further. Entry node: " - << entry_node.Name() << "(" << entry_node.OpType() << ")"; - continue; - } else { - early_stop = true; - LOGS(logger, VERBOSE) << "Node " << curr_node->Name() << "(" << curr_node->OpType() << ") is **NOT** in " - << "recompute op list, and its output [" << cur_output_arg_name - << "] does not exist in backward, search terminates. Entry node: " - << entry_node.Name() << "(" << entry_node.OpType() << ")"; - break; - } - } - - if (fw_op_output_arg_used_map.at(cur_output_arg_name).second) { - LOGS(logger, VERBOSE) << "Node " << curr_node->Name() << "(" << curr_node->OpType() << ") " - << "is in recompute op list, while its output [" << cur_output_arg_name - << "] is used in backward, we don't need trace bottom-up further. Entry node: " - << entry_node.Name() << "(" << entry_node.OpType() << ")"; - continue; - } - } - - // Append node to the selected graph. - if (std::find(nodes.begin(), nodes.end(), curr_node) == nodes.end()) { - nodes.push_back(curr_node); - LOGS(logger, VERBOSE) << "Node " << curr_node->Name() << "(" << curr_node->OpType() - << ") is added in selected subgraph "; - } - - // This check is not matured now, subject to be changed. - float ratio = InputOutputSizeRatio(curr_node); - float is_current_node_compromisable = (ratio < 1.f); - can_compromise_stashed_activation = can_compromise_stashed_activation || is_current_node_compromisable; - if (is_current_node_compromisable) { - LOGS(logger, VERBOSE) << "Node " << curr_node->Name() << "(" << curr_node->OpType() - << ") has input/output size " << ratio << " < 1.f, can compromise stashed activation"; - } - - if (is_current_node_compromisable && compromise_stashed_activation) { - LOGS(logger, VERBOSE) << "Node " << curr_node->Name() << "(" << curr_node->OpType() << ") is in " - << "recompute op list, and its output [" << cur_output_arg_name - << "] does not exist in backward, while it meet compromised check, we don't need trace " - << "bottom-up further."; - continue; - } - - // Iterate all input nodes according to allowed input arg index of the entry node. - const auto& input_arg_indices = op_recompute_config_it->second.input_arg_indices; - for (auto it = curr_node->InputEdgesBegin(), end = curr_node->InputEdgesEnd(); it != end; ++it) { - const Node::EdgeEnd& input_edge = *it; - const auto& parent_node = input_edge.GetNode(); - const auto parent_node_output_index = input_edge.GetSrcArgIndex(); - const auto current_node_input_index = input_edge.GetDstArgIndex(); - if (std::find(input_arg_indices.begin(), input_arg_indices.end(), current_node_input_index) != - input_arg_indices.end()) { - NodeOutputPort next_p = std::make_pair(&parent_node, parent_node_output_index); - - LOGS(logger, VERBOSE) << "Node " << parent_node.Name() << "(" << parent_node.OpType() << ")'s " - << parent_node_output_index - << "th output [" << parent_node.OutputDefs()[parent_node_output_index]->Name() - << "] is added in recompute search list "; - - q.push_back(next_p); - } - } - } - // After handle all entry node outputs, we set the flag to false. - is_first_queue_scan = false; - } - - // If input args are not found in bw, but op count exceed MAXIMUM_RECOMPUTE_NODE_COUNT, skip recompute. - if (!q.empty() || early_stop) { - LOGS(logger, VERBOSE) << "Fail to find a solution for recompute: current node count is " << nodes.size() - << ", queue size: " << q.size() << ", early stop: " << early_stop; - nodes.clear(); - } else { - // Re-order the nodes in topological order. - std::sort(nodes.begin(), nodes.end(), - [&node_index_to_its_order_in_topological_sort_map](const Node*& lhs, const Node*& rhs) { - return node_index_to_its_order_in_topological_sort_map.at(lhs->Index()) < - node_index_to_its_order_in_topological_sort_map.at(rhs->Index()); - }); - } - return Status::OK(); -} - -void MemoryOptimizer::CheckNodeForRecompute(const Node& node, - const ActivationUsedMap& fw_op_output_arg_used_map, - const InlinedHashMap& - node_index_to_its_order_in_topological_sort_map, - const InlinedHashMap>& - candidate_output_args_map, - SubGraphStores& subgraph_stores, - const logging::Logger& logger, - bool compromise_stashed_activation, - bool& can_compromise_stashed_activation) const { - if (recomputable_op_type_to_input_arg_index_map_.find(node.OpType()) == - recomputable_op_type_to_input_arg_index_map_.end()) { - return; - } - - InlinedVector nodes_in_topological_order; - ORT_ENFORCE(SelectRecomputeSubgraph(node, candidate_output_args_map.at(&node), - fw_op_output_arg_used_map, - node_index_to_its_order_in_topological_sort_map, - nodes_in_topological_order, logger, - compromise_stashed_activation, - can_compromise_stashed_activation) - .IsOK()); - if (nodes_in_topological_order.size() == 0) { - return; - } - - std::string subgraph_str_representation, log_info; - NodesInTopoOrderToString(nodes_in_topological_order, subgraph_str_representation, log_info); - LOGS(logger, VERBOSE) << "Node " << node.Name() << "(" << node.OpType() << ") can be recomputed" << log_info; - - // Update the subgraph optimization config map - key is the subgraph string representation, value is user config. - UserConfig user_config{OptimizationType::None, 0}; - if (pattern_subgraph_to_user_optimizer_config_map_.find(subgraph_str_representation) != - pattern_subgraph_to_user_optimizer_config_map_.end()) { - user_config = pattern_subgraph_to_user_optimizer_config_map_.at(subgraph_str_representation); - } - - SubGraphDesc& subgraph_desc = - subgraph_stores.Contains(subgraph_str_representation) - ? subgraph_stores.GetSubGraphDesc(subgraph_str_representation) - : subgraph_stores.CreateSubGraphDesc(subgraph_str_representation, user_config); - - subgraph_desc.total_frequency += 1; - - // Update the subgraph frequency map - key is the subgraph string representation, value is number of appearances. - for (size_t output_index : candidate_output_args_map.at(&node)) { - auto shape_str = TensorShapeProtoToString(node.OutputDefs()[output_index]->Shape()); - subgraph_desc.shape_str_frequency[shape_str]++; - } - - subgraph_stores.AddSubGraphInstance(&node, nodes_in_topological_order, subgraph_desc); - - return; -} - Status MemoryOptimizer::CreateRecomputeGraph(Graph& graph, const InlinedVector& nodes_in_topological_order, Node*& new_output_node_ptr) const { @@ -716,8 +217,8 @@ Status MemoryOptimizer::CreateRecomputeGraph(Graph& graph, // Check whether the node has been recomputed/offloaded or not. Simply check the existence of the first output // of the node has its corresponding recompute name or not. - // TODO: if there is more optimization types like offload added, we will add corresponding check whether the outputs - // already be offloaded or not. + // TODO: if there is more optimization types like offload added, we will add a corresponding check + // whether the outputs already be offloaded or not. if (graph.GetNodeArg(graph_utils::RecomputeName(node_to_duplicate->MutableOutputDefs()[0]->Name())) != nullptr) { continue; } diff --git a/orttraining/orttraining/core/optimizer/memory_optimizer.h b/orttraining/orttraining/core/optimizer/memory_optimizer.h index 1d21c9143f..13eb4cdb24 100644 --- a/orttraining/orttraining/core/optimizer/memory_optimizer.h +++ b/orttraining/orttraining/core/optimizer/memory_optimizer.h @@ -2,163 +2,39 @@ // Licensed under the MIT License. #pragma once -#include + #include "core/common/inlined_containers.h" #include "core/common/string_utils.h" #include "core/optimizer/graph_transformer.h" +#include "orttraining/core/optimizer/memory_optimizer/common.h" +#include "orttraining/core/optimizer/memory_optimizer/optimization_planner.h" +#include "orttraining/core/optimizer/memory_optimizer/recompute_analysis.h" +#include "orttraining/core/optimizer/memory_optimizer/memory_insight.h" namespace onnxruntime { /** @Class MemoryOptimizer -Find recomputable subgraphs and enable according to user configs. +(TODO) move to orttraining/orttraining/core/optimizer/memory_optimizer/ folder. + +Find recompute subgraphs and enable them according to user configs. The way we collect subgraphs +(in orttraining/orttraining/core/optimizer/memory_optimizer/recompute_analysis.h) in brief is: +1. Find all nodes that generate stashed activations. +2. For each node, check it data type is supported to recompute + a. If yes, add it in the subgraph, and append its input in the queue to scan next; + b. otherwise, stop collecting and return the subgraph (could be empty). +3. Pick up the input node from the queue, and do 2 again. The process ends when the queue is empty or 2.b happens. +4. Clone the recomputable subgraphs with lower node priority (to execute) and insert them back to the original graph. */ class MemoryOptimizer : public GraphTransformer { private: - using NodeOutputPort = std::pair; - using ActivationUsedMap = InlinedHashMap>; - - /** - * @brief Level to control allowed operations during subgraph detecting. - * Level 0: only allow cheap-to-compute operations. - * Level 1: allow more expensive operations. - */ - enum class ProbeLevel { - Basic = 0, - Advanced = 1, - LevelMax = 2, - }; - - /** - * @brief Type of memory reduction techniques. - */ - enum class OptimizationType { - None = 0, // Disabled. - Recompute = 1, - TypeMax = 2, - }; - - /** - * @brief Type of user config. - * type: type of memory reduction techniques. - * requested_count: the number of occurrences of a subgraph pattern for alleviation. -1 means apply all. - * One example: if a subgraph pattern is found 3 times, and requested_count is set 2, then the 1st and 2nd subgraph - * in topological order will be applied for alleviation. This is useful to avoid alleviating more memory than - * needed. - */ - struct UserConfig { - OptimizationType type; - int requested_count; - }; - - /** - * @brief Struct to store properties of a specific subgraph. - */ - struct SubGraphDesc { - SubGraphDesc() = default; - - // A string to represent the subgraph, used as a unique "ID" for a unique subgraph. - std::string subgraph_representative_str; - - InlinedHashMap shape_str_frequency; // shape string to frequency - UserConfig user_optimizer_config; - int total_frequency{0}; // The occurrence of this subgraph pattern in the graph. - - int applied_count{0}; // The number of times this subgraph pattern has been really applied in this transformer. - int skip_count{0}; // The number of times this subgraph instance has been skipped in reversed topological order. - float saving_ratio{1.0f}; // For compromised memory saving, the ratio of memory saving. - }; - - /** - * @brief A struct to maintain the information of target subgraphs to optimize. - * Imagine we loop all nodes finding recomputable/offload-able subgraphs, we want to store them first. - * Afterwards, we optionally pick up some of them to apply optimization according to user configs. - * - * subgraph_descs is a map from subgraph string representation to its subgraph related configurations. - * - * _optimization_target_graphs_ is a map from activation producer node pointers to its target optimization subgraph - * nodes. For example, if a subgraph Cast+Gelu can be recomputed, we may have a map like: - * key: node pointer of stashed activation producer Gelu; value: node vector {Cast, Gelu,}. - * - * When we AddSubGraphInstance, we must provider its corresponding subgraph desc in the parameter. - * Then we can know for each subgraph instance, what's the subgraph str representation, and what's the optimization - * config. - */ - struct SubGraphStores { - /********************************** - ** subgraph desc section starts ** - **********************************/ - - size_t SubGraphDescCount() const { - return subgraph_descs.size(); - } - - bool Contains(std::string_view subgraph_str) const { - return subgraph_descs.find(subgraph_str) != subgraph_descs.end(); - } - - SubGraphDesc& GetSubGraphDesc(std::string_view subgraph_string) { - ORT_ENFORCE(Contains(subgraph_string), "Subgraph string not found.", subgraph_string); - return subgraph_descs.at(subgraph_string); - } - - SubGraphDesc& CreateSubGraphDesc(const std::string& subgraph_string, - UserConfig& config) { - ORT_ENFORCE(!Contains(subgraph_string), "Subgraph string already exists.", subgraph_string); - subgraph_descs[subgraph_string].user_optimizer_config = config; - subgraph_descs[subgraph_string].subgraph_representative_str = subgraph_string; - return subgraph_descs[subgraph_string]; - } - - /********************************************************************** - ** subgraph desc section ends, and subgraph instance section starts. ** - ***********************************************************************/ - - // Pair of . - using GraphInstanceInfo = std::pair, std::string>; - - void AddSubGraphInstance(const Node* node, - const InlinedVector& nodes_in_topological_order, - const SubGraphDesc& subgraph_desc) { - ORT_ENFORCE(_optimization_target_graphs_.find(node) == _optimization_target_graphs_.end()); - _optimization_target_graphs_[node] = std::make_pair(nodes_in_topological_order, - subgraph_desc.subgraph_representative_str); - } - - bool ContainsSubGraphInstance(const Node* node) const { - return _optimization_target_graphs_.find(node) != _optimization_target_graphs_.end(); - } - - GraphInstanceInfo& GetSubGraphInstance(const Node* node) { - ORT_ENFORCE(_optimization_target_graphs_.find(node) != _optimization_target_graphs_.end()); - return _optimization_target_graphs_[node]; - } - - /*********************************** - ** subgraph instance section ends ** - ***********************************/ - - InlinedHashMap subgraph_descs; - InlinedHashMap _optimization_target_graphs_; - }; - - /** - * @brief Used to define per-op recompute config. - * - */ - struct AllowedRecomputeNodeConfig { - InlinedVector input_arg_indices; // input index to iterate further (bottom up) - }; - public: - MemoryOptimizer(const std::string& enable_memory_optimizer, const std::string& level) + MemoryOptimizer(const std::string& memory_optimizer_config, const std::string& level) : GraphTransformer("MemoryOptimizer") { // Parse user defined configs. - ORT_ENFORCE(ParseConfigFromString(enable_memory_optimizer, level).IsOK()); - - RegisterAllowedRecomputeOps(); + ORT_ENFORCE(ParseConfigFromString(memory_optimizer_config, level).IsOK()); } Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; @@ -166,35 +42,7 @@ class MemoryOptimizer : public GraphTransformer { bool ShouldOnlyApplyOnce() const override { return true; } private: - Status ParseConfigFromString(const std::string& enable_memory_optimizer, const std::string& level); - - /** - * @brief Prepare info including activation usage, node usage in fw and bw. - * - * @param graph Graph to iterate. - * @param fw_op_output_arg_used_map Collected activation usage mapping. - * - key: node arg name - * - value: a pair of bool, representing whether the activation is used by forward nodes or by backward nodes. - * @return int64_t value The boundary op (for example YieldOp) order in topological order. If no boundary op found, - * return -1; - */ - int64_t PrepareForTransformation(const Graph& graph, - ActivationUsedMap& fw_op_output_arg_used_map, - InlinedHashMap& - node_index_to_its_order_in_topological_sort_map) const; - /** - * @brief Find all stashed activations, e.g. activations used by forward operators and backward operators. - * - * @param graph Graph to iterate. - * @param fw_op_output_arg_used_map Activation usage mapping. - * @param candidate_output_args_map Candidate activations, which are consumed by both fw and bw ops. - * @return Status - */ - Status GetStashedActivationCandidates( - const Graph& graph, - const InlinedHashMap>& fw_op_output_arg_used_map, - InlinedHashMap>& candidate_output_args_map, - const logging::Logger& logger) const; + Status ParseConfigFromString(const std::string& memory_optimizer_config, const std::string& level); /** * @brief Apply graph modifications based on user configs. @@ -212,28 +60,15 @@ class MemoryOptimizer : public GraphTransformer { * @return false */ bool ModifyGraph(Graph& graph, - const InlinedHashMap& node_index_to_its_order_in_topological_sort_map, - const InlinedHashMap>& candidate_output_args_map, + const InlinedHashMap& + node_index_to_its_order_in_topological_sort_map, + const InlinedHashMap>& + candidate_output_args_map, const logging::Logger& logger, - int64_t boundary_op_order_in_topological_sort, - SubGraphStores& subgraph_stores, - Node* node) const; - - /** - * @brief Convert the recompute subgraph to its string representation. - * - * @param nodes_in_topological_order The subgraph nodes in topological order. - * @param subgraph_string_representation Returns subgraph string representation. - * @param log_info Returns log info for users. - */ - void NodesInTopoOrderToString(const InlinedVector& nodes_in_topological_order, - std::string& subgraph_string_representation, - std::string& log_info) const; - - /** - * @brief Convert optimization type to string. - */ - std::string UserConfigToString(const UserConfig& config) const; + ptrdiff_t boundary_op_order_in_topological_sort, + Node* node, + std::shared_ptr& node_plan, + std::shared_ptr& apply_context) const; /** * @brief Summarize transformation details. @@ -241,72 +76,16 @@ class MemoryOptimizer : public GraphTransformer { * @param stashed_activation_statistics statistics around stashed activation memory saving. * @return void */ - void PrintSummary(const SubGraphStores& recompute_stores, - const SubGraphStores& recompute_with_compromise_stores, + void PrintSummary(const optimizer::memory_optimizer::MemoryOptimizationPlanner& mem_opt_stats, + const InlinedHashMap>& + node_to_apply_contexts_map, const logging::Logger& logger) const; /************************************************** ** Recompute related function definition starts ** *************************************************/ - void RegisterAllowedRecomputeOps(); - - /** - * @brief Find recomputable subgraphs (has at least one nodes, at most MAXIMUM_RECOMPUTE_NODE_COUNT nodes). - * - * @param node The entry node to start the subgraph matching (bottom-up), usually the last node of found subgraphs. - * @param node_output_index_candidates Candidate output indices of "node", which are consumed by both fw and bw ops. - * @param fw_op_output_arg_used_map The activation usage (in fw and bw) mapping. - * @param node_index_to_its_order_in_topological_sort_map The mapping of node index to its order in topological sort. - * Used to re-order the collected subgraph nodes. - * @param nodes_in_topological_order Collected vector of nodes of found subgraph, in the order of the topological - * sorted. - * @param logger Logger. - * @param compromise_stashed_activation Whether to compromise stashed activation, e.g. if we cannot find a - * recomputable subgraph to save a stashed activation, we can compromise to find a recomputable subgraph to reduce the - * size of stashed activation. - * @param can_compromise_stashed_activation A bool return value, to indicate there is opportunaties for finding a - * compromised subgraph. - * @return Status - */ - Status SelectRecomputeSubgraph(const Node& node, - const InlinedVector& node_output_index_candidates, - const ActivationUsedMap& fw_op_output_arg_used_map, - const InlinedHashMap& - node_index_to_its_order_in_topological_sort_map, - InlinedVector& nodes_in_topological_order, - const logging::Logger& logger, - bool compromise_stashed_activation, - bool& can_compromise_stashed_activation) const; - - /** - * @brief For the node producing stashed activation, check whether a recomputable subgraph can be found or not. - * - * @param node The entry node to start the subgraph matching (bottom-up), usually the last node of found subgraphs. - * @param fw_op_output_arg_used_map The activation usage (in fw and bw) mapping. - * @param node_index_to_its_order_in_topological_sort_map The mapping of node index to its order in topological sort. - * Used to re-order the collected subgraph nodes. - * @param candidate_output_args_map A map from node to its candidate activations, which are consumed by both fw and - * bw ops. - * @param subgraph_stores A store to maintain all found subgraphs. - * @param logger Logger. - * @param compromise_stashed_activation Whether to compromise stashed activation, e.g. if we cannot find a - * recomputable subgraph to save a stashed activation, we can compromise to find a recomputable subgraph to reduce the - * size of stashed activation. - * @param can_compromise_stashed_activation A bool return value, to indicate there is opportunaties for finding a - * compromised subgraph. - */ - void CheckNodeForRecompute(const Node& node, - const ActivationUsedMap& fw_op_output_arg_used_map, - const InlinedHashMap& - node_index_to_its_order_in_topological_sort_map, - const InlinedHashMap>& - candidate_output_args_map, - SubGraphStores& subgraph_stores, - const logging::Logger& logger, - bool compromise_stashed_activation, - bool& can_compromise_stashed_activation) const; - /** * @brief Duplicate nodes to create a recompute subgraph. * @@ -323,12 +102,10 @@ class MemoryOptimizer : public GraphTransformer { ** Recompute related function definition ends ** *************************************************/ - // The op types that are supported predefined. - InlinedHashMap recomputable_op_type_to_input_arg_index_map_; // User enabled map of the subgraph string representation to the alleviation type. - InlinedHashMap pattern_subgraph_to_user_optimizer_config_map_; + InlinedHashMap pattern_subgraph_to_user_optimizer_config_map_; std::string optimizer_config_; - ProbeLevel recompute_probe_level_; + optimizer::memory_optimizer::ProbeLevel recompute_probe_level_; }; } // namespace onnxruntime diff --git a/orttraining/orttraining/core/optimizer/memory_optimizer/common.cc b/orttraining/orttraining/core/optimizer/memory_optimizer/common.cc new file mode 100644 index 0000000000..2291d7e4f3 --- /dev/null +++ b/orttraining/orttraining/core/optimizer/memory_optimizer/common.cc @@ -0,0 +1,149 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include + +#include "orttraining/core/optimizer/memory_optimizer/common.h" +#include "core/graph/graph_utils.h" +#include "core/optimizer/utils.h" +#include "core/graph/graph_viewer.h" +#include "core/framework/tensorprotoutils.h" + +#include "core/common/string_utils.h" + +namespace onnxruntime::optimizer::memory_optimizer { + +namespace { + +constexpr const char empty_dim_param_placeholder[] = "empty_dim_param"; +static size_t index_empty_dim = 0; + +bool TensorShapeProtoToDimParamVector(const ONNX_NAMESPACE::TensorShapeProto* shape, + std::vector& dim_params) { + bool has_unknown_dim = false; + for (int dim_index = 0; dim_index < shape->dim_size(); dim_index++) { + auto dim = shape->dim(dim_index); + if (utils::HasDimValue(dim)) { + dim_params.push_back(std::to_string(dim.dim_value())); + } else { + std::string trimmed_dim_param = utils::TrimString(dim.dim_param()); + if (trimmed_dim_param.empty()) { + has_unknown_dim = true; + dim_params.push_back(empty_dim_param_placeholder + std::to_string(index_empty_dim++)); + } else { + dim_params.push_back(trimmed_dim_param); + } + } + } + + if (shape->dim_size() == 0) { + dim_params.push_back("(1)"); // Scalar + } + + return has_unknown_dim; +} + +bool HasUnknowDimension(const ONNX_NAMESPACE::TensorShapeProto* shape) { + if (shape == nullptr) { + return true; + } + + std::vector dim_params; + return TensorShapeProtoToDimParamVector(shape, dim_params); +} + +std::string TensorShapeProtoToString(const ONNX_NAMESPACE::TensorShapeProto* shape) { + if (shape == nullptr) { + return "unknown"; + } + + std::vector dim_params; + TensorShapeProtoToDimParamVector(shape, dim_params); + + std::ostringstream oss; + oss << "("; + for (auto it = dim_params.begin(); it != dim_params.end(); ++it) { + oss << "(" << *it << ")"; + if (it != (dim_params.end() - 1)) { + oss << "*"; + } + } + oss << ")"; + + return oss.str(); +} + +} // namespace + +std::string GetTensorElemCountInSymbolicString(const Node* node, size_t output_index) { + const auto& output_def = node->OutputDefs()[output_index]; + const auto shape = output_def->Shape(); + + std::string shape_str = TensorShapeProtoToString(shape); + + // If the output shape contains unknown dimension, we try to get the shape from input. + // though the input shape might be different, but its elem size and count should be the same + // with the output. + if (node->OpType() == "Reshape" && HasUnknowDimension(shape) && + !HasUnknowDimension(node->InputDefs()[0]->Shape())) { + shape_str = TensorShapeProtoToString(node->InputDefs()[0]->Shape()); + } + + return shape_str; +} + +std::string OptimizationTypeToString(OptimizationType type) { + switch (type) { + case OptimizationType::None: + return "None"; + case OptimizationType::Recompute: + return "Recompute"; + case OptimizationType::RecomputeWithCompromise: + return "RecomputeWithCompromise"; + default: + ORT_THROW("Unknown optimization type."); + } +} + +int ParseIntValueFromString(std::string_view str) { + int int_value = 0; + auto result = std::from_chars(str.data(), str.data() + str.size(), int_value); + ORT_ENFORCE(result.ec != std::errc::invalid_argument, "Fail to convert to int from string: ", str); + return int_value; +} + +Status ParseConfigFromString(std::string_view memory_optimization_config, + InlinedHashMap& cluster_id_to_config_map) { + if (!memory_optimization_config.empty()) { + const auto user_config_strs = utils::SplitString(memory_optimization_config, ","); + for (const auto& user_config_str : user_config_strs) { + const auto user_config = utils::SplitString(user_config_str, ":"); + ORT_RETURN_IF_NOT(user_config.size() == 3, + "User config should be in format of SubgraphStr:OptimizationType:RequestApplyCount."); + + const std::string subgraph_string_representation(user_config[0]); + int optimization_type_int = ParseIntValueFromString(user_config[1]); + int requested_apply_count = ParseIntValueFromString(user_config[2]); + ORT_RETURN_IF_NOT(optimization_type_int < + static_cast(OptimizationType::TypeMax) && + optimization_type_int >= 0, + "Invalid optimization type specified for subgraph: ", + subgraph_string_representation); + + ORT_RETURN_IF_NOT(requested_apply_count == -1 || requested_apply_count >= 0, + "Invalid requested_apply_count specified for subgraph: ", requested_apply_count); + + // At this point, subgraph_string_representation is a pattern graph string representation. + // If duplicated subgraph_string_representation is found in user config, the last one will be used. + cluster_id_to_config_map[subgraph_string_representation] = UserConfig{ + static_cast(optimization_type_int), + requested_apply_count}; + } + } + + return Status::OK(); +} + +} // namespace onnxruntime::optimizer::memory_optimizer diff --git a/orttraining/orttraining/core/optimizer/memory_optimizer/common.h b/orttraining/orttraining/core/optimizer/memory_optimizer/common.h new file mode 100644 index 0000000000..85e2bf4f5d --- /dev/null +++ b/orttraining/orttraining/core/optimizer/memory_optimizer/common.h @@ -0,0 +1,76 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include + +#include "core/common/common.h" +#include "core/common/logging/logging.h" +#include "core/common/inlined_containers_fwd.h" +#include "core/graph/basic_types.h" +#include "core/framework/data_types.h" +#include "core/graph/graph_viewer.h" + +namespace onnxruntime::optimizer::memory_optimizer { + +// Uncomment for debugging Memory optimizer (MO). +// #define MO_NEED_LOG_DEBUG_INFO 1 + +#ifndef MO_LOG_DEBUG_INFO +#ifdef MO_NEED_LOG_DEBUG_INFO +#define MO_LOG_DEBUG_INFO(logger, message) LOGS(logger, WARNING) << message +#else +#define MO_LOG_DEBUG_INFO(logger, message) \ + ORT_UNUSED_PARAMETER(logger); \ + do { \ + } while (0) +#endif +#endif + +using NodeOutputPort = std::pair; +using ActivationUsedMap = InlinedHashMap>; + +/** + * @brief Type of memory reduction techniques. + */ +enum class OptimizationType { + None = 0, // Disabled. + Recompute = 1, + RecomputeWithCompromise = 2, + TypeMax = 3, +}; + +std::string OptimizationTypeToString(OptimizationType type); + +/** + * @brief Type of user config. + * type: type of memory reduction techniques. + * requested_count: the number of occurrences of a subgraph pattern for alleviation. -1 means apply all. + * One example: if a subgraph pattern is found 3 times, and requested_count is set 2, then the 1st and 2nd subgraph + * in topological order will be applied for alleviation. This is useful to avoid alleviating more memory than + * needed. + */ +struct UserConfig { + OptimizationType type; + int requested_count; +}; + +/** + * @brief Get total element count inn format of a symbolic string. + * + * @param node The node to get element count. + * @param output_index The output index of the node. + * @return std::string + */ +std::string GetTensorElemCountInSymbolicString(const Node* node, size_t output_index); + +int ParseIntValueFromString(std::string_view str); + +Status ParseConfigFromString(std::string_view memory_optimization_config, + InlinedHashMap& cluster_id_to_config_map); + +} // namespace onnxruntime::optimizer::memory_optimizer diff --git a/orttraining/orttraining/core/optimizer/memory_optimizer/memory_insight.cc b/orttraining/orttraining/core/optimizer/memory_optimizer/memory_insight.cc new file mode 100644 index 0000000000..60f62a9881 --- /dev/null +++ b/orttraining/orttraining/core/optimizer/memory_optimizer/memory_insight.cc @@ -0,0 +1,763 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include +#include + +#include "core/graph/graph_utils.h" +#include "core/graph/graph_viewer.h" +#include "orttraining/core/optimizer/memory_optimizer/common.h" +#include "orttraining/core/optimizer/memory_optimizer/optimization_planner.h" +#include "orttraining/core/optimizer/memory_optimizer/recompute_analysis.h" +#include "orttraining/core/optimizer/memory_optimizer/memory_insight.h" + +namespace onnxruntime::optimizer::memory_optimizer { + +// Placeholder string for table row separator, which is used to be replaced by table row separator finally. +constexpr const char kTableRowSeparator[] = "TABLE_SEPARATOR_PLACEHOLDER"; +// Placeholder string for table border, which is used to be replaced by table border finally. +constexpr const char kTableBorder[] = "TABLE_BORDER_PLACEHOLDER"; + +// The max length of the first column in the table. +constexpr const int kFirstColumnWidth = 7; +// The max length of left part (e.g. title) in the second column. +constexpr const int kTitleWidthInSecondColumn = 15; + +/** + * @brief Prepare info including activation usage, node usage in fw and bw. + * + * @param graph Graph to iterate. + * @param boundary_op_order_in_topological_sort index of the boundary op between fw and bw. + * @param node_index_to_its_order_in_topological_sort_map The mapping of node index to its order in topological sort. + * @param fw_op_output_arg_used_map Collected activation usage mapping. + * - key: node arg name + * - value: a pair of bool, representing whether the activation is used by forward nodes or by backward nodes. + * @param is_forward_nodes Collected node is forward pass op mapping. + */ +void GetForwardOutputUsageMap(const GraphViewer& graph_viewer, + const ptrdiff_t boundary_op_order_in_topological_sort, + const InlinedHashMap& + node_index_to_its_order_in_topological_sort_map, + ActivationUsedMap& fw_op_output_arg_used_map, + InlinedHashMap& is_forward_nodes) { + ORT_ENFORCE(boundary_op_order_in_topological_sort >= 0); + const auto& node_ids = graph_viewer.GetNodesInTopologicalOrder(); + is_forward_nodes.clear(); + is_forward_nodes.reserve(node_ids.size()); + + auto is_forward_pass_operator = [](ptrdiff_t op_order_in_topological_sort, + ptrdiff_t boundary_op_order_in_topological_sort) -> bool { + return op_order_in_topological_sort <= boundary_op_order_in_topological_sort; + }; + + fw_op_output_arg_used_map.clear(); + fw_op_output_arg_used_map.reserve(node_ids.size()); + for (size_t i = 0; i < node_ids.size(); ++i) { + const Node* p_node = graph_viewer.GetNode(node_ids[i]); + if (p_node == nullptr /* skip removed nodes*/) { + continue; + } + + const Node& node = *p_node; + + bool is_forward_op = is_forward_pass_operator(static_cast(i), boundary_op_order_in_topological_sort); + if (!is_forward_op) { + is_forward_nodes[p_node] = false; + continue; + } + + is_forward_nodes[p_node] = true; + + for (auto& output_arg : node.OutputDefs()) { + if (!output_arg->Exists() || output_arg->Name().empty()) { + continue; + } + + bool used_in_fw = false; + bool used_in_bw = false; + for (auto& consumer_node : graph_viewer.GetConsumerNodes(output_arg->Name())) { + ORT_ENFORCE(consumer_node != nullptr, "Consumer node should not be null."); + auto it = node_index_to_its_order_in_topological_sort_map.find(consumer_node->Index()); + ORT_ENFORCE(it != + node_index_to_its_order_in_topological_sort_map.end(), + "Consumer node should be in topological order map."); + size_t consumer_node_index_in_topological_order = it->second; + if (is_forward_pass_operator(static_cast(consumer_node_index_in_topological_order), + boundary_op_order_in_topological_sort)) { + used_in_fw = true; + } else { + used_in_bw = true; + } + } + + ORT_ENFORCE(fw_op_output_arg_used_map.find(output_arg->Name()) == fw_op_output_arg_used_map.end(), + "Duplicated output arg found named: ", output_arg->Name()); + fw_op_output_arg_used_map.insert({{output_arg->Name(), std::make_pair(used_in_fw, used_in_bw)}}); + } + } +} + +/** + * @brief Find all stashed activations, e.g. activations used by forward operators and backward operators. + * + * @param graph_viewer Graph to iterate. + * @param boundary_op_order_in_topological_sort The order of the boundary op in the topological sort. + * @param fw_op_output_arg_used_map Activation usage mapping. + * @param candidate_output_args_map Candidate activations, which are consumed by both fw and bw ops. + * @param is_forward_nodes Whether a node is a forward node. + * @param logger Logger. + * @return Status + */ + +Status GetStashedActivationCandidates(const GraphViewer& graph_viewer, + const ptrdiff_t boundary_op_order_in_topological_sort, + ActivationUsedMap& fw_op_output_arg_used_map, + InlinedHashMap>& + candidate_output_args_map, + InlinedHashMap& is_forward_nodes, + const logging::Logger& logger) { + if (boundary_op_order_in_topological_sort < 0) { + LOGS(logger, VERBOSE) << "No boundary op found. Skip memory optimization."; + return Status::OK(); + } + + const auto& node_ids = graph_viewer.GetNodesInTopologicalOrder(); + + InlinedHashMap node_index_to_its_order_in_topological_sort_map; + for (size_t i = 0; i < node_ids.size(); ++i) { + const Node* p_node = graph_viewer.GetNode(node_ids[i]); + if (p_node == nullptr) { /* skip removed nodes*/ + continue; + } + + node_index_to_its_order_in_topological_sort_map[p_node->Index()] = i; + } + + GetForwardOutputUsageMap(graph_viewer, boundary_op_order_in_topological_sort, + node_index_to_its_order_in_topological_sort_map, + fw_op_output_arg_used_map, + is_forward_nodes); + + for (auto& kv : fw_op_output_arg_used_map) { + // used by fw and bw, then it is a candidate. + if (kv.second.first && kv.second.second) { + const Node* n = graph_viewer.GetProducerNode(kv.first); + ORT_ENFORCE(n, "Activation should have a producer node"); + size_t k = 0; + for (k = 0; k < n->OutputDefs().size(); ++k) { + if (n->OutputDefs()[k]->Name().compare(kv.first) == 0) { + break; + } + } + + if (std::find(candidate_output_args_map[n].begin(), candidate_output_args_map[n].end(), k) != + candidate_output_args_map[n].end()) { + ORT_ENFORCE(false, "Duplicated candidate output found."); + } + + candidate_output_args_map[n].push_back(k); + LOGS(logger, VERBOSE) << "Find candidate output named [" << kv.first << "] of Node " << n->Name() << "(" + << n->OpType() << ")"; + } + } + + return Status::OK(); +} + +Status FindORTModuleMemoryOpportunity(const GraphViewer& graph_viewer, + const ProbeLevel probe_level, + const logging::Logger& logger, + InlinedHashMap& + node_index_to_its_order_in_topological_sort_map, + ptrdiff_t& yield_op_order_in_topological_sort, + InlinedHashMap>& + candidate_output_args_map, + MemoryOptimizationPlanner& memory_opt_planner) { + const auto& node_ids = graph_viewer.GetNodesInTopologicalOrder(); + + // Find boundary ops between forward and backward pass, currently, it's limited to YieldOp. + yield_op_order_in_topological_sort = -1; + for (size_t i = 0; i < node_ids.size(); ++i) { + const Node* p_node = graph_viewer.GetNode(node_ids[i]); + if (p_node == nullptr) { /* skip removed nodes*/ + continue; + } + + if (p_node->OpType() == "YieldOp") { + if (yield_op_order_in_topological_sort != -1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "There are multiple YieldOps in the graph, node: ", + p_node->Name(), " is the second one."); + } + yield_op_order_in_topological_sort = static_cast(i); + } + + node_index_to_its_order_in_topological_sort_map[p_node->Index()] = static_cast(i); + } + + ActivationUsedMap fw_op_output_arg_used_map; + + InlinedHashMap is_forward_nodes; + ORT_RETURN_IF_ERROR(GetStashedActivationCandidates(graph_viewer, + yield_op_order_in_topological_sort, + fw_op_output_arg_used_map, + candidate_output_args_map, + is_forward_nodes, + logger)); + + // The first pass - find the candidate subgraphs. + for (int i = static_cast(node_ids.size()) - 1; i >= 0; --i) { + const Node* p_node = graph_viewer.GetNode(node_ids[i]); + if (p_node == nullptr) { + continue; + } + + if (candidate_output_args_map.find(p_node) == candidate_output_args_map.end()) { + continue; + } + + bool can_compromise_stashed_activation = false; + std::unique_ptr recompute_plan = + CheckNodeForRecompute(*p_node, + probe_level, + fw_op_output_arg_used_map, + node_index_to_its_order_in_topological_sort_map, + candidate_output_args_map, + logger, false, + can_compromise_stashed_activation); + if (recompute_plan != nullptr) { + memory_opt_planner.AddNodeOptimizationPlan(p_node, std::move(recompute_plan)); + } + + if (can_compromise_stashed_activation) { + LOGS(logger, VERBOSE) << "Searching Node " << p_node->Name() << "(" << p_node->OpType() + << ") for compromised recompute"; + // If the subgraph recompute can save memory by comprising the assumption - recompute graphs' input must exist + // during backward pass, then we can consider to recompute them. + std::unique_ptr recompute_with_compromise_plan = + CheckNodeForRecompute(*p_node, probe_level, fw_op_output_arg_used_map, + node_index_to_its_order_in_topological_sort_map, + candidate_output_args_map, + logger, true, + can_compromise_stashed_activation); + if (recompute_with_compromise_plan != nullptr) { + memory_opt_planner.AddNodeOptimizationPlan(p_node, std::move(recompute_with_compromise_plan)); + } + } + } + + return Status::OK(); +} + +void GetMemoryRecordsGroupedByNodeClusterId(const MemoryOptimizationPlanner& memory_opt_planner, + const NodeToClusterApplyContextMap& node_to_apply_contexts_map, + std::vector>& generated_records) { + // Group by node cluster id, generate memory record. + InlinedHashMap records; + const auto& node_to_optimization_plan_map = memory_opt_planner.GetNodeToOptimizationPlanMap(); + for (const auto& node_to_optimization_plan : node_to_optimization_plan_map) { + const auto& node = node_to_optimization_plan.first; + const auto& node_plans = node_to_optimization_plan.second; + const std::string node_cluster_id = memory_opt_planner.GenerateNodeClusterId(node); + + std::pair::iterator, bool> insert_result = + records.insert({node_cluster_id, MemoryRecord()}); + bool already_exist = !insert_result.second; + auto& record = insert_result.first->second; + record.freq++; + + // Collect more information for display. + for (auto& plan : node_plans) { + // Same node cluster id, plans might still have different reuse_buffer pattern, so we need to collect all of them. + if (plan->reuse_buffers.size() > 0) { + gsl::span output_indices = plan->GetActivationOutputIndices(); + for (auto output_index : output_indices) { + bool is_output_reusing_buffers = plan->reuse_buffers.find(output_index) != plan->reuse_buffers.end(); + if (plan->GetOptimizationType() == OptimizationType::RecomputeWithCompromise) { + if (is_output_reusing_buffers) { + record.output_port_reuse_recompute_with_compromise_count[output_index] += 1; + } + } else if (plan->GetOptimizationType() == OptimizationType::Recompute) { + if (is_output_reusing_buffers) { + record.output_port_reuse_recompute_count[output_index] += 1; + } + } + } + } + + // For other infos that are guaranteed identity by cluster id, just skip collecting. + if (already_exist) { + continue; + } + + if (plan->GetOptimizationType() == OptimizationType::RecomputeWithCompromise) { + record.recompute_with_compromise_subgraph_str = + dynamic_cast(plan.get())->GetNodesInTopoOrderStr(); + } else if (plan->GetOptimizationType() == OptimizationType::Recompute) { + record.recompute_subgraph_str = dynamic_cast(plan.get())->GetNodesInTopoOrderStr(); + } + + gsl::span output_indices = plan->GetActivationOutputIndices(); + for (auto output_index : output_indices) { + const auto& output_def = node->OutputDefs()[output_index]; + MLDataType ml_data_type = DataTypeImpl::TypeFromProto(*output_def->TypeAsProto()); + ORT_ENFORCE(ml_data_type->IsTensorType(), "ml_type must be a tensor type, but it is ", + DataTypeImpl::ToString(ml_data_type)); + const TensorTypeBase* tensor_type_base = ml_data_type->AsTensorType(); + ORT_ENFORCE(nullptr != tensor_type_base); + MLDataType elt_type = tensor_type_base->GetElementType(); + + const auto byte_count_per_element = elt_type->Size(); + if (plan->GetOptimizationType() == OptimizationType::RecomputeWithCompromise) { + record.compromise_recomputed_outputs.emplace_back( + output_index, + GetTensorElemCountInSymbolicString(node, output_index), + byte_count_per_element, + plan->GetSaveRatio()); + + } else if (plan->GetOptimizationType() == OptimizationType::Recompute) { + record.recomputed_outputs.emplace_back(output_index, + GetTensorElemCountInSymbolicString(node, output_index), + byte_count_per_element, + plan->GetSaveRatio()); + } + } + } + } + + // Sort by feq and then by record key, to make sure the output is deterministic. + InlinedVector> freq_to_record_key; + for (const auto& p : records) { + freq_to_record_key.push_back({p.second.freq, p.first}); + } + + std::sort(freq_to_record_key.begin(), freq_to_record_key.end(), [](auto& left, auto& right) { + if (left.first == right.first) { + return left.second.compare(right.second) > 0; + } + return left.first > right.first; + }); + + for (const auto& p : freq_to_record_key) { + const std::string record_key = p.second; + generated_records.push_back({record_key, records[record_key]}); + } + + // If apply context is provided, also update the actual applied count. + if (node_to_apply_contexts_map.size() > 0) { + InlinedHashMap node_cluster_id_to_record_map; + for (auto& p : generated_records) { + node_cluster_id_to_record_map[p.first] = &p.second; + } + + for (const auto& p : node_to_apply_contexts_map) { + const auto& node = p.first; + const auto& apply_context = p.second; + std::string node_cluster_id = memory_opt_planner.GenerateNodeClusterId(node); + if (apply_context->type == OptimizationType::Recompute) { + node_cluster_id_to_record_map[node_cluster_id]->actual_recompute_count += 1; + node_cluster_id_to_record_map[node_cluster_id]->request_recompute_count = apply_context->requested_count; + } else if (apply_context->type == OptimizationType::RecomputeWithCompromise) { + node_cluster_id_to_record_map[node_cluster_id]->actual_recompute_with_compromise_count += 1; + node_cluster_id_to_record_map[node_cluster_id]->request_recompute_with_compromise_count = + apply_context->requested_count; + } else { + ORT_THROW("Unsupported optimization type found."); + } + } + } +} + +// Function declare to make it compile. +void IterateNodeOptimizationPlan(const std::shared_ptr& plan, + const InlinedHashMap>>& + node_to_optimization_plans_map, + const InlinedVector>& + current_combination, + const logging::Logger& logger, + InlinedVector>>& + all_combinations); + +/* + * Iterate from a node, generate combinations for each optimization plan for it. + */ +void IterateNode(const Node* node, + const InlinedHashMap>>& + node_to_optimization_plans_map, + const InlinedVector>& + current_combination, + const logging::Logger& logger, + InlinedVector>>& + all_combinations) { + MO_LOG_DEBUG_INFO(logger, "Enter IterateNode: " + node->Name()); + if (node_to_optimization_plans_map.find(node) == node_to_optimization_plans_map.end()) { + MO_LOG_DEBUG_INFO(logger, "Exit IterateNode since reused node don't have optimization plans: " + node->Name()); + return; + } + + for (const std::shared_ptr& plan : node_to_optimization_plans_map.at(node)) { + if (std::find(current_combination.begin(), current_combination.end(), plan) != + current_combination.end()) { + continue; + } + InlinedVector> new_combination = current_combination; + new_combination.push_back(plan); + IterateNodeOptimizationPlan(plan, node_to_optimization_plans_map, new_combination, logger, all_combinations); + } + MO_LOG_DEBUG_INFO(logger, "Exit IterateNode: " + node->Name()); +} + +void ListAllCombinations(const InlinedVector>>>& + all_possible_node_optimization_plans, + int index, + const InlinedVector>& current_combination, + const logging::Logger& logger, + InlinedVector>>& + all_combinations) { + MO_LOG_DEBUG_INFO(logger, "Enter ListAllCombinations"); + if (index == static_cast(all_possible_node_optimization_plans.size())) { + if (std::find(all_combinations.begin(), all_combinations.end(), current_combination) == + all_combinations.end()) { + all_combinations.push_back(current_combination); + } + MO_LOG_DEBUG_INFO(logger, "Exit ListAllCombinations after finding a new combination"); + return; + } + + for (const auto& plans : all_possible_node_optimization_plans[index]) { + for (const auto& plan : plans) { + InlinedVector> new_combination = current_combination; + new_combination.push_back(plan); + ListAllCombinations(all_possible_node_optimization_plans, index + 1, new_combination, logger, all_combinations); + } + } + + MO_LOG_DEBUG_INFO(logger, "Exit ListAllCombinations"); +} + +/** + * Iterate from a node optimization plan, if there is any buffer reuse in its node outputs, + * iterate all possible reuse buffer plan combinations. + */ +void IterateNodeOptimizationPlan(const std::shared_ptr& plan, + const InlinedHashMap>>& + node_to_optimization_plans_map, + const InlinedVector>& + current_combination, + const logging::Logger& logger, + InlinedVector>>& + all_combinations) { + MO_LOG_DEBUG_INFO(logger, "Enter IterateNodeOptimizationPlan: " + plan->GetClusterId()); + + // No reuse buffer, don't need to iterate further, we found a plan combination already. + if (plan->reuse_buffers.size() == 0) { + MO_LOG_DEBUG_INFO(logger, "length of current_combination: " + + std::to_string(current_combination.size()) + ", " + plan->GetClusterId()); + all_combinations.push_back(current_combination); + MO_LOG_DEBUG_INFO(logger, "Exit IterateNodeOptimizationPlan"); + return; + } + + InlinedVector>>> + all_possible_node_optimization_plans; + all_possible_node_optimization_plans.resize(plan->reuse_buffers.size()); + + size_t i = 0; + for (const auto& p : plan->reuse_buffers) { + MO_LOG_DEBUG_INFO(logger, ">>>reuse buffer: " + std::to_string(p.first)); + IterateNode(p.second.first, node_to_optimization_plans_map, {}, logger, all_possible_node_optimization_plans[i]); + ++i; + } + + ListAllCombinations(all_possible_node_optimization_plans, 0, current_combination, logger, all_combinations); + + MO_LOG_DEBUG_INFO(logger, "Exit IterateNodeOptimizationPlan: " + plan->GetClusterId()); +} + +// Return a deterministic string for multiple plans combinations. +std::string GetMultiplePlanClusterId(const InlinedVector>& plans) { + constexpr const int request_count = -1; // -1 means apply optimization to all appearances. + + std::ostringstream oss; + InlinedVector sorted_plans; + for (const auto& plan : plans) { + sorted_plans.push_back(plan->GetClusterId() + ":" + std::to_string(static_cast(plan->GetOptimizationType())) + + ":" + std::to_string(request_count)); + } + + std::sort(sorted_plans.begin(), sorted_plans.end()); + + for (const auto& plan : sorted_plans) { + if (oss.str().size() > 0) { + oss << ","; + } + oss << plan; + } + return oss.str(); +} + +void GetMemorySavingSymbolicString(const MemoryOptimizationPlanner& memory_opt_planner, + const logging::Logger& logger, + std::map>& + combination_cluster_ids_to_saved_symbolic_byte_map) { + // Group by "ClusterId:OptimizationType:RequestCount". + InlinedVector>> all_combinations; + + combination_cluster_ids_to_saved_symbolic_byte_map.clear(); + const auto& node_to_optimization_plan_map = memory_opt_planner.GetNodeToOptimizationPlanMap(); + for (const auto& node_to_optimization_plan : node_to_optimization_plan_map) { + const auto& node = node_to_optimization_plan.first; + InlinedVector> current_combination; + MO_LOG_DEBUG_INFO(logger, ">>>Start looping node: " + node->Name()); + IterateNode(node, node_to_optimization_plan_map, current_combination, logger, all_combinations); + MO_LOG_DEBUG_INFO(logger, "<<Name()); + } + + for (const auto& combination : all_combinations) { + std::string combination_cluster_id = GetMultiplePlanClusterId(combination); + std::string symbolic_byte_count = ""; + for (const auto& plan : combination) { + if (symbolic_byte_count.size() > 0) { + symbolic_byte_count += " + "; + } + symbolic_byte_count += plan->GetMemorySavingSymbolicString(); + } + + if (symbolic_byte_count.size() > 0) { + symbolic_byte_count = "(" + symbolic_byte_count + ")"; + } + auto& p = combination_cluster_ids_to_saved_symbolic_byte_map[combination_cluster_id]; + const auto& original = p.first; + if (original.size() > 0) { + symbolic_byte_count = original + " + " + symbolic_byte_count; + } + + MO_LOG_DEBUG_INFO(logger, "combination_cluster_id: " + combination_cluster_id + + ", symbolic_byte_count: " + symbolic_byte_count); + + p.first = symbolic_byte_count; + p.second += 1; + } +} + +namespace { + +template +std::string ToFixedLengthString(T value, int length) { + std::ostringstream oss; + oss << std::setw(length) << std::left; + oss << value; + return oss.str(); +} + +void FormatRecomputeMemoryRecords(int option_index, + const MemoryRecord& record, + bool compromise_recompute, + InlinedVector& rows) { + const auto subgraph_str = compromise_recompute ? record.recompute_with_compromise_subgraph_str + : record.recompute_subgraph_str; + const auto opt_type = compromise_recompute ? OptimizationType::RecomputeWithCompromise + : OptimizationType::Recompute; + const auto request_count = compromise_recompute ? record.request_recompute_with_compromise_count + : record.request_recompute_count; + const auto actual_count = compromise_recompute ? record.actual_recompute_with_compromise_count + : record.actual_recompute_count; + + const std::string empty_first_col = "|" + ToFixedLengthString(std::string(), kFirstColumnWidth) + "|"; + + rows.push_back(empty_first_col); + rows.push_back(empty_first_col + + ToFixedLengthString(">>Option " + std::to_string(option_index), kTitleWidthInSecondColumn) + ": " + + OptimizationTypeToString(opt_type) + " subgraph " + subgraph_str); + + if (request_count) { + // Only show this if user requested it. + rows.push_back( + empty_first_col + + ToFixedLengthString(" Status", kTitleWidthInSecondColumn) + ": " + "Enabled, requested count=" + + std::to_string(request_count) + + ", actual applied count=" + std::to_string(actual_count)); + } else { + rows.push_back(empty_first_col + ToFixedLengthString(" Status", kTitleWidthInSecondColumn) + + ": Disabled. Enable with export ORTMODULE_MEMORY_OPT_CONFIG=" + + subgraph_str + ":" + std::to_string(static_cast(opt_type)) + ":-1"); + } + + std::string activation_str = empty_first_col + " Stashed Activations: "; + rows.push_back(activation_str); + + const auto& reused_buffers = compromise_recompute ? record.output_port_reuse_recompute_with_compromise_count + : record.output_port_reuse_recompute_count; + if (reused_buffers.size() > 0) { + std::string reused_buffers_summary = empty_first_col + ToFixedLengthString(" - ReuseFreq", kTitleWidthInSecondColumn) + ": "; + for (const auto& p : reused_buffers) { + reused_buffers_summary += " Output " + std::to_string(p.first) + "(" + std::to_string(p.second) + "),"; + } + + rows.push_back(reused_buffers_summary); + } + + const auto activation_count = compromise_recompute ? record.compromise_recomputed_outputs.size() + : record.recomputed_outputs.size(); + for (size_t i = 0; i < activation_count; ++i) { + const MemoryRecord::OutputStat* stat; + if (compromise_recompute) { + stat = &record.compromise_recomputed_outputs[i]; + } else { + stat = &record.recomputed_outputs[i]; + } + + rows.push_back(empty_first_col + + ToFixedLengthString(" - Output " + std::to_string(stat->output_index), kTitleWidthInSecondColumn) + + ": [" + stat->output_shape_str + "], byte/elem: " + + std::to_string(stat->output_byte_count_per_element) + + ", " + std::to_string(static_cast(stat->saving_ratio * 100)) + + "% saved"); + } +} +} // namespace + +std::string SerializeMemoryRecords( + const std::vector>& records_grouped_by_node_cluster_id, + std::string_view user_config) { + InlinedVector rows; + rows.push_back(kTableBorder); + rows.push_back("|" + ToFixedLengthString("Freq", kFirstColumnWidth) + + "| Memory Optimization Opportunities (Clustered by node-level activation patterns)"); + rows.push_back(kTableRowSeparator); + + for (const auto& p : records_grouped_by_node_cluster_id) { + const auto& record = p.second; + rows.push_back("|" + ToFixedLengthString(record.freq, kFirstColumnWidth) + + "|For each row options are mutually exclusive, only one of them can be enabled."); + + int option_index = 1; + if (record.recomputed_outputs.size() > 0) { + FormatRecomputeMemoryRecords(option_index, record, false, rows); + option_index++; + } + + if (record.compromise_recomputed_outputs.size() > 0) { + FormatRecomputeMemoryRecords(option_index, record, true, rows); + option_index++; + } + rows.push_back(kTableRowSeparator); + } + + rows.push_back(kTableBorder); + + size_t max_length = 0; + for (auto& row : rows) { + max_length = std::max(max_length, row.length()); + } + + // Example is: + // static const std::string row_separator = + // "|_ _ _ _|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _|\n"; + static const std::string kTableRowSeparatorStart = "|_ _ _ _|"; + size_t second_row_length = max_length - kTableRowSeparatorStart.length(); + if (second_row_length % 2 == 0) { + second_row_length += 2; + max_length += 2; + } else { + second_row_length += 3; // add 3 to make it even + max_length += 3; + } + std::string row_separator_full(second_row_length, ' '); + for (size_t i = 0; i < row_separator_full.size() - 1; ++i) { + if (i % 2 == 0) { + row_separator_full[i] = '_'; + } + } + row_separator_full[row_separator_full.size() - 1] = '|'; + row_separator_full = kTableRowSeparatorStart + row_separator_full; + + std::string table_border_full(max_length, '='); + std::ostringstream summary; + summary << std::endl; + summary << MakeString("MemoryInsight Summary - User config: ", (user_config.empty() ? "not provided" : user_config)) + << std::endl; + for (auto& row : rows) { + if (row == kTableRowSeparator) { + summary << row_separator_full << std::endl; + } else if (row == kTableBorder) { + summary << table_border_full << std::endl; + } else { + std::string filled_up = std::string(max_length - row.length(), ' '); + filled_up[filled_up.length() - 1] = '|'; + summary << row << filled_up << std::endl; + } + } + summary << "Note: use comma as a separator for enabling more than one subgraphs." << std::endl; + return summary.str(); +} + +std::string GetSerializedORTModuleMemoryStat(const GraphViewer& graph_viewer, + std::string_view memory_optimization_config, + std::string_view recompute_probe_level, + const logging::Logger& logger, + std::map>& + cluster_id_combinations_to_saved_symbolic_byte_map, + const OrtValueNameIdxMap* ortvalue_name_to_idx_map, + const SequentialExecutionPlan* p_seq_exec_plan) { + ProbeLevel probe_level = ProbeLevel::Advanced; + if (!recompute_probe_level.empty()) { + int probe_level_int = ParseIntValueFromString(recompute_probe_level); + ORT_ENFORCE(probe_level_int < static_cast(ProbeLevel::LevelMax) && + probe_level_int >= 0, + "Invalid probe level specified: ", recompute_probe_level); + probe_level = static_cast(probe_level); + } + + ptrdiff_t yield_op_order_in_topological_sort; + InlinedHashMap> candidate_output_args_map; + InlinedHashMap node_index_to_its_order_in_topological_sort_map; + + // The first pass - find the candidate subgraphs. + MemoryOptimizationPlanner memory_opt_planner; + ORT_ENFORCE(FindORTModuleMemoryOpportunity( + graph_viewer, + probe_level, + logger, + node_index_to_its_order_in_topological_sort_map, + yield_op_order_in_topological_sort, + candidate_output_args_map, + memory_opt_planner) + .IsOK()); + + InlinedHashMap cluster_id_to_config_map; + // Finalize the plan according to user config, + // then create a ClusterApplyContext for each unique cluster (having the same node pattern) + + NodeToClusterApplyContextMap node_to_apply_context_map; + + if (!memory_optimization_config.empty()) { + ORT_ENFORCE(ParseConfigFromString(memory_optimization_config, cluster_id_to_config_map) + .IsOK()); + InlinedHashMap> node_to_opt_plan_map; + ORT_ENFORCE(memory_opt_planner.FinalizeNodePlansFromUserConfig(cluster_id_to_config_map, + node_to_opt_plan_map, + node_to_apply_context_map) + .IsOK()); + } + + if (ortvalue_name_to_idx_map != nullptr && p_seq_exec_plan != nullptr) { + ORT_ENFORCE(memory_opt_planner.UpdateNodePlansFromExecutionPlan(graph_viewer, + *ortvalue_name_to_idx_map, + *p_seq_exec_plan) + .IsOK()); + } + + std::vector> records; + GetMemoryRecordsGroupedByNodeClusterId(memory_opt_planner, node_to_apply_context_map, records); + + GetMemorySavingSymbolicString(memory_opt_planner, logger, cluster_id_combinations_to_saved_symbolic_byte_map); + + return SerializeMemoryRecords(records, memory_optimization_config); +} + +} // namespace onnxruntime::optimizer::memory_optimizer diff --git a/orttraining/orttraining/core/optimizer/memory_optimizer/memory_insight.h b/orttraining/orttraining/core/optimizer/memory_optimizer/memory_insight.h new file mode 100644 index 0000000000..c4267efdbe --- /dev/null +++ b/orttraining/orttraining/core/optimizer/memory_optimizer/memory_insight.h @@ -0,0 +1,129 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include + +#include "orttraining/core/optimizer/memory_optimizer/common.h" +#include "orttraining/core/optimizer/memory_optimizer/optimization_planner.h" +#include "orttraining/core/optimizer/memory_optimizer/recompute_analysis.h" + +namespace onnxruntime::optimizer::memory_optimizer { + +/** + * @brief A data structure to store memory optimization statistics for a specific node cluster id. + * + * We will collect statistics for each node cluster id. + * The node cluster id is generated from all possible optimization plans for a specific node, plus shape, data type, + * outputs, etc. For the nodes have the same node cluster id, they will have one single MemoryRecord, displayed + * as a row in the final memory optimization statistics table. + */ +class MemoryRecord { + public: + class OutputStat { + public: + OutputStat(size_t output_index, std::string_view output_shape, size_t output_byte_count_per_element, + float saving_ratio) + : output_index(output_index), + output_shape_str(output_shape), + output_byte_count_per_element(output_byte_count_per_element), + saving_ratio(saving_ratio) {} + + // output index, shape, byte count per element, saving ratio + size_t output_index; + std::string output_shape_str; + size_t output_byte_count_per_element; + float saving_ratio; + }; + + // Recompute Column + std::string recompute_subgraph_str; + InlinedVector recomputed_outputs; + int request_recompute_count = 0; + int actual_recompute_count = 0; + InlinedHashMap output_port_reuse_recompute_count; + + // RecomputeWithCompromise Column + std::string recompute_with_compromise_subgraph_str; + InlinedVector compromise_recomputed_outputs; + int request_recompute_with_compromise_count = 0; + int actual_recompute_with_compromise_count = 0; + InlinedHashMap output_port_reuse_recompute_with_compromise_count; + + // Frequency Column + int freq = 0; +}; + +/** + * @brief Iterate the graph and find all possible memory optimization opportunities for related nodes. + * + * @param graph_viewer The graph to iterate. + * @param probe_level The level to control allowed operations during recomputable subgraph detecting. + * @param logger Logger. + * @param node_index_to_its_order_in_topological_sort_map The mapping of node index to its order in topological sort. + * @param yield_op_order_in_topological_sort The order of the boundary op in the topological sort. + * @param candidate_output_args_map A map from node to its candidate activations, which are consumed by both fw and + * @param mem_opt_stats A store to maintain all found optimization plans for related nodes. + * @return Status + */ +Status FindORTModuleMemoryOpportunity(const GraphViewer& graph_viewer, + const ProbeLevel probe_level, + const logging::Logger& logger, + InlinedHashMap& + node_index_to_its_order_in_topological_sort_map, + ptrdiff_t& yield_op_order_in_topological_sort, + InlinedHashMap>& candidate_output_args_map, + MemoryOptimizationPlanner& mem_opt_stats); + +/** + * @brief From the optimization plans, generate the memory optimization statistics table containing many MemoryRecords, + * each represents one node cluster id. + * + * @param memory_opt_planner The optimization planner to get optimization plans. + * @param node_to_apply_contexts_map The optimization applying information. + * @param generated_records Returns the generated memory optimization statistics table. + * (for example, how many are actually applied) to each MemoryRecord. + */ +void GetMemoryRecordsGroupedByNodeClusterId(const MemoryOptimizationPlanner& memory_opt_planner, + const NodeToClusterApplyContextMap& + node_to_apply_contexts_map, + std::vector>& generated_records); + +/** + * @brief Serialize the memory optimization statistics table to a string. + * + * @param records_grouped_by_node_cluster_id The memory optimization statistics table. + * @param user_config The user configuration to the serialized string. + * @return std::string + */ +std::string SerializeMemoryRecords(const std::vector>& + records_grouped_by_node_cluster_id, + std::string_view user_config); + +/** + * @brief A public API exposed to retrieve the memory optimization statistics table, given a graph. + * + * If possible, session's allocation plans and execution plan will also be available to help the analysis. + * + * @param graph_viewer The graph to analyze. + * @param memory_optimization_config The user configuration to control the memory optimization. + * @param recompute_probe_level The level to control allowed operations during recomputable subgraph detecting. + * @param logger Logger. + * @param ortvalue_name_to_idx_map Optional. If provided, we will use it to map ort value name to index. + * @param p_seq_exec_plan Optional. If provided, we will use it to get allocation plans. + * @return std::string + */ +std::string GetSerializedORTModuleMemoryStat(const GraphViewer& graph_viewer, + std::string_view memory_optimization_config, + std::string_view recompute_probe_level, + const logging::Logger& logger, + // used as Python binding, so used std::map instead of InlinedHashMap + std::map>& + cluster_id_combinations_to_saved_symbolic_byte_map, + const OrtValueNameIdxMap* ortvalue_name_to_idx_map = nullptr, + const SequentialExecutionPlan* p_seq_exec_plan = nullptr); + +} // namespace onnxruntime::optimizer::memory_optimizer diff --git a/orttraining/orttraining/core/optimizer/memory_optimizer/optimization_planner.cc b/orttraining/orttraining/core/optimizer/memory_optimizer/optimization_planner.cc new file mode 100644 index 0000000000..7e042031f6 --- /dev/null +++ b/orttraining/orttraining/core/optimizer/memory_optimizer/optimization_planner.cc @@ -0,0 +1,140 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include + +#include "core/graph/graph_utils.h" +#include "core/optimizer/utils.h" +#include "core/framework/ort_value_name_idx_map.h" +#include "core/framework/sequential_execution_plan.h" + +#include "orttraining/core/optimizer/memory_optimizer/common.h" +#include "orttraining/core/optimizer/memory_optimizer/optimization_planner.h" + +namespace onnxruntime::optimizer::memory_optimizer { + +std::string NodeOptimizationPlanBase::GetMemorySavingSymbolicString() const { + std::string saving_str; + for (auto output_index : activation_output_indices_) { + // If the output is reusing other node's buffer, then no memory saving. + if (reuse_buffers.find(output_index) != reuse_buffers.end()) { + continue; + } + + const auto& output_def = node->OutputDefs()[output_index]; + MLDataType ml_data_type = DataTypeImpl::TypeFromProto(*output_def->TypeAsProto()); + ORT_ENFORCE(ml_data_type->IsTensorType(), "ml_type must be a tensor type, but it is ", + DataTypeImpl::ToString(ml_data_type)); + const TensorTypeBase* tensor_type_base = ml_data_type->AsTensorType(); + ORT_ENFORCE(nullptr != tensor_type_base); + MLDataType elt_type = tensor_type_base->GetElementType(); + const auto byte_count_per_element = elt_type->Size(); + if (!saving_str.empty()) { + saving_str += " + "; + } + saving_str = "(" + GetTensorElemCountInSymbolicString(node, output_index) + " * " + + std::to_string(byte_count_per_element) + " * " + + std::to_string(GetSaveRatio()) + ")"; + } + if (saving_str.empty()) { + return saving_str; + } + return "(" + saving_str + ")"; +} + +Status MemoryOptimizationPlanner::UpdateNodePlansFromExecutionPlan(const GraphViewer& graph_viewer, + const OrtValueNameIdxMap& ortvalue_name_to_idx_map, + const SequentialExecutionPlan& p_seq_exec_plan) { + InlinedHashMap idx_to_ortvalue_name_map; + for (const auto& entry : ortvalue_name_to_idx_map) { + idx_to_ortvalue_name_map[entry.second] = entry.first; + } + + for (const auto& node_to_optimization_plan : node_to_optimization_plans_map) { + const auto& node_plans = node_to_optimization_plan.second; + + for (auto& node_plan : node_plans) { + const std::string cluster_id = node_plan->GetClusterId(); + const Node* node = node_plan->node; + for (auto& output_index : node_plan->GetActivationOutputIndices()) { + const NodeArg* node_arg = node->OutputDefs()[output_index]; + const auto& ort_value_name = node_arg->Name(); + int ort_value_idx; + ORT_ENFORCE(ortvalue_name_to_idx_map.GetIdx(ort_value_name, ort_value_idx).IsOK()); + const auto& alloc_plan = p_seq_exec_plan.allocation_plan; + ORT_ENFORCE(ort_value_idx >= 0 && static_cast(ort_value_idx) < alloc_plan.size()); + const auto& per_alloc_plan = alloc_plan[ort_value_idx]; + if (per_alloc_plan.alloc_kind != AllocKind::kReuse) { + continue; + } + int reused_ort_value_idx = per_alloc_plan.reused_buffer; + const auto& reused_ort_value_name = idx_to_ortvalue_name_map.at(reused_ort_value_idx); + + const Node* p_node = graph_viewer.GetProducerNode(reused_ort_value_name); + if (p_node == nullptr) { + // This is a graph input. + continue; + } + + int src_op_output_index = optimizer_utils::IndexOfNodeOutput(*p_node, *node_arg); + node_plan->reuse_buffers[output_index] = std::make_pair(p_node, src_op_output_index); + } + } + } + + return Status::OK(); +} + +Status MemoryOptimizationPlanner::FinalizeNodePlansFromUserConfig( + const InlinedHashMap& cluster_id_to_user_configs, + InlinedHashMap>& node_to_opt_plan_map, + NodeToClusterApplyContextMap& node_to_apply_context_map) const { + if (cluster_id_to_user_configs.size() == 0) { + return Status::OK(); + } + + // Create a temporary map to store the apply context for each cluster pattern. + InlinedHashMap> cluster_id_to_apply_contexts_map; + + // We loop all nodes' optimization plans and find the match in user configs. + // If found in user configs, we finalize the plan and create/update the apply context for this node. + // If not found in user configs, we will not include the node in the returned result. + for (const auto& node_to_optimization_plan : node_to_optimization_plans_map) { + const auto& node = node_to_optimization_plan.first; + const auto& node_plans = node_to_optimization_plan.second; + + for (auto& node_plan : node_plans) { + const std::string cluster_id = node_plan->GetClusterId(); + if (cluster_id_to_user_configs.find(cluster_id) == cluster_id_to_user_configs.end()) { + continue; + } + + const auto& user_config = cluster_id_to_user_configs.at(cluster_id); + if (node_plan->GetOptimizationType() == user_config.type) { + // First finalize the plan for this node. + node_to_opt_plan_map[node] = node_plan; + + // Create/Update the apply context for this node. + if (cluster_id_to_apply_contexts_map.find(cluster_id) == cluster_id_to_apply_contexts_map.end()) { + std::shared_ptr apply_context = std::make_shared(); + apply_context->requested_count = user_config.requested_count; + apply_context->type = user_config.type; + apply_context->total_frequency++; + cluster_id_to_apply_contexts_map.insert({cluster_id, apply_context}); + } + + node_to_apply_context_map[node] = cluster_id_to_apply_contexts_map.at(cluster_id); + + // If different plans for the same node have same cluster id, we only need to finalize the first one. + // The rest of them will be ignored. + break; + } + } + } + + return Status::OK(); +} + +} // namespace onnxruntime::optimizer::memory_optimizer diff --git a/orttraining/orttraining/core/optimizer/memory_optimizer/optimization_planner.h b/orttraining/orttraining/core/optimizer/memory_optimizer/optimization_planner.h new file mode 100644 index 0000000000..0e5e2967ec --- /dev/null +++ b/orttraining/orttraining/core/optimizer/memory_optimizer/optimization_planner.h @@ -0,0 +1,133 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include "orttraining/core/optimizer/memory_optimizer/common.h" +#include "core/framework/ort_value_name_idx_map.h" +#include "core/framework/sequential_execution_plan.h" + +namespace onnxruntime::optimizer::memory_optimizer { + +/** + * @brief Struct to store properties of a specific subgraph. + */ +class ClusterApplyContext { + public: + ClusterApplyContext() = default; + + OptimizationType type; + int requested_count{0}; + int total_frequency{0}; // The occurrence of this subgraph pattern in the graph. + + int applied_count{0}; // The number of times this subgraph pattern has been really applied in this transformer. + int skip_count{0}; // The number of times this subgraph instance has been skipped in reversed topological order. +}; + +/** + * @brief Base class for a concrete optimization plan. + * + */ +class NodeOptimizationPlanBase { + public: + NodeOptimizationPlanBase(const Node* node, + gsl::span activation_output_indices, + float save_ratio) + : node(node), + activation_output_indices_(activation_output_indices.begin(), activation_output_indices.end()), + save_ratio_(save_ratio) { + } + + virtual ~NodeOptimizationPlanBase() = default; + + virtual OptimizationType GetOptimizationType() const = 0; + + /** + * Get the cluster id for this optimization plan. + * This cluster id is used to enable the optimization as a unique identity, for example, for recompute it is a + * subgraph string representation. + * @return std::string + */ + virtual std::string GetClusterId() const = 0; + + /** + * Get a string used to generate node cluster id for this optimization plan. + * Node cluster id is on Node level, each node can have multiple optimization plans, each plan generates its + * normalization string. Once combined we get Node cluster id. This id is used to categorize nodes into different + * groups, showing them as one row in memory optimization opportunity table. + * @return std::string + */ + virtual std::string NormalizeForNodeClusterId() const = 0; + + /** + * Return all output indices that are used as activation buffers. + */ + gsl::span GetActivationOutputIndices() const { return activation_output_indices_; } + + /** + * Return the saving ratio for this optimization plan. + */ + float GetSaveRatio() const { return save_ratio_; } + + /** + * Get a symbolic string to represent the memory saving for this optimization plan. + */ + std::string GetMemorySavingSymbolicString() const; + + const Node* node; + // A map: output index reusing other node's output (other_node, output index) + InlinedHashMap reuse_buffers; + + private: + InlinedVector activation_output_indices_; + float save_ratio_ = 1.0f; +}; + +using NodeToClusterApplyContextMap = InlinedHashMap>; + +class MemoryOptimizationPlanner { + public: + void AddNodeOptimizationPlan(const Node* node, + std::shared_ptr plan) { + if (node_to_optimization_plans_map.find(node) == node_to_optimization_plans_map.end()) { + node_to_optimization_plans_map.insert({node, {}}); + } + + node_to_optimization_plans_map[node].emplace_back(plan); + } + + Status UpdateNodePlansFromExecutionPlan(const GraphViewer& graph_viewer, + const OrtValueNameIdxMap& ortvalue_name_to_idx_map, + const SequentialExecutionPlan& p_seq_exec_plan); + + Status FinalizeNodePlansFromUserConfig( + const InlinedHashMap& cluster_id_to_user_configs, + InlinedHashMap>& node_to_opt_plan_map, + NodeToClusterApplyContextMap& node_to_apply_context_map) const; + + std::string GenerateNodeClusterId(const Node* node) const { + ORT_ENFORCE(node_to_optimization_plans_map.find(node) != node_to_optimization_plans_map.end(), + "Node not found in node_to_optimization_plans_map."); + std::ostringstream oss; + const auto& node_plans = node_to_optimization_plans_map.at(node); + for (auto& plan : node_plans) { + oss << plan->NormalizeForNodeClusterId(); + } + + return oss.str(); + } + + const InlinedHashMap>>& + GetNodeToOptimizationPlanMap() const { + return node_to_optimization_plans_map; + } + + private: + InlinedHashMap>> node_to_optimization_plans_map; +}; + +} // namespace onnxruntime::optimizer::memory_optimizer diff --git a/orttraining/orttraining/core/optimizer/memory_optimizer/recompute_analysis.cc b/orttraining/orttraining/core/optimizer/memory_optimizer/recompute_analysis.cc new file mode 100644 index 0000000000..0782cbdae2 --- /dev/null +++ b/orttraining/orttraining/core/optimizer/memory_optimizer/recompute_analysis.cc @@ -0,0 +1,405 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include + +#include "orttraining/core/optimizer/memory_optimizer/common.h" +#include "orttraining/core/optimizer/memory_optimizer/recompute_analysis.h" +#include "core/framework/data_types.h" + +namespace onnxruntime::optimizer::memory_optimizer { + +namespace { + +constexpr int32_t MAXIMUM_RECOMPUTE_NODE_COUNT = 15; + +static size_t GetElementSize(const ONNX_NAMESPACE::DataType& tensor_type) { + const ONNX_NAMESPACE::TypeProto& type_proto = ONNX_NAMESPACE::Utils::DataTypeUtils::ToTypeProto(tensor_type); + MLDataType ml_data_type = DataTypeImpl::TypeFromProto(type_proto); + const TensorTypeBase* tensor_type_base = ml_data_type->AsTensorType(); + ORT_ENFORCE(nullptr != tensor_type_base); + MLDataType elt_type = tensor_type_base->GetElementType(); + return elt_type->Size(); +} + +// TODO(pengwa): extent this function to be more general. +float InputOutputSizeRatio(const Node* node) { + if (node->OpType().compare("Cast") == 0) { + const NodeArg* input = node->InputDefs()[0]; + const NodeArg* output = node->OutputDefs()[0]; + if (input->TypeAsProto()->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_STRING || + output->TypeAsProto()->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_STRING) { + return 1.0f; + } + const auto& ptype1 = input->Type(); + const auto& ptype2 = output->Type(); + float ratio = static_cast(GetElementSize(ptype1)) / static_cast(GetElementSize(ptype2)); + return ratio; + } + + return 1.0f; +} + +/** + * @brief Used to define per-op recompute config. + * + */ +struct AllowedRecomputeNodeConfig { + InlinedVector input_arg_indices; // input index to iterate further (bottom up) +}; + +// The op types that are supported predefined. + +const InlinedHashMap& GetAllowedRecomputeOps(int probe_op_level) { + static InlinedHashMap> recomputable_op_table_map; + if (recomputable_op_table_map.find(probe_op_level) != recomputable_op_table_map.end()) { + return recomputable_op_table_map.at(probe_op_level); + } + + recomputable_op_table_map.insert({probe_op_level, InlinedHashMap()}); + auto& recomputable_op_table = recomputable_op_table_map.at(probe_op_level); + if (probe_op_level >= static_cast(ProbeLevel::Basic)) { + recomputable_op_table.insert({ + // Binary elementwise + {"Add", AllowedRecomputeNodeConfig{{0, 1}}}, + {"BiasGelu", AllowedRecomputeNodeConfig{{0, 1}}}, + {"Div", AllowedRecomputeNodeConfig{{0, 1}}}, + {"Mul", AllowedRecomputeNodeConfig{{0, 1}}}, + {"Sub", AllowedRecomputeNodeConfig{{0, 1}}}, + + // Data layout + /// The shape input is trivial whether it exists or not in backward. + {"Reshape", AllowedRecomputeNodeConfig{{0}}}, + {"Squeeze", AllowedRecomputeNodeConfig{{0}}}, + {"Unsqueeze", AllowedRecomputeNodeConfig{{0}}}, + + // Unary elementwise + /// The ratio and mode input are trivial whether they exist or not in backward + {"BitmaskDropout", AllowedRecomputeNodeConfig{{0}}}, + /// The axis input is trivial whether it exists or not in backward + {"CumSum", AllowedRecomputeNodeConfig{{0}}}, + {"Dropout", AllowedRecomputeNodeConfig{{0}}}, + {"Gelu", AllowedRecomputeNodeConfig{{0}}}, + {"FastGelu", AllowedRecomputeNodeConfig{{0}}}, + + // Ternary elementwise + {"Where", AllowedRecomputeNodeConfig{{0, 1, 2}}}, + + // Data copy + {"Tile", AllowedRecomputeNodeConfig{{0}}}, + {"Cast", AllowedRecomputeNodeConfig{{0}}}, + }); + } + + if (probe_op_level >= static_cast(ProbeLevel::Advanced)) { + recomputable_op_table.insert({ + {"MatMul", AllowedRecomputeNodeConfig{{0, 1}}}, + {"FusedMatMul", AllowedRecomputeNodeConfig{{0, 1}}}, + {"Softmax", AllowedRecomputeNodeConfig{{0}}}, + {"BiasSoftmax", AllowedRecomputeNodeConfig{{0, 1}}}, + {"BiasSoftmaxDropout", AllowedRecomputeNodeConfig{{0, 1}}}, + }); + } + + return recomputable_op_table; +} + +/** + * @brief Check whether a node is a recomputable node at given probe level. + */ +bool IsRecomputable(const Node& node, ProbeLevel probe_level) { + const auto& op_table = GetAllowedRecomputeOps(static_cast(probe_level)); + return op_table.find(node.OpType()) != op_table.end(); +} + +/** + * @brief Find recomputable subgraphs (has at least one nodes, at most MAXIMUM_RECOMPUTE_NODE_COUNT nodes). + * + * @param node The entry node to start the subgraph matching (bottom-up), usually the last node of found subgraphs. + * @param node_output_index_candidates Candidate output indices of "node", which are consumed by both fw and bw ops. + * @param fw_op_output_arg_used_map The activation usage (in fw and bw) mapping. + * @param node_index_to_its_order_in_topological_sort_map The mapping of node index to its order in topological sort. + * Used to re-order the collected subgraph nodes. + * @param nodes_in_topological_order Collected vector of nodes of found subgraph, in the order of the topological + * sorted. + * @param logger Logger. + * @param compromise_stashed_activation Whether to compromise stashed activation, e.g. if we cannot find a + * recomputable subgraph to save a stashed activation, we can compromise to find a recomputable subgraph to reduce the + * size of stashed activation. + * @param can_compromise_stashed_activation A bool return value, to indicate there is opportunaties for finding a + * compromised subgraph. + * @param save_ratio The ratio of memory saving if we can find a recomputable subgraph. + * @return Status + */ +Status SelectRecomputeSubgraph(const Node& entry_node, + const ProbeLevel probe_level, + const InlinedVector& node_output_index_candidates, + const ActivationUsedMap& fw_op_output_arg_used_map, + const InlinedHashMap& + node_index_to_its_order_in_topological_sort_map, + const logging::Logger& logger, + InlinedVector& nodes, + bool compromise_stashed_activation, + bool& can_compromise_stashed_activation, + float& save_ratio) { + const auto& recomputable_op_table = GetAllowedRecomputeOps(static_cast(probe_level)); + + can_compromise_stashed_activation = false; + + LOGS(logger, VERBOSE) << "Enter SelectRecomputeSubgraph for Node " << entry_node.Name() << "(" + << entry_node.OpType() << ")"; + nodes.clear(); + + std::deque q; + for (auto output_index : node_output_index_candidates) { + q.push_back(NodeOutputPort(&entry_node, output_index)); + } + + bool early_stop = false; + std::set visited_output_arg_set; + std::set visited_node_set; + + // For the initial activations in queue, they are stashed ones, so we do differently when scanning the queue for them. + bool is_first_queue_scan = true; + while (nodes.size() < MAXIMUM_RECOMPUTE_NODE_COUNT && !q.empty() && !early_stop) { + // Loop all candidate NodeOutputPort, and find the next layer of input nodes. + size_t current_queue_size = q.size(); + for (size_t i = 0; i < current_queue_size; ++i) { + NodeOutputPort p = q.front(); + q.pop_front(); + const Node* curr_node = p.first; + + // Skip if the node output is already visited. + if (std::find(visited_output_arg_set.begin(), visited_output_arg_set.end(), p) != + visited_output_arg_set.end()) { + continue; + } + + visited_output_arg_set.insert({p}); + + // If the node is already visited by from its other output index, skip it. + if (visited_node_set.find(curr_node) != visited_node_set.end()) { + continue; + } + + visited_node_set.insert(curr_node); + + // Bottom-up search rules. + // If current op is entry output node (that generates stashed activations): + // 1. If the op is not in recomputable_op_table, skip it. + // Otherwise: + // If current op is in allowed list, check its input args, and append the producers' NodeOutputPorts to next_q. + // If current op is NOT in allowed list: + // 1). the output does not exist in backward, we cannot find a good solution for so, the search terminates. + // 2). the output is used in backward, we don't need to trace back further, so continue searching. + auto op_recompute_config_it = recomputable_op_table.find(curr_node->OpType()); + auto cur_output_arg_name = curr_node->OutputDefs()[p.second]->Name(); + if (is_first_queue_scan) { + // We handle the entry node outputs differently because, we don't want this case falls into and succeed one of + // the checks in the other branch + // 1. "op is not in recompute op list, but its output is used in backward" + // 2. "op is in recompute op list, but its output is used in backward" + // (either of the above checks is true for entry node outputs) + if (op_recompute_config_it == recomputable_op_table.end()) { + early_stop = true; + LOGS(logger, VERBOSE) << "Entry Node " << curr_node->Name() << "(" << curr_node->OpType() << ") is **NOT** " + << "in recompute op list, search terminates."; + break; + } + } else { + if (op_recompute_config_it == recomputable_op_table.end()) { + if (fw_op_output_arg_used_map.at(cur_output_arg_name).second) { + LOGS(logger, VERBOSE) << "Node " << curr_node->Name() << "(" << curr_node->OpType() << ") is **NOT** in " + << "recompute op list, but its output [" << cur_output_arg_name << "] is used in " + << "backward, we don't need trace bottom-up further. Entry node: " + << entry_node.Name() << "(" << entry_node.OpType() << ")"; + continue; + } else { + early_stop = true; + LOGS(logger, VERBOSE) << "Node " << curr_node->Name() << "(" << curr_node->OpType() << ") is **NOT** in " + << "recompute op list, and its output [" << cur_output_arg_name + << "] does not exist in backward, search terminates. Entry node: " + << entry_node.Name() << "(" << entry_node.OpType() << ")"; + break; + } + } + + if (fw_op_output_arg_used_map.at(cur_output_arg_name).second) { + LOGS(logger, VERBOSE) << "Node " << curr_node->Name() << "(" << curr_node->OpType() << ") " + << "is in recompute op list, while its output [" << cur_output_arg_name + << "] is used in backward, we don't need trace bottom-up further. Entry node: " + << entry_node.Name() << "(" << entry_node.OpType() << ")"; + continue; + } + } + + // Append node to the selected graph. + if (std::find(nodes.begin(), nodes.end(), curr_node) == nodes.end()) { + nodes.push_back(curr_node); + LOGS(logger, VERBOSE) << "Node " << curr_node->Name() << "(" << curr_node->OpType() + << ") is added in selected subgraph "; + } + + // This check is not matured now, subject to change. + float ratio = InputOutputSizeRatio(curr_node); + float saving_ratio = 1.0f - ratio; + float is_current_node_compromisable = (ratio < 1.f); + can_compromise_stashed_activation = can_compromise_stashed_activation || is_current_node_compromisable; + if (is_current_node_compromisable) { + LOGS(logger, VERBOSE) << "Node " << curr_node->Name() << "(" << curr_node->OpType() + << ") has input/output size " << ratio << " < 1.f, can compromise stashed activation"; + } + + if (is_current_node_compromisable && compromise_stashed_activation) { + LOGS(logger, VERBOSE) << "Node " << curr_node->Name() << "(" << curr_node->OpType() << ") is in " + << "recompute op list, and its output [" << cur_output_arg_name + << "] does not exist in backward, while it meets compromised check, we don't need trace " + << "bottom-up further."; + save_ratio = saving_ratio; + continue; + } + + // Iterate all input nodes according to allowed input arg index of the entry node. + const auto& input_arg_indices = op_recompute_config_it->second.input_arg_indices; + for (auto it = curr_node->InputEdgesBegin(), end = curr_node->InputEdgesEnd(); it != end; ++it) { + const Node::EdgeEnd& input_edge = *it; + const auto& parent_node = input_edge.GetNode(); + const auto parent_node_output_index = input_edge.GetSrcArgIndex(); + const auto current_node_input_index = input_edge.GetDstArgIndex(); + if (std::find(input_arg_indices.begin(), input_arg_indices.end(), current_node_input_index) != + input_arg_indices.end()) { + NodeOutputPort next_p = std::make_pair(&parent_node, parent_node_output_index); + + LOGS(logger, VERBOSE) << "Node " << parent_node.Name() << "(" << parent_node.OpType() << ")'s " + << parent_node_output_index + << "th output [" << parent_node.OutputDefs()[parent_node_output_index]->Name() + << "] is added in recompute search list "; + + q.push_back(next_p); + } + } + } + // After handling all entry node outputs, we set the flag to false. + is_first_queue_scan = false; + } + + // If input args are not found in bw, but op count exceed MAXIMUM_RECOMPUTE_NODE_COUNT, skip recompute. + if (!q.empty() || early_stop) { + LOGS(logger, VERBOSE) << "Fail to find a solution for recompute: current node count is " << nodes.size() + << ", queue size: " << q.size() << ", early stop: " << early_stop; + nodes.clear(); + } else { + // Re-order the nodes in topological order. + std::sort(nodes.begin(), nodes.end(), + [&node_index_to_its_order_in_topological_sort_map](const Node*& lhs, const Node*& rhs) { + return node_index_to_its_order_in_topological_sort_map.at(lhs->Index()) < + node_index_to_its_order_in_topological_sort_map.at(rhs->Index()); + }); + } + return Status::OK(); +} + +/** + * @brief Convert the recompute subgraph to its string representation. + * + * @param nodes_in_topological_order The subgraph nodes in topological order. + * @param subgraph_string_representation Returns subgraph string representation. + * @param log_info Returns log info for users. + */ +void NodesInTopoOrderToString(gsl::span nodes_in_topological_order, + std::string& subgraph_string_representation, + std::string& log_info) { + std::ostringstream oss; + std::ostringstream subgraph_string_representation_oss; + size_t node_count = nodes_in_topological_order.size(); + for (size_t i = 0; i < node_count; ++i) { + if (i < node_count - 1) { // Ignore the last node. + oss << "(name:" << nodes_in_topological_order[i]->Name() << ", type:" << nodes_in_topological_order[i]->OpType() + << "),"; + } + + subgraph_string_representation_oss << nodes_in_topological_order[i]->OpType() << "+"; + } + + subgraph_string_representation = subgraph_string_representation_oss.str(); + log_info = oss.str(); + if (log_info.size() > 0) { + log_info = " with its precedent nodes: " + log_info; + } +} + +} // namespace + +std::unique_ptr CheckNodeForRecompute(const Node& node, + const ProbeLevel probe_level, + const ActivationUsedMap& fw_op_output_arg_used_map, + const InlinedHashMap& + node_index_to_its_order_in_topological_sort_map, + const InlinedHashMap>& + candidate_output_args_map, + const logging::Logger& logger, + bool compromise_stashed_activation, + bool& can_compromise_stashed_activation) { + if (!IsRecomputable(node, probe_level)) { + return nullptr; + } + + InlinedVector nodes_in_topological_order; + float save_ratio = 1.f; + ORT_ENFORCE(SelectRecomputeSubgraph(node, + probe_level, + candidate_output_args_map.at(&node), + fw_op_output_arg_used_map, + node_index_to_its_order_in_topological_sort_map, + logger, + nodes_in_topological_order, + compromise_stashed_activation, + can_compromise_stashed_activation, + save_ratio) + .IsOK()); + if (nodes_in_topological_order.size() == 0) { + return nullptr; + } + + std::string subgraph_str_representation, log_info; + NodesInTopoOrderToString(nodes_in_topological_order, subgraph_str_representation, log_info); + + LOGS(logger, VERBOSE) << "Node " << node.Name() << "(" << node.OpType() << ") can be recomputed" << log_info; + + return std::make_unique(&node, candidate_output_args_map.at(&node), + nodes_in_topological_order, + compromise_stashed_activation, + save_ratio); +} + +std::string NodeRecomputePlan::GetClusterId() const { + std::ostringstream oss; + oss << GetNodesInTopoOrderStr(); + return oss.str(); +} + +std::string NodeRecomputePlan::NormalizeForNodeClusterId() const { + std::ostringstream oss; + oss << "recompute:" << node->OpType() << "-" + << compromise_recompute_ << "-"; + for (auto& output_index : GetActivationOutputIndices()) { + oss << output_index << ":" << GetTensorElemCountInSymbolicString(node, output_index); + oss << ":" << node->OutputDefs()[output_index]->TypeAsProto()->tensor_type().elem_type() << "-"; + } + + oss << GetNodesInTopoOrderStr(); + return oss.str(); +} + +std::string NodeRecomputePlan::GetNodesInTopoOrderStr() const { + std::string subgraph_str_representation, log_info; + NodesInTopoOrderToString(nodes_in_topological_order_, subgraph_str_representation, log_info); + return subgraph_str_representation; +} + +} // namespace onnxruntime::optimizer::memory_optimizer diff --git a/orttraining/orttraining/core/optimizer/memory_optimizer/recompute_analysis.h b/orttraining/orttraining/core/optimizer/memory_optimizer/recompute_analysis.h new file mode 100644 index 0000000000..9211e5044c --- /dev/null +++ b/orttraining/orttraining/core/optimizer/memory_optimizer/recompute_analysis.h @@ -0,0 +1,104 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include "orttraining/core/optimizer/memory_optimizer/common.h" +#include "orttraining/core/optimizer/memory_optimizer/optimization_planner.h" + +namespace onnxruntime::optimizer::memory_optimizer { + +/** + * @brief Level to control allowed operations during subgraph detecting. + * Level 0: only allow cheap-to-compute operations. + * Level 1: allow more expensive operations. + */ +enum class ProbeLevel { + Basic = 0, + Advanced = 1, + LevelMax = 2, +}; + +/** + * @brief A child class used for Recompute/RecomputeWithCompromise optimization plan. + * + * For each node generating stashed activations, a recompute plan can be created for it. + */ +class NodeRecomputePlan : public NodeOptimizationPlanBase { + public: + NodeRecomputePlan(const Node* node, + const InlinedVector& activation_output_indices, + const InlinedVector& nodes_in_topological_order, + bool compromise_recompute = false, + float save_ratio = 1.0f) : NodeOptimizationPlanBase(node, activation_output_indices, save_ratio) { + compromise_recompute_ = compromise_recompute; + // Be noted, recompute is node level, each node arg should have the same optimization type. + nodes_in_topological_order_ = nodes_in_topological_order; + } + + const InlinedVector& GetNodesInTopoOrder() const { return nodes_in_topological_order_; } + + bool IsCompromiseRecompute() const { return compromise_recompute_; } + + OptimizationType GetOptimizationType() const override { + return compromise_recompute_ ? OptimizationType::RecomputeWithCompromise + : OptimizationType::Recompute; + } + + /** + * @brief Get the cluster id for this recompute plan. + * The cluster id is used to identify a unique subgraph. + * User can pass such cluster id to enable specific memory optimization for some subgraph. + */ + std::string GetClusterId() const override; + + /** + * @brief Get the serialized string for this recompute plan to create Node-level cluster id. + * Imagine, a Node can have multiple optimization plans, each plan generates its normalization string. + * Once combined we get Node cluster id. + * + * Node cluster id is used to categorize nodes into different groups, showing them as one row in memory + * optimization opportunity table. + */ + std::string NormalizeForNodeClusterId() const override; + + std::string GetNodesInTopoOrderStr() const; + + private: + bool compromise_recompute_; + InlinedVector nodes_in_topological_order_; +}; + +/** + * @brief For the node producing stashed activation, check whether a recomputable subgraph can be found or not. + * + * @param node The entry node to start the subgraph matching (bottom-up), usually the last node of found subgraphs. + * @param probe_level The level to control allowed operations during subgraph detecting. + * @param fw_op_output_arg_used_map The activation usage (in fw and bw) mapping. + * @param node_index_to_its_order_in_topological_sort_map The mapping of node index to its order in topological sort. + * Used to re-order the collected subgraph nodes. + * @param candidate_output_args_map A map from node to its candidate activations, which are consumed by both fw and + * bw ops. + * @param subgraph_stores A store to maintain all found subgraphs. + * @param logger Logger. + * @param compromise_stashed_activation Whether to compromise stashed activation, e.g. if we cannot find a + * recomputable subgraph to save a stashed activation, we can compromise to find a recomputable subgraph to reduce the + * size of stashed activation. + * @param can_compromise_stashed_activation A bool return value, to indicate there is opportunaties for finding a + * compromised subgraph. + */ +std::unique_ptr CheckNodeForRecompute(const Node& node, + const ProbeLevel probe_level, + const ActivationUsedMap& fw_op_output_arg_used_map, + const InlinedHashMap& + node_index_to_its_order_in_topological_sort_map, + const InlinedHashMap>& + candidate_output_args_map, + const logging::Logger& logger, + bool compromise_stashed_activation, + bool& can_compromise_stashed_activation); + +} // namespace onnxruntime::optimizer::memory_optimizer diff --git a/orttraining/orttraining/core/optimizer/scaled_sum_fusion.cc b/orttraining/orttraining/core/optimizer/scaled_sum_fusion.cc index dcb3abf247..e719a21118 100644 --- a/orttraining/orttraining/core/optimizer/scaled_sum_fusion.cc +++ b/orttraining/orttraining/core/optimizer/scaled_sum_fusion.cc @@ -254,7 +254,9 @@ Status ScaledSumFusion::ApplyImpl(Graph& graph, bool& modified, int /*graph_leve handled_scaled_sum_count += 1; } - LOGS(logger, INFO) << "Total fused ScaledSum node count: " << handled_scaled_sum_count; + if (handled_scaled_sum_count > 0) { + LOGS(logger, INFO) << "Total fused ScaledSum node count: " << handled_scaled_sum_count; + } return Status::OK(); } diff --git a/orttraining/orttraining/python/orttraining_pybind_state.cc b/orttraining/orttraining/python/orttraining_pybind_state.cc index bb1cb4bbd3..a5f46d88e4 100644 --- a/orttraining/orttraining/python/orttraining_pybind_state.cc +++ b/orttraining/orttraining/python/orttraining_pybind_state.cc @@ -433,7 +433,20 @@ void addObjectMethodsForTraining(py::module& m) { if (!status.IsOK()) { throw std::runtime_error("Error in backward pass execution: " + status.ErrorMessage()); } - }); + }) + .def("get_serialized_ortmodule_memory_stat", // for memory optimization + [](TrainingAgent* agent, // agent + const std::string& memory_optimization_config, // user config string + const std::string& recompute_probe_level // user config string for probe level + ) -> std::tuple>> { + std::map> cluster_id_combinations_to_saved_symbolic_byte_map; + std::string opportunity_table = + agent->GetSerializedORTModuleMemoryStat(memory_optimization_config, + recompute_probe_level, + cluster_id_combinations_to_saved_symbolic_byte_map); + return std::tuple>>( + opportunity_table, cluster_id_combinations_to_saved_symbolic_byte_map); + }); py::enum_(m, "PropagateCastOpsStrategy", py::module_local(), py::arithmetic{}) .value("NONE", GraphTransformerConfiguration::PropagateCastOpsConfiguration::Strategy::None) diff --git a/orttraining/orttraining/python/training/ortmodule/_execution_agent.py b/orttraining/orttraining/python/training/ortmodule/_execution_agent.py index 533fea5a0a..7a89aadee9 100644 --- a/orttraining/orttraining/python/training/ortmodule/_execution_agent.py +++ b/orttraining/orttraining/python/training/ortmodule/_execution_agent.py @@ -3,6 +3,8 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- +from typing import Tuple + import onnxruntime from onnxruntime.capi import _pybind_state as C from onnxruntime.capi._pybind_state import TrainingAgent as C_TrainingAgent @@ -161,3 +163,13 @@ class TrainingAgent: :param state: State of the graph that is used for executing partial graph runs. """ self._training_agent.run_backward(feeds, fetches, state) + + def get_serialized_ortmodule_memory_stat( + self, memory_optimization_config: str, recompute_probe_level: str + ) -> Tuple[str, dict]: + """ + Get serialized memory stats for OrtModule. + """ + return self._training_agent.get_serialized_ortmodule_memory_stat( + memory_optimization_config, recompute_probe_level + ) diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py index 5eb1d9f382..26993dec17 100755 --- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py @@ -19,7 +19,7 @@ from torch.utils.cpp_extension import ROCM_HOME import onnxruntime from onnxruntime.capi import _pybind_state as C from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference -from onnxruntime.training.utils import ORTModelInputOutputSchemaType, onnx_dtype_to_pytorch_dtype +from onnxruntime.training.utils import ORTModelInputOutputSchemaType, PTable, onnx_dtype_to_pytorch_dtype from onnxruntime.training.utils.hooks import configure_ort_compatible_zero_stage3 from . import _are_deterministic_algorithms_enabled, _io, _logger, _onnx_models, _utils @@ -91,7 +91,8 @@ class GraphExecutionManager(GraphExecutionInterface): self._first_skip_check_warning = True # Inspector for runtime information, for example input data, memory usage, etc. - self._runtime_inspector = RuntimeInspector(self._logger) + self._runtime_inspector = RuntimeInspector(self._logger, self._original_module) + self._runtime_inspector.memory_ob.enable_memory_stats_by_step(self._runtime_options.print_memory_stat_by_step) # Tracker for ORTModule model export, session creation overhead. self.time_tracker = _logger.TimeTracker() @@ -242,12 +243,6 @@ class GraphExecutionManager(GraphExecutionInterface): # 0:Verbose, 1:Info, 2:Warning. 3:Error, 4:Fatal. Default is 2. session_options.log_severity_level = int(self._debug_options.logging.log_level) - session_options.add_session_config_entry( - "optimization.enable_memory_optimizer", self._runtime_options.memory_optimizer_config - ) - session_options.add_session_config_entry( - "optimization.enable_memory_probe_recompute_level", self._runtime_options.probe_level - ) # Disable weight prepacking session_options.add_session_config_entry("session.disable_prepacking", "1") @@ -318,7 +313,8 @@ class GraphExecutionManager(GraphExecutionInterface): """ # VERBOSE -> FULL export verbose log + FULL torch other logs from stdout and stderr (C++ backend) - # INFO -> FULL export verbose log + FILTERED torch other logs from stdout and stderr (C++ backend) + # DEVINFO -> FULL export verbose log + FULL torch other logs from stdout and stderr (C++ backend) + # INFO -> [Rank 0] FULL export verbose log + FILTERED torch other logs from stdout and stderr (C++ backend) # WARNING/ERROR -> [Rank 0] NO export verbose log + FILTERED torch other logs from stdout and stderr (C++ backend) # Be noted: rank 0 log only is controlled by logger configured in _logger.py torch_exporter_verbose_log = self._debug_options.logging.log_level <= LogLevel.INFO @@ -565,7 +561,6 @@ class GraphExecutionManager(GraphExecutionInterface): enable sparsity-based optimization. """ - # Enable data sparsity inspection if sparse optimizer is ON or user wants to print input density. if self._runtime_options.enable_sparse_optimizer or self._runtime_options.print_input_density: self._runtime_inspector.enable_input_inspector( @@ -612,9 +607,6 @@ class GraphExecutionManager(GraphExecutionInterface): if not self._runtime_options.print_input_density: self._runtime_inspector.disable_input_inspector() - if self._runtime_options.print_memory_stat: - self._runtime_inspector.enable_memory_inspector(self._original_module) - def _append_pull_weight_trigger_as_input(self, kwargs: Dict, device: torch.device): from ._zero_stage3_compatibility import ( STAGE3_PULL_WEIGHT_TRIGGER_NAME, @@ -634,105 +626,141 @@ class GraphExecutionManager(GraphExecutionInterface): if get_rank() != 0: return - feature_map: List[Tuple[str, bool, str]] = [ - ("ATen Executor", True, "Dispatch ATen operators to ORT's ATen executor"), - ( + if self._runtime_inspector.memory_ob.is_enabled() and self._debug_options.log_level <= LogLevel.DEVINFO: + self._logger.info(self._runtime_inspector.memory_ob.memory_optimization_opportunity_table_str) + + tbl = PTable() + + def _add_record(tbl, columns): + return tbl.add_row([columns[0], ":", "ON" if columns[1] else "OFF", ":", columns[2]]) + + notes = [] + + _add_record(tbl, ["ATen Executor", True, "Dispatch ATen operators to ORT's ATen executor"]) + _add_record( + tbl, + [ "Cast Propagation", self._runtime_options.propagate_cast_ops_level > 0, f"Level {self._runtime_options.propagate_cast_ops_level} enabled", - ), - ( + ], + ) + _add_record( + tbl, + [ "Custom Function", self._runtime_options.enable_custom_autograd_function, "Support custom torch.autograd.Function export and execution", - ), - ( + ], + ) + + output_memory_optimization_details = self._debug_options.log_level <= LogLevel.INFO + mem_row = _add_record( + tbl, + [ "Memory Optimizer", len(self._runtime_options.memory_optimizer_config) > 0, - "Enable with env ORTMODULE_MEMORY_OPT_CONFIG=", - ), - ] + ( + f"User config: {self._runtime_options.memory_optimizer_config}, probe level: {self._runtime_options.probe_level}" + if len(self._runtime_options.memory_optimizer_config) > 0 + else "Enable with env ORTMODULE_MEMORY_OPT_CONFIG=" + ), + ], + ) - # Add compute optimizer - feature_map.extend( + if self._runtime_inspector.memory_ob.is_enabled() and output_memory_optimization_details: + mem_notes, mem_tbl = self._runtime_inspector.memory_ob.display_memory_optimization_plans( + self._runtime_options.memory_optimizer_config + ) + if mem_tbl is not None: + mem_row.append_annotation_table(mem_tbl) + notes.extend(mem_notes) + + _add_record( + tbl, [ - ( - "Compute Optimizer", - self._runtime_options.enable_compute_optimizer, - "Enable/Disable with env ORTMODULE_ENABLE_COMPUTE_OPTIMIZER=1/0", - ), - ( - " -FLOPReduction", - self._runtime_options.enable_compute_optimizer, - "Reduce FLOPs by upstreaming shrinking-sized ops", - ), - ] + "Compute Optimizer", + self._runtime_options.enable_compute_optimizer, + "Enable/Disable with env ORTMODULE_ENABLE_COMPUTE_OPTIMIZER=1/0", + ], + ) + _add_record( + tbl, + [ + " - FLOPReduction", + self._runtime_options.enable_compute_optimizer, + "Reduce FLOPs by upstreaming shrinking-sized ops", + ], ) if self._runtime_options.enable_compute_optimizer: if len(self._runtime_options.label_sparsity_ratio) > 0: - feature_map.append( - (" -LabelSparsityOpt", True, f"Input density: {self._runtime_options.label_sparsity_ratio}") + _add_record( + tbl, [" - LabelSparsityOpt", True, f"Input density: {self._runtime_options.label_sparsity_ratio}"] ) if len(self._runtime_options.embed_sparsity_ratio) > 0: - feature_map.append( - (" -EmbedSparsityOpt", True, f"Input density: {self._runtime_options.embed_sparsity_ratio}") + _add_record( + tbl, [" - EmbedSparsityOpt", True, f"Input density: {self._runtime_options.embed_sparsity_ratio}"] ) # Add fallback - feature_map.append( - ( + _add_record( + tbl, + [ "Auto Fallback", self._runtime_options.fallback_policy is not _FallbackPolicy.FALLBACK_DISABLE, "Fallback to PyTorch when encountering unsupported ops", - ) + ], ) - if self._runtime_options.enable_triton: - feature_map.append( - ( - "TritonOp Enabled", - True, - "ORT will switch to Triton for executing some ops to further accelerate training.", - ) - ) + # Add Triton + _add_record( + tbl, + [ + "TritonOp Enabled", + self._runtime_options.enable_triton, + "ORT will switch to Triton for executing some ops to further accelerate training.", + ], + ) if self._runtime_options.enable_tuning: desc = "Enable tunning Ops online" if self._runtime_options.tuning_results_path: desc += f", save tuning results to {self._runtime_options.tuning_results_path}" - feature_map.append(("Online Op Tuning", True, desc)) + _add_record(tbl, ["Online Op Tuning", True, desc]) elif self._runtime_options.tuning_results_path: - feature_map.append( - ( + _add_record( + tbl, + [ "Offline Op Tuning", True, f"Use offline tuning results from {self._runtime_options.tuning_results_path}", - ) + ], ) - feature_map.append( - ( + _add_record( + tbl, + [ "ZeRO Stage3 Support", self._runtime_options.enable_zero_stage3_support, "Enable/Disable with env ORTMODULE_ENABLE_ZERO_STAGE3=1/0", - ) + ], ) mode = "training" if self._export_mode == torch.onnx.TrainingMode.TRAINING else "inference" mode = f"{_logger.LogColor.UNDERLINE}{mode}{_logger.LogColor.ENDC}" - - stat = f"\n\n{_logger.LogColor.HEADER}***** ONNX Runtime Training (ORTModule) is accelerating your model *****{_logger.LogColor.ENDC}\n\n" + stat = f"\n{_logger.LogColor.HEADER}***** ONNX Runtime Training (ORTModule) is accelerating your model *****{_logger.LogColor.ENDC}\n\n" stat += f"ORTModule is enabled with following features ON/OFF for [{mode}] mode:\n\n" - for feature_tuple in feature_map: - switch_str = "ON" if feature_tuple[1] else "OFF" - stat += f"{feature_tuple[0]:<20}:\t{switch_str:<10}:\t{feature_tuple[2]:<80}\n" + stat += tbl.get_string() + "\n" # Collect ORTModule overheads for different phases. stat += f"\n{self.time_tracker.to_string(self._debug_options.logging.log_level < LogLevel.WARNING)}\n" - stat += f"Versions: ONNX Runtime - {onnxruntime.__version__}, ONNX - {onnx.__version__}\n\n" - stat += f"{_logger.LogColor.HEADER}************************************************************************{_logger.LogColor.ENDC}\n\n" + # Add notes + for index, note in enumerate(notes): + stat += f"Note {index + 1}: {note}\n" + + stat += f"\n{_logger.LogColor.HEADER}************************************************************************{_logger.LogColor.ENDC}\n\n" self._logger.warning(stat) diff --git a/orttraining/orttraining/python/training/ortmodule/_io.py b/orttraining/orttraining/python/training/ortmodule/_io.py index 1b6e2df9d2..f5fbd5093f 100644 --- a/orttraining/orttraining/python/training/ortmodule/_io.py +++ b/orttraining/orttraining/python/training/ortmodule/_io.py @@ -210,6 +210,7 @@ def _combine_input_buffers_initializers( result = [] embed_sparsity_results = OrderedDict() label_sparsity_results = OrderedDict() + onnx_input_to_value_map = OrderedDict() for input_idx, name in enumerate(onnx_input_names): inp = None @@ -251,6 +252,8 @@ def _combine_input_buffers_initializers( if label_density < 100: label_sparsity_results[name] = label_density result.append(inp) + + onnx_input_to_value_map[name] = inp else: raise wrap_exception( ORTModuleONNXModelException, RuntimeError(f"Input is present in ONNX graph but not provided: {name}.") @@ -264,6 +267,10 @@ def _combine_input_buffers_initializers( else: result.extend(params) + if rt_inspector.memory_ob.is_enabled() and not rt_inspector.memory_ob.symbolic_dim_collecting_completed: + rt_inspector.memory_ob.collect_symbolic_dim_values(input_info.dynamic_axes, onnx_input_to_value_map) + rt_inspector.memory_ob.symbolic_dim_collecting_completed = True + return result, embed_sparsity_results, label_sparsity_results diff --git a/orttraining/orttraining/python/training/ortmodule/_logger.py b/orttraining/orttraining/python/training/ortmodule/_logger.py index 0728ebdf19..a01db28374 100644 --- a/orttraining/orttraining/python/training/ortmodule/_logger.py +++ b/orttraining/orttraining/python/training/ortmodule/_logger.py @@ -263,7 +263,7 @@ class SuppressLogs: raise RuntimeError("The class of the function to be tracked must have a '_debug_options' attribute.") with _suppress_os_stream_output( - enable=graph_execution_manager._debug_options.log_level >= LogLevel.INFO, + enable=graph_execution_manager._debug_options.log_level >= LogLevel.DEVINFO, on_exit=partial( _log_with_filter, graph_execution_manager._logger, diff --git a/orttraining/orttraining/python/training/ortmodule/_runtime_inspector.py b/orttraining/orttraining/python/training/ortmodule/_runtime_inspector.py index dda909e8cb..cfd2e25e13 100644 --- a/orttraining/orttraining/python/training/ortmodule/_runtime_inspector.py +++ b/orttraining/orttraining/python/training/ortmodule/_runtime_inspector.py @@ -5,12 +5,18 @@ from enum import IntEnum from logging import Logger -from typing import List, Tuple, Union +from typing import Dict, List, Optional, Tuple, Union import onnx import torch from onnx import ModelProto, helper from onnx import onnx_pb as onnx_proto +from sympy import Symbol, simplify +from sympy.parsing.sympy_parser import parse_expr + +from onnxruntime.training.utils import PTable + +from ._execution_agent import TrainingAgent class Phase(IntEnum): @@ -39,11 +45,11 @@ class RuntimeInspector: Runtime inspector for ORTModule. """ - def __init__(self, logger: Logger): + def __init__(self, logger: Logger, module: torch.nn.Module): self._logger = logger self.input_density_ob: Union[InputDensityObserver, None] = None - self.memory_ob: Union[MemoryObserver, None] = None + self.memory_ob = MemoryObserver(module, self._logger) def enable_input_inspector(self, model: ModelProto, user_input_names: List[str]) -> None: """Initialize input inspector from the given ONNX model and user input names. @@ -82,26 +88,6 @@ class RuntimeInspector: """Disable input density inspector.""" self.input_density_ob = None - def enable_memory_inspector(self, module: torch.nn.Module): - """Enable memory inspector for ORTModule. - - Args: - module: ORTModule. - """ - if self.memory_ob is None: - self.memory_ob = MemoryObserver(module, self._logger) - else: - raise RuntimeError("Memory observer is already enabled.") - - def inspect_memory(self, phase: Phase) -> None: - """Inspect memory usage and print statistics. - - Args: - phase: Phase to inspect. - """ - if self.memory_ob is not None: - self.memory_ob.inspect_memory(phase) - class InputDensityObserver: """Training input data observer for ORTModule. @@ -460,6 +446,16 @@ class InputDensityObserver: return value +class MemoryOptimizationSummary: + """Memory optimization summary for a cluster id combination.""" + + def __init__(self, saving_str="", simplified_saving_expr=None, evaluated_saving=None, freq=0): + self.raw_symbolic_saving_str = saving_str + self.simplified_symbolic_saving_expr: Optional[Symbol] = simplified_saving_expr + self.evaluated_saving: Union[str, int, None] = evaluated_saving + self.freq = freq + + class MemoryObserver: """Memory inspector across the training lifetime. @@ -472,6 +468,19 @@ class MemoryObserver: def __init__(self, m: torch.nn.Module, logger: Logger): self._logger = logger + self._is_enabled = True + + # Memory optimization related. + self.memory_optimization_opportunity_table_str = None + self.cluster_id_combination_to_saving_symbolics_map: Dict[str, MemoryOptimizationSummary] = {} + ## The value is a list of symbolic dim values parsed from the first batch. + self.symbolic_dim_name_to_value_map: Dict = {} + + ## Used to control only the first batch is used to collect symbolic dim values. + self.symbolic_dim_collecting_completed = False + + # For per-step memory inspection. + self._print_memory_stats_by_step = False self._current_step = 0 self._rank = 0 self._world_size = 1 @@ -485,8 +494,77 @@ class MemoryObserver: self._is_first_inspect = True + def is_enabled(self) -> bool: + """Check if memory inspector is enabled.""" + return self._is_enabled + + def enable_memory_stats_by_step(self, print_memory_stats_by_step: bool): + # For per-step memory inspection. + self._print_memory_stats_by_step = print_memory_stats_by_step + + def collect_symbolic_dim_values( + self, + onnx_input_name_to_dynamic_axes_map: Dict[str, Dict[int, str]], + onnx_input_to_value_map: Dict[str, torch.Tensor], + ): + """Collect symbolic dim values.""" + for input_name, dynamic_axes in onnx_input_name_to_dynamic_axes_map.items(): + if input_name in onnx_input_to_value_map: + for dim_idx, dim_name in dynamic_axes.items(): + self.symbolic_dim_name_to_value_map[Symbol(dim_name)] = onnx_input_to_value_map[input_name].size()[ + dim_idx + ] + + def find_memory_optimization_opportunity( + self, execution_agent: TrainingAgent, memory_optimizer_config, probe_level + ): + """Find memory optimization opportunity. + + Args: + execution_agent: TrainingAgent. + memory_optimizer_config: Memory optimization config. + probe_level: Memory probe level. + """ + ( + self.memory_optimization_opportunity_table_str, + memory_optimization_saving_symbolics, + ) = execution_agent.get_serialized_ortmodule_memory_stat(memory_optimizer_config, probe_level) + + cluster_id_to_saving_symbol_map: Dict[str, MemoryOptimizationSummary] = {} + for cluster_id, memory_saving_stat in memory_optimization_saving_symbolics.items(): + memory_saving_symbolic = memory_saving_stat[0] + freq = memory_saving_stat[1] + expr = parse_expr(memory_saving_symbolic) + simplified_expr = simplify(expr) + r = simplified_expr.evalf(subs=self.symbolic_dim_name_to_value_map) + evaluated_saving = None + if r.is_number: + evaluated_saving = float(r) + else: + evaluated_saving = r + + cluster_id_to_saving_symbol_map[cluster_id] = MemoryOptimizationSummary( + memory_saving_symbolic, simplified_expr, evaluated_saving, freq + ) + + # Sorted by evaluated_saving if it is a float + sorted_list = sorted( + cluster_id_to_saving_symbol_map.items(), + key=lambda x: x[1].evaluated_saving if isinstance(x[1].evaluated_saving, float) else 0, + reverse=True, + ) + + for cluster_id, values in sorted_list: + self.cluster_id_combination_to_saving_symbolics_map[cluster_id] = values + def inspect_memory(self, cur_phase: Phase): - if not torch.cuda.is_available(): + """Inspect memory usage and print statistics. + + Args: + phase: Phase to inspect. + """ + + if not torch.cuda.is_available() or not self._print_memory_stats_by_step: return if self._is_first_inspect: @@ -498,36 +576,38 @@ class MemoryObserver: if self._rank != 0: return - if cur_phase < Phase.PRE_FORWARD or cur_phase > self._last_phase: - raise RuntimeError(f"Invalid phase detected: {cur_phase}") + if cur_phase < Phase.PRE_FORWARD or (cur_phase <= self._last_phase): + raise RuntimeError(f"Invalid phase detected: {cur_phase}, last_phase: {self._last_phase}") if (cur_phase - self._pre_phase) != 1: raise RuntimeError(f"Invalid phase transition detected: {self._pre_phase} -> {cur_phase}") - cur_mem_allocated = self._normalize(torch.cuda.memory_allocated()) - max_mem_allocated = self._normalize(torch.cuda.max_memory_allocated()) - cur_mem_cached = self._normalize(torch.cuda.memory_reserved()) - max_mem_cached = self._normalize(torch.cuda.max_memory_reserved()) - torch_mem_stat = torch.cuda.memory_stats() - cur_mem_inactive = self._normalize(torch_mem_stat.get("inactive_split_bytes.all.current", 0)) - max_mem_inactive = self._normalize(torch_mem_stat.get("inactive_split_bytes.all.peak", 0)) - - mem_stats = [ - ["phase", _convert_phase_to_string(cur_phase)], - ["allocated", cur_mem_allocated], # current memory alloeated for tensors - ["max allocated", max_mem_allocated], # peak memory allocated for tensors - ["cached", cur_mem_cached], # current memory cached for caching allocator - ["max cached", max_mem_cached], # peak memory cached for caching allocator. - ["inactive", cur_mem_inactive], # amount of inactive, non-releasable memory - ["max inactive", max_mem_inactive], # peak of inactive, non-releasable memory - ] - - summ = f"{self._rank_info} step {self._current_step} memory ({MemoryObserver.NORMALIZER_UNIT})" - for stat in mem_stats: - summ += f" | {stat[0]}: {stat[1]}" - # For the 10+ steps, only print when it is power of 2. - if self._current_step < 10 or (self._current_step & (self._current_step - 1) == 0): + need_print = self._current_step < 10 or (self._current_step & (self._current_step - 1) == 0) + + if need_print: + cur_mem_allocated = self._normalize(torch.cuda.memory_allocated()) + max_mem_allocated = self._normalize(torch.cuda.max_memory_allocated()) + cur_mem_cached = self._normalize(torch.cuda.memory_reserved()) + max_mem_cached = self._normalize(torch.cuda.max_memory_reserved()) + torch_mem_stat = torch.cuda.memory_stats() + cur_mem_inactive = self._normalize(torch_mem_stat.get("inactive_split_bytes.all.current", 0)) + max_mem_inactive = self._normalize(torch_mem_stat.get("inactive_split_bytes.all.peak", 0)) + + mem_stats = [ + ["phase", _convert_phase_to_string(cur_phase)], + ["allocated", cur_mem_allocated], # current memory allocated for tensors + ["max allocated", max_mem_allocated], # peak memory allocated for tensors + ["cached", cur_mem_cached], # current memory cached for the caching allocator + ["max cached", max_mem_cached], # peak memory cached for caching allocator. + ["inactive", cur_mem_inactive], # amount of inactive, non-releasable memory + ["max inactive", max_mem_inactive], # peak of inactive, non-releasable memory + ] + + summ = f"{self._rank_info} step {self._current_step} memory ({MemoryObserver.NORMALIZER_UNIT})" + for stat in mem_stats: + summ += f" | {stat[0]}: {stat[1]}" + self._logger.info(summ) if cur_phase == self._last_phase: @@ -542,3 +622,72 @@ class MemoryObserver: def _normalize(self, mem_size_in_bytes: Union[float, int]) -> str: return f"{float(mem_size_in_bytes) / MemoryObserver.NORMALIZER_FACTOR:.0f}" + + def display_memory_optimization_plans(self, memory_optimizer_config) -> Tuple[List[str], PTable]: + mem_plan_count = len(self.cluster_id_combination_to_saving_symbolics_map) + + if mem_plan_count > 0: + mem_tbl = PTable() + mem_tbl.add_row(["", "", "", "", "Configs", "Freq", "Max Saving(Bytes)", "Saving Symbolic(Bytes)"]) + + index = 1 + + def _get_user_config_without_freq(configs: str): + if len(configs) == 0: + return [] + config_list = configs.split(",") + configs_with_out_freq = [] + for config in config_list: + config_values = config.split(":") + freq = int(config_values[2]) + if freq == 0: + continue + configs_with_out_freq.append(config_values[0] + ":" + config_values[1]) + + return configs_with_out_freq + + user_configs_with_out_freq = _get_user_config_without_freq(memory_optimizer_config) + + for ( + cluster_id, + saving_symbolic, + ) in self.cluster_id_combination_to_saving_symbolics_map.items(): + saving_bytes = saving_symbolic.evaluated_saving + if isinstance(saving_bytes, float): + saving_bytes = f"{saving_bytes:,.0f}" + + cluster_ids_without_freq = _get_user_config_without_freq(cluster_id) + + mem_tbl.add_row( + [ + f" - Plan {index}", + ":", + "ON" + if all(cluster_id in user_configs_with_out_freq for cluster_id in cluster_ids_without_freq) + else "OFF", + ":", + cluster_id, + saving_symbolic.freq, + saving_bytes, + saving_symbolic.simplified_symbolic_saving_expr, + ] + ) + + index += 1 + + saving_recommendation = ( + "use comma as delimiter to enable multiple memory optimization plans at the same time:\n" + ) + saving_recommendation += " export ORTMODULE_MEMORY_OPT_CONFIG=,,..." + + notes = [] + notes.append(saving_recommendation) + + saving_recommendation = "memory saving is calculated based on the 1st batch symbolic dim values:\n" + for dim_param, dim_value in self.symbolic_dim_name_to_value_map.items(): + saving_recommendation += f" {dim_param}={dim_value}," + notes.append(saving_recommendation) + + return notes, mem_tbl + + return [], None diff --git a/orttraining/orttraining/python/training/ortmodule/_training_manager.py b/orttraining/orttraining/python/training/ortmodule/_training_manager.py index bafb642355..96a95557bb 100644 --- a/orttraining/orttraining/python/training/ortmodule/_training_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_training_manager.py @@ -18,7 +18,7 @@ from ._fallback import ORTModuleFallbackException, _FallbackManager, _FallbackPo from ._gradient_accumulation_manager import GradientAccumulationManager from ._graph_execution_manager import GraphExecutionManager, _RunStateInfo from ._io import _FlattenedModule, _InputInfo, unflatten_user_output -from ._logger import ORTModuleInitPhase, TrackTime +from ._logger import LogLevel, ORTModuleInitPhase, TrackTime from ._runtime_inspector import Phase from ._utils import save_tuning_results, set_tuning_results from .graph_optimizer_registry import GraphOptimizerRegistry @@ -111,7 +111,7 @@ class TrainingManager(GraphExecutionManager): Module outputs are returned to the user """ - self._runtime_inspector.inspect_memory(Phase.PRE_FORWARD) + self._runtime_inspector.memory_ob.inspect_memory(Phase.PRE_FORWARD) if self._runtime_options.skip_check.is_set(_SkipCheck.SKIP_CHECK_DEVICE) is False: # Assert that the input and model device match @@ -146,7 +146,7 @@ class TrainingManager(GraphExecutionManager): for idx in self._graph_info.output_grad_indices_non_differentiable: ctx.mark_non_differentiable(user_outputs[idx]) - self._runtime_inspector.inspect_memory(Phase.POST_FORWARD) + self._runtime_inspector.memory_ob.inspect_memory(Phase.POST_FORWARD) return user_outputs @@ -154,7 +154,7 @@ class TrainingManager(GraphExecutionManager): def backward(ctx, *grad_outputs): """Performs backward pass based on grad wrt module output""" - self._runtime_inspector.inspect_memory(Phase.PRE_BACKWARD) + self._runtime_inspector.memory_ob.inspect_memory(Phase.PRE_BACKWARD) assert ctx.run_info is not None, "forward() or __call__() methods must be called before backward()" if self._runtime_options.skip_check.is_set(_SkipCheck.SKIP_CHECK_DEVICE) is False: @@ -205,7 +205,7 @@ class TrainingManager(GraphExecutionManager): # This version only works if backward_outputs is an OrtValueVector. transferred_backward_outputs = _utils._ortvalues_to_torch_tensor(backward_outputs, self._device) - self._runtime_inspector.inspect_memory(Phase.POST_BACKWARD) + self._runtime_inspector.memory_ob.inspect_memory(Phase.POST_BACKWARD) return tuple(transferred_backward_outputs[idx] if idx != -1 else None for idx in self._gradient_map) @@ -242,7 +242,6 @@ class TrainingManager(GraphExecutionManager): self._runtime_options.skip_check.is_set(_SkipCheck.SKIP_CHECK_EXECUTION_AGENT), self._runtime_options.skip_check.is_set(_SkipCheck.SKIP_CHECK_DEVICE), ) - # If exporting module to ONNX for the first time, this skip check will not take effect. # It will only take effect on subsequent forward calls. build_gradient_graph = False @@ -433,6 +432,39 @@ class TrainingManager(GraphExecutionManager): local_device_rank = self._device.index if device_type == "ort" else _utils.get_device_index(self._device) + # When log level is <= INFO, we would collect memory optimization opportunities. + # (TODO: consider to enable by default once memory optimization feature is stable and well improved.) + # Create a training agent without enabling memory optimization here is beneficial for memory analyzing + # when we have an allocation plan in place, and reuse information is available. + if self._runtime_inspector.memory_ob.is_enabled() and self._debug_options.log_level <= LogLevel.INFO: + # Create a training agent without enabling memory optimization. + execution_agent = TrainingAgent( + self._onnx_models.optimized_model.SerializeToString(), + fw_feed_names, + fw_outputs_device_info, + bw_fetches_names, + bw_outputs_device_info, + session_options, + providers, + provider_options, + local_device_rank, + ) + + self._runtime_inspector.memory_ob.find_memory_optimization_opportunity( + execution_agent, self._runtime_options.memory_optimizer_config, self._runtime_options.probe_level + ) + + # Release it as early as possible. + del execution_agent + + # Enable memory optimization if it is enabled in the session options. + session_options.add_session_config_entry( + "optimization.memory_optimizer_config", self._runtime_options.memory_optimizer_config + ) + session_options.add_session_config_entry( + "optimization.enable_memory_probe_recompute_level", self._runtime_options.probe_level + ) + self._execution_agent = TrainingAgent( self._onnx_models.optimized_model.SerializeToString(), fw_feed_names, diff --git a/orttraining/orttraining/python/training/ortmodule/options.py b/orttraining/orttraining/python/training/ortmodule/options.py index cddd9cd440..77022f86d3 100644 --- a/orttraining/orttraining/python/training/ortmodule/options.py +++ b/orttraining/orttraining/python/training/ortmodule/options.py @@ -137,7 +137,7 @@ class DebugOptions: def torch_exporter_filter(self): """Accessor for the filter export logs configuration.""" torch_version = get_runtime_pytorch_version() - if self.log_level >= LogLevel.INFO: + if self.log_level > LogLevel.DEVINFO: if torch_version < version.parse("2.0"): return [ # WARNING: The shape inference of com.microsoft::SoftmaxCrossEntropyLossInternal type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. @@ -262,7 +262,7 @@ class _RuntimeOptions: # Configuration for dev tools. self.print_input_density = False - self.print_memory_stat = False + self.print_memory_stat_by_step = False # Configuration for fallback. self.fallback_policy = ortmodule.ORTMODULE_FALLBACK_POLICY @@ -321,7 +321,7 @@ class _RuntimeOptions: if "ORTMODULE_PRINT_INPUT_DENSITY" in os.environ: self.print_input_density = int(os.getenv("ORTMODULE_PRINT_INPUT_DENSITY")) == 1 if "ORTMODULE_PRINT_MEMORY_STATS" in os.environ: - self.print_memory_stat = int(os.getenv("ORTMODULE_PRINT_MEMORY_STATS")) == 1 + self.print_memory_stat_by_step = int(os.getenv("ORTMODULE_PRINT_MEMORY_STATS")) == 1 # Configuration for fallback. if "ORTMODULE_FALLBACK_POLICY" in os.environ: diff --git a/orttraining/orttraining/python/training/utils/__init__.py b/orttraining/orttraining/python/training/utils/__init__.py index d40a6ddf7d..244557c3c1 100644 --- a/orttraining/orttraining/python/training/utils/__init__.py +++ b/orttraining/orttraining/python/training/utils/__init__.py @@ -2,6 +2,7 @@ # Licensed under the MIT License. # __init__.py +from onnxruntime.training.utils.ptable import PTable from onnxruntime.training.utils.torch_io_helper import ( ORTModelInputOutputSchemaType, ORTModelInputOutputType, @@ -24,4 +25,5 @@ __all__ = [ "pytorch_type_to_onnx_dtype", "onnx_dtype_to_pytorch_dtype", "pytorch_scalar_type_to_pytorch_dtype", + "PTable", ] diff --git a/orttraining/orttraining/python/training/utils/hooks/_zero_offload_subscriber.py b/orttraining/orttraining/python/training/utils/hooks/_zero_offload_subscriber.py index 0d268a7a4a..61f3b20224 100644 --- a/orttraining/orttraining/python/training/utils/hooks/_zero_offload_subscriber.py +++ b/orttraining/orttraining/python/training/utils/hooks/_zero_offload_subscriber.py @@ -291,7 +291,7 @@ class ORTZeROOffloadPreForwardFunction(torch.autograd.Function): raise RuntimeError(f"param {p} has no grad, this should not happen.") # Param gradient accumulation is triggered here, along with the attached hooks, done by PyTorch. assert p.shape == g.shape, f"param_index: {param_index} - param shape {p.shape} != grad shape {g.shape}" - p.backward(g) + # p.backward(g) # At this point, the **real** param grads are already updated, the following grads are only used for # completing the full backward propagation, will not affect parameter updates. diff --git a/orttraining/orttraining/python/training/utils/ptable.py b/orttraining/orttraining/python/training/utils/ptable.py new file mode 100644 index 0000000000..3b3b80d29e --- /dev/null +++ b/orttraining/orttraining/python/training/utils/ptable.py @@ -0,0 +1,64 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +from typing import List + + +class Row: + """A row in a PTable""" + + def __init__(self, columns: List[str]) -> None: + self._columns: List[str] = columns # List of strings + self._annotation_table = None # Optional PTable used for displaying detailed information about the feature row. + + def append_annotation_table(self, ptable) -> None: + self._annotation_table = ptable + + +class PTable: + """A table that can be printed to the console.""" + + def __init__(self) -> None: + self._rows: List[Row] = [] + self._column_count = None + + def add_row(self, columns: List[str]) -> Row: + """Add a row to the table. The number of columns must match the number of columns in the table.""" + if self._column_count is None: + self._column_count = len(columns) + assert self._column_count == len(columns) + row = Row(columns) + self._rows.append(row) + return row + + def get_string(self, first_column_width=None, second_column_width=None) -> str: + """Serialize the table to a string.""" + # Collect the max width of each column + column_widths = [] + for row in self._rows: + if column_widths: + assert len(column_widths) == len(row._columns) + else: + column_widths = [0] * len(row._columns) + for i, column in enumerate(row._columns): + column_widths[i] = max(column_widths[i], len(str(column))) + + if first_column_width: + column_widths[0] = max(first_column_width, column_widths[0]) + + if second_column_width: + column_widths[2] = max(second_column_width, column_widths[2]) + + serialized_table = "" + for row in self._rows: + for i, column in enumerate(row._columns): + serialized_table += f"{str(column).ljust(column_widths[i] + 2)}" + serialized_table += "\n" + if row._annotation_table: + serialized_table += row._annotation_table.get_string( + first_column_width=column_widths[0], second_column_width=column_widths[2] + ) + + return serialized_table