Reproduce “TerraceMoE: A Cost Model for Hierarchical MoE All-to-All Communication” on OpenResearch

About this paper

Weicheng Xue, Bingqiang Wang, Li Yuan, Huihui Zhou, Yonghong Tian

TerraceMoE introduces an accountable cost model for hierarchical Mixture-of-Experts (MoE) all-to-all communication, enabling prediction of dispatch restructuring viability from microbenchmarks to mitigate high empirical experimentation costs. The model, validated at the communication-call level, quantitatively demonstrates that implementation details like the 'arrival chain' significantly impact the effective breakeven hierarchy ratio, while its end-to-end step time predictions were explicitly withdrawn due to gate failures.

Key insights

• Implementation overhead, particularly the 'arrival chain' (local tensor work at representative ranks), is a dominant factor, significantly increasing the effective breakeven hierarchy ratio from a byte-only prediction of 1.31 to 3.98 with the measured PyTorch chain.
• The sum of fixed costs for the two collectives in hierarchical dispatch, along with the non-data-movement portions of the arrival chain, are crucial 'invisible' terms often missed by simple hierarchy ratio comparisons.
• A gate-based validation framework provides explicit accountability for performance models by defining and enforcing their predictive capabilities and limitations, which is crucial for making high-stakes 'go/no-go' engineering decisions based on negative predictions.

Summary by alphaXiv.