Add Gradient Checkpointing about lightning HOT 10 CLOSED

I would be really interested in this feature! What is the current status?

from lightning.

huggingface has this built-in at this point. I think it's important to have this feature in lightning as well.

from lightning.

any updates? i still can't find how to apply gradient checkpointing with lightning

from lightning.

@sidhanthholalkere @Borda https://gitter.im/PyTorch-Lightning/community?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge

from lightning.

would be really interested in this Gradient Checkpointing (https://github.com/cybertronai/gradient-checkpointing)
! What is the current status?

from lightning.

huggingface has this built-in at this point. I think it's important to have this feature in lightning as well.

yeah, would be interested in implementing it? 🐰

from lightning.

@TianHongZXY and future googlers. This is implemented via FSDP https://lightning.ai/docs/pytorch/stable/advanced/model_parallel/fsdp.html#activation-checkpointing

from lightning.

I'm not sure that this is really necessary(maybe just a link to the tutorial instead of implementing it).
For sequential nets, its as simple as

    def forward(self, input_var, chunks=3):
        modules = [module for k, module in self._modules.items()][0]
        input_var = checkpoint_sequential(modules, chunks, input_var)
        input_var = input_var.view(input_var.size(0), -1)
        input_var = self.fc(input_var)
        return input_var

which is basically just adding 3 extra lines to the forward().
For other models(RNNs, LSTMs, etc.) it doesn't seem like there's a nice generalization and can vary between model structure.

from lightning.

it's very tempting, but unfortunately I don't think I will be able to commit to that at this point, due to time constraints. it's a shame that I can't because I use lightning and I think it's a good project. if something changes I will let you know.

from lightning.

Is there a reason that the standard Pytorch checkpointing (https://pytorch.org/docs/stable/checkpoint.html) cannot be used with lightning? Why would this need to be reimplemented in lightning?

from lightning.