DeepSpeedPlugin with activation checkpoint fails
Hi there,
I'm trying to run pytorch-lightning training with deepspeed plugin and activation checkpoints to support bigger batch sizes, based on https://pytorch-lightning.readthedocs.io/en/stable/advanced/advanced_gpu.html#deepspeed-activation-checkpointing.
As specified in the docs, running the model should be done using the checkpoint function. However, this function seems to return a tensor without gradients. When computing loss based on this value and returning from training_step, I'm getting an error
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
Minimal code to reproduce
import os

import deepspeed
import pytorch_lightning as pl
import torch
from deepspeed.ops.adam import FusedAdam
from pytorch_lightning.plugins import DeepSpeedPlugin
from torch import nn
from pytorch_lightning.utilities.types import STEP_OUTPUT
from torch.utils.data import DataLoader, RandomSampler


class PlModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = nn.Linear(1, 1)

    def forward(self, batch):
        return self.model(batch)

    def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
        res = deepspeed.checkpointing.checkpoint(self.model, batch)
        return nn.MSELoss()(res, torch.zeros_like(res, device=res.device))

    def configure_optimizers(self):
        return FusedAdam(self.parameters(), lr=0.1)


if __name__ == '__main__':
    trainer = pl.Trainer(gpus=-1, precision=16, plugins=DeepSpeedPlugin(stage=3, partition_activations=True))
    model = PlModel()
    dataset = torch.rand(100, 1)
    dl = torch.utils.data.DataLoader(dataset, batch_size=1, num_workers=os.cpu_count(),
                                     sampler=RandomSampler(dataset))
    trainer.fit(model, dl)
pytorch-lightning version: 1.3.3
deepspeed version: 0.5.0
Thanks!
Thanks @nachshonc!
I've managed to reproduce the same case without Deepspeed using torch.utils.checkpoint and our bug report model:
import deepspeed
import torch
from pytorch_lightning import LightningModule, Trainer
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 BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return torch.utils.checkpoint.checkpoint(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():
    train_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    val_data = DataLoader(RandomDataset(32, 64), batch_size=2)

    model = BoringModel()
    trainer = Trainer(
        max_epochs=1,
    )
    trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)


if __name__ == "__main__":
    run()
I think the issue arises from the fact that the entire model's activations have been removed, with the input tensors not requiring any gradients, thus the autograd engine not being able to infer any gradients.
For activation checkpointing, it only makes sense to include it if you have intermediate layers which can create expensive activations. For example, swap the model out to look like this:
class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer_h = torch.nn.Linear(32, 32)
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        x = torch.utils.checkpoint.checkpoint(self.layer_h, x)
        return self.layer(x)
Activation checkpointing just means on the backwards, we'll need to re-compute the activations (unless you do CPU checkpointing with Deepspeed or something, where activations are just transferred to the CPU memory). In this case, there is no point checkpointing the final layer, as the final layer will instantly need to be re-computed.
class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer_h = torch.nn.Linear(32, 32)
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        x = self.layer_h(x)
        return torch.utils.checkpoint.checkpoint(self.layer, x) # no point doing this!
We should definitely make the docs clearer for this, I'll make this an issue :)