パラメータ

• context_parallel_config (ContextParallelConfig, 任意) — コンテキスト並列化の設定です。
• tensor_parallel_config (TensorParallelConfig, 任意) — テンソル並列化の設定です。

パラメータ

• ring_degree (int、任意、デフォルトは 1) — Ring Attention に使用するデバイス数です。シーケンスはデバイス間で分割されます。各デバイスは、リングを順に回されて渡されるローカルな Q と KV チャンク間のアテンションを計算します。メモリ使用量は低くなります（一度に KV の 1/N しか保持しない）、計算と通信をオーバーラップできますが、すべてのトークンを見るには N 回の反復が必要です。メモリや帯域幅が限られている長いシーケンスに最適です。コンテキスト並列領域内のリングアテンションに使用するデバイス数です。コンテキスト並列メッシュ内のデバイス総数の約数である必要があります。
• ulysses_degree (int、任意、デフォルトは 1) — Ulysses Attention に使用するデバイス数です。シーケンスはデバイス間で分割されます。各デバイスがローカル QKV を計算し、その後、すべての KV チャンクを all-gather して、1 パスで完全なアテンションを計算します。メモリ使用量が高くなります（すべての KV を保存するため）、高帯域の all-to-all 通信が必要ですが、レイテンシは低くなります。インターコネクト帯域が十分に確保できる中程度の長さのシーケンスに最適です。
• convert_to_fp32 (bool、省略可能、デフォルトは True) — リングアテンションの数値的安定性を確保するために、出力と LSE を float32 に変換するかどうかを指定します。
• rotate_method (str、任意、デフォルトは "allgather") — リングアテンションにおいて、デバイス間でキー/バリュー状態をローテーションするために使用するメソッド。現在サポートされているのは "allgather" のみです。
• ulysses_anything (bool、省略可能、デフォルトは False) — ulysses_degree で割り切れない任意のシーケンス長やヘッド数に対応する「Ulysses Anything」モードを有効にするかどうかを指定します。有効にした場合、ulysses_degree は 1 より大きく、ring_degree は 1 でなければなりません。
• ring_anything (bool、省略可能、デフォルトは False) — 任意のシーケンス長をサポートする「Ring Anything」モードを有効にするかどうか。有効にした場合、ring_degree は 1 より大きく、ulysses_degree は 1 でなければなりません。
• mesh (torch.distributed.device_mesh.DeviceMesh, 任意) — コンテキスト並列性で使用するカスタムのデバイスメッシュです。指定された場合、新しいメッシュを作成する代わりに、このメッシュが使用されます。これは、コンテキスト並列性を、同じデバイスメッシュを共有する他の並列化戦略（FSDP やテンソル並列性など）と組み合わせる場合に有用です。メッシュには「ring」と「ulysses」の両方の次元が必要です。使用しない次元にはサイズ 1 を指定します（例: FSDP を使用したリングアテンションのみの場合、mesh_shape=(2, 1, 4) と mesh_dim_names=("ring", "ulysses", "fsdp") を指定します）。

パラメータ

• tp_degree (int、デフォルトは 1) — シャードに分割するデバイスの数です。並列化対象のモデルのアテンションヘッド数（および FFN の隠れ層次元数）の約数である必要があります。
• mesh (torch.distributed.device_mesh.DeviceMesh、省略可) — 使用するカスタムのデバイスメッシュです。指定された場合、tp_degree は mesh.size() から推論され、この引数は無視されます。同じメッシュを共有する他の並列化戦略（例: CP）と TP を組み合わせる場合に有用です。

パラメータ

• model (torch.nn.Module) — インプレースでシャードするモデル。
• config (TensorParallelConfig) — テンソル並列の構成です。そのデバイスメッシュは、config.setup(...) で既に設定されている必要があります。
• tp_plan (dict) — モジュール名のグロブから"colwise"/"rowwise"スタイル(またはパックされたバリアント)へのフラットなマッピング。通常はモデルの_tp_planです。
• weights_already_sharded (bool、デフォルトはFalse) — 計画されたパラメータが、ストリーミングによるfrom_pretrained読み込み後にそうなるように、すでにDTensorのシャードであるかどうかを示します。Trueの場合、フォワードフックのみが登録されます。これは検出ではなく明示的に渡す必要があります。チェックポイントに存在しない計画されたパラメータは依然としてメタテンソルであり、実際には半分がシャード化されているモデルに対して「シャード化されていない」と誤って検出されてしまうためです。