tensor_parallel_example.py 和 sequence_parallel_example.py

FORWARD PASS [12/1864] *c10d_functional.all_reduce: 1 BACKWARD PASS ToyModel *module type: class 'main.ToyModel' FORWARD PASS *c10d_functional.all_reduce: 1 ToyModel.in_proj *module type: class 'torch.nn.modules.linear.Linear' *Parameter List *weight: (Shard(dim=0),) *bias: (Shard(dim=0),) FORWARD PASS **aten.addmm.default shape: [torch.Size([32]), torch.Size([4, 10]), torch.Size([10, 32])] sharding: [(Shard(dim=0),), (Replicate(),), (Shard(dim=1),)] device mesh: DeviceMesh('cuda', [0, 1, 2, 3]) BACKWARD PASS **aten.mm.default shape: [torch.Size([32, 4]), torch.Size([4, 10])] sharding: [(Shard(dim=0),), (Replicate(),)] device mesh: DeviceMesh('cuda', [0, 1, 2, 3]) **aten.sum.dim_IntList shape: [torch.Size([4, 32])] sharding: [(Shard(dim=1),)] device mesh: DeviceMesh('cuda', [0, 1, 2, 3]) **aten.add_.Tensor shape: [torch.Size([32]), torch.Size([32])] sharding: [(Shard(dim=0),), (Shard(dim=0),)] device mesh: DeviceMesh('cuda', [0, 1, 2, 3]) **aten.add_.Tensor shape: [torch.Size([32, 10]), torch.Size([32, 10])] sharding: [(Shard(dim=0),), (Shard(dim=0),)] device mesh: DeviceMesh('cuda', [0, 1, 2, 3]) ToyModel.relu *module type: class 'torch.nn.modules.activation.ReLU' FORWARD PASS BACKWARD PASS ToyModel.out_proj *module type: class 'torch.nn.modules.linear.Linear' *Parameter List *weight: (Shard(dim=1),) *bias: (Replicate(),) FORWARD PASS *c10d_functional.all_reduce: 1 **aten.addmm.default shape: [torch.Size([5]), torch.Size([4,