Skip to content

MergeKit

steerability.algorithms.structural_control.wrappers.mergekit

args

control

MergeKit

Bases: StructuralControl

Wrapper for merging models via MergeKit https://github.com/arcee-ai/mergekit.

MergeKit combines multiple language models using various merge strategies like linear interpolation, SLERP, and TIES. This wrapper integrates MergeKit's functionality to enable structural control through model composition.

The process involves loading a merge configuration (from YAML or dict), executing the merge operation, and optionally loading the resulting merged model. Supports caching to avoid redundant operations.

Reference:

  • "Arcee's MergeKit: A Toolkit for Merging Large Language Models" Charles Goddard, Shamane Siriwardhana, Malikeh Ehghaghi, Luke Meyers, Vladimir Karpukhin, Brian Benedict, Mark McQuade, Jacob Solawetz https://aclanthology.org/2024.emnlp-industry.36
Source code in steerability/algorithms/structural_control/wrappers/mergekit/control.py 
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
class MergeKit(StructuralControl):
    """
    Wrapper for merging models via MergeKit [https://github.com/arcee-ai/mergekit](https://github.com/arcee-ai/mergekit).

    MergeKit combines multiple language models using various merge strategies like linear interpolation, SLERP, and
    TIES. This wrapper integrates MergeKit's functionality to enable structural control through model composition.

    The process involves loading a merge configuration (from YAML or dict), executing the merge operation, and
    optionally loading the resulting merged model. Supports caching to avoid redundant operations.

    Reference:

    - "Arcee's MergeKit: A Toolkit for Merging Large Language Models"
      Charles Goddard, Shamane Siriwardhana, Malikeh Ehghaghi, Luke Meyers, Vladimir Karpukhin, Brian Benedict,
      Mark McQuade, Jacob Solawetz
      [https://aclanthology.org/2024.emnlp-industry.36](https://aclanthology.org/2024.emnlp-industry.36)
    """

    Args = MergeKitArgs

    def artifact_capability(self) -> Capability:
        """Merging always leaves a full-weights checkpoint at `out_path`."""
        return Capability.SERVE_CHECKPOINT

    def export_artifact(self) -> CheckpointArtifact:
        """The merged checkpoint directory written (or reused) by `steer()`."""
        return CheckpointArtifact(path=str(self.args.out_path))

    def export_state(self) -> dict:
        """The merged checkpoint under the `"artifact"` key, for freezing."""
        return {"artifact": self.export_artifact()}

    def frozen_form(self, state: dict) -> tuple[str, dict]:
        """The merged checkpoint as a `load_checkpoint` entry."""
        return "structural_control/load_checkpoint", {"path": state["artifact"]}

    def fit_identity(self) -> dict:
        """The merge-relevant args projection: the merge configuration and the dtype."""
        return {
            "config_path": str(self.args.config_path) if self.args.config_path else None,
            "config_dict": self.args.config_dict,
            "dtype": self.args.dtype,
        }

    def steer(
            self,
            model: PreTrainedModel,
            tokenizer: PreTrainedTokenizerBase = None,
            **_
    ) -> PreTrainedModel:
        """Execute model merging via MergeKit and optionally return the merged model.

        Performs structural steering by merging multiple models according to a configuration file or dictionary.
        Supports caching to avoid redundant merge operations and can either return the merged model or the original
        model based on configuration.

        The method follows this logic:

        1. Load merge configuration from YAML file or dictionary
        2. Check if merged model already exists (skip if `force_remerge=False`)
        3. Execute merge if needed using MergeKit
        4. Optionally load and return the merged model

        Args:
            model (PreTrainedModel): The base model (potentially unused depending on the method).
            tokenizer (PreTrainedTokenizerBase, optional): Base tokenizer (currently unused).

        Returns:
            PreTrainedModel: Either the merged model (if `load_merged=True`) or the original model. When returning
            merged model, attempts to attach a new tokenizer if one was created during merging.

        Note:

        - If out_path exists and `force_remerge=False`, skips merging and loads cached result
        - Merged model saved to `out_path` directory with full weights and config
        - If `load_merged=False`, performs merge but returns original model
        """
        args: MergeKitArgs = self.args

        if args.config_path:
            with open(args.config_path, "r", encoding="utf-8") as config_file:
                config = mk_config.MergeConfiguration.model_validate(yaml.safe_load(config_file))
        else:
            config = mk_config.MergeConfiguration.model_validate(args.config_dict)

        # find merged weights
        out_path = Path(args.out_path)
        if out_path.exists() and not args.force_remerge:
            if args.load_merged:
                merged = AutoModelForCausalLM.from_pretrained(
                    pretrained_model_name_or_path=str(out_path),
                    device_map=args.device_map,
                    trust_remote_code=args.trust_remote_code,
                    dtype=getattr(torch, args.dtype)
                )
                return merged
            return model

        # merge
        # with FileLock(str(out_path) + ".lock"):
        mk_merge.run_merge(
            merge_config=config,
            out_path=str(out_path),
            options=mk_merge.MergeOptions(
                cuda=args.allow_cuda,
                trust_remote_code=args.trust_remote_code,
                **args.extra_merge_options,
            )
        )

        # load merged checkpoint (and check if merge returned new tokenizer)
        if args.load_merged:
            merged = AutoModelForCausalLM.from_pretrained(
                out_path,
                dtype=getattr(torch, args.dtype),
                device_map=args.device_map,
                trust_remote_code=args.trust_remote_code,
            )
            try:
                merged.tokenizer = AutoTokenizer.from_pretrained(
                    out_path,
                    trust_remote_code=args.trust_remote_code
                )
            except Exception:
                pass
            return merged

        return model
