fusion-bench 0.2.23__py3-none-any.whl → 0.2.25__py3-none-any.whl

fusion_bench/__init__.py

@@ -5,46 +5,156 @@
 # ██║     ╚██████╔╝███████║██║╚██████╔╝██║ ╚████║      ██████╔╝███████╗██║ ╚████║╚██████╗██║  ██║
 # ╚═╝      ╚═════╝ ╚══════╝╚═╝ ╚═════╝ ╚═╝  ╚═══╝      ╚═════╝ ╚══════╝╚═╝  ╚═══╝ ╚═════╝╚═╝  ╚═╝
 # flake8: noqa: F401
-from . import (
-    constants,
-    dataset,
-    method,
-    metrics,
-    mixins,
-    modelpool,
-    models,
-    optim,
-    programs,
-    taskpool,
-    tasks,
-    utils,
-)
+import sys
+from typing import TYPE_CHECKING
+
+from fusion_bench.utils.lazy_imports import LazyImporter
+
+from . import constants, metrics, optim, tasks
 from .constants import RuntimeConstants
-from .method import BaseAlgorithm, BaseModelFusionAlgorithm
-from .mixins import auto_register_config
-from .modelpool import BaseModelPool
-from .models import (
-    create_default_model_card,
-    load_model_card_template,
-    save_pretrained_with_remote_code,
-    separate_io,
-)
-from .programs import BaseHydraProgram
-from .taskpool import BaseTaskPool
-from .utils import (
-    BoolStateDictType,
-    LazyStateDict,
-    StateDictType,
-    TorchModelType,
-    cache_with_joblib,
-    get_rankzero_logger,
-    import_object,
-    instantiate,
-    parse_dtype,
-    print_parameters,
-    seed_everything_by_time,
-    set_default_cache_dir,
-    set_print_function_call,
-    set_print_function_call_permeanent,
-    timeit_context,
-)
+from .method import _available_algorithms
+
+_extra_objects = {
+    "RuntimeConstants": RuntimeConstants,
+    "constants": constants,
+    "metrics": metrics,
+    "optim": optim,
+    "tasks": tasks,
+}
+_import_structure = {
+    "dataset": ["CLIPDataset"],
+    "method": _available_algorithms,
+    "mixins": [
+        "CLIPClassificationMixin",
+        "FabricTrainingMixin",
+        "HydraConfigMixin",
+        "LightningFabricMixin",
+        "OpenCLIPClassificationMixin",
+        "PyinstrumentProfilerMixin",
+        "SimpleProfilerMixin",
+        "YAMLSerializationMixin",
+        "auto_register_config",
+    ],
+    "modelpool": [
+        "AutoModelPool",
+        "BaseModelPool",
+        "CausalLMBackbonePool",
+        "CausalLMPool",
+        "CLIPVisionModelPool",
+        "GPT2ForSequenceClassificationPool",
+        "HuggingFaceGPT2ClassificationPool",
+        "NYUv2ModelPool",
+        "OpenCLIPVisionModelPool",
+        "PeftModelForSeq2SeqLMPool",
+        "ResNetForImageClassificationPool",
+        "Seq2SeqLMPool",
+        "SequenceClassificationModelPool",
+    ],
+    "models": [
+        "create_default_model_card",
+        "load_model_card_template",
+        "save_pretrained_with_remote_code",
+        "separate_load",
+        "separate_save",
+    ],
+    "programs": ["BaseHydraProgram", "FabricModelFusionProgram"],
+    "taskpool": [
+        "BaseTaskPool",
+        "CLIPVisionModelTaskPool",
+        "DummyTaskPool",
+        "GPT2TextClassificationTaskPool",
+        "LMEvalHarnessTaskPool",
+        "OpenCLIPVisionModelTaskPool",
+        "NYUv2TaskPool",
+    ],
+    "utils": [
+        "ArithmeticStateDict",
+        "BoolStateDictType",
+        "LazyStateDict",
+        "StateDictType",
+        "TorchModelType",
+        "cache_with_joblib",
+        "get_rankzero_logger",
+        "import_object",
+        "instantiate",
+        "parse_dtype",
+        "print_parameters",
+        "seed_everything_by_time",
+        "set_default_cache_dir",
+        "set_print_function_call",
+        "set_print_function_call_permeanent",
+        "timeit_context",
+    ],
+}
+
+if TYPE_CHECKING:
+    from .dataset import CLIPDataset
+    from .method import BaseAlgorithm, BaseModelFusionAlgorithm
+    from .mixins import (
+        CLIPClassificationMixin,
+        FabricTrainingMixin,
+        HydraConfigMixin,
+        LightningFabricMixin,
+        OpenCLIPClassificationMixin,
+        PyinstrumentProfilerMixin,
+        SimpleProfilerMixin,
+        YAMLSerializationMixin,
+        auto_register_config,
+    )
+    from .modelpool import (
+        AutoModelPool,
+        BaseModelPool,
+        CausalLMBackbonePool,
+        CausalLMPool,
+        CLIPVisionModelPool,
+        GPT2ForSequenceClassificationPool,
+        HuggingFaceGPT2ClassificationPool,
+        NYUv2ModelPool,
+        OpenCLIPVisionModelPool,
+        PeftModelForSeq2SeqLMPool,
+        ResNetForImageClassificationPool,
+        Seq2SeqLMPool,
+        SequenceClassificationModelPool,
+    )
+    from .models import (
+        create_default_model_card,
+        load_model_card_template,
+        save_pretrained_with_remote_code,
+        separate_load,
+        separate_save,
+    )
+    from .programs import BaseHydraProgram, FabricModelFusionProgram
+    from .taskpool import (
+        BaseTaskPool,
+        CLIPVisionModelTaskPool,
+        DummyTaskPool,
+        GPT2TextClassificationTaskPool,
+        LMEvalHarnessTaskPool,
+        NYUv2TaskPool,
+        OpenCLIPVisionModelTaskPool,
+    )
+    from .utils import (
+        ArithmeticStateDict,
+        BoolStateDictType,
+        LazyStateDict,
+        StateDictType,
+        TorchModelType,
+        cache_with_joblib,
+        get_rankzero_logger,
+        import_object,
+        instantiate,
+        parse_dtype,
+        print_parameters,
+        seed_everything_by_time,
+        set_default_cache_dir,
+        set_print_function_call,
+        set_print_function_call_permeanent,
+        timeit_context,
+    )
+else:
+    sys.modules[__name__] = LazyImporter(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        extra_objects=_extra_objects,
+    )

