Distillation Trainer

개요

Distillation Trainer는 Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist, Olivier Bachem의 On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes에 설명된 on-policy knowledge distillation을 구현합니다.

해당 논문의 초록은 다음과 같습니다:

Knowledge distillation(KD)은 더 작은 student model을 학습시켜 teacher model의 추론 비용과 메모리 사용량을 줄이기 위해 teacher model을 압축하는 데 널리 사용됩니다. 그러나 auto-regressive sequence model을 위한 현재의 KD 방법은 학습 중에 접하는 출력 시퀀스와 student가 추론 중에 생성하는 시퀀스 사이의 분포 불일치 문제를 겪습니다. 이 문제를 해결하기 위해 Generalized Knowledge Distillation(GKD)을 소개합니다. GKD는 고정된 출력 시퀀스 세트에만 의존하는 대신, teacher가 해당 시퀀스에 대해 제공하는 피드백을 활용하여 student가 자체적으로 생성한 출력 시퀀스로 student를 학습시킵니다. supervised KD 접근법과 달리, GKD는 student와 teacher 사이에 대안적인 loss function을 사용할 수 있는 유연성을 제공하며, 이는 teacher의 분포를 모방할 표현력이 부족한 student의 경우 유용할 수 있습니다.

DistillationTrainer는 더 작은 student 모델이 student 모델 자체의 on-policy 생성 결과에 대해 teacher의 next-token 분포를 일치하도록 학습합니다. student의 completions을 on-policy로 생성하고(선택적으로 vLLM 기반), memory-efficient한 chunked Jensen-Shannon divergence loss를 사용하여 teacher의 전체 next-token 분포와 일치시킵니다. 이를 통해 teacher의 dense 분포가 전체적으로 materialize되지 않습니다.

이 trainer는 Carlos Miguel Patiño님이 기여했습니다.

생성을 학습과 분리하고, 로컬 forward pass 대신 HTTP로 제공되는 teacher에 대해 점수를 매기는 비동기 counterpart에 대해서는 AsyncDistillationTrainer를 참조하세요.

빠른 시작

이 예제는 증류(distillation) 방법으로 모델을 학습하는 방법을 보여줍니다. UltraFeedback 프롬프트 데이터셋의 프롬프트를 사용하여 Qwen 2.5 1.5B Instruct teacher에서 Qwen 2.5 0.5B Instruct model을 증류합니다. 데이터셋의 데이터는 여기서 확인할 수 있습니다:

아래는 모델을 학습하는 스크립트입니다.

# train_distillation.py
from datasets import load_dataset
from trl import DistillationTrainer

dataset = load_dataset("trl-lib/ultrafeedback-prompt", split="train")

trainer = DistillationTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    teacher_model="Qwen/Qwen2.5-1.5B-Instruct",
    train_dataset=dataset,
)
trainer.train()

다음 명령을 사용하여 스크립트를 실행합니다:

accelerate launch train_distillation.py

증류 방법 자세히 살펴보기

On-policy knowledge distillation은 고정된 teacher 출력 세트가 아니라, student가 스스로 생성한 completions에서 teacher의 next-token 분포를 재현하도록 student를 학습시킵니다. 자신의 생성 결과로부터 학습함으로써 student는 자신의 실수를 수정할 수 있으며, 이는 일반적으로 off-policy distillation보다 우수한 성능을 보입니다. 이 섹션에서는 completions 생성과 loss 계산이라는 두 가지 핵심 단계를 다루며, 실제 작동 방식을 자세히 설명합니다.

완성 결과 생성

각 학습 단계에서 student는 샘플링된 프롬프트에 대해 배치의 completions를 생성합니다.

손실 계산

손실은 일반화된 Jensen-Shannon 발산(JSD) 학생 분포 $p_S$ 와 교사 분포 $p_T$ 사이의 생성된 completion 토큰에 대해 beta 로 보간되어 정의됩니다: $\mathcal{L}_\beta = \beta \, \mathbb{D}_{\mathrm{KL}}\!\left[ p_T \| p_M \right] + (1 - \beta) \, \mathbb{D}_{\mathrm{KL}}\!\left[ p_S \| p_M \right], \qquad p_M = (1 - \beta) \, p_S + \beta \, p_T,$

여기서 $p_M$ 은 두 분포의 $\beta$ -혼합입니다. 양 끝값은 순수한 발산으로 축소됩니다: beta=0.0 은 순방향 $\mathbb{D}_{\mathrm{KL}}\!\left[ p_T \| p_S \right]$ 이고 beta=1.0 은 역방향 $\mathbb{D}_{\mathrm{KL}}\!\left[ p_S \| p_T \right]$ .

실제로 어휘 logits로의 투사(projection)와 발산(divergence)은 청크(chunk) 단위로 계산되므로, 피크 활성화 메모리는 전체 어휘 × 시퀀스 길이의 logits 텐서에 비례하여 증가하지 않습니다. 메모리 사용량 줄이기를 참조하세요.

기대되는 데이터셋 유형

