Flux 로 위한 Diffusers 와 PEFT 를 사용한 빠른 LoRA 추론

이번 글에서는 텍스트 생성을 위한 Flux.1-Dev 모델을 널리 쓰이고 채택된다는 점과 LoRA 를 사용할 때 추론 속도를 최적화하는 방법 (~2.3 배) 을 다룹니다. Hugging Face Hub 플랫폼에서 보고된 바와 같이, 이 모델은 3 만 개 이상의 어댑터가 훈련되었습니다. 따라서 커뮤니티에 대한 중요성이 큽니다.

Flux 를 사용하여 속도 향상을 보여주는 것은 물론이지만, 우리의 신념은 우리가 제시한 레시피가 다른 모델에도 적용될 수 있는 만큼 일반적이라는 것입니다.

코드를 바로 시작하고 싶으시다면, 동반 코드 저장소를 확인하세요.

• LoRA 추론 최적화의 장애물
• 최적화 레시피
• 소비자 GPU 의 최적화된 LoRA 추론
• 결론

LoRA 를 서비스할 때, 다른 LoRA 를 교체 (hotswap) 하는 것은 일반적입니다. LoRA 는 기본 모델 아키텍처를 변경합니다. 또한 LoRA 들은 서로 다를 수 있습니다 – 각 LoRA 는 다른 랭크와 적응을 위한 다른 레이어를 가질 수 있습니다. 이러한 LoRA 의 동적 속성을 고려하기 위해, 우리가 적용하는 최적화가 견고하도록 필요한 단계를 취해야 합니다.

예를 들어, 우리는 특정 LoRA 로 로드된 모델에 torch.compile 를 적용하여 추론 지연 시간을 개선할 수 있습니다. 그러나 다른 LoRA (잠재적으로 다른 구성을 가진) 로 교체하는 순간, 재컴파일 문제가 발생하여 추론이 느려집니다.

LoRA 파라미터를 기본 모델 파라미터와 융합하고, 컴파일을 실행한 후 새로운 ones 을 로드할 때 LoRA 파라미터를 다시 융합 해제 (unfuse) 할 수도 있습니다. 그러나 이 접근법은 잠재적인 아키텍처 수준의 변경으로 인해 추론이 실행될 때마다 재컴파일 문제를 다시 마주하게 됩니다.

우리의 최적화 레시피는 위의 상황을 고려하여 최대한 현실적입니다. 아래에 우리의 최적화 레시피의 핵심 구성 요소가 나열됩니다:

• Flash Attention 3 (FA3)
  torch.compile

• TorchAO 의 FP8 양자화

• Hotswapping-ready

위의 것들 중에서, FP8 양자화는 손실적이지만 종종 가장 강력한 속도 - 메모리 트레이드오프를 제공합니다. 우리는 NVIDIA GPU 를 사용하여 레시피를 테스트했지만, AMD GPU 에서도 작동할 것입니다.

우리의 이전 블로그 글 (1 번 글과 2 번 글) 에서는 우리의 최적화 레시피의 첫 세 가지 구성 요소를 사용하는 이점을 이미 논의했습니다. 하나씩 적용하는 것은 몇 줄의 코드입니다:

from diffusers import DiffusionPipeline, TorchAoConfig
from diffusers.quantizers import PipelineQuantizationConfig
from utils.fa3_processor import FlashFluxAttnProcessor3_0
...

FA3 프로세서 는 여기에서 나옵니다.

문제는 컴파일된 확산 트랜스포머 (pipe.transformer) 에 LoRA 를 교체하고, 재컴파일을 트리거하지 않고 LoRA 를 교체할 때 나타납니다.

일반적으로 LoRA 를 로드하고 해제하는 것은 재컴파일을 필요로 하며, 이는 컴파일에서 얻은 속도 이점을 무효화합니다. 다행히도 재컴파일의 필요성을 피하는 방법이 있습니다. hotswap=True 를 전달하면 diffusers 는 모델 아키텍처를 변경하지 않고 LoRA 어댑터 자체의 가중치만 교환하며, 이는 재컴파일을 필요로 하지 않습니다.

pipe.enable_lora_hotswap(target_rank=max_rank)
pipe.load_lora_weights(<lora-adapter-name1>)
# 첫 번째 LoRA 를 로드한 후 컴파일 *
...

(기억해 두세요, pipe 의 첫 번째 호출은 torch.compile 가 just-in-time 컴파일러이므로 느릴 수 있습니다. 그러나 이후 호출들은 훨씬 더 빨라야 합니다.)

이것은 일반적으로 LoRA 를 재컴파일 없이 교체할 수 있게 하지만, 다음과 같은 제한 사항이 있습니다:

