Inférence distribuée

L'inférence distribuée répartit la charge de travail sur plusieurs GPU. C'est une technique utile pour faire tenir des modèles plus volumineux en mémoire et elle permet de traiter plusieurs requêtes pour un débit plus élevé.

Ce guide vous montre comment utiliser Accelerate et PyTorch Distributed pour l'inférence distribuée.

Accelerate

Accelerate est une bibliothèque conçue pour simplifier l'inférence et l'entraînement sur plusieurs accélérateurs en gérant la configuration, ce qui permet aux utilisateurs de se concentrer sur leur code PyTorch.

Installez Accelerate avec la commande suivante.

uv pip install accelerate

Initialisez une classe accelerate.PartialState dans un fichier Python pour créer un environnement distribué. La classe accelerate.PartialState gère la gestion des processus, le contrôle et la distribution des appareils, ainsi que la coordination des processus.

Déplacez le DiffusionPipeline vers accelerate.PartialState.device pour affecter un GPU à chaque processus.

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)

Utilisez l'utilitaire split_between_processes comme gestionnaire de contexte pour distribuer automatiquement les prompts entre le nombre de processus.

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")

Appel accelerate launch pour exécuter le script et utilisez l’argument --num_processes pour définir le nombre de GPU à utiliser.

accelerate launch run_distributed.py --num_processes=2

Reportez-vous à ce script minimal pour exécuter une inférence sur plusieurs GPU. Pour en savoir plus, consultez le guide Distributed Inference with 🤗 Accelerate.

PyTorch Distributed

PyTorch DistributedDataParallel permet le parallélisme de données, qui réplique le même modèle sur chaque appareil afin de traiter différents lots de données en parallèle.

Importez torch.distributed et torch.multiprocessing dans un fichier Python afin de configurer le groupe de processus distribués et de générer les processus pour l'inférence sur chaque 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,
)

Créez une fonction pour l’inférence avec init_process_group. Cette méthode crée un environnement distribué avec le type de backend, le rank du processus actuel et le world_size ou nombre de processus participants (par exemple, 2 GPU correspondraient à world_size=2).

Déplacez le pipeline sur le rank et utilisez get_rank pour attribuer un GPU à chaque processus. Chaque processus traite un prompt différent.

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")

Utilisez mp.spawn pour créer le nombre de processus défini dans world_size.

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


if __name__ == "__main__":
    main()

Appel de torchrun pour exécuter le script d’inférence et utilisation de l’argument --nproc_per_node pour définir le nombre de GPU à utiliser.

torchrun --nproc_per_node=2 run_distributed.py

device_map

L'argument device_map permet une inférence distribuée en plaçant automatiquement les composants du modèle sur des GPU distincts. C'est particulièrement utile lorsqu'un modèle ne tient pas sur un seul GPU. Vous pouvez utiliser device_map pour charger et décharger sélectivement les composants du modèle nécessaires à une étape donnée, comme le montre l'exemple ci-dessous (en supposant que deux GPU sont disponibles).

Définissez device_map="balanced" pour répartir uniformément les encodeurs de texte sur toutes les GPU disponibles. Vous pouvez utiliser l’argument max_memory pour allouer une quantité maximale de mémoire pour chaque encodeur de texte. Ne chargez aucun autre composant du pipeline afin d’éviter une consommation de mémoire.

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
    )

Une fois les embeddings de texte calculés, retirez-les du GPU pour libérer de l’espace pour le diffusion transformer.

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()

Définissez device_map="auto" pour distribuer automatiquement le modèle sur les deux GPU. Cette stratégie place d'abord le modèle sur l'appareil le plus rapide, puis, si nécessaire, sur un appareil plus lent comme un CPU ou un disque dur. Le compromis du stockage des paramètres du modèle sur des appareils plus lents est une latence d'inférence plus élevée.

from diffusers import AutoModel
import torch 

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

