training.py 37 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 torch
from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP
from apex.optimizers import FusedAdam as Adam

Neel Kant's avatar
Neel Kant committed
25
from megatron import get_args
Mohammad's avatar
Mohammad committed
26
27
from megatron import get_timers
from megatron import get_tensorboard_writer
mohammad's avatar
mohammad committed
28
29
from megatron import get_num_microbatches
from megatron import update_num_microbatches
30
from megatron import mpu
Neel Kant's avatar
Neel Kant committed
31
from megatron import print_rank_0
32
from megatron import print_rank_last
Mohammad's avatar
Mohammad committed
33
34
from megatron.checkpointing import load_checkpoint
from megatron.checkpointing import save_checkpoint
35
36
from megatron.fp16 import FP16_Module
from megatron.fp16 import FP16_Optimizer
Mohammad's avatar
Mohammad committed
37
from megatron.initialize import initialize_megatron
38
39
40
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
41
from megatron.model.realm_model import ICTBertModel
42
from megatron.utils import check_adlr_autoresume_termination
43
from megatron.data.data_loaders import build_pretraining_data_loader
44
from megatron.utils import report_memory
45
46


47
def pretrain(train_valid_test_dataset_provider, model_provider,
48
             forward_step_func, extra_args_provider=None, args_defaults={}):
49
50
51
    """Main training program.

    This function will run the followings in the order provided:
Mohammad's avatar
Mohammad committed
52
53
        1) initialize Megatron.
        2) setup model, optimizer and lr schedule using the model_provider.
54
        3) call train_val_test_data_provider to get train/val/test datasets.
Mohammad's avatar
Mohammad committed
55
        4) train the modle using the forward_step_func.
56
57

    Arguments:
58
59
60
        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
61
62
63
64
65
66
67
68
69
70
            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.
71
72
    """

73
    # Initalize and get arguments, timers, and Tensorboard writer.
74
75
    initialize_megatron(extra_args_provider=extra_args_provider,
                        args_defaults=args_defaults)
76

77
    args = get_args()
Mohammad's avatar
Mohammad committed
78
    timers = get_timers()
79
80

    # Model, optimizer, and learning rate.
Mohammad's avatar
Mohammad committed
81
82
83
    timers('model and optimizer').start()
    model, optimizer, lr_scheduler = setup_model_and_optimizer(model_provider)
    timers('model and optimizer').stop()
84
85

    # Data stuff.
86
87
88
89
90
    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()
Mohammad's avatar
Mohammad committed
91
92
93

    # Print setup timing.
    print_rank_0('done with setups ...')
94
    timers.log(['model and optimizer', 'train/valid/test data iterators'])
Mohammad's avatar
Mohammad committed
95
    print_rank_0('training ...')
96
97

    iteration = 0
98
    if args.do_train and args.train_iters > 0:
mohammad's avatar
mohammad committed
99
100
101
        iteration = train(forward_step_func,
                          model, optimizer, lr_scheduler,
                          train_data_iterator, valid_data_iterator)
Mohammad's avatar
Mohammad committed
102

103
104
105
    if args.do_valid:
        prefix = 'the end of training for val data'
        evaluate_and_print_results(prefix, forward_step_func,
106
                                   valid_data_iterator, model,
Mohammad's avatar
Mohammad committed
107
                                   iteration, False)
108
109

    if args.save and iteration != 0:
110
        save_checkpoint(iteration, model, optimizer, lr_scheduler)
111
112
113
114
115
116

    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
117
                                   0, True)
118

119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
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]):
            update_num_microbatches(consumed_samples)
            consumed_samples += get_num_microbatches() * \
                                args.micro_batch_size * \
                                args.data_parallel_size
            iterations += 1
        # Reset
        update_num_microbatches(0)
        # 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))

150

Mohammad's avatar
Mohammad committed
151
def get_model(model_provider_func):
152
    """Build the model."""
Mohammad's avatar
Mohammad committed
153
    args = get_args()
154
155

    # Build model on cpu.
Mohammad's avatar
Mohammad committed
156
    model = model_provider_func()
