[RFC] Scan & Gradient checkpointing in Flax
Author: patrickvonplaten
Pull request: no
Labels: name: WIP

### Feature request

We should add scan and remat (gradient checkpointing) to the most important Flax/JAX models (BERT, GPT2, OPT, T5, BART, Wav2Vec2).

### Motivation

Scan allows for much faster compilation and memory savings and `remat` is the equivalent of `gradient_checkpointing` in PyTorch.

@sanchit-gandhi already uses both features in the Flax Seq2Seq Speech project - see: https://github.com/sanchit-gandhi/seq2seq-speech so it'd be quite trivial to get them working.

**Implementation details:**

Given that both `scan` and `remat` are not related to the model architecture, they should IMO **not** be in the model's config (We've done this mistake in PyTorch and don't want to repeat it here).

I would advocate for the following API:

```python
model = FlaxBertForMaskedLM.from_pretrained("bert-base-cased")
model.scan()  # or model.scan_enable()
model.unscan()  # or model.scan_disable()
```

and 

```python
model = FlaxBertForMaskedLM.from_pretrained("bert-base-cased")
model.gradient_checkpoint_enable()
model.gradient_checkpoint_disable()
```

As can be seen here: https://github.com/sanchit-gandhi/seq2seq-speech/blob/b28d0c25c8fad0f9ffa6707f91f7aba320d44a4b/models/modeling_flax_wav2vec2.py#L504

We'll need to re-initialize the `flax.linen.module` inside the model. However this should be fine since it just means that we do

```
self.module = self.module_class(config=config, dtype=dtype, use_scan=True, **kwargs)
 self. _is_scan_enabled = True
```

similar to this line: https://github.com/huggingface/transformers/blob/71e602725b90f63f404109bae9f72cbdf755477b/src/transformers/models/wav2vec2/modeling_flax_wav2vec2.py#L868

We can see along the PR how much logic can reside in `modeling_flax_utils.py` and how much would go into the specific models, *e.g.* `modeling_flax_wav2vec2.py`.

The same API / logic could be used for the `gradient_checkpointing`.

### Your contribution

Happy to give this implementation a shot with @sanchit-gandhi  and @patil-suraj .

Also would love to hear feedback from @borisdayma @marcvanzee about the API

Comments:
I'm not sure you would need both versions within a same script (scan and unscanned, or with and without checkpointing which affects only training anyway).

Then maybe you could just add it directly as an arg to `model.from_pretrained(..., scan=False, gradient_checkpointing=False)`

You would just have to use some naming conventions on your params to see if you need to scan/unscan when loading a checkpoint.

Suppose you have a training script, it would be useful to be able to use `scan` and `remat` during training for faster compile times and larger batch sizes, and then switch to `unscan` and no `remat` during eval for faster inference?

I'm not sure it would be worth it:
* Most of the time evaluation is relatively fast
* You would have to reformat your parameters each time between eval and train, potentially leading to memory fragmentation

Hey @patrickvonplaten, I'm keen to get gradient checkpointing working in JAX for [long-t5](https://huggingface.co/google/long-t5-tglobal-xl/tree/main). If this is not on the cards to be added soon happy to work on a PR for it if that works with you all?

Hey @KMFODA! There's a PR that is close to being merged: https://github.com/huggingface/transformers/pull/17843 I'll let you know once it's complete, and you can copy the logic across to Flax T5 in a new PR if that sounds good to you!

This RFC targeted the Flax/JAX modeling backend (`scan` / `remat` on `FlaxBertForMaskedLM`, `modeling_flax_wav2vec2.py`, etc.), but that entire backend has since been removed from the library: TensorFlow and Jax support were fully dropped in https://github.com/huggingface/transformers/pull/40760 (merged Sep 2025). There are no longer any `modeling_flax_*.py` files or a `modeling_flax_utils.py` to hang this on, so the RFC no longer has a surface to land in. Closing as outdated