| 15 | init_process_group(backend="nccl") |
| 16 | |
| 17 | class Trainer: |
| 18 | def __init__( |
| 19 | self, |
| 20 | model: torch.nn.Module, |
| 21 | train_data: DataLoader, |
| 22 | optimizer: torch.optim.Optimizer, |
| 23 | save_every: int, |
| 24 | snapshot_path: str, |
| 25 | ) -> None: |
| 26 | self.local_rank = int(os.environ["LOCAL_RANK"]) |
| 27 | self.global_rank = int(os.environ["RANK"]) |
| 28 | self.model = model.to(self.local_rank) |
| 29 | self.train_data = train_data |
| 30 | self.optimizer = optimizer |
| 31 | self.save_every = save_every |
| 32 | self.epochs_run = 0 |
| 33 | self.snapshot_path = snapshot_path |
| 34 | if os.path.exists(snapshot_path): |
| 35 | print("Loading snapshot") |
| 36 | self._load_snapshot(snapshot_path) |
| 37 | |
| 38 | self.model = DDP(self.model, device_ids=[self.local_rank]) |
| 39 | |
| 40 | def _load_snapshot(self, snapshot_path): |
| 41 | loc = f"cuda:{self.local_rank}" |
| 42 | snapshot = torch.load(snapshot_path, map_location=loc) |
| 43 | self.model.load_state_dict(snapshot["MODEL_STATE"]) |
| 44 | self.epochs_run = snapshot["EPOCHS_RUN"] |
| 45 | print(f"Resuming training from snapshot at Epoch {self.epochs_run}") |
| 46 | |
| 47 | def _run_batch(self, source, targets): |
| 48 | self.optimizer.zero_grad() |
| 49 | output = self.model(source) |
| 50 | loss = F.cross_entropy(output, targets) |
| 51 | loss.backward() |
| 52 | self.optimizer.step() |
| 53 | |
| 54 | def _run_epoch(self, epoch): |
| 55 | b_sz = len(next(iter(self.train_data))[0]) |
| 56 | print(f"[GPU{self.global_rank}] Epoch {epoch} | Batchsize: {b_sz} | Steps: {len(self.train_data)}") |
| 57 | self.train_data.sampler.set_epoch(epoch) |
| 58 | for source, targets in self.train_data: |
| 59 | source = source.to(self.local_rank) |
| 60 | targets = targets.to(self.local_rank) |
| 61 | self._run_batch(source, targets) |
| 62 | |
| 63 | def _save_snapshot(self, epoch): |
| 64 | snapshot = { |
| 65 | "MODEL_STATE": self.model.module.state_dict(), |
| 66 | "EPOCHS_RUN": epoch, |
| 67 | } |
| 68 | torch.save(snapshot, self.snapshot_path) |
| 69 | print(f"Epoch {epoch} | Training snapshot saved at {self.snapshot_path}") |
| 70 | |
| 71 | def train(self, max_epochs: int): |
| 72 | for epoch in range(self.epochs_run, max_epochs): |
| 73 | self._run_epoch(epoch) |
| 74 | if self.global_rank == 0 and epoch % self.save_every == 0: |