A smarter, dynamic, quant and model agnostic Memory Manager.

[Zironic]

Core features of the Arena Memory Manager.

Noted incompatibility: Current iteration of Arena Memory Manager does not support live Text Encoder and requires cache_text_embeddings: true and unload_text_encoder: true. If requested support could be added but I personally felt if you're already offloading, you shouldn't make your life harder by having two models in memory at the same time.

1. Canonical, transactional host storage. The purpose of this is to ensure blocks can be moved in and out of VRAM without actually changing the block.
2. Generic saved-forward execution. The purpose of this is to ensure that models can be offloaded generically in a way that's transparent to model code and torch.compile compatible
3. Opaque quantized storage. The purpose of this is to make the offloader Quant agnostic and makes it trivial in most cases to support any new quant.
4. Compact asynchronous block transfers.
5. Checkpoint-aware lifetime and backward re-fetch. This is essentially an optimization that makes offloaded backwards faster.
6. Automatic block residency. This is one of the primary reasons I had this Arena Memory manager made. Rather then the way the current Memory Manager or the Diffusers variant work, establish a single static residency plan at the start of the job. The Arena Memory Manager dynamically establishes how much of the model can be kept resident and how much of the model should be offloaded, making the same code work for any amount of VRAM without the user having to change any settings*. It'll even adapt to outside memory pressure mid-run.
7. Separate training and sampling lifecycle. Since training has a much larger working set then inference does, it adapts automatically to move more of the model resident during inference.
8. Explicit ownership and cleanup
9. Early rejection of unsupported training modes
10. Low VRAM simulator for testing low VRAM support.
11. Shared host-memory safety. Because the shared GPU memory budget which memory pinning relies on is entirely invisible to Pytorch and CUDA. Arena Memory Manager implments the DXGI API to read the memory budget so it can avoid over-pinning memory and thus avoid CUDA memory allocation crashes other more naive offloaders suffer from. Also second major reason I had this Arena Memory Manager made, because I kept getting crashes from overpinned RAM.
12. Destination-first checkpoint loading. VRAM pressure can be reduced by supported models by loading straight into the canonical arena storage and avoid the ram spike from normal loading.
13. Native FP8 support.
14. Torch.Compile(fullgraph=True) compatible and supports Pytorch MegaCache.

*Provided the working set fits. Arena Memory Manager is smart memory management, not magic.

The Arena Memory Manager should teoretically support any model and any quant that supports the capability contract. It has been tested to work with:

Krea2
Z-Image
Anima

TorchAO FP8
Quanto qfloat8
OstrisLinear Convrot8
OstrisLinear Convrotint4

For a model to work, it needs discoverable blocks, frozen weights, model-owned checkpointing and compatible block inputs/outputs. Quantizers need to declare their physical storage through the storage-binding interface.

Example config:

train:
  train_text_encoder: false
  cache_text_embeddings: true
  unload_text_encoder: true

model:
  layer_offloading: true
  layer_offloading_smart: true

  # Automatic training policy.
  layer_offloading_smart_working_reserve_gb: -1
  layer_offloading_smart_physical_vram_headroom_gb: -1
  layer_offloading_smart_wddm_hard_gb: 1.0
  layer_offloading_smart_cap_calibration: true  # Windows-only

  # Checkpoint every selected block, ensuring backward re-fetch is tested.
  layer_offloading_checkpoint_keep_last: 0
  layer_offloading_prefetch_depth: 3

  compile: true
  compile_dynamic: true

Force low vram:

model:
  layer_offloading_simulated_vram_gb: 8
  layer_offloading_strict_vram_cap: true

The Arena Memory Manager was written for torchao==0.17.0 and updates the requirement to that. It will still run with 0.10.0 but with native TorchAO FP8 forward/backward support disabled.