training.py 21.1 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
25

import torch
from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP
from apex.optimizers import FusedAdam as Adam

Mohammad's avatar
Mohammad committed
26
27
28
from megatron import get_args
from megatron import get_timers
from megatron import get_tensorboard_writer
29
from megatron import mpu
Mohammad's avatar
Mohammad committed
30
31
32
from megatron import print_rank_0
from megatron.checkpointing import load_checkpoint
from megatron.checkpointing import save_checkpoint
33
34
from megatron.fp16 import FP16_Module
from megatron.fp16 import FP16_Optimizer
Mohammad's avatar
Mohammad committed
35
from megatron.initialize import initialize_megatron
36
37
38
from megatron.learning_rates import AnnealingLR
from megatron.model import DistributedDataParallel as LocalDDP
from megatron.model import get_params_for_weight_decay_optimization
39
from megatron.mpu.initialize import get_index_ready, get_train_group
40
from megatron.utils import check_adlr_autoresume_termination
41
from megatron.utils import make_data_loader
42
from megatron.utils import report_memory
43
44
45


INDEX_READY = None
46
47


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

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

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

75
    # Initalize and get arguments, timers, and Tensorboard writer.
76
77
78
79
80
81
82
83
84
    if initializer_func is None:
        initialize_megatron(extra_args_provider=extra_args_provider,
                            args_defaults=args_defaults)
    else:
        initializer_func(extra_args_provider=extra_args_provider,
                         args_defaults=args_defaults)
        global INDEX_READY
        INDEX_READY = get_index_ready()

85
    args = get_args()
Mohammad's avatar
Mohammad committed
86
    timers = get_timers()
87
88

    # Model, optimizer, and learning rate.
Mohammad's avatar
Mohammad committed
89
90
91
    timers('model and optimizer').start()
    model, optimizer, lr_scheduler = setup_model_and_optimizer(model_provider)
    timers('model and optimizer').stop()
92
93

    # Data stuff.
94
95
96
97
98
    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
99
100
101

    # Print setup timing.
    print_rank_0('done with setups ...')
102
    timers.log(['model and optimizer', 'train/valid/test data iterators'])
Mohammad's avatar
Mohammad committed
103
    print_rank_0('training ...')
104
105

    iteration = 0
106
    if args.do_train and args.train_iters > 0:
107
108
        iteration, _ = train(forward_step_func,
                             model, optimizer, lr_scheduler,
Neel Kant's avatar
Neel Kant committed
109
                             train_data_iterator, valid_data_iterator)
Mohammad's avatar
Mohammad committed
110

111
112
113
    if args.do_valid:
        prefix = 'the end of training for val data'
        evaluate_and_print_results(prefix, forward_step_func,
114
                                   valid_data_iterator, model,
Mohammad's avatar
Mohammad committed
115
                                   iteration, False)
116
117

    if args.save and iteration != 0:
118
        save_checkpoint(iteration, model, optimizer, lr_scheduler)
119
120
121
122
123
124

    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
125
                                   0, True)
126
127


Mohammad's avatar
Mohammad committed
128
def get_model(model_provider_func):
129
    """Build the model."""
Mohammad's avatar
Mohammad committed
130
    args = get_args()
131
132

    # Build model on cpu.
Mohammad's avatar
Mohammad committed
133
    model = model_provider_func()
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

    # Print number of parameters.
    if mpu.get_data_parallel_rank() == 0:
        print(' > number of parameters on model parallel rank {}: {}'.format(
            mpu.get_model_parallel_rank(),
            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)

    # Wrap model for distributed training."""
    if args.DDP_impl == 'torch':
        i = torch.cuda.current_device()
Mohammad's avatar
Mohammad committed
151
152
        model = torchDDP(model, device_ids=[i], output_device=i,
                         process_group=mpu.get_data_parallel_group())
153
154
        return model
    if args.DDP_impl == 'local':
Mohammad's avatar
Mohammad committed
155
        model = LocalDDP(model)
156
157
        return model

158
    raise NotImplementedError('Unknown DDP implementation specified: {}. '
159
                              'Exiting.'.format(args.DDP_impl))
160
161


Mohammad's avatar
Mohammad committed
162
def get_optimizer(model):
163
    """Set up the optimizer."""
Mohammad's avatar
Mohammad committed
164
    args = get_args()
165
166

    # Build parameter groups (weight decay and non-decay).
Mohammad's avatar
Mohammad committed
167
    while isinstance(model, (torchDDP, LocalDDP, FP16_Module)):
168
169
170
171
172
173
174
175
176
177
        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']:
            if not hasattr(param, 'model_parallel'):
                param.model_parallel = False

    # Use Adam.
Mohammad's avatar
Mohammad committed
178
    optimizer = Adam(param_groups, lr=args.lr, weight_decay=args.weight_decay)
179
180
181
182
183
184
185
186

    # 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
187
                                       'min_scale': args.min_scale,
188
189
190
191
192
                                       'delayed_shift': args.hysteresis})

    return optimizer


