training.py 41.8 KB
Newer Older
1
# coding=utf-8
Mohammad's avatar
Mohammad committed
2
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
3
4
5
6
7
8
9
10
11
12
13
14
15
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Mohammad's avatar
Mohammad committed
16
"""Pretrain utilities."""
17
18
19

from datetime import datetime
import math
Mohammad's avatar
Mohammad committed
20
import sys
21
22
23
24
import time
# The earliest we can measure the start time.
_TRAIN_START_TIME = time.time()

25
26
27
28
import torch
from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP
from apex.optimizers import FusedAdam as Adam

Neel Kant's avatar
Neel Kant committed
29
from megatron import get_args
Mohammad's avatar
Mohammad committed
30
31
from megatron import get_timers
from megatron import get_tensorboard_writer
32
from megatron import get_current_global_batch_size
mohammad's avatar
mohammad committed
33
from megatron import get_num_microbatches
mohammad's avatar
mohammad committed
34
from megatron import is_last_rank
mohammad's avatar
mohammad committed
35
from megatron import update_num_microbatches
36
from megatron import mpu
Neel Kant's avatar
Neel Kant committed
37
from megatron import print_rank_0
38
from megatron import print_rank_last
Mohammad's avatar
Mohammad committed
39
40
from megatron.checkpointing import load_checkpoint
from megatron.checkpointing import save_checkpoint
41
from megatron.fp16 import FP16_Module
mohammad's avatar
mohammad committed
42
43
44
#from megatron.fp16 import FP16_Optimizer
from megatron.optimizer.optimizer import get_megatron_optimizer

Mohammad's avatar
Mohammad committed
45
from megatron.initialize import initialize_megatron
46
from megatron.initialize import write_args_to_tensorboard
47
48
49
from megatron.learning_rates import AnnealingLR
from megatron.model import DistributedDataParallel as LocalDDP
from megatron.model import get_params_for_weight_decay_optimization
Neel Kant's avatar
Neel Kant committed
50
from megatron.model.realm_model import ICTBertModel
51
from megatron.utils import check_adlr_autoresume_termination
52
from megatron.data.data_loaders import build_pretraining_data_loader
53
from megatron.utils import report_memory
54
55


56
57
58
59
60
61
62
def print_datetime(string):
    """Note that this call will sync across all ranks."""
    torch.distributed.barrier()
    time_str = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    print_rank_0('[' + string + '] datetime: {} '.format(time_str))


63
def pretrain(train_valid_test_dataset_provider, model_provider,
64
             forward_step_func, extra_args_provider=None, args_defaults={}):
65
66
67
    """Main training program.

    This function will run the followings in the order provided:
Mohammad's avatar
Mohammad committed
68
69
        1) initialize Megatron.
        2) setup model, optimizer and lr schedule using the model_provider.
70
        3) call train_val_test_data_provider to get train/val/test datasets.
Mohammad's avatar
Mohammad committed
71
        4) train the modle using the forward_step_func.
72
73

    Arguments:
74
75
76
        train_valid_test_dataset_provider: a function that takes the size of
            train/valid/test dataset and returns `train, valid, test` datasets.
        model_provider: a function that returns a vanilla version of the
Mohammad's avatar
Mohammad committed
77
78
79
80
81
82
83
84
85
86
            model. By vanilla we mean a simple model on cpu with no fp16 or ddp.
        forward_step_func: a function that takes a `data iterator` and `model`,
            and returns a `loss` scalar with a dictionary with key:values being
            the info we would like to monitor during training, for example
            `lm-loss: value`. We also require that this function add
            `batch generator` to the timers class.
        extra_args_provider: a function that takes a parser and adds arguments
            to it. It is used for programs to add their own arguments.
        args_defaults: a dictionary from argument-name to argument-value. It
            to set already parse arguments.
87
88
    """

89
    # Initalize and get arguments, timers, and Tensorboard writer.
90
91
    initialize_megatron(extra_args_provider=extra_args_provider,
                        args_defaults=args_defaults)
92

93
94
95
96
97
98
99
100
    # Adjust the startup time so it reflects the largest value.
    # This will be closer to what scheduler will see (outside of
    # image ... launches.
    global _TRAIN_START_TIME
    start_time_tensor = torch.cuda.FloatTensor([_TRAIN_START_TIME])
    torch.distributed.all_reduce(start_time_tensor,
                                 op=torch.distributed.ReduceOp.MIN)
    _TRAIN_START_TIME = start_time_tensor.item()
mshoeybi's avatar
mshoeybi committed
101
    print_rank_0('time to initialize megatron (seconds): {:.3f}'.format(
102
103
104
        time.time() - _TRAIN_START_TIME))
    print_datetime('after megatron is initialized')

105
    args = get_args()
Mohammad's avatar
Mohammad committed
106
    timers = get_timers()
107
108

    # Model, optimizer, and learning rate.
Mohammad's avatar
Mohammad committed
109
110
111
    timers('model and optimizer').start()
    model, optimizer, lr_scheduler = setup_model_and_optimizer(model_provider)
    timers('model and optimizer').stop()
112
113
    print_datetime('after model, optimizer, and learning rate '
                   'scheduler are built')
114
115

    # Data stuff.
116
117
118
119
120
    timers('train/valid/test data iterators').start()
    train_data_iterator, valid_data_iterator, test_data_iterator \
        = build_train_valid_test_data_iterators(
            train_valid_test_dataset_provider)
    timers('train/valid/test data iterators').stop()
mshoeybi's avatar
mshoeybi committed
121
    print_datetime('after dataloaders are built')
Mohammad's avatar
Mohammad committed
122
123
124

    # Print setup timing.
    print_rank_0('done with setups ...')
125
    timers.log(['model and optimizer', 'train/valid/test data iterators'])
Mohammad's avatar
Mohammad committed
126
    print_rank_0('training ...')
127
128

    iteration = 0
129
    if args.do_train and args.train_iters > 0:
mohammad's avatar
mohammad committed
130
131
132
        iteration = train(forward_step_func,
                          model, optimizer, lr_scheduler,
                          train_data_iterator, valid_data_iterator)
