ml-training

ML トレーニング

概要

PyTorch、TensorFlow/Keras、JAX/Flax、Hugging Face Transformers/Diffusers/Accelerate/PEFT、scikit-learn、XGBoost、LightGBM、CatBoost、Spark MLlib、Ray によるモデルのトレーニングに、このスキルを使用します。まず正確性を重視してください。データの検証、データ漏洩を防ぐ分割、再現可能な設定、有意義な指標、そして複雑な分散処理やアクセラレーターを多用する実行に進む前のシンプルなベースラインが重要です。

トレーニング準備チェックリスト

1. タスク、目的変数、指標、ベースライン、合格基準を定義します。
2. データのスキーマ、ラベルの品質、欠損値、重複、クラスバランス、トレーニング時とサービング時の特徴量の整合性を検証します。
3. 分割方法を選択します。iid 分類にはランダムな層化分割、相関のあるエンティティにはグループ分割、時系列データには時系列分割、モデル選択に関する主張にはネストされた CV を使用します。
4. 環境を固定します。フレームワークのバージョン、CUDA/cuDNN、ドライバー、データセットのスナップショット、前処理コード、モデル設定、ハードウェアを記録します。
5. 有効な範囲でシードを設定し、非決定的な処理を記録します。決定論モードを検証していない限り、GPU 間や分散カーネル間でビット単位の再現性を約束しないでください。
6. 小さなバッチで過学習テストを行い、次に完全なベースラインを実行してから、チューニングやスケールアウトに進みます。

トレーニングループの基本

ディープラーニングでは、すべてのトレーニングループで、順伝播、損失計算、逆伝播、オプティマイザーステップ、スケジューラーステップ、勾配のゼロ化、指標の記録、検証、チェックポイントの保存、早期停止を明示してください。

AMP、勾配累積、勾配クリッピングを含む本番環境向け PyTorch トレーニングループの概要

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.cuda.amp import autocast, GradScaler

def train_one_epoch(
    model: nn.Module, 
    dataloader: DataLoader, 
    optimizer: torch.optim.Optimizer, 
    criterion: nn.Module, 
    device: torch.device,
    scaler: GradScaler,
    scheduler = None,
    grad_accum_steps: int = 1,
    max_grad_norm: float = 1.0
):
    model.train()
    optimizer.zero_grad(set_to_none=True)
    total_loss = 0.0
    
    for step, (inputs, targets) in enumerate(dataloader):
        inputs, targets = inputs.to(device, non_blocking=True), targets.to(device, non_blocking=True)
        
        # Mixed precision forward pass
        with autocast(dtype=torch.float16): # or torch.bfloat16 if hardware supports it
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            # Scale loss for gradient accumulation
            loss = loss / grad_accum_steps
            
        # Scaled backpropagation
        scaler.scale(loss).backward()
        
        if (step + 1) % grad_accum_steps == 0 or (step + 1) == len(dataloader):
            # Unscales gradients before clipping
            scaler.unscale_(optimizer)
            # Gradient clipping to prevent exploding gradients
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
            
            # Step optimizer and update scale factor
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad(set_to_none=True)
            
            if scheduler is not None:
                scheduler.step()
                
        total_loss += loss.item() * grad_accum_steps
        
    return total_loss / len(dataloader)

@torch.inference_mode()
def validate(model: nn.Module, dataloader: DataLoader, criterion: nn.Module, device: torch.device):
    model.eval()
    total_loss = 0.0
    correct = 0
    total = 0
    
    for inputs, targets in dataloader:
        inputs, targets = inputs.to(device, non_blocking=True), targets.to(device, non_blocking=True)
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        total_loss += loss.item()
        
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
        
    return total_loss / len(dataloader), correct / total

TensorFlow/Keras では、コールバックと分散ストラテジーで十分な場合は、組み込みの fit を優先してください。独自の制御が必要な場合は tf.GradientTape を使用します。JAX/Flax では、状態を明示的に保ち、純粋な更新関数を使用し、ドロップアウト、拡張、サンプリングにはそれぞれ別の PRNG キーを割り当てます。

古典的な機械学習では、交差検証で各変換がトレーニング用フォールドだけに適用されるよう、前処理をパイプラインオブジェクト内に構築してください。

from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.impute import SimpleImputer
from xgboost import XGBClassifier

numeric_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='median')),
    ('scaler', StandardScaler())
])

categorical_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='most_frequent')),
    ('onehot', OneHotEncoder(handle_unknown='ignore'))
])

preprocessor = ColumnTransformer(transformers=[
    ('num', numeric_transformer, numeric_cols),
    ('cat', categorical_transformer, categorical_cols)
])

