# 機能設計書 67-分散ストラテジーAPI

## 概要

本ドキュメントは、TensorFlowにおけるMirroredStrategy、MultiWorkerMirroredStrategy、TPUStrategyなどの分散学習戦略を提供するtf.distribute.Strategy APIの設計を記述する。

### 本機能の処理概要

分散ストラテジーAPIは、TensorFlowの分散学習を抽象化するフレームワークである。tf.distribute.Strategyベースクラスを中心に、データ並列学習のためのモデル複製、勾配集約、変数同期の仕組みを提供する。ユーザーは同じモデル定義コードを異なるストラテジーと組み合わせることで、単一GPU、マルチGPU、マルチワーカー、TPUなど様々な分散環境に対応できる。

**業務上の目的・背景**：大規模データセットや大規模モデルの学習には複数のアクセラレータを活用した分散学習が不可欠である。しかし、分散学習のプログラミングは複雑であり、デバイス配置、データ分割、勾配集約、変数同期など多くの技術的課題がある。分散ストラテジーAPIはこれらの複雑さをフレームワークレベルで抽象化し、ユーザーが最小限のコード変更で分散学習を実現できるようにする。

**機能の利用シーン**：単一マシンのマルチGPU学習、複数マシンにまたがるマルチワーカー学習、TPUクラスタでの学習、Keras model.fitでの自動分散、カスタムトレーニングループでの分散。

**主要な処理内容**：
1. Strategy.scope()によるコンテキスト管理（変数作成の制御）
2. Strategy.run()によるレプリカ単位の関数実行
3. データセットの分散配布（experimental_distribute_dataset等）
4. レプリカコンテキストとクロスレプリカコンテキストの管理
5. DistributedValues（PerReplica、Mirrored等）の管理
6. 勾配の集約（reduce、all_reduce）

**関連システム・外部連携**：Keras API、tf.data Dataset API、チェックポイントシステム、SavedModelシステム。

**権限による制御**：特になし。

## 関連画面

| 画面No | 画面名 | 関連種別 | 関連する操作・処理 |
|--------|--------|----------|------------------|
| - | - | - | 画面機能マッピングに該当なし |

## 機能種別

計算処理（分散制御フレームワーク）

## 入力仕様

### 入力パラメータ

| パラメータ名 | 型 | 必須 | 説明 | バリデーション |
|-------------|-----|-----|------|---------------|
| fn (run時) | Callable | Yes | レプリカで実行する関数 | callableであること |
| args (run時) | tuple | No | fnに渡す引数 | - |
| kwargs (run時) | dict | No | fnに渡すキーワード引数 | - |
| reduce_op (reduce時) | ReduceOp | Yes | 集約操作（SUM/MEAN） | ReduceOpインスタンス |
| dataset (分散配布時) | tf.data.Dataset | Yes | 分散するデータセット | Datasetインスタンス |

### 入力データソース

Python関数、テンソル、tf.data.Dataset。

## 出力仕様

### 出力データ

| 項目名 | 型 | 説明 |
|--------|-----|------|
| run結果 | PerReplica/Mirrored | 各レプリカでの実行結果 |
| reduce結果 | Tensor/Mirrored | 集約後のテンソル |
| 分散Dataset | DistributedDataset | 分散されたデータセット |

### 出力先

Python上のDistributedValuesオブジェクト。

## 処理フロー

### 処理シーケンス

```
1. Strategyのインスタンス化
   └─ 利用可能なデバイスの検出とExtendedの初期化
2. strategy.scope() でコンテキスト入場
   └─ 変数作成がMirroredVariable等になるよう設定
3. モデルとオプティマイザの構築（scope内）
   └─ 変数は自動的に全レプリカに複製される
4. データセットの分散配布
   └─ experimental_distribute_dataset でデバイスごとにプリフェッチ
5. strategy.run(fn, args) で学習ステップ実行
   └─ 5-1. 入力のアンラップ（DistributedValues → 各レプリカの値）
   └─ 5-2. 各レプリカでfnを実行
   └─ 5-3. 結果のマージ（PerReplicaとして返却）
6. 勾配の集約
   └─ replica_context.all_reduce または strategy.reduce で集約
```

### フローチャート

```mermaid
flowchart TD
    A[Strategy インスタンス化] --> B[strategy.scope 入場]
    B --> C[モデル・オプティマイザ構築]
    C --> D[データセット分散配布]
    D --> E[strategy.run で学習ステップ実行]
    E --> F[各レプリカでfn実行]
    F --> G[結果のマージ PerReplica]
    G --> H[勾配の集約 all_reduce]
    H --> I[変数更新]
    I --> E
```

## ビジネスルール

### 業務ルール

| ルールNo | ルール名 | 内容 | 適用条件 |
|---------|---------|------|---------|
| BR-67-01 | コンテキスト制約 | Strategy.reduce等のメソッドはクロスレプリカコンテキストでのみ呼び出し可能 | 常時 |
| BR-67-02 | scope入れ子禁止 | 異なるStrategyのscope()を入れ子にすることはできない | 常時 |
| BR-67-03 | MirroredVariable同期 | MirroredVariableはON_WRITE同期で全レプリカに同じ更新が適用される | 同期変数の場合 |
| BR-67-04 | SyncOnRead集約 | SyncOnReadVariableはクロスレプリカコンテキストでの読み取り時に集約される | ON_READ同期変数の場合 |
| BR-67-05 | デフォルトストラテジー | スコープ外ではデフォルトの単一レプリカストラテジーが使用される | strategy.scope()外 |