133
    print_datetime('after training is done')
Mohammad's avatar
Mohammad committed
134

135
136
137
    if args.do_valid:
        prefix = 'the end of training for val data'
        evaluate_and_print_results(prefix, forward_step_func,
138
                                   valid_data_iterator, model,
Mohammad's avatar
Mohammad committed
139
                                   iteration, False)
140
141

    if args.save and iteration != 0:
142
        save_checkpoint(iteration, model, optimizer, lr_scheduler)
143
144
145
146
147
148

    if args.do_test:
        # Run on test data.
        prefix = 'the end of training for test data'
        evaluate_and_print_results(prefix, forward_step_func,
                                   test_data_iterator, model,
Mohammad's avatar
Mohammad committed
149
                                   0, True)
150

151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
def update_train_iters(args):

    # For iteration-based training, we don't need to do anything
    if args.train_iters:
        return

    # Constant batch size with sample-based training.
    if args.rampup_batch_size is None:
        args.train_iters = args.train_samples // args.global_batch_size

    else:
        # Sample based training with rampup batch size.
        iterations = 0
        consumed_samples = 0
        # Rampup phase.
        while consumed_samples <= int(args.rampup_batch_size[2]):
167
168
            update_num_microbatches(consumed_samples, consistency_check=False)
            consumed_samples += get_current_global_batch_size()
169
170
            iterations += 1
        # Reset
171
        update_num_microbatches(0, consistency_check=False)
172
173
174
175
176
177
178
179
        # Constant phase
        # Note that we throw away any partial last batch.
        iterations += (args.train_samples - consumed_samples) // \
                      args.global_batch_size
        args.train_iters = iterations

    print_rank_0('setting training iterations to {}'.format(args.train_iters))

180

Mohammad's avatar
Mohammad committed
181
def get_model(model_provider_func):
182
    """Build the model."""
Mohammad's avatar
Mohammad committed
183
    args = get_args()
184
185

    # Build model on cpu.
Mohammad's avatar
Mohammad committed
186
    model = model_provider_func()
187
188
189

    # Print number of parameters.
    if mpu.get_data_parallel_rank() == 0:
190
        print(' > number of parameters on (tensor, pipeline) '
191
              'model parallel rank ({}, {}): {}'.format(
192
193
            mpu.get_tensor_model_parallel_rank(),
            mpu.get_pipeline_model_parallel_rank(),
194
195
196
197
198
199
200
201
202
203
204
            sum([p.nelement() for p in model.parameters()])), flush=True)

    # GPU allocation.
    model.cuda(torch.cuda.current_device())

    # Fp16 conversion.
    if args.fp16:
        model = FP16_Module(model)

    if args.DDP_impl == 'torch':
        i = torch.cuda.current_device()
Mohammad's avatar
Mohammad committed
205
206
        model = torchDDP(model, device_ids=[i], output_device=i,
                         process_group=mpu.get_data_parallel_group())
207
208
        return model
    if args.DDP_impl == 'local':
Mohammad's avatar
Mohammad committed
209
        model = LocalDDP(model)
210
211
        return model

212
    raise NotImplementedError('Unknown DDP implementation specified: {}. '
213
                              'Exiting.'.format(args.DDP_impl))
214
215


Mohammad's avatar
Mohammad committed
216
def get_optimizer(model):
217
    """Set up the optimizer."""
Mohammad's avatar
Mohammad committed
218
    args = get_args()
219
220

    # Build parameter groups (weight decay and non-decay).
Mohammad's avatar
Mohammad committed
221
    while isinstance(model, (torchDDP, LocalDDP, FP16_Module)):
222
223
224
225
226
227
        model = model.module
    param_groups = get_params_for_weight_decay_optimization(model)

    # Add model parallel attribute if it is not set.
    for param_group in param_groups:
        for param in param_group['params']:
228
229
            if not hasattr(param, 'tensor_model_parallel'):
                param.tensor_model_parallel = False
230
231

    # Use Adam.
232
233
    optimizer = Adam(param_groups, lr=args.lr, weight_decay=args.weight_decay,
        betas=(args.adam_beta1, args.adam_beta2), eps=args.adam_eps)
234
235

    # Wrap into fp16 optimizer.
mohammad's avatar
mohammad committed
236
237
    optimizer = get_megatron_optimizer(optimizer, model)
    '''
238
239
240
241
242
        optimizer = FP16_Optimizer(optimizer,
                                   static_loss_scale=args.loss_scale,
                                   dynamic_loss_scale=args.dynamic_loss_scale,
                                   dynamic_loss_args={
                                       'scale_window': args.loss_scale_window,
Neel Kant's avatar
Neel Kant committed
243
                                       'min_scale': args.min_scale,
244
                                       'delayed_shift': args.hysteresis})
mohammad's avatar
mohammad committed
245
    '''
246
247
248
    return optimizer


Mohammad's avatar
Mohammad committed
249
def get_learning_rate_scheduler(optimizer):
250
    """Build the learning rate scheduler."""
Mohammad's avatar
Mohammad committed
251
    args = get_args()
252

253
254
255
256
257
    # Iteration-based training.
    if args.train_iters:
        if args.lr_decay_iters is None:
            args.lr_decay_iters = args.train_iters
        decay_steps = args.lr_decay_iters * args.global_batch_size
258
259
        if args.lr_warmup_fraction is not None:
            warmup_steps = args.lr_warmup_fraction * decay_steps
260
261
        else:
            warmup_steps = args.lr_warmup_iters * args.global_batch_size
262
263
264
265
266
    # Sample-based training.
    elif args.train_samples:
        # We need to set training iters for later use. Technically
        # we need to adjust the training samples too (due to last
        # batch being incomplete) but we leave it as is for now.
267
        update_train_iters(args)
268
269
270
        if args.lr_decay_samples is None:
            args.lr_decay_samples = args.train_samples
        decay_steps = args.lr_decay_samples
271
272
        if args.lr_warmup_fraction is not None:
            warmup_steps = args.lr_warmup_fraction * decay_steps
