From 8028c4c35250fc530663d5edb58914b8583d0665 Mon Sep 17 00:00:00 2001 From: Chaoyang He Date: Sat, 25 Sep 2021 12:57:56 -0700 Subject: [PATCH 1/7] Update engine.py --- deepspeed/runtime/engine.py | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index b2d9c83eb2b3..b8ba9786c99b 100755 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -138,6 +138,7 @@ def __init__(self, self.block_eigenvalue = None self.gas_boundary_ctr = 0 self.dist_backend = "nccl" + self.excluded_weights_for_ddp_sync = set() self.has_moe_layers = False self.num_experts = None @@ -246,6 +247,9 @@ def __init__(self, self.flatten = util_ops.flatten self.unflatten = util_ops.unflatten + if mpu is not None and hasattr(mpu, "get_excluded_weights_for_ddp_sync"): + self.excluded_weights_for_ddp_sync = mpu.get_excluded_weights_for_ddp_sync(model) + def get_batch_info(self): """ Get all training batch related settings. @@ -1687,17 +1691,19 @@ def buffered_allreduce_fallback(self, grads=None, elements_per_buffer=500000000) grads.append(param.grad.data) else: grad_data = param.grad.data - if self.sparse_gradients_enabled( - ) and param_name in self.csr_tensor_module_names: - if is_moe_param: - expert_grads.append(CSRTensor(grad_data)) - else: - grads.append(CSRTensor(grad_data)) - else: - if is_moe_param: - expert_grads.append(grad_data) + if not self.excluded_weights_for_ddp_sync.__contains__(param_name): + grads.append(grad_data) + if self.sparse_gradients_enabled( + ) and param_name in self.csr_tensor_module_names: + if is_moe_param: + expert_grads.append(CSRTensor(grad_data)) + else: + grads.append(CSRTensor(grad_data)) else: - grads.append(grad_data) + if is_moe_param: + expert_grads.append(grad_data) + else: + grads.append(grad_data) split_buckets = split_half_float_double_csr(grads) for i, bucket_tuple in enumerate(split_buckets): From 2918360e93dc928f95a019c1673e59d9d199c0ac Mon Sep 17 00:00:00 2001 From: Chaoyang He Date: Sat, 25 Sep 2021 13:26:29 -0700 Subject: [PATCH 2/7] add more arguments: nnodes/num_local_procs/node_rank --- deepspeed/launcher/launch.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/deepspeed/launcher/launch.py b/deepspeed/launcher/launch.py index 2ad98ea358a5..702cd8422e19 100755 --- a/deepspeed/launcher/launch.py +++ b/deepspeed/launcher/launch.py @@ -122,7 +122,10 @@ def main(): sys.executable, "-u", args.training_script, - "--local_rank={}".format(local_rank) + "--local_rank={}".format(local_rank), + "--nnodes={}".format(args.nnodes), + "--num_local_procs={}".format(num_local_procs), + "--node_rank={}".format(args.node_rank), ] + args.training_script_args sig_names = {2: "SIGINT", 15: "SIGTERM"} From 796070faf9add60121056b0a4467e56ac2431c7b Mon Sep 17 00:00:00 2001 From: Chaoyang He Date: Sat, 25 Sep 2021 15:19:50 -0700 Subject: [PATCH 3/7] add group branch without MoE --- deepspeed/runtime/engine.py | 4 +-- deepspeed/utils/groups.py | 56 ++++++++++++++++++++++++++++++++----- 2 files changed, 51 insertions(+), 9 deletions(-) diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index b8ba9786c99b..43b52029aae5 100755 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -750,11 +750,11 @@ def _configure_distributed_model(self, model): assert self.mpu.get_model_parallel_world_size() == groups.get_model_parallel_world_size(), "mpu object provided must match mpu object provided to groups.initialize()" else: # Scenario 3 - groups.initialize(mpu=self.mpu) + groups.initialize(mpu=self.mpu, self.has_moe_layers) else: if not groups.is_initialized(): # Scenario 1 - groups.initialize() + groups.initialize(self.has_moe_layers) #else: # Scenario 2 # Scenario 4 - Case 2 diff --git a/deepspeed/utils/groups.py b/deepspeed/utils/groups.py index dcd774d76501..242d5dec0bf6 100644 --- a/deepspeed/utils/groups.py +++ b/deepspeed/utils/groups.py @@ -17,6 +17,8 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import logging + """ Support, expert, data, and model (only megatron-style) parallelism in DeepSpeed @@ -68,7 +70,7 @@ def ensure_divisibility(numerator, denominator): numerator, denominator) -def initialize(ep_size=1, mpu=None): +def initialize(ep_size=1, mpu=None, has_moe_layers=False): """ Process groups initialization supporting expert (E), data (D), and model (M) parallelism. DeepSpeed considers the following scenarios w.r.t. process group creation. @@ -103,13 +105,21 @@ def initialize(ep_size=1, mpu=None): that descibes model/data parallel ranks. """ - if mpu is not None: - log_dist(message="initializing deepspeed groups using mpu", ranks=[0]) - initialize_model_and_expert_parallel(ep_size, mpu) + if has_moe_layers: + if mpu is not None: + log_dist(message="initializing deepspeed groups using mpu", ranks=[0]) + initialize_model_and_expert_parallel(ep_size, mpu) + else: + log_dist(message="initializing deepspeed groups", ranks=[0]) + initialize_model_parallel(1) + initialize_expert_parallel(ep_size) else: - log_dist(message="initializing deepspeed groups", ranks=[0]) - initialize_model_parallel(1) - initialize_expert_parallel(ep_size) + if mpu is not None: + log_dist(message="initializing data and model parallel groups using mpu", ranks=[0]) + initialize_data_and_model_parallel(mpu) + else: + log_dist(message="initializing deepspeed groups", ranks=[0]) + initialize_model_parallel(1) def initialize_model_parallel(model_parallel_size_): @@ -216,6 +226,35 @@ def initialize_expert_parallel(expert_parallel_size_): _EXPERT_PARALLEL_GROUP = group +def initialize_data_and_model_parallel(mpu): + """ + * S3: There is model parallelism but no expert parallelism (M):: + + mpu.init() # client initializes its model parallel unit + model = my_model(args) + engine = deepspeed.initialize(model, mpu=mpu) # init w. mpu but ep_size = dp_world_size + """ + assert torch.distributed.is_initialized(), "torch distributed is not initialized" + assert mpu.model_parallel_is_initialized(), "model parallel group is not initialized" + model_parallel_size_ = mpu.get_model_parallel_world_size() + + world_size = torch.distributed.get_world_size() + rank = torch.distributed.get_rank() + dp_world_size = mpu.get_data_parallel_world_size() + dp_rank = mpu.get_data_parallel_rank() + + log_dist( + f"Initializing deepspeed groups with model parallel size {model_parallel_size_}, expert parallel size {expert_parallel_size_}, and data parallel size {world_size}", + [0]) + + global _DATA_PARALLEL_GROUP, _MODEL_PARALLEL_GROUP + global _EXPERT_PARALLEL_GROUP, _EXPERT_DATA_PARALLEL_GROUP + + # Get world size and rank. Ensure some consistencies. + _DATA_PARALLEL_GROUP = mpu.get_data_parallel_group() + _MODEL_PARALLEL_GROUP = mpu.get_model_parallel_group() + + def initialize_model_and_expert_parallel(expert_parallel_size_, mpu): """ Initialize Expert groups based on MPU groups. @@ -251,6 +290,7 @@ def initialize_model_and_expert_parallel(expert_parallel_size_, mpu): expert_parallel_size_ = min(expert_parallel_size_, dp_world_size) ensure_divisibility(world_size, expert_parallel_size_) + logging.info("expert_parallel_size_ = {}".format(expert_parallel_size_)) # Build the expert data parallel groups. assert _EXPERT_DATA_PARALLEL_GROUP is None, \ @@ -272,6 +312,7 @@ def initialize_model_and_expert_parallel(expert_parallel_size_, mpu): [0]) if rank in list(ranks): _EXPERT_DATA_PARALLEL_GROUP = group + logging.info("create expert data parallel process group. MP {}, EP {}, Ranks = {}".format(j, i, ranks)) for i in range(dp_world_size // expert_parallel_size_): ranks = range(i * expert_parallel_size_ * model_parallel_size_ + j, @@ -284,6 +325,7 @@ def initialize_model_and_expert_parallel(expert_parallel_size_, mpu): [0]) if rank in list(ranks): _EXPERT_PARALLEL_GROUP = group + logging.info("creating expert parallel process group. MP {}, EP {}, Ranks = {}".format(j, i, ranks)) def is_initialized(): From b7ac656f43475958530034c4fcb5640c74cb876a Mon Sep 17 00:00:00 2001 From: Chaoyang He Date: Sat, 25 Sep 2021 15:25:00 -0700 Subject: [PATCH 4/7] Update engine.py --- deepspeed/runtime/engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 43b52029aae5..74117d8fb0e1 100755 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -750,7 +750,7 @@ def _configure_distributed_model(self, model): assert self.mpu.get_model_parallel_world_size() == groups.get_model_parallel_world_size(), "mpu object provided must match mpu object provided to groups.initialize()" else: # Scenario 3 - groups.initialize(mpu=self.mpu, self.has_moe_layers) + groups.initialize(mpu=self.mpu, has_moe_layers=self.has_moe_layers) else: if not groups.is_initialized(): # Scenario 1 From de9a084fb28061c6506ec0a8070f701ddad25697 Mon Sep 17 00:00:00 2001 From: Chaoyang He Date: Sat, 25 Sep 2021 15:37:02 -0700 Subject: [PATCH 5/7] add API "add_excluded_weights_for_ddp_sync" --- deepspeed/runtime/engine.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 74117d8fb0e1..73464d4af39d 100755 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -138,7 +138,6 @@ def __init__(self, self.block_eigenvalue = None self.gas_boundary_ctr = 0 self.dist_backend = "nccl" - self.excluded_weights_for_ddp_sync = set() self.has_moe_layers = False self.num_experts = None @@ -247,8 +246,7 @@ def __init__(self, self.flatten = util_ops.flatten self.unflatten = util_ops.unflatten - if mpu is not None and hasattr(mpu, "get_excluded_weights_for_ddp_sync"): - self.excluded_weights_for_ddp_sync = mpu.get_excluded_weights_for_ddp_sync(model) + self.excluded_weights_for_ddp_sync = set() def get_batch_info(self): """ Get all training batch related settings. @@ -754,7 +752,7 @@ def _configure_distributed_model(self, model): else: if not groups.is_initialized(): # Scenario 1 - groups.initialize(self.has_moe_layers) + groups.initialize(self.has_moe_layers, has_moe_layers=self.has_moe_layers) #else: # Scenario 2 # Scenario 4 - Case 2 @@ -1669,6 +1667,10 @@ def allreduce_no_retain(self, bucket, dp_group, numel_per_bucket=500000000): if len(small_bucket) > 0: self.allreduce_and_copy(small_bucket, dp_group) + def add_excluded_weights_for_ddp_sync(self, params_name): + logger.info("add_excluded_weights_for_ddp_sync: {}".format(params_name)) + self.excluded_weights_for_ddp_sync.add(params_name) + def buffered_allreduce_fallback(self, grads=None, elements_per_buffer=500000000): grads, expert_grads = [], [] for param_name, param in self.module.named_parameters(): @@ -1704,6 +1706,8 @@ def buffered_allreduce_fallback(self, grads=None, elements_per_buffer=500000000) expert_grads.append(grad_data) else: grads.append(grad_data) + else: + logger.info("weights {} are excluded from DDP sync".format(param_name)) split_buckets = split_half_float_double_csr(grads) for i, bucket_tuple in enumerate(split_buckets): From e880ce7743a99403b7141194c02b7ae2beb06635 Mon Sep 17 00:00:00 2001 From: Chaoyang He Date: Sun, 26 Sep 2021 09:38:39 -0700 Subject: [PATCH 6/7] skip params --- deepspeed/runtime/engine.py | 35 +++++++++++++++-------------------- 1 file changed, 15 insertions(+), 20 deletions(-) diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 73464d4af39d..5b0f5400ccfe 100755 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -246,8 +246,6 @@ def __init__(self, self.flatten = util_ops.flatten self.unflatten = util_ops.unflatten - self.excluded_weights_for_ddp_sync = set() - def get_batch_info(self): """ Get all training batch related settings. @@ -1297,7 +1295,6 @@ def allreduce_gradients(self, bucket_size=MEMORY_OPT_ALLREDUCE_SIZE): # Pass (PP) gas boundary flag to optimizer (required for zero) self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary( ) - # ZeRO stage 2 communicates during non gradient accumulation boundaries as well if self.zero_optimization_partition_gradients(): self.optimizer.overlapping_partition_gradients_reduce_epilogue() @@ -1667,13 +1664,14 @@ def allreduce_no_retain(self, bucket, dp_group, numel_per_bucket=500000000): if len(small_bucket) > 0: self.allreduce_and_copy(small_bucket, dp_group) - def add_excluded_weights_for_ddp_sync(self, params_name): - logger.info("add_excluded_weights_for_ddp_sync: {}".format(params_name)) - self.excluded_weights_for_ddp_sync.add(params_name) - def buffered_allreduce_fallback(self, grads=None, elements_per_buffer=500000000): grads, expert_grads = [], [] for param_name, param in self.module.named_parameters(): + # if hasattr(param, 'expert'): + if 'expert' in param_name: + # Skip gradient sync for unshared parameters + logger.info("Skip param_name {}'s gradient sync for unshared parameters".format(param_name)) + continue if hasattr(param, 'allreduce') and not param.allreduce: is_moe_param = True else: @@ -1693,21 +1691,18 @@ def buffered_allreduce_fallback(self, grads=None, elements_per_buffer=500000000) grads.append(param.grad.data) else: grad_data = param.grad.data - if not self.excluded_weights_for_ddp_sync.__contains__(param_name): - grads.append(grad_data) - if self.sparse_gradients_enabled( - ) and param_name in self.csr_tensor_module_names: - if is_moe_param: - expert_grads.append(CSRTensor(grad_data)) - else: - grads.append(CSRTensor(grad_data)) + grads.append(grad_data) + if self.sparse_gradients_enabled( + ) and param_name in self.csr_tensor_module_names: + if is_moe_param: + expert_grads.append(CSRTensor(grad_data)) else: - if is_moe_param: - expert_grads.append(grad_data) - else: - grads.append(grad_data) + grads.append(CSRTensor(grad_data)) else: - logger.info("weights {} are excluded from DDP sync".format(param_name)) + if is_moe_param: + expert_grads.append(grad_data) + else: + grads.append(grad_data) split_buckets = split_half_float_double_csr(grads) for i, bucket_tuple in enumerate(split_buckets): From edb6964603b04cdd8c8fe2d1134ee83e110f13e6 Mon Sep 17 00:00:00 2001 From: Chaoyang He Date: Sun, 26 Sep 2021 14:42:59 -0700 Subject: [PATCH 7/7] Update groups.py --- deepspeed/utils/groups.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/utils/groups.py b/deepspeed/utils/groups.py index 242d5dec0bf6..db620a0d1777 100644 --- a/deepspeed/utils/groups.py +++ b/deepspeed/utils/groups.py @@ -244,7 +244,7 @@ def initialize_data_and_model_parallel(mpu): dp_rank = mpu.get_data_parallel_rank() log_dist( - f"Initializing deepspeed groups with model parallel size {model_parallel_size_}, expert parallel size {expert_parallel_size_}, and data parallel size {world_size}", + f"Initializing deepspeed groups with model parallel size {model_parallel_size_}, and data parallel size {world_size}", [0]) global _DATA_PARALLEL_GROUP, _MODEL_PARALLEL_GROUP