157
158
159

    # Print number of parameters.
    if mpu.get_data_parallel_rank() == 0:
160
        print(' > number of parameters on (tensor, pipeline) '
161
              'model parallel rank ({}, {}): {}'.format(
162
163
            mpu.get_tensor_model_parallel_rank(),
            mpu.get_pipeline_model_parallel_rank(),
164
165
166
167
168
169
170
171
172
173
174
            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
175
176
        model = torchDDP(model, device_ids=[i], output_device=i,
                         process_group=mpu.get_data_parallel_group())
177
178
        return model
    if args.DDP_impl == 'local':
Mohammad's avatar
Mohammad committed
179
        model = LocalDDP(model)
180
181
        return model

182
    raise NotImplementedError('Unknown DDP implementation specified: {}. '
183
                              'Exiting.'.format(args.DDP_impl))
184
185


Mohammad's avatar
Mohammad committed
186
def get_optimizer(model):
187
    """Set up the optimizer."""
Mohammad's avatar
Mohammad committed
188
    args = get_args()
189
190

    # Build parameter groups (weight decay and non-decay).
Mohammad's avatar
Mohammad committed
191
    while isinstance(model, (torchDDP, LocalDDP, FP16_Module)):
192
193
194
195
196
197
        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']:
198
199
            if not hasattr(param, 'tensor_model_parallel'):
                param.tensor_model_parallel = False
200
201

    # Use Adam.
202
203
    optimizer = Adam(param_groups, lr=args.lr, weight_decay=args.weight_decay,
        betas=(args.adam_beta1, args.adam_beta2), eps=args.adam_eps)
204
205
206
207
208
209
210
211

    # Wrap into fp16 optimizer.
    if args.fp16:
        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
212
                                       'min_scale': args.min_scale,
213
214
215
216
217
                                       'delayed_shift': args.hysteresis})

    return optimizer


Mohammad's avatar
Mohammad committed
218
def get_learning_rate_scheduler(optimizer):
219
    """Build the learning rate scheduler."""
Mohammad's avatar
Mohammad committed
220
    args = get_args()
221

222
223
224
225
226
    # 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
227
228
229
230
        if args.lr_warmup_percent is not None:
            warmup_steps = args.lr_warmup_percent * decay_steps
        else:
            warmup_steps = args.lr_warmup_iters * args.global_batch_size
231
232
233
234
235
    # 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.
236
        update_train_iters(args)
237
238
239
        if args.lr_decay_samples is None:
            args.lr_decay_samples = args.train_samples
        decay_steps = args.lr_decay_samples
240
241
242
243
        if args.lr_warmup_percent is not None:
            warmup_steps = args.lr_warmup_percent * decay_steps
        else:
            warmup_steps = args.lr_warmup_samples
244
    else:
245
246
247
        raise Exception(
            'either train-iters or train-samples should be provided.')

248
249
    lr_scheduler = AnnealingLR(
        optimizer,
250
        max_lr=args.lr,
251
        min_lr=args.min_lr,
252
253
        warmup_steps=warmup_steps,
        decay_steps=decay_steps,
254
        decay_style=args.lr_decay_style,
255
256
257
258
259
260
        use_checkpoint_lr_scheduler=args.use_checkpoint_lr_scheduler,
        override_lr_scheduler=args.override_lr_scheduler)

    return lr_scheduler


Mohammad's avatar
Mohammad committed
261
def setup_model_and_optimizer(model_provider_func):
262
    """Setup model and optimizer."""
Mohammad's avatar
Mohammad committed
263
    args = get_args()
264

Mohammad's avatar
Mohammad committed
265
266
267
    model = get_model(model_provider_func)
    optimizer = get_optimizer(model)
    lr_scheduler = get_learning_rate_scheduler(optimizer)
268
269

    if args.load is not None:
270
        args.iteration = load_checkpoint(model, optimizer, lr_scheduler)
271
272
273
    else:
        args.iteration = 0

mohammad's avatar
mohammad committed
274
    # We only support local DDP with multiple micro-batches.
mohammad's avatar
mohammad committed
275
276
277
    if get_num_microbatches() > 1:
        assert args.DDP_impl == 'local'

