Remove parameters from autograd backward hook

[Honzys]

Hello,

I am trying to remove some layers from DistributedDataParallel to prevent them being synchronized between devices.

I spent last 6 hours googling, and I have found out, that there's a attribute _ddp_params_and_buffers_to_ignore which can be set to module that is passed to DistributedDataParallel constructor. I've implemented custom strategy plugin to Trainer, I have checked that the parameters are then passed to a parameters_to_ignore attribute of the DistributedDataParallel but somehow if I check gradients, of the layer, they are always the same.

Is there some simpler way to remove some layer / module from being synchronized between more devices in DDP strategy?

Thank you in advance for any help!

[rohitgr7]

hey @Honzys !

can you confirm whether the list of the param names you have set to be ignored are same as that of then your model is initialized in LightningModule:

class LitModel(LightningModule):
    def __init__(self, ...):
        super().__init__()
        self.model = ...
        self.named_parameters() # extract key names from this and check whether the list of keys are exactly same

for reference, you can check out the test here where this feature was integrated: https://github.com/pytorch/pytorch/pull/44826/files

[Honzys]

Hi!
First of all thank you for your response!
The trick is, I have to create the model after the DDP is initialized. So right now I am constructing the models right before configure_optimizers is called. I didn't find the optimal hook which is called right after the DDP is optimized.

But I can confirm that. The parameter is called head.kernel.
To make sure it's getting ignored I added multiple module. prefixes in _setup_model() method from DDPPlugin, like this:

        model._ddp_params_and_buffers_to_ignore = [
            "head.kernel",
            "module.head.kernel",
            "module.module.head.kernel",
            "module.module.module.head.kernel",
        ]
        return DistributedDataParallel(module=model, device_ids=self.determine_ddp_device_ids(), **self._ddp_kwargs)

In my use-case I have one layer in the training, which I want to treat as Normal model and don't want that layer to be synced during the backward pass. Apparently the ignore is passed correctly, I tried to remove setting the ignore parameters and it crashed, because the layer is not the same size at each device (which is what I want).
But still, the gradient is the same (except for the extra parameters which are present due to mentioned different size of each linear layer).

Do you have some pointers, clues where to watch? I also checked the hooks, if they are registered to the layers, but I guess it is done using the Reducer and not by the specific hooks in each layer.

Very ugly workaround for me would be to rename the layer on each device to make sure there's no connection between the layers in each device, but that's quite ugly solution and I guess there must by some simple way which I am not seeing.

[Honzys]

Okay, It was my mistake, I deeply apologize for wasting your time there. The layer indeeds gets removed from the DistributedDataParallel (or rather not even getting there).

But I've found another error when trying to set the _ddp_params_and_buffers_to_ignore inside the LightningModule, so I've created issue here - #11844 .

Thank you anyway!