Achieving Near-Linear Training Scalability for Pinterest’s Foundation Models

What this post added

This post details the journey to achieve near-linear training scalability for Pinterest's foundation models. It identifies the primary bottleneck as distributed embedding collective communication, exacerbated by large embedding tables. The key contributions include: 1. Quantized Communications (QComms) to reduce payload size by compressing FP32 to FP8. 2. Balanced Sharding to ensure even workload distribution across GPUs. 3. Bandwidth-Aware Embedding Optimization by reshaping embedding tables (halving dimension, doubling rows) to reduce data transfer. 4. A refined 2D Parallel topology, specifically the 'All-to-All Optimized' approach, which localizes expensive all-to-all operations within nodes and uses cheaper all-reduce for inter-node synchronization. These optimizations collectively improved 2-node scaling from 1.13x to 2.0x and 4-node scaling from 1.21x to 3.9x, extending to 8 nodes at 7.5x.