support bf16 loss in static (PaddlePaddle#7874)

Author: heavyrain-lzy

llm/llama/auto_parallel/run_pretrain_auto.py

@@ -27,6 +27,7 @@
 import paddle
 import paddle.distributed as dist
 import paddle.distributed.auto_parallel as auto
+from paddle.base.data_feeder import convert_uint16_to_float
 from paddle.profiler.utils import job_schedule_profiler_range
 
 from paddlenlp.ops import Topology
@@ -668,7 +669,10 @@ def loss_func(loss, outputs):
                 outs = engine.run(micro_batch, mode="train")
 
                 if "loss" in outs:
-                    tr_loss_step = np.sum(outs["loss"])
+                    if outs["loss"].dtype == np.uint16:
+                        tr_loss_step = np.sum(convert_uint16_to_float(outs["loss"]))
+                    else:
+                        tr_loss_step = np.sum(outs["loss"])
                 else:
                     tr_loss_step = float(0)