dFactory: 쉽고 효율적인 Diffusion LLM의 파인튜닝을 위한 오픈소스 프레임워크

읽을거리&정보공유
, , , , , ,
1

dFactory: 쉽고 효율적인 Diffusion LLM의 파인튜닝을 위한 오픈소스 프레임워크
dFactory: 쉽고 효율적인 Diffusion LLM의 파인튜닝을 위한 오픈소스 프레임워크1098×376 34.3 KB

dFactory 소개

dFactory는 최근 새롭게 떠오르고 있는 확산 대형 언어 모델(Diffusion Large Language Model, dLLM) 의 미세 조정(Fine-Tuning)을 쉽고 효율적으로 수행할 수 있도록 돕는 오픈소스 프레임워크입니다. 앤트그룹(Ant Group)의 InclusionAI 팀이 주도하여 개발했으며, 자사의 강력한 100B 매개변수 규모 모델인 LLaDA 2.0 시리즈와 함께 dLLM 생태계 확장을 위해 공개되었습니다.

기존의 대형 언어 모델(LLM) 생태계는 주로 왼쪽에서 오른쪽으로 한 단어씩 순차적으로 예측하는 자기회귀(Autoregressive) 방식에 맞춰져 있었습니다. 그러나 최근 확산 모델(Diffusion Model)이 텍스트 생성 영역에서도 강력한 성능과 압도적인 병렬 디코딩의 이점을 증명하면서, 이를 커스텀 데이터로 훈련하려는 개발자들의 수요가 급증하고 있습니다. 텍스트 토큰은 이미지 픽셀과 달리 딱 떨어지는 이산적(Discrete) 성질을 가지기 때문에, dFactory는 텍스트에 연속적인 노이즈를 주는 대신 토큰을 무작위로 가리고(Masking) 복원하는 특수한 '이산 확산(Discrete Diffusion)' 과정을 최적화하여 학습할 수 있는 환경을 제공합니다.

특히 dFactory 프레임워크는 대규모 분산 학습 프레임워크인 VeOmni를 기반으로 구축되어, FSDP2(Fully Sharded Data Parallel)와 같은 복잡한 병렬화 기법을 매끄럽게 지원합니다. 덕분에 개발자와 연구자들은 소비자용 혹은 기업용 GPU 클러스터 환경에서 LLaDA2.0-mini(16B)나 LLaDA2.0-flash(100B)와 같은 거대한 MoE(Mixture-of-Experts) 기반의 확산 모델을 상대적으로 적은 리소스로 빠르고 안정적으로 학습시킬 수 있습니다.

dFactory vs. 기존 SFT 프레임워크와 비교

기존의 Hugging Face transformers나 TRL 같은 프레임워크들은 일반적인 LLM, 즉, 자기회귀 언어 모델(Autoregressive LLM)의 SFT(Supervised Fine-Tuning)에 고도로 최적화되어 있습니다. 하지만 확산 언어 모델을 학습하기 위해서는 텍스트 블록 전체를 한 번에 처리하는 'Block Diffusion' 방식이나 랜덤 마스킹 기반의 손실 함수 등 완전히 다른 학습 레시피가 필요합니다.

dFactory는 dLLM을 위한 전용 학습 알고리즘을 내장하고 있다는 점에서 기존 도구들과 차별화됩니다. 또한 가장 큰 아키텍처적 차이는 MoE(Mixture-of-Experts) 가중치 병합 및 분리(Merging and Splitting) 시스템에 있습니다. 일반적인 Hugging Face 모델은 각 전문가(Expert)의 가중치를 독립적인 텐서로 저장하지만, dFactory는 GPU 상에서 고도의 배치 행렬 곱셈(Batched Matrix Multiplication)을 활용해 연산 속도를 극대화할 수 있도록 학습 전 가중치를 하나의 텐서로 병합하고, 추론 단계에서 이를 다시 분리하는 독자적인 유틸리티를 제공합니다.

dFactory의 지원 모델 및 핵심 알고리즘

dFactory는 현재 InclusionAI가 공개한 최신 확산 언어 모델(dLLM)인 LLaDA2.0을 지원하며, 확산 텍스트 생성에 필요한 다양한 알고리즘을 내장하고 있습니다.

  • 지원 모델 라인업:

  • Discrete Diffusion Model (랜덤 마스킹): 정방향 프로세스(Forward Process)에서 텍스트 토큰을 무작위로 마스킹하고, 역방향 프로세스에서 이를 복원하는 사전 학습 및 SFT 오브젝티브를 제공합니다.

  • Block Diffusion: SFT 시 연속된 텍스트 블록 전체를 한 번에 처리하는 기법으로, 자기회귀 모델의 학습 방식을 확산 모델에 맞게 차용하여 학습 안정성과 효율성을 높입니다.

  • Trainable Parallel Decoding & Path Distillation: 여러 토큰을 병렬로 생성하여 추론 속도를 혁신적으로 높이는 궤적 압축(Trajectory Compression) 및 학습 가능한 병렬 디코딩 모듈을 지원합니다.

