Added concatenated layout support - #46
IlyasMoutawwakil wants to merge 2 commits into
Conversation
|
cc @tridao |
|
this fix is pretty interesting, it's simpler than I thought (by just changing the layout of the weight tensor). |
|
we'll need to think through how this affects the rest of the kernels (esp bwd). But if it's just about changing the tensor layout (and eating the extra TMA instruction) I think we should support this concat layout! |
|
@GarlGuo is about to land a big refactor in the next couple of days so if we're to merge it it'll be after the refactor. |
|
confirming that concat layout doesnt' even generate extra TMA instruction, it just changes TMA 3D to TMA 4D. That's why perf stays the same. This is great! |
|
I will block this PR until #39 is merged (which might happen tmrw). After that I will work on the concatenated layout support. |
|
i have sth in quack to support concat layout, will open PR soon |
|
@IlyasMoutawwakil i adapted this to quack (Dao-AILab/quack#107), adding you as co-author if that's ok. |
|
@IlyasMoutawwakil I have created a new PR as #47. I will also add you as a co-author if that's ok. |
|
awesome ! works for me and thank you both 🙏 |
|
@IlyasMoutawwakil #47 has been merged and is available on the latest pypi package: https://pypi.org/project/sonic-moe/ |
|
Awesome!! We will start integrating |
|
@IlyasMoutawwakil @vasqu I will close this PR and issue #45 for now. Please let me know if more help is needed. |
|
Sounds good, can you also ref the PR in the issue for others that might stumble on it? |
|
@vasqu Sounds good. I have left a note. |
We wanna add sonic-moe support in Transformers, but the experts layout we are using in our modeling across MoEs is concatenated [gate;up] and not the interleaved one (see #45) changing the weight layout to interleaved breaks a couple things for us and seems to be an adoption pain point in other projects as well so we are suggesting this PR 🤗