code: def __init__(
      self,
      total_samples: int,
      consumed_samples: int,
      local_minibatch_size: int,
      data_parallel_rank: int,
      data_parallel_size: int,
      drop_last: bool = True,
  ):
      # Sanity checks.
      if total_samples <= 0:
          raise RuntimeError(
              "no sample to consume: {}".format(self.total_samples)
          )
      if consumed_samples >= total_samples:
          raise RuntimeError(
              "no samples left to consume: {}, {}".format(
                  self.consumed_samples, self.total_samples
              )
          )
      if local_minibatch_size <= 0:
          raise RuntimeError(
              f"local minibatch size must be greater than 0: {local_minibatch_size}"
          )
      if data_parallel_size <= 0:
          raise RuntimeError(
              f"data parallel size must be greater than 0: {data_parallel_size}"
          )
      if data_parallel_rank >= data_parallel_size:
          raise RuntimeError(
              "data_parallel_rank should be smaller than data size: {}, {}".format(
                  self.data_parallel_rank, data_parallel_size
              )
          )
      # Keep a copy of input params for later use.
      self.total_samples = total_samples
      self.consumed_samples = consumed_samples
      self._local_minibatch_size = local_minibatch_size
      self.data_parallel_rank = data_parallel_rank
      self.data_parallel_size = data_parallel_size
      self.local_minibatch_times_data_parallel_size = (
          self._local_minibatch_size * data_parallel_size
      )
      self.drop_last = drop_last
  
imports: import abc
  import torch
label_desc: Use image node_id 1 to create a new MegatronPretrainingSampler object from inherited base classes: _Base with example: obj = MegatronPretrainingSampler(total_samples, consumed_samples, local_minibatch_size, data_parallel_rank, data_parallel_size, drop_last)