FusionBench

https://tanganke.github.io/fusion_bench
Admin
FusionBench is a library and toolkit for deep model fusion that merges, ensembles, or fuses multiple...
Tokens:579,709
Snippets:3,675
Trust Score:9.7
Update:3 months ago
Show doc for...
Tokens:
Raw
### Run First Fusion Experiment with CLI

Source: https://tanganke.github.io/fusion_bench/get_started/basic_examples

This is the quickest way to get started with FusionBench by running a simple average fusion experiment using the command-line interface. Ensure you have FusionBench installed and its prerequisites met.

```bash
fusion_bench \
    method=simple_average \
    modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TA8 \
    taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TA8

```

--------------------------------

### Complete Configuration Example

Source: https://tanganke.github.io/fusion_bench/get_started/basic_examples/structured_configs

This is a full configuration example combining multiple components for a CLIP task arithmetic.

```yaml
dir: ${oc.env:FUSION_BENCH_PROJECT_ROOT,"."}/outputs/multirun/${hydra.job.config_name}/${now:%Y-%m-%d_%H-%M-%S}
subdir: ${hydra.job.num}
job:
  env_set:
    HYDRA_FULL_ERROR: ${oc.env:HYDRA_FULL_ERROR,1}
output_subdir: ""

_target_: fusion_bench.programs.FabricModelFusionProgram
_recursive_: false
method:
  _target_: fusion_bench.method.TaskArithmeticAlgorithm
  scaling_factor: 0.7
modelpool:
  _target_: fusion_bench.modelpool.CLIPVisionModelPool
  models:
    _pretrained_: openai/clip-vit-base-patch32
    sun397: tanganke/clip-vit-base-patch32_sun397
    stanford-cars: tanganke/clip-vit-base-patch32_stanford-cars
taskpool:
  _target_: fusion_bench.taskpool.CLIPVisionModelTaskPool
  test_datasets:
    sun397:
      _target_: datasets.load_dataset
      path: tanganke/sun397
      split: test
    stanford-cars:
      _target_: datasets.load_dataset
      path: tanganke/stanford_cars
      split: test
  clip_model: openai/clip-vit-base-patch32
  processor: openai/clip-vit-base-patch32

```

--------------------------------

### Install Optuna

Source: https://tanganke.github.io/fusion_bench/get_started/advanced_examples/hyperparameter_optimization_with_optuna

Install Optuna and the Optuna dashboard for hyperparameter optimization and visualization.

```bash
pip install optuna
pip install optuna-dashboard

```

--------------------------------

### FusionBench CLI Example: Run with Specific Configuration

Source: https://tanganke.github.io/fusion_bench/get_started

Example of running FusionBench using a named configuration file.

```bash
# Run with a specific configuration
fusion_bench --config-name custom_config
```

--------------------------------

### Setup Model, Data, and Optimizer

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/training

Orchestrates the setup process by calling methods to configure the model, data, and optimizer, then integrates them with the fabric environment.

```python
fabric = self.fabric

        self.setup_model()
        self.setup_data()

        optimizer = self.configure_optimizer()
        optimizer, lr_scheduler = optimizer["optimizer"], optimizer["lr_scheduler"]

        self.model, self.optimizer = fabric.setup(self.model, optimizer)
        self.lr_scheduler = lr_scheduler
```

--------------------------------

### Install SwanLab

Source: https://tanganke.github.io/fusion_bench/get_started/intermediate_examples/select_logger

Install the SwanLab library to enable its logging capabilities.

```bash
pip install swanlab

```

--------------------------------

### FusionBench CLI Example: Use Custom Config Path

Source: https://tanganke.github.io/fusion_bench/get_started

Example of running FusionBench with a custom configuration file path and name.

```bash
# Use a custom config path
fusion_bench --config-path ./my_configs --config-name custom_fusion
```

--------------------------------

### Fusion-Bench CLI Usage Example

Source: https://tanganke.github.io/fusion_bench/guides/custom_modelpool

Example command to run Fusion-Bench using a custom ModelPool and a dummy taskpool.

```bash
fusion_bench \
  method=simple_average \
  modelpool=BertClassifierPool/glue_tasks \
  taskpool=dummy

```

--------------------------------

### Main Setup Method

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/training

Orchestrates the entire setup process for training. It calls methods to set up the model, data, and optimizer, then integrates them with the fabric instance. This is the primary entry point for initializing the training environment.

```python
    def setup(self):
        fabric = self.fabric

        self.setup_model()
        self.setup_data()

        optimizer = self.configure_optimizer()
        optimizer, lr_scheduler = optimizer["optimizer"], optimizer["lr_scheduler"]

        self.model, self.optimizer = fabric.setup(self.model, optimizer)
        self.lr_scheduler = lr_scheduler
```

