feat: 量子化チェックポイント/レジューム機能の追加

[sotanengel]

概要

長時間かかる量子化処理が途中で中断された場合に、最初からやり直さずにチェックポイントから再開できる機能を提案します。

背景・動機

Llama-70B などの大規模モデルを GPTQ で量子化する場合、単一 GPU では数時間かかることがあります。この間にプロセスが OOM・ネットワーク障害・クラッシュ等で終了すると、すべての計算結果が失われます。チェックポイント/レジューム機能があれば、最後に保存したブロック境界から処理を再開できるため、コストと時間を大幅に削減できます。

提案内容

新規クラス CheckpointConfig

from onecomp import CheckpointConfig, Runner, ModelConfig
from onecomp.quantizer.gptq import GPTQ

runner = Runner(
    model_config=ModelConfig(model_id="meta-llama/Llama-2-7b-hf"),
    quantizer=GPTQ(wbits=4, groupsize=128),
    checkpoint_config=CheckpointConfig(
        checkpoint_dir="./ckpt",
        interval_blocks=1,   # N ブロックごとに保存
        resume=True,         # 既存チェックポイントがあれば自動レジューム
    ),
)
runner.run()

CLI オプション

# チェックポイントを有効化して量子化
onecomp meta-llama/Llama-2-7b-hf \
    --wbits 4 \
    --checkpoint-dir ./ckpt \
    --checkpoint-interval-blocks 1

# 中断後に再開（同じコマンドを再実行するだけ）
onecomp meta-llama/Llama-2-7b-hf \
    --wbits 4 \
    --checkpoint-dir ./ckpt

# チェックポイントがあっても最初からやり直す場合
onecomp meta-llama/Llama-2-7b-hf --wbits 4 --checkpoint-dir ./ckpt --no-resume

実装詳細

コンポーネント	変更内容
onecomp/checkpoint.py (新規)	CheckpointConfig データクラス・CheckpointManager の実装
onecomp/runner.py	checkpoint_config パラメータ追加、レジューム時のモデル重み適用
onecomp/cli.py	--checkpoint-dir・--checkpoint-interval-blocks・--no-resume オプション追加
onecomp/quantizer/_quantizer.py	on_layer_quantized コールバックフックの追加
onecomp/runner_methods/chunked_quantization.py	レジューム済みレイヤーのスキップ対応

チェックポイントの仕組み

• Transformer ブロック境界ごとに checkpoint_<NNNN>.pt（量子化結果）と checkpoint_meta.json（互換性情報）を保存
• 古いチェックポイントファイルは自動削除（最新 + checkpoint_final.pt のみ保持）
• レジューム時は、完了済みレイヤーの重みを脱量子化してモデルに適用することで、後続レイヤーの活性化が正しく計算される
• モデル ID・量子化器クラス・主要パラメータの互換性チェックつき

サポート範囲

• 通常のキャリブレーション付き量子化
• チャンク分割キャリブレーション（calibration_config.batch_size 使用時）
• キャリブレーションなし量子化
• Llama・Qwen・Gemma など主要アーキテクチャのレイヤー名パターンに対応

参考実装

実装済みのプロトタイプが以下のブランチにあります:
sotanengel/OneCompression の feature/checkpoint