fusion-bench 0.2.6__py3-none-any.whl → 0.2.8__py3-none-any.whl

fusion_bench/compat/method/__init__.py

@@ -21,6 +21,7 @@ class AlgorithmFactory:
         "clip_task_wise_adamerging": ".adamerging.clip_task_wise_adamerging.CLIPTaskWiseAdaMergingAlgorithm",
         "clip_layer_wise_adamerging": ".adamerging.clip_layer_wise_adamerging.CLIPLayerWiseAdaMergingAlgorithm",
         "singular_projection_merging": "fusion_bench.method.smile_upscaling.singular_projection_merging.SingularProjectionMergingAlgorithm",
+        "clip_layer_wise_adamerging_surgery": ".surgery.clip_layer_wise_adamerging_surgery.CLIPLayerWiseAdaMergingSurgeryAlgorithm",
         # plug-and-play model merging methods
         "clip_concrete_task_arithmetic": ".concrete_subspace.clip_concrete_task_arithmetic.ConcreteTaskArithmeticAlgorithmForCLIP",
         "clip_concrete_task_wise_adamerging": ".concrete_subspace.clip_concrete_adamerging.ConcreteTaskWiseAdaMergingForCLIP",

fusion_bench/compat/method/base_algorithm.py

@@ -1,8 +1,11 @@
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 from omegaconf import DictConfig
 
+if TYPE_CHECKING:
+    from fusion_bench.programs.base_program import BaseHydraProgram
+
 __all__ = ["ModelFusionAlgorithm"]
 
 
@@ -18,6 +21,9 @@ class ModelFusionAlgorithm(ABC):
         config (DictConfig): Configuration for the algorithm.
     """
 
+    _program: "BaseHydraProgram" = None
+    """A reference to the program that is running the algorithm."""
+
     def __init__(self, algorithm_config: Optional[DictConfig] = None):
         """
         Initialize the model fusion algorithm with the given configuration.

fusion_bench/compat/modelpool/__init__.py