Neel Kant's avatar
Neel Kant committed
278
279
280
281
282
    # get model without FP16 and/or TorchDDP wrappers
    unwrapped_model = model
    while hasattr(unwrapped_model, 'module'):
        unwrapped_model = unwrapped_model.module

283
284
    if args.iteration == 0 and hasattr(unwrapped_model,
                                       'init_state_dict_from_bert'):
285
        print("Initializing ICT from pretrained BERT model", flush=True)
286
        unwrapped_model.init_state_dict_from_bert()
Neel Kant's avatar
Neel Kant committed
287

288
289
290
    return model, optimizer, lr_scheduler


291
292
293
294
295
296
297
298
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
299
    tensor_shape = (args.seq_length, args.micro_batch_size, args.hidden_size)
300
301
302
    if recv_forward:
        tensor_recv_prev = torch.empty(tensor_shape,
                                       requires_grad=True,
303
304
                                       device=torch.cuda.current_device(),
                                       dtype=args.params_dtype)
305
306
307
    if recv_backward:
        tensor_recv_next = torch.empty(tensor_shape,
                                       requires_grad=True,
308
309
                                       device=torch.cuda.current_device(),
                                       dtype=args.params_dtype)
310
311
312
313
314
315

    # 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,
316
                                    group=mpu.get_pipeline_model_parallel_group())
317
318
319
320
321

    return tensor_recv_prev, tensor_recv_next


def backward_step(optimizer, model, input_tensor, output_tensor, output_tensor_grad):
322
    """Backward step."""
Mohammad's avatar
Mohammad committed
323
324
    args = get_args()
    timers = get_timers()
325

326
327
328
329
    # Retain the grad on the input_tensor.
    if input_tensor is not None:
        input_tensor.retain_grad()

330
    # Backward pass.
331
332
333
334
335
336
337
338
339
340
341
342
343
344
    if args.fp16:
        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)

    # 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


345
346
347
def forward_step_with_communication(forward_step_func, data_iterator, model,
                                    input_tensors, output_tensors,
                                    losses_reduced, timers):
348
349
    args = get_args()

350
    if not mpu.is_pipeline_first_stage():
351
        timers('forward-recv').start()
352
353
354
355
356
        input_tensor, _ = communicate(
            tensor_send_next=None,
            tensor_send_prev=None,
            recv_forward=True,
            recv_backward=False)
357
        timers('forward-recv').stop()
358
359
360
361
    else:
        input_tensor = None

    # Forward model for one step.
362
    timers('forward-compute').start()
363
    output_tensor = forward_step_func(data_iterator, model, input_tensor)
364
    timers('forward-compute').stop()
365
366
367

    if mpu.is_pipeline_last_stage():
        loss, loss_reduced = output_tensor
mohammad's avatar
mohammad committed
368
        output_tensor = loss / get_num_microbatches()
369
370
        losses_reduced.append(loss_reduced)
    else:
371
        timers('forward-send').start()
372
373
374
375
376
        communicate(
            tensor_send_next=output_tensor,
            tensor_send_prev=None,
            recv_forward=False,
            recv_backward=False)
377
        timers('forward-send').stop()
378
379
380
381
382
383
384
385
386
387
388
389

    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:
390
        timers('backward-recv').start()
391
392
393
394
395
        _, output_tensor_grad = communicate(
            tensor_send_next=None,
            tensor_send_prev=None,
            recv_forward=False,
            recv_backward=True)
396
        timers('backward-recv').stop()
397
398

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

    if not mpu.is_pipeline_first_stage():
405
        timers('backward-send').start()
406
407
408
409
410
        communicate(
            tensor_send_next=None,
            tensor_send_prev=input_grad_tensor,
            recv_forward=False,
            recv_backward=False)
411
        timers('backward-send').stop()
412
413


414
415
416
417
418
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):
419
420
    args = get_args()

421
422
423
424
425
426
427
    # 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
428
        output_tensor = loss / get_num_microbatches()
429
430
431
        output_tensor_grad = None
        losses_reduced.append(loss_reduced)
    else:
