Importing `torchdata.stateful_dataloader` hides `torch` RandomSampler and BatchSampler about data HOT 8 CLOSED

Thanks for flagging @byi8220 ! We put this patch in to give default samplers the faster behaviour by default. I'll look at pulling this out to a non-monkey patch solution

from data.

cc @gokulavasan see the original issue in HF Accelerate on assumptions on number of calls to iter: huggingface/accelerate#2894

from data.

I see 3 approaches currently to fix this,

1. we fork the dataloader init code, this is maybe the best way but introduces more forked code
2. we repro the sampler/batch_sampler set up in init before we call super().init(...), this seems hacky and not significantly better than 1)
3. we check isinstance after super().init() and replace, but that heading down old-school Lightning territory and surely will lead to head-scratching and further problems down the road.

I'm going to go with 1)

from data.

Thanks for resolving this! This appears to have fixed the breaking tests in accelerate, and the repro above shows there is no more monkey patching going on.

However, it might be worth mentioning that if one passes in an existing non-stateful BatchSampler to the StatefulDataLoader constructor, then the constructed StatefulDataLoader will use the provided sampler. Only pointing this out since it wasn't clear to me if this is the intended behavior.

Repro output with nightly 0.7.1.dev20240703+cpu (important stuff in green diff):

--------------------------------------------------------------------------------
BatchSampler before importing `stateful_dataloader`: <class 'torch.utils.data.sampler.BatchSampler'>
RandomSampler before importing `stateful_dataloader`: <class 'torch.utils.data.sampler.RandomSampler'>
--------------------------------------------------------------------------------
type(non_stateful_dataloader.batch_sampler): <class 'torch.utils.data.sampler.BatchSampler'>
type(non_stateful_dataloader.batch_sampler.sampler): <class '__main__.MyRandomSamplerWrapper'>
type(non_stateful_dataloader.batch_sampler.sampler.original_sampler): <class 'torch.utils.data.sampler.RandomSampler'>
--------------------------------------------------------------------------------
+ BatchSampler after importing `stateful_dataloader`: <class 'torch.utils.data.sampler.BatchSampler'>
+ RandomSampler after importing `stateful_dataloader`: <class 'torch.utils.data.sampler.RandomSampler'>
--------------------------------------------------------------------------------
+ # Even after the fix, stateful_dataloader has a non-stateful sampler
type(stateful_dataloader.batch_sampler): <class 'torch.utils.data.sampler.BatchSampler'>
type(stateful_dataloader.batch_sampler.sampler): <class '__main__.MyRandomSamplerWrapper'>
type(stateful_dataloader.batch_sampler.sampler.original_sampler): <class 'torch.utils.data.sampler.RandomSampler'>
--------------------------------------------------------------------------------
type(stateful_dataloader_2.batch_sampler): <class 'torchdata.stateful_dataloader.sampler.BatchSampler'>
type(stateful_dataloader_2.batch_sampler.sampler): <class 'torch.utils.data.sampler.SequentialSampler'>
--------------------------------------------------------------------------------

from data.

@byi8220 thanks for the details! If I understand correctly: user has explicitly passed in a non-stateful Batch Sampler to DataLoader constructor? In this case, I think this is correct, we should respect what the user has passed in and not try to override it under the hood. I've seen code that does this before in other libraries and it can cause some nasty surprises and un-intuitive behaviour, and can be very hard to debug.

from data.

user has explicitly passed in a non-stateful Batch Sampler to DataLoader constructor?

Yes

In this case, I think this is correct, we should respect what the user has passed in and not try to override it under the hood.

This makes sense, I was just curious if this would cause an issue with saving or loading a state dict into this dataloader

from data.

That's a great call, yes it might cause an issue if checkpoint was saved before and then loaded with the new code-change, but it might also just try to fast-forward, I haven't attempted. We'll be cutting a release this month, so once that's out it should be easier to manage these types of changes.

For the case where users are explicitly passing in a BatchSampler, they can import it from torchdata.stateful_dataloader.samplers instead of from torch.utils.data

from data.

I'm not super familiar with the stateful_dataloader code, but I agree with your hunch that it should fall back to fast forwarding.

Still, it might be worth adding a warning? Worst case, it might help catch a hard to find bug. But best case, isn't fast forwarding an iterator costly?

from data.