--------------------------------

### Install FusionBench from PyPI

Source: https://tanganke.github.io/fusion_bench

Install the fusion-bench library and toolkit from PyPI. This is the standard method for using FusionBench as a dependency in other projects.

```bash
pip install fusion-bench

# you can also install a specific version
# pip install fusion-bench==0.1.6

```

--------------------------------

### Using an Algorithm Example

Source: https://tanganke.github.io/fusion_bench/api

This example demonstrates how to instantiate and use a `SimpleAverageAlgorithm` from FusionBench. It shows the process of creating an algorithm instance with specified parameters and then running the fusion process with a model pool.

```APIDOC
## Using an Algorithm

### Description
This section provides an example of how to use a pre-defined algorithm, `SimpleAverageAlgorithm`, for model fusion.

### Usage
1. Instantiate the algorithm with necessary parameters like `weight_key` and `scale_factor`.
2. Call the `run` method with a `modelpool` object to perform the fusion.

### Example
```python
from fusion_bench import SimpleAverageAlgorithm

# Create algorithm instance
algorithm = SimpleAverageAlgorithm(
    weight_key="task_weight",
    scale_factor=1.0
)

# Run fusion
merged_model = algorithm.run(modelpool)
```
```

--------------------------------

### Install and Login to Wandb

Source: https://tanganke.github.io/fusion_bench/get_started/intermediate_examples/select_logger

Install the Weights & Biases library and log in to your account to enable logging.

```bash
pip install wandb
wandb login

```

--------------------------------

### MagMax Configuration Example

Source: https://tanganke.github.io/fusion_bench/algorithms/magmax

This is an example of a configuration file for the MagMax merging method in FusionBench. It specifies the method to be used.

```yaml
# =============================================================================
# FusionBench Method Configuration: MagMax

```

--------------------------------

### FusionBench CLI Example: Override Specific Parameters

Source: https://tanganke.github.io/fusion_bench/get_started

Example of running FusionBench and overriding specific configuration parameters like `method.alpha` on the fly.

```bash
# Override specific parameters
fusion_bench --config-name custom_config method.alpha=0.5
```

--------------------------------

### Install FusionBench from GitHub

Source: https://tanganke.github.io/fusion_bench

Install the latest version of fusion-bench from the GitHub repository in editable mode. This method is useful for development or when needing the absolute latest changes.

```bash
git clone https://github.com/tanganke/fusion_bench.git
cd fusion_bench

# checkout to use a specific version. for example, v0.1.6
# git checkout v0.1.6

pip install -e . # install the package in editable mode

```

--------------------------------

### Install Documentation Dependencies

Source: https://tanganke.github.io/fusion_bench/guides/docs

Install the necessary Python dependencies for documentation development.

```bash
pip install -e ".[docs]"
```

--------------------------------

### Example: Merging Eight CLIP-ViT-L/14 Models with Custom Batch Settings

Source: https://tanganke.github.io/fusion_bench/algorithms/fisher_merging

Command-line example for merging eight CLIP-ViT-L/14 models with custom dataloader batch size and number of workers.

```bash
fusion_bench \
  method=fisher_merging/clip_fisher_merging \
    method.dataloader_kwargs.batch_size=8 \
    method.dataloader_kwargs.num_workers=4 \
  modelpool=CLIPVisionModelPool/clip-vit-large-patch14_TA8 \
  taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TA8 \
    taskpool.clip_model=openai/clip-vit-large-patch14

```

--------------------------------

### Using Model Pool Example

Source: https://tanganke.github.io/fusion_bench/api

This example illustrates how to load models using the `CLIPVisionModelPool`. It shows the process of initializing the model pool from a configuration file and then loading the associated models.

```APIDOC
## Using Model Pool

### Description
This section demonstrates how to load models using a specific model pool implementation, `CLIPVisionModelPool`.

### Usage
1. Initialize the `CLIPVisionModelPool` by providing a path to its configuration file.
2. Call the `load_models` method to retrieve the models managed by the pool.

### Example
```python
from fusion_bench import CLIPVisionModelPool

# Load models from pool
modelpool = CLIPVisionModelPool.from_config("config/modelpool/CLIPVisionModelPool/clip-vit-base-patch32_TA8.yaml")
models = modelpool.load_models()
```
```

--------------------------------

### DAWE CLI Usage Example

Source: https://tanganke.github.io/fusion_bench/algorithms/dawe

Example of how to run DAWE using the command-line interface, specifying merge mode, training steps, learning rate, and model/task pools.

```bash
fusion_bench \
    method=dawe/dawe_for_clip \
    method.merge_mode=task_wise \
    method.max_steps=1000 \
    method.learning_rate=1e-5 \
    modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TA8 \
    taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TA8

```