Deepak Narayanan's avatar
Deepak Narayanan committed
432
        timers('forward-send-backward-recv').start()
433
434
435
436
437
        _, 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
438
        timers('forward-send-backward-recv').stop()
439
440
441
442
443
444
445
446
447
448
449
450
451
452

    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
453
        timers('backward-send-forward-recv').start()
454
455
456
457
458
        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
459
        timers('backward-send-forward-recv').stop()
460
461
462
463
464
465
    else:
        input_tensor = None

    return input_tensor


466
467
468
def forward_backward_no_pipelining(forward_step_func, data_iterator, model,
                                   optimizer, timers):
    """Run forward and backward passes without inter-stage communication."""
469
470
    args = get_args()

471
    losses_reduced = []
mohammad's avatar
mohammad committed
472
    for i in range(get_num_microbatches()):
473
474
        timers('forward-compute').start()
        loss, loss_reduced = forward_step_func(data_iterator, model, input_tensor=None)
mohammad's avatar
mohammad committed
475
        output_tensor = loss / get_num_microbatches()
476
477
478
479
480
481
482
483
484
485
        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
486

487
488
489
490
491
492
493

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
494
    num_microbatches = get_num_microbatches()
495
496
497
498
499
    num_warmup_microbatches = \
        (mpu.get_pipeline_model_parallel_world_size() -
         mpu.get_pipeline_model_parallel_rank() - 1)
    num_warmup_microbatches = min(
        num_warmup_microbatches,
500
501
502
        num_microbatches)
    num_microbatches_remaining = \
        num_microbatches - num_warmup_microbatches
503
504
505
506
507

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

508
509
    # Run warmup forward passes.
    for i in range(num_warmup_microbatches):
510
511
512
513
        forward_step_with_communication(
            forward_step_func, data_iterator, model,
            input_tensors, output_tensors,
            losses_reduced, timers)
514

515
    # Before running 1F1B, need to receive first forward tensor.
516
517
    # If all microbatches are run in warmup / cooldown phase, then no need to
    # receive this tensor here.
518
    if num_microbatches_remaining > 0:
519
520
521
        if mpu.is_pipeline_first_stage():
            input_tensor = None
        else:
522
            timers('forward-recv').start()
523
524
525
526
            input_tensor, _ = communicate(tensor_send_next=None,
                                          tensor_send_prev=None,
                                          recv_forward=True,
                                          recv_backward=False)
527
            timers('forward-recv').stop()
528
529

    # Run 1F1B.
530
531
    for i in range(num_microbatches_remaining):
        last_iteration = (i == (num_microbatches_remaining - 1))
532
533
534
535
536
537
538
        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)

539
540
    # Run cooldown backward passes.
    for i in range(num_warmup_microbatches):
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
        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.
    if args.fp16:
        optimizer.zero_grad(set_grads_to_None=True)
    else:
        optimizer.zero_grad()

    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)
565
566
567

    # All-reduce if needed.
    if args.DDP_impl == 'local':
568
        timers('backward-params-all-reduce').start()
569
570
        model.allreduce_params(reduce_after=False,
                               fp32_allreduce=args.fp32_allreduce)
571
        timers('backward-params-all-reduce').stop()
572

573
574
575
576
    # 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).
577
    timers('backward-embedding-all-reduce').start()
578
    if (mpu.is_pipeline_first_stage() or mpu.is_pipeline_last_stage()) and \
579
            mpu.get_pipeline_model_parallel_world_size() > 1:
580
581
582
583
        unwrapped_model = model
        while isinstance(unwrapped_model, (torchDDP, LocalDDP, FP16_Module)):
            unwrapped_model = unwrapped_model.module

584
585
586
587
        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())
588
    timers('backward-embedding-all-reduce').stop()
589

590
591
592
593
594
595
    # Update master gradients.
    timers('backward-master-grad').start()
    if args.fp16:
        optimizer.update_master_grads()
    timers('backward-master-grad').stop()

596
    # Clipping gradients helps prevent the exploding gradient.
597
    timers('backward-clip-grad').start()
