Skip to main content

AI Model Dynamic Offloader

This project is a pytorch VRAM allocator that implements on-demand offloading of model weights when the primary pytorch VRAM allocator comes under pressure.

Support:

  • Nvidia GPUs (CUDA) and AMD GPUs (ROCm/HIP)
  • PyTorch 2.8+
  • CUDA 12.8+ (Nvidia) / ROCm 7+ (AMD)
  • Windows 11+ / Linux as per python ManyLinux support

How it works:

  • The pytorch application creates a Virtual Base Address Register (VBAR) for a model. Creating a VBAR doesn't cost any VRAM, only GPU virtual address space (which is pretty much free).
  • The pytorch application allocates tensors for model weights within the VBAR. These tensors are initially un-allocated and will segfault if touched.
  • The pytorch application faults in the tensors using the fault() API at the time the tensor is needed. This is where VRAM actually gets allocated.
If the fault() is successful (sufficient VRAM for this tensor):
  1. If the fault() resultant signature is changed or unknown:
    • The application uses tensor::_copy() to populate the weight data on the GPU.
    • The application saves the returned signature against this weight for future comparison
  2. The layer uses the weight tensor.
  3. The application calls unpin() on the tensor to allow it to be freed under pressure later if needed.
If the fault() is unsuccessful (offloaded weight):
  1. The application allocates a temporary regular GPU tensor.
  2. Uses _copy to populate weight data on the GPU.
  3. The layer uses the temporary as the weight.
  4. Pytorch garbage collects the temp when the layer is finished.

see examples/example.py


Priorities:

  • The most recent VBARs are the highest priority and lower addresses in the VBAR take priority over higher addresses.
  • Applications should order their tensor allocations in the VBAR in load-priority order with the lowest addresses for the highest priority weights.
  • Calling fault() on a weight that is higher priority than other weights will cause those lower priority weights to get freed to make space.
  • Having a weight evicted sets that VBAR's watermark to that weight's level. Any weights in the same VBAR above the watermark automatically fail the fault() API. This avoids constantly faulting in all weights each model iteration while allowing the application to just blindly call fault() every layer and check the results. There is no need for the application to manage any VRAM quotas or watermarks.
  • Existing VBARs can be pushed to top priority with the prioritize() API. This allows use of an already loaded or partially model (e.g. using the same model twice in a complex workflow). Using prioritize resets the offload watermark of that model to no offloading, giving its weights priority over any other currently loaded models.

Backend:

  • VBAR allocation is done with cuMemAddressReserve(), faulting with cuMemCreate() and cuMemMap() and all frees done with appropriate converse APIs.
  • For consistency with VBAR memory management, main pytorch allocator plugin is also implemented with cuMemAddressReserve -> cuMemCreate -> cuMemMap. This also behaves a lot better on Windows systems with System Memory fallback.

On AMD, the equivalent HIP APIs (hipMemAddressReserve -> hipMemCreate -> hipMemMap, and their converse calls) are used throughout via the same flow.

Caveats:

  • There is no real way for this allocator to tell the difference between high usage and bad fragmentation in the pytorch caching allocator. As we always return success to the pytorch caching allocator it experiences no pressure while weights are being offloaded which means it can run in an extremely fragmented mode. The assumption is model weight access patterns are reasonably regular over blocks or iterations and it finds a good set of sizes to cache. What you should generally do though, is completely flush the pytorch caching allocator before each new model run, which avoids completely un-used reservations from taking priority over the next models weights.

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distributions

No source distribution files available for this release.See tutorial on generating distribution archives.

Built Distributions

If you're not sure about the file name format, learn more about wheel file names.

comfy_aimdo-0.5.2-py3-none-any.whl (24.0 kB view details)

Uploaded Python 3

comfy_aimdo-0.5.2-cp39-abi3-win_arm64.whl (248.0 kB view details)

Uploaded CPython 3.9+Windows ARM64

comfy_aimdo-0.5.2-cp39-abi3-win_amd64.whl (273.4 kB view details)

Uploaded CPython 3.9+Windows x86-64

comfy_aimdo-0.5.2-cp39-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (429.0 kB view details)

Uploaded CPython 3.9+manylinux: glibc 2.17+ x86-64

comfy_aimdo-0.5.2-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl (655.8 kB view details)

Uploaded CPython 3.9+manylinux: glibc 2.17+ ARM64

File details

Details for the file comfy_aimdo-0.5.2-py3-none-any.whl.

File metadata

  • Download URL: comfy_aimdo-0.5.2-py3-none-any.whl
  • Upload date:
  • Size: 24.0 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for comfy_aimdo-0.5.2-py3-none-any.whl
Algorithm Hash digest
SHA256 3ac0cdbbe3451e5d6e2295b1ea5db5f0b4bc8a33396082739cbcfa3d7df1f8ca
MD5 0cb3e177b35f21aa889a7fc1b8b9a7b8
BLAKE2b-256 f8f8b7d4e4e1002790dbd44386d89eb59c376f18b398ee2c52865f7174a7c40e

