optimizer.py 53.3 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# coding=utf-8
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# 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 optimizer."""
mohammad's avatar
mohammad committed
17
18
19
20
21
22
23
24
25

from abc import ABC
from abc import abstractmethod

import torch

from apex.multi_tensor_apply import multi_tensor_applier
import amp_C

mohammad's avatar
mohammad committed
26
27
from megatron import get_timers
from megatron import mpu
mohammad's avatar
mohammad committed
28
29
from megatron import print_rank_0

Rewon Child's avatar
Rewon Child committed
30
from .clip_grads import clip_grad_norm_fp32, count_zeros_fp32
mohammad's avatar
mohammad committed
31

32
33
34
# >>>
from lutil import pax, tp
# <<<
mohammad's avatar
mohammad committed
35

Lawrence McAfee's avatar
Lawrence McAfee committed
36

mohammad's avatar
mohammad committed
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def _zero_grad_group_helper(group, set_to_none):
    """Zero out the gradient for a group of parameters.
    Note: copied from torch.optim.optimizer."""
    for param in group:
        if param.grad is not None:
            if set_to_none:
                param.grad = None
            else:
                if param.grad.grad_fn is not None:
                    param.grad.detach_()
                else:
                    param.grad.requires_grad_(False)
                param.grad.zero_()


52
def _multi_tensor_copy_this_to_that(this, that, overflow_buf=None):
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
53
54
55
56
    """Use multi-tensor-applier to copy values from one list to another.
    We don't have a blfoat16 implementation so for now if the overflow_buf
    is not provided, we default back to simple loop copy to be compatible
    with bfloat16."""
57
58
    if overflow_buf:
        overflow_buf.fill_(0)
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
59
60
61
62
63
        # Scaling with factor `1.0` is equivalent to copy.
        multi_tensor_applier(amp_C.multi_tensor_scale,
                             overflow_buf,
                             [this, that],
                             1.0)
64
    else:
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
65
66
67
        for this_, that_ in zip(this, that):
            that_.copy_(this_)

68

mohammad's avatar
mohammad committed
69
70
71

class MegatronOptimizer(ABC):

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
72
73
74

    def __init__(self, optimizer, clip_grad,
                 log_num_zeros_in_grad,
75
                 params_have_main_grad,
76
                 use_contiguous_buffers_in_local_ddp):
77

mohammad's avatar
mohammad committed
78
79
80
        """Input optimizer is the base optimizer for example Adam."""
        self.optimizer = optimizer
        assert self.optimizer, 'no optimizer is provided.'
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
81
82
83
84
        # Set gradient clipping and logging params.
        self.clip_grad = clip_grad
        self.log_num_zeros_in_grad = log_num_zeros_in_grad
        self.params_have_main_grad = params_have_main_grad
85
        self.use_contiguous_buffers_in_local_ddp = use_contiguous_buffers_in_local_ddp
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
86

87
        if self.use_contiguous_buffers_in_local_ddp:
88
89
            assert self.params_have_main_grad, \
                "use of contiguous buffer requires that params have main grad"
mohammad's avatar
mohammad committed
90

Rewon Child's avatar
Rewon Child committed
91
    def get_parameters(self):
92
93
94
95
        params = []
        for param_group in self.optimizer.param_groups:
            for param in param_group['params']:
                params.append(param)
Rewon Child's avatar
Rewon Child committed
96
97
        return params

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
98

Rewon Child's avatar
Rewon Child committed
99
100
    def clip_grad_norm(self, clip_grad):
        params = self.get_parameters()
Lawrence McAfee's avatar
Lawrence McAfee committed
101
102
103
104
105
106
107
        # >>>
        # pax(0, {
        #     "clip_grad" : clip_grad,
        #     "params": [ (p.tensor_model_parallel, tp(p)) for p in params ],
        #     "grads" : [ p.grad for p in params ],
        # })
        # <<<
108
        return clip_grad_norm_fp32(params, clip_grad)
109

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
110

Rewon Child's avatar
Rewon Child committed
111
112
113
114
    def count_zeros(self):
        params = self.get_parameters()
        return count_zeros_fp32(params)

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
115

mohammad's avatar
mohammad committed
116
117
118
119
    @abstractmethod
    def zero_grad(self, set_to_none=True):
        pass

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
120

mohammad's avatar
mohammad committed
121
122
    @abstractmethod
    def get_loss_scale(self):
123
        """The output should be a cuda tensor of size 1."""
mohammad's avatar
mohammad committed
124
125
        pass

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
126

mohammad's avatar
mohammad committed
127
128
129
130
    def scale_loss(self, loss):
        """Simple scaling."""
        return self.get_loss_scale() * loss

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
131

Lawrence McAfee's avatar
Lawrence McAfee committed
132
133
134
135
136
    @abstractmethod
    def reduce_gradients(self):
        pass


mohammad's avatar
mohammad committed
137
138
139
140
    @abstractmethod
    def step(self):
        pass

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
141

Lawrence McAfee's avatar
Lawrence McAfee committed
142
143
144
145
146
    @abstractmethod
    def gather_params(self):
        pass


147
148
    @abstractmethod
    def reload_model_params(self):
149
150
151
152
153
        """Refreshes any internal state from the current model parameters.
        Call whenever the parameters are changed outside of the optimizer.
        For example, when we load a model from a checkpoint  without loading
        the optimizer, the model parameters are updated but for fp16 optimizer
        with main parameters, the main parameters need to also be updated."""
154
155
        pass

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
156

mohammad's avatar
mohammad committed
157
158
159
160
    @abstractmethod
    def state_dict(self):
        pass

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
161

mohammad's avatar
mohammad committed
162
163
164
165
    @abstractmethod
    def load_state_dict(self, state_dict):
        pass

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
166

mohammad's avatar
mohammad committed
167
168
169
170
171
172
173
174
175
176
    # Promote state so it can be retrieved or set via
    # "optimizer_instance.state"
    def _get_state(self):
        return self.optimizer.state

    def _set_state(self, value):
        self.optimizer.state = value

    state = property(_get_state, _set_state)

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
177

mohammad's avatar
mohammad committed
178
179
180
181
182
183
184
185
186
187
188
189
    # Promote param_groups so it can be retrieved or set via
    # "optimizer_instance.param_groups"
    # (for example, to adjust the learning rate)
    def _get_param_groups(self):
        return self.optimizer.param_groups

    def _set_param_groups(self, value):
        self.optimizer.param_groups = value

    param_groups = property(_get_param_groups, _set_param_groups)


Lawrence McAfee's avatar
Lawrence McAfee committed
190
class BaseFloat16Optimizer(MegatronOptimizer):
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
191
192

    def __init__(self, optimizer, clip_grad, log_num_zeros_in_grad,
193
                 params_have_main_grad, use_contiguous_buffers_in_local_ddp,
194
195
                 bf16, grad_scaler,
                 models):
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
196

Lawrence McAfee's avatar
Lawrence McAfee committed
197
        super().__init__(
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
198
            optimizer, clip_grad, log_num_zeros_in_grad,
199
            params_have_main_grad, use_contiguous_buffers_in_local_ddp)
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
200

201
202
203
        # >>>
        self.models = models
        # <<<
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
204
        self.bf16 = bf16
mohammad's avatar
mohammad committed
205
        self.grad_scaler = grad_scaler
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
206
207
208
        # None grad scaler is only supported for bf16.
        if self.grad_scaler is None:
            assert self.bf16, 'fp16 expects a grad scaler.'
mohammad's avatar
mohammad committed
209
210
211

        # Tensor used to determine if a nan/if has happend.
        # Any non-zero value indicates inf/nan.
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
212
213
214
215
        # Note that we keep this for the cases that grad scaler is none.
        # We still record nan/inf if we have a bfloat16 with a grad scaler.
        if self.grad_scaler:
            self.found_inf = torch.cuda.FloatTensor([0.0])