598
    if args.clip_grad > 0.:
599
        if not args.fp16:
600
601
602
603
604
605
606
607
            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)
608
609
        else:
            optimizer.clip_master_grads(args.clip_grad)
610
    timers('backward-clip-grad').stop()
611
612
613
614
615
616
617
618
619

    # Update parameters.
    timers('optimizer').start()
    optimizer.step()
    timers('optimizer').stop()

    # Update learning rate.
    skipped_iter = 0
    if not (args.fp16 and optimizer.overflow):
620
621
622
623
        increment = get_num_microbatches() * \
                    args.micro_batch_size * \
                    args.data_parallel_size
        lr_scheduler.step(increment=increment)
624
625
626
    else:
        skipped_iter = 1

627
    if mpu.is_pipeline_last_stage():
628
629
630
631
        # Average loss across microbatches.
        loss_reduced = {}
        for key in losses_reduced[0]:
            losses_reduced_for_key = [x[key] for x in losses_reduced]
632
            loss_reduced[key] = sum(losses_reduced_for_key) / len(losses_reduced_for_key)
633
634
        return loss_reduced, skipped_iter
    return {}, skipped_iter
635
636


Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
637
def training_log(loss_dict, total_loss_dict, learning_rate, iteration,
mohammad's avatar
mohammad committed
638
                 loss_scale, report_memory_flag, skipped_iter):
Mohammad's avatar
Mohammad committed
639
640
641
642
    """Log training information such as losses, timing, ...."""
    args = get_args()
    timers = get_timers()
    writer = get_tensorboard_writer()
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
643
644

    # Update losses.
mohammad's avatar
mohammad committed
645
646
647
    skipped_iters_key = 'skipped iterations'
    total_loss_dict[skipped_iters_key] = total_loss_dict.get(
        skipped_iters_key, 0) + skipped_iter
mohammad's avatar
mohammad committed
648
    got_nan_key = 'got nan'
mohammad's avatar
mohammad committed
649
650

    got_nan = False
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
651
    for key in loss_dict:
mohammad's avatar
mohammad committed
652
        if not skipped_iter:
653
654
            total_loss_dict[key] = total_loss_dict.get(
                key, torch.cuda.FloatTensor([0.0])) + loss_dict[key]
mohammad's avatar
mohammad committed
655
656
657
658
659
        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
660
661
662
663
            got_nan = got_nan or is_nan

    total_loss_dict[got_nan_key] = total_loss_dict.get(
        got_nan_key, 0) + int(got_nan)
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
664
665
666

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

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
668
669
670
    def add_to_logging(name):
        if name in timers.timers:
            timers_to_log.append(name)
671
672
673
    add_to_logging('forward-compute')
    add_to_logging('forward-recv')
    add_to_logging('forward-send')
Deepak Narayanan's avatar
Deepak Narayanan committed
674
    add_to_logging('forward-send-backward-recv')
675
676
677
    add_to_logging('backward-compute')
    add_to_logging('backward-recv')
    add_to_logging('backward-send')
Deepak Narayanan's avatar
Deepak Narayanan committed
678
    add_to_logging('backward-send-forward-recv')
679
    add_to_logging('backward-master-grad')
680
    add_to_logging('backward-params-all-reduce')
681
    add_to_logging('backward-embedding-all-reduce')
682
    add_to_logging('backward-clip-grad')
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
    add_to_logging('optimizer')
    add_to_logging('batch generator')

    # Tensorboard values.
    if writer and torch.distributed.get_rank() == 0:
        writer.add_scalar('learning_rate', learning_rate, iteration)
        for key in loss_dict:
            writer.add_scalar(key, loss_dict[key], iteration)
        if args.fp16:
            writer.add_scalar('loss_scale', loss_scale, iteration)
        normalizer = iteration % args.log_interval
        if normalizer == 0:
            normalizer = args.log_interval
        timers.write(timers_to_log, writer, iteration,
                     normalizer=normalizer)

    if iteration % args.log_interval == 0:
        elapsed_time = timers('interval time').elapsed()
        if writer and torch.distributed.get_rank() == 0:
            writer.add_scalar('iteration_time',
                              elapsed_time / args.log_interval, iteration)