데이터셋은 conversational prompt-only 데이터셋 형식으로 포맷되어야 합니다. 학생 모델이 온-폴리시로 자체 완성물을 생성하므로 프롬프트만 필요합니다:

{"prompt": [{"role": "user", "content": "What color is the sky?"}]}

기록된 지표

학습 및 평가 중에 다음 메트릭을 기록합니다:

num_tokens: 프롬프트와 완성 결과를 모두 포함하여 지금까지 처리된 토큰의 총 수입니다. 도구를 사용할 경우 도구 토큰은 제외하고 계산합니다.
step_time: 학습 단계당 평균 소요 시간(초 단위, 생성 포함).
completions/mean_length: 생성된 completions의 평균 길이입니다. 도구를 사용할 경우, 도구 토큰을 제외한 토큰만 계산됩니다.
completions/min_length: 생성된 completions의 최소 길이입니다. 도구를 사용할 경우, 도구 토큰이 아닌 토큰만 계산됩니다.
completions/max_length: 생성된 completions의 최대 길이. 도구를 사용할 경우 도구 토큰은 제외하고 카운트합니다.
completions/mean_terminated_length: EOS로 종료되는 생성된 completions의 평균 길이입니다. 도구를 사용할 경우, 도구 토큰은 제외하고 계산합니다.
completions/min_terminated_length: EOS로 종료되는 생성된 completions의 최소 길이입니다. 도구를 사용할 경우, 도구 토큰이 아닌 토큰만 계산됩니다.
completions/max_terminated_length: EOS로 종료되는 생성된 completions의 최대 길이입니다. 도구를 사용할 경우, 도구 토큰이 아닌 토큰만 계산됩니다.
completions/clipped_ratio: 잘림(클립)된 completions의 비율.
tools/call_frequency: 생성 배치당 완료(completion)별 평균 도구 호출 횟수입니다. tools가 제공된 경우에만 기록됩니다.
tools/failure_frequency: 실패한 도구 호출의 비율(도구를 찾을 수 없거나, 예외가 발생하거나, 호출 유형이 지원되지 않는 경우). 도구가 호출되지 않은 경우 0.0입니다. tools가 제공된 경우에만 기록됩니다.
entropy: 생성된 완성본 전체에 걸친 토큰 예측의 평균 엔트로피(nats 단위).

커스터마이징

vLLM 기반 생성으로 학습 속도 향상

on-policy 방식으로 학습할 때 생성은 종종 주요 병목이 됩니다. 생성을 가속화하려면 LLM용 고처리량, 저지연 추론 엔진인 vLLM을 사용할 수 있습니다. 이를 활성화하려면 먼저 패키지를 설치해야 합니다.

pip install trl[vllm]

학습 중 vLLM을 사용하는 두 가지 방식을 지원합니다: colocate 모드와 server 모드입니다.

옵션 1: Colocate 모드

이 모드에서는 vLLM이 트레이너 프로세스 내에서 실행되며 학습 모델과 GPU 메모리를 공유합니다. 이를 통해 별도의 서버를 시작할 필요가 없으며 GPU 활용도를 높일 수 있지만, 학습 GPU에서 메모리 경쟁이 발생할 수 있습니다. 이 모드가 기본 모드입니다.

from trl import DistillationConfig

training_args = DistillationConfig(
    ...,
    use_vllm=True,  # vllm_mode="colocate" by default
)

옵션 2: 서버 모드

이 모드에서는 vLLM이 별도의 프로세스(그리고 별도의 GPU)에서 실행되며 HTTP를 통해 trainer와 통신합니다. 추론에 전용 GPU가 있는 경우 이상적입니다.

vLLM 서버 시작:

VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --max-logprobs -1

학습 스크립트에서 서버 모드를 활성화합니다:

from trl import DistillationConfig

training_args = DistillationConfig(
    ...,
    use_vllm=True,
    vllm_mode="server",
)

서버가 trainer와 다른 GPU를 사용하도록 확인하세요. 그렇지 않으면 NCCL 오류가 발생할 수 있습니다. 사용할 GPU는 CUDA_VISIBLE_DEVICES 환경 변수로 지정할 수 있습니다.

모델 크기와 학습에 필요한 전체 GPU 메모리 요구 사항에 따라, 리소스 과소 사용 또는 메모리 부족 오류를 방지하기 위해 DistillationConfig의 vllm_gpu_memory_utilization 매개변수를 조정해야 할 수 있습니다.

자세한 내용은 vLLM을 활용한 학습 속도 향상을 참조하세요.

PEFT로 어댑터 학습

🤗 PEFT 라이브러리와 긴밀한 통합을 지원하여, 전체 student를 학습하는 대신 어댑터를 학습하고 Hub에서 공유할 수 있습니다.

from datasets import load_dataset
from trl import DistillationTrainer
from peft import LoraConfig

dataset = load_dataset("trl-lib/ultrafeedback-prompt", split="train")

trainer = DistillationTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    teacher_model="Qwen/Qwen2.5-1.5B-Instruct",
    train_dataset=dataset,
    peft_config=LoraConfig(),
)
trainer.train()

