Does HvdAllToAllEmbedding shard embeddings across ranks? Or is it copied on all ranks?

[alykhantejani]

hi,

I'm using

  self.kv_creator = de.CuckooHashTableCreator(
            saver=de.FileSystemSaver(proc_size=1, proc_rank=hvd.local_rank())
        )

        self.emb = de.keras.layers.HvdAllToAllEmbedding(..., devices=['CPU'])

And I notice in your keras callback for horovod you remove the TrainableWrapper and DEResourceVariables from the horovod broadcats.

This leads me to believe that dense parts of the models are broadcasted and then synchronously updated, but the embeddings are just updated on each device? This would mean the embedding ids are sharded (via the CuckooHashTableCreator) to a specific rank

The reason Im curious is because each of my worker ranks are using the same ammount of host memory which makes me think they are all sharing the same embeddings?

[alykhantejani]

Hi, just following up on this?

[MoFHeka]

shard embeddings across ranks according partition_fn