__init__.py 4.25 KB
Newer Older
mohammad's avatar
mohammad committed
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.

from apex.optimizers import FusedAdam as Adam
17
from apex.optimizers import FusedSGD as SGD
mohammad's avatar
mohammad committed
18

mohammad's avatar
mohammad committed
19
from megatron import get_args
20
from megatron.model import LayerNorm
mohammad's avatar
mohammad committed
21
22

from .grad_scaler import ConstantGradScaler, DynamicGradScaler
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
23
from .optimizer import Float16OptimizerWithFloat16Params, FP32Optimizer
mohammad's avatar
mohammad committed
24
25


26
def _get_params_for_weight_decay_optimization(modules):
mohammad's avatar
mohammad committed
27
28
29
30
31
32
    """Divide params into with-weight-decay and without-weight-decay groups.
    Layernorms and baises will have no weight decay but the rest will.
    """

    weight_decay_params = {'params': []}
    no_weight_decay_params = {'params': [], 'weight_decay': 0.0}
33
34
35
36
37
38
39
40
41
42
43
44
45
    for module in modules:
        for module_ in module.modules():
            if isinstance(module_, LayerNorm):
                no_weight_decay_params['params'].extend(
                    [p for p in list(module_._parameters.values())
                     if p is not None])
            else:
                weight_decay_params['params'].extend(
                    [p for n, p in list(module_._parameters.items())
                     if p is not None and n != 'bias'])
                no_weight_decay_params['params'].extend(
                    [p for n, p in list(module_._parameters.items())
                     if p is not None and n == 'bias'])
mohammad's avatar
mohammad committed
46
47
48
49
50
51
52
53
54

    return weight_decay_params, no_weight_decay_params


def get_megatron_optimizer(model):
    args = get_args()

    # Base optimizer.
    param_groups = _get_params_for_weight_decay_optimization(model)
55
56
57
58
59
60
    if args.optimizer == 'adam':
        optimizer = Adam(param_groups,
                         lr=args.lr,
                         weight_decay=args.weight_decay,
                         betas=(args.adam_beta1, args.adam_beta2),
                         eps=args.adam_eps)
Vijay Korthikanti's avatar
Vijay Korthikanti committed
61
    elif args.optimizer == 'sgd':
62
63
64
65
        optimizer = SGD(param_groups,
                        lr=args.lr,
                        weight_decay=args.weight_decay,
                        momentum=args.sgd_momentum)
Vijay Korthikanti's avatar
Vijay Korthikanti committed
66
67
    else:
        raise Exception('{} optimizer is not supported.'.format(
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
68
            args.optimizer))
mohammad's avatar
mohammad committed
69

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
70
71
72
73
74
75
76
77
78
79
80
81
82
83
    # Determine whether the params have main-grad field.
    params_have_main_grad = False
    if args.DDP_impl == 'local':
        params_have_main_grad = True

    if args.fp16 or args.bf16:

        # Grad scaler:
        #    if loss-scale is provided, instantiate the constant scaler.
        #    if we are using fp16 and loss-scale is not present, use a
        #       dynamic scaler.
        #    otherwise we are running in bf16 with no loss-scale so
        #       leave it as None.
        grad_scaler = None
mohammad's avatar
mohammad committed
84
85
86
87
88
        # Constant loss scale.
        if args.loss_scale:
            grad_scaler = ConstantGradScaler(args.loss_scale)
        # Dynamic loss scale.
        else:
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
89
90
91
92
93
94
95
96
97
            if args.fp16:
                grad_scaler = DynamicGradScaler(
                    initial_scale=args.initial_loss_scale,
                    min_scale=args.min_loss_scale,
                    growth_factor=2.0,
                    backoff_factor=0.5,
                    growth_interval=args.loss_scale_window,
                    hysteresis=args.hysteresis)

mohammad's avatar
mohammad committed
98
        # Megatron optimizer.
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
99
100
101
102
103
104
        return Float16OptimizerWithFloat16Params(optimizer,
                                                 args.clip_grad,
                                                 args.log_num_zeros_in_grad,
                                                 params_have_main_grad,
                                                 args.bf16,
                                                 grad_scaler)
mohammad's avatar
mohammad committed
105
106

    # FP32.
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
107
108
109
    return FP32Optimizer(optimizer, args.clip_grad,
                         args.log_num_zeros_in_grad,
                         params_have_main_grad)