add offload to disk #124
add offload to disk #124
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @tenderness-git, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request introduces a robust system for predicting GPU memory usage of deep learning models, particularly VRAM. It provides the necessary infrastructure to profile a model's memory footprint under varying input conditions, train a linear regression model on this collected data, and then leverage the trained model to accurately estimate memory requirements for new, unseen inputs. This capability is crucial for optimizing resource management and preventing out-of-memory errors in large-scale deep learning deployments.
Highlights
- New Memory Prediction Module: Introduced a new
memoryutility package underdiffsynth_engine/utilsto provide tools for GPU memory profiling and prediction for deep learning models. - Custom Linear Regression Implementation: Added a
LinearRegressionclass inlinear_regression.pythat leverages PyTorch'storch.linalg.lstsqfor ordinary least squares fitting. This class includes methods for prediction, as well as serialization and deserialization for model persistence. - GPU Memory Profiling Utilities: Implemented functions in
memory_predcit_model.pyto accurately profile both activation memory (memory used during forward pass) and weight memory (memory consumed by model parameters) on CUDA devices. - Predictive Model for Memory Usage: Developed a
MemoryPredictModelclass that utilizes the customLinearRegressionmodel. This class can be trained on collected memory profiling data and then used to predict the total GPU memory consumption for a given model based on its input parameters. - Example Usage Scripts: Provided example scripts (
profile_model.pyandtrain_model.py) demonstrating the end-to-end workflow: profiling a model (e.g.,FluxDiT) to gather memory data, training theMemoryPredictModelwith this data, and subsequently using the trained model for memory prediction.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
tenderness-git
force-pushed
the
wangguan
branch
from
228dac6 to
fd3a489
Compare
| vae_tile_stride: int | Tuple[int, int] = 256 | ||
| device: str = "cuda" | ||
| offload_mode: Optional[str] = None | ||
| offload_to_disk: bool = False |
There was a problem hiding this comment.
感觉应该换个名字，看了一下实现我觉得跟offload是两个正交的能力，这个名字太困惑了。一般理解上的offload_to_disk是真的在磁盘上存点数据，然后通过一个方法能还原回来，但是现在实现上没有这么一个方法；而且我们的目的是节省一次性运行过程中的内存占用，应该也不需要这样的能力，结合model lifecycle换给名字会更好理解一些
| logger.warning("must set an non cpu device for pipeline before calling enable_cpu_offload") | ||
| return | ||
| if offload_mode == "cpu_offload": | ||
| if offload_mode is None or offload_mode == "disable": |
There was a problem hiding this comment.
都是在初始化的时候设置的，感觉disable这个选项没啥用呀
|
|
||
| def enable_cpu_offload(self, offload_mode: str): | ||
| valid_offload_mode = ("cpu_offload", "sequential_cpu_offload") | ||
| def enable_cpu_offload(self, offload_mode: str | None, offload_to_disk:bool = False): |
There was a problem hiding this comment.
另外是正交方法的话应该拆成另外一个函数，而不是作为这个函数的参数会更好一些
No description provided.