CCCMKホールディングス TECH LABの Tech Blog

TECH LABのエンジニアが技術情報を発信しています

ブログタイトル

databricksでgpt-oss-20bのFine-Tuningをする。

こんにちは、CCCMKホールディングスAIエンジニアの三浦です。

OpenAIがリリースしたオープンウェイトの言語モデルgpt-ossをカスタマイズ、つまりFine-Tuningするにはどうするんだろう?と最近調べていました。色々情報が見つかり、トライしてみた中で一番わかりやすかったのがOpenAIのCookbookで紹介されている方法でした。

cookbook.openai.com

gpt-oss-20bはリーズニングモデルで回答の前にリーズニング(論理的な思考・推論の過程)が実行されます。gpt-oss-20bはデフォルトの状態だと入力が英語以外の言語であってもリーズニングは英語で実行されます。先ほどのCookbookの中ではFine-Tuningによってgpt-oss-20bが任意の言語でリーズニング出来るようにする手順が紹介されています。シンプルな内容ですが、この内容を応用することで様々な形でgpt-oss-20bをカスタマイズすることが出来るといえます。

今回はOpenAIのCookbookの手順をAzure databricksで実行し、gpt-oss-20bのFine-Tuningを行ってみました。一部databricks特有の設定が必要な箇所もありましたので、databricksでgpt-oss-20bのFine-Tuningをする際の参考になれば幸いです。

使用したデータセット

Cookbookで使われているものと同じデータセットを使用しました。Hugging Faceで公開されているこちらのデータセットです。

huggingface.co

このデータセットは"reasoning_language", "developer", "user", "analysis", "final"というカラムで構成されています。これはOpenAIの"Harmony Response Format"というメッセージフォーマットに従っていて、"developer"にシステムプロンプト、"user"にユーザーの入力、"analysis"にリーズニング、"final"に応答が格納されています。"reasoning_language"には"analysis", つまりリーズニングの言語種が格納されていて、"Spanish", "French", "Italian","German"のいずれかが格納されます。

このデータを読み込み、Harmony Response Formatのテンプレートを適用すると以下のような文字列に変換されます。この文字列がgpt-oss-20bに入力されるデータになります。

<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI.
Knowledge cutoff: 2024-06
Current date: 2025-08-26

Reasoning: medium

# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>developer<|message|># Instructions

reasoning language: French

You are an AI chatbot with a lively and energetic personality.

<|end|><|start|>user<|message|>Can you show me the latest trends on Twitter right now?<|end|><|start|>assistant<|channel|>analysis<|message|>D'accord, l'utilisateur demande les tendances Twitter les plus récentes. Tout d'abord, ...(中略)... techniques.<|end|><|start|>assistant<|channel|>final<|message|>Hey there!  While I can't check Twitter (X) in real-time or access live data, ...(中略)... brainstorm *what* might trend next? I’ve got ideas!<|return|>

databricksの環境

今回の処理を実行したdatabricksの環境は以下の通りです。

実行したdatabricksのCompute

Standard_NC24ads_A100_v4はA100GPUが1台搭載されています。

実装内容

ライブラリのインストール

ライブラリのインストールです。

%pip install torch --index-url https://download.pytorch.org/whl/cu128
%pip install "trl>=0.20.0" "peft>=0.17.0" "transformers>=4.55.0" "mlflow>=3.2"
dbutils.library.restartPython()

データセットのダウンロード

データセットのダウンロードです。

from datasets import load_dataset

dataset = load_dataset("HuggingFaceH4/Multilingual-Thinking", split="train")

モデルのダウンロード

モデルやトークナイザをロードします。model_kwargsの内容はOpenAIのCookbookに従っています。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, Mxfp4Config

# トークナイザのロード
tokenizer = AutoTokenizer.from_pretrained("openai/gpt-oss-20b")

# モデルのロード設定
quantization_config = Mxfp4Config(dequantize=True)
model_kwargs = dict(
    attn_implementation="eager",
    torch_dtype=torch.bfloat16,
    quantization_config=quantization_config,
    use_cache=False,
    device_map="auto",
)

# モデルのロード
model = AutoModelForCausalLM.from_pretrained(
    "openai/gpt-oss-20b", 
    **model_kwargs
)

LoRAの設定

今回のFine-Tuningはパラメータのフル学習ではなく、軽量なアダプタをいくつかの層に接続し、そのアダプタを学習対象にする"LoRA"というテクニックを使用します。

LoRAの設定を行い、最後に学習対象のパラメータのボリュームを表示します。

from peft import LoraConfig, get_peft_model

peft_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules="all-linear",
    target_parameters=[
        "7.mlp.experts.gate_up_proj",
        "7.mlp.experts.down_proj",
        "15.mlp.experts.gate_up_proj",
        "15.mlp.experts.down_proj",
        "23.mlp.experts.gate_up_proj",
        "23.mlp.experts.down_proj",
    ],
)
peft_model = get_peft_model(model, peft_config)
peft_model.print_trainable_parameters()

学習対象のパラメータ数は1,500万程度でパラメータ全体のおよそ0.07%程度のボリュームです。

trainable params: 15,040,512 || all params: 20,929,797,696 || trainable%: 0.0719

