Thanks for posting @Maxwell_Albert Looking into your problem, I couldn’t find anything suspicious yet. But I think potentially this might potentially be related to how the backward runs across two GPUs.
Since you want to do model parallel across two GPUs and split two models. Maybe you want to consider using RPC to do model parallel? RPC and Distributed Autograd could hook up the autograd graph implicitly without user writing custom send/recv functions, this might can resolve the nan loss issue that you met. Getting Started with Distributed RPC Framework — PyTorch Tutorials 1.10.1+cu102 documentation