• 모든 LoRA 어댑터의 최대 랭크를 미리 제공해야 합니다. 따라서 랭크 16 의 어댑터 하나와 랭크 32 의 어댑터 하나를 가진 경우, max_rank=32 를 전달해야 합니다.
• Hotswapped 된 LoRA 어댑터는 첫 번째 LoRA 가 타겟팅한 레이어와 동일한 레이어 또는 그 하위 집합만 타겟팅할 수 있습니다.
• Text encoder 를 타겟팅하는 것은 아직 지원되지 않습니다.

Diffusers 의 hotswapping 과 그 제한 사항에 대한 자세한 정보는 문서를 방문하세요.

이 워크플로의 이점은 hotswapping 을 사용하지 않고 컴파일 없이 추론 지연 시간을 볼 때 명확해집니다.

Option	Time (s) ⬇️	Speedup (vs baseline) ⬆️	Notes
baseline	7.8910	–	Baseline
...
Key takeaways:

• "regular + compile" 옵션은 regular 옵션에 비해 합리적인 속도 향상을 제공하지만, 재컴파일 문제를 유발하여 전체 실행 시간을 증가시킵니다. 우리의 벤치마크에서는 컴파일 시간을 제시하지 않습니다.
• Hotswapping 을 통해 (또는 "optimized" 옵션으로 알려진) 재컴파일 문제가 제거되면 가장 높은 속도 향상을 달성합니다.
• "optimized" 옵션에서 FP8 양자화는 활성화되어 품질 손실을 초래할 수 있습니다. FP8 을 사용하지 않더라도 합리적인 양의 속도 향상 ("no_fp8" 옵션) 을 얻습니다.
• 데모 목적으로는 컴파일과 함께 hotswapping 을 위해 두 개의 LoRA 풀을 사용합니다. 전체 코드는 동반 코드 저장소를 참조하세요.

지금까지 논의한 최적화 레시피는 H100 과 같은 강력한 GPU 에 대한 액세스를 가정합니다. 그러나 RTX 4090 과 같은 소비자용 GPU 를 사용해야 하는 경우 어떻게 할 수 있을까요? 확인해 봅시다.

Flux.1-Dev (어떤 LoRA 없이), Bfloat16 데이터 타입을 사용하여 실행하면 약 33GB 의 메모리를 사용합니다. LoRA 모듈의 크기에 따라 최적화를 사용하지 않는 경우, 이 메모리 발량은 더욱 증가할 수 있습니다. 많은 소비자용 GPU 는 RTX 4090 만이 24GB 를 가지고 있습니다. 이 섹션의 나머지에서는 RTX 4090 머신을 테스트베드 (testbed) 로 고려합니다.

먼저 Flux.1-Dev 의 엔드 투 엔드 실행을 활성화하려면, 현재 계산에 필요하지 않은 컴포넌트를 CPU 로 오프로드하여 더 많은 가속기 메모리를 확보할 수 있습니다. 이를 통해 RTX 4090 에서 전체 파이프라인을 약 35.403 초 내에 실행할 수 있습니다. 컴파일을 활성화하면 지연 시간을 31.205 초로 줄일 수 있습니다 (1.12x 속도 향상). 코드 측면에서는 몇 줄만입니다:

pipe = DiffusionPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16,
)
...

여기서 FP8 양자화를 적용하지 않았음을 주목하세요. CPU 오프로드와 컴파일 (지원 이슈 스레드) 과 호환되지 않기 때문입니다. 따라서 Flux Transformer 에 FP8 양자화를 적용하는 것만으로는 메모리 고갈 문제를 완화하기에는 충분하지 않습니다. 이 경우 우리는 이를 제거했습니다.

따라서 FP8 양자화 스키마를 활용하려면 CPU 오프로드 없이 수행하는 방법을 찾아야 합니다. Flux.1-Dev 에서 T5 text encoder 에 추가로 양자화를 적용하면 24GB 로 전체 파이프라인을 로드하고 실행할 수 있습니다. 아래는 T5 text encoder 를 양자화하지 않은 경우와 양자화된 경우 (bitsandbytes 의 NF4 양자화) 의 결과 비교입니다.

위 그림에서 알 수 있듯이 T5 텍스트 인코더를 양자화 (quantizing) 하더라도 품질 손실이 크게 발생하지 않습니다. 양자화된 T5 텍스트 인코더와 FP8-양자화된 Flux Transformer 를 torch.compile 와 함께 결합하면 32.27 초 에서 9.668 초 로 (약 3.3 배의 속도 향상) 질감 감지 없이 합리적인 결과를 얻을 수 있습니다.

