checkpointing.py 22 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
16
17
18
19
20
21
22
23
24
#
# 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.

"""Input/output checkpointing."""

import os
import random
import sys
import numpy as np

import torch

25
26
27
from megatron import (get_args,
                      mpu,
                      print_rank_0,
28
29
                      update_num_microbatches,
                      utils)
30

31
32
33
34
# >>>
from lutil import pax
# <<<

Vijay Korthikanti's avatar
Vijay Korthikanti committed
35
36
37
38
_CHECKPOINT_VERSION = None

def set_checkpoint_version(value):
    global _CHECKPOINT_VERSION
Jared Casper's avatar
Jared Casper committed
39
40
41
    if _CHECKPOINT_VERSION is not None:
        assert _CHECKPOINT_VERSION == value, \
            "checkpoint versions do not match"
Vijay Korthikanti's avatar
Vijay Korthikanti committed
42
43
44
45
46
    _CHECKPOINT_VERSION = value

def get_checkpoint_version():
    global _CHECKPOINT_VERSION
    return _CHECKPOINT_VERSION
47
48
49

def check_checkpoint_args(checkpoint_args):
    """Ensure fixed arguments for a model are the same for the input
50
    arguments and the one retrieved from checkpoint."""
51
52
    args = get_args()

53
54
55
56
57
    def _compare(arg_name, old_arg_name=None):
        if old_arg_name is not None:
            checkpoint_value = getattr(checkpoint_args, old_arg_name)
        else:
            checkpoint_value = getattr(checkpoint_args, arg_name)
58
59
60
61
62
63
64
65
66
        args_value = getattr(args, arg_name)
        error_message = '{} value from checkpoint ({}) is not equal to the ' \
                        'input argument value ({}).'.format(
                            arg_name, checkpoint_value, args_value)
        assert checkpoint_value == args_value, error_message

    _compare('num_layers')
    _compare('hidden_size')
    _compare('num_attention_heads')
Vijay Korthikanti's avatar
Vijay Korthikanti committed
67
    if args.vocab_file:
Vijay Korthikanti's avatar
Vijay Korthikanti committed
68
        _compare('max_position_embeddings')
69
70
71
        _compare('make_vocab_size_divisible_by')
        _compare('padded_vocab_size')
        _compare('tokenizer_type')
Vijay Korthikanti's avatar
Vijay Korthikanti committed
72
73
    if args.data_parallel_random_init:
        _compare('data_parallel_random_init')
74
75
76
77
78
79
    if get_checkpoint_version() < 3.0:
        _compare('tensor_model_parallel_size',
                 old_arg_name='model_parallel_size')
    if get_checkpoint_version() >= 3.0:
        _compare('tensor_model_parallel_size')
        _compare('pipeline_model_parallel_size')
80
81
82
83
84
85
86
87

def ensure_directory_exists(filename):
    """Build filename's path if it does not already exists."""
    dirname = os.path.dirname(filename)
    if not os.path.exists(dirname):
        os.makedirs(dirname)


88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
# >>
# def get_checkpoint_name(checkpoints_path, iteration,
#                         release=False):
#     """A unified checkpoint name."""
#     if release:
#         directory = 'release'
#     else:
#         directory = 'iter_{:07d}'.format(iteration)
#     # Use both the tensor and pipeline MP rank.
#     if mpu.get_pipeline_model_parallel_world_size() == 1:
#         return os.path.join(checkpoints_path, directory,
#                             'mp_rank_{:02d}'.format(
#                                 mpu.get_tensor_model_parallel_rank()),
#                             'model_optim_rng.pt')
#     return os.path.join(checkpoints_path, directory,
#                         'mp_rank_{:02d}_{:03d}'.format(
#                             mpu.get_tensor_model_parallel_rank(),
#                             mpu.get_pipeline_model_parallel_rank()),
#                         'model_optim_rng.pt')
107
108
def get_checkpoint_names(checkpoints_path, iteration, use_distributed_optimizer,
                         release=False):
109
110
111
112
113
    """A unified checkpoint name."""
    if release:
        directory = 'release'
    else:
        directory = 'iter_{:07d}'.format(iteration)
