pytorch-precision-converter 2.0.0__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.
- pytorch_precision_converter-2.0.0/LICENSE +9 -0
- pytorch_precision_converter-2.0.0/PKG-INFO +191 -0
- pytorch_precision_converter-2.0.0/README.md +165 -0
- pytorch_precision_converter-2.0.0/precision_converter/__init__.py +27 -0
- pytorch_precision_converter-2.0.0/precision_converter/__main__.py +5 -0
- pytorch_precision_converter-2.0.0/precision_converter/cli.py +181 -0
- pytorch_precision_converter-2.0.0/precision_converter/core.py +222 -0
- pytorch_precision_converter-2.0.0/pyproject.toml +44 -0
- pytorch_precision_converter-2.0.0/pytorch_precision_converter.egg-info/PKG-INFO +191 -0
- pytorch_precision_converter-2.0.0/pytorch_precision_converter.egg-info/SOURCES.txt +15 -0
- pytorch_precision_converter-2.0.0/pytorch_precision_converter.egg-info/dependency_links.txt +1 -0
- pytorch_precision_converter-2.0.0/pytorch_precision_converter.egg-info/entry_points.txt +2 -0
- pytorch_precision_converter-2.0.0/pytorch_precision_converter.egg-info/requires.txt +6 -0
- pytorch_precision_converter-2.0.0/pytorch_precision_converter.egg-info/top_level.txt +1 -0
- pytorch_precision_converter-2.0.0/setup.cfg +4 -0
- pytorch_precision_converter-2.0.0/tests/test_cli.py +116 -0
- pytorch_precision_converter-2.0.0/tests/test_core.py +165 -0
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
# MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2023 Angelo Lamonaca
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
|
|
6
|
+
|
|
7
|
+
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
|
|
8
|
+
|
|
9
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
|
@@ -0,0 +1,191 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: pytorch-precision-converter
|
|
3
|
+
Version: 2.0.0
|
|
4
|
+
Summary: Convert PyTorch checkpoints and safetensors files to fp16, bf16 or fp8 without breaking them.
|
|
5
|
+
Author: Angelo Lamonaca
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/angelolamonaca/PyTorch-Precision-Converter
|
|
8
|
+
Project-URL: Issues, https://github.com/angelolamonaca/PyTorch-Precision-Converter/issues
|
|
9
|
+
Project-URL: Changelog, https://github.com/angelolamonaca/PyTorch-Precision-Converter/blob/main/CHANGELOG.md
|
|
10
|
+
Keywords: pytorch,safetensors,fp16,bf16,fp8,checkpoint,quantization,stable-diffusion,ema
|
|
11
|
+
Classifier: Development Status :: 5 - Production/Stable
|
|
12
|
+
Classifier: Environment :: Console
|
|
13
|
+
Classifier: Intended Audience :: Developers
|
|
14
|
+
Classifier: Intended Audience :: Science/Research
|
|
15
|
+
Classifier: Programming Language :: Python :: 3
|
|
16
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
17
|
+
Requires-Python: >=3.9
|
|
18
|
+
Description-Content-Type: text/markdown
|
|
19
|
+
License-File: LICENSE
|
|
20
|
+
Requires-Dist: torch>=2.1
|
|
21
|
+
Requires-Dist: safetensors>=0.4.3
|
|
22
|
+
Requires-Dist: numpy
|
|
23
|
+
Provides-Extra: test
|
|
24
|
+
Requires-Dist: pytest>=7; extra == "test"
|
|
25
|
+
Dynamic: license-file
|
|
26
|
+
|
|
27
|
+
# PyTorch Precision Converter
|
|
28
|
+
|
|
29
|
+
[](https://pypi.org/project/pytorch-precision-converter/)
|
|
30
|
+
[](https://github.com/angelolamonaca/PyTorch-Precision-Converter/actions/workflows/ci.yml)
|
|
31
|
+

|
|
32
|
+

|
|
33
|
+
[](https://github.com/angelolamonaca/PyTorch-Precision-Converter/blob/main/LICENSE)
|
|
34
|
+
|
|
35
|
+
Make PyTorch checkpoints and `safetensors` files **half the size (a quarter with fp8)** by converting them to
|
|
36
|
+
`fp16`, `bf16` or `fp8`, and strip the EMA weights you don't need for inference. One command, no model code
|
|
37
|
+
needed, runs on the CPU.
|
|
38
|
+
|
|
39
|
+
```console
|
|
40
|
+
$ precision-convert model.ckpt -p fp16 -t ema-only -st
|
|
41
|
+
Converted 280 tensors to fp16, left 1 unchanged, dropped 242 EMA entries
|
|
42
|
+
Used the EMA weights for 240 tensors
|
|
43
|
+
Saved model-ema-only-fp16.safetensors (29.93 MB, 27% of 109.23 MB)
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
## Why not just call `.half()`?
|
|
47
|
+
|
|
48
|
+
A one-line cast quietly breaks things. This tool handles the details:
|
|
49
|
+
|
|
50
|
+
- **Only floating-point tensors are cast.** Integer and bool tensors like `position_ids` or step counters keep
|
|
51
|
+
their dtype instead of turning into floats.
|
|
52
|
+
- **No silent `inf` or `NaN`.** fp16 tops out at 65,504 and fp8 at 448. Values beyond the range are clamped
|
|
53
|
+
to it and you get a warning, rather than infinities appearing in your weights.
|
|
54
|
+
- **Real EMA extraction.** `-t ema-only` puts the EMA weights where the model loads them from, for Stable
|
|
55
|
+
Diffusion and other latent-diffusion checkpoints (`model_ema.*`), then drops the duplicates.
|
|
56
|
+
- **safetensors metadata survives.** LoRA training info (`ss_*` keys) and other metadata is copied to the output.
|
|
57
|
+
- **Safe by default.** PyTorch checkpoints load with `weights_only=True`, so a malicious `.ckpt` can't run code.
|
|
58
|
+
Files that need full unpickling require an explicit `--unsafe-load`.
|
|
59
|
+
- **Your files are safe too.** Output is written to a temporary file and moved into place when complete, the
|
|
60
|
+
input is never replaced, and existing files are only overwritten with `--overwrite`.
|
|
61
|
+
- **Mixed precision when you need it.** `--keep-original "first_stage_model.*"` keeps matching tensors (here the
|
|
62
|
+
VAE) in their original dtype.
|
|
63
|
+
|
|
64
|
+
## Install
|
|
65
|
+
|
|
66
|
+
```bash
|
|
67
|
+
pip install pytorch-precision-converter
|
|
68
|
+
```
|
|
69
|
+
|
|
70
|
+
For the latest unreleased changes, install from GitHub instead:
|
|
71
|
+
`pip install git+https://github.com/angelolamonaca/PyTorch-Precision-Converter`.
|
|
72
|
+
|
|
73
|
+
Either way you get the `precision-convert` command (`python -m precision_converter` works too). It needs Python 3.9+,
|
|
74
|
+
PyTorch 2.1+ and `safetensors`.
|
|
75
|
+
|
|
76
|
+
## Usage
|
|
77
|
+
|
|
78
|
+
```bash
|
|
79
|
+
precision-convert INPUT [-p PRECISION] [-t TYPE] [-st] [-o OUTPUT] [options]
|
|
80
|
+
```
|
|
81
|
+
|
|
82
|
+
| Option | What it does |
|
|
83
|
+
|---|---|
|
|
84
|
+
| `INPUT` | `.safetensors` file or PyTorch checkpoint (`.ckpt`, `.pt`, `.pth`, `.bin`). `-f INPUT` works too. |
|
|
85
|
+
| `-p`, `--precision` | `fp32`, `fp16`, `bf16`, `fp8_e4m3fn`, `fp8_e5m2`, or `keep` (default: change no dtypes). |
|
|
86
|
+
| `-t`, `--type` | `full` keeps everything (default). `ema-only` uses the EMA weights. `no-ema` drops the EMA weights. |
|
|
87
|
+
| `-st`, `--safe-tensors` | Save as `.safetensors` instead of a PyTorch `.ckpt`. |
|
|
88
|
+
| `-o`, `--output` | Output path; its extension picks the format. Default: `<input>-<type>-<precision>` next to the input. |
|
|
89
|
+
| `--keep-original PATTERN` | Leave tensors whose key matches the glob in their original dtype. Repeatable. |
|
|
90
|
+
| `--info` | Show the tensors, dtypes, sizes, EMA and metadata in the file, then exit. |
|
|
91
|
+
| `--unsafe-load` | Fully unpickle PyTorch checkpoints. Only for files you trust. |
|
|
92
|
+
| `--overwrite` | Replace the output file if it exists. |
|
|
93
|
+
| `-v`, `--verbose` | List each tensor that had values clamped. |
|
|
94
|
+
|
|
95
|
+
## Recipes
|
|
96
|
+
|
|
97
|
+
**Shrink a Stable Diffusion training checkpoint for inference** (EMA weights, fp16, safetensors):
|
|
98
|
+
|
|
99
|
+
```bash
|
|
100
|
+
precision-convert model.ckpt -p fp16 -t ema-only -st
|
|
101
|
+
```
|
|
102
|
+
|
|
103
|
+
**Convert a LoRA to fp16 and keep its training metadata:**
|
|
104
|
+
|
|
105
|
+
```bash
|
|
106
|
+
precision-convert my-lora.safetensors -p fp16 -o my-lora-fp16.safetensors
|
|
107
|
+
```
|
|
108
|
+
|
|
109
|
+
**Store a large model in fp8** for UIs that load fp8 weights, such as ComfyUI, keeping normalization layers
|
|
110
|
+
precise:
|
|
111
|
+
|
|
112
|
+
```bash
|
|
113
|
+
precision-convert model.safetensors -p fp8_e4m3fn -st --keep-original "*norm*"
|
|
114
|
+
```
|
|
115
|
+
|
|
116
|
+
**Keep the VAE in fp32** (fp16 VAEs can produce NaNs and black images in some models):
|
|
117
|
+
|
|
118
|
+
```bash
|
|
119
|
+
precision-convert model.safetensors -p fp16 -st --keep-original "first_stage_model.*"
|
|
120
|
+
```
|
|
121
|
+
|
|
122
|
+
**Turn a pickled checkpoint into safetensors without changing precision:**
|
|
123
|
+
|
|
124
|
+
```bash
|
|
125
|
+
precision-convert model.ckpt -st
|
|
126
|
+
```
|
|
127
|
+
|
|
128
|
+
**Check what's inside a file before converting:**
|
|
129
|
+
|
|
130
|
+
```console
|
|
131
|
+
$ precision-convert model.ckpt --info
|
|
132
|
+
model.ckpt: 523 tensors, 27,274,319 values, 109.10 MB
|
|
133
|
+
float32 521 tensors 109.10 MB
|
|
134
|
+
int64 1 tensor 616 B
|
|
135
|
+
int32 1 tensor 4 B
|
|
136
|
+
EMA: 242 model_ema.* tensors (-t ema-only or -t no-ema removes them)
|
|
137
|
+
```
|
|
138
|
+
|
|
139
|
+
## Which precision should I pick?
|
|
140
|
+
|
|
141
|
+
| Precision | Bytes per value | Largest value | Good for |
|
|
142
|
+
|---|---|---|---|
|
|
143
|
+
| `fp32` | 4 | 3.4 × 10³⁸ | Training, reference copies, upcasting half-precision files |
|
|
144
|
+
| `bf16` | 2 | 3.4 × 10³⁸ | The safe 2-byte choice: same range as fp32, so nothing overflows |
|
|
145
|
+
| `fp16` | 2 | 65,504 | GPUs without fast bf16, and tools that expect fp16 |
|
|
146
|
+
| `fp8_e4m3fn` | 1 | 448 | Smallest files, for inference in tools that load fp8 weights; some quality loss |
|
|
147
|
+
| `fp8_e5m2` | 1 | 57,344 | fp8 with more range and less precision than `e4m3fn` |
|
|
148
|
+
|
|
149
|
+
The fp8 options cast each value directly, without per-tensor scaling factors.
|
|
150
|
+
|
|
151
|
+
## Python API
|
|
152
|
+
|
|
153
|
+
```python
|
|
154
|
+
from precision_converter import convert_state_dict, load_checkpoint, save_checkpoint
|
|
155
|
+
|
|
156
|
+
state_dict, metadata = load_checkpoint("model.safetensors")
|
|
157
|
+
converted, report = convert_state_dict(state_dict, precision="bf16", conversion="no-ema")
|
|
158
|
+
save_checkpoint(converted, "model-bf16.safetensors", metadata=metadata)
|
|
159
|
+
|
|
160
|
+
print(f"{report.converted} tensors converted, {len(report.clamped)} had values clamped")
|
|
161
|
+
```
|
|
162
|
+
|
|
163
|
+
`convert_state_dict` takes any mapping of tensors, so `model.state_dict()` works too.
|
|
164
|
+
|
|
165
|
+
## Upgrading from 1.x
|
|
166
|
+
|
|
167
|
+
`python convert.py ...` still works, but 2.0 fixes bugs that changed what earlier versions wrote:
|
|
168
|
+
|
|
169
|
+
- `-t ema-only` used to keep the regular weights and store the EMA weights under keys no model loads. It now
|
|
170
|
+
outputs the actual EMA weights.
|
|
171
|
+
- Integer tensors are no longer cast to floating point, and overflowing values are clamped instead of
|
|
172
|
+
becoming `inf`.
|
|
173
|
+
- The default precision is now `keep`. `-p fp32` used to change nothing; it now converts fp16 and bf16 tensors
|
|
174
|
+
to fp32. With the default precision the output name no longer ends in `-fp32`.
|
|
175
|
+
- Checkpoints are loaded with `weights_only=True`. Use `--unsafe-load` for trusted files that need more.
|
|
176
|
+
- Existing output files are no longer overwritten unless you pass `--overwrite`.
|
|
177
|
+
|
|
178
|
+
See [CHANGELOG.md](https://github.com/angelolamonaca/PyTorch-Precision-Converter/blob/main/CHANGELOG.md) for details.
|
|
179
|
+
|
|
180
|
+
## Contributing
|
|
181
|
+
|
|
182
|
+
Issues and pull requests are welcome. To run the tests:
|
|
183
|
+
|
|
184
|
+
```bash
|
|
185
|
+
pip install -e ".[test]"
|
|
186
|
+
pytest
|
|
187
|
+
```
|
|
188
|
+
|
|
189
|
+
## License
|
|
190
|
+
|
|
191
|
+
MIT. See [LICENSE](https://github.com/angelolamonaca/PyTorch-Precision-Converter/blob/main/LICENSE).
|
|
@@ -0,0 +1,165 @@
|
|
|
1
|
+
# PyTorch Precision Converter
|
|
2
|
+
|
|
3
|
+
[](https://pypi.org/project/pytorch-precision-converter/)
|
|
4
|
+
[](https://github.com/angelolamonaca/PyTorch-Precision-Converter/actions/workflows/ci.yml)
|
|
5
|
+

|
|
6
|
+

|
|
7
|
+
[](https://github.com/angelolamonaca/PyTorch-Precision-Converter/blob/main/LICENSE)
|
|
8
|
+
|
|
9
|
+
Make PyTorch checkpoints and `safetensors` files **half the size (a quarter with fp8)** by converting them to
|
|
10
|
+
`fp16`, `bf16` or `fp8`, and strip the EMA weights you don't need for inference. One command, no model code
|
|
11
|
+
needed, runs on the CPU.
|
|
12
|
+
|
|
13
|
+
```console
|
|
14
|
+
$ precision-convert model.ckpt -p fp16 -t ema-only -st
|
|
15
|
+
Converted 280 tensors to fp16, left 1 unchanged, dropped 242 EMA entries
|
|
16
|
+
Used the EMA weights for 240 tensors
|
|
17
|
+
Saved model-ema-only-fp16.safetensors (29.93 MB, 27% of 109.23 MB)
|
|
18
|
+
```
|
|
19
|
+
|
|
20
|
+
## Why not just call `.half()`?
|
|
21
|
+
|
|
22
|
+
A one-line cast quietly breaks things. This tool handles the details:
|
|
23
|
+
|
|
24
|
+
- **Only floating-point tensors are cast.** Integer and bool tensors like `position_ids` or step counters keep
|
|
25
|
+
their dtype instead of turning into floats.
|
|
26
|
+
- **No silent `inf` or `NaN`.** fp16 tops out at 65,504 and fp8 at 448. Values beyond the range are clamped
|
|
27
|
+
to it and you get a warning, rather than infinities appearing in your weights.
|
|
28
|
+
- **Real EMA extraction.** `-t ema-only` puts the EMA weights where the model loads them from, for Stable
|
|
29
|
+
Diffusion and other latent-diffusion checkpoints (`model_ema.*`), then drops the duplicates.
|
|
30
|
+
- **safetensors metadata survives.** LoRA training info (`ss_*` keys) and other metadata is copied to the output.
|
|
31
|
+
- **Safe by default.** PyTorch checkpoints load with `weights_only=True`, so a malicious `.ckpt` can't run code.
|
|
32
|
+
Files that need full unpickling require an explicit `--unsafe-load`.
|
|
33
|
+
- **Your files are safe too.** Output is written to a temporary file and moved into place when complete, the
|
|
34
|
+
input is never replaced, and existing files are only overwritten with `--overwrite`.
|
|
35
|
+
- **Mixed precision when you need it.** `--keep-original "first_stage_model.*"` keeps matching tensors (here the
|
|
36
|
+
VAE) in their original dtype.
|
|
37
|
+
|
|
38
|
+
## Install
|
|
39
|
+
|
|
40
|
+
```bash
|
|
41
|
+
pip install pytorch-precision-converter
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
For the latest unreleased changes, install from GitHub instead:
|
|
45
|
+
`pip install git+https://github.com/angelolamonaca/PyTorch-Precision-Converter`.
|
|
46
|
+
|
|
47
|
+
Either way you get the `precision-convert` command (`python -m precision_converter` works too). It needs Python 3.9+,
|
|
48
|
+
PyTorch 2.1+ and `safetensors`.
|
|
49
|
+
|
|
50
|
+
## Usage
|
|
51
|
+
|
|
52
|
+
```bash
|
|
53
|
+
precision-convert INPUT [-p PRECISION] [-t TYPE] [-st] [-o OUTPUT] [options]
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
| Option | What it does |
|
|
57
|
+
|---|---|
|
|
58
|
+
| `INPUT` | `.safetensors` file or PyTorch checkpoint (`.ckpt`, `.pt`, `.pth`, `.bin`). `-f INPUT` works too. |
|
|
59
|
+
| `-p`, `--precision` | `fp32`, `fp16`, `bf16`, `fp8_e4m3fn`, `fp8_e5m2`, or `keep` (default: change no dtypes). |
|
|
60
|
+
| `-t`, `--type` | `full` keeps everything (default). `ema-only` uses the EMA weights. `no-ema` drops the EMA weights. |
|
|
61
|
+
| `-st`, `--safe-tensors` | Save as `.safetensors` instead of a PyTorch `.ckpt`. |
|
|
62
|
+
| `-o`, `--output` | Output path; its extension picks the format. Default: `<input>-<type>-<precision>` next to the input. |
|
|
63
|
+
| `--keep-original PATTERN` | Leave tensors whose key matches the glob in their original dtype. Repeatable. |
|
|
64
|
+
| `--info` | Show the tensors, dtypes, sizes, EMA and metadata in the file, then exit. |
|
|
65
|
+
| `--unsafe-load` | Fully unpickle PyTorch checkpoints. Only for files you trust. |
|
|
66
|
+
| `--overwrite` | Replace the output file if it exists. |
|
|
67
|
+
| `-v`, `--verbose` | List each tensor that had values clamped. |
|
|
68
|
+
|
|
69
|
+
## Recipes
|
|
70
|
+
|
|
71
|
+
**Shrink a Stable Diffusion training checkpoint for inference** (EMA weights, fp16, safetensors):
|
|
72
|
+
|
|
73
|
+
```bash
|
|
74
|
+
precision-convert model.ckpt -p fp16 -t ema-only -st
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
**Convert a LoRA to fp16 and keep its training metadata:**
|
|
78
|
+
|
|
79
|
+
```bash
|
|
80
|
+
precision-convert my-lora.safetensors -p fp16 -o my-lora-fp16.safetensors
|
|
81
|
+
```
|
|
82
|
+
|
|
83
|
+
**Store a large model in fp8** for UIs that load fp8 weights, such as ComfyUI, keeping normalization layers
|
|
84
|
+
precise:
|
|
85
|
+
|
|
86
|
+
```bash
|
|
87
|
+
precision-convert model.safetensors -p fp8_e4m3fn -st --keep-original "*norm*"
|
|
88
|
+
```
|
|
89
|
+
|
|
90
|
+
**Keep the VAE in fp32** (fp16 VAEs can produce NaNs and black images in some models):
|
|
91
|
+
|
|
92
|
+
```bash
|
|
93
|
+
precision-convert model.safetensors -p fp16 -st --keep-original "first_stage_model.*"
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
**Turn a pickled checkpoint into safetensors without changing precision:**
|
|
97
|
+
|
|
98
|
+
```bash
|
|
99
|
+
precision-convert model.ckpt -st
|
|
100
|
+
```
|
|
101
|
+
|
|
102
|
+
**Check what's inside a file before converting:**
|
|
103
|
+
|
|
104
|
+
```console
|
|
105
|
+
$ precision-convert model.ckpt --info
|
|
106
|
+
model.ckpt: 523 tensors, 27,274,319 values, 109.10 MB
|
|
107
|
+
float32 521 tensors 109.10 MB
|
|
108
|
+
int64 1 tensor 616 B
|
|
109
|
+
int32 1 tensor 4 B
|
|
110
|
+
EMA: 242 model_ema.* tensors (-t ema-only or -t no-ema removes them)
|
|
111
|
+
```
|
|
112
|
+
|
|
113
|
+
## Which precision should I pick?
|
|
114
|
+
|
|
115
|
+
| Precision | Bytes per value | Largest value | Good for |
|
|
116
|
+
|---|---|---|---|
|
|
117
|
+
| `fp32` | 4 | 3.4 × 10³⁸ | Training, reference copies, upcasting half-precision files |
|
|
118
|
+
| `bf16` | 2 | 3.4 × 10³⁸ | The safe 2-byte choice: same range as fp32, so nothing overflows |
|
|
119
|
+
| `fp16` | 2 | 65,504 | GPUs without fast bf16, and tools that expect fp16 |
|
|
120
|
+
| `fp8_e4m3fn` | 1 | 448 | Smallest files, for inference in tools that load fp8 weights; some quality loss |
|
|
121
|
+
| `fp8_e5m2` | 1 | 57,344 | fp8 with more range and less precision than `e4m3fn` |
|
|
122
|
+
|
|
123
|
+
The fp8 options cast each value directly, without per-tensor scaling factors.
|
|
124
|
+
|
|
125
|
+
## Python API
|
|
126
|
+
|
|
127
|
+
```python
|
|
128
|
+
from precision_converter import convert_state_dict, load_checkpoint, save_checkpoint
|
|
129
|
+
|
|
130
|
+
state_dict, metadata = load_checkpoint("model.safetensors")
|
|
131
|
+
converted, report = convert_state_dict(state_dict, precision="bf16", conversion="no-ema")
|
|
132
|
+
save_checkpoint(converted, "model-bf16.safetensors", metadata=metadata)
|
|
133
|
+
|
|
134
|
+
print(f"{report.converted} tensors converted, {len(report.clamped)} had values clamped")
|
|
135
|
+
```
|
|
136
|
+
|
|
137
|
+
`convert_state_dict` takes any mapping of tensors, so `model.state_dict()` works too.
|
|
138
|
+
|
|
139
|
+
## Upgrading from 1.x
|
|
140
|
+
|
|
141
|
+
`python convert.py ...` still works, but 2.0 fixes bugs that changed what earlier versions wrote:
|
|
142
|
+
|
|
143
|
+
- `-t ema-only` used to keep the regular weights and store the EMA weights under keys no model loads. It now
|
|
144
|
+
outputs the actual EMA weights.
|
|
145
|
+
- Integer tensors are no longer cast to floating point, and overflowing values are clamped instead of
|
|
146
|
+
becoming `inf`.
|
|
147
|
+
- The default precision is now `keep`. `-p fp32` used to change nothing; it now converts fp16 and bf16 tensors
|
|
148
|
+
to fp32. With the default precision the output name no longer ends in `-fp32`.
|
|
149
|
+
- Checkpoints are loaded with `weights_only=True`. Use `--unsafe-load` for trusted files that need more.
|
|
150
|
+
- Existing output files are no longer overwritten unless you pass `--overwrite`.
|
|
151
|
+
|
|
152
|
+
See [CHANGELOG.md](https://github.com/angelolamonaca/PyTorch-Precision-Converter/blob/main/CHANGELOG.md) for details.
|
|
153
|
+
|
|
154
|
+
## Contributing
|
|
155
|
+
|
|
156
|
+
Issues and pull requests are welcome. To run the tests:
|
|
157
|
+
|
|
158
|
+
```bash
|
|
159
|
+
pip install -e ".[test]"
|
|
160
|
+
pytest
|
|
161
|
+
```
|
|
162
|
+
|
|
163
|
+
## License
|
|
164
|
+
|
|
165
|
+
MIT. See [LICENSE](https://github.com/angelolamonaca/PyTorch-Precision-Converter/blob/main/LICENSE).
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
"""Convert the tensor precision of PyTorch checkpoints and safetensors files."""
|
|
2
|
+
|
|
3
|
+
__version__ = "2.0.0"
|
|
4
|
+
|
|
5
|
+
from .core import ( # noqa: E402
|
|
6
|
+
CONVERSION_TYPES,
|
|
7
|
+
PRECISIONS,
|
|
8
|
+
CheckpointInfo,
|
|
9
|
+
ConversionReport,
|
|
10
|
+
UnsafeCheckpointError,
|
|
11
|
+
convert_state_dict,
|
|
12
|
+
describe,
|
|
13
|
+
load_checkpoint,
|
|
14
|
+
save_checkpoint,
|
|
15
|
+
)
|
|
16
|
+
|
|
17
|
+
__all__ = [
|
|
18
|
+
"CONVERSION_TYPES",
|
|
19
|
+
"PRECISIONS",
|
|
20
|
+
"CheckpointInfo",
|
|
21
|
+
"ConversionReport",
|
|
22
|
+
"UnsafeCheckpointError",
|
|
23
|
+
"convert_state_dict",
|
|
24
|
+
"describe",
|
|
25
|
+
"load_checkpoint",
|
|
26
|
+
"save_checkpoint",
|
|
27
|
+
]
|
|
@@ -0,0 +1,181 @@
|
|
|
1
|
+
"""Command-line interface: ``precision-convert``."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import argparse
|
|
6
|
+
import sys
|
|
7
|
+
from pathlib import Path
|
|
8
|
+
|
|
9
|
+
import torch
|
|
10
|
+
|
|
11
|
+
from . import __version__
|
|
12
|
+
from .core import (
|
|
13
|
+
CONVERSION_TYPES,
|
|
14
|
+
PRECISIONS,
|
|
15
|
+
UnsafeCheckpointError,
|
|
16
|
+
convert_state_dict,
|
|
17
|
+
describe,
|
|
18
|
+
load_checkpoint,
|
|
19
|
+
save_checkpoint,
|
|
20
|
+
)
|
|
21
|
+
|
|
22
|
+
EXAMPLES = """\
|
|
23
|
+
examples:
|
|
24
|
+
precision-convert model.safetensors -p fp16 -st
|
|
25
|
+
precision-convert sd-model.ckpt -p fp16 -t ema-only -st
|
|
26
|
+
precision-convert lora.safetensors -p bf16 -o lora-bf16.safetensors
|
|
27
|
+
precision-convert flux.safetensors -p fp8_e4m3fn -st --keep-original "*norm*"
|
|
28
|
+
precision-convert model.ckpt --info
|
|
29
|
+
"""
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def build_parser() -> argparse.ArgumentParser:
|
|
33
|
+
parser = argparse.ArgumentParser(
|
|
34
|
+
prog="precision-convert",
|
|
35
|
+
description="Convert the tensor precision of PyTorch checkpoints and safetensors files.",
|
|
36
|
+
epilog=EXAMPLES,
|
|
37
|
+
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
38
|
+
)
|
|
39
|
+
parser.add_argument(
|
|
40
|
+
"input", nargs="?", help="Model file: .safetensors or a PyTorch checkpoint (.ckpt, .pt, .pth, .bin)."
|
|
41
|
+
)
|
|
42
|
+
parser.add_argument("-f", "--file", help="Same as INPUT, kept for compatibility with 1.x.")
|
|
43
|
+
parser.add_argument(
|
|
44
|
+
"-p",
|
|
45
|
+
"--precision",
|
|
46
|
+
default="keep",
|
|
47
|
+
choices=PRECISIONS,
|
|
48
|
+
help="Target precision for floating-point tensors (default: keep, which changes no dtype).",
|
|
49
|
+
)
|
|
50
|
+
parser.add_argument(
|
|
51
|
+
"-t",
|
|
52
|
+
"--type",
|
|
53
|
+
default="full",
|
|
54
|
+
choices=CONVERSION_TYPES,
|
|
55
|
+
help="full: keep everything (default). ema-only: use the EMA weights and drop the EMA copies. "
|
|
56
|
+
"no-ema: drop the EMA copies.",
|
|
57
|
+
)
|
|
58
|
+
parser.add_argument(
|
|
59
|
+
"-st", "--safe-tensors", action="store_true", help="Save in safetensors format (default: PyTorch .ckpt)."
|
|
60
|
+
)
|
|
61
|
+
parser.add_argument(
|
|
62
|
+
"-o",
|
|
63
|
+
"--output",
|
|
64
|
+
help="Output path. Its extension picks the format (.safetensors or anything else for PyTorch). "
|
|
65
|
+
"Default: <input>-<type>-<precision> next to the input.",
|
|
66
|
+
)
|
|
67
|
+
parser.add_argument(
|
|
68
|
+
"--keep-original",
|
|
69
|
+
metavar="PATTERN",
|
|
70
|
+
action="append",
|
|
71
|
+
default=[],
|
|
72
|
+
help='Leave tensors whose key matches this glob in their original dtype, e.g. "first_stage_model.*". '
|
|
73
|
+
"Repeatable.",
|
|
74
|
+
)
|
|
75
|
+
parser.add_argument("--overwrite", action="store_true", help="Replace the output file if it already exists.")
|
|
76
|
+
parser.add_argument(
|
|
77
|
+
"--unsafe-load",
|
|
78
|
+
action="store_true",
|
|
79
|
+
help="Load PyTorch checkpoints with full unpickling (weights_only=False). Unpickling can run code "
|
|
80
|
+
"hidden in the file: only use this for files you trust.",
|
|
81
|
+
)
|
|
82
|
+
parser.add_argument("--info", action="store_true", help="Describe the input file and exit without converting.")
|
|
83
|
+
parser.add_argument("-v", "--verbose", action="store_true", help="List the tensors whose values were clamped.")
|
|
84
|
+
parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
|
|
85
|
+
return parser
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
def format_bytes(n: float) -> str:
|
|
89
|
+
for unit in ("B", "KB", "MB", "GB"):
|
|
90
|
+
if n < 1000 or unit == "GB":
|
|
91
|
+
return f"{n:.0f} {unit}" if unit == "B" else f"{n:.2f} {unit}"
|
|
92
|
+
n /= 1000
|
|
93
|
+
raise AssertionError("unreachable")
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
def default_output(input_path: Path, conversion: str, precision: str, safetensors: bool) -> Path:
|
|
97
|
+
parts = [input_path.stem, conversion] + ([precision] if precision != "keep" else [])
|
|
98
|
+
return input_path.with_name("-".join(parts) + (".safetensors" if safetensors else ".ckpt"))
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
def print_info(path: Path, state_dict: dict, metadata: dict | None) -> None:
|
|
102
|
+
info = describe(state_dict)
|
|
103
|
+
print(f"{path}: {info.tensors:,} tensors, {info.parameters:,} values, {format_bytes(info.bytes)}")
|
|
104
|
+
for dtype, (count, nbytes) in sorted(info.dtypes.items(), key=lambda item: -item[1][1]):
|
|
105
|
+
print(f" {dtype:<16} {count:>7,} {'tensor ' if count == 1 else 'tensors'} {format_bytes(nbytes):>10}")
|
|
106
|
+
if info.ema_tensors:
|
|
107
|
+
print(f" EMA: {info.ema_tensors:,} model_ema.* tensors (-t ema-only or -t no-ema removes them)")
|
|
108
|
+
if info.other_keys:
|
|
109
|
+
print(f" Non-tensor keys: {', '.join(info.other_keys)}")
|
|
110
|
+
if metadata:
|
|
111
|
+
print(f" Metadata: {len(metadata)} entries ({', '.join(sorted(metadata)[:8])}{', ...' if len(metadata) > 8 else ''})")
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
def main(argv: list[str] | None = None) -> int:
|
|
115
|
+
parser = build_parser()
|
|
116
|
+
args = parser.parse_args(argv)
|
|
117
|
+
|
|
118
|
+
if args.input and args.file:
|
|
119
|
+
parser.error("give the input file either as INPUT or with -f, not both")
|
|
120
|
+
input_path = Path(args.input or args.file or "model.ckpt")
|
|
121
|
+
if not input_path.is_file():
|
|
122
|
+
parser.error(f"input file not found: {input_path}")
|
|
123
|
+
|
|
124
|
+
if args.output:
|
|
125
|
+
output_path = Path(args.output)
|
|
126
|
+
if args.safe_tensors and output_path.suffix != ".safetensors":
|
|
127
|
+
parser.error("--safe-tensors needs an --output path ending in .safetensors")
|
|
128
|
+
else:
|
|
129
|
+
output_path = default_output(input_path, args.type, args.precision, args.safe_tensors)
|
|
130
|
+
|
|
131
|
+
try:
|
|
132
|
+
state_dict, metadata = load_checkpoint(input_path, unsafe=args.unsafe_load)
|
|
133
|
+
except (UnsafeCheckpointError, ValueError, OSError, RuntimeError) as e:
|
|
134
|
+
print(f"error: {e}", file=sys.stderr)
|
|
135
|
+
return 1
|
|
136
|
+
|
|
137
|
+
if args.info:
|
|
138
|
+
print_info(input_path, state_dict, metadata)
|
|
139
|
+
return 0
|
|
140
|
+
|
|
141
|
+
if output_path.resolve() == input_path.resolve():
|
|
142
|
+
parser.error("the output would replace the input file; choose another --output")
|
|
143
|
+
if output_path.exists() and not args.overwrite:
|
|
144
|
+
parser.error(f"{output_path} already exists; pass --overwrite to replace it")
|
|
145
|
+
|
|
146
|
+
converted, report = convert_state_dict(state_dict, args.precision, args.type, args.keep_original)
|
|
147
|
+
|
|
148
|
+
if args.precision == "keep":
|
|
149
|
+
summary = f"Kept every dtype as it was ({report.unchanged:,} entries)"
|
|
150
|
+
else:
|
|
151
|
+
summary = f"Converted {report.converted:,} tensors to {args.precision}, left {report.unchanged:,} unchanged"
|
|
152
|
+
if report.dropped:
|
|
153
|
+
summary += f", dropped {report.dropped:,} EMA entries"
|
|
154
|
+
print(summary)
|
|
155
|
+
if args.type == "ema-only":
|
|
156
|
+
if report.ema_swapped:
|
|
157
|
+
print(f"Used the EMA weights for {report.ema_swapped:,} tensors")
|
|
158
|
+
else:
|
|
159
|
+
print("warning: no EMA weights found; the output keeps the regular weights", file=sys.stderr)
|
|
160
|
+
if report.clamped:
|
|
161
|
+
limit = torch.finfo(PRECISIONS[args.precision]).max
|
|
162
|
+
print(
|
|
163
|
+
f"warning: {len(report.clamped):,} tensors had values outside the {args.precision} range "
|
|
164
|
+
f"(±{limit:g}); they were clamped instead of becoming inf/NaN",
|
|
165
|
+
file=sys.stderr,
|
|
166
|
+
)
|
|
167
|
+
if args.verbose:
|
|
168
|
+
for key, count in report.clamped.items():
|
|
169
|
+
print(f" {key}: {count:,} values", file=sys.stderr)
|
|
170
|
+
|
|
171
|
+
skipped = save_checkpoint(converted, output_path, metadata=metadata)
|
|
172
|
+
if skipped:
|
|
173
|
+
print(f"warning: safetensors only stores tensors, left out: {', '.join(skipped)}", file=sys.stderr)
|
|
174
|
+
|
|
175
|
+
before, after = input_path.stat().st_size, output_path.stat().st_size
|
|
176
|
+
print(f"Saved {output_path} ({format_bytes(after)}, {after / before:.0%} of {format_bytes(before)})")
|
|
177
|
+
return 0
|
|
178
|
+
|
|
179
|
+
|
|
180
|
+
if __name__ == "__main__":
|
|
181
|
+
sys.exit(main())
|