分散推論

分散推論は、ワークロードを複数のGPUに分散して処理します。これは、より大きなモデルをメモリに収めるための有用な技術であり、複数のプロンプトを処理することでスループットを向上させることができます。

このガイドでは、AccelerateとPyTorch Distributedを使用して分散推論を行う方法を紹介します。

Accelerate

Accelerate は、セットアップを処理することで、複数のアクセラレーター上での推論と学習を簡素化するために設計されたライブラリであり、ユーザーが PyTorch コードに集中できるようにします。

次のコマンドでAccelerateをインストールします。

uv pip install accelerate

Pythonファイル内でaccelerate.PartialStateクラスを初期化して、分散環境を作成します。accelerate.PartialStateクラスは、プロセス管理、デバイスの制御と分散、プロセス間の調整を管理します。

DiffusionPipeline を accelerate.PartialState.device に移動し、各プロセスに GPU を割り当てます。

import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.float16
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)

split_between_processes ユーティリティをコンテキストマネージャーとして使用し、プロンプトをプロセス数に自動で分配します。

with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
    result = pipeline(prompt).images[0]
    result.save(f"result_{distributed_state.process_index}.png")

accelerate launch を呼び出してスクリプトを実行し、--num_processes 引数を使用して使用する GPU の数を設定します。

accelerate launch run_distributed.py --num_processes=2

複数GPUで推論を実行するための最小限の例として、スクリプトを参照してください。詳細については、Distributed Inference with 🤗 Accelerateガイドをご覧ください。

PyTorch Distributed

PyTorch の DistributedDataParallel は データ並列性 を有効にし、各デバイスに同じモデルを複製して、異なるバッチのデータを並列で処理します。

Pythonファイルにtorch.distributedとtorch.multiprocessingをインポートし、分散プロセスグループを設定して、各GPUでの推論用のプロセスを生成します。

import torch
import torch.distributed as dist
import torch.multiprocessing as mp

from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.float16,
)

init_process_group を用いた推論用の関数を作成します。このメソッドは、バックエンドの種類、現在のプロセスのrank、および参加するプロセス数であるworld_size（例えば、2つのGPUの場合はworld_size=2）を指定して、分散環境を作成します。

パイプラインをrankに移動し、get_rankを使用して各プロセスにGPUを割り当てます。各プロセスは異なるプロンプトを処理します。

def run_inference(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

    pipeline.to(rank)

    if torch.distributed.get_rank() == 0:
        prompt = "a dog"
    elif torch.distributed.get_rank() == 1:
        prompt = "a cat"

    image = sd(prompt).images[0]
    image.save(f"./{'_'.join(prompt)}.png")

mp.spawn を使用して、world_size で定義された数のプロセスを作成します。

def main():
    world_size = 2
    mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)


if __name__ == "__main__":
    main()

torchrun を呼び出して推論スクリプトを実行し、--nproc_per_node 引数を使用して使用する GPU の数を設定します。

torchrun --nproc_per_node=2 run_distributed.py

device_map

device_map 引数を使用すると、モデルのコンポーネントを自動的に複数の GPU に配置することで、分散推論を有効にできます。これは、モデルが単一の GPU に収まらない場合に特に有用です。以下の例（2 つの GPU が利用可能であると仮定）に示すように、device_map を使用して、特定の段階で必要なモデルコンポーネントを選択的にロードおよびアンロードできます。

device_map="balanced" を設定すると、テキストエンコーダーが利用可能なすべての GPU に均等に分散配置されます。各テキストエンコーダーに割り当てる最大メモリ量を指定するには、max_memory 引数を使用できます。メモリ使用量を避けるため、他のパイプラインコンポーネントは読み込まないでください。

from diffusers import FluxPipeline
import torch

prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=None,
    vae=None,
    device_map="balanced",
    max_memory={0: "16GB", 1: "16GB"},
    dtype=torch.bfloat16
)
with torch.no_grad():
    print("Encoding prompts.")
    prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
        prompt=prompt, prompt_2=None, max_sequence_length=512
    )

