hpcaitech
diff --git a/‎colossalai/booster/plugin/hybrid_parallel_plugin.py‎
Lines changed: 27 additions & 5 deletions b/‎colossalai/booster/plugin/hybrid_parallel_plugin.py‎
Lines changed: 27 additions & 5 deletions
diff --git a/‎colossalai/booster/plugin/low_level_zero_plugin.py‎
Lines changed: 24 additions & 1 deletion b/‎colossalai/booster/plugin/low_level_zero_plugin.py‎
Lines changed: 24 additions & 1 deletion
diff --git a/‎colossalai/cluster/process_group_mesh.py‎
Lines changed: 6 additions & 1 deletion b/‎colossalai/cluster/process_group_mesh.py‎
Lines changed: 6 additions & 1 deletion
diff --git a/‎colossalai/device/device_mesh.py‎
Lines changed: 3 additions & 3 deletions b/‎colossalai/device/device_mesh.py‎
Lines changed: 3 additions & 3 deletions
diff --git a/‎colossalai/interface/optimizer.py‎
Lines changed: 24 additions & 1 deletion b/‎colossalai/interface/optimizer.py‎
Lines changed: 24 additions & 1 deletion
diff --git a/‎colossalai/nn/optimizer/__init__.py‎
Lines changed: 28 additions & 1 deletion b/‎colossalai/nn/optimizer/__init__.py‎
Lines changed: 28 additions & 1 deletion
@@ -1,7 +1,9 @@
 import ctypes
 import random
 import warnings
+from collections import defaultdict
 from contextlib import contextmanager
+from copy import deepcopy
 from functools import partial
 from types import MethodType
 from typing import Any, Callable, Dict, Iterator, List, Optional, OrderedDict, Tuple, Union
@@ -24,6 +26,8 @@
 from colossalai.checkpoint_io import CheckpointIO, HybridParallelCheckpointIO
 from colossalai.cluster import ProcessGroupMesh
 from colossalai.interface import AMPModelMixin, ModelWrapper, OptimizerWrapper
+from colossalai.interface.optimizer import DistributedOptim
+from colossalai.nn.optimizer import DistGaloreAwamW
 from colossalai.pipeline.schedule import InterleavedSchedule, OneForwardOneBackwardSchedule
 from colossalai.pipeline.stage_manager import PipelineStageManager
 from colossalai.shardformer import GradientCheckpointConfig, ShardConfig, ShardFormer
@@ -1171,6 +1175,15 @@ def configure(
         lr_scheduler: Optional[LRScheduler] = None,
     ) -> Tuple[Module, OptimizerWrapper, Callable, DataLoader, LRScheduler]:
         param_info = get_param_info(optimizer)
