add option merged_dtype to clip_concrete_adamerging method

config/method/clip_concrete_layer_wise_adamerging.yaml

@@ -5,6 +5,7 @@ name: clip_concrete_layer_wise_adamerging
 batch_size: 16
 num_workers: 8
 
+merge_dtype: null
 optimizer: adam
 lr: 1e-3
 base_lr: 1

config/method/clip_concrete_task_wise_adamerging.yaml

@@ -5,6 +5,7 @@ name: clip_concrete_task_wise_adamerging
 batch_size: 16
 num_workers: 8
 
+merge_dtype: null
 optimizer: adam
 lr: 1e-3
 base_lr: 1

fusion_bench/method/concrete_subspace/clip_concrete_adamerging.py

@@ -42,6 +42,7 @@
     get_task_wise_weights,
 )
 from fusion_bench.tasks.clip_classification.clip_mixin import CLIPClassificationMixin
+from fusion_bench.utils.dtype import parse_dtype
 from fusion_bench.utils.parameters import print_parameters
 from fusion_bench.utils.type import StateDictType
 
@@ -55,6 +56,8 @@ class ConcreteTaskWiseAdaMergingForCLIP(
 ):
     @torch.no_grad()
     def setup_models(self):
+        config = self.config
+        self.merge_dtype = parse_dtype(config.get("merge_dtype", None))
         modelpool = self.modelpool
 
         # Load the pretrained model
@@ -66,6 +69,8 @@ def setup_models(self):
             ignore_untrained_params=True,
             parameter_type="logits",
         )
+        if self.merge_dtype is not None:
+            mask_model.to(self.merge_dtype)
         mask_model.fill_(self.config.initial_logits)
         # TODO: ablation study for the initialization of mask model
         # for param in mask_model.parameters():
@@ -92,6 +97,7 @@ def setup_models(self):
             clamp_weights=self.config.clamp_weights,
             tie_weights=self.config.tie_weights,
             strict=self.config.strict,
+            task_vector_dtype=self.merge_dtype,
         )
         return module, mask_model
 
@@ -321,6 +327,8 @@ class ConcreteLayerWiseAdaMergingForCLIP(
 ):
     @torch.no_grad()
     def setup_models(self):
+        config = self.config
+        self.merge_dtype = parse_dtype(config.get("merge_dtype", None))
         modelpool = self.modelpool
 
         # Load the pretrained model
@@ -332,6 +340,8 @@ def setup_models(self):
             ignore_untrained_params=True,
             parameter_type="logits",
         )
+        if self.merge_dtype is not None:
+            mask_model.to(self.merge_dtype)
         mask_model.fill_(self.config.initial_logits)
         # TODO: ablation study for the initialization of mask model
         # for param in mask_model.parameters():
@@ -361,6 +371,7 @@ def setup_models(self):
             clamp_weights=self.config.clamp_weights,
             tie_weights=self.config.tie_weights,
             strict=self.config.strict,
+            layer_vector_dtype=self.merge_dtype,
         )
         return module, mask_model