`batch_sampler.batch_size` is None with deepspeed and `DataLoader(batch_size=None)`

[olegsinavski]

Bug description

Hello,

After Lightning 2.2.0 upgrade we experience a crash when using deepspeed with DataLoader(batch_size=None, ...):

  File ".../pip-ai-experimental_pytorch_lightning/site-packages/pytorch_lightning/trainer/trainer.py", line 579, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File ".../pip-ai-experimental_pytorch_lightning/site-packages/pytorch_lightning/trainer/trainer.py", line 962, in _run
    self.strategy.setup(self)
  File ".../pip-ai-experimental_pytorch_lightning/site-packages/pytorch_lightning/strategies/deepspeed.py", line 335, in setup
    self._init_config_if_needed()
  File ".../pip-ai-experimental_pytorch_lightning/site-packages/pytorch_lightning/strategies/deepspeed.py", line 804, in _init_config_if_needed
    self._format_config()
  File ".../pip-ai-experimental_pytorch_lightning/site-packages/pytorch_lightning/strategies/deepspeed.py", line 813, in _format_config
    self._format_batch_size_and_grad_accum_config()
  File ".../pip-ai-experimental_pytorch_lightning/site-packages/pytorch_lightning/strategies/deepspeed.py", line 906, in _format_batch_size_and_grad_accum_config
    batch_size = self._auto_select_batch_size()
  File ".../pip-ai-experimental_pytorch_lightning/site-packages/pytorch_lightning/strategies/deepspeed.py", line 920, in _auto_select_batch_size
    batch_size = train_dataloader.batch_sampler.batch_size
AttributeError: 'NoneType' object has no attribute 'batch_size'

This was working on 2.1 versions

What version are you seeing the problem on?

v2.2

How to reproduce the bug

class DataModule(LightningDataModule):
    def train_dataloader(self) -> DataLoader:
        # Don't set batch size here, it's done in the datapipe
        return DataLoader(
            self.train_dp, batch_size=None,...
        )

trainer = Trainer(strategy='deepspeed', ...)
trainer.fit(model, DataModule())

Error messages and logs

  File ".../pip-ai-experimental_pytorch_lightning/site-packages/pytorch_lightning/trainer/trainer.py", line 579, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File ".../pip-ai-experimental_pytorch_lightning/site-packages/pytorch_lightning/trainer/trainer.py", line 962, in _run
    self.strategy.setup(self)
  File ".../pip-ai-experimental_pytorch_lightning/site-packages/pytorch_lightning/strategies/deepspeed.py", line 335, in setup
    self._init_config_if_needed()
  File ".../pip-ai-experimental_pytorch_lightning/site-packages/pytorch_lightning/strategies/deepspeed.py", line 804, in _init_config_if_needed
    self._format_config()
  File ".../pip-ai-experimental_pytorch_lightning/site-packages/pytorch_lightning/strategies/deepspeed.py", line 813, in _format_config
    self._format_batch_size_and_grad_accum_config()
  File ".../pip-ai-experimental_pytorch_lightning/site-packages/pytorch_lightning/strategies/deepspeed.py", line 906, in _format_batch_size_and_grad_accum_config
    batch_size = self._auto_select_batch_size()
  File ".../pip-ai-experimental_pytorch_lightning/site-packages/pytorch_lightning/strategies/deepspeed.py", line 920, in _auto_select_batch_size
    batch_size = train_dataloader.batch_sampler.batch_size
AttributeError: 'NoneType' object has no attribute 'batch_size'

Environment

Unfortunately, I can't collect the environment because we use a custom build system :(

lightning==2.2.0
torch==2.1.2

More info

No response

cc @awaelchli

[olegsinavski]

Workaround is strategy=pl.strategies.DeepSpeedStrategy(..., logging_batch_size_per_gpu=batch_size) actually.

Seems like this was broken in this commit:
3518f9e

I actually kinda agree that its better to have a hard crash, than warning. But maybe exception should be more verbose.

[awaelchli]

@olegsinavski Do you know what value for the batch size the DeepSpeed strategy determined internally prior to 2.2? It might have just been wrong, and since the batch size can't be known from looking at the datapipe itself (can it?), it's probably required that the users passes logging_batch_size_per_gpu. If this is true, then the error we get now just surfaces a bug we had prior and I agree with your suggestion to make it an error with good user instructions.

[anthonyhu]

Hey @awaelchli, I've experienced this bug with Oleg so I can answer the question about the batchsize.

Looking at the previous code to infer automatically the batchsize:

        import deepspeed

        # train_micro_batch_size_per_gpu is used for throughput logging purposes
        # by default we try to use the batch size of the loader
        assert self.lightning_module is not None
        batch_size = 1
        data_source = self.lightning_module.trainer.fit_loop._data_source
        if data_source.is_defined():
            try:
                train_dataloader = data_source.dataloader()
                if hasattr(train_dataloader, "batch_sampler"):
                    batch_size = train_dataloader.batch_sampler.batch_size
            # broad exception on purpose as `source.dataloader()` will fail if the dataloader requires `setup`
            # to have been called before
            except Exception:
                if self.global_rank == 0:
                    deepspeed.utils.logging.logger.warning(
                        "Tried to infer the batch size for internal deepspeed logging from the `train_dataloader()`. "
                        "To ensure DeepSpeed logging remains correct, please manually pass the strategy with the "
                        "batch size, `Trainer(strategy=DeepSpeedStrategy(logging_batch_size_per_gpu=batch_size))`."
                    )
        return batch_size

Previously, the deepspeed strategy internally determined the batch size to be equal to 1. As it would run into the exception, and then return the default batch size (which is 1). Perhaps we could add back this logging warning before crashing? It would have facilitated finding the source of the error. :)

[awaelchli]

Perhaps yes, we should add the warning back. But we should do an explicit None-check for the sampler instead of the previous try-except IMO. If you have the time, a PR for this would be appreciated.