LLM事前学習とスケーリング則：Chinchillaからフラッシュアテンション・MoEまで

import math

def chinchilla_optimal(compute_budget_flops: float):
    """
    Chinchilla則に基づく最適N、Dの計算
    compute_budget_flops: 総FLOPs (6 * N * Dで近似)
    Returns: (optimal_params, optimal_tokens)
    """
    # Hoffmann et al. Table A3の係数を使用
    # C = 6 * N * D の近似から:
    # N_opt = sqrt(C / (6 * 20)), D_opt = 20 * N_opt
    N_opt = math.sqrt(compute_budget_flops / (6 * 20))
    D_opt = 20 * N_opt
    return N_opt, D_opt

# 例: A100 1000台、30日間の学習
# A100: ~312 TFLOPS (bf16)、稼働率40%
flops_per_sec = 1000 * 312e12 * 0.4
duration_sec = 30 * 24 * 3600
total_flops = flops_per_sec * duration_sec

N_opt, D_opt = chinchilla_optimal(total_flops)
print(f"総FLOPs: {total_flops:.2e}")
print(f"最適パラメータ数: {N_opt/1e9:.1f}B")
print(f"最適トークン数: {D_opt/1e12:.1f}T")

from datasketch import MinHash, MinHashLSH
import re

def get_shingles(text: str, k: int = 5):
    """テキストから文字kシングル集合を生成"""
    text = re.sub(r'\s+', ' ', text.lower())
    return {text[i:i+k] for i in range(len(text) - k + 1)}

def build_minhash(text: str, num_perm: int = 128) -> MinHash:
    """テキストからMinHashシグネチャを作成"""
    m = MinHash(num_perm=num_perm)
    for shingle in get_shingles(text):
        m.update(shingle.encode('utf-8'))
    return m

def deduplicate_corpus(documents: list, threshold: float = 0.8):
    """
    MinHash LSHで類似文書を除去
    threshold: 重複とみなすJaccard類似度の閾値
    """
    lsh = MinHashLSH(threshold=threshold, num_perm=128)
    unique_docs = []

    for idx, doc in enumerate(documents):
        mh = build_minhash(doc)
        result = lsh.query(mh)

        if len(result) == 0:
            lsh.insert(f"doc_{idx}", mh)
            unique_docs.append(doc)
        # 類似文書がすでに存在する場合はスキップ

    print(f"元の文書数: {len(documents)} -> 重複排除後: {len(unique_docs)}")
    return unique_docs

from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.processors import ByteLevel as ByteLevelProcessor

def train_bpe_tokenizer(
    corpus_files: list,
    vocab_size: int = 32000,
    save_path: str = "tokenizer.json"
):
    """GPT-2スタイルのByte-level BPEトークナイザーを学習"""
    tokenizer = Tokenizer(BPE(unk_token=None))
    tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)
    tokenizer.post_processor = ByteLevelProcessor(trim_offsets=False)

    trainer = BpeTrainer(
        vocab_size=vocab_size,
        min_frequency=2,
        special_tokens=["<pad>", "<s>", "</s>", "<unk>"],
        show_progress=True,
    )

    tokenizer.train(files=corpus_files, trainer=trainer)
    tokenizer.save(save_path)
    print(f"トークナイザー保存完了: {save_path} (vocab_size={vocab_size})")
    return tokenizer

import torch
import torch.nn as nn
import torch.nn.functional as F

class MoELayer(nn.Module):
    """シンプルなMixture of Expertsレイヤー"""
    def __init__(self, d_model: int, d_ff: int, num_experts: int = 8, top_k: int = 2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.router = nn.Linear(d_model, num_experts, bias=False)
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_ff),
                nn.SiLU(),
                nn.Linear(d_ff, d_model),
            )
            for _ in range(num_experts)
        ])

    def forward(self, x: torch.Tensor):
        # x: (batch, seq_len, d_model)
        B, T, D = x.shape
        x_flat = x.view(-1, D)  # (B*T, D)

        # ルーター計算
        router_logits = self.router(x_flat)  # (B*T, num_experts)
        router_probs = F.softmax(router_logits, dim=-1)

        # Top-K専門家の選択
        topk_probs, topk_idx = router_probs.topk(self.top_k, dim=-1)
        topk_probs = topk_probs / topk_probs.sum(dim=-1, keepdim=True)  # 再正規化

        # 負荷分散損失（専門家のバランス維持）
        expert_load = router_probs.mean(0)  # (num_experts,)
        load_balance_loss = self.num_experts * (expert_load * expert_load.mean()).sum()

        # 専門家の出力計算
        output = torch.zeros_like(x_flat)
        for k in range(self.top_k):
            expert_indices = topk_idx[:, k]
            expert_weights = topk_probs[:, k].unsqueeze(-1)
            for e_idx in range(self.num_experts):
                mask = (expert_indices == e_idx)
                if mask.any():
                    expert_out = self.experts[e_idx](x_flat[mask])
                    output[mask] += expert_weights[mask] * expert_out

        return output.view(B, T, D), load_balance_loss

