Efficient and Automatic Rematerialization for Pytorch training
Project description
Rockmate
The Rockmate framework is designed for training a PyTorch neural network within a given GPU budget
constraint using automatic re-materialization (activation checkpointing) technique.
Given a PyTorch model, a sample input, and a GPU memory budget, Rockmate builds a new
torch.nn.Module, which performs forward and backward pass keeping activations under the given
budget.
- The new model produces the same outputs and gradients as the original one.
- Model training with a budget constraint, which is lower than the one required by PyTorch Autodiff, is achieved by re-computing some of the activations instead of storing them for gradient calculation.
- Depending on the budget,
Rockmatedefines automatically which activations should be recomputed.
More information on our repository.
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
rockmate-2.0.0.tar.gz
(78.8 kB
view details)
File details
Details for the file rockmate-2.0.0.tar.gz.
File metadata
- Download URL: rockmate-2.0.0.tar.gz
- Upload date:
- Size: 78.8 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.0.1 CPython/3.10.12
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
68c41a21858121c569ea6dbe973c0208b14579c4c5a360fdbcf00e8d7b1ffbe6
|
|
| MD5 |
0bc63c8746da0fe09a3254063c8c5653
|
|
| BLAKE2b-256 |
6c8f7052feecb918599323be3fad744b3ac53acea483bb577a907ba9241a6f84
|