114
    # Use both the tensor and pipeline MP rank.
115
116
117
    common_path = os.path.join(
        checkpoints_path,
        directory,
118
        "mp_rank_%02d_%03d" % (
119
            mpu.get_tensor_model_parallel_rank(),
120
            mpu.get_pipeline_model_parallel_rank()))
121
    model_name = os.path.join(common_path, "model_rng.pt")
122
123
124
125
126
127
    if use_distributed_optimizer:
        optim_name = os.path.join(
            common_path + "_%03d" % mpu.get_data_parallel_rank(),
            "optim.pt")
    else:
        optim_name = os.path.join(common_path, "optim.pt")
128
129
    return model_name, optim_name
# <<<
130
131
132
133
134
135
136
137


def get_checkpoint_tracker_filename(checkpoints_path):
    """Tracker file rescords the latest chckpoint during
    training to restart from."""
    return os.path.join(checkpoints_path, 'latest_checkpointed_iteration.txt')


138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
def read_metadata(tracker_filename):
    # Read the tracker file and either set the iteration or
    # mark it as a release checkpoint.
    iteration = 0
    release = False
    with open(tracker_filename, 'r') as f:
        metastring = f.read().strip()
        try:
            iteration = int(metastring)
        except ValueError:
            release = metastring == 'release'
            if not release:
                print_rank_0('ERROR: Invalid metadata file {}. Exiting'.format(
                    tracker_filename))
                sys.exit()
    assert iteration > 0 or release, 'error parsing metadata file {}'.format(
        tracker_filename)

156
157
158
159
    # Get the max iteration retrieved across the ranks.
    iters_cuda = torch.cuda.LongTensor([iteration])
    torch.distributed.all_reduce(iters_cuda, op=torch.distributed.ReduceOp.MAX)
    max_iter = iters_cuda[0].item()
160
161
162
163

    # We should now have all the same iteration.
    # If not, print a warning and chose the maximum
    # iteration across all ranks.
164
165
166
167
168
    if iteration != max_iter:
        print('WARNING: on rank {} found iteration {} in the '
              'metadata while max iteration across the ranks '
              'is {}, replacing it with max iteration.'.format(
                  rank, iteration, max_iter), flush=True)
169
170
171
    return max_iter, release


172
173
def get_rng_state():
    """ collect rng state across data parallel ranks """
174
    args = get_args()
175
176
177
178
179
180
181
182
183
    rng_state = {
        'random_rng_state': random.getstate(),
        'np_rng_state': np.random.get_state(),
        'torch_rng_state': torch.get_rng_state(),
        'cuda_rng_state': torch.cuda.get_rng_state(),
        'rng_tracker_states': mpu.get_cuda_rng_tracker().get_states()}

    rng_state_list = None
    if torch.distributed.is_initialized() and \
184
185
            mpu.get_data_parallel_world_size() > 1 and \
            args.data_parallel_random_init:
186
187
188
        rng_state_list = \
            [None for i in range(mpu.get_data_parallel_world_size())]
        torch.distributed.all_gather_object(
189
            rng_state_list,
190
            rng_state,
191
192
193
194
195
196
197
            group=mpu.get_data_parallel_group())
    else:
        rng_state_list = [rng_state]

    return rng_state_list


198
def save_checkpoint(iteration, model, optimizer, opt_param_scheduler):
199
200
201
202
    """Save a model checkpoint."""
    args = get_args()

    # Only rank zero of the data parallel writes to the disk.
203
    model = utils.unwrap_model(model)
204

Jared Casper's avatar
Jared Casper committed
205
206
    print_rank_0('saving checkpoint at iteration {:7d} to {}'.format(
        iteration, args.save))
207

208
    # Collect rng state across data parallel ranks.
209
210
    rng_state = get_rng_state()

211
212
    # Checkpoint file names.
    model_checkpoint_name, optim_checkpoint_name = \
213
214
215
216
217
218
        get_checkpoint_names(args.save, iteration, args.use_distributed_optimizer)

    pax(0, {
        "model_checkpoint_name" : model_checkpoint_name,
        "optim_checkpoint_name" : optim_checkpoint_name,
    })
