fix for nemo: do not initialize mpu if it is already initialized

6e433055 · mohammad · c63906a6 · 6e433055
Commit 6e433055 authored Oct 01, 2020 by mohammad
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

megatron/initialize.py megatron/initialize.py +4 -1

No files found.
--- a/megatron/initialize.py
+++ b/megatron/initialize.py
@@ -123,7 +123,10 @@ def _initialize_distributed():
    # Set the model-parallel / data-parallel communicators.
    if device_count > 0:
-        mpu.initialize_model_parallel(args.model_parallel_size)
+        if mpu.model_parallel_is_initialized():
+            print('model parallel is already initialized')
+        else:
+            mpu.initialize_model_parallel(args.model_parallel_size)
 def _init_autoresume():