fixed incorrect order of fp8 metadata initialization

Signed-off-by: Alp Dener <adener@nvidia.com>
NVIDIA · Jan 21, 2024 · a868a5c · a868a5c
1 parent ce2b738
commit a868a5c
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -860,12 +860,12 @@ def __init__(
                 del self.weight_tensor
                 del self.bias_tensor
 
-        self.reset_parameters(defer_init=(device == 'meta'))
-
         if self.primary_weights_in_fp8:
             self.init_fp8_metadata()
             self.fp8_meta["update_amax_and_scale_fwd"] = True
 
+        self.reset_parameters(defer_init=(device == 'meta'))
+
         self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
         # For RPL, bias has to be added after TP collectives

diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -1226,12 +1226,12 @@ def __init__(
         else:
             self.fc2_bias = torch.Tensor().to(dtype=params_dtype, device=device)
 
-        self.reset_parameters(defer_init=(device == 'meta'))
-
         if self.primary_weights_in_fp8:
             self.init_fp8_metadata(num_gemms=2)
             self.fp8_meta["update_amax_and_scale_fwd"] = True
 
+        self.reset_parameters(defer_init=(device == 'meta'))
+
         # For RPL, bias has to be added after TP collectives
         # So it cannot be fused with the GEMM
         if self.set_parallel_mode and self.apply_bias:

diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
@@ -754,12 +754,12 @@ def __init__(
                 del self.weight_tensor
                 del self.bias_tensor
 
-        self.reset_parameters(defer_init=(device == 'meta'))
-
         if self.primary_weights_in_fp8:
             self.init_fp8_metadata()
             self.fp8_meta["update_amax_and_scale_fwd"] = True
 
+        self.reset_parameters(defer_init=(device == 'meta'))
+
         self.fp8_weight_shapes.append(torch.Size((self.out_features, self.in_features)))
 
         # For RPL, bias has to be added after TP collectives