Traceback 12b

🚀 TraceBack 12b Release

TraceBackは、「推論トレースデータの生成を効果的にスケーリングするにはどうすればよいか」という疑問から生まれたものです。推論トレースの作成には、推論モデル（r1、o1、o3など）だけに依存する必要はないことがわかりました。

このモデルには多くの目標がありますが、主なものは以下の通りです。

• 小さなモデル（r1などよりも小さい）を使用するため、推論が高速で、合成推論データセットの生成をより迅速に行えるため、スケーリングが容易です。
• ドメイン外の検証不能な問題に対して、合成トレースを用いた知識蒸留が可能です。
• 入力として使用することで、非推論モデルの出力やデータセットを推論合成データセットに変換できます。

現時点での概念実証では、1と3の目標を達成しており、今後さらにスケーリングする予定です。

• 現在はMistral Nemo 12bをベースモデルとして使用しています。
• 2エポックのみのトレーニングを行っています。
• 微調整（Qlora）には200kのサンプルのみを使用しており、データセットはこちらです。

したがって、まだ改善の余地がたくさんあります。

このモデルは、命令と解決策を入力として使用し、それに基づいた妥当な推論トレースを出力します。私はこれが推論データ生成の未来だと信じています。評価版のリリースにご期待ください。

🚀 クイックスタート

このセクションでは、TraceBack 12bの基本的な使い方を説明します。

✨ 主な機能

• 高速な合成推論データセットの生成を可能にします。
• ドメイン外の非検証可能な問題に対する知識蒸留ができます。
• 非推論モデルの出力やデータセットを推論合成データセットに変換できます。

📦 インストール

このモデルを使用するには、必要なライブラリをインストールする必要があります。以下のコードを参考にしてください。

# 必要なライブラリのインストール
pip install transformers torch unsloth axolotl deepspeed

💻 使用例

基本的な使用法

以下は、ChatGPTの命令と解決策を入力として使用した推論の例です。

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# トークナイザーとモデルの読み込み
model_name = "secemp9/TraceBack-12b"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# モデルを目的のデバイスに移動
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)

# メッセージの定義
messages = [
    {"role": "user", "content": """Instruction:
how many r in strawberry

Solution:
There are **three** "r"s in "strawberry."
"""}
]

# ステップ1: チャットテンプレートを適用して、フォーマットされたテキストを文字列として取得
formatted_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

# ステップ2: フォーマットされたテキストをテンソルの辞書にトークナイズ
inputs = tokenizer(formatted_text, return_tensors="pt").to(device)

# 応答の生成
outputs = model.generate(**inputs, max_new_tokens=32000)

# 出力をデコードして表示
generated_text = tokenizer.decode(outputs[0])
print(generated_text)

高度な使用法

unslothを使用した例です。

from unsloth import FastLanguageModel

# モデルとトークナイザーの読み込み
model, tokenizer = FastLanguageModel.from_pretrained("secemp9/TraceBack-12b")

# メッセージの定義（"stuff_here"を実際の入力に置き換えてください）
messages = [
    {"role": "user", "content": """Instruction:
how many r in strawberry

Solution:
There are **three** "r"s in "strawberry."
"""}
]

# ステップ1: チャットテンプレートを適用して、フォーマットされたテキストを文字列として取得
formatted_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

# ステップ2: フォーマットされたテキストをテンソルの辞書にトークナイズ
inputs = tokenizer(formatted_text, return_tensors="pt").to(model.device)

# 応答の生成
outputs = model.generate(**inputs, max_new_tokens=32000)

# 出力をデコードして表示
generated_text = tokenizer.decode(outputs[0])
print(generated_text)

📚 ドキュメント

推論例

ここでは、ChatGPTの命令と解決策をモデルに入力する簡単な例を示します。

データセット例

データセットの形式は、命令 + 解決策: 推論トレースのペアに従っています。以下はサンプル会話です。

{
  "messages": [
    {
      "role": "user",
      "content": "Instruction:
      text_here

      Solution:
      text_here
    },
    {
      "role": "assistant",
      "content": "text_here"
    }
  ]
}

これは以下のように表示されます。

プロンプト形式

現時点でのプロンプト形式は以下の通りです。

Instruction:

Solution:

モデルの出力は、現時点では特にフォーマットされておらず、推論結果がそのまま出力されます。

🔧 技術詳細

Axolotl設定

UnslothのコードをAxolotlの設定ファイルに変換し、DeepSpeedを使用しています。以下が設定内容です。

config.yml

# ベースモデルの設定
base_model: unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit
load_in_4bit: true

# データセットの設定
datasets:
  - path: instruction_solution_to_thought_dataset.jsonl
    type: chat_template

# チャットテンプレート
chat_template: chatml

# LoRAアダプターの設定
adapter: lora
lora_r: 16
lora_alpha: 16
lora_dropout: 0
lora_target_modules:
  - q_proj
  - k_proj
  - v_proj
  - o_proj
  - gate_proj
  - up_proj
  - down_proj

# トレーニングのハイパーパラメータ
max_seq_length: 128000
micro_batch_size: 2
gradient_accumulation_steps: 8
learning_rate: 3e-5
num_epochs: 3
warmup_steps: 100
optimizer: adamw_8bit
weight_decay: 0.01
lr_scheduler_type: cosine
max_grad_norm: 1.0
output_dir: ./outputs_solution_to_thought
seed: 3407
merge_lora: true
hf_upload: true
hf_repo: secemp9/TraceBack-12b
xformers_attention:
flash_attention: True
bf16: true          # BF16混合精度を有効にする
# DeepSpeedによるマルチGPUトレーニング
deepspeed: deepspeed_configs/zero2.json

# オプション: 勾配チェックポインティングを有効にする
gradient_checkpointing: true

deepspeed_configs/zero2.json

{
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  },
  "bf16": {
    "enabled": true
  },
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": "auto",
      "weight_decay": "auto",
      "betas": [0.9, 0.999],
      "eps": 1e-8
    }
  },
  "scheduler": {
    "type": "WarmupLR",
    "params": {
      "warmup_min_lr": 0,
      "warmup_max_lr": "auto",
      "warmup_num_steps": "auto"
    }
  },
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto",
  "steps_per_print": 10,
  "wandb": {
    "enabled": true
  }
}

📄 ライセンス

このプロジェクトは、Apache-2.0ライセンスの下でライセンスされています。

📋 モデル情報