mohammad's avatar
mohammad committed
216
217

        # Dummy tensor needed for apex multi-apply tensor.
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
218
219
220
221
222
223
224
225
226
227
        # For bfloat, we don't have multi-tensor apply and for now
        # we set it to none so the multi-tensor apply gets ignored.
        if bf16:
            self._dummy_overflow_buf = None
        else:
            self._dummy_overflow_buf = torch.cuda.IntTensor([0])

        # In case grad scaler is not passed, define the unity scale.
        if self.grad_scaler is None:
            self._scale_one = torch.cuda.FloatTensor([1.0])
mohammad's avatar
mohammad committed
228

Lawrence McAfee's avatar
Lawrence McAfee committed
229
230
231
232
233
234
235

    def get_loss_scale(self):
        if self.grad_scaler is None:
            return self._scale_one
        return self.grad_scaler.scale


Lawrence McAfee's avatar
Lawrence McAfee committed
236
237
238
239
    def reload_model_params(self):
        self._copy_model_params_to_main_params()


Lawrence McAfee's avatar
Lawrence McAfee committed
240
241
242
243
244
245
246
247
248
249
250
251
252
253
    def _unscale_main_grads_and_check_for_nan(self):

        # Collect main grads.
        main_grads = self._collect_main_grad_data_for_unscaling()
        # pax(1, {"main_grads": main_grads})

        # Reset found inf.
        self.found_inf.fill_(0.0)

        # Unscale and set found inf/nan
        torch._amp_foreach_non_finite_check_and_unscale_(
            main_grads, self.found_inf, self.grad_scaler.inv_scale)

        # Update across all model parallel instances.
254
255
256
257
258
        # >>>
        # torch.distributed.all_reduce(self.found_inf,
        #                              op=torch.distributed.ReduceOp.MAX,
        #                              group=mpu.get_model_parallel_group())
        # +++
Lawrence McAfee's avatar
Lawrence McAfee committed
259
        torch.distributed.all_reduce(self.found_inf,
260
261
                                     op=torch.distributed.ReduceOp.MAX)
        # <<<
Lawrence McAfee's avatar
Lawrence McAfee committed
262
263
264
265
266
267
268
269

        # Check for nan.
        found_inf_flag = (self.found_inf.item() > 0)

        return found_inf_flag


    @torch.no_grad()
270
    def step(self, ITERATION):
Lawrence McAfee's avatar
Lawrence McAfee committed
271
272
273
274
275

        timers = get_timers()

        # Copy gradients from model params to main params.
        timers('optimizer-copy-to-main-grad').start()
276
        self._copy_model_grads_to_main_grads(ITERATION)
Lawrence McAfee's avatar
Lawrence McAfee committed
277
278
        timers('optimizer-copy-to-main-grad').stop()

Lawrence McAfee's avatar
Lawrence McAfee committed
279
280
281
282
283
        # pax(0, {
        #     "params" : self.get_parameters(), # self.main_param_shards,
        #     "grads" : [ p.grad for p in self.get_parameters() ], # self.main_param_shards ],
        # })

Lawrence McAfee's avatar
Lawrence McAfee committed
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
        # Do unscale, check for inf, and update grad scaler only for
        # the case that grad scaler is provided.
        if self.grad_scaler:

            # Unscale and check for inf/nan.
            timers('optimizer-unscale-and-check-inf').start()
            found_inf_flag = self._unscale_main_grads_and_check_for_nan()
            timers('optimizer-unscale-and-check-inf').stop()

            # We are done with scaling gradients
            # so we can update the loss scale.
            self.grad_scaler.update(found_inf_flag)

            # If we found inf/nan, skip the update.
            if found_inf_flag:
299
300
301
302
303
                pax(0, {
                    "main params" : self.get_main_params(),
                    "main grads" : self.get_main_grads(),
                    "found_inf_flag" : found_inf_flag,
                })
Lawrence McAfee's avatar
Lawrence McAfee committed
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
                return False, None, None

        # Clip the main gradients.
        timers('optimizer-clip-main-grad').start()
        grad_norm = None
        if self.clip_grad > 0.0:
            grad_norm = self.clip_grad_norm(self.clip_grad)
        timers('optimizer-clip-main-grad').stop()

        # count the zeros in the grads
        num_zeros_in_grad = self.count_zeros() if \
                            self.log_num_zeros_in_grad else None

        # Step the optimizer.
        self.optimizer.step()

        # >>>
        # pax(0, {
322
323
        #     "main params" : self.get_main_params(),
        #     "main grads" : self.get_main_grads(),
Lawrence McAfee's avatar
Lawrence McAfee committed
324
325
326
327
328
        # })
        # <<<

        # Update params from main params.
        timers('optimizer-copy-main-to-model-params').start()
329
        self._copy_main_params_to_model_params(ITERATION)
Lawrence McAfee's avatar
Lawrence McAfee committed
330
331
        timers('optimizer-copy-main-to-model-params').stop()

332
333
334
335
336
337
338
        # >>>
        # pax(1, {
        #     "ITERATION" : ITERATION,
        #     "model_params" : [ p for m in self.models for p in m.parameters() ],
        # })
        # <<<

Lawrence McAfee's avatar
Lawrence McAfee committed
339
340
341
342
        # Successful update.
        return True, grad_norm, num_zeros_in_grad


Lawrence McAfee's avatar
Lawrence McAfee committed
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
# class Float16OptimizerWithFloat16Params(MegatronOptimizer):
class Float16OptimizerWithFloat16Params(BaseFloat16Optimizer):
    """Float16 optimizer for fp16 and bf16 data types.

    Arguments:
        optimizer: base optimizer such as Adam or SGD
        clip_grad: clip gradeints with this global L2 norm. Note
            that clipping is ignored if clip_grad == 0
        log_num_zeros_in_grad: return number of zeros in the gradients.
        params_have_main_grad: flag indicating if parameters have
            a `main_grad` field. If this is set, we are assuming
            that the model parameters are store in the `main_grad`
            field instead of the typical `grad` field. This happens
            for the DDP cases where there is a continuous buffer
            holding the gradients. For example for bfloat16, we want
            to do gradient accumulation and all-reduces in float32
            and as a result we store those gradients in the main_grad.
            Note that main grad is not necessarily in float32.
        bf16: if true, the model is running in bfloat16.
        grad_scaler: used for scaling gradients. Note that this can be
            None. This case happens when `bf16 = True` and we don't
            use any loss scale. Note that for `bf16 = True`, we can have
            a constnat gradient scaler. Also for `bf16 = False`, we
            always require a grad scaler.
    """

    def __init__(self, optimizer, clip_grad, log_num_zeros_in_grad,
                 params_have_main_grad, use_contiguous_buffers_in_local_ddp,
Lawrence McAfee's avatar
Lawrence McAfee committed
371
                 bf16, grad_scaler, models):
Lawrence McAfee's avatar
Lawrence McAfee committed
372
373
374
375

        super().__init__(
            optimizer, clip_grad, log_num_zeros_in_grad,
            params_have_main_grad, use_contiguous_buffers_in_local_ddp,
Lawrence McAfee's avatar
Lawrence McAfee committed
376
            bf16, grad_scaler, models)
Lawrence McAfee's avatar
Lawrence McAfee committed
377

mohammad's avatar
mohammad committed
378
        # ======================
379
        # main parameter stuff
mohammad's avatar
mohammad committed
380
381
382
        # ======================

        # Three groups of parameters:
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
383
384
        #   float16_groups: original float16 parameters
        #   fp32_from_float16_groups: fp32 copy of float16 parameters
mohammad's avatar
mohammad committed
385
        #   fp32_from_fp32_groups: original fp32 parameters
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
386
387
        self.float16_groups = []
        self.fp32_from_float16_groups = []
mohammad's avatar
mohammad committed
388
389
390
391
        self.fp32_from_fp32_groups = []

        # For all the groups in the original optimizer:
        for param_group in self.optimizer.param_groups:
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
392
            float16_params_this_group = []