디스틸레이션 손실은 lm_head.weight을 직접 읽고 PeftModel.forward()를 거치지 않고 학생 백본을 실행합니다. 따라서 lm_head에 적용된 어댑터(target_modules를 통해)와 프롬프트 학습 방법(PromptTuning, PrefixTuning, P-Tuning)은 묵시적으로 무시되므로 거부됩니다. 헤드를 학습하려면 대신 modules_to_save=["lm_head"]를 사용하세요.

에이전트 학습

DistillationTrainer는 에이전트 학습을 지원합니다: 학생 모델은 생성 과정에서 도구를 호출하고, 그 결과로 생성된 궤적에 대해 증류됩니다. 도구 결과 토큰은 손실 함수에서 마스킹되므로, 학생 모델은 자신이 생성한 토큰에만 대해 학습됩니다.

도구

tools 인수는 에이전트가 사용할 수 있는 도구를 정의하는 Python 함수 목록을 기대합니다:

from trl import DistillationTrainer

trainer = DistillationTrainer(
    tools=[tool1, tool2],
    ...,
)

각 도구는 타입 힌트가 지정된 인자와 반환 타입을 가진 표준 Python 함수여야 하며, 목적, 인자, 반환 값을 설명하는 Google 스타일 docstring을 포함해야 합니다. 자세한 내용은 Passing tools 가이드를 참조하세요.

Tool call 루프는 채팅 템플릿이 prefix-preserving이어야 합니다(도구 메시지를 추가해도 이전 메시지의 렌더링 방식이 변경되어서는 안 됩니다). 알려진 모델 계열(예: Qwen3, DeepSeek-V3)의 경우, TRL은 도구가 활성화되면 패치된 학습 템플릿으로 자동으로 교체합니다. 전체 목록은 Chat Templates을 참조하세요.

DistillationConfig에서 max_tool_calling_iterations를 사용하여 도구 호출 턴의 수를 제한할 수 있습니다. 기본적으로 제한이 없으며, 학생 모델이 도구 호출이 없는 응답 턴을 생성하면 생성이 중단됩니다.

예:

from trl import DistillationTrainer

def multiply(a: int, b: int) -> int:
    """
    Multiplies two integers.

    Args:
        a: The first integer.
        b: The second integer.

    Returns:
        The product of the two integers.
    """
    return a * b

trainer = DistillationTrainer(
    tools=[multiply],
    ...,
)

Async tools은 아직 지원되지 않습니다 — 동기 함수를 전달하세요.

멀티모달 도구 응답

Tools는 텍스트와 함께 이미지를 반환할 수 있으며, 이를 위해 콘텐츠 블록의 리스트를 반환합니다. 이는 도구가 시각적 피드백(예: 스크린샷, 플롯, 카메라 캡처)을 제공하는 VLM 에이전트 학습에 유용합니다.

from PIL import Image

def take_screenshot() -> list:
    """
    Takes a screenshot of the current screen.

    Returns:
        The screenshot image with a description.
    """
    img = Image.open("screenshot.png")
    return [{"type": "image", "image": img}, {"type": "text", "text": "Here is the screenshot."}]

반환된 이미지는 자동으로 대화에 삽입되어 VLM에 전달되며, 이후의 생성 턴에 사용됩니다.

비전 언어 모델 학습

DistillationTrainer는 텍스트와 이미지를 모두 포함하는 멀티모달 데이터셋에서 Vision-Language Models(VLM)의 증류를 지원합니다. 학생 모델과 교사 모델 모두에 VLM을 전달하고, image 컬럼(샘플당 단일 이미지) 또는 images 컬럼(샘플당 이미지 목록)을 가진 프롬프트 전용 데이터셋을 제공하세요. 예상되는 데이터셋 구조에 대한 자세한 내용은 Dataset Format — Vision datasets 섹션을 참조하세요.

테스트 환경:

Gemma 3 — 예: google/gemma-3-4b-it
LLaVA-NeXT — 예: llava-hf/llava-v1.6-mistral-7b-hf
Qwen2-VL — 예: Qwen/Qwen2-VL-2B-Instruct
Qwen2.5-VL — 예: Qwen/Qwen2.5-VL-3B-Instruct

모든 VLM과의 호환성이 보장되지는 않습니다. 특정 모델이 지원되어야 한다고 생각하시면 GitHub에서 이슈를 열어 주세요. 더 좋다면, 필요한 변경 사항이 포함된 pull request를 제출해 주세요.

명령줄 인터페이스

명령줄에서 증류 학습을 시작하려면 trl distillation CLI를 사용하세요. 표준 ModelConfig 플래그를 통해 전체 학습 및 LoRA를 지원합니다.

# Full training:
trl distillation \
    --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
    --teacher_model_name_or_path Qwen/Qwen2.5-1.5B-Instruct \
    --dataset_name trl-lib/ultrafeedback-prompt \
    --learning_rate 2e-5 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 8 \
    --output_dir distilled-model \
    --num_train_epochs 1

