dLLM: 확산 언어 모델의 학습 및 평가를 위한 통합 오픈소스 라이브러리

섹션별 상세

dLLM은 확산 언어 모델의 전체 개발 수명 주기를 지원하는 통합 라이브러리이다. Transformers Trainer를 기반으로 설계되어 확장성이 뛰어나며, LoRA를 통한 효율적인 파인튜닝과 DeepSpeed, FSDP를 활용한 대규모 분산 학습 환경을 기본적으로 제공한다.

python

import transformers
import dllm

model = dllm.utils.get_model(model_args=model_args)
tokenizer = dllm.utils.get_tokenizer(model_args=model_args)

trainer = dllm.core.trainers.MDLMTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    args=training_args,
    data_collator=transformers.DataCollatorForSeq2Seq(
        tokenizer, return_tensors="pt", padding=True,
        label_pad_token_id=tokenizer.pad_token_id,
    ),
)
trainer.train()

dLLM 라이브러리를 사용하여 Masked Diffusion Language Model(MDLM)을 학습시키는 기본 코드 구조

LLaDA와 Dream 같은 최신 확산 언어 모델을 위한 최소 학습 및 추론 레시피를 포함한다. Masked Diffusion(MDLM), Block Diffusion(BD3LM), Edit Flows 등 다양한 확산 알고리즘의 구현체를 제공하여 연구자가 목적에 맞는 기법을 선택하여 적용할 수 있다.

Fast-dLLM 기능을 통해 확산 모델의 고질적인 문제인 추론 속도를 개선했다. 캐싱 메커니즘과 신뢰도 임계값 기반 디코딩(Confidence-threshold decoding) 기술을 도입하여 LLaDA와 Dream 모델의 추론 및 평가 속도를 획기적으로 가속화한다.

python

import dllm

sampler = dllm.core.samplers.MDLMSampler(model=model, tokenizer=tokenizer)
messages = [
    [{"role": "user", "content": "Lily runs 12 km/h for 4 hours. How far in 8 hours?"}],
]
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True)
outputs = sampler.sample(inputs, return_dict=True)

학습된 확산 언어 모델을 사용하여 추론(Sampling)을 수행하는 예시

dLLM을 통해 구현된 확산 언어 모델의 대화형 인터페이스 시연 화면 — Screenshot확산 모델이 텍스트를 생성하는 과정을 시각적으로 보여주며, 멀티턴 대화가 가능함을 입증한다. 특히 LLaDA-8B-Instruct 모델의 실제 작동 모습을 통해 라이브러리의 실용성을 강조한다.

기존 자기회귀 모델을 확산 모델로 변환하는 Tiny-A2D 레시피를 제공한다. Qwen, Llama, GPT-2와 같은 일반적인 AR 모델을 Masked Diffusion 방식으로 전환하여 0.5B~0.6B 규모의 소형 확산 모델로 재탄생시키는 전체 과정을 공개했다.

BERT와 같은 인코더 모델을 확산 기반 챗봇으로 변환하는 BERT-Chat 기능을 지원한다. RoBERTa, ModernBERT 등 다양한 인코더 모델을 경량 챗봇으로 파인튜닝할 수 있는 레시피를 통해 확산 모델의 응용 범위를 넓혔다.

용어 해설

확산 언어 모델(Diffusion Language Model): — 텍스트 생성을 확산 프로세스로 모델링하는 기법으로, 노이즈가 섞인 상태에서 점진적으로 원래 데이터를 복원하며 문장을 생성한다. 기존 자기회귀(Autoregressive) 모델과 달리 병렬 생성이나 비순차적 편집이 가능하다는 장점이 있다.
자기회귀 모델(Autoregressive Model): — 이전 토큰들을 기반으로 다음 토큰을 순차적으로 예측하는 방식의 언어 모델이다. GPT 시리즈가 대표적이며, 문맥 이해 능력이 뛰어나지만 생성 속도가 시퀀스 길이에 비례하여 느려지는 특성이 있다.
완전 샤딩 데이터 병렬 처리(FSDP): — 모델 파라미터, 그래디언트, 옵티마이저 상태를 여러 GPU에 분산하여 저장함으로써 대규모 모델을 효율적으로 학습시키는 기법이다. 메모리 사용량을 획기적으로 줄여 단일 GPU 메모리를 초과하는 모델 학습을 가능하게 한다.
언어 모델 평가 프레임워크(LM Evaluation Harness): — 다양한 벤치마크 데이터셋을 사용하여 언어 모델의 성능을 표준화된 방식으로 측정하는 도구이다. MMLU, GSM8K 등 수백 개의 평가 태스크를 지원하여 모델 간의 객관적인 비교를 돕는다.

기술

PyTorch
Transformers
DeepSpeed
FSDP
LoRA
Accelerate

활용 사례

확산 언어 모델 학습 및 파인튜닝
자기회귀 모델의 확산 모델 변환
경량 BERT 기반 챗봇 구축
DLM 성능 벤치마킹

언급된 리소스

GitHubdLLM GitHub Repository

논문dLLM: Simple Diffusion Language Modeling Paper