@@ -22,7 +22,7 @@ class ModelPoolFactory:
     """
 
     _modelpool = {
-        "NYUv2ModelPool": ".nyuv2_modelpool.NYUv2ModelPool",
+        "NYUv2ModelPool": "fusion_bench.modelpool.nyuv2_modelpool.NYUv2ModelPool",
         "huggingface_clip_vision": HuggingFaceClipVisionPool,
         "HF_GPT2ForSequenceClassification": GPT2ForSequenceClassificationPool,
         "AutoModelPool": ".huggingface_automodel.AutoModelPool",

fusion_bench/compat/taskpool/__init__.py

@@ -20,7 +20,7 @@ class TaskPoolFactory:
         "dummy": DummyTaskPool,
         "clip_vit_classification": ".clip_image_classification.CLIPImageClassificationTaskPool",
         "FlanT5GLUETextGenerationTaskPool": ".flan_t5_glue_text_generation.FlanT5GLUETextGenerationTaskPool",
-        "NYUv2TaskPool": ".nyuv2_taskpool.NYUv2TaskPool",
+        "NYUv2TaskPool": "fusion_bench.taskpool.nyuv2_taskpool.NYUv2TaskPool",
     }
 
     @staticmethod

fusion_bench/dataset/arc_agi/arc.py

@@ -152,6 +152,11 @@ class Task:
         tasks = []
         for test_data in data["test"]:
             task = cls.deserialize(
+                {
+                    "train": data["train"],
+                    "test": [test_data],
+                    "name": data.get("name", ""),
+                },
                 {
                     "train": data["train"],
                     "test": [test_data],

fusion_bench/dataset/arc_agi/preprocess.py

@@ -143,7 +143,7 @@ def format_and_filter(
     data = {
         "input_ids": prompt_tokens + output_tokens,
         "attention_mask": [1] * len(prompt_tokens) + [1] * len(output_tokens),
-        "labels": [-100] * len(prompt_tokens) + output_tokens,
+        "labels": prompt_tokens + output_tokens,
         "task_id": task_id,
         "num_prompt_tokens": len(prompt_tokens),
         "num_output_tokens": len(output_tokens),

fusion_bench/dataset/clip_dataset.py

@@ -65,4 +65,7 @@ class CLIPDataset(torch.utils.data.Dataset):
         else:
             # if processor is None, return the raw image directly
             inputs = image
+        # convert boolean label to int, this is for the case when the label is a binary classification task
+        if isinstance(item["label"], bool):
+            item["label"] = 1 if item["label"] else 0
         return inputs, item["label"]

fusion_bench/dataset/fer2013.py

@@ -0,0 +1,12 @@
+from datasets import load_dataset
+
+
+def load_fer2013(path: str = "clip-benchmark/wds_fer2013", split: str = "train"):
+    dataset = load_dataset(path, split=split)
+    dataset = dataset.remove_columns(["__key__", "__url__"])
+    dataset = dataset.rename_columns({"jpg": "image", "cls": "label"})
+    return dataset
+
+if __name__ == "__main__":
+    dataset = load_fer2013(split="test")
+    print(dataset)

fusion_bench/dataset/llama/__init__.py

@@ -0,0 +1 @@
+from . import collate

fusion_bench/dataset/llama/alpaca.py

@@ -1,16 +1,100 @@
 import logging
 import os
+import warnings
 from typing import Any, Dict, List, Optional
 
 from datasets import Dataset, load_dataset, load_from_disk
+from lightning.fabric.utilities import rank_zero_only
+from tqdm.auto import tqdm
 from transformers import PreTrainedTokenizer
 
 import fusion_bench
+from fusion_bench.utils import timeit_context
 
 log = logging.getLogger(__name__)
 
 
-def tokenize_alpaca_dataset(
+def convert_alpaca_to_conversation(alpaca_data: List[Dict[str, str]]):
+    """
+    Convert Alpaca format data to conversation format.
+
+    Args:
+        alpaca_data (list): List of dictionaries in Alpaca format with
+            'instruction', 'input', and 'output' keys
+
+    Returns:
+        list: List of conversations in ChatML format
+    """
+    conversations = []
+
+    for item in tqdm(
+        alpaca_data,
+        "Converting Alpaca to conversations",
+        disable=not rank_zero_only.rank == 0,
+    ):
+        # Skip if required fields are missing
+        if not item.get("instruction") or not item.get("output"):
+            continue
+
+        conversation = []
+
+        # Create user message
+        user_content = item["instruction"]
+        if item.get("input") and item["input"].strip():
+            user_content += f"\n\n{item['input']}"
+
+        conversation.append({"role": "user", "content": user_content})
+
+        # Create assistant message
+        conversation.append({"role": "assistant", "content": item["output"]})
+
+        conversations.append(conversation)
+
+    return conversations
+
+
+def load_tokenized_alpaca_dataset(
+    tokenizer: PreTrainedTokenizer,
+    path: str = "yahma/alpaca-cleaned",
+    split: str = "train",
+    cache_path: Optional[str] = None,
+):
+    """
+    Load and tokenized Alpaca dataset and Alpaca-like dataset.
+
+    Args:
+        tokenizer (PreTrainedTokenizer): The tokenizer to use for tokenizing the dataset.
+        path (str, optional): The path to the Alpaca dataset. Defaults to "yahma/alpaca-cleaned".
+        split (str, optional): The dataset split to load (e.g., "train", "test"). Defaults to "train".
+        cache_path (Optional[str], optional): The path to cache the tokenized dataset. If provided and the cache exists,
+            the dataset will be loaded from the cache. Defaults to None.
+
+    Returns:
+        Dataset: The tokenized dataset.
+    """
+    if cache_path is not None and os.path.exists(cache_path):
+        dataset = load_from_disk(cache_path)
+        if split is not None and split in dataset:
+            return dataset[split]
+        else:
+            return dataset
+
+    dataset = load_dataset(path, split=split)
+
+    alpaca_data = dataset.to_list()
+    conversations = convert_alpaca_to_conversation(alpaca_data)
+    with timeit_context("Tokenizing dataset"):
+        tokenized_dataset = tokenizer.apply_chat_template(
+            conversations, return_dict=True
+        )
+    tokenized_dataset = Dataset.from_dict(tokenized_dataset)
+
+    if cache_path is not None and rank_zero_only.rank == 0:
+        tokenized_dataset.save_to_disk(cache_path)
+    return tokenized_dataset
+
+
+def _tokenize_alpaca_dataset_with_template(
     dataset: Dataset,
     tokenizer: PreTrainedTokenizer,
     max_length: int = 2048,
@@ -32,6 +116,10 @@ def tokenize_alpaca_dataset(
     Returns:
         Tokenized dataset
     """
