training.py 41.7 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
236

    # Wrap into fp16 optimizer.
    if args.fp16:
mohammad's avatar
mohammad committed
237
238
        optimizer = get_megatron_optimizer(optimizer)
        '''
239
240
241
242
243
        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
244
                                       'min_scale': args.min_scale,
245
                                       'delayed_shift': args.hysteresis})
mohammad's avatar
mohammad committed
246
        '''
247
248
249
    return optimizer


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

254
255
256
257
258
    # 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
259
260
        if args.lr_warmup_fraction is not None:
            warmup_steps = args.lr_warmup_fraction * decay_steps
261
262
        else:
            warmup_steps = args.lr_warmup_iters * args.global_batch_size
263
264
265
266
267
    # 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.
268
        update_train_iters(args)
269
270
271
        if args.lr_decay_samples is None:
            args.lr_decay_samples = args.train_samples
        decay_steps = args.lr_decay_samples
272
273
        if args.lr_warmup_fraction is not None:
            warmup_steps = args.lr_warmup_fraction * decay_steps
274
275
        else:
            warmup_steps = args.lr_warmup_samples
276
    else:
277
278
279
        raise Exception(
            'either train-iters or train-samples should be provided.')

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

    return lr_scheduler


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

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

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

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

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

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

328
329
330
    return model, optimizer, lr_scheduler


331
332
333
334
335
336
337
338
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
339
    tensor_shape = (args.seq_length, args.micro_batch_size, args.hidden_size)
340
341
342
    dtype = args.params_dtype
    if args.fp32_residual_connection:
        dtype = torch.float
343
344
345
    if recv_forward:
        tensor_recv_prev = torch.empty(tensor_shape,
                                       requires_grad=True,
346
                                       device=torch.cuda.current_device(),
347
                                       dtype=dtype)
348
349
350
    if recv_backward:
        tensor_recv_next = torch.empty(tensor_shape,
                                       requires_grad=True,
351
                                       device=torch.cuda.current_device(),
352
                                       dtype=dtype)
353
354
355
356
357
358

    # 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,
359
                                    group=mpu.get_pipeline_model_parallel_group())
360
361
362
363
364

    return tensor_recv_prev, tensor_recv_next


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

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

373
    # Backward pass.
mohammad's avatar
mohammad committed
374
375
376
377
378
    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:
379
380
381
382
        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
383
    '''
384
385
386
387
388
389
390
391
    # 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


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

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

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

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

    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:
437
        timers('backward-recv').start()
438
439
440
441
442
        _, output_tensor_grad = communicate(
            tensor_send_next=None,
            tensor_send_prev=None,
            recv_forward=False,
            recv_backward=True)
443
        timers('backward-recv').stop()
444
445

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

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


461
462
463
464
465
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):
466
467
    args = get_args()

468
469
470
471
472
473
474
    # 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
475
        output_tensor = loss / get_num_microbatches()
476
477
478
        output_tensor_grad = None
        losses_reduced.append(loss_reduced)
    else:
Deepak Narayanan's avatar
Deepak Narayanan committed
479
        timers('forward-send-backward-recv').start()
480
481
482
483
484
        _, 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
485
        timers('forward-send-backward-recv').stop()
486
487
488
489
490
491
492
493
494
495
496
497
498
499

    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
500
        timers('backward-send-forward-recv').start()
501
502
503
504
505
        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
506
        timers('backward-send-forward-recv').stop()
507
508
509
510
511
512
    else:
        input_tensor = None

    return input_tensor


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

518
    losses_reduced = []
mohammad's avatar
mohammad committed
519
    for i in range(get_num_microbatches()):
520
521
        timers('forward-compute').start()
        loss, loss_reduced = forward_step_func(data_iterator, model, input_tensor=None)
mohammad's avatar
mohammad committed
522
        output_tensor = loss / get_num_microbatches()
523
524
525
526
527
528
529
530
531
532
        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
533

534
535
536
537
538
539
540

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
541
    num_microbatches = get_num_microbatches()
542
543
544
545
546
    num_warmup_microbatches = \
        (mpu.get_pipeline_model_parallel_world_size() -
         mpu.get_pipeline_model_parallel_rank() - 1)
    num_warmup_microbatches = min(
        num_warmup_microbatches,
547
548
549
        num_microbatches)
    num_microbatches_remaining = \
        num_microbatches - num_warmup_microbatches
550
551
552
553
554

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

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

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

    # Run 1F1B.
577
578
    for i in range(num_microbatches_remaining):
        last_iteration = (i == (num_microbatches_remaining - 1))
579
580
581
582
583
584
585
        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)

586
587
    # Run cooldown backward passes.
    for i in range(num_warmup_microbatches):
588
589
590
591
592
593
594
595
596
597
598
599
600
        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
601
602
    optimizer.zero_grad()
    '''
