Diffusers와 PEFT를 사용한 Flux의 빠른 LoRA 추론
Hugging Face는 Flux.1-Dev 텍스트‑투‑이미지 모델에 대한 최적화 레시피를 개발했으며, LoRA 추론 속도를 약 2.3배 향상시킵니다. 이 접근 방식은 서로 다른 LoRA 어댑터를 교체할 때 발생하는 재컴파일 정체 문제를 해결하여, 지연 시간을 희생하지 않고 고성능 맞춤화를 가능하게 합니다.
LoRA 추론 병목 현상 극복
LoRA 추론을 최적화하는 것은 어려운 일입니다. 서로 다른 어댑터를 hotswap하면(랭크가 다르고 대상 레이어가 다를 수 있음) 일반적으로 모델 아키텍처가 변경되기 때문입니다. 표준 워크플로에서는 특정 LoRA가 적용된 모델에 torch.compile을 사용하면 속도가 향상되지만, 해당 LoRA를 다른 것으로 교체하면 그래프가 재컴파일되어 추론 속도가 크게 저하됩니다.
이를 해결하기 위해 Hugging Face는 “hotswapping” 메커니즘을 사용합니다. Diffusers에서 hotswap=True를 설정하면 모델 아키텍처는 그대로 유지되고 LoRA 어댑터의 가중치만 교체됩니다. 다음 조건이 충족되면 재컴파일이 필요하지 않습니다:
- Maximum Rank Definition:
max_rank는 풀에서 가장 큰 어댑터를 수용하도록 미리 지정되어야 합니다. - Layer Consistency: 이후 LoRA는 첫 번째 로드된 LoRA가 대상한 동일한 레이어 또는 그 하위 집합을 대상으로 해야 합니다.
- Text Encoder Limitation: 현재 hotswapping은 텍스트 인코더를 대상으로 하는 것을 지원하지 않습니다.
고성능 GPU용 최적화 레시피
NVIDIA H100 GPU와 같은 고성능 하드웨어의 경우, 최적화된 추론 파이프라인은 네 가지 핵심 요소를 결합합니다:
- Flash Attention 3 (FA3): 어텐션 메커니즘 효율성을 향상시킵니다.
torch.compile: 실행 그래프를 최적화하는 JIT 컴파일러입니다.- FP8 Quantization: TorchAO를 통해 제공되며, 속도와 메모리 사용량 사이의 좋은 균형을 제공하지만 손실이 있습니다.
- Hotswapping: 재컴파일을 일으키지 않고 어댑터 전환을 가능하게 합니다.
성능 벤치마크 (H100)
| 옵션 | 시간 (초) | 속도 향상 (기준 대비) | 비고 |
|---|---|---|---|
| 기준 | 7.8910 | – | 기준 |
| 최적화 | 3.5464 | 2.23× | Hotswapping + 컴파일 + FP8 |
| FP8 없음 | 4.3520 | 1.81× | FP8 양자화 없이 최적화 |
| FA3 없음 | 4.3020 | 1.84× | Flash Attention 3 없이 최적화 |
| 기준 + 컴파일 | 5.0920 | 1.55× | 컴파일은 활성화되었지만 재컴파일 정체가 발생합니다 |
소비자 GPU (RTX 4090) 최적화
Flux.1-Dev를 Bfloat16으로 실행하려면 약 33GB VRAM이 필요하며, 이는 RTX 4090과 같은 소비자 GPU의 24GB 용량을 초과합니다. 이 하드웨어에서 모델을 사용할 수 있도록 Hugging Face는 특정 메모리 감소 전략을 사용합니다:
- T5 Text Encoder Quantization:
bitsandbytes의 NF4 양자화를 사용하여 텍스트 인코더의 메모리 사용량을 줄입니다. - FP8 Quantization: Flux Transformer에 적용됩니다.
- Regional Compilation:
compile_repeated_blocks를 사용하여 컴파일 시간과 메모리 사용량을 줄입니다.
RTX 4090에서 FP8 양자화, torch.compile, 그리고 T5 양자화(NF4)를 결합하면 2.04배의 속도 향상이 이루어졌으며, 추론 시간이 23.6060초(기준)에서 11.5715초로 감소했습니다.
Hotswapping 기술 구현
어댑터 교체 시 재컴파일을 방지하기 위해 구현에서는 두 가지 기술적 과제를 해결합니다:
- Tensor Conversion: LoRA 스케일링 팩터를 float에서 torch 텐서로 변환합니다.
- Weight Padding: LoRA 가중치를 가장 큰 필요 형태(
max_rank로 정의)로 패딩합니다. 이를 통해 전체 속성을 재할당하지 않고 가중치 데이터를 교체할 수 있습니다. 패딩은 수학적 결과가 변하지 않도록 0으로 채워지며, 과도한 패딩은 계산 속도에 약간 영향을 줄 수 있습니다.
이 워크플로를 구현하는 경우, LoRA 로드 순서는 매우 중요합니다. 어댑터가 서로 다른 레이어를 대상으로 하면, 모든 필요한 레이어의 합집합을 대상으로 하는 더미 LoRA를 먼저 생성해야 합니다.