File size: 4,465 Bytes
a164e13 | 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 134 135 136 137 138 139 140 141 142 143 144 145 | # Copyright (C) 2024 Charles O. Goddard
#
# This software is free software: you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This software is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program. If not, see http://www.gnu.org/licenses/.
import random
from typing import List
import click
import yaml
from mergekit.architecture import get_architecture_info
from mergekit.common import ModelReference
from mergekit.config import (
InputSliceDefinition,
MergeConfiguration,
OutputSliceDefinition,
)
from mergekit.merge import run_merge
from mergekit.options import MergeOptions, add_merge_options
@click.command("mergekit-layershuffle")
@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 = 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()
|