From 37ad5bbb889779211647dcef08b0e683a9d8dce7 Mon Sep 17 00:00:00 2001 From: Wanqiang Chen Date: Wed, 5 Aug 2026 15:53:52 -0700 Subject: [PATCH] Make BatchedUnaryEmbeddingBag constructible on meta/fake tensors Summary: `BatchedUnaryEmbeddingBag.split_embedding_weights()` sliced `self.weight` using tensor-valued bounds read out of the `table_offsets_tensor` buffer. Meta and fake tensors carry shape but no data, so those reads lower to `aten::item`. Under a dispatch mode that fabricates placeholder scalars for value-less reads, both slice bounds come back as the same constant, so every slice collapses to shape `(0, 1)` and the `assert param.shape == (num_emb, 1)` in `init_parameters()` fails during construction. Because a bare `assert` carries no message, this surfaces as a blank-message exception. `__init__` now computes the cumulative offsets once as a Python `list[int]` and derives `table_offsets_tensor` from it, giving a single source of truth. `split_embedding_weights()` slices with the Python ints, keeping shapes static. Behavior on real devices is unchanged - same offsets, same `uniform_` bounds, same per-table slices, bit-identical initialization for a given seed. The change additionally removes `2*N*T` device-to-host syncs per `split_embedding_weights()` call on CUDA, and makes the module constructible under `torch.device("meta")` and FakeTensor, which model-analysis and export tooling depend on. Differential Revision: D114802150 --- .../batched_unary_embeddings_ops.py | 22 ++++----- .../test/batched_unary_embeddings_test.py | 46 +++++++++++++++++++ 2 files changed, 57 insertions(+), 11 deletions(-) diff --git a/fbgemm_gpu/fbgemm_gpu/batched_unary_embeddings_ops.py b/fbgemm_gpu/fbgemm_gpu/batched_unary_embeddings_ops.py index 9835e15e23..bbc0049b8c 100644 --- a/fbgemm_gpu/fbgemm_gpu/batched_unary_embeddings_ops.py +++ b/fbgemm_gpu/fbgemm_gpu/batched_unary_embeddings_ops.py @@ -38,17 +38,17 @@ def __init__(self, num_tasks: int, hash_sizes: list[int], long_index: bool = Fal # [N][sum(E)][1] embedding_data = torch.randn(size=(num_tasks, sum(self.hash_sizes), 1)) self.weight = torch.nn.Parameter(embedding_data) + # Keep the offsets as Python ints as well as a buffer. Slicing the + # weight with the buffer's elements would read values off a tensor, + # which yields empty slices on meta / fake tensors (they carry shape + # but no data); the Python copy keeps split_embedding_weights() + # statically shaped so the module is constructible on those devices. + table_offsets = [0] + for hash_size in self.hash_sizes: + table_offsets.append(table_offsets[-1] + hash_size) + self.table_offsets: list[int] = table_offsets index_dtype = torch.int64 if long_index else torch.int32 - table_offsets_tensor = torch.cat( - [ - torch.tensor([0], dtype=index_dtype), - torch.cumsum( - torch.tensor(hash_sizes), - dim=0, - dtype=index_dtype, - ), - ] - ) + table_offsets_tensor = torch.tensor(table_offsets, dtype=index_dtype) self.register_buffer("table_offsets_tensor", table_offsets_tensor) self.init_parameters() @@ -71,7 +71,7 @@ def split_embedding_weights(self): embedding_weights.append( self.weight.detach()[ n, - self.table_offsets_tensor[t] : self.table_offsets_tensor[t + 1], + self.table_offsets[t] : self.table_offsets[t + 1], :, ] ) diff --git a/fbgemm_gpu/test/batched_unary_embeddings_test.py b/fbgemm_gpu/test/batched_unary_embeddings_test.py index 7f8cf5bc88..b5d0453974 100644 --- a/fbgemm_gpu/test/batched_unary_embeddings_test.py +++ b/fbgemm_gpu/test/batched_unary_embeddings_test.py @@ -239,6 +239,52 @@ def test_gpu(self) -> None: def test_cpu(self) -> None: self._test_main(gpu_infer=False) + def test_meta_device_construction(self) -> None: + """ + Constructing on meta must work: split_embedding_weights() slices the + weight per table, and meta tensors carry shape but no data. Slicing + with elements of table_offsets_tensor would read values off a + value-less tensor, collapsing every slice to length 0 and tripping the + shape assert in init_parameters(). The Python-int offsets keep the + slices statically shaped. + + Model-analysis tooling constructs modules under torch.device("meta") + to estimate FLOPs and parameter counts without allocating, which is + the path that motivated this test. + """ + hash_sizes = [100, 200] + num_tasks = 3 + with torch.device("meta"): + unary_emb = batched_unary_embeddings_ops.BatchedUnaryEmbeddingBag( + num_tasks=num_tasks, hash_sizes=hash_sizes, long_index=True + ) + + self.assertEqual(unary_emb.weight.shape, (num_tasks, sum(hash_sizes), 1)) + split_weights = unary_emb.split_embedding_weights() + self.assertEqual(len(split_weights), num_tasks * len(hash_sizes)) + # Order matches init_parameters()'s `hash_sizes * num_tasks` zip. + for i, param in enumerate(split_weights): + self.assertTrue(param.is_meta) + self.assertEqual(param.shape, (hash_sizes[i % len(hash_sizes)], 1)) + + def test_torchscript_scriptable(self) -> None: + """ + split_embedding_weights() and init_parameters() are @torch.jit.export, + so the module must stay scriptable. Guards the Python-int + table_offsets attribute against TorchScript inference regressions. + """ + hash_sizes = [100, 200] + num_tasks = 3 + unary_emb = batched_unary_embeddings_ops.BatchedUnaryEmbeddingBag( + num_tasks=num_tasks, hash_sizes=hash_sizes, long_index=True + ) + scripted = torch.jit.script(unary_emb) + + split_weights = scripted.split_embedding_weights() + self.assertEqual(len(split_weights), num_tasks * len(hash_sizes)) + for i, param in enumerate(split_weights): + self.assertEqual(param.shape, (hash_sizes[i % len(hash_sizes)], 1)) + @unittest.skipIf(*gpu_unavailable) # This test exercises the HIP launch-side limit and requires a large # output tensor (~17 GiB) plus offsets (~4 GiB) — total ~22 GiB GPU