Add Gradient Checkpointing

[williamFalcon]

Super useful to minimize RAM (trade-off speed).

Anyone interested in implementing?
(Instructions are here).

Probably needs to be generalized though.

[sholalkere]

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.

[williamFalcon]

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

[lukasfolle]

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

[wxj630]

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

[levhaikin]

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

[Borda]

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? 🐰

[levhaikin]

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.