「非決定モード → 決定モード」に切り替える例

 

前提

• Python + PyTorch 環境がある（GPU 使用可能、CUDA 有効）。

• thinking-machines-lab/batch_invariant_ops ライブラリをインストール済みであること。 GitHub

• モデルが FlashAttention など高速な Attention バックエンドを使う場合、Attention 部分を固定戦略 (fixed splits, 固定チャンク・KV キャッシュの統一レイアウト等) に対応できることが望ましい。

---

セットアップ

まず、batch_invariant_ops を使い始めるための準備。

# リポジトリをクローン or pip 経由でインストール
git clone https://github.com/thinking-machines-lab/batch_invariant_ops.git
cd batch_invariant_ops
pip install -e .

または、PyPI で公開されていれば:

pip install batch_invariant_ops

---

サンプルコード：非決定モード → 決定モード切り替え

以下は、簡単な Transformer モデル（あるいは vLLM 等）で、標準モードとバッチ不変モードを切り替えて推論を行い、応答が一貫するかを比較する例です。

import torch
from batch_invariant_ops import set_batch_invariant_mode
# 仮に FlashAttention を使うライブラリ
from your_model_lib import TransformerModel, FlashAttentionBackend

def inference(model, input_ids, attention_mask):
    """
    単純な推論関数
    """
    # モデルの順伝播 (推論モード)
    with torch.no_grad():
        logits = model(input_ids=input_ids, attention_mask=attention_mask)
    return logits

def test_determinism(input_ids, attention_mask, runs=5, batch_mode=False):
    """
    同じ入力で複数回推論して結果が一致するかどうかチェックする
    batch_mode=False: 標準モード（非決定モード）
    batch_mode=True: バッチ不変モードを有効にした状態
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = TransformerModel(backend=FlashAttentionBackend()).to(device)
    input_ids = input_ids.to(device)
    attention_mask = attention_mask.to(device)

   _outputs = []
    for i in range(runs):
        if batch_mode:
            # 決定性モードを有効にする
            with set_batch_invariant_mode():
                out = inference(model, input_ids, attention_mask)
        else:
            out = inference(model, input_ids, attention_mask)
        _outputs.append(out.cpu())

    # 比較
    first = _outputs[0]
    all_same = all(torch.allclose(first, o, atol=0, rtol=0) for o in _outputs[1:])
    return all_same, _outputs

if __name__ == "__main__":
    # 適当な入力を作成（例：バッチサイズ 2、シーケンス長 128 等）
    batch_size = 2
    seq_len = 128
    vocab_size = 30522  # 例
    # ダミー input_ids, attention_mask
    input_ids = torch.randint(low=0, high=vocab_size, size=(batch_size, seq_len), dtype=torch.long)
    attention_mask = torch.ones((batch_size, seq_len), dtype=torch.long)

    # 標準モードでテスト
    same_std, outs_std = test_determinism(input_ids, attention_mask, runs=3, batch_mode=False)
    print(f"Standard mode deterministic? {same_std}")

    # バッチ不変モードでテスト
    same_bi, outs_bi = test_determinism(input_ids, attention_mask, runs=3, batch_mode=True)
    print(f"Batch‐invariant mode deterministic? {same_bi}")

---

注意点・補足実装

この切り替えを実用レベルで使うには、以下の点を押さえておく必要があります。

Attention / FlashAttention に関する扱い

• FlashAttention や類似バックエンドでは、Attention の計算中にチャンク分割 (chunking) や Split‑KV, カーネル最適化などが入っていることが多い。これらの最適化戦略がバッチサイズやシーケンス長、KV キャッシュ状態等に応じて動的に変化する場合、非決定性の原因となる。

• 決定モードにおいては、チャンク／分割サイズを固定する、KV キャッシュの内部レイアウトを一貫させるなどの設計が必要。

例えば：

# 仮想コードスニペット：FlashAttention バックエンドの Attention 呼び出しを修正する例

class FixedAttention(FlashAttentionBackend):
    def forward(self, query, key, value, mask=None):
        # 固定チャンクサイズと分割戦略を明示
        chunk_size = 64  # 例：固定
        # KV キャッシュがあるかどうかに関わらず、キャッシュと現在入力を統一したレイアウトで結合
        if hasattr(self, 'kv_cache') and self.kv_cache is not None:
            # レイアウト統一処理
            key = torch.cat([self.kv_cache['key'], key], dim=1)
            value = torch.cat([self.kv_cache['value'], value], dim=1)
        # 固定チャンク毎に処理
        # 注意：この部分はライブラリ／GPU カーネル設定によっては細部が異なる
        out = super().attend(query, key, value, mask, chunk_size=chunk_size)
        return out

他の演算子（RMSNorm, MatMul, mean/reduction etc.）

• batch_invariant_ops ライブラリでは、torch.mm(), torch.addmm(), torch.log_softmax(), torch.mean() など、標準的な行列演算や削減 (reduction) 演算をバッチ不変なカーネルに置き換える仕組みが用意されている。 GitHub

• 例：torch.mean() を使うとき、通常はバッチ軸または特徴軸での reduce の順序がバッチサイズ等に依存することがあるので、決定モードではその順序を固定するか、reduce を分割しない／一定の分割サイズでのみ行うようにする。

---

テスト・検証

決定性モードを導入したあと、次のようなテストを行うと良い：

1. 繰り返し推論テスト：同じ入力を多回（例 10～100 回）推論し、それぞれの出力ロジットがビット単位で一致するかをチェック。

2. バッチサイズ変更テスト：入力のバッチサイズを変えても、個々の要素の出力が一致するかどうか（例、バッチサイズ = 1 vs バッチサイズ = N）を比較。

3. シーケンス長・KV キャッシュ状態のテスト：アテンションのプリフィル／デコード／キャッシュ未使用・使用時で、出力が変わらないか確認。

4. 異なる GPU / 同じモデルでも異なるカードでのテスト：異なるハードウェア構成で同じコード・ライブラリを使ったときの再現性を確認。

5. 性能測定：標準モードと決定性モードの速度・メモリ消費・レイテンシを比較し、「どの程度のトレードオフがあるか」を把握。

---

トレードオフと実用上の判断基準

• 決定性モードを有効にすると、特に小バッチや短シーケンス入力で GPU のコア利用効率が下がることがあり、レイテンシが増える可能性が高い。

• 通常のモデル提供時には、応答の多様性や速度が重視されるケースもあるので、用途によって標準モード／決定モードを切り替えるオプションを持たせるのが望ましい。

• モデルが大きく・入力が長くなるほど、非決定性が蓄積されやすいため、決定モードの効果を実感しやすいが、そのぶんオーバーヘッドも出やすい。