fix: fuse optimizer oom and make it stateless

bagua/torch_api/algorithms/async_model_average.py

@@ -78,15 +78,19 @@ def __init__(
             process_ranks, stream=torch.cuda.Stream(priority=-1)
         )
 
-    def tensors_to_buckets(self, tensors: List[List[BaguaTensor]]) -> List[BaguaBucket]:
+    def tensors_to_buckets(
+        self, tensors: List[List[BaguaTensor]], do_flatten: bool
+    ) -> List[BaguaBucket]:
+        # TODO: async algorithm conflict with fused optimizer, can only support flattened inplace bucket.
+        assert do_flatten, "async does not support `do_flatten=False` at present."
         if self.step_id < self.warmup_steps:
-            return super().tensors_to_buckets(tensors)
+            return super().tensors_to_buckets(tensors, do_flatten)
 
         all_tensors = []
         for idx, bucket in enumerate(tensors):
             all_tensors.extend(bucket)
 
-        bagua_bucket = BaguaBucket(all_tensors, flatten=True, name=str(0))
+        bagua_bucket = BaguaBucket(all_tensors, flatten=do_flatten, name=str(0))
 
         return [bagua_bucket]
 

bagua/torch_api/algorithms/base.py

@@ -73,7 +73,9 @@ def init_tensors(self, bagua_module: BaguaModule) -> List[BaguaTensor]:
         ), "tensor names should be unique"
         return tensors
 
-    def tensors_to_buckets(self, tensors: List[List[BaguaTensor]]) -> List[BaguaBucket]:
+    def tensors_to_buckets(
+        self, tensors: List[List[BaguaTensor]], do_flatten: bool
+    ) -> List[BaguaBucket]:
         """
         Given the bucketing suggestion from Bagua, return the actual Bagua buckets.
         The default implementation follows the suggestion to do the bucketing.
@@ -82,14 +84,15 @@ def tensors_to_buckets(self, tensors: List[List[BaguaTensor]]) -> List[BaguaBuck
             tensors: Bagua tensors grouped in different
                 lists, representing Bagua's suggestion on how to bucketing the
                 tensors.
+            do_flatten: Whether to flatten the Bagua buckets.
 
         Returns:
             A list of Bagua buckets.
         """
         bagua_buckets = []
         for idx, bucket in enumerate(tensors):
             bagua_bucket = BaguaBucket(
-                bucket, flatten=True, name=str(idx)
+                bucket, flatten=do_flatten, name=str(idx)
             )  # TODO: check duplicated names
             bagua_buckets.append(bagua_bucket)
         return bagua_buckets

bagua/torch_api/algorithms/bytegrad.py

@@ -30,24 +30,14 @@ def __init__(
         self.hierarchical = hierarchical
         self.average = average
 
-    def tensors_to_buckets(self, tensors: List[List[BaguaTensor]]) -> List[BaguaBucket]:
-        """
-        Given the bucketing suggestion from Bagua, return the actual Bagua buckets.
-        The default implementation follows the suggestion to do the bucketing.
-
-        Args:
-            tensors: Bagua tensors grouped in different
-                lists, representing Bagua's suggestion on how to bucketing the
-                tensors.
-
-        Returns:
-            A list of Bagua buckets.
-        """
+    def tensors_to_buckets(
+        self, tensors: List[List[BaguaTensor]], do_flatten: bool
+    ) -> List[BaguaBucket]:
         bagua_buckets = []
         for idx, bucket in enumerate(tensors):
             bagua_bucket = BaguaBucket(
                 bucket,
-                flatten=True,
+                flatten=do_flatten,
                 name=str(idx),
                 alignment=self.process_group.get_global_communicator().nranks(),
             )

bagua/torch_api/algorithms/decentralized.py

@@ -48,12 +48,14 @@ def init_tensors(self, bagua_module: BaguaModule) -> List[BaguaTensor]:
         ]
         return self.tensors
 
-    def tensors_to_buckets(self, tensors: List[List[BaguaTensor]]) -> List[BaguaBucket]:
+    def tensors_to_buckets(
+        self, tensors: List[List[BaguaTensor]], do_flatten: bool
+    ) -> List[BaguaBucket]:
         all_tensors = []
         for idx, bucket in enumerate(tensors):
             all_tensors.extend(bucket)
 
