Spaces:
Running
on
Zero
Running
on
Zero
<!--Copyright 2024 The HuggingFace Team. All rights reserved. | |
Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with | |
the License. You may obtain a copy of the License at | |
http://www.apache.org/licenses/LICENSE-2.0 | |
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on | |
an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the | |
specific language governing permissions and limitations under the License. | |
--> | |
# ๋ฉ๋ชจ๋ฆฌ์ ์๋ | |
๋ฉ๋ชจ๋ฆฌ ๋๋ ์๋์ ๋ํด ๐ค Diffusers *์ถ๋ก *์ ์ต์ ํํ๊ธฐ ์ํ ๋ช ๊ฐ์ง ๊ธฐ์ ๊ณผ ์์ด๋์ด๋ฅผ ์ ์ํฉ๋๋ค. | |
์ผ๋ฐ์ ์ผ๋ก, memory-efficient attention์ ์ํด [xFormers](https://github.com/facebookresearch/xformers) ์ฌ์ฉ์ ์ถ์ฒํ๊ธฐ ๋๋ฌธ์, ์ถ์ฒํ๋ [์ค์น ๋ฐฉ๋ฒ](xformers)์ ๋ณด๊ณ ์ค์นํด ๋ณด์ธ์. | |
๋ค์ ์ค์ ์ด ์ฑ๋ฅ๊ณผ ๋ฉ๋ชจ๋ฆฌ์ ๋ฏธ์น๋ ์ํฅ์ ๋ํด ์ค๋ช ํฉ๋๋ค. | |
| | ์ง์ฐ์๊ฐ | ์๋ ํฅ์ | | |
| ---------------- | ------- | ------- | | |
| ๋ณ๋ ์ค์ ์์ | 9.50s | x1 | | |
| cuDNN auto-tuner | 9.37s | x1.01 | | |
| fp16 | 3.61s | x2.63 | | |
| Channels Last ๋ฉ๋ชจ๋ฆฌ ํ์ | 3.30s | x2.88 | | |
| traced UNet | 3.21s | x2.96 | | |
| memory-efficient attention | 2.63s | x3.61 | | |
<em> | |
NVIDIA TITAN RTX์์ 50 DDIM ์คํ ์ "a photo of an astronaut riding a horse on mars" ํ๋กฌํํธ๋ก 512x512 ํฌ๊ธฐ์ ๋จ์ผ ์ด๋ฏธ์ง๋ฅผ ์์ฑํ์์ต๋๋ค. | |
</em> | |
## cuDNN auto-tuner ํ์ฑํํ๊ธฐ | |
[NVIDIA cuDNN](https://developer.nvidia.com/cudnn)์ ์ปจ๋ณผ๋ฃจ์ ์ ๊ณ์ฐํ๋ ๋ง์ ์๊ณ ๋ฆฌ์ฆ์ ์ง์ํฉ๋๋ค. Autotuner๋ ์งง์ ๋ฒค์น๋งํฌ๋ฅผ ์คํํ๊ณ ์ฃผ์ด์ง ์ ๋ ฅ ํฌ๊ธฐ์ ๋ํด ์ฃผ์ด์ง ํ๋์จ์ด์์ ์ต๊ณ ์ ์ฑ๋ฅ์ ๊ฐ์ง ์ปค๋์ ์ ํํฉ๋๋ค. | |
**์ปจ๋ณผ๋ฃจ์ ๋คํธ์ํฌ**๋ฅผ ํ์ฉํ๊ณ ์๊ธฐ ๋๋ฌธ์ (๋ค๋ฅธ ์ ํ๋ค์ ํ์ฌ ์ง์๋์ง ์์), ๋ค์ ์ค์ ์ ํตํด ์ถ๋ก ์ ์ cuDNN autotuner๋ฅผ ํ์ฑํํ ์ ์์ต๋๋ค: | |
```python | |
import torch | |
torch.backends.cudnn.benchmark = True | |
``` | |
### fp32 ๋์ tf32 ์ฌ์ฉํ๊ธฐ (Ampere ๋ฐ ์ดํ CUDA ์ฅ์น๋ค์์) | |
Ampere ๋ฐ ์ดํ CUDA ์ฅ์น์์ ํ๋ ฌ๊ณฑ ๋ฐ ์ปจ๋ณผ๋ฃจ์ ์ TensorFloat32(TF32) ๋ชจ๋๋ฅผ ์ฌ์ฉํ์ฌ ๋ ๋น ๋ฅด์ง๋ง ์ฝ๊ฐ ๋ ์ ํํ ์ ์์ต๋๋ค. | |
๊ธฐ๋ณธ์ ์ผ๋ก PyTorch๋ ์ปจ๋ณผ๋ฃจ์ ์ ๋ํด TF32 ๋ชจ๋๋ฅผ ํ์ฑํํ์ง๋ง ํ๋ ฌ ๊ณฑ์ ์ ํ์ฑํํ์ง ์์ต๋๋ค. | |
๋คํธ์ํฌ์ ์์ ํ float32 ์ ๋ฐ๋๊ฐ ํ์ํ ๊ฒฝ์ฐ๊ฐ ์๋๋ฉด ํ๋ ฌ ๊ณฑ์ ์ ๋ํด์๋ ์ด ์ค์ ์ ํ์ฑํํ๋ ๊ฒ์ด ์ข์ต๋๋ค. | |
์ด๋ ์ผ๋ฐ์ ์ผ๋ก ๋ฌด์ํ ์ ์๋ ์์น์ ์ ํ๋ ์์ค์ด ์์ง๋ง, ๊ณ์ฐ ์๋๋ฅผ ํฌ๊ฒ ๋์ผ ์ ์์ต๋๋ค. | |
๊ทธ๊ฒ์ ๋ํด [์ฌ๊ธฐ](https://huggingface.co/docs/transformers/v4.18.0/en/performance#tf32)์ ๋ ์ฝ์ ์ ์์ต๋๋ค. | |
์ถ๋ก ํ๊ธฐ ์ ์ ๋ค์์ ์ถ๊ฐํ๊ธฐ๋ง ํ๋ฉด ๋ฉ๋๋ค: | |
```python | |
import torch | |
torch.backends.cuda.matmul.allow_tf32 = True | |
``` | |
## ๋ฐ์ ๋ฐ๋ ๊ฐ์ค์น | |
๋ ๋ง์ GPU ๋ฉ๋ชจ๋ฆฌ๋ฅผ ์ ์ฝํ๊ณ ๋ ๋น ๋ฅธ ์๋๋ฅผ ์ป๊ธฐ ์ํด ๋ชจ๋ธ ๊ฐ์ค์น๋ฅผ ๋ฐ์ ๋ฐ๋(half precision)๋ก ์ง์ ๋ถ๋ฌ์ค๊ณ ์คํํ ์ ์์ต๋๋ค. | |
์ฌ๊ธฐ์๋ `fp16`์ด๋ผ๋ ๋ธ๋์น์ ์ ์ฅ๋ float16 ๋ฒ์ ์ ๊ฐ์ค์น๋ฅผ ๋ถ๋ฌ์ค๊ณ , ๊ทธ ๋ `float16` ์ ํ์ ์ฌ์ฉํ๋๋ก PyTorch์ ์ง์ํ๋ ์์ ์ด ํฌํจ๋ฉ๋๋ค. | |
```Python | |
pipe = StableDiffusionPipeline.from_pretrained( | |
"runwayml/stable-diffusion-v1-5", | |
torch_dtype=torch.float16, | |
) | |
pipe = pipe.to("cuda") | |
prompt = "a photo of an astronaut riding a horse on mars" | |
image = pipe(prompt).images[0] | |
``` | |
<Tip warning={true}> | |
์ด๋ค ํ์ดํ๋ผ์ธ์์๋ [`torch.autocast`](https://pytorch.org/docs/stable/amp.html#torch.autocast) ๋ฅผ ์ฌ์ฉํ๋ ๊ฒ์ ๊ฒ์์ ์ด๋ฏธ์ง๋ฅผ ์์ฑํ ์ ์๊ณ , ์์ํ float16 ์ ๋ฐ๋๋ฅผ ์ฌ์ฉํ๋ ๊ฒ๋ณด๋ค ํญ์ ๋๋ฆฌ๊ธฐ ๋๋ฌธ์ ์ฌ์ฉํ์ง ์๋ ๊ฒ์ด ์ข์ต๋๋ค. | |
</Tip> | |
## ์ถ๊ฐ ๋ฉ๋ชจ๋ฆฌ ์ ์ฝ์ ์ํ ์ฌ๋ผ์ด์ค ์ดํ ์ | |
์ถ๊ฐ ๋ฉ๋ชจ๋ฆฌ ์ ์ฝ์ ์ํด, ํ ๋ฒ์ ๋ชจ๋ ๊ณ์ฐํ๋ ๋์ ๋จ๊ณ์ ์ผ๋ก ๊ณ์ฐ์ ์ํํ๋ ์ฌ๋ผ์ด์ค ๋ฒ์ ์ ์ดํ ์ (attention)์ ์ฌ์ฉํ ์ ์์ต๋๋ค. | |
<Tip> | |
Attention slicing์ ๋ชจ๋ธ์ด ํ๋ ์ด์์ ์ดํ ์ ํค๋๋ฅผ ์ฌ์ฉํ๋ ํ, ๋ฐฐ์น ํฌ๊ธฐ๊ฐ 1์ธ ๊ฒฝ์ฐ์๋ ์ ์ฉํฉ๋๋ค. | |
ํ๋ ์ด์์ ์ดํ ์ ํค๋๊ฐ ์๋ ๊ฒฝ์ฐ *QK^T* ์ดํ ์ ๋งคํธ๋ฆญ์ค๋ ์๋นํ ์์ ๋ฉ๋ชจ๋ฆฌ๋ฅผ ์ ์ฝํ ์ ์๋ ๊ฐ ํค๋์ ๋ํด ์์ฐจ์ ์ผ๋ก ๊ณ์ฐ๋ ์ ์์ต๋๋ค. | |
</Tip> | |
๊ฐ ํค๋์ ๋ํด ์์ฐจ์ ์ผ๋ก ์ดํ ์ ๊ณ์ฐ์ ์ํํ๋ ค๋ฉด, ๋ค์๊ณผ ๊ฐ์ด ์ถ๋ก ์ ์ ํ์ดํ๋ผ์ธ์์ [`~StableDiffusionPipeline.enable_attention_slicing`]๋ฅผ ํธ์ถํ๋ฉด ๋ฉ๋๋ค: | |
```Python | |
import torch | |
from diffusers import StableDiffusionPipeline | |
pipe = StableDiffusionPipeline.from_pretrained( | |
"runwayml/stable-diffusion-v1-5", | |
torch_dtype=torch.float16, | |
) | |
pipe = pipe.to("cuda") | |
prompt = "a photo of an astronaut riding a horse on mars" | |
pipe.enable_attention_slicing() | |
image = pipe(prompt).images[0] | |
``` | |
์ถ๋ก ์๊ฐ์ด ์ฝ 10% ๋๋ ค์ง๋ ์ฝ๊ฐ์ ์ฑ๋ฅ ์ ํ๊ฐ ์์ง๋ง ์ด ๋ฐฉ๋ฒ์ ์ฌ์ฉํ๋ฉด 3.2GB ์ ๋์ ์์ VRAM์ผ๋ก๋ Stable Diffusion์ ์ฌ์ฉํ ์ ์์ต๋๋ค! | |
## ๋ ํฐ ๋ฐฐ์น๋ฅผ ์ํ sliced VAE ๋์ฝ๋ | |
์ ํ๋ VRAM์์ ๋๊ท๋ชจ ์ด๋ฏธ์ง ๋ฐฐ์น๋ฅผ ๋์ฝ๋ฉํ๊ฑฐ๋ 32๊ฐ ์ด์์ ์ด๋ฏธ์ง๊ฐ ํฌํจ๋ ๋ฐฐ์น๋ฅผ ํ์ฑํํ๊ธฐ ์ํด, ๋ฐฐ์น์ latent ์ด๋ฏธ์ง๋ฅผ ํ ๋ฒ์ ํ๋์ฉ ๋์ฝ๋ฉํ๋ ์ฌ๋ผ์ด์ค VAE ๋์ฝ๋๋ฅผ ์ฌ์ฉํ ์ ์์ต๋๋ค. | |
์ด๋ฅผ [`~StableDiffusionPipeline.enable_attention_slicing`] ๋๋ [`~StableDiffusionPipeline.enable_xformers_memory_efficient_attention`]๊ณผ ๊ฒฐํฉํ์ฌ ๋ฉ๋ชจ๋ฆฌ ์ฌ์ฉ์ ์ถ๊ฐ๋ก ์ต์ํํ ์ ์์ต๋๋ค. | |
VAE ๋์ฝ๋๋ฅผ ํ ๋ฒ์ ํ๋์ฉ ์ํํ๋ ค๋ฉด ์ถ๋ก ์ ์ ํ์ดํ๋ผ์ธ์์ [`~StableDiffusionPipeline.enable_vae_slicing`]์ ํธ์ถํฉ๋๋ค. ์๋ฅผ ๋ค์ด: | |
```Python | |
import torch | |
from diffusers import StableDiffusionPipeline | |
pipe = StableDiffusionPipeline.from_pretrained( | |
"runwayml/stable-diffusion-v1-5", | |
torch_dtype=torch.float16, | |
) | |
pipe = pipe.to("cuda") | |
prompt = "a photo of an astronaut riding a horse on mars" | |
pipe.enable_vae_slicing() | |
images = pipe([prompt] * 32).images | |
``` | |
๋ค์ค ์ด๋ฏธ์ง ๋ฐฐ์น์์ VAE ๋์ฝ๋๊ฐ ์ฝ๊ฐ์ ์ฑ๋ฅ ํฅ์์ด ์ด๋ฃจ์ด์ง๋๋ค. ๋จ์ผ ์ด๋ฏธ์ง ๋ฐฐ์น์์๋ ์ฑ๋ฅ ์ํฅ์ ์์ต๋๋ค. | |
<a name="sequential_offloading"></a> | |
## ๋ฉ๋ชจ๋ฆฌ ์ ์ฝ์ ์ํด ๊ฐ์ ๊ธฐ๋ฅ์ ์ฌ์ฉํ์ฌ CPU๋ก ์คํ๋ก๋ฉ | |
์ถ๊ฐ ๋ฉ๋ชจ๋ฆฌ ์ ์ฝ์ ์ํด ๊ฐ์ค์น๋ฅผ CPU๋ก ์คํ๋ก๋ํ๊ณ ์๋ฐฉํฅ ์ ๋ฌ์ ์ํํ ๋๋ง GPU๋ก ๋ก๋ํ ์ ์์ต๋๋ค. | |
CPU ์คํ๋ก๋ฉ์ ์ํํ๋ ค๋ฉด [`~StableDiffusionPipeline.enable_sequential_cpu_offload`]๋ฅผ ํธ์ถํ๊ธฐ๋ง ํ๋ฉด ๋ฉ๋๋ค: | |
```Python | |
import torch | |
from diffusers import StableDiffusionPipeline | |
pipe = StableDiffusionPipeline.from_pretrained( | |
"runwayml/stable-diffusion-v1-5", | |
torch_dtype=torch.float16, | |
) | |
prompt = "a photo of an astronaut riding a horse on mars" | |
pipe.enable_sequential_cpu_offload() | |
image = pipe(prompt).images[0] | |
``` | |
๊ทธ๋ฌ๋ฉด ๋ฉ๋ชจ๋ฆฌ ์๋น๋ฅผ 3GB ๋ฏธ๋ง์ผ๋ก ์ค์ผ ์ ์์ต๋๋ค. | |
์ฐธ๊ณ ๋ก ์ด ๋ฐฉ๋ฒ์ ์ ์ฒด ๋ชจ๋ธ์ด ์๋ ์๋ธ๋ชจ๋ ์์ค์์ ์๋ํฉ๋๋ค. ์ด๋ ๋ฉ๋ชจ๋ฆฌ ์๋น๋ฅผ ์ต์ํํ๋ ๊ฐ์ฅ ์ข์ ๋ฐฉ๋ฒ์ด์ง๋ง ํ๋ก์ธ์ค์ ๋ฐ๋ณต์ ํน์ฑ์ผ๋ก ์ธํด ์ถ๋ก ์๋๊ฐ ํจ์ฌ ๋๋ฆฝ๋๋ค. ํ์ดํ๋ผ์ธ์ UNet ๊ตฌ์ฑ ์์๋ ์ฌ๋ฌ ๋ฒ ์คํ๋ฉ๋๋ค('num_inference_steps' ๋งํผ). ๋งค๋ฒ UNet์ ์๋ก ๋ค๋ฅธ ์๋ธ๋ชจ๋์ด ์์ฐจ์ ์ผ๋ก ์จ๋ก๋๋ ๋ค์ ํ์์ ๋ฐ๋ผ ์คํ๋ก๋๋๋ฏ๋ก ๋ฉ๋ชจ๋ฆฌ ์ด๋ ํ์๊ฐ ๋ง์ต๋๋ค. | |
<Tip> | |
๋ ๋ค๋ฅธ ์ต์ ํ ๋ฐฉ๋ฒ์ธ <a href="#model_offloading">๋ชจ๋ธ ์คํ๋ก๋ฉ</a>์ ์ฌ์ฉํ๋ ๊ฒ์ ๊ณ ๋ คํ์ญ์์ค. ์ด๋ ํจ์ฌ ๋น ๋ฅด์ง๋ง ๋ฉ๋ชจ๋ฆฌ ์ ์ฝ์ด ํฌ์ง๋ ์์ต๋๋ค. | |
</Tip> | |
๋ํ ttention slicing๊ณผ ์ฐ๊ฒฐํด์ ์ต์ ๋ฉ๋ชจ๋ฆฌ(< 2GB)๋ก๋ ๋์ํ ์ ์์ต๋๋ค. | |
```Python | |
import torch | |
from diffusers import StableDiffusionPipeline | |
pipe = StableDiffusionPipeline.from_pretrained( | |
"runwayml/stable-diffusion-v1-5", | |
torch_dtype=torch.float16, | |
) | |
prompt = "a photo of an astronaut riding a horse on mars" | |
pipe.enable_sequential_cpu_offload() | |
pipe.enable_attention_slicing(1) | |
image = pipe(prompt).images[0] | |
``` | |
**์ฐธ๊ณ **: 'enable_sequential_cpu_offload()'๋ฅผ ์ฌ์ฉํ ๋, ๋ฏธ๋ฆฌ ํ์ดํ๋ผ์ธ์ CUDA๋ก ์ด๋ํ์ง **์๋** ๊ฒ์ด ์ค์ํฉ๋๋ค.๊ทธ๋ ์ง ์์ผ๋ฉด ๋ฉ๋ชจ๋ฆฌ ์๋น์ ์ด๋์ด ์ต์ํ๋ฉ๋๋ค. ๋ ๋ง์ ์ ๋ณด๋ฅผ ์ํด [์ด ์ด์](https://github.com/huggingface/diffusers/issues/1934)๋ฅผ ๋ณด์ธ์. | |
<a name="model_offloading"></a> | |
## ๋น ๋ฅธ ์ถ๋ก ๊ณผ ๋ฉ๋ชจ๋ฆฌ ๋ฉ๋ชจ๋ฆฌ ์ ์ฝ์ ์ํ ๋ชจ๋ธ ์คํ๋ก๋ฉ | |
[์์ฐจ์ CPU ์คํ๋ก๋ฉ](#sequential_offloading)์ ์ด์ ์น์ ์์ ์ค๋ช ํ ๊ฒ์ฒ๋ผ ๋ง์ ๋ฉ๋ชจ๋ฆฌ๋ฅผ ๋ณด์กดํ์ง๋ง ํ์์ ๋ฐ๋ผ ์๋ธ๋ชจ๋์ GPU๋ก ์ด๋ํ๊ณ ์ ๋ชจ๋์ด ์คํ๋ ๋ ์ฆ์ CPU๋ก ๋ฐํ๋๊ธฐ ๋๋ฌธ์ ์ถ๋ก ์๋๊ฐ ๋๋ ค์ง๋๋ค. | |
์ ์ฒด ๋ชจ๋ธ ์คํ๋ก๋ฉ์ ๊ฐ ๋ชจ๋ธ์ ๊ตฌ์ฑ ์์์ธ _modules_์ ์ฒ๋ฆฌํ๋ ๋์ , ์ ์ฒด ๋ชจ๋ธ์ GPU๋ก ์ด๋ํ๋ ๋์์ ๋๋ค. ์ด๋ก ์ธํด ์ถ๋ก ์๊ฐ์ ๋ฏธ์น๋ ์ํฅ์ ๋ฏธ๋ฏธํ์ง๋ง(ํ์ดํ๋ผ์ธ์ 'cuda'๋ก ์ด๋ํ๋ ๊ฒ๊ณผ ๋น๊ตํ์ฌ) ์ฌ์ ํ ์ฝ๊ฐ์ ๋ฉ๋ชจ๋ฆฌ๋ฅผ ์ ์ฝํ ์ ์์ต๋๋ค. | |
์ด ์๋๋ฆฌ์ค์์๋ ํ์ดํ๋ผ์ธ์ ์ฃผ์ ๊ตฌ์ฑ ์์ ์ค ํ๋๋ง(์ผ๋ฐ์ ์ผ๋ก ํ ์คํธ ์ธ์ฝ๋, unet ๋ฐ vae) GPU์ ์๊ณ , ๋๋จธ์ง๋ CPU์์ ๋๊ธฐํ ๊ฒ์ ๋๋ค. | |
์ฌ๋ฌ ๋ฐ๋ณต์ ์ํด ์คํ๋๋ UNet๊ณผ ๊ฐ์ ๊ตฌ์ฑ ์์๋ ๋ ์ด์ ํ์ํ์ง ์์ ๋๊น์ง GPU์ ๋จ์ ์์ต๋๋ค. | |
์ด ๊ธฐ๋ฅ์ ์๋์ ๊ฐ์ด ํ์ดํ๋ผ์ธ์์ `enable_model_cpu_offload()`๋ฅผ ํธ์ถํ์ฌ ํ์ฑํํ ์ ์์ต๋๋ค. | |
```Python | |
import torch | |
from diffusers import StableDiffusionPipeline | |
pipe = StableDiffusionPipeline.from_pretrained( | |
"runwayml/stable-diffusion-v1-5", | |
torch_dtype=torch.float16, | |
) | |
prompt = "a photo of an astronaut riding a horse on mars" | |
pipe.enable_model_cpu_offload() | |
image = pipe(prompt).images[0] | |
``` | |
์ด๋ ์ถ๊ฐ์ ์ธ ๋ฉ๋ชจ๋ฆฌ ์ ์ฝ์ ์ํ attention slicing๊ณผ๋ ํธํ๋ฉ๋๋ค. | |
```Python | |
import torch | |
from diffusers import StableDiffusionPipeline | |
pipe = StableDiffusionPipeline.from_pretrained( | |
"runwayml/stable-diffusion-v1-5", | |
torch_dtype=torch.float16, | |
) | |
prompt = "a photo of an astronaut riding a horse on mars" | |
pipe.enable_model_cpu_offload() | |
pipe.enable_attention_slicing(1) | |
image = pipe(prompt).images[0] | |
``` | |
<Tip> | |
์ด ๊ธฐ๋ฅ์ ์ฌ์ฉํ๋ ค๋ฉด 'accelerate' ๋ฒ์ 0.17.0 ์ด์์ด ํ์ํฉ๋๋ค. | |
</Tip> | |
## Channels Last ๋ฉ๋ชจ๋ฆฌ ํ์ ์ฌ์ฉํ๊ธฐ | |
Channels Last ๋ฉ๋ชจ๋ฆฌ ํ์์ ์ฐจ์ ์์๋ฅผ ๋ณด์กดํ๋ ๋ฉ๋ชจ๋ฆฌ์์ NCHW ํ ์ ๋ฐฐ์ด์ ๋์ฒดํ๋ ๋ฐฉ๋ฒ์ ๋๋ค. | |
Channels Last ํ ์๋ ์ฑ๋์ด ๊ฐ์ฅ ์กฐ๋ฐํ ์ฐจ์์ด ๋๋ ๋ฐฉ์์ผ๋ก ์ ๋ ฌ๋ฉ๋๋ค(์ผ๋ช ํฝ์ ๋น ์ด๋ฏธ์ง๋ฅผ ์ ์ฅ). | |
ํ์ฌ ๋ชจ๋ ์ฐ์ฐ์ Channels Last ํ์์ ์ง์ํ๋ ๊ฒ์ ์๋๋ผ ์ฑ๋ฅ์ด ์ ํ๋ ์ ์์ผ๋ฏ๋ก, ์ฌ์ฉํด๋ณด๊ณ ๋ชจ๋ธ์ ์ ์๋ํ๋์ง ํ์ธํ๋ ๊ฒ์ด ์ข์ต๋๋ค. | |
์๋ฅผ ๋ค์ด ํ์ดํ๋ผ์ธ์ UNet ๋ชจ๋ธ์ด channels Last ํ์์ ์ฌ์ฉํ๋๋ก ์ค์ ํ๋ ค๋ฉด ๋ค์์ ์ฌ์ฉํ ์ ์์ต๋๋ค: | |
```python | |
print(pipe.unet.conv_out.state_dict()["weight"].stride()) # (2880, 9, 3, 1) | |
pipe.unet.to(memory_format=torch.channels_last) # in-place ์ฐ์ฐ | |
# 2๋ฒ์งธ ์ฐจ์์์ ์คํธ๋ผ์ด๋ 1์ ๊ฐ์ง๋ (2880, 1, 960, 320)๋ก, ์ฐ์ฐ์ด ์๋ํจ์ ์ฆ๋ช ํฉ๋๋ค. | |
print(pipe.unet.conv_out.state_dict()["weight"].stride()) | |
``` | |
## ์ถ์ (tracing) | |
์ถ์ ์ ๋ชจ๋ธ์ ํตํด ์์ ์ ๋ ฅ ํ ์๋ฅผ ํตํด ์คํ๋๋๋ฐ, ํด๋น ์ ๋ ฅ์ด ๋ชจ๋ธ์ ๋ ์ด์ด๋ฅผ ํต๊ณผํ ๋ ํธ์ถ๋๋ ์์ ์ ์บก์ฒํ์ฌ ์คํ ํ์ผ ๋๋ 'ScriptFunction'์ด ๋ฐํ๋๋๋ก ํ๊ณ , ์ด๋ just-in-time ์ปดํ์ผ๋ก ์ต์ ํ๋ฉ๋๋ค. | |
UNet ๋ชจ๋ธ์ ์ถ์ ํ๊ธฐ ์ํด ๋ค์์ ์ฌ์ฉํ ์ ์์ต๋๋ค: | |
```python | |
import time | |
import torch | |
from diffusers import StableDiffusionPipeline | |
import functools | |
# torch ๊ธฐ์ธ๊ธฐ ๋นํ์ฑํ | |
torch.set_grad_enabled(False) | |
# ๋ณ์ ์ค์ | |
n_experiments = 2 | |
unet_runs_per_experiment = 50 | |
# ์ ๋ ฅ ๋ถ๋ฌ์ค๊ธฐ | |
def generate_inputs(): | |
sample = torch.randn((2, 4, 64, 64), device="cuda", dtype=torch.float16) | |
timestep = torch.rand(1, device="cuda", dtype=torch.float16) * 999 | |
encoder_hidden_states = torch.randn((2, 77, 768), device="cuda", dtype=torch.float16) | |
return sample, timestep, encoder_hidden_states | |
pipe = StableDiffusionPipeline.from_pretrained( | |
"runwayml/stable-diffusion-v1-5", | |
torch_dtype=torch.float16, | |
).to("cuda") | |
unet = pipe.unet | |
unet.eval() | |
unet.to(memory_format=torch.channels_last) # Channels Last ๋ฉ๋ชจ๋ฆฌ ํ์ ์ฌ์ฉ | |
unet.forward = functools.partial(unet.forward, return_dict=False) # return_dict=False์ ๊ธฐ๋ณธ๊ฐ์ผ๋ก ์ค์ | |
# ์๋ฐ์ | |
for _ in range(3): | |
with torch.inference_mode(): | |
inputs = generate_inputs() | |
orig_output = unet(*inputs) | |
# ์ถ์ | |
print("tracing..") | |
unet_traced = torch.jit.trace(unet, inputs) | |
unet_traced.eval() | |
print("done tracing") | |
# ์๋ฐ์ ๋ฐ ๊ทธ๋ํ ์ต์ ํ | |
for _ in range(5): | |
with torch.inference_mode(): | |
inputs = generate_inputs() | |
orig_output = unet_traced(*inputs) | |
# ๋ฒค์น๋งํน | |
with torch.inference_mode(): | |
for _ in range(n_experiments): | |
torch.cuda.synchronize() | |
start_time = time.time() | |
for _ in range(unet_runs_per_experiment): | |
orig_output = unet_traced(*inputs) | |
torch.cuda.synchronize() | |
print(f"unet traced inference took {time.time() - start_time:.2f} seconds") | |
for _ in range(n_experiments): | |
torch.cuda.synchronize() | |
start_time = time.time() | |
for _ in range(unet_runs_per_experiment): | |
orig_output = unet(*inputs) | |
torch.cuda.synchronize() | |
print(f"unet inference took {time.time() - start_time:.2f} seconds") | |
# ๋ชจ๋ธ ์ ์ฅ | |
unet_traced.save("unet_traced.pt") | |
``` | |
๊ทธ ๋ค์, ํ์ดํ๋ผ์ธ์ `unet` ํน์ฑ์ ๋ค์๊ณผ ๊ฐ์ด ์ถ์ ๋ ๋ชจ๋ธ๋ก ๋ฐ๊ฟ ์ ์์ต๋๋ค. | |
```python | |
from diffusers import StableDiffusionPipeline | |
import torch | |
from dataclasses import dataclass | |
@dataclass | |
class UNet2DConditionOutput: | |
sample: torch.Tensor | |
pipe = StableDiffusionPipeline.from_pretrained( | |
"runwayml/stable-diffusion-v1-5", | |
torch_dtype=torch.float16, | |
).to("cuda") | |
# jitted unet ์ฌ์ฉ | |
unet_traced = torch.jit.load("unet_traced.pt") | |
# pipe.unet ์ญ์ | |
class TracedUNet(torch.nn.Module): | |
def __init__(self): | |
super().__init__() | |
self.in_channels = pipe.unet.config.in_channels | |
self.device = pipe.unet.device | |
def forward(self, latent_model_input, t, encoder_hidden_states): | |
sample = unet_traced(latent_model_input, t, encoder_hidden_states)[0] | |
return UNet2DConditionOutput(sample=sample) | |
pipe.unet = TracedUNet() | |
with torch.inference_mode(): | |
image = pipe([prompt] * 1, num_inference_steps=50).images[0] | |
``` | |
## Memory-efficient attention | |
์ดํ ์ ๋ธ๋ก์ ๋์ญํญ์ ์ต์ ํํ๋ ์ต๊ทผ ์์ ์ผ๋ก GPU ๋ฉ๋ชจ๋ฆฌ ์ฌ์ฉ๋์ด ํฌ๊ฒ ํฅ์๋๊ณ ํฅ์๋์์ต๋๋ค. | |
@tridao์ ๊ฐ์ฅ ์ต๊ทผ์ ํ๋์ ์ดํ ์ : [code](https://github.com/HazyResearch/flash-attention), [paper](https://arxiv.org/pdf/2205.14135.pdf). | |
๋ฐฐ์น ํฌ๊ธฐ 1(ํ๋กฌํํธ 1๊ฐ)์ 512x512 ํฌ๊ธฐ๋ก ์ถ๋ก ์ ์คํํ ๋ ๋ช ๊ฐ์ง Nvidia GPU์์ ์ป์ ์๋ ํฅ์์ ๋ค์๊ณผ ๊ฐ์ต๋๋ค: | |
| GPU | ๊ธฐ์ค ์ดํ ์ FP16 | ๋ฉ๋ชจ๋ฆฌ ํจ์จ์ ์ธ ์ดํ ์ FP16 | | |
|------------------ |--------------------- |--------------------------------- | | |
| NVIDIA Tesla T4 | 3.5it/s | 5.5it/s | | |
| NVIDIA 3060 RTX | 4.6it/s | 7.8it/s | | |
| NVIDIA A10G | 8.88it/s | 15.6it/s | | |
| NVIDIA RTX A6000 | 11.7it/s | 21.09it/s | | |
| NVIDIA TITAN RTX | 12.51it/s | 18.22it/s | | |
| A100-SXM4-40GB | 18.6it/s | 29.it/s | | |
| A100-SXM-80GB | 18.7it/s | 29.5it/s | | |
์ด๋ฅผ ํ์ฉํ๋ ค๋ฉด ๋ค์์ ๋ง์กฑํด์ผ ํฉ๋๋ค: | |
- PyTorch > 1.12 | |
- Cuda ์ฌ์ฉ ๊ฐ๋ฅ | |
- [xformers ๋ผ์ด๋ธ๋ฌ๋ฆฌ๋ฅผ ์ค์นํจ](xformers) | |
```python | |
from diffusers import StableDiffusionPipeline | |
import torch | |
pipe = StableDiffusionPipeline.from_pretrained( | |
"runwayml/stable-diffusion-v1-5", | |
torch_dtype=torch.float16, | |
).to("cuda") | |
pipe.enable_xformers_memory_efficient_attention() | |
with torch.inference_mode(): | |
sample = pipe("a small cat") | |
# ์ ํ: ์ด๋ฅผ ๋นํ์ฑํ ํ๊ธฐ ์ํด ๋ค์์ ์ฌ์ฉํ ์ ์์ต๋๋ค. | |
# pipe.disable_xformers_memory_efficient_attention() | |
``` | |