(models: list[LazyModel])
| 865 | |
| 866 | |
| 867 | def merge_sharded(models: list[LazyModel]) -> LazyModel: |
| 868 | # Original LLaMA models have each file contain one part of each tensor. |
| 869 | # Use a dict instead of a set to preserve order. |
| 870 | names = {name: None for model in models for name in model} |
| 871 | |
| 872 | def convert(name: str) -> LazyTensor: |
| 873 | lazy_tensors = [model[name] for model in models] |
| 874 | if len(lazy_tensors) == 1: |
| 875 | # only one file; don't go through this procedure since there might |
| 876 | # be quantized tensors |
| 877 | return lazy_tensors[0] |
| 878 | if len(lazy_tensors[0].shape) == 1: |
| 879 | # the tensor is just duplicated in every file |
| 880 | return lazy_tensors[0] |
| 881 | if name.startswith('tok_embeddings.') or \ |
| 882 | name.endswith('.attention.wo.weight') or \ |
| 883 | name.endswith('.feed_forward.w2.weight'): |
| 884 | # split by columns |
| 885 | axis = 1 |
| 886 | else: |
| 887 | # split by rows |
| 888 | axis = 0 |
| 889 | concatenated_shape = list(lazy_tensors[0].shape) |
| 890 | concatenated_shape[axis] = sum(tensor.shape[axis] for tensor in lazy_tensors) |
| 891 | |
| 892 | def load() -> UnquantizedTensor: |
| 893 | ndarrays = [load_unquantized(tensor) for tensor in lazy_tensors] |
| 894 | concatenated = np.concatenate(ndarrays, axis=axis) |
| 895 | return UnquantizedTensor(concatenated) |
| 896 | description = 'concatenated[[' + '] | ['.join(lt.description for lt in lazy_tensors) + ']]' |
| 897 | return LazyTensor(load, concatenated_shape, lazy_tensors[0].data_type, description) |
| 898 | return {name: convert(name) for name in names} |
| 899 | |
| 900 | |
| 901 | def merge_multifile_models(models_plus: list[ModelPlus]) -> ModelPlus: |
no test coverage detected