mohammad's avatar
mohammad committed
393
            fp32_params_this_group = []
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
394
            fp32_from_float16_params_this_group = []
mohammad's avatar
mohammad committed
395
396
397
398
            # For all the parameters in this group:
            for i, param in enumerate(param_group['params']):
                if param.requires_grad:

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
399
400
401
402
                    # float16 params:
                    if param.type() in ['torch.cuda.HalfTensor',
                                        'torch.cuda.BFloat16Tensor']:
                        float16_params_this_group.append(param)
mohammad's avatar
mohammad committed
403
                        # Create a copy
404
                        main_param = param.detach().clone().float()
mohammad's avatar
mohammad committed
405
                        # Copy tensor model parallel attributes.
406
                        mpu.copy_tensor_model_parallel_attributes(main_param,
mohammad's avatar
mohammad committed
407
                                                                  param)
408
                        if hasattr(param, 'shared'):
409
                            main_param.shared = param.shared
mohammad's avatar
mohammad committed
410
                        # Replace the optimizer params with the new fp32 copy.
411
                        param_group['params'][i] = main_param
412

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
413
                        fp32_from_float16_params_this_group.append(main_param)
414
                        # Reset existing state dict key to the new main param.
mohammad's avatar
mohammad committed
415
                        if param in self.optimizer.state:
416
                            self.optimizer.state[main_param] \
mohammad's avatar
mohammad committed
417
418
419
420
                                = self.optimizer.state.pop(param)

                    # fp32 params.
                    elif param.type() == 'torch.cuda.FloatTensor':
Lawrence McAfee's avatar
Lawrence McAfee committed
421
422
423
                        # >>>
                        pax(0, {"param": param})
                        # <<<
mohammad's avatar
mohammad committed
424
425
426
427
                        fp32_params_this_group.append(param)
                        param_group['params'][i] = param

                    else:
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
428
429
430
431
432
433
434
435
436
                        raise TypeError('Wrapped parameters must be one of '
                                        'torch.cuda.FloatTensor,  '
                                        'torch.cuda.HalfTensor, or '
                                        'torch.cuda.BFloat16Tensor. '
                                        'Received {}'.format(param.type()))

            self.float16_groups.append(float16_params_this_group)
            self.fp32_from_float16_groups.append(
                fp32_from_float16_params_this_group)
mohammad's avatar
mohammad committed
437
438
439
440
441
442
            self.fp32_from_fp32_groups.append(fp32_params_this_group)

        # Leverage state_dict() and load_state_dict() to
        # recast preexisting per-param state tensors
        self.optimizer.load_state_dict(self.optimizer.state_dict())

Lawrence McAfee's avatar
Lawrence McAfee committed
443
444
445
446
447
448
449
450
451
452
        # >>>
        # from megatron.mpu.layers import param_is_not_tensor_parallel_duplicate
        # params = self.get_parameters()
        # pax(0, {
        #     # "params / 0" : params[0],
        #     "params" : [ (p.tensor_model_parallel, tp(p)) for p in params ],
        #     "grads" : [ (param_is_not_tensor_parallel_duplicate(p.grad), tp(p.grad)) for p in params ],
        # })
        # <<<

mohammad's avatar
mohammad committed
453
454
455

    def zero_grad(self, set_to_none=True):
        """We only need to zero the model related parameters, i.e.,
456
457
458
459
        float16_groups & fp32_from_fp32_groups. We additionally zero
        fp32_from_float16_groups as a memory optimization to reduce
        fragmentation; in the case of set_to_none==True, the space
        used by this field can be safely deallocated at this point."""
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
460
        for group in self.float16_groups:
mohammad's avatar
mohammad committed
461
            _zero_grad_group_helper(group, set_to_none)
462
463
        for group in self.fp32_from_float16_groups:
            _zero_grad_group_helper(group, set_to_none)
mohammad's avatar
mohammad committed
464
465
466
467
        for group in self.fp32_from_fp32_groups:
            _zero_grad_group_helper(group, set_to_none)


468
    # >>>
469
470
471
472
473
474
475
476
477
478
479
480
481
482
    def reduce_gradients(self, model):

        # >>>
        from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP

        from megatron import get_args
        from megatron import get_timers
        from megatron.model import DistributedDataParallel as LocalDDP
        from megatron.model import Float16Module
        from megatron.utils import unwrap_model

        args = get_args()
        timers = get_timers()
        # <<<
483
484
485
486
487
488
489
490
491
492
493
494
495
496
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

        # >>>
        # if not args.use_distributed_optimizer:

        # All-reduce if needed.
        # >>>
        # if args.DDP_impl == 'local' and not args.use_distributed_optimizer:
        if args.DDP_impl == 'local':
        # <<<
            timers('backward-params-all-reduce').start()
            for model_module in model:
                # >>>
                # from lutil import pax, tp
                # pax(0, {
                #     "model" : model,
                #     "model_module" : model_module,
                # })
                # <<<
                # >>>
                # e.g., grad_shard = optimizer.get_grad_shard()
                # <<<
                model_module.allreduce_gradients()
            timers('backward-params-all-reduce').stop()

        # All-reduce word_embeddings' grad across first and last stages to ensure
        # that word_embeddings parameters stay in sync.
        # This should only run for models that support pipelined model parallelism
        # (BERT and GPT-2).
        timers('backward-embedding-all-reduce').start()
        if mpu.is_rank_in_embedding_group(ignore_virtual=True) and \
                mpu.get_pipeline_model_parallel_world_size() > 1:
            if mpu.is_pipeline_first_stage(ignore_virtual=True):
                unwrapped_model = model[0]
            elif mpu.is_pipeline_last_stage(ignore_virtual=True):
                unwrapped_model = model[-1]
            else:  # We do not support the interleaved schedule for T5 yet.
                unwrapped_model = model[0]
            unwrapped_model = unwrap_model(
                unwrapped_model, (torchDDP, LocalDDP, Float16Module))

            if unwrapped_model.share_word_embeddings:
                word_embeddings_weight = unwrapped_model.word_embeddings_weight()
                # >>>
526
527
528
529
530
                if args.DDP_impl == 'local':
                    grad = word_embeddings_weight.main_grad
                else:
                    grad = word_embeddings_weight.grad
                torch.distributed.all_reduce(grad, group=mpu.get_embedding_group())
531
                # +++
532
533
534
                # grad_shard = optimizer.get_grad_shard(word_embeddings)
                # torch.distributed.all_reduce(grad_shard,
                #                              group=mpu.get_embedding_group())
535
536
537
538
539
540
541
542
543
544
545
546
547
548
                # <<<

        # All-reduce position_embeddings grad across first (encoder) and split (decoder) 
        # stages to ensure that position embeddings parameters stay in sync.
        # This should only run for T5 models with pipeline parallelism
        if mpu.is_rank_in_position_embedding_group() and \
                mpu.get_pipeline_model_parallel_world_size() > 1 and \
                args.pipeline_model_parallel_split_rank is not None:
            unwrapped_model = model[0]
            unwrapped_model = unwrap_model(
                unwrapped_model, (torchDDP, LocalDDP, Float16Module))
            assert args.DDP_impl == 'local', \
                'T5 model is only supported with local DDP mode'
            # >>>
549
550
            grad = unwrapped_model.language_model.embedding.position_embeddings.weight.main_grad
            torch.distributed.all_reduce(grad, group=mpu.get_position_embedding_group())
551
            # +++
552
553
554
555
            # grad_shard = optimizer.get_grad_shard(
            #     unwrapped_model.language_model.embedding.position_embeddings.weight)
            # torch.distributed.all_reduce(grad_shard,
            #                              group=mpu.get_position_embedding_group())
556
557
558
            # <<<
        timers('backward-embedding-all-reduce').stop()

Lawrence McAfee's avatar
Lawrence McAfee committed
559
    def gather_params(self):