+
+        # TODO: Support Galore + ZeRO
+        zero_stage = self.zero_stage
+        zero_config = deepcopy(self.zero_config)
+        if isinstance(optimizer, DistGaloreAwamW) and zero_stage > 0 and self.dp_size > 0:
+            warnings.warn("Galore is only supported for Tensor Parallel and vanilla Data Parallel yet. Disabling ZeRO.")
+            zero_config["partition_grad"] = False
+            zero_stage = 0
+
         if not isinstance(model, ModelWrapper):
             use_ddp = (self.dp_size > 1 and self.pp_size == 1 and self.zero_stage == 0) or (
                 self.dp_size == 1
@@ -1194,7 +1207,8 @@ def configure(
                 custom_policy=self.custom_policy,
             )
         if optimizer is not None and not isinstance(optimizer, OptimizerWrapper):
-            if self.zero_stage == 0:
+            if zero_stage == 0:
+                is_zero = False
                 if self.precision in ["fp16", "bf16"]:
                     optimizer = HybridParallelAMPOptimizer(
                         optimizer,
@@ -1218,11 +1232,11 @@ def configure(
                         tp_process_group=self.tp_group,
                     )
             else:
-                zero_dp_size = dist.get_world_size(dp_group)
-                if zero_dp_size == 1:
+                is_zero = self.dp_size > 1
+                if self.dp_size == 1:
                     warnings.warn(
                         "Use Zero Optimizer when data parallel size is 1 may introduce unnecessary overhead. "
-                        "If you are not intended to use cpu_offload, please consider set zero_stage=0."
+                        "If you do not intend to use cpu_offload, please consider set zero_stage=0."
                     )
 
                 assert self.precision != "fp32", "Please set precision to 'fp16' or 'bf16' when using ZeRO."
@@ -1236,11 +1250,19 @@ def configure(
                     pp_process_group=self.pp_group,
                     verbose=True,
                     clip_grad_norm=self.max_norm,
-                    **self.zero_config,
+                    **zero_config,
                     **self.amp_config,
                 )
             # inject update_master_params
             model.update_master_params = MethodType(optimizer.update_master_params, model)
+
+            # Setup optimizers that require global states
+            optim = optimizer.optim
+            if isinstance(optim, DistributedOptim):
+                shard_to_param = optimizer.get_master_to_working_map() if is_zero else {}
+                padding_map = optimizer.get_param_padding_map() if is_zero else defaultdict(int)
+                optim.setup_distributed(self.tp_group, self.dp_group, shard_to_param, padding_map, is_zero)
+
         return model, optimizer, criterion, dataloader, lr_scheduler
 
     def execute_pipeline(
 
@@ -8,7 +8,10 @@
 from typing import Callable, Dict, Iterator, List, Optional, Tuple
 
 import torch
+import torch.distributed
+import torch.distributed as dist
 import torch.nn as nn
+from torch.distributed.distributed_c10d import _get_default_group
 from torch.nn import Parameter
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
@@ -28,6 +31,8 @@
     sharded_optimizer_loading_epilogue,
 )
 from colossalai.interface import AMPModelMixin, ModelWrapper, OptimizerWrapper
+from colossalai.interface.optimizer import DistributedOptim
+from colossalai.nn.optimizer import DistGaloreAwamW
 from colossalai.quantization import BnbQuantizationConfig, quantize_model
 from colossalai.zero import LowLevelZeroOptimizer
 
@@ -428,13 +433,31 @@ def configure(
         if not isinstance(model, ModelWrapper):
             model = LowLevelZeroModel(model, self.precision)
 
+        # TODO: Support Galore + ZeRO
+        zero_stage = self.stage
+        zero_optim_kwargs = {**self.zero_optim_kwargs}
+        dp_size = dist.get_world_size()
+        if isinstance(optimizer, DistGaloreAwamW) and zero_stage > 0 and dp_size > 0:
+            warnings.warn("Galore is only supported for Tensor Parallel and vanilla Data Parallel yet. Disabling ZeRO.")
+            zero_optim_kwargs["partition_grad"] = False
+            zero_stage = 0
+
         if optimizer is not None and not isinstance(optimizer, OptimizerWrapper):
             optimizer: LowLevelZeroOptimizer = LowLevelZeroOptimizer(
-                optimizer, **self.zero_optim_kwargs, verbose=self.verbose
+                optimizer, **zero_optim_kwargs, verbose=self.verbose
             )
             # inject update_master_params
             model.update_master_params = MethodType(optimizer.update_master_params, model)
 
+            # Setup optimizers that require global states
+            optim = optimizer.optim
+            is_zero = dp_size > 1 and zero_stage > 0
+            dp_group = _get_default_group()  # Use the whole world
+            if isinstance(optim, DistributedOptim):
+                shard_to_param = optimizer.get_master_to_working_map()
+                padding_map = optimizer.get_param_padding_map()
+                optim.setup_distributed(None, dp_group, shard_to_param, padding_map, is_zero)
+
         return model, optimizer, criterion, dataloader, lr_scheduler
 
     def control_checkpoint_io(self) -> bool:
 
@@ -38,7 +38,12 @@ class ProcessGroupMesh:
 
     def __init__(self, *size: int) -> None:
         assert dist.is_initialized(), "Please initialize torch.distributed first."
-        assert prod(size) == dist.get_world_size(), "The product of the size must be equal to the world size."
+        world_size = dist.get_world_size()
+        prod_size = prod(size)
+        assert (
+            prod_size == world_size
+        ), f"The product of the size({prod_size}) must be equal to the world size({world_size})."
+
         self._shape = size
         self._rank = dist.get_rank()
         self._coord = ProcessGroupMesh.unravel(self._rank, self._shape)
 
@@ -306,9 +306,8 @@ def _init_global_to_logical_rank_mapping(
             # index means the local rank in the current axis
             # inner_tensor refers to the processes with the same local rank
 
-            if inner_tensor.numel() == 1:
-                # if the inner_tensor only has one element, it means that
-                # it already reaches the last axis
+            if inner_tensor.dim() == 0:
+                # if the inner_tensor already reaches the last axis,
                 # we append its local_rank in the last axis to the index_list
                 # and assign to the mapping
                 # the value of the mapping is the the local rank at the indexed axis of the device mesh
@@ -459,6 +458,7 @@ def _collate_global_ranks_in_same_process_group(self, global_rank):
 
                 # replace the local rank in the given dimension with the
                 # local rank of the current process iterated
+
                 process_coordinates[dim] = _local_rank
                 processes_in_the_same_process_group[dim].append(process_coordinates)
 
 
@@ -1,6 +1,7 @@
-from typing import Union
+from typing import Dict, Optional, Union
 
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 from torch import Tensor
 from torch.optim import Optimizer
@@ -133,3 +134,25 @@ def unwrap(self):
         Unwrap the optimizer for checkpoint saving/loading.
         """
         return self.optim
+
+
+class DistributedOptim(Optimizer):
+    def setup_distributed(
+        self,
+        tp_group: Optional[dist.ProcessGroup] = None,
+        dp_group: Optional[dist.ProcessGroup] = None,
+        shard_to_working_param: Optional[Dict] = {},
+        padding_map: Optional[Dict] = None,
+        is_zero: Optional[bool] = False,
+    ):
+        """Assign process groups for TP and ZeRO 2.
+        Arguments:
+            tp_group (dist.ProcessGroup): Tensor Parallel process group
+            dp_group (dist.ProcessGroup): ZeRO stage 2 process group
+            shard_to_working_param (Dict): ZeRO stage 2 feeds the optimizer a sharded param view to match grad shape.
+                This maps from id(view) to model params used in forward & backward.
+            padding_map (Dict): Per-param padding from ZeRO stage 2
+            is_zero (bool): Whether to use ZeRO stage 2.
+        """
+
+        raise NotImplementedError("setup_distributed for TP/DP isn't supported by this optimizer yet!")
@@ -1,9 +1,36 @@
+from galore_torch import GaLoreAdafactor, GaLoreAdamW
+
+from .came import CAME
 from .cpu_adam import CPUAdam
+from .distributed_adafactor import DistributedAdaFactor
+from .distributed_came import DistributedCAME
+from .distributed_galore import DistGaloreAwamW
+from .distributed_lamb import DistributedLamb
 from .fused_adam import FusedAdam
 from .fused_lamb import FusedLAMB
 from .fused_sgd import FusedSGD
+from .galore import GaLoreAdamW8bit
 from .hybrid_adam import HybridAdam
 from .lamb import Lamb
 from .lars import Lars
 
-__all__ = ["FusedLAMB", "FusedAdam", "FusedSGD", "Lamb", "Lars", "CPUAdam", "HybridAdam"]
+from .adafactor import Adafactor  # noqa
+
+__all__ = [
+    "FusedLAMB",
+    "FusedAdam",
+    "FusedSGD",
+    "Lamb",
+    "Lars",
+    "CPUAdam",
+    "HybridAdam",
+    "DistributedLamb",
+    "DistGaloreAwamW",
+    "GaLoreAdamW",
+    "GaLoreAdafactor",
+    "GaLoreAdamW8bit",
+    "CAME",
+    "DistributedCAME",
+    "Adafactor",
+    "DistributedAdaFactor",
+]