Issue with DistributedSampler and Custom Batch Sampler in PyTorch Lightning

Uma_Sushmitha_Guntur: <p>Hi everyone,</p>
<p>I’m working on a Pytorch Lightning pipeline (on a machine with 4 H100 GPUs) where I need to pass a sorted dataset (audio here) into a custom batch sampler to (1) bucket these samples and (2) batch them down the line. My goal is to use the sorted indices for custom batch sampling to minimize padding (since sorting can bring similar lengths together and can reduce padding within a batch) during data loading in the <code>collate_fn</code> (which currently uses a zero padding technique which pads all samples to the max sequence).<br>
Below is my get_dataloader code which is called in train code:</p>
<pre><code class="lang-auto">if len(dataset_list) == 1:
    dataset = dataset_list[0]
else:
    dataset = torch.utils.data.ConcatDataset(dataset_list)

# Sort dataset by audio length in decreasing order
sorted_indices = sorted(range(len(dataset)), key=lambda i: dataset[i].audio.shape[1], reverse=True)
sorted_dataset = [dataset[i] for i in sorted_indices]

dataset = sorted_dataset

# Use DistributedSampler without shuffling
sampler = DistributedSampler(dataset, shuffle=False)

# Custom BucketBatchSampler with bucketing logic
bucket_sampler = BucketBatchSampler(
    sampler, 
    batch_size=self.config.physical_batch_size,
    drop_last=train,
    buckets=25,
)

# Create DataLoader with the custom batch sampler
dataloader = torch.utils.data.DataLoader(
    dataset=dataset,
    batch_sampler=bucket_sampler,
    collate_fn=dataset_list[0].collate_fn,
    num_workers=4,
    pin_memory=True,
    shuffle=train,
)

</code></pre>
<p>The issue I’m encountering is that when I manually initialize DistributedSampler (also turned off use_distributed_sampler=False in Trainer initalization), the below error occurs. If I let Lightning handle the distributed setup without manually changing the argument shuffle=False, no error occurs but it disrupts my sorted indices (shuffles it) and my bucketing logic needs the sampler to output sorted indices.</p>
<p>Interestingly, If i use a sampler =<code>SequentialSampler(dataset)</code>, it overrides this with the DistributedSampler again…</p>
<p>Here’s the error I receive when using <code>DistributedSampler(dataset, shuffle=False)</code>:</p>
<pre><code class="lang-auto">ValueError: Default process group has not been initialized, please make sure to call init_process_group.
</code></pre>
<p>Has anyone faced a similar challenge, or does anyone know how to prevent the dataset from being shuffled while using DistributedSampler in this context? I’m specifically looking for a way to maintain the sorted order for my custom batch sampling logic.</p>
<p>Also providing my bucket_sampler logic for more context:</p>
<pre><code class="lang-auto">class BucketBatchSampler(BatchSampler):
    def __init__(self, sampler, batch_size, drop_last=False, buckets=25, lengths_dict=None):
        self.sampler = sampler
        self.batch_size = batch_size
        self.drop_last = drop_last
        self.buckets = buckets
    
    def __iter__(self):
        # Get indices from the sampler (this will be different for each GPU)
        # print("SAMPLER NAME within bucket batch", self.sampler)
        indices = list(self.sampler)
        print("sorted_indices", indices)
        


        if self.drop_last or len(indices) % self.buckets == 0:
            bucket_size = len(indices) // self.buckets
        else:
            bucket_size = (len(indices) // self.buckets) + 1

        # Create buckets
        buckets = [
            indices[i:i + bucket_size]
            for i in range(0, len(indices), bucket_size)
        ]
        print("bucket", buckets)

        # Shuffle buckets
        random.shuffle(buckets)
        
        for bucket in buckets:
            # Shuffle samples within the bucket
            random.shuffle(bucket)
            
            # Yield batches from the bucket
            for i in range(0, len(bucket), self.batch_size):
                batch = bucket[i:i + self.batch_size]
                if len(batch) == self.batch_size or not self.drop_last:
                    print("length of each batch should be 8", len(batch))
                    print("batch post bucketing", batch)
                    yield batch

    def __len__(self):
        if self.drop_last:
            length = len(self.sampler) // self.batch_size
            # print("total number of batches (drop_last=true)", length)
            return length
        else:
            length = (len(self.sampler) + self.batch_size - 1) // self.batch_size
            # print("total number of batches (drop_last=False)", length)
            return length
</code></pre>
<p>this is a good answer: <a href="https://github.com/Lightning-AI/pytorch-lightning/discussions/7573" class="inline-onebox" rel="noopener nofollow ugc">Error while using custom DistributedSampler · Lightning-AI/pytorch-lightning · Discussion #7573 · GitHub</a> but I still don’t understand where to initialize the process group?</p>
<p>I understand that I need to follow the rules mentioned here (as suggested by <a class="mention" href="/u/ptrblck">@ptrblck</a> in one of the discussions, but what should I initialize for these parameters?</p>
<p>os.environ[‘MASTER_ADDR’] = ‘localhost’<br>
os.environ[‘MASTER_PORT’] = ‘12355’</p>
<p>dist.init_process_group(“nccl”, rank=rank, world_size=world_size)</p>
<p>Thanks in advance!</p>