273
274
        else:
            warmup_steps = args.lr_warmup_samples
275
    else:
276
277
278
        raise Exception(
            'either train-iters or train-samples should be provided.')

279
280
    lr_scheduler = AnnealingLR(
        optimizer,
281
        max_lr=args.lr,
282
        min_lr=args.min_lr,
283
284
        warmup_steps=warmup_steps,
        decay_steps=decay_steps,
285
        decay_style=args.lr_decay_style,
286
287
288
289
290
291
        use_checkpoint_lr_scheduler=args.use_checkpoint_lr_scheduler,
        override_lr_scheduler=args.override_lr_scheduler)

    return lr_scheduler


Mohammad's avatar
Mohammad committed
292
def setup_model_and_optimizer(model_provider_func):
293
    """Setup model and optimizer."""
Mohammad's avatar
Mohammad committed
294
    args = get_args()
295

Mohammad's avatar
Mohammad committed
296
297
298
    model = get_model(model_provider_func)
    optimizer = get_optimizer(model)
    lr_scheduler = get_learning_rate_scheduler(optimizer)
299
300

    if args.load is not None:
301
302
303
304
305
        timers = get_timers()
        # Extra barrier is added to make sure all ranks report the
        # max time.
        torch.distributed.barrier()
        timers('load checkpoint').start()
306
        args.iteration = load_checkpoint(model, optimizer, lr_scheduler)
307
308
309
        torch.distributed.barrier()
        timers('load checkpoint').stop()
        timers.log(['load checkpoint'])
310
311
312
    else:
        args.iteration = 0

mohammad's avatar
mohammad committed
313
    # We only support local DDP with multiple micro-batches.
mohammad's avatar
mohammad committed
314
315
316
    if get_num_microbatches() > 1:
        assert args.DDP_impl == 'local'

Neel Kant's avatar
Neel Kant committed
317
318
319
320
321
    # get model without FP16 and/or TorchDDP wrappers
    unwrapped_model = model
    while hasattr(unwrapped_model, 'module'):
        unwrapped_model = unwrapped_model.module

322
323
    if args.iteration == 0 and hasattr(unwrapped_model,
                                       'init_state_dict_from_bert'):
324
        print("Initializing ICT from pretrained BERT model", flush=True)
325
        unwrapped_model.init_state_dict_from_bert()
Neel Kant's avatar
Neel Kant committed
326

327
328
329
    return model, optimizer, lr_scheduler


330
331
332
333
334
335
336
337
def communicate(tensor_send_next, tensor_send_prev, recv_forward, recv_backward):
    """Communicate tensors between stages using torch.distributed.ring_exchange(.) API."""
    args = get_args()

    # Create placeholder tensors for receive in forward and backward directions
    # if needed.
    tensor_recv_prev = None
    tensor_recv_next = None
338
    tensor_shape = (args.seq_length, args.micro_batch_size, args.hidden_size)
339
340
341
    dtype = args.params_dtype
    if args.fp32_residual_connection:
        dtype = torch.float
342
343
344
    if recv_forward:
        tensor_recv_prev = torch.empty(tensor_shape,
                                       requires_grad=True,
345
                                       device=torch.cuda.current_device(),
346
                                       dtype=dtype)
347
348
349
    if recv_backward:
        tensor_recv_next = torch.empty(tensor_shape,
                                       requires_grad=True,
350
                                       device=torch.cuda.current_device(),
351
                                       dtype=dtype)
352
353
354
355
356
357

    # Send tensors in both the forward and backward directions as appropriate.
    torch.distributed.ring_exchange(tensor_send_prev=tensor_send_prev,
                                    tensor_recv_prev=tensor_recv_prev,
                                    tensor_send_next=tensor_send_next,
                                    tensor_recv_next=tensor_recv_next,
358
                                    group=mpu.get_pipeline_model_parallel_group())
359
360
361
362
363

    return tensor_recv_prev, tensor_recv_next


def backward_step(optimizer, model, input_tensor, output_tensor, output_tensor_grad):
364
    """Backward step."""
Mohammad's avatar
Mohammad committed
365
366
    args = get_args()
    timers = get_timers()
367

368
369
370
371
    # Retain the grad on the input_tensor.
    if input_tensor is not None:
        input_tensor.retain_grad()

372
    # Backward pass.
mohammad's avatar
mohammad committed
373
374
375
376
377
    if output_tensor_grad is None:
        output_tensor = optimizer.scale_loss(output_tensor)
    torch.autograd.backward(output_tensor, grad_tensors=output_tensor_grad)
    '''
    if args.fp16 and output_tensor_grad is None:
378
379
380
381
        optimizer.backward(output_tensor, update_master_grads=False,
                           output_tensor_grad=output_tensor_grad)
    else:
        torch.autograd.backward(output_tensor, grad_tensors=output_tensor_grad)
mohammad's avatar
mohammad committed
382
    '''
383
384
385
386
387
388
389
390
    # Collect the grad of the input_tensor.
    input_tensor_grad = None
    if input_tensor is not None:
        input_tensor_grad = input_tensor.grad

    return input_tensor_grad


391
392
393
def forward_step_with_communication(forward_step_func, data_iterator, model,
                                    input_tensors, output_tensors,
                                    losses_reduced, timers):
394
395
    args = get_args()

396
    if not mpu.is_pipeline_first_stage():
397
        timers('forward-recv').start()
398
399
400
401
402
        input_tensor, _ = communicate(
            tensor_send_next=None,
            tensor_send_prev=None,
            recv_forward=True,
            recv_backward=False)
403
        timers('forward-recv').stop()
404
405
406
407
    else:
        input_tensor = None

    # Forward model for one step.
408
    timers('forward-compute').start()
409
    output_tensor = forward_step_func(data_iterator, model, input_tensor)
410
    timers('forward-compute').stop()
411
412
413

    if mpu.is_pipeline_last_stage():
        loss, loss_reduced = output_tensor
mohammad's avatar
mohammad committed
414
        output_tensor = loss / get_num_microbatches()
415
416
        losses_reduced.append(loss_reduced)
    else:
