saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-27 20:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Support more ops for recompute (#20234)

Browse source 
### Support more ops for recompute

To cover Mistral model, and support padding elimination ops.

### 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 2024-04-09 09:24:48 +08:00 committed by GitHub
parent 22a61a3cf5
commit 41acd8c543
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 59 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

80
orttraining/orttraining/core/optimizer/memory_optimizer/recompute_analysis.cc
View file

@ -144,6 +144,20 @@ const InlinedHashMap<std::string, OpsetToIgnorableIndicesMap>& GetAllowedRecompu
{20, {0}},
},
},
{
utils::GetFullQualifiedOpName("Cos", kOnnxDomain),
{
{7, {}},
},
},
{
utils::GetFullQualifiedOpName("CumSum", kOnnxDomain),
{
// The axis input is trivial
{11, {1}},
{14, {1}},
},
},
{
utils::GetFullQualifiedOpName("Dropout", kOnnxDomain),
{
@ -162,27 +176,6 @@ const InlinedHashMap<std::string, OpsetToIgnorableIndicesMap>& GetAllowedRecompu
{14, {}},
},
},
{
utils::GetFullQualifiedOpName("Expand", kOnnxDomain),
{
{8, {1}}, // Ignore the shape.
{13, {1}},
},
},
{
utils::GetFullQualifiedOpName("Cos", kOnnxDomain),
{
{7, {}},
},
},
{
utils::GetFullQualifiedOpName("CumSum", kOnnxDomain),
{
// The axis input is trivial
{11, {1}},
{14, {1}},
},
},
{
utils::GetFullQualifiedOpName("Einsum", kOnnxDomain),
{
@ -199,12 +192,25 @@ const InlinedHashMap<std::string, OpsetToIgnorableIndicesMap>& GetAllowedRecompu
{19, {}},
},
},
{
utils::GetFullQualifiedOpName("Expand", kOnnxDomain),
{
{8, {1}}, // Ignore the shape.
{13, {1}},
},
},
{
utils::GetFullQualifiedOpName("FastGelu", kMSDomain),
{
{1, {}},
},
},
{
utils::GetFullQualifiedOpName("FlattenAndUnpad", kMSDomain),
{
{1, {1}}, // ignore the indices
},
},
{
utils::GetFullQualifiedOpName("Gather", kOnnxDomain),
{
@ -225,6 +231,17 @@ const InlinedHashMap<std::string, OpsetToIgnorableIndicesMap>& GetAllowedRecompu
{1, {}},
},
},
{
utils::GetFullQualifiedOpName("Gemm", kOnnxDomain),
{
{1, {}},
{6, {}},
{7, {}},
{9, {}},
{11, {}},
{13, {}},
},
},
{
utils::GetFullQualifiedOpName("Less", kOnnxDomain),
{
@ -244,6 +261,27 @@ const InlinedHashMap<std::string, OpsetToIgnorableIndicesMap>& GetAllowedRecompu
{14, {}},
},
},
{
utils::GetFullQualifiedOpName("Neg", kOnnxDomain),
{
{1, {}},
{6, {}},
{13, {}},
},
},
{
utils::GetFullQualifiedOpName("NonZero", kOnnxDomain),
{
{9, {}},
{13, {}},
},
},
{
utils::GetFullQualifiedOpName("PadAndUnflatten", kMSDomain),
{
{1, {1, 2}}, // ignore the indices and unflatten_dims
},
},
{
utils::GetFullQualifiedOpName("Range", kOnnxDomain),
{