テキスト埋め込みが計算された後、拡散トランスフォーマーのためにスペースを確保するため、GPUから削除します。

import gc 

def flush():
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()

del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline

flush()

device_map="auto"を設定すると、モデルが2つのGPUに自動で分散配置されます。この戦略では、まず最も高速なデバイスにモデルを配置し、必要に応じてCPUやハードドライブといった低速なデバイスに配置します。モデルのパラメータを低速なデバイスに保存するトレードオフは、推論のレイテンシが長くなることです。

from diffusers import AutoModel
import torch 

transformer = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev", 
    subfolder="transformer",
    device_map="auto",
    dtype=torch.bfloat16
)

pipeline.hf_device_map を実行すると、各モデルがデバイスにどのように分散配置されているかを確認できます。これは、モデルのデバイス配置を追跡するのに便利です。また、transformer モデルに対して hf_device_map を呼び出すことで、その分散状況を確認することもできます。

パイプラインに transformer モデルを追加し、output_type="latent" を設定して latent を生成します。

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    text_encoder=None,
    text_encoder_2=None,
    tokenizer=None,
    tokenizer_2=None,
    vae=None,
    transformer=transformer,
    dtype=torch.bfloat16
)

print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
    prompt_embeds=prompt_embeds,
    pooled_prompt_embeds=pooled_prompt_embeds,
    num_inference_steps=50,
    guidance_scale=3.5,
    height=height,
    width=width,
    output_type="latent",
).images

パイプラインとtransformerをメモリから解放し、VAEを読み込んで潜変数（latents）をデコードします。VAEは通常、単一のデバイスに読み込むのに十分なほど小さいサイズです。

import torch
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor

vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", dtype=torch.bfloat16).to("cuda")  # or "mps", "xpu", "cpu"
vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)

with torch.no_grad():
    print("Running decoding.")
    latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
    latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor

    image = vae.decode(latents, return_dict=False)[0]
    image = image_processor.postprocess(image, output_type="pil")
    image[0].save("split_transformer.png")

特定の段階で必要なモデルを選択的にロードおよびアンロードし、最大サイズのモデルを複数のGPUにシャーディングすることで、コンシューマー向けGPU上で大規模モデルの推論を実行することが可能になります。

コンテキスト並列

Context parallelismは、入力シーケンスを複数のGPUに分割してメモリ使用量を削減します。各GPUは、シーケンスの自分のスライスを処理します。

より最適化されたアテンションバックエンドに切り替えるには、set_attention_backend()を使用してください。利用可能なバックエンドの完全なリストについては、この表を参照してください。

ほとんどのアテンションバックエンドはコンテキスト並列処理と互換性があります。互換性のないバックエンドがある場合は、issueを作成してください。

リングアテンション

Key (K) と value (V) の表現は、Ring Attention を用いてデバイス間で通信します。これにより、各分割が他のすべてのトークンの K/V を参照できるようになります。各 GPU はローカルの K/V に対する attention を計算し、リング内の次の GPU に渡します。単一の GPU がシーケンス全体を保持しないため、通信レイテンシが低減されます。

transformer モデルの parallel_config 引数に ContextParallelConfig を渡します。この設定は、Ring Attention に使用するデバイス数を決定する ring_degree 引数をサポートしています。

import torch
from torch import distributed as dist
from diffusers import DiffusionPipeline, ContextParallelConfig

def setup_distributed():
    if not dist.is_initialized():
        dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    device = torch.device(f"cuda:{rank}")
    torch.cuda.set_device(device)
    return device

def main():
    device = setup_distributed()
    world_size = dist.get_world_size()

    pipeline = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-dev", dtype=torch.bfloat16
    ).to(device)
    pipeline.transformer.set_attention_backend("_native_cudnn")

    cp_config = ContextParallelConfig(ring_degree=world_size)
    pipeline.transformer.enable_parallelism(config=cp_config)

    prompt = """
    cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
    highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
    """

    # Must specify generator so all ranks start with same latents (or pass your own)
    generator = torch.Generator().manual_seed(42)
    image = pipeline(
        prompt,
        guidance_scale=3.5,
        num_inference_steps=50,
        generator=generator,
    ).images[0]

    if dist.get_rank() == 0:
        image.save(f"output.png")

    if dist.is_initialized():
        dist.destroy_process_group()


