[RFC] Make more operations inplace (GELU, BatchNorm, LayerNorm)
Labels: module: nn, triaged, enhancement, actionable

### 🚀 The feature, motivation and pitch

# Make more operations inplace (GELU, BatchNorm, LayerNorm)

## **Summary**
Hi PyTorch team,

We would like to enable users to make the following operations inplace: LayerNorm, BatchNorm and GELU.

## **Motivation**
In-place operations can lower the amount of memory used by operators in specific cases with minimal overhead
compared to other methods of reducing the memory footprint of models such as activation checkpointing, this
can then lead to increased performance in the form of throughput gains through batch size improvements. In-place
GELU and In-place LayerNorm combined for example, leads to an activation memory footprint savings of over
10% for Huggingface BERT at a sequence length of 512.

These benefits and prior art are described in the papers listed in the prior arts. Currently, in the
github repositories associated with these papers these In-place versions of these operators are implemented as 
PyTorch modules that must be imported.

The value add of this feature would be allowing users to take advantage of these operators directly from PyTorch. 
This would have no impact to other users of the library.

## **Proposed Implementation**
Our proposed implementation is to add an inplace flag to existing modules, which can then be set to True to enable 
the in-place version, such as with dropout for example.

`torch.nn.Dropout1d(p=0.5, inplace=False)`

The existing LayerNorm module, for example, would change from

`torch.nn.LayerNorm(normalized_shape, eps=1e-05, elementwise_affine=True, device=None, dtype=None)`
to 
`torch.nn.LayerNorm(normalized_shape, eps=1e-05, elementwise_affine=True, inplace=False, device=None, dtype=None)`.

Internally, this can then be handled as other inplace operators with a custom implementation in ATen, and adding 
the appropriate flags in the Autograd engine. Suggestions on how to do this would be welcome.

## **Metrics**
The main metric to evaluate whether this change is successful is to evaluate the memory usage of a model with
and without this change.

## **Drawbacks**
There are little to no drawbacks considering that the implementation cost is low, it does not introduce any breaking changes,
and this is not a torch wide change that would affect all users.

However, we note the In-place GELU is lossy, and this needs to be communicated/handled correctly.


## **Alternatives**
An alternative design would be to add these In-place versions as separate operators. For example, add an
`InPlaceLayerNorm` layer as its own separate module.


## **Prior Art**
These ideas are introduced in the following papers:

[In-Place Activated BatchNorm for Memory-Optimized Training of DNNs (CVPR '18)](https://openaccess.thecvf.com/content_cvpr_2018/papers/Bulo_In-Place_Activated_BatchNorm_CVPR_2018_paper.pdf)
Github repository: https://github.com/mapillary/inplace_abn

[Tempo: Accelerating Transformer-Based Model Training through Memory Footprint Reduction (NeurIPS '22)](https://openreview.net/forum?id=xqyEG7EhTZ)
Github repository: https://github.com/UofT-EcoSystem/Tempo

Example usage of Tempo:

```
from tempo.inplace_gelu import InplaceGelu
from tempo.inplace_layernorm import InplaceLayerNorm
from tempo.combined import Combined

# Original LayerNorm in Transformer model
# self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=1e-12)
# Replace it with Tempo's InplaceLayerNorm
self.LayerNorm = InplaceLayerNorm(config.hidden_size, eps=1e-12)
```

## **Unresolved questions**
* How to deal with the lossy nature of GELU
* Exact implementation details

### Alternatives

See above

cc @albanD @mruberry @jbschlosser @walterddr @saketh-are