Add OPCM configuration and implementation for distributed model merging (#205)

* Add OPCM configuration and implementation for distributed model merging

* Update fusion_bench/method/opcm/opcm_general.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update fusion_bench/method/opcm/opcm_general.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Reorder class inheritance for OPCMForCLIP to improve clarity

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: tanganke

config/method/opcm/opcm_general.yaml

@@ -0,0 +1,18 @@
+# =============================================================================
+# FusionBench Method Configuration: OPCM
+# =============================================================================
+# Incrementally merges models via SVD projection and evaluation per step.
+# =============================================================================
+_target_: fusion_bench.method.opcm.opcm_general.OPCM
+# shuffle the order of the models
+shuffle_order: true
+# the scaling factor for the SVD projection
+alpha: 0.5
+# the random seed to use
+seed: null
+# save the merged model on every step
+save_on_every_step: true
+# evaluate the merged model on every step
+evaluate_on_every_step: true
+# the number of ray actors to use for distributed merging
+num_ray_actors: 0

examples/opcm/distributed_opcm.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.opcm.opcm_general import OPCM
+from fusion_bench.utils import timeit_context
+
+config_file = DEFAULT_CONFIG_PATH / "method/opcm/opcm_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: OPCM = OPCM.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/opcm/opcm.py

@@ -30,9 +30,9 @@
 
 
 class OPCMForCLIP(
-    BaseAlgorithm,
     LightningFabricMixin,
     SimpleProfilerMixin,
+    BaseAlgorithm,
 ):
     def __init__(
         self,
@@ -220,6 +220,9 @@ def run(self, modelpool: BaseModelPool):
         return merged_model
 
     def save_merged_model(self, merged_model: CLIPVisionModel, step: int):
+        if self.log_dir is None:
+            print("Log dir is None, skip saving merged model.")
+            return
         os.makedirs(Path(self.log_dir) / "checkpoints", exist_ok=True)
         merged_model.save_pretrained(
             Path(self.log_dir) / "checkpoints" / f"merged_model_{step}"