704
705
        log_string = ' iteration {:8d}/{:8d} |'.format(
            iteration, args.train_iters)
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
706
707
708
        log_string += ' elapsed time per iteration (ms): {:.1f} |'.format(
            elapsed_time * 1000.0 / args.log_interval)
        log_string += ' learning rate: {:.3E} |'.format(learning_rate)
mohammad's avatar
mohammad committed
709
710
        num_iterations = max(
            1, args.log_interval - total_loss_dict[skipped_iters_key])
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
711
        for key in total_loss_dict:
mohammad's avatar
mohammad committed
712
            if key not in [skipped_iters_key, got_nan_key]:
mohammad's avatar
mohammad committed
713
                avg = total_loss_dict[key].item() / float(num_iterations)
714
715
716
                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
717
718
        if args.fp16:
            log_string += ' loss scale: {:.1f} |'.format(loss_scale)
mohammad's avatar
mohammad committed
719
720
        log_string += ' number of skipped iterations: {:3d} |'.format(
            total_loss_dict[skipped_iters_key])
mohammad's avatar
mohammad committed
721
722
        log_string += ' number of nan iterations: {:3d} |'.format(
            total_loss_dict[got_nan_key])
mohammad's avatar
mohammad committed
723
        total_loss_dict[skipped_iters_key] = 0
mohammad's avatar
mohammad committed
724
        total_loss_dict[got_nan_key] = 0
725
        print_rank_last(log_string)
726
727
728
        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
729
730
731
732
733
734
            report_memory_flag = False
        timers.log(timers_to_log, normalizer=args.log_interval)

    return report_memory_flag


735
def train(forward_step_func, model, optimizer, lr_scheduler,
736
          train_data_iterator, valid_data_iterator):
737
    """Train the model function."""
Mohammad's avatar
Mohammad committed
738
739
    args = get_args()
    timers = get_timers()
740
741
742
743
744
745
746
747
748
749
750
751
752

    # Turn on training mode which enables dropout.
    model.train()

    # Tracking loss.
    total_loss_dict = {}

    # Iterations.
    iteration = args.iteration

    timers('interval time').start()
    report_memory_flag = True
    while iteration < args.train_iters:
mohammad's avatar
mohammad committed
753
        update_num_microbatches(args.consumed_train_samples)
754
755
756
757
        loss_dict, skipped_iter = train_step(forward_step_func,
                                             train_data_iterator,
                                             model,
                                             optimizer,
Mohammad's avatar
Mohammad committed
758
                                             lr_scheduler)
759
        iteration += 1
760
        args.consumed_train_samples += mpu.get_data_parallel_world_size() * \
761
                                       args.micro_batch_size * \
mohammad's avatar
mohammad committed
762
                                       get_num_microbatches()
763
764

        # Logging.
Mohammad's avatar
Mohammad committed
765
766
767
        loss_scale = None
        if args.fp16:
            loss_scale = optimizer.loss_scale
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
768
769
        report_memory_flag = training_log(loss_dict, total_loss_dict,
                                          optimizer.param_groups[0]['lr'],
Mohammad's avatar
Mohammad committed
770
                                          iteration, loss_scale,
mohammad's avatar
mohammad committed
771
                                          report_memory_flag, skipped_iter)
772
773

        # Autoresume
774
775
        if args.adlr_autoresume and \
           (iteration % args.adlr_autoresume_interval == 0):
776
            check_adlr_autoresume_termination(iteration, model, optimizer,
777
                                              lr_scheduler)
778
779
780
781

        # Checkpointing
        if args.save and args.save_interval and \
           iteration % args.save_interval == 0:
782
            save_checkpoint(iteration, model, optimizer, lr_scheduler)
783
784
785
786
787
788

        # 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,
789
                                       valid_data_iterator, model,
Mohammad's avatar
Mohammad committed
790
                                       iteration, False)
791
792

        if args.exit_interval and iteration % args.exit_interval == 0:
793
            torch.distributed.barrier()
794
795
            time_str = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            rank = torch.distributed.get_rank()
