class_docstr: Megatron style BatchSampler.
  
  Let mbs, gbs, tp, pp, and dp stand for "micro batch size", "global batch size",
  "tensor model parallel world size", "pipeline model parallel world size", and
  "data parallel world size", the number of micro batches (hereafter, nmb) is defined as
  :math:`nmb = gbs \div (mbs \times dp)`.
  
  See `apex/transformer/microbatches.py#L91-L98 <https://github.com/NVIDIA/apex/blob/
  44c3043685b6115e7b81b3458a6c76601b1e55b4/apex/transformer/microbatches.py#L91-L98>`_
  for the initial settings of the number of micro batches and
  `apex/transformer/microbatches.py#L160-L177 <https://github.com/NVIDIA/apex/blob/
  44c3043685b6115e7b81b3458a6c76601b1e55b4/apex/transformer/microbatches.py#L160-L177>_`.
  for warming up of global batch size.
  
  e.g.) `(mbs, gbs, tp, pp, dp) = (1, 16, 1, 1, 2)`, then the number of micro batches is
  :math:`gbs \div (mbs \times dp) = 16 \div (1 \times 2) = 8`.
  In this case, an instance of Megatron Batch Sampler on each data parallel rank is expected
  returns :math:`nmb \times mbs = 8` indices.
class_docstr_tok: ["Megatron","style","BatchSampler",".","Let","mbs",",","gbs",",","tp",",","pp",",","and","dp","stand","for","``","micro","batch","size","''",",","``","global","batch","size","''",",","''","tensor","model","parallel","world","size","''",",","``","pipeline","model","parallel","world","size","''",",","and","''","data","parallel","world","size","''",",","the","number","of","micro","batches","(","hereafter",",","nmb",")","is","defined","as",":","math",":","`","nmb","=","gbs","\\div","(","mbs","\\times","dp",")","`",".","See","`","apex\/transformer\/microbatches.py","#","L91-L98","<","https",":","\/\/github.com\/NVIDIA\/apex\/blob\/","44c3043685b6115e7b81b3458a6c76601b1e55b4\/apex\/transformer\/microbatches.py","#","L91-L98",">","`","_","for","the","initial","settings","of","the","number","of","micro","batches","and","`","apex\/transformer\/microbatches.py","#","L160-L177","<","https",":","\/\/github.com\/NVIDIA\/apex\/blob\/","44c3043685b6115e7b81b3458a6c76601b1e55b4\/apex\/transformer\/microbatches.py","#","L160-L177",">","_","`",".","for","warming","up","of","global","batch","size",".","e.g",".",")","`","(","mbs",",","gbs",",","tp",",","pp",",","dp",")","=","(","1",",","16",",","1",",","1",",","2",")","`",",","then","the","number","of","micro","batches","is",":","math",":","`","gbs","\\div","(","mbs","\\times","dp",")","=","16","\\div","(","1","\\times","2",")","=","8","`",".","In","this","case",",","an","instance","of","Megatron","Batch","Sampler","on","each","data","parallel","rank","is","expected","returns",":","math",":","`","nmb","\\times","mbs","=","8","`","indices","."]
docstr: Update the global batch size.
docstr_tok: ["Update","the","global","batch","size","."]
code: def update_global_batch_size(
      self, new_global_batch_size: int
  ) -> None:
      """Update the global batch size."""
      self._global_batch_size = new_global_batch_size
      if (
          self._global_batch_size
          % (self.micro_batch_size * self.data_parallel_size)
          != 0
      ):
          raise RuntimeError(
              f"`global_batch_size` ({self._global_batch_size}) is not divisible by "
              f"`micro_batch_size ({self.micro_batch_size}) x data_parallel_size "
              f"({self.data_parallel_size})`"
          )
      self._num_micro_batches = self._global_batch_size // (
          self.micro_batch_size * self.data_parallel_size
      )
      self._global_batch_size_on_this_data_parallel_rank = (
          self._num_micro_batches * self.micro_batch_size
      )
  
imports: import abc
  import warnings
  from typing import Tuple
  import torch
  from nemo.utils.decorators import experimental
label_desc: Use image node_id 2 for calling the BaseMegatronBatchSampler obj's underlying member method code with example usage: obj.update_global_batch_size(new_global_batch_size) without return types