class_docstr: Megatron Bert pretraining.
  Model returns [batch, seq, hidden] shape
class_docstr_tok: ["Megatron","Bert","pretraining",".","Model","returns","[","batch",",","seq",",","hidden","]","shape"]
code: def allreduce_first_last_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 and last stages to ensure
      # that word_embeddings parameters stay in sync.
      # This should only run for models that support pipelined model parallelism
      # (BERT and GPT-2).
      if (
          parallel_state.get_pipeline_model_parallel_world_size() > 1
          and (
              parallel_state.is_pipeline_first_stage(
                  ignore_virtual=True
              )
              or parallel_state.is_pipeline_last_stage(
                  ignore_virtual=True
              )
          )
      ):
          if parallel_state.is_pipeline_first_stage(
              ignore_virtual=True
          ):
              if isinstance(self.model, list):
                  module = self.model[
                      0
                  ]  # only the first virtual rank has the embeddings
              else:
                  module = self.model
          if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
              if isinstance(self.model, list):
                  module = self.model[
                      -1
                  ]  # only the last virtual rank has the embeddings
              else:
                  module = self.model
          if module.share_token_embeddings:
              word_embeddings_weight = (
                  module.shared_embedding_or_output_weight()
                  if self.mcore_bert
                  else module.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()
              )
  
imports: import itertools
  from typing import Any, Dict, List, Optional
  import torch
  import torch.nn.functional
  from omegaconf.dictconfig import DictConfig
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.data.language_modeling.megatron import dataset_utils
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
  from nemo.collections.nlp.models.language_modeling.megatron.bert_model import BertModel
  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.utils import ApexGuardDefaults, average_losses_across_data_parallel_group, get_params_for_weight_decay_optimization
  from nemo.collections.nlp.parts.nlp_overrides import GradScaler
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.core.classes.common import PretrainedModelInfo
  from nemo.core.neural_types import ChannelType, MaskType, NeuralType
  from nemo.utils import AppState, logging
label_desc: Use image node_id 7 for calling the MegatronBertModel obj's underlying member method code with example usage: obj.allreduce_first_last_embeddings() without return types