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 configure_optimizers(self):
      if self.with_distributed_adam:
          # Disable overlapped grad sync for embedding grad when
          # pipeline parallelism is enabled
          if (
              parallel_state.get_pipeline_model_parallel_world_size()
              > 1
          ):
              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 module.share_token_embeddings:
                      param = module.word_embeddings_weight()
                      param._disable_greedy_grad_copy = (
                          not self.megatron_amp_O2
                      )
                      param._disable_overlap_grad_sync = True
              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:
                      param = module.word_embeddings_weight()
                      param._disable_greedy_grad_copy = (
                          not self.megatron_amp_O2
                      )
                      param._disable_overlap_grad_sync = True
  
          # Disable overlapped grad sync for layer norm grads when
          # sequence parallelism is enabled
          for param in self.parameters():
              if getattr(param, "sequence_parallel", False):
                  param._disable_greedy_grad_copy = (
                      not self.megatron_amp_O2
                  )
                  param._disable_overlap_grad_sync = True
  
          # sequence parallelism is enabled
          for param in self.parameters():
              if getattr(param, "sequence_parallel", False):
                  param._disable_greedy_grad_copy = (
                      not self.megatron_amp_O2
                  )
                  param._disable_overlap_grad_sync = True
  
          # Initialize parameter buckets for overlapped grad and param syncs
          # Note: Params with disabled overlapping are put in the
          # last param bucket
          buckets = []
          if (
              self.cfg.get("virtual_pipeline_model_parallel_size", None)
              is not None
          ):
              # Initialize a bucket for each virtual pipeline stage
              for module in self.model:
                  if isinstance(
                      module, (Float16Module, MCoreFloat16Module)
                  ):
                      module = module.module
                  stage_bucket = []
                  layers = (
                      module.encoder.layers
                      if self.mcore_bert
                      else module.language_model.encoder.layers
                  )
                  for layer in layers:
                      stage_bucket.extend(
                          p
                          for p in layer.parameters()
                          if not getattr(
                              p, "_disable_overlap_grad_sync", False
                          )
                      )
                  buckets.append(stage_bucket)
          else:
              # Initialize a bucket for each Transformer layer
              modules = (
                  self.model
                  if isinstance(self.model, list)
                  else [self.model]
              )
              for module in modules:
                  if isinstance(
                      module, (Float16Module, MCoreFloat16Module)
                  ):
                      module = module.module
                  layers = (
                      module.encoder.layers
                      if self.mcore_bert
                      else module.language_model.encoder.layers
                  )
                  for layer in layers:
                      buckets.append(
                          [
                              p
                              for p in layer.parameters()
                              if not getattr(
                                  p, "_disable_overlap_grad_sync", False
                              )
                          ]
                      )
          buckets.reverse()
          used_params = set()
          for bucket in buckets:
              used_params.update(bucket)
          buckets[-1].extend(
              p for p in self.parameters() if p not in used_params
          )
          self.distributed_adam_buckets = buckets
  
      return super().configure_optimizers()
  
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 30 for calling the MegatronBertModel obj's underlying member method code with example usage: obj.configure_optimizers() and returns: super