mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
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=<plan1 config>,<plan2 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 <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. -->
This commit is contained in:
pengwa
GitHub
parent
64dacc2892
commit
43a5147e01
31 changed files with 2637 additions and 1038 deletions
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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=<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=<plan1 config>,<plan2 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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
// <subgraph string : optimization strategy : number of subgraph to apply>.
|
||||
// <subgraph string: optimization strategy: number of subgraph to apply>.
|
||||
// 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.
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
|
||||
#pragma once
|
||||
|
||||
#include <string>
|
||||
#include <string_view>
|
||||
#include <vector>
|
||||
|
||||
|
|
@ -37,5 +38,32 @@ inline InlinedVector<std::string_view> 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
|
||||
|
|
|
|||
|
|
@ -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<std::unique_ptr<GraphTransformer>> GenerateTransformers(
|
|||
// fusions might be prevented if this one removes a Q/DQ node too early.
|
||||
transformers.emplace_back(std::make_unique<QDQFinalCleanupTransformer>(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<MemoryOptimizer>(enable_memory_optimizer, probe_level));
|
||||
#endif
|
||||
|
||||
} break;
|
||||
|
||||
case TransformerLevel::Level3: {
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -1,11 +1,17 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include <string>
|
||||
|
||||
#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<std::string> 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<std::string>& feed_names,
|
||||
const std::vector<std::string>& fetches_names,
|
||||
const std::vector<OrtDevice>& outputs_device_info,
|
||||
std::unique_ptr<FeedsFetchesManager>& feeds_fetches_manager) {
|
||||
std::unique_ptr<FeedsFetchesManager>&
|
||||
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<std::string, std::pair<std::string, int>>&
|
||||
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
|
||||
|
|
|
|||
|
|
@ -5,11 +5,15 @@
|
|||
|
||||
#include <thread>
|
||||
#include <future>
|
||||
#include <map>
|
||||
#include <utility>
|
||||
#include <string>
|
||||
|
||||
#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<OrtDevice>& outputs_device_info,
|
||||
std::unique_ptr<FeedsFetchesManager>& feeds_fetches_manager);
|
||||
|
||||
std::string GetSerializedORTModuleMemoryStat(std::string_view memory_optimization_config,
|
||||
std::string_view recompute_probe_level,
|
||||
std::map<std::string, std::pair<std::string, int>>&
|
||||
cluster_id_combinations_to_saved_symbolic_byte_map) const;
|
||||
|
||||
private:
|
||||
// TrainingAgent runs on a InferenceSession under the hood
|
||||
InferenceSession& inference_session_;
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -1,233 +1,84 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include <algorithm>
|
||||
#include <iomanip>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#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<int>(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<OptimizationType>(optimization_type_int), requested_apply_count};
|
||||
}
|
||||
}
|
||||
|
||||
int probe_level = ParseIntValueFromString(level);
|
||||
ORT_RETURN_IF_NOT(probe_level < static_cast<int>(ProbeLevel::LevelMax) && probe_level >= 0,
|
||||
int probe_level = optimizer::memory_optimizer::ParseIntValueFromString(level);
|
||||
ORT_RETURN_IF_NOT(probe_level < static_cast<int>(optimizer::memory_optimizer::ProbeLevel::LevelMax) &&
|
||||
probe_level >= 0,
|
||||
"Invalid probe level specified: ", level);
|
||||
recompute_probe_level_ = static_cast<ProbeLevel>(probe_level);
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
int64_t MemoryOptimizer::PrepareForTransformation(const Graph& graph,
|
||||
ActivationUsedMap& fw_op_output_arg_used_map,
|
||||
InlinedHashMap<NodeIndex, size_t>&
|
||||
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<ptrdiff_t>(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<ptrdiff_t>(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<ptrdiff_t>(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<std::string, std::pair<bool, bool>>&
|
||||
fw_op_output_arg_used_map,
|
||||
InlinedHashMap<const Node*, InlinedVector<size_t>>&
|
||||
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<optimizer::memory_optimizer::ProbeLevel>(probe_level);
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
bool MemoryOptimizer::ModifyGraph(Graph& graph,
|
||||
const InlinedHashMap<NodeIndex, size_t>&
|
||||
const InlinedHashMap<NodeIndex, ptrdiff_t>&
|
||||
node_index_to_its_order_in_topological_sort_map,
|
||||
const InlinedHashMap<const Node*, InlinedVector<size_t>>&
|
||||
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<optimizer::memory_optimizer::NodeOptimizationPlanBase>& node_plan,
|
||||
std::shared_ptr<optimizer::memory_optimizer::ClusterApplyContext>& 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<optimizer::memory_optimizer::NodeRecomputePlan*>(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<int>(recompute_probe_level_);
|
||||
|
||||
InlinedHashMap<std::string, std::pair<bool, bool>> fw_op_output_arg_used_map;
|
||||
InlinedHashMap<NodeIndex, size_t> 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<const Node*, InlinedVector<size_t>> 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<NodeIndex, ptrdiff_t> node_index_to_its_order_in_topological_sort_map;
|
||||
|
||||
// The first pass - find the candidate subgraphs.
|
||||
for (int i = static_cast<int>(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<const Node*, std::shared_ptr<optimizer::memory_optimizer::NodeOptimizationPlanBase>>
|
||||
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<int>(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<const Node*>& 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<optimizer::memory_optimizer::ClusterApplyContext>>&
|
||||
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<std::pair<std::string, optimizer::memory_optimizer::MemoryRecord>> 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<int>(recompute_probe_level_) >= static_cast<int>(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<int>(recompute_probe_level_) >= static_cast<int>(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<size_t>& node_output_index_candidates,
|
||||
const ActivationUsedMap& fw_op_output_arg_used_map,
|
||||
const InlinedHashMap<NodeIndex, size_t>&
|
||||
node_index_to_its_order_in_topological_sort_map,
|
||||
InlinedVector<const Node*>& 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<NodeOutputPort> q;
|
||||
for (auto output_index : node_output_index_candidates) {
|
||||
q.push_back(NodeOutputPort(&entry_node, static_cast<int>(output_index)));
|
||||
}
|
||||
|
||||
bool early_stop = false;
|
||||
std::set<NodeOutputPort> visited_output_arg_set;
|
||||
std::set<const Node*> 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<NodeIndex, size_t>&
|
||||
node_index_to_its_order_in_topological_sort_map,
|
||||
const InlinedHashMap<const Node*, InlinedVector<size_t>>&
|
||||
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<const Node*> 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<const Node*>& 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;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,163 +2,39 @@
|
|||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
#include <charconv>
|
||||
|
||||
#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<const Node*, int>;
|
||||
using ActivationUsedMap = InlinedHashMap<std::string, std::pair<bool, bool>>;
|
||||
|
||||
/**
|
||||
* @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<std::string, int> 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 <nodes in topological order, a string to represent the subgraph>.
|
||||
using GraphInstanceInfo = std::pair<InlinedVector<const Node*>, std::string>;
|
||||
|
||||
void AddSubGraphInstance(const Node* node,
|
||||
const InlinedVector<const Node*>& 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<std::string /*subgraph_representative_str*/, SubGraphDesc> subgraph_descs;
|
||||
InlinedHashMap<const Node*, GraphInstanceInfo> _optimization_target_graphs_;
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Used to define per-op recompute config.
|
||||
*
|
||||
*/
|
||||
struct AllowedRecomputeNodeConfig {
|
||||
InlinedVector<int> 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<NodeIndex, size_t>&
|
||||
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<std::string, std::pair<bool, bool>>& fw_op_output_arg_used_map,
|
||||
InlinedHashMap<const Node*, InlinedVector<size_t>>& 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<NodeIndex, size_t>& node_index_to_its_order_in_topological_sort_map,
|
||||
const InlinedHashMap<const Node*, InlinedVector<size_t>>& candidate_output_args_map,
|
||||
const InlinedHashMap<NodeIndex, ptrdiff_t>&
|
||||
node_index_to_its_order_in_topological_sort_map,
|
||||
const InlinedHashMap<const Node*, InlinedVector<size_t>>&
|
||||
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<const Node*>& 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<optimizer::memory_optimizer::NodeOptimizationPlanBase>& node_plan,
|
||||
std::shared_ptr<optimizer::memory_optimizer::ClusterApplyContext>& 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<const Node*,
|
||||
std::shared_ptr<optimizer::memory_optimizer::ClusterApplyContext>>&
|
||||
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<size_t>& node_output_index_candidates,
|
||||
const ActivationUsedMap& fw_op_output_arg_used_map,
|
||||
const InlinedHashMap<NodeIndex, size_t>&
|
||||
node_index_to_its_order_in_topological_sort_map,
|
||||
InlinedVector<const Node*>& 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<NodeIndex, size_t>&
|
||||
node_index_to_its_order_in_topological_sort_map,
|
||||
const InlinedHashMap<const Node*, InlinedVector<size_t>>&
|
||||
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<std::string, AllowedRecomputeNodeConfig> recomputable_op_type_to_input_arg_index_map_;
|
||||
// User enabled map of the subgraph string representation to the alleviation type.
|
||||
InlinedHashMap<std::string, UserConfig> pattern_subgraph_to_user_optimizer_config_map_;
|
||||
InlinedHashMap<std::string, optimizer::memory_optimizer::UserConfig> pattern_subgraph_to_user_optimizer_config_map_;
|
||||
std::string optimizer_config_;
|
||||
ProbeLevel recompute_probe_level_;
|
||||
optimizer::memory_optimizer::ProbeLevel recompute_probe_level_;
|
||||
};
|
||||
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -0,0 +1,149 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include <charconv>
|
||||
#include <vector>
|
||||
#include <utility>
|
||||
|
||||
#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<std::string>& 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<std::string> dim_params;
|
||||
return TensorShapeProtoToDimParamVector(shape, dim_params);
|
||||
}
|
||||
|
||||
std::string TensorShapeProtoToString(const ONNX_NAMESPACE::TensorShapeProto* shape) {
|
||||
if (shape == nullptr) {
|
||||
return "unknown";
|
||||
}
|
||||
|
||||
std::vector<std::string> 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<std::string, UserConfig>& 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<int>(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<OptimizationType>(optimization_type_int),
|
||||
requested_apply_count};
|
||||
}
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
} // namespace onnxruntime::optimizer::memory_optimizer
|
||||
|
|
@ -0,0 +1,76 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <utility>
|
||||
|
||||
#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<const Node*, size_t>;
|
||||
using ActivationUsedMap = InlinedHashMap<std::string, std::pair<bool, bool>>;
|
||||
|
||||
/**
|
||||
* @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<std::string, UserConfig>& cluster_id_to_config_map);
|
||||
|
||||
} // namespace onnxruntime::optimizer::memory_optimizer
|
||||
|
|
@ -0,0 +1,763 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include <algorithm>
|
||||
#include <iomanip>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#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<NodeIndex, size_t>&
|
||||
node_index_to_its_order_in_topological_sort_map,
|
||||
ActivationUsedMap& fw_op_output_arg_used_map,
|
||||
InlinedHashMap<const Node*, bool>& 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<ptrdiff_t>(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<ptrdiff_t>(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<const Node*, InlinedVector<size_t>>&
|
||||
candidate_output_args_map,
|
||||
InlinedHashMap<const Node*, bool>& 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<NodeIndex, size_t> 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<NodeIndex, ptrdiff_t>&
|
||||
node_index_to_its_order_in_topological_sort_map,
|
||||
ptrdiff_t& yield_op_order_in_topological_sort,
|
||||
InlinedHashMap<const Node*, InlinedVector<size_t>>&
|
||||
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<ptrdiff_t>(i);
|
||||
}
|
||||
|
||||
node_index_to_its_order_in_topological_sort_map[p_node->Index()] = static_cast<ptrdiff_t>(i);
|
||||
}
|
||||
|
||||
ActivationUsedMap fw_op_output_arg_used_map;
|
||||
|
||||
InlinedHashMap<const Node*, bool> 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<int>(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<NodeRecomputePlan> 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<NodeRecomputePlan> 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<std::pair<std::string, MemoryRecord>>& generated_records) {
|
||||
// Group by node cluster id, generate memory record.
|
||||
InlinedHashMap<std::string, MemoryRecord> 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<InlinedHashMap<std::string, MemoryRecord>::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<const size_t> 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<NodeRecomputePlan*>(plan.get())->GetNodesInTopoOrderStr();
|
||||
} else if (plan->GetOptimizationType() == OptimizationType::Recompute) {
|
||||
record.recompute_subgraph_str = dynamic_cast<NodeRecomputePlan*>(plan.get())->GetNodesInTopoOrderStr();
|
||||
}
|
||||
|
||||
gsl::span<const size_t> 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<std::pair<int, std::string>> 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<std::string, MemoryRecord*> 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<NodeOptimizationPlanBase>& plan,
|
||||
const InlinedHashMap<const Node*, InlinedVector<std::shared_ptr<NodeOptimizationPlanBase>>>&
|
||||
node_to_optimization_plans_map,
|
||||
const InlinedVector<std::shared_ptr<NodeOptimizationPlanBase>>&
|
||||
current_combination,
|
||||
const logging::Logger& logger,
|
||||
InlinedVector<InlinedVector<std::shared_ptr<NodeOptimizationPlanBase>>>&
|
||||
all_combinations);
|
||||
|
||||
/*
|
||||
* Iterate from a node, generate combinations for each optimization plan for it.
|
||||
*/
|
||||
void IterateNode(const Node* node,
|
||||
const InlinedHashMap<const Node*, InlinedVector<std::shared_ptr<NodeOptimizationPlanBase>>>&
|
||||
node_to_optimization_plans_map,
|
||||
const InlinedVector<std::shared_ptr<NodeOptimizationPlanBase>>&
|
||||
current_combination,
|
||||
const logging::Logger& logger,
|
||||
InlinedVector<InlinedVector<std::shared_ptr<NodeOptimizationPlanBase>>>&
|
||||
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<NodeOptimizationPlanBase>& plan : node_to_optimization_plans_map.at(node)) {
|
||||
if (std::find(current_combination.begin(), current_combination.end(), plan) !=
|
||||
current_combination.end()) {
|
||||
continue;
|
||||
}
|
||||
InlinedVector<std::shared_ptr<NodeOptimizationPlanBase>> 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<InlinedVector<InlinedVector<std::shared_ptr<NodeOptimizationPlanBase>>>>&
|
||||
all_possible_node_optimization_plans,
|
||||
int index,
|
||||
const InlinedVector<std::shared_ptr<NodeOptimizationPlanBase>>& current_combination,
|
||||
const logging::Logger& logger,
|
||||
InlinedVector<InlinedVector<std::shared_ptr<NodeOptimizationPlanBase>>>&
|
||||
all_combinations) {
|
||||
MO_LOG_DEBUG_INFO(logger, "Enter ListAllCombinations");
|
||||
if (index == static_cast<int>(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<std::shared_ptr<NodeOptimizationPlanBase>> 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<NodeOptimizationPlanBase>& plan,
|
||||
const InlinedHashMap<const Node*, InlinedVector<std::shared_ptr<NodeOptimizationPlanBase>>>&
|
||||
node_to_optimization_plans_map,
|
||||
const InlinedVector<std::shared_ptr<NodeOptimizationPlanBase>>&
|
||||
current_combination,
|
||||
const logging::Logger& logger,
|
||||
InlinedVector<InlinedVector<std::shared_ptr<NodeOptimizationPlanBase>>>&
|
||||
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<InlinedVector<InlinedVector<std::shared_ptr<NodeOptimizationPlanBase>>>>
|
||||
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<std::shared_ptr<NodeOptimizationPlanBase>>& plans) {
|
||||
constexpr const int request_count = -1; // -1 means apply optimization to all appearances.
|
||||
|
||||
std::ostringstream oss;
|
||||
InlinedVector<std::string> sorted_plans;
|
||||
for (const auto& plan : plans) {
|
||||
sorted_plans.push_back(plan->GetClusterId() + ":" + std::to_string(static_cast<int>(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<std::string, std::pair<std::string, int>>&
|
||||
combination_cluster_ids_to_saved_symbolic_byte_map) {
|
||||
// Group by "ClusterId:OptimizationType:RequestCount".
|
||||
InlinedVector<InlinedVector<std::shared_ptr<NodeOptimizationPlanBase>>> 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<std::shared_ptr<NodeOptimizationPlanBase>> 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, "<<<End looping node: " + node->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 <typename T>
|
||||
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<std::string>& 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<int>(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<int>(stat->saving_ratio * 100)) +
|
||||
"% saved");
|
||||
}
|
||||
}
|
||||
} // namespace
|
||||
|
||||
std::string SerializeMemoryRecords(
|
||||
const std::vector<std::pair<std::string, MemoryRecord>>& records_grouped_by_node_cluster_id,
|
||||
std::string_view user_config) {
|
||||
InlinedVector<std::string> 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<std::string, std::pair<std::string, int>>&
|
||||
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<int>(ProbeLevel::LevelMax) &&
|
||||
probe_level_int >= 0,
|
||||
"Invalid probe level specified: ", recompute_probe_level);
|
||||
probe_level = static_cast<ProbeLevel>(probe_level);
|
||||
}
|
||||
|
||||
ptrdiff_t yield_op_order_in_topological_sort;
|
||||
InlinedHashMap<const Node*, InlinedVector<size_t>> candidate_output_args_map;
|
||||
InlinedHashMap<NodeIndex, ptrdiff_t> 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<std::string, UserConfig> 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<const Node*, std::shared_ptr<NodeOptimizationPlanBase>> 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<std::pair<std::string, MemoryRecord>> 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
|
||||
|
|
@ -0,0 +1,129 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <map>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
|
||||
#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<OutputStat> recomputed_outputs;
|
||||
int request_recompute_count = 0;
|
||||
int actual_recompute_count = 0;
|
||||
InlinedHashMap<size_t, int> output_port_reuse_recompute_count;
|
||||
|
||||
// RecomputeWithCompromise Column
|
||||
std::string recompute_with_compromise_subgraph_str;
|
||||
InlinedVector<OutputStat> compromise_recomputed_outputs;
|
||||
int request_recompute_with_compromise_count = 0;
|
||||
int actual_recompute_with_compromise_count = 0;
|
||||
InlinedHashMap<size_t, int> 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<NodeIndex, ptrdiff_t>&
|
||||
node_index_to_its_order_in_topological_sort_map,
|
||||
ptrdiff_t& yield_op_order_in_topological_sort,
|
||||
InlinedHashMap<const Node*, InlinedVector<size_t>>& 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<std::pair<std::string, MemoryRecord>>& 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<std::pair<std::string, MemoryRecord>>&
|
||||
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<std::string, std::pair<std::string, int>>&
|
||||
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
|
||||
|
|
@ -0,0 +1,140 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
|
||||
#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<int, std::string> 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<size_t>(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<std::string, UserConfig>& cluster_id_to_user_configs,
|
||||
InlinedHashMap<const Node*, std::shared_ptr<NodeOptimizationPlanBase>>& 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<std::string, std::shared_ptr<ClusterApplyContext>> 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<ClusterApplyContext> apply_context = std::make_shared<ClusterApplyContext>();
|
||||
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
|
||||
|
|
@ -0,0 +1,133 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
#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<const size_t> 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<const size_t> 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<size_t, NodeOutputPort> reuse_buffers;
|
||||
|
||||
private:
|
||||
InlinedVector<size_t> activation_output_indices_;
|
||||
float save_ratio_ = 1.0f;
|
||||
};
|
||||
|
||||
using NodeToClusterApplyContextMap = InlinedHashMap<const Node*, std::shared_ptr<ClusterApplyContext>>;
|
||||
|
||||
class MemoryOptimizationPlanner {
|
||||
public:
|
||||
void AddNodeOptimizationPlan(const Node* node,
|
||||
std::shared_ptr<NodeOptimizationPlanBase> 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<std::string, UserConfig>& cluster_id_to_user_configs,
|
||||
InlinedHashMap<const Node*, std::shared_ptr<NodeOptimizationPlanBase>>& 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<const Node*,
|
||||
InlinedVector<std::shared_ptr<NodeOptimizationPlanBase>>>&
|
||||
GetNodeToOptimizationPlanMap() const {
|
||||
return node_to_optimization_plans_map;
|
||||
}
|
||||
|
||||
private:
|
||||
InlinedHashMap<const Node*, InlinedVector<std::shared_ptr<NodeOptimizationPlanBase>>> node_to_optimization_plans_map;
|
||||
};
|
||||
|
||||
} // namespace onnxruntime::optimizer::memory_optimizer
|
||||
|
|
@ -0,0 +1,405 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include <algorithm>
|
||||
#include <deque>
|
||||
#include <memory>
|
||||
#include <set>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
|
||||
#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<float>(GetElementSize(ptype1)) / static_cast<float>(GetElementSize(ptype2));
|
||||
return ratio;
|
||||
}
|
||||
|
||||
return 1.0f;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Used to define per-op recompute config.
|
||||
*
|
||||
*/
|
||||
struct AllowedRecomputeNodeConfig {
|
||||
InlinedVector<int> input_arg_indices; // input index to iterate further (bottom up)
|
||||
};
|
||||
|
||||
// The op types that are supported predefined.
|
||||
|
||||
const InlinedHashMap<std::string, AllowedRecomputeNodeConfig>& GetAllowedRecomputeOps(int probe_op_level) {
|
||||
static InlinedHashMap<int, InlinedHashMap<std::string, AllowedRecomputeNodeConfig>> 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<std::string, AllowedRecomputeNodeConfig>()});
|
||||
auto& recomputable_op_table = recomputable_op_table_map.at(probe_op_level);
|
||||
if (probe_op_level >= static_cast<int>(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<int>(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<int>(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<size_t>& node_output_index_candidates,
|
||||
const ActivationUsedMap& fw_op_output_arg_used_map,
|
||||
const InlinedHashMap<NodeIndex, ptrdiff_t>&
|
||||
node_index_to_its_order_in_topological_sort_map,
|
||||
const logging::Logger& logger,
|
||||
InlinedVector<const Node*>& nodes,
|
||||
bool compromise_stashed_activation,
|
||||
bool& can_compromise_stashed_activation,
|
||||
float& save_ratio) {
|
||||
const auto& recomputable_op_table = GetAllowedRecomputeOps(static_cast<int>(probe_level));
|
||||
|
||||
can_compromise_stashed_activation = false;
|
||||
|
||||
LOGS(logger, VERBOSE) << "Enter SelectRecomputeSubgraph for Node " << entry_node.Name() << "("
|
||||
<< entry_node.OpType() << ")";
|
||||
nodes.clear();
|
||||
|
||||
std::deque<NodeOutputPort> q;
|
||||
for (auto output_index : node_output_index_candidates) {
|
||||
q.push_back(NodeOutputPort(&entry_node, output_index));
|
||||
}
|
||||
|
||||
bool early_stop = false;
|
||||
std::set<NodeOutputPort> visited_output_arg_set;
|
||||
std::set<const Node*> 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<const Node* const> 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<NodeRecomputePlan> CheckNodeForRecompute(const Node& node,
|
||||
const ProbeLevel probe_level,
|
||||
const ActivationUsedMap& fw_op_output_arg_used_map,
|
||||
const InlinedHashMap<NodeIndex, ptrdiff_t>&
|
||||
node_index_to_its_order_in_topological_sort_map,
|
||||
const InlinedHashMap<const Node*, InlinedVector<size_t>>&
|
||||
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<const Node*> 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<NodeRecomputePlan>(&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
|
||||
|
|
@ -0,0 +1,104 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
#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<size_t>& activation_output_indices,
|
||||
const InlinedVector<const Node*>& 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<const Node*>& 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<const Node*> 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<NodeRecomputePlan> CheckNodeForRecompute(const Node& node,
|
||||
const ProbeLevel probe_level,
|
||||
const ActivationUsedMap& fw_op_output_arg_used_map,
|
||||
const InlinedHashMap<NodeIndex, ptrdiff_t>&
|
||||
node_index_to_its_order_in_topological_sort_map,
|
||||
const InlinedHashMap<const Node*, InlinedVector<size_t>>&
|
||||
candidate_output_args_map,
|
||||
const logging::Logger& logger,
|
||||
bool compromise_stashed_activation,
|
||||
bool& can_compromise_stashed_activation);
|
||||
|
||||
} // namespace onnxruntime::optimizer::memory_optimizer
|
||||
|
|
@ -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();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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::string, std::map<std::string, std::pair<std::string, int>>> {
|
||||
std::map<std::string, std::pair<std::string, int>> 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<std::string, std::map<std::string, std::pair<std::string, int>>>(
|
||||
opportunity_table, cluster_id_combinations_to_saved_symbolic_byte_map);
|
||||
});
|
||||
|
||||
py::enum_<GraphTransformerConfiguration::PropagateCastOpsConfiguration::Strategy>(m, "PropagateCastOpsStrategy", py::module_local(), py::arithmetic{})
|
||||
.value("NONE", GraphTransformerConfiguration::PropagateCastOpsConfiguration::Strategy::None)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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=<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=<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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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=<plan1 config>,<plan2 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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
64
orttraining/orttraining/python/training/utils/ptable.py
Normal file
64
orttraining/orttraining/python/training/utils/ptable.py
Normal file
|
|
@ -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
|
||||
Loading…
Reference in a new issue