KOEN
Read Original
#AI

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

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

Summary

How Pinterest Achieved 93.75% Training Efficiency on 8-Node Clusters

Breaking the communication bottleneck in embedding-heavy foundation models with 2D Parallel and QComms.

Scaling recommendation models is notoriously difficult due to massive embedding tables. This article explores how Pinterest transformed a failing multi-node setup into a near-linear scaling powerhouse. Learn how they optimized every layer of the communication stack to unlock larger, more capable models.

  • 01Recovered baseline performance using AWS EFA and OS-bypass networking
  • 02Reduced NCCL communication payload by 75% via FP8 quantization (QComms)
  • 03Implemented All-to-All optimized 2D Parallel topology to utilize intra-node NVLink
  • 04Reshaped embedding dimensions to minimize bytes-on-the-wire without losing capacity
  • 05Migrated to Distributed Checkpoint (DCP) for flexible load-time resharding

RECOMMENDATION

Essential reading for ML engineers working on large-scale recommendation systems. If your distributed training throughput isn't scaling with more GPUs, the communication-centric optimizations detailed here provide a proven playbook.

The Problem

Pinterest's foundation models, characterized by embedding-heavy architectures, exceeded single-GPU memory and faced severe performance degradation in multi-node setups due to network communication bottlenecks during distributed embedding lookups.

The Solution

The team implemented a multi-layered optimization strategy: enabling AWS EFA for OS-bypass, adopting Quantized Communications (FP8), reshaping embedding dimensions for bandwidth efficiency, and redesigning the 2D Parallel topology to keep All-to-All communication intra-node.

The Result

The optimizations achieved near-linear scaling factors of 3.9x on 4 nodes (97.5% efficiency) and 7.5x on 8 nodes (93.75% efficiency), resulting in a 13x total throughput increase from the initial pre-EFA baseline.

Trade-off

Certain optimizations like Balanced Sharding were constrained by serving-side requirements, and frequent PyTorch/TorchRec upgrades introduced significant operational overhead due to library version conflicts and memory profiler issues.

03

Key Concepts

Concept · 01

TorchRec & DistributedModelParallel

A PyTorch domain library designed for large-scale recommendation models, enabling efficient sharding of massive embedding tables across multiple GPUs.

  • Used to manage embedding tables that constitute 99% of Pinterest's foundation model parameters across GPU clusters.
Concept · 02

2D Parallel (All-to-All Optimized)

A hybrid parallelism strategy that organizes GPUs into groups to optimize specific communication patterns based on hardware interconnects.

  • Pinterest flipped the topology to confine expensive All-to-All lookups within nodes (NVLink) while syncing replicas via cross-node AllReduce.
Concept · 03

Quantized Communications (QComms)

A technique to compress communication payloads into lower-precision formats like FP8 to reduce network bandwidth usage.

  • Utilized FBGEMM to shrink NCCL SendRecv operations by over 75%, significantly boosting the scaling factor for 4-node clusters.