219
220
221
222

    # Save args, model, RNG.
    if not torch.distributed.is_initialized() \
       or mpu.get_data_parallel_rank() == 0:
223
224
225
226

        # Arguments, iteration, and model.
        state_dict = {}
        state_dict['args'] = args
227
        state_dict['checkpoint_version'] = 3.0
228
        state_dict['iteration'] = iteration
229
230
231
232
233
234
        if len(model) == 1:
            state_dict['model'] = model[0].state_dict_for_save_checkpoint()
        else:
            for i in range(len(model)):
                mpu.set_virtual_pipeline_model_parallel_rank(i)
                state_dict['model%d' % i] = model[i].state_dict_for_save_checkpoint()
235

236
237
238
239
240
241
242
243
        # >>>
        # # Optimizer stuff.
        # if not args.no_save_optim:
        #     if optimizer is not None:
        #         state_dict['optimizer'] = optimizer.state_dict()
        #     if opt_param_scheduler is not None:
        #         state_dict['opt_param_scheduler'] = opt_param_scheduler.state_dict()
        # <<<
244
245
246

        # RNG states.
        if not args.no_save_rng:
247
            state_dict["rng_state"] = rng_state
248
249

        # Save.
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
        ensure_directory_exists(model_checkpoint_name)
        torch.save(state_dict, model_checkpoint_name)

    # >>>
    # Save optimizer state.
    if not args.no_save_optim \
       and (not torch.distributed.is_initialized()
            or mpu.get_data_parallel_rank() == 0
            or args.use_distributed_optimizer):

        # Optimizer stuff.
        state_dict = {}
        if optimizer is not None:
            state_dict['optimizer'] = optimizer.state_dict()
        if opt_param_scheduler is not None:
            state_dict['opt_param_scheduler'] = opt_param_scheduler.state_dict()

        # Save.
        ensure_directory_exists(optim_checkpoint_name)
        torch.save(state_dict, optim_checkpoint_name)
        # >>>
        # pax({
        #     "model_checkpoint_name" : model_checkpoint_name,
        #     "optim_checkpoint_name" : optim_checkpoint_name,
        #     "state_dict" : state_dict,
        # })
        # <<<
    # <<<
278
279

    # Wait so everyone is done (necessary)
Jared Casper's avatar
Jared Casper committed
280
281
282
283
284
285
    if torch.distributed.is_initialized():
        torch.distributed.barrier()

    print_rank_0('  successfully saved checkpoint at iteration {:7d} to {}'.format(
        iteration, args.save))

286
    # And update the latest iteration
Jared Casper's avatar
Jared Casper committed
287
    if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
288
289
290
291
292
        tracker_filename = get_checkpoint_tracker_filename(args.save)
        with open(tracker_filename, 'w') as f:
            f.write(str(iteration))

    # Wait so everyone is done (not necessary)
Jared Casper's avatar
Jared Casper committed
293
294
    if torch.distributed.is_initialized():
        torch.distributed.barrier()
295

296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
def _transpose_first_dim(t, num_splits, num_splits_first, model):
    input_shape = t.size()
    # We use a self_attention module but the values extracted aren't
    # specific to self attention so should work for cross attention as well
    while hasattr(model, 'module'):
        model = model.module
    attention_module = model.language_model.encoder.layers[0].self_attention
    hidden_size_per_attention_head = attention_module.hidden_size_per_attention_head
    num_attention_heads_per_partition = attention_module.num_attention_heads_per_partition
    if num_splits_first:
        """[num_splits * np * hn, h]
        -->(view) [num_splits, np, hn, h]
        -->(tranpose) [np, num_splits, hn, h]
        -->(view) [np * num_splits * hn, h] """

        intermediate_shape = \
            (num_splits, num_attention_heads_per_partition,
             hidden_size_per_attention_head) + input_shape[1:]

        t = t.view(*intermediate_shape)
        t = t.transpose(0, 1).contiguous()
    else:
        """[np * hn * num_splits, h]
        -->(view) [np, hn, num_splits, h]
        -->(tranpose) [np, num_splits, hn, h]
        -->(view) [np * num_splits * hn, h] """

        intermediate_shape = \
            (num_attention_heads_per_partition,
             hidden_size_per_attention_head, num_splits) +\
             input_shape[1:]

        t = t.view(*intermediate_shape)
        t = t.transpose(1, 2).contiguous()
    t = t.view(*input_shape)

    return t