Mohammad's avatar
Mohammad committed
193
def get_learning_rate_scheduler(optimizer):
194
    """Build the learning rate scheduler."""
Mohammad's avatar
Mohammad committed
195
    args = get_args()
196
197
198
199
200
201
202

    # Add linear learning rate scheduler.
    if args.lr_decay_iters is not None:
        num_iters = args.lr_decay_iters
    else:
        num_iters = args.train_iters
    num_iters = max(1, num_iters)
Mohammad's avatar
Mohammad committed
203
    init_step = 0
204
205
206
207
208
    warmup_iter = args.warmup * num_iters
    lr_scheduler = AnnealingLR(
        optimizer,
        start_lr=args.lr,
        warmup_iter=warmup_iter,
Mohammad's avatar
Mohammad committed
209
        total_iters=num_iters,
210
211
212
213
214
215
216
217
218
        decay_style=args.lr_decay_style,
        last_iter=init_step,
        min_lr=args.min_lr,
        use_checkpoint_lr_scheduler=args.use_checkpoint_lr_scheduler,
        override_lr_scheduler=args.override_lr_scheduler)

    return lr_scheduler


Mohammad's avatar
Mohammad committed
219
def setup_model_and_optimizer(model_provider_func):
220
    """Setup model and optimizer."""
Mohammad's avatar
Mohammad committed
221
    args = get_args()
222

Mohammad's avatar
Mohammad committed
223
224
225
    model = get_model(model_provider_func)
    optimizer = get_optimizer(model)
    lr_scheduler = get_learning_rate_scheduler(optimizer)
226
227

    if args.load is not None:
228
        args.iteration = load_checkpoint(model, optimizer, lr_scheduler)
229
230
231
232
233
234
    else:
        args.iteration = 0

    return model, optimizer, lr_scheduler


Mohammad's avatar
Mohammad committed
235
def backward_step(optimizer, model, loss):
236
    """Backward step."""
Mohammad's avatar
Mohammad committed
237
238
    args = get_args()
    timers = get_timers()
Neel Kant's avatar
Neel Kant committed
239
    torch.cuda.synchronize()
240
241

    # Backward pass.
Mohammad's avatar
Mohammad committed
242
    optimizer.zero_grad(set_grads_to_None=True)
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
    if args.fp16:
        optimizer.backward(loss, update_master_grads=False)
    else:
        loss.backward()

    # All-reduce if needed.
    if args.DDP_impl == 'local':
        timers('allreduce').start()
        model.allreduce_params(reduce_after=False,
                               fp32_allreduce=args.fp32_allreduce)
        timers('allreduce').stop()
    # Update master gradients.
    if args.fp16:
        optimizer.update_master_grads()
    # Clipping gradients helps prevent the exploding gradient.
    if args.clip_grad > 0:
        if not args.fp16:
            mpu.clip_grad_norm(model.parameters(), args.clip_grad)
        else:
            optimizer.clip_master_grads(args.clip_grad)


Mohammad's avatar
Mohammad committed
265
266
def train_step(forward_step_func, data_iterator,
               model, optimizer, lr_scheduler):
267
    """Single training step."""
Mohammad's avatar
Mohammad committed
268
269
    args = get_args()
    timers = get_timers()
270
271
272

    # Forward model for one step.
    timers('forward').start()
Mohammad's avatar
Mohammad committed
273
    loss, loss_reduced = forward_step_func(data_iterator, model)
274
275
276
    timers('forward').stop()

    timers('backward').start()
Mohammad's avatar
Mohammad committed
277
    backward_step(optimizer, model, loss)
278
279
    timers('backward').stop()

Neel Kant's avatar
Neel Kant committed
280
281
    # Calculate gradients, reduce across processes, and clip.

