| 600 | |
| 601 | |
| 602 | def merge_multifile_models(models_plus: list[ModelPlus]) -> ModelPlus: |
| 603 | formats = set(mp.format for mp in models_plus) |
| 604 | # assert len(formats) == 1, "different formats?" |
| 605 | format = formats.pop() |
| 606 | paths = [path for mp in models_plus for path in mp.paths] |
| 607 | # Use the first non-None vocab, if any. |
| 608 | try: |
| 609 | vocab = next(mp.vocab for mp in models_plus if mp.vocab is not None) |
| 610 | except StopIteration: |
| 611 | vocab = None |
| 612 | |
| 613 | if any("model.embed_tokens.weight" in mp.model for mp in models_plus) or \ |
| 614 | any("model.layers.0.fc1.weight" in mp.model for mp in models_plus): |
| 615 | # Transformers models put different tensors in different files, but |
| 616 | # don't split indivdual tensors between files. |
| 617 | model: LazyModel = {} |
| 618 | for mp in models_plus: |
| 619 | model.update(mp.model) |
| 620 | else: |
| 621 | model = merge_sharded([mp.model for mp in models_plus]) |
| 622 | |
| 623 | return ModelPlus(model, paths, format, vocab) |
| 624 | |
| 625 | |
| 626 | def permute_lazy(lazy_tensor: LazyTensor, n_head: int, n_head_kv: int) -> LazyTensor: |