tensorplay.distributed.tensor.parallel API

Functions 2

tensorplay.distributed.tensor.parallel.parallelize_module(module: Any, device_mesh: Any = None, parallelize_plan: ParallelStyle | dict[str, ParallelStyle] | None = None, *, src_data_rank: int | None = 0) → Any[source]

class tensorplay.distributed.tensor.parallel.ColwiseParallel(*, input_layouts: Placement | None = None, output_layouts: Placement | None = None, use_local_output: bool = True)[source]

class tensorplay.distributed.tensor.parallel.PrepareModuleInput(*, input_layouts: Placement | tuple[Placement | None, ...] | None = None, desired_input_layouts: Placement | tuple[Placement | None, ...] | None = None, input_kwarg_layouts: dict[str, Placement] | None = None, desired_input_kwarg_layouts: dict[str, Placement] | None = None, use_local_output: bool = False)[source]

class tensorplay.distributed.tensor.parallel.PrepareModuleInputOutput(*, input_layouts: Placement | tuple[Placement | None, ...] | None = None, desired_input_layouts: Placement | tuple[Placement | None, ...] | None = None, input_kwarg_layouts: dict[str, Placement] | None = None, desired_input_kwarg_layouts: dict[str, Placement] | None = None, use_local_input: bool = False, output_layouts: Placement | tuple[Placement | None, ...], desired_output_layouts: Placement | tuple[Placement, ...], use_local_output: bool = True)[source]

class tensorplay.distributed.tensor.parallel.PrepareModuleOutput(*, output_layouts: Placement | tuple[Placement | None, ...], desired_output_layouts: Placement | tuple[Placement, ...], use_local_output: bool = True)[source]

class tensorplay.distributed.tensor.parallel.RowwiseParallel(*, input_layouts: Placement | None = None, output_layouts: Placement | None = None, use_local_output: bool = True)[source]

class tensorplay.distributed.tensor.parallel.SequenceParallel(*, sequence_dim: int = 1, use_local_output: bool = False)[source]