Adding support for checkpointing about transformers.jl HOT 12 OPEN

That doesn't sounds the same. One could always completely turn off all dropouts, but normally we would want the checkpoint computed with the same dropout state as the first forward call so that the gradient with or without checkpoint are the same.

from transformers.jl.

These are good points @ToucheSir. I will come to this in two weeks timeframe, I am a bit busy now with academic stuff.

from transformers.jl.

Sounds good to have!

HF handle it in the forward method of hf-models (equiv. Layers.Transformer). I'm not sure Checkedpointed as AbstractTransformerBlock is the best place to add the checkpoint functionality. Some alternative ideas I currently have in mind:

1. Generalized Checkedpointed{S} <: LayerStruct and overload Checkpointed{<:Transformer} to add checkpoint per blocks.
2. Modify Layers.applyblocks to allow hooks and use Zygote.checkpointed as the hook function.
3. Similar to 2. but provide a HookedTransformerBlock <: AbstractTransformerBlock.

The wrapping function can be implemented with postwalk like the Layers.set_dropout.

from transformers.jl.

I will look at your suggestions. Checkpointed as a AbstractTransformerBlock was quick and dirty trick. I like the postwalk trick.

from transformers.jl.

One thing you'll want to think about is stateful layers like Dropout and BatchNorm which would not behave the same in subsequent calls. For the former I think some mechanism to snapshot RNG state would be required, and for the latter maybe an explicit overload?

from transformers.jl.

It seems the problem is that we cannot know if a Dropout or BatchNorm is executed under checkpointed environment?

from transformers.jl.

@ToucheSir I have not thought about this. Is there still switch to toggle train and test mode? That would effectively solve the problem.

from transformers.jl.

If the pullback is only called once, I believe BatchNorm and co should actually not require any special handling. Otherwise, the approach would be to traverse the model looking for these layers, saving their current train/test status, doing the checkpointing and then restoring the saved status.

As Peter notes, Dropout is trickier because you still need the RNG state around to create a mask. The most straightforward solution using struct Checkpointed would be to recurse through the model looking for Dropout layers and snapshotting their RNG state beforehand. Then that state can be restored whenever the checkpointing runs. I haven't quite thought about how this interacts with RNGs shared between layers (as is the default), but that should be solvable.

Medium-long term, we may want to consider a mechanism like https://github.com/vchuravy/ScopedValues.jl for exposing whether checkpointing is currently active in Flux itself. Then layers can query that info and change their behaviour accordingly without a wrapper.

from transformers.jl.

@ToucheSir I wonder if we could subtype the AContext in Zygote for a CheckpointedContext and overload the pullback behavior for dropout or so?

from transformers.jl.

Maybe, but generally we'd like to avoid coupling Flux to Zygote wherever possible (e.g. no custom pullbacks).

from transformers.jl.

I would say that only need to couple NNlib to Zygote since dropout is moved out from Flux.

from transformers.jl.

Yeah, NNlib has no dep (hard or weak) on Zygote right now and it'd be better to keep it that way. Porting Zygote.checkpoint to use the ChainRules API shouldn't be an issue, just need to decide if it lives in Flux or NNlib.

from transformers.jl.