Lawrence McAfee's avatar
Lawrence McAfee committed
560
        pass
Lawrence McAfee's avatar
Lawrence McAfee committed
561

562
    def _copy_model_grads_to_main_grads(self):
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
563
564
565
        # This only needs to be done for the float16 group.
        for model_group, main_group in zip(self.float16_groups,
                                           self.fp32_from_float16_groups):
566
            for model_param, main_param in zip(model_group, main_group):
567
                if self.params_have_main_grad and hasattr(model_param, 'main_grad'):
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
568
569
570
571
                    main_param.grad = model_param.main_grad.float()
                else:
                    if model_param.grad is not None:
                        main_param.grad = model_param.grad.float()
572
573
574
575
576

                # Safe to deallocate model's grad/main_grad after copying.
                # (If using contiguous buffers, main_grad's memory should
                # persist and therefore should not be deallocated.)
                model_param.grad = None
577
                if self.params_have_main_grad and \
578
                   not self.use_contiguous_buffers_in_local_ddp:
579
580
                    model_param.main_grad = None

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
581
582
583
584
585
        # For fp32 grads, we need to reset the grads to main grad.
        if self.params_have_main_grad:
            for model_group in self.fp32_from_fp32_groups:
                for model_param in model_group:
                    model_param.grad = model_param.main_grad
mohammad's avatar
mohammad committed
586

587
588
589
                    # Safe to de-reference model's main_grad after copying.
                    # (If using contiguous buffers, main_grad's memory should
                    # persist and therefore should not be deallocated.)
590
                    if not self.use_contiguous_buffers_in_local_ddp:
591
                        model_param.main_grad = None
mohammad's avatar
mohammad committed
592

Lawrence McAfee's avatar
Lawrence McAfee committed
593
594
    def _collect_main_grad_data_for_unscaling(self):

595
        main_grads = []
Lawrence McAfee's avatar
Lawrence McAfee committed
596
597

        # fp32 params from float16 ones.
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
598
        for main_group in self.fp32_from_float16_groups:
599
600
601
            for main_param in main_group:
                if main_param.grad is not None:
                    main_grads.append(main_param.grad.data)
Lawrence McAfee's avatar
Lawrence McAfee committed
602
603
604

        # pax(1, {"main_grads": main_grads})

mohammad's avatar
mohammad committed
605
        # Append fp32 parameters.
606
607
608
609
        for main_group in self.fp32_from_fp32_groups:
            for main_param in main_group:
                if main_param.grad is not None:
                    main_grads.append(main_param.grad.data)
Lawrence McAfee's avatar
Lawrence McAfee committed
610
611
612
613
614
        
        # >>>
        # from megatron.mpu.layers import param_is_not_tensor_parallel_duplicate
        # pax(1, {"main_grads": [ (param_is_not_tensor_parallel_duplicate(t), tp(t)) for t in main_grads ]})
        # <<<
mohammad's avatar
mohammad committed
615

Lawrence McAfee's avatar
Lawrence McAfee committed
616
        return main_grads
mohammad's avatar
mohammad committed
617
618


Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
619
    def _get_model_and_main_params_data_float16(self):
mohammad's avatar
mohammad committed
620
        model_data = []
621
        main_data = []
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
622
623
        for model_group, main_group in zip(self.float16_groups,
                                           self.fp32_from_float16_groups):
624
            for model_param, main_param in zip(model_group, main_group):
mohammad's avatar
mohammad committed
625
                model_data.append(model_param.data)
626
627
                main_data.append(main_param.data)
        return model_data, main_data
628
629


630
    def _copy_main_params_to_model_params(self):
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
631
632
        # Only needed for the float16 params.
        model_data, main_data = self._get_model_and_main_params_data_float16()
633
634
635
636
637
        _multi_tensor_copy_this_to_that(this=main_data, that=model_data,
                                        overflow_buf=self._dummy_overflow_buf)


    def _copy_model_params_to_main_params(self):
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
638
639
        # Only needed for the float16 params.
        model_data, main_data = self._get_model_and_main_params_data_float16()
640
641
        _multi_tensor_copy_this_to_that(this=model_data, that=main_data,
                                        overflow_buf=self._dummy_overflow_buf)
642
643


mohammad's avatar
mohammad committed
644
645
646
    def state_dict(self):
        state_dict = {}
        state_dict['optimizer'] = self.optimizer.state_dict()
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
647
648
649
        if self.grad_scaler:
            state_dict['grad_scaler'] = self.grad_scaler.state_dict()
        state_dict['fp32_from_fp16_params'] = self.fp32_from_float16_groups
mohammad's avatar
mohammad committed
650
651
652
653
        return state_dict


    def load_state_dict(self, state_dict):
mohammad's avatar
mohammad committed
654
655
656
657
658
659
660
661
662
663
664
665
666
        # Optimizer.
        optimizer_key = 'optimizer'
        if optimizer_key not in state_dict:
            optimizer_key = 'optimizer_state_dict'
            print_rank_0('***WARNING*** loading optimizer from '
                         'an old checkpoint ...')
        self.optimizer.load_state_dict(state_dict[optimizer_key])

        # Grad scaler.
        if 'grad_scaler' not in state_dict:
            print_rank_0('***WARNING*** found an old checkpoint, will not '
                         'load grad scaler ...')
        else:
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
667
668
669
670
671
672
            if self.grad_scaler:
                self.grad_scaler.load_state_dict(state_dict['grad_scaler'])
            else:
                print_rank_0('***WARNING*** fould the grad scaler in the '
                             'checkpoint but it is None in the class. '
                             'Skipping loading grad scaler ...')
mohammad's avatar
mohammad committed
673

674
        # Copy data for the main params.
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
675
676
677
        fp32_from_float16_params_key = 'fp32_from_fp16_params'
        if fp32_from_float16_params_key not in state_dict:
            fp32_from_float16_params_key = 'fp32_from_fp16'
mohammad's avatar
mohammad committed
678
        for current_group, saved_group in zip(
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
679
680
                self.fp32_from_float16_groups,
                state_dict[fp32_from_float16_params_key]):
mohammad's avatar
mohammad committed
681
682
683
684
            for current_param, saved_param in zip(current_group, saved_group):
                current_param.data.copy_(saved_param.data)


685
# >>>
686
import math
687

688
from megatron import get_args
689

Lawrence McAfee's avatar
Lawrence McAfee committed
690
691
692
693
694
695
696
697
698
699
700
# class ShardIndex:
class Shard:
    def __init__(self, start, end):
        self.start = start
        self.end = end
        self.size = end - start
    def normalize(self, start = 0):
        return Shard(start, start + self.size)
    def __str__(self):
        return "%d,%d [%d]" % (self.start, self.end, self.size)

701
# class Float16DistributedOptimizer(Float16OptimizerWithFloat16Params):
Lawrence McAfee's avatar
Lawrence McAfee committed
702
703
# class Float16DistributedOptimizer(MegatronOptimizer):
class Float16DistributedOptimizer(BaseFloat16Optimizer):
704

705
    # >>>
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
    # @classmethod
    # def test_reduce_scatter(cls):

    #     torch.manual_seed(mpu.get_data_parallel_rank())
    #     size = (20,)
    #     dtype = torch.float
    #     device = torch.cuda.current_device()
    #     data_parallel_world_size = mpu.get_data_parallel_world_size()
    #     data_parallel_group = mpu.get_data_parallel_group()

    #     input_list = [
    #         # torch.randn(size, dtype = dtype, device = device)
    #         5 * torch.randint(low = 1, high = 3, size = size, dtype = dtype, device = device)
    #         for _ in range(data_parallel_world_size)
    #     ]
    #     output = torch.empty(size, dtype = dtype, device = device)

    #     torch.distributed.reduce_scatter(
    #         output,
    #         input_list,
    #         group = data_parallel_group,
    #     )

    #     if torch.distributed.get_rank() == 0:
    #         print(output)
    #     pax(0, {
    #         "data_parallel_world_size" : data_parallel_world_size,
    #         "data_parallel_group" : data_parallel_group,
    #         "input_list" : input_list,
    #         "output" : tp(output),
    #     })