--------------------------------

### Model Stock Configuration Example

Source: https://tanganke.github.io/fusion_bench/algorithms/model_stock

Example configuration for the Model Stock algorithm, specifying the target class, keys to ignore, and model save path.

```yaml
_target_: fusion_bench.method.model_stock.ModelStock
ignore_keys:
  [
    "model.positional_embedding",
    "model.text_projection",
    "model.logit_scale",
    "model.token_embedding.weight",
    "model.ln_final.weight",
    "model.ln_final.bias",
  ]
model_save_path: ${path.log_dir}/checkpoint
model_save_kwargs: null

```

--------------------------------

### Full Training Setup

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/training

Orchestrates the complete setup for training by calling methods to set up the model, data, optimizer, and learning rate scheduler. It then uses the fabric utility to set up the model and optimizer.

```python
    def setup(self):
        fabric = self.fabric

        self.setup_model()
        self.setup_data()

        optimizer = self.configure_optimizer()
        optimizer, lr_scheduler = optimizer["optimizer"], optimizer["lr_scheduler"]

        self.model = self.fabric.setup_module(self.model)
        self.optimizer = self.fabric.setup_optimizers(optimizer)
        self.lr_scheduler = lr_scheduler
```

--------------------------------

### Setup and Configuration for Optuna Optimization

Source: https://tanganke.github.io/fusion_bench/get_started/advanced_examples/hyperparameter_optimization_with_optuna

Initializes Lightning Fabric and loads configuration using Hydra for model fusion tasks. This setup is required before defining the optimization objective.

```python
import os
import lightning as L
import optuna
from hydra import compose, initialize

from fusion_bench import instantiate
from fusion_bench.method import TaskArithmeticAlgorithm
from fusion_bench.modelpool import CLIPVisionModelPool
from fusion_bench.scripts.cli import get_default_config_path
from fusion_bench.taskpool import CLIPVisionModelTaskPool

# Initialize Lightning Fabric for efficient computation
fabric = L.Fabric(accelerator="auto", devices=1)

# Load configuration using Hydra
with initialize(
    version_base=None,
    config_path=os.path.relpath(
        get_default_config_path(), start=os.path.dirname(__file__)
    ),
):
    cfg = compose(
        config_name="fabric_model_fusion",
        overrides=[
            "modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TA8",
            "taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TA8",
        ],
    )
    modelpool: CLIPVisionModelPool = instantiate(cfg.modelpool)
    taskpool: CLIPVisionModelTaskPool = instantiate(cfg.taskpool)
    taskpool._fabric_instance = fabric

```

--------------------------------

### Example: Move Module to CUDA

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.utils/torch

Demonstrates moving a PyTorch nn.Module to a CUDA device.

```python
>>> model = torch.nn.Linear(2, 2)
>>> to_device(model, torch.device('cuda'))
```

--------------------------------

### RegMeanAlgorithmForCLIPPlusPlus Get Regression Mean Weights Hook Setup

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/merging

Sets up a hook function to compute regression mean weights for each linear module's input. It initializes dictionaries to store weights, computed example counts, and actual example counts.

```python
def get_regmean_weights(
    self,
    model_name: str,
    layer: Module,
    batches_input: List[Tensor],
    linear_modules_to_merge: Dict[str, Module],
):
    layer = self.fabric.setup(layer)

    def compute_regmean_weights(module_name: str):
        """
        compute the regmean weights, a hook function to deal with each module's input
        :param module_name: str, module name
        :return:
        """

        def hook(module: nn.Module, input: tuple, output: torch.Tensor):
            # Tensor, shape (batch_size, sequence_length, hidden_dim)
            x = cast(Tensor, input[0]).detach()
            batch_num_actual_examples = x.shape[0]
            # Tensor, shape (batch_size * sequence_length, hidden_dim)
            x = x.reshape(-1, x.shape[-1])
            # Tensor, shape (hidden_dim, hidden_dim)
            xtx = torch.matmul(x.transpose(0, 1), x)
            # store the averaged weights in regmean_weights
            if module_name not in regmean_weights.keys():
                regmean_weights[module_name] = xtx / x.shape[0]
                num_computed_examples[module_name] = x.shape[0]
                num_actual_examples[module_name] = batch_num_actual_examples
            else:
                regmean_weights[module_name] = (
                    regmean_weights[module_name]
                    * num_computed_examples[module_name]
                    + xtx
                ) / (num_computed_examples[module_name] + x.shape[0])
                num_computed_examples[module_name] += x.shape[0]
                num_actual_examples[module_name] += batch_num_actual_examples

        return hook

    handles = []
    # dictionary, regmean matrices for each linear module inputs
    regmean_weights = {}
    # dictionary, number of examples (multiplied the sequence length) used for computing regmean matrices
    num_computed_examples = {}
    # dictionary, number of actual examples used for computing regmean matrices
    num_actual_examples = {}
```

