Repository: Lightning-AI/pytorch-lightning
Pull request: Add a warning to deepspeed when inferring batch size
## What does this PR do?

Closes #8872 

Adds a warning for users, as this may crash their code if they require setup to be called first.

## Before submitting

- [x] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- [x] Did you make sure to **update the documentation** with your changes? (if necessary)
- [x] Did you write any **new necessary tests**? (not for typos and docs)
- [x] Did you verify new and **existing tests pass** locally with your changes?
- [x] Did you list all the **breaking changes** introduced by this pull request?
- [x] Did you **update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)**? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [x] Is this pull request ready for review? (if not, please submit in draft mode)
- [x] Check that all items from **Before submitting** are resolved
- [x] Make sure the title is self-explanatory and the description concisely explains the PR
- [x] 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 🙃


Problem statement:
Deepspeed accelerator calls datamodule.train_dataloader() prior to setup()
## 🐛 Bug

Deepspeed accelerator calls datamodule.train_dataloader() prior to setup(). This does not happen with other accelerators.

### To Reproduce

```py
import os

import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer, LightningDataModule


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class PlDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self._setup = False

    def setup(self, stage):
        self._setup = True

    def train_dataloader(self):
        assert self._setup
        return DataLoader(RandomDataset(32, 64), batch_size=2)

    def val_dataloader(self):
        assert self._setup
        return DataLoader(RandomDataset(32, 64), batch_size=2)

    def test_dataloader(self):
        assert self._setup
        return DataLoader(RandomDataset(32, 64), batch_size=2)


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("train_loss", loss)
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("valid_loss", loss)

    def test_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("test_loss", loss)

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


def run():

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_train_batches=1,
        limit_val_batches=1,
        num_sanity_val_steps=0,
        max_epochs=1,
        weights_summary=None,
        accelerator="deepspeed",
    )
    trainer.fit(model, datamodule=PlDataModule())


if __name__ == "__main__":
    run()
```

### Expected behavior

`train_dataloader` is never called before `setup`.

### Additional context

Backtrace

```
  /home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py(553)fit()
-> self._run(model)
  /home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py(865)_run()
-> self.accelerator.setup_environment()
  /home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/accelerators/gpu.py(30)setup_environment()
-> super().setup_environment()
  /home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py(76)setup_environment()
-> self.training_type_plugin.setup_environment()
  /home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp.py(166)setup_environment()
-> self.setup_distributed()
  /home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/deepspeed.py(341)setup_distributed()
-> self._format_config()
  /home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/deepspeed.py(545)_format_config()
-> self._format_batch_size_and_grad_accum_config()
  /home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/deepspeed.py(555)_format_batch_size_and_grad_accum_config()
-> batch_size = self._auto_select_batch_size()
  /home/ubuntu/.local/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/deepspeed.py(566)_auto_select_batch_size()
-> train_dataloader = self.lightning_module.train_dataloader()
```

Hints from issue discussion:
Dear @leezu,

The `DeepSpeed Plugin` is doing a hack to automatically resolved your batch_size, but won't actually use this dataloader for training.
And there is no simple way to improve this unless you provide the batch_size within the deepspeed config. 

And, the DataLoader used for training will be defined within a distributed setting.

My advice is to use `is_distributed_available()` function if you have custom samplers.

I will be closing this issue for now.

Best,
T.C
@tchaton thank you for looking into a fix. 

> And there is no simple way to improve this unless you provide the batch_size within the deepspeed config.

Even if `setup` can't be called automatically, how about adding sanity assertions that ensure train_dataloader is never called before `setup`? In that case users can at least be informed about the issue
@SeanNaren Any idea ?
Thanks for the issue @leezu!

To skip the auto-infer of the batch size for logging, you can pass the batch size directly to the plugin like such:

```python
Trainer(plugins=DeepSpeedPlugin(logging_batch_size_per_gpu=32))
```

If I modify the case, it works :)

```python
import os

import torch
from pytorch_lightning import LightningModule, Trainer, LightningDataModule
from pytorch_lightning.plugins import DeepSpeedPlugin
from torch.utils.data import DataLoader, Dataset


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class PlDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self._setup = False

    def setup(self, stage):
        self._setup = True

    def train_dataloader(self):
        assert self._setup
        return DataLoader(RandomDataset(32, 64), batch_size=2)

    def val_dataloader(self):
        assert self._setup
        return DataLoader(RandomDataset(32, 64), batch_size=2)

    def test_dataloader(self):
        assert self._setup
        return DataLoader(RandomDataset(32, 64), batch_size=2)


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("train_loss", loss)
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("valid_loss", loss)

    def test_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("test_loss", loss)

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


def run():

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_train_batches=1,
        limit_val_batches=1,
        num_sanity_val_steps=0,
        max_epochs=1,
        weights_summary=None,
        plugins=DeepSpeedPlugin(logging_batch_size_per_gpu=32),
        gpus=1
    )
    trainer.fit(model, datamodule=PlDataModule())


if __name__ == "__main__":
    run()
```

