Problem
Piper does not correctly support tied embeddings in the presence of pipeline parallelism. Currently, when tied embeddings are placed on separate devices, the embedding parameters are duplicated on each actor and their gradients are not synchronized before the optimizer step.
Expected behavior
Tied embeddings should be treated as though there is a single copy of the parameters, regardless of the distribution strategy. The solution should add minimal network overhead while avoiding duplication, where possible.
Implementation notes
The suggested solution is to duplicate tied embeddings which appear on distinct actors, and insert a gradient reduction on the duplicates after both have completed the backward pass and before the optimizer step.
Tied embeddings which appear in different stages may not appear on distinct devices (e.g. in a V-shaped stage placement, tied embeddings can appear in the first and last stages which belong to the same device). In this case, duplication and synchronization can be avoided (the current Piper behavior will create a duplicate for each stage, regardless of device placement).
The solution can likely be implemented entirely in the backend (transparent to the user). The most straightforward implementation path is probably to place tied embeddings in their own buckets because TrainingDAG nodes operate on buckets, so this should make it easy to e.g. add a reduction node for the embedding buckets after the embeddings' backward compute nodes.
You will likely need to make a new process group for each pair of devices that share tied embeddings.
Problem
Piper does not correctly support tied embeddings in the presence of pipeline parallelism. Currently, when tied embeddings are placed on separate devices, the embedding parameters are duplicated on each actor and their gradients are not synchronized before the optimizer step.
Expected behavior
Tied embeddings should be treated as though there is a single copy of the parameters, regardless of the distribution strategy. The solution should add minimal network overhead while avoiding duplication, where possible.
Implementation notes
The suggested solution is to duplicate tied embeddings which appear on distinct actors, and insert a gradient reduction on the duplicates after both have completed the backward pass and before the optimizer step.
Tied embeddings which appear in different stages may not appear on distinct devices (e.g. in a V-shaped stage placement, tied embeddings can appear in the first and last stages which belong to the same device). In this case, duplication and synchronization can be avoided (the current Piper behavior will create a duplicate for each stage, regardless of device placement).
The solution can likely be implemented entirely in the backend (transparent to the user). The most straightforward implementation path is probably to place tied embeddings in their own buckets because
TrainingDAGnodes operate on buckets, so this should make it easy to e.g. add a reduction node for the embedding buckets after the embeddings' backward compute nodes.You will likely need to make a new process group for each pair of devices that share tied embeddings.