2025년, NVIDIA는 Black Forest Labs(BFL)와 협력하여 FLUX.1 텍스트-이미지 모델 시리즈를 최적화했으며, 이를 통해 NVIDIA Blackwell GeForce RTX 50 시리즈 GPU에서 FP4 이미지 생성 성능을 구현했습니다.
잠재 확산 모델의 자연스러운 확장선인 FLUX.1 Kontext [dev]는 인컨텍스트 학습이 거대 언어 모델(LLM)뿐만 아니라 시각 생성 모델에서도 유효한 기술임을 입증한 것인데요. NVIDIA는 이러한 경험을 더욱 널리 확산시키기 위해 BFL과 협력하여 저정밀도 양자화를 활용한 실시간에 가까운 편집 경험을 구현했습니다.
FLUX.2는 멀티 이미지 참조 기능을 제공하며 최고 수준의 기업용 모델에 필적하는 품질을 갖춘 중대한 도약의 결과물입니다. 하지만 FLUX.2 [dev]는 상당한 컴퓨팅 자원을 요구하기 때문에, BFL과 Comfy, NVIDIA는 협력을 통해 중대한 돌파구를 마련했습니다. 바로 FLUX.2 [dev]의 메모리 요구 사항을 40% 이상 절감하고 ComfyUI를 통한 로컬 배포를 가능하게 한 것입니다. FP8 정밀도를 사용한 이 최적화 덕분에 FLUX.2 [dev]는 이미지 생성 분야에서 가장 인기 있는 모델 중 하나가 되었습니다.
FLUX.2 [dev]가 오픈 웨이트 모델의 표준으로 자리 잡음에 따라, NVIDIA 팀은 BFL과 협력하여 다음 단계의 성능 도약을 공유합니다. 바로 NVIDIA DGX B200 및 NVIDIA DGX B300을 포함한 가장 강력한 데이터 센터용 NVIDIA Blackwell GPU에서 구현되는 FLUX.2 [dev]용 4비트 가속 기술입니다.
본 포스팅에서는 코드 스니펫과 시작 단계를 포함하여, 해당 NVIDIA 데이터 센터 아키텍처에서 FLUX.2 [dev]를 가속하기 위해 사용된 다양한 추론 최적화 기술을 살펴봅니다. 이러한 최적화 기술들이 결합되어 지연 시간(Latency)을 획기적으로 줄였으며, 데이터 센터 GPU에서의 효율적인 배포를 가능하게 했습니다.
FLUX.2 [dev] 기반 BF16과 NVFP4의 시각적 비교
세부 사항을 논하기에 앞서, 기본 BF16 정밀도에서의 FLUX.2 [dev] 출력 품질과 NVFP4를 사용하여 얻은 결과가 매우 유사함을 확인해 보시기 바랍니다 (그림 1 및 그림 2 참조).
그림 1에 사용된 프롬프트는 다음과 같습니다. “포근한 소파 위에서 평온하게 낮잠을 자고 있는 고양이. 소파는 달 표면에서 자라난 높은 나무 꼭대기에 놓여 있다. 멀리 보이는 지구는 우주의 어둠 속에서 활기찬 파란색과 초록색 보석처럼 빛난다. 매끈한 우주선이 근처에 떠서 장면에 부드러운 빛을 비추고 있으며, 전체적인 디지털 아트 구성은 꿈결 같은 분위기를 자아낸다.”
그림 2에 사용된 프롬프트는 다음과 같습니다. “정장을 차려입고 귀가하던 중 우산도 없이 폭우를 만난 커플을 묘사한 유화.” 이 사례의 경우 두 결과물 간의 차이를 식별하기가 더욱 어렵습니다. 가장 눈에 띄는 차이점은 BF16 이미지 속 남성의 미소와 NVFP4 이미지 배경에 나타난 여러 개의 우산 정도입니다. 이를 제외하면 전경과 배경의 세밀한 디테일 대부분이 두 이미지 모두에서 잘 보존되어 있습니다.
FLUX.2 [dev] 최적화
FLUX.2 [dev] 모델은 텍스트 임베딩 모델(Mistral Small 3), 확산 트랜스포머(Diffusion Transformer) 모델, 그리고 오토인코더라는 세 가지 핵심 구성 요소로 이루어져 있습니다. NVIDIA 팀은 TensorRT-LLM/feat/visual_gen 브랜치에 마련된 프로토타입 런타임을 사용하여, 오픈 소스 디퓨저(Diffusers) 구현체에 다음과 같은 최적화 기술을 적용했습니다.
- NVFP4 양자화
- 타임스텝 임베딩 인식 캐싱 (TeaCache)
- CUDA 그래프 (CUDA Graphs)
- Torch compile
- 멀티 GPU 지원
NVFP4 양자화
NVFP4는 2단계 마이크로블록 스케일링 전략을 도입하여 마이크로스케일링 데이터 포맷(Microscaling Data Formats)의 개념을 한 단계 발전시켰습니다. 이 방식은 정확도 저하를 최소화하도록 설계되었으며, 텐서별 스케일링(Per-tensor scaling)과 블록별 스케일링(Per-block scaling)이라는 두 가지 차별화된 메커니즘을 특징으로 합니다.
텐서별 스케일링은 FP32 정밀도로 저장되는 값으로, 전체 텐서 분포를 조정하며 정적 또는 동적으로 계산될 수 있습니다. 이와 대조적으로 블록별 스케일링은 텐서를 16개 요소 단위의 블록으로 나누어 실시간으로 동적으로 계산됩니다.
사용자는 유연성을 극대화하기 위해 특정 레이어를 고정밀도로 유지하거나, 아래의 FLUX.2 [dev] 적용 예시와 같이 동적 양자화를 적용하도록 선택할 수 있습니다.
exclude_pattern =
r"^(?!.*(embedder|norm_out|proj_out|to_add_out|to_added_qkv|stream)).*"
NVFP4 연산은 다음과 같은 구문을 통해 적용됩니다.
from visual_gen.layers import apply_visual_gen_linear
apply_visual_gen_linear(
model,
load_parameters=True,
quantize_weights=True,
exclude_pattern=exclude_pattern,
)
TeaCache
추론 프로세스의 가속화를 위해 TeaCache 기술이 사용되었습니다. TeaCache는 확산 과정 중 이전에 생성된 잠재 변수(Latent)를 활용하여 특정 확산 단계를 조건부로 건너뛰는 방식입니다. 이 효과를 정량화하기 위해 테스트를 진행한 결과, 20개의 프롬프트와 50단계의 추론 프로세스 시나리오에서 TeaCache는 평균 16단계를 우회하였으며, 이는 추론 지연 시간을 약 30% 단축하는 결과로 이어졌습니다.
TeaCache 하이퍼파라미터의 최적 구성을 결정하기 위해 그리드 탐색(Grid search) 방법론을 도입하였습니다. 이 구성은 연산 속도와 생성 품질 사이의 최적의 균형점을 제공합니다.
dit_configs = {
...
"teacache": {
"enable_teacache": True,
"use_ret_steps": True,
"teacache_thresh": 0.05,
"ret_steps": 10,
"cutoff_steps": 50,
},
...
}
캐싱 메커니즘을 위한 스케일링 인자는 경험적으로 결정되었으며, 3차 다항식을 통해 근사화되었습니다. 이 다항식은 텍스트-이미지 및 멀티 참조 이미지 생성 사례를 포함하는 캘리브레이션 데이터셋을 사용하여 피팅되었습니다.
그림 3은 이러한 경험적 접근 방식을 시각화한 것으로, 원시 캘리브레이션 데이터 포인트와 함께 변조된 입력 차이와 모델 출력 차이 간의 관계를 모델링한 3차 다항식 곡선(붉은색)을 보여줍니다.
CUDA 그래프
NVIDIA TensorRT-LLM visual_gen은 CUDA 그래프 캡처를 지원하기 위해 즉시 사용 가능한 래퍼(Wrapper)를 제공합니다. 래퍼를 임포트하고 forward 함수를 교체하는 것만으로 간단히 적용할 수 있습니다.
from visual_gen.utils.cudagraph import cudagraph_wrapper
model.forward = cudagraph_wrapper(model.forward)
Torch compile
FLUX.2 [dev]는 기본적으로 torch.compile이 활성화되어 있지 않으나, 본 팀은 베이스라인 실행을 제외한 모든 실험에서 이를 활성화하였습니다.
model = torch.compile(model)
멀티 GPU 지원
TensorRT-LLM visual_gen을 사용하여 멀티 GPU를 활성화하는 과정은 다음 네 단계로 구성됩니다.
- GPU 간 통신을 처리하는 코드를 삽입하도록
model.forward함수 수정 - 모델의 어텐션 구현을
ditAttnProcessor로 교체 - 설정 파일에서 병렬 알고리즘 선택 및 병렬 처리 크기 설정
- torchrun으로 실행
다음 스니펫은 그 예시입니다. 입력 데이터를 여러 GPU에 분산시키기 위해 model.forward 시작 부분에 분할 코드를 삽입하세요.
from visual_gen.utils import (
dit_sp_gather,
dit_sp_split,
)
# ...
hidden_states = dit_sp_split(hidden_states, dim=1)
encoder_hidden_states = dit_sp_split(encoder_hidden_states, dim=1)
img_ids = dit_sp_split(img_ids, dim=1)
txt_ids = dit_sp_split(txt_ids, dim=1)
이어서, 반환(return) 직전인 model.forward 끝부분에 수집(gather) 코드를 삽입하세요.
output = dit_sp_gather(output, dim=1)
그다음, 기존의 어텐션 구현을 여러 GPU 간의 적절한 통신을 보장하는 제공된 오리지널 어텐션 프로세서로 교체하세요.
from visual_gen.layers import ditAttnProcessor
# ...
def attention(...):
# ...
x = ditAttnProcessor().visual_gen_attn(q, k, v, tensor_layout="HND")
# ...
설정 파일에서 올바른 병렬 크기를 설정하세요. 예를 들어, 4개의 GPU에서 Ulysses 병렬 처리를 사용하려면 다음과 같이 설정합니다.
dit_config = {
...
"parallel": {
"dit_ulysses_size": 4,
}
...
}
마지막으로, setup_configs API를 호출하여 설정을 활성화하십시오.
visual_gen.setup_configs(**dit_configs)
멀티 GPU를 사용할 때는 반드시 torchrun으로 스크립트를 실행해야 합니다. TensorRT-LLM visual_gen은 torchrun의 랭크(rank) 정보를 활용하여 모든 통신과 작업 분할을 정확하게 처리합니다.
성능 분석
저정밀도 커널, 캐싱 기술, 멀티 GPU 추론을 포함한 모든 추론 최적화 기법이 엔드 투 엔드(End-to-End) FLUX.2 [dev] 예제에 포함되었습니다.
그림 4에 나타난 바와 같이, NVIDIA DGX B200 아키텍처는 기본 BF16 정밀도를 사용할 때도 NVIDIA H200 대비 1.7배의 세대 간 성능 도약을 보여줍니다. 나아가 CUDA 그래프, torch.compile, NVFP4 정밀도, TeaCache 등의 추론 최적화를 단계적으로 적용하면 단일 B200의 성능이 베이스라인 대비 무려 6.3배까지 향상됩니다.
최종적으로 2개의 B200 구성을 통한 멀티 GPU 추론은 현재 업계 표준인 H200 대비 10.2배의 성능 향상을 달성하였습니다.
베이스라인은 어떠한 최적화나 torch.compile 활성화도 적용되지 않은 순정 상태의 FLUX.2 [dev]입니다. 최적화 시리즈에는 torch.compile, CUDA 그래프, NVFP4 및 TeaCache 활성화가 포함되었습니다. 벤치마크에 사용된 확산 단계는 50단계입니다.
단일 GPU에서 본 팀은 NVFP4와 TeaCache가 속도 향상과 출력 품질 사이에서 훌륭한 균형을 제공하며, 각각 약 2배의 속도 향상을 구현함을 확인하였습니다. torch.compile은 대부분의 개발자에게 익숙한 손실 없는 가속 기술이지만 그 혜택은 제한적이었습니다. CUDA 그래프는 주로 멀티 GPU 추론에서 유익하며, NVIDIA B200에서 여러 GPU를 사용한 점진적 스케일링을 가능하게 합니다. 마지막으로 전체 파이프라인은 텍스트 인코더의 FP8 양자화와 함께 견고함을 입증하였으며, 대규모 배포 시 추가적인 이점을 제공합니다.
멀티 GPU 환경에서 TensorRT-LLM visual_gen 시퀀스 병렬 처리는 GPU를 추가함에 따라 선형에 가까운 스케일링 성능을 보여줍니다. 이러한 효과는 NVIDIA Blackwell B200 및 GB200은 물론, NVIDIA Blackwell Ultra B300 및 GB300 GPU에서도 동일하게 관찰됩니다. 현재 NVIDIA Blackwell Ultra GPU를 위한 추가 최적화 작업이 진행 중입니다.
NVIDIA Blackwell GPU 기반 FLUX.2 시작하기
FLUX.2는 고품질 출력과 사용자 친화적인 배포 옵션을 성공적으로 결합한 이미지 생성 분야의 혁신입니다. NVIDIA 팀은 BFL과의 협력을 통해 가장 강력한 NVIDIA 데이터 센터 GPU에서 FLUX.2 [dev]의 비약적인 가속을 실현하였습니다.
NVFP4 양자화 및 TeaCaching을 포함한 새로운 기술을 FLUX.2 [dev] 모델에 적용함으로써 추론 속도에서 강력한 세대적 도약을 이루어냈습니다. 이러한 최적화의 결합은 지연 시간을 획기적으로 단축하여 NVIDIA 데이터 센터 GPU에서의 효율적인 배포를 가능하게 합니다.
최신 최적화 기술이 적용된 추론 파이프라인 구축을 시작하시려면, NVIDIA/TensorRT-LLM/visual_gen GitHub 리포지토리에서 엔드 투 엔드 FLUX.2 예제와 관련 코드를 확인해 보시기 바랍니다.
