global_vars.py 9.6 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
import time
21
22
from functools import reduce
import operator
Mohammad's avatar
Mohammad committed
23
import torch
Mohammad's avatar
Mohammad committed
24

25
from megatron import dist_signal_handler
26
from megatron.tokenizer import build_tokenizer
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
36
_GLOBAL_MEMORY_BUFFER = None
Mohammad's avatar
Mohammad committed
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
def get_signal_handler():
    _ensure_var_is_initialized(_GLOBAL_SIGNAL_HANDLER, 'signal handler')
    return _GLOBAL_SIGNAL_HANDLER

85
86
87
88
89
90

def get_global_memory_buffer():
    _ensure_var_is_initialized(_GLOBAL_MEMORY_BUFFER, 'global memory buffer')
    return _GLOBAL_MEMORY_BUFFER


91
92
93
94
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
95

96

97

98
def set_global_variables(args):
Mohammad's avatar
Mohammad committed
99
    """Set args, tokenizer, tensorboard-writer, adlr-autoresume, and timers."""
100
101
102
103
104
105

    assert args is not None

    _ensure_var_is_not_initialized(_GLOBAL_ARGS, 'args')
    set_args(args)

mohammad's avatar
mohammad committed
106
    _build_num_microbatches_calculator(args)
107
108
    if args.vocab_file:
        _ = _build_tokenizer(args)
Mohammad's avatar
Mohammad committed
109
110
    _set_tensorboard_writer(args)
    _set_adlr_autoresume(args)
Mohammad's avatar
Mohammad committed
111
    _set_timers()
112
    _set_global_memory_buffer()
Mohammad's avatar
Mohammad committed
113

114
115
    if args.exit_signal_handler:
        _set_signal_handler()
116
117
    

118
119
120
def set_args(args):
    global _GLOBAL_ARGS
    _GLOBAL_ARGS = args
Mohammad's avatar
Mohammad committed
121
122


123
124
def _parse_args(extra_args_provider=None, defaults={},
                ignore_unknown_args=False):
Mohammad's avatar
Mohammad committed
125
126
127
    """Parse entire arguments."""
    global _GLOBAL_ARGS
    _ensure_var_is_not_initialized(_GLOBAL_ARGS, 'args')
128
129
130


    _GLOBAL_ARGS = args
Mohammad's avatar
Mohammad committed
131
    return _GLOBAL_ARGS
Mohammad's avatar
Mohammad committed
132
133


mohammad's avatar
mohammad committed
134
135
136
137
138
139
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
140
141
    _GLOBAL_NUM_MICROBATCHES_CALCULATOR = build_num_microbatches_calculator(
        args)
mohammad's avatar
mohammad committed
142
143


Mohammad's avatar
Mohammad committed
144
def _build_tokenizer(args):
Mohammad's avatar
Mohammad committed
145
146
147
    """Initialize tokenizer."""
    global _GLOBAL_TOKENIZER
    _ensure_var_is_not_initialized(_GLOBAL_TOKENIZER, 'tokenizer')
Mohammad's avatar
Mohammad committed
148
    _GLOBAL_TOKENIZER = build_tokenizer(args)
Mohammad's avatar
Mohammad committed
149
150
151
152
153
154
155
    return _GLOBAL_TOKENIZER


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


Mohammad's avatar
Mohammad committed
158
def _set_tensorboard_writer(args):
Mohammad's avatar
Mohammad committed
159
160
161
162
163
164
    """Set tensorboard writer."""
    global _GLOBAL_TENSORBOARD_WRITER
    _ensure_var_is_not_initialized(_GLOBAL_TENSORBOARD_WRITER,
                                   'tensorboard writer')

    if hasattr(args, 'tensorboard_dir') and \
165
       args.tensorboard_dir and args.rank == (args.world_size - 1):
Mohammad's avatar
Mohammad committed
166
167
168
169
        try:
            from torch.utils.tensorboard import SummaryWriter
            print('> setting tensorboard ...')
            _GLOBAL_TENSORBOARD_WRITER = SummaryWriter(
170
171
                log_dir=args.tensorboard_dir,
                max_queue=args.tensorboard_queue_size)
Mohammad's avatar
Mohammad committed
172
173
174
175
176
177
        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
178
def _set_adlr_autoresume(args):
Mohammad's avatar
Mohammad committed
179
180
181
182
183
184
185
186
187
188
    """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
189
        except BaseException:
Mohammad's avatar
Mohammad committed
190
191
192
193
194
195
196
197
198
199
200
201
            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()

202
203
204
205
206
207
def _set_global_memory_buffer():
    """Initialize global buffer"""
    global _GLOBAL_MEMORY_BUFFER
    _ensure_var_is_not_initialized(_GLOBAL_MEMORY_BUFFER, 'global memory buffer')
    _GLOBAL_MEMORY_BUFFER = GlobalMemoryBuffer()

Mohammad's avatar
Mohammad committed
208
209
210
211
212
213
214
215
216

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
217
218


219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
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
264
265
266
267
268
269
270
271
class Timers:
    """Group of timers."""

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

    def __call__(self, name):
        if name not in self.timers:
272
            self.timers[name] = _Timer(name)
Mohammad's avatar
Mohammad committed
273
274
275
276
277
278
279
280
281
282
        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
283
            writer.add_scalar(name + '-time', value, iteration)
Mohammad's avatar
Mohammad committed
284
285
286
287
288
289
290

    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(
291
                reset=reset) * 1000.0 / normalizer
Mohammad's avatar
Mohammad committed
292
293
            string += ' | {}: {:.2f}'.format(name, elapsed_time)
        if torch.distributed.is_initialized():
mohammad's avatar
mohammad committed
294
295
            if torch.distributed.get_rank() == (
                    torch.distributed.get_world_size() - 1):
Mohammad's avatar
Mohammad committed
296
297
298
                print(string, flush=True)
        else:
            print(string, flush=True)
299
300
301


class GlobalMemoryBuffer:
Vijay Korthikanti's avatar
Vijay Korthikanti committed
302
303
304
    """Global buffer to avoid dynamic memory allocations.
    Caller should ensure that buffers of the same name 
    are not used concurrently."""
305
306
307
308

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

Vijay Korthikanti's avatar
Vijay Korthikanti committed
309
    def get_tensor(self, tensor_shape, dtype, name):
310
        required_len = reduce(operator.mul, tensor_shape, 1)
Vijay Korthikanti's avatar
Vijay Korthikanti committed
311
312
313
314
315
316
317
318
319
        if self.buffer.get((name, dtype), None) is None or \
                self.buffer[(name, dtype)].numel() < required_len:
            self.buffer[(name, dtype)] = \
                torch.empty(required_len,
                            dtype=dtype,
                            device=torch.cuda.current_device(),
                            requires_grad=False)

        return self.buffer[(name, dtype)][0:required_len].view(*tensor_shape)