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 only
  • Pytorch 2.8+
  • Cuda 12.8+
  • 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.

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.4.12-py3-none-any.whl (22.8 kB view details)

Uploaded Python 3

comfy_aimdo-0.4.12-cp39-abi3-win_arm64.whl (227.5 kB view details)

Uploaded CPython 3.9+Windows ARM64

comfy_aimdo-0.4.12-cp39-abi3-win_amd64.whl (253.6 kB view details)

Uploaded CPython 3.9+Windows x86-64

comfy_aimdo-0.4.12-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl (567.0 kB view details)

Uploaded CPython 3.9+manylinux: glibc 2.17+ ARM64

comfy_aimdo-0.4.12-cp39-abi3-manylinux2010_x86_64.manylinux2014_x86_64.manylinux_2_12_x86_64.manylinux_2_17_x86_64.whl (342.2 kB view details)

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

File details

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

File metadata

  • Download URL: comfy_aimdo-0.4.12-py3-none-any.whl
  • Upload date:
  • Size: 22.8 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.4.12-py3-none-any.whl
Algorithm Hash digest
SHA256 346c597012a7db53ee472a6675f8cdd6714f4e3c88e372ccbb881df16c78ed7c
MD5 0cec58bb756b49d9f430629fdfa6a204
BLAKE2b-256 e258e0d0c30a502fb77a8beff00142c779b08244bf997dfde09fd9754f099f60

See more details on using hashes here.

Provenance

The following attestation bundles were made for comfy_aimdo-0.4.12-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.4.12-cp39-abi3-win_arm64.whl.

File metadata

  • Download URL: comfy_aimdo-0.4.12-cp39-abi3-win_arm64.whl
  • Upload date:
  • Size: 227.5 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.4.12-cp39-abi3-win_arm64.whl
Algorithm Hash digest
SHA256 8b3cc36b62429ac3d8e0b6548a06611bead49ad5de3a9a92944760c2d100c402
MD5 aa1e42271ac771124885ceca84ce4f5e
BLAKE2b-256 3b4a0fefdf4041378ef123f1531378930095461780b282bfea36d0bb4dabb671

See more details on using hashes here.

Provenance

The following attestation bundles were made for comfy_aimdo-0.4.12-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.4.12-cp39-abi3-win_amd64.whl.

File metadata

  • Download URL: comfy_aimdo-0.4.12-cp39-abi3-win_amd64.whl
  • Upload date:
  • Size: 253.6 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.4.12-cp39-abi3-win_amd64.whl
Algorithm Hash digest
SHA256 9d1a48e08e287c07cb26a40f7b93376558f4f5b35f6300e3f707a568d2af7fc3
MD5 2d0792c8c85077ea9828004e7edfe872
BLAKE2b-256 084a2d9191a32df7262c63e16a1f37cb63d71523b4d00f57c9945224131234b9

See more details on using hashes here.

Provenance

The following attestation bundles were made for comfy_aimdo-0.4.12-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.4.12-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl.

File metadata

File hashes

Hashes for comfy_aimdo-0.4.12-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl
Algorithm Hash digest
SHA256 fa3dd26e5bf51e31e1b53762a565490f6571fc6b7b52aeaa4f45f422b50ace1e
MD5 3ed962d9e828652027bb41e85daff2ca
BLAKE2b-256 ef1b8691b070a9e38c258780bb68332c439f3232e274094374bdf3ef7bdc49c8

See more details on using hashes here.

Provenance

The following attestation bundles were made for comfy_aimdo-0.4.12-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.

File details

Details for the file comfy_aimdo-0.4.12-cp39-abi3-manylinux2010_x86_64.manylinux2014_x86_64.manylinux_2_12_x86_64.manylinux_2_17_x86_64.whl.

File metadata

File hashes

Hashes for comfy_aimdo-0.4.12-cp39-abi3-manylinux2010_x86_64.manylinux2014_x86_64.manylinux_2_12_x86_64.manylinux_2_17_x86_64.whl
Algorithm Hash digest
SHA256 5c9ec713b86e47e8a889e9f2e1a83670f93a24b58890517ea8ede55984f7954a
MD5 05c059edc5ace586d0de76444483662d
BLAKE2b-256 59e177a6cf7add393f4fad35e38ec6d1e0cfb51ca416efa66fa75b51de00b0df

See more details on using hashes here.

Provenance

The following attestation bundles were made for comfy_aimdo-0.4.12-cp39-abi3-manylinux2010_x86_64.manylinux2014_x86_64.manylinux_2_12_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.

Release history Release notifications | RSS feed

0.4.13

5 files

This release

0.4.12 This release

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

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page