Some adjustment for supporting Deepspeed-Ulysses

src/accelerate/accelerator.py

@@ -103,6 +103,7 @@
     save_fsdp_optimizer,
     wait_for_everyone,
 )
+from .utils import parallel_state as mpu
 from .utils.constants import FSDP_PYTORCH_VERSION, PROFILE_PATTERN_NAME
 from .utils.modeling import get_state_dict_offloaded_model
 from .utils.other import is_compiled_module
@@ -1634,11 +1635,16 @@ def _prepare_deepspeed(self, *args):
             gradient_accumulation_steps=self.gradient_accumulation_steps,
         )
 
+        if mpu.model_parallel_is_initialized():
+            world_size = mpu.get_data_parallel_world_size()
+        else:
+            world_size = self.num_processes
+
         config_kwargs = {
             "train_micro_batch_size_per_gpu": batch_size_per_device,
             "train_batch_size": batch_size_per_device
             * deepspeed_plugin.get_value("gradient_accumulation_steps")
-            * self.num_processes,
+            * world_size,
             "gradient_clipping": 1.0,
             "zero_optimization.stage3_gather_16bit_weights_on_model_save": False,
         }
@@ -1758,6 +1764,8 @@ def _prepare_deepspeed(self, *args):
             deepspeed_plugin.deepspeed_config_process(must_match=False, **config_kwargs)
             self.deepspeed_config = deepspeed_plugin.deepspeed_config
             kwargs = dict(model=model, config_params=self.deepspeed_config)
+            if mpu.model_parallel_is_initialized():
+                kwargs["mpu"] = mpu
             if optimizer is not None:
                 if isinstance(optimizer, (DummyOptim)):
                     kwargs["model_parameters"] = optimizer.params

src/accelerate/data_loader.py

@@ -17,7 +17,7 @@
 from typing import Callable, List, Optional, Union
 
 import torch
-from torch.utils.data import BatchSampler, DataLoader, IterableDataset, RandomSampler
+from torch.utils.data import BatchSampler, DataLoader, DistributedSampler, IterableDataset, RandomSampler
 
 from .logging import get_logger
 from .state import AcceleratorState, DistributedType, GradientState, is_torch_xla_available
@@ -941,8 +941,20 @@ def prepare_data_loader(
         generator = torch.Generator().manual_seed(42)
         dataloader.generator = generator
         dataloader.sampler.generator = generator
+
+    is_distributed_sampler = isinstance(
+        dataloader.sampler.sampler if sampler_is_batch_sampler else dataloader.sampler, DistributedSampler
+    )
+
+    if is_distributed_sampler and split_batches:
+        raise ValueError("Using `split_batches=True` with a `DistributedSampler` is not supported.")
+
     # No change if no multiprocess
-    if (num_processes != 1 or state.distributed_type == DistributedType.MEGATRON_LM) and not dispatch_batches:
+    if (
+        (num_processes != 1 or state.distributed_type == DistributedType.MEGATRON_LM)
+        and not dispatch_batches
+        and not is_distributed_sampler
+    ):
         if isinstance(new_dataset, IterableDataset):
             if getattr(dataloader.dataset, "generator", None) is not None:
                 synchronized_generator = dataloader.dataset.generator