ADLR/megatron-lm!2316 - respect perform_initialization

Author: akoumpa

megatron/core/extensions/transformer_engine.py

@@ -338,7 +338,7 @@ def __init__(
                 input_size,
                 output_size_per_partition,
                 0,
-                init_method,
+                init_method=condition_init_method(config, init_method),
                 stride=1,
                 return_master_weight=False,
                 rank=rank,
@@ -427,7 +427,7 @@ def __init__(
                 input_size,
                 output_size_per_partition,
                 0,
-                init_method,
+                init_method=condition_init_method(config, init_method),
                 stride=1,
                 return_master_weight=False,
                 rank=rank,
@@ -501,7 +501,7 @@ def __init__(
                 input_size,
                 input_size_per_partition,
                 1,
-                init_method,
+                init_method=condition_init_method(config, init_method),
                 stride=1,
                 return_master_weight=False,
                 params_dtype=config.params_dtype,