diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index ccf898a8323b..1d363e1bad97 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -280,15 +280,15 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page ### Feature Matrix -| Model | FP8 blockwise | NVFP4 | TeaCache | CFG Parallelism | Ulysses Parallelism | Parallel VAE | CUDA Graph | torch.compile | trtllm-serve | Attention2D | Ring Attention | Tensor Parallelism | -|---|---|---|---|---|---|---|---|---|---|--|--|--| -| **FLUX.1** | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | -| **FLUX.2** | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | -| **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | -| **Wan 2.2** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | -| **LTX-2** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | -| **MiniMax-H3** | Yes | Yes | No | No | Yes | Yes | No | Yes | Yes | No | No | Yes | -| **Qwen-Image** | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | -| **Qwen-Image-Layered** | Yes | No | No | No | No | No | Yes | Yes | Yes | No | No | No | -| **Qwen-Image-Edit-2511** | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | No | No | No | -| **Cosmos3** | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | +| Model | FP8 blockwise | NVFP4 | TeaCache | CFG Parallelism | Ulysses Parallelism | Parallel VAE | CUDA Graph | torch.compile | trtllm-serve | Attention2D | Ring Attention | Tensor Parallelism | Text Encoder TP | +|---|---|---|---|---|---|---|---|---|---|--|--|--|---| +| **FLUX.1** | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | +| **FLUX.2** | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | +| **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | +| **Wan 2.2** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | +| **LTX-2** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | +| **MiniMax-H3** | Yes | Yes | No | No | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes | +| **Qwen-Image** | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No | +| **Qwen-Image-Layered** | Yes | No | No | No | No | No | Yes | Yes | Yes | No | No | No | No | +| **Qwen-Image-Edit-2511** | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | No | No | No | No | +| **Cosmos3** | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | No | diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 863314f7a5fe..19553ecab5a4 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -56,22 +56,22 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models ### Feature Matrix -| Model | FP8 blockwise | NVFP4 | TeaCache | Cache-DiT | CPU Offloading | CFG Parallelism | Ulysses Parallelism | Parallel VAE | CUDA Graph | torch.compile | trtllm-serve | Attention2D | Ring Attention | Tensor Parallelism\* | SOL-Attn | -|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---| -| **FLUX.1** | Yes | Yes | Yes | Yes | No | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | -| **FLUX.2** | Yes | Yes | Yes | Yes | No | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | -| **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | -| **Wan 2.1 VSA** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | Yes | -| **Wan 2.2** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | -| **FastWan 2.2** | Yes | Yes | No | No | No | No | No | No | Yes | Yes | Yes | No | No | No | No | -| **LTX-2** | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | Yes | -| **MiniMax-H3** | Yes | Yes | No | No | No | No | Yes | Yes | No | Yes | Yes | No | No | Yes | No | -| **Qwen-Image** | Yes | Yes | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No | -| **Qwen-Image-Layered** | Yes | No | No | Yes | No | No | No | No | Yes | Yes | Yes | No | No | No | No | -| **Qwen-Image-Edit-2511** | Yes | Yes | No | No | No | Yes | Yes | No | Yes | Yes | Yes | No | No | No | No | -| **Cosmos3** | Yes | Yes | No | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | No | -| **HunyuanVideo 1.5** | Yes | Yes | No | No | No | No | No | No | No | No | Yes | No | No | No | No | -| **GlmImage** | Yes | Yes | No | No | No | No | No | No | No | No | Yes | No | No | No | No | +| Model | FP8 blockwise | NVFP4 | TeaCache | Cache-DiT | CPU Offloading | CFG Parallelism | Ulysses Parallelism | Parallel VAE | CUDA Graph | torch.compile | trtllm-serve | Attention2D | Ring Attention | Tensor Parallelism\* | Text Encoder TP | SOL-Attn | +|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---| +| **FLUX.1** | Yes | Yes | Yes | Yes | No | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No | +| **FLUX.2** | Yes | Yes | Yes | Yes | No | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No | +| **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | Yes | +| **Wan 2.1 VSA** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | No | Yes | +| **Wan 2.2** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | Yes | +| **FastWan 2.2** | Yes | Yes | No | No | No | No | No | No | Yes | Yes | Yes | No | No | No | No | No | +| **LTX-2** | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | +| **MiniMax-H3** | Yes | Yes | No | No | No | No | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes | No | +| **Qwen-Image** | Yes | Yes | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No | No | +| **Qwen-Image-Layered** | Yes | No | No | Yes | No | No | No | No | Yes | Yes | Yes | No | No | No | No | No | +| **Qwen-Image-Edit-2511** | Yes | Yes | No | No | No | Yes | Yes | No | Yes | Yes | Yes | No | No | No | No | No | +| **Cosmos3** | Yes | Yes | No | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | No | No | +| **HunyuanVideo 1.5** | Yes | Yes | No | No | No | No | No | No | No | No | Yes | No | No | No | No | No | +| **GlmImage** | Yes | Yes | No | No | No | No | No | No | No | No | Yes | No | No | No | No | No | \* Wan 2.1 and Wan 2.2 (not VSA) also support a token-sharded TP layout (prototype), enabled with `parallel_config.tp_layout: token_sharded`; see [Multi-GPU Parallelism](#multi-gpu-parallelism). @@ -406,6 +406,10 @@ Configured under `VisualGenArgs.parallel_config`. Modes can be combined: - **Ring Attention** (`ring_size: N`): Shards the sequence axis across a 1D ring of `N` ranks, streaming K/V blocks (CP degree = `N`; total SP degree = `N ยท ulysses_size`; mutually exclusive with Attention2D). - **Tensor Parallelism** (`tp_size: N`): Splits attention heads and transformer MLPs across GPUs for faster compute and reduced memory usage. - **Token-sharded TP** (`tp_layout: token_sharded` in `parallel_config`, prototype): Keeps the residual stream sharded by tokens across the TP group, so each all-reduce becomes a reduce-scatter plus an all-gather and the norms, modulation and quantization between projections run on `1/tp_size` of the tokens. Requires `tp_size > 1`; composes with CFG parallelism, not with Ulysses, Ring Attention, Attention2D or VSA. Supported for Wan 2.1 and Wan 2.2. +- **Text encoder tensor parallelism** (`text_encoder_tp_size: N`): MiniMax-H3 + only. Shards the Qwen3-VL text encoder over the first N ranks for text-only + prompt encoding, then broadcasts prompt embeddings to the remaining VisualGen + ranks. This does not enable DiT tensor parallelism for MiniMax-H3. For multi-node execution, VisualGen relies on the external launcher to terminate the remaining ranks when any rank exits. `torchrun` provides this behavior. With diff --git a/tensorrt_llm/_torch/attention/attention.py b/tensorrt_llm/_torch/attention/attention.py index 858c92a44ad2..9dae20042e9b 100644 --- a/tensorrt_llm/_torch/attention/attention.py +++ b/tensorrt_llm/_torch/attention/attention.py @@ -33,6 +33,16 @@ from .rotary_embedding import MRotaryEmbedding, RotaryEmbedding +def _inherit_process_group_attrs(mapping: Mapping, source: Mapping) -> Mapping: + """Preserve out-of-band process group routing on derived TP mappings.""" + tp_group_pg = getattr(source, "_tp_group_pg", None) + if tp_group_pg is not None: + mapping._tp_group_pg = tp_group_pg + if getattr(source, "_use_torch_distributed_allreduce", False): + mapping._use_torch_distributed_allreduce = True + return mapping + + def extract_extra_attrs(layer_idx: str, attn_type: str): assert attn_type in ["mla", "attn"], "Invalid attention type" extra_attrs = get_model_extra_attrs() @@ -614,6 +624,7 @@ def __init__( gpus_per_node=self.mapping.gpus_per_node, enable_attention_dp=self.mapping.enable_attention_dp, ) + mapping = _inherit_process_group_attrs(mapping, self.mapping) self.tp_size = tp_size self.cp_size = cp_size self.tp_rank = mapping.tp_rank @@ -672,6 +683,7 @@ def __init__( gpus_per_node=self.mapping.gpus_per_node, enable_attention_dp=self.mapping.enable_attention_dp, ) + mapping_o = _inherit_process_group_attrs(mapping_o, self.mapping) self.mapping_o = mapping_o self.o_proj = Linear( diff --git a/tensorrt_llm/_torch/device_mesh.py b/tensorrt_llm/_torch/device_mesh.py index 257f72402a63..64a69273da1c 100644 --- a/tensorrt_llm/_torch/device_mesh.py +++ b/tensorrt_llm/_torch/device_mesh.py @@ -79,7 +79,9 @@ class DeviceMeshTopologyImpl(_MappingBaseForTypeCheck): # attribute loads, matching MpiTopology's precomputed rank lists. @torch.compiler.disable def _resolve_mesh_dim(self, name: str) -> tuple: - pg = self._get_mesh_dim_by_name(name).get_group() + pg = getattr(self, '_tp_group_pg', None) if name == 'tp' else None + if pg is None: + pg = self._get_mesh_dim_by_name(name).get_group() group_name = pg.group_name if hasattr(pg, 'group_name') else '' resolved = (pg, pg.rank(), self._get_group_ranks(pg), group_name) # Before torch.distributed is initialized (single-process runs) the diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 1ea2b401e9a9..6c5e5ec4d69a 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -1310,6 +1310,25 @@ def forward( if all_reduce_params is None: all_reduce_params = AllReduceParams() + use_torch_distributed_allreduce = self._disable_mpi and ( + getattr(self.mapping, "_use_torch_distributed_allreduce", False) + or getattr(self.mapping, "_tp_group_pg", None) is not None) + + if use_torch_distributed_allreduce: + if all_reduce_params.fusion_op != AllReduceFusionOp.NONE or any( + value is not None for value in ( + all_reduce_params.residual, + all_reduce_params.norm_weight, + all_reduce_params.scale, + all_reduce_params.bias, + )): + raise NotImplementedError( + "Torch distributed all-reduce fallback supports only " + "plain all-reduce.") + output = input.clone() + torch.distributed.all_reduce(output, group=self.mapping.tp_group_pg) + return output + # Try Symmetric Memory AllReduce first if available # Note: Currently only supports NONE fusion op (plain allreduce) if self.symm_mem_allreduce and all_reduce_params.fusion_op == AllReduceFusionOp.NONE: diff --git a/tensorrt_llm/_torch/models/modeling_qwen3.py b/tensorrt_llm/_torch/models/modeling_qwen3.py index 50474525a6b8..24cd6d8d418d 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3.py @@ -215,6 +215,7 @@ def __init__(self, mapping=config.mapping, tensor_parallel_mode=TensorParallelMode.COLUMN, gather_output=True, + allreduce_strategy=config.allreduce_strategy, ) self.layers = nn.ModuleList([ Qwen3DecoderLayer( diff --git a/tensorrt_llm/_torch/modules/embedding.py b/tensorrt_llm/_torch/modules/embedding.py index 33dd576dc683..b89250bc6fad 100644 --- a/tensorrt_llm/_torch/modules/embedding.py +++ b/tensorrt_llm/_torch/modules/embedding.py @@ -8,7 +8,7 @@ import torch.nn.functional as F from torch.nn.parameter import Parameter -from tensorrt_llm.functional import AllReduceParams +from tensorrt_llm.functional import AllReduceParams, AllReduceStrategy from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from tensorrt_llm.math_utils import ceil_div @@ -62,6 +62,7 @@ def __init__( use_custom_cublas_mm: bool = False, quant_config=None, output_dtype: Optional[torch.dtype] = None, + allreduce_strategy: AllReduceStrategy = AllReduceStrategy.AUTO, ): local_in_features = embedding_dim local_out_features = num_embeddings @@ -105,6 +106,7 @@ def __init__( reduce_output=reduce_output, use_custom_cublas_mm=use_custom_cublas_mm, quant_config=quant_config, + allreduce_strategy=allreduce_strategy, ) if tensor_parallel_mode == TensorParallelMode.ROW: @@ -327,6 +329,7 @@ def __init__( reduce_output: bool = True, enable_torch_compile_for_embedding: Optional[bool] = False, use_custom_cublas_mm: bool = False, + allreduce_strategy: AllReduceStrategy = AllReduceStrategy.AUTO, ): super().__init__( embedding_dim=embedding_dim, @@ -337,6 +340,7 @@ def __init__( gather_output=gather_output, reduce_output=reduce_output, use_custom_cublas_mm=use_custom_cublas_mm, + allreduce_strategy=allreduce_strategy, ) self.enable_torch_compile_for_embedding = enable_torch_compile_for_embedding diff --git a/tensorrt_llm/_torch/visual_gen/checkpoints/weight_loader.py b/tensorrt_llm/_torch/visual_gen/checkpoints/weight_loader.py index e729e4f4d21c..4e69bab6bd2c 100644 --- a/tensorrt_llm/_torch/visual_gen/checkpoints/weight_loader.py +++ b/tensorrt_llm/_torch/visual_gen/checkpoints/weight_loader.py @@ -72,6 +72,8 @@ def load_weights( Args: checkpoint_dir: Path to checkpoint (pipeline root or component dir) mapping: Distributed mapping (for future TP/PP support) + kwargs: Optional loading controls. ``prefetch=False`` disables + safetensors host page-cache prefetch before loading. Returns: - If single component: Dict mapping weight names to tensors @@ -79,6 +81,7 @@ def load_weights( Example: {"transformer": {...}, "transformer_2": {...}} """ checkpoint_path = Path(checkpoint_dir) + prefetch = kwargs.get("prefetch", True) # Standard and modular Diffusers manifests use the same component # subdirectory layout for weight-bearing models. @@ -109,7 +112,7 @@ def load_weights( if not weight_files: raise ValueError(f"No weight files found in {weight_dir}") - if all(wf.endswith(".safetensors") for wf in weight_files): + if prefetch and all(wf.endswith(".safetensors") for wf in weight_files): prefetch_files_to_host_cache( weight_files, description="visual-gen checkpoint", diff --git a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py index 6ef5e86db6bf..76751febecd2 100644 --- a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py +++ b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/pipeline_minimax_h3.py @@ -25,6 +25,7 @@ import numpy as np import torch +import torch.distributed as dist from diffusers import AutoencoderKLMiniMaxH3Audio, MiniMaxH3Scheduler from diffusers.utils.torch_utils import randn_tensor from PIL import Image, ImageOps @@ -40,8 +41,10 @@ RoleSpec, ) from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline +from tensorrt_llm._utils import mpi_disabled from tensorrt_llm.inputs.utils import load_image from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import MpiTopology from .packing import ( MINIMAX_H3_AUDIO_CHANNELS, @@ -70,6 +73,7 @@ ) from .parallel_vae import TiledAutoencoderKLMiniMaxH3 from .ref2va import load_references, prepare_references, validate_reference_order +from .text_encoder_tp import MiniMaxH3TensorParallelTextEncoder from .transformer_minimax_h3 import MiniMaxH3Transformer3DModel if TYPE_CHECKING: @@ -135,6 +139,29 @@ def _check_denoise_step(velocity: torch.Tensor, name: str, step: int) -> None: ) +class _MiniMaxH3TextEncoderMapping(MpiTopology): + """Mapping for the MiniMax-H3 text-encoder-only TP subgroup.""" + + def __new__(cls, *args, **kwargs): + del args, kwargs + return object.__new__(cls) + + def __init__(self, *args, tp_group_pg=None, **kwargs) -> None: + self._tp_group_pg = tp_group_pg + self._use_torch_distributed_allreduce = True + super().__init__(*args, **kwargs) + + @property + def tp_group_pg(self): + if self._tp_group_pg is None: + raise RuntimeError("MiniMax-H3 text encoder TP process group is not initialized.") + return self._tp_group_pg + + @property + def tp_group_name(self) -> str: + return self.tp_group_pg.group_name + + @register_pipeline( "MiniMaxH3ModularPipeline", hf_ids=["MiniMaxAI/MiniMax-H3"], @@ -176,8 +203,16 @@ def _request_generator(seed: int) -> torch.Generator: return torch.Generator().manual_seed(seed) def __init__(self, pipeline_config: DiffusionPipelineConfig) -> None: - vgm = pipeline_config.visual_gen_mapping + self.text_encoder_tp_size = int( + getattr(getattr(pipeline_config, "parallel", None), "text_encoder_tp_size", 1) + ) + vgm = getattr(pipeline_config, "visual_gen_mapping", None) world_size = vgm.world_size if vgm is not None else pipeline_config.mapping.world_size + if self.text_encoder_tp_size > world_size: + raise ValueError( + f"text_encoder_tp_size ({self.text_encoder_tp_size}) exceeds " + f"world_size ({world_size})." + ) if world_size != 1: supports_topology = ( vgm is not None @@ -200,6 +235,29 @@ def __init__(self, pipeline_config: DiffusionPipelineConfig) -> None: "Other attention backends do not support the required " "packed-row padding mask." ) + self._text_encoder_tp_group = None + if ( + self._uses_tensor_parallel_text_encoder + and self.text_encoder_tp_size < world_size + and not mpi_disabled() + ): + raise RuntimeError( + "MiniMax-H3 subset text_encoder_tp_size requires TLLM_DISABLE_MPI=1 " + "so Qwen text encoder TP uses a torch process group instead of MPI collectives." + ) + if self._uses_tensor_parallel_text_encoder and dist.is_initialized(): + text_encoder_tp_ranks = list(range(self.text_encoder_tp_size)) + rank = int(getattr(pipeline_config.mapping, "rank", 0)) + if self.text_encoder_tp_size == world_size: + group = dist.group.WORLD + else: + group = dist.new_group( + ranks=text_encoder_tp_ranks, + use_local_synchronization=False, + ) + if rank in text_encoder_tp_ranks: + self._text_encoder_tp_group = group + if ( pipeline_config.attention.backend == "TRTLLM" and torch.cuda.get_device_capability() not in ((10, 0), (10, 3)) @@ -398,6 +456,36 @@ def load_weights(self, weights: dict) -> None: self.transformer.load_weights(transformer_weights) self.transformer.eval() + def _load_tensor_parallel_text_encoder( + self, + checkpoint_dir: str, + device: torch.device, + ) -> MiniMaxH3TensorParallelTextEncoder: + if not self._loads_text_encoder_on_this_rank: + raise RuntimeError( + "MiniMax-H3 tensor-parallel text encoder is loaded only on " + f"ranks < {self.text_encoder_tp_size}; rank {self.rank} " + "must receive broadcast prompt embeddings." + ) + text_encoder_mapping = _MiniMaxH3TextEncoderMapping( + world_size=self.text_encoder_tp_size, + rank=self.rank, + tp_size=self.text_encoder_tp_size, + gpus_per_node=self.text_encoder_tp_size, + tp_group_pg=self._text_encoder_tp_group, + ) + text_encoder = MiniMaxH3TensorParallelTextEncoder.from_pretrained_config( + checkpoint_dir, text_encoder_mapping, device + ) + # Safetensors prefetch synchronizes on the default process group. Disable + # it when only a subset of ranks loads the text encoder. + prefetch = self.text_encoder_tp_size == self.world_size + weights = WeightLoader(components=PipelineComponent.TEXT_ENCODER).load_weights( + checkpoint_dir, text_encoder_mapping, prefetch=prefetch + ) + text_encoder.load_weights(weights) + return text_encoder.eval() + def load_standard_components( self, checkpoint_dir: str, @@ -419,12 +507,26 @@ def load_standard_components( subfolder=PipelineComponent.PROCESSOR, ) if not _component_skipped(skip_components, PipelineComponent.TEXT_ENCODER): - self.text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( - checkpoint_dir, - subfolder=PipelineComponent.TEXT_ENCODER, - torch_dtype=torch.bfloat16, - ).to(device) - self.text_encoder.eval() + if self._loads_text_encoder_on_this_rank: + if self._uses_tensor_parallel_text_encoder: + self.text_encoder = self._load_tensor_parallel_text_encoder( + checkpoint_dir, device + ) + else: + self.text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( + checkpoint_dir, + subfolder=PipelineComponent.TEXT_ENCODER, + torch_dtype=torch.bfloat16, + ).to(device) + self.text_encoder.eval() + else: + self.text_encoder = None + logger.info( + "MiniMax-H3 text encoder is loaded only on ranks < %d; " + "rank %d will receive broadcast prompt embeddings.", + self.text_encoder_tp_size, + self.rank, + ) if not _component_skipped(skip_components, PipelineComponent.VAE): self.vae = TiledAutoencoderKLMiniMaxH3.from_pretrained( checkpoint_dir, @@ -638,6 +740,100 @@ def _encode_prompt( prompt: str, keyframes: list[Image.Image], ) -> tuple[torch.Tensor, torch.Tensor]: + if self._uses_tensor_parallel_text_encoder and keyframes: + raise NotImplementedError( + "MiniMax-H3 tensor-parallel text encoder currently supports " + "text-only prompts; set text_encoder_tp_size=1 for keyframe requests." + ) + if self.world_size == 1: + return self._encode_prompt_local(prompt, keyframes) + + prompt_embeds: torch.Tensor | None = None + text_token_tags: torch.Tensor | None = None + with self._coordinate_text_encoder_errors(): + if self._loads_text_encoder_on_this_rank: + prompt_embeds, text_token_tags = self._encode_prompt_local(prompt, keyframes) + + prompt_embeds = self._broadcast_tensor_from_rank0(prompt_embeds, dtype=self.dtype) + text_token_tags = self._broadcast_tensor_from_rank0(text_token_tags, dtype=torch.long) + return prompt_embeds, text_token_tags.cpu() + + @property + def _uses_tensor_parallel_text_encoder(self) -> bool: + return self.text_encoder_tp_size > 1 + + @property + def _loads_text_encoder_on_this_rank(self) -> bool: + return self.rank < self.text_encoder_tp_size + + @contextmanager + def _coordinate_text_encoder_errors(self) -> Iterator[None]: + """Propagate caught encoding failures before peers enter payload broadcasts.""" + if self.world_size == 1 or not dist.is_initialized(): + yield + return + + # Allocate before encoding so an encoder OOM does not prevent reporting it. + failed_rank = torch.zeros(1, dtype=torch.int32, device=self.device) + local_error: Exception | None = None + try: + yield + except Exception as error: + local_error = error + failed_rank.fill_(self.rank + 1) + + # This coordinates Python failures only; failed TP collectives still + # require distributed worker teardown. + dist.all_reduce(failed_rank, op=dist.ReduceOp.MAX) + failed_rank_id = int(failed_rank.item()) - 1 + if failed_rank_id >= 0: + message = f"MiniMax-H3 text encoding failed on rank {failed_rank_id}" + if local_error is not None: + message += f": {local_error}" + raise RuntimeError(message) from local_error + + def _broadcast_tensor_from_rank0( + self, + tensor: torch.Tensor | None, + *, + dtype: torch.dtype, + ) -> torch.Tensor: + if not dist.is_initialized(): + if tensor is None: + raise ValueError("rank 0 tensor is required when distributed is not initialized.") + return tensor.to(device=self.device, dtype=dtype) + + shape = torch.zeros(8, dtype=torch.long, device=self.device) + if self.rank == 0: + if tensor is None: + raise ValueError("rank 0 tensor is required for MiniMax-H3 broadcast.") + tensor = tensor.to(device=self.device, dtype=dtype).contiguous() + shape[0] = tensor.ndim + shape[1 : tensor.ndim + 1] = torch.tensor(tensor.shape, device=self.device) + dist.broadcast(shape, src=0) + ndim = int(shape[0].item()) + tensor_shape = tuple(int(value) for value in shape[1 : ndim + 1].tolist()) + if self.rank != 0: + tensor = torch.empty(tensor_shape, dtype=dtype, device=self.device) + dist.broadcast(tensor, src=0) + return tensor + + def _encode_prompt_local( + self, + prompt: str, + keyframes: list[Image.Image], + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.text_encoder is None: + raise RuntimeError("MiniMax-H3 text encoder is not loaded on this rank.") + if not self._uses_tensor_parallel_text_encoder: + num_layers = self.text_encoder.config.text_config.num_hidden_layers + if num_layers <= MINIMAX_H3_TEXT_ENCODER_LAYER: + raise ValueError( + "MiniMax-H3 requires the unnormalized hidden state after Qwen3-VL " + f"layer {MINIMAX_H3_TEXT_ENCODER_LAYER}, but the encoder has " + f"{num_layers} layers." + ) + pixel_values = None image_grid_thw = None token_ids = [] @@ -677,10 +873,41 @@ def _encode_prompt( token_tags, dtype=torch.long ) + def _encode_reference_text_tokens( + self, + token_ids: list[int], + vision_inputs: dict[str, torch.Tensor], + ) -> torch.Tensor: + """Encode Ref2VA presentation tokens on owner ranks and broadcast.""" + if self._uses_tensor_parallel_text_encoder and vision_inputs: + raise NotImplementedError( + "MiniMax-H3 tensor-parallel text encoder currently supports " + "text-only prompts; set text_encoder_tp_size=1 for Ref2VA " + "requests with image or video references." + ) + if self.world_size == 1: + return self._encode_text_tokens(token_ids, vision_inputs) + + prompt_embeds: torch.Tensor | None = None + with self._coordinate_text_encoder_errors(): + if self._loads_text_encoder_on_this_rank: + prompt_embeds = self._encode_text_tokens(token_ids, vision_inputs) + return self._broadcast_tensor_from_rank0(prompt_embeds, dtype=self.dtype) + def _encode_text_tokens( self, token_ids: list[int], vision_inputs: dict[str, torch.Tensor] ) -> torch.Tensor: """Encode one H3 presentation, shared by keyframes and mixed references.""" + input_ids = torch.tensor([token_ids], dtype=torch.long, device=self.device) + if self._uses_tensor_parallel_text_encoder: + if vision_inputs: + raise NotImplementedError( + "MiniMax-H3 tensor-parallel text encoder currently supports " + "text-only prompts; set text_encoder_tp_size=1 for requests " + "with Qwen3-VL vision inputs." + ) + return self.text_encoder.encode(input_ids).to(device=self.device, dtype=self.dtype) + num_layers = self.text_encoder.config.text_config.num_hidden_layers if num_layers <= MINIMAX_H3_TEXT_ENCODER_LAYER: raise ValueError( @@ -689,7 +916,6 @@ def _encode_text_tokens( f"{num_layers} layers." ) - input_ids = torch.tensor([token_ids], dtype=torch.long, device=self.device) mm_token_type_ids = torch.tensor( self.processor.create_mm_token_type_ids([token_ids]), dtype=torch.long, diff --git a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/ref2va.py b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/ref2va.py index e1c537f54159..dc515e1168f7 100644 --- a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/ref2va.py +++ b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/ref2va.py @@ -353,7 +353,7 @@ def prepare_references( ids, tags = _build_presentation( pipeline.tokenizer, prompt, references, image_counts, video_counts, timestamps ) - prompt_embeds = pipeline._encode_text_tokens(ids, vision) + prompt_embeds = pipeline._encode_reference_text_tokens(ids, vision) visual_latents, audio_latents = [], [] audio_mean = torch.tensor(pipeline.audio_vae.config.latents_mean).view(1, 1, -1) audio_std = torch.tensor(pipeline.audio_vae.config.latents_std).view(1, 1, -1) diff --git a/tensorrt_llm/_torch/visual_gen/models/minimax_h3/text_encoder_tp.py b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/text_encoder_tp.py new file mode 100644 index 000000000000..a0f9b1099595 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/minimax_h3/text_encoder_tp.py @@ -0,0 +1,161 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tensor-parallel Qwen text encoder used by MiniMax-H3 visual generation.""" + +from __future__ import annotations + +import copy +import re +from pathlib import Path +from typing import Any + +import torch +from transformers import AutoConfig + +from tensorrt_llm._torch.attention.backends.vanilla import VanillaAttentionMetadata +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_qwen3 import Qwen3ForTextEmbedding +from tensorrt_llm.functional import AllReduceStrategy +from tensorrt_llm.mapping import Mapping + +from .packing import MINIMAX_H3_TEXT_ENCODER_LAYER + + +class MiniMaxH3TensorParallelTextEncoder(torch.nn.Module): + """TRTLLM Qwen3 backbone that returns MiniMax-H3's layer-50 state. + + MiniMax-H3 consumes the unnormalized hidden state from the Qwen3-VL text + backbone rather than final logits. The regular TRTLLM Qwen3 model already + shards attention/MLP weights when ``mapping.tp_size > 1``; this wrapper only + adapts that backbone to the visual-gen prompt-encoding contract. + """ + + _LANGUAGE_MODEL_PREFIX = "model.language_model." + _BARE_LANGUAGE_MODEL_PREFIX = "language_model." + _LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.") + + def __init__( + self, + text_config: Any, + mapping: Mapping, + device: torch.device, + *, + target_layer: int = MINIMAX_H3_TEXT_ENCODER_LAYER, + max_num_tokens: int | None = None, + ) -> None: + super().__init__() + if text_config.num_hidden_layers <= target_layer: + raise ValueError( + "MiniMax-H3 requires the unnormalized hidden state after Qwen3-VL " + f"layer {target_layer}, but the encoder has " + f"{text_config.num_hidden_layers} layers." + ) + + self.target_layer = target_layer + self.mapping = mapping + self.device = device + self.dtype = torch.bfloat16 + + config = copy.deepcopy(text_config) + config.architectures = ["Qwen3ForTextEmbedding"] + config.torch_dtype = self.dtype + config.num_hidden_layers = target_layer + config.disable_fuse_rope = True + if getattr(config, "rope_scaling", None) is None: + config.rope_scaling = {} + config.rope_scaling["type"] = "mrope" + + model_config = ModelConfig( + pretrained_config=config, + mapping=mapping, + attn_backend="VANILLA", + # MiniMax-H3 can load its Qwen text encoder on a subset of ranks. + # NCCL keeps all-reduce scoped to that TP group and avoids the + # custom workspace setup that synchronizes the full world. + allreduce_strategy=AllReduceStrategy.NCCL, + max_num_tokens=max_num_tokens or getattr(config, "max_position_embeddings", 8192), + ) + model_config._frozen = True + self.model = Qwen3ForTextEmbedding(model_config).to(device).eval() + + @classmethod + def from_pretrained_config( + cls, + checkpoint_dir: str | Path, + mapping: Mapping, + device: torch.device, + *, + target_layer: int = MINIMAX_H3_TEXT_ENCODER_LAYER, + ) -> "MiniMaxH3TensorParallelTextEncoder": + config = AutoConfig.from_pretrained( + checkpoint_dir, + subfolder="text_encoder", + ) + if not hasattr(config, "text_config"): + raise ValueError("MiniMax-H3 text_encoder config must contain text_config.") + return cls( + config.text_config, + mapping, + device, + target_layer=target_layer, + ) + + def load_weights(self, weights: dict[str, torch.Tensor]) -> None: + mapped_weights = {} + for key, value in weights.items(): + mapped_key = self._map_weight_key(key) + if mapped_key is None: + continue + mapped_weights[mapped_key] = value + self.model.load_weights(mapped_weights) + + def _map_weight_key(self, key: str) -> str | None: + if key.startswith("model.visual."): + return None + if key.startswith(self._LANGUAGE_MODEL_PREFIX): + key = f"model.{key[len(self._LANGUAGE_MODEL_PREFIX) :]}" + elif key.startswith(self._BARE_LANGUAGE_MODEL_PREFIX): + key = f"model.{key[len(self._BARE_LANGUAGE_MODEL_PREFIX) :]}" + + if key.startswith(("lm_head.", "model.lm_head.")): + return None + + lookup_key = key if key.startswith("model.") else f"model.{key}" + layer_match = self._LAYER_RE.match(lookup_key) + if layer_match is not None and int(layer_match.group(1)) >= self.target_layer: + return None + return key + + @torch.inference_mode() + def encode(self, input_ids: torch.Tensor) -> torch.Tensor: + input_ids = input_ids.reshape(-1).to(device=self.device, dtype=torch.long) + if self.device.type == "cuda": + torch.cuda.set_device(self.device) + seq_len = input_ids.shape[0] + attn_metadata = VanillaAttentionMetadata( + max_num_requests=1, + max_num_tokens=seq_len, + seq_lens=torch.tensor([seq_len], dtype=torch.int32), + num_contexts=1, + request_ids=[0], + mapping=self.mapping, + ) + attn_metadata.prepare_encoder_only() + + position_ids = torch.arange(seq_len, device=self.device, dtype=torch.long) + position_ids = position_ids.view(1, 1, seq_len).expand(3, 1, seq_len).contiguous() + + hidden_states = self.model.model.embed_tokens(input_ids) + residual = None + for decoder_layer in self.model.model.layers: + hidden_states, residual = decoder_layer( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + residual=residual, + ) + + if residual is not None: + hidden_states = hidden_states + residual + return hidden_states.unsqueeze(0) diff --git a/tensorrt_llm/mapping.py b/tensorrt_llm/mapping.py index 8a38434438b3..450f8e713e5a 100644 --- a/tensorrt_llm/mapping.py +++ b/tensorrt_llm/mapping.py @@ -661,6 +661,9 @@ def repurpose_helix_cp_to_tp(self): # DeviceMesh specific methods @property def tp_group_pg(self) -> ProcessGroup: + tp_group_pg = getattr(self, "_tp_group_pg", None) + if tp_group_pg is not None: + return tp_group_pg raise NotImplementedError("tp_group_pg is not implemented.") @property diff --git a/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json index b9c2b97f4b5c..66e5b4812bab 100644 --- a/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json @@ -397,6 +397,14 @@ "title": "Ring Size", "type": "integer" }, + "text_encoder_tp_size": { + "default": 1, + "description": "Text encoder tensor-parallel group size. Currently MiniMax-H3 only: text-only prompts shard the Qwen3-VL text encoder across ranks [0, text_encoder_tp_size), while Ref2VA or FL2VA requests with Qwen vision inputs must use text_encoder_tp_size=1. The rank 0 worker broadcasts prompt embeddings to the remaining ranks.", + "minimum": 1, + "status": "prototype", + "title": "Text Encoder Tp Size", + "type": "integer" + }, "tp_layout": { "anyOf": [ { diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index 374d8563196a..d829f1647563 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -413,6 +413,20 @@ def token_sharded_tp(self) -> bool: """True iff ``tp_layout`` is ``'token_sharded'`` (test this, not the field's truthiness).""" return self.tp_layout == "token_sharded" + text_encoder_tp_size: int = Field( + 1, + ge=1, + status="prototype", + description=( + "Text encoder tensor-parallel group size. Currently MiniMax-H3 " + "only: text-only prompts shard the Qwen3-VL text encoder across " + "ranks [0, text_encoder_tp_size), while Ref2VA or FL2VA requests " + "with Qwen vision inputs must use text_encoder_tp_size=1. " + "The rank 0 worker broadcasts prompt " + "embeddings to the remaining ranks." + ), + ) + @property def seq_parallel_size(self) -> int: """Parallelism degree over the sequence/context axis. @@ -480,6 +494,11 @@ def validate_world_size(self, world_size: int) -> None: f"total_parallel_size ({self.total_parallel_size}) " f"exceeds world_size ({world_size})" ) + if self.text_encoder_tp_size > world_size: + raise ValueError( + f"text_encoder_tp_size ({self.text_encoder_tp_size}) " + f"exceeds world_size ({world_size})" + ) class BaseCacheConfig(StrictBaseModel): diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py index 700a76642044..06bd8dfc616a 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_multi_gpu.py @@ -448,6 +448,63 @@ def test_wan22_t2v_lpips_against_golden_tp( ) +@pytest.mark.parametrize("world_size", [2, 4]) +def test_minimax_h3_text_encoder_tp_lpips(tmp_path, monkeypatch, world_size): + """Compare encoder TP with HF encoding while keeping DiT parallelism fixed.""" + from defs.examples.visual_gen.test_minimax_h3_e2e import ( + MINIMAX_H3_AUDIO_LOG_STFT_THRESHOLD, + _mean_lpips_distance, + _minimax_h3_checkpoint_path, + _multi_resolution_log_stft_distance, + ) + + from tensorrt_llm import VisualGen, VisualGenArgs, VisualGenParams + + if torch.cuda.device_count() < world_size: + pytest.skip(f"Requires {world_size} CUDA GPUs") + monkeypatch.setenv("TLLM_DISABLE_MPI", "1") + checkpoint = _minimax_h3_checkpoint_path() + videos, audios = [], [] + for encoder_tp in (1, 2): + args = VisualGenArgs( + model=checkpoint, + parallel_config={"ulysses_size": world_size, "text_encoder_tp_size": encoder_tp}, + attention_config={"backend": "VANILLA"}, + torch_compile_config={"enable": False}, + cuda_graph_config={"enable": False}, + compilation_config={"skip_warmup": True}, + ) + engine = VisualGen(model=checkpoint, args=args) + try: + output = engine.generate( + inputs="A woman smiles and waves in a sunlit park, with birds chirping.", + params=VisualGenParams( + height=128, width=128, num_frames=124, num_inference_steps=4, seed=42 + ), + ) + video = torch.as_tensor(output.video).cpu() + if video.ndim == 4: + video = video.unsqueeze(0) + audio = torch.as_tensor(output.audio).cpu() + assert video.shape == (1, 124, 128, 128, 3) + assert video.float().std() > 1 + assert torch.isfinite(audio).all() + videos.append(video) + audios.append(audio) + torch.save( + {"video": video, "audio": audio}, + tmp_path / f"h3_world{world_size}_encoder_tp{encoder_tp}.pt", + ) + finally: + engine.shutdown() + reference = videos[0].permute(0, 1, 4, 2, 3).float().div(255) + score = _mean_lpips_distance(videos[1], reference) + audio_score = _multi_resolution_log_stft_distance(audios[1], audios[0]) + print(f"H3 encoder TP=2/{world_size}: LPIPS={score:.6f}, audio log-STFT={audio_score:.6f}") + assert score < 0.05, f"Encoder TP changed video: LPIPS={score:.6f}" + assert audio_score < MINIMAX_H3_AUDIO_LOG_STFT_THRESHOLD + + @pytest.mark.parametrize( "ulysses_size,vae_size,backend", [(2, 2, "VANILLA"), (4, 3, "VANILLA"), (2, 2, "FA4")] ) diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 3f7f679f90d9..a347984ea457 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -41,6 +41,8 @@ l0_dgx_b200: - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[FA4-2-1] - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_ulysses_and_parallel_vae - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_tp2_example_lpips TIMEOUT (30) + - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_text_encoder_tp.py::test_minimax_h3_text_encoder_tp[tp2_world2] + - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_text_encoder_tp_lpips[2] TIMEOUT (30) - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_ulysses_parallel_vae_lpips[2-2-VANILLA] TIMEOUT (30) - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_ulysses_parallel_vae_lpips[2-2-FA4] TIMEOUT (30) - unittest/_torch/visual_gen/test_flux_pipeline.py::TestFluxParallelism::test_ulysses_2gpu_correctness @@ -66,6 +68,8 @@ l0_dgx_b200: - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[VANILLA-2-2] - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_tensor_parallel[FA4-2-2] - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_parallel.py::test_h3_partial_vae_groups + - unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_text_encoder_tp.py::test_minimax_h3_text_encoder_tp[tp2_world4] + - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_text_encoder_tp_lpips[4] TIMEOUT (30) - examples/visual_gen/test_visual_gen_multi_gpu.py::test_minimax_h3_ulysses_parallel_vae_lpips[4-3-VANILLA] TIMEOUT (30) # Synthetic weights cover Engram head/row shards and TP/EP shared-expert collectives. - unittest/_torch/multi_gpu/test_engram_sharding.py::test_engram_shard_roundtrip -k "world:4" diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_text_encoder_tp.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_text_encoder_tp.py new file mode 100644 index 000000000000..0a0bb0f8b8bc --- /dev/null +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_minimax_h3_text_encoder_tp.py @@ -0,0 +1,196 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Real Qwen3-VL layer-50 parity with full-world and subset text-encoder TP.""" + +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLTextConfig +from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextModel + +from tensorrt_llm._torch.visual_gen.models.minimax_h3.packing import MINIMAX_H3_TEXT_ENCODER_LAYER +from tensorrt_llm._torch.visual_gen.models.minimax_h3.pipeline_minimax_h3 import ( + _MiniMaxH3TextEncoderMapping, +) +from tensorrt_llm._torch.visual_gen.models.minimax_h3.text_encoder_tp import ( + MiniMaxH3TensorParallelTextEncoder, +) + + +def _assert_projection_shards( + candidate: MiniMaxH3TensorParallelTextEncoder, + reference: Qwen3VLTextModel, + group: dist.ProcessGroup, + rank: int, +) -> None: + embedding = candidate.model.model.embed_tokens + mapping = embedding.mapping + assert mapping.tp_size == 2 + assert mapping.tp_rank == rank + assert mapping.tp_group == [0, 1] + assert mapping.tp_group_pg is group + assert mapping.tp_group_name == group.group_name + assert dist.get_world_size(mapping.tp_group_pg) == 2 + assert dist.get_rank(mapping.tp_group_pg) == rank + assert mapping._use_torch_distributed_allreduce + torch.testing.assert_close( + embedding.weight, reference.embed_tokens.weight.chunk(2, dim=0)[rank], rtol=0, atol=0 + ) + assert len(candidate.model.model.layers) == MINIMAX_H3_TEXT_ENCODER_LAYER == 50 + for layer, hf_layer in zip(candidate.model.model.layers, reference.layers): + projections = ( + layer.self_attn.qkv_proj, + layer.self_attn.o_proj, + layer.mlp.gate_up_proj, + layer.mlp.down_proj, + ) + for projection in projections: + mapping = projection.mapping + assert mapping.tp_size == 2 + assert mapping.tp_rank == rank + assert mapping.tp_group == [0, 1] + assert mapping.tp_group_pg is group + assert mapping.tp_group_name == group.group_name + assert dist.get_world_size(mapping.tp_group_pg) == 2 + assert dist.get_rank(mapping.tp_group_pg) == rank + assert mapping._use_torch_distributed_allreduce + + expected_qkv = torch.cat( + [ + projection.weight.chunk(2, dim=0)[rank] + for projection in ( + hf_layer.self_attn.q_proj, + hf_layer.self_attn.k_proj, + hf_layer.self_attn.v_proj, + ) + ], + dim=0, + ) + expected_gate_up = torch.cat( + [ + projection.weight.chunk(2, dim=0)[rank] + for projection in (hf_layer.mlp.gate_proj, hf_layer.mlp.up_proj) + ], + dim=0, + ) + expected_weights = ( + expected_qkv, + hf_layer.self_attn.o_proj.weight.chunk(2, dim=1)[rank], + expected_gate_up, + hf_layer.mlp.down_proj.weight.chunk(2, dim=1)[rank], + ) + for projection, expected in zip(projections, expected_weights): + torch.testing.assert_close(projection.weight, expected, rtol=0, atol=0) + + +def _worker(rank: int, port: int, world_size: int) -> None: + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=180), + ) + try: + # Every rank creates the subgroup; only ranks 0 and 1 encode. Idle ranks + # never enter TP collectives, exposing accidental world-sized reductions. + group = ( + dist.group.WORLD + if world_size == 2 + else dist.new_group(ranks=[0, 1], timeout=timedelta(seconds=180)) + ) + if rank < 2: + device = torch.device("cuda", rank) + torch.manual_seed(19807) + config = Qwen3VLTextConfig( + vocab_size=128, + hidden_size=256, + intermediate_size=512, + num_hidden_layers=51, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=64, + max_position_embeddings=128, + attention_dropout=0.0, + rope_scaling={ + "rope_type": "default", + "mrope_section": [16, 8, 8], + "mrope_interleaved": True, + }, + ) + config._attn_implementation = "eager" + reference = Qwen3VLTextModel(config).to(device=device, dtype=torch.bfloat16).eval() + mapping = _MiniMaxH3TextEncoderMapping( + world_size=2, rank=rank, tp_size=2, gpus_per_node=2, tp_group_pg=group + ) + candidate = MiniMaxH3TensorParallelTextEncoder(config, mapping, device) + # Use the released conditional-generation checkpoint's language-model + # prefix and the real loader, including the unused final decoder layer. + candidate.load_weights( + { + f"model.language_model.{name}": value + for name, value in reference.state_dict().items() + } + ) + _assert_projection_shards(candidate, reference, group, rank) + with torch.inference_mode(): + for seq_len in (1, 7, 19): + input_ids = ((torch.arange(seq_len, device=device) * 17 + 3) % 128)[None] + position_ids = torch.arange(seq_len, device=device).view(1, 1, -1) + outputs = reference( + input_ids=input_ids, + position_ids=position_ids.expand(3, 1, -1), + use_cache=False, + output_hidden_states=True, + return_dict=True, + ) + # With 51 HF layers, hidden_states[50] is the unnormalized + # output of the first 50, not the final RMSNorm output. + expected = outputs.hidden_states[MINIMAX_H3_TEXT_ENCODER_LAYER] + first_layer_outputs: list[torch.Tensor] = [] + handle = candidate.model.model.layers[0].register_forward_hook( + lambda module, args, output: first_layer_outputs.append( + output[0] + output[1] + ) + ) + try: + actual = candidate.encode(input_ids) + finally: + handle.remove() + torch.testing.assert_close( + first_layer_outputs[0].unsqueeze(0), + outputs.hidden_states[1], + rtol=2e-2, + atol=5e-3, + ) + assert actual.shape == (1, seq_len, config.hidden_size) + assert actual.dtype == torch.bfloat16 + assert torch.isfinite(actual).all() + # BF16 GEMM and fused residual/RMSNorm rounding accumulates + # across 50 layers, even without TP. Keep a tighter first-layer + # check above and bound both outliers and aggregate error here. + torch.testing.assert_close(actual, expected, rtol=2e-2, atol=8e-2) + relative_error = ( + actual.float() - expected.float() + ).norm() / expected.float().norm() + assert relative_error < 3e-2, relative_error.item() + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("world_size", [2, 4], ids=["tp2_world2", "tp2_world4"]) +def test_minimax_h3_text_encoder_tp(monkeypatch: pytest.MonkeyPatch, world_size: int) -> None: + from ._visual_gen_dist_utils import spawn_with_retry + + if torch.cuda.device_count() < world_size: + pytest.skip(f"Requires {world_size} CUDA GPUs") + monkeypatch.setenv("TLLM_DISABLE_MPI", "1") + spawn_with_retry( + lambda port: mp.spawn(_worker, args=(port, world_size), nprocs=world_size, join=True) + ) diff --git a/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py b/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py index a968da0e04e8..92f229135c1c 100644 --- a/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_minimax_h3_pipeline.py @@ -17,14 +17,19 @@ import itertools import json +from datetime import timedelta from pathlib import Path from types import SimpleNamespace import pytest import torch +import torch.distributed as dist +import torch.multiprocessing as mp from diffusers import MiniMaxH3Scheduler from PIL import Image +from tensorrt_llm._torch.attention.attention import _inherit_process_group_attrs +from tensorrt_llm._torch.device_mesh import DeviceMeshTopologyImpl from tensorrt_llm._torch.visual_gen.config import ( DiffusionPipelineConfig, discover_pipeline_components, @@ -38,11 +43,25 @@ from tensorrt_llm._torch.visual_gen.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm._torch.visual_gen.pipeline_registry import AutoPipeline, PipelineComponent -from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.visual_gen.args import ParallelConfig, TorchCompileConfig, VisualGenArgs pytestmark = pytest.mark.cpu_only +def _ulysses_mapping(world_size: int) -> SimpleNamespace: + return SimpleNamespace( + world_size=world_size, + tp_size=1, + cfg_size=1, + ulysses_size=world_size, + cp_size=1, + ring_size=1, + attn2d_row_size=1, + attn2d_col_size=1, + ) + + class _FakeMiniMaxH3Transformer: def __init__(self) -> None: self.device = torch.device("cpu") @@ -349,9 +368,9 @@ def test_pipeline_keeps_torch_compile_enabled( def _config(torch_compile: TorchCompileConfig) -> SimpleNamespace: return SimpleNamespace( extra_attrs={}, - mapping=SimpleNamespace(world_size=1), + mapping=SimpleNamespace(world_size=1, tp_size=1), visual_gen_mapping=None, - parallel=SimpleNamespace(parallel_vae_size=1), + parallel=ParallelConfig(), attention=SimpleNamespace(backend="VANILLA"), cache=None, cpu_offload_config=SimpleNamespace(enable=False), @@ -430,9 +449,9 @@ def test_trtllm_attention_rejects_unvalidated_gpu( monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: capability) config = SimpleNamespace( extra_attrs={}, - mapping=SimpleNamespace(world_size=1), + mapping=SimpleNamespace(world_size=1, tp_size=1), visual_gen_mapping=None, - parallel=SimpleNamespace(parallel_vae_size=1), + parallel=ParallelConfig(), attention=SimpleNamespace(backend="TRTLLM"), ) with pytest.raises(NotImplementedError, match="SM100"): @@ -449,9 +468,9 @@ def test_trtllm_attention_accepts_supported_gpu( ) config = SimpleNamespace( extra_attrs={}, - mapping=SimpleNamespace(world_size=1), + mapping=SimpleNamespace(world_size=1, tp_size=1), visual_gen_mapping=None, - parallel=SimpleNamespace(parallel_vae_size=1), + parallel=ParallelConfig(), attention=SimpleNamespace(backend="TRTLLM"), cache=None, cpu_offload_config=SimpleNamespace(enable=False), @@ -460,6 +479,469 @@ def test_trtllm_attention_accepts_supported_gpu( MiniMaxH3Pipeline(config) +def test_text_encoder_residency_skips_non_owner_rank( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + h3_pipeline.BasePipeline, + "__init__", + lambda self, config: torch.nn.Module.__init__(self), + ) + monkeypatch.setattr(MiniMaxH3Pipeline, "rank", property(lambda self: 1)) + monkeypatch.setattr(MiniMaxH3Pipeline, "world_size", property(lambda self: 2)) + config = SimpleNamespace( + extra_attrs={}, + mapping=SimpleNamespace(world_size=2, tp_size=1), + visual_gen_mapping=_ulysses_mapping(2), + parallel=ParallelConfig(ulysses_size=2, text_encoder_tp_size=1), + attention=SimpleNamespace(backend="VANILLA"), + cache=None, + cpu_offload_config=SimpleNamespace(enable=False), + cuda_graph=SimpleNamespace(enable=False), + ) + pipeline = MiniMaxH3Pipeline(config) + + pipeline.load_standard_components( + "/tmp/checkpoint", + torch.device("cpu"), + skip_components=[ + PipelineComponent.TOKENIZER, + PipelineComponent.PROCESSOR, + PipelineComponent.VAE, + PipelineComponent.AUDIO_VAE, + PipelineComponent.SCHEDULER, + PipelineComponent.AUDIO_SCHEDULER, + ], + ) + + assert pipeline.text_encoder is None + + +def test_subset_tensor_parallel_text_encoder_loads_on_participant_rank( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(h3_pipeline, "mpi_disabled", lambda: True) + monkeypatch.setattr( + h3_pipeline.BasePipeline, + "__init__", + lambda self, config: torch.nn.Module.__init__(self), + ) + monkeypatch.setattr(MiniMaxH3Pipeline, "rank", property(lambda self: 1)) + monkeypatch.setattr(MiniMaxH3Pipeline, "world_size", property(lambda self: 4)) + sentinel = object() + calls = [] + + def load_text_encoder(self, checkpoint_dir: str, device: torch.device) -> object: + calls.append((checkpoint_dir, device)) + return sentinel + + monkeypatch.setattr( + MiniMaxH3Pipeline, + "_load_tensor_parallel_text_encoder", + load_text_encoder, + ) + config = SimpleNamespace( + extra_attrs={}, + mapping=SimpleNamespace(world_size=4, tp_size=1), + visual_gen_mapping=_ulysses_mapping(4), + parallel=ParallelConfig(ulysses_size=4, text_encoder_tp_size=2), + attention=SimpleNamespace(backend="VANILLA"), + cache=None, + cpu_offload_config=SimpleNamespace(enable=False), + cuda_graph=SimpleNamespace(enable=False), + ) + pipeline = MiniMaxH3Pipeline(config) + + pipeline.load_standard_components( + "/tmp/checkpoint", + torch.device("cpu"), + skip_components=[ + PipelineComponent.TOKENIZER, + PipelineComponent.PROCESSOR, + PipelineComponent.VAE, + PipelineComponent.AUDIO_VAE, + PipelineComponent.SCHEDULER, + PipelineComponent.AUDIO_SCHEDULER, + ], + ) + + assert pipeline.text_encoder is sentinel + assert calls == [("/tmp/checkpoint", torch.device("cpu"))] + + +@pytest.mark.parametrize( + "world_size, expected_prefetch", + [(2, True), (4, False)], +) +def test_tensor_parallel_text_encoder_loader_prefetch_matches_rank_residency( + monkeypatch: pytest.MonkeyPatch, + world_size: int, + expected_prefetch: bool, +) -> None: + monkeypatch.setattr( + h3_pipeline.BasePipeline, + "__init__", + lambda self, config: torch.nn.Module.__init__(self), + ) + monkeypatch.setenv("TLLM_DISABLE_MPI", "1") + monkeypatch.setattr(MiniMaxH3Pipeline, "rank", property(lambda self: 1)) + monkeypatch.setattr(MiniMaxH3Pipeline, "world_size", property(lambda self: world_size)) + calls = [] + + class FakeTextEncoder: + def load_weights(self, weights: dict) -> None: + self.weights = weights + + def eval(self) -> "FakeTextEncoder": + return self + + def from_pretrained_config(checkpoint_dir, mapping, device): + calls.append(("config", checkpoint_dir, mapping, device)) + return FakeTextEncoder() + + class FakeWeightLoader: + def __init__(self, components) -> None: + self.components = components + + def load_weights(self, checkpoint_dir, mapping, **kwargs): + calls.append(("weights", checkpoint_dir, mapping, kwargs)) + return {"weight": torch.ones(1)} + + monkeypatch.setattr( + h3_pipeline.MiniMaxH3TensorParallelTextEncoder, + "from_pretrained_config", + from_pretrained_config, + ) + monkeypatch.setattr(h3_pipeline, "WeightLoader", FakeWeightLoader) + config = SimpleNamespace( + extra_attrs={}, + mapping=SimpleNamespace(world_size=world_size, tp_size=1), + visual_gen_mapping=_ulysses_mapping(world_size), + parallel=ParallelConfig(ulysses_size=world_size, text_encoder_tp_size=2), + attention=SimpleNamespace(backend="VANILLA"), + cache=None, + cpu_offload_config=SimpleNamespace(enable=False), + cuda_graph=SimpleNamespace(enable=False), + ) + pipeline = MiniMaxH3Pipeline(config) + + text_encoder = pipeline._load_tensor_parallel_text_encoder( + "/tmp/checkpoint", torch.device("cpu") + ) + + assert isinstance(text_encoder, FakeTextEncoder) + _, _, mapping, kwargs = calls[1] + assert mapping.world_size == 2 + assert mapping.tp_size == 2 + assert kwargs == {"prefetch": expected_prefetch} + + +def test_subset_tensor_parallel_text_encoder_skips_nonparticipant_rank( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(h3_pipeline, "mpi_disabled", lambda: True) + monkeypatch.setattr( + h3_pipeline.BasePipeline, + "__init__", + lambda self, config: torch.nn.Module.__init__(self), + ) + monkeypatch.setattr(MiniMaxH3Pipeline, "rank", property(lambda self: 3)) + monkeypatch.setattr(MiniMaxH3Pipeline, "world_size", property(lambda self: 4)) + config = SimpleNamespace( + extra_attrs={}, + mapping=SimpleNamespace(world_size=4, tp_size=1), + visual_gen_mapping=_ulysses_mapping(4), + parallel=ParallelConfig(ulysses_size=4, text_encoder_tp_size=2), + attention=SimpleNamespace(backend="VANILLA"), + cache=None, + cpu_offload_config=SimpleNamespace(enable=False), + cuda_graph=SimpleNamespace(enable=False), + ) + pipeline = MiniMaxH3Pipeline(config) + + pipeline.load_standard_components( + "/tmp/checkpoint", + torch.device("cpu"), + skip_components=[ + PipelineComponent.TOKENIZER, + PipelineComponent.PROCESSOR, + PipelineComponent.VAE, + PipelineComponent.AUDIO_VAE, + PipelineComponent.SCHEDULER, + PipelineComponent.AUDIO_SCHEDULER, + ], + ) + + assert pipeline.text_encoder is None + + +@pytest.mark.parametrize("rank", range(4)) +def test_subset_text_encoder_tp_rejects_mpi_before_group_creation( + monkeypatch: pytest.MonkeyPatch, rank: int +) -> None: + monkeypatch.setattr(h3_pipeline, "mpi_disabled", lambda: False) + monkeypatch.setattr(h3_pipeline.dist, "is_initialized", lambda: True) + + def unexpected_group_creation(*args, **kwargs): + pytest.fail("Unsupported MPI subset must fail before creating a process group") + + monkeypatch.setattr(h3_pipeline.dist, "new_group", unexpected_group_creation) + config = SimpleNamespace( + mapping=SimpleNamespace(world_size=4, tp_size=1, rank=rank), + visual_gen_mapping=_ulysses_mapping(4), + parallel=ParallelConfig(ulysses_size=4, text_encoder_tp_size=2), + attention=SimpleNamespace(backend="VANILLA"), + ) + with pytest.raises(RuntimeError, match="TLLM_DISABLE_MPI=1"): + MiniMaxH3Pipeline(config) + + +@pytest.mark.parametrize("rank", [0, 1]) +@pytest.mark.parametrize("mesh_initialized", [False, True]) +def test_text_encoder_attention_mapping_preserves_explicit_tp_group( + monkeypatch: pytest.MonkeyPatch, rank: int, mesh_initialized: bool +) -> None: + monkeypatch.setenv("TLLM_DISABLE_MPI", "1") + monkeypatch.setattr( + DeviceMeshTopologyImpl, "device_mesh", object() if mesh_initialized else None + ) + group = SimpleNamespace(rank=lambda: rank) + + def unexpected_mesh_access(*args, **kwargs): + pytest.fail("Explicit encoder group must bypass the VisualGen device mesh") + + monkeypatch.setattr(DeviceMeshTopologyImpl, "build_mesh", unexpected_mesh_access) + monkeypatch.setattr(DeviceMeshTopologyImpl, "_get_mesh_dim_by_name", unexpected_mesh_access) + monkeypatch.setattr( + DeviceMeshTopologyImpl, "_get_group_ranks", lambda self, pg: [0, 1] if pg is group else [] + ) + source = SimpleNamespace(_tp_group_pg=group, _use_torch_distributed_allreduce=True) + mapping = _inherit_process_group_attrs(Mapping(world_size=2, rank=rank, tp_size=2), source) + assert mapping.tp_group_pg is group + assert mapping.tp_rank == rank + assert mapping.tp_group == [0, 1] + + +@pytest.mark.parametrize("rank", range(4)) +def test_tp_keyframes_rejected_before_encode_or_broadcast( + monkeypatch: pytest.MonkeyPatch, rank: int +) -> None: + pipeline = _SyntheticMiniMaxH3Pipeline() + pipeline.text_encoder_tp_size = 2 + monkeypatch.setattr(MiniMaxH3Pipeline, "rank", property(lambda self: rank)) + monkeypatch.setattr(MiniMaxH3Pipeline, "world_size", property(lambda self: 4)) + + def unexpected_call(*args, **kwargs): + pytest.fail("Unsupported keyframes must be rejected before encoding or broadcasting") + + monkeypatch.setattr(pipeline, "_encode_prompt_local", unexpected_call) + monkeypatch.setattr(pipeline, "_broadcast_tensor_from_rank0", unexpected_call) + with pytest.raises(NotImplementedError, match="text_encoder_tp_size=1"): + MiniMaxH3Pipeline._encode_prompt(pipeline, "test", [Image.new("RGB", (16, 16))]) + + +@pytest.mark.parametrize("rank", range(4)) +@pytest.mark.parametrize("encoder_tp,with_keyframe", [(2, False), (1, True)]) +def test_prompt_encoding_keeps_supported_owner_and_broadcast_paths( + monkeypatch: pytest.MonkeyPatch, rank: int, encoder_tp: int, with_keyframe: bool +) -> None: + pipeline = _SyntheticMiniMaxH3Pipeline() + pipeline.pipeline_config.torch_dtype = torch.float32 + pipeline.text_encoder_tp_size = encoder_tp + monkeypatch.setattr(MiniMaxH3Pipeline, "rank", property(lambda self: rank)) + monkeypatch.setattr(MiniMaxH3Pipeline, "world_size", property(lambda self: 4)) + keyframes = [Image.new("RGB", (16, 16))] if with_keyframe else [] + embeds = torch.ones(1, 2, 4) + tags = torch.zeros(2, dtype=torch.long) + encoded = [] + broadcasts = [] + broadcast_dtypes = [] + + def encode(prompt, images): + encoded.append((prompt, images)) + return embeds, tags + + def broadcast(tensor, *, dtype): + broadcasts.append(tensor) + broadcast_dtypes.append(dtype) + return tags if dtype == torch.long else embeds + + monkeypatch.setattr(pipeline, "_encode_prompt_local", encode) + monkeypatch.setattr(pipeline, "_broadcast_tensor_from_rank0", broadcast) + actual_embeds, actual_tags = MiniMaxH3Pipeline._encode_prompt(pipeline, "test", keyframes) + assert encoded == ([("test", keyframes)] if rank < encoder_tp else []) + assert len(broadcasts) == 2 + assert broadcast_dtypes == [torch.float32, torch.long] + assert broadcasts[0] is (embeds if rank < encoder_tp else None) + assert broadcasts[1] is (tags if rank < encoder_tp else None) + torch.testing.assert_close(actual_embeds, embeds) + torch.testing.assert_close(actual_tags, tags) + + +def test_single_rank_keyframe_encoding_does_not_broadcast(monkeypatch: pytest.MonkeyPatch) -> None: + pipeline = _SyntheticMiniMaxH3Pipeline() + pipeline.text_encoder_tp_size = 1 + monkeypatch.setattr(MiniMaxH3Pipeline, "world_size", property(lambda self: 1)) + keyframes = [Image.new("RGB", (16, 16))] + expected = (torch.ones(1, 2, 4), torch.zeros(2, dtype=torch.long)) + calls = [] + + def encode(prompt, images): + calls.append((prompt, images)) + return expected + + def unexpected_broadcast(*args, **kwargs): + pytest.fail("Single-rank prompt encoding must not broadcast") + + monkeypatch.setattr(pipeline, "_encode_prompt_local", encode) + monkeypatch.setattr(pipeline, "_broadcast_tensor_from_rank0", unexpected_broadcast) + assert MiniMaxH3Pipeline._encode_prompt(pipeline, "test", keyframes) is expected + assert calls == [("test", keyframes)] + + +def _text_encoder_failure_worker(rank: int, init_method: str) -> None: + dist.init_process_group( + "gloo", init_method=init_method, rank=rank, world_size=4, timeout=timedelta(seconds=60) + ) + try: + pipeline = _SyntheticMiniMaxH3Pipeline() + pipeline.pipeline_config.torch_dtype = torch.float32 + original_broadcast = pipeline._broadcast_tensor_from_rank0 + expected = torch.full((1, 2, 4), 3.0) + tags = torch.zeros(2, dtype=torch.long) + + for reference_path in (False, True): + for encoder_tp, failing_rank in ((1, 0), (2, 1)): + pipeline.text_encoder_tp_size = encoder_tp + local_error = ValueError("injected encoder failure") + failing = True + encode_calls = [] + broadcast_calls = [] + + def encode(*args): + encode_calls.append(rank) + if failing and rank == failing_rank: + raise local_error + return expected if reference_path else (expected, tags) + + def broadcast(tensor, *, dtype): + broadcast_calls.append(dtype) + return original_broadcast(tensor, dtype=dtype) + + pipeline._encode_prompt_local = encode + pipeline._encode_text_tokens = encode + pipeline._broadcast_tensor_from_rank0 = broadcast + + def run_encode(): + if reference_path: + return pipeline._encode_reference_text_tokens([11, 12], {}) + return MiniMaxH3Pipeline._encode_prompt(pipeline, "test", [])[0] + + with pytest.raises( + RuntimeError, match=f"text encoding failed on rank {failing_rank}" + ) as error: + run_encode() + assert error.value.__cause__ is (local_error if rank == failing_rank else None) + assert not broadcast_calls + assert encode_calls == ([rank] if rank < encoder_tp else []) + + # Request handling catches the error and keeps workers alive. + # A second request must complete on the same process group. + failing = False + torch.testing.assert_close(run_encode(), expected) + assert broadcast_calls == ( + [torch.float32] if reference_path else [torch.float32, torch.long] + ) + dist.barrier() + finally: + dist.destroy_process_group() + + +def test_text_encoder_failures_reach_all_ranks_and_allow_next_request(tmp_path: Path) -> None: + mp.spawn( + _text_encoder_failure_worker, + args=((tmp_path / "encoder-error-group").as_uri(),), + nprocs=4, + join=True, + ) + + +def test_single_rank_text_encoder_error_is_preserved() -> None: + pipeline = _SyntheticMiniMaxH3Pipeline() + error = ValueError("injected tokenizer failure") + with pytest.raises(ValueError) as raised: + with pipeline._coordinate_text_encoder_errors(): + raise error + assert raised.value is error + + +def test_tensor_parallel_text_encoder_encode_does_not_require_hf_config() -> None: + pipeline = _SyntheticMiniMaxH3Pipeline() + pipeline.pipeline_config.torch_dtype = torch.float32 + pipeline.text_encoder_tp_size = 2 + calls = [] + + class FakeTensorParallelTextEncoder: + def encode(self, input_ids: torch.Tensor) -> torch.Tensor: + calls.append(input_ids.clone()) + return torch.ones(1, input_ids.numel(), 4) + + pipeline.text_encoder = FakeTensorParallelTextEncoder() + + encoded = pipeline._encode_text_tokens([11, 12], {}) + + assert encoded.shape == (1, 2, 4) + assert torch.equal(calls[0], torch.tensor([[11, 12]])) + + +def test_tensor_parallel_text_encoder_rejects_qwen_vision_inputs() -> None: + pipeline = _SyntheticMiniMaxH3Pipeline() + pipeline.pipeline_config.torch_dtype = torch.float32 + pipeline.text_encoder_tp_size = 2 + pipeline.text_encoder = SimpleNamespace(encode=lambda input_ids: torch.ones(1, 1, 4)) + + with pytest.raises(NotImplementedError, match="vision inputs"): + pipeline._encode_text_tokens([11], {"pixel_values": torch.zeros(1)}) + + +def test_ref2va_text_encoding_uses_broadcast_on_non_owner_rank( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pipeline = _SyntheticMiniMaxH3Pipeline() + pipeline.pipeline_config.torch_dtype = torch.float32 + pipeline.text_encoder_tp_size = 1 + sentinel = torch.full((1, 2, 4), 3.0) + calls = [] + + monkeypatch.setattr(MiniMaxH3Pipeline, "rank", property(lambda self: 1)) + monkeypatch.setattr(MiniMaxH3Pipeline, "world_size", property(lambda self: 2)) + monkeypatch.setattr( + pipeline, + "_encode_text_tokens", + lambda token_ids, vision_inputs: calls.append((token_ids, vision_inputs)), + ) + + def broadcast(tensor: torch.Tensor | None, *, dtype: torch.dtype) -> torch.Tensor: + calls.append(("broadcast", tensor, dtype)) + return sentinel + + monkeypatch.setattr(pipeline, "_broadcast_tensor_from_rank0", broadcast) + + encoded = pipeline._encode_reference_text_tokens([11, 12], {}) + + assert encoded is sentinel + assert calls == [("broadcast", None, torch.float32)] + + +def test_ref2va_tensor_parallel_text_encoder_rejects_vision_inputs() -> None: + pipeline = _SyntheticMiniMaxH3Pipeline() + pipeline.pipeline_config.torch_dtype = torch.float32 + pipeline.text_encoder_tp_size = 2 + + with pytest.raises(NotImplementedError, match="Ref2VA"): + pipeline._encode_reference_text_tokens([11], {"pixel_values": torch.zeros(1)}) + + def test_default_generation_steps_match_reference_app() -> None: pipeline = _SyntheticMiniMaxH3Pipeline() assert pipeline.default_generation_params["num_inference_steps"] == 28 @@ -1119,15 +1601,7 @@ def _ulysses_pipeline_config(world_size: int = 8) -> SimpleNamespace: return SimpleNamespace( mapping=SimpleNamespace(world_size=1, tp_size=1), parallel=SimpleNamespace(parallel_vae_size=1), - visual_gen_mapping=SimpleNamespace( - world_size=world_size, - cfg_size=1, - ulysses_size=world_size, - cp_size=1, - ring_size=1, - attn2d_row_size=1, - attn2d_col_size=1, - ), + visual_gen_mapping=_ulysses_mapping(world_size), attention=SimpleNamespace(backend="VANILLA"), cache=None, cpu_offload_config=SimpleNamespace(enable=False), diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py index 85d2c29d4478..39c742234b9b 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py @@ -595,6 +595,16 @@ def test_parallel_vae_size_must_be_positive(self): with pytest.raises(ValidationError): ParallelConfig(parallel_vae_size=0) + def test_text_encoder_tp_size_does_not_change_worker_count(self): + pc = ParallelConfig(ulysses_size=2, text_encoder_tp_size=2) + assert pc.n_workers == 2 + pc.validate_world_size(2) + + def test_text_encoder_tp_size_validate_world_size_fails(self): + pc = ParallelConfig(text_encoder_tp_size=2) + with pytest.raises(ValueError, match="text_encoder_tp_size"): + pc.validate_world_size(1) + def test_attn2d_and_ulysses_seq_parallel_size(self): pc = ParallelConfig( attn2d_size=(2, 2), diff --git a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py index 38cf3edc82ba..5b9dbaa91d93 100644 --- a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py +++ b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py @@ -166,6 +166,9 @@ def schemas(tmp_path_factory: pytest.TempPathFactory) -> dict[str, dict]: (generator.VISUAL_GEN_SCHEMA, {"parallel_config": {"cfg_size": 1}}, True), (generator.VISUAL_GEN_SCHEMA, {"parallel_config": {"cfg_szie": 1}}, False), (generator.VISUAL_GEN_SCHEMA, {"parallel_config": {"cfg_size": 0}}, False), + (generator.VISUAL_GEN_SCHEMA, {"parallel_config": {"text_encoder_tp_size": 1}}, True), + (generator.VISUAL_GEN_SCHEMA, {"parallel_config": {"text_encoder_tp_size": 2}}, True), + (generator.VISUAL_GEN_SCHEMA, {"parallel_config": {"text_encoder_tp_size": 0}}, False), (generator.VISUAL_GEN_SCHEMA, {"attention_config": {"backend": "not-a-backend"}}, False), ], )