clf_pipeline = Pipeline(steps=[
    ('preprocessor', preprocessor),
    ('classifier', XGBClassifier(use_label_encoder=False, eval_metric='logloss'))
])

オプティマイザー、スケジュール、正則化

重み減衰を分離する必要がある場合、ディープラーニングでは AdamW を堅実なデフォルトとして使用します。一部のビジョン領域ではモーメンタム付き SGD を使用します。メモリの制約がある場合に限り、Adafactor または 8-bit オプティマイザーを使用します。Transformer 系のモデルでは、アーキテクチャのドキュメントで別途指定されていない限り、バイアスと正規化パラメーターを重み減衰の対象から除外してください。

学習率の戦略は、通常、オプティマイザーの細かな調整より重要です。一般的なデフォルトは次のとおりです。

• Transformer と大規模バッチのトレーニングには、ウォームアップとコサイン減衰を組み合わせます。
• 試験的なトレーニングには、One Cycle またはコサイン・リスタートを使用します。
• 収束が不安定な小規模の教師ありタスクには、性能停滞時に学習率を下げます。
• グローバルバッチを増やす場合は、実効バッチサイズに応じて線形にウォームアップします。

過学習は最終指標だけでなく、トレーニング曲線と検証曲線で検出します。データ拡張、ドロップアウト、ラベルスムージング、重み減衰、早期停止、MixUp/CutMix、より強固な検証分割、小規模なモデルを使用してください。過少適合の場合は、複雑さを加える前に、ラベル、特徴量のシグナル、損失のスケーリング、オプティマイザーの学習率、モデル容量、前処理のバグを確認します。

バッチサイズ、勾配累積、精度

実効バッチサイズは、デバイスごとのバッチサイズにデバイス数と勾配累積ステップ数を掛けたものです。検証品質が安定する場合に限り、スループット向上のためにバッチサイズを増やします。メモリが限られている場合は、勾配累積、活性化チェックポイント、シーケンスの短縮、画像の縮小、バケット化、パラメーター効率の高いファインチューニング、オプティマイザー状態のシャーディング、オフロードを使用します。

最新のアクセラレーターでは混合精度を使用します。ネイティブの bf16 に対応したハードウェアでは、損失スケーリングに関する問題を減らせるため bf16 を優先します。それ以外では、動的損失スケーリングを用いた fp16 AMP を使用します。必要に応じて、一部のリダクション、ソフトマックス／ロジットの処理、指標の累積など、数値的に繊細な処理は fp32 で行います。NaN、オーバーフロー、アンダーフロー、不安定な正規化に注意してください。

分散トレーニングのパターン

ボトルネックを解消できる、最もシンプルな分散モードを選択します。

パターン	使用する状況	注意点
データ並列／DDP	モデルが1台のデバイスに収まり、スループットを上げたい場合	グローバルバッチ、サンプラーのシャーディング、メトリクスの同期を適切に設定する必要があります
FSDP／ZeRO	パラメーターやオプティマイザーの状態をシャーディングしないとモデルが収まらない場合	チェックポイントが複雑になります。ラッピングとオフロードを調整してください
テンソル並列	個々のレイヤーが1台のデバイスに収まらない場合、または大規模なトランスフォーマーを扱う場合	フレームワーク／ランタイムのサポートが必要で、通信負荷が高くなります
パイプライン並列	非常に深いモデルをレイヤーごとに分割する必要がある場合	バブルによるオーバーヘッドが発生します。マイクロバッチのスケジューリングが重要です
エキスパート並列／モデル並列	混合エフェクトモデルや巨大なモデルを扱う場合	ルーティングと負荷分散が複雑になります

PyTorch DDP／FSDP 初期化のひな形

import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

def setup_distributed():
    # Environment variables set by torchrun
    world_size = int(os.environ["WORLD_SIZE"])
    rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])
    
    dist.init_process_group(
        backend="nccl", 
        init_method="env://", 
        world_size=world_size, 
        rank=rank
    )
    torch.cuda.set_device(local_rank)
    return rank, local_rank, world_size

def cleanup():
    dist.destroy_process_group()

# Usage DDP:
# rank, local_rank, _ = setup_distributed()
# model = MyModel().to(local_rank)
# model = DDP(model, device_ids=[local_rank])

# Usage FSDP:
# rank, local_rank, _ = setup_distributed()
# model = MyModel().to(local_rank)
# model = FSDP(model)

DeepSpeed JSON 設定のひな形（deepspeed_config.json）

