36 lines
746 B
Python
36 lines
746 B
Python
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
|
|
|
|
def main():
|
|
dist.init_process_group("nccl")
|
|
|
|
rank = dist.get_rank()
|
|
torch.cuda.set_device(0)
|
|
|
|
model = nn.Linear(10, 10).cuda()
|
|
ddp = torch.nn.parallel.DistributedDataParallel(model, device_ids=[0])
|
|
|
|
opt = optim.SGD(ddp.parameters(), lr=0.01)
|
|
loss_fn = nn.MSELoss()
|
|
|
|
for i in range(5):
|
|
x = torch.randn(32, 10).cuda()
|
|
y = torch.randn(32, 10).cuda()
|
|
|
|
opt.zero_grad()
|
|
out = ddp(x)
|
|
loss = loss_fn(out, y)
|
|
loss.backward()
|
|
opt.step()
|
|
|
|
print(f"rank {rank} step {i} loss {loss.item()}")
|
|
|
|
dist.destroy_process_group()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|