Repository: Lightning-AI/pytorch-lightning
Pull request: Better error message and instructions for handling custom batch samplers
## What does this PR do?

Fixes #18391
Fixes #17793
Partially addresses #18023

<!-- Does your PR introduce any breaking changes? If yes, please list them. -->

<details>
  <summary><b>Before submitting</b></summary>

- Was this **discussed/agreed** via a GitHub issue? (not for typos and docs)
- [ ] Did you read the [contributor guideline](https://github.com/Lightning-AI/lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [ ] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- Did you make sure to **update the documentation** with your changes? (if necessary)
- Did you write any **new necessary tests**? (not for typos and docs)
- [ ] Did you verify new and **existing tests pass** locally with your changes?
- Did you list all the **breaking changes** introduced by this pull request?
- Did you **update the CHANGELOG**? (not for typos, docs, test updates, or minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

</details>

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/Review-guidelines). In short, see the following bullet-list:

<details>
  <summary>Reviewer checklist</summary>

- [ ] Is this pull request ready for review? (if not, please submit in draft mode)
- [ ] Check that all items from **Before submitting** are resolved
- [ ] Make sure the title is self-explanatory and the description concisely explains the PR
- [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified

</details>

<!--

Did you have fun?

Make sure you had fun coding 🙃

-->


cc @borda @justusschock @awaelchli @carmocca

Problem statement:
Custom batch sampler leads to an error
### Bug description

I am using a custom batch sampler with variable batch sizes. Lightning tries to read `batch_sampler.batch_size` parameter and fails

### What version are you seeing the problem on?

v2.0

### How to reproduce the bug

```python
class BatchSamplerForPaddingHaters(Sampler):
    def __init__(self, lengths, sampler):
        self.sampler = sampler
        self.lengths = lengths
        self.batch_sizes = []

        cur_len = lengths[0]
        cur_bs = 0
        for l in lengths:
            if l != cur_len:
                self.batch_sizes.append(cur_bs)
                cur_bs = 0
            cur_bs += 1
            cur_len = l

    def __iter__(self):
        sampler_iter = iter(self.sampler)
        for bs in self.batch_sizes:
            yield [next(sampler_iter) for _ in range(bs)]

    def __len__(self):
        return len(self.batch_sizes)

class MyDataModule(L.LightningDataModule):
    def train_dataloader(self):
        train_sampler = BatchSamplerForPaddingHaters(self.train_data.lengths, 
                                                          SequentialSampler(self.train_data))
        return DataLoader(self.train_data, num_workers=12, 
                          prefetch_factor=256, batch_sampler=train_sampler)

    def val_dataloader(self):
        test_sampler = BatchSamplerForPaddingHaters(self.test_data.lengths,
                                                         SequentialSampler(self.test_data))
        return DataLoader(self.test_data, num_workers=2, 
                          prefetch_factor=16, batch_sampler=test_sampler)

trainer_opts = {
        'accelerator': 'gpu',
        'log_every_n_steps': 1,
        'gradient_clip_val': GRAD_CLIP_VAL,
        'logger': loggers.WandbLogger(project='icat', log_model=True),
        'callbacks': [callbacks.ModelCheckpoint(monitor='val_loss',mode='min')],
        'devices': 2,
        'num_nodes': 1,
        'strategy': 'ddp'
    }

  trainer = L.Trainer(**trainer_opts)
  trainer.fit(model, data)
```


### Error messages and logs

```
Traceback (most recent call last):
  File "/home/mcs001/20194474/code/mimicseq/learn.py", line 93, in <module>
    train_icat()
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/click/core.py", line 1130, in __call__
    return self.main(*args, **kwargs)
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/click/core.py", line 1055, in main
    rv = self.invoke(ctx)
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/click/core.py", line 1404, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/click/core.py", line 760, in invoke
    return __callback(*args, **kwargs)
  File "/cm/shared/apps/Anaconda/2022.10/pth3.9/lib/python3.9/site-packages/tenacity/__init__.py", line 324, in wrapped_f
    return self(f, *args, **kw)
  File "/cm/shared/apps/Anaconda/2022.10/pth3.9/lib/python3.9/site-packages/tenacity/__init__.py", line 404, in __call__
    do = self.iter(retry_state=retry_state)
  File "/cm/shared/apps/Anaconda/2022.10/pth3.9/lib/python3.9/site-packages/tenacity/__init__.py", line 349, in iter
    return fut.result()
  File "/cm/shared/apps/Anaconda/2022.10/pth3.9/lib/python3.9/concurrent/futures/_base.py", line 439, in result
    return self.__get_result()
  File "/cm/shared/apps/Anaconda/2022.10/pth3.9/lib/python3.9/concurrent/futures/_base.py", line 391, in __get_result
    raise self._exception
  File "/cm/shared/apps/Anaconda/2022.10/pth3.9/lib/python3.9/site-packages/tenacity/__init__.py", line 407, in __call__
    result = fn(*args, **kwargs)
  File "/home/mcs001/20194474/code/mimicseq/learn.py", line 88, in train_icat
    trainer.fit(model, data)
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/lightning/pytorch/trainer/trainer.py", line 531, in fit
    call._call_and_handle_interrupt(
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/lightning/pytorch/trainer/call.py", line 41, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/lightning/pytorch/strategies/launchers/subprocess_script.py", line 91, in launch
    return function(*args, **kwargs)
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/lightning/pytorch/trainer/trainer.py", line 570, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/lightning/pytorch/trainer/trainer.py", line 975, in _run
    results = self._run_stage()
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/lightning/pytorch/trainer/trainer.py", line 1016, in _run_stage
    self._run_sanity_check()
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/lightning/pytorch/trainer/trainer.py", line 1045, in _run_sanity_check
    val_loop.run()
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/lightning/pytorch/loops/utilities.py", line 177, in _decorator
    return loop_run(self, *args, **kwargs)
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 98, in run
    self.setup_data()
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 169, in setup_data
    dl = _process_dataloader(trainer, dl)
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/lightning/pytorch/trainer/connectors/data_connector.py", line 489, in _process_dataloader
    dataloader = trainer._data_connector._prepare_dataloader(dataloader, shuffle=is_shuffled, mode=stage)
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/lightning/pytorch/trainer/connectors/data_connector.py", line 200, in _prepare_dataloader
    return _update_dataloader(dataloader, sampler, mode=mode)
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/lightning/pytorch/utilities/data.py", line 132, in _update_dataloader
    dl_args, dl_kwargs = _get_dataloader_init_args_and_kwargs(dataloader, sampler, mode)
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/lightning/pytorch/utilities/data.py", line 192, in _get_dataloader_init_args_and_kwargs
    dl_kwargs.update(_dataloader_init_kwargs_resolve_sampler(dataloader, sampler, mode, disallow_batch_sampler))
  File "/home/mcs001/20194474/.local/lib/python3.9/site-packages/lightning/pytorch/utilities/data.py", line 296, in _dataloader_init_kwargs_resolve_sampler
    batch_size=batch_sampler.batch_size,
AttributeError: 'BatchSamplerForPaddingHaters' object has no attribute 'batch_size'
```


### Environment

<details>
  <summary>Current environment</summary>

```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow): Trainer, LightiningDataModule
#- PyTorch Lightning Version (e.g., 1.5.0): 2.0.3
#- Lightning App Version (e.g., 0.5.2):2.0.3
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9): 3.9
#- OS (e.g., Linux):  Linux-3.10.0-1160.83.1.el7.x86_64-x86_64-with-glibc2.17
#- CUDA/cuDNN version: 11.8.0
#- GPU models and configuration: 2 Tesla V100-PCIE-16GB
#- How you installed Lightning(`conda`, `pip`, source): pip
#- Running environment of LightningApp (e.g. local, cloud): slurm cluster
```

</details>


### More info

_No response_

cc @justusschock @awaelchli

How can I use a custom distributed batch sampler?
### Bug description

i want to use custom batch sampler like this 
``` python3
class DistributedBucketSampler(torch.utils.data.distributed.DistributedSampler):
    """
    Maintain similar input lengths in a batch.
    Length groups are specified by boundaries.
    Ex) boundaries = [b1, b2, b3] -> any batch is included either {x | b1 < length(x) <=b2} or {x | b2 < length(x) <= b3}.

    It removes samples which are not included in the boundaries.
    Ex) boundaries = [b1, b2, b3] -> any x s.t. length(x) <= b1 or length(x) > b3 are discarded.
    """

    def __init__(
        self,
        dataset,
        batch_size,
        boundaries,
        num_replicas=None,
        rank=None,
        shuffle=True,
        expand_new_speaker=False,
        new_speaker_ratio=None,
        speaker_uniform_batch=False,
        drop_last=False,
    ):
        if not shuffle and expand_new_speaker:
            raise NotImplementedError("expand_new_speaker is supported only in case of shuffle=True")

        super().__init__(dataset, num_replicas=num_replicas, rank=rank, shuffle=shuffle, drop_last=drop_last)
        self.lengths = dataset.lengths
        self.expand_new_speaker = expand_new_speaker
        self.batch_size = batch_size
        self.boundaries = boundaries
        self.speaker_uniform_batch = speaker_uniform_batch

        if expand_new_speaker:
            self.is_new_speaker = dataset.is_new_speaker
            assert new_speaker_ratio <= 1
            self.new_speaker_ratio = new_speaker_ratio
            self.buckets, self.num_samples_per_bucket, self.prev_buckets = self._create_buckets()
        else:
            self.buckets, self.num_samples_per_bucket = self._create_buckets()

        self.total_size = sum(self.num_samples_per_bucket)
        self.num_samples = self.total_size // self.num_replicas

        if speaker_uniform_batch:
            # it doesn't work well if expand_new_speaker==True
            self.speaker_names = [one_item["speaker_name"] for one_item in dataset.filelist]

    def _create_buckets(self):
        buckets = [[] for _ in range(len(self.boundaries) - 1)]
        if self.expand_new_speaker:
            prev_buckets = [[] for _ in range(len(self.boundaries) - 1)]

        n_out_of_boundaries = 0
        for idx_data in range(len(self.lengths)):
            length = self.lengths[idx_data]
            idx_bucket = self._bisect(length)
            if idx_bucket != -1:
                if self.expand_new_speaker:
                    if self.is_new_speaker[idx_data]:
                        buckets[idx_bucket].append(idx_data)
                    else:
                        prev_buckets[idx_bucket].append(idx_data)
                else:
                    buckets[idx_bucket].append(idx_data)
            else:
                n_out_of_boundaries += 1
                # print(f"Data is out of boundaries. wav_size // (2 * hop_length): {length}")
        print(
            f"Sampler filtering result (by wav length): {len(self.lengths)} -> {len(self.lengths) - n_out_of_boundaries}"
        )

        for idx_bucket in range(len(buckets) - 1, 0, -1):
            if len(buckets[idx_bucket]) == 0:
                buckets.pop(idx_bucket)
                if self.expand_new_speaker:
                    prev_buckets.pop(idx_bucket)

        num_samples_per_bucket = []
        for idx_bucket in range(len(buckets)):
            if self.expand_new_speaker:
                len_bucket = int(len(buckets[idx_bucket]) / self.new_speaker_ratio)
            else:
                len_bucket = len(buckets[idx_bucket])
            total_batch_size = self.num_replicas * self.batch_size  # total_batch_size = n_gpus * batch_size
            remainder = (total_batch_size - (len_bucket % total_batch_size)) % total_batch_size
            num_samples_per_bucket.append(len_bucket + remainder)  # num_samples_per_bucket을 total_batch_size의 배수로 맞춤

        if self.expand_new_speaker:
            return buckets, num_samples_per_bucket, prev_buckets
        else:
            return buckets, num_samples_per_bucket

    def __iter__(self):
        # deterministically shuffle based on epoch
        g = torch.Generator()
        g.manual_seed(self.epoch)

        buckets = copy.deepcopy(self.buckets)
        if self.expand_new_speaker:
            for idx_bucket in range(len(buckets)):
                n_sample_from_pre = int(len(buckets[idx_bucket]) * (1 / self.new_speaker_ratio - 1))
                sampled_indices = torch.randperm(len(self.prev_buckets[idx_bucket]), generator=g)[:n_sample_from_pre]
                sample_from_pre = [self.prev_buckets[idx_bucket][i] for i in sampled_indices]
                if len(sampled_indices) < n_sample_from_pre:
                    assert len(sampled_indices) != 0, "prev_bucket must not be empty."
                    sample_from_pre = sample_from_pre * (n_sample_from_pre // len(sampled_indices) + 1)
                    sample_from_pre = sample_from_pre[:n_sample_from_pre]
                buckets[idx_bucket].extend(sample_from_pre)

        indices = []
        if self.shuffle:
            for bucket in buckets:
                indices.append(torch.randperm(len(bucket), generator=g).tolist())
        else:
            for bucket in buckets:
                indices.append(list(range(len(bucket))))

        batches = []
        for i in range(len(buckets)):
            bucket = buckets[i]
            len_bucket = len(bucket)
            ids_bucket = indices[i]
            num_samples_bucket = self.num_samples_per_bucket[i]

            # speaker uniform batch
            if self.speaker_uniform_batch:
                speaker_idx_dict = defaultdict(list)
                for id_bucket in ids_bucket:
                    speaker_idx_dict[self.speaker_names[bucket[id_bucket]]].append(id_bucket)
                for speaker_name, idx_list in list(speaker_idx_dict.items()):
                    if len(idx_list) < self.num_replicas:
                        del speaker_idx_dict[speaker_name]

                new_ids_bucket = []
                while len(speaker_idx_dict.keys()) >= self.batch_size:
                    sampled_speaker_names = random.sample(speaker_idx_dict.keys(), self.batch_size)
                    for sampled_speaker_name in sampled_speaker_names:
                        new_ids_bucket.extend(speaker_idx_dict[sampled_speaker_name][: self.num_replicas])
                        speaker_idx_dict[sampled_speaker_name] = speaker_idx_dict[sampled_speaker_name][
                            self.num_replicas :
                        ]
                        if len(speaker_idx_dict[sampled_speaker_name]) < self.num_replicas:
                            del speaker_idx_dict[sampled_speaker_name]
                ids_bucket = new_ids_bucket

            # not speaker uniform batch - add extra samples to make it evenly divisible
            else:
                remainder = num_samples_bucket - len_bucket
                ids_bucket = (
                    ids_bucket + ids_bucket * (remainder // len_bucket) + ids_bucket[: (remainder % len_bucket)]
                )

            # subsample
            assert len(ids_bucket) % (self.batch_size * self.num_replicas) == 0
            ids_bucket = ids_bucket[self.rank :: self.num_replicas]

            # batching
            for j in range(len(ids_bucket) // self.batch_size):
                batch = [bucket[idx] for idx in ids_bucket[j * self.batch_size : (j + 1) * self.batch_size]]
                batches.append(batch)

        if self.shuffle:
            batch_ids = torch.randperm(len(batches), generator=g).tolist()
            batches = [batches[i] for i in batch_ids]
        self.batches = batches

        assert self.speaker_uniform_batch or len(self.batches) * self.batch_size == self.num_samples
        return iter(self.batches)

    def _bisect(self, x, lo=0, hi=None):
        """data length를 input으로 받아 해당 data가 어떤 boundary에 속하는지 return"""
        if hi is None:
            hi = len(self.boundaries) - 1

        if hi > lo:
            mid = (hi + lo) // 2
            if self.boundaries[mid] < x and x <= self.boundaries[mid + 1]:
                return mid
            elif x <= self.boundaries[mid]:
                return self._bisect(x, lo, mid)
            else:
                return self._bisect(x, mid + 1, hi)
        else:
            return -1

    def __len__(self):
        return self.num_samples // self.batch_size
```

but i think `fabric.utilities.data._dataloader_init_kwargs_resolve_sampler`'s 215 lines that can not used additional argument(like `boundaries`) and is that fully considered on DistributedBatchSampler? (why after re init batchsampler, batch size is 1 on return dict??)

please help me!

### What version are you seeing the problem on?

v2.0, master

### How to reproduce the bug

_No response_

### Error messages and logs

```
# Error messages and logs here please
```
`TypeError: DistributedBucketSampler.__init__() missing 1 required positional argument: 'boundaries'`

### Environment

<details>
  <summary>Current environment</summary>

```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):
```

</details>


### More info

_No response_

cc @justusschock @awaelchli

Hints from issue discussion:
Similar issue in 1.8.4 when I try to use a custom batch sampler in `predict_dataloader()`.
Interesting! 
Did you fix this issue @vadim0x60  ?
Can I know what model your training it ? (Application)

I think No one is working on it !.  

Can I work on this issue with LSTM model ?

@awaelchli  should I open PR for this issue  (user-specific issue)? 



@ManoBharathi93 
I am afraid I don't understand the source code of lightning well enough to fix it myself at this point.
The model is `torch.TransformerEncoder`
Given that @ZeguanXiao ran into the same, this is not a user-specific issue.
@vadim0x60 can you share the full code that you implemented  ?
I believe lightning doesn't support *any* batch sampler with non-standard constructor arguments, my particular use case is probably a distraction. You can check it [here](https://github.com/vadim0x60/mimicseq) if you insist.
Found it! .. You didn't inherit from base class to `MyDataModule` class. You missed constructor which requires `train_data`, `test_data`, and `batch_size` arguments.

I have trained on LSTM it's work fine for me using batch sampler.
```import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, Sampler, SequentialSampler, DataLoader
from typing import List, Tuple
import pytorch_lightning as L
from pytorch_lightning import loggers, callbacks

GRAD_CLIP_VAL = 1.0

class MyDataset(Dataset):
    def __init__(self, data: List[Tuple[List[int], int]]):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        x, y = self.data[idx]
        return torch.LongTensor(x), torch.LongTensor([y])

class BatchSamplerForPaddingHaters(Sampler):
    def __init__(self, lengths, sampler):
        self.sampler = sampler
        self.lengths = lengths
        self.batch_sizes = []

        cur_len = lengths[0]
        cur_bs = 0
        for l in lengths:
            if l != cur_len:
                self.batch_sizes.append(cur_bs)
                cur_bs = 0
            cur_bs += 1
            cur_len = l

    def __iter__(self):
        sampler_iter = iter(self.sampler)
        for bs in self.batch_sizes:
            yield [next(sampler_iter) for _ in range(bs)]

    def __len__(self):
        return len(self.batch_sizes)

class MyDataModule(L.LightningDataModule):
    def __init__(self, train_data, test_data, batch_size=32):
        super().__init__()
        self.train_data = train_data
        self.test_data = test_data
        self.batch_size = batch_size

    def prepare_data(self):
        pass

    def setup(self, stage=None):
        pass

    def train_dataloader(self):
        train_sampler = BatchSamplerForPaddingHaters(self.train_data.lengths, 
                                                          SequentialSampler(self.train_data))
        return DataLoader(self.train_data, num_workers=12, 
                          prefetch_factor=256, batch_sampler=train_sampler)

    def val_dataloader(self):
        test_sampler = BatchSamplerForPaddingHaters(self.test_data.lengths,
                                                         SequentialSampler(self.test_data))
        return DataLoader(self.test_data, num_workers=2, 
                          prefetch_factor=16, batch_sampler=test_sampler)


class MyModel(L.LightningModule):
    def __init__(self, vocab_size: int, embedding_dim: int, hidden_dim: int, num_classes: int):
        super().__init__()
        self.embeddings = torch.nn.Embedding(vocab_size, embedding_dim)
        self.lstm = torch.nn.LSTM(embedding_dim, hidden_dim, num_layers=2, batch_first=True)
        self.fc = torch.nn.Linear(hidden_dim, num_classes)

    def forward(self, inputs):
        x = self.embeddings(inputs)
        lstm_out, _ = self.lstm(x)
        out = self.fc(lstm_out[:, -1, :])
        return out

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self.forward(x)
        loss = F.cross_entropy(logits, y.squeeze())
        self.log('train_loss', loss)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self.forward(x)
        loss = F.cross_entropy(logits, y.squeeze())
        self.log('val_loss', loss)
        preds = torch.argmax(logits, dim=1)
        acc = torch.sum(preds == y.squeeze()).float() / y.size(0)
        self.log('val_acc', acc)
        return {'val_loss': loss, 'val_acc': acc}

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters())
        return optimizer

train_data = MyDataset([(list(range(i)), i % 3) for i in range(1, 1001)])
test_data = MyDataset([(list(range(i)), i % 3) for i in range(1001, 2001)])

data = MyDataModule(train_data, test_data)
model = MyModel(vocab_size=1000, embedding_dim=32, hidden_dim=64, num_classes=3)

trainer_opts = {
        'accelerator': 'gpu',
        'log_every_n_steps': 1,
        'gradient_clip_val': GRAD_CLIP_VAL,
        'logger': loggers.WandbLogger(project='icat', log_model=True),
        'callbacks': [callbacks.ModelCheckpoint(monitor='val_loss',mode='min')],
        'devices': 2,
        'num_nodes': 1,
        'distributed_backend': 'ddp'
    }

trainer = L.Trainer(**trainer_opts)
trainer.fit(model, data)
```
It works fine..

@vadim0x60 
Correct way of Implementing `MyDataModule` : 
```
class MyDataModule(L.LightningDataModule):
def __init__(self, train_data, test_data, batch_size=32):
        super().__init__()
        self.train_data = train_data
        self.test_data = test_data
        self.batch_size = batch_size
    def train_dataloader(self):
        train_sampler = BatchSamplerForPaddingHaters(self.train_data.lengths, 
                                                          SequentialSampler(self.train_data))
        return DataLoader(self.train_data, num_workers=12, 
                          prefetch_factor=256, batch_sampler=train_sampler)

    def val_dataloader(self):
        test_sampler = BatchSamplerForPaddingHaters(self.test_data.lengths,
                                                         SequentialSampler(self.test_data))
        return DataLoader(self.test_data, num_workers=2, 
                          prefetch_factor=16, batch_sampler=test_sampler)
```
If it fixes your problem kindly close this issue ! 
Thank you

I do inherit from `LightningDataModule`. And I actually have `super().__init__()` in my code (though not in the simplified snippet I posted here), it doesn't help. I cannot include `batch_size` in the constructor, because my custom batch sampler has a variable batch size.
Facing the same issue. I also have a dynamic batch size 😟 
Hi @YooSungHyun 
Could you please post the full error stack trace?

Also it would be helpful if you could share the dataset code too. This sampler seems to require a specific dataset that has "lengths" defined, and it wasn't possible for me to run this in an example since I don't know what the expected format is. 
If you are using the Lightning Trainer, you can set
```py
trainer = Trainer(use_distributed_sampler=False) 
```
to [disable Lightning attempting to inject the distributed sampler](https://lightning.ai/docs/pytorch/stable/common/trainer.html#use-distributed-sampler).

And likewise, in Fabric [it would be done in setup](https://lightning.ai/docs/fabric/stable/api/fabric_methods.html#setup-dataloaders):
```py
fabric.setup_dataloaders(..., use_distributed_sampler=False)
```

Commit message: [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Patch to review:
diff --git a/src/lightning/fabric/CHANGELOG.md b/src/lightning/fabric/CHANGELOG.md
index 8ded9d4cb160e..ae35727286dff 100644
--- a/src/lightning/fabric/CHANGELOG.md
+++ b/src/lightning/fabric/CHANGELOG.md
@@ -118,6 +118,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added validation of user input for `devices` and `num_nodes` when running with `SLURM` or `TorchElastic` ([#18292](https://github.com/Lightning-AI/lightning/pull/18292))
 
 
+- Improved the error messaging and instructions when handling custom batch samplers in distributed settings ([#18402](https://github.com/Lightning-AI/lightning/pull/18402))
+
+
 ### Changed
 
 - Allow using iterable-style datasets with TPUs ([#17331](https://github.com/Lightning-AI/lightning/pull/17331))
diff --git a/src/lightning/fabric/utilities/data.py b/src/lightning/fabric/utilities/data.py
index de11d8a303c32..001c8bc571d28 100644
--- a/src/lightning/fabric/utilities/data.py
+++ b/src/lightning/fabric/utilities/data.py
@@ -180,6 +180,7 @@ def _dataloader_init_kwargs_resolve_sampler(
     if batch_sampler is not None and type(batch_sampler) is not BatchSampler:
         batch_sampler_cls = type(batch_sampler)
         if hasattr(batch_sampler, "__pl_saved_args"):
+            # This is a PyTorch `BatchSampler` for which we captured the init args
             args = batch_sampler.__pl_saved_args
             kwargs = batch_sampler.__pl_saved_kwargs
             default_kwargs = batch_sampler.__pl_saved_default_kwargs
@@ -196,7 +197,8 @@ def _dataloader_init_kwargs_resolve_sampler(
                 )
 
             batch_sampler = _reinstantiate_wrapped_cls(batch_sampler, *args, **kwargs)
-        else:
+        elif isinstance(batch_sampler, BatchSampler):
+            # This is a PyTorch `BatchSampler` for which we could not capture the init args
             try:
                 batch_sampler = batch_sampler_cls(
                     sampler,
@@ -218,6 +220,14 @@ def _dataloader_init_kwargs_resolve_sampler(
                     "To mitigate this, either follow the API of `BatchSampler` or instantiate "
                     "your custom batch sampler inside `*_dataloader` hooks of your module."
                 ) from ex
+        else:
+            # The sampler is not a PyTorch `BatchSampler`, we don't know how to inject a custom sampler
+            raise TypeError(
+                " Lightning can't inject a (distributed) sampler into your batch sampler, because it doesn't subclass"
+                " PyTorch's `BatchSampler`. To mitigate this, either follow the API of `BatchSampler` or set"
+                " `.setup_dataloaders(..., use_distributed_sampler=False)`. If you choose the latter, you will be "
+                " responsible for handling the distributed sampling within your batch sampler."
+            )
 
         return {
             "sampler": None,
diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md
index 3afc2a79b0484..18dc1ddcda65d 100644
--- a/src/lightning/pytorch/CHANGELOG.md
+++ b/src/lightning/pytorch/CHANGELOG.md
@@ -115,6 +115,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Made the text delimiter in the rich progress bar configurable ([#18372](https://github.com/Lightning-AI/lightning/pull/18372))
 
 
+- Improved the error messaging and instructions when handling custom batch samplers in distributed settings ([#18402](https://github.com/Lightning-AI/lightning/pull/18402))
+
+
+
 ### Changed
 
 - Removed the limitation to call `self.trainer.model.parameters()` in `LightningModule.configure_optimizers()` ([#17309](https://github.com/Lightning-AI/lightning/pull/17309))
diff --git a/src/lightning/pytorch/overrides/distributed.py b/src/lightning/pytorch/overrides/distributed.py
index 9b86b5db3055d..3ecc58f56b9fa 100644
--- a/src/lightning/pytorch/overrides/distributed.py
+++ b/src/lightning/pytorch/overrides/distributed.py
@@ -17,11 +17,13 @@
 import torch
 from torch import Tensor
 from torch.nn.parallel.distributed import DistributedDataParallel
-from torch.utils.data import BatchSampler, DistributedSampler, Sampler
+from torch.utils.data import DistributedSampler, Sampler
+from typing_extensions import Self
 
 from lightning.fabric.utilities.distributed import _DatasetSamplerWrapper
 from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_1_12
 from lightning.pytorch.utilities.rank_zero import rank_zero_debug, rank_zero_info
+from lightning.pytorch.utilities.types import _SizedIterable
 
 
 def _find_tensors(
@@ -245,10 +247,10 @@ def __iter__(self) -> Iterator:
         return (self.dataset[index] for index in super().__iter__())
 
 
-class _IndexBatchSamplerWrapper(BatchSampler):
+class _IndexBatchSamplerWrapper:
     """This class is used to wrap a :class:`torch.utils.data.BatchSampler` and capture its indices."""
 
-    def __init__(self, batch_sampler: BatchSampler) -> None:
+    def __init__(self, batch_sampler: _SizedIterable) -> None:
         # do not call super().__init__() on purpose
         self.seen_batch_indices: List[List[int]] = []
 
@@ -266,7 +268,7 @@ def __next__(self) -> List[int]:
         self.seen_batch_indices.append(batch)
         return batch
 
-    def __iter__(self) -> Iterator[List[int]]:
+    def __iter__(self) -> Self:
         self.seen_batch_indices = []
         self._iterator = iter(self._batch_sampler)
         return self
diff --git a/src/lightning/pytorch/utilities/data.py b/src/lightning/pytorch/utilities/data.py
index cbfae478503de..b9978ef38d499 100644
--- a/src/lightning/pytorch/utilities/data.py
+++ b/src/lightning/pytorch/utilities/data.py
@@ -245,6 +245,7 @@ def _dataloader_init_kwargs_resolve_sampler(
 
     if batch_sampler is not None and (batch_sampler_cls is not BatchSampler or is_predicting):
         if hasattr(batch_sampler, "__pl_saved_args"):
+            # This is a PyTorch `BatchSampler` for which we captured the init args
             args = batch_sampler.__pl_saved_args
             kwargs = batch_sampler.__pl_saved_kwargs
             default_kwargs = batch_sampler.__pl_saved_default_kwargs
@@ -273,7 +274,8 @@ def _dataloader_init_kwargs_resolve_sampler(
                 )
 
             batch_sampler = _reinstantiate_wrapped_cls(batch_sampler, *args, **kwargs)
-        else:
+        elif isinstance(batch_sampler, BatchSampler):
+            # This is a PyTorch `BatchSampler` for which we could not capture the init args
             try:
                 batch_sampler = batch_sampler_cls(
                     sampler,
@@ -295,6 +297,15 @@ def _dataloader_init_kwargs_resolve_sampler(
                     "To mitigate this, either follow the API of `BatchSampler` or instantiate "
                     "your custom batch sampler inside `*_dataloader` hooks of your module."
                 ) from ex
+        else:
+            # The sampler is not a PyTorch `BatchSampler`, we don't know how to inject a custom sampler or
+            # how to adjust the `drop_last` value
+            raise TypeError(
+                " Lightning can't inject a (distributed) sampler into your batch sampler, because it doesn't subclass"
+                " PyTorch's `BatchSampler`. To mitigate this, either follow the API of `BatchSampler` or set"
+                " `Trainer(use_distributed_sampler=False)`. If you choose the latter, you will be responsible for"
+                " handling the distributed sampling within your batch sampler."
+            )
 
         if is_predicting:
             batch_sampler = _IndexBatchSamplerWrapper(batch_sampler)
diff --git a/src/lightning/pytorch/utilities/types.py b/src/lightning/pytorch/utilities/types.py
index b394eacb2b335..d933e371d6c5e 100644
--- a/src/lightning/pytorch/utilities/types.py
+++ b/src/lightning/pytorch/utilities/types.py
@@ -18,7 +18,7 @@
 """
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Any, Generator, List, Mapping, Optional, Protocol, runtime_checkable, Type, Union
+from typing import Any, Generator, Iterator, List, Mapping, Optional, Protocol, runtime_checkable, Type, Union
 
 import torch
 from torch import Tensor
@@ -82,3 +82,11 @@ class LRSchedulerConfig:
     monitor: Optional[str] = None
     # enforce that the monitor exists for ReduceLROnPlateau
     strict: bool = True
+
+
+class _SizedIterable(Protocol):
+    def __len__(self) -> int:
+        pass
+
+    def __iter__(self) -> Iterator:
+        pass
diff --git a/tests/tests_fabric/utilities/test_data.py b/tests/tests_fabric/utilities/test_data.py
index 1f1316eb9bd9a..56ad900f12c4c 100644
--- a/tests/tests_fabric/utilities/test_data.py
+++ b/tests/tests_fabric/utilities/test_data.py
@@ -412,7 +412,7 @@ def __init__(self, randomize, *args, **kwargs):
     assert isinstance(new_dataloader, GoodImpl)
 
 
-def test_custom_batch_sampler():
+def test_custom_torch_batch_sampler():
     """This test asserts, that custom `BatchSampler`, with all the arguments, that are required in order to properly
     reinstantiate the class, is invoked properly.
 
@@ -456,6 +456,20 @@ def __init__(self, sampler, extra_arg, drop_last=True):
     assert not hasattr(batch_sampler, "__pl_saved_default_kwargs")
 
 
+def test_custom_batch_sampler():
+    """Test that a custom (non-PyTorch) batch sampler requires the user to set `use_distributed_sampler=False`."""
+
+    class CustomBatchSampler:  # not inheriting from `BatchSampler`
+        def __iter__(self):
+            while True:
+                yield [0, 1, 2, 3]
+
+    batch_sampler = CustomBatchSampler()
+    dataloader = DataLoader(range(100), batch_sampler=batch_sampler)
+    with pytest.raises(TypeError, match=r"can't inject a \(distributed\) sampler into your batch sampler"):
+        _ = _update_dataloader(dataloader, sampler=Mock())
+
+
 def test_custom_batch_sampler_no_sampler():
     """Tests whether appropriate error is raised when the custom `BatchSampler` does not support sampler argument."""
 
diff --git a/tests/tests_pytorch/overrides/test_distributed.py b/tests/tests_pytorch/overrides/test_distributed.py
index c221238e4fbe8..84a4f03d85325 100644
--- a/tests/tests_pytorch/overrides/test_distributed.py
+++ b/tests/tests_pytorch/overrides/test_distributed.py
@@ -78,7 +78,6 @@ def test_index_batch_sampler():
     batch_sampler = BatchSampler(sampler, 3, False)
     index_batch_sampler = _IndexBatchSamplerWrapper(batch_sampler)
 
-    assert isinstance(index_batch_sampler, BatchSampler)
     assert batch_sampler.batch_size == index_batch_sampler.batch_size
     assert batch_sampler.drop_last == index_batch_sampler.drop_last
     assert batch_sampler.sampler is sampler
diff --git a/tests/tests_pytorch/utilities/test_data.py b/tests/tests_pytorch/utilities/test_data.py
index 2093cf3da4c1f..9df42a6225f86 100644
--- a/tests/tests_pytorch/utilities/test_data.py
+++ b/tests/tests_pytorch/utilities/test_data.py
@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from unittest.mock import Mock
 
 import numpy as np
 import pytest
@@ -137,7 +138,7 @@ def __init__(self, randomize, *args, **kwargs):
 
 
 @pytest.mark.parametrize("predicting", [True, False])
-def test_custom_batch_sampler(predicting):
+def test_custom_torch_batch_sampler(predicting):
     """This test asserts, that custom `BatchSampler`, with all the arguments, that are required in order to properly
     reinstantiate the class, is invoked properly.
 
@@ -188,6 +189,20 @@ def __init__(self, sampler, extra_arg, drop_last=True):
     assert not hasattr(batch_sampler, "__pl_saved_default_kwargs")
 
 
+def test_custom_batch_sampler():
+    """Test that a custom (non-PyTorch) batch sampler requires the user to set `use_distributed_sampler=False`."""
+
+    class CustomBatchSampler:  # not inheriting from `BatchSampler`
+        def __iter__(self):
+            while True:
+                yield [0, 1, 2, 3]
+
+    batch_sampler = CustomBatchSampler()
+    dataloader = DataLoader(range(100), batch_sampler=batch_sampler)
+    with pytest.raises(TypeError, match=r"can't inject a \(distributed\) sampler into your batch sampler"):
+        _ = _update_dataloader(dataloader, sampler=Mock())
+
+
 def test_custom_batch_sampler_no_drop_last():
     """Tests whether appropriate warning is raised when the custom `BatchSampler` does not support `drop_last` and we
     want to reset it."""
@@ -213,7 +228,7 @@ def __init__(self, sampler, extra_arg):
 
     # Assert that warning is raised
     with pytest.warns(UserWarning, match="drop_last=False"):
-        dataloader = _update_dataloader(dataloader, dataloader.sampler, mode=RunningStage.PREDICTING)
+        _ = _update_dataloader(dataloader, dataloader.sampler, mode=RunningStage.PREDICTING)
 
 
 def test_custom_batch_sampler_no_sampler():


Code review comments:
diff_hunk: @@ -196,7 +197,8 @@ def _dataloader_init_kwargs_resolve_sampler(
                 )
 
             batch_sampler = _reinstantiate_wrapped_cls(batch_sampler, *args, **kwargs)
-        else:
+        elif isinstance(batch_sampler, BatchSampler):
+            # This is a PyTorch `BatchSampler` for which we could not capture the init args
path: src/lightning/fabric/utilities/data.py
text: @user1:
I think this is a breaking change. Doesn't this change mean that objects that implement the exact batch sampler API but don't subclass it are now unsupported?

@author:
Is that what this else was handling? This code path is completely untested and I struggled looking at the history what this is for. 

@user1:
I believe so, and https://github.com/Lightning-AI/lightning/pull/18402/files#diff-092f8fcfded5ca85ce38376661e86413badaa7454530cc9496f7bdaaa655b9dbL216-L220 suggests this.

The best way to check would be to write a test for this and see if it fails now

@author:
I can't determine where your link should be pointing. Could you tell me which file?

@author:
@user1 I was able to add that test and support the case you were requesting. In summary, we handle these cases

- the sampler is a subclass of BatchSampler: We capture the init args and re-instantiate with updated args
- the sampler is not a subclass of BatchSampler but looks like it conforms to the BatchSampler interface (having batch_size and drop_last defined): We re-instantiate based on these args only and hope for the best
- the sampler is not a subclass of BatchSampler and doesn't look like it conforms to any interface we can handle: This PR raises a new error now with the suggestion to set `use_distributed_sampler=False`

@user1:
Thanks