--------------------------------

### on_fisher_merging_start

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/merging

Setup the zero-shot classification head before starting the Fisher merging process.

```APIDOC
## on_fisher_merging_start

### Description
Setup the zero-shot classification head before starting the Fisher merging process.

### Method
on_fisher_merging_start

### Parameters
None

### Request Example
None

### Response
#### Success Response (200)
None

#### Response Example
None
```

--------------------------------

### Setup for Test-Time Adaptation

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/merging

Prepares the necessary components before starting test-time adaptation, such as loading task-specific scores.

```python
    def on_test_time_adaptation_start(self):
        """
        Something to do before the test-time adaptation starts. Such as setting up the task-specific heads.
        """
        self.scores = {}
        for model_name in self.modelpool.model_names:
            score = cast(
                GPT2ForSequenceClassification,
                self.modelpool.load_classifier(model_name),
            ).score.requires_grad_(False)
            score = score.to(self.fabric.device)
            self.scores[model_name] = score
```

--------------------------------

### On Test Time Adaptation Start Hook

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/ensemble

A placeholder method to be called before test-time adaptation begins, intended for setup tasks like initializing task-specific heads.

```python
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
```

--------------------------------

### Setup Zero-Shot Classification Head for Fisher Merging

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/merging

Initializes the zero-shot classification head, which is a prerequisite for starting the Fisher merging process. This method should be called before any Fisher merging operations.

```python
def on_fisher_merging_start(self):
    """
    Setup the zero-shot classification head before starting the Fisher merging process.
    """
    self.setup_zero_shot_classification_head()
```

--------------------------------

### get_current_device

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.utils/torch

Gets the current available device for PyTorch operations, prioritizing XPU, NPU, MPS, CUDA, and falling back to CPU. It respects the LOCAL_RANK environment variable for distributed setups.

```APIDOC
## get_current_device()

### Description
Gets the current available device for PyTorch operations. This is used for distributed training.

This function checks the availability of various types of devices in the following order: 1. XPU (Intel's AI accelerator) 2. NPU (Neural Processing Unit) 3. MPS (Metal Performance Shaders, for Apple devices) 4. CUDA (NVIDIA's GPU) 5. CPU (Central Processing Unit, used as a fallback)

The function returns the first available device found in the above order. If none of the specialized devices are available, it defaults to the CPU.

### Returns
* `device` (torch.device): The current available device for PyTorch operations.

### Environment Variables
* `LOCAL_RANK`: This environment variable is used to specify the device index for multi-device setups. If not set, it defaults to "0".

### Example
```python
>>> device = get_current_device()
>>> print(device)
# xpu:0  # or npu:0, mps:0, cuda:0, cpu depending on availability
```

### Method
```python
def get_current_device() -> torch.device:
    R"""
    Gets the current available device for PyTorch operations.
    This is used for distributed training.

    This function checks the availability of various types of devices in the following order:
    1. XPU (Intel's AI accelerator)
    2. NPU (Neural Processing Unit)
    3. MPS (Metal Performance Shaders, for Apple devices)
    4. CUDA (NVIDIA's GPU)
    5. CPU (Central Processing Unit, used as a fallback)

    The function returns the first available device found in the above order. If none of the specialized devices
    are available, it defaults to the CPU.

    Returns:
        torch.device: The current available device for PyTorch operations.

    Environment Variables:
        LOCAL_RANK: This environment variable is used to specify the device index for multi-device setups.
                    If not set, it defaults to "0".

    Example:

        >>> device = get_current_device()
        >>> print(device)
        xpu:0  # or npu:0, mps:0, cuda:0, cpu depending on availability
    """

    if is_torch_xpu_available():
        device = "xpu:{}".format(os.environ.get("LOCAL_RANK", "0"))
    elif is_torch_npu_available():
        device = "npu:{}".format(os.environ.get("LOCAL_RANK", "0"))
    elif is_torch_mps_available():
        device = "mps:{}".format(os.environ.get("LOCAL_RANK", "0"))
    elif is_torch_cuda_available():
        device = "cuda:{}".format(os.environ.get("LOCAL_RANK", "0"))
    else:
        device = "cpu"

    return torch.device(device)
```
```

--------------------------------

### Get Current PyTorch Device

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.utils/torch

Determines the best available device (XPU, NPU, MPS, CUDA, CPU) for PyTorch operations, respecting the LOCAL_RANK environment variable for distributed setups.

