minimax_h3_parallel ¶

def decode_to_pixels_parallel(
    vae: AutoencoderKLMiniMaxH3,
    z: torch.Tensor,
    output: torch.Tensor | None,
    group: "GroupCoordinator",
    strategy: str = DEFAULT_DECODE_GATHER_STRATEGY,
) -> torch.Tensor | None:
    """Chunk-parallel ``decode_to_pixels`` across a sequence-parallel group.

    All group ranks call this together with identical ``z``. Temporal chunks
    are decoded round-robin across the group and their segments move to the
    group's first rank, which assembles bitwise the serial
    ``decode_to_pixels`` result into ``output``. Only the first rank passes
    ``output`` (validated exactly like the serial API); other ranks pass
    ``None`` and receive ``None``.
    """
    if strategy not in DECODE_GATHER_STRATEGIES:
        raise ValueError(f"Unknown parallel-decode strategy {strategy!r}; expected one of {DECODE_GATHER_STRATEGIES}.")
    is_leader = group.rank_in_group == 0
    if is_leader:
        if output is None:
            raise ValueError("The first sequence-parallel rank must provide the CPU output buffer.")
        expected_shape = vae.decoded_pixel_shape(z.shape)
        if output.device.type != "cpu" or output.dtype != torch.float32 or tuple(output.shape) != expected_shape:
            raise ValueError(
                "`output` must be a CPU float32 tensor with shape "
                f"{expected_shape}, got device={output.device}, dtype={output.dtype}, shape={tuple(output.shape)}.")
    elif output is not None:
        raise ValueError("Only the first sequence-parallel rank may provide an output buffer.")
    if group.world_size == 1:
        return vae.decode_to_pixels(z, output)

    try:
        if vae.use_slicing and z.shape[0] > 1:
            for batch_index, z_slice in enumerate(z.split(1)):
                slice_output = output[batch_index:batch_index + 1] if output is not None else None
                _decode_single_parallel(vae, z_slice, slice_output, group, strategy)
        else:
            _decode_single_parallel(vae, z, output, group, strategy)
    finally:
        # Drain the leader's async chunk copies before the caller (or an
        # exception handler) can read or release the pinned buffer.
        if output is not None and vae._streams_chunk_copies(z, output):
            torch.cuda.current_stream(z.device).synchronize()
    return output

def encode_pixels_parallel(
    vae: AutoencoderKLMiniMaxH3,
    pixels: torch.Tensor,
    group: "GroupCoordinator",
) -> AutoencoderKLOutput:
    """Clip-parallel ``encode_pixels`` across a sequence-parallel group.

    Encoder clips have no cross-clip dependency (no overlap, no blending), so
    ranks encode disjoint clips and all-gather the per-clip moment tensors.
    Every rank returns the identical full posterior — preserving the serial
    contract that all ranks hold the same encoded latents — bitwise equal to
    ``vae.encode_pixels(pixels)``. Moments are latent-sized (a few MB per
    clip), so the all-gather is negligible next to the clip forwards.
    """
    if pixels.ndim != 5 or pixels.shape[1] != vae.config.in_channels or pixels.shape[2] <= 0:
        raise ValueError(
            f"`pixels` must have shape [B, {vae.config.in_channels}, T, H, W] with T > 0, "
            f"got {tuple(pixels.shape)}.")
    if pixels.device.type != "cpu":
        raise ValueError(f"`pixels` must remain on CPU, got device={pixels.device}.")
    if pixels.dtype != torch.uint8 and not pixels.is_floating_point():
        raise TypeError(f"`pixels` must use uint8 or a floating-point dtype, got {pixels.dtype}.")
    if group.world_size == 1:
        return vae.encode_pixels(pixels)
    if vae.use_slicing and pixels.shape[0] > 1:
        moments = torch.cat([_encode_single_parallel(vae, pixel_slice, group) for pixel_slice in pixels.split(1)])
    else:
        moments = _encode_single_parallel(vae, pixels, group)
    return AutoencoderKLOutput(latent_dist=DiagonalGaussianDistribution(moments))

def parallel_chunk_indices(num_chunks: int, world_size: int, rank_in_group: int) -> list[int]:
    """Round-robin chunk ownership: chunk ``i`` belongs to rank ``i % world_size``."""
    if num_chunks < 0:
        raise ValueError(f"num_chunks must be non-negative, got {num_chunks}.")
    if world_size < 1:
        raise ValueError(f"world_size must be positive, got {world_size}.")
    if not 0 <= rank_in_group < world_size:
        raise ValueError(f"rank_in_group {rank_in_group} out of range for world_size {world_size}.")
    return list(range(rank_in_group, num_chunks, world_size))