333

Mostofa Patwary's avatar
Mostofa Patwary committed
334
335
336
337
338
def fix_query_key_value_ordering(model, checkpoint_version):
    """Fix up query/key/value matrix ordering if checkpoint
    version is smaller than 2.0
    """
    if checkpoint_version < 2.0:
339
340
341
        if isinstance(model, list):
            assert len(model)==1
            model = model[0]
Mostofa Patwary's avatar
Mostofa Patwary committed
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
        for name, param in model.named_parameters():
            if name.endswith(('.query_key_value.weight', '.query_key_value.bias')):
                if checkpoint_version == 0:
                    fixed_param = _transpose_first_dim(param.data, 3, True, model)
                elif checkpoint_version == 1.0:
                    fixed_param = _transpose_first_dim(param.data, 3, False, model)
                else:
                    print_rank_0(f"Invalid checkpoint version {checkpoint_version}.")
                    sys.exit()
                param.data.copy_(fixed_param)
            if name.endswith(('.key_value.weight', '.key_value.bias')):
                if checkpoint_version == 0:
                    fixed_param = _transpose_first_dim(param.data, 2, True, model)
                elif checkpoint_version == 1.0:
                    fixed_param = _transpose_first_dim(param.data, 2, False, model)
                else:
                    print_rank_0(f"Invalid checkpoint version {checkpoint_version}.")
                    sys.exit()
                param.data.copy_(fixed_param)
        print_rank_0(" succesfully fixed query-key-values ordering for"
                    " checkpoint version {}".format(checkpoint_version))

364
def load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', strict=True):
365
366
367
368
369
    """Load a model checkpoint and return the iteration.
    strict (bool): whether to strictly enforce that the keys in
        :attr:`state_dict` of the checkpoint match the names of
        parameters and buffers in model.
    """
370
    args = get_args()
371
    load_dir = getattr(args, load_arg)
372

373
    model = utils.unwrap_model(model)
374

375
    # Read the tracker file and set the iteration.
376
    tracker_filename = get_checkpoint_tracker_filename(load_dir)
377
378
379
380
381
382
383
384
385
386
387

    # If no tracker file, return iretation zero.
    if not os.path.isfile(tracker_filename):
        print_rank_0('WARNING: could not find the metadata file {} '.format(
            tracker_filename))
        print_rank_0('    will not load any checkpoints and will start from '
                     'random')
        return 0

    # Otherwise, read the tracker file and either set the iteration or
    # mark it as a release checkpoint.
388
    iteration, release = read_metadata(tracker_filename)
389
390

    # Checkpoint.
391
    model_checkpoint_name, optim_checkpoint_name = \
392
393
394
        get_checkpoint_names(load_dir, iteration,
                             args.use_distributed_optimizer,
                             release)
Jared Casper's avatar
Jared Casper committed
395
    print_rank_0(f' loading checkpoint from {args.load} at iteration {iteration}')
396
397
398

    # Load the checkpoint.
    try:
399
400
        model_state_dict = torch.load(model_checkpoint_name, map_location='cpu')
        optim_state_dict = torch.load(optim_checkpoint_name, map_location='cpu')
401
    except ModuleNotFoundError:
mohammad's avatar
mohammad committed
402
        from megatron.fp16_deprecated import loss_scaler
403
404
405
        # For backward compatibility.
        print_rank_0(' > deserializing using the old code structure ...')
        sys.modules['fp16.loss_scaler'] = sys.modules[
mohammad's avatar
mohammad committed
406
407
408
            'megatron.fp16_deprecated.loss_scaler']
        sys.modules['megatron.fp16.loss_scaler'] = sys.modules[
            'megatron.fp16_deprecated.loss_scaler']
409
410
        model_state_dict = torch.load(model_checkpoint_name, map_location='cpu')
        optim_state_dict = torch.load(optim_checkpoint_name, map_location='cpu')