417
        timers('forward-send').start()
418
419
420
421
422
        communicate(
            tensor_send_next=output_tensor,
            tensor_send_prev=None,
            recv_forward=False,
            recv_backward=False)
423
        timers('forward-send').stop()
424
425
426
427
428
429
430
431
432
433
434
435

    input_tensors.append(input_tensor)
    output_tensors.append(output_tensor)


def backward_step_with_communication(optimizer, model, input_tensors, output_tensors, timers):
    input_tensor = input_tensors.pop(0)
    output_tensor = output_tensors.pop(0)

    if mpu.is_pipeline_last_stage():
        output_tensor_grad = None
    else:
436
        timers('backward-recv').start()
437
438
439
440
441
        _, output_tensor_grad = communicate(
            tensor_send_next=None,
            tensor_send_prev=None,
            recv_forward=False,
            recv_backward=True)
442
        timers('backward-recv').stop()
443
444

    # Backward pass for one step.
445
    timers('backward-compute').start()
446
447
    input_grad_tensor = \
        backward_step(optimizer, model, input_tensor, output_tensor, output_tensor_grad)
448
    timers('backward-compute').stop()
449
450

    if not mpu.is_pipeline_first_stage():
451
        timers('backward-send').start()
452
453
454
455
456
        communicate(
            tensor_send_next=None,
            tensor_send_prev=input_grad_tensor,
            recv_forward=False,
            recv_backward=False)
457
        timers('backward-send').stop()
458
459


460
461
462
463
464
def forward_and_backward_steps_with_communication(forward_step_func, data_iterator, model,
                                                  optimizer,
                                                  input_tensor, last_microbatch,
                                                  input_tensors, output_tensors,
                                                  losses_reduced, timers):
465
466
    args = get_args()

467
468
469
470
471
472
473
    # Forward model for one step.
    timers('forward-compute').start()
    output_tensor = forward_step_func(data_iterator, model, input_tensor)
    timers('forward-compute').stop()

    if mpu.is_pipeline_last_stage():
        loss, loss_reduced = output_tensor
mohammad's avatar
mohammad committed
474
        output_tensor = loss / get_num_microbatches()
475
476
477
        output_tensor_grad = None
        losses_reduced.append(loss_reduced)
    else:
Deepak Narayanan's avatar
Deepak Narayanan committed
478
        timers('forward-send-backward-recv').start()
479
480
481
482
483
        _, output_tensor_grad = communicate(
            tensor_send_next=output_tensor,
            tensor_send_prev=None,
            recv_forward=False,
            recv_backward=True)
Deepak Narayanan's avatar
Deepak Narayanan committed
484
        timers('forward-send-backward-recv').stop()
485
486
487
488
489
490
491
492
493
494
495
496
497
498

    input_tensors.append(input_tensor)
    output_tensors.append(output_tensor)

    input_tensor = input_tensors.pop(0)
    output_tensor = output_tensors.pop(0)

    # Backward pass for one step.
    timers('backward-compute').start()
    input_grad_tensor = \
        backward_step(optimizer, model, input_tensor, output_tensor, output_tensor_grad)
    timers('backward-compute').stop()

    if not mpu.is_pipeline_first_stage():
Deepak Narayanan's avatar
Deepak Narayanan committed
499
        timers('backward-send-forward-recv').start()
500
501
502
503
504
        input_tensor, _ = communicate(
            tensor_send_next=None,
            tensor_send_prev=input_grad_tensor,
            recv_forward=(not last_microbatch),
            recv_backward=False)
Deepak Narayanan's avatar
Deepak Narayanan committed
505
        timers('backward-send-forward-recv').stop()
506
507
508
509
510
511
    else:
        input_tensor = None

    return input_tensor


512
513
514
def forward_backward_no_pipelining(forward_step_func, data_iterator, model,
                                   optimizer, timers):
    """Run forward and backward passes without inter-stage communication."""
515
516
    args = get_args()

517
    losses_reduced = []
mohammad's avatar
mohammad committed
518
    for i in range(get_num_microbatches()):
519
520
        timers('forward-compute').start()
        loss, loss_reduced = forward_step_func(data_iterator, model, input_tensor=None)
mohammad's avatar
mohammad committed
521
        output_tensor = loss / get_num_microbatches()
522
523
524
525
526
527
528
529
530
531
        losses_reduced.append(loss_reduced)
        timers('forward-compute').stop()

        timers('backward-compute').start()
        output_tensor_grad = None
        backward_step(optimizer, model, input_tensor=None,
                      output_tensor=output_tensor, output_tensor_grad=None)
        timers('backward-compute').stop()

    return losses_reduced
532

533
534
535
536
537
538
539

def forward_backward_pipelining(forward_step_func, data_iterator, model,
                                optimizer, timers):
    """Run 1F1B schedule, with communication and warmup + cooldown microbatches as needed."""
    args = get_args()

    # Compute number of warmup microbatches.
mohammad's avatar
mohammad committed
540
    num_microbatches = get_num_microbatches()
541
542
543
544
545
    num_warmup_microbatches = \
        (mpu.get_pipeline_model_parallel_world_size() -
         mpu.get_pipeline_model_parallel_rank() - 1)
    num_warmup_microbatches = min(
        num_warmup_microbatches,
546
547
548
        num_microbatches)
    num_microbatches_remaining = \
        num_microbatches - num_warmup_microbatches
549
550
551
552
553

    input_tensors = []
    output_tensors = []
    losses_reduced = []

554
555
    # Run warmup forward passes.
    for i in range(num_warmup_microbatches):
556
557
558
559
        forward_step_with_communication(
            forward_step_func, data_iterator, model,
            input_tensors, output_tensors,
            losses_reduced, timers)
560

561
    # Before running 1F1B, need to receive first forward tensor.
562
563
    # If all microbatches are run in warmup / cooldown phase, then no need to
    # receive this tensor here.
564
    if num_microbatches_remaining > 0:
565
566
567
        if mpu.is_pipeline_first_stage():
            input_tensor = None
        else:
568
            timers('forward-recv').start()