-        bagua_bucket = BaguaBucket(all_tensors, flatten=True, name=str(0))
+        bagua_bucket = BaguaBucket(all_tensors, flatten=do_flatten, name=str(0))
 
         return [bagua_bucket]
 

bagua/torch_api/algorithms/q_adam.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 from bagua.torch_api.bucket import BaguaBucket
 from bagua.torch_api.tensor import BaguaTensor
-from bagua.torch_api import get_world_size
 from bagua.torch_api.distributed import BaguaModule
 from bagua.torch_api.algorithms import Algorithm, AlgorithmImpl
 from bagua.torch_api.communication import BaguaProcessGroup
@@ -45,7 +44,7 @@ def __init__(
             raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
         defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
         super(QAdamOptimizer, self).__init__(params, defaults)
-
+        # TODO: qadam optimizer maintain `step_id` in its state
         self.step_id = 0
         self.warmup_steps = warmup_steps
 
@@ -162,12 +161,14 @@ def set_momentum_fn(param, t):
         tensor_groups.sort(key=lambda x: x._q_adam_idx)
         return tensor_groups
 
-    def tensors_to_buckets(self, tensors: List[List[BaguaTensor]]) -> List[BaguaBucket]:
+    def tensors_to_buckets(
+        self, tensors: List[List[BaguaTensor]], do_flatten: bool
+    ) -> List[BaguaBucket]:
         bagua_buckets = []
         for idx, bucket in enumerate(tensors):
             bagua_bucket = BaguaBucket(
                 bucket,
-                flatten=True,
+                flatten=do_flatten,
                 name=str(idx),
                 alignment=self.process_group.get_global_communicator().nranks(),
             )

bagua/torch_api/bucket.py

@@ -8,7 +8,7 @@
 import torch
 
 from bagua.torch_api.tensor import BaguaTensor
-from bagua.torch_api.utils import check_contiguous
+from bagua.torch_api.utils import check_contiguous, get_flattened_tensor
 from bagua.torch_api.communication import (
     BaguaProcessGroup,
     _bagua_backend_comm,
@@ -87,25 +87,10 @@ def flattened_tensor(self) -> torch.Tensor:
         :attr:`self` tensors and padding tensor (if exists).
         """
 
-        all_registered_tensors = [
+        all_effective_tensors = [
             tensor.bagua_getter_closure() for tensor in self._all_tensors
         ]
-        total_size = 0
-        for tensor in all_registered_tensors:
-            total_size += tensor.numel()
-
-        flatten_tensor = torch.zeros(
-            total_size,
-            dtype=all_registered_tensors[0].dtype,
-            device=all_registered_tensors[0].device,
-        )
-
-        offset = 0
-        for tensor in all_registered_tensors:
-            # copy data
-            flatten_tensor[offset : offset + tensor.numel()] = tensor.reshape(-1)
-            offset += tensor.numel()
-        return flatten_tensor
+        return get_flattened_tensor(all_effective_tensors)
 
     def _flatten_(self):
         """
@@ -372,7 +357,7 @@ def clear_ops(self) -> BaguaBucket:
 
     def bytes(self) -> int:
         """Returns the total number of bytes occupied by the bucket."""
-        registered_tensors = [tensor.bagua_getter_closure() for tensor in self.tensors]
+        effective_tensors = [tensor.bagua_getter_closure() for tensor in self.tensors]
         return sum(
-            tensor.numel() * tensor.element_size() for tensor in registered_tensors
+            tensor.numel() * tensor.element_size() for tensor in effective_tensors
         )

bagua/torch_api/contrib/__init__.py

@@ -1,4 +1,4 @@
-from .fused_optimizer import FusedOptimizer  # noqa: F401
+from .fuse.optimizer import fuse_optimizer  # noqa: F401
 from .load_balancing_data_loader import (  # noqa: F401
     LoadBalancingDistributedSampler,
     LoadBalancingDistributedBatchSampler,