[Pipe] Tied layers

[stas00]

🚀 Feature

Continuing the requests to support various needs of the models in the new Pipe pytorch feature, this one brings up

Tied layers

For models that have tied weights (e.g. encoder/decoder in transformer-type of models. It's important to:

• avoid memory duplication
• handle the gradients correctly (have to be reduced twice in the typical encoder+decoder transformer scenario)

You can see this feature provided in DeepSpeed Tied Layers and its discussion.

Thank you.

@pritamdamania87

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @agolynski @SciPioneer @H-Huang @mrzzd @cbalioglu @gcramer23

[stas00]

@msbaines wrote in #50693 (comment)

Pipe supports shared-parameters by default as long the partitions with the shared-parameter are mapped to the same device:

pytorch/torch/distributed/pipeline/sync/pipe.py

Line 87 in 2303c24

raise ValueError("module with duplicate parameters on distinct devices is not supported")

Exactly, which makes it very difficult to accomplish for the user. It's not possible to accomplish with automatic partitioning, other than doing it totally manually and trying to handpick which layer should go where.

That's why I've opened this issue - there must be an easier way to tell the Pipe that this is a shared weight and let the Pipe figure out the details.

What type of models are you using shared-parameters in? So far, we've only seen shared-parameters being used for the shared-embedding table used in Machine Translation models.

all of HF transformers, so yes MT is one of them.

[msbaines]

@msbaines wrote in #50693 (comment)

Pipe supports shared-parameters by default as long the partitions with the shared-parameter are mapped to the same device:

pytorch/torch/distributed/pipeline/sync/pipe.py

Line 87 in 2303c24

raise ValueError("module with duplicate parameters on distinct devices is not supported")

Exactly, which makes it very difficult to accomplish for the user. It's not possible to accomplish with automatic partitioning, other than doing it totally manually and trying to handpick which layer should go where.

Pipe currently does not perform automatic partitioning but it already knows when two layers are sharing weights. So I'm not sure this mechanism is needed.

That's why I've opened this issue - there must be an easier way to tell the Pipe that this is a shared weight and let the Pipe figure out the details.

What type of models are you using shared-parameters in? So far, we've only seen shared-parameters being used for the shared-embedding table used in Machine Translation models.

all of HF transformers, so yes MT is one of them.

[msbaines]

I think this feature would be useful in the context of adding something like Memory-Efficient Model Construction.

[mrzzd]

I think support for Tied layers does not need to come with the pipeline code. All-reduce happens after processing each batch, so it can be invoked in main trainer loop, which is part of the user code. Is my understanding correct?

[stas00]

Tied layers are part of the model and are not part of the trainer loop. Users never interact with those directly.

[msbaines]

Just to clarify. Pipe already knows which parameters are shared since a complete input model is passed in the constructor. So no additional annotation is necessary. If support for a deferred Module allocation mechanism like "Memory-Efficient Model Construction" were supported, where we don't have the complete model available, annotations such as TiedLayer may be necessary. But with LazyModuleMixin and UninitializedParameter it might be possible to implement something where you can access the whole model in a memory-efficient way.

[mrzzd]

Tied layers are part of the model and are not part of the trainer loop. Users never interact with those directly

Totally agreed, All I was saying is that in the absence of support for Tied layers in pipeline framework, the user can invoke necessary means to make it functional in the trainer loop.

[stas00]

All I was saying is that in the absence of support for Tied layers in pipeline framework, the user can invoke necessary means to make it functional in the trainer loop.

If I'm understanding you correctly, you're saying that there is a technical way to solve this issue during training. i.e. you're saying it's feasible.

But even so how would a user know that they need to do that? General users deal with models on the opaque level and (almost?) never interact with its internals.