569
570
571
572
            input_tensor, _ = communicate(tensor_send_next=None,
                                          tensor_send_prev=None,
                                          recv_forward=True,
                                          recv_backward=False)
573
            timers('forward-recv').stop()
574
575

    # Run 1F1B.
576
577
    for i in range(num_microbatches_remaining):
        last_iteration = (i == (num_microbatches_remaining - 1))
578
579
580
581
582
583
584
        input_tensor = \
            forward_and_backward_steps_with_communication(forward_step_func, data_iterator, model,
                                                          optimizer,
                                                          input_tensor, last_iteration,
                                                          input_tensors, output_tensors,
                                                          losses_reduced, timers)

585
586
    # Run cooldown backward passes.
    for i in range(num_warmup_microbatches):
587
588
589
590
591
592
593
594
595
596
597
598
599
        backward_step_with_communication(
            optimizer, model, input_tensors, output_tensors, timers)

    return losses_reduced


def train_step(forward_step_func, data_iterator,
               model, optimizer, lr_scheduler):
    """Single training step."""
    args = get_args()
    timers = get_timers()

    # Set grad to zero.
mohammad's avatar
mohammad committed
600
601
    optimizer.zero_grad()
    '''
602
603
604
605
    if args.fp16:
        optimizer.zero_grad(set_grads_to_None=True)
    else:
        optimizer.zero_grad()
mohammad's avatar
mohammad committed
606
    '''
607
608
609
610
611
612
613

    if mpu.get_pipeline_model_parallel_world_size() > 1:
        losses_reduced = forward_backward_pipelining(
            forward_step_func, data_iterator, model, optimizer, timers)
    else:
        losses_reduced = forward_backward_no_pipelining(
            forward_step_func, data_iterator, model, optimizer, timers)
614
615
616

    # All-reduce if needed.
    if args.DDP_impl == 'local':
617
        timers('backward-params-all-reduce').start()
618
619
        model.allreduce_params(reduce_after=False,
                               fp32_allreduce=args.fp32_allreduce)
620
        timers('backward-params-all-reduce').stop()
621

622
623
624
625
    # All-reduce word_embeddings' grad across first and last stages to ensure
    # that word_embeddings parameters stay in sync.
    # This should only run for models that support pipelined model parallelism
    # (BERT and GPT-2).
626
    timers('backward-embedding-all-reduce').start()
627
    if (mpu.is_pipeline_first_stage() or mpu.is_pipeline_last_stage()) and \
628
            mpu.get_pipeline_model_parallel_world_size() > 1:
629
630
631
632
        unwrapped_model = model
        while isinstance(unwrapped_model, (torchDDP, LocalDDP, FP16_Module)):
            unwrapped_model = unwrapped_model.module

633
634
635
636
        if unwrapped_model.share_word_embeddings:
            word_embeddings_weight = unwrapped_model.word_embeddings_weight()
            torch.distributed.all_reduce(word_embeddings_weight.grad,
                                         group=mpu.get_embedding_group())
637
    timers('backward-embedding-all-reduce').stop()
638

639
    # Update master gradients.
mohammad's avatar
mohammad committed
640
    '''
641
642
643
644
    timers('backward-master-grad').start()
    if args.fp16:
        optimizer.update_master_grads()
    timers('backward-master-grad').stop()
mohammad's avatar
mohammad committed
645
    '''
646
    # Clipping gradients helps prevent the exploding gradient.
mohammad's avatar
mohammad committed
647
    '''
648
    timers('backward-clip-grad').start()
649
    if args.clip_grad > 0.:
650
        if not args.fp16:
651
652
653
654
655
656
657
658
            named_parameters = model.named_parameters()
            parameters = []
            parameter_names = []
            for parameter_name, parameter in model.named_parameters():
                parameters.append(parameter)
                parameter_names.append(parameter_name)
            mpu.clip_grad_norm(parameters, args.clip_grad,
                               parameter_names=parameter_names)
659
660
        else:
            optimizer.clip_master_grads(args.clip_grad)
661
    timers('backward-clip-grad').stop()
mohammad's avatar
mohammad committed
662
    '''
663
664
665

    # Update parameters.
    timers('optimizer').start()
mohammad's avatar
mohammad committed
666
    update_successfull = optimizer.step()
667
668
669
    timers('optimizer').stop()

    # Update learning rate.
mohammad's avatar
mohammad committed
670
    if update_successfull:
671
672
673
674
        increment = get_num_microbatches() * \
                    args.micro_batch_size * \
                    args.data_parallel_size
        lr_scheduler.step(increment=increment)
mohammad's avatar
mohammad committed
675
        skipped_iter = 0
676
677
678
    else:
        skipped_iter = 1

679
    if mpu.is_pipeline_last_stage():
680
681
682
683
        # Average loss across microbatches.
        loss_reduced = {}
        for key in losses_reduced[0]:
            losses_reduced_for_key = [x[key] for x in losses_reduced]
684
            loss_reduced[key] = sum(losses_reduced_for_key) / len(losses_reduced_for_key)
685
686
        return loss_reduced, skipped_iter
    return {}, skipped_iter
687
688


Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
689
def training_log(loss_dict, total_loss_dict, learning_rate, iteration,
mohammad's avatar
mohammad committed
690
                 loss_scale, report_memory_flag, skipped_iter):
Mohammad's avatar
Mohammad committed
691
692
693
694
    """Log training information such as losses, timing, ...."""
    args = get_args()
    timers = get_timers()
    writer = get_tensorboard_writer()
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
695

mohammad's avatar
mohammad committed
696
697
    # Advanced, skipped, and Nan iterations.
    advanced_iters_key = 'advanced iterations'
mohammad's avatar
mohammad committed
698
    skipped_iters_key = 'skipped iterations'
mohammad's avatar
mohammad committed
699
700
701
702
703
704
705
706
707
    nan_iters_key = 'nan iterations'
    # Advanced iterations.
    if not skipped_iter:
        total_loss_dict[advanced_iters_key] = total_loss_dict.get(
            advanced_iters_key, 0) + 1
    else:
        if advanced_iters_key not in total_loss_dict:
            total_loss_dict[advanced_iters_key] = 0
    # Skipped iterations.
