PartialFC: Distributing only part of the model

[w2ex]

Hi,
I am willing to implement PartialFC based on InsightFace implementation, which is fully Pytorch.
The main idea is to distribute the backbone of the model with torch.nn.distributed.DistributedDataParallel (see here ), but not the head. This is straightforward in plain Pytorch, but not in Pytorch Lightning.

If I understand correctly, in Pytorch Lightning, the parallelized strategy has to be specified through the strategy argument of the Trainer.

My question is : if i create a new strategy inherited from DDPStrategy, and override the _setup_model() function as follow (only apply DDP on model.backbone and not the full model), will it be enough ?

    def _setup_model(self, model: Module) -> Module:
        device_ids = self.determine_ddp_device_ids()
        log.detail(f"setting up Partial-FC DDP model with device ids: {device_ids}, kwargs: {self._ddp_kwargs}")
        model.backbone = DistributedDataParallel(module=model.backbone, device_ids=device_ids, **self._ddp_kwargs)
        return model

I know I also have to deal with some instance check in _register_ddp_hooks and teardown, but I want to know if I am missing anything else.

Last, will I be able to launch my script using torchrun or is there any incompatibility with Pytorch Lightning ?

Thank you very much.

[homomorfism]

Hi, @w2ex! Have you implemented PartialFC in pytorch lightning? Or did you stick to torch for following architecture?