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

[MoFHeka]

Because TF can not train a Hash Table directly, so I put the TrainableWrapper and DEResourceVariables as agent. When forward running, looking up values from hash table and fill the value into TrainableWrapper, besides storing IDs in DEResourceVariables for backward running. So each TrainableWrapper and DEResourceVariable is different from one worker rank to another rank.

It's OK? May I close this issue?