```python
import os
import torch

def is_torch_xpu_available():
    # Placeholder for actual check
    return False

def is_torch_npu_available():
    # Placeholder for actual check
    return False

def is_torch_mps_available():
    # Placeholder for actual check
    return False

def is_torch_cuda_available():
    return torch.cuda.is_available()

def get_current_device() -> torch.device:
    R"""
    Gets the current available device for PyTorch operations.
    This is used for distributed training.

    This function checks the availability of various types of devices in the following order:
    1. XPU (Intel's AI accelerator)
    2. NPU (Neural Processing Unit)
    3. MPS (Metal Performance Shaders, for Apple devices)
    4. CUDA (NVIDIA's GPU)
    5. CPU (Central Processing Unit, used as a fallback)

    The function returns the first available device found in the above order. If none of the specialized devices
    are available, it defaults to the CPU.

    Returns:
        torch.device: The current available device for PyTorch operations.

    Environment Variables:
        LOCAL_RANK: This environment variable is used to specify the device index for multi-device setups.
                    If not set, it defaults to "0".

    Example:

        >>> device = get_current_device()
        >>> print(device)
        xpu:0  # or npu:0, mps:0, cuda:0, cpu depending on availability
    """

    if is_torch_xpu_available():
        device = "xpu:{}".format(os.environ.get("LOCAL_RANK", "0"))
    elif is_torch_npu_available():
        device = "npu:{}".format(os.environ.get("LOCAL_RANK", "0"))
    elif is_torch_mps_available():
        device = "mps:{}".format(os.environ.get("LOCAL_RANK", "0"))
    elif is_torch_cuda_available():
        device = "cuda:{}".format(os.environ.get("LOCAL_RANK", "0"))
    else:
        device = "cpu"

    return torch.device(device)
```

--------------------------------

### Instantiate TaskPool from Configuration

Source: https://tanganke.github.io/fusion_bench/taskpool

Demonstrates how to create a taskpool instance by loading a configuration file using OmegaConf and instantiating it with fusion_bench.utils.instantiate.

```python
# Create from configuration file
from fusion_bench.utils import instantiate
from omegaconf import OmegaConf

config = OmegaConf.load("path/to/taskpool/config.yaml")
taskpool = instantiate(config)

```

--------------------------------

### Example: Merging GPT-2 Models for Text Classification

Source: https://tanganke.github.io/fusion_bench/algorithms/fisher_merging

Command-line example for merging GPT-2 models for text classification, specifying custom Fisher example count, batch size, and workers.

```bash
fusion_bench \
  method=fisher_merging/gpt2_fisher_merging \
    method.num_fisher_examples=512 \
    method.batch_size=8 \
    method.num_workers=2 \
  modelpool=gpt-2_glue \
  taskpool=gpt-2_glue

```

--------------------------------

### setup_model

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/training

Sets up the tokenizer and the model for training. It loads the tokenizer, sets the padding token, and loads the pre-trained model, configuring its padding token if necessary.

```APIDOC
## setup_model

### Description
This method is responsible for initializing and configuring the tokenizer and the model. It loads the tokenizer from the `modelpool`, ensures a padding token is set (using the end-of-sequence token if necessary), and then loads the pre-trained model, also ensuring its padding token is correctly configured.

### Method Signature
`setup_model(self)`

### Internal State Updates
- **self.tokenizer**: Loaded tokenizer from `modelpool`.
- **self.model**: Loaded pre-trained model (e.g., `LlamaForSequenceClassification`).

### Side Effects
- Sets `self.tokenizer.pad_token_id` if it is `None`.
- Sets `model.config.pad_token_id` if it is `None`.
```

--------------------------------

### FusionBench Model and Tokenizer Setup

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/training

Sets up the model and tokenizer for training. It loads the tokenizer and ensures a pad token is set, then loads the pre-trained model.

```python
def setup_model(self):
        # https://github.com/Lightning-AI/litgpt/blob/main/litgpt/finetune/lora.py
        self.tokenizer = self.modelpool.load_tokenizer()
        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id

        model = self.modelpool.load_pretrained_model()
```

--------------------------------

### __init__

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/mixing

Initializes the SparseWeightEnsemblingMoEAlgorithm with the provided configuration. It sets up a profiler based on the cache directory specified in the configuration.

```APIDOC
## __init__

### Description
Initialize the SparseWeightEnsemblingMoEAlgorithm with the given configuration.

### Parameters
#### Path Parameters
- **`algorithm_config`** (DictConfig) - Required - The configuration for the algorithm.
```

--------------------------------

### Get Number of Models

Source: https://tanganke.github.io/fusion_bench/modelpool

Get the total count of models available in the ModelPool.

```python
# Get number of models
num_models = len(modelpool)
```

--------------------------------

### Fusion-Bench Model Setup

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/training

Initializes the model, tokenizer, and sets up training parameters. Handles tokenizer padding token configuration and optionally freezes token embeddings.

