[docs] Memory optims #11385
Open
[docs] Memory optims #11385
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
AutoencoderKLWan and AsymmetricAutoencoderKL does not support tiling or slicing (asymmetric have just unused flags), it most likely should be mentioned. |
sayakpaul
reviewed
Apr 25, 2025
| Use the [`~DiffusionPipeline.reset_device_map`] method to reset the `device_map`. This is necessary if you want to use methods like `.to()`, [`~DiffusionPipeline.enable_sequential_cpu_offload`], and [`~DiffusionPipeline.enable_model_cpu_offload`] on a pipeline that was device-mapped. | ||
|
|
||
| ```py | ||
| pipeline.reset_device_map |
There was a problem hiding this comment.
Suggested change
| pipeline.reset_device_map | |
| pipeline.reset_device_map() |
|
|
||
| Model offloading moves entire models to the GPU instead of selectively moving *some* layers or model components. One of the main pipeline models, usually the text encoder, UNet, and VAE, is placed on the GPU while the other components are held on the CPU. Components like the UNet that run multiple times stays on the GPU until its completely finished and no longer needed. This eliminates the communication overhead of [CPU offloading](#cpu-offloading) and makes model offloading a faster alternative. The tradeoff is memory savings won't be as large. | ||
|
|
||
| > [!WARNING] |
There was a problem hiding this comment.
Do we want to add a warning after enable_sequential_cpu_offload() that it's terribly slow and can often appear to be impractical?
| Model offloading moves entire models to the GPU instead of selectively moving *some* layers or model components. One of the main pipeline models, usually the text encoder, UNet, and VAE, is placed on the GPU while the other components are held on the CPU. Components like the UNet that run multiple times stays on the GPU until its completely finished and no longer needed. This eliminates the communication overhead of [CPU offloading](#cpu-offloading) and makes model offloading a faster alternative. The tradeoff is memory savings won't be as large. | ||
|
|
||
| > [!WARNING] | ||
| > To properly offload models after they're called, it is required to run the entire pipeline and models in the expected order. Keep this in mind if models are reused outside the pipeline context after hooks have been installed (see [Removing Hooks](https://huggingface.co/docs/accelerate/en/package_reference/big_modeling#accelerate.hooks.remove_hook_from_module) for more details). This is a stateful operation that installs hooks on the model. |
There was a problem hiding this comment.
This reads a bit incomplete:
Keep this in mind if models are reused outside the pipeline context after hooks have been installed
There was a problem hiding this comment.
Also, not sure if it would make sense to include but users can still benefit from pipeline.enable_model_cpu_offload() when doing stuff like pipeline.encode_prompt(). See #11376
Comment on lines
+214
to
+215
| > [!WARNING] | ||
| > Group offloading may not work with all models if the forward implementation contains weight-dependent device casting of inputs because it may clash with group offloading's device casting mechanism. |
There was a problem hiding this comment.
Providing some example models would be helpful here. Cc: @a-r-r-o-w
| apply_group_offloading(pipe.vae, onload_device=onload_device, offload_type="leaf_level") | ||
| # Use the apply_group_offloading method for other model components | ||
| apply_group_offloading(pipeline.text_encoder, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=2) | ||
| apply_group_offloading(pipeline.vae, onload_device=onload_device, offload_type="leaf_level") |
There was a problem hiding this comment.
vae is a subclass of ModelMixin so, we should be able to use enable_group_offload().
|
|
||
| PyTorch supports `torch.float8_e4m3fn` and `torch.float8_e5m2` as weight storage dtypes, but they can't be used for computation in many different tensor operations due to unimplemented kernel support. However, you can use these dtypes to store model weights in fp8 precision and upcast them on-the-fly when the layers are used in the forward pass. This is known as layerwise weight-casting. | ||
| ## FP8 layerwise casting |
There was a problem hiding this comment.
Suggested change
| ## FP8 layerwise casting | |
| ## Layerwise casting |
As it's not specific to FP8.
|
|
||
| Typically, inference on most models is done with `torch.float16` or `torch.bfloat16` weight/computation precision. Layerwise weight-casting cuts down the memory footprint of the model weights by approximately half. | ||
| Layerwise casting stores weights in a smaller data format (`torch.float8_e4m3fn` and `torch.float8_e5m2`) to use less memory and upcasts those weights to `torch.float16` or `torch.bfloat16` for computation. Certain layers (normalization and modulation related weights) are skipped because storing them in fp8 can degrade generation quality. |
There was a problem hiding this comment.
Suggested change
| Layerwise casting stores weights in a smaller data format (`torch.float8_e4m3fn` and `torch.float8_e5m2`) to use less memory and upcasts those weights to `torch.float16` or `torch.bfloat16` for computation. Certain layers (normalization and modulation related weights) are skipped because storing them in fp8 can degrade generation quality. | |
| Layerwise casting stores weights in a smaller data format (for example: `torch.float8_e4m3fn` and `torch.float8_e5m2`) to use less memory and upcasts those weights to a higher precision, e.g., `torch.float16` or `torch.bfloat16` for computation. Certain layers (normalization and modulation related weights) are skipped because storing them in fp8 can degrade generation quality. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Refactors the memory optimization docs and combines it with working with big models (distributed setups).
Let me know if I'm missing anything!