Commit 945ece94 authored by Jared Casper's avatar Jared Casper
Browse files

Merge branch 't5_pipeline_fixes' into 'main' 

t5_pipeline_fix

See merge request ADLR/megatron-lm!353
parents d4169684 ea128da5

megatron/mpu/initialize.py

+1 −1
Original line number Diff line number Diff line
@@ -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]]