Train the model function.
(forward_step_func, model, optimizer, opt_param_scheduler,
train_data_iterator, valid_data_iterator,
process_non_loss_data_func, config)
| 564 | timers.log(['save-checkpoint']) |
| 565 | |
| 566 | def train(forward_step_func, model, optimizer, opt_param_scheduler, |
| 567 | train_data_iterator, valid_data_iterator, |
| 568 | process_non_loss_data_func, config): |
| 569 | """Train the model function.""" |
| 570 | args = get_args() |
| 571 | timers = get_timers() |
| 572 | |
| 573 | # Write args to tensorboard |
| 574 | write_args_to_tensorboard() |
| 575 | |
| 576 | # Turn on training mode which enables dropout. |
| 577 | for model_module in model: |
| 578 | model_module.train() |
| 579 | |
| 580 | # Tracking loss. |
| 581 | total_loss_dict = {} |
| 582 | |
| 583 | # Iterations. |
| 584 | iteration = args.iteration |
| 585 | |
| 586 | # Setup some training config params |
| 587 | config.grad_scale_func = optimizer.scale_loss |
| 588 | config.timers = timers |
| 589 | # TODO: Remove this once we move DDP to Core. |
| 590 | |
| 591 | |
| 592 | |
| 593 | |
| 594 | if len(model) == 1 and isinstance(model[0], DDP) and \ |
| 595 | args.overlap_grad_reduce: |
| 596 | assert config.no_sync_func is None, \ |
| 597 | ('When overlap_grad_reduce is True, config.no_sync_func must be None; ' |
| 598 | 'a custom no_sync_func is not supported when overlapping grad-reduce') |
| 599 | if args.delay_grad_reduce: |
| 600 | config.grad_sync_func = model[0].grad_sync |
| 601 | config.no_sync_func = model[0].no_sync |
| 602 | |
| 603 | timers('interval-time', log_level=0).start(barrier=True) |
| 604 | print_datetime('before the start of training step') |
| 605 | report_memory_flag = True |
| 606 | while iteration < args.train_iters: |
| 607 | if args.profile and \ |
| 608 | iteration == args.profile_step_start and \ |
| 609 | torch.distributed.get_rank() in args.profile_ranks: |
| 610 | torch.cuda.cudart().cudaProfilerStart() |
| 611 | torch.autograd.profiler.emit_nvtx(record_shapes=True).__enter__() |
| 612 | |
| 613 | update_num_microbatches(args.consumed_train_samples) |
| 614 | args.curr_iteration = iteration |
| 615 | loss_dict, skipped_iter, grad_norm, num_zeros_in_grad = \ |
| 616 | train_step(forward_step_func, |
| 617 | train_data_iterator, |
| 618 | model, |
| 619 | optimizer, |
| 620 | opt_param_scheduler, |
| 621 | config) |
| 622 | iteration += 1 |
| 623 | args.consumed_train_samples += mpu.get_data_parallel_world_size() * \ |
no test coverage detected