{
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto",
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 5e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 5e8,
    "contiguous_gradients": true,
    "cpu_offload": false
  },
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "initial_scale_power": 16,
    "hysteresis": 2,
    "min_loss_scale": 1
  },
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": "auto",
      "betas": [0.9, 0.999],
      "eps": 1e-8,
      "weight_decay": "auto"
    }
  }
}

チェックポイントとフォールトトレランス

復旧可能なチェックポイントには、モデルの重み、オプティマイザーの状態、スケジューラーの状態、スケーラーの状態、エポック／ステップ、可能な場合は RNG の状態、トークナイザー／前処理器、ラベルの対応表、設定、メトリクスの履歴、コードまたはコミット識別子を含めます。検証データに基づくベストチェックポイントと最新のチェックポイントは、別々に保存します。長時間のトレーニングを始める前に、再開できることをテストしてください。成果物は、揮発性のワーカーディスクではなく、永続ストレージに保存します。

PyTorch の堅牢なチェックポイントの保存と読み込み

def save_checkpoint(state, is_best, checkpoint_dir="checkpoints"):
    os.makedirs(checkpoint_dir, exist_ok=True)
    filepath = os.path.join(checkpoint_dir, "last_checkpoint.pt")
    torch.save(state, filepath)
    if is_best:
        shutil.copyfile(filepath, os.path.join(checkpoint_dir, "best_model.pt"))

def load_checkpoint(filepath, model, optimizer=None, scheduler=None, scaler=None):
    checkpoint = torch.load(filepath, map_location="cpu")
    model.load_state_dict(checkpoint["state_dict"])
    if optimizer is not None and "optimizer" in checkpoint:
        optimizer.load_state_dict(checkpoint["optimizer"])
    if scheduler is not None and "scheduler" in checkpoint:
        scheduler.load_state_dict(checkpoint["scheduler"])
    if scaler is not None and "scaler" in checkpoint:
        scaler.load_state_dict(checkpoint["scaler"])
    return checkpoint.get("epoch", 0), checkpoint.get("best_metric", float("inf"))

不均衡なデータとノイズの多いデータ

クラス不均衡に対しては、まず目的を適切に反映するメトリクスを選択します。PR-AUC、F1 の各種指標、バランス精度、適合率を指定した再現率、コスト重み付き効用、またはキャリブレーションなどがあります。クラス重み、フォーカル損失、バランスサンプリング、しきい値の調整、またはデータ収集を活用します。分割前のオーバーサンプリングは避けてください。ラベルノイズに対しては、損失の大きいサンプルを調査し、ロバストな損失関数を使用し、重要なサンプルのラベルを付け直し、不確実性を報告します。

スケーリングと高度なトレーニング

スケーリング則は実践的に適用します。スケールを拡大する前に、品質がデータ、モデル、計算リソースのどれに制約されているかを見積もります。損失を計算量、データサイズ、パラメーター数に対して追跡します。大規模モデルでは、カリキュラム学習、データ品質フィルタリング、重複排除、シーケンスパッキング、トークン予算の管理、計算量に対して最適なトレードオフを検討します。RLHF または強化学習では、オンライン最適化の前に、報酬の品質、オフライン評価、安全性の制約、分布シフトを検証します。

よくある失敗の診断

• トレーニング損失が減少しない：ラベルの誤り、パラメーターの凍結、学習率の不適切さ、損失関数と活性化関数の不一致、dtype のオーバーフロー、train() モードの未設定、入力パイプラインの破損。
• 検証結果がトレーニング結果より大幅に悪い：データリークのない分割によって、過学習、分布シフト、データ拡張の不一致、または正則化の不足が明らかになっている可能性があります。
• NaN：学習率が高すぎる、fp16 のオーバーフロー、無効な入力、不安定な損失、ゼロ除算、勾配の爆発。
• トレーニングが遅い：データローダーのボトルネック、小さなバッチ、CPU での変換、過剰なログ記録、同期、精度設定の未最適化、ネットワークストレージ。
• 結果を再現できない：データ／環境のバージョンが固定されていない、非決定的なカーネル、ランダムなデータ拡張、分散サンプラーの違い、隠れた前処理状態。

参考資料

• PyTorch 分散処理および AMP のドキュメント：https://pytorch.org/docs/stable/
• TensorFlow のトレーニングおよび分散処理ガイド：https://www.tensorflow.org/guide
• JAX および Flax のドキュメント：https://jax.readthedocs.io/ and https://flax.readthedocs.io/
• Hugging Face Transformers、Accelerate、PEFT のドキュメント：https://huggingface.co/docs
• scikit-learn のモデル評価およびパイプラインのドキュメント：https://scikit-learn.org/stable/user_guide.html