Exécutez pipeline.hf_device_map pour voir comment les différents modèles sont répartis sur les appareils. Cela est utile pour suivre l’emplacement des modèles sur les appareils. Vous pouvez également appeler hf_device_map sur le modèle transformer pour voir comment il est réparti.

Ajoutez le modèle transformer au pipeline et définissez output_type="latent" pour générer les latents.

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

Supprimez le pipeline et le transformer de la mémoire et chargez un VAE pour décoder les latents. Le VAE est généralement assez petit pour être chargé sur un seul appareil.

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")

En chargeant et déchargeant sélectivement les modèles nécessaires à une étape donnée et en répartissant les plus grands modèles sur plusieurs GPU, il est possible d’exécuter l’inférence avec de grands modèles sur des GPU grand public.

Parallélisme de contexte

Le parallélisme de contexte répartit les séquences d'entrée sur plusieurs GPU afin de réduire l'utilisation de la mémoire. Chaque GPU traite sa propre portion de la séquence.

Utilisez set_attention_backend() pour passer à un backend d'attention plus optimisé. Consultez ce tableau pour obtenir la liste complète des backends disponibles.

La plupart des backends d'attention sont compatibles avec le parallélisme de contexte. Ouvrez un issue si un backend n'est pas compatible.

Ring Attention

Les représentations de clé (K) et de valeur (V) communiquent entre les appareils à l'aide de Ring Attention. Cela garantit que chaque partition voit les K/V de tous les autres jetons. Chaque GPU calcule l'attention pour ses K/V locaux et les transmet au GPU suivant dans l'anneau. Aucun GPU ne détient la séquence complète, ce qui réduit la latence de communication.

Passez un ContextParallelConfig à l'argument parallel_config du modèle transformer. La configuration prend en charge l'argument ring_degree qui détermine le nombre d'appareils à utiliser pour Ring Attention.

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()

Le script ci-dessus doit être exécuté avec un lanceur distribué, tel que torchrun, compatible avec PyTorch. --nproc-per-node est défini sur le nombre de GPU disponibles.

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

Ulysses Attention

Ulysses Attention répartit une séquence sur plusieurs GPU et effectue une communication all-to-all (chaque appareil envoie/reçoit des données à tous les autres appareils). Chaque GPU se retrouve avec tous les jetons pour un seul sous-ensemble de têtes d'attention. Chaque GPU calcule l'attention localement sur tous les jetons pour sa tête, puis effectue un autre all-to-all pour regrouper les résultats par jetons pour la couche suivante.

ContextParallelConfig prend en charge Ulysses Attention via l'argument ulysses_degree. Celui-ci détermine le nombre d'appareils à utiliser pour Ulysses Attention.

Passez ContextParallelConfig à enable_parallelism().

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

Attention unifiée

Unified Sequence Parallelism combine Ring Attention et Ulysses Attention en une seule approche pour un traitement efficace des séquences longues. Il applique d’abord la communication all-to-all d’Ulysses pour redistribuer les têtes et les jetons de séquence, puis utilise Ring Attention pour traiter les données redistribuées, et inverse enfin le all-to-all pour restaurer la disposition d’origine.

Cette approche hybride exploite les forces des deux méthodes :

• Ulysses Attention parallélise efficacement sur les têtes d'attention
• Ring Attention gère les séquences très longues avec un surcoût mémoire minimal
• Ensemble, ils permettent une parallélisation 2D à la fois sur les dimensions des têtes et de la séquence

ContextParallelConfig prend en charge Unified Attention en spécifiant à la fois ulysses_degree et ring_degree. Le nombre total d’appareils utilisés est ulysses_degree * ring_degree, organisés dans une grille 2D où les groupes Ulysses et Ring sont orthogonaux (sans chevauchement). Transmettez le ContextParallelConfig avec ulysses_degree et ring_degree définis sur des valeurs supérieures à 1 à enable_parallelism().

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

Unified Attention doit être utilisé lorsque le nombre d’appareils est suffisant pour les organiser en une grille 2D (au moins 4 appareils).

