Lightning-AI / Lightning-AI/pytorch-lightning

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

Open
#19,460 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug help wanted strategy: deepspeed
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### 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

```python
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

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in pytorch_lightning/strategies/deepspeed.py, especially _format_batch_size_and_grad_accum_config and _auto_select_batch_size. Reproduce the issue with a DataLoader(batch_size=None) and the deepspeed strategy, then trace how the missing batch_sampler is handled. Done means this configuration no longer raises AttributeError during setup and has regression coverage.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.