Deepspeed accelerator calls datamodule.train_dataloader() prior to setup()

[leezu]

🐛 Bug

Deepspeed accelerator calls datamodule.train_dataloader() prior to setup(). This does not happen with other accelerators.

To Reproduce

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()

[tchaton]

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

[leezu]

@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

[tchaton]

@SeanNaren Any idea ?

[SeanNaren]

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:

Trainer(plugins=DeepSpeedPlugin(logging_batch_size_per_gpu=32))

If I modify the case, it works :)

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