mohammad's avatar
mohammad committed
708
709
    total_loss_dict[skipped_iters_key] = total_loss_dict.get(
        skipped_iters_key, 0) + skipped_iter
mohammad's avatar
mohammad committed
710
    # Update losses and set nan iterations
mohammad's avatar
mohammad committed
711
    got_nan = False
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
712
    for key in loss_dict:
mohammad's avatar
mohammad committed
713
        if not skipped_iter:
714
715
            total_loss_dict[key] = total_loss_dict.get(
                key, torch.cuda.FloatTensor([0.0])) + loss_dict[key]
mohammad's avatar
mohammad committed
716
717
718
719
720
        else:
            value = loss_dict[key].float().sum().item()
            is_nan = value == float('inf') or \
                     value == -float('inf') or \
                     value != value
mohammad's avatar
mohammad committed
721
            got_nan = got_nan or is_nan
mohammad's avatar
mohammad committed
722
723
    total_loss_dict[nan_iters_key] = total_loss_dict.get(
        nan_iters_key, 0) + int(got_nan)
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
724
725
726

    # Logging.
    timers_to_log = []
Neel Kant's avatar
Neel Kant committed
727

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
728
729
730
    def add_to_logging(name):
        if name in timers.timers:
            timers_to_log.append(name)
731
732
733
    add_to_logging('forward-compute')
    add_to_logging('forward-recv')
    add_to_logging('forward-send')
Deepak Narayanan's avatar
Deepak Narayanan committed
734
    add_to_logging('forward-send-backward-recv')
735
736
737
    add_to_logging('backward-compute')
    add_to_logging('backward-recv')
    add_to_logging('backward-send')
Deepak Narayanan's avatar
Deepak Narayanan committed
738
    add_to_logging('backward-send-forward-recv')
739
    add_to_logging('backward-params-all-reduce')
740
    add_to_logging('backward-embedding-all-reduce')
mohammad's avatar
mohammad committed
741
742
743
744
    add_to_logging('optimizer-copy-to-master-grad')
    add_to_logging('optimizer-unscale-and-check-inf')
    add_to_logging('optimizer-clip-master-grad')
    add_to_logging('optimizer-copy-master-to-model-params')
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
745
    add_to_logging('optimizer')
mohammad's avatar
mohammad committed
746
    add_to_logging('batch-generator')
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
747

mohammad's avatar
mohammad committed
748
    # Calculate batch size.
mshoeybi's avatar
mshoeybi committed
749
750
751
    batch_size = args.micro_batch_size * args.data_parallel_size * \
        get_num_microbatches()

mohammad's avatar
mohammad committed
752
753
754
    total_iterations = total_loss_dict[advanced_iters_key] + \
                       total_loss_dict[skipped_iters_key]

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
755
    # Tensorboard values.
mohammad's avatar
mohammad committed
756
757
758
    if writer and is_last_rank():
        writer.add_scalar('learning-rate', learning_rate, iteration)
        writer.add_scalar('learning-rate vs samples', learning_rate,
759
                          args.consumed_train_samples)
mohammad's avatar
mohammad committed
760
761
        writer.add_scalar('batch-size', batch_size, iteration)
        writer.add_scalar('batch-size vs samples', batch_size,
762
                          args.consumed_train_samples)
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
763
        for key in loss_dict:
mohammad's avatar
mohammad committed
764
765
            writer.add_scalar(key , loss_dict[key], iteration)
            writer.add_scalar(key + ' vs samples', loss_dict[key],
766
                              args.consumed_train_samples)
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
767
        if args.fp16:
mohammad's avatar
mohammad committed
768
769
            writer.add_scalar('loss-scale', loss_scale, iteration)
            writer.add_scalar('loss-scale vs samples', loss_scale,
770
                              args.consumed_train_samples)
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
771
        timers.write(timers_to_log, writer, iteration,
mohammad's avatar
mohammad committed
772
                     normalizer=total_iterations)
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
773
774
775

    if iteration % args.log_interval == 0:
        elapsed_time = timers('interval time').elapsed()
mohammad's avatar
mohammad committed
776
        elapsed_time_per_iteration = elapsed_time / total_iterations
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
777
        if writer and torch.distributed.get_rank() == 0:
mohammad's avatar
mohammad committed
778
779
            writer.add_scalar('iteration-time',
                              elapsed_time_per_iteration, iteration)
780
781
        log_string = ' iteration {:8d}/{:8d} |'.format(
            iteration, args.train_iters)
mshoeybi's avatar
mshoeybi committed
782
        log_string += ' consumed samples: {:12d} |'.format(
783
            args.consumed_train_samples)
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
784
        log_string += ' elapsed time per iteration (ms): {:.1f} |'.format(
mohammad's avatar
mohammad committed
785
            elapsed_time_per_iteration * 1000.0)
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
786
        log_string += ' learning rate: {:.3E} |'.format(learning_rate)
mohammad's avatar
mohammad committed
787
        log_string += ' global batch size: {:5d} |'.format(batch_size)
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
788
        for key in total_loss_dict:
mohammad's avatar
mohammad committed
789
790
791
792
            if key not in [advanced_iters_key, skipped_iters_key,
                           nan_iters_key]:
                avg = total_loss_dict[key].item() / \
                      float(max(1, total_loss_dict[advanced_iters_key]))
793
794
795
                if avg > 0.0:
                    log_string += ' {}: {:.6E} |'.format(key, avg)
                total_loss_dict[key] = torch.cuda.FloatTensor([0.0])
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
796
797
        if args.fp16:
            log_string += ' loss scale: {:.1f} |'.format(loss_scale)
mohammad's avatar
mohammad committed
798
799
        log_string += ' number of skipped iterations: {:3d} |'.format(
            total_loss_dict[skipped_iters_key])
mohammad's avatar
mohammad committed
800
        log_string += ' number of nan iterations: {:3d} |'.format(
mohammad's avatar
mohammad committed
801
802
            total_loss_dict[nan_iters_key])
        total_loss_dict[advanced_iters_key] = 0
