class_docstr: Megatron style Random Batch Sampler.
  
  Major difference is that `__iter__` yields a local minibatch, not a microbatch.
  A local minibatch consists of `global_batch_size / data_parallel_size`
  
  Args:
      total_samples: The number of data samples, i.e. ``len(dataset)``.
      consumed_samples: The number of samples already consumed in pretraining.
      local_minibatch_size: The number of data in each batch returned from `__iter__`. Basically
          `local_minibatch_size = global_batch_size / data_parallel_size`.
      data_parallel_rank:
      data_parallel_size:
class_docstr_tok: ["Megatron","style","Random","Batch","Sampler",".","Major","difference","is","that","`","__iter__","`","yields","a","local","minibatch",",","not","a","microbatch",".","A","local","minibatch","consists","of","`","global_batch_size","\/","data_parallel_size","`","Args",":","total_samples",":","The","number","of","data","samples",",","i.e",".","``","len","(","dataset",")","``",".","consumed_samples",":","The","number","of","samples","already","consumed","in","pretraining",".","local_minibatch_size",":","The","number","of","data","in","each","batch","returned","from","`","__iter__","`",".","Basically","`","local_minibatch_size","=","global_batch_size","\/","data_parallel_size","`",".","data_parallel_rank",":","data_parallel_size",":"]
code: def __iter__(self):
      active_total_samples = self.total_samples - self.last_batch_size
      self.epoch = self.consumed_samples // active_total_samples
      current_epoch_samples = (
          self.consumed_samples % active_total_samples
      )
      # note(mkozuki): might be better to uncomment
      # assert current_epoch_samples % (self.data_parallel_size * apex.transformer.pipeline_parallel.utils.get_micro_batch_size()) == 0
  
      # data sharding and random sampling
      bucket_size = (
          self.total_samples
          // self.local_minibatch_times_data_parallel_size
      ) * self.local_minibatch_size
      bucket_offset = current_epoch_samples // self.data_parallel_size
      start_idx = self.data_parallel_rank * bucket_size
  
      g = torch.Generator()
      g.manual_seed(self.epoch)
      random_idx = torch.randperm(bucket_size, generator=g).tolist()
      idx_range = [start_idx + x for x in random_idx[bucket_offset:]]
  
      batch = []
      # Last batch if not complete will be dropped.
      for idx in idx_range:
          batch.append(idx)
          if len(batch) == self.local_minibatch_size:
              self.consumed_samples += (
                  self.local_minibatch_times_data_parallel_size
              )
              yield batch
              batch = []
  
imports: import abc
  import torch
label_desc: Use image node_id 5 for calling the MegatronPretrainingRandomSampler obj's underlying member method code with example usage: obj.__iter__() without return types