import math

def cosine_lr_with_warmup(
    optimizer,
    step: int,
    warmup_steps: int,
    total_steps: int,
    max_lr: float,
    min_lr: float = 0.0,
):
    """線形ウォームアップ付きコサインアニーリング"""
    if step < warmup_steps:
        lr = max_lr * step / warmup_steps
    else:
        progress = (step - warmup_steps) / (total_steps - warmup_steps)
        lr = min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * progress))

    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
    return lr

# 典型的な事前学習設定
# warmup: 全ステップの1~2%
# max_lr: 3e-4 (7Bモデル向け)
# min_lr: max_lr * 0.1

import torch

def train_step_with_stability(model, optimizer, batch, grad_clip: float = 1.0):
    """勾配ノルムクリッピングとスパイク検出を備えた学習ステップ"""
    optimizer.zero_grad()
    loss = model(batch)
    loss.backward()

    # 勾配ノルムの計算とクリッピング
    grad_norm = torch.nn.utils.clip_grad_norm_(
        model.parameters(), max_norm=grad_clip
    )

    # 異常に大きな勾配の検出
    if grad_norm > 100 * grad_clip:
        print(f"警告: 異常なgrad_norm={grad_norm:.2f}、ステップをスキップ")
        optimizer.zero_grad()
        return None, grad_norm

    optimizer.step()
    return loss.item(), grad_norm

# Flash Attention 2の使用例
# pip install flash-attn --no-build-isolation

import torch
from flash_attn import flash_attn_func

def flash_attention_forward(
    q: torch.Tensor,  # (batch, seqlen, nheads, headdim)
    k: torch.Tensor,
    v: torch.Tensor,
    causal: bool = True,
    softmax_scale: float = None,
):
    """
    Flash Attention 2の順伝播
    causal=True: 自己回帰言語モデル用の因果マスクを適用
    """
    if softmax_scale is None:
        softmax_scale = q.shape[-1] ** (-0.5)

    out = flash_attn_func(
        q, k, v,
        dropout_p=0.0,
        softmax_scale=softmax_scale,
        causal=causal,
    )
    return out  # (batch, seqlen, nheads, headdim)

import torch
import math

def compute_perplexity(model, dataloader, device: str = "cuda"):
    """検証セットのパープレキシティを計算"""
    model.eval()
    total_loss = 0.0
    total_tokens = 0

    with torch.no_grad():
        for batch in dataloader:
            input_ids = batch["input_ids"].to(device)
            labels = batch["labels"].to(device)

            outputs = model(input_ids=input_ids, labels=labels)
            batch_tokens = (labels != -100).sum().item()
            total_loss += outputs.loss.item() * batch_tokens
            total_tokens += batch_tokens

    avg_loss = total_loss / total_tokens
    perplexity = math.exp(avg_loss)
    return perplexity

# 典型的な基準値
# GPT-2レベル: PPL 約20 (WebText)
# 良好な7Bモデル: PPL 約6〜8 (検証セット)

# lm-evaluation-harnessのインストールと実行
pip install lm-eval

# Llama 3 8Bの標準ベンチマーク評価
lm_eval --model hf \
    --model_args pretrained=meta-llama/Meta-Llama-3-8B \
    --tasks mmlu,hellaswag,arc_challenge,winogrande \
    --device cuda:0 \
    --batch_size 8 \
    --output_path results/llama3-8b/

# 学習中の特定チェックポイントを評価
lm_eval --model hf \
    --model_args pretrained=./checkpoints/step-50000 \
    --tasks hellaswag \
    --num_fewshot 10 \
    --device cuda:0