RUNTIME_KWARGS_SCHEMA = [] class-attribute instance-attribute
args = self.Args.validate(*args, **kwargs) instance-attribute
enabled = True class-attribute instance-attribute
supports_batching = True class-attribute instance-attribute
artifact_capability()

Merging always leaves a full-weights checkpoint at out_path.

Source code in steerability/algorithms/structural_control/wrappers/mergekit/control.py 
38
39
40
def artifact_capability(self) -> Capability:
    """Merging always leaves a full-weights checkpoint at `out_path`."""
    return Capability.SERVE_CHECKPOINT
cleanup()

Release resources allocated during steer().

Override in subclasses that allocate GPU memory or other resources during steering. Default no-op.

Source code in steerability/algorithms/core/base_control.py 
195
196
197
198
199
200
201
def cleanup(self) -> None:
    """Release resources allocated during `steer()`.

    Override in subclasses that allocate GPU memory or other resources during steering. Default
    no-op.
    """
    pass
clone_for_call(seed=None)

A configuration-preserving shallow clone for one generation call.

The clone shares steer-time artifacts (memories, steering vectors, attached tokenizers) with the original but has its own attribute namespace, so per-call attribute mutation on the clone never races another call using the original. When seed is given and the control defines reseed(seed), the clone's client-side RNG is re-seeded.

Parameters:

Name Type Description Default
seed int | None

Optional seed forwarded to the clone's reseed().

 None

Returns:

Type Description
BaseControl

The clone.
Source code in steerability/algorithms/core/base_control.py 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def clone_for_call(self, seed: int | None = None) -> "BaseControl":
    """A configuration-preserving shallow clone for one generation call.

    The clone shares steer-time artifacts (memories, steering vectors, attached tokenizers)
    with the original but has its own attribute namespace, so per-call attribute mutation on
    the clone never races another call using the original. When `seed` is given and the
    control defines `reseed(seed)`, the clone's client-side RNG is re-seeded.

    Args:
        seed: Optional seed forwarded to the clone's `reseed()`.

    Returns:
        The clone.
    """
    clone = copy.copy(self)
    if seed is not None:
        reseed = getattr(clone, "reseed", None)
        if callable(reseed):
            reseed(seed)
    return clone
