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: Constructor of Megatron-LM style Batch Sampler.
  
  Args:
      total_samples: The size of dataset.
      consumed_samples: The number of samples that have been used.
      micro_batch_size: The size of each micro batch.
      global_batch_size: The size of global batch.
      data_parallel_rank: The value you can obtain via
          `parallel_state.get_data_parallel_rank()` of megatron.core.
      data_parallel_size: The value you can obtain via
          `parallel_state.get_data_parallel_world_size()` of megatron.core.
docstr_tok: ["Constructor","of","Megatron-LM","style","Batch","Sampler",".","Args",":","total_samples",":","The","size","of","dataset",".","consumed_samples",":","The","number","of","samples","that","have","been","used",".","micro_batch_size",":","The","size","of","each","micro","batch",".","global_batch_size",":","The","size","of","global","batch",".","data_parallel_rank",":","The","value","you","can","obtain","via","`","parallel_state.get_data_parallel_rank","(",")","`","of","megatron.core",".","data_parallel_size",":","The","value","you","can","obtain","via","`","parallel_state.get_data_parallel_world_size","(",")","`","of","megatron.core","."]
code: def __init__(
      self,
      total_samples: int,
      consumed_samples: int,
      micro_batch_size: int,
      global_batch_size: int,
      data_parallel_rank: int,
      data_parallel_size: int,
      drop_last: bool,
      pad_samples_to_global_batch_size=False,
  ) -> None:
      """Constructor of Megatron-LM style Batch Sampler.
  
      Args:
          total_samples: The size of dataset.
          consumed_samples: The number of samples that have been used.
          micro_batch_size: The size of each micro batch.
          global_batch_size: The size of global batch.
          data_parallel_rank: The value you can obtain via
              `parallel_state.get_data_parallel_rank()` of megatron.core.
          data_parallel_size: The value you can obtain via
              `parallel_state.get_data_parallel_world_size()` of megatron.core.
      """
      # Sanity checks.
      if total_samples <= 0:
          raise RuntimeError(
              "no sample to consume: {}".format(total_samples)
          )
      if micro_batch_size <= 0:
          raise RuntimeError(
              f"micro_batch_size size must be greater than 0, but {micro_batch_size}"
          )
      if data_parallel_size <= 0:
          raise RuntimeError(
              f"data parallel size must be greater than 0, but {data_parallel_size}"
          )
      if data_parallel_rank >= data_parallel_size:
          raise RuntimeError(
              "data_parallel_rank should be smaller than data size, but {} >= {}".format(
                  data_parallel_rank, data_parallel_size
              )
          )
      # Keep a copy of input params for later use.
      self.total_samples: int = total_samples
      self.consumed_samples: int = consumed_samples
      self.micro_batch_size: int = micro_batch_size
      self.data_parallel_rank: int = data_parallel_rank
      self.data_parallel_size: int = data_parallel_size
      self.drop_last: bool = drop_last
      self.pad_samples_to_global_batch_size = (
          pad_samples_to_global_batch_size
      )
  
      self.update_global_batch_size(global_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 1 to create a new BaseMegatronBatchSampler object  with example: obj = BaseMegatronBatchSampler(total_samples, consumed_samples, micro_batch_size, global_batch_size, data_parallel_rank, data_parallel_size, drop_last, pad_samples_to_global_batch_size)