if __name__ == "__main__":
    main()

上記のスクリプトは、PyTorch と互換性のある分散ランチャー、例えば torchrun を使用して実行する必要があります。--nproc-per-node は、利用可能な GPU の数に設定されます。

torchrun --nproc-per-node 2 above_script.py

Ulysses Attention

Ulysses Attention splits a sequence across GPUs and performs an all-to-all communication (every device sends/receives data to every other device). Each GPU ends up with all tokens for only a subset of attention heads. Each GPU computes attention locally on all tokens for its head, then performs another all-to-all to regroup results by tokens for the next layer.

ContextParallelConfigは、ulysses_degree引数を通じてUlysses Attentionをサポートしています。この引数は、Ulysses Attentionに使用するデバイスの数を決定します。

enable_parallelism() に ContextParallelConfig を渡します。

# Depending on the number of GPUs available.
pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2))

統一アテンション

Unified Sequence Parallelismは、Ring AttentionとUlysses Attentionを単一の手法に統合し、効率的な長シーケンス処理を実現します。まずUlyssesのall-to-all通信を適用してヘッドとシーケンストークンを再分配し、次にRing Attentionを用いて再分配されたデータを処理し、最後にall-to-allを逆転させて元のレイアウトを復元します。

このハイブリッドアプローチは、両方の手法の強みを活用します：

• Ulysses Attention は、アテンションヘッド間で効率的に並列化します
• Ring Attentionは、最小限のメモリオーバーヘッドで非常に長いシーケンスを処理します
• これらにより、ヘッド次元とシーケンス次元の両方に対する2D並列化が可能になります。

ContextParallelConfig は、ulysses_degree と ring_degree の両方を指定することで Unified Attention をサポートします。使用するデバイスの総数は ulysses_degree * ring_degree となり、Ulysses グループと Ring グループが直交（重複しない）する 2 次元グリッドとして配置されます。 ulysses_degree と ring_degree の両方を 1 より大きい値に設定した ContextParallelConfig を enable_parallelism() に渡します。

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ring_degree=2))

Unified Attention は、2D グリッドに配置できるだけのデバイス（少なくとも 4 台）がある場合に使用します。

4つのH100 GPUで構成されるノード上で、このスクリプトを使用してUlysess、Ring、Unified Attentionのベンチマークを実行しました。結果は以下の通りです。

上記の表から、Ulyssesの方がスループットに優れていることが明らかですが、使用可能なデバイスの数はアテンションヘッドの数に制限されるという制約が残っています。この制約はunified attentionによって解決されます。

Ulysses Anything Attention

デフォルトのUlysses Attentionメカニズムでは、hidden statesのシーケンス長がデバイス数で割り切れる必要があります。これはUlysses Attentionの実用的な応用において大きな制約を課します。Ulysses Anything Attentionは、任意のシーケンス長と任意のattention heads数をサポートするUlysses Attentionのバリアントであり、実用上の汎用性を高めています。

ContextParallelConfig は、ulysses_degree と ulysses_anything の両方を指定することで、Ulysses Anything Attention をサポートします。なお、Ulysses Anything Attention は現在、Unified Attention によってサポートされていません。ulysses_degree を 1 より大きい値に設定し、ulysses_anything=True を設定した ContextParallelConfig を enable_parallelism() に渡してください。

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ulysses_anything=True))

H2DおよびD2H転送によって引き起こされる複数の強制CUDA同期を回避するために、init_process_groupにglooバックエンドを追加してください。これにより、通信レイテンシが大幅に低減されます。

4つのL20 GPUで構成されるノード上で、このスクリプトを使用して、Ulysses、Ring、Unified Attention、Ulysses Anything AttentionでFLUX.1-devのベンチマークを実行しました。結果は以下の通りです。

