Unverified Commit a454734f authored by Hongzhi (Steve), Chen's avatar Hongzhi (Steve), Chen Committed by GitHub
Browse files

[Misc] Follow up on the missing comment in PR#4989. (#5437)


Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-28-63.ap-northeast-1.compute.internal>
Co-authored-by: default avatarXin Yao <yaox12@outlook.com>
parent 8d5d8962
......@@ -168,11 +168,14 @@ class NodeEmbedding: # NodeEmbedding
Target device to put the collected embeddings.
"""
if not self._comm:
# embeddings are stored on the CPU
# For embeddings stored on the CPU.
emb = self._tensor[node_ids].to(device)
else:
# embeddings are stored on the GPU
# the following method also covers self._world_size = 0 or 1
# For embeddings stored on the GPU.
# The following method is designed to perform communication
# across multiple GPUs and can handle situations where only one GPU
# is present gracefully, a.k.a. self._world_size == 1 or
# 0 (when th.distributed.is_initialized() is false).
emb = nccl.sparse_all_to_all_pull(
node_ids, self._tensor, self._partition
)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment