[GPT-3] Support Grad Merge with FP32 main grad for BF16 training of GPT-3 model

benchmarks/test_tipc/ernie/dygraph/hybrid_parallel/benchmark_common/run_benchmark.sh

@@ -87,7 +87,7 @@ function _train(){
                -o Global.micro_batch_size=${micro_batch_size} \
                -o Engine.max_steps=${max_iter} \
                -o Engine.eval_freq=${eval_freq} \
-               -o Engine.mix_precision.use_pure_fp16=${use_pure_fp16} \
+               -o Engine.mix_precision.enable=${use_pure_fp16} \
                -o Engine.save_load.save_steps=100000 \
                -o Model.hidden_size=1024 \
                -o Model.num_hidden_layers=${num_layers} \

benchmarks/test_tipc/gpt/dygraph/hybrid_parallel/benchmark_common/run_benchmark.sh

@@ -86,7 +86,7 @@ function _train(){
                -o Global.micro_batch_size=${micro_batch_size} \
                -o Engine.max_steps=${max_iter} \
                -o Engine.eval_freq=${eval_freq} \
-               -o Engine.mix_precision.use_pure_fp16=${use_pure_fp16} \
+               -o Engine.mix_precision.enable=${use_pure_fp16} \
                -o Engine.save_load.save_steps=100000 \
                -o Model.hidden_size=1024 \
                -o Model.num_layers=${num_layers} \

benchmarks/test_tipc/gpt/dygraph/sharding/benchmark_common/run_benchmark.sh

@@ -81,7 +81,7 @@ function _train(){
                -o Global.micro_batch_size=${micro_batch_size} \
                -o Engine.max_steps=${max_iter} \
                -o Engine.eval_freq=${eval_freq} \
-               -o Engine.mix_precision.use_pure_fp16=${use_pure_fp16} \
+               -o Engine.mix_precision.enable=${use_pure_fp16} \
                -o Engine.save_load.save_steps=100000 \
                -o Model.use_recompute=${use_recompute} \
                -o Distributed.dp_degree=${dp_degree} \

docs/standard.md

@@ -102,7 +102,9 @@ Engine:
   eval_iters: 10
   test_iters:
   mix_precision:
-    use_pure_fp16: True
+    enable: True
+    dtype: "float16"
+    level: "O2"
     scale_loss: 32768.0
     custom_black_list: ["reduce_sum", "c_softmax_with_cross_entropy", "elementwise_div"]
     custom_white_list: ["lookup_table", "lookup_table_v2"]
@@ -123,10 +125,12 @@ Engine:
 | logging_freq      | 训练日志打印的频率                            |
 | eval_freq         | 模型评估间隔                               |
 | eval_iters        | 模型评估时训练评估测试集的轮数                      |
-| use_pure_fp16     | 是否使用purefp16精度训练                     |
-| scale_loss        | 使用fp16精度下，loss的放缩比例                  |
-| custom_black_list | 自定义算子黑名单。这个名单中的算子在支持float16计算时会被认为是数值危险的，它们的影响也可能会在下游操作中观察到。这些算子通常不会转为float16计算。 |
-| custom_white_list | 自定义算子白名单。这个名单中的算子在支持float16计算时会被认为是数值安全的，并且对性能至关重要。如果设置了白名单，该名单中的算子会使用float16计算。|
+| enable            | 是否使用混合精度策略进行训练                     |
+| dtype             | 混合精度训练数据类型使用float16还是bfloat16，默认为float16类型 |
+| level             | 混合精度训练模式，默认``O2``模式                 |
+| scale_loss        | 使用fp16混合精度策略下，loss的放缩比例                  |
+| custom_black_list | 自定义算子黑名单。这个名单中的算子在支持混合精度计算时会被认为是数值危险的，它们的影响也可能会在下游操作中观察到。这些算子通常不会转为float16/bfloat16计算 |
+| custom_white_list | 自定义算子白名单。这个名单中的算子在支持混合精度计算时会被认为是数值安全的，并且对性能至关重要。如果设置了白名单，该名单中的算子会使用float16/bfloat16计算 |
 | save_steps        | 保存模型间隔                               |
 | save_epoch        | 保存模型epoch间隔                               |
 | output_dir        | 指定输出文件                               |

examples/transformer/models/GPT/docs/README.md

@@ -102,7 +102,9 @@ cd .. # 回到 GPT 目录下
     eval_iters: 10
     test_iters:
     mix_precision:
-      use_pure_fp16: True
+      enable: True
+      dtype: "float16"
+      level: "O2"
       scale_loss: 32768.0
       custom_black_list: ["reduce_sum", "c_softmax_with_cross_entropy", "elementwise_div"]
       custom_white_list: ["lookup_table", "lookup_table_v2"]
@@ -128,10 +130,12 @@ cd .. # 回到 GPT 目录下
 | eval_freq         | 模型评估间隔                               |
 | eval_iters        | 模型评估时训练评估测试集的轮数                      |
 | test_iters        | 模型测试或推理时的轮数                      |
-| use_pure_fp16     | 是否使用purefp16精度训练                     |
-| scale_loss        | 使用fp16精度下，loss的放缩比例                  |
-| custom_black_list | 自定义算子黑名单。这个名单中的算子在支持float16计算时会被认为是数值危险的，它们的影响也可能会在下游操作中观察到。这些算子通常不会转为float16计算。 |
-| custom_white_list | 自定义算子白名单。这个名单中的算子在支持float16计算时会被认为是数值安全的，并且对性能至关重要。如果设置了白名单，该名单中的算子会使用float16计算。|
+| enable            | 是否使用混合精度策略进行训练                     |
+| dtype             | 混合精度训练数据类型使用float16还是bfloat16，默认为float16类型 |
+| level             | 混合精度训练模式，默认``O2``模式                 |
+| scale_loss        | 使用fp16混合精度策略下，loss的放缩比例                  |
+| custom_black_list | 自定义算子黑名单。这个名单中的算子在支持混合精度计算时会被认为是数值危险的，它们的影响也可能会在下游操作中观察到。这些算子通常不会转为float16/bfloat16计算 |
+| custom_white_list | 自定义算子白名单。这个名单中的算子在支持混合精度计算时会被认为是数值安全的，并且对性能至关重要。如果设置了白名单，该名单中的算子会使用float16/bfloat16计算 |
 | save_steps        | 保存模型间隔step数                         |
 | save_epoch        | 保存模型间隔epoch数                        |
 | output_dir        | 指定输出文件                              |

examples/transformer/models/GPT/finetune/configs/finetune_gpt_345M_single_card_glue.yaml

@@ -11,7 +11,7 @@ Global:
   logging_freq: 10
   eval_freq: 1
   mix_precision:
-    use_pure_fp16: True
+    enable: True
     scale_loss: 32768.0
     custom_black_list: ["reduce_sum", "c_softmax_with_cross_entropy", "elementwise_div", "reduce_mean"]
     custom_white_list: ["lookup_table", "lookup_table_v2"]

examples/transformer/models/GPT/finetune/impls.py

@@ -190,7 +190,7 @@ def fit_impl(config, batch, forward_func, **kwargs):
 def eval_impl(config, batch, model, loss_fn, eval_metric):
     model.eval()
 
-    use_fp16 = config.Global.mix_precision.use_pure_fp16
+    use_fp16 = config.Global.mix_precision.enable
     black_list = config.Global.mix_precision.custom_black_list
     white_list = config.Global.mix_precision.custom_white_list
 

examples/transformer/models/GPT/finetune/run.py

@@ -62,7 +62,7 @@
     # build GPT model
     model, tokenizer, train_loss_fn, eval_loss_fn = impls.build_model(config)
 
-    if config.Global.mix_precision.use_pure_fp16:
+    if config.Global.mix_precision.enable:
         scaler = paddle.amp.GradScaler(
             init_loss_scaling=config.Global.mix_precision.scale_loss)
         # Note: Save dtype is the same as model dtype. Also can set save_dtype='float32' when 
@@ -98,14 +98,14 @@
 
     if 'multi_precision' in config.Optimizer:
         assert config.Optimizer.pop('multi_precision') \
-            == config.Global.mix_precision.use_pure_fp16
+            == config.Global.mix_precision.enable
 
     lr_scheduler = cpn.build_lr_scheduler(config.Optimizer.lr)
     optimizer = cpn.build_optimizer(
         config.Optimizer,
         model,
         lr_scheduler,
-        multi_precision=config.Global.mix_precision.use_pure_fp16)
+        multi_precision=config.Global.mix_precision.enable)
 
     # call fleet wrapper
     if nranks > 1:

examples/transformer/models/GPT/generation/configs/generation_gpt_base.yaml

@@ -14,7 +14,7 @@ Global:
   eval_iters: 10
   test_iters:
   mix_precision:
-    use_pure_fp16: True
+    enable: True
     scale_loss: 32768.0
     custom_black_list: ["reduce_sum", "c_softmax_with_cross_entropy", "elementwise_div"]
     custom_white_list: ["lookup_table", "lookup_table_v2"]

examples/transformer/models/GPT/generation/export.py

@@ -51,8 +51,8 @@
     cfg.process_configs(config)
     cfg.print_config(config)
 
-    if config.Global.mix_precision.use_pure_fp16:
-        logger.info("NOTE: disable use_pure_fp16 in export mode")
+    if config.Global.mix_precision.enable:
+        logger.info("NOTE: disable mix_precision in export mode")
 
     # build GPT model
     model, _ = impls.build_model(config)

examples/transformer/models/GPT/offline-eval/configs/eval_gpt_base.yaml

@@ -14,7 +14,7 @@ Global:
   eval_iters: 10
   test_iters:
   mix_precision:
-    use_pure_fp16: True
+    enable: True
     scale_loss: 32768.0
     custom_black_list: ["reduce_sum", "c_softmax_with_cross_entropy", "elementwise_div"]
     custom_white_list: ["lookup_table", "lookup_table_v2"]

examples/transformer/models/GPT/offline-eval/impls.py

@@ -61,7 +61,7 @@ def build_model(config):
 def eval_impl(config, batch, model):
     model.eval()
 
-    use_fp16 = config.Global.mix_precision.use_pure_fp16
+    use_fp16 = config.Global.mix_precision.enable
     black_list = config.Global.mix_precision.custom_black_list
     white_list = config.Global.mix_precision.custom_white_list
 

examples/transformer/models/GPT/offline-eval/run.py

@@ -77,7 +77,7 @@
         ]
         model, quanter = qat.compress_model(config, model, input_spec)
 
-    if config.Global.mix_precision.use_pure_fp16:
+    if config.Global.mix_precision.enable:
         scaler = paddle.amp.GradScaler(
             init_loss_scaling=config.Global.mix_precision.scale_loss)
         # Note: Save dtype is the same as model dtype. Also can set save_dtype='float32' when 

examples/transformer/models/GPT/pretrain/configs/pretrain_gpt_base.yaml

@@ -14,7 +14,7 @@ Global:
   eval_iters: 10
   test_iters:
   mix_precision:
-    use_pure_fp16: True
+    enable: True
     scale_loss: 32768.0
     custom_black_list: ["reduce_sum", "c_softmax_with_cross_entropy", "elementwise_div"]
     custom_white_list: ["lookup_table", "lookup_table_v2"]

examples/transformer/models/GPT/pretrain/export.py

@@ -51,8 +51,8 @@
     cfg.process_configs(config)
     cfg.print_config(config)
 
-    if config.Global.mix_precision.use_pure_fp16:
-        logger.info("NOTE: disable use_pure_fp16 in export mode")
+    if config.Global.mix_precision.enable:
+        logger.info("NOTE: disable mix_precision in export mode")
 
     # build GPT model
     model, _, _ = impls.build_model(config)

examples/transformer/models/GPT/pretrain/impls.py

@@ -101,7 +101,7 @@ def build_model(config):
 
 def model_forward_backward(config, batch, forward_func, **kwargs):
     acc_steps = config.Global.accumulate_steps
-    use_fp16 = config.Global.mix_precision.use_pure_fp16
+    use_fp16 = config.Global.mix_precision.enable
     black_list = config.Global.mix_precision.custom_black_list
     white_list = config.Global.mix_precision.custom_white_list
 
@@ -165,7 +165,7 @@ def model_forward_backward(config, batch, forward_func, **kwargs):
 
 def optim_update_params(config, **kwargs):
     hcg = env.get_hcg()
-    use_fp16 = config.Global.mix_precision.use_pure_fp16
+    use_fp16 = config.Global.mix_precision.enable
 
     dp_degree = config.Distributed.dp_degree
     sharding_stage = config.Distributed.sharding.sharding_stage
@@ -221,7 +221,7 @@ def fit_impl(config, batch, forward_func, **kwargs):
 def eval_impl(config, batch, model, loss_fn):
     model.eval()
 
-    use_fp16 = config.Global.mix_precision.use_pure_fp16
+    use_fp16 = config.Global.mix_precision.enable
     black_list = config.Global.mix_precision.custom_black_list
     white_list = config.Global.mix_precision.custom_white_list
 

examples/transformer/models/GPT/pretrain/run.py

@@ -83,7 +83,7 @@
         ]
         model, quanter = qat.compress_model(config, model, input_spec)
 
-    if config.Global.mix_precision.use_pure_fp16:
+    if config.Global.mix_precision.enable:
         scaler = paddle.amp.GradScaler(
             init_loss_scaling=config.Global.mix_precision.scale_loss)
         # Note: Save dtype is the same as model dtype. Also can set save_dtype='float32' when 
@@ -104,7 +104,7 @@
         config.Optimizer,
         model,
         lr_scheduler,
-        multi_precision=config.Global.mix_precision.use_pure_fp16)
+        multi_precision=config.Global.mix_precision.enable)
 
     # call fleet wrapper
     if nranks > 1:

examples/transformer/models/GPT/pretrain_moe/configs/pretrain_moe_345M_single_card.yaml

@@ -7,7 +7,7 @@ Global:
   max_steps: 20000
   logging_freq: 10
   mix_precision:
-    use_pure_fp16: True
+    enable: True
 
 Data:
   Train:

examples/transformer/models/GPT/pretrain_moe/configs/pretrain_moe_base.yaml

@@ -14,7 +14,7 @@ Global:
   eval_iters: 10
   test_iters:
   mix_precision:
-    use_pure_fp16: True
+    enable: True
     scale_loss: 32768.0
     custom_black_list: ["reduce_sum", "c_softmax_with_cross_entropy", "elementwise_div"]
     custom_white_list: ["lookup_table", "lookup_table_v2"]

examples/transformer/models/GPT/pretrain_moe/impls.py

@@ -59,7 +59,7 @@ def _get_model_size(l, h, v, s, ne, ei):
             # gate
             P += (h * nei + nei)
             # experts
-            P += nei * (8 * h * h + 5 * h) 
+            P += nei * (8 * h * h + 5 * h)
         # FFN Layer
         else:
             P += 8 * h * h + 5 * h
@@ -120,7 +120,7 @@ def build_model(config):
 
 def model_forward_backward(config, batch, forward_func, **kwargs):
     acc_steps = config.Global.accumulate_steps
-    use_fp16 = config.Global.mix_precision.use_pure_fp16
+    use_fp16 = config.Global.mix_precision.enable
     black_list = config.Global.mix_precision.custom_black_list
     white_list = config.Global.mix_precision.custom_white_list
 
@@ -199,7 +199,7 @@ def model_forward_backward(config, batch, forward_func, **kwargs):
 
 def optim_update_params(config, **kwargs):
     hcg = env.get_hcg()
-    use_fp16 = config.Global.mix_precision.use_pure_fp16
+    use_fp16 = config.Global.mix_precision.enable
 
     dp_degree = config.Distributed.dp_degree
     sharding_stage = config.Distributed.sharding.sharding_stage