export_artifact()

The merged checkpoint directory written (or reused) by steer().

Source code in steerability/algorithms/structural_control/wrappers/mergekit/control.py 
42
43
44
def export_artifact(self) -> CheckpointArtifact:
    """The merged checkpoint directory written (or reused) by `steer()`."""
    return CheckpointArtifact(path=str(self.args.out_path))
export_state()

The merged checkpoint under the "artifact" key, for freezing.

Source code in steerability/algorithms/structural_control/wrappers/mergekit/control.py 
46
47
48
def export_state(self) -> dict:
    """The merged checkpoint under the `"artifact"` key, for freezing."""
    return {"artifact": self.export_artifact()}
fit_identity()

The merge-relevant args projection: the merge configuration and the dtype.

Source code in steerability/algorithms/structural_control/wrappers/mergekit/control.py 
54
55
56
57
58
59
60
def fit_identity(self) -> dict:
    """The merge-relevant args projection: the merge configuration and the dtype."""
    return {
        "config_path": str(self.args.config_path) if self.args.config_path else None,
        "config_dict": self.args.config_dict,
        "dtype": self.args.dtype,
    }
frozen_form(state)

The merged checkpoint as a load_checkpoint entry.

Source code in steerability/algorithms/structural_control/wrappers/mergekit/control.py 
50
51
52
def frozen_form(self, state: dict) -> tuple[str, dict]:
    """The merged checkpoint as a `load_checkpoint` entry."""
    return "structural_control/load_checkpoint", {"path": state["artifact"]}
requirements()

Backend requirements computed from this instance's configuration, per phase.

The generate phase requires Capability.IN_PROCESS_TORCH for in-process adoption of the returned model; when the configuration produces an on-disk artifact (artifact_capability()), serving that artifact is an alternative, so a backend advertising the matching serve capability also supports the generate phase.

Returns:

Type Description
Requirements

The control's phase-keyed requirements.
Source code in steerability/algorithms/structural_control/base.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def requirements(self) -> Requirements:
    """Backend requirements computed from this instance's configuration, per phase.

    The generate phase requires `Capability.IN_PROCESS_TORCH` for in-process adoption of
    the returned model; when the configuration produces an on-disk artifact
    (`artifact_capability()`), serving that artifact is an alternative, so a backend
    advertising the matching serve capability also supports the generate phase.

    Returns:
        The control's phase-keyed requirements.
    """
    generate = needs(Capability.IN_PROCESS_TORCH)
    capability = self.artifact_capability()
    if capability is not None:
        generate = any_of(
            generate,
            needs(capability, hint="serve the steer-time artifact on a vLLM backend"),
        )
    return Requirements(generate=generate)
steer(model, tokenizer=None, **_)

Execute model merging via MergeKit and optionally return the merged model.

Performs structural steering by merging multiple models according to a configuration file or dictionary. Supports caching to avoid redundant merge operations and can either return the merged model or the original model based on configuration.

The method follows this logic:

  1. Load merge configuration from YAML file or dictionary
  2. Check if merged model already exists (skip if force_remerge=False)
  3. Execute merge if needed using MergeKit
  4. Optionally load and return the merged model

Parameters:

Name Type Description Default
model PreTrainedModel

The base model (potentially unused depending on the method).

 required
tokenizer PreTrainedTokenizerBase

Base tokenizer (currently unused).

 None

Returns:

Name Type Description
PreTrainedModel PreTrainedModel

Either the merged model (if load_merged=True) or the original model. When returning
PreTrainedModel

merged model, attempts to attach a new tokenizer if one was created during merging.

Note:

  • If out_path exists and force_remerge=False, skips merging and loads cached result
  • Merged model saved to out_path directory with full weights and config
  • If load_merged=False, performs merge but returns original model
