MemScale

Drop-in memory optimizer for PyTorch training. Cut VRAM up to ~76% — and train models that otherwise don't fit — with 1 line of code.

The problem

Training large models on GPUs hits a wall: VRAM.

• GPT-2 Large training → 14.9 GB on a 24 GB RTX 3090, with little headroom
• 1.5B parameter model → out of memory on single 24GB GPU
• DeepSpeed ZeRO setup → 2 weeks of configuration

MemScale solves this. Wrap your model in 1 line, cut VRAM up to ~76%, and train models that otherwise run out of memory — no code changes.

Benchmarks

Validated on RTX 3090 24GB (PyTorch 2.12, CUDA 13) Reproducible: python -m memscale.benchmarks --output results.json

Configuration: AGGRESSIVE mode (8-bit Adam + BF16 + checkpointing). Reduction % scales with workload size — larger batches and longer sequences typically yield higher percentages. See BENCHMARK_REPORT.md for methodology details.

Quick start

pip install memscale
pip install bitsandbytes  # optional, for additional 8-bit Adam savings

import memscale
from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(per_device_train_batch_size=16),
    train_dataset=dataset,
)

# Add this one line:
trainer = memscale.wrap(trainer)

trainer.train()  # Up to ~76% less VRAM, same speed

That's it. MemScale automatically:

1. Profiles your model's memory usage per layer
2. Decides which optimization technique fits each layer best
3. Applies boundary checkpointing, 8-bit Adam, mixed precision
4. Reports memory savings and throughput in real time

No API key required. Library works fully offline. Anonymous telemetry is enabled by default to help improve the decision engine — opt out anytime with memscale.disable_telemetry() or MEMSCALE_TELEMETRY=0. See Telemetry below for what's collected.

Experimental: async CPU offload (v1.1+)

v1.1 adds an experimental async CPU offload engine. It overlaps GPU→CPU transfers with compute using tier-aware CUDA streams, while keeping the same VRAM savings as the default sync path.

import memscale

model = memscale.wrap(model, async_offload=True)
# Same VRAM savings, async transfer.
# Note: speedup is not yet benchmarked — see the v1.2 roadmap.

async_offload defaults to False; existing users are unaffected. Enabling it emits an ExperimentalFeatureWarning. Numerical equivalence with the sync path is validated (rtol 1e-6 / atol 1e-7 on RTX 3090); the training-loop speedup is deferred to a dedicated v1.2 benchmarking phase. For production, use the default sync offload. See the async offload user guide for details.

Maximum reduction (combined techniques)

For maximum savings, enable all techniques:

import torch
from memscale import Config, OptimizationMode
from memscale.phase_f import apply_all_optimizations

model = YourModel()
optimizer = torch.optim.AdamW(model.parameters())

config = Config(
    mode=OptimizationMode.AGGRESSIVE,
    use_8bit_optimizer=True,    # bitsandbytes 8-bit Adam
    use_mixed_precision=True,   # BF16 on Ampere+, FP16 fallback
)

# One call applies all techniques
model, optimizer = apply_all_optimizations(model, optimizer, config)

# Train normally
for batch in dataloader:
    loss = model(batch).loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

This stack achieved 73.4% reduction on BERT-Large (7.60 GB → 2.02 GB) in the v1.1 benchmarks — see BENCHMARK_REPORT.md.

How it works

MemScale combines proven memory optimization techniques and chooses what fits each layer:

The decision engine analyzes your model and picks the right technique per layer — you don't need to configure individual layers.

HuggingFace integration

For HuggingFace autoregressive models (GPT-2, Llama, Mistral, T5), MemScale automatically disables config.use_cache when checkpointing is enabled. This prevents the CheckpointError that occurs when KV-cache concatenation conflicts with backward recompute. No code changes needed — just memscale.wrap().

Multi-GPU support

Multi-GPU training works via standard PyTorch DDP. MemScale's per-GPU optimizations apply on each GPU:

torchrun --nproc_per_node=2 your_training_script.py

import memscale
import torch.nn.parallel as parallel

model = YourModel().to(local_rank)
model, optimizer = apply_all_optimizations(model, optimizer, config)
model = parallel.DistributedDataParallel(model, device_ids=[local_rank])

# Train normally - 87% per-GPU reduction with 2x throughput

Validated on 2x RTX 3090: 1.69 GB per GPU (vs 13 GB baseline single-GPU).

Distributed sharding (research preview)

memscale.distributed provides ZeRO-3 inspired parameter and optimizer sharding building blocks. Full integration with model forward/backward hooks is planned for v1.1. For production multi-GPU training requiring 95%+ reduction today, FSDP or DeepSpeed remain the recommended choice.

Usage modes

HuggingFace Trainer

import memscale
trainer = memscale.wrap(your_hf_trainer)
trainer.train()

PyTorch Lightning

from lightning import Trainer
from memscale.integrations.lightning import MemScaleLightningCallback

trainer = Trainer(
    callbacks=[MemScaleLightningCallback()],
    max_epochs=10,
)
trainer.fit(model, dataloader)

Custom training loop

import memscale

with memscale.optimize(model, optimizer) as ms:
    for batch in dataloader:
        loss = model(batch).loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

Configuration

Most users don't need this. Defaults work for 90% of cases.

from memscale import wrap, Config, OptimizationMode

config = Config(
    mode=OptimizationMode.AGGRESSIVE,  # or BALANCED (default), CONSERVATIVE
    enable_checkpointing=True,
    enable_offloading=True,
    use_8bit_optimizer=False,    # set True for max reduction
    use_mixed_precision=False,   # set True for max reduction
    target_gpu_utilization=0.85,
)