T5 텍스트 인코더를 양자화하지 않더라도 24 GB VRAM 으로 이미지를 생성하는 것은 가능하지만, 이는 생성 파이프라인을 다소 복잡하게 만들었을 것입니다.

이제 우리는 RTX 4090 에서 FP8 양자화를 사용하여 전체 Flux.1-Dev 파이프라인을 실행할 수 있는 방법을 가졌습니다. 이전으로 확립된 LoRA 추론 최적화 레시피를 동일한 하드웨어에 적용할 수 있습니다. FA3 는 RTX 4090 에서 지원되지 않으므로, T5 양자화가 추가된 다음 최적화 레시피를 유지하겠습니다:

• FP8 양자화
  torch.compile

• Hotswapping-ready (핫스왑 준비)

• T5 양자화 (NF4 포함)

아래 표에서는 위 구성 요소의 다른 조합을 적용한 추론 지연 시간 수치를 보여줍니다.

옵션	주요 아규스 플래그	시간 (초) ⬇️	속도 향상 (기준 대비) ⬆️
baseline	disable_fp8=False disable_compile=True quantize_t5=True offload=False
23.6060	–
optimized	disable_fp8=False disable_compile=False quantize_t5=True offload=False
11.5715	2.04×

빠른 노트:

• 컴파일레이션은 기준에 비해 약 2 배의 속도 향상을 제공합니다.
• 다른 옵션들은 오프로딩이 활성화되어도 OOM 오류를 발생시켰습니다.

핫스왑을 트리거링하지 않고 활성화하려면 두 가지 장벽을 극복해야 합니다. 첫째, LoRA 스케일링 인자는 부동 소수점 (floats) 에서 torch 텐서로 변환되어야 하며, 이는 비교적 쉽게 달성됩니다. 둘째, LoRA 가중치의 모양이 필요한 최대 모양으로 패딩되어야 합니다. 이를 통해 데이터가 전체 속성을 재할당할 필요가 없이 교체될 수 있습니다. 이것이 위에서 논의된 max_rank 인자가 중요한 이유입니다. 우리는 값을 0 으로 패딩하므로 결과는 변하지 않지만, 패딩 크기에 따라 계산은 약간 느려집니다.

새로운 LoRA 속성이 추가되지 않으므로, 이는 첫 번째 LoRA 이후의 각 LoRA 가 첫 번째 LoRA 가 타겟으로 하는 레이어 또는 그 하위 집합만 대상으로 할 수 있음을 의미합니다. 따라서 로딩 순서를 지혜롭게 선택하세요. LoRA 가 불연속적인 레이어를 타겟으로 한다면, 모든 타겟 레이어의 합집합을 타겟으로 하는 dummy LoRA 를 생성할 가능성이 있습니다.

이 구현의 세부 사항을 보려면 PEFT 의 hotswap.py 파일을 방문하세요.

이 포스트는 Flux 와 빠른 LoRA 추론을 위한 최적화 레시피를 개요로 제시하며, 상당한 속도 향상을 보여줍니다. 우리의 접근 방식은 Flash Attention 3, torch.compile, 및 FP8 양자화를 결합하고 재컴파일레이션 문제를 해결하지 않고 핫스왑 기능을 보장합니다. H100 과 같은 고성능 GPU 에서는 이 최적화된 설정이 기준에 비해 2.23 배의 속도 향상을 제공합니다.

소비자 GPU, 특히 RTX 4090 의 경우, 우리는 T5 텍스트 인코더 양자화 (NF4) 를 도입하고 지역 컴파일레이션을 활용하여 메모리 제한을 극복했습니다. 이 포괄적인 레시피는 상당한 2.04 배의 속도 향상을 달성하여 VRAM 이 제한되어 있어도 Flux 로 LoRA 추론이 가능하고 성능을 발휘하게 했습니다. 핵심 통찰은 컴파일레이션과 양자화를 신중하게 관리함으로써 LoRA 의 이점을 다른 하드웨어 구성에 걸쳐 완전히 실현할 수 있다는 것입니다.

이 포스트의 레시피가 LoRA 기반 사용 사례를 최적화하고 빠른 추론을 누릴 수 있도록 영감을 주기를 바랍니다.

다음은 이 포스트 전반에 걸쳐 인용한 중요한 리소스 목록입니다:

• Flux Fast 소개: H100에서 Flux 가속하기
• torch.compile 과 Diffusers: 최고 성능을 위한 실전 가이드
• Diffusers 내 LoRA 가이드

이 접근법을 시도해 보려는 독자를 위해 Parag Ekbote 가 Replicate 에서 배포한 데모가 있습니다. 이 데모는 NVIDIA L40 및 A100 GPU 에서 설계되고 테스트되었으며, 비교 가능한 성능을 입증했습니다.