Skip to content

Switch to use or not einsum op. Needed for ORT#1456

Merged
jeffra merged 12 commits into
deepspeedai:masterfrom
ganik:ganik/ort
Oct 21, 2021
Merged

Switch to use or not einsum op. Needed for ORT#1456
jeffra merged 12 commits into
deepspeedai:masterfrom
ganik:ganik/ort

Conversation

@ganik
Copy link
Copy Markdown
Contributor

@ganik ganik commented Oct 14, 2021

Replace einsum op as its not supported by onnxruntime.
Branch by USE_EINSUM flag.
Example:

...
deepspeed.moe.sharded_moe.USE_EINSUM = False
...
self.fc3 = deepspeed.moe.layer.MoE(..)
...

@ganik
Copy link
Copy Markdown
Contributor Author

ganik commented Oct 14, 2021

Trying draft PR first to see if there are any build breaks

@jeffra
Copy link
Copy Markdown
Collaborator

jeffra commented Oct 14, 2021

Thanks for the contribution @ganik, the formatting tests that are failing can be fixed via these instructions: https://github.com/microsoft/DeepSpeed/blob/master/CONTRIBUTING.md#prerequisites

@ykim362
Copy link
Copy Markdown
Contributor

ykim362 commented Oct 15, 2021

@ganik Can ORT support einsum? Einsum is becoming more universal especially in the context of MoE.

@ganik
Copy link
Copy Markdown
Contributor Author

ganik commented Oct 19, 2021

@ganik Can ORT support einsum? Einsum is becoming more universal especially in the context of MoE.

Yes, we are looking into this, its not easy to make it performant. Einsum is a sort of UBER op where depending on rule it branches out to different translations.

@ganik ganik changed the title [WIP] Replace einsum op Replace einsum op Oct 21, 2021
@ganik ganik changed the title Replace einsum op Switch to use or not einsum op. Needed for ORT Oct 21, 2021
@jeffra
Copy link
Copy Markdown
Collaborator

jeffra commented Oct 21, 2021

Ready for review, tagging @samyam, @awan-10 , @ykim362

Copy link
Copy Markdown
Collaborator

@jeffra jeffra left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, waiting for review from MoE folks though.

@awan-10
Copy link
Copy Markdown
Contributor

awan-10 commented Oct 21, 2021

Thank you for this PR @ganik.

@jeffra - it looks good to me.

@ykim362
Copy link
Copy Markdown
Contributor

ykim362 commented Oct 21, 2021

LGTM!

@jeffra jeffra enabled auto-merge (squash) October 21, 2021 19:59
@jeffra jeffra merged commit 20bf1cc into deepspeedai:master Oct 21, 2021
Dipet pushed a commit to Dipet/DeepSpeed that referenced this pull request Oct 26, 2021
Co-authored-by: Gani Nazirov <ganaziro@microsoft.com>
Co-authored-by: Gani Nazirov <ganaziro@OrtTrainingDev4.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
Co-authored-by: Jeff Rasley <jerasley@microsoft.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants