Single training step.
(forward_step_func, data_iterator, model, optimizer, lr_scheduler)
| 475 | |
| 476 | |
| 477 | def train_step(forward_step_func, data_iterator, model, optimizer, lr_scheduler): |
| 478 | """Single training step.""" |
| 479 | args = get_args() |
| 480 | timers = get_timers() |
| 481 | |
| 482 | if args.deepspeed and args.ds_pipeline_enabled: |
| 483 | skipped_iter = 0 |
| 484 | num_zeros_in_grad = 0 |
| 485 | assert isinstance(model[0], deepspeed.PipelineEngine) |
| 486 | loss = model[0].train_batch(data_iter=data_iterator) |
| 487 | grad_norm = model[0].get_global_grad_norm() |
| 488 | return {"lm loss": loss}, skipped_iter, grad_norm, num_zeros_in_grad |
| 489 | |
| 490 | # Set grad to zero. |
| 491 | if not args.deepspeed: |
| 492 | if args.DDP_impl == "local" and args.use_contiguous_buffers_in_ddp: |
| 493 | for partition in model: |
| 494 | partition.zero_grad_buffer() |
| 495 | else: |
| 496 | optimizer.zero_grad() |
| 497 | |
| 498 | if mpu.get_pipeline_model_parallel_world_size() > 1: |
| 499 | if args.virtual_pipeline_model_parallel_size is not None: |
| 500 | # print_rank_0("===> fb_func = w/ interleaving") |
| 501 | forward_backward_func = forward_backward_pipelining_with_interleaving |
| 502 | assert get_num_microbatches() % args.pipeline_model_parallel_size == 0, ( |
| 503 | "number of microbatches is not divisible by pipeline-parallel " |
| 504 | "size when using interleaved schedule" |
| 505 | ) |
| 506 | else: |
| 507 | # print_rank_0("===> fb_func = w/o interleaving") |
| 508 | forward_backward_func = forward_backward_pipelining_without_interleaving |
| 509 | else: |
| 510 | # print_rank_0("===> fb_func = no_pp") |
| 511 | forward_backward_func = forward_backward_no_pipelining |
| 512 | # print_rank_0("===> running fb_func") |
| 513 | losses_reduced = forward_backward_func( |
| 514 | forward_step_func, data_iterator, model, optimizer, timers, forward_only=False |
| 515 | ) |
| 516 | |
| 517 | # All-reduce if needed. |
| 518 | if not args.deepspeed and args.DDP_impl == "local": |
| 519 | timers("backward-params-all-reduce").start() |
| 520 | for model_module in model: |
| 521 | model_module.allreduce_gradients() |
| 522 | timers("backward-params-all-reduce").stop() |
| 523 | |
| 524 | # All-reduce word_embeddings' grad across first and last stages to ensure |
| 525 | # that word_embeddings parameters stay in sync. |
| 526 | # This should only run for models that support pipelined model parallelism |
| 527 | # (BERT and GPT-2). |
| 528 | if not args.deepspeed: |
| 529 | timers("backward-embedding-all-reduce").start() |
| 530 | if ( |
| 531 | mpu.is_pipeline_first_stage(ignore_virtual=True) |
| 532 | or mpu.is_pipeline_last_stage(ignore_virtual=True) |
| 533 | ) and mpu.get_pipeline_model_parallel_world_size() > 1: |
| 534 | if mpu.is_pipeline_first_stage(ignore_virtual=True): |
no test coverage detected