initialize.py 4.74 KB
Newer Older
Mohammad's avatar
Mohammad committed
1
# coding=utf-8
Mohammad's avatar
Mohammad committed
2
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
Mohammad's avatar
Mohammad committed
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
#
# 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.

"""Megatron initialization."""

18
import datetime
Mohammad's avatar
Mohammad committed
19
20
21
import random
import os

22
import numpy as np
Mohammad's avatar
Mohammad committed
23
24
import torch

25
26
27
from megatron import get_adlr_autoresume
from megatron import get_args
from megatron import get_tensorboard_writer
Mohammad's avatar
Mohammad committed
28
from megatron import mpu
29
from megatron.global_vars import set_global_variables
Mohammad's avatar
Mohammad committed
30
31


32
def initialize_megatron(extra_args_provider=None, args_defaults={},
Raul Puri's avatar
Raul Puri committed
33
                        ignore_unknown_args=False, allow_no_cuda=False):
Mohammad's avatar
Mohammad committed
34
    """Set global variables, initialize distributed, and
Raul Puri's avatar
Raul Puri committed
35
36
37
38
    set autoresume and random seeds.
    `allow_no_cuda` should not be set unless using megatron for cpu only 
    data processing. In general this arg should not be set unless you know 
    what you are doing."""
Raul Puri's avatar
Raul Puri committed
39
40
41
    if not allow_no_cuda:
        # Make sure cuda is available.
        assert torch.cuda.is_available(), 'Megatron requires CUDA.'
Mohammad's avatar
Mohammad committed
42

Mohammad's avatar
Mohammad committed
43
44
    # Parse args, build tokenizer, and set adlr-autoresume,
    # tensorboard-writer, and timers.
Mohammad's avatar
Mohammad committed
45
    set_global_variables(extra_args_provider=extra_args_provider,
46
47
                         args_defaults=args_defaults,
                         ignore_unknown_args=ignore_unknown_args)
Mohammad's avatar
Mohammad committed
48
49
50
51
52
53
54

    # Pytorch distributed.
    _initialize_distributed()

    # Autoresume.
    _init_autoresume()

55
    # Random seeds for reproducibility.
Mohammad's avatar
Mohammad committed
56
57
58
59
60
    args = get_args()
    if args.rank == 0:
        print('> setting random seeds to {} ...'.format(args.seed))
    _set_random_seed(args.seed)

Mohammad's avatar
Mohammad committed
61
62
63
    # Write arguments to tensorboard.
    _write_args_to_tensorboard()

Mohammad's avatar
Mohammad committed
64

Neel Kant's avatar
Neel Kant committed
65
def init_distributed():
Mohammad's avatar
Mohammad committed
66
67
    args = get_args()

Raul Puri's avatar
Raul Puri committed
68
    device_count = torch.cuda.device_count()
Mohammad's avatar
Mohammad committed
69
70
71
72
73
74
75
    if torch.distributed.is_initialized():

        if args.rank == 0:
            print('torch distributed is already initialized, '
                  'skipping initialization ...', flush=True)
        args.rank = torch.distributed.get_rank()
        args.world_size = torch.distributed.get_world_size()
76
77
        if device_count > 0:
            device = torch.cuda.current_device()
Raul Puri's avatar
Raul Puri committed
78
            local_rank = args.rank % device_count
79
80
            assert local_rank == device, \
                'expected local-rank to be the same as rank % device-count.'
Mohammad's avatar
Mohammad committed
81
82
83
84
85
86

    else:

        if args.rank == 0:
            print('> initializing torch distributed ...', flush=True)
        # Manually set the device ids.
87
        if device_count > 0:
Raul Puri's avatar
Raul Puri committed
88
            device = args.rank % device_count
89
90
91
92
93
94
            if args.local_rank is not None:
                assert args.local_rank == device, \
                    'expected local-rank to be the same as rank % device-count.'
            else:
                args.local_rank = device
            torch.cuda.set_device(device)
Mohammad's avatar
Mohammad committed
95
96
97
98
99
100
101
102
103
104
        # Call the init process
        init_method = 'tcp://'
        master_ip = os.getenv('MASTER_ADDR', 'localhost')
        master_port = os.getenv('MASTER_PORT', '6000')
        init_method += master_ip + ':' + master_port
        torch.distributed.init_process_group(
            backend=args.distributed_backend,
            world_size=args.world_size, rank=args.rank,
            init_method=init_method)

Neel Kant's avatar
Neel Kant committed
105
106
107
108
109
110
111

def _initialize_distributed():
    """Initialize torch.distributed and mpu."""
    init_distributed()
    args = get_args()
    device_count = torch.cuda.device_count()

Mohammad's avatar
Mohammad committed
112
    # Set the model-parallel / data-parallel communicators.
113
114
    if device_count > 0:
        mpu.initialize_model_parallel(args.model_parallel_size)
Mohammad's avatar
Mohammad committed
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131


def _init_autoresume():
    """Set autoresume start time."""
    autoresume = get_adlr_autoresume()
    if autoresume:
        torch.distributed.barrier()
        autoresume.init()
        torch.distributed.barrier()


def _set_random_seed(seed):
    """Set random seed for reproducability."""
    if seed is not None and seed > 0:
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
132
133
        if torch.cuda.device_count() > 0:
            mpu.model_parallel_cuda_manual_seed(seed)
Mohammad's avatar
Mohammad committed
134
135
    else:
        raise ValueError('Seed ({}) should be a positive integer.'.format(seed))
Mohammad's avatar
Mohammad committed
136
137
138
139
140
141
142
143
144


def _write_args_to_tensorboard():
    """Write arguments to tensorboard."""
    args = get_args()
    writer = get_tensorboard_writer()
    if writer:
        for arg in vars(args):
            writer.add_text(arg, str(getattr(args, arg)))