README.md

---
language:
  - en
  - ja
license: apache-2.0
tags:
  - olmo
  - causal-lm
  - bilingual
  - japanese
  - from-scratch
pipeline_tag: text-generation
---

# OLMo-2-386M-JA

英日バイリンガルの小規模言語モデル。OLMo-2 アーキテクチャをベースに、**ランダム初期化から** Pretrain → CPT → SFT → DPO の4ステージパイプラインで学習しました。

事前学習済みモデルの重みをダウンロードせず、AllenAI が公開しているアーキテクチャ設定（`config.json`）のみを使用してゼロから学習しています。

---

## モデル概要

| 項目 | 値 |
|------|-----|
| アーキテクチャ | OLMo-2（縮小版） |
| パラメータ数 | ~386M |
| 語彙数 | 65,536（英日 BPE） |
| コンテキスト長 | 2,048 tokens |
| 精度 | bfloat16 |
| 対応言語 | 英語 / 日本語 |

### アーキテクチャ詳細

| ハイパーパラメータ | 値 |
|---|---|
| hidden_size | 1,024 |
| num_hidden_layers | 16 |
| num_attention_heads | 16 |
| num_key_value_heads | 8（GQA） |
| intermediate_size | 4,096 |

---

## 学習パイプライン

```
ランダム重み初期化
      ↓
Stage 1: Pretraining    (~2B tokens, EN 70% / JA 30%)
      ↓
Stage 2: CPT / Annealing  (高品質データで継続学習)
      ↓
Stage 3: SFT            (~115K instruction samples)
      ↓
Stage 4: DPO            (UltraFeedback preference alignment)
      ↓
  本モデル
```

### Stage 1: Pretraining

| 項目 | 詳細 |
|------|------|
| データ | FineWeb (EN 70%) + Wikipedia JA (JA 30%) |
| トークン数 | ~2B tokens |
| LR | 3e-4（cosine + min-LR 10%） |
| Effective batch | 128 seq ≈ 262K tokens |
| Optimizer | AdamW (weight_decay=0.1) |

### Stage 2: CPT（Continued Pretraining）

| 項目 | 詳細 |
|------|------|
| データ | FineWeb-Edu (EN 60%) + CulturaX JA (JA 40%) |
| ドキュメント数 | 500K docs |
| LR | 1e-4（cosine） |

### Stage 3: SFT（Supervised Fine-Tuning）

| データセット | サンプル数 | 言語 |
|---|---|---|
| allenai/tulu-3-sft-mixture | 70,000 | 英語 |
| kunishou/databricks-dolly-15k-ja | ~15,000 | 日本語 |
| lightblue/tagengo-gpt4（日本語のみ） | 10,000 | 日本語 |
| izumi-lab/llm-japanese-dataset | 20,000 | 日本語 |
| **合計** | **~115,000** | |

| 項目 | 値 |
|---|---|
| LR | 5e-6（linear） |
| Epochs | 2 |
| Effective batch | 64 |
| Loss | Cross-Entropy（assistant turns のみ） |

### Stage 4: DPO（Direct Preference Optimization）

| 項目 | 詳細 |
|------|------|
| データ | allenai/ultrafeedback_binarized_cleaned |
| Loss | DPO sigmoid（β=0.1） |
| LR | 5e-7（linear） |
| Epochs | 1 |

---

## トークナイザー

英日バイリンガル BPE トークナイザーをゼロから学習しました。

| 項目 | 値 |
|---|---|
| 種類 | ByteLevel BPE |
| 語彙数 | 65,536（= 2¹⁶） |
| 学習データ | FineWeb 60K docs + Wikipedia JA 60K docs |
| 特殊トークン | `<\|endoftext\|>`, `<\|padding\|>`, `<\|user\|>`, `<\|assistant\|>`, `<\|system\|>` |

OLMo-2 のオリジナルトークナイザー（vocab=50,280）と比べて日本語トークン効率が大幅に向上しています。

---

## 使い方

```python
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_id = "itacora/olmo2-386m-ja"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

messages = [{"role": "user", "content": "機械学習とは何ですか？"}]
input_text = tokenizer.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True
)
inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
inputs.pop("token_type_ids", None)

with torch.inference_mode():
    output = model.generate(
        **inputs,
        max_new_tokens=200,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        repetition_penalty=1.3,
    )

response = tokenizer.decode(output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
print(response)
```

### チャットテンプレート

```
<|user|>
{ユーザーの入力}
<|assistant|>
{モデルの返答}<|endoftext|>
```

---

## 学習環境

| 項目 | 詳細 |
|------|------|
| GPU | NVIDIA RTX 4090 (24GB VRAM) |
| フレームワーク | PyTorch 2.4 + Transformers + TRL |
| 精度 | bfloat16 + gradient checkpointing |
| プラットフォーム | RunPod |

---

## 限界と注意事項

- **Pretrain データ量**: Chinchilla 最適値（386M × 20 = ~8B tokens）に対して約 2B tokens のみ。言語能力・知識量ともに制限があります。
- **繰り返し**: 長い生成では同じ文が繰り返されることがあります。`repetition_penalty=1.3` の使用を推奨します。
- **事実の誤り**: 知識が不十分なため、誤った情報を生成する場合があります。重要な用途には使用しないでください。
- **日本語品質**: SFT データの一部（izumi-lab）に `【P】` 等のフォーマットマーカーが含まれており、日本語生成に影響することがあります。
- **研究・実験目的**: 本モデルはフルパイプラインの動作確認・学習実験を目的として作成されました。プロダクション用途には適しません。

---

## ライセンス

Apache 2.0

学習データに使用した各データセットのライセンスも確認してください：
- [FineWeb](https://huggingface.co/datasets/HuggingFaceFW/fineweb): ODC-By
- [Wikipedia](https://huggingface.co/datasets/wikimedia/wikipedia): CC BY-SA 4.0
- [Tülu-3 SFT mixture](https://huggingface.co/datasets/allenai/tulu-3-sft-mixture): ODC-BY
- [UltraFeedback](https://huggingface.co/datasets/allenai/ultrafeedback_binarized_cleaned): MIT

---

## 引用

本モデルは以下の研究・成果物をベースにしています：

```bibtex
@article{olmo2,
  title={OLMo 2: The best fully open language model to date},
  author={Team OLMo},
  year={2025},
  url={https://arxiv.org/abs/2501.00656}
}
```