[AMP] Support pure fp16 training mode for dygraph

[zhangbo9674]

PR types

New features

PR changes

Others

Describe

1. Background:

The paddy static graph training mode has developed the pure fp16 training mode, but dynamic graph lacks this training mode.

• static graph pure fp16 pr：Support pure fp16 training for AMP API. #29544

It is verified on GPT-2 117M model that there is a large gap on training speed between Paddle AMP and Megatron (dygraph mode), Megatron uses fp16 training, all model parameters use fp16, all op executed in fp16.

To sum up, we developed the dynamic graph pure fp16 training mode。

2. Pure fp16 API: (dygraph mode)

# 1. Decorate the models and optimizers for mixed accuracy training, include amp and pure fp16:
# paddle.amp.decorator()
def decorate(models, optimizers=None, level='O1', master_weight=True, save_dtype=None)

#2. The context manager of mixed accuracy training, include amp and pure fp16:
# paddle.amp.auto_cast()
def auto_cast(enable=True, custom_white_list=None, custom_black_list=None, level='O1')

2.1. Rewrite network parameters from fp32 to fp16 by decorate:

In amp training mode, it uses the blac&white list to control the fp16 computation, so it will insert many cast OPs. In pure fp16 training mode, all OPs are executed in fp16, unless this OP does not support fp16.

Therefore, under the dygraph training, it is necessary to rewrite all the network parameters to fp16, so there is no need to insert cast OP for data conversion in the process of execution.

2.2. Optimzier update the fp16 parameters of the network by decorate:

Rewriting the network parameters from fp32 to fp16 in 2.1 is not an inplace operation. Therefore, optimzier needs to update the fp16 network parameters.

2.3. level parameter in decorate and auto_cast:

The level parameter is set to unify the training results of AMP and pure fp16. The level accept values are O1 and O2:

• O1 represent amp, the input data type of each operator will be casted by white_list and black_list;

• O2 represent pure fp16, all OPs parameters and input data will be casted to fp16, except OPs in black_list, don't support fp16 kernel and batchnorm parameters.

• In decorate, default value for level is O1. In amp training mode, the decorate will do nothing, so you do not need call this api, but in pure fp16, you need call decorate explicitly.

• In auto_cast, default value for level is O1. This is to be compatible with the original amp training mode of Paddle.

2.4. master_weight in decorate:

Until now, Momentum, Adam and AdamW support the float16 computation. All three of them have the multi_precision parameters, which can avoid poor accuracy or slow convergence in a way.

In decorate, the default master_weight is None. If master_weight is None or True, in pure fp16, multi_precision will be set to True, if user do not want use this strategy, user should set to False.

2.5. save_dtype in decorate:

In pure fp16, the model parameters will rewrite from fp32 to fp16. For Inference, it usually need to save the model of fp32 data type, so we provide an interface save_dtype, when save_dtype is not None, we will register a data type conversion function hook for Layer.state_dict(). So that, all parameters in state_dict will cast to save_dtype. Finally, the data type of the model parameters saved through paddle.save and paddle.jit.save will be save_dtype.

save_dtype support fp16、fp32、fp64 or None. If save_dtype is None, we will not register data type conversion function hook.

2.6. custom_white_list and custom_black_list:

• In amp, the input data type of each operator will be casted by white_list and black_list.

• In pure fp16, only black_list is effective.

3. Use example:

import paddle
import paddle.nn as nn

# define a fp32 model:
class SimpleNet(nn.Layer):
    def __init__(self, input_size, output_size):
        super(SimpleNet, self).__init__()
        self.linear1 = nn.Linear(input_size, output_size)
        self.relu1 = nn.ReLU()
        self.linear2 = nn.Linear(input_size, output_size)
    def forward(self, x):
        x = self.linear1(x)
        x = self.relu1(x)
        x = self.linear2(x)
        return x

paddle.seed(10)
epochs = 20
input_size = 4096   
output_size = 4096  
batch_size = 512   
nums_batch = 50

# training data:
train_data = [paddle.randn((batch_size, input_size)) for _ in range(nums_batch)]
label_data = [paddle.randn((batch_size, output_size)) for _ in range(nums_batch)]

# define model:
model = SimpleNet(input_size, output_size)  
# define optimizer：
optimizer = paddle.optimizer.SGD(learning_rate=0.0001, parameters=model.parameters())  
# define GradScaler：
scaler = paddle.amp.GradScaler(init_loss_scaling=1024)

# step1：decorate model & optimizer for pure fp16 training mode:
model, optimizer = paddle.amp.decorate(models=model, optimizers=optimizer, level='O2', save_dtype='float32')

# train:
for epoch in range(epochs):
    datas = zip(train_data, label_data)
    for batch_id, (data, label) in enumerate(datas):
        # step2：open pure fp16 context by auto_cast:
        with paddle.amp.auto_cast(enable=True, custom_white_list=None, custom_black_list=None, level='O2'):
            output = model(data)
            loss = paddle.nn.MSELoss(output, label)  
        scaled = scaler.scale(loss)
        scaled.backward()
        scaler.minimize(optimizer, scaled)
        optimizer.clear_grad()

4. Performance testing：

4.1. Test GPT2 performance by pure fp16:

• batch size = 8
• repo: https://github.com/PaddlePaddle/PaddleNLP/tree/develop/examples/language_model/gpt-3/dygraph

**4.2. Result：

• loss：

• speed：

5、文档预览：

中文文档预览：

• overview: http://10.136.157.23:8090/documentation/docs/zh/api/paddle/amp/Overview_cn.html?reviewVersion=jenkins-doc-review-py3.8-252

• auto_cast: http://10.136.157.23:8090/documentation/docs/zh/api/paddle/amp/auto_cast_cn.html?reviewVersion=jenkins-doc-review-py3.8-252

• decorate: http://10.136.157.23:8090/documentation/docs/zh/api/paddle/amp/decorate_cn.html?reviewVersion=jenkins-doc-review-py3.8-252

英文文档预览：

• auto_cast: http://10.136.157.23:8090/documentation/docs/en/api/paddle/amp/auto_cast_en.html?reviewVersion=jenkins-doc-review-py3.8-252

• decorate: http://10.136.157.23:8090/documentation/docs/en/api/paddle/amp/decorate_en.html?reviewVersion=jenkins-doc-review-py3.8-252

[paddle-bot-old]

Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[GuoxiaWang]

下次提交的时候，帮忙把 sparse_momentum 也加上 MasterParam 吧，谢谢

[zhangbo9674]

目前框架中还没有找到使用sparse_momentum的优化器，以及动态图调用sparse_momentum的地方，所以pure fp16的pr中暂时先不加入了。

[zhiqiu]

change tracer.enebale_amp to tracer.amp_level

[zhangbo9674]

Done，thanks！

[TCChenlong]

LGTM for Docs

[zhiqiu]

LGTM!

[ForFishes]

LGTM!!

[TCChenlong]

LGTM