411
        sys.modules.pop('fp16.loss_scaler', None)
mohammad's avatar
mohammad committed
412
        sys.modules.pop('megatron.fp16.loss_scaler', None)
413
    except BaseException as e:
414
        print_rank_0('could not load the checkpoint')
415
        print_rank_0(e)
416
417
        sys.exit()

418
419
420
421
    # >>>
    pax({"hi.": "there."})
    # <<<

Vijay Korthikanti's avatar
Vijay Korthikanti committed
422
    # set checkpoint version
423
    set_checkpoint_version(model_state_dict.get('checkpoint_version', 0))
Vijay Korthikanti's avatar
Vijay Korthikanti committed
424

425
426
427
428
429
    # Set iteration.
    if args.finetune or release:
        iteration = 0
    else:
        try:
430
            iteration = model_state_dict['iteration']
431
        except KeyError:
Neel Kant's avatar
Neel Kant committed
432
            try:  # Backward compatible with older checkpoints
433
                iteration = model_state_dict['total_iters']
434
435
436
437
438
439
440
            except KeyError:
                print_rank_0('A metadata file exists but unable to load '
                             'iteration from checkpoint {}, exiting'.format(
                                 checkpoint_name))
                sys.exit()

    # Check arguments.
mohammad's avatar
mohammad committed
441
442
    assert args.consumed_train_samples == 0
    assert args.consumed_valid_samples == 0
443
444
    if 'args' in model_state_dict:
        checkpoint_args = model_state_dict['args']
445
        check_checkpoint_args(checkpoint_args)
446
447
        args.consumed_train_samples = getattr(checkpoint_args,
                                              'consumed_train_samples', 0)
mohammad's avatar
mohammad committed
448
        update_num_microbatches(consumed_samples=args.consumed_train_samples)
449
450
        args.consumed_valid_samples = getattr(checkpoint_args,
                                              'consumed_valid_samples', 0)
451
452
453
454
    else:
        print_rank_0('could not find arguments in the checkpoint ...')

    # Model.
455
    if len(model) == 1:
456
        model[0].load_state_dict(model_state_dict['model'], strict=strict)
457
458
459
    else:
        for i in range(len(model)):
            mpu.set_virtual_pipeline_model_parallel_rank(i)
460
            model[i].load_state_dict(model_state_dict['model%d' % i], strict=strict)
461

Mostofa Patwary's avatar
Mostofa Patwary committed
462
463
464
465
    # Fix up query/key/value matrix ordering if needed
    checkpoint_version = get_checkpoint_version()
    print_rank_0(f' checkpoint version {checkpoint_version}')
    fix_query_key_value_ordering(model, checkpoint_version)
466

467
468
469
470
471
472
473
    # >>>
    # pax(0, {
    #     "model_state_dict" : model_state_dict,
    #     "optim_state_dict" : optim_state_dict,
    # })
    # <<<

474
    # Optimizer.
475
476
477
478
479
    pax({
        "release" : release,
        "finetune" : args.finetune,
        "no_load_optim" : args.no_load_optim,
    })
480
481
482
    if not release and not args.finetune and not args.no_load_optim:
        try:
            if optimizer is not None:
483
                optimizer.load_state_dict(optim_state_dict['optimizer'])
484
            if opt_param_scheduler is not None:
485
                if 'lr_scheduler' in optim_state_dict: # backward compatbility
486
                    opt_param_scheduler.load_state_dict(optim_state_dict['lr_scheduler'])
487
                else:
488
                    opt_param_scheduler.load_state_dict(optim_state_dict['opt_param_scheduler'])
489
490
491
492
493
494
495
496
497
498
        except KeyError:
            print_rank_0('Unable to load optimizer from checkpoint {}. '
                         'Specify --no-load-optim or --finetune to prevent '
                         'attempting to load the optimizer state, '
                         'exiting ...'.format(checkpoint_name))
            sys.exit()

    # rng states.
    if not release and not args.finetune and not args.no_load_rng:
        try:
499
            if 'rng_state' in model_state_dict:
500
                # access rng_state for data parallel rank
501
                if args.data_parallel_random_init:
