PyTorch多GPU模型并行：跨GPU计算的可行性与常见实践咨询

首先可以明确：你给出的代码是完全可行的，完全符合你的预期——X @ A的计算会在GPU 0上执行，得到的结果转移到GPU 1后，再与GPU 1上的参数B完成矩阵乘法，整个流程不会有设备不兼容的问题。

关于这种多设备模型分布方式是否属于常见实践：

• 这种手动拆分模型参数到不同GPU、手动控制张量设备流转的方式，属于早期多GPU计算的手动并行方案，目前已经不是主流的常见实践。
• 当前PyTorch生态中，更常用的是框架原生的并行方案：
  • 数据并行（DataParallel/DistributedDataParallel）：将完整模型复制到多个GPU，每个GPU处理部分数据，适合模型可单GPU容纳的场景。
  • 模型并行（如FSDP、Megatron-LM等）：针对大模型无法单GPU容纳的场景，自动拆分模型的层或参数到不同GPU，自动处理设备间的数据传输，无需手动编写设备转移逻辑。
• 手动拆分方式的缺点很突出：需要开发者手动管理所有参数和张量的设备位置，代码维护成本高，易出现设备不匹配错误，且难以适配复杂模型和大规模计算场景，仅在简单自定义场景或需要极致手动控制流程的情况下才会偶尔使用。

内容的提问来源于stack exchange，提问作者Matthew