fusion-bench 0.2.9__py3-none-any.whl

fusion_bench/method/linear/expo.py

@@ -0,0 +1,118 @@
+"""
+This module contains the implementation of ExPO merge for general nn.Modules.
+
+Reference:
+- Zheng et al. Weak-to-Strong Extrapolation Expedites Alignment.
+"""
+
+import logging
+from copy import deepcopy
+
+import torch
+from torch import nn
+
+from fusion_bench import BaseAlgorithm, BaseModelPool
+from fusion_bench.method import SimpleAverageAlgorithm
+from fusion_bench.utils.state_dict_arithmetic import (
+    state_dict_add,
+    state_dict_mul,
+    state_dict_sub,
+)
+
+log = logging.getLogger(__name__)
+
+
+def expo_merge(
+    sft_model: nn.Module,
+    rlhf_model: nn.Module,
+    extrapolation_factor: float,
+    inplace: bool = True,
+    enable_grad: bool = False,
+):
+    """
+    Minimal implementation of ExPO merge.
+
+    Args:
+        sft_model (nn.Module): The pretrained model (base model).
+        rlhf_model (nn.Module): The finetuned model (medium-aligned model).
+        extrapolation_factor (float): The extrapolation factor.
+        inplace (bool): Whether to perform the merge in-place. If not, the rlhf_model will be copied before merging.
+        enable_grad (bool): Whether to enable gradient computation during the merge.
+
+    Returns:
+        nn.Module: The merged model.
+    """
+
+    if not inplace:
+        rlhf_model = deepcopy(rlhf_model)
+
+    with torch.set_grad_enabled(enable_grad):
+        for (sft_name, sft_param), (rlhf_name, rlhf_param) in zip(
+            sft_model.named_parameters(), rlhf_model.named_parameters()
+        ):
+            assert sft_name == rlhf_name, f"Model mismatch: {sft_name} != {rlhf_name}"
+            rlhf_param.data = rlhf_param.data + extrapolation_factor * (
+                rlhf_param.data - sft_param.data
+            )
+    return rlhf_model
+
+
+class ExPOAlgorithm(BaseAlgorithm):
+    R"""
+    ExPO merge algorithm.
+
+    This algorithm merges a pretrained model with a finetuned model.
+
+    $$\theta_{merged} = \theta_{sft} + \alpha (\theta_{rlhf} - \theta_{sft})$$
+
+    where $\theta_{merged}$ is the merged model, $\theta_{rlhf}$ is the finetuned model (medium-aligned model),
+    $\theta_{sft}$ is the pretrained model (base model), and $\alpha$ is the extrapolation factor.
+
+    In the configuration, the SFT model should have name `_pretrained_` and the rlhf name can be set arbitarily.
+    """
+
+    _config_mapping = BaseAlgorithm._config_mapping | {
+        "extrapolation_factor": "extrapolation_factor"
+    }
+
+    def __init__(self, extrapolation_factor: float, **kwargs):
+        self.extrapolation_factor = extrapolation_factor
+        super().__init__(**kwargs)
+
+    def run(self, modelpool: BaseModelPool):
+        """
+        Run the ExPO merge algorithm.
+
+        Args:
+            modelpool (BaseModelPool): The pool of models to merge.
+
+        Returns:
+            nn.Module: The merged model.
+        """
+        if not isinstance(modelpool, BaseModelPool):
+            modelpool = BaseModelPool(modelpool)
+
+        assert len(modelpool.model_names) >= 1, "ExPO requires at least one model."
+        assert modelpool.has_pretrained, "ExPO requires pretrained models (base model)."
+
+        sft_model = modelpool.load_pretrained_model()
+        if len(modelpool) == 1:
+            rlhf_model = modelpool.load_model(modelpool.model_names[0])
+        else:
+            # if there are multiple RLHF models, use simple average to merge them before running ExPO
+            log.info(
+                f"There are {len(modelpool)} models in the model pool, averaging them first..."
+            )
+            rlhf_model = SimpleAverageAlgorithm().run(modelpool)
+
+        # merge the pretrained model and the finetuned model
+        delta_parameters = state_dict_sub(
+            rlhf_model.state_dict(), sft_model.state_dict()
+        )
+        merged_sd = state_dict_add(
+            rlhf_model.state_dict(),
+            state_dict_mul(delta_parameters, scalar=self.extrapolation_factor),
+        )
+
+        rlhf_model.load_state_dict(merged_sd)
+        return rlhf_model