dFactory 설치 및 사용

dFactory 설치 가이드

dFactory는 서브모듈로 포함된 ByteDance의 VeOmni 분산 학습 엔진을 사용하므로, 환경을 구성할 때 종속성을 올바르게 설치해야 합니다. 이 때, 파이썬 패키지 관리자인 uv를 사용하는 것을 권장합니다.

# 1. uv 설치 (없는 경우)
curl -LsSf https://astral.sh/uv/install.sh | sh

# 2. dFactory 저장소 및 서브모듈 클론
git clone https://github.com/inclusionAI/dFactory.git --recursive
cd dFactory/VeOmni

# 3. 종속성 설치 및 가상환경 활성화
uv sync --extra gpu
source .venv/bin/activate
cd ..

또는, pip를 사용할 경우 pip install -e VeOmni/ 명령어로 설치가 가능합니다.

MoE 모델 가중치 병합 (Expert Merging)

최적의 훈련 성능을 위해 분리된 전문가(Separate-Expert) 형식의 모델 가중치를 하나로 합치는 과정이 필수적입니다. 예를 들어 8개의 전문가가 [4096, 14336] 형태의 개별 텐서로 존재했다면, 병합 후에는 [8, 4096, 14336] 형태의 단일 텐서 구조로 변경되어 연산 효율이 극대화됩니다.

# 1. Hugging Face에서 모델 다운로드
python ./scripts/download_hf_model.py \
  --repo_id inclusionAI/LLaDA2.0-mini-preview \
  --local_dir /path/to/separate_expert_model

# 2. 가중치 병합 유틸리티 실행
python scripts/moe_convertor.py \
  --input-path /path/to/separate_expert_model \
  --output-path /path/to/save/merged_model \
  --mode merge

학습 데이터 준비 및 학습 실행

학습 데이터는 대화형(Conversational) 메시지 포맷의 JSONL 파일로 구성해야 합니다. 저장소 내에 openai/gsm8k 데이터셋을 변환하는 예제 스크립트(build_gsm8k_dataset.py)가 포함되어 있으므로 이를 참고하여 커스텀 데이터를 준비할 수 있습니다.

데이터가 준비되면 configs/sft/llada2_mini_bd_sft.yaml 파일을 열어 model_path, train_path, output_dir 경로를 내 환경에 맞게 수정한 후, 아래의 명령어로 SFT를 시작합니다.

PYTHONPATH=$(pwd)/VeOmni:$PYTHONPATH sh train.sh tasks/train_llada2_bd.py configs/sft/llada2_mini_bd_sft.yaml

학습 후 가중치 분리 및 SGLang 추론

학습이 완료된 모델을 일반적인 프레임워크(Hugging Face 또는 SGLang)에서 추론하려면, 병합했던 가중치를 다시 분리(Split)하는 작업을 거쳐야 합니다.

# 가중치 분리 (Split) - 학습 아웃풋 디렉토리 내 hf_ckpt 폴더 지정
python scripts/moe_convertor.py \
  --input-path /path/to/merged_model/checkpoints/global_step_XXX/hf_ckpt/ \
  --output-path /path/to/save/separate_expert_model \
  --mode split

# 원본 베이스 모델의 아키텍처 모델링 파일 수동 복사
cp /path/to/original_base_model/modeling_llada2_moe.py /path/to/save/separate_expert_model/

이후, dFactory 공식 문서에 안내된 SGLang 추론 가이드에 따라 로컬 서버를 구동하면, 확산 언어 모델 특유의 병렬 디코딩 이점을 살려 기존 대비 매우 빠른 속도로 텍스트를 생성할 수 있습니다.

라이선스

dFactory 프로젝트는 Apache License 2.0으로 공개 및 배포되고 있습니다.

:books: dFactory 설치 및 사용 관련 문서 사이트

:github: dFactory 프로젝트 GitHub 저장소



이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있습니다. 관심있는 내용이시라면 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다. :hugs:

:pytorch:파이토치 한국 사용자 모임:south_korea:이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일:love_letter:로 보내드립니다! (기본은 Weekly지만 Daily로 변경도 가능합니다.)

:wrapped_gift: 아래:down_right_arrow:쪽에 좋아요:+1:를 눌러주시면 새로운 소식들을 정리하고 공유하는데 힘이 됩니다~ :star_struck:

1개의 좋아요