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

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
+47 -52
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));