See more details on using hashes here.

Provenance

The following attestation bundles were made for comfy_aimdo-0.5.2-py3-none-any.whl:

Publisher: build-wheels.yml on Comfy-Org/comfy-aimdo

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file comfy_aimdo-0.5.2-cp39-abi3-win_arm64.whl.

File metadata

  • Download URL: comfy_aimdo-0.5.2-cp39-abi3-win_arm64.whl
  • Upload date:
  • Size: 248.0 kB
  • Tags: CPython 3.9+, Windows ARM64
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for comfy_aimdo-0.5.2-cp39-abi3-win_arm64.whl
Algorithm Hash digest
SHA256 7cb6a7560ab827ddac96e9f6fd8aede58b3cf92a20800eb48916dc3be21c0a3e
MD5 092c237c452595d6b337250c166711ab
BLAKE2b-256 2f26d874fdcf8840a17475f12728fe1dff09a34926ff23bc20cdc019542ba835

See more details on using hashes here.

Provenance

The following attestation bundles were made for comfy_aimdo-0.5.2-cp39-abi3-win_arm64.whl:

Publisher: build-wheels.yml on Comfy-Org/comfy-aimdo

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file comfy_aimdo-0.5.2-cp39-abi3-win_amd64.whl.

File metadata

  • Download URL: comfy_aimdo-0.5.2-cp39-abi3-win_amd64.whl
  • Upload date:
  • Size: 273.4 kB
  • Tags: CPython 3.9+, Windows x86-64
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for comfy_aimdo-0.5.2-cp39-abi3-win_amd64.whl
Algorithm Hash digest
SHA256 5017987269c319dce8c05ac4ce72a1f18284922de56a40be7248cdd1ba22b74c
MD5 703b8687b823f4d79d32f033a1b95f0e
BLAKE2b-256 36d34fa4b3b5c533087ac5c5af0bd69b35b13c89e4abbcc5c164142d9cb4c727

See more details on using hashes here.

Provenance

The following attestation bundles were made for comfy_aimdo-0.5.2-cp39-abi3-win_amd64.whl:

Publisher: build-wheels.yml on Comfy-Org/comfy-aimdo

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file comfy_aimdo-0.5.2-cp39-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.

File metadata

File hashes

Hashes for comfy_aimdo-0.5.2-cp39-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl
Algorithm Hash digest
SHA256 56965d183b501c0ca0c7aac9336e7079ec0066211a402ac23ad64f40afde32b2
MD5 3f784e30ef240e9ed8330fd023eee446
BLAKE2b-256 9a84f7562dcf3be4eaf2bca67ecba451b01905e4c8d6b560e9a9f52d17b622aa

See more details on using hashes here.

Provenance

The following attestation bundles were made for comfy_aimdo-0.5.2-cp39-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl:

Publisher: build-wheels.yml on Comfy-Org/comfy-aimdo

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file comfy_aimdo-0.5.2-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl.

File metadata

File hashes

Hashes for comfy_aimdo-0.5.2-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl
Algorithm Hash digest
SHA256 abf86a49d5b86d7ede0a4a39f3f9e1ddff4efe60f6b34da2bb4df43c398a4ac6
MD5 0680f8e4abf742c7d486d3fa116c6066
BLAKE2b-256 8e67197d1f55b7a53c074f2027d892943c8c457a86535a99c7b37ceb126fa3be

See more details on using hashes here.

Provenance

The following attestation bundles were made for comfy_aimdo-0.5.2-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl:

Publisher: build-wheels.yml on Comfy-Org/comfy-aimdo

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Release history Release notifications | RSS feed

0.5.3

5 files

This release

0.5.2 This release

5 files

0.5.1

5 files

0.5.0

5 files

0.4.15

5 files

0.4.14

5 files

0.4.13

5 files

0.4.12

5 files

0.4.11

5 files

0.4.10

5 files

0.4.9

5 files

0.4.8

5 files

0.4.7

5 files

0.4.6

5 files

0.4.5

5 files

0.4.4

5 files

0.4.3

5 files

0.4.2

5 files

0.4.1

5 files

0.4.0

5 files

0.3.0

5 files

0.2.14

3 files

0.2.12

3 files

0.2.11

3 files

0.2.10

3 files

0.2.9

3 files

0.2.8

3 files

0.2.7

3 files

0.2.6

3 files

0.2.5

3 files

0.2.4

3 files

0.2.3

3 files

0.2.2

3 files

0.2.1

3 files

0.2.0

3 files

0.1.8

3 files

0.1.7

3 files

0.1.6

3 files

0.1.5

3 files

0.1.4

3 files

0.1.3

3 files

0.1.2

11 files

0.1.1

11 files

0.1.0

11 files

0.0.214

3 files

0.0.213

3 files

0.0.4

11 files

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