上記の表から、Ulysses Anything Attention は、標準の Ulysses Attention と同等のパフォーマンスを維持しながら、任意のシーケンス長に対する互換性が優れていることが明らかです。

Ring Anything Attention

デフォルトのRing Attentionでは、hidden states のシーケンス長がリングの次数で均等に分かれる必要があります。Ring Anything Attentionは、Ring Attention のバリアントであり、任意の（均等に分けられない）シーケンス長をサポートします。各ランクのローカル KV をグローバルな最大シーケンス長までパディングし、パディング済みの KV バッファを all-gather してから、attention を実行する前に各ランクの実際の長さにスライスし直します。

Ring Anything Attention は Unified Attention によってサポートされていません。Ring Anything Attention を有効にするには、ring_degree > 1 および ring_anything=True を設定してください。

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=2, ring_anything=True))

H2DおよびD2H転送による複数の強制CUDA同期を回避するために、init_process_groupにglooバックエンドを追加します。

import torch.distributed as dist

dist.init_process_group(backend="cpu:gloo,cuda:nccl")

Ring Anything Attention は現在、推論のみをサポートしています。また、None 以外の attention mask はサポートされていません。attn_mask は None でなければなりません。

4つのRTX 4090（48GB）GPUを備えたノードでのFLUX.1-devのベンチマーク結果は以下の通りです。

上記の表から、Ring Anything Attention は、標準的な Ring Attention と同等のパフォーマンスを維持しながら、任意のシーケンス長との互換性を提供していることがわかります。 Ring Anything Attention の動機とトレードオフの詳細については、このコメントを参照してください。

parallel_config

モデルの初期化時にparallel_configを渡すことで、コンテキスト並列化を有効にできます。

CKPT_ID = "black-forest-labs/FLUX.1-dev"

cp_config = ContextParallelConfig(ring_degree=2)
transformer = AutoModel.from_pretrained(
    CKPT_ID, 
    subfolder="transformer", 
    dtype=torch.bfloat16,
    parallel_config=cp_config
)

pipeline = DiffusionPipeline.from_pretrained(
    CKPT_ID, transformer=transformer, dtype=torch.bfloat16,
).to(device)

テンソル並列

テンソル並列では、モデルの重み行列を複数のデバイスに分割して配置します。各デバイスは、各レイヤーの列方向("colwise")または行方向("rowwise")のスライスを保持し、部分結果を計算します。レイヤーの境界でAllReduce/AllGatherを実行することで、完全な出力が再構築されます。コンテキスト並列とは異なり、重みのデバイスあたりメモリを削減できるため、単一のデバイスに収まらないモデルに有用です。

モデルの from_pretrained() の parallel_config 引数に TensorParallelConfig を渡します。tp_degree はシャードに分割するデバイスの数であり、モデルのアテンションヘッド数で割り切れる必要があります。モデルは _tp_plan（モジュール名のグロブから "colwise"/"rowwise" スタイルへのフラットなマッピング）を定義する必要があります。

この方法で読み込むと、読み込み中にチェックポイントがシャード化されます。各ランクは、シャード化された重みのうち自分のスライスのみを読み取り、それをそのまま自分のデバイスに配置します。フルサイズの状態がメモリ上に生成されることは一切ないため、tp_degreeが上昇するにつれて、ランクごとのメモリ使用量は減少します。

enable_parallelism() を呼び出す前にフルモデルを読み込む方法と比較すると、tp_degree が大きくなるにつれてその差は拡大しますが、読み込みが速く、ランクあたりの CPU メモリ使用量も少なくなります。以下の数値は、4x A10G (23GB) 上で black-forest-labs/FLUX.2-dev (transformer のみ、32B パラメータ、bf16) を使用した場合のものです。

import torch
from torch import distributed as dist
from diffusers import DiffusionPipeline, Flux2Transformer2DModel, TensorParallelConfig