```python
def setup_model(self):
        self.tokenizer = self.modelpool.load_tokenizer()
        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id

        model = self.modelpool.load_pretrained_model()
        self.model: "LlamaForCausalLM" = model

        if self.fix_token_embedding:
            self.model.model.embed_tokens.requires_grad_(False)
```

--------------------------------

### Install Pyinstrument

Source: https://tanganke.github.io/fusion_bench/guides/fusion_bench/mixins/pyinstrument_profiling

Install the pyinstrument library using pip. This is a prerequisite for using the PyinstrumentProfilerMixin.

```bash
pip install pyinstrument

```

--------------------------------

### setup_model

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/training

Sets up the model, optimizer, and learning rate scheduler. Initializes the CLIP model, applies LoRA if specified, and configures the optimizer and learning rate scheduler.

```APIDOC
## setup_model()

### Description
Sets up the model, optimizer, and learning rate scheduler. This method initializes the CLIP model, applies LoRA if specified, and configures the optimizer and learning rate scheduler.

### Returns
- **Tuple** - A tuple containing the processor, classifier, optimizer, and learning rate scheduler.
```

--------------------------------

### setup_model

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/training

Sets up the model, optimizer, and learning rate scheduler. This method initializes the CLIP model, applies LoRA if specified, and configures the optimizer and learning rate scheduler.

```APIDOC
## setup_model()

### Description
Sets up the model, optimizer, and learning rate scheduler. This method initializes the CLIP model, applies LoRA if specified, and configures the optimizer and learning rate scheduler.

### Returns
* **Tuple** - A tuple containing the processor, classifier, optimizer, and learning rate scheduler.
```

--------------------------------

### PostDefenseSAUAlgorithmForCLIP Model Setup

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/merging

Initializes models for the PostDefenseSAUAlgorithmForCLIP, loading pretrained, merge, and reference merge models.

```python
class PostDefenseSAUAlgorithmForCLIP(
    CLIPClassificationMixin,
    SimpleProfilerMixin,
    ModelFusionAlgorithm,
):
    @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
        pretrained_model = modelpool.load_model("_pretrained_")
        merge_model = modelpool.load_model("merge")
        merge_model_ref = modelpool.load_model("merge")

```

--------------------------------

### Single Model Configuration Example

Source: https://tanganke.github.io/fusion_bench/modelpool/llm

Example YAML configuration for a single Llama model within CausalLMPool.

```yaml
_target_: fusion_bench.modelpool.CausalLMPool
_recursive_: false

```

--------------------------------

### Install LM Evaluation Harness

Source: https://tanganke.github.io/fusion_bench/taskpool/lm_eval_harness_cli

Installs the LM Evaluation Harness library with optional dependencies for enhanced functionality.

```bash
pip install -e '.[lm-eval-harness]'
```

--------------------------------

### Setup Model, Optimizer, and Scheduler

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/training

Initializes the CLIP model, optionally applies LoRA, and configures the optimizer and learning rate scheduler. Prints model summary for the vision model if running on the global zero process.

```python
def setup_model(self):
    """
    Sets up the model, optimizer, and learning rate scheduler.

    This method initializes the CLIP model, applies LoRA if specified, and configures the optimizer and learning rate scheduler.

    Returns:
        Tuple: A tuple containing the processor, classifier, optimizer, and learning rate scheduler.
    """
    config = self.config
    modelpool = self.modelpool

    clip_model: CLIPModel = modelpool.load_clip_model("_pretrained_")
    processor = modelpool.load_processor()

    self.finetune_method = "full fine-tune"
    if self.use_lora:
        self.finetune_method = "lora fine-tune"
        lora_config = LoraConfig(
            **OmegaConf.to_container(
                self.lora_config, resolve=True, enum_to_str=True
            )
        )
        clip_model.vision_model = get_peft_model(
            clip_model.vision_model, lora_config
        )

    classifier = HFCLIPClassifier(clip_model, processor=processor)

    if self.fabric.is_global_zero:
        print("=== Model Summary (For Vision Model Only) ===")
        print_parameters(classifier.clip_model.vision_model)
    # configure optimizers
    optimizer = torch.optim.Adam(
        [
            p
            for p in classifier.clip_model.vision_model.parameters()
            if p.requires_grad
        ],
        lr=self.learning_rate,
        weight_decay=self.weight_decay,
    )
    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer=optimizer, T_max=self.num_steps
    )

    return processor, classifier, optimizer, lr_scheduler
```

--------------------------------

### ImageClassificationFineTuningForCLIP.setup_model

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/training

Sets up the CLIP model, including applying LoRA if configured, and initializes the optimizer and learning rate scheduler.