737
738
    # <<<

739
    @classmethod
Lawrence McAfee's avatar
Lawrence McAfee committed
740
    def get_model_gbuf_param_shard_map(cls, model, dtype, gbuf_world_shard):
741

Lawrence McAfee's avatar
Lawrence McAfee committed
742
743
        # Param shard map.
        param_world_index_map = model._grad_buffer_param_index_map[dtype]
744
        param_shard_map = {}
Lawrence McAfee's avatar
Lawrence McAfee committed
745
        for param, param_world_indexes in param_world_index_map.items():
746

Lawrence McAfee's avatar
Lawrence McAfee committed
747
748
749
            # Shard range.
            param_world_start, param_world_end = param_world_indexes
            param_local_start = max(
750
                0,
Lawrence McAfee's avatar
Lawrence McAfee committed
751
752
753
754
755
756
757
758
759
                param_world_start - gbuf_world_shard.start)
            param_local_end = min(
                gbuf_world_shard.size,
                param_world_end - gbuf_world_shard.start)

            # Add shard, if within range.
            if param_local_end > param_local_start:
                param_local_shard = Shard(param_local_start, param_local_end)
                param_world_shard = param_local_shard.normalize(param_world_start)
760
761
                sub_param_start = max(0, gbuf_world_shard.start-param_world_start)
                sub_param_shard = param_local_shard.normalize(sub_param_start)
Lawrence McAfee's avatar
Lawrence McAfee committed
762
                param_shard_map[param] = {
763
764
765
                    "gbuf_world" : param_world_shard,
                    "gbuf_local" : param_local_shard,
                    "param" : sub_param_shard,
766
                }
767
                # >>>
768
769
                # if param_world_start < gbuf_world_shard.start:
                #     pax({"param shards": param_shard_map[param]})
770
                # <<<
771

Lawrence McAfee's avatar
Lawrence McAfee committed
772
        # pax(0, {"param_shard_map": [ str((str(p.shape), s)) for p,s in param_shard_map.items() ]})
773
774
775
776

        return param_shard_map

    @classmethod
Lawrence McAfee's avatar
Lawrence McAfee committed
777
    def get_model_gbuf_shard(cls, model, dtype):
778

Lawrence McAfee's avatar
Lawrence McAfee committed
779
780
        data_parallel_rank = mpu.get_data_parallel_rank()
        data_parallel_world_size = mpu.get_data_parallel_world_size()
781
782

        # Grad buffer shard.
Lawrence McAfee's avatar
Lawrence McAfee committed
783
784
785
786
        grad_buffer = model._grad_buffers[dtype]
        gbuf_size = grad_buffer.numel
        max_gbuf_shard_size = int(math.ceil(gbuf_size / data_parallel_world_size))

787
788
789
790
791
792
793
        gbuf_world_all_shards = []
        for r in range(data_parallel_world_size):
            gbuf_world_start = r * max_gbuf_shard_size
            gbuf_world_end = min(gbuf_size, gbuf_world_start+max_gbuf_shard_size)
            gbuf_world_shard = Shard(gbuf_world_start, gbuf_world_end)
            gbuf_world_all_shards.append(gbuf_world_shard)
        gbuf_world_shard = gbuf_world_all_shards[data_parallel_rank]
Lawrence McAfee's avatar
Lawrence McAfee committed
794
795
796
797
798
799
800
801
802
803
804
        gbuf_local_shard = gbuf_world_shard.normalize()

        # Param shards.
        param_shard_map = cls.get_model_gbuf_param_shard_map(model,
                                                             dtype,
                                                             gbuf_world_shard)

        # Altogether.
        data = {
            "local" : gbuf_local_shard,
            "world" : gbuf_world_shard,
805
            "world_all" : gbuf_world_all_shards,
Lawrence McAfee's avatar
Lawrence McAfee committed
806
            "param_map" : param_shard_map,
807
808
        }

Lawrence McAfee's avatar
Lawrence McAfee committed
809
        # pax(0, {"data": data})
810

Lawrence McAfee's avatar
Lawrence McAfee committed
811
        return data
812
813

    @classmethod
Lawrence McAfee's avatar
Lawrence McAfee committed
814
    def get_model_gbuf_shard_map(cls, model):
815
        return {
Lawrence McAfee's avatar
Lawrence McAfee committed
816
            dtype : cls.get_model_gbuf_shard(model, dtype)
817
818
819
            for dtype in model._grad_buffers
        }

Lawrence McAfee's avatar
Lawrence McAfee committed
820
821
    @classmethod
    def get_param_gbuf_map(cls, model_gbuf_shards):
822

Lawrence McAfee's avatar
Lawrence McAfee committed
823
824
825
826
827
828
829
830
831
832
833
834
835
        param_gbuf_map = {}
        for model_index, model_gbuf_shard_map in enumerate(model_gbuf_shards):
            for dtype, gbuf_shard_map in model_gbuf_shard_map.items():
                for param, param_shard_map in gbuf_shard_map["param_map"].items():
                    # assert param not in param_size_map
                    # param_size_map[param] = param_shard_map["local"].size
                    param_gbuf_map[param] = (model_index, dtype)
                    # pax(0, {
                    #     "dtype" : dtype,
                    #     "gbuf_shard_map" : gbuf_shard_map,
                    #     "param" : tp(param),
                    #     "param_shard_map" : param_shard_map,
                    # })
836

Lawrence McAfee's avatar
Lawrence McAfee committed
837
838
839
840
841
842
        # pax(0, {
        #     "model_gbuf_shards" : model_gbuf_shards,
        #     # "param_size_map" :
        #     # [ (str(p.shape), s) for p, s in param_size_map.items() ],
        #     "param_gbuf_map" : param_gbuf_map,
        # })
843

Lawrence McAfee's avatar
Lawrence McAfee committed
844
        return param_gbuf_map
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865

    @classmethod
    def get_optimizer_group_shards(cls, param_groups, model_gbuf_shards):

        num_groups = len(param_groups)

        # Param group map.
        param_group_map = {}
        for group_index, group in enumerate(param_groups):
            for param in group["params"]:
                assert param.requires_grad
                param_group_map[param] = group_index

        # Optimizer group shards.
        group_shards = [ {"size": 0, "param_map": {}} for _ in param_groups ]
        for model_gbuf_shard_map in model_gbuf_shards:
            for dtype, gbuf_shard_map in model_gbuf_shard_map.items():
                for param in gbuf_shard_map["param_map"]:
                    
                    group_index = param_group_map[param]
                    group_shard = group_shards[group_index]
866
                    param_size = gbuf_shard_map["param_map"][param]["param"].size
867
868
869
870
871
872
873
874

                    param_group_start = group_shard["size"]
                    param_group_end = param_group_start + param_size
                    param_group_shard = Shard(param_group_start, param_group_end)

                    group_shard["size"] += param_size
                    group_shard["param_map"][param] = param_group_shard

875
876
877
878
879
880
881
882
883
884
                    # >>>
                    # if torch.distributed.get_rank() == 1:
                    #     print(">>> [%d] ... group %d, size %d, param %s. <<<" % (
                    #         torch.distributed.get_rank(),
                    #         group_index,
                    #         param_size,
                    #         str(tuple(param.shape)),
                    #     ))
                    # <<<

885
886
887
888
889
890
        # Squeeze zero-size group shards.
        for group_index, group_shard in enumerate(group_shards):
            group_shard["orig_group"] = param_groups[group_index]
        group_shards = [ g for g in group_shards if g["size"] > 0 ]

        # pax(0, {
891
892
893
894
895
896
        #     "param_group_map": [
        #         (g, str(p.shape))
        #         for p, g in param_group_map.items()
        #     ],
        #     "group_shards" : group_shards,
        # })
897
898
899

        return group_shards

