encoder_split ¶

def h3_broadcast_condition(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> None:
    """Send the conditioning output to every rank over the world NCCL group.

    Only world rank 0 encodes a request (one prompt presentation per request),
    so this runs on rank 0 as the source. Other encoder ranks mirror the same
    collectives through ``h3_receive_condition`` and drop the payload; skipping
    them entirely would desynchronize the world-group broadcasts for the
    denoise ranks.
    """
    # Imported lazily: stages/__init__ imports the decoding stage, which imports
    # this module, so a top-level import would close that cycle mid-init.
    from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_conditioning import MINIMAX_H3_TEXT_TOKEN_TAGS_KEY

    world_group = get_world_group()
    if world_group.rank != 0:
        raise RuntimeError("h3_broadcast_condition must run on the primary encoder worker (rank 0).")
    if not batch.prompt_embeds:
        raise RuntimeError("MiniMax-H3 encoder worker has no prompt_embeds to send.")
    tags = batch.extra.get(MINIMAX_H3_TEXT_TOKEN_TAGS_KEY)
    if tags is None:
        raise RuntimeError("MiniMax-H3 encoder worker has no text token tags to send.")

    device = get_local_torch_device()
    meta = {
        "embeds": _embed_specs(batch.prompt_embeds),
        "num_tags": int(tags.numel()),
    }
    world_group.broadcast_object(meta, src=0)
    for tensor in batch.prompt_embeds:
        world_group.broadcast(tensor.contiguous(), src=0)
    tags_on_device = tags.to(device=device, dtype=torch.long).contiguous()
    world_group.broadcast(tags_on_device, src=0)
    logger.info("MiniMax-H3 encoder worker broadcast prompt_embeds: %s", _embed_specs(batch.prompt_embeds))

def h3_denoise_device_mesh(fastvideo_args: FastVideoArgs) -> DeviceMesh:
    """FSDP mesh covering exactly the denoise ranks.

    ``init_device_mesh`` always lays a mesh over ranks ``0..numel-1``, but the
    encoder group owns the leading ranks, so the mesh is built from an explicit
    rank map. Only the denoise ranks build it: the DiT exists nowhere else, and
    the groups it creates are local to those ranks.
    """
    from fastvideo.platforms import current_platform

    start = h3_encoder_worker_count(fastvideo_args)
    world_size = get_world_group().world_size
    device_type = "npu" if current_platform.is_npu() else "cuda"
    ranks = torch.arange(start, world_size, dtype=torch.int).view(1, -1)
    logger.info("MiniMax-H3 encoder split: FSDP mesh over denoise ranks %d..%d", start, world_size - 1)
    return DeviceMesh(device_type, ranks, mesh_dim_names=("replicate", "shard"))

def h3_encoder_worker_count(fastvideo_args: FastVideoArgs) -> int:
    """Number of world ranks reserved for the encoder group."""
    world_size = get_world_group().world_size
    count, unit = _resolve_encoder_count(fastvideo_args)
    if not 1 <= count < world_size:
        raise ValueError(f"MiniMax-H3 encoder split needs 1 <= encoder {unit} < world size, got "
                         f"encoder={count}, world_size={world_size}.")
    return count

def h3_is_output_worker(fastvideo_args: FastVideoArgs) -> bool:
    """Whether this rank owns the decoded output (see ``h3_output_worker_rank``)."""
    return get_world_group().rank == h3_output_worker_rank(fastvideo_args)

def h3_is_primary_encoder_worker(fastvideo_args: FastVideoArgs) -> bool:
    """Rank 0 encodes every request; ranks 1..n-1 only mirror the broadcast."""
    return get_world_group().rank == 0

def h3_output_worker_rank(fastvideo_args: FastVideoArgs) -> int:
    """World rank that owns the decoded output: world rank 0 normally, first denoise rank in split."""
    if not h3_encoder_split_enabled(fastvideo_args):
        return get_world_group().first_rank
    return h3_encoder_worker_count(fastvideo_args)

def h3_prepare_split_worker_parallelism(fastvideo_args: FastVideoArgs, rank: int,
                                        world_size: int) -> tuple[int, list[list[int]], list[list[int]]]:
    """Resolve per-rank SP size and the custom SP/DP group layout for split mode.

    Called by the worker before ``maybe_init_distributed_environment_and_model_parallel``,
    and stamped onto the worker-local ``fastvideo_args`` so the consistency check in
    ``ComposedPipelineBase.__init__`` sees the same per-rank SP size.
    """
    if fastvideo_args.distributed_executor_backend != "ray":
        raise ValueError("MiniMax-H3 encoder split (h3_encoder_split) requires distributed_executor_backend='ray'.")
    count, unit = _resolve_encoder_count(fastvideo_args)
    if not 1 <= count < world_size:
        raise ValueError(f"MiniMax-H3 encoder split needs 1 <= encoder {unit} < num_gpus, got "
                         f"encoder={count}, world={world_size}.")
    tp_size = max(1, int(getattr(fastvideo_args, "tp_size", 1) or 1))
    if tp_size > 1 and (count % tp_size or (world_size - count) % tp_size):
        # TP groups are consecutive ranks, so a group straddling the boundary
        # would all-reduce over ranks that never run the TP-sharded conditioner.
        raise ValueError(f"MiniMax-H3 encoder split: the encoder ({count}) and denoise ({world_size - count}) "
                         f"group sizes must both be multiples of tp_size ({tp_size}).")
    from fastvideo.fastvideo_args import h3_split_sp_error, probe_h3_attention_heads

    heads = probe_h3_attention_heads(fastvideo_args.model_path)
    if heads and heads % (world_size - count):
        raise ValueError(h3_split_sp_error(heads, world_size - count, world_size, unit))
    if rank < count:
        sp_size = 1
    else:
        sp_size = world_size - count
        if int(getattr(fastvideo_args, "h3_sequential_load", False) or False):
            logger.info("MiniMax-H3 encoder split: sequential module load is redundant on denoise ranks and stays off")
    sp_group_ranks = [[r] for r in range(count)] + [list(range(count, world_size))]
    dp_group_ranks = [[r] for r in range(world_size)]
    fastvideo_args.sp_size = sp_size
    fastvideo_args.h3_encoder_workers = count
    # The denoise ranks are the FSDP group when FSDP inference is on: the encoder
    # ranks never load the DiT, so the mesh must not span the world group.
    # ``h3_denoise_device_mesh`` builds it from an explicit rank map.
    fastvideo_args.hsdp_shard_dim = world_size - count
    fastvideo_args.hsdp_replicate_dim = 1
    logger.info("MiniMax-H3 split: rank %d is an %s worker (sp_size=%d, denoise ranks %d..%d)",
                rank,
                "encoder" if rank < count else "denoise",
                sp_size,
                count,
                world_size - 1,
                local_main_process_only=False)
    return sp_size, sp_group_ranks, dp_group_ranks

def h3_receive_condition(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> None:
    """Materialize the encoder worker's conditioning output on a denoise rank."""
    from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_conditioning import MINIMAX_H3_TEXT_TOKEN_TAGS_KEY

    world_group = get_world_group()
    device = get_local_torch_device()
    meta = world_group.broadcast_object(None, src=0)
    embeds: list[torch.Tensor] = []
    for spec in meta["embeds"]:
        dtype = _DTYPE_BY_NAME.get(spec["dtype"])
        if dtype is None:
            raise RuntimeError(f"Unsupported prompt_embeds dtype from encoder worker: {spec['dtype']}")
        tensor = torch.empty(tuple(spec["shape"]), dtype=dtype, device=device)
        world_group.broadcast(tensor, src=0)
        embeds.append(tensor)
    tags = torch.empty(int(meta["num_tags"]), dtype=torch.long, device=device)
    world_group.broadcast(tags, src=0)
    batch.prompt_embeds = embeds
    batch.extra[MINIMAX_H3_TEXT_TOKEN_TAGS_KEY] = tags.cpu()
    logger.info("MiniMax-H3 split worker received prompt_embeds: %s",
                _embed_specs(embeds),
                local_main_process_only=False)