Mohammad's avatar
Mohammad committed
796
797
798
            print_rank_0('rank: {} | time: {} | exiting the program at '
                         'iteration {}'.format(rank, time_str, iteration))
            sys.exit()
799

mohammad's avatar
mohammad committed
800
    return iteration
801
802


Mohammad's avatar
Mohammad committed
803
def evaluate(forward_step_func, data_iterator, model, verbose=False):
804
    """Evaluation."""
Mohammad's avatar
Mohammad committed
805
    args = get_args()
806
807
808
809
810
811
812
813
814
815
816
817
818

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

mohammad's avatar
mohammad committed
820
            for _ in range(get_num_microbatches()):
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
                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)
845

846
            args.consumed_valid_samples += mpu.get_data_parallel_world_size() \
847
                                           * args.micro_batch_size \
mohammad's avatar
mohammad committed
848
                                           * get_num_microbatches()
849
850
851
852
    # Move model back to the train mode.
    model.train()

    for key in total_loss_dict:
mohammad's avatar
mohammad committed
853
        total_loss_dict[key] /= args.eval_iters * get_num_microbatches()
854
855
856
857
858

    return total_loss_dict

def evaluate_and_print_results(prefix, forward_step_func,
                               data_iterator, model,
Mohammad's avatar
Mohammad committed
859
                               iteration, verbose=False):
860
    """Helper function to evaluate and dump results on screen."""
Mohammad's avatar
Mohammad committed
861
862
863
    writer = get_tensorboard_writer()

    total_loss_dict = evaluate(forward_step_func, data_iterator, model, verbose)
864
865
866
867
868
869
870
871
872
873
874
875
    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
876
877
878
    print_rank_last('-' * length)
    print_rank_last(string)
    print_rank_last('-' * length)
879
880


881
882
883
def build_train_valid_test_data_iterators(
        build_train_valid_test_datasets_provider):
    """XXX"""
Mohammad's avatar
Mohammad committed
884
    args = get_args()
885

886
887
888
    (train_dataloader, valid_dataloader, test_dataloader) = (None, None, None)

    print_rank_0('> building train, validation, and test datasets ...')
889
890
891

    # Backward compatibility, assume fixed batch size.
    if args.iteration > 0 and args.consumed_train_samples == 0:
892
893
        assert args.train_samples is None, \
            'only backward compatiblity support for iteration-based training'
mohammad's avatar
mohammad committed
894
        args.consumed_train_samples = args.iteration * args.global_batch_size
895
    if args.iteration > 0 and args.consumed_valid_samples == 0:
896
897
        assert args.train_samples is None, \
            'only backward compatiblity support for iteration-based training'
898
        args.consumed_valid_samples = (args.iteration // args.eval_interval) * \
mohammad's avatar
mohammad committed
899
            args.eval_iters * args.global_batch_size
900

901
    # Data loader only on rank 0 of each model parallel group.
902
    if mpu.get_tensor_model_parallel_rank() == 0:
903
904

        # Number of train/valid/test samples.
905
906
907
908
909
910
        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
911
        test_iters = args.eval_iters
912
        train_val_test_num_samples = [train_samples,
mohammad's avatar
mohammad committed
913
914
                                      eval_iters * args.global_batch_size,
                                      test_iters * args.global_batch_size]
915
916
917
918
919
920
921
922
923
924
        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.
925
926
927
928
929
        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)
930
931
932
933
934
935
936
937
938
939
940
941
942

        # 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,
943
944
                                mpu.get_tensor_model_parallel_src_rank(),
                                group=mpu.get_tensor_model_parallel_group())
945
946
947
948
949
950
951
    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)
952
953
954
    else:
        train_data_iterator = None

955
956
    if valid_dataloader is not None:
        valid_data_iterator = iter(valid_dataloader)
957
    else:
958
        valid_data_iterator = None
959

960
961
    if test_dataloader is not None:
        test_data_iterator = iter(test_dataloader)
962
963
964
    else:
        test_data_iterator = None

965
    return train_data_iterator, valid_data_iterator, test_data_iterator