282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
    # Update parameters.
    timers('optimizer').start()
    optimizer.step()
    timers('optimizer').stop()

    # Update learning rate.
    skipped_iter = 0
    if not (args.fp16 and optimizer.overflow):
        lr_scheduler.step()
    else:
        skipped_iter = 1

    return loss_reduced, skipped_iter


Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
297
def training_log(loss_dict, total_loss_dict, learning_rate, iteration,
Mohammad's avatar
Mohammad committed
298
299
300
301
302
                 loss_scale, report_memory_flag):
    """Log training information such as losses, timing, ...."""
    args = get_args()
    timers = get_timers()
    writer = get_tensorboard_writer()
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
303
304
305
306
307
308
309

    # Update losses.
    for key in loss_dict:
        total_loss_dict[key] = total_loss_dict.get(key, 0.) + loss_dict[key]

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

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
    def add_to_logging(name):
        if name in timers.timers:
            timers_to_log.append(name)
    add_to_logging('forward')
    add_to_logging('backward')
    add_to_logging('allreduce')
    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)
        log_string = ' iteration {:8d}/{:8d} |'.format(iteration,
                                                       args.train_iters)
        log_string += ' elapsed time per iteration (ms): {:.1f} |'.format(
            elapsed_time * 1000.0 / args.log_interval)
        log_string += ' learning rate: {:.3E} |'.format(learning_rate)
        for key in total_loss_dict:
            avg = total_loss_dict[key].item() / args.log_interval
            log_string += ' {}: {:.6E} |'.format(key, avg)
            total_loss_dict[key] = 0.0
        if args.fp16:
            log_string += ' loss scale: {:.1f} |'.format(loss_scale)
        print_rank_0(log_string)
        if report_memory_flag:
            report_memory('after {} iterations'.format(iteration))
            report_memory_flag = False
        timers.log(timers_to_log, normalizer=args.log_interval)

    return report_memory_flag


358
def train(forward_step_func, model, optimizer, lr_scheduler,
359
          train_data_iterator, valid_data_iterator):
360
    """Train the model function."""
Mohammad's avatar
Mohammad committed
361
362
    args = get_args()
    timers = get_timers()
363
364
365
366
367
368
369
370
371
372
373
374
375

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

    # Tracking loss.
    total_loss_dict = {}

    # Iterations.
    iteration = args.iteration
    skipped_iters = 0

    timers('interval time').start()
    report_memory_flag = True
376
377
378

    global INDEX_READY
    recv_handle = torch.distributed.broadcast(INDEX_READY, args.max_training_rank, async_op=True)
379
    while iteration < args.train_iters:
380
381
382
383
384
385
386
387
388
389
        if hasattr(model, 'retriever') and INDEX_READY == 1:
            model.retriever.reload_index()
            save_checkpoint(iteration, model, optimizer, lr_scheduler)

            if args.rank == 0:
                INDEX_READY = 1 - INDEX_READY
                print("Switched index ready", flush=True)
            send_handle = torch.distributed.broadcast(INDEX_READY, 0, async_op=True)
            torch.distributed.barrier(get_train_group())
            recv_handle = torch.distributed.broadcast(INDEX_READY, args.max_training_rank, async_op=True)
390

391
392
393
394
        loss_dict, skipped_iter = train_step(forward_step_func,
                                             train_data_iterator,
                                             model,
                                             optimizer,
Mohammad's avatar
Mohammad committed
395
                                             lr_scheduler)
396
397
398
399
        skipped_iters += skipped_iter
        iteration += 1

        # Logging.
Mohammad's avatar
Mohammad committed
400
401
402
        loss_scale = None
        if args.fp16:
            loss_scale = optimizer.loss_scale
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
403
404
        report_memory_flag = training_log(loss_dict, total_loss_dict,
                                          optimizer.param_groups[0]['lr'],
Mohammad's avatar
Mohammad committed
405
                                          iteration, loss_scale,
Mohammad's avatar
Mohammad committed
406
                                          report_memory_flag)
407
408

        # Autoresume
409
410
        if args.adlr_autoresume and \
           (iteration % args.adlr_autoresume_interval == 0):
411
            check_adlr_autoresume_termination(iteration, model, optimizer,
412
                                              lr_scheduler)
413
414
415
416

        # Checkpointing
        if args.save and args.save_interval and \
           iteration % args.save_interval == 0:
417
            save_checkpoint(iteration, model, optimizer, lr_scheduler)
418
419
420
421
422
423

        # 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,
424
                                       valid_data_iterator, model,
Mohammad's avatar
Mohammad committed
425
                                       iteration, False)
426
427
428
429
430

        if args.exit_interval and iteration % args.exit_interval == 0:
            torch.distributed.barrier()
            time_str = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            rank = torch.distributed.get_rank()
