[QUESTION]您好，请教一下如果我需要在moe_ag的allgather通信之后对ctx.inputs增加一个transpose(0,1)，需要修改哪处代码？

Open jinchen89 opened this issue 8 months ago • 1 comments

我理解allgather得到ctx.inputs后第一个维度应该是tpsize，如果直接reshape成第一个维度是ctx.nexperts_ep组内的专家数的话是不是不太合逻辑，是否需要先把组内专家数那个维度先转置到第一个维度？

flux/src/moe_ag_scatter/ths_op/gemm_grouped_v3_ag_scatter.cc 我能不能直接在这个文件的all_gather_all2all函数的最后对input_buffer做transpose，我测下来推理第二次就会挂掉？

Apr 15 '25 02:04 jinchen89

是要在op内做啥操作吗？我看op->run之前所有tp rank的数据还没allgather完成@wenlei-bao

Apr 24 '25 05:04 jinchen89