Nous avons effectué un benchmark avec Ulysess, Ring et Unified Attention à l'aide de ce script sur un nœud de 4 GPU H100. Les résultats sont résumés comme suit :

D’après le tableau ci-dessus, il est clair que Ulysses offre un meilleur débit, mais le nombre d’appareils qu’il peut utiliser reste limité au nombre de têtes d’attention, une limitation qui est résolue par l’attention unifiée.

Ulysses Anything Attention

Le mécanisme d'attention Ulysses par défaut exige que la longueur de séquence des états cachés soit divisible par le nombre d'appareils. Cela impose des limitations importantes à l'application pratique de l'attention Ulysses. Ulysses Anything Attention est une variante de l'attention Ulysses qui prend en charge des longueurs de séquence arbitraires et un nombre arbitraire de têtes d'attention, améliorant ainsi la polyvalence de l'attention Ulysses dans un contexte d'utilisation pratique.

ContextParallelConfig prend en charge Ulysses Anything Attention en spécifiant à la fois ulysses_degree et ulysses_anything. Veuillez noter que Ulysses Anything Attention n'est actuellement pas pris en charge par Unified Attention. Transmettez le ContextParallelConfig avec ulysses_degree défini sur une valeur supérieure à 1 et ulysses_anything=True à enable_parallelism().

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

Pour éviter de multiples synchronisations CUDA forcées causées par les transferts H2D et D2H, veuillez ajouter le backend gloo dans init_process_group. Cela réduira considérablement la latence de communication.

Nous avons effectué un benchmark de FLUX.1-dev avec Ulysses, Ring, Unified Attention et Ulysses Anything Attention à l'aide de ce script sur un nœud de 4 GPU L20. Les résultats sont résumés comme suit :

D'après le tableau ci-dessus, il est clair que Ulysses Anything Attention offre une meilleure compatibilité avec des longueurs de séquence arbitraires tout en maintenant les mêmes performances que l'Ulysses Attention standard.

Ring Anything Attention

Le Ring Attention par défaut exige que la longueur de séquence des états cachés soit divisible sans reste par le degré de l’anneau. Ring Anything Attention est une variante du Ring Attention qui prend en charge des longueurs de séquence arbitraires (non divisibles sans reste). Elle remplit le KV local de chaque rang jusqu’à la longueur de séquence maximale globale, effectue un all-gather sur le tampon KV rembourré, puis recoupe le tampon à la longueur réelle de chaque rang avant d’exécuter l’attention.

Ring Anything Attention n'est pas pris en charge par Unified Attention. Définissez ring_degree > 1 et ring_anything=True pour activer Ring Anything Attention.

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

Ajoutez le backend gloo à init_process_group pour éviter de multiples synchronisations CUDA forcées dues aux transferts H2D et D2H.

import torch.distributed as dist

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

Ring Anything Attention ne prend actuellement en charge que l’inférence et les masques d’attention non-None ne sont pas pris en charge. attn_mask doit être None.

Consultez les benchmarks de FLUX.1-dev ci-dessous sur un nœud de 4 GPU RTX 4090 (48 Go).

D'après le tableau ci-dessus, Ring Anything Attention offre une compatibilité avec des longueurs de séquence arbitraires tout en maintenant des performances comparables à celles du Ring Attention standard. Pour plus de détails sur la motivation et les compromis de Ring Anything Attention, voir ce commentaire.

parallel_config

Passez parallel_config lors de l'initialisation du modèle pour activer le parallélisme de contexte.

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)

Parallélisme de tenseurs

Le parallélisme de tenseurs répartit les matrices de poids d'un modèle sur plusieurs appareils. Chaque appareil détient une tranche par colonne ("colwise") ou par ligne ("rowwise") de chaque couche, calcule un résultat partiel, et un AllReduce/AllGather à la frontière de la couche reconstruit la sortie complète. Contrairement au parallélisme de contexte, il réduit la mémoire de poids par appareil, ce qui est utile pour les modèles qui ne tiennent pas sur un seul appareil.

