튜토리얼: OpenAI GPT-OSS 미세 조정

pip install --user --upgrade \
  torch --extra-index-url https://download.pytorch.org/whl/cu128 \
  "transformers>=4.55.0" "trl>=0.20.0" "peft>=0.17.0" \
  datasets trackio wandb weave huggingface_hub ipywidgets "jinja2>=3.1.0"

import wandb

import os

# 0.1 Setup WANDB for training metrics
os.environ["WANDB_PROJECT"] = "YOUR_PROJECT_NAME"  # Change to your desired project!
wandb.init(project=os.environ["WANDB_PROJECT"])

from huggingface_hub import notebook_login

notebook_login()

# 1. LOAD DATASET
from datasets import load_dataset

dataset = load_dataset("HuggingFaceH4/Multilingual-Thinking", split="train")
print(f"Total examples: {len(dataset)}")
print("Sample:", dataset[0])

# 2. LOAD TOKENIZER & PREVIEW TEMPLATE
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("openai/gpt-oss-20b")

messages = dataset[0]["messages"]
conversation = tokenizer.apply_chat_template(messages, tokenize=False)
print(conversation)  # Shows how dialogue is formatted

# 3. LOAD BASE MODEL WITH QUANTIZATION
import torch
from transformers import AutoModelForCausalLM, Mxfp4Config

quant_config = Mxfp4Config(dequantize=True)
model_kwargs = dict(
    attn_implementation="eager",
    torch_dtype=torch.bfloat16,
    quantization_config=quant_config,
    use_cache=False,
    device_map="auto",
)
model = AutoModelForCausalLM.from_pretrained("openai/gpt-oss-20b", **model_kwargs)

# 4. WRAP WITH LORA (PEFT)
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()

# 5. CONFIGURE TRAINING (& use wandb)
from trl import SFTConfig

training_args = SFTConfig(
    learning_rate=2e-4,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    gradient_checkpointing=True,
    num_train_epochs=1,
    logging_steps=1,
    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",
    report_to="wandb",  # ← use wandb
    push_to_hub=False,  # ← DO NOT push to hub
)

# 6. START TRAINING
from trl import SFTTrainer

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

# 7. SAVE MODEL LOCALLY (no pushing to hub)
trainer.save_model(training_args.output_dir)

del trainer
del peft_model
del model

# 8. INFERENCE: LOAD FINE-TUNED MODEL & GENERATE OUTPUT

from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import weave; weave.init("gpt-oss")


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()
peft_model_id = "gpt-oss-20b-multilingual-reasoner"
model = PeftModel.from_pretrained(base_model, peft_model_id)
model = model.merge_and_unload()

# Example generation: User asks in Spanish, wants reasoning in German
import torch

@weave.op
def run_inference(system_prompt, user_prompt, reasoning_language="German"):
    messages = [
        {"role": "system", "content": f"reasoning language: {reasoning_language}"},
        {"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": 512, "do_sample": True, "temperature": 0.6}
    output_ids = model.generate(input_ids, **gen_kwargs)
    response = tokenizer.batch_decode(output_ids)[0]
    return response



# Inference with German reasoning
out = run_inference("reasoning language: German", "¿Cuál es el capital de Australia?", "German")
print("\nGerman reasoning output:\n", out)

# Inference with Chinese reasoning (not seen in finetuning)
out_chinese = run_inference("reasoning language: Chinese", "What is the national symbol of Canada?", "Chinese")
print("\nChinese reasoning output:\n", out_chinese)