Mohammad's avatar
Mohammad committed
431
432
433
            print_rank_0('rank: {} | time: {} | exiting the program at '
                         'iteration {}'.format(rank, time_str, iteration))
            sys.exit()
434
435
436
437

    return iteration, skipped_iters


Mohammad's avatar
Mohammad committed
438
def evaluate(forward_step_func, data_iterator, model, verbose=False):
439
    """Evaluation."""
Mohammad's avatar
Mohammad committed
440
    args = get_args()
441
442
443
444
445
446
447
448
449
450
451
452
453
454

    # 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))
            # Forward evaluation.
Mohammad's avatar
Mohammad committed
455
            _, loss_dict = forward_step_func(data_iterator, model)
456
457
458
            # Reduce across processes.
            for key in loss_dict:
                total_loss_dict[key] = total_loss_dict.get(key, 0.) + \
Neel Kant's avatar
Neel Kant committed
459
                    loss_dict[key]
460
461
462
463
464
465
466
467
468
469
470
    # Move model back to the train mode.
    model.train()

    for key in total_loss_dict:
        total_loss_dict[key] /= args.eval_iters

    return total_loss_dict


def evaluate_and_print_results(prefix, forward_step_func,
                               data_iterator, model,
Mohammad's avatar
Mohammad committed
471
                               iteration, verbose=False):
472
    """Helper function to evaluate and dump results on screen."""
Mohammad's avatar
Mohammad committed
473
474
475
    writer = get_tensorboard_writer()

    total_loss_dict = evaluate(forward_step_func, data_iterator, model, verbose)
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
    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
    print_rank_0('-' * length)
    print_rank_0(string)
    print_rank_0('-' * length)


493
494
495
def build_train_valid_test_data_iterators(
        build_train_valid_test_datasets_provider):
    """XXX"""
Mohammad's avatar
Mohammad committed
496
    args = get_args()
497

498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
    (train_dataloader, valid_dataloader, test_dataloader) = (None, None, None)

    print_rank_0('> building train, validation, and test datasets ...')
    # Data loader only on rank 0 of each model parallel group.
    if mpu.get_model_parallel_rank() == 0:
        # Rank, size, and global batch size.
        data_parallel_size = mpu.get_data_parallel_world_size()
        global_batch_size = args.batch_size * data_parallel_size

        # Number of train/valid/test samples.
        train_iters = args.train_iters
        eval_iters = (train_iters // args.eval_interval + 1) * args.eval_iters
        test_iters = args.eval_iters
        train_val_test_num_samples = [train_iters * global_batch_size,
                                      eval_iters * global_batch_size,
                                      test_iters * global_batch_size]
        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.
        train_dataloader = make_data_loader(train_ds)
        valid_dataloader = make_data_loader(valid_ds)
        test_dataloader = make_data_loader(test_ds)

        # 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,
                                mpu.get_model_parallel_src_rank(),
                                group=mpu.get_model_parallel_group())
    args.do_train = flags[0].item()
    args.do_valid = flags[1].item()
    args.do_test = flags[2].item()

Mohammad's avatar
Mohammad committed
546
    # Shift the start iterations.
547
548
    if train_dataloader is not None:
        train_dataloader.batch_sampler.start_iter = args.iteration % \
Neel Kant's avatar
Neel Kant committed
549
            len(train_dataloader)
Mohammad's avatar
Mohammad committed
550
        print_rank_0('setting training data start iteration to {}'.
551
552
                     format(train_dataloader.batch_sampler.start_iter))
    if valid_dataloader is not None:
Mohammad's avatar
Mohammad committed
553
        start_iter_val = (args.iteration // args.eval_interval) * \
Neel Kant's avatar
Neel Kant committed
554
            args.eval_iters
555
        valid_dataloader.batch_sampler.start_iter = start_iter_val % \
Neel Kant's avatar
Neel Kant committed
556
            len(valid_dataloader)
Mohammad's avatar
Mohammad committed
557
        print_rank_0('setting validation data start iteration to {}'.
558
                     format(valid_dataloader.batch_sampler.start_iter))
559

560
561
562
    # Build iterators.
    if train_dataloader is not None:
        train_data_iterator = iter(train_dataloader)
563
564
565
    else:
        train_data_iterator = None

566
567
    if valid_dataloader is not None:
        valid_data_iterator = iter(valid_dataloader)
568
    else:
569
        valid_data_iterator = None
570

571
572
    if test_dataloader is not None:
        test_data_iterator = iter(test_dataloader)
573
574
575
    else:
        test_data_iterator = None

576
    return train_data_iterator, valid_data_iterator, test_data_iterator