Skip to content

fix: prevent torch checkpoint offset overflow #1832

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
leejet merged 1 commit into from
Jul 29, 2026
Merged

fix: prevent torch checkpoint offset overflow #1832

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 60 additions & 5 deletions src/model_io/pickle_io.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

#include <cstdlib>
#include <cstring>
#include <limits>
#include <string>
#include <unordered_map>
#include <utility>
Expand Down Expand Up @@ -512,8 +513,51 @@ static bool parse_storage_type(const std::string& global_name, PickleStorageInfo
return false;
}

static bool tensor_is_contiguous(const PickleTensorInfo& tensor) {
if (tensor.tensor_storage.nelements() == 0) {
static bool checked_pickle_byte_count(int64_t element_count,
uint64_t element_nbytes,
uint64_t* byte_count) {
if (element_count < 0 || element_nbytes == 0) {
return false;
}

uint64_t count = static_cast<uint64_t>(element_count);
if (count > std::numeric_limits<uint64_t>::max() / element_nbytes) {
return false;
}

*byte_count = count * element_nbytes;
return true;
}

static bool tensor_layout_is_valid(const PickleTensorInfo& tensor, uint64_t raw_element_nbytes) {
if (raw_element_nbytes == 0) {
return false;
}

bool has_zero_dimension = false;
uint64_t element_count = 1;
for (int i = 0; i < tensor.tensor_storage.n_dims; ++i) {
int64_t dimension = tensor.tensor_storage.ne[i];
if (dimension < 0) {
return false;
}
if (dimension == 0) {
has_zero_dimension = true;
continue;
}

uint64_t size = static_cast<uint64_t>(dimension);
if (element_count > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) / size) {
return false;
}
element_count *= size;
}

if (!has_zero_dimension &&
element_count > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) / raw_element_nbytes) {
return false;
}
if (has_zero_dimension) {
return true;
}
if (tensor.stride_n_dims != tensor.tensor_storage.n_dims) {
Expand Down Expand Up @@ -932,7 +976,12 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
if (storage.key.empty() || !parse_storage_type(pid.items[1].str_value, &storage)) {
return false;
}
storage.nbytes = (uint64_t)pid.items[4].int_value * storage.raw_element_nbytes;
if (!checked_pickle_byte_count(pid.items[4].int_value,
storage.raw_element_nbytes,
&storage.nbytes)) {
set_error(error, "invalid storage size in torch pickle");
return false;
}
storage_nbytes[storage.key] = storage.nbytes;
stack.push_back(make_storage_value(storage));
} break;
Expand Down Expand Up @@ -963,7 +1012,12 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
tensor.tensor_storage.is_f64 = args.items[0].storage.is_f64;
tensor.tensor_storage.is_i64 = args.items[0].storage.is_i64;
tensor.tensor_storage.storage_key = args.items[0].storage.key;
tensor.tensor_storage.offset = (uint64_t)args.items[1].int_value * args.items[0].storage.raw_element_nbytes;
if (!checked_pickle_byte_count(args.items[1].int_value,
args.items[0].storage.raw_element_nbytes,
&tensor.tensor_storage.offset)) {
set_error(error, "invalid tensor storage offset in torch pickle");
return false;
}

for (const auto& item : args.items[2].items) {
if (item.kind != PickleValue::INT || tensor.tensor_storage.n_dims >= SD_MAX_DIMS) {
Expand All @@ -979,7 +1033,8 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
tensor.stride[tensor.stride_n_dims++] = item.int_value;
}

if (!tensor_is_contiguous(tensor)) {
if (!tensor_layout_is_valid(tensor, args.items[0].storage.raw_element_nbytes)) {
set_error(error, "invalid tensor shape or stride in torch pickle");
return false;
}
stack.push_back(make_tensor_value(tensor));
Expand Down
17 changes: 12 additions & 5 deletions src/model_io/torch_legacy_io.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,11 +139,16 @@ bool read_torch_legacy_file(const std::string& file_path,
if (it == legacy_storage_map.end()) {
return false;
}
if (current_offset + LEGACY_STORAGE_HEADER_SIZE + it->second > file_size) {
if (current_offset > file_size ||
LEGACY_STORAGE_HEADER_SIZE > file_size - current_offset) {
return false;
}
storage_offsets[storage_key] = current_offset + LEGACY_STORAGE_HEADER_SIZE;
current_offset += LEGACY_STORAGE_HEADER_SIZE + it->second;
uint64_t storage_offset = current_offset + LEGACY_STORAGE_HEADER_SIZE;
if (it->second > file_size - storage_offset) {
return false;
}
storage_offsets[storage_key] = storage_offset;
current_offset = storage_offset + it->second;
}

for (auto& tensor_storage : tensor_storages) {
Expand All @@ -159,8 +164,10 @@ bool read_torch_legacy_file(const std::string& file_path,

uint64_t base_offset = it_offset->second;
uint64_t storage_nbytes = it_size->second;
uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
if (tensor_storage.offset + tensor_nbytes > storage_nbytes) {
int64_t tensor_nbytes = tensor_storage.nbytes_to_read();
if (tensor_nbytes < 0 ||
tensor_storage.offset > storage_nbytes ||
static_cast<uint64_t>(tensor_nbytes) > storage_nbytes - tensor_storage.offset) {
return false;
}

Expand Down
6 changes: 4 additions & 2 deletions src/model_io/torch_zip_io.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,8 +76,10 @@ static bool parse_zip_data_pkl(const uint8_t* buffer,
return false;
}

uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
if (tensor_storage.offset + tensor_nbytes > entry_size) {
int64_t tensor_nbytes = tensor_storage.nbytes_to_read();
if (tensor_nbytes < 0 ||
tensor_storage.offset > entry_size ||
static_cast<uint64_t>(tensor_nbytes) > entry_size - tensor_storage.offset) {
set_error(error, "tensor '" + tensor_storage.name + "' exceeds storage entry '" + entry_name + "'");
return false;
}
Expand Down
48 changes: 39 additions & 9 deletions src/model_loader.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1077,6 +1077,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,

std::vector<uint8_t> read_buffer;
std::vector<uint8_t> convert_buffer;
std::vector<uint8_t> zip_entry_buffer;

while (true) {
int64_t t0, t1;
Expand Down Expand Up @@ -1115,34 +1116,60 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,

size_t nbytes_to_read = tensor_storage.nbytes_to_read();

auto read_data = [&](char* buf, size_t n) {
auto read_data = [&](char* buf, size_t n) -> bool {
if (zip != nullptr) {
zip_entry_openbyindex(zip, tensor_storage.index_in_zip);
if (zip_entry_openbyindex(zip, tensor_storage.index_in_zip) != 0) {
LOG_ERROR("failed to open zip entry for tensor '%s'", tensor_storage.name.c_str());
return false;
}
size_t entry_size = zip_entry_size(zip);
if (tensor_storage.offset > entry_size) {
LOG_ERROR("tensor '%s' exceeds its zip storage entry", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
size_t tensor_offset = static_cast<size_t>(tensor_storage.offset);
if (n > entry_size - tensor_offset) {
LOG_ERROR("tensor '%s' exceeds its zip storage entry", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}

if (entry_size != n) {
int64_t t_memcpy_start;
read_buffer.resize(entry_size);
zip_entry_noallocread(zip, (void*)read_buffer.data(), entry_size);
zip_entry_buffer.resize(entry_size);
auto bytes_read = zip_entry_noallocread(zip, (void*)zip_entry_buffer.data(), entry_size);
if (bytes_read < 0 || static_cast<size_t>(bytes_read) != entry_size) {
LOG_ERROR("failed to read zip entry for tensor '%s'", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
t_memcpy_start = ggml_time_ms();
memcpy((void*)buf, (void*)(read_buffer.data() + tensor_storage.offset), n);
memcpy((void*)buf, (void*)(zip_entry_buffer.data() + tensor_offset), n);
memcpy_time_ms.fetch_add(ggml_time_ms() - t_memcpy_start);
} else {
zip_entry_noallocread(zip, (void*)buf, n);
auto bytes_read = zip_entry_noallocread(zip, (void*)buf, n);
if (bytes_read < 0 || static_cast<size_t>(bytes_read) != n) {
LOG_ERROR("failed to read zip entry for tensor '%s'", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
}
zip_entry_close(zip);
} else if (mmapped) {
if (!mmapped->copy_data(buf, n, tensor_storage.offset)) {
LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
failed = true;
return false;
}
} else {
file.seekg(tensor_storage.offset);
file.read(buf, n);
if (!file) {
LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
failed = true;
return false;
}
}
return true;
};

char* read_buf = nullptr;
Expand Down Expand Up @@ -1176,7 +1203,10 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
}

t0 = ggml_time_ms();
read_data(read_buf, nbytes_to_read);
if (!read_data(read_buf, nbytes_to_read)) {
failed = true;
break;
}
t1 = ggml_time_ms();
read_time_ms.fetch_add(t1 - t0);

Expand Down
6 changes: 3 additions & 3 deletions src/model_manager.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,9 +54,9 @@ static bool backend_supports_host_buffer(ggml_backend_t backend) {
}

static bool device_supports_param_op(ggml_backend_dev_t device,
ggml_tensor* weight,
enum ggml_op op,
ggml_backend_buffer_type_t buft) {
ggml_tensor* weight,
enum ggml_op op,
ggml_backend_buffer_type_t buft) {
if (op == GGML_OP_NONE) {
return true;
}
Expand Down
6 changes: 3 additions & 3 deletions src/model_manager.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,9 +131,9 @@ class ModelManager : public RunnerWeightManager {
ResidencyMode residency_mode,
ggml_backend_t compute_backend,
ggml_backend_t params_backend,
size_t* registered_tensor_size = nullptr,
bool allow_split_buffer = false,
bool params_follow_compute_backend = false,
size_t* registered_tensor_size = nullptr,
bool allow_split_buffer = false,
bool params_follow_compute_backend = false,
const std::map<ggml_tensor*, enum ggml_op>* tensor_ops = nullptr);

bool unregister_param_tensors(const std::string& desc,
Expand Down
Loading