fusion_bench/method/linear/linear_interpolation.py

@@ -0,0 +1,60 @@
+import logging
+
+import torch
+
+from fusion_bench import BaseAlgorithm, BaseModelPool
+from fusion_bench.utils.state_dict_arithmetic import state_dict_weighted_sum
+
+log = logging.getLogger(__name__)
+
+
+class LinearInterpolationAlgorithm(BaseAlgorithm):
+    R"""
+    LinearInterpolationAlgorithm performs linear interpolation between two models.
+    Returns a model with the state dict that is a linear interpolation of the state dicts of the two models.
+    $\theta = (1-t) \theta_1 + t \theta_2$
+    """
+
+    _config_mapping = BaseAlgorithm._config_mapping | {
+        "t": "t",
+    }
+
+    def __init__(self, t: float, **kwargs):
+        """
+        Initialize the LinearInterpolationAlgorithm with the given interpolation parameter.
+
+        Args:
+            t (float): The interpolation parameter, should be in the range [0, 1].
+            **kwargs: Additional keyword arguments.
+        """
+        assert 0 <= t <= 1, "t should be in the range [0, 1]"
+        self.t = t
+        super().__init__(**kwargs)
+
+    def run(self, modelpool: BaseModelPool):
+        """
+        Run the linear interpolation algorithm on the given model pool.
+
+        This method performs linear interpolation between two models in the model pool
+        and returns a model with the interpolated state dict.
+
+        Args:
+            modelpool (BaseModelPool): The pool of models to interpolate. Must contain exactly two models.
+
+        Returns:
+            nn.Module: The model with the interpolated state dict.
+        """
+        assert (
+            modelpool.all_model_names == 2
+        ), "linear interpolation expect exactly 2 models"
+        primary_model = modelpool.load_model(modelpool.all_model_names[0])
+        secondary_model = modelpool.load_model(modelpool.all_model_names[1])
+
+        with torch.no_grad():
+            primary_state_dict = primary_model.state_dict()
+            secondary_state_dict = secondary_model.state_dict()
+            state_dict = state_dict_weighted_sum(
+                [primary_state_dict, secondary_state_dict], [1 - self.t, self.t]
+            )
+        primary_model.load_state_dict(state_dict)
+        return primary_model

fusion_bench/method/linear/llama_expo.py

