Revise LayerDrop implementation #484
Conversation
facebook-github-bot
added
the
CLA Signed
| drop_p: float = 0.0, | ||
| generator: Optional[Generator] = None, |
There was a problem hiding this comment.
Shouldn't we remove these as well?
There was a problem hiding this comment.
You mean the generator? It should be still around in case someone wants to provide a different RNG for layerdrop
|
|
||
| # compat | ||
| @final | ||
| class ModuleList(TorchModuleList): |
There was a problem hiding this comment.
What is the purpose of this class? Why don't we just use torch.nn.ModuleList instead of this? I don't see this even used anywhere, I propose for this file to be removed.
There was a problem hiding this comment.
We have several teams using this module for now. Everything tagged with “# compat” will eventually get removed once we migrate those uses (before v0.3 release)
| def backward(ctx: Any, grad_output: Tensor) -> Tuple[Tensor, Tensor]: | ||
| return grad_output, torch.zeros_like(grad_output) |
There was a problem hiding this comment.
Smartly done! The gradient with respect to x is going to be just the same as the gradient of the output, since this is just the identity function when you drop layers.
There was a problem hiding this comment.
Also, what is the advantage of using this method over PyTorch hooks?
Did this happen because rng across the different gpus was different - causing different layers to be dropped across gpus ? |
| @@ -204,6 +221,19 @@ def forward( | |||
|
|
|||
| return seqs, padding_mask | |||
|
|
|||
| def _drop_iter(self) -> Iterator[Tuple[Module, bool]]: | |||
There was a problem hiding this comment.
I think this logic (and some other pieces in the forward pass) are used in decoder as well. Maybe we could consider adding a base component in the future (that both encoder and decoder would inherit from)
This PR revises the LayerDrop implementation and moves its from
ModuleListtoStandardTransformerEncoderandStandardTransformerDecoder. Although the original implementation was ideal, both DDP and FSDP had trouble correctly handling forward/backward passes and either silently failed to sync gradients (DDP), or failed with some cryptic error. This implementation causes redundant computation of dropped layers, but since the autograd graph stays constant, both DDP and FSDP can handle it correctly.