ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter8/MultilingualReasoning/gpt_oss_20b_sft.py

@@ -0,0 +1,554 @@
+"""
+多语言推理模型微调脚本
+
+本脚本展示如何使用 Hugging Face 的 TRL 库对 OpenAI 的 gpt-oss-20b 模型进行微调，
+使其能够在多种语言中进行有效推理。
+
+基于 OpenAI Cookbook 教程：
+https://cookbook.openai.com/articles/gpt-oss/fine-tune-transfomers
+
+作者: Edward Beeching, Quentin Gallouédec, Lewis Tunstall
+修改: 适配为完整的 Python 脚本
+
+⚠️  硬件要求（重要！）:
+- GPU: H100（80GB 显存）或更高配置
+- 训练时间: H100 上约 18 分钟
+- 使用 Mxfp4Config 量化和 LoRA 进行内存高效训练
+
+功能特性:
+- 使用 Mxfp4Config（针对 OpenAI 模型优化的 4-bit 浮点格式）
+- 使用 LoRA 进行内存高效的微调（包括 MoE 专家层）
+- 支持多语言推理（英语、西班牙语、法语、德语、意大利语等）
+- 可以混合语言（用一种语言提问，用另一种语言推理）
+- 所有超参数与 OpenAI Cookbook 教程完全一致
+"""
+
+import os
+import argparse
+try:
+    import torch
+    from datasets import load_dataset
+    from transformers import AutoModelForCausalLM, AutoTokenizer, Mxfp4Config
+    from peft import LoraConfig, PeftModel, get_peft_model
+    from trl import SFTTrainer, SFTConfig
+except ImportError:
+    torch = None
+    load_dataset = None
+    AutoModelForCausalLM = AutoTokenizer = Mxfp4Config = None
+    LoraConfig = PeftModel = get_peft_model = None
+    SFTTrainer = SFTConfig = None
+
+
+# ============================================================================
+# 第一部分：数据集准备
+# ============================================================================
+
+def load_and_prepare_dataset():
+    """
+    加载并准备多语言推理数据集
+    
+    使用 HuggingFaceH4/Multilingual-Thinking 数据集，该数据集包含：
+    - 多种语言的推理链（思维链）
+    - 支持英语、西班牙语、法语、德语、意大利语等
+    
+    Returns:
+        Dataset: 格式化后的训练数据集
+    """
+    print("=" * 80)
+    print("步骤 1: 加载数据集")
+    print("=" * 80)
+    
+    # 从 Hugging Face Hub 加载数据集
+    dataset = load_dataset("HuggingFaceH4/Multilingual-Thinking")
+    
+    print(f"数据集加载完成！")
+    print(f"训练样本数: {len(dataset['train'])}")
+    print(f"数据集列: {dataset['train'].column_names}")
+    print(f"\n示例数据:")
+    print(dataset['train'][0])
+    
+    return dataset['train']
+
+
+def format_chat_template(example, tokenizer):
+    """
+    格式化对话模板
+    
+    将数据集中的消息格式化为模型可以理解的对话格式
+    
+    Args:
+        example: 数据集中的一个样本
+        tokenizer: 分词器
+        
+    Returns:
+        dict: 格式化后的样本
+    """
+    # 应用聊天模板（带 messages 类型校验）
+    messages = example.get("messages")
+    if not isinstance(messages, list):
+        messages = []
+    example["text"] = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+    )
+    return example
+
+
+# ============================================================================
+# 第二部分：模型准备
+# ============================================================================
+
+def load_base_model(model_name="openai/gpt-oss-20b"):
+    """
+    加载基础模型和分词器
+    
+    使用 Mxfp4Config 进行量化，这是专门为 OpenAI 模型优化的 4-bit 浮点格式。
+    
+    Args:
+        model_name: 模型名称或路径
+        
+    Returns:
+        tuple: (model, tokenizer)
+    """
+    print("\n" + "=" * 80)
+    print("步骤 2: 加载基础模型")
+    print("=" * 80)
+    
+    # 加载分词器
+    print(f"加载分词器: {model_name}")
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    
+    # 设置 pad token（如果不存在）
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
+    # 配置 Mxfp4 量化（针对 OpenAI 模型优化）
+    print("使用 Mxfp4Config 量化...")
+    quantization_config = Mxfp4Config(dequantize=True)
+    
+    # 配置模型加载参数
+    model_kwargs = {
+        "attn_implementation": "eager",      # 注意力实现方式
+        "torch_dtype": torch.bfloat16,       # 使用 bfloat16 提高效率
+        "quantization_config": quantization_config,  # Mxfp4 量化配置
+        "use_cache": False,                   # 训练时禁用 KV 缓存
+        "device_map": "auto",                 # 自动分配设备
+    }
+    
+    # 加载模型
+    print(f"加载模型: {model_name}")
+    print("这可能需要几分钟时间...")
+    model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
+    
+    print(f"模型加载完成！")
+    print(f"模型参数量: {model.num_parameters() / 1e9:.2f}B")
+    
+    return model, tokenizer
+
+
+def prepare_model_for_lora(model, lora_rank=8, lora_alpha=16):
+    """
+    配置 LoRA（低秩适应）进行高效微调
+    
+    LoRA 只训练少量参数，大大减少内存使用和训练时间。
+    针对 openai/gpt-oss-20b 的 MoE（混合专家）架构，除了注意力层外，
+    还需要特别指定 MLP 专家层进行训练。
+    
+    Args:
+        model: 基础模型
+        lora_rank: LoRA 秩（默认 8，与官方教程一致）
+        lora_alpha: LoRA 缩放参数（默认 16）
+        
+    Returns:
+        PeftModel: 配置了 LoRA 的模型
+    """
+    print("\n" + "=" * 80)
+    print("步骤 3: 配置 LoRA")
+    print("=" * 80)
+    
+    # LoRA 配置（与 OpenAI Cookbook 一致）
+    peft_config = LoraConfig(
+        r=lora_rank,                     # LoRA 秩
+        lora_alpha=lora_alpha,           # LoRA 缩放参数
+        target_modules="all-linear",     # 目标所有线性层
+        target_parameters=[              # MoE 专家层的特定参数
+            "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",
+        ],
+    )
+    
+    print("LoRA 配置:")
+    print(f"  - Rank: {lora_rank}")
+    print(f"  - Alpha: {lora_alpha}")
+    print(f"  - 目标模块: {peft_config.target_modules}")
+    print(f"  - MoE 专家层参数: {len(peft_config.target_parameters)} 个")
+    
+    # 应用 LoRA
+    model = get_peft_model(model, peft_config)
+    
+    # 打印可训练参数统计
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_percent = 100 * trainable_params / total_params
+    
+    print(f"\n可训练参数统计:")
+    print(f"  - 可训练参数: {trainable_params:,} ({trainable_percent:.2f}%)")
+    print(f"  - 总参数: {total_params:,}")
+    
+    return model
+
+
+# ============================================================================
+# 第三部分：训练
+# ============================================================================
+
+def train_model(model, tokenizer, dataset, output_dir="./gpt-oss-20b-multilingual-reasoner", 
+                batch_size=4, num_epochs=1, learning_rate=2e-4, max_seq_length=2048):
+    """
+    使用 SFTTrainer 训练模型
+    
+    Args:
+        model: 配置了 LoRA 的模型
+        tokenizer: 分词器
+        dataset: 训练数据集
+        output_dir: 输出目录
+        batch_size: 批次大小（根据 GPU 显存调整，默认 4）
+        num_epochs: 训练轮数（默认 1）
+        learning_rate: 学习率（默认 2e-4）
+        max_seq_length: 最大序列长度
+        
+    Returns:
+        SFTTrainer: 训练好的 trainer 对象
+    """
+    print("\n" + "=" * 80)
+    print("步骤 4: 开始训练")
+    print("=" * 80)
+    
+    # 训练参数配置（与 OpenAI Cookbook 完全一致）
+    training_args = SFTConfig(
+        learning_rate=learning_rate,
+        gradient_checkpointing=True,
+        num_train_epochs=num_epochs,
+        logging_steps=1,
+        per_device_train_batch_size=batch_size,
+        gradient_accumulation_steps=4,
+        max_length=max_seq_length,
+        warmup_ratio=0.03,
+        lr_scheduler_type="cosine_with_min_lr",
+        lr_scheduler_kwargs={"min_lr_rate": 0.1},
+        output_dir=output_dir,
+        report_to="trackio",  # 设为 "trackio" 以启用实验跟踪
+        push_to_hub=False,  # 设为 True 以自动推送到 Hub
+    )
+    
+    print("训练配置:")
+    print(f"  - 批次大小: {batch_size}")
+    print(f"  - 梯度累积步数: {training_args.gradient_accumulation_steps}")
+    print(f"  - 有效批次大小: {batch_size * training_args.gradient_accumulation_steps}")
+    print(f"  - 训练轮数: {num_epochs}")
+    print(f"  - 学习率: {learning_rate}")
+    print(f"  - 学习率调度: {training_args.lr_scheduler_type}")
+    print(f"  - 最大序列长度: {max_seq_length}")
+    print(f"  - 输出目录: {output_dir}")
+    
+    # 初始化 SFTTrainer
+    trainer = SFTTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset,
+        processing_class=tokenizer,
+    )
+    
+    # 开始训练
+    print("\n开始训练...")
+    print("⚠️  在 H100 GPU 上训练约需 18 分钟")
+    print("-" * 80)
+    
+    trainer.train()
+    
+    print("\n" + "=" * 80)
+    print("训练完成！")
+    print("=" * 80)
+    
+    return trainer
+
+
+# ============================================================================
+# 第四部分：保存和推送模型
+# ============================================================================
+
+def save_and_push_model(trainer, output_dir, push_to_hub=False, hub_model_id=None):
+    """
+    保存模型并可选择推送到 Hugging Face Hub
+    
+    Args:
+        trainer: 训练好的 trainer 对象
+        output_dir: 输出目录
+        push_to_hub: 是否推送到 Hub
+        hub_model_id: Hub 上的模型 ID
+    """
+    print("\n" + "=" * 80)
+    print("步骤 5: 保存模型")
+    print("=" * 80)
+    
+    # 保存模型到本地
+    print(f"保存模型到: {output_dir}")
+    trainer.save_model(output_dir)
+    print("模型保存完成！")
+    
+    # 可选：推送到 Hugging Face Hub
+    if push_to_hub:
+        if hub_model_id is None:
+            raise ValueError("需要提供 hub_model_id 才能推送到 Hub")
+
+        print(f"\n推送模型到 Hugging Face Hub: {hub_model_id}")
+        # Trainer.push_to_hub 从 args.hub_model_id 取仓库名；不设置的话会
+        # 忽略用户传入的 --hub_model_id，推到 output_dir 同名的默认仓库。
+        trainer.args.hub_model_id = hub_model_id
+        trainer.push_to_hub(
+            dataset_name="HuggingFaceH4/Multilingual-Thinking",
+        )
+        print("模型已成功推送到 Hub！")
+
+
+# ============================================================================
+# 第五部分：推理
+# ============================================================================
+
+def load_trained_model(base_model_name, peft_model_path):
+    """
+    加载训练好的模型进行推理
+    
+    Args:
+        base_model_name: 基础模型名称
+        peft_model_path: LoRA 权重路径
+        
+    Returns:
+        tuple: (model, tokenizer)
+    """
+    print("\n" + "=" * 80)
+    print("加载训练好的模型进行推理")
+    print("=" * 80)
+    
+    # 加载分词器
+    print(f"加载分词器: {base_model_name}")
+    tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+    
+    # 加载基础模型
+    print(f"加载基础模型: {base_model_name}")
+    model_kwargs = {
+        "attn_implementation": "eager",
+        "torch_dtype": "auto",
+        "use_cache": True,  # 推理时启用 KV 缓存
+        "device_map": "auto",
+    }
+    base_model = AutoModelForCausalLM.from_pretrained(base_model_name, **model_kwargs)
+    
+    # 加载并合并 LoRA 权重
+    print(f"加载 LoRA 权重: {peft_model_path}")
+    model = PeftModel.from_pretrained(base_model, peft_model_path)
+    
+    print("合并 LoRA 权重与基础模型...")
+    model = model.merge_and_unload()
+    
+    print("模型加载完成！")
+    
+    return model, tokenizer
+
+
+def generate_response(model, tokenizer, reasoning_language, user_prompt, 
+                     max_new_tokens=512, temperature=0.6, format_output=True):
+    """
+    生成多语言推理响应
+    
+    Args:
+        model: 训练好的模型
+        tokenizer: 分词器
+        reasoning_language: 推理使用的语言
+        user_prompt: 用户提问
+        max_new_tokens: 最大生成 token 数
+        temperature: 采样温度（越高越随机）
+        format_output: 是否格式化输出（使用明显的标记）
+        
+    Returns:
+        str: 生成的完整响应
+    """
+    # 构建消息
+    system_prompt = f"reasoning language: {reasoning_language}"
+    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": max_new_tokens,
+        "do_sample": True,
+        "temperature": temperature,
+        "top_p": None,
+        "top_k": None,
+    }
+    
+    # 生成响应
+    print(f"\n生成响应...")
+    print(f"推理语言: {reasoning_language}")
+    print(f"用户提问: {user_prompt}")
+    
+    with torch.no_grad():
+        output_ids = model.generate(input_ids, **gen_kwargs)
+    
+    # 解码输出 - 保留特殊标记以便解析
+    response_with_tokens = tokenizer.batch_decode(output_ids, skip_special_tokens=False)[0]
+    print("-" * 80)
+    print(response_with_tokens)
+    print("-" * 80)
+
+
+def run_inference_examples(model, tokenizer):
+    """
+    运行多个推理示例
+    
+    Args:
+        model: 训练好的模型
+        tokenizer: 分词器
+    """
+    print("\n" + "=" * 80)
+    print("推理示例")
+    print("=" * 80)
+    
+    # 示例 1: 西班牙语提问，德语推理
+    print("\n[示例 1: 西班牙语提问 + 德语推理]")
+    generate_response(
+        model, tokenizer,
+        reasoning_language="German",
+        user_prompt="¿Cuál es el capital de Australia?",  # 澳大利亚的首都是什么？
+        format_output=True,
+    )
+    
+    # 示例 2: 英语提问，中文推理
+    print("\n\n[示例 2: 英语提问 + 中文推理]")
+    generate_response(
+        model, tokenizer,
+        reasoning_language="Chinese",
+        user_prompt="What is the national symbol of Canada?",
+        format_output=True,
+    )
+    
+    # 示例 3: 中文提问，中文推理
+    print("\n\n[示例 3: 中文提问 + 中文推理]")
+    generate_response(
+        model, tokenizer,
+        reasoning_language="Chinese",
+        user_prompt="求解 x^2 - 2x + 1 = 0 的根",
+        format_output=True,
+    )
+
+
+# ============================================================================
+# 主函数
+# ============================================================================
+
+def main():
+    """主函数：完整的训练流程"""
+    
+    parser = argparse.ArgumentParser(description="多语言推理模型微调")
+    parser.add_argument(
+        "--mode", 
+        type=str, 
+        choices=["train", "inference", "full"],
+        default="full",
+        help="运行模式: train（仅训练）, inference（仅推理）, full（完整流程）"
+    )
+    parser.add_argument("--model_name", type=str, default="openai/gpt-oss-20b", help="基础模型名称")
+    parser.add_argument("--output_dir", type=str, default="./gpt-oss-20b-multilingual-reasoner", help="输出目录")
+    parser.add_argument("--batch_size", type=int, default=4, help="训练批次大小（默认 4，与官方教程一致）")
+    parser.add_argument("--num_epochs", type=int, default=1, help="训练轮数（默认 1，与官方教程一致）")
+    parser.add_argument("--learning_rate", type=float, default=2e-4, help="学习率（默认 2e-4，与官方教程一致）")
+    parser.add_argument("--max_seq_length", type=int, default=2048, help="最大序列长度")
+    parser.add_argument("--lora_rank", type=int, default=8, help="LoRA 秩（默认 8，与官方教程一致）")
+    parser.add_argument("--lora_alpha", type=int, default=16, help="LoRA alpha")
+    parser.add_argument("--push_to_hub", action="store_true", default=False, help="推送模型到 Hugging Face Hub")
+    parser.add_argument("--hub_model_id", type=str, default=None, help="Hub 模型 ID")
+    
+    args = parser.parse_args()
+    
+    print("=" * 80)
+    print("多语言推理模型微调")
+    print("=" * 80)
+    print(f"模式: {args.mode}")
+    print(f"基础模型: {args.model_name}")
+    print(f"输出目录: {args.output_dir}")
+    
+    # 训练模式
+    if args.mode in ["train", "full"]:
+        # 1. 加载数据集
+        dataset = load_and_prepare_dataset()
+        
+        # 2. 加载基础模型（使用 Mxfp4Config 量化）
+        model, tokenizer = load_base_model(args.model_name)
+        
+        # 3. 配置 LoRA
+        model = prepare_model_for_lora(model, args.lora_rank, args.lora_alpha)
+        
+        # 4. 训练模型
+        trainer = train_model(
+            model, 
+            tokenizer, 
+            dataset,
+            output_dir=args.output_dir,
+            batch_size=args.batch_size,
+            num_epochs=args.num_epochs,
+            learning_rate=args.learning_rate,
+            max_seq_length=args.max_seq_length,
+        )
+        
+        # 5. 保存模型
+        save_and_push_model(
+            trainer, 
+            args.output_dir,
+            push_to_hub=args.push_to_hub,
+            hub_model_id=args.hub_model_id,
+        )
+        
+        if args.mode == "full":
+            # full 模式继续跑推理：先释放训练占用的显存，
+            # 再按推理路径从 output_dir 重新加载已保存的模型。
+            del trainer
+            del model
+            torch.cuda.empty_cache()
+            print("\n训练完成！已释放训练显存，继续运行推理示例。")
+        else:
+            print("\n训练完成！建议重启内核以释放 GPU 显存后再进行推理。")
+
+    # 推理模式（inference 单独运行；full 在训练后接着运行）
+    if args.mode in ["inference", "full"]:
+        if not os.path.exists(args.output_dir):
+            print(f"错误: 未找到模型目录 {args.output_dir}")
+            print("请先运行训练或指定正确的模型路径")
+            return
+        
+        # 加载训练好的模型
+        model, tokenizer = load_trained_model(args.model_name, args.output_dir)
+        
+        # 运行推理示例
+        run_inference_examples(model, tokenizer)
+    
+    print("\n" + "=" * 80)
+    print("完成！")
+    print("=" * 80)
+
+
+if __name__ == "__main__":
+    main()
+