Allowed custom BatchSamplers when instantiated in *_dataloader hook

[otaj]

What does this PR do?

Allowed custom BatchSamplers when instantiated in *_dataloader hook

Fixes #13007, #11807

Does your PR introduce any breaking changes? If yes, please list them.

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, 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)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:

• 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

Did you have fun?

Make sure you had fun coding 🙃

cc @Borda @justusschock @awaelchli @ninginthecloud @rohitgr7 @otaj @akihironitta

[carmocca]

This is getting really complex 😂

Posting a partial review

[otaj]

Thank you for the review! I addressed all your concerns 👍

Btw, regarding the complexity, it surprisingly wasn't too bad for me, since it's basically just logic from dataloader handling made generic

[carmocca]

Great! Minor comments

[otaj]

@carmocca, @awaelchli I finally managed to properly debug the cause of the IPU regression. The problem was, that I was passing all default kwargs from torch.utils.data.DataLoader to poptorch.DataLoader which didn't play nicely with it. I don't think there's a substantial change from when you reviewed it (just split __pl_saved_kwargs into __pl_saved_kwargs and __pl_saved_default_kwargs).

[psridhar-asapp]

Hey @otaj I am confused on how to update my sampler to fit this API.

We are dealing with audio data and we want to batch samples by sequence length. In order to do this, we have shape files that are CSVs mapping between each audio sample and its length. Our BatchSampler reads these files, sorts utterances by length, and then constructs fixed batches of similarly long utterances such that all batches are similar in size.

I am not sure what to move to a Sampler here. We could move the loading of this shapes file to the sampler backed by a new dataset, but we do not sample from the shapes file to build batches, we consume the entire file at once.

I think part of my doubt is a lack of understanding on what exactly you do to the Sampler passed into the BatchSampler. I could move the loading of the shapes file inside a Dataset, pass in a Sampler to this dataset as a hacky way of giving the dataset to the BatchSampler, and totally ignore the Sampler in my BatchSampler logic while just reading from its member Dataset, but I imagine this will not solve the problem I'm having.

[otaj]

Hi @psridhar-asapp. First things first, this doesn't really add new API. It keeps the functionality as before, just now it allows your custom BatchSampler to have also more arguments, that don't appear in a signature of regular BatchSampler.

I remember we spoke on this topic in some other issue (I think it was #11807). I don't think we have a good resolution to it yet, honestly.

We don't really do much to your sampler. Everything important regarding Sampler and DataLoader is happening in this method (and of course, all the other methods it calls). https://github.com/Lightning-AI/lightning/blob/caaf35689c585fba1bb33b70243e7ddd2cf5d13c/src/pytorch_lightning/trainer/connectors/data_connector.py#L238

Currently, it sets shuffling if it was set to None, and we're training, wraps it in a distributed wrapper and does a bit more of a magic in _update_dataloader, which is just juggling arguments. The most important bit of call tree in _update_dataloader is honestly this function https://github.com/Lightning-AI/lightning/blob/caaf35689c585fba1bb33b70243e7ddd2cf5d13c/src/pytorch_lightning/utilities/data.py#L315

That is where we wrap either BatchSampler or Sampler in structures used for prediction or fault tolerance. I think (and think is the operative word here) it might be possible to remove the Sampler from being passed around and you could be happy with your BatchSampler only solution, but I don't really have a clear vision about that yet and it will require non-trivial engineering work.

If you don't care about either of these things, you can just expose argument sampler in your custom BatchSampler, (which needs to be a proper subclass of it and instantiated in *_dataloader hook) and don't use that sampler at all. I'm sorry, I don't have a better solution for you