# LoRA:
trl distillation \
    --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
    --teacher_model_name_or_path Qwen/Qwen2.5-1.5B-Instruct \
    --dataset_name trl-lib/ultrafeedback-prompt \
    --learning_rate 2e-4 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 8 \
    --output_dir distilled-model \
    --num_train_epochs 1 \
    --use_peft \
    --lora_r 64 \
    --lora_alpha 16

DistillationTrainer

클래스 trl.DistillationTrainer

< source >

( model: str | PreTrainedModel | PeftModelteacher_model: str | transformers.modeling_utils.PreTrainedModel | None = Noneargs: trl.trainer.distillation_config.DistillationConfig | None = Nonetrain_dataset: datasets.arrow_dataset.Dataset | None = Noneeval_dataset: datasets.arrow_dataset.Dataset | dict[str, datasets.arrow_dataset.Dataset] | None = Noneprocessing_class: transformers.tokenization_utils_base.PreTrainedTokenizerBase | transformers.processing_utils.ProcessorMixin | None = Nonecallbacks: list[transformers.trainer_callback.TrainerCallback] | None = Noneoptimizers: tuple[typing.Optional[torch.optim.Optimizer], typing.Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None)quantization_config: BitsAndBytesConfig | None = Nonepeft_config: PeftConfig | None = Nonetools: list[collections.abc.Callable] | None = None )

매개변수

model (str 또는 PreTrainedModel 또는 PeftModel) — 학습할 모델. 다음 중 하나여야 합니다:
- huggingface.co의 모델 저장소에 호스팅된 사전 학습 모델의 model id인 문자열, 또는 save_pretrained를 사용하여 모델 가중치를 저장한 디렉터리의 경로, 예: './my_model_directory/'입니다. 모델은 args.model_init_kwargs의 키워드 인수를 사용하여 <ModelArchitecture>.from_pretrained(여기서 <ModelArchitecture>는 모델 구성에서 파생됨)로 로드됩니다. args.model_init_kwargs에 dtype이 지정되지 않은 경우 기본값은 float32입니다. 이는 from_pretrained와 다르며, 후자의 경우(Transfomers v5부터) dtype은 모델 구성에서 추론됩니다.
- PreTrainedModel 객체입니다. 인과적 언어 모델만 지원됩니다.
- PeftModel 객체입니다. 인과 언어 모델만 지원됩니다.
teacher_model (str 또는 PreTrainedModel, optional) — Student가 다음 토큰 분포를 일치하도록 학습하는 Teacher 모델입니다. model id / 경로 (model과 유사하게 로드되며, args.teacher_model_init_kwargs 사용) 또는 인스턴스화된 PreTrainedModel일 수 있습니다. Student의 어휘를 공유해야 합니다. Teacher를 다른 방식으로 공급하는 서브클래스 (예: 원격 서버)에서는 생략할 수 있습니다.
args (DistillationConfig, 선택 사항) — 이 trainer의 설정입니다. None인 경우 기본 설정이 사용됩니다.
train_dataset (Dataset 또는 IterableDataset) — 학습에 사용할 Dataset입니다. 반드시 "prompt". 컬럼을 포함해야 합니다. Dataset의 추가 컬럼은 무시됩니다. 샘플 형식은 다음과 같습니다:
- Standard: 각 샘플에는 일반 텍스트가 포함됩니다.
- 대화형: 각 샘플에는 구조화된 메시지(예: 역할 및 내용)가 포함됩니다.
train_dataset이 IterableDataset(예: 스트리밍 데이터셋)인 경우, 그 길이를 추론할 수 없으므로 학습 루프를 제한하고 학습률 스케줄러를 구성하는 데 필요한 총 학습 단계 수를 지정하기 위해 학습 인자에 max_steps를 설정해야 합니다.
eval_dataset (Dataset, IterableDataset, DatasetDict, IterableDatasetDict 또는 dict[str, Dataset | IterableDataset]) — 평가에 사용할 Dataset입니다. train_dataset과 동일한 요구 사항을 충족해야 합니다.
processing_class (PreTrainedTokenizerBase, ProcessorMixin, 선택 사항) — 데이터를 처리하는 데 사용되는 처리 클래스입니다. 패딩 쪽은 “left”로 설정해야 합니다. None인 경우, from_pretrained를 사용하여 모델 이름에서 처리 클래스를 로드합니다. 패딩 토큰인 tokenizer.pad_token이 설정되어 있어야 합니다. 처리 클래스에 패딩 토큰이 설정되지 않은 경우, tokenizer.eos_token이 기본값으로 사용됩니다.
callbacks (의 목록 TrainerCallback, 선택 사항) — 학습 루프를 사용자 정의할 콜백 목록입니다. 여기.
기본으로 사용되는 콜백 중 하나를 제거하려면 remove_callback 메서드를 사용하세요.
optimizers (tuple[torch.optim.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None], 선택 사항, 기본값 (None, None)) — 사용할 옵티마이저와 스케줄러를 포함하는 튜플입니다. 기본적으로 모델에 AdamW 인스턴스와 args에 의해 제어되는 get_linear_schedule_with_warmup로 생성된 스케줄러를 사용합니다.
quantization_config (BitsAndBytesConfig, 선택 사항) — 모델 식별자로부터 모델을 로드할 때 사용되는 양자화 구성입니다. QLoRA 학습을 위해 peft_config와 함께 사용하세요. 모델이 이미 인스턴스화되어 있으면 무시됩니다.
peft_config (PeftConfig, 선택 사항) — 모델을 래핑하는 데 사용되는 PEFT 구성입니다. None인 경우 모델은 래핑되지 않습니다.
tools (Callable의 목록, 선택 사항) — 모델이 생성 과정에서 호출할 수 있는 호출 가능한 도구 함수의 목록입니다. 각 도구는 적절히 타입 힌트가 지정된 인자와 반환 값을 가진 표준 Python 함수여야 하며, 목적, 인자, 반환 값을 설명하는 Google 스타일 docstring을 포함해야 합니다. 자세한 내용은 https://huggingface.co/docs/transformers/en/chat_extras#passing-tools를 참조하세요. 모델은 함수의 이름, 타입 힌트, docstring을 기반으로 호출 방법을 결정합니다. 모델의 chat template이 도구 사용을 지원하고 도구 호출을 위해 fine-tuning이 되었는지 확인하세요.

