Skip to content

Commit 9493c2c

Browse files
houseroadfacebook-github-bot
authored andcommitted
Treat meta as CPU device, so we can run the meta on CPU with MetaSparseDispatchMode (#1617)
Summary: Pull Request resolved: #1617 When we use meta, it's not treated as CPU tensor. But the use case is that we want to process the weights on CPU first. So we set use_cpu as true when device is true in `IntNBitTableBatchedEmbeddingBagsCodegen` Reviewed By: jianyuh, nanoax Differential Revision: D43633336 fbshipit-source-id: bfb5f52c17f6ecf8bedb4f47a85aabcf35b87416
1 parent 8083b9b commit 9493c2c

File tree

1 file changed

+4
-0
lines changed

1 file changed

+4
-0
lines changed

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1799,6 +1799,10 @@ def __init__(
17991799
dims: List[int] = [e[2] for e in embedding_specs]
18001800
weights_tys: List[SparseType] = [e[3] for e in embedding_specs]
18011801
locations: List[EmbeddingLocation] = [e[4] for e in embedding_specs]
1802+
# if target device is meta then we set use_cpu based on the embedding location
1803+
# information in embedding_specs.
1804+
if self.current_device.type == "meta":
1805+
self.use_cpu = all(loc == EmbeddingLocation.HOST for loc in locations)
18021806

18031807
if row_alignment is None:
18041808
self.row_alignment: int = 1 if self.use_cpu else 16

0 commit comments

Comments
 (0)