def METHOD_NAME(self, dataset, consumed_samples):
    if isinstance(dataset, BlendableDataset):
        collate_fun = dataset.datasets[0].collate_fn
    else:
        collate_fun = dataset.collate_fn
    collate_fn = partial(collate_fun, tp_workers=0)
    global_batch_size = self.trainer.world_size * self.cfg.micro_batch_size // self.cfg.tensor_model_parallel_size
    batch_sampler = MegatronPretrainingBatchSampler(
        total_samples=len(dataset),
        consumed_samples=consumed_samples,
        micro_batch_size=self.cfg.micro_batch_size,
        global_batch_size=global_batch_size,
        data_parallel_rank=parallel_state.get_data_parallel_rank(),
        data_parallel_size=parallel_state.get_data_parallel_world_size(),
        drop_last=True,
    )
    return torch.utils.data.DataLoader(
        dataset, batch_sampler=batch_sampler, collate_fn=collate_fn, num_workers=0, pin_memory=True,
    )