(LoRA) Fine-Tuning FLUX.1-dev on Consumer Hardware

🖼️ 인포그래픽

(LoRA) Fine-Tuning FLUX.1-dev on Consumer Hardware 내용을 설명하는 본문 이미지

🖼️ 4컷 인포그래픽

💡 한 줄 요약

FLUX.1-dev의 트랜스포머를 4비트로 양자화하고 LoRA, 8비트 옵티마이저, 체크포인팅, 잠재 표현·텍스트 임베딩 캐시를 결합해 단일 GPU에서 약 10GB 미만의 VRAM으로 미세 조정하는 방법을 설명한다.

📌 핵심 요약

목표는 소규모 데이터셋을 이용해 FLUX.1-dev가 알폰스 무하의 예술적 스타일을 따르도록 미세 조정하면서, 최대 GPU 메모리 사용량을 약 10GB 미만으로 낮추는 것이다.
FLUX.1-dev의 CLIP·T5 텍스트 인코더와 VAE는 고정하고, 핵심 트랜스포머만 4비트 NF4 형식으로 불러온 뒤 LoRA 어댑터를 학습한다.
QLoRA는 양자화된 기반 모델 위에서 FP16 또는 BF16 LoRA 어댑터만 갱신하며, 예시 설정에서는 약 119억 개 전체 파라미터 가운데 약 467만 개만 학습 대상으로 삼는다.
8비트 AdamW, 그래디언트 체크포인팅, VAE 잠재 표현 캐시를 적용해 옵티마이저 상태와 활성화, 반복적인 이미지 인코딩에 필요한 메모리와 계산량을 줄인다.
CLIP·T5의 텍스트 임베딩까지 사전에 계산해 Parquet 파일로 저장하면 학습 중에는 양자화된 트랜스포머와 LoRA 어댑터만 주로 GPU에 남길 수 있으며, 제시된 실행 예시는 RTX 4090에서 BF16과 배치 크기 1을 사용한다.

🧩 주요 포인트

목표는 소규모 데이터셋을 이용해 FLUX.1-dev가 알폰스 무하의 예술적 스타일을 따르도록 미세 조정하면서, 최대 GPU 메모리 사용량을 약 10GB 미만으로 낮추는 것이다.
FLUX.1-dev의 CLIP·T5 텍스트 인코더와 VAE는 고정하고, 핵심 트랜스포머만 4비트 NF4 형식으로 불러온 뒤 LoRA 어댑터를 학습한다.
QLoRA는 양자화된 기반 모델 위에서 FP16 또는 BF16 LoRA 어댑터만 갱신하며, 예시 설정에서는 약 119억 개 전체 파라미터 가운데 약 467만 개만 학습 대상으로 삼는다.
8비트 AdamW, 그래디언트 체크포인팅, VAE 잠재 표현 캐시를 적용해 옵티마이저 상태와 활성화, 반복적인 이미지 인코딩에 필요한 메모리와 계산량을 줄인다.
CLIP·T5의 텍스트 임베딩까지 사전에 계산해 Parquet 파일로 저장하면 학습 중에는 양자화된 트랜스포머와 LoRA 어댑터만 주로 GPU에 남길 수 있으며, 제시된 실행 예시는 RTX 4090에서 BF16과 배치 크기 1을 사용한다.

🧠 상세 정리

1. 추론 최적화에서 미세 조정으로

글은 앞선 양자화 백엔드 비교가 FLUX.1-dev 같은 확산 모델의 추론 메모리를 줄이는 데 초점을 맞췄다는 설명에서 출발한다. bitsandbytes와 torchao 등의 양자화 기법을 이용하면 이미지 생성 성능을 크게 훼손하지 않으면서 모델의 메모리 점유를 낮출 수 있지만, 모델을 사용자의 목적과 스타일에 맞게 만들려면 추론만으로는 충분하지 않다고 본다. 이에 이번 글은 diffusers 기반 QLoRA를 이용해 단일 GPU에서 FLUX.1-dev를 효율적으로 미세 조정하고, 최대 VRAM 사용량을 약 10GB 미만으로 억제하는 구성을 다룬다. 실험 장비로는 VRAM 24GB의 NVIDIA RTX 4090을 사용하며, 호환 하드웨어에서 torchao의 FP8 학습으로 속도를 더 최적화할 수 있다는 방향도 서두에서 제시한다.

2. 데이터셋과 학습 범위