900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
    @classmethod
    def allocate_main_param_shards(cls, opt_group_shards):

        # Allocate main param/grad shard.
        # ** torch.nn.Parameter ??
        # ** MemoryBuffer ??
        allocate_shard = lambda shard_size, dtype : torch.empty(
            (shard_size,),
            dtype = dtype,
            device = torch.cuda.current_device(),
            requires_grad = True)
        
        # main_param_shards = []
        for group_index, group_shard in enumerate(opt_group_shards):

            group_size = group_shard["size"]
            assert group_size != 0, "temporary check ... remove me."

            # ** todo: for dtype in model_main_dtypes ........ **

            # Allocate shard.
            # if group_size == 0:
            #     main_param = None
            # else:
            main_param = allocate_shard(group_size, torch.float)
            main_param.grad = allocate_shard(group_size, torch.float)
            mpu.set_tensor_model_parallel_attributes(main_param, True, 0, 1)

            # main_param_shards.append(main_param)
            group_shard["orig_group"]["params"] = [ main_param ]

            # # Update optimizer group.
            # self.optimizer.param_groups[group_index]["params"] = [ main_param ]

        # pax(1, {
        #     "opt_group_shards" : opt_group_shards,
        #     "main_param_shards" : main_param_shards,
        # })

        # return main_param_shards

941
942
    def __init__(self, optimizer, clip_grad, log_num_zeros_in_grad,
                 params_have_main_grad, use_contiguous_buffers_in_local_ddp,
943
                 bf16, grad_scaler, models):
944
945
946

        super().__init__(
            optimizer, clip_grad, log_num_zeros_in_grad,
Lawrence McAfee's avatar
Lawrence McAfee committed
947
            params_have_main_grad, use_contiguous_buffers_in_local_ddp,
948
            bf16, grad_scaler, models)
949

950
951
        # >>>
        args = get_args()
952
        assert args.use_contiguous_buffers_in_local_ddp # already checked in args
953
        # <<<
954

Lawrence McAfee's avatar
Lawrence McAfee committed
955
956
957
958
        # # Data parallel info.
        # self.data_parallel_group = mpu.get_data_parallel_group()
        # self.data_parallel_rank = mpu.get_data_parallel_rank()
        # self.data_parallel_world_size = mpu.get_data_parallel_world_size()
959

960
961
962
963
        # Model grad buffer shards.
        self.model_gbuf_shards = []
        for model_index, model in enumerate(self.models):
            self.model_gbuf_shards.append(self.get_model_gbuf_shard_map(model))
Lawrence McAfee's avatar
Lawrence McAfee committed
964
        self.param_gbuf_map = self.get_param_gbuf_map(self.model_gbuf_shards)
965

966
967
        # pax(0, {"param_gbuf_map": [ (str(tuple(p.shape)), d) for p, d in self.param_gbuf_map.items() ]})

968
969
970
971
972
        # Optimizer shards.
        self.opt_group_shards = self.get_optimizer_group_shards(
            self.optimizer.param_groups,
            self.model_gbuf_shards)

973
        # pax(0, {**{"opt_group_shards / %d" % i : g for i, g in enumerate(self.opt_group_shards)}})
Lawrence McAfee's avatar
Lawrence McAfee committed
974

975
976
977
978
        # Allocate main param shards.
        # self.main_param_shards = \
        #     self.allocate_main_param_shards(self.opt_group_shards)
        self.allocate_main_param_shards(self.opt_group_shards)
979

980
        # >>>
981
982
983
984
985
        # pax(0, {
        #     "model_gbuf_shards" : self.model_gbuf_shards,
        #     "opt_group_shards" : self.opt_group_shards,
        #     "main_param_shards" : self.main_param_shards,
        # })
986
987
        # <<<

988
989
990
991
992
        # Update optimizer groups.
        # - Also, leverage state_dict() and load_state_dict() to
        #   recast preexisting per-param state tensors.
        self.optimizer.param_groups = \
            [ g["orig_group"] for g in self.opt_group_shards ]
Lawrence McAfee's avatar
Lawrence McAfee committed
993
994
        self.optimizer.load_state_dict(self.optimizer.state_dict())

995
996
997
998
999
1000
1001
        # pax(1, {
        #     "opt_group_shards" : self.opt_group_shards,
        #     "param_groups" : self.optimizer.param_groups,
        # })

        # Initialize main params.
        self._copy_model_params_to_main_params()
1002

1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
    @staticmethod
    def has_nan_debug(tensors):
        if isinstance(tensors, torch.Tensor):
            tensors = [ tensors ]
        assert isinstance(tensors, list)
        has_nans = [ (not torch.all(torch.isfinite(t)).item()) for t in tensors ]
        has_nan = any(has_nans)
        return has_nan
    def get_local_model_param_views(self):
        '''** FOR DEBUGGING. **'''
        model_param_views = []
        for group_index, opt_group_shard in enumerate(self.opt_group_shards):
            for param, opt_shard in opt_group_shard["param_map"].items():
                model_index, dtype = self.param_gbuf_map[param]
                gbuf_shard_map = \
                    self.model_gbuf_shards[model_index][dtype]["param_map"][param]
                model_param_shard = gbuf_shard_map["param"]
                model_param_views.append(
                    param.view(-1)[model_param_shard.start:model_param_shard.end])
        return model_param_views
    def get_local_model_grad_views(self):
        '''** FOR DEBUGGING. **'''
        model_grad_views = []
        for group_index, opt_group_shard in enumerate(self.opt_group_shards):
            for param, opt_shard in opt_group_shard["param_map"].items():
                model_index, dtype = self.param_gbuf_map[param]
                gbuf = self.models[model_index]._grad_buffers[dtype].data
                gbuf_shard_map = \
                    self.model_gbuf_shards[model_index][dtype]["param_map"][param]
                gbuf_world_shard = gbuf_shard_map["gbuf_world"]
                model_grad_views.append(
                    gbuf[gbuf_world_shard.start:gbuf_world_shard.end])
        return model_grad_views
    def get_world_model_params(self):
        '''** FOR DEBUGGING. **'''
        return [ p for m in self.models for p in m.parameters() ]

    def get_main_params(self):
        return [ g["params"][0] for g in self.optimizer.param_groups ]
    def get_main_grads(self):
        return [ p.grad for p in self.get_main_params() ]
1044
    def get_main_param(self, group_index):
1045
1046
        # return self.optimizer.param_groups[group_index]["params"][0]
        return self.get_main_params()[group_index]
1047
1048
1049
    def get_main_grad(self, group_index):
        return self.get_main_param(group_index).grad

1050
1051
1052
1053
1054
1055
    def load_state_dict(self):
        raise Exception("hi.")
    def reload_model_params(self):
        raise Exception("hi.")
    def state_dict(self):
        raise Exception("hi.")
Lawrence McAfee's avatar
Lawrence McAfee committed
1056
1057
1058

    def zero_grad(self, set_to_none=True):

Lawrence McAfee's avatar
Lawrence McAfee committed
1059
        model_params = []
Lawrence McAfee's avatar
Lawrence McAfee committed
1060
1061
        for model in self.models:
            for dtype, param_map in model._grad_buffer_param_index_map.items():
Lawrence McAfee's avatar
Lawrence McAfee committed
1062
1063
1064
1065
                model_params.extend(param_map.keys())
        # main_params = []
        # for main_group in self.optimizer.param_groups:
        #     main_params.extend(main_group["params"])
Lawrence McAfee's avatar
Lawrence McAfee committed
1066

Lawrence McAfee's avatar
Lawrence McAfee committed
1067
1068
        _zero_grad_group_helper(model_params, set_to_none)
        # _zero_grad_group_helper(params, set_to_none = False)
Lawrence McAfee's avatar
Lawrence McAfee committed
1069

Lawrence McAfee's avatar
Lawrence McAfee committed
1070
        # pax(0, {"params": params})
1071

