Initialize Paddle Fleet and get rank of worker
(seed: int = 20)
| 108 | |
| 109 | |
| 110 | def init_distributed_environment(seed: int = 20) -> Tuple[int, int]: |
| 111 | """Initialize Paddle Fleet and get rank of worker""" |
| 112 | # Global rank |
| 113 | ranks = dist.get_world_size() |
| 114 | dist_strategy = fleet.DistributedStrategy() |
| 115 | if ranks > 0: |
| 116 | dist_strategy.hybrid_configs = { |
| 117 | "dp_degree": 1, |
| 118 | "mp_degree": ranks, |
| 119 | "pp_degree": 1, |
| 120 | "sharding_degree": 1, |
| 121 | } |
| 122 | |
| 123 | # Set control in tensor parallel |
| 124 | dist_strategy.tensor_parallel_configs = {"tensor_init_seed": seed} |
| 125 | fleet.init(is_collective=True, strategy=dist_strategy) |
| 126 | |
| 127 | # Local rank |
| 128 | local_rank = fleet.worker_index() |
| 129 | else: |
| 130 | local_rank = 0 |
| 131 | return ranks, local_rank |
| 132 | |
| 133 | |
| 134 | def update_fd_config_for_mm(fd_config: FDConfig) -> None: |
no outgoing calls