Per-sharding-type EBC benchmark workloads

[siqihuang]

Summary:
The EBC benchmark ran all sharding types on one table-wise shape (256 tables x 1M rows x 256 dim). That shape is unrepresentative for row-wise and column-wise: row-wise reduce-scatters partial sums for all 256 tables (8x table-wise's output bytes on 8 ranks), and column-wise splits 256 dims into 32-dim shards, below the planner's MIN_CW_DIM (128).

This diff gives each sharding type its own default shape (_DEFAULT_SHAPES). Explicit kwargs still override every field.

How the defaults are chosen

1. Use case: each shape models the case that sharding type is picked for.
   • Table-wise: many medium tables.
   • Row-wise: a few large, lookup-heavy tables, i.e. long-sequence features with high pooling.
   • Column-wise: a few wide tables, where 4096 dims keeps every shard at or above MIN_CW_DIM on up to 32 ranks.
   • Row-wise and column-wise keep fewer tables than ranks; with more tables, table-wise would be the better choice.
2. Matched load: the table count is fixed by the use case, so batch, pooling and dim make up the load. Every shape looks up the same embedding bytes per iteration per rank (batch * tables * pooling * dim * 4 B = 21.5 GB), so latencies compare directly.
3. Memory: row count does not change traffic, so it is sized to keep peak memory under about 43 GB per GPU on 8 ranks, which fits 80 GB H100/A100.

Sharding type	Tables x rows x dim	Pooling	Batch
table_wise (unchanged)	256 x 1M x 256	20	4096
row_wise	4 x 50M x 256	320	16384
column_wise	4 x 2.5M x 4096	20	16384

Results (8 x H100, unpipelined). Normalized throughput is QPS x tables * pooling * dim * 4 B, the embedding bytes served per second per rank. Raw QPS counts samples, and a sample costs different work in each shape, so only the normalized figure compares across sharding types.

Sharding type	GPU time (P50)	Normalized throughput
table_wise	51.1 ms	0.42 TB/s
row_wise	42.2 ms	0.51 TB/s
column_wise	49.5 ms	0.44 TB/s

It also doubles unpipelined num_benchmarks to 200 and raises the primitive benchmarks' default to 100.

Differential Revision: D123169610

[meta-codesync]

@siqihuang has exported this pull request. If you are a Meta employee, you can view the originating Diff in D123169610.