1072
1073
    def get_model_grad_buffer_dp_views(self):

Lawrence McAfee's avatar
Lawrence McAfee committed
1074
        # >>>
1075
1076
1077
        # ** only contiguous grad buffer supported, for now [ TEMPORARY ] **
        args = get_args()
        assert args.use_contiguous_buffers_in_local_ddp
Lawrence McAfee's avatar
Lawrence McAfee committed
1078
        # <<<
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096

        # Grad buffer views.
        gbuf_view_items = []
        for model_index, model in enumerate(self.models):
            for dtype, gbuf_shard in self.model_gbuf_shards[model_index].items():
                world_shards = gbuf_shard["world_all"]

                gbuf = model._grad_buffers[dtype]
                gbuf_views = []
                for shard in world_shards:
                    gbuf_views.append(gbuf.data[shard.start:shard.end])

                gbuf_view_items.append((model_index, dtype, gbuf_views))

        # pax(0, {"gbuf_view_items": gbuf_view_items})

        return gbuf_view_items

1097
    def reduce_gradients(self, model):
1098

Lawrence McAfee's avatar
Lawrence McAfee committed
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
        # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
        # Sync word embedding params.

        # ... todo ...

        # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
        # Sync T5 position embedding params.

        # ... todo ...

1109
1110
        # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
        # Reduce-scatter.
Lawrence McAfee's avatar
Lawrence McAfee committed
1111
1112
        data_parallel_rank = mpu.get_data_parallel_rank()
        data_parallel_group = mpu.get_data_parallel_group()
1113

1114
        gbuf_view_items = self.get_model_grad_buffer_dp_views()
Lawrence McAfee's avatar
Lawrence McAfee committed
1115

1116
1117
1118
1119
1120
1121
1122
1123
        for model_index, dtype, gbuf_views in gbuf_view_items:
            torch.distributed.reduce_scatter(
                gbuf_views[data_parallel_rank],
                gbuf_views,
                group = data_parallel_group,
            )
            
        # pax(0, {"gbuf_view_items": gbuf_view_items})
Lawrence McAfee's avatar
Lawrence McAfee committed
1124

1125
    def gather_params(self):
Lawrence McAfee's avatar
Lawrence McAfee committed
1126

1127
1128
        data_parallel_rank = mpu.get_data_parallel_rank()
        data_parallel_group = mpu.get_data_parallel_group()
1129

1130
1131
        gbuf_view_items = self.get_model_grad_buffer_dp_views()

Lawrence McAfee's avatar
Lawrence McAfee committed
1132
        # All-gather updated main params.
1133
1134
1135
1136
1137
1138
1139
        for model_index, dtype, gbuf_views in gbuf_view_items:
            torch.distributed.all_gather(
                gbuf_views,
                gbuf_views[data_parallel_rank],
                group = data_parallel_group,
            )

Lawrence McAfee's avatar
Lawrence McAfee committed
1140
1141
        # Each model param now contains its updated values in it's
        # '.main_grad' field.
1142
1143
1144
1145
1146
1147
1148
        for param in self.param_gbuf_map:
            param.detach().copy_(param.main_grad)
            # pax(0, {
            #     "param" : tp(param),
            #     "main_grad" : tp(param.main_grad),
            #     # "grad" : tp(param.grad),
            # })
1149

1150
1151
1152
1153
1154
1155
1156
        # pax(0, {
        #     "gbuf_view_items" : gbuf_view_items,
        #     "param_gbuf_map" : [
        #         (str(tuple(p.shape)), d)
        #         for p, d in self.param_gbuf_map.items()
        #     ],
        # })
1157
1158
1159
1160
1161
1162
        pax(1, {
            "main params" : self.get_main_params(),
            "model params / world" : self.get_world_model_params(),
            "gbuf_view_item" : tp(gbuf_view[data_parallel_rank]),
            # "model params / local" : self.get_local_model_param_views(),
        })
1163

Lawrence McAfee's avatar
Lawrence McAfee committed
1164
    def _collect_main_grad_data_for_unscaling(self):
1165
        # return [ p.grad.data for p in self.main_param_shards ]
1166
        # return [ p.grad.data for p in self.main_param_shards if p is not None ]
1167
1168
1169
        # return [ self.get_main_grad(gi).data
        #          for gi in range(len(self.opt_group_shards)) ]
        return [ g.data for g in self.get_main_grads() ]
Lawrence McAfee's avatar
Lawrence McAfee committed
1170

1171
1172
1173
    def _copy_model_params_to_main_params(self):

        for group_index, group_shard in enumerate(self.opt_group_shards):
1174
            main_param = self.get_main_param(group_index)
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
            for model_param, main_shard in group_shard["param_map"].items():

                # Model shard.
                model_index, dtype = self.param_gbuf_map[model_param]
                model_shard = self.model_gbuf_shards \
                    [model_index][dtype]["param_map"][model_param]["param"]

                assert main_shard.size == model_shard.size

                # Copy shard data.
                main_view = main_param[main_shard.start:main_shard.end]
1186
1187
                model_view = model_param.view(-1)[model_shard.start:model_shard.end]
                # try:
1188
                main_view.detach().copy_(model_view)
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
                # except:
                #     pax({
                #         "main_param" : tp(main_param),
                #         "model_param" : tp(model_param),
                #         "main_view" : tp(main_view),
                #         "model_view" : tp(model_view),
                #         "main_shard" : str(main_shard),
                #         "model_shard" : str(model_shard),
                #     })

1199
        # pax(0, {
1200
1201
1202
1203
        #     **{
        #         "opt_group_shards / %d" % i : s
        #         for i, s in enumerate(self.opt_group_shards)
        #     },
1204
        #     "main_params" : self.get_main_params(),
1205
        # })
1206

1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
    def _copy_model_grads_to_main_grads(self, ITERATION):

        # >>>
        model_grads = self.get_local_model_grad_views()
        model_has_nan = self.has_nan_debug(model_grads)
        if model_has_nan:
            pax(1, {
                "ITERATION" : ITERATION,
                "model grads" : model_grads,
                "model_has_nan" : model_has_nan,
                "model params / local" : self.get_local_model_param_views(),
                # "model params / world" : [ list(self.param_gbuf_map),
                # "main grads" : self.get_main_grads(),
            })
        # <<<
Lawrence McAfee's avatar
Lawrence McAfee committed
1222
1223

        for group_index, group_shard in enumerate(self.opt_group_shards):
1224
            for model_param, main_shard in group_shard["param_map"].items():
Lawrence McAfee's avatar
Lawrence McAfee committed
1225

1226
                model_index, dtype = self.param_gbuf_map[model_param]
Lawrence McAfee's avatar
Lawrence McAfee committed
1227
                model_shard = self.model_gbuf_shards \
1228
                    [model_index][dtype]["param_map"][model_param]["gbuf_world"]
Lawrence McAfee's avatar
Lawrence McAfee committed
1229
1230
1231

                assert main_shard.size == model_shard.size

Lawrence McAfee's avatar
Lawrence McAfee committed
1232
                # Copy from DDP's contiguous buffer to main shard's grad.
1233
                model_grad = self.models[model_index]._grad_buffers[dtype].data
1234
                main_grad = self.get_main_grad(group_index)
Lawrence McAfee's avatar
Lawrence McAfee committed
1235

Lawrence McAfee's avatar
Lawrence McAfee committed
1236
                # Copy sub-range within tensor.
1237
1238
                model_view = model_grad[model_shard.start:model_shard.end]
                main_view = main_grad[main_shard.start:main_shard.end]
Lawrence McAfee's avatar
Lawrence McAfee committed
1239

1240
                main_view.detach().copy_(model_view)
