Class ParallelConfig

src/diffusers/models/_modeling_parallel.py:158–187

Configuration for applying different parallelisms. Args: context_parallel_config (`ContextParallelConfig`, *optional*): Configuration for context parallelism.

@dataclass
class ParallelConfig:
    """
    Configuration for applying different parallelisms.

    Args:
        context_parallel_config (`ContextParallelConfig`, *optional*):
            Configuration for context parallelism.
    """

    context_parallel_config: ContextParallelConfig | None = None

    _rank: int = None
    _world_size: int = None
    _device: torch.device = None
    _mesh: torch.distributed.device_mesh.DeviceMesh = None

    def setup(
        self,
        rank: int,
        world_size: int,
        device: torch.device,
        *,
        mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    ):
        self._rank = rank
        self._world_size = world_size
        self._device = device
        self._mesh = mesh
        if self.context_parallel_config is not None:
            self.context_parallel_config.setup(rank, world_size, device, mesh)


@dataclass(frozen=True)