미세 조정의 구체적인 목표는 소규모 데이터셋을 사용해 black-forest-labs/FLUX.1-dev가 알폰스 무하의 예술적 스타일을 따르도록 만드는 것이다. FLUX 모델은 크게 CLIP과 T5로 구성된 텍스트 인코더, 주 모델인 Flux 트랜스포머, 이미지와 잠재 표현 사이를 변환하는 VAE의 세 부분으로 나뉜다. 이 QLoRA 구성에서는 세 요소를 모두 학습하지 않고 트랜스포머에만 LoRA 어댑터를 추가해 최적화한다. 텍스트 인코더와 VAE는 학습 내내 고정되며, 이후에는 두 구성 요소가 매번 GPU에서 실행되지 않도록 출력 결과까지 미리 계산하고 캐시하는 방식으로 메모리 사용을 더 줄인다.

3. LoRA와 QLoRA의 작동 원리

LoRA는 기존의 전체 가중치 행렬을 직접 갱신하는 대신, 가중치 변화량을 두 개의 작은 저랭크 행렬 A와 B의 곱으로 표현한다. 변화량은 ΔW=BA로 계산되며, 랭크 r을 원래 행렬 차원보다 훨씬 작게 설정하기 때문에 학습해야 할 파라미터 수가 크게 감소한다. LoRA의 알파 값은 어댑터 활성의 크기를 조절해 사전 학습 모델과 새 어댑터가 미치는 영향을 균형 있게 맞추며, 보통 랭크와 같거나 그 배수로 설정된다. QLoRA는 여기서 한 단계 더 나아가 기반 모델을 일반적으로 4비트 형식으로 불러오고, 그 위에 FP16 또는 BF16 정밀도의 LoRA 어댑터만 학습함으로써 거대한 기반 모델을 GPU에 보관하는 비용을 크게 낮춘다.

4. 옵티마이저와 활성화 메모리 절감

표준 AdamW는 각 학습 파라미터에 대한 1차·2차 모멘트 추정치를 FP32로 유지하므로 옵티마이저 상태가 상당한 메모리를 차지한다. 글에서 사용하는 8비트 AdamW는 블록 단위 양자화로 이 상태를 8비트 정밀도로 저장하면서 학습 안정성을 유지하며, 표준 FP32 AdamW와 비교해 옵티마이저 메모리를 약 75% 줄일 수 있다고 설명한다. 실행 시 use_8bit_adam 플래그가 지정되면 bitsandbytes의 AdamW8bit를 선택하고, 그렇지 않으면 일반 torch.optim.AdamW를 사용한다. 그래디언트 체크포인팅도 함께 적용해 순전파 중 모든 중간 활성화를 저장하지 않고 일부 체크포인트만 남긴 뒤, 역전파 시 필요한 값을 다시 계산함으로써 추가 연산과 더 낮은 메모리 사용량을 교환한다.

5. VAE 잠재 표현 사전 캐시

잠재 표현 캐시는 학습을 시작하기 전에 모든 훈련 이미지를 VAE 인코더에 통과시켜 잠재 분포를 미리 계산하고 보관하는 최적화다. 이후 각 학습 단계에서는 원본 이미지를 VAE로 반복 인코딩하지 않고 캐시된 잠재 표현을 직접 사용하므로, 중복 계산을 없애 단계별 처리 시간을 단축할 수 있다. 모든 잠재 표현을 만든 뒤에는 VAE 객체를 삭제하고 메모리를 해제할 수 있어 학습 중 GPU에 유지해야 하는 모델 구성 요소도 줄어든다. 대신 캐시 데이터를 저장하기 위한 시스템 RAM 사용량이 증가하지만, 글은 소규모 데이터셋을 대상으로 할 때 일반적으로 감당할 수 있는 수준이라고 설명하며 cache_latents 플래그로 이 절차를 활성화한다.

6. 4비트 NF4 기반 모델과 LoRA 어댑터 구성

기반 트랜스포머는 BitsAndBytesConfig를 통해 4비트 NF4 형식으로 로드하며, 실제 연산 자료형은 실행 옵션에 따라 FP32, FP16 또는 BF16으로 선택한다. FluxTransformer2DModel.from_pretrained 호출에서 transformer 하위 모델만 불러오고 해당 양자화 설정을 전달한 뒤, prepare_model_for_kbit_training으로 저비트 학습에 맞게 준비한다. LoRA 어댑터는 FLUX 어텐션 블록의 to_k, to_q, to_v, to_out.0 모듈을 대상으로 추가하며, 랭크와 알파에는 같은 값을 사용하고 가중치는 가우시안 방식으로 초기화한다. 제시된 출력에서는 전체 파라미터가 11,906,077,760개인 반면 실제로 학습되는 LoRA 파라미터는 4,669,440개이므로, 전체 모델을 갱신하지 않는 방식의 규모 차이를 확인할 수 있다.