+    warnings.warn(
+        "This function is deprecated. Use `apply_chat_template` from `transformers` instead.",
+        DeprecationWarning,
+    )
 
     def prepare_samples(samples: Dict[str, List[str]]) -> Dict[str, List[List[int]]]:
         # Format prompts based on whether input field exists
@@ -115,7 +203,7 @@ def tokenize_alpaca_dataset(
     return tokenized_dataset
 
 
-def load_tokenized_alpaca_dataset_from_json(
+def load_tokenized_alpaca_dataset_from_json_with_prompt(
     data_files: str,
     tokenizer: PreTrainedTokenizer,
     max_length: int,
@@ -138,5 +226,7 @@ def load_tokenized_alpaca_dataset_from_json(
     dataset = load_dataset("json", data_files=data_files)
     if split is not None:
         dataset = dataset[split]
-    dataset = tokenize_alpaca_dataset(dataset, tokenizer, max_length=max_length)
+    dataset = _tokenize_alpaca_dataset_with_template(
+        dataset, tokenizer, max_length=max_length
+    )
     return dataset

fusion_bench/dataset/llama/collate.py

@@ -7,7 +7,7 @@ from torch.nn.utils.rnn import pad_sequence
 
 def padded_collate_sft(
     batch: List[Dict[str, List[int]]],
-    padding_idx: int = 0,
+    pad_token_id: int = 0,
     input_ids_key: str = "input_ids",
     attention_mask_key: Optional[str] = "attention_mask",
     labels_key: Optional[str] = "labels",
@@ -28,7 +28,7 @@ def padded_collate_sft(
     input_ids = pad_sequence(
         [torch.tensor(x[input_ids_key]) for x in batch],
         batch_first=True,
-        padding_value=padding_idx,
+        padding_value=pad_token_id,
     )
     if attention_mask_key is not None and attention_mask_key in batch[0]:
         attention_mask = pad_sequence(
@@ -38,6 +38,12 @@ def padded_collate_sft(
         )
     else:
         attention_mask = None
+
+    for i, item in enumerate(batch):
+        # if labels_key not in item, copy input_ids to labels_key
+        if labels_key not in item:
+            item[labels_key] = item[input_ids_key]
+
     labels = pad_sequence(
         [torch.tensor(x[labels_key]) for x in batch],
         batch_first=True,
@@ -58,3 +64,57 @@ def padded_collate_sft(
             collated_batch[key] = [x[key] for x in batch]
 
     return collated_batch
+
+
+def bradley_terry_rm_collate(
+    batch: List[Dict[str, List[int]]],
+    pad_token_id: int = 0,
+    padding_side="right",
+):
+    """
+    Collate function for Bradley-Terry reward modeling.
+
+    Args:
+        batch (List[Dict[str, List[int]]]): A list of dictionaries containing input, label pairs.
+        pad_token_id (int): Padding index for input ids. Defaults to 0.
+
+    Returns:
+        Dict[str, torch.Tensor]: Collated input and label tensors. The first half of the batch is the winner, and the second half is the loser.
+    """
+    converted_batch = []
+    for item in batch:
+        new_item = {
+            "input_ids": item["chosen_input_ids"],
+            "attention_mask": item["chosen_attention_mask"],
+        }
+        converted_batch.append(new_item)
+    for item in batch:
+        new_item = {
+            "input_ids": item["rejected_input_ids"],
+            "attention_mask": item["rejected_attention_mask"],
+        }
+        converted_batch.append(new_item)
+
+    input_ids = pad_sequence(
+        [torch.tensor(x["input_ids"]) for x in converted_batch],
+        batch_first=True,
+        padding_value=pad_token_id,
+        padding_side=padding_side,
+    )
+    attention_mask = pad_sequence(
+        [torch.tensor(x["attention_mask"]) for x in converted_batch],
+        batch_first=True,
+        padding_value=0,
+        padding_side=padding_side,
+    )
+
+    collated_batch = {"input_ids": input_ids, "attention_mask": attention_mask}
+    for key in batch[0]:
+        if key not in [
+            "chosen_input_ids",
+            "chosen_attention_mask",
+            "rejected_input_ids",
+            "rejected_attention_mask",
+        ]:
+            collated_batch[key] = [x[key] for x in batch]
+    return collated_batch

fusion_bench/dataset/llama/metamathqa.py

@@ -0,0 +1,50 @@
+import os
+from typing import TYPE_CHECKING, Optional
+
+from datasets import Dataset, load_dataset, load_from_disk
+from lightning.fabric.utilities import rank_zero_only
+from tqdm.auto import tqdm
+
+from fusion_bench.utils import timeit_context
+
+from .alpaca import convert_alpaca_to_conversation
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizer
+
+
+def load_tokenized_metamathqa(
+    tokenizer: "PreTrainedTokenizer",
+    path: str = "meta-math/MetaMathQA",
+    split: str = "train",
+    cache_path: Optional[str] = None,
+):
+    if cache_path is not None and os.path.exists(cache_path):
+        dataset = load_from_disk(cache_path)
+        if split is not None and split in dataset:
+            return dataset[split]
+        else:
+            return dataset
+
+    dataset = load_dataset(path, split=split)
+
+    # convert dataset to alpaca format and save to ../data/MetaMathQA.json
+    alpaca_dataset = []
+    for example in tqdm(dataset, disable=not rank_zero_only.rank == 0):
+        alpaca_example = {
+            "instruction": example["query"],
+            "input": "",
+            "output": example["response"],
+        }
+        alpaca_dataset.append(alpaca_example)
+
+    conversations = convert_alpaca_to_conversation(alpaca_dataset)
+    with timeit_context("Tokenizing dataset"):
+        tokenized_dataset = tokenizer.apply_chat_template(
+            conversations, return_dict=True
+        )
+    tokenized_dataset = Dataset.from_dict(tokenized_dataset)
+
+    if cache_path is not None and rank_zero_only.rank == 0:
+        tokenized_dataset.save_to_disk(cache_path)
+    return tokenized_dataset

fusion_bench/dataset/llama/preference_700k.py

@@ -0,0 +1,70 @@
+import logging
+import os
+from copy import deepcopy
+from typing import TYPE_CHECKING, Optional
+
+from datasets import Dataset, load_dataset, load_from_disk
+from lightning.fabric.utilities import rank_zero_only
+from tqdm.auto import tqdm
+
+from fusion_bench.utils import timeit_context
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizer
+
+log = logging.getLogger(__name__)
+
+
+def load_tokenized_preference_700k_for_rlhf(
+    tokenizer: "PreTrainedTokenizer",
+    path: str = "hendrydong/preference_700K",
+    split: str = "train",
+    num_proc: int = 8,
+    cache_path: Optional[str] = None,
+):
+    R"""
+    Load and tokenized Preference 700k dataset for Bradley-Terry ranking model.
+
+    The returned dataset contains the following fields:
+
+    - chosen_input_ids: The input token ids for the winner.
+    - chosen_attention_mask: The attention mask for the winner.
+    - rejected_input_ids: The input token ids for the loser.
+    - rejected_attention_mask: The attention mask for the loser.
+    """
+    if cache_path is not None and os.path.exists(cache_path):
+        dataset = load_from_disk(cache_path)
+        return dataset
+
+    dataset = load_dataset(path, split=split)
+
+    def tokenize(sample):
+        sample["chosen_chat"] = tokenizer.apply_chat_template(
+            sample["chosen"], tokenize=False, add_generation_prompt=False
+        )
+        sample["rejected_chat"] = tokenizer.apply_chat_template(
+            sample["rejected"], tokenize=False, add_generation_prompt=False
+        )
+
+        tokenized_pos = tokenizer(sample["chosen_chat"], truncation=True)
+        tokenized_neg = tokenizer(sample["rejected_chat"], truncation=True)
+
+        # Ensure that the chosen response does not contain an PAD token
+        sample["chosen_input_ids"] = tokenized_pos["input_ids"]
+        sample["chosen_attention_mask"] = tokenized_pos["attention_mask"]
+        if tokenizer.pad_token_id in tokenized_pos["input_ids"]:
+            log.warning(f"Prompt contains PAD token: {sample['chosen_chat']}")
+
+        sample["rejected_input_ids"] = tokenized_neg["input_ids"]
+        sample["rejected_attention_mask"] = tokenized_neg["attention_mask"]
+        # Ensure that the rejected response does not contain an PAD token
+        if tokenizer.pad_token_id in tokenized_neg["input_ids"]:
+            log.warning(f"Prompt contains PAD token: {sample['rejected_chat']}")
+
+        return sample
+
+    dataset = dataset.map(tokenize, num_proc=num_proc)
+
+    if cache_path is not None and rank_zero_only.rank == 0:
+        dataset.save_to_disk(cache_path)
+    return dataset

fusion_bench/dataset/llama/stanford_shp.py

@@ -0,0 +1,90 @@
+import os
+from copy import deepcopy
+from typing import TYPE_CHECKING, Optional
+
+from datasets import Dataset, load_dataset, load_from_disk
+from lightning.fabric.utilities import rank_zero_only
+from tqdm.auto import tqdm
+
+from fusion_bench.utils import timeit_context
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizer
+
+
+def load_tokenized_stanford_shp_for_rlhf(
+    tokenizer: "PreTrainedTokenizer",
+    path: str = "stanfordnlp/SHP",
+    split: str = "train",
+    num_proc: int = 8,
+    cache_path: Optional[str] = None,
+):
+    if cache_path is not None and os.path.isdir(cache_path):
+        dataset = load_from_disk(cache_path)
+        return dataset
+
+    dataset = load_dataset(path, split=split)
+
+    def tokenize(sample):
+        """
+        - history: the post title concatented to the post body (string)
+        - human_ref_A: text of comment A (string)
+        - human_ref_B: text of comment B (string)
+        - labels: the preference label -- it is 1 if A is preferred to B; 0 if B is preferred to A. This was randomized such that the label distribution is roughly 50/50. (integer)
+        """
+        # Create a conversation with the post title and body, followed by comments
+        conversation = [{"role": "user", "content": sample["history"]}]
+        if sample["labels"] == 0:
+            sample["chosen"] = deepcopy(conversation).append(
+                {"role": "assistant", "content": sample["human_ref_B"]}
+            )
+            sample["rejected"] = deepcopy(conversation).append(
+                {"role": "assistant", "content": sample["human_ref_A"]}
+            )
+        else:
+            sample["chosen"] = deepcopy(conversation).append(
+                {"role": "assistant", "content": sample["human_ref_A"]}
+            )
+            sample["rejected"] = deepcopy(conversation).append(
+                {"role": "assistant", "content": sample["human_ref_B"]}
+            )
+
+        # apply chat template
+        sample["chosen_chat"] = tokenizer.apply_chat_template(
+            sample["chosen"], tokenize=False, add_generation_prompt=False
+        )
+        sample["rejected_chat"] = tokenizer.apply_chat_template(
+            sample["rejected"], tokenize=False, add_generation_prompt=False
+        )
+
+        # tokenize the conversation
+        tokenized_pos = tokenizer(sample["chosen_chat"], truncation=True)
+        tokenized_neg = tokenizer(sample["rejected_chat"], truncation=True)
+
+        # Ensure that the chosen response does not contain an EOS token
+        sample["chosen_input_ids"] = tokenized_pos["input_ids"]
+        sample["chosen_attention_mask"] = tokenized_pos["attention_mask"]
+        assert (
+            tokenizer.eos_token_id not in tokenized_pos["input_ids"][:-1]
+        ), f"Prompt contains EOS token: {sample['positive']}"
+        if sample["chosen_input_ids"][-1] != tokenizer.eos_token_id:
+            sample["chosen_input_ids"].append(tokenizer.eos_token_id)
+            sample["chosen_attention_mask"].append(1)
+
+        sample["rejected_input_ids"] = tokenized_neg["input_ids"]
+        sample["rejected_attention_mask"] = tokenized_neg["attention_mask"]
+        # Ensure that the rejected response does not contain an EOS token
+        assert (
+            tokenizer.eos_token_id not in tokenized_neg["input_ids"][:-1]
+        ), f"Prompt contains EOS token: {sample['rejected']}"
+        if sample["rejected_input_ids"][-1] != tokenizer.eos_token_id:
+            sample["rejected_input_ids"].append(tokenizer.eos_token_id)
+            sample["rejected_attention_mask"].append(1)
+
+        return sample
+
+    dataset = dataset.map(tokenize, num_proc=num_proc)
+
+    if cache_path is not None and rank_zero_only.rank == 0:
+        dataset.save_to_disk(cache_path)
+    return dataset

fusion_bench/dataset/llama/ultrachat.py

@@ -0,0 +1,58 @@
+import os
+from typing import TYPE_CHECKING, Optional
+
+from datasets import Dataset, load_dataset, load_from_disk
+from lightning.fabric.utilities import rank_zero_only
+from tqdm.auto import tqdm
+
+from fusion_bench.utils import timeit_context
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizer
+
+
+def load_tokenized_ultrachat_200k(
+    tokenizer: "PreTrainedTokenizer",
+    path: str = "HuggingFaceH4/ultrachat_200k",
+    split: str = "train_sft",
+    num_proc: int = 8,
+    cache_path: Optional[str] = None,
+):
+    R"""
+    Load and tokenized Ultrachat 200k dataset for Bradley-Terry ranking model.
+
+    The returned dataset contains the following fields:
+
+    - input_ids: The input token ids for the winner.
+    - attention_mask: The attention mask for the winner.
+    """
+    if cache_path is not None and os.path.exists(cache_path):
+        dataset = load_from_disk(cache_path)
+        return dataset
+
+    dataset = load_dataset(path, split=split)
+
+    def tokenize(sample):
+
+        # ? is it necessary to `.replace(tokenizer.bos_token, "")`?
+        sample["input_ids"] = tokenizer.apply_chat_template(
+            sample["messages"], tokenize=True, add_generation_prompt=False
+        )
+        sample["attention_mask"] = [1] * len(sample["input_ids"])
+
+        return sample
+
+    dataset = dataset.map(tokenize, num_proc=num_proc)
+
+    if cache_path is not None and rank_zero_only.rank == 0:
+        dataset.save_to_disk(cache_path)
+    return dataset
+
+
+if __name__ == "__main__":
+    # Example usage and testing
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")
+    dataset = load_tokenized_ultrachat_200k(tokenizer)
+    print(dataset)

fusion_bench/method/__init__.py

@@ -10,7 +10,7 @@ _import_structure = {
     "dummy": ["DummyAlgorithm"],
     # single task learning (fine-tuning)
     "classification": ["ImageClassificationFineTuningForCLIP"],
-    "lm_finetune": ["FullFinetuneSFT", "PeftFinetuneSFT"],
+    "lm_finetune": ["FullFinetuneSFT", "PeftFinetuneSFT", "BradleyTerryRewardModeling"],
     # analysis
     "analysis": ["TaskVectorCosSimilarity", "TaskVectorViolinPlot"],
     # model ensemble methods
@@ -49,6 +49,7 @@ _import_structure = {
         "PWEMoExactParetoOptimalForCLIP",
     ],
     "ada_svd": ["AdaSVDMergingForCLIPVisionModel"],
+    "task_singular_vector": ["TaskSingularVectorMerging"],
     # plug-and-play model merging methods
     "concrete_subspace": [
         "ConcreteTaskArithmeticAlgorithmForCLIP",
@@ -153,6 +154,7 @@ if TYPE_CHECKING:
         SparseLoForLlama,
     )
     from .task_arithmetic import TaskArithmeticAlgorithm
+    from .task_singular_vector import TaskSingularVectorMerging
     from .ties_merging import TiesMergingAlgorithm
     from .we_moe import CLIPWeightEnsemblingMoEAlgorithm
     from .weighted_average import WeightedAverageAlgorithm, WeightedAverageForLLama

fusion_bench/method/adamerging/layer_wise_adamerging.py

@@ -1,12 +1,12 @@
 import logging
 import os
 from abc import abstractmethod
-from typing import Any, List, Mapping, Union, cast  # noqa: F401
+from typing import TYPE_CHECKING, Any, List, Mapping, TypeVar, Union, cast  # noqa: F401
 
 import torch
 from lightning.fabric.utilities.rank_zero import rank_zero_only
 from omegaconf import DictConfig
-from torch import Tensor
+from torch import Tensor, nn
 from torch.utils.data import DataLoader
 from tqdm.autonotebook import tqdm
 
@@ -19,10 +19,14 @@ from fusion_bench.models.wrappers.layer_wise_fusion import (
     get_layer_wise_weights,
 )
 from fusion_bench.utils.data import load_tensor_from_file
+from fusion_bench.utils.type import TorchModelType
 
 from .entropy_loss import entropy_loss
 from .utils import get_memory_usage
 
+if TYPE_CHECKING:
+    from fusion_bench.programs.fabric_fusion_program import FabricModelFusionProgram
+
 log = logging.getLogger(__name__)
 
 
@@ -31,6 +35,9 @@ class LayerWiseAdaMergingAlgorithm(
     LightningFabricMixin,
     SimpleProfilerMixin,
 ):
+    _program: "FabricModelFusionProgram"
+    """The program that this algorithm is running on."""
+
     """
     Implements the Layer-Wise AdaMerging Algorithm.
 
@@ -48,7 +55,7 @@ class LayerWiseAdaMergingAlgorithm(
         super().__init__(algorithm_config)
 
     @torch.no_grad()
-    def construct_layer_wise_merged_model(self, modelpool: ModelPool):
+    def construct_layer_wise_merged_model(self, modelpool: "ModelPool"):
         """
         Constructs a wrapped layer-wise merged model from model pool.
 
@@ -183,7 +190,7 @@ class LayerWiseAdaMergingAlgorithm(
         """
         pass
 
-    def test_time_adaptation(self, module: LayerWiseMergedModel):
+    def test_time_adaptation(self, module: "LayerWiseMergedModel[TorchModelType]"):
         """
         Perform test-time adaptation on the merged model.