Hey @Rizhiy
You might not be able to pass the optimizer as that, because every subprocess needs its own dedicated optimizer. And not sure about how dateset/criterion/scheduler in the code behave in multiprocessing use cases. If it is just the model, the following might work.
def train(rank, model):
    model = DistributedDataParallel(model.to(rank), device_ids=[rank], output_device=rank)
    ...

def main():
    model.share_memory() 
    mp.spawn(
        train, 
        args=(model, dataset, ...),
        nprocs=world_size,
        join=True)

if __name__=="__main__":
    main()

Rank is the subprocess id, which is provided my mp.spawn as the first argument to the target function.
If the reason for this is to keep the train() signature intact, is it possible to create another wrapper function to wrap train(), configure everything in wrapper, call train in wrapper, and use wrapper as the spawn target?