Single training step.
(data_iterator, model, optimizer, lr_scheduler,
args, timers, hooks=None, single_step=False, **kwargs)
| 431 | |
| 432 | |
| 433 | def train_step(data_iterator, model, optimizer, lr_scheduler, |
| 434 | args, timers, hooks=None, single_step=False, **kwargs): |
| 435 | """Single training step.""" |
| 436 | if hooks is None: |
| 437 | hooks = {} |
| 438 | lm_loss_total, metrics_total, count, metrics_count = 0.0, {}, 0, {} |
| 439 | forward_step = hooks['forward_step'] |
| 440 | |
| 441 | while True: |
| 442 | profiling_flag = (args.profiling != -1 and args.iteration >= args.profiling) |
| 443 | # Forward model for one step. |
| 444 | if profiling_flag: |
| 445 | torch.cuda.nvtx.range_push("forward") |
| 446 | timers('forward').start() |
| 447 | forward_ret = forward_step(data_iterator, model, args, timers, **kwargs) |
| 448 | if isinstance(forward_ret, tuple): |
| 449 | lm_loss, metrics = forward_ret |
| 450 | else: |
| 451 | lm_loss, metrics = forward_ret, {} |
| 452 | timers('forward').stop() |
| 453 | if profiling_flag: |
| 454 | torch.cuda.nvtx.range_pop() |
| 455 | |
| 456 | # Check nan or inf in forward, preventing it from interfering loss scaler, |
| 457 | # and all reduce metrics by the way |
| 458 | if profiling_flag: |
| 459 | torch.cuda.nvtx.range_push("loss_and_metrics") |
| 460 | lm_loss_reduced = lm_loss.detach().clone() |
| 461 | torch.distributed.all_reduce(lm_loss_reduced.data) |
| 462 | lm_loss_reduced.data = lm_loss_reduced.data / args.world_size |
| 463 | |
| 464 | loss_checker = lm_loss_reduced |
| 465 | for name in metrics: |
| 466 | if not 'eval' in name: |
| 467 | metrics[name] = metrics[name].detach().clone() |
| 468 | if metrics[name].data.item() == -100: |
| 469 | cnt = torch.zeros(1, dtype=torch.int64, device=metrics[name].data.device) |
| 470 | metrics[name].data = torch.tensor(0., device=metrics[name].data.device) |
| 471 | else: |
| 472 | cnt = torch.ones(1, dtype=torch.int64, device=metrics[name].data.device) |
| 473 | torch.distributed.all_reduce(metrics[name].data) |
| 474 | torch.distributed.all_reduce(cnt) |
| 475 | if cnt.item() == 0: |
| 476 | metrics[name].data = torch.tensor(-100, device=metrics[name].data.device) |
| 477 | else: |
| 478 | metrics[name].data /= cnt.cpu().item() # args.world_size |
| 479 | loss_checker = loss_checker + metrics[name] |
| 480 | if loss_checker.isnan().any() or loss_checker.isinf().any(): |
| 481 | print_all('Skipping backward and optimizer step for nan or inf in forwarding metrics/loss!') |
| 482 | return lm_loss.detach(), 1, metrics |
| 483 | |
| 484 | # Accumulate the statistics |
| 485 | lm_loss_total += lm_loss_reduced |
| 486 | for name in metrics: |
| 487 | if name not in metrics_total: |
| 488 | metrics_total[name] = torch.tensor(0.0, device=metrics[name].data.device) |
| 489 | if name not in metrics_count: |
| 490 | metrics_count[name] = 0 |
no test coverage detected