@@ -255,7 +255,7 @@ def fit_impl(config, batch, forward_func, **kwargs):
 def eval_impl(config, batch, model, loss_fn):
     model.eval()
 
-    use_fp16 = config.Global.mix_precision.use_pure_fp16
+    use_fp16 = config.Global.mix_precision.enable
     black_list = config.Global.mix_precision.custom_black_list
     white_list = config.Global.mix_precision.custom_white_list
 

examples/transformer/models/GPT/pretrain_moe/run.py

@@ -84,7 +84,7 @@
         ]
         model, quanter = qat.compress_model(config, model, input_spec)
 
-    if config.Global.mix_precision.use_pure_fp16:
+    if config.Global.mix_precision.enable:
         scaler = paddle.amp.GradScaler(
             init_loss_scaling=config.Global.mix_precision.scale_loss)
         # Note: Save dtype is the same as model dtype. Also can set save_dtype='float32' when 
@@ -105,7 +105,7 @@
         config.Optimizer,
         model,
         lr_scheduler,
-        multi_precision=config.Global.mix_precision.use_pure_fp16)
+        multi_precision=config.Global.mix_precision.enable)
 
     # call fleet wrapper
     if nranks > 1:

examples/transformer/utils/config.py

@@ -413,7 +413,7 @@ def process_global_configs(config):
     global_cfg['mix_precision'] = global_cfg.get('mix_precision', {})
     amp_cfg = global_cfg.mix_precision
 
