search

tensorflow / recommenders-addons

Additional utils and helpers to extend TensorFlow when build recommendation systems, contributed and maintained by SIG Recommenders.
Apache License 2.0
587 stars 132 forks source link

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

Open alykhantejani opened 2 months ago

alykhantejani commented 2 months ago

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 commented 1 month ago

Hi, just following up on this?