Searched refs:ParallelStyle (Results 1 – 7 of 7) sorted by relevance
| /external/pytorch/torch/distributed/tensor/parallel/ |
| D | api.py | 14 from torch.distributed.tensor.parallel.style import ParallelStyle
25 parallelize_plan: Union[ParallelStyle, Dict[str, ParallelStyle]], argument
84 if isinstance(parallelize_plan, ParallelStyle):
|
| D | style.py | 22 class ParallelStyle(ABC): class
35 class ColwiseParallel(ParallelStyle):
145 class RowwiseParallel(ParallelStyle):
251 class SequenceParallel(ParallelStyle):
333 class PrepareModuleInput(ParallelStyle):
463 class PrepareModuleOutput(ParallelStyle):
|
| D | __init__.py | 6 ParallelStyle,
|
| /external/pytorch/test/distributed/_tensor/experimental/ |
| D | test_tp_transform.py | 11 ParallelStyle,
104 parallel_strategies: Dict[str, ParallelStyle] = {
142 parallel_strategies: Dict[str, ParallelStyle] = {
|
| /external/pytorch/docs/source/ |
| D | distributed.tensor.parallel.rst | 38 parameters to DTensors, the following ``ParallelStyle`` s can be used in
50 ``ParallelStyle`` s, we assume the input/output activation tensors are evenly sharded on
55 and use ``use_local_output=False`` to return DTensor after each ``ParallelStyle``, where
|
| /external/pytorch/torch/distributed/tensor/experimental/ |
| D | _tp_transform.py | 17 from torch.distributed.tensor.parallel.style import ColwiseParallel, ParallelStyle
39 parallel_strategies: Dict[str, ParallelStyle], argument
82 parallel_strategies: Dict[str, ParallelStyle], argument
109 parallel_strategies: Dict[str, ParallelStyle], argument
|
| D | _attention.py | 29 from torch.distributed.tensor.parallel.style import ParallelStyle
626 class _AttentionContextParallel(ParallelStyle):
|