7. CLIP·T5 텍스트 임베딩 사전 계산

학습 중 CLIP과 T5 텍스트 인코더를 GPU에 계속 올려 두지 않도록, 캡션에 대한 프롬프트 임베딩과 풀링된 임베딩, 텍스트 ID를 미리 계산해 저장한다. 예시 스크립트는 T5 인코더를 8비트로 불러오고 트랜스포머와 VAE를 제외한 FluxPipeline을 구성한 뒤, 각 프롬프트에 encode_prompt를 적용한다. 데이터셋의 각 이미지는 픽셀 데이터에 대한 SHA-256 해시로 캡션 및 임베딩과 연결되며, 계산 결과는 CPU 배열과 목록으로 변환되어 Parquet 파일에 직렬화된다. 학습 데이터 로더는 이후 이 파일에서 캐시된 결과를 읽기 때문에 같은 캡션을 단계마다 다시 인코딩할 필요가 없고, 텍스트 인코더도 학습 중 GPU 메모리에 상주하지 않아도 된다.

8. 학습 실행 구성과 메모리 절감의 결합

예시에서는 먼저 최대 시퀀스 길이 77로 텍스트 임베딩을 계산해 embeddings_alphonse_mucha.parquet 파일을 만든 다음, accelerate로 DreamBooth 방식의 FLUX LoRA 학습 스크립트를 실행한다. 학습 명령은 black-forest-labs/FLUX.1-dev와 사전 계산한 Parquet 파일을 입력으로 사용하고, 결과를 alphonse_mucha_lora_flux_nf4 디렉터리에 저장하도록 구성한다. 공개된 명령 부분에서는 BF16 혼합 정밀도, 8비트 AdamW, 가중치 방식 none, 512×768 해상도, 배치 크기 1, 반복 횟수 1, 학습률 1e-4를 지정한다. 글의 핵심은 4비트 트랜스포머와 LoRA만으로 끝내지 않고 VAE 잠재 표현 및 텍스트 임베딩 캐시를 함께 적용해, 실제 학습 시 활성 모델을 사실상 양자화된 트랜스포머와 LoRA 어댑터 중심으로 축소하고 전체 미세 조정을 약 10GB 미만의 GPU 메모리 범위에 두는 것이다.

🧾 핵심 주장 / 시사점

메모리 절감은 하나의 기법이 아니라 기반 모델, 옵티마이저 상태, 중간 활성화, VAE 연산, 텍스트 인코더라는 서로 다른 점유 원인을 각각 4비트 양자화, 8비트 AdamW, 체크포인팅, 잠재 표현 캐시, 임베딩 캐시로 제거하는 조합에서 나온다.
FLUX.1-dev 전체를 업데이트하지 않고 어텐션의 일부 투영 모듈에 추가된 약 467만 개의 LoRA 파라미터만 학습하는 것이 소비자용 GPU에서의 미세 조정을 가능하게 하는 핵심 구조다.
VAE 잠재 표현과 CLIP·T5 임베딩을 학습 전에 계산하면 GPU 메모리뿐 아니라 반복 연산도 줄일 수 있지만, 그 대가로 사전 처리 단계와 캐시 파일 및 시스템 RAM 관리가 필요하다.

✅ 액션 아이템

FLUX.1-dev 미세조정 실험에서는 CLIP·T5 텍스트 인코더와 VAE를 동결하고 트랜스포머만 4비트 NF4로 로드해 LoRA 어댑터만 업데이트하도록 구성한다.
8비트 AdamW, 그래디언트 체크포인팅, VAE 잠재 표현 캐시를 결합해 옵티마이저 상태·활성화 메모리와 반복 인코딩 비용을 줄여 단일 GPU에서 10GB 미만을 유지한다.
학습 파라미터를 119억 중 약 467만개 LoRA로 한정하고 BF16 또는 FP16를 정한 뒤 Parquet 텍스트 임베딩 저장과 배치 1 조건으로 동일 조건의 재현 실험을 맞춘다.