| 334 | |
| 335 | |
| 336 | def evaluate(model, val_dloader, device, distributed): |
| 337 | model.eval() |
| 338 | total_loss = 0.0 |
| 339 | total_tokens = 0 |
| 340 | with torch.no_grad(): |
| 341 | for batch in val_dloader: |
| 342 | input_ids = batch["input_ids"].to(device) |
| 343 | labels = batch["labels"].to(device) |
| 344 | with autocast(enabled=False): |
| 345 | outputs = model(input_ids=input_ids, labels=labels) |
| 346 | loss = outputs.loss |
| 347 | num_tokens = input_ids.ne(tokenizer.pad_token_id).sum().item() |
| 348 | total_loss += loss.item() * num_tokens |
| 349 | total_tokens += num_tokens |
| 350 | if distributed: |
| 351 | loss_tensor = torch.tensor([total_loss, total_tokens], dtype=torch.float32, device=device) |
| 352 | dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM) |
| 353 | total_loss, total_tokens = loss_tensor[0].item(), loss_tensor[1].item() |
| 354 | model.train() |
| 355 | return total_loss / total_tokens if total_tokens > 0 else 0.0 |
| 356 | |
| 357 | |
| 358 | def save_full_checkpoint(model, optimizer, scheduler, epoch, step_in_epoch, global_step, output_dir): |