Passez un TensorParallelConfig à l’argument parallel_config de la méthode from_pretrained() du modèle. tp_degree est le nombre d’appareils sur lesquels effectuer le sharding et doit diviser le nombre de têtes d’attention du modèle. Le modèle doit définir un _tp_plan (un mapping plat des glob de noms de modules vers un style "colwise"/"rowwise").

Charger de cette manière sharde le checkpoint durant sa lecture : chaque rang ne lit que sa propre tranche de chaque poids shardé et le place directement sur son propre appareil. Rien de taille complète n'est jamais matérialisé, de sorte que la mémoire par rang diminue lorsque tp_degree augmente.

Par rapport au chargement du modèle complet suivi de l’appel à enable_parallelism(), cette méthode charge plus rapidement et utilise moins de mémoire CPU par rang, avec un écart qui augmente à mesure que tp_degree augmente. Les chiffres ci-dessous concernent black-forest-labs/FLUX.2-dev (transformer uniquement, 32 Md de paramètres, bf16) sur 4x A10G (23 Go).

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 est dérivé de world_size ci-dessus, de sorte que --nproc-per-node 4 répartit le transformer sur 4 appareils.

Le chargement avec un parallel_config parallèle par tenseurs n’est pas encore pris en charge avec device_map, quantization_config, low_cpu_mem_usage=False, use_flashpack=True ou des poids non-safetensors ; chacun provoque une erreur plutôt que de basculer silencieusement sur le chargement du point de contrôle complet.

La combinaison de la parallélisation des tenseurs avec la quantification, le déchargement ou les adaptateurs LoRA n’est pas encore prise en charge non plus, de sorte que ceux-ci lèvent une erreur quelle que soit la façon dont le modèle est réparti.

Nous recommandons de charger les modèles parallélisés par tenseurs de cette manière. Vous pouvez également fractionner un modèle déjà en mémoire en appelant enable_parallelism() avec la même configuration, mais cela charge d’abord l’ensemble du modèle puis le refractionne, ce qui entraîne un coût mémoire égal à celui du point de contrôle complet sur chaque rang.

La sauvegarde d'un modèle parallélisé par tenseurs n'est pas encore prise en charge, et save_pretrained() lève une erreur sur ce type de modèle. Enregistrez le modèle avant de le fragmenter.

Rédiger un plan de parallélisme de tenseurs

Le parallélisme tensoriel ne fonctionne que sur les modèles qui définissent un _tp_plan, un attribut de classe plat associant des motifs de noms de modules à un style de partitionnement. L’écriture d’un tel plan consiste principalement à associer chaque projection qui étend la dimension cachée à la projection qui la réduit à nouveau.

Chaque clé peut contenir au plus un *, et le préfixe qui la précède doit résoudre vers un nn.ModuleList afin qu'une seule entrée couvre tous les blocs. Une clé sans * s'applique au modèle lui-même. Les chemins sont relatifs au modèle.

Colwise et rowwise

Associez-les toujours dans cet ordre. Une projection "colwise" laisse sa sortie fragmentée, la projection "rowwise" suivante consomme directement ce fragment, et un unique AllReduce à la frontière du bloc reconstruit le résultat. Fragmenter la paire de toute autre manière impose un regroupement intermédiaire et communique beaucoup plus.

Pour l'attention, cela signifie que chaque rang possède un sous-ensemble de têtes, c'est pourquoi tp_degree doit diviser le nombre de têtes. Les duplications du flux encodeur (add_q_proj, to_add_out, ff_context) suivent le même schéma que leurs équivalents du flux image.

Projections fusionnées

Lorsqu’un Linear regroupe plusieurs tenseurs logiques le long de la dimension qui est partagée, des tranches simples "colwise"/"rowwise" traversent la concaténation et désalignent les morceaux. Utilisez plutôt PackedColwiseParallel/PackedRowwiseParallel, qui partagent chaque bloc regroupé indépendamment.

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

