global_vars.py 8.63 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
18
19
#
# 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 global variables."""

import os
import sys
Mohammad's avatar
Mohammad committed
20
21
22
import time

import torch
Mohammad's avatar
Mohammad committed
23

24
from megatron import dist_signal_handler
25
from megatron.tokenizer import build_tokenizer
Mohammad's avatar
Mohammad committed
26
from .arguments import parse_args
mohammad's avatar
mohammad committed
27
from .microbatches import build_num_microbatches_calculator
Mohammad's avatar
Mohammad committed
28
29

_GLOBAL_ARGS = None
mohammad's avatar
mohammad committed
30
_GLOBAL_NUM_MICROBATCHES_CALCULATOR = None
Mohammad's avatar
Mohammad committed
31
32
33
34
_GLOBAL_TOKENIZER = None
_GLOBAL_TENSORBOARD_WRITER = None
_GLOBAL_ADLR_AUTORESUME = None
_GLOBAL_TIMERS = None
35
_GLOBAL_SIGNAL_HANDLER = None
Mohammad's avatar
Mohammad committed
36
37
38
39
40
41
42
43


def get_args():
    """Return arguments."""
    _ensure_var_is_initialized(_GLOBAL_ARGS, 'args')
    return _GLOBAL_ARGS


mohammad's avatar
mohammad committed
44
45
46
47
def get_num_microbatches():
    return _GLOBAL_NUM_MICROBATCHES_CALCULATOR.get()


48
49
50
51
52
53
54
def get_current_global_batch_size():
    return _GLOBAL_NUM_MICROBATCHES_CALCULATOR.get_current_global_batch_size()


def update_num_microbatches(consumed_samples, consistency_check=True):
    _GLOBAL_NUM_MICROBATCHES_CALCULATOR.update(consumed_samples,
                                               consistency_check)
mohammad's avatar
mohammad committed
55
56


Mohammad's avatar
Mohammad committed
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def get_tokenizer():
    """Return tokenizer."""
    _ensure_var_is_initialized(_GLOBAL_TOKENIZER, 'tokenizer')
    return _GLOBAL_TOKENIZER


def get_tensorboard_writer():
    """Return tensorboard writer. It can be None so no need
    to check if it is initialized."""
    return _GLOBAL_TENSORBOARD_WRITER


def get_adlr_autoresume():
    """ADLR autoresume object. It can be None so no need
    to check if it is initialized."""
    return _GLOBAL_ADLR_AUTORESUME


def get_timers():
    """Return timers."""
    _ensure_var_is_initialized(_GLOBAL_TIMERS, 'timers')
    return _GLOBAL_TIMERS

80
81
82
83
84
85
86
87
def get_signal_handler():
    _ensure_var_is_initialized(_GLOBAL_SIGNAL_HANDLER, 'signal handler')
    return _GLOBAL_SIGNAL_HANDLER

def _set_signal_handler():
    global _GLOBAL_SIGNAL_HANDLER
    _ensure_var_is_not_initialized(_GLOBAL_SIGNAL_HANDLER, 'signal handler')
    _GLOBAL_SIGNAL_HANDLER = dist_signal_handler.DistributedSignalHandler().__enter__()
Mohammad's avatar
Mohammad committed
88

89
90
def set_global_variables(extra_args_provider=None, args_defaults={},
                         ignore_unknown_args=False):
Mohammad's avatar
Mohammad committed
91
    """Set args, tokenizer, tensorboard-writer, adlr-autoresume, and timers."""
Mohammad's avatar
Mohammad committed
92
    args = _parse_args(extra_args_provider=extra_args_provider,
93
94
                       defaults=args_defaults,
                       ignore_unknown_args=ignore_unknown_args)
mohammad's avatar
mohammad committed
95
    _build_num_microbatches_calculator(args)
96
97
    if args.vocab_file:
        _ = _build_tokenizer(args)
