| 29 | |
| 30 | |
| 31 | def demo_basic(rank): |
| 32 | |
| 33 | print( |
| 34 | f"[{os.getpid()}] rank = {dist.get_rank()}, " |
| 35 | + f"world_size = {dist.get_world_size()}" |
| 36 | ) |
| 37 | |
| 38 | model = ToyModel().to(rank) |
| 39 | ddp_model = DDP(model, device_ids=[rank]) |
| 40 | |
| 41 | loss_fn = nn.MSELoss() |
| 42 | optimizer = optim.SGD(ddp_model.parameters(), lr=0.001) |
| 43 | |
| 44 | optimizer.zero_grad() |
| 45 | outputs = ddp_model(torch.randn(20, 10)) |
| 46 | labels = torch.randn(20, 5).to(rank) |
| 47 | loss_fn(outputs, labels).backward() |
| 48 | optimizer.step() |
| 49 | |
| 50 | print(f"training completed in rank {rank}!") |
| 51 | |
| 52 | |
| 53 | def main(): |