mohammad's avatar
mohammad committed
803
        total_loss_dict[skipped_iters_key] = 0
mohammad's avatar
mohammad committed
804
        total_loss_dict[nan_iters_key] = 0
805
        print_rank_last(log_string)
806
807
808
        if report_memory_flag and learning_rate > 0.:
            # Report memory after optimizer state has been initialized.
            report_memory('(after {} iterations)'.format(iteration))
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
809
810
811
812
813
814
            report_memory_flag = False
        timers.log(timers_to_log, normalizer=args.log_interval)

    return report_memory_flag


815
816
817
818
819
820
821
822
823
824
825
826
def save_checkpoint_and_time(iteration, model, optimizer, lr_scheduler):
    timers = get_timers()
    # Extra barrier is added to make sure
    # all ranks report the max time.
    torch.distributed.barrier()
    timers('save checkpoint').start()
    save_checkpoint(iteration, model, optimizer, lr_scheduler)
    torch.distributed.barrier()
    timers('save checkpoint').stop()
    timers.log(['save checkpoint'])


827
def train(forward_step_func, model, optimizer, lr_scheduler,
828
          train_data_iterator, valid_data_iterator):
829
    """Train the model function."""
Mohammad's avatar
Mohammad committed
830
831
    args = get_args()
    timers = get_timers()
832

833
834
835
    # Write args to tensorboard
    write_args_to_tensorboard()

836
837
838
839
840
841
842
843
844
845
    # Turn on training mode which enables dropout.
    model.train()

    # Tracking loss.
    total_loss_dict = {}

    # Iterations.
    iteration = args.iteration

    timers('interval time').start()
846
    print_datetime('before the start of training step')
847
848
    report_memory_flag = True
    while iteration < args.train_iters:
mohammad's avatar
mohammad committed
849
        update_num_microbatches(args.consumed_train_samples)
850
851
852
853
        loss_dict, skipped_iter = train_step(forward_step_func,
                                             train_data_iterator,
                                             model,
                                             optimizer,
Mohammad's avatar
Mohammad committed
854
                                             lr_scheduler)
855
        iteration += 1
856
        args.consumed_train_samples += mpu.get_data_parallel_world_size() * \
857
                                       args.micro_batch_size * \
mohammad's avatar
mohammad committed
858
                                       get_num_microbatches()
859
860

        # Logging.
Mohammad's avatar
Mohammad committed
861
862
        loss_scale = None
        if args.fp16:
mohammad's avatar
mohammad committed
863
            loss_scale = optimizer.get_loss_scale().item()
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
864
865
        report_memory_flag = training_log(loss_dict, total_loss_dict,
                                          optimizer.param_groups[0]['lr'],
Mohammad's avatar
Mohammad committed
866
                                          iteration, loss_scale,
mohammad's avatar
mohammad committed
867
                                          report_memory_flag, skipped_iter)
868
869

        # Autoresume
870
871
        if args.adlr_autoresume and \
           (iteration % args.adlr_autoresume_interval == 0):
872
            check_adlr_autoresume_termination(iteration, model, optimizer,
873
                                              lr_scheduler)
874
875
876
877
878
879

        # Evaluation
        if args.eval_interval and iteration % args.eval_interval == 0 and \
           args.do_valid:
            prefix = 'iteration {}'.format(iteration)
            evaluate_and_print_results(prefix, forward_step_func,
880
                                       valid_data_iterator, model,
Mohammad's avatar
Mohammad committed
881
                                       iteration, False)
882

883
884
885
886
887
888
889
890
        # Checkpointing
        saved_checkpoint = False
        if args.save and args.save_interval and \
           iteration % args.save_interval == 0:
            save_checkpoint_and_time(iteration, model, optimizer,
                                     lr_scheduler)
            saved_checkpoint = True

891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
        # Exiting based on duration
        if args.exit_duration_in_mins:
            train_time = (time.time() - _TRAIN_START_TIME) / 60.0
            done_cuda = torch.cuda.IntTensor(
                [train_time > args.exit_duration_in_mins])
            torch.distributed.all_reduce(
                done_cuda, op=torch.distributed.ReduceOp.MAX)
            done = done_cuda.item()
            if done:
                if not saved_checkpoint:
                    save_checkpoint_and_time(iteration, model, optimizer,
                                             lr_scheduler)
                print_datetime('exiting program after {} minutes'.format(train_time))                
                sys.exit()

        # Exiting based on iterations        
907
        if args.exit_interval and iteration % args.exit_interval == 0:
908
909
910
            if not saved_checkpoint:
                save_checkpoint_and_time(iteration, model, optimizer,
                                         lr_scheduler)
911
            torch.distributed.barrier()
912
            print_datetime('exiting program at iteration {}'.format(iteration))                
Mohammad's avatar
Mohammad committed
913
            sys.exit()
914

915

mohammad's avatar
mohammad committed
916
    return iteration
917
918


Mohammad's avatar
Mohammad committed
919
def evaluate(forward_step_func, data_iterator, model, verbose=False):
920
    """Evaluation."""
Mohammad's avatar
Mohammad committed
921
    args = get_args()
922
923
924
925
926
927
928
929
930
931
932
933
934

    # Turn on evaluation mode which disables dropout.
    model.eval()

    total_loss_dict = {}

    with torch.no_grad():
        iteration = 0
        while iteration < args.eval_iters:
            iteration += 1
            if verbose and iteration % args.log_interval == 0:
                print_rank_0('Evaluating iter {}/{}'.format(iteration,
                                                            args.eval_iters))
935

mohammad's avatar
mohammad committed
936
            for _ in range(get_num_microbatches()):
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
                if not mpu.is_pipeline_first_stage():
                    input_tensor, _ = communicate(
                        tensor_send_next=None,
                        tensor_send_prev=None,
                        recv_forward=True,
                        recv_backward=False)
                else:
                    input_tensor = None

                # Forward evaluation.
                output_tensor = forward_step_func(data_iterator, model, input_tensor)

                if mpu.is_pipeline_last_stage():
                    _, loss_dict = output_tensor
                    # Reduce across processes.
                    for key in loss_dict:
                        total_loss_dict[key] = total_loss_dict.get(key, torch.cuda.FloatTensor([0.0])) + \
                            loss_dict[key]
                else:
                    communicate(
                        tensor_send_next=output_tensor,
                        tensor_send_prev=None,
                        recv_forward=False,
                        recv_backward=False)