603
604
605
606
    if args.fp16:
        optimizer.zero_grad(set_grads_to_None=True)
    else:
        optimizer.zero_grad()
mohammad's avatar
mohammad committed
607
    '''
608
609
610
611
612
613
614

    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)
615
616
617

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

623
624
625
626
    # 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).
627
    timers('backward-embedding-all-reduce').start()
628
    if (mpu.is_pipeline_first_stage() or mpu.is_pipeline_last_stage()) and \
629
            mpu.get_pipeline_model_parallel_world_size() > 1:
630
631
632
633
        unwrapped_model = model
        while isinstance(unwrapped_model, (torchDDP, LocalDDP, FP16_Module)):
            unwrapped_model = unwrapped_model.module

634
635
636
637
        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())
638
    timers('backward-embedding-all-reduce').stop()
639

640
    # Update master gradients.
mohammad's avatar
mohammad committed
641
    '''
642
643
644
645
    timers('backward-master-grad').start()
    if args.fp16:
        optimizer.update_master_grads()
    timers('backward-master-grad').stop()
mohammad's avatar
mohammad committed
646
    '''
647
    # Clipping gradients helps prevent the exploding gradient.
mohammad's avatar
mohammad committed
648
    '''
649
    timers('backward-clip-grad').start()
650
    if args.clip_grad > 0.:
651
        if not args.fp16:
652
653
654
655
656
657
658
659
            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)
660
661
        else:
            optimizer.clip_master_grads(args.clip_grad)
662
    timers('backward-clip-grad').stop()
mohammad's avatar
mohammad committed
663
    '''
664
665
666

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

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

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


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

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

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

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

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

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

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

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

    return report_memory_flag


814
815
816
817
818
819
820
821
822
823
824
825
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'])


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

832
833
834
    # Write args to tensorboard
    write_args_to_tensorboard()

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

    # Tracking loss.
    total_loss_dict = {}

    # Iterations.
    iteration = args.iteration

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

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

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

        # 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,
879
                                       valid_data_iterator, model,
Mohammad's avatar
Mohammad committed
880
                                       iteration, False)
881

882
883
884
885
886
887
888
889
        # 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

890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
        # 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        
906
        if args.exit_interval and iteration % args.exit_interval == 0:
907
908
909
            if not saved_checkpoint:
                save_checkpoint_and_time(iteration, model, optimizer,
                                         lr_scheduler)
910
            torch.distributed.barrier()
911
            print_datetime('exiting program at iteration {}'.format(iteration))                
Mohammad's avatar
Mohammad committed
912
            sys.exit()
913

914

mohammad's avatar
mohammad committed
915
    return iteration
916
917


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

    # 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))
934

mohammad's avatar
mohammad committed
935
            for _ in range(get_num_microbatches()):
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
                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)
960

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

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

    return total_loss_dict

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

    total_loss_dict = evaluate(forward_step_func, data_iterator, model, verbose)
979
980
981
982
983
984
985
986
987
988
989
990
    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
991
992
993
    print_rank_last('-' * length)
    print_rank_last(string)
    print_rank_last('-' * length)
994
995


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

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

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

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

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

        # Number of train/valid/test samples.
1020
1021
1022
1023
1024
1025
        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
1026
        test_iters = args.eval_iters
1027
        train_val_test_num_samples = [train_samples,
mohammad's avatar
mohammad committed
1028
1029
                                      eval_iters * args.global_batch_size,
                                      test_iters * args.global_batch_size]
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
        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.
1040
1041
1042
1043
1044
        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)
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057

        # 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,
1058
1059
                                mpu.get_tensor_model_parallel_src_rank(),
                                group=mpu.get_tensor_model_parallel_group())
1060
1061
1062
1063
1064
1065
1066
    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)
1067
1068
1069
    else:
        train_data_iterator = None

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

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

1080
    return train_data_iterator, valid_data_iterator, test_data_iterator