support different entry size for different ranks

rapidsai / wholegraph

WholeGraph - large scale Graph Neural Networks

https://docs.rapids.ai/api/cugraph/stable/wholegraph/

Apache License 2.0

100 stars 38 forks source link

support different entry size for different ranks #194

Closed zhuofan1123 closed 3 months ago

zhuofan1123 commented 4 months ago

Allow users to specify the entry size on each rank.

    node_feat_wm_embedding = wgth.create_embedding(
        ...
        embedding_entry_partition=[283071, 401722, 356680, 329221, 238065, 238060, 217897, 384313]
    )

embedding_entry_partition[i] indicates the number of embedding entries stored on the rank i.
If embedding_entry_partition is None, embedding will be partitioned equally.
Only chunked device and distributed host/device are supported.

copy-pr-bot[bot] commented 4 months ago

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.