Loading…
Achieving Near-Linear Training Scalability for Pinterest’s Foundation Models
PinterestPinterest Engineering
Summary
Pinterest’s foundation ranking models store approximately 99% of their parameters in embedding tables, causing cross-node NCCL all-to-all communication to severely bottleneck multi-node distributed training. Initial multi-node scaling stalled at 1.21x on four nodes even after adopting AWS Elastic Fabric Adapter for OS-bypass networking. To resolve this, engineers profiled GPU activity with PyTorch Profiler and systematically addressed communication bottlenecks across multiple layers. The team implemented FP8 quantized communications using FBGEMM, balanced table sharding across GPUs, halved embedding dimensions while doubling row counts, and adopted a 2D parallel communication topology alongside PyTorch Distributed Checkpoint. These combined interventions achieved near-linear scaling factors of 3.9x on four nodes and 7.5x on eight nodes across 64 GPUs.
Context
Multi-node distributed training for Pinterest's embedding-heavy foundation ranking models, where roughly 99% of parameters reside in embedding tables, initially suffered severe communication bottlenecks. Adding a second node degraded training throughput to 0.2x of a single node without OS-bypass networking. Even after enabling AWS Elastic Fabric Adapter (EFA), multi-node scaling remained inefficient at 1.13x for two nodes and 1.21x for four nodes because GPUs spent substantial time idle while waiting for cross-node NCCL all-to-all embedding exchanges.
Approach / What changed
Engineers used PyTorch Profiler and NCCL traces to isolate communication delays and deployed several layered optimizations: compressing embedding wire payloads from FP32 to FP8 using FBGEMM quantized communications (QComms), balancing table-wise sharding across GPUs, halving embedding dimensions while doubling row counts to reduce bytes per collective, and implementing 2D parallel training topology. Additionally, the infrastructure upgraded from TorchSnapshot to PyTorch Distributed Checkpoint (DCP) to support load-time resharding across varying world sizes.
Takeaways
- Compressing embedding communication from FP32 to FP8 with FBGEMM QComms reduced the largest NCCL SendRecv operation by over 75% without compromising model quality or training loss convergence.
- Halving embedding dimensions while doubling row counts maintained model capacity while reducing communication payload volume, elevating scaling factors to 1.78x on two nodes and 2.8x on four nodes.
- Combining quantized communication, balanced sharding, payload reshaping, and 2D parallelism increased distributed scaling from 1.21x to 3.9x on 4 nodes and reached 7.5x on 8 nodes (64 GPUs).
Related reading
Scaling Conditional Learned Retrieval for Pinterest Home Feed
Pinterest home feed candidate generation relies on large-scale retrieval, but standard two-tower models produce single embeddings that struggle to capture diverse concurrent user intentions. To solve this, Pinterest expanded Conditional Learned Retrieval (CLR) to generate condition-aware embeddings conditioned on explicit interest, Pin, and Board contexts. The team evolved model representations by appending condition tokens to user sequences in transformers, integrating the PinFM foundation model, and introducing hierarchical semantic IDs to improve cold-start coverage. On the serving side, engineers consolidated duplicate feature requests and shifted inference to GPU instances using NVEmbed under a one-query-to-many-documents paradigm. These infrastructure and modeling changes unlocked new condition types, achieved seven-figure cost savings, and reduced p90 model retrieval latency from 80ms down to 12ms.
Pinterest Engineering