Merge branch 'akoumparouli/fix_te_skip_init' into 'main'

respect perform_initialization

See merge request ADLR/megatron-lm!2316

Author: ko3n1g

megatron/core/extensions/transformer_engine.py

@@ -348,7 +348,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,
@@ -441,7 +441,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,
@@ -519,7 +519,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,