MLPerf DLRM-v2 train-perf benchmark on torch_tpu (meta-pytorch#4587)

Summary:

Adds `train_perf_mlperf_tpu.py`, a steady-state training-throughput benchmark for the faithful MLPerf DLRM-v2 model on the torch_tpu stack (torch side only, no JAX). Same model as the GPU `mlperf_dlrm` runner in D110238830: torchrec `DLRM_DCN` (DCNv2, 3 layers / low-rank 512) over a multi-hot `EmbeddingBagCollection` (pool sum 214), `embedding_dim` 128, dense arch `[512,256,128]`, over arch `[1024,1024,512,256,1]`, per-element Adagrad. Cardinalities selectable via `--cardinality` (`shrunk` sum 30M / `canonical` sum 228M).

Embeddings are row-wise sharded on the `UNFUSED_TPU` compute kernel under `DistributedModelParallel` over `tpu_dist`; the dense DCN/MLP path runs on TPU via torch_tpu. Reports steady-state ms/step and K samples/s/chip, the same metric as the jte TPU side, so it is a direct third comparison point for the GPU vs TPU per-chip gap-closure work (B200 GPU vs v7x jte-JAX vs v7x torch_tpu).

Faithful-model caveat: the full MLPerf model uses pooled, multi-hot embeddings, which produce an uneven row-wise all2all. The torch_tpu RW path was first brought up for `EmbeddingCollection` / 1-hot / even-split only, so running this faithful config may require pooled-multi-hot / uneven-all2all support that is still landing.

Differential Revision: D110794150

Author: kausv

torchrec/distributed/batched_embedding_kernel.py

@@ -4795,3 +4795,136 @@ def forward(self, features: KeyedJaggedTensor) -> torch.Tensor:
                 self._emb_modules[self._feature_table_map[feature_idx]](feature_values)
             )
         return outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=0)
+
+
+class BatchedTPUEmbeddingBag(BaseBatchedEmbeddingBag[torch.Tensor]):
+    """Pooled TPU compute kernel (`UNFUSED_TPU`).
+
+    Pooled sibling of `BatchedTPUEmbedding`, and a drop-in alongside
+    `BatchedDenseEmbeddingBag` in the pooled dispatch
+    (`GroupedPooledEmbeddingsLookup._create_embedding_kernel`). Backs the lookup with one
+    `TPUEmbeddingUnfused` table per table in the group; `forward` routes each feature to
+    its table, gathers its ids, and pools per sample (SUM or MEAN per `config.pooling`).
+
+    Pooling is done in torch (scatter-add over per-id sample indices) on top of the
+    unfused gather, so no dedicated pooled kernel is required and the backward flows
+    through the gather op's autograd.
+
+    Args:
+        config (GroupedEmbeddingConfig): grouped table config (one or more tables).
+        pg (Optional[dist.ProcessGroup]): process group (unused locally).
+        device (Optional[torch.device]): compute device (e.g. ``"tpu"``).
+        sharding_type (Optional[ShardingType]): sharding type (for the pooling mode).
+
+    Example::
+
+        kernel = BatchedTPUEmbeddingBag(grouped_config, device="tpu")
+        out = kernel(features)  # [batch, sum(embedding_dim over features)]
+    """
+
+    def __init__(
+        self,
+        config: GroupedEmbeddingConfig,
+        pg: Optional[dist.ProcessGroup] = None,
+        device: Optional[torch.device] = None,
+        sharding_type: Optional[ShardingType] = None,
+    ) -> None:
+        super().__init__(config, pg, device, sharding_type)
+        # Lazy import to avoid pulling in experimental TPU code at module load time.
+        from torchrec.experimental.torch_tpu.modules.embedding_modules import (
+            TPUEmbeddingUnfused,
+        )
+
+        dtype = data_type_to_sparse_type(config.data_type).as_dtype()
+        # One TPUEmbeddingUnfused per table in the group (mirrors BatchedTPUEmbedding);
+        # the base class fills _local_rows / _local_cols / _feature_table_map.
+        self._emb_modules: nn.ModuleList = nn.ModuleList()
+        for local_rows, local_cols in zip(self._local_rows, self._local_cols):
+            self._emb_modules.append(
+                TPUEmbeddingUnfused(
+                    num_embeddings=local_rows,
+                    embedding_dim=local_cols,
+                    device=device,
+                    dtype=dtype,
+                )
+            )
+        self.init_parameters()
+
+    @property
+    # pyrefly: ignore [bad-override]  # TPUEmbeddingUnfused is not one of the
+    # fbgemm codegen types the base class enumerates (same as the Triton kernel).
+    def emb_module(self) -> "TPUEmbeddingUnfused":  # noqa: F821
+        # pyre-ignore[16]: ModuleList returns Module, pyre thinks Union
+        return self._emb_modules[0]
+
+    def split_embedding_weights(self) -> List[torch.Tensor]:
+        # pyre-ignore[16]
+        return [emb_module.weight for emb_module in self._emb_modules]
+
+    def named_split_embedding_weights(
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
+    ) -> Iterator[Tuple[str, torch.Tensor]]:
+        assert (
+            remove_duplicate
+        ), "remove_duplicate=False not supported in named_split_embedding_weights"
+        for table, emb_module in zip(self._config.embedding_tables, self._emb_modules):
+            # pyre-ignore[16]
+            yield append_prefix(prefix, f"{table.name}.weight"), emb_module.weight
+
+    def named_parameters(
+        self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True
+    ) -> Iterator[Tuple[str, nn.Parameter]]:
+        for table, emb_module in zip(self._config.embedding_tables, self._emb_modules):
+            # pyre-ignore[7]
+            yield append_prefix(prefix, f"{table.name}.weight"), emb_module.weight
+
+    def _pool(
+        self, gathered: torch.Tensor, lengths: torch.Tensor, batch_size: int
+    ) -> torch.Tensor:
+        """Pool a feature's per-id embeddings [L, dim] into [batch, dim].
+
+        Sum-pool by scatter-adding each id's row into its sample slot; MEAN divides by
+        the bag length. `lengths` are this feature's per-sample bag sizes (may be jagged
+        after the row-wise all2all), summing to L. Uses only arange / repeat_interleave /
+        index_add_ / div, so it lowers on TPU without a bespoke pooled kernel.
+        """
+        dim = gathered.shape[1]
+        device = gathered.device
+        pooled = torch.zeros(batch_size, dim, dtype=gathered.dtype, device=device)
+        # sample index for each gathered row: sample b repeated lengths[b] times.
+        sample_idx = torch.repeat_interleave(
+            torch.arange(batch_size, device=device), lengths
+        )
+        pooled.index_add_(0, sample_idx, gathered)
+        if self._pooling == PoolingMode.MEAN:
+            denom = lengths.clamp(min=1).unsqueeze(1).to(pooled.dtype)
+            pooled = pooled / denom
+        return pooled
+
+    # pyrefly: ignore [bad-override]  # same signature shape as the sequence kernel
+    def forward(self, features: KeyedJaggedTensor) -> torch.Tensor:
+        batch_size = features.stride()
+        length_per_key = features.length_per_key()
+        # KJT lengths are feature-major: feature i's per-sample bag sizes are the
+        # i-th block of batch_size entries.
+        all_lengths = features.lengths()
+        per_feature_values = (
+            torch.split(features.values(), length_per_key)
+            if len(length_per_key) > 1
+            else [features.values()]
+        )
+        outputs: List[torch.Tensor] = []
+        for feature_idx, feature_values in enumerate(per_feature_values):
+            emb_module = self._emb_modules[self._feature_table_map[feature_idx]]
+            # pyre-ignore[9, 16]: ModuleList returns Module, pyre thinks Union;
+            # the table's `weight` is a Tensor so `.device` is a torch.device.
+            device: torch.device = emb_module.weight.device
+            gathered = emb_module(
+                feature_values.to(device=device, dtype=torch.int32)
+            )  # [L_i, dim]
+            lengths_i = all_lengths[
+                feature_idx * batch_size : (feature_idx + 1) * batch_size
+            ].to(device)
+            outputs.append(self._pool(gathered, lengths_i, batch_size))
+        # Pooled features concatenate along the embedding dim -> [batch, sum(dim)].
+        return outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=1)