def main():
    dist.init_process_group(backend="nccl")
    rank, world_size = dist.get_rank(), dist.get_world_size()
    device = torch.device(f"cuda:{rank}")
    torch.cuda.set_device(device)

    # Each rank reads only its own shard of every planned weight, straight onto `cuda:rank`.
    transformer = Flux2Transformer2DModel.from_pretrained(
        "black-forest-labs/FLUX.2-dev",
        subfolder="transformer",
        torch_dtype=torch.bfloat16,
        parallel_config=TensorParallelConfig(tp_degree=world_size),
    )

    pipeline = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.2-dev", transformer=transformer, torch_dtype=torch.bfloat16
    )
    # The transformer is already on its device; move the remaining components individually. Do not call
    # `pipeline.to(device)` — that would move every rank's shards onto the same device.
    pipeline.text_encoder.to(device)
    pipeline.vae.to(device)

    generator = torch.Generator().manual_seed(42)
    image = pipeline(prompt="a cat holding a sign that says hello", generator=generator).images[0]
    if rank == 0:
        image.save("output.png")
    dist.destroy_process_group()

if __name__ == "__main__":
    main()

torchrun --nproc-per-node 4 tensor_parallel_flux.py

tp_degree は上記の world_size から取得されるため、--nproc-per-node 4 を指定すると、transformer が 4 つのデバイスにシャードされます。

parallel_config を使用したテンソル並列での読み込みは、device_map、quantization_config、low_cpu_mem_usage=False、use_flashpack=True、または非 safetensors 形式の重みではまだサポートされていません。これらのいずれかを指定すると、完全なチェックポイントを読み込むように静かにフォールバックするのではなく、例外が発生します。

テンソル並列化を量子化、オフロード、または LoRA アダプターと組み合わせることもまだサポートされていないため、モデルがどのようにシャード化されていても、これらの機能は例外を発生させます。

テンソル並列モデルの読み込みには、この方法をおすすめします。すでにメモリ上に読み込まれたモデルをシャーディングすることも可能ですが、その場合は同じ設定でenable_parallelism()を呼び出します。ただし、この方法ではまずすべてを読み込んでから再シャーディングを行うため、すべてのランクでチェックポイント全体分のメモリコストが発生します。

テンソル並列モデルの保存は現在サポートされておらず、save_pretrained() はテンソル並列モデルに対して例外を発生させます。モデルをシャーディングする前に保存してください。

テンソル並列化プランの作成

Tensor parallelism は、モジュール名のグロブをシャーディング方式にマッピングするフラットなクラス属性である _tp_plan を定義しているモデルでのみ機能します。この定義を作成する際は、主に、隠れ次元を 拡張 する投影と、それを元に戻して 縮小 する投影をペアにする必要があります。

各キーには*を最大1つ含めることができ、その接頭辞はnn.ModuleListに解決され、単一のエントリですべてのブロックをカバーする必要があります。*を含まないキーはモデル自体に適用されます。パスはモデルを基準とした相対パスです。

Colwise と rowwise

常にこの順序でペアにしてください。"colwise" 投影は出力をシャードしたまま残し、続く "rowwise" 投影はそのシャードを直接消費し、ブロック境界での単一の AllReduce によって結果が再構築されます。このペアを他の方法でシャードすると、中間で gather が発生し、通信量が大幅に増加します。

attention においては、各 rank が head の部分集合を所有するため、tp_degree は head 数で割り切れる必要があります。Encoder-stream の重複部分（add_q_proj、to_add_out、ff_context）は、image-stream 側の対応する部分と同じパターンに従います。

融合された投影

1つのLinearがシャード対象の次元に沿って複数の論理テンソルをパックしている場合、単純な"colwise"/"rowwise"スライスでは連結を横断してそのままスライスされてしまい、各パーツがずれてしまいます。代わりにPackedColwiseParallel/PackedRowwiseParallelを使用してください。これにより、各パックされたブロックが独立してシャードされます。

# in src/diffusers/models/transformers/your_model.py
from ...hooks.tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel

blocks は、パックされた次元を割り切る整数のリストです。例えば、等しい 2 分割の SwiGLU gate+up 投影では [1, 1]、mlp_ratio=3 の Q+K+V+gate+up 融合投影では [1, 1, 1, 3, 3] となります。