Vijay Korthikanti's avatar
Vijay Korthikanti committed
502

503
                    rng_state = model_state_dict['rng_state'][mpu.get_data_parallel_rank()]
504
                else:
505
                    rng_state = model_state_dict['rng_state'][0]
506
507
508
509
510
511
512
513
                random.setstate(rng_state['random_rng_state'])
                np.random.set_state(rng_state['np_rng_state'])
                torch.set_rng_state(rng_state['torch_rng_state'])
                torch.cuda.set_rng_state(rng_state['cuda_rng_state'])
                # Check for empty states array
                if not rng_state['rng_tracker_states']:
                    raise KeyError
                mpu.get_cuda_rng_tracker().set_states(
Vijay Korthikanti's avatar
Vijay Korthikanti committed
514
                    rng_state['rng_tracker_states'])
515
            else:  # backward compatability
516
517
518
519
                random.setstate(model_state_dict['random_rng_state'])
                np.random.set_state(model_state_dict['np_rng_state'])
                torch.set_rng_state(model_state_dict['torch_rng_state'])
                torch.cuda.set_rng_state(model_state_dict['cuda_rng_state'])
520
                # Check for empty states array
521
                if not model_state_dict['rng_tracker_states']:
522
523
                    raise KeyError
                mpu.get_cuda_rng_tracker().set_states(
524
                    model_state_dict['rng_tracker_states'])
525
        except KeyError:
526
            print_rank_0('Unable to load rng state from checkpoint {}. '
527
                         'Specify --no-load-rng or --finetune to prevent '
528
                         'attempting to load the rng state, '
529
530
531
                         'exiting ...'.format(checkpoint_name))
            sys.exit()

Jared Casper's avatar
Jared Casper committed
532
    # Some utilities want to load a checkpoint without distributed being initialized
533
    # pax({"hi.": "there."})
Jared Casper's avatar
Jared Casper committed
534
535
536
537
538
    if torch.distributed.is_initialized():
        torch.distributed.barrier()

    print_rank_0(f'  successfully loaded checkpoint from {args.load} '
                 f'at iteration {iteration}')
539
540

    return iteration
Neel Kant's avatar
Neel Kant committed
541
542


543
544
545
546
547
548
def load_biencoder_checkpoint(model, only_query_model=False,
        only_context_model=False, custom_load_path=None):
    """
    selectively load retrieval models for indexing/retrieving 
    from saved checkpoints
    """
Neel Kant's avatar
Neel Kant committed
549
550
551

    args = get_args()

552
    model = utils.unwrap_model(model)
Neel Kant's avatar
Neel Kant committed
553

554
    load_path = custom_load_path if custom_load_path is not None else args.load
Neel Kant's avatar
Neel Kant committed
555
556
557
558
559

    tracker_filename = get_checkpoint_tracker_filename(load_path)
    with open(tracker_filename, 'r') as f:
        iteration = int(f.read().strip())

560
561
562
    checkpoint_name, _ = get_checkpoint_names(load_path, iteration,
                                              args.use_distributed_optimizer,
                                              False)
Neel Kant's avatar
Neel Kant committed
563
564
565
566
    if mpu.get_data_parallel_rank() == 0:
        print('global rank {} is loading checkpoint {}'.format(
            torch.distributed.get_rank(), checkpoint_name))

567
    state_dict = torch.load(model_checkpoint_name, map_location='cpu')
568
    ret_state_dict = state_dict['model']
Neel Kant's avatar
Neel Kant committed
569
570

    if only_query_model:
571
        ret_state_dict.pop('context_model')
Mostofa Patwary's avatar
Mostofa Patwary committed
572
    if only_context_model:
573
        ret_state_dict.pop('query_model')
Neel Kant's avatar
Neel Kant committed
574

575
576
    assert len(model) == 1
    model[0].load_state_dict(ret_state_dict)
Neel Kant's avatar
Neel Kant committed
577
578
579
580
581
    torch.distributed.barrier()

    if mpu.get_data_parallel_rank() == 0:
        print(' successfully loaded {}'.format(checkpoint_name))

Neel Kant's avatar
Neel Kant committed
582
    return model
583