class_docstr: Megatron encoder-decoder base class
class_docstr_tok: ["Megatron","encoder-decoder","base","class"]
code: def allreduce_word_and_position_embeddings(self):
      # Modified from megatron-lm: https://github.com/NVIDIA/Megatron-LM/blob/d41696840ed0a7edb7e0499eb82a48ae112d9bb3/megatron/training.py#L407
      # All-reduce word_embeddings' grad across first, last stages to ensure that word_embeddings parameters stay in sync.
      if (
          parallel_state.get_pipeline_model_parallel_world_size() > 1
          and (parallel_state.is_rank_in_embedding_group())
      ):
          if self.cfg.get(
              "share_token_embeddings", True
          ) and self.cfg.get(
              "share_decoder_tokens_head_embeddings", True
          ):
              word_embeddings_weight = (
                  self.enc_dec_model.word_embeddings_weight()
              )
              if self.megatron_amp_O2:
                  # O2 recipe stores a "main" copy of weights and grads
                  grad = word_embeddings_weight.main_grad
              else:
                  grad = word_embeddings_weight.grad
              torch.distributed.all_reduce(
                  grad, group=parallel_state.get_embedding_group()
              )
          else:
              raise ValueError(
                  f"Attempting to allreduce word_embeddings for pipeline parallel size > 1, but found untied word embeddings or token head embeddings. This is not supported yet."
              )
  
      # All-reduce position embeddings for T5.
      if (
          parallel_state.is_rank_in_position_embedding_group()
          and parallel_state.get_pipeline_model_parallel_world_size()
          > 1
          and parallel_state.get_pipeline_model_parallel_split_rank()
          is not None
          and self.cfg.encoder.get("position_embedding_type")
          == "learned_absolute"
          and self.cfg.decoder.get("position_embedding_type")
          == "learned_absolute"
      ):
          if self.cfg.get("share_token_embeddings", True):
              position_embeddings_weight = (
                  self.enc_dec_model.position_embeddings_weight()
              )
              if self.megatron_amp_O2:
                  grad = position_embeddings_weight.main_grad
              else:
                  grad = position_embeddings_weight.grad
              torch.distributed.all_reduce(
                  grad,
                  group=parallel_state.get_position_embedding_group(),
              )
  
      # All-reduce relative position embeddings for T5.
      if (
          parallel_state.get_pipeline_model_parallel_world_size()
          > 2  # This > 2 and not > 1 since with PP=2 encoder RPE can live only on one rank.
          and parallel_state.get_pipeline_model_parallel_split_rank()
          is not None
      ):
          # For split rank = 1, we have only one encoder rank and so we don't need to allreduce.
          if (
              self.cfg.encoder.get("position_embedding_type")
              == "relative"
              and parallel_state.is_rank_in_encoder_relative_position_embedding_group()
              and parallel_state.get_pipeline_model_parallel_split_rank()
              > 1
          ):
              position_embeddings_weight = (
                  self.enc_dec_model.encoder_relative_position_embeddings_weight()
              )
              if self.megatron_amp_O2:
                  grad = position_embeddings_weight.main_grad
              else:
                  grad = position_embeddings_weight.grad
              torch.distributed.all_reduce(
                  grad,
                  group=parallel_state.get_encoder_relative_position_embedding_group(),
              )
  
          # For split rank == pipeline_world_size - 1, we have only one decoder rank and so we don't need to allreduce.
          if (
              self.cfg.decoder.get("position_embedding_type")
              == "relative"
              and parallel_state.is_rank_in_decoder_relative_position_embedding_group()
          ):
              position_embeddings_weight = (
                  self.enc_dec_model.decoder_relative_position_embeddings_weight()
              )
              if self.megatron_amp_O2:
                  grad = position_embeddings_weight.main_grad
              else:
                  grad = position_embeddings_weight.grad
              torch.distributed.all_reduce(
                  grad,
                  group=parallel_state.get_decoder_relative_position_embedding_group(),
              )
  
              # If the model also has separate RPE weights for decoder cross-attention, allreduce those as well.
              if not self.cfg.decoder.get(
                  "relative_position_bias_self_attention_only", True
              ):
                  position_embeddings_weight = (
                      self.enc_dec_model.decoder_cross_attention_relative_position_embeddings_weight()
                  )
                  if self.megatron_amp_O2:
                      grad = position_embeddings_weight.main_grad
                  else:
                      grad = position_embeddings_weight.grad
                  torch.distributed.all_reduce(
                      grad,
                      group=parallel_state.get_decoder_relative_position_embedding_group(),
                  )
  
imports: import copy
  import functools
  import inspect
  from typing import Any, Dict, List, Optional
  import torch
  from omegaconf import OmegaConf, open_dict
  from omegaconf.dictconfig import DictConfig
  from pytorch_lightning.accelerators import CPUAccelerator
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.build_model import build_model
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module
  from nemo.collections.nlp.modules.common.megatron.token_level_encoder_decoder import MegatronTokenLevelEncoderDecoderModule
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, average_losses_across_data_parallel_group, get_params_for_weight_decay_optimization
  from nemo.collections.nlp.modules.common.text_generation_utils import compute_beam_search_len_penalty, get_sampling_token_fn
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.utils import AppState, logging
label_desc: Use image node_id 16 for calling the MegatronLMEncoderDecoderModel obj's underlying member method code with example usage: obj.allreduce_word_and_position_embeddings() without return types