Skip to content

npu: a config.json vocab_size the container's embedding table disagrees with is refused (fixes #1843) #1876

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

Closed
bong-water-water-bong wants to merge 1 commit into from
Closed

npu: a config.json vocab_size the container's embedding table disagrees with is refused (fixes #1843) #1876

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
34 changes: 32 additions & 2 deletions npu/model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,17 +79,47 @@ Model::Model(const std::string& dir) {
// 1.9/1.0.12 Decision packs do), so `tensor("model.embed_tokens.weight")` threw "model has no tensor …";
// (2) a tiled entry's shape is [tile_rows, tile_cols, tile_bytes], so `shape[0]` is a TILE count -- a 248320-row
// table is [7760, 4, 8704] and was reported as vocab 7760, hidden 4.
// 3-D is one tile grid (npu/q4nx.h: 32 rows x 256 columns per tile); 2-D is the plain row shape.
static constexpr int kTileRows = 32;
if (c.contains("vocab_size")) {
dims_.vocab = c.at("vocab_size");
// engine#1843: when the container also carries the table, it is evidence, and a table that disagrees with the
// config is a damaged or mislabelled pack -- neither number can be trusted, so it is refused rather than one
// being picked. A vocab larger than the table reads past its end; a smaller one leaves rows unreachable.
// Compared only where the row count is unambiguous: a float [V, H] table (exact), and a 3-D tile grid, whose
// last tile-row may be padding (V <= tile_rows * 32 < V + 32). A 2-D I8 table is either plain int8 [V, H] or
// [ntiles, chunk_bytes] tile chunks -- bonsai-27B's is [155200, 5120], where 5120 is both the chunk size and
// the hidden width -- so its rows cannot be read off the shape, and it is left to the forward's own checks.
// (Measured 11-10 on the 40 containers on strixhalo: every float and 3-D table matches its vocab_size exactly.)
if (const Tensor* emb = find("model.embed_tokens.weight")) {
const std::string& dt = emb->dtype;
const bool flt = dt == "BF16" || dt == "F16" || dt == "F32";
int rows = -1;
bool ok = true;
if (emb->shape.size() == 3) {
rows = int(emb->shape[0] * kTileRows);
ok = dims_.vocab > 0 && dims_.vocab <= rows && rows - dims_.vocab < kTileRows;
} else if (emb->shape.size() == 2 && flt) {
rows = int(emb->shape[0]);
ok = rows == dims_.vocab;
}
if (!ok)
throw std::runtime_error(dir + "/config.json states vocab_size " + std::to_string(dims_.vocab) +
" but the container's model.embed_tokens.weight has " +
std::to_string(rows) + " rows" +
(emb->shape.size() == 3 ? " (" + std::to_string(emb->shape[0]) +
" tile-rows of 32)"
: std::string()) +
"; one of the two is wrong, so the container is refused rather than "
"served with either");
}
} else {
const Tensor* emb = find("model.embed_tokens.weight");
if (!emb && c.value("tie_word_embeddings", false)) emb = find("lm_head.weight");
if (!emb)
throw std::runtime_error(dir + "/config.json states no vocab_size and the container has no embedding "
"table (`model.embed_tokens.weight`, or `lm_head.weight` for a model "
"that declares tie_word_embeddings)");
// 3-D is one tile grid (npu/q4nx.h: 32 rows x 256 columns per tile); 2-D is the plain row shape.
static constexpr int kTileRows = 32;
dims_.vocab = emb->shape.size() >= 3 ? int(emb->shape[0] * kTileRows) : int(emb->shape.at(0));
}
dims_.rope_theta = c.value("rope_theta", 10000.0f);
Expand Down
33 changes: 32 additions & 1 deletion tests/npu_model_bounds_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ struct Entry {
std::string name;
std::vector<int64_t> shape;
uint64_t bytes;
std::string dtype = "U8";
};

// Writes dir/model.q4nx (tensors packed in order, zero-filled) and dir/config.json.
Expand All @@ -75,7 +76,7 @@ void write_model(const fs::path& dir, const std::vector<Entry>& entries, uint64_
json table = json::object();
uint64_t off = 0;
for (const auto& e : entries) {
table[e.name] = {{"dtype", "U8"}, {"shape", e.shape}, {"data_offsets", {off, off + e.bytes}}};
table[e.name] = {{"dtype", e.dtype}, {"shape", e.shape}, {"data_offsets", {off, off + e.bytes}}};
off += e.bytes;
}
if (!bad_name.empty()) table[bad_name]["data_offsets"] = bad_offsets;
Expand Down Expand Up @@ -160,6 +161,36 @@ int main() {
check(why.empty(), "a tiled embedding without vocab_size derives rows, not tiles" +
(why.empty() ? "" : " (" + why + ")"));
}
{ // engine#1843: vocab_size from config.json is compared with the table when the table is there
struct Case { const char* tag; std::vector<int64_t> shape; std::string dtype; int vocab; bool loads; };
const std::vector<Case> cases = {
{"bf16-match", {2, 128}, "BF16", 2, true},
{"bf16-config-larger", {2, 128}, "BF16", 3, false}, // the embedding read would run past the table
{"bf16-config-smaller", {2, 128}, "BF16", 1, false}, // a row the config says does not exist
{"tiled-exact", {2, 4, 8704}, "I8", 64, true},
{"tiled-padded", {2, 4, 8704}, "I8", 40, true}, // the last tile-row is partly padding
{"tiled-config-larger", {2, 4, 8704}, "I8", 65, false},
{"tiled-whole-spare-tile-row", {2, 4, 8704}, "I8", 32, false},
{"i8-2d-not-compared", {2, 128}, "I8", 999, true}, // plain int8 or tile chunks: rows not on the shape
};
for (const auto& k : cases) {
std::vector<Entry> e = one_layer();
for (auto& x : e)
if (x.name == "model.embed_tokens.weight") {
x.shape = k.shape;
x.dtype = k.dtype;
}
write_model(root / k.tag, e, 0, nullptr, "", json{{"vocab_size", k.vocab}});
const std::string why = thrown([&] { Model m((root / k.tag).string()); });
if (k.loads)
check(why.empty(), std::string("vocab check ") + k.tag + ": loads" +
(why.empty() ? "" : " (" + why + ")"));
else
check(why.find("vocab_size " + std::to_string(k.vocab)) != std::string::npos &&
why.find("model.embed_tokens.weight") != std::string::npos,
std::string("vocab check ") + k.tag + ": refused, naming both numbers (" + why + ")");
}
}
{ // engine#1782: neither -> refused by name, not by "model has no tensor"
std::vector<Entry> e = one_layer();
e.erase(std::remove_if(e.begin(), e.end(),
Expand Down
Loading