diff --git a/npu/model.cpp b/npu/model.cpp index d5a312b4a..855cedf86 100644 --- a/npu/model.cpp +++ b/npu/model.cpp @@ -79,8 +79,40 @@ 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"); @@ -88,8 +120,6 @@ Model::Model(const std::string& dir) { 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); diff --git a/tests/npu_model_bounds_test.cpp b/tests/npu_model_bounds_test.cpp index d76e7805a..b80392306 100644 --- a/tests/npu_model_bounds_test.cpp +++ b/tests/npu_model_bounds_test.cpp @@ -64,6 +64,7 @@ struct Entry { std::string name; std::vector shape; uint64_t bytes; + std::string dtype = "U8"; }; // Writes dir/model.q4nx (tensors packed in order, zero-filled) and dir/config.json. @@ -75,7 +76,7 @@ void write_model(const fs::path& dir, const std::vector& 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; @@ -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 shape; std::string dtype; int vocab; bool loads; }; + const std::vector 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 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 e = one_layer(); e.erase(std::remove_if(e.begin(), e.end(),