Ulysses SP for HF Integration (deepspeedai#7268)

This is the Deepspeed counterpart of
snowflakedb/ArcticTraining#45 - as the new
feature(s) require changes on both sides.


For PR reviewers: 

Readiness status:
- [x] Code
- [x] Tests
- [ ] Docs - working on it


Features:

- [x] add support for delaying grad addition via
`param.ds_grad_is_ready` flag (used when performing tiled compute in an
autograd function)
- [x] add light sp-only mpu version (Jeff Rasley)
- [x] improved debug
- [x] added `all_gather_object` to `dist`
- [x] `UlyssesSPAttentionHF` (port of UlyssesAttention from
Megatron-Deepspeed plus modern MHA-variations)
- [x] `UlyssesSPDataLoaderAdapter` - DL adapter to shard the normal DL
batches to be used by `UlyssesSPAttentionHF`
- [x] `SequenceTiledCompute` - generic autograd function to perform
compute after tiling on the sequence dimension
- [x] `TiledMLP` - a specific autograd function to perform tiled MLP
(it's much easier to understand before trying to grok
`SequenceTiledCompute`)
- [x] added a differentiable `_DimZeroAllToAll` (Samyam Rajbhandari)
- [x] torch-dist-check now allows `torch.distributed.nn` (which is
needed since deepspeed's dist is not up to date with
`torch.distributed.nn`)

---------

Signed-off-by: Stas Bekman <[email protected]>
Signed-off-by: Stas Bekman <[email protected]>
Co-authored-by: Stas Bekman <[email protected]>
Co-authored-by: Jeff Rasley <[email protected]>
Co-authored-by: Olatunji Ruwase <[email protected]>

Author: 4 people

.github/workflows/nv-ds-chat.yml

@@ -43,8 +43,8 @@ jobs:
 
       - name: Install deepspeed
         run: |
-          pip install transformers==4.48.3
           pip install .[dev]
+          pip install transformers==4.48.3
           ds_report
 
       - name: Install deepspeed-chat

deepspeed/comm/comm.py

@@ -242,6 +242,12 @@ def all_gather(tensor_list,
     return cdb.all_gather(tensor_list=tensor_list, tensor=tensor, group=group, async_op=async_op)
 
 
+@timed_op
+def all_gather_object(object_list, obj, group=None, prof=False, log_name='all_gather_object', debug=get_caller_func()):
+    global cdb
+    return cdb.all_gather_object(object_list=object_list, obj=obj, group=group)
+
+
 def has_reduce_scatter_tensor():
     global cdb
     assert cdb is not None and cdb.is_initialized(

deepspeed/comm/torch.py

@@ -268,6 +268,10 @@ def all_gather_coalesced(self, output_tensors, input_tensors, group=None, async_
             else:
                 reqs[-1].wait()
 
+    @disable_compiler_collective
+    def all_gather_object(self, object_list, obj, group=None):
+        return torch.distributed.all_gather_object(object_list=object_list, obj=obj, group=group)
+
     @disable_compiler_collective
     def reduce_scatter_tensor(self, output_tensor, input_tensor, op=ReduceOp.SUM, group=None, async_op=False):
         if self.has_reduce_scatter_tensor():

deepspeed/runtime/config.py

@@ -721,14 +721,23 @@ def __init__(self, config: Union[str, dict], mpu=None, mesh_device=None):
                 raise ValueError(
                     f"Expected a string path to an existing deepspeed config, or a dictionary or a valid base64. Received: {config}"
                 )
+
         try:
             self.global_rank = dist.get_rank()
             if mpu is not None:
-                self.world_size = mpu.get_data_parallel_world_size()
+                # Ulysses SP
+                if not hasattr(mpu, "get_data_parallel_world_size"):
+                    self.world_size = dist.get_world_size() / mpu.get_sequence_parallel_world_size()
+                else:
+                    self.world_size = mpu.get_data_parallel_world_size()
             elif mesh_device is not None:
                 self.world_size = dist.get_world_size(mesh_device.get_group(mesh_dim="data_parallel"))
             else:
-                self.world_size = dist.get_world_size()
+                # HF zero.init case where there is no mpu
+                if "sequence_parallel_size" in config:
+                    self.world_size = dist.get_world_size() / config["sequence_parallel_size"]
+                else:
+                    self.world_size = dist.get_world_size()
         except:
             self.global_rank = 0
             self.world_size = 1
@@ -941,7 +950,7 @@ def _set_batch_related_parameters(self):
         micro_batch = self.train_micro_batch_size_per_gpu
         grad_acc = self.gradient_accumulation_steps
 
-        #print(f"train_batch = {train_batch}, micro_batch={micro_batch}")
+        #print(f"in: train_batch = {train_batch}, micro_batch={micro_batch}")
 
         # all values are provided nothing needs to be set
         if train_batch is not None and micro_batch is not None and grad_acc is not None:
@@ -980,6 +989,8 @@ def _set_batch_related_parameters(self):
             assert False, \
                 'Either train_batch_size or train_micro_batch_size_per_gpu needs to be provided'
 
+        #print(f"final: {self.train_batch_size=} {self.train_micro_batch_size_per_gpu=} {self.gradient_accumulation_steps=}")
+
     def _configure_train_batch_size(self):
         self._set_batch_related_parameters()
         self._batch_assertion()

deepspeed/runtime/engine.py

@@ -1354,6 +1354,15 @@ def _configure_distributed_model(self, model):
             self.communication_data_type = self._config.seq_parallel_communication_data_type
             self.seq_parallel_group = groups._get_sequence_parallel_group()
 
+        if dist.get_rank() == 0:
+            summary = "********** distributed groups summary **********\n"
+            summary += f"\t {self.dp_world_size=}\n"
+            summary += f"\t {self.mp_world_size=}\n"
+            summary += f"\t {self.seq_dp_world_size=}\n"
+            summary += f"\t {self.sequence_parallel_size=}\n"
+            summary += "***********************************************"
+            logger.info(summary)
+
         if not (self.amp_enabled() or is_zero_init_model):
             self._broadcast_model()
 

deepspeed/runtime/sequence_parallel/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) The DeepSpeed Contributors
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team

deepspeed/runtime/sequence_parallel/parallel_state_sp.py

@@ -0,0 +1,90 @@
+# Copyright (c) The DeepSpeed Contributors
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""
+This is a slimmed-down version of parallel_state.py (mpu) from Megatron-Deepspeed
+"""
+
+from deepspeed import comm as dist
+
+# Sequence parallel groups to handle both data and sequence parallelisms.
+# These groups are used to reduce gradients and shard parameters and optimizer stages for ZeRO.
+_SEQUENCE_PARALLEL_GROUP = None
+_SEQUENCE_DATA_PARALLEL_GROUP = None
+
+
+def initialize_sequence_parallel(sequence_parallel_size: int) -> None:
+    """Initialize sequence parallel groups."""
+
+    assert dist.is_initialized()
+    world_size: int = dist.get_world_size()
+
+    if world_size < sequence_parallel_size:
+        raise RuntimeError(f"world_size ({world_size}) is less than sequence_parallel_size {sequence_parallel_size}")
+
+    if sequence_parallel_size <= 1:
+        raise ValueError(f"sequence_parallel_size must be greater than 1, got {sequence_parallel_size}")
+
+    if world_size % sequence_parallel_size != 0:
+        raise RuntimeError(
+            f"world_size ({world_size}) is not divisible by sequence_parallel_size {sequence_parallel_size})")
+
+    data_parallel_size: int = world_size // sequence_parallel_size
+    sequence_data_parallel_size: int = sequence_parallel_size * data_parallel_size
+    num_sequence_parallel_groups: int = world_size // sequence_parallel_size
+    num_sequence_data_parallel_groups: int = world_size // sequence_parallel_size // data_parallel_size
+
+    rank = dist.get_rank()
+
+    # Build the sequence parallel groups.
+    global _SEQUENCE_PARALLEL_GROUP
+    assert _SEQUENCE_PARALLEL_GROUP is None, "sequence parallel group is already initialized"
+    for i in range(num_sequence_parallel_groups):
+        ranks = range(i * sequence_parallel_size, (i + 1) * sequence_parallel_size)
+        group = dist.new_group(ranks)
+        if rank in ranks:
+            _SEQUENCE_PARALLEL_GROUP = group
+
+    # Build the sequence data parallel groups.
+    global _SEQUENCE_DATA_PARALLEL_GROUP
+    assert _SEQUENCE_DATA_PARALLEL_GROUP is None, "sequence data parallel group is already initialized"
+    all_data_sequence_parallel_group_ranks = []
+    for i in range(num_sequence_data_parallel_groups):
+        ranks = range(i * sequence_data_parallel_size, (i + 1) * sequence_data_parallel_size)
+        group = dist.new_group(ranks)
+        all_data_sequence_parallel_group_ranks.append(list(ranks))
+        if rank in ranks:
+            _SEQUENCE_DATA_PARALLEL_GROUP = group
+
+
+def get_sequence_parallel_group():
+    """Get the sequence parallel group the caller rank belongs to."""
+    assert _SEQUENCE_PARALLEL_GROUP is not None, "sequence parallel group is not initialized"
+    return _SEQUENCE_PARALLEL_GROUP
+
+
+def get_sequence_data_parallel_group():
+    """Get the sequence parallel group the caller rank belongs to."""
+    assert _SEQUENCE_DATA_PARALLEL_GROUP is not None, "sequence data parallel group is not initialized"
+    return _SEQUENCE_DATA_PARALLEL_GROUP
+
+
+def get_sequence_parallel_world_size():
+    """Return world size for the sequence parallel group."""
+    return dist.get_world_size(group=get_sequence_parallel_group())
+
+
+def get_sequence_data_parallel_world_size():
+    """Return world size for the sequence parallel group."""
+    return dist.get_world_size(group=get_sequence_data_parallel_group())
+
+
+def get_sequence_parallel_rank():
+    """Return my rank for the sequence parallel group."""
+    return dist.get_rank(group=get_sequence_parallel_group())
+
+
+def get_sequence_data_parallel_rank():
+    """Return my rank for the sequence data parallel group."""
+    return dist.get_rank(group=get_sequence_data_parallel_group())