fusion-bench 0.2.21__py3-none-any.whl → 0.2.22__py3-none-any.whl

fusion_bench/method/smile_upscaling/smile_qwen2_upscaling.py

@@ -16,10 +16,11 @@ from transformers.models.qwen2.modeling_qwen2 import Qwen2DecoderLayer
 
 from fusion_bench import BaseAlgorithm, BaseModelPool
 from fusion_bench.compat.modelpool import to_modelpool
+from fusion_bench.constants import RuntimeConstants
 from fusion_bench.mixins import SimpleProfilerMixin, auto_register_config
 from fusion_bench.modelpool import CausalLMPool
 from fusion_bench.models.hf_utils import (
-    generate_complete_readme,
+    create_default_model_card,
     save_pretrained_with_remote_code,
 )
 from fusion_bench.models.modeling_smile_qwen2 import (
@@ -41,7 +42,10 @@ log = logging.getLogger(__name__)
 
 
 @auto_register_config
-class SmileQwen2UpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
+class SmileQwen2UpscalingAlgorithm(
+    SimpleProfilerMixin,
+    BaseAlgorithm,
+):
     R"""
     SmileQwen2UpscalingAlgorithm is a model fusion algorithm designed to upscale
     a pretrained Qwen2 model using a set of fine-tuned expert models. The algorithm
@@ -62,7 +66,7 @@ class SmileQwen2UpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
         self,
         device,
         accelerator,
-        model_path,
+        model_save_path,
         model_dtype,
         num_experts_per_tok,
         rank_of_router,
@@ -71,6 +75,11 @@ class SmileQwen2UpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
         **kwargs,
     ):
         super().__init__(**kwargs)
+        if not torch.cuda.is_available():
+            if "cuda" in self.device:
+                self.device = "cpu"
+            if "cuda" in self.accelerator:
+                self.accelerator = "cpu"
 
     @torch.no_grad()
     def run(self, modelpool) -> SmileQwen2ForCausalLM:
@@ -86,13 +95,6 @@ class SmileQwen2UpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
         self.modelpool = modelpool = to_modelpool(modelpool)
         config = self.config
 
-        # load model from path if provided and return directly
-        if config.model_path is not None and os.path.exists(config.model_path):
-            log.info(f"Loading model from {config.model_path}")
-            model = AutoModelForCausalLM.from_pretrained(config.model_path)
-            print_parameters(model)
-            return model
-
         with self.profile("load pretrained model"):
             pretrained_model = modelpool.load_pretrained_model()
         with self.profile("load fine-tuned model"):
@@ -100,7 +102,7 @@ class SmileQwen2UpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
                 m for m in tqdm(modelpool.models(), total=len(modelpool.model_names))
             ]
 
-        if config.device == "cuda" and torch.cuda.is_available():
+        if self.device == "cuda" and torch.cuda.is_available():
             pretrained_model = pretrained_model.cuda()
             print("parameter count of pretrained model:")
             print_parameters(pretrained_model)
@@ -114,17 +116,17 @@ class SmileQwen2UpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
         print_parameters(model)
         print(model)
 
-        if config.model_dtype is not None:
-            model.to(dtype=parse_dtype(config.model_dtype))
+        if self.model_dtype is not None:
+            model.to(dtype=parse_dtype(self.model_dtype))
 
-        if config.model_path is not None:
-            if os.path.dirname(config.model_path):
-                os.makedirs(os.path.dirname(config.model_path), exist_ok=True)
-            log.info(f"Saving model to {config.model_path}")
+        if self.model_save_path is not None:
+            if os.path.dirname(self.model_save_path):
+                os.makedirs(os.path.dirname(self.model_save_path), exist_ok=True)
+            log.info(f"Saving model to {self.model_save_path}")
             tokenizer = self.modelpool.load_tokenizer()
-            tokenizer.save_pretrained(config.model_path)
+            tokenizer.save_pretrained(self.model_save_path)
             if not self.save_with_remote_code:
-                model.save_pretrained(config.model_path)
+                model.save_pretrained(self.model_save_path)
             else:
                 save_pretrained_with_remote_code(
                     model,
@@ -133,17 +135,18 @@ class SmileQwen2UpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
                         "AutoModel": SmileQwen2Model,
                         "AutoModelForCausalLM": SmileQwen2ForCausalLM,
                     },
-                    save_directory=config.model_path,
+                    save_directory=self.model_save_path,
                 )
 
             # save readme
-            complete_readme = generate_complete_readme(
-                algorithm=self,
-                modelpool=modelpool,
+            model_card_str = create_default_model_card(
                 models=[modelpool.get_model_path(m) for m in modelpool.all_model_names],
+                description="Merged Qwen model using SMILE Upscaling",
+                algorithm_config=self.config,
+                modelpool_config=modelpool.config,
             )
-            with open(os.path.join(config.model_path, "README.md"), "w") as f:
-                f.write(complete_readme)
+            with open(os.path.join(self.model_save_path, "README.md"), "w") as f:
+                f.write(model_card_str)
 
         return model
 
@@ -174,9 +177,9 @@ class SmileQwen2UpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
                 )
             base_config = AutoConfig.from_pretrained(pretrained_path)
             model_config = SmileQwen2Config(
-                num_experts_per_tok=config.num_experts_per_tok,
-                rank_of_router=config.rank_of_router,
-                rank_of_expert=config.rank_of_expert,
+                num_experts_per_tok=self.num_experts_per_tok,
+                rank_of_router=self.rank_of_router,
+                rank_of_expert=self.rank_of_expert,
                 num_local_experts=len(finetuned_models),
                 **base_config.to_dict(),
             )
@@ -186,7 +189,7 @@ class SmileQwen2UpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
 
         # copy pretrained model weights
         state_dict = model.state_dict()
-        pretrained_state_dict = dict(pretrained_model.state_dict())
+        pretrained_state_dict = pretrained_model.state_dict()
         for key in list(pretrained_state_dict.keys()):
             if key not in state_dict:
                 pretrained_state_dict.pop(key)
@@ -198,6 +201,12 @@ class SmileQwen2UpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
             "Upscaling Modules (layer)",
             dynamic_ncols=True,
         ):
+            if RuntimeConstants.debug and layer_idx > 0:
+                log.info(
+                    "Debug mode enabled: processing only the first layer, skipping remaining layers"
+                )
+                break
+
             pretrained_layer: Qwen2DecoderLayer = pretrained_model.model.layers[
                 layer_idx
             ]
@@ -213,7 +222,7 @@ class SmileQwen2UpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
                         base=getattr(pretrained_layer.self_attn, n),
                         experts=[getattr(m.self_attn, n) for m in finetuned_layers],
                         target=getattr(target_layer.self_attn, n),
-                        accelerator=config.accelerator,
+                        accelerator=self.accelerator,
                     )
                 except ExpertNotTrainedError:
                     setattr(
@@ -228,7 +237,7 @@ class SmileQwen2UpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
                         base=getattr(pretrained_layer.mlp, n),
                         experts=[getattr(m.mlp, n) for m in finetuned_layers],
                         target=getattr(target_layer.mlp, n),
-                        accelerator=config.accelerator,
+                        accelerator=self.accelerator,
                     )
                 except ExpertNotTrainedError:
                     setattr(

fusion_bench/method/smile_upscaling/smile_upscaling.py

@@ -20,8 +20,8 @@ from fusion_bench.models.smile_moe.linear_from_module import (
     SmileMoELinear,
 )
 from fusion_bench.models.utils import get_attr, set_attr
-from fusion_bench.utils.parameters import print_parameters
 from fusion_bench.utils.devices import get_device
+from fusion_bench.utils.parameters import print_parameters
 
 log = logging.getLogger(__name__)
 

fusion_bench/method/we_moe/__init__.py

@@ -1,2 +1,3 @@
 # flake8: noqa F401
 from .clip_we_moe import CLIPWeightEnsemblingMoEAlgorithm
+from .flan_t5_we_moe import FlanT5WeightEnsemblingMoEAlgorithm

fusion_bench/method/we_moe/entropy_loss.py

@@ -0,0 +1,25 @@
+import torch
+from torch import Tensor
+
+
+def entropy_loss(logits: Tensor, eps: float = 1e-8) -> Tensor:
+    """
+    Compute the entropy loss of a set of logits.
+
+    Args:
+        logits (Tensor): The logits to compute the entropy loss of.
+        eps (float): A small value to avoid log(0). Default is 1e-8.
+
+    Returns:
+        Tensor: The entropy loss of the logits.
+    """
+    # Ensure the logits tensor has 2 dimensions
+    assert (
+        logits.dim() == 2
+    ), f"Expected logits to have 2 dimensions, found {logits.dim()}, {logits.size()=}"
+
+    # Compute the softmax probabilities
+    probs = torch.softmax(logits, dim=-1)
+
+    # Compute the entropy loss
+    return -torch.sum(probs * torch.log(probs + eps), dim=-1).mean()

fusion_bench/method/we_moe/flan_t5_we_moe.py

@@ -0,0 +1,331 @@
+import functools
+import logging
+import os
+from copy import deepcopy
+from typing import Any, Dict, List, Mapping, Optional, Union, cast  # noqa: F401
+
+import lightning
+import lightning as L
+import lightning.fabric.wrappers
+import torch
+from torch import Tensor
+from torch.utils.data import DataLoader
+from tqdm.autonotebook import tqdm
+from transformers import T5ForConditionalGeneration
+from transformers.data import default_data_collator
+
+from fusion_bench.method import BaseAlgorithm
+from fusion_bench.method.task_arithmetic.task_arithmetic import task_arithmetic_merge
+from fusion_bench.mixins.lightning_fabric import LightningFabricMixin
+from fusion_bench.mixins.simple_profiler import SimpleProfilerMixin
+from fusion_bench.modelpool import Seq2SeqLMPool
+from fusion_bench.models.we_moe import WeightEnsemblingMoE
+from fusion_bench.utils import timeit_context
+from fusion_bench.utils.data import InfiniteDataLoader, load_tensor_from_file
+from fusion_bench.utils.instantiate_utils import instantiate
+from fusion_bench.utils.parameters import print_parameters
+
+from .entropy_loss import entropy_loss
+from .utils import get_memory_usage
+
+log = logging.getLogger(__name__)
+
+
+class FlanT5WeightEnsemblingMoEAlgorithm(
+    BaseAlgorithm,
+    LightningFabricMixin,
+    SimpleProfilerMixin,
+):
+    """
+    FlanT5WeightEnsemblingMoEAlgorithm is a class that implements the WeightEnsemblingMoEAlgorithm
+    for FlanT5 models. It extends the WeightEnsemblingMoEAlgorithm and CLIPClassificationMixin classes.
+
+    Attributes:
+        modelpool (Seq2SeqLMPool): The model pool containing the FlanT5 models.
+    """
+
+    modelpool: Seq2SeqLMPool = None
+
+    def __init__(
+        self,
+        checkpoint: bool = False,
+        save_checkpoint: bool = False,
+        router_hidden_layers: int = 2,
+        init_lambda: float = 0.3,
+        batch_reduce: bool = True,
+        lr: float = 1e-4,
+        optimizer: str = "adam",
+        devices: int = 1,
+        batch_size: int = 16,
+        num_workers: int = 0,
+        max_steps: int = 1000,
+        use_grad_accumulate: bool = True,
+        cache_dir: bool = "outputs",
+        fast_dev_run: bool = False,
+        **kwargs,
+    ):
+        """
+        Initialize the WeightEnsemblingMoEAlgorithm with the given configuration.
+
+        Args:
+            algorithm_config (DictConfig): The configuration for the algorithm.
+        """
+        self.checkpoint = checkpoint
+        self.save_checkpoint = save_checkpoint
+        self.router_hidden_layers = router_hidden_layers
+        self.init_lambda = init_lambda
+        self.batch_reduce = batch_reduce
+        self.lr = lr
+        self.optimizer = optimizer
+        self.devices = devices
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.max_steps = max_steps
+        self.use_grad_accumulate = use_grad_accumulate
+        self.cache_dir = cache_dir
+        self.fast_dev_run = fast_dev_run
+        super().__init__(**kwargs)
+
+    def construct_moe_model(self) -> WeightEnsemblingMoE:
+        """
+        Construct the Mixture of Experts (MoE) model using the models in the model pool.
+
+        Returns:
+            WeightEnsemblingMoE: The constructed MoE model.
+        """
+        base_model = self.modelpool.load_model("_pretrained_")
+        expert_models = [
+            self.modelpool.load_model(name) for name in self.modelpool.model_names
+        ]
+
+        # Merge the models using task arithmetic
+        moe_model = task_arithmetic_merge(
+            # This function modifies the model in place, so we need to pass a deepcopy
+            deepcopy(base_model),
+            expert_models,
+            scaling_factor=self.init_lambda,
+        ).requires_grad_(False)
+
+        print(base_model)
+
+        # Up-scale MLP modules
+        num_layer = 12
+        encoder_mlp_index = 1
+        base_encoder = base_model.encoder
+        moe_encoder = moe_model.encoder
+        expert_encoders = [m.encoder for m in expert_models]
+
+        for layer_idx in range(num_layer):
+            base_mlp = (
+                base_encoder.block[layer_idx].layer[encoder_mlp_index].DenseReluDense
+            )
+            expert_mlps = [
+                e.block[layer_idx].layer[encoder_mlp_index].DenseReluDense
+                for e in expert_encoders
+            ]
+
+            moe_encoder.block[layer_idx].layer[encoder_mlp_index].DenseReluDense = (
+                WeightEnsemblingMoE(
+                    hidden_size=base_encoder.config.hidden_size,
+                    base_model=base_mlp,
+                    expert_models=expert_mlps,
+                    init_lambda=self.init_lambda,
+                    batch_first=True,
+                    router_hidden_layers=self.router_hidden_layers,
+                    batch_reduce=self.batch_reduce,
+                )
+            )
+
+        decoder_mlp_index = 2
+        base_decoder = base_model.decoder
+        moe_decoder = moe_model.decoder
+        expert_decoders = [m.decoder for m in expert_models]
+
+        for layer_idx in range(num_layer):
+            base_mlp = (
+                base_decoder.block[layer_idx].layer[decoder_mlp_index].DenseReluDense
+            )
+            expert_mlps = [
+                e.block[layer_idx].layer[decoder_mlp_index].DenseReluDense
+                for e in expert_decoders
+            ]
+
+            moe_decoder.block[layer_idx].layer[decoder_mlp_index].DenseReluDense = (
+                WeightEnsemblingMoE(
+                    hidden_size=base_decoder.config.hidden_size,
+                    base_model=base_mlp,
+                    expert_models=expert_mlps,
+                    init_lambda=self.init_lambda,
+                    batch_first=True,
+                    router_hidden_layers=self.router_hidden_layers,
+                    batch_reduce=self.batch_reduce,
+                )
+            )
+
+        print(moe_model)
+        return moe_model
+
+    @functools.cache
+    def get_shuffled_test_loader_iter(self, task: str) -> DataLoader:
+        """
+        Loader of test dataset for test-time adaptation. labels are not needed.
+
+        Args:
+            task (str): The name of the task.
+
+        Returns:
+            DataLoader: The data loader for the test dataset.
+        """
+        # dataloader_kwargs = dict(self.dataloader_kwargs)
+        # dataloader_kwargs.update(dict(shuffle=True, collate_fn=default_data_collator))
+
+        dataset = self.modelpool.load_test_dataset(task)
+        log.info("get_shuffled_test_loader_iter")
+        loader = DataLoader(
+            dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            collate_fn=default_data_collator,
+        )
+        # loader = DataLoader(dataset, **dataloader_kwargs)
+        if self.fabric is not None:
+            loader = self.fabric.setup_dataloaders(loader)
+        return iter(InfiniteDataLoader(loader))
+
+    def compute_logits(
+        self,
+        module: Union[T5ForConditionalGeneration],
+        batch,
+        task: str,
+    ) -> Tensor:
+        """
+        Compute the logits for the given images and task.
+
+        Args:
+            module: The model module.
+            images (Tensor): The input images.
+            task (str): The name of the task.
+
+        Returns:
+            Tensor: The computed logits.
+        """
+        input_ids: Tensor = batch["input_ids"]
+        attention_mask: Tensor = batch["attention_mask"]
+
+        # remove padding tokens from the input
+        while attention_mask[:, -1].eq(0).all():
+            input_ids = input_ids[:, :-1]
+            attention_mask = attention_mask[:, :-1]
+
+        outputs = module(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            decoder_input_ids=torch.ones(
+                input_ids.size(0), 1, dtype=torch.long, device=input_ids.device
+            ),
+        )
+        logits = outputs.logits[:, 0, :]
+        return logits
+
+    def test_time_adaptation(self, module):
+        """
+        Perform test-time adaptation for the given module.
+
+        Args:
+            module (WeightEnsemblingMoE): The MoE module to adapt.
+
+        Returns:
+            WeightEnsemblingMoE: The adapted MoE module.
+        """
+        self.on_test_time_adaptation_start()
+
+        # configure optimizer
+        if self.optimizer == "adam":
+            print([name for name, p in module.named_parameters() if p.requires_grad])
+            optimizer = torch.optim.Adam(
+                [p for p in module.parameters() if p.requires_grad], lr=self.lr
+            )
+        else:
+            raise ValueError(f"Unsupported optimizer: {self.optimizer}")
+
+        module, optimizer = self.fabric.setup(module, optimizer)
+
+        module.train()
+        # module.merge_weights()
+        for step_idx in (
+            pbar := tqdm(
+                range(self.max_steps if not self.is_debug_mode else 1),
+                ("[DEBUG MODE] " if self.is_debug_mode else "")
+                + "WEMoE Test-time adaptation",
+                dynamic_ncols=True,
+            )
+        ):
+            total_loss = 0
+            for task in self.modelpool.model_names:
+                with self.profile("data loading"):
+                    batch = next(self.get_shuffled_test_loader_iter(task))
+                with self.profile("forward pass"):
+                    logits = self.compute_logits(module, batch, task)
+                    logits = logits.mean(dim=0, keepdim=True)
+                    loss = entropy_loss(logits)
+                    total_loss += loss
+                with self.profile("backward pass"):
+                    self.fabric.backward(loss, retain_graph=True)
+
+            with self.profile("optimizer step"):
+                optimizer.step()
+                optimizer.zero_grad()
+
+            metrics = {
+                "train/loss": total_loss.item(),
+            }
+            self.fabric.log_dict(metrics, step=step_idx)
+            pbar.set_postfix(metrics)
+
+        log.info(get_memory_usage(f"after adamerging, the memory usage of GPU is:"))
+        self.print_profile_summary()
+        return module
+
+    def on_test_time_adaptation_start(self):
+        """
+        Something to do before the test-time adaptation starts. Such as setting up the task-specific heads.
+        """
+        pass
+
+    def run(self, modelpool: Seq2SeqLMPool, **kwargs):
+        """
+        Run the WeightEnsemblingMoEAlgorithm to fuse models using Weight Ensembling Mixture of Experts.
+
+        Args:
+            modelpool (ModelPool): The pool of models to be fused.
+
+        Returns:
+            WeightEnsemblingMoE: The fused MoE model.
+        """
+        log.info("Fusing models using layer-wise adaptive merging.")
+        self.modelpool = modelpool
+
+        with timeit_context("upscaling models to a weight-ensembling MoE model"):
+            moe_model = self.construct_moe_model()
+            print_parameters(moe_model)
+
+        if self.checkpoint != False:
+            log.info(
+                f"load checkpoint from {self.checkpoint}, test-time adaptation will be skipped."
+            )
+            self.load_checkpoint(moe_model, self.checkpoint)
+        else:
+            with self.profile("test-time adaptation"):
+                moe_model = self.test_time_adaptation(moe_model)
+            if self.save_checkpoint != False:
+                log.info(f"save checkpoint to {self.save_checkpoint}")
+                self.save_checkpoint(moe_model, self.save_checkpoint)
+
+            if lightning.fabric.wrappers.is_wrapped(moe_model):
+                moe_model = lightning.fabric.wrappers._unwrap_objects(moe_model)
+
+        # enable sample-wise adaptation
+        moe_model.batch_reduce = False
+        self.print_profile_summary()
+        return moe_model

fusion_bench/method/we_moe/utils.py

@@ -0,0 +1,15 @@
+import torch
+
+
+def get_memory_usage(desc):
+    """
+    obtain the current GPU memory usage
+
+    Returns:
+        str: A string containing the allocated and cached memory in MB.
+    """
+    allocated = torch.cuda.memory_allocated() / 1024**2  # 转换为 MB
+    cached = torch.cuda.memory_reserved() / 1024**2  # 转换为 MB
+    return (
+        f"{desc}\nAllocated Memory: {allocated:.2f} MB\nCached Memory: {cached:.2f} MB"
+    )

fusion_bench/method/weighted_average/llama.py

@@ -7,11 +7,11 @@ from transformers import PreTrainedModel
 from typing_extensions import override
 
 from fusion_bench.method import BaseAlgorithm
+from fusion_bench.mixins import auto_register_config
 from fusion_bench.modelpool import CausalLMPool
 from fusion_bench.utils import timeit_context
 from fusion_bench.utils.state_dict_arithmetic import state_dict_add, state_dict_mul
 from fusion_bench.utils.type import StateDictType
-from fusion_bench.mixins import auto_register_config
 
 log = logging.getLogger(__name__)
 

fusion_bench/mixins/clip_classification.py

@@ -22,6 +22,7 @@ from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 from transformers import CLIPModel, CLIPProcessor, CLIPVisionModel
 
+from fusion_bench import cache_with_joblib
 from fusion_bench.dataset.clip_dataset import CLIPDataset
 from fusion_bench.mixins import LightningFabricMixin
 from fusion_bench.modelpool import CLIPVisionModelPool
@@ -46,7 +47,6 @@ class CLIPClassificationMixin(LightningFabricMixin):
 
     - `_dataloader_kwargs` (Dict[str, Any]): Keyword arguments for the dataloader.
     - `modelpool` (CLIPVisionModelPool): The model pool containing the CLIP models.
-    - `zeroshot_weights_cache_dir` (Optional[str]): The directory to cache the zero-shot weights.
     """
 
     dataloader_kwargs: Dict[str, Any] = {}
@@ -54,7 +54,6 @@ class CLIPClassificationMixin(LightningFabricMixin):
     modelpool: CLIPVisionModelPool = None
     _clip_processor: CLIPProcessor = None
     # a dict of zeroshot weights for each task, each key is the task name
-    zeroshot_weights_cache_dir: str = "outputs/cache/clip_zeroshot_weights"
     zeroshot_weights: Dict[str, torch.Tensor] = {}
     whether_setup_zero_shot_classification_head = False
 
@@ -131,26 +130,16 @@ class CLIPClassificationMixin(LightningFabricMixin):
         self.visual_projection = self.fabric.to_device(self.visual_projection)
         self.logit_scale_exp = self.fabric.to_device(self.logit_scale_exp)
 
-        # get cache directory
-        if self.modelpool.has_pretrained:
-            model_name = self.modelpool.get_model_config("_pretrained_")
-            if not isinstance(model_name, str):
-                model_name = model_name.pretrained_model_name_or_path
-        else:
-            model_name = self.modelpool.get_model_config(self.modelpool.model_names[0])
-            if not isinstance(model_name, str):
-                model_name = model_name.pretrained_model_name_or_path
-        cache_dir = os.path.join(
-            self.zeroshot_weights_cache_dir,
-            os.path.normpath(model_name.split("/")[-1]),
-        )
-        if not os.path.exists(cache_dir):
-            log.info(
-                f"Creating cache directory for zero-shot classification head at {cache_dir}"
-            )
-            os.makedirs(cache_dir)
+        @cache_with_joblib()
+        def construct_classification_head(task: str):
+            nonlocal clip_classifier
+
+            classnames, templates = get_classnames_and_templates(task)
+            clip_classifier.set_classification_task(classnames, templates)
+            zeroshot_weights = clip_classifier.zeroshot_weights.detach().clone()
+
+            return zeroshot_weights
 
-        log.info(f"cache directory for zero-shot classification head: {cache_dir}")
         for task in tqdm(
             self.modelpool.model_names if task_names is None else task_names,
             "Setting up zero-shot classification head",
@@ -158,27 +147,7 @@ class CLIPClassificationMixin(LightningFabricMixin):
         ):
             zeroshot_weights = None
             if self.fabric.is_global_zero:
-                cache_file = os.path.join(
-                    cache_dir, os.path.normpath(f"{task}_zeroshot_weights.pt")
-                )
-                if os.path.exists(cache_file):
-                    zeroshot_weights = torch.load(
-                        cache_file,
-                        map_location="cpu",
-                        weights_only=True,
-                    ).detach()
-                    log.info(
-                        f"Loadded cached zeroshot weights for task: {task}, shape: {zeroshot_weights.shape}"
-                    )
-                else:
-                    log.info(
-                        f"Construct zero shot classification head for task: {task}"
-                    )
-                    classnames, templates = get_classnames_and_templates(task)
-                    clip_classifier.set_classification_task(classnames, templates)
-                    zeroshot_weights = clip_classifier.zeroshot_weights.detach().clone()
-                    log.info(f"save zeroshot weights to {cache_file}")
-                    torch.save(zeroshot_weights, cache_file)
+                zeroshot_weights = construct_classification_head(task)
 
             self.fabric.barrier()
             self.zeroshot_weights[task] = self.fabric.broadcast(zeroshot_weights, src=0)