Lawrence McAfee's avatar
Lawrence McAfee committed
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255

                # pax(0, {
                #     "group_index" : group_index,
                #     "group_shard" : group_shard,
                #     "param" : tp(param),
                #     "model_index" : model_index,
                #     "gbuf_dtype" : str(gbuf_dtype),
                #     "model_grad_tensor" : tp(model_grad_tensor),
                #     "main_grad_tensor" : tp(main_grad_tensor),
                #     "model_grad_view" : tp(model_grad_view),
                #     "main_grad_view" : tp(main_grad_view),
                #     "model_shard" : str(model_shard),
                #     "main_shard" : str(main_shard),
                # })

Lawrence McAfee's avatar
Lawrence McAfee committed
1256
        # >>>
1257
1258
1259
1260
1261
1262
1263
        # pax(1, {
        #     # "model_gbuf_shards" : self.model_gbuf_shards,
        #     **{
        #         "opt_group_shards / %d" % i : s
        #         for i, s in enumerate(self.opt_group_shards)
        #     },
        #     "main_grads" : self.get_main_grads(),
Lawrence McAfee's avatar
Lawrence McAfee committed
1264
        # })
1265
1266
        # for group_index, main_grad in enumerate(self.get_main_grads()):
        #     # is_nan = torch.any(torch.isnan(main_grad)).item()
1267
        #     if is_nan:
1268
1269
1270
1271
        #         # opt_group_shard = self.opt_group_shards[group_index]
        #         # param_views = []
        #         # for param, shard in opt_group_shard["param_map"].items():
        #         #     ddd
1272
        #         pax(0, {
1273
1274
1275
1276
1277
1278
        #             "opt_group_shard" : self.opt_group_shards[group_index],
        #             "param_map" : [ (str(p.shape), str(d)) for p, d in self.opt_group_shards[group_index]["param_map"].items() ],
        #             "gbufs" : [ b.data for m in self.models for d, b in m._grad_buffers.items() ],
        #             "group_index" : group_index,
        #             "main_param" : tp(self.get_main_param(group_index)),
        #             "main_grad" : tp(main_grad),
1279
1280
        #             "is_nan" : is_nan,
        #         })
1281
1282
1283
1284
1285
1286
1287
1288
        main_grads = self.get_main_grads()
        main_has_nan = self.has_nan_debug(main_grads)
        if main_has_nan:
            raise Exception("hi.")

        # pax(1, {
        #     "model grads" : self.get_local_model_grad_views(),
        # })
Lawrence McAfee's avatar
Lawrence McAfee committed
1289
        # <<<
Lawrence McAfee's avatar
Lawrence McAfee committed
1290

1291

1292
    def _copy_main_params_to_model_params(self, ITERATION):
1293
1294

        for group_index, group_shard in enumerate(self.opt_group_shards):
1295
            for model_param, main_shard in group_shard["param_map"].items():
1296

1297
                model_index, dtype = self.param_gbuf_map[model_param]
1298
                model_shard = self.model_gbuf_shards \
1299
                    [model_index][dtype]["param_map"][model_param]["gbuf_world"]
1300
1301
1302
1303

                assert main_shard.size == model_shard.size

                # Use DDP's contiguous buffer to temporarily hold params.
1304
                model_param = self.models[model_index]._grad_buffers[dtype].data
1305
                main_param = self.get_main_param(group_index)
1306
1307

                # Copy sub-range within tensor.
1308
1309
                model_view = model_param[model_shard.start:model_shard.end]
                main_view = main_param[main_shard.start:main_shard.end]
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331

                model_view.detach().copy_(main_view)

                # Debug.
                # pax(0, {
                #     "group_index" : group_index,
                #     "group_shard" : group_shard,
                #     "param" : tp(param),
                #     "model_index" : model_index,
                #     "gbuf_dtype" : str(gbuf_dtype),
                #     "model_grad_tensor" : tp(model_grad_tensor),
                #     "main_grad_tensor" : tp(main_grad_tensor),
                #     "model_grad_view" : tp(model_grad_view),
                #     "main_grad_view" : tp(main_grad_view),
                #     "model_shard" : str(model_shard),
                #     "main_shard" : str(main_shard),
                # })

        # pax(0, {
        #     "model_gbuf_shards" : self.model_gbuf_shards,
        #     "opt_group_shards" : self.opt_group_shards,
        # })
Lawrence McAfee's avatar
Lawrence McAfee committed
1332
1333
        # >>>
        for param in self.param_gbuf_map:
1334
1335
            # is_nan = torch.any(torch.isnan(param)).item()
            is_nan = not torch.all(torch.isfinite(param)).item()
Lawrence McAfee's avatar
Lawrence McAfee committed
1336
            if is_nan:
1337
                pax({
Lawrence McAfee's avatar
Lawrence McAfee committed
1338
1339
1340
                    "param" : tp(param),
                    "is_nan" : is_nan,
                })
1341
1342
1343
1344
1345
1346
1347

        # model_param_views = self.get_local_model_param_views()
        # pax(1, {
        #     "ITERATION" : ITERATION,
        #     "main params" : self.get_main_params(),
        #     "model params / local" : self.get_local_model_param_views(),
        # })
Lawrence McAfee's avatar
Lawrence McAfee committed
1348
        # <<<
1349

1350
1351
# <<<

mohammad's avatar
mohammad committed
1352

mohammad's avatar
mohammad committed
1353
1354
class FP32Optimizer(MegatronOptimizer):

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
1355
1356
    def __init__(self, optimizer, clip_grad,
                 log_num_zeros_in_grad,
1357
                 params_have_main_grad,
1358
                 use_contiguous_buffers_in_local_ddp):
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
1359
1360
1361

        super(FP32Optimizer, self).__init__(
            optimizer, clip_grad, log_num_zeros_in_grad,
1362
            params_have_main_grad, use_contiguous_buffers_in_local_ddp)
mohammad's avatar
mohammad committed
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380

        self._scale = torch.cuda.FloatTensor([1.0])


    def zero_grad(self, set_to_none=True):
        """Copied from torch.optim.optimizer"""
        for group in self.optimizer.param_groups:
            _zero_grad_group_helper(group['params'], set_to_none)


    def get_loss_scale(self):
        """FP32 optimizer does not do any scaling."""
        return self._scale


    @torch.no_grad()
    def step(self):
        """Clip gradients (if needed) and step the base optimizer.
mohammad's avatar
mohammad committed
1381
        Always return successful since there is no overflow."""
mohammad's avatar
mohammad committed
1382

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
1383
1384
1385
1386
1387
1388
        # Copy main_grads to grads.
        if self.params_have_main_grad:
            for param_group in self.optimizer.param_groups:
                for param in param_group['params']:
                    param.grad = param.main_grad

1389
1390
1391
                    # Safe to de-reference model's main_grad after copying.
                    # (If using contiguous buffers, main_grad's memory should
                    # persist and therefore should not be deallocated.)
1392
                    if not self.use_contiguous_buffers_in_local_ddp:
1393
1394
                        param.main_grad = None

mohammad's avatar
mohammad committed
1395
        # Clip gradients.
1396
        grad_norm = None
mohammad's avatar
mohammad committed
1397
        if self.clip_grad > 0.0:
1398
            grad_norm = self.clip_grad_norm(self.clip_grad)
mohammad's avatar
mohammad committed
1399

Rewon Child's avatar
Rewon Child committed
1400
        # count the zeros in the grads
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
1401
1402
        num_zeros_in_grad = self.count_zeros() if \
                            self.log_num_zeros_in_grad else None
Rewon Child's avatar
Rewon Child committed
1403

mohammad's avatar
mohammad committed
1404
1405
1406
1407
        # Update parameters.
        self.optimizer.step()

        # No overflow for FP32 optimizer.
1408
        return True, grad_norm, num_zeros_in_grad
mohammad's avatar
mohammad committed
1409
1410


1411
1412
1413
1414
    def reload_model_params(self):
        pass


mohammad's avatar
mohammad committed
1415
1416
1417
1418
1419
1420
    def state_dict(self):
        return self.optimizer.state_dict()


    def load_state_dict(self, state_dict):
        self.optimizer.load_state_dict(state_dict)