API Reference — Overview

MemScale’s public API is intentionally small. Everything below is exported from the top-level memscale package.

Public functions

Function	Purpose	Returns
`wrap(target, config=None, ...)`	Optimize a model or `Trainer`	the same object, with hooks attached
`optimize(model, optimizer=None, ...)`	Context manager for custom training loops	yields the executor
`detach(model)`	Remove MemScale from a model	`None`

Configuration

Object	Purpose
`Config`	All optimization settings
`OptimizationMode`	Enum: `CONSERVATIVE`, `BALANCED`, `AGGRESSIVE`

Importing

import memscale
 
# functions
from memscale import wrap, optimize, detach
 
# configuration
from memscale import Config, OptimizationMode

wrap() vs apply_all_optimizations. wrap() applies checkpointing and CPU offload — the lightweight path. For the headline VRAM reductions, use apply_all_optimizations(model, optimizer), which adds BF16 mixed precision and 8-bit Adam on top; it returns an optimized (model, optimizer) pair. It is exported from the package and is what the benchmark suite measures. To optimize a Hugging Face Trainer, pass it straight to wrap() — see the Hugging Face guide.

Two ways in

wrap() — the common path. Works on a raw nn.Module or a Hugging Face Trainer. Attaches hooks and returns the object; you keep using it as before.
optimize() — a context manager for hand-written training loops where you want the optimization scoped to a with block.

Both are documented in full on their own pages.

Decision Engine wrap()