### 計算ロジック

該当なし。

## データベース操作仕様

### 操作別データベース影響一覧

データベース操作は行わない。

## エラー処理

### エラーケース一覧

| エラーコード | エラー種別 | 発生条件 | 対処方法 |
|------------|----------|---------|---------|
| RuntimeError | コンテキストエラー | クロスレプリカコンテキスト外でStrategy.reduceを呼び出し | strategy.scope()内で呼び出す |
| RuntimeError | ストラテジー混在エラー | 異なるStrategyのscope()内でメソッドを呼び出し | 正しいStrategyのscope()を使用 |

### リトライ仕様

該当なし。

## トランザクション仕様

該当なし。

## パフォーマンス要件

- レプリカ間の通信オーバーヘッドを最小化する設計
- batch_reduceで複数のreduceをバッチ化して通信効率を向上
- デフォルトストラテジーの場合、オーバーヘッドは最小限

## セキュリティ考慮事項

マルチワーカー環境ではワーカー間通信のセキュリティ設定（TLS等）が必要。

## 備考

- `distribution_strategy_gauge`と`distribution_strategy_replica_gauge`でストラテジーの使用状況がモニタリングされる
- スレッドローカルの`_update_replica_id`でUpdate Contextが管理される

---

## コードリーディングガイド

本機能を理解するために参照すべきファイルと、推奨する読み解き順序を以下に示す。

### 推奨読解順序

#### Step 1: データ構造を理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 1-1 | distribute_lib.py | `tensorflow/python/distribute/distribute_lib.py` | モジュールドキュメント（16-189行目）：分散学習の用語と概念 |
| 1-2 | distribute_lib.py | `tensorflow/python/distribute/distribute_lib.py` | `_ThreadMode`クラス（346-349行目）：スレッドごとのコンテキスト管理 |
| 1-3 | distribute_lib.py | `tensorflow/python/distribute/distribute_lib.py` | `UpdateContext`クラス（254-269行目）：更新コンテキスト |

**読解のコツ**: distribute_lib.pyの冒頭には非常に詳細な用語集（Data parallelism、Replica、Host、Worker等）があり、これを理解することが分散学習機能全体の理解に不可欠。

#### Step 2: コンテキスト管理APIを理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 2-1 | distribute_lib.py | `tensorflow/python/distribute/distribute_lib.py` | `get_replica_context`関数（453-498行目） |
| 2-2 | distribute_lib.py | `tensorflow/python/distribute/distribute_lib.py` | `get_strategy`関数（542-559行目） |
| 2-3 | distribute_lib.py | `tensorflow/python/distribute/distribute_lib.py` | `has_strategy`関数（562-575行目） |

**主要処理フロー**:
- **246-251行目**: `get_update_replica_id`でスレッドローカルのレプリカIDを取得
- **276-293行目**: `_require_cross_replica_or_default_context_extended`でコンテキスト検証
- **308-320行目**: `require_replica_context`でレプリカコンテキスト検証

#### Step 3: Strategyベースクラス

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 3-1 | distribute_lib.py | `tensorflow/python/distribute/distribute_lib.py` | `Strategy`/`StrategyBase`クラスの定義 |

### プログラム呼び出し階層図

```
tf.distribute.Strategy
    │
    ├─ scope()  [コンテキスト管理]
    │      └─ _ThreadMode 設定
    │
    ├─ run(fn, args)  [レプリカ実行]
    │      ├─ DistributedValues のアンラップ
    │      ├─ 各レプリカでfn実行
    │      └─ 結果のマージ (PerReplica)
    │
    ├─ reduce(reduce_op, value)  [集約]
    │      └─ CrossDeviceOps.reduce()
    │
    ├─ experimental_distribute_dataset(dataset)
    │      └─ DistributedDataset 構築
    │
    └─ StrategyExtended  [内部実装]
           ├─ _create_variable()
           ├─ reduce_to()
           └─ batch_reduce_to()
```

### データフロー図

```
[入力]                          [処理]                              [出力]

Python fn + Dataset ─────────▶ strategy.scope() ──▶ MirroredVariables
                                │
Dataset ─────────────────────▶ distribute_dataset ──▶ DistributedDataset
                                │
DistributedDataset ──────────▶ strategy.run(fn) ──▶ PerReplica results
                                │
PerReplica gradients ────────▶ all_reduce ──▶ Mirrored gradients
                                │
Mirrored gradients ──────────▶ optimizer.apply ──▶ Updated MirroredVariables
```

### 関連ファイル一覧

| ファイル | パス | 種別 | 役割 |
|---------|------|------|------|
| distribute_lib.py | `tensorflow/python/distribute/distribute_lib.py` | ソース | Strategyベースクラスとコンテキスト管理 |
| values.py | `tensorflow/python/distribute/values.py` | ソース | DistributedValues（PerReplica、Mirrored等） |
| reduce_util.py | `tensorflow/python/distribute/reduce_util.py` | ソース | ReduceOp定義 |
| distribute_utils.py | `tensorflow/python/distribute/distribute_utils.py` | ソース | 分散ユーティリティ |
| input_lib.py | `tensorflow/python/distribute/input_lib.py` | ソース | 分散データセット入力 |
| device_util.py | `tensorflow/python/distribute/device_util.py` | ソース | デバイスユーティリティ |