Mohammad's avatar
Mohammad committed
98
99
    _set_tensorboard_writer(args)
    _set_adlr_autoresume(args)
Mohammad's avatar
Mohammad committed
100
101
    _set_timers()

102
103
104
    if args.exit_signal_handler:
        _set_signal_handler()

Mohammad's avatar
Mohammad committed
105

106
107
def _parse_args(extra_args_provider=None, defaults={},
                ignore_unknown_args=False):
Mohammad's avatar
Mohammad committed
108
109
110
    """Parse entire arguments."""
    global _GLOBAL_ARGS
    _ensure_var_is_not_initialized(_GLOBAL_ARGS, 'args')
Mohammad's avatar
Mohammad committed
111
    _GLOBAL_ARGS = parse_args(extra_args_provider=extra_args_provider,
112
113
                              defaults=defaults,
                              ignore_unknown_args=ignore_unknown_args)
Mohammad's avatar
Mohammad committed
114
    return _GLOBAL_ARGS
Mohammad's avatar
Mohammad committed
115
116


mohammad's avatar
mohammad committed
117
118
119
120
121
122
def _build_num_microbatches_calculator(args):

    global _GLOBAL_NUM_MICROBATCHES_CALCULATOR
    _ensure_var_is_not_initialized(_GLOBAL_NUM_MICROBATCHES_CALCULATOR,
                                   'num microbatches calculator')

mohammad's avatar
mohammad committed
123
124
    _GLOBAL_NUM_MICROBATCHES_CALCULATOR = build_num_microbatches_calculator(
        args)
mohammad's avatar
mohammad committed
125
126


Mohammad's avatar
Mohammad committed
127
def _build_tokenizer(args):
Mohammad's avatar
Mohammad committed
128
129
130
    """Initialize tokenizer."""
    global _GLOBAL_TOKENIZER
    _ensure_var_is_not_initialized(_GLOBAL_TOKENIZER, 'tokenizer')
Mohammad's avatar
Mohammad committed
131
    _GLOBAL_TOKENIZER = build_tokenizer(args)
Mohammad's avatar
Mohammad committed
132
133
134
135
136
137
138
    return _GLOBAL_TOKENIZER


def rebuild_tokenizer(args):
    global _GLOBAL_TOKENIZER
    _GLOBAL_TOKENIZER = None
    return _build_tokenizer(args)
Mohammad's avatar
Mohammad committed
139
140


Mohammad's avatar
Mohammad committed
141
def _set_tensorboard_writer(args):
Mohammad's avatar
Mohammad committed
142
143
144
145
146
147
    """Set tensorboard writer."""
    global _GLOBAL_TENSORBOARD_WRITER
    _ensure_var_is_not_initialized(_GLOBAL_TENSORBOARD_WRITER,
                                   'tensorboard writer')

    if hasattr(args, 'tensorboard_dir') and \
148
       args.tensorboard_dir and args.rank == (args.world_size - 1):
Mohammad's avatar
Mohammad committed
149
150
151
152
        try:
            from torch.utils.tensorboard import SummaryWriter
            print('> setting tensorboard ...')
            _GLOBAL_TENSORBOARD_WRITER = SummaryWriter(
153
154
                log_dir=args.tensorboard_dir,
                max_queue=args.tensorboard_queue_size)
Mohammad's avatar
Mohammad committed
155
156
157
158
159
160
        except ModuleNotFoundError:
            print('WARNING: TensorBoard writing requested but is not '
                  'available (are you using PyTorch 1.1.0 or later?), '
                  'no TensorBoard logs will be written.', flush=True)


Mohammad's avatar
Mohammad committed
161
def _set_adlr_autoresume(args):
Mohammad's avatar
Mohammad committed
162
163
164
165
166
167
168
169
170
171
    """Initialize ADLR autoresume."""
    global _GLOBAL_ADLR_AUTORESUME
    _ensure_var_is_not_initialized(_GLOBAL_ADLR_AUTORESUME, 'adlr autoresume')

    if args.adlr_autoresume:
        if args.rank == 0:
            print('enabling autoresume ...', flush=True)
        sys.path.append(os.environ.get('SUBMIT_SCRIPTS', '.'))
        try:
            from userlib.auto_resume import AutoResume
