upstream-archive/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-08-05 09:30:49 -05:00
Code Issues Packages Projects Releases Wiki Activity

model : allow reshape of tensors during load (#26531)

Browse Source
This commit is contained in:
Georgi Gerganov
2026-08-04 09:06:44 +03:00
committed by GitHub
parent 935cad6497
commit 1269cb1ff1
5 changed files with 60 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

99
src/llama-model-loader.cpp
View File

@@ -857,7 +857,11 @@ struct ggml_tensor * llama_model_loader::require_tensor_meta(const std::string &
return tensor;
}
const struct ggml_tensor * llama_model_loader::check_tensor_dims(const std::string & name, const std::vector<int64_t> & ne, bool required) const {
const struct ggml_tensor * llama_model_loader::check_tensor_dims(
const std::string & name,
const std::vector<int64_t> & ne,
bool required,
bool allow_reshape) const {
const struct ggml_tensor * cur = get_tensor_meta(name.c_str());
if (cur == NULL) {
@@ -867,21 +871,33 @@ const struct ggml_tensor * llama_model_loader::check_tensor_dims(const std::stri
throw std::runtime_error(format("%s: tensor '%s' not found", __func__, name.c_str()));
}
{
bool is_ok = true;
bool is_ok = true;
if (allow_reshape) {
// check total number of elements only
const int64_t ncur = ggml_nelements(cur);
int64_t nexp = 1;
for (size_t i = 0; i < ne.size(); ++i) {
nexp *= ne[i];
}
if (ncur != nexp) {
is_ok = false;
}
} else {
for (size_t i = 0; i < GGML_MAX_DIMS; ++i) {
if ((i < ne.size() && ne[i] != cur->ne[i]) || (i >= ne.size() && cur->ne[i] != 1)) {
is_ok = false;
break;
}
}
if (!is_ok) {
throw std::runtime_error(
format("%s: tensor '%s' has wrong shape; expected %s, got %s",
__func__, name.c_str(),
llama_format_tensor_shape(ne).c_str(),
llama_format_tensor_shape(cur).c_str()));
}
}
if (!is_ok) {
throw std::runtime_error(
format("%s: tensor '%s' has wrong shape; expected %s, got %s",
__func__, name.c_str(),
llama_format_tensor_shape(ne).c_str(),
llama_format_tensor_shape(cur).c_str()));
}
return cur;
@@ -1246,11 +1262,25 @@ struct ggml_tensor * llama_model_loader::create_tensor(
return ret;
}
ggml_tensor * t_meta = get_tensor_meta(tn.str().c_str());
ggml_backend_buffer_type_t buft = buft_for_tensor(t_meta);
if (buft == nullptr) {
return nullptr; // return type is ggml_tensor *
LLAMA_LOG_DEBUG("%s: loading tensor %s\n", __func__, tn.str().c_str());
const struct ggml_tensor * cur = check_tensor_dims(tn.str(), ne, !(flags & TENSOR_NOT_REQUIRED), flags & TENSOR_ALLOW_RESHAPE);
if (cur == NULL) {
return NULL;
}
ggml_tensor t_meta = *cur;
if (flags & TENSOR_ALLOW_RESHAPE) {
for (size_t dim = 0; dim < GGML_MAX_DIMS; dim++) {
t_meta.ne[dim] = dim < ne.size() ? ne.begin()[dim] : 1;
t_meta.nb[dim] = dim == 0 ? ggml_type_size(t_meta.type) : t_meta.ne[dim-1]*t_meta.nb[dim-1];
}
}
ggml_backend_buffer_type_t buft = buft_for_tensor(&t_meta);
if (buft == nullptr) {
return nullptr;
}
ggml_context * ctx = ctx_for_buft(buft);
// if duplicated, check if the original tensor was allocated in the same buffer type context and avoid creating a new one
@@ -1261,20 +1291,13 @@ struct ggml_tensor * llama_model_loader::create_tensor(
}
}
LLAMA_LOG_DEBUG("%s: loading tensor %s\n", __func__, tn.str().c_str());
const struct ggml_tensor * cur = check_tensor_dims(tn.str(), ne, !(flags & TENSOR_NOT_REQUIRED));
if (cur == NULL) {
return NULL;
}
const bool duplicated = flags & TENSOR_DUPLICATED;
struct ggml_tensor * tensor = ggml_dup_tensor(ctx, cur);
ggml_set_name(tensor, ggml_get_name(cur));
struct ggml_tensor * tensor = ggml_dup_tensor(ctx, &t_meta);
ggml_set_name(tensor, ggml_get_name(&t_meta));
if (duplicated) {
size_data += ggml_nbytes(cur);
size_data += ggml_nbytes(&t_meta);
} else {
n_created++;
}
@@ -1282,34 +1305,6 @@ struct ggml_tensor * llama_model_loader::create_tensor(
return tensor;
}
struct ggml_tensor * llama_model_loader::create_tensor_as_view(struct ggml_context * ctx, struct ggml_tensor * base, const std::string & name, const std::initializer_list<int64_t> & ne, size_t offset, bool required) {
const struct ggml_tensor * cur = check_tensor_dims(name, ne, required);
if (cur == NULL) {
return NULL;
}
if (cur->type != base->type) {
throw std::runtime_error(format("%s: tensor '%s' has wrong type; expected %s, got %s", __func__, name.c_str(), ggml_type_name(base->type), ggml_type_name(cur->type)));
}
std::array<int64_t, GGML_MAX_DIMS> dims;
for (size_t i = 0; i < GGML_MAX_DIMS; ++i) {
dims[i] = i < ne.size() ? ne.begin()[i] : 1;
}
struct ggml_tensor * tensor = ggml_view_4d(ctx, base,
dims[0], dims[1], dims[2], dims[3],
cur->nb[1], cur->nb[2], cur->nb[3],
offset);
ggml_set_name(tensor, name.c_str());
n_created++;
return tensor;
}
void llama_model_loader::done_getting_tensors(bool partial) const {
if (n_created > n_tensors) {
throw std::runtime_error(format("%s: too many tensors created; expected %d, got %d", __func__, n_tensors, n_created));

9
src/llama-model-loader.h
View File

@@ -67,6 +67,7 @@ struct llama_model_loader {
static const int TENSOR_DUPLICATED = 1 << 1;
static const int TENSOR_SKIP = 1 << 2;
static const int TENSOR_SKIP_IF_VIRTUAL = 1 << 3;
static const int TENSOR_ALLOW_RESHAPE = 1 << 4;
int n_kv = 0;
int n_tensors = 0;
@@ -177,14 +178,16 @@ struct llama_model_loader {
struct ggml_tensor * require_tensor_meta(const std::string & name) const;
const struct ggml_tensor * check_tensor_dims(const std::string & name, const std::vector<int64_t> & ne, bool required) const;
const struct ggml_tensor * check_tensor_dims(
const std::string & name,
const std::vector<int64_t> & ne,
bool required,
bool allow_reshape) const;
struct ggml_tensor * create_tensor(
const llama_hparams & hparams, const buft_list_t * buft_list_cpu, const buft_list_t * buft_list_input, const buft_list_t * buft_list_output,
const buft_list_t * buft_list_layer, const LLM_TN_IMPL & tn, const std::initializer_list<int64_t> & ne, int flags);
struct ggml_tensor * create_tensor_as_view(struct ggml_context * ctx, struct ggml_tensor * base, const std::string & name, const std::initializer_list<int64_t> & ne, size_t offset, bool required = true);
void done_getting_tensors(bool partial = false) const;
void init_mappings(bool prefetch = true, llama_mlocks * mlock_mmaps = nullptr);

3
src/llama-model.cpp
View File

@@ -2867,7 +2867,8 @@ llama_model_base::llama_model_base(const struct llama_model_params & params) : l
TENSOR_DUPLICATED (llama_model_loader::TENSOR_DUPLICATED),
TENSOR_NOT_REQUIRED (llama_model_loader::TENSOR_NOT_REQUIRED),
TENSOR_SKIP (llama_model_loader::TENSOR_SKIP),
TENSOR_SKIP_IF_VIRTUAL(llama_model_loader::TENSOR_SKIP_IF_VIRTUAL) {}
TENSOR_SKIP_IF_VIRTUAL(llama_model_loader::TENSOR_SKIP_IF_VIRTUAL),
TENSOR_ALLOW_RESHAPE (llama_model_loader::TENSOR_ALLOW_RESHAPE) {}
ggml_tensor * llama_model_base::create_tensor(const LLM_TN_IMPL & tn, const std::initializer_list<int64_t> & ne, int flags) {
GGML_ASSERT(ml != nullptr);

1
src/llama-model.h
View File

@@ -719,6 +719,7 @@ struct llama_model_base : public llama_model {
const int TENSOR_NOT_REQUIRED;
const int TENSOR_SKIP;
const int TENSOR_SKIP_IF_VIRTUAL;
const int TENSOR_ALLOW_RESHAPE;
explicit llama_model_base(const llama_model_params & params);
virtual ~llama_model_base() = default;

6
src/models/deepseek4.cpp
View File

@@ -114,7 +114,9 @@ void llama_model_deepseek4::load_arch_tensors(llama_model_loader & ml) {
layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head}, flags);
layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", i), {n_embd, n_embd_head}, flags);
layer.attn_kv_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM, "weight", i), {n_embd_head}, flags);
layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank * o_groups}, flags);
// for wo_a, the shape in the file is (n_head * n_embd_head / o_groups, o_lora_rank*o_groups)
// so we reshape here, to avoid reshaping the tensor in the graph
layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank, o_groups}, flags | TENSOR_ALLOW_RESHAPE);
layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), {o_groups * o_lora_rank, n_embd}, flags);
layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc_dim, hc_mix_dim}, flags);
@@ -1258,7 +1260,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention_impl(
out = ggml_reshape_3d(ctx0, out, o_group_dim, n_groups, nt);
out = ggml_permute(ctx0, out, 0, 2, 1, 3);
ggml_tensor * oa = ggml_mul_mat(ctx0, ggml_reshape_3d(ctx0, layer.wo_a, layer.wo_a->ne[0], o_lora_rank, n_groups), out);
ggml_tensor * oa = ggml_mul_mat(ctx0, layer.wo_a, out);
cb(oa, "attn_wo_a", il);
oa = ggml_permute(ctx0, oa, 0, 2, 1, 3);
oa = ggml_cont_2d(ctx0, oa, o_lora_rank*n_groups, nt);