File size: 3,896 Bytes
656b04b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 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 | # Copyright (C) 2025 Arcee AI
# SPDX-License-Identifier: BUSL-1.1
import random
from typing import List
import click
import yaml
from mergekit.architecture import ArchitectureInfoUtils
from mergekit.common import ModelReference
from mergekit.config import (
InputSliceDefinition,
MergeConfiguration,
OutputSliceDefinition,
)
from mergekit.merge import run_merge
from mergekit.options import MergeOptions, PrettyPrintHelp, add_merge_options
@click.command("mergekit-layershuffle", cls=PrettyPrintHelp)
@click.argument("out_path", type=str)
@click.option("--model", "-m", multiple=True, type=str, help="Add a model to the merge")
@click.option(
"--weight",
"-w",
multiple=True,
type=float,
default=[],
show_default=False,
help="Weighting for a model",
)
@click.option(
"--print-yaml/--no-print-yaml",
is_flag=True,
help="Print YAML merge config for resulting model",
)
@click.option(
"--write-yaml",
type=click.Path(writable=True),
help="Path to write YAML merge config to",
)
@click.option(
"--dry-run", is_flag=True, help="Generate a config but do not run the merge"
)
@click.option("--fp16/--no-fp16", is_flag=True, help="Use FP16 precision")
@click.option(
"--full-random/--no-full-random",
is_flag=True,
help="Randomize layer index as well as source model",
)
@add_merge_options
def main(
out_path: str,
model: List[str],
weight: List[float],
print_yaml: bool,
write_yaml: bool,
dry_run: bool,
fp16: bool,
full_random: bool,
merge_options: MergeOptions,
):
models = [ModelReference.parse(m) for m in model]
m0_cfg = models[0].config()
arch_info = ArchitectureInfoUtils.get_architecture_info(m0_cfg)
total_num_layers = arch_info.num_layers(m0_cfg)
out_slices: List[OutputSliceDefinition] = []
if full_random:
for model, frac in zip(models, weight):
cfg = model.config()
num_layers = int(arch_info.num_layers(cfg) * frac)
for _ in range(num_layers):
src_idx = random.randrange(0, num_layers)
out_slices.append(
OutputSliceDefinition(
sources=[
InputSliceDefinition(
model=str(model),
layer_range=(src_idx, src_idx + 1),
)
]
)
)
random.shuffle(out_slices)
else:
for layer_idx in range(total_num_layers):
src_model = random.choices(models, weights=weight, k=1)[0]
if out_slices and out_slices[-1].sources[0].model == str(src_model):
out_slices[-1].sources[0].layer_range = (
out_slices[-1].sources[0].layer_range[0],
layer_idx + 1,
)
else:
out_slices.append(
OutputSliceDefinition(
sources=[
InputSliceDefinition(
model=str(src_model),
layer_range=(layer_idx, layer_idx + 1),
)
]
)
)
merge_config = MergeConfiguration(
merge_method="passthrough", slices=out_slices, dtype="float16" if fp16 else None
)
if print_yaml or write_yaml:
yaml_str = yaml.dump(merge_config.model_dump(exclude_none=True, mode="json"))
if print_yaml:
print(yaml_str)
if write_yaml:
with open(write_yaml, "w", encoding="utf-8") as file:
file.write(yaml_str)
if dry_run:
return
run_merge(
merge_config,
out_path,
options=merge_options,
)
if __name__ == "__main__":
main()
|