Neel Kant's avatar
Neel Kant committed
172
        except BaseException:
Mohammad's avatar
Mohammad committed
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
            print('ADLR autoresume is not available, exiting ...')
            sys.exit()

        _GLOBAL_ADLR_AUTORESUME = AutoResume


def _set_timers():
    """Initialize timers."""
    global _GLOBAL_TIMERS
    _ensure_var_is_not_initialized(_GLOBAL_TIMERS, 'timers')
    _GLOBAL_TIMERS = Timers()


def _ensure_var_is_initialized(var, name):
    """Make sure the input variable is not None."""
    assert var is not None, '{} is not initialized.'.format(name)


def _ensure_var_is_not_initialized(var, name):
    """Make sure the input variable is not None."""
    assert var is None, '{} is already initialized.'.format(name)
Mohammad's avatar
Mohammad committed
194
195


196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
class _Timer:
    """Timer."""

    def __init__(self, name):
        self.name_ = name
        self.elapsed_ = 0.0
        self.started_ = False
        self.start_time = time.time()

    def start(self):
        """Start the timer."""
        assert not self.started_, 'timer has already been started'
        torch.cuda.synchronize()
        self.start_time = time.time()
        self.started_ = True

    def stop(self):
        """Stop the timer."""
        assert self.started_, 'timer is not started'
        torch.cuda.synchronize()
        self.elapsed_ += (time.time() - self.start_time)
        self.started_ = False

    def reset(self):
        """Reset timer."""
        self.elapsed_ = 0.0
        self.started_ = False

    def elapsed(self, reset=True):
        """Calculate the elapsed time."""
        started_ = self.started_
        # If the timing in progress, end it first.
        if self.started_:
            self.stop()
        # Get the elapsed time.
        elapsed_ = self.elapsed_
        # Reset the elapsed time
        if reset:
            self.reset()
        # If timing was in progress, set it back.
        if started_:
            self.start()
        return elapsed_


Mohammad's avatar
Mohammad committed
241
242
243
244
245
246
247
248
class Timers:
    """Group of timers."""

    def __init__(self):
        self.timers = {}

    def __call__(self, name):
        if name not in self.timers:
249
            self.timers[name] = _Timer(name)
Mohammad's avatar
Mohammad committed
250
251
252
253
254
255
256
257
258
259
        return self.timers[name]

    def write(self, names, writer, iteration, normalizer=1.0, reset=False):
        """Write timers to a tensorboard writer"""
        # currently when using add_scalars,
        # torch.utils.add_scalars makes each timer its own run, which
        # polutes the runs list, so we just add each as a scalar
        assert normalizer > 0.0
        for name in names:
            value = self.timers[name].elapsed(reset=reset) / normalizer
mohammad's avatar
mohammad committed
260
            writer.add_scalar(name + '-time', value, iteration)
Mohammad's avatar
Mohammad committed
261
262
263
264
265
266
267

    def log(self, names, normalizer=1.0, reset=True):
        """Log a group of timers."""
        assert normalizer > 0.0
        string = 'time (ms)'
        for name in names:
            elapsed_time = self.timers[name].elapsed(
268
                reset=reset) * 1000.0 / normalizer
Mohammad's avatar
Mohammad committed
269
270
            string += ' | {}: {:.2f}'.format(name, elapsed_time)
        if torch.distributed.is_initialized():
mohammad's avatar
mohammad committed
271
272
            if torch.distributed.get_rank() == (
                    torch.distributed.get_world_size() - 1):
Mohammad's avatar
Mohammad committed
273
274
275
                print(string, flush=True)
        else:
            print(string, flush=True)