961

962
            args.consumed_valid_samples += mpu.get_data_parallel_world_size() \
963
                                           * args.micro_batch_size \
mohammad's avatar
mohammad committed
964
                                           * get_num_microbatches()
965
966
967
968
    # Move model back to the train mode.
    model.train()

    for key in total_loss_dict:
mohammad's avatar
mohammad committed
969
        total_loss_dict[key] /= args.eval_iters * get_num_microbatches()
970
971
972
973
974

    return total_loss_dict

def evaluate_and_print_results(prefix, forward_step_func,
                               data_iterator, model,
Mohammad's avatar
Mohammad committed
975
                               iteration, verbose=False):
976
    """Helper function to evaluate and dump results on screen."""
Mohammad's avatar
Mohammad committed
977
978
979
    writer = get_tensorboard_writer()

    total_loss_dict = evaluate(forward_step_func, data_iterator, model, verbose)
980
981
982
983
984
985
986
987
988
989
990
991
    string = ' validation loss at {} | '.format(prefix)
    for key in total_loss_dict:
        string += '{} value: {:.6E} | '.format(key, total_loss_dict[key].item())
        ppl = math.exp(min(20, total_loss_dict[key].item()))
        string += '{} PPL: {:.6E} | '.format(key, ppl)
        if writer and torch.distributed.get_rank() == 0:
            writer.add_scalar('{} value'.format(key),
                              total_loss_dict[key].item(),
                              iteration)
            writer.add_scalar('{} ppl'.format(key), ppl, iteration)

    length = len(string) + 1
992
993
994
    print_rank_last('-' * length)
    print_rank_last(string)
    print_rank_last('-' * length)
995
996


997
998
999
def build_train_valid_test_data_iterators(
        build_train_valid_test_datasets_provider):
    """XXX"""
Mohammad's avatar
Mohammad committed
1000
    args = get_args()
1001

1002
1003
1004
    (train_dataloader, valid_dataloader, test_dataloader) = (None, None, None)

    print_rank_0('> building train, validation, and test datasets ...')
1005
1006
1007

    # Backward compatibility, assume fixed batch size.
    if args.iteration > 0 and args.consumed_train_samples == 0:
1008
1009
        assert args.train_samples is None, \
            'only backward compatiblity support for iteration-based training'
mohammad's avatar
mohammad committed
1010
        args.consumed_train_samples = args.iteration * args.global_batch_size
1011
    if args.iteration > 0 and args.consumed_valid_samples == 0:
1012
1013
        assert args.train_samples is None, \
            'only backward compatiblity support for iteration-based training'
1014
        args.consumed_valid_samples = (args.iteration // args.eval_interval) * \
mohammad's avatar
mohammad committed
1015
            args.eval_iters * args.global_batch_size
1016

1017
    # Data loader only on rank 0 of each model parallel group.
1018
    if mpu.get_tensor_model_parallel_rank() == 0:
1019
1020

        # Number of train/valid/test samples.
1021
1022
1023
1024
1025
1026
        if args.train_samples:
            train_samples = args.train_samples
        else:
            train_samples = args.train_iters * args.global_batch_size
        eval_iters = (args.train_iters // args.eval_interval + 1) * \
                     args.eval_iters
1027
        test_iters = args.eval_iters
1028
        train_val_test_num_samples = [train_samples,
mohammad's avatar
mohammad committed
1029
1030
                                      eval_iters * args.global_batch_size,
                                      test_iters * args.global_batch_size]
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
        print_rank_0(' > datasets target sizes (minimum size):')
        print_rank_0('    train:      {}'.format(train_val_test_num_samples[0]))
        print_rank_0('    validation: {}'.format(train_val_test_num_samples[1]))
        print_rank_0('    test:       {}'.format(train_val_test_num_samples[2]))

        # Build the datasets.
        train_ds, valid_ds, test_ds = build_train_valid_test_datasets_provider(
            train_val_test_num_samples)

        # Build dataloders.
1041
1042
1043
1044
1045
        train_dataloader = build_pretraining_data_loader(
            train_ds, args.consumed_train_samples)
        valid_dataloader = build_pretraining_data_loader(
            valid_ds, args.consumed_valid_samples)
        test_dataloader = build_pretraining_data_loader(test_ds, 0)
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058

        # Flags to know if we need to do training/validation/testing.
        do_train = train_dataloader is not None and args.train_iters > 0
        do_valid = valid_dataloader is not None and args.eval_iters > 0
        do_test = test_dataloader is not None and args.eval_iters > 0
        # Need to broadcast num_tokens and num_type_tokens.
        flags = torch.cuda.LongTensor(
            [int(do_train), int(do_valid), int(do_test)])
    else:
        flags = torch.cuda.LongTensor([0, 0, 0])

    # Broadcast num tokens.
    torch.distributed.broadcast(flags,
1059
1060
                                mpu.get_tensor_model_parallel_src_rank(),
                                group=mpu.get_tensor_model_parallel_group())
1061
1062
1063
1064
1065
1066
1067
    args.do_train = flags[0].item()
    args.do_valid = flags[1].item()
    args.do_test = flags[2].item()

    # Build iterators.
    if train_dataloader is not None:
        train_data_iterator = iter(train_dataloader)
1068
1069
1070
    else:
        train_data_iterator = None

1071
1072
    if valid_dataloader is not None:
        valid_data_iterator = iter(valid_dataloader)
1073
    else:
1074
        valid_data_iterator = None
1075

1076
1077
    if test_dataloader is not None:
        test_data_iterator = iter(test_dataloader)
1078
1079
1080
    else:
        test_data_iterator = None

1081
    return train_data_iterator, valid_data_iterator, test_data_iterator