poc ple split

common/arg.cpp

@@ -3062,6 +3062,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.model.path = value;
         }
     ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_EXPORT_LORA, LLAMA_EXAMPLE_DOWNLOAD, LLAMA_EXAMPLE_TOKENIZE}).set_env("LLAMA_ARG_MODEL"));
+    add_opt(common_arg(
+        {"--shard"}, "FNAME,...",
+        "comma-separated paths of the model splits following the one given by -m, in order\n"
+        "use it when the splits are not named <name>-%05d-of-%05d.gguf, e.g. to share a split between models",
+        [](common_params & params, const std::string & value) {
+            for (const auto & item : parse_csv_row(value)) {
+                params.model.shards.push_back(item);
+            }
+        }
+    ).set_env("LLAMA_ARG_SHARD"));
     add_opt(common_arg(
         {"-mu", "--model-url"}, "MODEL_URL",
         "model download url (default: unused)",

common/common.cpp

@@ -1276,13 +1276,13 @@ common_init_result::common_init_result(common_params & params, bool model_only)
         cparams_dft.n_rs_seq = 0;
 
         const common_fit_extra_model extra = {
-            /*.path_model   =*/ params_dft.model.path.c_str(),
+            /*.paths_model  =*/ params_dft.model.get_paths(),
             /*.mparams      =*/ &mparams_dft,
             /*.cparams      =*/ &cparams_dft,
             /*.shares_model =*/ !has_draft, // an MTP context runs on the weights of the main model
         };
 
-        common_fit_params(params.model.path.c_str(), &mparams, &cparams,
+        common_fit_params(params.model.get_paths(), &mparams, &cparams,
             params.tensor_split,
             params.tensor_buft_overrides.data(),
             params.fit_params_target.data(),
@@ -1291,7 +1291,7 @@ common_init_result::common_init_result(common_params & params, bool model_only)
             params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR);
     }
 
-    llama_model * model = llama_model_load_from_file(params.model.path.c_str(), mparams);
+    llama_model * model = common_model_load(params.model.get_paths(), mparams);
     if (model == NULL) {
         return;
     }
@@ -1723,6 +1723,19 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
     return mparams;
 }
 
+struct llama_model * common_model_load(const std::vector<std::string> & paths, const struct llama_model_params & mparams) {
+    if (paths.size() == 1) {
+        return llama_model_load_from_file(paths[0].c_str(), mparams);
+    }
+
+    std::vector<const char *> c_paths;
+    c_paths.reserve(paths.size());
+    for (const auto & path : paths) {
+        c_paths.push_back(path.c_str());
+    }
+    return llama_model_load_from_splits(c_paths.data(), c_paths.size(), mparams);
+}
+
 struct llama_context_params common_context_params_to_llama(const common_params & params) {
     auto cparams = llama_context_default_params();
 

common/common.h

@@ -311,6 +311,14 @@ struct common_params_model {
     std::string hf_file     = ""; // HF file
     std::string docker_repo = ""; // Docker repo
 
+    std::vector<std::string> shards; // shards following path, for splits not named <name>-%05d-of-%05d.gguf
+
+    std::vector<std::string> get_paths() const {
+        std::vector<std::string> paths = { path };
+        paths.insert(paths.end(), shards.begin(), shards.end());
+        return paths;
+    }
+
     std::string get_name() const {
         if (!hf_repo.empty()) {
             return hf_repo;
@@ -1017,6 +1025,9 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode
 struct llama_model_params   common_model_params_to_llama  (      common_params & params);
 struct llama_context_params common_context_params_to_llama(const common_params & params);
 
+// a single path also finds the other splits by file name
+struct llama_model * common_model_load(const std::vector<std::string> & paths, const struct llama_model_params & mparams);
+
 // clear LoRA adapters from context, then apply new list of adapters
 void common_set_adapter_lora(struct llama_context * ctx, std::vector<common_adapter_lora_info> & lora);
 

common/fit.cpp

@@ -1,5 +1,7 @@
 #include "fit.h"
 
+#include "common.h"
+
 #include "json.h"
 #include "log.h"
 
@@ -28,7 +30,7 @@ class common_params_fit_exception : public std::runtime_error {
 };
 
 static std::vector<llama_device_memory_data> common_get_device_memory_data_impl(
-        const char * path_model,
+        const std::vector<std::string> & paths_model,
         const llama_model_params * mparams,
         const llama_context_params * cparams,
         std::vector<ggml_backend_dev_t> & devs,
@@ -57,7 +59,7 @@ static std::vector<llama_device_memory_data> common_get_device_memory_data_impl(
     mparams_copy.no_alloc  = true;
     mparams_copy.load_mode = LLAMA_LOAD_MODE_NONE;
 
-    llama_model * model = llama_model_load_from_file(path_model, mparams_copy);
+    llama_model * model = common_model_load(paths_model, mparams_copy);
     if (model == nullptr) {
         llama_log_set(ud.original_logger.callback, ud.original_logger.user_data);
         throw std::runtime_error("failed to load model");
@@ -154,7 +156,7 @@ static std::vector<llama_device_memory_data> common_get_device_memory_data_impl(
 }
 
 common_device_memory_data_vec common_get_device_memory_data(
-        const char * path_model,
+        const std::vector<std::string> & paths_model,
         const llama_model_params * mparams,
         const llama_context_params * cparams,
         std::vector<ggml_backend_dev_t> & devs,
@@ -163,7 +165,7 @@ common_device_memory_data_vec common_get_device_memory_data(
         uint32_t & hp_n_expert,
         ggml_log_level log_level) {
     std::vector<llama_device_memory_data> impl = common_get_device_memory_data_impl(
-            path_model, mparams, cparams, devs, hp_ngl, hp_n_ctx_train, hp_n_expert, log_level);
+            paths_model, mparams, cparams, devs, hp_ngl, hp_n_ctx_train, hp_n_expert, log_level);
 
     common_device_memory_data_vec ret(impl.size());
     for (size_t i = 0; i < impl.size(); i++) {
@@ -177,7 +179,7 @@ common_device_memory_data_vec common_get_device_memory_data(
 }
 
 static void common_params_fit_impl(
-        const char * path_model, struct llama_model_params * mparams, struct llama_context_params * cparams,
+        const std::vector<std::string> & paths_model, struct llama_model_params * mparams, struct llama_context_params * cparams,
         float * tensor_split, struct llama_model_tensor_buft_override * tensor_buft_overrides,
         size_t * margins_s, uint32_t n_ctx_min, const common_fit_extra_model * extra, enum ggml_log_level log_level) {
     if (mparams->split_mode == LLAMA_SPLIT_MODE_TENSOR) {
@@ -221,7 +223,7 @@ static void common_params_fit_impl(
             dmds_t measured;
             try {
                 measured = common_get_device_memory_data_impl(
-                    extra->path_model, extra->mparams, extra->cparams, devs_extra, ngl_extra, nct_extra, nex_extra, log_level);
+                    extra->paths_model, extra->mparams, extra->cparams, devs_extra, ngl_extra, nct_extra, nex_extra, log_level);
             } catch (const std::runtime_error & e) {
                 // the extra model is optional, fit the main model alone rather than giving up
                 LOG_WRN("%s: failed to measure the memory of the extra model, fitting without it: %s\n", __func__, e.what());
@@ -261,7 +263,7 @@ static void common_params_fit_impl(
     // step 1: get data for default parameters and check whether any changes are necessary in the first place
 
     LOG_TRC("%s: getting device memory data for initial parameters:\n", __func__);
-    dmds_t dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+    dmds_t dmds_full = common_get_device_memory_data_impl(paths_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
 
     // saturate instead of overflowing, this also preserves the UINT32_MAX sentinel of n_ctx_min:
     const uint32_t n_ctx_max       = (uint32_t) std::min<uint64_t>(uint64_t(hp_nct)    * n_streams, UINT32_MAX);
@@ -273,7 +275,7 @@ static void common_params_fit_impl(
         if (n_streams > 1) {
             LOG_TRC("%s: context size unset and KV cache not unified -> using %" PRIu32 " for %" PRIu32 " sequences:\n",
                 __func__, n_ctx_max, n_streams);
-            dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+            dmds_full = common_get_device_memory_data_impl(paths_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
         }
     }
     add_extra_memory(dmds_full);
@@ -412,7 +414,7 @@ static void common_params_fit_impl(
 
                     int64_t sum_projected_used_min_ctx = 0;
                     cparams->n_ctx = n_ctx_min_total;
-                    dmds_t dmds_min_ctx = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+                    dmds_t dmds_min_ctx = common_get_device_memory_data_impl(paths_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
                     add_extra_memory(dmds_min_ctx);
                     if (nd == 0) {
                         sum_projected_used_min_ctx = dmds_min_ctx.back().mb.total();
@@ -594,7 +596,7 @@ static void common_params_fit_impl(
         set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, mparams_copy);
 
         dmds_t dmd_nl = common_get_device_memory_data_impl(
-            path_model, &mparams_copy, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+            paths_model, &mparams_copy, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
         add_extra_memory(dmd_nl);
 
         LOG_TRC("%s: memory for test allocation by device:\n", func_name);
@@ -623,7 +625,7 @@ static void common_params_fit_impl(
 
         LOG_TRC("%s: getting device memory data with all MoE tensors moved to system memory:\n", __func__);
         dmds_t dmds_cpu_moe = common_get_device_memory_data_impl(
-            path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+            paths_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
         add_extra_memory(dmds_cpu_moe);
 
         for (size_t id = 0; id < nd; id++) {
@@ -877,7 +879,7 @@ static void common_params_fit_impl(
 }
 
 enum common_params_fit_status common_fit_params(
-        const char * path_model,
+        const std::vector<std::string> & paths_model,
         llama_model_params * mparams,
         llama_context_params * cparams,
         float * tensor_split,
@@ -889,7 +891,7 @@ enum common_params_fit_status common_fit_params(
     const int64_t t0_us = llama_time_us();
     common_params_fit_status status = COMMON_PARAMS_FIT_STATUS_SUCCESS;
     try {
-        common_params_fit_impl(path_model, mparams, cparams, tensor_split, tensor_buft_overrides, margins, n_ctx_min, extra, log_level);
+        common_params_fit_impl(paths_model, mparams, cparams, tensor_split, tensor_buft_overrides, margins, n_ctx_min, extra, log_level);
         LOG_TRC("%s: successfully fit params to free device memory\n", __func__);
     } catch (const common_params_fit_exception & e) {
         LOG_WRN("%s: failed to fit params to free device memory: %s\n", __func__, e.what());
@@ -1086,7 +1088,7 @@ void common_memory_breakdown_print(const struct llama_context * ctx) {
 }
 
 void common_fit_print(
-        const char * path_model,
+        const std::vector<std::string> & paths_model,
         llama_model_params * mparams,
         llama_context_params * cparams) {
     std::vector<ggml_backend_dev_t> devs;
@@ -1094,7 +1096,7 @@ void common_fit_print(
     uint32_t hp_nct = 0; // hparams.n_ctx_train
     uint32_t hp_nex = 0; // hparams.n_expert
 
-    auto dmd = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, GGML_LOG_LEVEL_ERROR);
+    auto dmd = common_get_device_memory_data_impl(paths_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, GGML_LOG_LEVEL_ERROR);
     GGML_ASSERT(dmd.size() == devs.size() + 1);
 
     for (size_t id = 0; id < devs.size(); id++) {

common/fit.h

@@ -3,6 +3,7 @@
 #include "ggml.h"
 #include "llama.h"
 
+#include <string>
 #include <vector>
 
 enum common_params_fit_status {
@@ -15,7 +16,7 @@ enum common_params_fit_status {
 //   - its context follows the context of the main model, so its memory is measured again whenever that context changes
 //   - shares_model tells the fit that the weights are already counted in the main model, as for an MTP context
 struct common_fit_extra_model {
-    const char * path_model;
+    std::vector<std::string> paths_model;
     llama_model_params * mparams;
     llama_context_params * cparams;
     bool shares_model;
@@ -27,7 +28,7 @@ struct common_fit_extra_model {
 //   - only parameters that have the same value as in llama_default_model_params are modified
 //     with the exception of the context size which is modified if and only if equal to 0
 common_params_fit_status common_fit_params(
-                         const char * path_model,
+    const std::vector<std::string> & paths_model,
                  llama_model_params * mparams,
                llama_context_params * cparams,
                               float * tensor_split,          // writable buffer for tensor split, needs at least llama_max_devices elements
@@ -39,7 +40,7 @@ common_params_fit_status common_fit_params(
 
 // print estimated memory to stdout
 void common_fit_print(
-                         const char * path_model,
+    const std::vector<std::string> & paths_model,
                  llama_model_params * mparams,
                llama_context_params * cparams);
 
@@ -57,7 +58,7 @@ using common_device_memory_data_vec = std::vector<common_device_memory_data>;
 
 // Load a model + context with no_alloc and return the per-device memory breakdown.
 common_device_memory_data_vec common_get_device_memory_data(
-                         const char * path_model,
+    const std::vector<std::string> & paths_model,
            const llama_model_params * mparams,
          const llama_context_params * cparams,
     std::vector<ggml_backend_dev_t> & devs,

conversion/base.py

@@ -132,7 +132,8 @@ class ModelBase:
                  target_model_dir: Path | None = None,
                  fuse_gate_up_exps: bool = False,
                  fp8_as_q8: bool = False,
-                 fuse_qkv: bool = False):
+                 fuse_qkv: bool = False,
+                 ple_split: bool = False):
         if type(self) is ModelBase or \
                 type(self) is TextModel or \
                 type(self) is MmprojModel:
@@ -197,8 +198,10 @@ class ModelBase:
                 logger.info("heuristics unable to detect tensor dtype, defaulting to --outtype f16")
 
         # Configure GGUF Writer
+        second_shard_tensors = [gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.PER_LAYER_TOKEN_EMBD] + ".weight"] if ple_split else []
         self.gguf_writer = gguf.GGUFWriter(path=None, arch=gguf.MODEL_ARCH_NAMES[self.model_arch], endianess=self.endianess, use_temp_file=self.use_temp_file,
-                                           split_max_tensors=split_max_tensors, split_max_size=split_max_size, dry_run=dry_run, small_first_shard=small_first_shard)
+                                           split_max_tensors=split_max_tensors, split_max_size=split_max_size, dry_run=dry_run, small_first_shard=small_first_shard,
+                                           second_shard_tensors=second_shard_tensors)
 
         # Mistral specific
         self.disable_mistral_community_chat_template = disable_mistral_community_chat_template

convert_hf_to_gguf.py

@@ -101,6 +101,10 @@ def parse_args() -> argparse.Namespace:
         "--no-tensor-first-split", action="store_true",
         help="do not add tensors to the first split (disabled by default)"
     )
+    parser.add_argument(
+        "--ple-split", action="store_true",
+        help="put the per-layer embedding table alone in the second split, so it can be shared between quants (load with --shard)"
+    )
     parser.add_argument(
         "--metadata", type=Path,
         help="Specify the path for an authorship metadata override file"
@@ -219,7 +223,7 @@ def main() -> None:
         "auto": gguf.LlamaFileType.GUESSED,
     }
 
-    is_split = args.split_max_tensors > 0 or args.split_max_size != "0"
+    is_split = args.split_max_tensors > 0 or args.split_max_size != "0" or args.ple_split
     if args.use_temp_file and is_split:
         logger.error("Error: Cannot use temp file when splitting")
         sys.exit(1)
@@ -299,6 +303,7 @@ def main() -> None:
                                      fuse_gate_up_exps=args.fuse_gate_up_exps,
                                      fp8_as_q8=args.fp8_as_q8,
                                      fuse_qkv=args.fuse_qkv,
+                                     ple_split=args.ple_split,
                                      )
 
         if args.vocab_only:

gguf-py/gguf/gguf_writer.py

@@ -11,7 +11,7 @@ from enum import Enum, auto
 from math import prod
 from pathlib import Path
 from io import BufferedWriter
-from typing import IO, Any, Sequence, Mapping
+from typing import IO, Any, Collection, Sequence, Mapping
 from string import ascii_letters, digits
 
 import numpy as np
@@ -86,7 +86,8 @@ class GGUFWriter:
 
     def __init__(
         self, path: os.PathLike[str] | str | None, arch: str, use_temp_file: bool = False, endianess: GGUFEndian = GGUFEndian.LITTLE,
-        split_max_tensors: int = 0, split_max_size: int = 0, dry_run: bool = False, small_first_shard: bool = False
+        split_max_tensors: int = 0, split_max_size: int = 0, dry_run: bool = False, small_first_shard: bool = False,
+        second_shard_tensors: Collection[str] = (),
     ):
         self.fout = None
         self.path = Path(path) if path else None
@@ -101,13 +102,21 @@ class GGUFWriter:
         self.split_max_size = split_max_size
         self.dry_run = dry_run
         self.small_first_shard = small_first_shard
+        self.second_shard_tensors = set(second_shard_tensors)
         logger.info("gguf: This GGUF file is for {0} Endian only".format(
             "Big" if self.endianess == GGUFEndian.BIG else "Little",
         ))
         self.state = WriterState.NO_FILE
 
+        # a fixed shard index lets this shard be shared by models with a different shard count
+        if self.second_shard_tensors:
+            if self.use_temp_file:
+                raise ValueError("Cannot use temp file with second_shard_tensors")
+            self.tensors.append({})
+
         if self.small_first_shard:
             self.tensors.append({})
+        self.cur_shard = len(self.tensors) - 1 if self.small_first_shard else 0
 
         self.add_architecture()
 
@@ -178,6 +187,9 @@ class GGUFWriter:
         if path is not None:
             self.path = path
 
+        if self.second_shard_tensors and not self.tensors[1]:
+            raise ValueError(f"None of the second shard tensors were added: {sorted(self.second_shard_tensors)}")
+
         if self.path is not None:
             filenames = self.print_plan()
             self.fout = [open(filename, "wb") for filename in filenames]
@@ -364,18 +376,23 @@ class GGUFWriter:
             if tensor_dtype == np.uint8:
                 tensor_shape = quant_shape_from_byte_shape(tensor_shape, raw_dtype)
 
+        if name in self.second_shard_tensors:
+            self.tensors[1][name] = TensorInfo(shape=tensor_shape, dtype=dtype, nbytes=tensor_nbytes)
+            return
+
         # make sure there is at least one tensor before splitting
-        if len(self.tensors[-1]) > 0:
+        if len(self.tensors[self.cur_shard]) > 0:
             if (  # split when over tensor limit
                 self.split_max_tensors != 0
-                and len(self.tensors[-1]) >= self.split_max_tensors
+                and len(self.tensors[self.cur_shard]) >= self.split_max_tensors
             ) or (   # split when over size limit
                 self.split_max_size != 0
-                and sum(ti.nbytes for ti in self.tensors[-1].values()) + tensor_nbytes > self.split_max_size
+                and sum(ti.nbytes for ti in self.tensors[self.cur_shard].values()) + tensor_nbytes > self.split_max_size
             ):
                 self.tensors.append({})
+                self.cur_shard = len(self.tensors) - 1
 
-        self.tensors[-1][name] = TensorInfo(shape=tensor_shape, dtype=dtype, nbytes=tensor_nbytes)
+        self.tensors[self.cur_shard][name] = TensorInfo(shape=tensor_shape, dtype=dtype, nbytes=tensor_nbytes)
 
     def add_tensor(
         self, name: str, tensor: np.ndarray[Any, Any], raw_shape: Sequence[int] | None = None,
@@ -397,7 +414,8 @@ class GGUFWriter:
         self.add_tensor_info(name, shape, tensor.dtype, tensor.nbytes, raw_dtype=raw_dtype)
 
         if self.temp_file is None:
-            self.tensors[-1][name].tensor = tensor
+            shard = self.tensors[1] if name in self.second_shard_tensors else self.tensors[self.cur_shard]
+            shard[name].tensor = tensor
             return
 
         tensor.tofile(self.temp_file)

tools/fit-params/fit-params.cpp

@@ -31,7 +31,7 @@ int llama_fit_params(int argc, char ** argv) {
     auto cparams = common_context_params_to_llama(params);
 
     if (!params.fit_params_print) {
-        const common_params_fit_status status = common_fit_params(params.model.path.c_str(), &mparams, &cparams,
+        const common_params_fit_status status = common_fit_params(params.model.get_paths(), &mparams, &cparams,
                 params.tensor_split, params.tensor_buft_overrides.data(), params.fit_params_target.data(), params.fit_params_min_ctx,
                 nullptr,
                 params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR);
@@ -71,7 +71,7 @@ int llama_fit_params(int argc, char ** argv) {
         LOG_INF("%s: printing estimated memory in MiB to stdout (device, model, context, compute) ...\n", __func__);
         common_log_flush(common_log_main());
 
-        common_fit_print(params.model.path.c_str(), &mparams, &cparams);
+        common_fit_print(params.model.get_paths(), &mparams, &cparams);
     }
 
     return 0;

tools/llama-bench/llama-bench.cpp

@@ -2352,7 +2352,7 @@ int llama_bench(int argc, char ** argv) {
             uint32_t n_ctx_needed = std::max<uint32_t>(inst.n_prompt + inst.n_gen + inst.n_depth, inst.fit_min_ctx);
             cparams.n_ctx = std::max(cparams.n_ctx, n_ctx_needed);
 
-            common_fit_params(inst.model.c_str(), &mparams, &cparams,
+            common_fit_params({ inst.model }, &mparams, &cparams,
                 fit_tensor_split.data(),
                 fit_overrides.data(),
                 margins.data(),