ハイパーパラメータ

学習時のハイパーパラメータは以下のようになります。

gradient_checkpointing=Trueは必須で、これがないとOut-Of-Memoryエラーが発生してしまいます。また、ddp_find_unused_parameters=FalseはOpenAIのCookbookにはない設定ですが、databricksで実行する場合、これを設定しないとエラーが発生してしまいました。

from trl import SFTConfig

training_args = SFTConfig(
    learning_rate=2e-4,
    gradient_checkpointing=True,
    logging_steps=1,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    max_length=2048,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine_with_min_lr",
    lr_scheduler_kwargs={"min_lr_rate": 0.1},
    output_dir="gpt-oss-20b-multilingual-reasoner",
    ddp_find_unused_parameters=False,
)

学習の実行

以下で学習を開始します。

from trl import SFTTrainer


trainer = SFTTrainer(
    model=peft_model,
    args=training_args,
    train_dataset=dataset,
    processing_class=tokenizer,
)
trainer.train()
trainer.save_model(training_args.output_dir)

学習曲線は以下のようになりました。

約1,000件のデータセットを3epoch実行で約1.6時間で完了しました。

推論

推論を行う手順は次の通りです。

まずベースとなるモデルをダウンロードします。ロードする時の設定値のuse_cacheは学習時はFalseに設定していましたが推論時はTrueにすることで高速に推論を実行することが出来るようになります。

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# トークナイザ
tokenizer = AutoTokenizer.from_pretrained("openai/gpt-oss-20b")

# モデル
model_kwargs = dict(
    attn_implementation="eager", 
    torch_dtype="auto", 
    use_cache=True, 
    device_map="auto"
)
base_model = AutoModelForCausalLM.from_pretrained(
    "openai/gpt-oss-20b", 
    **model_kwargs
).cuda()

次にベースのモデルに学習したアダプタを接続します。

model = PeftModel.from_pretrained(
    base_model, 
    "./gpt-oss-20b-multilingual-reasoner"
)
model = model.merge_and_unload()

次のようなコードで推論を実行することが出来ます。

REASONING_LANGUAGE = "Japanese"
#システムプロンプトでリーズニングの言語指定
SYSTEM_PROMPT = f"reasoning language: {REASONING_LANGUAGE}"
USER_PROMPT = "カレーの美味しい作り方を教えて。"
messages = [
    {"role": "system", "content": SYSTEM_PROMPT},
    {"role": "user", "content": USER_PROMPT},
]

# テンプレート適用
input_ids = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    return_tensors="pt",
).to(model.device)

# 推論実行
gen_kwargs = {"max_new_tokens": 5120, "do_sample": True, "temperature": 0.6, "top_p": None, "top_k": None}
output_ids = model.generate(input_ids, **gen_kwargs)

# 文字列にデコード
response = tokenizer.batch_decode(output_ids)[0]
print(response)

結果は以下のようになりました。確かにリーズニング(<|channel|>analysis<|message|>以降)は日本語で生成出来るようになりました。ただ、回答(<|channel|>final<|message|>以降)が英語で生成されるようになってしまいました。

<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI.
Knowledge cutoff: 2024-06
Current date: 2025-08-27

Reasoning: medium

# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>developer<|message|># Instructions

reasoning language: Japanese

<|end|><|start|>user<|message|>カレーの美味しい作り方を教えて。<|end|><|start|>assistant<|channel|>analysis<|message|>わかりました、ユーザーにカレーのレシピを教えてほしいということですね。まず、...(中略)...これらを踏まえて、レシピを作成します。<|end|><|start|>assistant<|channel|>final<|message|>Here’s a basic recipe for a delicious curry that’s adaptable to your taste...(中略)...Enjoy your homemade curry! Let me know if you’d like variations (e.g., Thai, Jamaican, or vegetarian).<|return|>

比較のため、Fine-Tuning前の回答結果も掲載します。こちらはリーズニングや英語で出力されますが、回答は日本語で生成されます。

<|start|>user<|message|>カレーの美味しい作り方を教えて。<|end|><|start|>assistant<|channel|>analysis<|message|>The user wants a recipe for making delicious curry, ...(中略)...We'll produce a nice, thorough recipe.<|end|><|start|>assistant<|channel|>final<|message|>## ざっくりした作り方(簡単日本風カレー)...(中略)...もし**インド風カレー**や**タイ風カレー**に挑戦したい場合は、スパイスミックスを変えたり、ココナッツミルクを加えてアレンジしてみてください。Happy cooking!<|return|>

要因は今回使用した学習データにあると考えられます。今回のデータの回答はすべて英語で記載されていたため、回答を英語で出力するよう学習してしまったのではないかと思われます。日本語で回答しているデータをいくつか学習データに含めたほうが良さそうだな、と感じました。

まとめ

いかがでしょうか。結果はちょっと意図したものと違う方向にいってしまいましたが、gpt-oss-20bをdatabricksでFine-Tuningするための大筋の手順は掴むことが出来たと思います。今回使用したデータと同じ形式のデータを用意し、同様の手順を実行することで様々なカスタマイズをgpt-oss-20bに施すことが出来ることが分かり、今後色々やってみたいことが増えたように感じました。ぜひトライしてみてください!