```APIDOC
## ImageClassificationFineTuningForCLIP.setup_model

### Description
Sets up the model, optimizer, and learning rate scheduler. This method initializes the CLIP model, applies LoRA if specified, and configures the optimizer and learning rate scheduler.

### Method Signature
`setup_model(self)`

### Returns
- **Tuple**: A tuple containing the processor, classifier, optimizer, and learning rate scheduler.

### Usage Example
```python
# Assuming 'trainer' is an instance of ImageClassificationFineTuningForCLIP
processor, classifier, optimizer, lr_scheduler = trainer.setup_model()
```
```

--------------------------------

### Python: Get TensorBoard SummaryWriter

Source: https://tanganke.github.io/fusion_bench/guides/fusion_bench/mixins/lightning_fabric

Get the TensorBoard SummaryWriter for detailed logging. Raises an AttributeError if the logger is not a TensorBoardLogger.

```python
SummaryWriter(SummaryWriter) – 
The TensorBoard SummaryWriter instance.
```

--------------------------------

### Inspecting Configuration (Formatted)

Source: https://tanganke.github.io/fusion_bench/cli/fusion_bench

Use 'print_config=true' and 'dry_run=true' for a beautifully formatted output with syntax highlighting.

```bash
fusion_bench print_config=true dry_run=true
```

--------------------------------

### TaskVectorCosSimilarity Example

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/utility

Example of how to instantiate and run the TaskVectorCosSimilarity algorithm. This algorithm computes cosine similarity between task vectors of models.

```python
>>> algorithm = TaskVectorCosSimilarity(
...     plot_heatmap=True,
...     trainable_only=True,
...     output_path="/path/to/outputs"
... )
>>> result = algorithm.run(modelpool)
```

--------------------------------

### Setup Lightning Fabric

Source: https://tanganke.github.io/fusion_bench/guides/fusion_bench/mixins/lightning_fabric

Initializes and launches a PyTorch Lightning Fabric instance with optional TensorBoard logging. Requires a configuration dictionary specifying fabric and logger settings.

```python
def setup_lightning_fabric(self, config: DictConfig):
    """
    Initializes and launches the Lightning Fabric with optional logging.

    This method sets up the Lightning Fabric for distributed computing based on the provided configuration. If a fabric
    configuration is not found, it logs a warning and exits. Optionally, if a fabric logger configuration is provided,
    it initializes a TensorBoardLogger with the specified settings.

    Expected configuration keys:
    - fabric: The configuration for the Lightning Fabric.
    - fabric.loggers: The configuration for the TensorBoardLogger.
    """
    if self._fabric_instance is None:
        if config.get("fabric", None) is None:
            log.warning("No fabric configuration found. use default settings. By default, use 1 device.")
            self._fabric_instance = L.Fabric(devices=1)
        else:
            self._fabric_instance = instantiate(config.fabric)
        if not _is_using_cli():  # if not using cli, launch the fabric
            self._fabric_instance.launch()
        # Set the log directory in config if it is not already set
        if (
            self.log_dir is not None
            and hasattr(config, "log_dir")
            and config.get("log_dir", None) is None
        ):
            if self._fabric_instance.is_global_zero:
                log.info(f"Setting log_dir to {self.log_dir}")
            config.log_dir = self.log_dir
```

--------------------------------

### Example Test with Specific Checkpoint Path

Source: https://tanganke.github.io/fusion_bench/guides/resnet/image_classification_finetune

An example of testing a fine-tuned ResNet model with a concrete checkpoint path provided.

```bash
fusion_bench --config-name model_fusion \
    method=classification/image_classification_finetune_test \
    method.checkpoint_path="outputs/resnet18/cifar10/version_0/checkpoints/epoch\=9-step\=1960.ckpt" \
    modelpool=ResNetForImageClassfication/transformers/resnet18_cifar10

```

--------------------------------

### Example: Merging Eight CLIP-ViT-B/32 Models

Source: https://tanganke.github.io/fusion_bench/algorithms/fisher_merging

Command-line example for merging eight CLIP-ViT-B/32 models using the fisher_merging/clip_fisher_merging method.

```bash
fusion_bench method=fisher_merging/clip_fisher_merging \
  modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TA8 \
  taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TA8

```

--------------------------------

### setup_lightning_fabric

Source: https://tanganke.github.io/fusion_bench/guides/fusion_bench/mixins/lightning_fabric

Initializes and launches the Lightning Fabric with optional logging. This method configures the fabric for distributed computing based on a provided configuration dictionary, optionally setting up a TensorBoardLogger.

