fix linear weight initilize in paddle and resume

Signed-off-by: Guoxia Wang <[email protected]>

transformer_engine/paddle/layer/base.py

@@ -187,6 +187,8 @@ def _get_fp8_state(self) -> paddle.Tensor:
             state = {}
             state["scaling_fwd"] = self.fp8_meta["scaling_fwd"].to_numpy()
             state["scaling_bwd"] = self.fp8_meta["scaling_bwd"].to_numpy()
+            get_global_fp8_state().get_fp8_fwd_buffer().wait()
+            get_global_fp8_state().get_fp8_bwd_buffer().wait()
             state["global_fp8_fwd_buffer"] = get_global_fp8_state().get_fp8_fwd_buffer().to_numpy()
             state["global_fp8_bwd_buffer"] = get_global_fp8_state().get_fp8_bwd_buffer().to_numpy()
             # Store other pickelable values.
@@ -224,7 +226,10 @@ def _set_fp8_state(self, state: paddle.Tensor) -> None:
         if state is None:
             return
 
-        state = pickle.loads(state.numpy().tobytes())
+        if isinstance(state, paddle.Tensor):
+            state = pickle.loads(state.numpy().tobytes())
+        else:
+            state = pickle.loads(state.tobytes())
         if state is None:
             return
 

transformer_engine/paddle/layer/linear.py

@@ -786,17 +786,19 @@ def __init__(
 
         # Initialize weight parameter
         with track_rng_state(enable=self.tensor_parallel):
-            # TE linear weight is in column major
+            # Ensure that the random seed state is consistent with Paddle Linear.
             self.weight = self.create_parameter(
-                shape=(
-                    [self.out_features, self.in_features]
-                    if self.backend == "transformer_engine"
-                    else [self.in_features, self.out_features]
-                ),
+                shape=([self.in_features, self.out_features]),
                 attr=self._weight_attr,
                 dtype=self._dtype,
                 is_bias=False,
             )
+            # TE linear weight is in column major
+            if self.backend == "transformer_engine":
+                with paddle.no_grad():
+                    clone_weight = self.weight.clone().t().contiguous()
+                    self.weight.get_tensor()._share_data_with(clone_weight.get_tensor())
+
         set_weight_tensor_dist_attr(
             self.weight, self.tensor_parallel, self.parallel_mode, self.backend
         )