t5_pipeline_fix

ea128da5 · Vijay Korthikanti · 3ae12a47 · ea128da5
Commit ea128da5 authored Nov 05, 2021 by Vijay Korthikanti
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

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

No files found.
--- a/megatron/mpu/initialize.py
+++ b/megatron/mpu/initialize.py
@@ -177,7 +177,7 @@ def initialize_model_parallel(tensor_model_parallel_size_=1,
        if len(ranks) > 1:
            embedding_ranks = [ranks[0], ranks[-1]]
            if pipeline_model_parallel_split_rank_ is not None and \
-                    pipeline_model_parallel_split_rank_ not in embedding_ranks:
+                    ranks[pipeline_model_parallel_split_rank_] not in embedding_ranks:
                embedding_ranks = [ranks[0],
                                   ranks[pipeline_model_parallel_split_rank_],
                                   ranks[-1]]