fit auto_parallel amp for llama (#59497)

zhiqiu · web-flow · commit 9fb7f90678bb · 2023-11-30T12:07:13.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -34,6 +34,7 @@
 )
 from paddle.metric import Metric
 from paddle.static import InputSpec, Operator, Variable, global_scope
+from paddle.static.amp.fp16_utils import _convert_float_to_bfloat16
 
 from ...utils.log_utils import get_logger
 from ..interface import CollectionNames, fetch, get_collection
@@ -833,25 +834,42 @@ def _initialize(self, mode, init_parameters=True):
             random.seed(self._strategy.seed + self._dp_ranks[0])
 
         dist_context = self._dist_contexts[mode]
+        dist_main_program = dist_context.dist_main_programs[self._cur_rank]
         if self._dygraph_mode:
             if not init_parameters:
                 self._share_parameters()
             else:
-                dist_main_program = dist_context.dist_main_programs[
-                    self._cur_rank
-                ]
                 self.program_helper.init(
                     dist_main_program, self._place, dist_context
                 )
             # The model's instance variables (not paramters), used in forward function,
             # have been initialized when initialize model in dynamic mode.
             if self._model and len(self._model.buffers()) > 0:
                 for buffer in self._model.buffers():
-                    scope_var = global_scope().find_var(buffer.name)
-                    buffer_tensor = global_scope().var(buffer.name).get_tensor()
-                    if scope_var and buffer_tensor._is_initialized():
-                        continue
-                    buffer_tensor.set(buffer.numpy(), self._place)
+                    if dist_main_program.global_block().has_var(buffer.name):
+                        dest_type = (
+                            dist_main_program.global_block()
+                            .var(buffer.name)
+                            .dtype
+                        )
+                        scope_var = global_scope().find_var(buffer.name)
+                        buffer_tensor = (
+                            global_scope().var(buffer.name).get_tensor()
+                        )
+                        if scope_var and buffer_tensor._is_initialized():
+                            continue
+                        # for amp
+                        if dest_type == core.VarDesc.VarType.BF16:
+                            buffer_tensor.set(
+                                _convert_float_to_bfloat16(buffer.numpy()),
+                                self._place,
+                            )
+                        elif dest_type == core.VarDesc.VarType.FP16:
+                            buffer_tensor.set(
+                                np.float16(buffer.numpy()), self._place
+                            )
+                        else:
+                            buffer_tensor.set(buffer.numpy(), self._place)
 
         if self._executor is None:
             self._executor = paddle.static.Executor(self._place)
diff --git a/python/paddle/distributed/passes/auto_parallel_fp16.py b/python/paddle/distributed/passes/auto_parallel_fp16.py
@@ -291,7 +291,10 @@ def resolute_tensor_dtype(self, block):
                         if out_var.dtype == __target_dtype__:
                             out_var.desc.set_dtype(core.VarDesc.VarType.FP32)
             elif is_backward_op(op):
-                if self._is_fp16_op(op.desc.original_id()) is True:
+                if (
+                    self._is_fp16_op(op.desc.original_id()) is True
+                    or op.type == "cast"
+                ):
                     for out_name in op.output_names:
                         if _keep_fp32_output(op, out_name):
                             continue