fusion_bench/dataset/__init__.py

@@ -1,16 +1,20 @@
 # flake8: noqa F401
-from datasets import load_dataset
-from omegaconf import DictConfig, open_dict
+import sys
+from typing import TYPE_CHECKING
 
-from fusion_bench.utils import instantiate
+from omegaconf import DictConfig, open_dict
 
-from .clip_dataset import CLIPDataset
+from fusion_bench.utils.lazy_imports import LazyImporter
 
 
 def load_dataset_from_config(dataset_config: DictConfig):
     """
     Load the dataset from the configuration.
     """
+    from datasets import load_dataset
+
+    from fusion_bench.utils import instantiate
+
     assert hasattr(dataset_config, "type"), "Dataset type not specified"
     if dataset_config.type == "instantiate":
         return instantiate(dataset_config.object)
@@ -27,3 +31,22 @@ def load_dataset_from_config(dataset_config: DictConfig):
         return dataset
     else:
         raise ValueError(f"Unknown dataset type: {dataset_config.type}")
+
+
+_extra_objects = {
+    "load_dataset_from_config": load_dataset_from_config,
+}
+_import_structure = {
+    "clip_dataset": ["CLIPDataset"],
+}
+
+if TYPE_CHECKING:
+    from .clip_dataset import CLIPDataset
+
+else:
+    sys.modules[__name__] = LazyImporter(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        extra_objects=_extra_objects,
+    )

fusion_bench/dataset/clip_dataset.py

@@ -6,7 +6,7 @@ from typing import Optional, Tuple
 
 import torch
 from torch.utils.data import Dataset
-from transformers import CLIPProcessor, ProcessorMixin
+from transformers import BaseImageProcessor, CLIPProcessor, ProcessorMixin
 
 __all__ = ["CLIPDataset"]
 
@@ -60,7 +60,7 @@ class CLIPDataset(torch.utils.data.Dataset):
             raise ValueError("Each item should be a dictionary or a tuple of length 2")
         image = item["image"]
         if self.processor is not None:
-            if isinstance(self.processor, ProcessorMixin):
+            if isinstance(self.processor, (ProcessorMixin, BaseImageProcessor)):
                 # Apply the processor to the image to get the input tensor
                 inputs = self.processor(images=[image], return_tensors="pt")[
                     "pixel_values"

fusion_bench/method/__init__.py

@@ -2,6 +2,7 @@
 import sys
 from typing import TYPE_CHECKING
 
+from fusion_bench.utils import join_lists
 from fusion_bench.utils.lazy_imports import LazyImporter
 
 _import_structure = {
@@ -12,6 +13,8 @@ _import_structure = {
     "classification": [
         "ImageClassificationFineTuningForCLIP",
         "ContinualImageClassificationFineTuningForCLIP",
+        "ImageClassificationFineTuning",
+        "ImageClassificationFineTuning_Test",
     ],
     "lm_finetune": ["FullFinetuneSFT", "PeftFinetuneSFT", "BradleyTerryRewardModeling"],
     # analysis
@@ -26,9 +29,12 @@ _import_structure = {
     "linear": [
         "ExPOAlgorithm",
         "ExPOAlgorithmForLlama",
+        "SimpleAverageForCausalLM",
         "SimpleAverageForLlama",
+        "TaskArithmeticForCausalLM",
         "TaskArithmeticForLlama",
         "LinearInterpolationAlgorithm",
+        "TiesMergingForCausalLM",
     ],
     "slerp": ["SlerpMergeAlgorithm", "SlerpForCausalLM"],
     "simple_average": ["SimpleAverageAlgorithm"],
@@ -72,6 +78,7 @@ _import_structure = {
     "fw_merging": ["FrankWolfeHardAlgorithm", "FrankWolfeSoftAlgorithm"],
     "tall_mask": ["TallMaskTaskArithmeticAlgorithm"],
     "model_stock": ["ModelStock"],
+    "wudi": ["wudi_merging", "WUDIMerging"],
     # plug-and-play model merging methods
     "concrete_subspace": [
         "ConcreteTaskArithmeticAlgorithmForCLIP",
@@ -127,7 +134,10 @@ _import_structure = {
         "ProgressivePruningForMixtral",
     ],
 }
-
+_available_algorithms = join_lists(list(_import_structure.values()))
+_extra_objects = {
+    "_available_algorithms": _available_algorithms,
+}
 
 if TYPE_CHECKING:
     from .ada_svd import AdaSVDMergingForCLIPVisionModel
@@ -137,6 +147,8 @@ if TYPE_CHECKING:
     from .bitdelta import BitDeltaAlgorithm
     from .classification import (
         ContinualImageClassificationFineTuningForCLIP,
+        ImageClassificationFineTuning,
+        ImageClassificationFineTuning_Test,
         ImageClassificationFineTuningForCLIP,
     )
     from .concrete_subspace import (
@@ -184,8 +196,11 @@ if TYPE_CHECKING:
         ExPOAlgorithm,
         ExPOAlgorithmForLlama,
         LinearInterpolationAlgorithm,
+        SimpleAverageForCausalLM,
         SimpleAverageForLlama,
+        TaskArithmeticForCausalLM,
         TaskArithmeticForLlama,
+        TiesMergingForCausalLM,
     )
     from .lm_finetune import *
     from .mixture_of_experts import (
@@ -238,10 +253,12 @@ if TYPE_CHECKING:
         FlanT5WeightEnsemblingMoEAlgorithm,
     )
     from .weighted_average import WeightedAverageAlgorithm, WeightedAverageForLLama
+    from .wudi import WUDIMerging, wudi_merging
 
 else:
     sys.modules[__name__] = LazyImporter(
         __name__,
         globals()["__file__"],
         _import_structure,
+        extra_objects=_extra_objects,
     )

fusion_bench/method/classification/__init__.py

@@ -1,3 +1,28 @@
 # flake8: noqa F401
-from .clip_finetune import ImageClassificationFineTuningForCLIP
-from .continual_clip_finetune import ContinualImageClassificationFineTuningForCLIP
+import sys
+from typing import TYPE_CHECKING
+
+from fusion_bench.utils.lazy_imports import LazyImporter
+
+_import_structure = {
+    "clip_finetune": ["ImageClassificationFineTuningForCLIP"],
+    "continual_clip_finetune": ["ContinualImageClassificationFineTuningForCLIP"],
+    "image_classification_finetune": [
+        "ImageClassificationFineTuning",
+        "ImageClassificationFineTuning_Test",
+    ],
+}
+
+if TYPE_CHECKING:
+    from .clip_finetune import ImageClassificationFineTuningForCLIP
+    from .continual_clip_finetune import ContinualImageClassificationFineTuningForCLIP
+    from .image_classification_finetune import (
+        ImageClassificationFineTuning,
+        ImageClassificationFineTuning_Test,
+    )
+else:
+    sys.modules[__name__] = LazyImporter(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+    )

fusion_bench/method/classification/image_classification_finetune.py

@@ -0,0 +1,214 @@
+import os
+from typing import Optional
+
+import lightning as L
+import lightning.pytorch.callbacks as pl_callbacks
+import torch
+from lightning.pytorch.loggers import TensorBoardLogger
+from lightning_utilities.core.rank_zero import rank_zero_only
+from lit_learn.lit_modules import ERM_LitModule
+from omegaconf import DictConfig
+from torch import nn
+from torch.utils.data import DataLoader
+from torchmetrics.classification import Accuracy
+
+from fusion_bench import (
+    BaseAlgorithm,
+    BaseModelPool,
+    RuntimeConstants,
+    auto_register_config,
+    get_rankzero_logger,
+    instantiate,
+)
+from fusion_bench.dataset import CLIPDataset
+from fusion_bench.modelpool import ResNetForImageClassificationPool
+from fusion_bench.tasks.clip_classification import get_num_classes
+
+log = get_rankzero_logger(__name__)
+
+
+@auto_register_config
+class ImageClassificationFineTuning(BaseAlgorithm):
+    def __init__(
+        self,
+        max_epochs: Optional[int],
+        max_steps: Optional[int],
+        label_smoothing: float,
+        optimizer: DictConfig,
+        lr_scheduler: DictConfig,
+        dataloader_kwargs: DictConfig,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        assert (max_epochs is None) or (
+            max_steps is None or max_steps < 0
+        ), "Only one of max_epochs or max_steps should be set."
+        self.training_interval = "epoch" if max_epochs is not None else "step"
+        if self.training_interval == "epoch":
+            self.max_steps = -1
+        log.info(f"Training interval: {self.training_interval}")
+        log.info(f"Max epochs: {max_epochs}, max steps: {max_steps}")
+
+    def run(self, modelpool: ResNetForImageClassificationPool):
+        # load model and dataset
+        model = modelpool.load_pretrained_or_first_model()
+        assert isinstance(model, nn.Module), "Loaded model is not a nn.Module."
+
+        assert (
+            len(modelpool.train_dataset_names) == 1
+        ), "Exactly one training dataset is required."
+        self.dataset_name = dataset_name = modelpool.train_dataset_names[0]
+        num_classes = get_num_classes(dataset_name)
+        train_dataset = modelpool.load_train_dataset(dataset_name)
+        train_dataset = CLIPDataset(
+            train_dataset, processor=modelpool.load_processor(stage="train")
+        )
+        train_loader = self.get_dataloader(train_dataset, stage="train")
+        if modelpool.has_val_dataset:
+            val_dataset = modelpool.load_val_dataset(dataset_name)
+            val_dataset = CLIPDataset(
+                val_dataset, processor=modelpool.load_processor(stage="val")
+            )
+            val_loader = self.get_dataloader(val_dataset, stage="val")
+
+        # configure optimizer
+        optimizer = instantiate(self.optimizer, params=model.parameters())
+        if self.lr_scheduler is not None:
+            lr_scheduler = instantiate(self.lr_scheduler, optimizer=optimizer)
+            optimizer = {
+                "optimizer": optimizer,
+                "lr_scheduler": {
+                    "scheduler": lr_scheduler,
+                    "interval": self.training_interval,
+                    "frequency": 1,
+                },
+            }
+        log.info(f"optimizer:\n{optimizer}")
+
+        lit_module = ERM_LitModule(
+            model,
+            optimizer,
+            objective=nn.CrossEntropyLoss(label_smoothing=self.label_smoothing),
+            metrics={
+                "acc@1": Accuracy(task="multiclass", num_classes=num_classes),
+                "acc@5": Accuracy(task="multiclass", num_classes=num_classes, top_k=5),
+            },
+        )
+
+        log_dir = (
+            self._program.path.log_dir
+            if self._program is not None
+            else "outputs/lightning_logs"
+        )
+        trainer = L.Trainer(
+            max_epochs=self.max_epochs,
+            max_steps=self.max_steps,
+            accelerator="auto",
+            devices="auto",
+            callbacks=[
+                pl_callbacks.LearningRateMonitor(logging_interval="step"),
+                pl_callbacks.DeviceStatsMonitor(),
+            ],
+            logger=TensorBoardLogger(
+                save_dir=log_dir,
+                name="",
+            ),
+            fast_dev_run=RuntimeConstants.debug,
+        )
+
+        trainer.fit(
+            lit_module, train_dataloaders=train_loader, val_dataloaders=val_loader
+        )
+        model = lit_module.model
+        if rank_zero_only.rank == 0:
+            log.info(f"Saving the final model to {log_dir}/raw_checkpoints/final")
+            modelpool.save_model(
+                model,
+                path=os.path.join(
+                    trainer.log_dir if trainer.log_dir is not None else log_dir,
+                    "raw_checkpoints",
+                    "final",
+                ),
+            )
+        return model
+
+    def get_dataloader(self, dataset, stage: str):
+        assert stage in ["train", "val", "test"], f"Invalid stage: {stage}"
+        dataloader_kwargs = dict(self.dataloader_kwargs)
+        if "shuffle" not in dataloader_kwargs:
+            dataloader_kwargs["shuffle"] = stage == "train"
+        return DataLoader(dataset, **dataloader_kwargs)
+
+
+@auto_register_config
+class ImageClassificationFineTuning_Test(BaseAlgorithm):
+    def __init__(self, checkpoint_path: str, dataloader_kwargs: DictConfig, **kwargs):
+        super().__init__(**kwargs)
+
+    def run(self, modelpool: BaseModelPool):
+        assert (
+            modelpool.has_val_dataset or modelpool.has_test_dataset
+        ), "No validation or test dataset found in the model pool."
+
+        # load model and dataset
+        model = modelpool.load_pretrained_or_first_model()
+        assert isinstance(model, nn.Module), "Loaded model is not a nn.Module."
+
+        if modelpool.has_test_dataset:
+            assert (
+                len(modelpool.test_dataset_names) == 1
+            ), "Exactly one test dataset is required."
+            self.dataset_name = dataset_name = modelpool.test_dataset_names[0]
+            dataset = modelpool.load_test_dataset(dataset_name)
+            dataset = CLIPDataset(
+                dataset, processor=modelpool.load_processor(stage="test")
+            )
+        else:
+            assert (
+                len(modelpool.val_dataset_names) == 1
+            ), "Exactly one validation dataset is required."
+            self.dataset_name = dataset_name = modelpool.val_dataset_names[0]
+            dataset = modelpool.load_val_dataset(dataset_name)
+            dataset = CLIPDataset(
+                dataset, processor=modelpool.load_processor(stage="test")
+            )
+        num_classes = get_num_classes(dataset_name)
+
+        test_loader = self.get_dataloader(dataset, stage="test")
+
+        if self.checkpoint_path is None:
+            lit_module = ERM_LitModule(
+                model,
+                metrics={
+                    "acc@1": Accuracy(task="multiclass", num_classes=num_classes),
+                    "acc@5": Accuracy(
+                        task="multiclass", num_classes=num_classes, top_k=5
+                    ),
+                },
+            )
+        else:
+            lit_module = ERM_LitModule.load_from_checkpoint(
+                checkpoint_path=self.checkpoint_path,
+                model=model,
+                metrics={
+                    "acc@1": Accuracy(task="multiclass", num_classes=num_classes),
+                    "acc@5": Accuracy(
+                        task="multiclass", num_classes=num_classes, top_k=5
+                    ),
+                },
+            )
+
+        trainer = L.Trainer(
+            devices=1, num_nodes=1, logger=False, fast_dev_run=RuntimeConstants.debug
+        )
+
+        test_metrics = trainer.test(lit_module, dataloaders=test_loader)
+        log.info(f"Test metrics: {test_metrics}")
+        return model
+
+    def get_dataloader(self, dataset, stage: str):
+        assert stage in ["train", "val", "test"], f"Invalid stage: {stage}"
+        dataloader_kwargs = dict(self.dataloader_kwargs)
+        if "shuffle" not in dataloader_kwargs:
+            dataloader_kwargs["shuffle"] = stage == "train"
+        return DataLoader(dataset, **dataloader_kwargs)

fusion_bench/method/ensemble.py

@@ -17,7 +17,21 @@ from fusion_bench.models.wrappers.ensemble import (
 log = logging.getLogger(__name__)
 
 
+@auto_register_config
 class SimpleEnsembleAlgorithm(BaseAlgorithm):
+    def __init__(
+        self,
+        device_map: Optional[Mapping[int, Union[str, torch.device]]] = None,
+        **kwargs,
+    ):
+        """
+        Initializes the SimpleEnsembleAlgorithm with an optional device map.
+
+        Args:
+            device_map (Optional[Mapping[int, Union[str, torch.device]]], optional): A mapping from model index to device. Defaults to None.
+        """
+        super().__init__(**kwargs)
+
     @torch.no_grad()
     def run(self, modelpool: BaseModelPool | List[nn.Module]) -> EnsembleModule:
         """
@@ -30,9 +44,10 @@ class SimpleEnsembleAlgorithm(BaseAlgorithm):
             EnsembleModule: The ensembled model.
         """
         log.info(f"Running ensemble algorithm with {len(modelpool)} models")
-
         models = [modelpool.load_model(m) for m in modelpool.model_names]
-        ensemble = EnsembleModule(models=models)
+
+        log.info("creating ensemble module")
+        ensemble = EnsembleModule(models=models, device_map=self.device_map)
         return ensemble
 
 

fusion_bench/method/linear/__init__.py

@@ -2,5 +2,9 @@
 from .expo import ExPOAlgorithm
 from .linear_interpolation import LinearInterpolationAlgorithm
 from .llama_expo import ExPOAlgorithmForLlama
-from .simple_average_for_llama import SimpleAverageForLlama
-from .task_arithmetic_for_llama import TaskArithmeticForLlama
+from .simple_average_for_causallm import SimpleAverageForCausalLM, SimpleAverageForLlama
+from .task_arithmetic_for_causallm import (
+    TaskArithmeticForCausalLM,
+    TaskArithmeticForLlama,
+)
+from .ties_merging_for_causallm import TiesMergingForCausalLM

fusion_bench/method/linear/{simple_average_for_llama.py → simple_average_for_causallm.py}

@@ -18,16 +18,16 @@ log = get_rankzero_logger(__name__)
 
 
 @auto_register_config
-class SimpleAverageForLlama(BaseAlgorithm):
+class SimpleAverageForCausalLM(BaseAlgorithm):
     R"""
     A simple averaging algorithm for LLama models. If `merge_backbone` is set to `True`, the backbone of the model will be averaged and the rest of the model will be loaded from the pre-trained model.
 
     Examples:
-        The following example demonstrates how to use the `SimpleAverageForLlama` algorithm to merge Mistral models.
+        The following example demonstrates how to use the `SimpleAverageForCausalLM` algorithm to merge Mistral models.
 
         ```bash
         fusion_bench \
-            method=linear/simple_average_for_llama \
+            method=linear/simple_average_for_causallm \
             method.model_save_path=outputs/simle_mixtral_exp_v4/simple_average \
             modelpool=CausalLMPool/simle_mixtral_exp_v4.yaml
         ```
@@ -35,7 +35,7 @@ class SimpleAverageForLlama(BaseAlgorithm):
 
     def __init__(
         self,
-        merge_backbone: bool,
+        merge_backbone: bool = False,
         model_save_path: Optional[str] = None,
         show_pbar: bool = False,
         **kwargs,
@@ -81,3 +81,7 @@ class SimpleAverageForLlama(BaseAlgorithm):
                 with open(os.path.join(self.model_save_path, "README.md"), "w") as f:
                     f.write(model_card_str)
         return model
+
+
+SimpleAverageForLlama = SimpleAverageForCausalLM
+"""Alias for SimpleAverageForCausalLM"""