__init__.py 5.89 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
20
from megatron import get_args

21
from .distrib_optimizer import DistributedOptimizer
mohammad's avatar
mohammad committed
22
from .grad_scaler import ConstantGradScaler, DynamicGradScaler
23
from .optimizer import Float16OptimizerWithFloat16Params, FP32Optimizer
24

mohammad's avatar
mohammad committed
25

26
27
28
29
30
31
32
33
def get_param_groups(modules,
                     no_weight_decay_cond,
                     scale_lr_cond,
                     lr_mult):
    """creates param groups based on weight decay condition (regularized vs non regularized)
       and learning rate scale condition (args.lr vs lr_mult * args.lr)
       scale_lr_cond is used during finetuning where head of the network requires a scaled
       version of the base learning rate. 
mohammad's avatar
mohammad committed
34
    """
35
36
37
38
    wd_no_scale_lr = []
    wd_scale_lr = []
    no_wd_no_scale_lr = []
    no_wd_scale_lr = []
39
    for module in modules:
40
41
42
43
44
45
        for name, param in module.named_parameters():
            if not param.requires_grad:
                continue

            if no_weight_decay_cond is not None:
                no_wd = no_weight_decay_cond(name, param)
46
            else:
Vijay Korthikanti's avatar
Vijay Korthikanti committed
47
                # do not regularize biases nor Norm parameters
48
                no_wd = name.endswith(".bias") or len(param.shape) == 1
mohammad's avatar
mohammad committed
49

50
51
52
53
            if scale_lr_cond is not None:
                scale_lr = scale_lr_cond(name, param)
            else:
                scale_lr = False
mohammad's avatar
mohammad committed
54

55
56
57
58
59
60
61
62
            if not no_wd and not scale_lr:
                wd_no_scale_lr.append(param)
            elif not no_wd and scale_lr:
                wd_scale_lr.append(param)
            elif no_wd and not scale_lr:
                no_wd_no_scale_lr.append(param)
            else:
                no_wd_scale_lr.append(param)
mohammad's avatar
mohammad committed
63

64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
    param_groups = []
    if len(wd_no_scale_lr):
        param_groups.append({'params': wd_no_scale_lr, 'wd_mult': 1.0, 'lr_mult': 1.0})
    if len(wd_scale_lr):
        param_groups.append({'params': wd_scale_lr, 'wd_mult': 1.0, 'lr_mult': lr_mult})
    if len(no_wd_no_scale_lr):
        param_groups.append({'params': no_wd_no_scale_lr, 'wd_mult': 0.0, 'lr_mult': 1.0})
    if len(no_wd_scale_lr):
        param_groups.append({'params': no_wd_scale_lr, 'wd_mult': 0.0, 'lr_mult': lr_mult})

    return param_groups

def get_megatron_optimizer(model,
                           no_weight_decay_cond=None,
                           scale_lr_cond=None,
                           lr_mult=1.0):
mohammad's avatar
mohammad committed
80
81
82
    args = get_args()

    # Base optimizer.
83
84
85
86
87
    param_groups = get_param_groups(model,
                                    no_weight_decay_cond,
                                    scale_lr_cond,
                                    lr_mult)

88
    if args.optimizer == 'adam':
89
90
91
92
93
        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
94
    elif args.optimizer == 'sgd':
95
96
97
98
        optimizer = SGD(param_groups,
                        lr=args.lr,
                        weight_decay=args.weight_decay,
                        momentum=args.sgd_momentum)
Vijay Korthikanti's avatar
Vijay Korthikanti committed
99
100
    else:
        raise Exception('{} optimizer is not supported.'.format(
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
101
            args.optimizer))
mohammad's avatar
mohammad committed
102

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
103
104
105
106
107
    # Determine whether the params have main-grad field.
    params_have_main_grad = False
    if args.DDP_impl == 'local':
        params_have_main_grad = True

108
109
110
111
    # Mixed precision optimizer.
    # - Note: both the Float16Optimizer and the DistributedOptimizer inherit
    #   from the MixedPrecisionOptimizer, which manages any optimizer where
    #   the model params and main params are distinct.
112
    if args.fp16 or args.bf16 or args.use_distributed_optimizer:
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
113
114
115
116
117
118
119
120

        # 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
121

mohammad's avatar
mohammad committed
122
123
124
        # Constant loss scale.
        if args.loss_scale:
            grad_scaler = ConstantGradScaler(args.loss_scale)
125

mohammad's avatar
mohammad committed
126
127
        # Dynamic loss scale.
        else:
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
128
129
130
131
132
133
134
135
136
            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
137
        # Megatron optimizer.
138
139
140
141
        opt_ty = DistributedOptimizer \
            if args.use_distributed_optimizer else \
            Float16OptimizerWithFloat16Params
        return opt_ty(optimizer,
142
143
144
145
                      args.clip_grad,
                      args.log_num_zeros_in_grad,
                      params_have_main_grad,
                      args.use_contiguous_buffers_in_local_ddp,
146
                      args.fp16,
147
                      args.bf16,
148
149
                      grad_scaler,
                      model)
mohammad's avatar
mohammad committed
150
151

    # FP32.
152
153
154
155
156
    return FP32Optimizer(optimizer, args.clip_grad,
                         args.log_num_zeros_in_grad,
                         params_have_main_grad,
                         args.use_contiguous_buffers_in_local_ddp,
                         model)