Sounds good! We are also working on a broader effort to make DDP more modular and configurable across the board. It seems like this definitely an interesting direction to support - feel free to make an issue on our GitHub repo or add to this RFC tracking the DDP configurability effort: [RFC] Modularize DistributedDataParallel · Issue #37002 · pytorch/pytorch · GitHub 8