TRL

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.17.1).

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

トレーニングのカスタマイズ

TRLはモジュール性を重視して設計されており、ユーザーがニーズに合わせてトレーニングループを効率的にカスタマイズできるようにしています。以下に、さまざまな手法を適用してテストする方法の例を示します。

これらの例ではDPOTrainerを使用していますが、これらのカスタマイズ手法はTRLのほとんどの（あるいはすべての）トレーナーに適用できます。

異なるオプティマイザーとスケジューラーを使用する

デフォルトでは、DPOTrainerはtorch.optim.AdamWオプティマイザーを作成します。別のオプティマイザーを作成して定義し、DPOTrainerに渡すには、次のようにします。

from datasets import load_dataset
from torch import optim
from transformers import AutoModelForCausalLM
from trl import DPOTrainer

dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
optimizer = optim.SGD(model.parameters(), lr=1e-6)

trainer = DPOTrainer(
    model=model,
    train_dataset=dataset,
    optimizers=(optimizer, None),
)
trainer.train()

学習率スケジューラーを追加します

との両方を渡すことで、学習率スケジューラーを追加することもできます。

from torch import optim

optimizer = optim.AdamW(model.parameters(), lr=1e-6)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

trainer = DPOTrainer(..., optimizers=(optimizer, lr_scheduler))

8ビット参照モデルを渡す

trlはtransformersからfrom_pretrainedを使用してモデルを読み込む際にすべてのキーワード引数をサポートしているため、よりメモリ効率の良いファインチューニングのためにtransformersのload_in_8bitも活用できます。

transformersの8bitまたは4bitでの読み込みで、8ビットモデルの読み込みについて詳しく説明しています。

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(load_in_8bit=True)
ref_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", quantization_config=quantization_config)

trainer = DPOTrainer(..., ref_model=ref_model)

カスタムコールバックの追加

ログ記録、監視、または早期停止用のコールバックを追加することで、トレーニングループをカスタマイズできます。コールバックを使用すると、トレーニング中の特定のタイミングでカスタムコードを実行できます。

from transformers import TrainerCallback


class CustomLoggingCallback(TrainerCallback):
    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is not None:
            print(f"Step {state.global_step}: {logs}")


trainer = DPOTrainer(..., callbacks=[CustomLoggingCallback()])

カスタム評価メトリクスの追加

トレーニング中に追跡するカスタム評価メトリクスを定義できます。これは、特定のタスクにおけるモデルの性能を監視するために役立ちます。

def compute_metrics(eval_preds):
    logits, labels = eval_preds
    # Add your metric computation here
    return {"custom_metric": 0.0}


training_args = DPOConfig(..., eval_strategy="steps", eval_steps=100)

trainer = DPOTrainer(..., eval_dataset=eval_dataset, compute_metrics=compute_metrics)

混合精度トレーニングを使用する

混合精度トレーニングにより、トレーニングの高速化とメモリ使用量の削減が大幅に実現できます。トレーニング設定でbf16=Trueまたはfp16=Trueを設定することで、これを有効にできます。

# Use bfloat16 precision (recommended for modern GPUs)
training_args = DPOConfig(..., bf16=True)

注: Ampere GPU (A100、RTX 30xx) 以降の GPU には bf16=True を、それ以前の GPU には fp16=True を使用してください。

勾配の蓄積を使用する

GPUメモリが限られている環境での学習において、勾配の蓄積（gradient accumulation）を使用すると、重みの更新を行う前に複数のステップにわたって勾配を蓄積することで、より大きなバッチサイズをシミュレートできます。

# Simulate a batch size of 32 with per_device_train_batch_size=4 and gradient_accumulation_steps=8
training_args = DPOConfig(
    ...,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
)

学習目的を変更する

トレーナーのサブクラス化は、トレーニングをカスタマイズする最もシンプルな方法です。損失はcompute_lossメソッドで定義されるため、別の目的関数でトレーニングするには、トレーナーをサブクラス化してこのメソッドをオーバーライドします。データの準備、生成、ログ出力、チェックポイントの保存は継承されるため、サブクラスには実際に変更される部分のみを保持します。

import torch

from trl import DPOTrainer
from trl.trainer.utils import selective_log_softmax


class MyDPOTrainer(DPOTrainer):
    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        # In this example: the hinge loss from https://huggingface.co/papers/2309.06657
        input_ids, attention_mask = inputs["input_ids"], inputs["attention_mask"]
        shift_labels, shift_completion_mask = input_ids[:, 1:], inputs["completion_mask"][:, 1:]
        logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
        with torch.no_grad():
            ref_logits = self.ref_model(input_ids=input_ids, attention_mask=attention_mask).logits

        # Sum the log-probs over the completion tokens, for the policy and for the reference model
        logps = (selective_log_softmax(logits[:, :-1], shift_labels) * shift_completion_mask).sum(dim=1)
        ref_logps = (selective_log_softmax(ref_logits[:, :-1], shift_labels) * shift_completion_mask).sum(dim=1)

        chosen_logps, rejected_logps = logps.chunk(2, dim=0)  # batch is [chosen, rejected]
        ref_chosen_logps, ref_rejected_logps = ref_logps.chunk(2, dim=0)
        delta = (chosen_logps - ref_chosen_logps) - (rejected_logps - ref_rejected_logps)
        return torch.relu(1 - self.beta * delta).mean()

設定にパラメータを追加する

損失に必要なパラメータを宣言するために、configをサブクラス化します:

from dataclasses import dataclass, field

from trl import DPOConfig


@dataclass
class MyDPOConfig(DPOConfig):
    my_coef: float = field(default=0.1, metadata={"help": "Coefficient of the custom term."})

バッチ形式を変更する

損失がデフォルトのcollatorでは生成されない入力が必要な場合は、独自のcollatorを渡し、_prepare_datasetをオーバーライドして、データセットがすでに期待される形式である場合はそのままにしてください。

class MyDPOTrainer(DPOTrainer):
    def _prepare_dataset(self, dataset, *args, **kwargs):
        return dataset


trainer = MyDPOTrainer(..., data_collator=my_collator)

完全な例

ブロック拡散SFTの例であるexamples/sft_diffusion_gemma/sft_diffusion_gemma.pyは、これら3つを組み合わせます。DiffusionGemmaSFTConfigはキャンバスと破損のパラメータでSFTConfigを拡張し、DiffusionGemmaSFTTrainerはSFTTrainerを拡張して、自己回帰クロスエントロピーをブロック拡散のノイズ除去目的関数に置き換えます。どちらもライブラリの変更を必要としません。

Antidoomは、Liquid AIによるオープンソースツールで、DPOTrainerをFinal Token Preference Optimization（Antislop）で拡張します。これは単一のトークン位置に対する選好損失であり、推論モデルにおける反復ループを軽減します。

Update on GitHub

←Jobs を使用したトレーニングメモリ使用量の削減→