Skip to main content

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, Rockmate defines automatically which activations should be recomputed.

More information on our repository.

Metadata

Release files for rockmate 2.0.0

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for rockmate 2.0.0
File Size Uploaded Provenance Actions
rockmate-2.0.0.tar.gz 78.8 kB Without provenance Details for rockmate-2.0.0.tar.gz

Release files / rockmate-2.0.0.tar.gz

Download URL rockmate-2.0.0.tar.gz
Size 78.8 kB
Tags Source
SHA-256 checksum
How to use checksums
68c41a21858121c569ea6dbe973c0208b14579c4c5a360fdbcf00e8d7b1ffbe6
BLAKE2b-256 checksum
How to use checksums
6c8f7052feecb918599323be3fad744b3ac53acea483bb577a907ba9241a6f84
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via twine/6.0.1 CPython/3.10.12

Security for rockmate 2.0.0

Provenance

Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.

No provenance found

None of the files in this release were uploaded with provenance metadata.

A note on package security

PyPI is a repository for distributing packages, not a service for auditing them. While third-party security researchers and our staff do their best to detect and remove malicious packages, PyPI cannot guarantee the quality, security or suitability of any given package.

To learn more about protecting yourself from malicious actors, visit our package security documentation.

Release history Release notifications | RSS feed

This release

2.0.0 This release

1 release file No provenance

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page