Skip to content

Fix GPU memory allocation issues #276

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Apr 19, 2025
Merged

Conversation

Ednaordinary
Copy link
Contributor

Previously, the following would not properly clear all memory:

import gc, torch
del pipeline, transformer
torch.cuda.empty_cache()

This PR fixes that by resetting the internal quantized model state on __del__ (please note this is not actually called until the ref count for the transformer is 0, but this matches the behaviour of other projects like diffusers)

Since this calls the full reset() method of ModuleWrapper, it may have unintended consequences if multiple models are loaded at once (it's probably fine though)

@Ednaordinary Ednaordinary changed the title Fix memory allocation issues Fix GPU memory allocation issues Apr 13, 2025
@Ednaordinary
Copy link
Contributor Author

This may solve the following:
image

nunchaku-tech/ComfyUI-nunchaku#65 should be fixed by this PR (I do not use comfy ui, dunno how it manages models)

I believe nunchaku-tech/ComfyUI-nunchaku#57 is caused by memory pinning, a separate issue, which was an issue I had where pinned memory could not be allocated but still took up memory (even after the process had exited, had to fully reboot).
image

Can be solved by telling nunchaku not to pin memory:

import os
os.environ['NUNCHAKU_LOAD_METHOD'] = 'READNOPIN'

I am less sure about comfyui specifically since I don't use it, but providing the variable on server start may work

export NUNCHAKU_LOAD_METHOD='READNOPIN'
python3 main.py

Copy link
Collaborator

@lmxyy lmxyy left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR. Looks good to me @sxtyzhangzk Could you take a look?

@lmxyy
Copy link
Collaborator

lmxyy commented Apr 19, 2025

Passing all tests, approved.

@lmxyy lmxyy merged commit 9912f57 into nunchaku-tech:main Apr 19, 2025
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants