Fix: resolve Qwen3 text encoder loading issues for FP8 and GGUF formats

개요

Qwen3 텍스트 인코더를 FP8 Safetensors 및 GGUF 형식의 양자화된 버전과 Z-Image 모델을 함께 로드할 때 발생하는 두 가지 별도의 크래싱 문제를 해결합니다.

주요 내용

  • FP8 Safetensors 형식의 lm_head.weight 누락 문제:
  • 원인: 많은 양자화된 FP8 체크포인트는 디스크 공간 절약을 위해 lm_head.weight 키를 의도적으로 제외하며, 이는 model.embed_tokens.weight와 연결되어 있습니다. fast_load_transformers_model 호출 시 mmgp/offload.py에서 엄격한 키 검사가 발생하여 크래시를 유발합니다.
  • 해결책: unified_preprocessor를 구현하여 로드 시 상태 사전을 가로챕니다. 엄격한 검사 전에 model.embed_tokens.weight 메모리 참조를 lm_head.weight 슬롯으로 동적으로 별칭 지정하여 오류를 우회하고 VRAM 낭비를 방지합니다.
  • GGUF 형식의 SDPA Dtype 불일치 문제:
  • 원인: GGUF 텍스트 인코더를 사용할 때, 백엔드 통합은 Value 텐서를 bfloat16으로 제공하는 반면, Hugging Face Transformers는 Query와 Key의 RoPE를 float32로 계산합니다. PyTorch의 네이티브 SDPA는 세 가지 모두 정확히 동일한 데이터 타입을 요구하여 즉시 크래시가 발생합니다.
  • 해결책: text_encoder.to(dtype) 호출을 통해 텍스트 인코더를 로드 직후 균일한 정밀도 상태로 강제 변환합니다. 이는 모든 어텐션 프로젝션을 정렬하여 핵심 어텐션 백엔드를 수정하지 않고도 HF 네이티브 SDPA가 원활하게 작동하도록 합니다.
  • models/z_image/z_image_main.py 변경 사항:
  • 즉석에서 연결된 임베딩 가중치 누락을 처리하기 위한 상태 사전 전처리기 주입.
  • offload.fast_load_transformers_model 완료 직후 .to(dtype) 캐스트 추가하여 엄격한 Dtype 균일성 강제 적용.

시사점

이 수정 사항은 Qwen3 텍스트 인코더의 다양한 양자화 형식 로딩 시 발생하는 치명적인 오류를 해결하여, Z-Image 모델과 함께 이러한 모델을 더 안정적이고 폭넓게 사용할 수 있도록 합니다.

원문 읽기 →
원문을 불러오는 중...

댓글

GitHub Discussions