지식 증류용 Trainer입니다. 학생 모델은 온-폴리시(on-policy) 방식으로 학습되며, 스스로 완성물을 생성하여 On-Policy Distillation of Language Models에서 소개된 일반화된 Jensen-Shannon 발산(전방 KL, 역방 KL, JSD를 beta로 보간) 하에서 교사의 다음 토큰 분포를 일치시킵니다.

예:

>>> from trl import DistillationTrainer
>>> from datasets import load_dataset

>>> dataset = load_dataset("trl-lib/tldr", split="train")

>>> trainer = DistillationTrainer(
...     model="Qwen/Qwen2.5-0.5B-Instruct",
...     teacher_model="Qwen/Qwen2.5-1.5B-Instruct",
...     train_dataset=dataset,
... )
>>> trainer.train()

train

< source >

매개변수

resume_from_checkpoint (str 또는 bool, 선택 사항) — str인 경우, 이전 Trainer 인스턴스가 저장한 체크포인트의 로컬 경로입니다. bool이고 True와 같으면, 이전 Trainer 인스턴스가 args.output_dir에 저장한 마지막 체크포인트를 로드합니다. 이 값이 지정되어 있으면, 여기서 로드된 모델/옵티마이저/스케줄러 상태로부터 학습이 재개됩니다.
trial (optuna.Trial 또는 dict[str, Any], 선택 사항) — 하이퍼파라미터 탐색을 위한 트라이얼 실행 또는 하이퍼파라미터 사전입니다.
ignore_keys_for_eval (list[str], 선택 사항) — 모델의 출력(사전에 해당하는 경우)에서 학습 중 평가 예측을 수집할 때 무시해야 할 키의 목록입니다.

반환값

~trainer_utils.TrainOutput

전체 스텝 수, 학습 손실, 메트릭을 포함하는 객체.

주요 학습 진입점.

save_model

< source >

( output_dir: str | None = None_internal_call: bool = False )

모델을 저장하므로 from_pretrained()를 사용하여 다시 로드할 수 있습니다.

main 프로세스에서만 저장합니다.

push_to_hub

< source >

( commit_message: str | None = 'End of training'blocking: bool = Truetoken: str | None = Nonerevision: str | None = None**kwargs )

매개변수

commit_message (str, 선택 사항, 기본값 "End of training") — 푸시할 때 커밋할 메시지입니다.
blocking (bool, 선택 사항, 기본값 True) — git push가 완료된 후에만 함수가 반환되어야 하는지 여부.
token (str, 선택 사항, 기본값 None) — Trainer의 원래 args를 덮어쓸 수 있는 쓰기 권한이 있는 토큰.
revision (str, 선택 사항) — 커밋할 git revision입니다. 기본값은 “main” 브랜치의 head입니다.
kwargs (dict[str, Any], 선택) — ~Trainer.create_model_card에 함께 전달되는 추가 키워드 인자입니다.

self.args.hub_model_id 저장소에 self.model과 self.processing_class를 🤗 모델 허브에 업로드합니다.

DistillationConfig

클래스 trl.DistillationConfig

< source >

