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