|
51 | 51 | from .collective import P2POp # noqa: F401
|
52 | 52 | from .collective import reduce_scatter # noqa: F401
|
53 | 53 |
|
54 |
| -from .communication import * # noqa: F401 |
| 54 | +from .communication import stream |
55 | 55 |
|
56 | 56 | from .auto_parallel import shard_op # noqa: F401
|
57 | 57 | from .auto_parallel import shard_tensor # noqa: F401
|
|
66 | 66 |
|
67 | 67 | from . import cloud_utils # noqa: F401
|
68 | 68 |
|
69 |
| -from .sharding import * # noqa: F401 |
| 69 | +from .sharding import group_sharded_parallel, save_group_sharded_model |
70 | 70 |
|
71 | 71 | __all__ = [ # noqa
|
72 |
| - "spawn", "launch", "scatter", "broadcast", "ParallelEnv", "new_group", |
73 |
| - "init_parallel_env", "gloo_init_parallel_env", "gloo_barrier", |
74 |
| - "gloo_release", "QueueDataset", "split", "CountFilterEntry", |
75 |
| - "ShowClickEntry", "get_world_size", "get_group", "all_gather", |
76 |
| - "all_gather_object", "InMemoryDataset", "barrier", "all_reduce", "alltoall", |
77 |
| - "alltoall_single", "send", "reduce", "recv", "ReduceOp", "wait", "get_rank", |
78 |
| - "ProbabilityEntry", "ParallelMode", "is_initialized", |
79 |
| - "destroy_process_group", "isend", "irecv", "reduce_scatter", "stream" |
| 72 | + "spawn", |
| 73 | + "launch", |
| 74 | + "scatter", |
| 75 | + "broadcast", |
| 76 | + "ParallelEnv", |
| 77 | + "new_group", |
| 78 | + "init_parallel_env", |
| 79 | + "gloo_init_parallel_env", |
| 80 | + "gloo_barrier", |
| 81 | + "gloo_release", |
| 82 | + "QueueDataset", |
| 83 | + "split", |
| 84 | + "CountFilterEntry", |
| 85 | + "ShowClickEntry", |
| 86 | + "get_world_size", |
| 87 | + "get_group", |
| 88 | + "all_gather", |
| 89 | + "all_gather_object", |
| 90 | + "InMemoryDataset", |
| 91 | + "barrier", |
| 92 | + "all_reduce", |
| 93 | + "alltoall", |
| 94 | + "alltoall_single", |
| 95 | + "send", |
| 96 | + "reduce", |
| 97 | + "recv", |
| 98 | + "ReduceOp", |
| 99 | + "wait", |
| 100 | + "get_rank", |
| 101 | + "ProbabilityEntry", |
| 102 | + "ParallelMode", |
| 103 | + "is_initialized", |
| 104 | + "destroy_process_group", |
| 105 | + "isend", |
| 106 | + "irecv", |
| 107 | + "reduce_scatter", |
80 | 108 | ]
|
0 commit comments