Source code in steerability/algorithms/structural_control/wrappers/mergekit/control.py 
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
def steer(
        self,
        model: PreTrainedModel,
        tokenizer: PreTrainedTokenizerBase = None,
        **_
) -> PreTrainedModel:
    """Execute model merging via MergeKit and optionally return the merged model.

    Performs structural steering by merging multiple models according to a configuration file or dictionary.
    Supports caching to avoid redundant merge operations and can either return the merged model or the original
    model based on configuration.

    The method follows this logic:

    1. Load merge configuration from YAML file or dictionary
    2. Check if merged model already exists (skip if `force_remerge=False`)
    3. Execute merge if needed using MergeKit
    4. Optionally load and return the merged model

    Args:
        model (PreTrainedModel): The base model (potentially unused depending on the method).
        tokenizer (PreTrainedTokenizerBase, optional): Base tokenizer (currently unused).

    Returns:
        PreTrainedModel: Either the merged model (if `load_merged=True`) or the original model. When returning
        merged model, attempts to attach a new tokenizer if one was created during merging.

    Note:

    - If out_path exists and `force_remerge=False`, skips merging and loads cached result
    - Merged model saved to `out_path` directory with full weights and config
    - If `load_merged=False`, performs merge but returns original model
    """
    args: MergeKitArgs = self.args

    if args.config_path:
        with open(args.config_path, "r", encoding="utf-8") as config_file:
            config = mk_config.MergeConfiguration.model_validate(yaml.safe_load(config_file))
    else:
        config = mk_config.MergeConfiguration.model_validate(args.config_dict)

    # find merged weights
    out_path = Path(args.out_path)
    if out_path.exists() and not args.force_remerge:
        if args.load_merged:
            merged = AutoModelForCausalLM.from_pretrained(
                pretrained_model_name_or_path=str(out_path),
                device_map=args.device_map,
                trust_remote_code=args.trust_remote_code,
                dtype=getattr(torch, args.dtype)
            )
            return merged
        return model

    # merge
    # with FileLock(str(out_path) + ".lock"):
    mk_merge.run_merge(
        merge_config=config,
        out_path=str(out_path),
        options=mk_merge.MergeOptions(
            cuda=args.allow_cuda,
            trust_remote_code=args.trust_remote_code,
            **args.extra_merge_options,
        )
    )

    # load merged checkpoint (and check if merge returned new tokenizer)
    if args.load_merged:
        merged = AutoModelForCausalLM.from_pretrained(
            out_path,
            dtype=getattr(torch, args.dtype),
            device_map=args.device_map,
            trust_remote_code=args.trust_remote_code,
        )
        try:
            merged.tokenizer = AutoTokenizer.from_pretrained(
                out_path,
                trust_remote_code=args.trust_remote_code
            )
        except Exception:
            pass
        return merged

    return model
steer_access()

ModelAccess.MODULE; training happens on live weights.

Source code in steerability/algorithms/structural_control/base.py 
112
113
114
def steer_access(self) -> ModelAccess:
    """`ModelAccess.MODULE`; training happens on live weights."""
    return ModelAccess.MODULE
steer_fits()

The fit artifacts this instance's steer step will produce, for the steer plan.

Each entry is (artifact, artifact_class), where artifact is the fit source or recipe class name and artifact_class is "direction" or "calibrated". The default is an empty tuple (no fits).

Returns:

Type Description
tuple[tuple[str, str], ...]

The declared fit artifacts, in declaration order.
Source code in steerability/algorithms/core/base_control.py 
105
106
107
108
109
110
111
112
113
114
115
def steer_fits(self) -> tuple[tuple[str, str], ...]:
    """The fit artifacts this instance's steer step will produce, for the steer plan.

    Each entry is `(artifact, artifact_class)`, where `artifact` is the fit source or
    recipe class name and `artifact_class` is `"direction"` or `"calibrated"`. The default
    is an empty tuple (no fits).

    Returns:
        The declared fit artifacts, in declaration order.
    """
    return ()