( output_dir: str | None = Noneper_device_train_batch_size: int = 8num_train_epochs: float = 3.0max_steps: int = -1learning_rate: float = 1e-06lr_scheduler_type: transformers.trainer_utils.SchedulerType | str = 'linear'lr_scheduler_kwargs: dict | str | None = Nonewarmup_steps: float = 0optim: transformers.training_args.OptimizerNames | str = 'adamw_torch_fused'optim_args: str | None = Noneweight_decay: float = 0.0adam_beta1: float = 0.9adam_beta2: float = 0.999adam_epsilon: float = 1e-08optim_target_modules: None | str | list[str] = Nonegradient_accumulation_steps: int = 1average_tokens_across_devices: bool = Truemax_grad_norm: float = 1.0label_smoothing_factor: float = 0.0bf16: bool | None = Nonefp16: bool = Falsebf16_full_eval: bool = Falsefp16_full_eval: bool = Falsetf32: bool | None = Nonegradient_checkpointing: bool = Truegradient_checkpointing_kwargs: dict[str, typing.Any] | str | None = Nonetorch_compile: bool = Falsetorch_compile_backend: str | None = Nonetorch_compile_mode: str | None = Noneuse_liger_kernel: bool = Falseliger_kernel_config: dict[str, bool] | None = Noneuse_cache: bool = Falseneftune_noise_alpha: float | None = Nonetorch_empty_cache_steps: int | None = Noneauto_find_batch_size: bool = Falselogging_strategy: transformers.trainer_utils.IntervalStrategy | str = 'steps'logging_steps: float = 10logging_first_step: bool = Falselog_on_each_node: bool = Truelogging_nan_inf_filter: bool = Trueinclude_num_input_tokens_seen: str | bool = 'no'log_level: str = 'passive'log_level_replica: str = 'warning'disable_tqdm: bool | None = Nonereport_to: None | str | list[str] = 'none'run_name: str | None = Noneproject: str = 'huggingface'trackio_space_id: str | None = Nonetrackio_bucket_id: str | None = Nonetrackio_static_space_id: typing.Union[str, NoneType, typing.Literal[False]] = Noneeval_strategy: transformers.trainer_utils.IntervalStrategy | str = 'no'eval_steps: float | None = Noneeval_delay: float = 0per_device_eval_batch_size: int = 8prediction_loss_only: bool = Falseeval_on_start: bool = Falseeval_do_concat_batches: bool = Trueeval_use_gather_object: bool = Falseeval_accumulation_steps: int | None = Noneinclude_for_metrics: list[str] = <factory>batch_eval_metrics: bool = Falsesave_only_model: bool = Falsesave_strategy: transformers.trainer_utils.SaveStrategy | str = 'steps'save_steps: float = 500save_on_each_node: bool = Falsesave_total_limit: int | None = Noneenable_jit_checkpoint: bool = Falsepush_to_hub: bool = Falsehub_token: str | None = Nonehub_private_repo: bool | None = Nonehub_model_id: str | None = Nonehub_strategy: transformers.trainer_utils.HubStrategy | str = 'every_save'hub_always_push: bool = Falsehub_revision: str | None = Noneload_best_model_at_end: bool = Falsemetric_for_best_model: str | None = Nonegreater_is_better: bool | None = Noneignore_data_skip: bool = Falserestore_callback_states_from_checkpoint: bool = Falsefull_determinism: bool = Falseseed: int = 42data_seed: int | None = Noneuse_cpu: bool = Falseaccelerator_config: dict | str | None = Noneparallelism_config: accelerate.parallelism_config.ParallelismConfig | None = Nonedataloader_drop_last: bool = Falsedataloader_num_workers: int = 0dataloader_pin_memory: bool = Truedataloader_persistent_workers: bool = Falsedataloader_prefetch_factor: int | None = Nonedataloader_multiprocessing_context: str | None = Nonedataloader_in_order: bool = Trueremove_unused_columns: bool | None = Falselabel_names: list[str] | None = Nonetrain_sampling_strategy: str = 'random'length_column_name: str = 'length'ddp_find_unused_parameters: bool | None = Noneddp_bucket_cap_mb: int | None = Noneddp_broadcast_buffers: bool | None = Noneddp_static_graph: bool | None = Noneddp_backend: str | None = Noneddp_timeout: int = 1800fsdp: str | None = Nonefsdp_config: dict[str, typing.Any] | str | None = Nonedeepspeed: dict | str | None = Nonedebug: str | list[transformers.debug_utils.DebugOption] = ''skip_memory_metrics: bool = Truedo_train: bool = Falsedo_eval: bool = Falsedo_predict: bool = Falseresume_from_checkpoint: str | None = Nonelocal_rank: int = -1model_init_kwargs: dict[str, typing.Any] | str | None = Nonetrust_remote_code: bool = Falseteacher_model_name_or_path: str | None = Noneteacher_model_revision: str | None = Noneteacher_model_init_kwargs: dict[str, typing.Any] | str | None = Nonedisable_dropout: bool = Falsemax_completion_length: int | None = 512ds3_gather_for_generation: bool = Trueshuffle_dataset: bool | None = Truepad_to_multiple_of: int | str | None = 'auto'temperature: float = 1.0top_p: float = 1.0top_k: int = 0min_p: float | None = Nonegeneration_kwargs: dict | str | None = Nonechat_template_kwargs: dict | str | None = Nonerepetition_penalty: float = 1.0cache_implementation: str | None = Noneuse_vllm: bool = Falsevllm_mode: str = 'colocate'vllm_model_impl: str = 'vllm'vllm_enable_sleep_mode: bool = Falsevllm_structured_outputs_regex: str | None = Nonevllm_server_base_url: str | None = Nonevllm_server_host: str = '0.0.0.0'vllm_server_port: int = 8000vllm_server_timeout: float = 240.0vllm_group_port: int = 51216vllm_gpu_memory_utilization: float = 0.3vllm_max_model_length: int | None = Nonevllm_tensor_parallel_size: int = 1beta: float = 1.0max_tool_calling_iterations: int | None = Nonelog_completions: bool = Falsenum_completions_to_print: int | None = Nonelog_unique_prompts: bool = False )