```APIDOC
## setup_lightning_fabric

### Description
Initializes and launches the Lightning Fabric with optional logging. This method sets up the Lightning Fabric for distributed computing based on the provided configuration. If a fabric configuration is not found, it logs a warning and exits. Optionally, if a fabric logger configuration is provided, it initializes a TensorBoardLogger with the specified settings.

### Method Signature
`setup_lightning_fabric(config)`

### Parameters
#### Parameters
- **`config`** (`DictConfig`) - Required - The configuration dictionary for setting up the Lightning Fabric. Expected keys include 'fabric' for fabric settings and 'fabric.loggers' for logger configuration.
```

--------------------------------

### Setup CLIP Model for Fine-tuning

Source: https://tanganke.github.io/fusion_bench/guides/clip_vit/finetune

Sets up the CLIP model, optimizer, and learning rate scheduler. Supports full fine-tuning, LoRA, and L-LoRA by applying PEFT configurations.

```python
def setup_model(self):
        """
        Sets up the model, optimizer, and learning rate scheduler.

        This method initializes the CLIP model, applies LoRA if specified, and configures the optimizer and learning rate scheduler.

        Returns:
            Tuple: A tuple containing the processor, classifier, optimizer, and learning rate scheduler.
        """
        config = self.config
        modelpool = self.modelpool

        clip_model: CLIPModel = modelpool.load_clip_model("_pretrained_")
        processor = modelpool.load_processor()

        self.finetune_method = "full fine-tune"
        if config.use_lora or config.use_l_lora:
            self.finetune_method = "lora fine-tune"
            lora_config = LoraConfig(
                **OmegaConf.to_container(
                    config.lora_config, resolve=True, enum_to_str=True
                )
            )
            clip_model.vision_model = get_peft_model(
                clip_model.vision_model, lora_config
            )

            if config.use_l_lora:
                # http://arxiv.org/abs/2310.04742
                # Anke Tang et al. Parameter Efficient Multi-task Model Fusion with Partial Linearization. ICLR 2024.
                self.finetune_method = "l-lora fine-tune"
                print("Linearizing Lora Layers")
                linearize_lora_model_(clip_model.vision_model)

        classifier = HFCLIPClassifier(clip_model, processor=processor)

        if self.fabric.is_global_zero:
            print("=== Model Summary (For Vision Model Only) ===")
            print_parameters(classifier.clip_model.vision_model)
        # configure optimizers
        optimizer = torch.optim.Adam(
            [
                p
                for p in classifier.clip_model.vision_model.parameters()
                if p.requires_grad
            ],
            lr=config.learning_rate,
            weight_decay=config.weight_decay,
        )
        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer=optimizer, T_max=config.num_steps
        )

        return processor, classifier, optimizer, lr_scheduler
```

--------------------------------

### Fish Shell Completion Installation

Source: https://tanganke.github.io/fusion_bench/cli/fusion_bench

Install Fish shell completion by piping the output of `fusion_bench -sc install=fish` to the source command.

```bash
fusion_bench -sc install=fish | source
```

--------------------------------

### Running FusionBench with Depth Upscaling

Source: https://tanganke.github.io/fusion_bench/algorithms/depth_upscaling

This command demonstrates how to execute the `fusion_bench` tool with the depth_upscaling method. Ensure the configuration file is correctly specified when running.

```bash
fusion_bench method=depth_upscaling ...

```

--------------------------------

### Bash Shell Completion Installation

Source: https://tanganke.github.io/fusion_bench/cli/fusion_bench

Install Bash shell completion for tab-completing arguments by evaluating the output of `fusion_bench -sc install=bash`.

```bash
eval "$(fusion_bench -sc install=bash)"
```

--------------------------------

### Basic Python Function Example

Source: https://tanganke.github.io/fusion_bench/guides/docs

Example of a Python function using code fences with language specification, as per Markdown standards.

```python
def example_function():
    return "Hello, World!"
```

--------------------------------

### __init__

Source: https://tanganke.github.io/fusion_bench/api/fusion_bench.method/mixing

Initializes the RankOneMoEAlgorithm with the provided configuration. It sets up the Fabric environment if CUDA is available and initializes a profiler.

```APIDOC
## __init__

### Description
Initialize the RankOneMoEAlgorithm with the given configuration.

### Parameters
#### Path Parameters
- **`algorithm_config`** (DictConfig) - Required - The configuration for the algorithm.

### Source Code
```python
def __init__(self, algorithm_config: DictConfig):
    """
    Initialize the RankOneMoEAlgorithm with the given configuration.

    Args:
        algorithm_config (DictConfig): The configuration for the algorithm.
    """
    super().__init__(algorithm_config)

    if self._fabric is None and torch.cuda.is_available():
        self._fabric = L.Fabric(
            devices=self.config.get("devices", 1),
        )
        self._fabric.launch()
    else:
        assert "No CUDA device available."
    self.profiler = SimpleProfiler(
        self.config.get("cache_dir", "outputs"), "we_moe_profiler.txt"
    )
```
```
Tokens:
Raw