Add distributed_ray example and enhance DOPMerging with GPU support

Author: tanganke

examples/dop/distributed_ray.py

@@ -0,0 +1,24 @@
+from transformers import CLIPVisionModel
+
+from fusion_bench import BaseModelPool
+from fusion_bench.constants.paths import DEFAULT_CONFIG_PATH
+from fusion_bench.method.dop.dop_general import DOPMerging
+from fusion_bench.utils import timeit_context
+
+config_file = DEFAULT_CONFIG_PATH / "method/dop/dop_general.yaml"
+
+
+with timeit_context("loading models"):
+    models = {
+        "_pretrained_": CLIPVisionModel.from_pretrained("openai/clip-vit-base-patch32"),
+        "sun397": CLIPVisionModel.from_pretrained(
+            "tanganke/clip-vit-base-patch32_sun397"
+        ),
+        "stanford-cars": CLIPVisionModel.from_pretrained(
+            "tanganke/clip-vit-base-patch32_stanford-cars"
+        ),
+    }
+
+algo: DOPMerging = DOPMerging.from_yaml(config_file)
+algo.num_ray_actors = 2  # set the number of ray actors to use for parallel merging
+algo.run(BaseModelPool(models))

fusion_bench/method/dop/dop_general.py

@@ -140,9 +140,13 @@ def run(self, modelpool: BaseModelPool):
                     ray.init()
 
                 # create actors
+                if self.fabric.device.type == "cuda":
+                    actor_options = {"num_gpus": 1}
+                else:
+                    actor_options = {}
                 self.ray_actor_pool = ActorPool(
                     [
-                        DOPMergingActor.remote(**self.config)
+                        DOPMergingActor.options(**actor_options).remote(**self.config)
                         for _ in range(self.num_ray_actors)
                     ]
                 )
@@ -250,7 +254,7 @@ def _layer_wise_optimize(
                         )
                     self.ray_actor_pool.submit(
                         lambda actor, kwargs: actor._optimize_linear_layer.remote(
-                            *kwargs
+                            **kwargs
                         ),
                         {
                             "module_name": module_name,

fusion_bench/models/utils.py

@@ -15,6 +15,7 @@ def is_leaf_module(module: nn.Module) -> bool:
 def named_leaf_modules(
     module: nn.Module,
     prefix: str = "",
+    ignore_empty: bool = True,
 ) -> Iterable[tuple[str, nn.Module]]:
     """
     Recursively find the leaf modules in a module.
@@ -28,6 +29,8 @@ def named_leaf_modules(
     """
     for name, submodule in module.named_modules(prefix=prefix):
         if is_leaf_module(submodule):
+            if ignore_empty and len(list(submodule.parameters())) == 0:
+                continue
             yield name, submodule