-    amp_cfg['use_pure_fp16'] = amp_cfg.get('use_pure_fp16', False)
+    amp_cfg['enable'] = amp_cfg.get('enable', False)
     amp_cfg['scale_loss'] = amp_cfg.get('scale_loss', 32768)
     amp_cfg['custom_black_list'] = amp_cfg.get('custom_black_list', None)
     amp_cfg['custom_white_list'] = amp_cfg.get('custom_white_list', None)

ppfleetx/configs/multimodal/imagen/imagen_base.yaml

@@ -15,7 +15,7 @@ Engine:
   eval_freq: 10000000
   eval_iters: 10000000
   mix_precision:
-    use_pure_fp16: False 
+    enable: False 
     scale_loss: 32768.0
     custom_black_list: ["reduce_sum", "c_softmax_with_cross_entropy", "elementwise_div"]
     custom_white_list: ["lookup_table", "lookup_table_v2"]

ppfleetx/configs/multimodal/imagen/imagen_super_resolution_1024.yaml

@@ -37,7 +37,7 @@ Engine:
   eval_freq: 10000000
   eval_iters: 10000000
   mix_precision:
-    use_pure_fp16: False 
+    enable: False 
     scale_loss: 32768.0
     custom_black_list: ["reduce_sum", "c_softmax_with_cross_entropy", "elementwise_div"]
     custom_white_list: ["lookup_table", "lookup_table_v2"]

ppfleetx/configs/nlp/ernie/auto/pretrain_ernie_base.yaml

@@ -17,7 +17,7 @@ Engine:
   eval_iters: 10
   test_iters: -1
   mix_precision:
-    use_pure_fp16: False
+    enable: False
     scale_loss: 32768.0
     custom_black_list: ["reduce_sum", "c_softmax_with_cross_entropy", "elementwise_div"]
     custom_white_list: ["lookup_table", "lookup_table_v2"]

ppfleetx/configs/nlp/ernie/finetune_ernie_base.yaml

@@ -17,7 +17,7 @@ Engine:
   eval_iters: 10
   test_iters: -1
   mix_precision:
-    use_pure_fp16: False
+    enable: False
     scale_loss: 32768.0
     custom_black_list: ["reduce_sum", "c_softmax_with_cross_entropy", "elementwise_div"]
     custom_white_list: ["lookup_table", "lookup_table_v2"]

ppfleetx/configs/nlp/ernie/pretrain_ernie_base.yaml

@@ -17,7 +17,7 @@ Engine:
   eval_iters: 10
   test_iters: -1
   mix_precision:
-    use_pure_fp16: False
+    enable: False
     scale_loss: 32768.0
     custom_black_list: ["reduce_sum", "c_softmax_with_cross_entropy", "elementwise_div"]
     custom_white_list: ["lookup_table", "lookup_table_v2"]

ppfleetx/configs/nlp/ernie/qat_ernie_base.yaml

@@ -17,7 +17,7 @@ Engine:
   eval_iters: 10
   test_iters: -1
   mix_precision:
-    use_pure_fp16: False
+    enable: False
     scale_loss: 32768.0
     custom_black_list: ["reduce_sum", "c_softmax_with_cross_entropy", "elementwise_div"]
     custom_white_list: ["lookup_table", "lookup_table_v2"]

ppfleetx/configs/nlp/gpt/finetune_gpt_345M_single_card_glue.yaml

@@ -13,7 +13,7 @@ Engine:
   logging_freq: 10
   eval_freq: 1
   mix_precision:
-    use_pure_fp16: True
+    enable: True
     scale_loss: 32768.0
     custom_black_list: ["reduce_sum", "c_softmax_with_cross_entropy", "elementwise_div", "reduce_mean"]
     custom_white_list: ["lookup_table", "lookup_table_v2"]