모델과 teacher model을 제어하는 매개변수

model_init_kwargs (str 또는 dict[str, Any], 선택 사항) — model 인수가 문자열로 제공될 때 AutoModelForCausalLM.from_pretrained에 사용되는 키워드 인수입니다. revision 값은 처리 클래스를 로드할 때에도 사용됩니다.
trust_remote_code (bool, optional, defaults to False) — Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to from_pretrained and from_pretrained, for both the student and teacher.
teacher_model_name_or_path (str, 선택 사항) — 교사 모델의 이름 또는 경로입니다. 교사가 로컬에서 로드될 때 사용됩니다.
teacher_model_revision (str, 선택 사항) — 교사 모델의 리비전(예: 브랜치 이름, 태그, 커밋 해시).
teacher_model_init_kwargs (str 또는 dict[str, Any], 선택 사항) — 문자열에서 teacher model을 인스턴스화할 때 AutoModelForCausalLM.from_pretrained에 전달되는 키워드 인자
disable_dropout (bool, 선택 사항, 기본값 False) — 학습 중 학생 모델에서 dropout을 비활성화할지 여부입니다.

데이터 전처리를 제어하는 매개변수

remove_unused_columns (bool, 선택 사항, 기본값 False) — 데이터셋에서 "prompt" 컬럼만 유지할지 여부입니다. 트레이너는 원본 프롬프트 컬럼을 소비하고 온-폴리시로 완성을 생성하므로 기본값은 False입니다.
max_completion_length (int 또는 None, 선택 사항, 기본값 512) — 온-폴리시(on-policy) 생성 과정에서 완료(completion)당 생성할 최대 토큰 수.
ds3_gather_for_generation (bool, 선택 사항, 기본값 True) — 이 설정은 DeepSpeed ZeRO-3에 적용됩니다. 활성화하면 정책 모델 가중치가 생성을 위해 수집되어 생성 속도가 향상됩니다. 그러나 이 옵션을 비활성화하면 단일 GPU의 VRAM 용량을 초과하는 모델을 학습할 수 있지만, 생성 속도가 느려집니다. 이 옵션을 비활성화하면 vLLM 생성과 호환되지 않습니다.
shuffle_dataset (bool, 선택 사항, 기본값 True) — 학습 데이터셋을 섞을지 여부입니다.
pad_to_multiple_of (int 또는 str, 선택 사항, 기본값 "auto") — 설정된 경우, 프롬프트 ID와 완성문 ID는 이 값의 배수로 패딩됩니다. "auto"는 Neuron(AWS Trainium)에서 입력 형식별로 한 번만 컴파일되도록 128의 배수로 패딩하며, 다른 환경에서는 패딩하지 않습니다.

생성을 제어하는 매개변수

temperature (float, 선택 사항, 기본값 1.0) — 생성 시 샘플링 및 증류 손실 계산에 사용되는 온도입니다. 값이 높을수록 더 부드러운 확률 분포가 생성됩니다.
top_p (float, 선택 사항, 기본값 1.0) — 온-폴리시(on-policy) 생성을 위한 Top-p (nucleus) 샘플링 파라미터.
top_k (int, 선택 사항, 기본값 0) — 온-폴리시(on-policy) 생성을 위한 Top-k 샘플링 매개변수입니다. 0은 top-k 필터링을 비활성화합니다.
min_p (float, 선택 사항) — 최대 확률을 가진 토큰의 확률로 스케일링되는 최소 토큰 확률입니다. 0.0과 1.0 사이의 값이어야 합니다. 일반적인 값은 0.01-0.2 범위입니다.
generation_kwargs (dict[str, Any], 선택 사항) — completions 샘플링 시 GenerationConfig(transformers 사용 시) 또는 SamplingParams(vLLM 사용 시)에 전달할 추가 키워드 인자입니다. suppress_tokens, num_beams 등을 설정하여 생성 동작을 더 자세히 커스터마이징하는 데 사용할 수 있습니다. 다른 생성 매개변수(예: min_p, top_p 등)와 충돌하는 키가 포함된 경우, 해당 키가 우선 적용됩니다.
chat_template_kwargs (dict[str, Any], 선택) — 완성 결과를 생성할 때 apply_chat_template 함수에 전달할 추가 키워드 인자입니다.
repetition_penalty (float, 선택 사항, 기본값 1.0) — 프롬프트와 지금까지 생성된 텍스트에 해당 토큰이 등장했는지에 따라 새로운 토큰에 페널티를 부과하는 값입니다. 1.0보다 큰 값은 모델이 새로운 토큰을 사용하도록 장려하고, 1.0보다 작은 값은 모델이 토큰을 반복하도록 장려합니다.
cache_implementation (str, 선택 사항) — use_vllm이 False로 설정된 경우 더 빠른 생성을 위한 캐시 메서드 구현입니다.