blocks est une liste d'entiers proportionnels dont la somme divise la dimension empilée — [1, 1] pour une projection gate+up SwiGLU de moitiés égales, ou [1, 1, 1, 3, 3] pour une projection Q+K+V+gate+up fusionnée avec mlp_ratio=3.

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

Lorsque les tailles de blocs ne sont connues que depuis la configuration, omettez l’argument et stockez plutôt les tailles absolues sur Linear lors de __init__, sous la forme de _tp_packed_col_blocks ou _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(),

Les deux formes apparaissent dans un même plan dans transformer_flux2.py : la paire SwiGLU double-flux transmet blocks littéralement, tandis que la projection QKV+MLP fusionnée en flux unique stocke ses tailles sur le Linear dans Flux2ParallelSelfAttention.__init__ car elles dépendent de mlp_ratio et mlp_mult_factor.

À exclure

Tout ce qui est absent du plan reste répliqué sur chaque rang, ce qui est le bon choix pour les couches de normalisation, la modulation AdaLN (img_mod/txt_mod), les embeddings de patch et de texte, ainsi que la couche finale norm_out/proj_out. Ces éléments étant petits, leur sharding n'économise que peu de mémoire tout en ajoutant de la communication.

Contraintes et vérification

• tp_degree doit diviser config.num_attention_heads.
• Chaque bloc compacté doit être individuellement divisible par tp_degree, et non seulement leur somme.

Les deux sont validés par from_pretrained() et enable_parallelism() avant que n'importe quel poids ne soit chargé ou fractionné.

Valider un nouveau plan numériquement plutôt qu’à l’œil : générer avec une graine fixe sur un seul appareil, puis à nouveau sous parallélisme de tenseurs, et comparer les sorties. Un "colwise"/"rowwise" mal placé s’exécute généralement tout de même et produit une image plausible mais incorrecte.

Partez d’un plan existant pour une architecture similaire. QwenImageTransformer2DModel est entièrement non fusionné et chaque entrée est un simple "colwise"/"rowwise", FluxTransformer2DModel ajoute une unique projection row-wise packée, et Flux2Transformer2DModel couvre les deux styles packés — voir respectivement transformer_qwenimage.py, transformer_flux.py et transformer_flux2.py.

Choisir une stratégie

Les stratégies ci-dessus résolvent des problèmes différents, et la question utile n'est pas laquelle est la plus rapide en théorie, mais ce qui vous manque.

Quelques conseils pratiques :

• Débit sur de nombreuses requêtes, le modèle tient déjà en mémoire. Utilisez le parallélisme de données. C'est la seule stratégie ici qui augmente le débit de manière linéaire sans modifier le modèle, et qui laisse la latence des requêtes individuelles inchangée.
• Les poids d’un composant unique ne tiennent pas. Tournez-vous d’abord vers le parallélisme de tenseurs, car il réduit à la fois la mémoire et la latence. Il communique à chaque frontière de bloc, il exige donc un interconnect rapide comme NVLink ; sur PCIe, ce trafic par couche peut dépasser le calcul qu’il économise, et device_map devient alors le meilleur choix. device_map gère également le cas où les composants sont individuellement corrects mais collectivement trop volumineux.
• Ce sont les activations, et non les poids, qui posent problème. C'est le régime des séquences longues — grandes images, nombreuses images — et le parallélisme de contexte est la réponse directe. Choisissez un backend parmi les benchmarks Ulysses/Ring ci-dessus.
• Les poids et la séquence sont tous deux trop volumineux. Combinez le parallélisme de tenseurs et le parallélisme de contexte. TensorParallelConfig accepte un argument mesh afin que les deux puissent partager un même maillage d’appareils. Cette combinaison n’est pas encore validée, il convient donc de la considérer comme expérimentale.