@@ -0,0 +1,229 @@
+"""
+This module contains an extention implementation of ExPO merge for LLAMA models by @tanganke.
+
+Reference:
+- Zheng et al. Weak-to-Strong Extrapolation Expedites Alignment.
+"""
+
+import logging
+from typing import Optional, cast
+
+import torch
+from torch import nn
+from transformers import LlamaForCausalLM, LlamaModel
+from typing_extensions import override
+
+from fusion_bench import BaseAlgorithm, BaseModelPool
+from fusion_bench.method import DareSimpleAverage, SimpleAverageAlgorithm
+from fusion_bench.method.pruning.prune_utils import unstructured_magnitude_prune_
+from fusion_bench.utils.state_dict_arithmetic import StateDictType
+
+log = logging.getLogger(__name__)
+
+
+def expo_(
+    sft_model: nn.Module,
+    rlhf_model: nn.Module,
+    extrapolation_factor: float,
+    merge_dtype: Optional[torch.dtype] = None,
+    magnitude_sparsity_ratio: Optional[float] = None,
+):
+    """
+    Applies extrapolation to the parameters of the RLHF model based on the SFT model.
+    The RLHF model is updated in place.
+
+    Args:
+        sft_model (nn.Module): The supervised fine-tuned model.
+        rlhf_model (nn.Module): The reinforcement learning from human feedback model.
+        extrapolation_factor (float): The factor by which to extrapolate the parameters.
+
+    Returns:
+        nn.Module: The RLHF model with updated parameters.
+    """
+    rlhf_state_dict: StateDictType = rlhf_model.state_dict()
+    sft_state_dict: StateDictType = sft_model.state_dict()
+
+    merged_state_dict = {}
+
+    for n in rlhf_state_dict:
+        rlhf_p = rlhf_state_dict[n]
+        sft_p = sft_state_dict[n]
+        if merge_dtype is not None:
+            orignal_dtype = rlhf_state_dict[n].dtype
+            rlhf_p = rlhf_state_dict[n].to(dtype=merge_dtype)
+            sft_p = sft_state_dict[n].to(dtype=merge_dtype)
+
+        delta_p = rlhf_p - sft_p
+        if magnitude_sparsity_ratio is not None:
+            delta_p = unstructured_magnitude_prune_(
+                delta_p, torch.abs, magnitude_sparsity_ratio, return_pruned_weight=False
+            )
+        sft_p = rlhf_p + extrapolation_factor * delta_p
+
+        if merge_dtype is not None:
+            merged_state_dict[n] = rlhf_p.to(dtype=orignal_dtype)
+        else:
+            merged_state_dict[n] = rlhf_p
+
+    rlhf_model.load_state_dict(merged_state_dict)
+    return rlhf_model
+
+
+def expo_linear_modules_(
+    sft_model: nn.Module,
+    rlhf_model: nn.Module,
+    extrapolation_factor: float,
+    merge_dtype: Optional[torch.dtype] = None,
+    magnitude_sparsity_ratio: Optional[float] = None,
+):
+    """
+    Applies extrapolation to the linear modules of the RLHF model based on the SFT model.
+    The RLHF model is updated in place.
+
+    Args:
+        sft_model (nn.Module): The supervised fine-tuned model.
+        rlhf_model (nn.Module): The reinforcement learning from human feedback model.
+        extrapolation_factor (float): The factor by which to extrapolate the parameters.
+
+    Returns:
+        nn.Module: The RLHF model with updated linear modules.
+    """
+    for name, module in sft_model.named_modules():
+        if isinstance(module, nn.Linear):
+            expo_(
+                module,
+                rlhf_model.get_submodule(name),
+                extrapolation_factor=extrapolation_factor,
+                merge_dtype=merge_dtype,
+                magnitude_sparsity_ratio=magnitude_sparsity_ratio,
+            )
+    return rlhf_model
+
+
+class ExPOAlgorithmForLlama(BaseAlgorithm):
+
+    def __init__(
+        self,
+        extrapolation_factor: float,
+        attention_scaling_factor: float = 0.5,
+        only_on_backbone: bool = True,
+        on_linear_weights: bool = True,
+        on_linear_bias: bool = False,
+        on_embedding: bool = False,
+        fix_last_n_layers: int = 0,
+        fix_first_n_layers: int = 0,
+        magnitude_sparsity_ratio: Optional[float] = None,
+        **kwargs,
+    ):
+        self.extrapolation_factor = extrapolation_factor
+        self.attention_scaling_factor = attention_scaling_factor
+        self.only_on_backbone = only_on_backbone
+        self.on_linear_weights = on_linear_weights
+        self.on_linear_bias = on_linear_bias
+        self.on_embedding = on_embedding
+        self.fix_last_n_layers = fix_last_n_layers
+        self.fix_first_n_layers = fix_first_n_layers
+        self.magnitude_sparsity_ratio = magnitude_sparsity_ratio
+        super().__init__(**kwargs)
+
+    def load_models(self, modelpool: BaseModelPool):
+        sft_model: LlamaForCausalLM = modelpool.load_pretrained_model()
+        if len(modelpool) == 1:
+            rlhf_model = modelpool.load_model(modelpool.model_names[0])
+        else:
+            # if there are multiple RLHF models, use simple average to merge them before running ExPO
+            log.info(
+                f"There are {len(modelpool)} models in the model pool, averaging them first..."
+            )
+            rlhf_model = SimpleAverageAlgorithm().run(modelpool)
+        rlhf_model = cast(LlamaForCausalLM, rlhf_model)
+        return sft_model, rlhf_model
+
+    def run(self, modelpool: BaseModelPool):
+        if not isinstance(modelpool, BaseModelPool):
+            modelpool = BaseModelPool(modelpool)
+
+        assert len(modelpool.model_names) >= 1, "ExPO requires at least one model."
+        assert modelpool.has_pretrained, "ExPO requires pretrained models (base model)."
+
+        sft_model, rlhf_model = self.load_models(modelpool)
+
+        if not self.on_linear_bias:
+            for name, module in sft_model.named_modules():
+                if isinstance(module, nn.Linear):
+                    module.bias = rlhf_model.get_submodule(name).bias
+        if not self.on_linear_weights:
+            for name, module in sft_model.named_modules():
+                if isinstance(module, nn.Linear):
+                    module.weight = rlhf_model.get_submodule(name).weight
+
+        if not self.only_on_backbone:
+            expo_(sft_model.lm_head, rlhf_model.lm_head, self.extrapolation_factor)
+
+        # expo on the backbone
+        self._expo_lm_model_(
+            sft_model.model, rlhf_model.model, self.extrapolation_factor
+        )
+        return rlhf_model
+
+    def _expo_lm_model_(
+        self,
+        sft_model: LlamaModel,
+        rlhf_model: LlamaModel,
+        extrapolation_factor: float,
+    ):
+        if self.on_embedding:
+            expo_(sft_model.embed_tokens, rlhf_model.embed_tokens, extrapolation_factor)
+
+        if self.fix_first_n_layers == "half":
+            self.fix_first_n_layers = len(sft_model.layers) // 2
+        if self.fix_last_n_layers == "half":
+            self.fix_last_n_layers = len(sft_model.layers) // 2
+
+        for layer_idx in range(
+            self.fix_first_n_layers, len(sft_model.layers) - self.fix_last_n_layers
+        ):
+            sft_layer = sft_model.layers[layer_idx]
+            expo_linear_modules_(
+                sft_layer.self_attn,
+                rlhf_model.layers[layer_idx].self_attn,
+                extrapolation_factor=extrapolation_factor
+                * self.attention_scaling_factor,
+                merge_dtype=torch.float32,
+                magnitude_sparsity_ratio=self.magnitude_sparsity_ratio,
+            )
+            expo_linear_modules_(
+                sft_layer.mlp,
+                rlhf_model.layers[layer_idx].mlp,
+                extrapolation_factor=extrapolation_factor,
+                merge_dtype=torch.float32,
+                magnitude_sparsity_ratio=self.magnitude_sparsity_ratio,
+            )
+
+
+class ExPOWithDareForLLama(ExPOAlgorithmForLlama):
+    def __init__(
+        self,
+        dare_sparsity_ratio: float,
+        dare_only_on_linear_weights: bool,
+        dare_rescale: bool = True,
+        **kwargs,
+    ):
+        self.dare_sparsity_ratio = dare_sparsity_ratio
+        self.dare_only_on_linear_weights = dare_only_on_linear_weights
+        self.dare_rescale = dare_rescale
+        super().__init__(**kwargs)
+
+    @override
+    def load_models(self, modelpool: BaseModelPool):
+        log.info(
+            f"There are {len(modelpool)} models in the model pool, averaging them first..."
+        )
+        rlhf_model = DareSimpleAverage(
+            sparsity_ratio=self.dare_sparsity_ratio,
+            only_on_linear_weights=self.dare_only_on_linear_weights,
+            rescale=self.dare_rescale,
+        ).run(modelpool)
+        rlhf_model = cast(LlamaForCausalLM, rlhf_model)
+        sft_model: LlamaForCausalLM = modelpool.load_pretrained_model()
+        return sft_model, rlhf_model

