Tinyllama V1.1

🚀 TinyLlama-1.1B-v1.1

• コードベース: github.com/jzhang38/TinyLlama
• 技術レポート: arxiv.org/pdf/2401.02385

このプロジェクトでは、Llama 2とまったく同じアーキテクチャとトークナイザーを採用しています。これにより、TinyLlamaはLlamaをベースに構築された多くのオープンソースプロジェクトですぐに利用できます。さらに、TinyLlamaはわずか11億個のパラメータで構成されており、計算とメモリ使用量が制限された多くのアプリケーションに適しています。

🚀 クイックスタート

概要

このプロジェクトでは、単一のTinyLlamaモデルをトレーニングするだけでなく、まず15兆トークンのコーパスでTinyLlamaをトレーニングして基本的な言語能力を獲得します。その後、このモデルを使用して、3つの異なるデータサンプリングを用いた継続的な事前学習により、3つの異なるモデルを作成します。このプロセスの視覚的な表現については、以下の図を参照してください。

事前学習

いくつかの問題（bug1、bug2）のため、より良いモデルを提供するためにTinyLlamaを再トレーニングしました。モデルは2兆トークンでトレーニングし、事前学習を3つの異なる段階に分けました：1) 基本的な事前学習、2) 特定のドメインでの継続的な事前学習、3) クールダウン。

基本的な事前学習

この初期段階では、slimpajamaのみを使用してモデルをトレーニングし、常識的な推論能力を開発しました。この基本的な事前学習期間中、モデルは1.5兆トークンでトレーニングされました。ノードあたり4つのA100-40Gを備えたクラスタを使用し、ノード内でのみモデルの重みを分割したため、このときのバッチサイズは約180万に設定できました。

特定のドメインでの継続的な事前学習

この事前学習では、3種類のコーパスを組み込みました：slimpajama（最初の段階と同じ）、Math&Code（starcoderとproof pile）、およびChinese（Skypile）。このアプローチにより、特殊な機能を持つ3つのバリアントモデルを開発することができました。

この段階の最初の約60億トークンでは、ドメイン固有のコーパス（Slimpajamaを除く、なぜなら段階1と比較して変更されていないため）のサンプリング割合を線形に増やしました。このウォームアップサンプリング増加戦略は、事前学習データの分布を徐々に調整し、より安定したトレーニングプロセスを確保するために設計されています。このサンプリング増加段階の後、安定したサンプリング戦略でモデルを約1.85兆トークンに達するまで継続的に事前学習しました。

クールダウン

事前学習の最後に、より良いモデルの収束を達成するために、クールダウンフェーズを実装することが重要な技術になっています。ただし、最初からコサイン学習率戦略を使用しているため、MiniCPMやdeepseekのようにクールダウンのために学習率を変更することは困難です。そのため、バッチサイズを調整してクールダウンを試みました。具体的には、クールダウン段階でバッチサイズを180万から720万に増やし、元のコサイン学習率スケジュールを維持しました。

Tinyllamaモデルファミリー

広範かつ詳細な事前学習プロセスの後、以下の3つの特殊バージョンのモデルをリリースしています：

1. TinyLlama_v1.1：標準バージョン、一般的な用途に使用。
2. TinyLlama_v1.1_Math&Code：数学とコードに関する能力が高い。
3. TinyLlama_v1.1_Chinese：中国語の理解能力が高い。

データ

各段階でのデータ分布を以下に示します：

TinyLlama_v1.1

TinyLlama_v1.1_math_code

TinyLlama_v1.1_chinese

使い方

transformers>=4.31が必要です。詳細情報については、TinyLlamaのGitHubページを確認してください。

from transformers import AutoTokenizer
import transformers 
import torch
model = "TinyLlama/TinyLlama_v1.1"
tokenizer = AutoTokenizer.from_pretrained(model)
pipeline = transformers.pipeline(
    "text-generation",
    model=model,
    torch_dtype=torch.float16,
    device_map="auto",
)

sequences = pipeline(
    'The TinyLlama project aims to pretrain a 1.1B Llama model on 3 trillion tokens. With some proper optimization, we can achieve this within a span of "just" 90 days using 16 A100-40G GPUs 🚀🚀. The training has started on 2023-09-01.',
    do_sample=True,
    top_k=10,
    num_return_sequences=1,
    repetition_penalty=1.5,
    eos_token_id=tokenizer.eos_token_id,
    max_length=500,
)
for seq in sequences:
    print(f"Result: {seq['generated_text']}")

評価

📄 ライセンス

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