With the current FSDP trainer, we can train large-scale knowledge graph embeddings on a single node with multiple GPUs — for example, a KG with roughly 200M entities and 1B triples on 4×H100 can be trained for 500 epochs in around 53 hours with an embedding dimension of 256. Extending FSDP to a multi-node setting would let us scale to even larger graphs. To do this, both the embedding tables and the optimiser states must be fully sharded across nodes, rather than each node handling its own copies of the embedding tables and optimiser states.
With the current FSDP trainer, we can train large-scale knowledge graph embeddings on a single node with multiple GPUs — for example, a KG with roughly 200M entities and 1B triples on 4×H100 can be trained for 500 epochs in around 53 hours with an embedding dimension of 256. Extending FSDP to a multi-node setting would let us scale to even larger graphs. To do this, both the embedding tables and the optimiser states must be fully sharded across nodes, rather than each node handling its own copies of the embedding tables and optimiser states.