hf-mem 0.5.6__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- hf_mem-0.5.6/LICENSE +25 -0
- hf_mem-0.5.6/PKG-INFO +131 -0
- hf_mem-0.5.6/README.md +116 -0
- hf_mem-0.5.6/pyproject.toml +40 -0
- hf_mem-0.5.6/pyproject.toml.orig +41 -0
- hf_mem-0.5.6/src/hf_mem/__init__.py +4 -0
- hf_mem-0.5.6/src/hf_mem/_fetch.py +14 -0
- hf_mem-0.5.6/src/hf_mem/_print.py +133 -0
- hf_mem-0.5.6/src/hf_mem/_types.py +9 -0
- hf_mem-0.5.6/src/hf_mem/_version.py +5 -0
- hf_mem-0.5.6/src/hf_mem/cli.py +148 -0
- hf_mem-0.5.6/src/hf_mem/gguf/__init__.py +35 -0
- hf_mem-0.5.6/src/hf_mem/gguf/fetch.py +77 -0
- hf_mem-0.5.6/src/hf_mem/gguf/kv_cache.py +38 -0
- hf_mem-0.5.6/src/hf_mem/gguf/metadata.py +187 -0
- hf_mem-0.5.6/src/hf_mem/gguf/print.py +251 -0
- hf_mem-0.5.6/src/hf_mem/gguf/types.py +206 -0
- hf_mem-0.5.6/src/hf_mem/py.typed +0 -0
- hf_mem-0.5.6/src/hf_mem/run.py +612 -0
- hf_mem-0.5.6/src/hf_mem/safetensors/__init__.py +41 -0
- hf_mem-0.5.6/src/hf_mem/safetensors/fetch.py +61 -0
- hf_mem-0.5.6/src/hf_mem/safetensors/kv_cache.py +180 -0
- hf_mem-0.5.6/src/hf_mem/safetensors/metadata.py +199 -0
- hf_mem-0.5.6/src/hf_mem/safetensors/print.py +277 -0
- hf_mem-0.5.6/src/hf_mem/safetensors/types.py +58 -0
hf_mem-0.5.6/LICENSE
ADDED
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
Copyright (c) 2026-present Alvaro Bartolome <alvarobartt@gmail.com>
|
|
2
|
+
|
|
3
|
+
Permission is hereby granted, free of charge, to any
|
|
4
|
+
person obtaining a copy of this software and associated
|
|
5
|
+
documentation files (the "Software"), to deal in the
|
|
6
|
+
Software without restriction, including without
|
|
7
|
+
limitation the rights to use, copy, modify, merge,
|
|
8
|
+
publish, distribute, sublicense, and/or sell copies of
|
|
9
|
+
the Software, and to permit persons to whom the Software
|
|
10
|
+
is furnished to do so, subject to the following
|
|
11
|
+
conditions:
|
|
12
|
+
|
|
13
|
+
The above copyright notice and this permission notice
|
|
14
|
+
shall be included in all copies or substantial portions
|
|
15
|
+
of the Software.
|
|
16
|
+
|
|
17
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
|
|
18
|
+
ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
|
|
19
|
+
TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
|
|
20
|
+
PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
|
|
21
|
+
SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
|
|
22
|
+
CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
|
|
23
|
+
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
|
|
24
|
+
IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
|
|
25
|
+
DEALINGS IN THE SOFTWARE.
|
hf_mem-0.5.6/PKG-INFO
ADDED
|
@@ -0,0 +1,131 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: hf-mem
|
|
3
|
+
Version: 0.5.6
|
|
4
|
+
Summary: A CLI to estimate inference memory requirements for Hugging Face models, written in Python.
|
|
5
|
+
Author: Alvaro Bartolome
|
|
6
|
+
Author-email: Alvaro Bartolome <alvarobartt@gmail.com>
|
|
7
|
+
License-Expression: MIT
|
|
8
|
+
License-File: LICENSE
|
|
9
|
+
Requires-Dist: httpx2[http2]>=2.12.0
|
|
10
|
+
Requires-Python: >=3.10
|
|
11
|
+
Project-URL: Homepage, https://github.com/alvarobartt/hf-mem
|
|
12
|
+
Project-URL: Repository, https://github.com/alvarobartt/hf-mem
|
|
13
|
+
Project-URL: Issues, https://github.com/alvarobartt/hf-mem/issues
|
|
14
|
+
Description-Content-Type: text/markdown
|
|
15
|
+
|
|
16
|
+
<img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/hf-mem.png" />
|
|
17
|
+
|
|
18
|
+
---
|
|
19
|
+
|
|
20
|
+
> [!WARNING]
|
|
21
|
+
> `hf-mem` is still experimental and therefore subject to major changes across releases, so please keep in mind that breaking changes may occur until v1.0.0.
|
|
22
|
+
|
|
23
|
+
`hf-mem` is a CLI to estimate inference memory requirements for Hugging Face models, written in Python. `hf-mem` is lightweight, only depends on [`httpx2`](https://github.com/pydantic/httpx2), as it pulls the [Safetensors](https://github.com/huggingface/safetensors) and / or [GGUF](https://github.com/ggml-org/ggml/blob/master/docs/gguf.md) metadata via [HTTP Range requests](https://developer.mozilla.org/en-US/docs/Web/HTTP/Guides/Range_requests). It's recommended to run with [`uv`](https://github.com/astral-sh/uv) for a better experience.
|
|
24
|
+
|
|
25
|
+
`hf-mem` lets you estimate the inference requirements to run any model from the Hugging Face Hub, including [Transformers](https://github.com/huggingface/transformers), [Diffusers](https://github.com/huggingface/diffusers) and [Sentence Transformers](https://github.com/huggingface/sentence-transformers) models, or really any model as long as it contains any of [Safetensors](https://github.com/huggingface/safetensors) or [GGUF](https://github.com/ggml-org/ggml/blob/master/docs/gguf.md) weights.
|
|
26
|
+
|
|
27
|
+
Read more information about `hf-mem` in [this short-form post](https://alvarobartt.com/hf-mem), but note it's not up-to-date as it was written in January 2026.
|
|
28
|
+
|
|
29
|
+
## Usage
|
|
30
|
+
|
|
31
|
+
### CLI (Recommended)
|
|
32
|
+
|
|
33
|
+
#### Transformers
|
|
34
|
+
|
|
35
|
+
```bash
|
|
36
|
+
uvx hf-mem --model-id MiniMaxAI/MiniMax-M2
|
|
37
|
+
```
|
|
38
|
+
|
|
39
|
+
<img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/transformers.png" />
|
|
40
|
+
|
|
41
|
+
#### Diffusers
|
|
42
|
+
|
|
43
|
+
```bash
|
|
44
|
+
uvx hf-mem --model-id Qwen/Qwen-Image
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
<img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/diffusers.png" />
|
|
48
|
+
|
|
49
|
+
#### Sentence Transformers
|
|
50
|
+
|
|
51
|
+
```bash
|
|
52
|
+
uvx hf-mem --model-id google/embeddinggemma-300m
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
<img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/sentence-transformers.png" />
|
|
56
|
+
|
|
57
|
+
### Python
|
|
58
|
+
|
|
59
|
+
You can also run it programmatically with Python as:
|
|
60
|
+
|
|
61
|
+
```python
|
|
62
|
+
from hf_mem import run
|
|
63
|
+
|
|
64
|
+
result = run(model_id="MiniMaxAI/MiniMax-M2", experimental=True)
|
|
65
|
+
print(result)
|
|
66
|
+
# Result(model_id='MiniMaxAI/MiniMax-M2', revision='main', filename=None, memory=230121630720, kv_cache=24964497408, total_memory=255086128128, details=False)
|
|
67
|
+
```
|
|
68
|
+
|
|
69
|
+
If you're already inside an async application, use `arun(...)` instead:
|
|
70
|
+
|
|
71
|
+
```python
|
|
72
|
+
from hf_mem import arun
|
|
73
|
+
|
|
74
|
+
result = await arun(model_id="MiniMaxAI/MiniMax-M2", experimental=True)
|
|
75
|
+
print(result)
|
|
76
|
+
# Result(model_id='MiniMaxAI/MiniMax-M2', revision='main', filename=None, memory=230121630720, kv_cache=24964497408, total_memory=255086128128, details=False)
|
|
77
|
+
```
|
|
78
|
+
|
|
79
|
+
## Experimental
|
|
80
|
+
|
|
81
|
+
By enabling the `--experimental` flag, you can enable the KV Cache memory estimation for LLMs (`...ForCausalLM`) and VLMs (`...ForConditionalGeneration`), even including a custom `--max-model-len` (defaults to the `config.json` default), `--batch-size` (defaults to 1), and the `--kv-cache-dtype` (defaults to `auto` which means it uses the default data type set in `config.json` under `torch_dtype` or `dtype`, or rather from `quantization_config` when applicable).
|
|
82
|
+
|
|
83
|
+
Additionally, for Mixture of Experts (MoEs), the `--experimental` flag also adds a weights-only breakdown of the base model and all the experts combined.
|
|
84
|
+
|
|
85
|
+
```bash
|
|
86
|
+
uvx hf-mem --model-id MiniMaxAI/MiniMax-M2 --experimental
|
|
87
|
+
```
|
|
88
|
+
|
|
89
|
+
<img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/experimental.png" />
|
|
90
|
+
|
|
91
|
+
## GGUF
|
|
92
|
+
|
|
93
|
+
If the repository contains GGUF model weights, those will be listed by default (only if there are no Safetensors weights, otherwise the GGUFs will be ignored) and the memory will be estimated for each one of those; whereas if a specific file is provided, then the memory estimation will be targeted for that given file instead.
|
|
94
|
+
|
|
95
|
+
```bash
|
|
96
|
+
uvx hf-mem --model-id TheBloke/deepseek-llm-7B-chat-GGUF --experimental
|
|
97
|
+
```
|
|
98
|
+
|
|
99
|
+
<img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/gguf.png" />
|
|
100
|
+
|
|
101
|
+
Or if you want to only get the estimation on a given file:
|
|
102
|
+
|
|
103
|
+
```bash
|
|
104
|
+
uvx hf-mem --model-id TheBloke/deepseek-llm-7B-chat-GGUF --gguf-file deepseek-llm-7b-chat.Q2_K.gguf --experimental
|
|
105
|
+
```
|
|
106
|
+
|
|
107
|
+
<img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/gguf-file.png" />
|
|
108
|
+
|
|
109
|
+
## Skills
|
|
110
|
+
|
|
111
|
+
Optionally, you can add `hf-mem` as an agent skill, which allows the underlying coding agent to discover and use it when provided as a [`SKILL.md`](skills/hf-mem/SKILL.md), e.g., `.claude/skills/hf-mem/SKILL.md`.
|
|
112
|
+
|
|
113
|
+
More information can be found at [Anthropic Agent Skills and how to use them](https://github.com/anthropics/skills).
|
|
114
|
+
|
|
115
|
+
## Extensions
|
|
116
|
+
|
|
117
|
+
Optionally, you can also add `hf-mem` as an extension to the Hugging Face Hub CLI under `hf` as `hf mem ...`. First you need to install `hf` as explained in [Hugging Face CLI - "Getting started"](https://huggingface.co/docs/huggingface_hub/en/guides/cli). To add `hf-mem` as an extension, all you need to do is run:
|
|
118
|
+
|
|
119
|
+
```bash
|
|
120
|
+
hf extensions add alvarobartt/hf-mem
|
|
121
|
+
```
|
|
122
|
+
More information can be found at [Hugging Face Hub CLI - "How Python extensions are installed"](https://huggingface.co/docs/huggingface_hub/en/guides/cli-extensions#how-python-extensions-are-installed).
|
|
123
|
+
|
|
124
|
+
Once installed, you can use `hf-mem` via `hf` as `hf mem ...` along with the rest of Hugging Face extensions you have installed (which you can list via `hf extensions list`).
|
|
125
|
+
|
|
126
|
+
## References
|
|
127
|
+
|
|
128
|
+
- [Safetensors Metadata parsing](https://huggingface.co/docs/safetensors/en/metadata_parsing)
|
|
129
|
+
- [GGUF File format](https://github.com/ggml-org/ggml/blob/master/docs/gguf.md)
|
|
130
|
+
- [GGUF on the Hugging Face Hub](https://huggingface.co/docs/hub/en/gguf)
|
|
131
|
+
- [usgraphics - TR-100 Machine Report](https://github.com/usgraphics/usgc-machine-report)
|
hf_mem-0.5.6/README.md
ADDED
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
<img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/hf-mem.png" />
|
|
2
|
+
|
|
3
|
+
---
|
|
4
|
+
|
|
5
|
+
> [!WARNING]
|
|
6
|
+
> `hf-mem` is still experimental and therefore subject to major changes across releases, so please keep in mind that breaking changes may occur until v1.0.0.
|
|
7
|
+
|
|
8
|
+
`hf-mem` is a CLI to estimate inference memory requirements for Hugging Face models, written in Python. `hf-mem` is lightweight, only depends on [`httpx2`](https://github.com/pydantic/httpx2), as it pulls the [Safetensors](https://github.com/huggingface/safetensors) and / or [GGUF](https://github.com/ggml-org/ggml/blob/master/docs/gguf.md) metadata via [HTTP Range requests](https://developer.mozilla.org/en-US/docs/Web/HTTP/Guides/Range_requests). It's recommended to run with [`uv`](https://github.com/astral-sh/uv) for a better experience.
|
|
9
|
+
|
|
10
|
+
`hf-mem` lets you estimate the inference requirements to run any model from the Hugging Face Hub, including [Transformers](https://github.com/huggingface/transformers), [Diffusers](https://github.com/huggingface/diffusers) and [Sentence Transformers](https://github.com/huggingface/sentence-transformers) models, or really any model as long as it contains any of [Safetensors](https://github.com/huggingface/safetensors) or [GGUF](https://github.com/ggml-org/ggml/blob/master/docs/gguf.md) weights.
|
|
11
|
+
|
|
12
|
+
Read more information about `hf-mem` in [this short-form post](https://alvarobartt.com/hf-mem), but note it's not up-to-date as it was written in January 2026.
|
|
13
|
+
|
|
14
|
+
## Usage
|
|
15
|
+
|
|
16
|
+
### CLI (Recommended)
|
|
17
|
+
|
|
18
|
+
#### Transformers
|
|
19
|
+
|
|
20
|
+
```bash
|
|
21
|
+
uvx hf-mem --model-id MiniMaxAI/MiniMax-M2
|
|
22
|
+
```
|
|
23
|
+
|
|
24
|
+
<img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/transformers.png" />
|
|
25
|
+
|
|
26
|
+
#### Diffusers
|
|
27
|
+
|
|
28
|
+
```bash
|
|
29
|
+
uvx hf-mem --model-id Qwen/Qwen-Image
|
|
30
|
+
```
|
|
31
|
+
|
|
32
|
+
<img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/diffusers.png" />
|
|
33
|
+
|
|
34
|
+
#### Sentence Transformers
|
|
35
|
+
|
|
36
|
+
```bash
|
|
37
|
+
uvx hf-mem --model-id google/embeddinggemma-300m
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
<img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/sentence-transformers.png" />
|
|
41
|
+
|
|
42
|
+
### Python
|
|
43
|
+
|
|
44
|
+
You can also run it programmatically with Python as:
|
|
45
|
+
|
|
46
|
+
```python
|
|
47
|
+
from hf_mem import run
|
|
48
|
+
|
|
49
|
+
result = run(model_id="MiniMaxAI/MiniMax-M2", experimental=True)
|
|
50
|
+
print(result)
|
|
51
|
+
# Result(model_id='MiniMaxAI/MiniMax-M2', revision='main', filename=None, memory=230121630720, kv_cache=24964497408, total_memory=255086128128, details=False)
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
If you're already inside an async application, use `arun(...)` instead:
|
|
55
|
+
|
|
56
|
+
```python
|
|
57
|
+
from hf_mem import arun
|
|
58
|
+
|
|
59
|
+
result = await arun(model_id="MiniMaxAI/MiniMax-M2", experimental=True)
|
|
60
|
+
print(result)
|
|
61
|
+
# Result(model_id='MiniMaxAI/MiniMax-M2', revision='main', filename=None, memory=230121630720, kv_cache=24964497408, total_memory=255086128128, details=False)
|
|
62
|
+
```
|
|
63
|
+
|
|
64
|
+
## Experimental
|
|
65
|
+
|
|
66
|
+
By enabling the `--experimental` flag, you can enable the KV Cache memory estimation for LLMs (`...ForCausalLM`) and VLMs (`...ForConditionalGeneration`), even including a custom `--max-model-len` (defaults to the `config.json` default), `--batch-size` (defaults to 1), and the `--kv-cache-dtype` (defaults to `auto` which means it uses the default data type set in `config.json` under `torch_dtype` or `dtype`, or rather from `quantization_config` when applicable).
|
|
67
|
+
|
|
68
|
+
Additionally, for Mixture of Experts (MoEs), the `--experimental` flag also adds a weights-only breakdown of the base model and all the experts combined.
|
|
69
|
+
|
|
70
|
+
```bash
|
|
71
|
+
uvx hf-mem --model-id MiniMaxAI/MiniMax-M2 --experimental
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
<img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/experimental.png" />
|
|
75
|
+
|
|
76
|
+
## GGUF
|
|
77
|
+
|
|
78
|
+
If the repository contains GGUF model weights, those will be listed by default (only if there are no Safetensors weights, otherwise the GGUFs will be ignored) and the memory will be estimated for each one of those; whereas if a specific file is provided, then the memory estimation will be targeted for that given file instead.
|
|
79
|
+
|
|
80
|
+
```bash
|
|
81
|
+
uvx hf-mem --model-id TheBloke/deepseek-llm-7B-chat-GGUF --experimental
|
|
82
|
+
```
|
|
83
|
+
|
|
84
|
+
<img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/gguf.png" />
|
|
85
|
+
|
|
86
|
+
Or if you want to only get the estimation on a given file:
|
|
87
|
+
|
|
88
|
+
```bash
|
|
89
|
+
uvx hf-mem --model-id TheBloke/deepseek-llm-7B-chat-GGUF --gguf-file deepseek-llm-7b-chat.Q2_K.gguf --experimental
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
<img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/gguf-file.png" />
|
|
93
|
+
|
|
94
|
+
## Skills
|
|
95
|
+
|
|
96
|
+
Optionally, you can add `hf-mem` as an agent skill, which allows the underlying coding agent to discover and use it when provided as a [`SKILL.md`](skills/hf-mem/SKILL.md), e.g., `.claude/skills/hf-mem/SKILL.md`.
|
|
97
|
+
|
|
98
|
+
More information can be found at [Anthropic Agent Skills and how to use them](https://github.com/anthropics/skills).
|
|
99
|
+
|
|
100
|
+
## Extensions
|
|
101
|
+
|
|
102
|
+
Optionally, you can also add `hf-mem` as an extension to the Hugging Face Hub CLI under `hf` as `hf mem ...`. First you need to install `hf` as explained in [Hugging Face CLI - "Getting started"](https://huggingface.co/docs/huggingface_hub/en/guides/cli). To add `hf-mem` as an extension, all you need to do is run:
|
|
103
|
+
|
|
104
|
+
```bash
|
|
105
|
+
hf extensions add alvarobartt/hf-mem
|
|
106
|
+
```
|
|
107
|
+
More information can be found at [Hugging Face Hub CLI - "How Python extensions are installed"](https://huggingface.co/docs/huggingface_hub/en/guides/cli-extensions#how-python-extensions-are-installed).
|
|
108
|
+
|
|
109
|
+
Once installed, you can use `hf-mem` via `hf` as `hf mem ...` along with the rest of Hugging Face extensions you have installed (which you can list via `hf extensions list`).
|
|
110
|
+
|
|
111
|
+
## References
|
|
112
|
+
|
|
113
|
+
- [Safetensors Metadata parsing](https://huggingface.co/docs/safetensors/en/metadata_parsing)
|
|
114
|
+
- [GGUF File format](https://github.com/ggml-org/ggml/blob/master/docs/gguf.md)
|
|
115
|
+
- [GGUF on the Hugging Face Hub](https://huggingface.co/docs/hub/en/gguf)
|
|
116
|
+
- [usgraphics - TR-100 Machine Report](https://github.com/usgraphics/usgc-machine-report)
|
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "hf-mem"
|
|
3
|
+
version = "0.5.6"
|
|
4
|
+
description = "A CLI to estimate inference memory requirements for Hugging Face models, written in Python."
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
license = "MIT"
|
|
7
|
+
license-files = ["LICENSE"]
|
|
8
|
+
requires-python = ">=3.10"
|
|
9
|
+
dependencies = ["httpx2[http2]>=2.12.0"]
|
|
10
|
+
|
|
11
|
+
[[project.authors]]
|
|
12
|
+
name = "Alvaro Bartolome"
|
|
13
|
+
email = "alvarobartt@gmail.com"
|
|
14
|
+
|
|
15
|
+
[project.urls]
|
|
16
|
+
Homepage = "https://github.com/alvarobartt/hf-mem"
|
|
17
|
+
Repository = "https://github.com/alvarobartt/hf-mem"
|
|
18
|
+
Issues = "https://github.com/alvarobartt/hf-mem/issues"
|
|
19
|
+
|
|
20
|
+
[project.scripts]
|
|
21
|
+
hf-mem = "hf_mem.cli:main"
|
|
22
|
+
|
|
23
|
+
[build-system]
|
|
24
|
+
requires = ["uv_build>=0.9.5,<0.10.0"]
|
|
25
|
+
build-backend = "uv_build"
|
|
26
|
+
|
|
27
|
+
[dependency-groups]
|
|
28
|
+
lint = [
|
|
29
|
+
"pre-commit>=4.5.1",
|
|
30
|
+
"ruff>=0.14.10",
|
|
31
|
+
]
|
|
32
|
+
|
|
33
|
+
[tool.ruff]
|
|
34
|
+
line-length = 112
|
|
35
|
+
|
|
36
|
+
[tool.ruff.lint]
|
|
37
|
+
select = ["I"]
|
|
38
|
+
|
|
39
|
+
[tool.ruff.lint.isort]
|
|
40
|
+
force-single-line = false
|
|
@@ -0,0 +1,41 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "hf-mem"
|
|
3
|
+
version = "0.5.6"
|
|
4
|
+
description = "A CLI to estimate inference memory requirements for Hugging Face models, written in Python."
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
authors = [
|
|
7
|
+
{ name = "Alvaro Bartolome", email = "alvarobartt@gmail.com" }
|
|
8
|
+
]
|
|
9
|
+
license = "MIT"
|
|
10
|
+
license-files = ["LICENSE"]
|
|
11
|
+
requires-python = ">=3.10"
|
|
12
|
+
dependencies = [
|
|
13
|
+
"httpx2[http2]>=2.12.0",
|
|
14
|
+
]
|
|
15
|
+
|
|
16
|
+
[project.urls]
|
|
17
|
+
Homepage = "https://github.com/alvarobartt/hf-mem"
|
|
18
|
+
Repository = "https://github.com/alvarobartt/hf-mem"
|
|
19
|
+
Issues = "https://github.com/alvarobartt/hf-mem/issues"
|
|
20
|
+
|
|
21
|
+
[build-system]
|
|
22
|
+
requires = ["uv_build>=0.9.5,<0.10.0"]
|
|
23
|
+
build-backend = "uv_build"
|
|
24
|
+
|
|
25
|
+
[dependency-groups]
|
|
26
|
+
lint = [
|
|
27
|
+
"pre-commit>=4.5.1",
|
|
28
|
+
"ruff>=0.14.10",
|
|
29
|
+
]
|
|
30
|
+
|
|
31
|
+
[tool.ruff]
|
|
32
|
+
line-length = 112
|
|
33
|
+
|
|
34
|
+
[tool.ruff.lint]
|
|
35
|
+
select = ["I"]
|
|
36
|
+
|
|
37
|
+
[tool.ruff.lint.isort]
|
|
38
|
+
force-single-line = false
|
|
39
|
+
|
|
40
|
+
[project.scripts]
|
|
41
|
+
hf-mem = "hf_mem.cli:main"
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
import os
|
|
2
|
+
from typing import Any, Dict
|
|
3
|
+
|
|
4
|
+
import httpx2
|
|
5
|
+
|
|
6
|
+
# NOTE: Shared across all HTTP fetch operations; overridable via the `REQUEST_TIMEOUT` environment variable
|
|
7
|
+
REQUEST_TIMEOUT = float(os.getenv("REQUEST_TIMEOUT", 30.0))
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
# NOTE: Return type-hint set to `Any`, but it will only be a JSON-compatible object
|
|
11
|
+
async def get_json_file(client: httpx2.AsyncClient, url: str, headers: Dict[str, str] | None = None) -> Any:
|
|
12
|
+
response = await client.get(url, headers=headers, timeout=REQUEST_TIMEOUT)
|
|
13
|
+
response.raise_for_status()
|
|
14
|
+
return response.json()
|
|
@@ -0,0 +1,133 @@
|
|
|
1
|
+
from typing import Literal
|
|
2
|
+
|
|
3
|
+
MIN_NAME_LEN = 5
|
|
4
|
+
MAX_NAME_LEN = 14
|
|
5
|
+
MAX_DATA_LEN = 64
|
|
6
|
+
BORDERS_AND_PADDING = 4
|
|
7
|
+
|
|
8
|
+
BOX = {
|
|
9
|
+
"tl": "┌",
|
|
10
|
+
"tr": "┐",
|
|
11
|
+
"bl": "└",
|
|
12
|
+
"br": "┘",
|
|
13
|
+
"ht": "─",
|
|
14
|
+
"vt": "│",
|
|
15
|
+
"tsep": "┬",
|
|
16
|
+
"bsep": "┴",
|
|
17
|
+
"lm": "├",
|
|
18
|
+
"rm": "┤",
|
|
19
|
+
"mm": "┼",
|
|
20
|
+
}
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def _print_with_color(content: str) -> None:
|
|
24
|
+
print(f"\x1b[38;2;244;183;63m{content}\x1b[0m")
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def _print_header(current_len: int, badge: str | None = None) -> None:
|
|
28
|
+
length = current_len + 2 * BORDERS_AND_PADDING + 2
|
|
29
|
+
if badge:
|
|
30
|
+
badge_text = f" {badge} "
|
|
31
|
+
fill_len = max(0, length - 2 - len(badge_text))
|
|
32
|
+
top = BOX["tl"] + (BOX["tsep"] * fill_len) + badge_text + BOX["tr"]
|
|
33
|
+
else:
|
|
34
|
+
top = BOX["tl"] + (BOX["tsep"] * (length - 2)) + BOX["tr"]
|
|
35
|
+
_print_with_color(top)
|
|
36
|
+
|
|
37
|
+
bottom = BOX["lm"] + (BOX["bsep"] * (length - 2)) + BOX["rm"]
|
|
38
|
+
_print_with_color(bottom)
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def _print_centered(text: str, current_len: int) -> None:
|
|
42
|
+
total_width = current_len + 2 * BORDERS_AND_PADDING
|
|
43
|
+
text_len = len(text)
|
|
44
|
+
pad_left = (total_width - text_len) // 2
|
|
45
|
+
pad_right = total_width - text_len - pad_left
|
|
46
|
+
_print_with_color(f"{BOX['vt']}{' ' * pad_left}{text}{' ' * pad_right}{BOX['vt']}")
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
def _print_divider(
|
|
50
|
+
current_len: int,
|
|
51
|
+
side: Literal["top", "top-continue", "mid", "bottom", "bottom-continue"] | None = None,
|
|
52
|
+
name_len: int = MAX_NAME_LEN,
|
|
53
|
+
) -> None:
|
|
54
|
+
match side:
|
|
55
|
+
case "top":
|
|
56
|
+
left, mid, right = BOX["lm"], BOX["tsep"], BOX["rm"]
|
|
57
|
+
case "top-continue":
|
|
58
|
+
left, mid, right = BOX["lm"], BOX["bsep"], BOX["rm"]
|
|
59
|
+
case "bottom":
|
|
60
|
+
left, mid, right = BOX["bl"], BOX["bsep"], BOX["br"]
|
|
61
|
+
case "bottom-continue":
|
|
62
|
+
left, mid, right = BOX["lm"], BOX["bsep"], BOX["rm"]
|
|
63
|
+
case "mid":
|
|
64
|
+
left, mid, right = BOX["lm"], BOX["ht"], BOX["rm"]
|
|
65
|
+
case _:
|
|
66
|
+
left, mid, right = BOX["lm"], BOX["mm"], BOX["rm"]
|
|
67
|
+
|
|
68
|
+
name_col_inner = name_len + 2
|
|
69
|
+
data_col_inner = current_len + 1
|
|
70
|
+
|
|
71
|
+
line = left
|
|
72
|
+
line += BOX["ht"] * name_col_inner
|
|
73
|
+
line += mid
|
|
74
|
+
line += BOX["ht"] * data_col_inner
|
|
75
|
+
line += right
|
|
76
|
+
_print_with_color(line)
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
def _print_full_divider(
|
|
80
|
+
current_len: int,
|
|
81
|
+
side: Literal["top", "top-continue", "bottom", "bottom-continue", "mid"] | None = None,
|
|
82
|
+
) -> None:
|
|
83
|
+
total_inner = current_len + 2 * BORDERS_AND_PADDING
|
|
84
|
+
|
|
85
|
+
match side:
|
|
86
|
+
case "top":
|
|
87
|
+
left, right = BOX["lm"], BOX["rm"]
|
|
88
|
+
case "top-continue":
|
|
89
|
+
left, right = BOX["lm"], BOX["rm"]
|
|
90
|
+
case "bottom":
|
|
91
|
+
left, right = BOX["bl"], BOX["br"]
|
|
92
|
+
case "bottom-continue":
|
|
93
|
+
left, right = BOX["lm"], BOX["rm"]
|
|
94
|
+
case _:
|
|
95
|
+
left, right = BOX["lm"], BOX["rm"]
|
|
96
|
+
|
|
97
|
+
_print_with_color(left + BOX["ht"] * total_inner + right)
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
def _format_name(name: str, max_len: int = MAX_NAME_LEN) -> str:
|
|
101
|
+
if len(name) < MIN_NAME_LEN:
|
|
102
|
+
return f"{name:<{MIN_NAME_LEN}}"
|
|
103
|
+
if len(name) > max_len:
|
|
104
|
+
return name[: max_len - 3] + "..."
|
|
105
|
+
return f"{name:<{max_len}}"
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
def _print_row(name: str, text: str, current_len: int, name_len: int = MAX_NAME_LEN) -> None:
|
|
109
|
+
name_fmt = _format_name(name, name_len)
|
|
110
|
+
data_fmt = f"{str(text):<{current_len}}"
|
|
111
|
+
_print_with_color(f"{BOX['vt']} {name_fmt} {BOX['vt']} {data_fmt} {BOX['vt']}")
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
def _make_bar(used: float, total: float, width: int) -> str:
|
|
115
|
+
if total <= 0:
|
|
116
|
+
return "░" * width
|
|
117
|
+
frac = min(max(used / total, 0.0), 1.0)
|
|
118
|
+
filled = int(round(frac * width))
|
|
119
|
+
filled = max(0, min(width, filled))
|
|
120
|
+
return "█" * filled + "░" * (width - filled)
|
|
121
|
+
|
|
122
|
+
|
|
123
|
+
def _format_short_number(n: float) -> str:
|
|
124
|
+
n = float(n)
|
|
125
|
+
for unit in ("", "K", "M", "B", "T"):
|
|
126
|
+
if abs(n) < 1000.0:
|
|
127
|
+
return f"{int(n)}" if unit == "" else f"{n:.2f}{unit}"
|
|
128
|
+
n /= 1000.0
|
|
129
|
+
return f"{n:.2f}P"
|
|
130
|
+
|
|
131
|
+
|
|
132
|
+
def _bytes_to_gib(nbytes: int) -> float:
|
|
133
|
+
return nbytes / (1024**3)
|
|
@@ -0,0 +1,148 @@
|
|
|
1
|
+
import argparse
|
|
2
|
+
import asyncio
|
|
3
|
+
import json
|
|
4
|
+
import warnings
|
|
5
|
+
|
|
6
|
+
from hf_mem.gguf.print import print_gguf_files_report, print_gguf_report
|
|
7
|
+
from hf_mem.gguf.types import GGUFDtype
|
|
8
|
+
from hf_mem.run import Result, arun
|
|
9
|
+
from hf_mem.safetensors.kv_cache import KV_CACHE_DTYPE_CHOICES
|
|
10
|
+
from hf_mem.safetensors.print import print_safetensors_report
|
|
11
|
+
|
|
12
|
+
VALID_KV_CACHE_DTYPES = set(KV_CACHE_DTYPE_CHOICES) | set(GGUFDtype.__members__.keys())
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def _print_result(result: Result) -> None:
|
|
16
|
+
if result.safetensors is not None:
|
|
17
|
+
print_safetensors_report(
|
|
18
|
+
model_id=result.model_id,
|
|
19
|
+
revision=result.revision,
|
|
20
|
+
metadata=result.safetensors,
|
|
21
|
+
kv_cache=result.kv_cache_metadata,
|
|
22
|
+
moe=result.moe_metadata,
|
|
23
|
+
)
|
|
24
|
+
return
|
|
25
|
+
|
|
26
|
+
if result.gguf_files is not None:
|
|
27
|
+
if result.filename is not None:
|
|
28
|
+
print_gguf_report(
|
|
29
|
+
model_id=result.model_id,
|
|
30
|
+
filename=result.filename,
|
|
31
|
+
revision=result.revision,
|
|
32
|
+
metadata=result.gguf_files[result.filename],
|
|
33
|
+
kv_cache=result.gguf_files[result.filename].kv_cache,
|
|
34
|
+
)
|
|
35
|
+
else:
|
|
36
|
+
print_gguf_files_report(
|
|
37
|
+
model_id=result.model_id,
|
|
38
|
+
revision=result.revision,
|
|
39
|
+
gguf_files=result.gguf_files,
|
|
40
|
+
memory=result.memory, # type: ignore[arg-type]
|
|
41
|
+
kv_cache=result.kv_cache, # type: ignore[arg-type]
|
|
42
|
+
)
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
def main() -> None:
|
|
46
|
+
parser = argparse.ArgumentParser()
|
|
47
|
+
|
|
48
|
+
parser.add_argument("--model-id", required=True, help="Model ID on the Hugging Face Hub")
|
|
49
|
+
parser.add_argument(
|
|
50
|
+
"--revision",
|
|
51
|
+
default="main",
|
|
52
|
+
help="Model revision on the Hugging Face Hub",
|
|
53
|
+
)
|
|
54
|
+
parser.add_argument(
|
|
55
|
+
"--hf-token",
|
|
56
|
+
type=str,
|
|
57
|
+
default=None,
|
|
58
|
+
required=False,
|
|
59
|
+
help="The Hugging Face Hub token to use when sending requests to the Hub. If not provided it will be retrieved from either the `HF_TOKEN` environment variable or from the `.cache/huggingface/token` file if applicable.",
|
|
60
|
+
)
|
|
61
|
+
|
|
62
|
+
parser.add_argument(
|
|
63
|
+
"--experimental",
|
|
64
|
+
action="store_true",
|
|
65
|
+
help=(
|
|
66
|
+
"Whether to enable the experimental KV Cache estimation or not. For Safetensors "
|
|
67
|
+
"MoE Transformers models it also includes a base-model-vs-experts weights "
|
|
68
|
+
"breakdown. Only applies to `...ForCausalLM` and `...ForConditionalGeneration` "
|
|
69
|
+
"models from Transformers."
|
|
70
|
+
),
|
|
71
|
+
)
|
|
72
|
+
parser.add_argument(
|
|
73
|
+
"--max-model-len",
|
|
74
|
+
type=int,
|
|
75
|
+
default=None,
|
|
76
|
+
# NOTE: https://docs.vllm.ai/en/stable/configuration/engine_args/#-max-model-len
|
|
77
|
+
help="Model context length (prompt and output). If unspecified, will be automatically derived from the model config.",
|
|
78
|
+
)
|
|
79
|
+
parser.add_argument(
|
|
80
|
+
"--batch-size",
|
|
81
|
+
type=int,
|
|
82
|
+
default=1,
|
|
83
|
+
help="Batch size to help estimate the required RAM for caching when running the inference. Defaults to 1.",
|
|
84
|
+
)
|
|
85
|
+
parser.add_argument(
|
|
86
|
+
"--kv-cache-dtype",
|
|
87
|
+
type=str,
|
|
88
|
+
default="auto",
|
|
89
|
+
# NOTE: https://docs.vllm.ai/en/stable/cli/serve/#-kv-cache-dtype
|
|
90
|
+
help=f"Data type for the KV cache storage. If `auto` is specified, it will use the default model dtype specified in the `config.json` (if available) or F16 for GGUF files. Despite the FP8 data types having different formats, all those take 1 byte, meaning that the calculation would lead to the same results. Valid values are {KV_CACHE_DTYPE_CHOICES} for safetensors files and {['auto'] + list(GGUFDtype.__members__.keys())} for GGUF files. Defaults to `auto`.",
|
|
91
|
+
)
|
|
92
|
+
|
|
93
|
+
parser.add_argument(
|
|
94
|
+
"--json-output",
|
|
95
|
+
action="store_true",
|
|
96
|
+
help="Whether to provide the output as a JSON instead of printed as table.",
|
|
97
|
+
)
|
|
98
|
+
parser.add_argument(
|
|
99
|
+
"--ignore-table-width",
|
|
100
|
+
action="store_true",
|
|
101
|
+
help="Deprecated: table width now auto-expands when needed, so this flag is ignored.",
|
|
102
|
+
)
|
|
103
|
+
parser.add_argument(
|
|
104
|
+
"--gguf-file",
|
|
105
|
+
type=str,
|
|
106
|
+
default=None,
|
|
107
|
+
help="Specific GGUF file to estimate. If not provided, all GGUF files found in the repo will be estimated. The path within the repo is required (e.g. Q4_K_M/model-Q4_K_M.gguf).",
|
|
108
|
+
)
|
|
109
|
+
parser.add_argument(
|
|
110
|
+
"--details",
|
|
111
|
+
action="store_true",
|
|
112
|
+
help="Include per-component dtype and parameter breakdowns in JSON output (`--json-output` only).",
|
|
113
|
+
)
|
|
114
|
+
args = parser.parse_args()
|
|
115
|
+
|
|
116
|
+
if args.experimental:
|
|
117
|
+
warnings.warn(
|
|
118
|
+
"`--experimental` is set, which means that models with an architecture as `...ForCausalLM` and `...ForConditionalGeneration` will include estimations for the KV Cache as well. You can also provide the args `--max-model-len` and `--batch-size` as part of the estimation. Note that enabling `--experimental` means that the output will be different both when displayed and when dumped as JSON with `--json-output`, so bear that in mind."
|
|
119
|
+
)
|
|
120
|
+
|
|
121
|
+
if args.ignore_table_width:
|
|
122
|
+
warnings.warn(
|
|
123
|
+
"`--ignore-table-width` is deprecated and has no effect; table width auto-expands when needed by default."
|
|
124
|
+
)
|
|
125
|
+
|
|
126
|
+
if args.kv_cache_dtype not in VALID_KV_CACHE_DTYPES:
|
|
127
|
+
raise RuntimeError(
|
|
128
|
+
f"--kv-cache-dtype={args.kv_cache_dtype} not recognized. Valid options: {VALID_KV_CACHE_DTYPES}."
|
|
129
|
+
)
|
|
130
|
+
|
|
131
|
+
result = asyncio.run(
|
|
132
|
+
arun(
|
|
133
|
+
model_id=args.model_id,
|
|
134
|
+
revision=args.revision,
|
|
135
|
+
hf_token=args.hf_token,
|
|
136
|
+
experimental=args.experimental,
|
|
137
|
+
max_model_len=args.max_model_len,
|
|
138
|
+
batch_size=args.batch_size,
|
|
139
|
+
kv_cache_dtype=args.kv_cache_dtype,
|
|
140
|
+
gguf_file=args.gguf_file,
|
|
141
|
+
details=args.details,
|
|
142
|
+
)
|
|
143
|
+
)
|
|
144
|
+
|
|
145
|
+
if args.json_output:
|
|
146
|
+
print(json.dumps(result.to_json()))
|
|
147
|
+
else:
|
|
148
|
+
_print_result(result)
|