vLLM 기반 생성 가속을 제어하는 매개변수

use_vllm (bool, 선택 사항, 기본값 False) — 학생 모델에서 온-폴리시(on-policy) 완성물을 생성할 때 vLLM을 사용할지 여부입니다.
vllm_mode (str, 선택 사항, 기본값 "colocate") — student vLLM 통합을 위한 모드. "server" 또는 "colocate" 중 하나입니다.
vllm_model_impl (str, 선택 사항, 기본값 "vllm") — vLLM의 모델 구현 백엔드입니다. "vllm" 또는 "transformers"를 사용하세요.
vllm_enable_sleep_mode (bool, 선택 사항, 기본값 False) — 옵티마이저 단계에서 학생 가중치를 오프로드하기 위해 vLLM sleep mode를 활성화합니다.
vllm_structured_outputs_regex (str, 선택 사항) — vLLM 구조화 출력을 위한 정규식 패턴입니다.

vLLM 서버를 제어하는 매개변수(`vllm_mode`가 `"server"`인 경우에만 사용)

vllm_server_base_url (str, 선택 사항) — student vLLM 서버의 기본 URL입니다. 이 값이 제공되면 vllm_server_host와 vllm_server_port는 무시됩니다.
vllm_server_host (str, 선택 사항, 기본값 "0.0.0.0") — 학생 vLLM 서버의 호스트입니다.
vllm_server_port (int, 선택 사항, 기본값 8000) — student vLLM 서버의 포트입니다.
vllm_server_timeout (float, 선택 사항, 기본값 240.0) — student vLLM 서버에 연결하는 데 걸리는 시간 제한입니다.
vllm_group_port (int, 선택 사항, 기본값 51216) — vLLM 가중치 업데이트 그룹(NCCL communicator)에 사용할 포트입니다.

colocate된 vLLM 실행을 제어하는 매개변수(`vllm_mode`가 `"colocate"`인 경우에만 사용)

vllm_gpu_memory_utilization (float, 선택 사항, 기본값 0.3) — 동일 위치에 배치된 student vLLM 엔진의 GPU 메모리 사용률입니다.
vllm_max_model_length (int, 선택 사항) — 동일 위치에 배치된 vLLM 엔진의 최대 모델 시퀀스 길이입니다.
vllm_tensor_parallel_size (int, 선택 사항, 기본값 1) — 동일 위치에 배치된 student vLLM 엔진의 Tensor parallel size입니다.

학습을 제어하는 매개변수

beta (float, 선택 사항, 기본값 1.0) — Generalized Jensen-Shannon Divergence 손실의 보간 계수입니다. 0.0이면 손실은 forward KL 발산이 됩니다. 1.0이면 손실은 reverse KL 발산이 됩니다. 0.5이면 표준 JSD가 됩니다. GRPO의 beta(참조 모델에 대한 KL 페널티 계수)와 달리, 여기서는 발산 자체를 선택하는 역할을 하며 참조 모델 KL 페널티가 없습니다.
max_tool_calling_iterations (int, 선택 사항) — 에이전트를 학습할 때의 최대 도구 호출 턴 수입니다. None인 경우 제한이 없으며, 모델이 도구 호출이 없는 응답 턴을 생성하거나 전체 응답 길이가 max_model_length에 도달하면 생성이 중단됩니다.

로깅을 제어하는 매개변수

log_completions (bool, 선택 사항, 기본값 False) — logging_steps 단계마다 (prompt, completion) 쌍의 샘플을 기록할지 여부입니다. rich가 설치되어 있으면 샘플을 출력합니다. wandb 및/또는 trackio 로깅이 활성화되어 있으면 wandb 및/또는 trackio에 기록합니다.
num_completions_to_print (int, 선택 사항) — rich로 출력할 completions의 수입니다. None인 경우 모든 completions이 기록됩니다.
log_unique_prompts (bool, 선택 사항, 기본값 False) — 고유한 프롬프트만 기록할지 여부입니다. True인 경우 고유한 프롬프트만 기록되고, False인 경우 모든 프롬프트가 기록됩니다.

DistillationTrainer의 구성 클래스입니다.

지식 증류에 특화된 매개변수로 TrainingArguments를 확장합니다. 필요한 모든 필드가 여기에 선언되어 있습니다.

HfArgumentParser를 사용하면 이 클래스를 커맨드 라인에서 지정할 수 있는 argparse 인수로 변환할 수 있습니다.

이러한 파라미터는 TrainingArguments와 다른 기본값을 가집니다:

logging_steps: 기본값이 500 대신 10입니다.

gradient_checkpointing: 기본값이 False가 아닌 True입니다.

bf16: fp16가 설정되지 않은 경우 기본값이 False가 아닌 True입니다.

learning_rate: 기본값이 5e-5 대신 1e-6입니다.

Update on GitHub