This should skip the auto-infer of the batch size. We should add a warning here so that users are aware! Will add that in a PR

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/CHANGELOG.md b/CHANGELOG.md
index aa3669c40ccf3..b39e4dab4a6f9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -106,6 +106,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added `on_exception` callback hook ([#9183](https://github.com/PyTorchLightning/pytorch-lightning/pull/9183))
 
+- Add a warning to deepspeed when inferring batch size ([#9221](https://github.com/PyTorchLightning/pytorch-lightning/pull/9221))
+
+
 ### Changed
 
 - Parsing of the `gpus` Trainer argument has changed: `gpus="n"` (str) no longer selects the GPU index n and instead selects the first n devices. ([#8770](https://github.com/PyTorchLightning/pytorch-lightning/pull/8770))
diff --git a/pytorch_lightning/plugins/training_type/deepspeed.py b/pytorch_lightning/plugins/training_type/deepspeed.py
index 36d487e3419c7..a4173b57f0b8d 100644
--- a/pytorch_lightning/plugins/training_type/deepspeed.py
+++ b/pytorch_lightning/plugins/training_type/deepspeed.py
@@ -546,6 +546,11 @@ def _format_batch_size_and_grad_accum_config(self):
                 " as this will be set via accumulate_grad_batches=x argument passed via the Lightning Trainer."
             )
         if "train_micro_batch_size_per_gpu" not in self.config:
+            rank_zero_warn(
+                "Inferring the batch size for internal deepspeed logging from the `train_dataloader()`. "
+                "If you require skipping this, please pass "
+                "`Trainer(plugins=DeepSpeedPlugin(logging_batch_size_per_gpu=batch_size)`"
+            )
             batch_size = self._auto_select_batch_size()
             self.config["train_micro_batch_size_per_gpu"] = batch_size
         self.config["gradient_accumulation_steps"] = self.lightning_module.trainer.accumulate_grad_batches
diff --git a/tests/plugins/test_deepspeed_plugin.py b/tests/plugins/test_deepspeed_plugin.py
index 29086bd72a053..de4bb3ea987f9 100644
--- a/tests/plugins/test_deepspeed_plugin.py
+++ b/tests/plugins/test_deepspeed_plugin.py
@@ -1,6 +1,6 @@
 import json
 import os
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 from unittest import mock
 
 import pytest
@@ -11,7 +11,7 @@
 from torch.utils.data import DataLoader
 from torchmetrics import Accuracy
 
-from pytorch_lightning import LightningModule, seed_everything, Trainer
+from pytorch_lightning import LightningDataModule, LightningModule, seed_everything, Trainer
 from pytorch_lightning.callbacks import Callback, LearningRateMonitor, ModelCheckpoint
 from pytorch_lightning.plugins import DeepSpeedPlugin, DeepSpeedPrecisionPlugin
 from pytorch_lightning.plugins.training_type.deepspeed import LightningDeepSpeedModule
@@ -808,3 +808,53 @@ def training_step(self, batch, batch_idx):
     trainer = Trainer(default_root_dir=tmpdir, plugins=[DeepSpeedPlugin()], gpus=1, fast_dev_run=True, precision=16)
     with pytest.raises(MisconfigurationException, match="returning `None` .* is not supported"):
         trainer.fit(model)
+
+
+@RunIf(min_gpus=1, deepspeed=True, special=True)
+def test_deepspeed_warn_train_dataloader_called(tmpdir):
+    """Test DeepSpeed warns when it calls ``lightning_module.train_dataloader`` internally for logging batch
+    size."""
+    model = BoringModel()
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        plugins=[DeepSpeedPlugin()],
+        gpus=1,
+        fast_dev_run=True,
+    )
+    with pytest.warns(UserWarning, match="Inferring the batch size for internal deepspeed logging"):
+        trainer.fit(model)
+
+
+@RunIf(min_gpus=1, deepspeed=True, special=True)
+def test_deepspeed_setup_train_dataloader(tmpdir):
+    """Test DeepSpeed works when setup is required to call, and the user passes the batch size manually."""
+
+    class TestSetupIsCalledDataModule(LightningDataModule):
+        def __init__(self):
+            super().__init__()
+            self._setup = False
+
+        def setup(self, stage: Optional[str] = None) -> None:
+            self._setup = True
+
+        def train_dataloader(self):
+            assert self._setup
+            return DataLoader(RandomDataset(32, 64), batch_size=2)
+
+        def val_dataloader(self):
+            assert self._setup
+            return DataLoader(RandomDataset(32, 64), batch_size=2)
+
+        def test_dataloader(self):
+            assert self._setup
+            return DataLoader(RandomDataset(32, 64), batch_size=2)
+
+    model = BoringModel()
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        plugins=[DeepSpeedPlugin(logging_batch_size_per_gpu=32)],
+        gpus=1,
+        fast_dev_run=True,
+    )
+    trainer.fit(model, datamodule=TestSetupIsCalledDataModule())
+    trainer.test(model)


Code review comments:
diff_hunk: @@ -546,6 +546,11 @@ def _format_batch_size_and_grad_accum_config(self):
                 " as this will be set via accumulate_grad_batches=x argument passed via the Lightning Trainer."
             )
         if "train_micro_batch_size_per_gpu" not in self.config:
+            rank_zero_warn(
+                "Inferring the batch size for internal deepspeed logging from the `train_dataloader()`. "
+                "If you require skipping this, please pass "
+                "`Trainer(plugins=DeepSpeedPlugin(logging_batch_size_per_gpu=batch_size)`"
path: pytorch_lightning/plugins/training_type/deepspeed.py
text: @user1:
```suggestion
                "`Trainer(plugins=DeepSpeedPlugin(logging_batch_size_per_gpu=batch_size))`"
```

@author:
great catch!