chore: import upstream snapshot with attribution

2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
@@ -0,0 +1,84 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+from dataclasses import dataclass
+from deepspeed import comm as dist
+from deepspeed.constants import CROSS_RANK, CROSS_SIZE, LOCAL_RANK
+from .data_parallel_writer_factory import DataParallelWriterFactory
+
+# TODO: parse socket number from env.
+SOCKETS_PER_MACHINE = 2
+
+
+@dataclass
+class MPUInfo(object):
+    pp_world_size: int
+    pp_rank: int
+    tp_world_size: int
+    tp_rank: int
+    dp_world_size: int
+    dp_peer_ranks: list
+    dp_rank: int
+
+
+def _create_model_parallel_info(mpu):
+    return MPUInfo(pp_world_size=mpu.get_pipeline_model_parallel_world_size(),
+                   pp_rank=mpu.get_pipeline_model_parallel_rank(),
+                   tp_world_size=mpu.get_tensor_model_parallel_world_size(),
+                   tp_rank=mpu.get_tensor_model_parallel_rank(),
+                   dp_world_size=mpu.get_data_parallel_world_size(),
+                   dp_peer_ranks=mpu.get_data_parallel_group_ranks(),
+                   dp_rank=mpu.get_data_parallel_rank())
+
+
+@dataclass
+class ExpertParallelInfo(object):
+    ep_world_size: int
+    ep_rank: int
+    dp_world_size: int
+    dp_peer_ranks: list
+    dp_rank: int
+
+
+def _create_expert_parallel_info(groups):
+    group_name = groups._get_max_expert_size_name()
+    return ExpertParallelInfo(ep_world_size=groups._get_expert_parallel_world_size(group_name),
+                              ep_rank=groups._get_expert_parallel_rank(group_name),
+                              dp_world_size=groups._get_expert_data_parallel_world_size(group_name),
+                              dp_peer_ranks=groups._get_expert_data_parallel_group_ranks(group_name),
+                              dp_rank=groups._get_expert_data_parallel_rank(group_name))
+
+
+@dataclass
+class UniversalParallelInfo(object):
+    global_world_size: int
+    global_rank: int
+    local_rank: int
+    mpu_info: MPUInfo
+    ep_info: ExpertParallelInfo
+    pure_dp: bool
+    num_machines: int
+    machine_rank: int
+    num_sockets: int
+
+
+def create_universal_parallel_info(groups, has_moe_layers):
+    return UniversalParallelInfo(global_world_size=dist.get_world_size(),
+                                 global_rank=dist.get_rank(),
+                                 local_rank=int(os.environ[LOCAL_RANK]),
+                                 mpu_info=None if groups.mpu is None else _create_model_parallel_info(groups.mpu),
+                                 ep_info=_create_expert_parallel_info(groups) if has_moe_layers else None,
+                                 pure_dp=groups.mpu is None
+                                 or groups.mpu.get_data_parallel_world_size() == dist.get_world_size(),
+                                 num_machines=int(os.environ[CROSS_SIZE]),
+                                 machine_rank=int(os.environ[CROSS_RANK]),
+                                 num_sockets=int(os.environ[CROSS_SIZE]) * SOCKETS_PER_MACHINE)
+
+
+def create_data_parallel_writer_config(groups, parallel_unit, zero_stage, has_moe_layers):
+    uni_parallel_info = create_universal_parallel_info(groups, has_moe_layers)
+    writer_factory = DataParallelWriterFactory(uni_parallel_info, parallel_unit)
+    return writer_factory.create_config(zero_stage, has_moe_layers)