"transformer_blocks.*.ff.linear_in": PackedColwiseParallel([1, 1]),

ブロックサイズが設定ファイルからのみ判明する場合、引数を省略し、代わりに__init__の時点でLinearに絶対サイズを_tp_packed_col_blocksまたは_tp_packed_row_blocksとして保存してください。

# in the attention module's __init__
self.to_out._tp_packed_row_blocks = [self.inner_dim, self.mlp_hidden_dim]

# in the model's _tp_plan
"single_transformer_blocks.*.attn.to_out": PackedRowwiseParallel(),

transformer_flux2.py には、両方の形式が 1 つのプランに登場します。ダブルストリームの SwiGLU ペアは blocks をそのまま渡すのに対し、フュージョンされたシングルストリームの QKV+MLP 投影は、サイズが mlp_ratio と mlp_mult_factor に依存するため、Flux2ParallelSelfAttention.__init__ 内の Linear にそのサイズを格納します。

除外する項目

計画に含まれていないものは、すべてのランクでレプリカ化されます。これは、正規化レイヤー、AdaLN 変調（img_mod/txt_mod）、パッチおよびテキストの埋め込み、そして最終的な norm_out/proj_out に対しては適切な選択です。これらのパラメータはサイズが小さいため、シャーディングしてもメモリ削減の効果は限定的であり、通信コストが増加するだけです。

制約と検証

• tp_degreeはconfig.num_attention_headsで割り切れる必要があります。
• すべてのパックされたブロックは、合計値だけでなく、それぞれ個別にtp_degreeで割り切れる必要があります。

両者は、重みの読み込みやシャーディングが行われる前に、from_pretrained() および enable_parallelism() によって検証されます。

目視ではなく数値的に新しいプランを検証します。固定のシードで単一デバイス上で生成し、次にテンソル並列環境下で再度生成して、出力を比較します。"colwise"/"rowwise"の配置が誤っていても、通常は実行され、妥当に見えますが誤った画像が生成されます。

類似アーキテクチャの既存プランを起点とします。QwenImageTransformer2DModel は完全に非融合であり、すべてのエントリがプレーンな"colwise"/"rowwise"です。FluxTransformer2DModel は単一のパックされた行方向射影を追加し、Flux2Transformer2DModel は両方のパックスタイルをカバーしています。それぞれtransformer_qwenimage.py、transformer_flux.py、transformer_flux2.pyを参照してください。

戦略の選択

上記の戦略はそれぞれ異なる問題を解決します。重要なのは、抽象的に見てどれが最も高速かではなく、リソースの不足がどの部分で起きているかということです。

実用的なガイドライン:

• 多数のプロンプトに対するスループット、モデルはすでに収まる。 データ並列性を使用します。ここでは、モデルを変更せずにスループットを線形にスケールさせ、単一プロンプトのレイテンシに影響を与えない唯一の戦略です。
• 単一コンポーネントの重みが収まらない場合は、まずテンソル並列化を検討してください。テンソル並列化はメモリとレイテンシの両方を削減できるためです。ブロック境界ごとに通信が発生するため、NVLinkのような高速なインターコネクトが望まれます。PCIe経由では、レイヤーごとの通信トラフィックが計算の削減効果を上回る可能性があり、その場合はdevice_mapの方が適切な選択となります。device_mapは、各コンポーネントが個別には問題ないものの、全体として大きすぎるケースにも対応します。
• 問題なのは重みではなく、アクティベーションです。 これは長シーケンスの領域（大きな画像、多数のフレーム）であり、コンテキスト並列性が直接的な解決策です。上記の Ulysses/Ring ベンチマーク からバックエンドを選択してください。
• 重みとシーケンスの両方が大きすぎる場合は、テンソル並列とコンテキスト並列を組み合わせます。TensorParallelConfigはmesh引数を受け取り、両者が1つのデバイスメッシュを共有できます。この組み合わせはまだ検証されていないため、実験的なものとみなしてください。