torchrec/distributed/embedding_lookup.py

@@ -49,6 +49,7 @@
     BatchedFusedEmbedding,
     BatchedFusedEmbeddingBag,
     BatchedTPUEmbedding,
+    BatchedTPUEmbeddingBag,
     KeyValueEmbedding,
     KeyValueEmbeddingBag,
     ShardedBatchedFusedEmbedding,
@@ -731,6 +732,13 @@ def _create_embedding_kernel(
                 pg=pg,
                 device=device,
             )
+        elif config.compute_kernel == EmbeddingComputeKernel.UNFUSED_TPU:
+            return BatchedTPUEmbeddingBag(
+                config=config,
+                pg=pg,
+                device=device,
+                sharding_type=sharding_type,
+            )
         elif config.compute_kernel in {
             EmbeddingComputeKernel.KEY_VALUE,
         }:

torchrec/distributed/embeddingbag.py

@@ -1381,6 +1381,24 @@ def _initialize_torch_state(self, skip_registering: bool = False) -> None:  # no
                                 self._table_name_to_config[table_name].data_type
                             )
                         ),
+                        # Set grad-ness so ShardedTensor init does not reject an
+                        # unfused (autograd-trained) kernel whose shard is a
+                        # grad-requiring nn.Parameter. Derive it from the compute
+                        # kernel rather than the local shard: the kernel comes from
+                        # the sharding plan and is identical on every rank, whereas
+                        # probing local_shards would record False on ranks that hold
+                        # no shard (e.g. table-wise), yielding globally inconsistent
+                        # TensorProperties for the same logical tensor. Only the
+                        # unfused kernels (DENSE, UNFUSED_TPU) keep grad-requiring
+                        # weights; fused/KV/quant kernels train via in-backward
+                        # optimizers on non-grad buffers.
+                        requires_grad=(
+                            _model_parallel_name_to_compute_kernel[table_name]
+                            in {
+                                EmbeddingComputeKernel.DENSE.value,
+                                EmbeddingComputeKernel.UNFUSED_TPU.value,
+                            }
+                        ),
                     )
 
                     sharded_tensor_metadata = sharding_spec.build_metadata(