comments

                grad = word_embeddings_weight.grad

            torch.distributed.all_reduce(grad, group=mpu.get_embedding_group())

    # All-reduce position_embeddings grad across first (encoder) and split (decoder) 

    # stages to ensure that position embeddings parameters stay in sync.

    # This should only run for T5 models with pipeline parallelism

    if mpu.is_rank_in_position_embedding_group() and \

            mpu.get_pipeline_model_parallel_world_size() > 1 and \

            args.pipeline_model_parallel_split_rank is not None:

        unwrapped_model = model[0]

        unwrapped_model = unwrap_model(

            unwrapped_model, (torchDDP, LocalDDP, Float16Module))

        assert args.DDP_impl == 'local', \

            'T5 model is only supported with local DDP mode'

        grad = unwrapped_model.language_model.embedding.position_embeddings.weight.main_grad

        torch.distributed.all_reduce(grad, group=mpu.get_position_embedding_group())

    timers('backward-embedding-all-reduce').stop()