DiffSynth-Studio/diffsynth/distributed/__init__.py at 52049293d2f23173bf7b1d48531c4b9fd7971235 · modelscope/DiffSynth-Studio · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
"""
Distributed computing support for DiffSynth-Studio.

Provides multi-GPU support through:
- Data Parallel (DP): Batch-level parallelism
- Tensor Parallel (TP): Layer-level parallelism
- Pipeline utilities for distributed inference
"""

from .parallel import (
    init_distributed,
    cleanup_distributed,
    get_rank,
    get_world_size,
    is_distributed,
    is_main_process,
    barrier,
    broadcast,
    all_reduce,
    all_gather,
)

from .tensor_parallel import (
    TensorParallelLinear,
    ColumnParallelLinear,
    RowParallelLinear,
    split_tensor_along_dim,
    gather_tensor_along_dim,
)

from .data_parallel import (
    DataParallelPipeline,
    scatter_batch,
    gather_outputs,
)

from .multi_gpu import (
    MultiGPUPipeline,
    auto_device_map,
    get_optimal_device_map,
)

__all__ = [
    # Initialization
    "init_distributed",
    "cleanup_distributed",
    "get_rank",
    "get_world_size",
    "is_distributed",
    "is_main_process",
    # Communication
    "barrier",
    "broadcast",
    "all_reduce",
    "all_gather",
    # Tensor Parallel
    "TensorParallelLinear",
    "ColumnParallelLinear",
    "RowParallelLinear",
    "split_tensor_along_dim",
    "gather_tensor_along_dim",
    # Data Parallel
    "DataParallelPipeline",
    "scatter_batch",
    "gather_outputs",
    # Multi-GPU
    "MultiGPUPipeline",
    "auto_device_map",
    "get_optimal_device_map",
]