fusion_bench/method/linear/simple_average_for_llama.py

@@ -0,0 +1,54 @@
+from typing import Optional
+
+from typing_extensions import override
+
+from fusion_bench import timeit_context
+from fusion_bench.method.base_algorithm import BaseAlgorithm
+from fusion_bench.method.simple_average import SimpleAverageAlgorithm
+from fusion_bench.modelpool import CausalLMBackbonePool, CausalLMPool
+
+
+class SimpleAverageForLlama(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.
+
+        ```bash
+        fusion_bench \
+            method=linear/simple_average_for_llama \
+            method.model_save_path=outputs/simle_mixtral_exp_v4/simple_average \
+            modelpool=CausalLMPool/simle_mixtral_exp_v4.yaml
+        ```
+    """
+
+    _config_mapping = BaseAlgorithm._config_mapping | {
+        "merge_backbone": "merge_backbone",
+    }
+
+    def __init__(self, merge_backbone: bool, model_save_path: Optional[str] = None):
+        super().__init__()
+        self.merge_backbone = merge_backbone
+        self.model_save_path = model_save_path
+
+    @override
+    def run(self, modelpool: CausalLMPool):
+        if self.model_save_path:
+            tokenizer = modelpool.load_tokenizer()
+
+        if self.merge_backbone:
+            assert modelpool.has_pretrained
+            backbone_modelpool = CausalLMBackbonePool(**modelpool.config)
+            model = modelpool.load_model("_pretrained_")
+            backbone_model = SimpleAverageAlgorithm().run(backbone_modelpool)
+            model.model.layers = backbone_model
+        else:
+            model = SimpleAverageAlgorithm().run()
+
+        if self.model_save_path is not None:
+            with timeit_context(f"Saving the model to {self.model_save_path}"):
+                tokenizer.save_pretrained(self.model_save_path)
+                model.save_pretrained(self.model_save_path)
+        return model

fusion_bench/method/linear/task_arithmetic_for_llama.py

@@ -0,0 +1,57 @@
+import logging
+from typing import Dict, List, Mapping, Optional, TypeVar, Union  # noqa: F401
+
+from typing_extensions import override
+
+from fusion_bench import timeit_context
+from fusion_bench.method import TaskArithmeticAlgorithm
+from fusion_bench.mixins.simple_profiler import SimpleProfilerMixin
+from fusion_bench.modelpool import CausalLMBackbonePool, CausalLMPool
+
+log = logging.getLogger(__name__)
+
+
+class TaskArithmeticForLlama(TaskArithmeticAlgorithm, SimpleProfilerMixin):
+    R"""
+    Examples:
+
+    fusion_bench \
+        method=linear/task_arithmetic_for_llama \
+            method.scaling_factor=0.3 \
+        method.model_save_path=outputs/simle_mixtral_exp_v4/task_arithmetic_0.3 \
+        modelpool=CausalLMPool/simle_mixtral_exp_v4.yaml
+    """
+
+    _config_mapping = TaskArithmeticAlgorithm._config_mapping | {
+        "merge_backbone": "merge_backbone",
+    }
+
+    def __init__(
+        self,
+        scaling_factor: float,
+        merge_backbone: bool,
+        model_save_path: Optional[str] = None,
+    ):
+        self.merge_backbone = merge_backbone
+        self.model_save_path = model_save_path
+        super().__init__(scaling_factor=scaling_factor)
+
+    @override
+    def run(self, modelpool: CausalLMPool):
+        if self.model_save_path:
+            tokenizer = modelpool.load_tokenizer()
+
+        if self.merge_backbone:
+            assert modelpool.has_pretrained
+            backbone_modelpool = CausalLMBackbonePool(**modelpool.config)
+            model = modelpool.load_model("_pretrained_")
+            backbone_model = super().run(backbone_modelpool)
+            model.model.layers = backbone_model
+        else:
+            model = super().run(modelpool)
+
+        if self.model_save_path is not None:
+            with timeit_context(f"Saving the model to {self.model_save_path}"):
+                tokenizer.save_pretrained(self.model_save_path)
+                model.save_pretrained(self.model_save_path)
+        return model

fusion_bench/method/lm_finetune/__init__.py

@@ -0,0 +1,3 @@
+from .bradley_terry_rm import BradleyTerryRewardModeling
+from .fullfinetune_sft import FullFinetuneSFT
+from .peftfinetune_sft import PeftFinetuneSFT