trainer = wrap(trainer, config=config)

Cost attribution

Track how much money MemScale saves on cloud GPU bills:

from memscale.cost_attribution import CostTracker, estimate_savings

# Quick estimate
report = estimate_savings(
    baseline_vram_gb=70.0,
    memscale_vram_gb=35.0,
    training_hours=10.0,
    gpu_type='A40',
    baseline_gpu_type='A100 80GB',  # GPU you'd need WITHOUT MemScale
)
print(report)
# Baseline (A100 80GB): $24.90
# MemScale (A40):       $15.30
# Savings:              $9.60 (38.6%)

Built-in pricing for 16 GPU types (V100, A100, H100, RTX series, AMD MI300X, etc.) plus auto-inference of the cheapest GPU sufficient for your workload.

OOM prediction

Catch out-of-memory before training starts:

from memscale import OOMPredictor

predictor = OOMPredictor(model_params_bytes=2_000_000_000)  # 2 GB params
risk = predictor.predict(batch_size=16, sequence_length=2048, optimizer='adamw')

if risk.level == 'CRITICAL':
    print(f"⚠️ {risk.message}")
    print(f"Recommendations: {risk.recommendations}")

Telemetry

MemScale ships with anonymous telemetry enabled by default to improve the decision engine across diverse hardware and workloads. To opt out:

import memscale
memscale.disable_telemetry()

Or via environment variable (set before importing MemScale):

export MEMSCALE_TELEMETRY=0

What's collected (~1 KB per training run)

• Anonymous client ID (random UUID, stored locally at ~/.memscale/client_id)
• Library version, Python version, PyTorch version, OS
• Hardware: GPU model, VRAM, CUDA version, number of GPUs
• Model architecture: layer types, parameter count (no weights)
• Optimization outcome: techniques applied, memory saved, throughput overhead

What's NEVER collected

• ❌ Model weights or training data
• ❌ Code or scripts
• ❌ File paths, hostnames, IP addresses
• ❌ Email or any identifying information
• ❌ Layer-level activations or gradients

Telemetry is fire-and-forget (silent failure on network error), sent over HTTPS to api.memscale.id/v1/telemetry. See our privacy policy for details.

Re-enable after opting out

memscale.enable_telemetry()

Or:

export MEMSCALE_TELEMETRY=1

Compatibility

AMD GPU support (ROCm) coming in a future release.

FAQ

Q: Does MemScale change my training results? The activation checkpointing and DDP techniques are mathematically lossless. BF16/FP16 mixed precision introduces small numerical differences — same as standard PyTorch AMP.

Q: How does this compare to DeepSpeed and FSDP? DeepSpeed and FSDP are powerful but require significant configuration and distributed training expertise. MemScale's value is plug-and-play: 1-line wrap with auto-detection. For 95%+ reduction in production multi-GPU setups, DeepSpeed ZeRO-3 is more mature. For single-GPU and DDP workloads, MemScale is competitive and easier to use.

Q: Will this slow down my training? Activation checkpointing adds 20-30% compute overhead (the standard tradeoff). 8-bit Adam adds ~2-5%. Net effect: training is slower per step, but you can use larger batches (better hardware utilization), so end-to-end time often improves.

Q: What if my model has custom architecture? The decision engine handles standard transformers (PyTorch native, HuggingFace BERT/GPT2/Llama/Mistral, vision transformers) automatically. Custom architectures fall back to per-module heuristics. Both are tested.

Q: Why a range instead of a flat number? Reduction depends on model architecture, batch size, sequence length, and which techniques you enable. The v1.1 benchmarks show roughly 68-76% on standard transformers at the default workload, and MemScale enables training models that otherwise run out of memory. Larger batches and longer sequences typically yield higher percentages.

Q: Is MemScale open source? Source code is currently proprietary. PyPI distribution is public (free to install and use). We may open source later based on community feedback. See memscale.id for licensing.

Roadmap

• v1.0.4 (current): 5 medium bug fixes (seq_len awareness, dtype-aware param count, tiling outputs, GPU downgrade cost, estimate_training_memory)
• v1.1 (Q3 2026): Reproducible benchmark CLI, bug fixes (F-1…F-4), honest cost attribution, and experimental async CPU offload (opt-in, correctness-validated; speedup deferred to v1.2)
• v1.2 (Q4 2026): Async offload speedup benchmarking + older-GPU validation, tensor splitting, ML-based decision policy trained on telemetry data
• v2.0 (2027): Multi-framework support (JAX, TensorFlow) + MemScale Serve (inference)

Architecture

MemScale's optimization happens in stages:

1. Profiling: Static analysis with empirical fallback for dynamic models
2. Decision engine: Per-layer technique selection based on memory profile, hardware budget, and configuration
3. Execution: Apply chosen techniques via PyTorch hooks
4. Observation: Track memory and throughput, report to user

Reporting Issues

For bug reports, please include:

1. Minimal reproducible example
2. Hardware (GPU model, VRAM)
3. PyTorch version
4. Output of memscale.profile_model(model) if relevant

Email: team@memscale.id

License

Proprietary. Full terms: contact team@memscale.id or visit memscale.id.

Citation

If you use MemScale in your research, please cite:

@software{memscale2026,
  title={MemScale: Drop-in Memory Optimization for PyTorch Training},
  author={MemScale Team},
  year={2026},
  url={https://memscale.id}
}

Built for ML practitioners. Questions? team@memscale.id