Thanks for posting the question @George_Diamantop From the code snippets it seems like you are doing some early exits, which means not all parameters will participate in the calculation of model output and loss right? Can you try use find_unused_parameters=True when initializing DDP?
Currently, find_unused_parameters=True must be passed into torch.nn.parallel.DistributedDataParallel() initialization if there are parameters that may be unused in the forward pass.