2D Parallelism for Recommendation Model Training

2D parallelism for recommendation model training changes the topology of distributed training so high-volume embedding lookup all-to-all stays inside local GPU groups, while cross-node communication synchronizes replicas. This note assumes Distributed Training for Embedding-Heavy Recommenders and GPU Collective Communication.

The problem it solves is not “how do we split a model at all?” The previous notes already covered sharding. The problem is: if global sharding makes every lookup exchange cross the whole cluster, can we arrange ranks so the expensive lookup shuffle stays local?

Topology Shift

This diagram contrasts global sharding with local replica groups. The left side has one sharded model spread across nodes. The right side has one complete replica per node, with each replica internally sharded across local GPUs.

Global sharding spreads one model replica across every GPU in the cluster. It maximizes aggregate memory for one replica, but lookup all-to-all can cross nodes for every training step.

2D parallelism uses two axes:

Axis	What it does	Communication shape
Model-parallel axis inside a group	shards one replica across local GPUs	local lookup all-to-all
Data-parallel axis across groups	runs several replicas on different batch slices	cross-group gradient synchronization

For a four-node job with eight GPUs per node, each node can be one group. The eight local GPUs shard one complete replica and serve lookup exchange locally. The four node-level replicas process different batch slices and synchronize updates across nodes.

Mapping Example

Without 2D parallelism, one training batch on node A can need rows owned by any GPU in the cluster:

global sharding:
  node A rank 0 needs board_789
  owner(board_789) = node C rank 18
  lookup vector crosses nodes during the forward pass

With one replica group per node, each node has local owners for the rows in that replica:

2D layout:
  node A rank 0 needs board_789
  owner_inside_node_A(board_789) = node A rank 6
  lookup vector stays inside node A
 
  node B processes a different batch slice
  node C processes a different batch slice
  replicas synchronize updates across nodes later

The lookup payload did not disappear. It moved from slow cross-node paths to faster local paths. Cross-node bandwidth is then used for replica synchronization, which is usually more regular than per-id lookup exchange.

Memory Tradeoff

2D parallelism spends memory to reduce cross-node lookup traffic. Each group must hold a complete model replica after internal sharding.

training-time model data: 640 GB
node shape: 8 GPUs x 80 GB = 640 GB
 
one node can hold one sharded replica before overhead
four nodes can hold four replicas

If one replica needs 1 TB and a node has only eight 80 GB GPUs, one-node groups do not fit. The system then needs larger groups, more compression, offload, smaller tables, or a different sharding plan.

Pinterest Example

Pinterest first used a 2D layout that separated local model-parallel sharding from cross-node data-parallel replication. They then optimized rank placement around all-to-all locality rather than all-reduce convenience. That means ranks exchanging high-volume embedding vectors were placed on local links even if another rank order would have made gradient synchronization simpler.

The article reports all-to-all latency dropping from 78 ms to 13 ms. Scaling reached 2.0x at two nodes and 3.9x at four nodes. The later eight-node run reached 7.5x, or 93.75% of ideal linear scaling.

Sources

• Pinterest Engineering - Achieving Near-Linear Training Scalability for Pinterest’s Foundation Models
• NVIDIA NCCL User Guide - Collective Operations

See also

• Distributed Training for Embedding-Heavy Recommenders
• GPU Collective Communication
• Training Scalability for Distributed ML
• Quantized Communication for Distributed Training