diff --git a/examples/gpt2_chat/CMakeLists.txt b/examples/gpt2_chat/CMakeLists.txt new file mode 100644 index 00000000..fc4485d4 --- /dev/null +++ b/examples/gpt2_chat/CMakeLists.txt @@ -0,0 +1,14 @@ +# Example: minimal GPT-2 chat demo. Disabled by default; enable with +# cmake -DFDEEP_BUILD_GPT2_CHAT=ON .. +# The example needs the rebuilt model JSON and the GPT-2 vocab/merges, see +# keras_export/save_gpt2_backbone_for_fdeep.py and the README in this directory. + +option(FDEEP_BUILD_GPT2_CHAT "Build the GPT-2 chat demo example" OFF) + +if(FDEEP_BUILD_GPT2_CHAT) + add_executable(gpt2_chat main.cpp) + target_link_libraries(gpt2_chat PRIVATE fdeep) + if(NOT MSVC) + target_compile_options(gpt2_chat PRIVATE -O3) + endif() +endif() diff --git a/examples/gpt2_chat/main.cpp b/examples/gpt2_chat/main.cpp new file mode 100644 index 00000000..6e8dee2f --- /dev/null +++ b/examples/gpt2_chat/main.cpp @@ -0,0 +1,158 @@ +// Minimal GPT-2 chat demo on top of frugally-deep. +// +// Usage: +// gpt2_chat [max_new_tokens] [temperature] [top_k] +// +// Reads a prompt per line from stdin and prints the continuation. A blank +// line exits. Uses a stateful key/value cache (see +// include/fdeep/llm/gpt2_cached.hpp), so each new token costs roughly +// constant time rather than re-encoding the whole prefix. +// +// The weights binary is produced by +// keras_export/save_gpt2_weights_bin.py --output weights.bin + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +int sample_logits(const std::vector& logits, float temperature, + std::size_t top_k, std::mt19937_64& rng) +{ + if (temperature <= 0.0f) { + std::size_t best = 0; + float best_val = logits[0]; + for (std::size_t i = 1; i < logits.size(); ++i) { + if (logits[i] > best_val) { + best_val = logits[i]; + best = i; + } + } + return static_cast(best); + } + + std::vector> cand; + cand.reserve(logits.size()); + const float inv_t = 1.0f / temperature; + for (std::size_t i = 0; i < logits.size(); ++i) { + cand.emplace_back(logits[i] * inv_t, static_cast(i)); + } + if (top_k > 0 && top_k < cand.size()) { + std::partial_sort(cand.begin(), + cand.begin() + static_cast(top_k), cand.end(), + [](const auto& a, const auto& b) { return a.first > b.first; }); + cand.resize(top_k); + } + float m = cand[0].first; + for (const auto& c : cand) + m = std::max(m, c.first); + float total = 0.0f; + for (auto& c : cand) { + c.first = std::exp(c.first - m); + total += c.first; + } + std::uniform_real_distribution dist(0.0f, total); + const float pick = dist(rng); + float acc = 0.0f; + for (const auto& c : cand) { + acc += c.first; + if (acc >= pick) + return c.second; + } + return cand.back().second; +} + +} // namespace + +int main(int argc, char** argv) +{ + if (argc < 4) { + std::cerr << "usage: " << argv[0] + << " " + " [max_new_tokens] [temperature] [top_k] [max_seq_len]\n"; + return 1; + } + const std::string weights_path = argv[1]; + const std::string vocab_path = argv[2]; + const std::string merges_path = argv[3]; + const std::size_t max_new = (argc > 4) ? std::stoul(argv[4]) : 100; + const float temperature = (argc > 5) ? std::stof(argv[5]) : 0.0f; + const std::size_t top_k = (argc > 6) ? std::stoul(argv[6]) : 0; + const std::size_t max_seq_len = (argc > 7) ? std::stoul(argv[7]) : 256; + + std::cerr << "loading tokenizer...\n"; + fdeep::llm::gpt2_bpe_tokenizer tok(vocab_path, merges_path); + + std::cerr << "loading weights...\n"; + const auto t0 = std::chrono::steady_clock::now(); + fdeep::llm::gpt2_cached_model gpt(weights_path, max_seq_len); + const auto t1 = std::chrono::steady_clock::now(); + std::cerr << "loaded in " + << std::chrono::duration(t1 - t0).count() << " s\n"; + + std::mt19937_64 rng(static_cast( + std::chrono::steady_clock::now().time_since_epoch().count())); + + std::cerr << "ready (max_seq_len=" << max_seq_len + << ", temperature=" << temperature + << ", top_k=" << top_k + << "). type a prompt and press enter; blank line quits.\n"; + + std::string line; + while (true) { + std::cout << "> " << std::flush; + if (!std::getline(std::cin, line)) + break; + if (line.empty()) + break; + + gpt.reset(); + const auto prompt_ids = tok.encode(line); + if (prompt_ids.empty()) + continue; + if (prompt_ids.size() >= max_seq_len) { + std::cerr << "[prompt is " << prompt_ids.size() + << " tokens; max_seq_len is " << max_seq_len + << "]\n"; + continue; + } + + std::cout << line; + std::cout.flush(); + + const auto t_pre0 = std::chrono::steady_clock::now(); + auto logits = gpt.prefill(prompt_ids); + const auto t_pre1 = std::chrono::steady_clock::now(); + + std::vector generated; + const auto t_gen0 = std::chrono::steady_clock::now(); + for (std::size_t step = 0; step < max_new; ++step) { + const int next_id = sample_logits(logits, temperature, top_k, rng); + generated.push_back(next_id); + std::cout << tok.decode({ next_id }) << std::flush; + if (next_id == tok.eos_token_id()) + break; + if (gpt.cur_len() >= gpt.max_seq_len()) + break; + logits = gpt.step(next_id); + } + const auto t_gen1 = std::chrono::steady_clock::now(); + std::cout << "\n[" + << "prefill " << prompt_ids.size() << " tok " + << std::chrono::duration(t_pre1 - t_pre0).count() << " s, " + << "decode " << generated.size() << " tok " + << std::chrono::duration(t_gen1 - t_gen0).count() << " s, " + << (generated.empty() ? 0.0 : 1000.0 * std::chrono::duration(t_gen1 - t_gen0).count() / static_cast(generated.size())) + << " ms/tok]\n"; + } + return 0; +} diff --git a/include/fdeep/llm/gpt2_bpe.hpp b/include/fdeep/llm/gpt2_bpe.hpp new file mode 100644 index 00000000..7882f892 --- /dev/null +++ b/include/fdeep/llm/gpt2_bpe.hpp @@ -0,0 +1,582 @@ +// Copyright 2026, Tobias Hermann. +// https://github.com/Dobiasd/frugally-deep +// Distributed under the MIT License. +// (See accompanying LICENSE file or at +// https://opensource.org/licenses/MIT) + +// GPT-2-compatible byte-level BPE tokenizer. +// +// Loads vocab.json (token-string -> id) and merges.txt (one merge rule per +// line, " ") as exported from keras_hub's GPT2Tokenizer or HuggingFace's +// gpt2 tokenizer. Provides encode(text) -> vector and decode(ids) -> string. +// +// Pre-tokenization uses GPT-2's regex pattern restricted to the ASCII subset. +// Non-ASCII text is byte-level safe (UTF-8 bytes round-trip via the byte +// encoder), but pre-tokenization will not split unicode words the way the +// reference Python implementation does. For the standard chat-demo use case +// this is sufficient. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace fdeep { +namespace llm { + + namespace internal { + + // Build the canonical GPT-2 byte<->unicode mapping. Returns a 256-entry + // table of UTF-8 strings, one per byte. The inverse mapping is built + // alongside. + inline std::array make_byte_to_unicode( + std::unordered_map& unicode_to_byte) + { + std::array table; + std::vector mapped(256, false); + + auto encode_codepoint = [](uint32_t cp) { + std::string out; + if (cp < 0x80) { + out.push_back(static_cast(cp)); + } else if (cp < 0x800) { + out.push_back(static_cast(0xC0 | (cp >> 6))); + out.push_back(static_cast(0x80 | (cp & 0x3F))); + } else if (cp < 0x10000) { + out.push_back(static_cast(0xE0 | (cp >> 12))); + out.push_back(static_cast(0x80 | ((cp >> 6) & 0x3F))); + out.push_back(static_cast(0x80 | (cp & 0x3F))); + } else { + out.push_back(static_cast(0xF0 | (cp >> 18))); + out.push_back(static_cast(0x80 | ((cp >> 12) & 0x3F))); + out.push_back(static_cast(0x80 | ((cp >> 6) & 0x3F))); + out.push_back(static_cast(0x80 | (cp & 0x3F))); + } + return out; + }; + + // The "printable" ranges that map to themselves. + auto add_range = [&](uint32_t lo, uint32_t hi) { + for (uint32_t b = lo; b <= hi; ++b) { + table[b] = encode_codepoint(b); + mapped[b] = true; + } + }; + add_range(0x21, 0x7E); // '!'..'~' + add_range(0xA1, 0xAC); + add_range(0xAE, 0xFF); + + // Remaining bytes get codepoints starting at 256, in byte order. + uint32_t next_cp = 256; + for (uint32_t b = 0; b < 256; ++b) { + if (!mapped[b]) { + table[b] = encode_codepoint(next_cp++); + mapped[b] = true; + } + } + + unicode_to_byte.clear(); + for (uint32_t b = 0; b < 256; ++b) { + unicode_to_byte[table[b]] = static_cast(b); + } + return table; + } + + // Strip surrounding whitespace from a string (in place). + inline void strip(std::string& s) + { + std::size_t i = 0; + while (i < s.size() && (s[i] == ' ' || s[i] == '\t' || s[i] == '\r' || s[i] == '\n')) { + ++i; + } + std::size_t j = s.size(); + while (j > i && (s[j - 1] == ' ' || s[j - 1] == '\t' || s[j - 1] == '\r' || s[j - 1] == '\n')) { + --j; + } + s = s.substr(i, j - i); + } + + // Minimal JSON string-key -> int reader for vocab.json. Assumes the + // file is the output of json.dump on a flat dict where keys are the + // BPE token strings (already byte-encoded) and values are non-negative + // integers. Handles standard JSON escapes (\\, \", \n, \t, \r, \b, \f, + // \/) and \uXXXX surrogate pairs. + inline std::unordered_map load_vocab_json(const std::string& path) + { + std::ifstream in(path); + if (!in) { + throw std::runtime_error("could not open " + path); + } + std::ostringstream ss; + ss << in.rdbuf(); + const std::string text = ss.str(); + + std::unordered_map vocab; + vocab.reserve(60000); + + auto encode_codepoint = [](uint32_t cp, std::string& out) { + if (cp < 0x80) { + out.push_back(static_cast(cp)); + } else if (cp < 0x800) { + out.push_back(static_cast(0xC0 | (cp >> 6))); + out.push_back(static_cast(0x80 | (cp & 0x3F))); + } else if (cp < 0x10000) { + out.push_back(static_cast(0xE0 | (cp >> 12))); + out.push_back(static_cast(0x80 | ((cp >> 6) & 0x3F))); + out.push_back(static_cast(0x80 | (cp & 0x3F))); + } else { + out.push_back(static_cast(0xF0 | (cp >> 18))); + out.push_back(static_cast(0x80 | ((cp >> 12) & 0x3F))); + out.push_back(static_cast(0x80 | ((cp >> 6) & 0x3F))); + out.push_back(static_cast(0x80 | (cp & 0x3F))); + } + }; + + std::size_t i = 0; + const std::size_t n = text.size(); + + auto skip_ws = [&]() { + while (i < n && (text[i] == ' ' || text[i] == '\t' || text[i] == '\n' || text[i] == '\r')) { + ++i; + } + }; + + auto read_string = [&](std::string& out) { + if (text[i] != '"') { + throw std::runtime_error("vocab.json: expected '\"' at offset " + std::to_string(i)); + } + ++i; + while (i < n && text[i] != '"') { + if (text[i] == '\\' && i + 1 < n) { + const char esc = text[++i]; + switch (esc) { + case '"': + out.push_back('"'); + break; + case '\\': + out.push_back('\\'); + break; + case '/': + out.push_back('/'); + break; + case 'b': + out.push_back('\b'); + break; + case 'f': + out.push_back('\f'); + break; + case 'n': + out.push_back('\n'); + break; + case 'r': + out.push_back('\r'); + break; + case 't': + out.push_back('\t'); + break; + case 'u': { + if (i + 4 >= n) { + throw std::runtime_error("vocab.json: bad \\u escape"); + } + uint32_t cp = 0; + for (int k = 0; k < 4; ++k) { + const char h = text[++i]; + cp <<= 4; + if (h >= '0' && h <= '9') + cp |= static_cast(h - '0'); + else if (h >= 'a' && h <= 'f') + cp |= static_cast(h - 'a' + 10); + else if (h >= 'A' && h <= 'F') + cp |= static_cast(h - 'A' + 10); + else + throw std::runtime_error("vocab.json: bad hex digit"); + } + if (cp >= 0xD800 && cp <= 0xDBFF && i + 6 < n + && text[i + 1] == '\\' && text[i + 2] == 'u') { + uint32_t lo = 0; + std::size_t j = i + 3; + for (int k = 0; k < 4; ++k) { + const char h = text[j++]; + lo <<= 4; + if (h >= '0' && h <= '9') + lo |= static_cast(h - '0'); + else if (h >= 'a' && h <= 'f') + lo |= static_cast(h - 'a' + 10); + else if (h >= 'A' && h <= 'F') + lo |= static_cast(h - 'A' + 10); + else + throw std::runtime_error("vocab.json: bad hex digit"); + } + if (lo >= 0xDC00 && lo <= 0xDFFF) { + cp = 0x10000u + ((cp - 0xD800u) << 10) + (lo - 0xDC00u); + i += 6; + } + } + encode_codepoint(cp, out); + break; + } + default: + throw std::runtime_error( + std::string("vocab.json: unknown escape \\") + esc); + } + ++i; + } else { + out.push_back(text[i++]); + } + } + if (i >= n) { + throw std::runtime_error("vocab.json: unterminated string"); + } + ++i; // consume closing quote + }; + + auto read_int = [&]() { + std::size_t start = i; + if (i < n && (text[i] == '-' || text[i] == '+')) + ++i; + while (i < n && text[i] >= '0' && text[i] <= '9') + ++i; + return std::stoi(text.substr(start, i - start)); + }; + + skip_ws(); + if (i >= n || text[i] != '{') { + throw std::runtime_error("vocab.json: expected '{'"); + } + ++i; + skip_ws(); + if (i < n && text[i] == '}') { + return vocab; + } + while (i < n) { + skip_ws(); + std::string key; + read_string(key); + skip_ws(); + if (i >= n || text[i] != ':') { + throw std::runtime_error("vocab.json: expected ':'"); + } + ++i; + skip_ws(); + vocab.emplace(std::move(key), read_int()); + skip_ws(); + if (i < n && text[i] == ',') { + ++i; + continue; + } + if (i < n && text[i] == '}') { + ++i; + break; + } + throw std::runtime_error("vocab.json: expected ',' or '}'"); + } + return vocab; + } + + inline std::vector> load_merges(const std::string& path) + { + std::ifstream in(path); + if (!in) { + throw std::runtime_error("could not open " + path); + } + std::vector> merges; + std::string line; + bool first = true; + while (std::getline(in, line)) { + if (first) { + first = false; + if (line.size() >= 1 && line[0] == '#') { + continue; // skip "#version: 0.2" header + } + } + if (line.empty()) + continue; + const std::size_t sp = line.find(' '); + if (sp == std::string::npos) + continue; + merges.emplace_back(line.substr(0, sp), line.substr(sp + 1)); + } + return merges; + } + + } // namespace internal + + class gpt2_bpe_tokenizer { + public: + gpt2_bpe_tokenizer(const std::string& vocab_path, const std::string& merges_path) + : byte_to_unicode_(internal::make_byte_to_unicode(unicode_to_byte_)) + , vocab_(internal::load_vocab_json(vocab_path)) + { + const auto merges = internal::load_merges(merges_path); + merge_ranks_.reserve(merges.size()); + for (std::size_t i = 0; i < merges.size(); ++i) { + merge_ranks_.emplace(merges[i], static_cast(i)); + } + inv_vocab_.resize(vocab_.size()); + for (const auto& kv : vocab_) { + if (kv.second < 0 || static_cast(kv.second) >= inv_vocab_.size()) { + throw std::runtime_error("vocab id out of range"); + } + inv_vocab_[static_cast(kv.second)] = kv.first; + } + } + + std::size_t vocab_size() const { return vocab_.size(); } + + std::vector encode(const std::string& text) const + { + std::vector out; + const auto pieces = pre_tokenize(text); + for (const auto& piece : pieces) { + std::string encoded; + encoded.reserve(piece.size() * 2); + for (unsigned char b : piece) { + encoded += byte_to_unicode_[b]; + } + bpe_encode(encoded, out); + } + return out; + } + + std::string decode(const std::vector& ids) const + { + std::string concat; + for (int id : ids) { + if (id < 0 || static_cast(id) >= inv_vocab_.size()) { + continue; // skip unknown + } + concat += inv_vocab_[static_cast(id)]; + } + // Decode the byte-encoded unicode string back to raw bytes. + std::string out; + out.reserve(concat.size()); + std::size_t i = 0; + while (i < concat.size()) { + const unsigned char c = static_cast(concat[i]); + std::size_t len = 1; + if ((c & 0x80) == 0) { + len = 1; + } else if ((c & 0xE0) == 0xC0) { + len = 2; + } else if ((c & 0xF0) == 0xE0) { + len = 3; + } else if ((c & 0xF8) == 0xF0) { + len = 4; + } + if (i + len > concat.size()) + break; + const std::string ch = concat.substr(i, len); + auto it = unicode_to_byte_.find(ch); + if (it != unicode_to_byte_.end()) { + out.push_back(static_cast(it->second)); + } else { + out += ch; + } + i += len; + } + return out; + } + + int eos_token_id() const + { + auto it = vocab_.find("<|endoftext|>"); + return it != vocab_.end() ? it->second : -1; + } + + private: + // Apply GPT-2 pre-tokenization to ASCII input. Splits the text into + // pieces that are then byte-encoded and BPE-merged independently. + // Pattern matched (in order): + // 's | 't | 're | 've | 'm | 'll | 'd + // ' ?' followed by [A-Za-z]+ + // ' ?' followed by [0-9]+ + // ' ?' followed by run of non-space, non-letter, non-digit + // whitespace runs (final-only or other) + // Non-ASCII bytes are passed through as 1-byte pieces. + std::vector pre_tokenize(const std::string& text) const + { + std::vector out; + const std::size_t n = text.size(); + std::size_t i = 0; + auto is_letter = [](unsigned char c) { + return (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z'); + }; + auto is_digit = [](unsigned char c) { + return c >= '0' && c <= '9'; + }; + auto is_space = [](unsigned char c) { + return c == ' ' || c == '\t' || c == '\n' || c == '\r' || c == '\v' || c == '\f'; + }; + while (i < n) { + // Contractions: 's, 't, 're, 've, 'm, 'll, 'd + if (text[i] == '\'') { + static const char* const contractions[] = { "'s", "'t", "'re", "'ve", "'m", "'ll", "'d" }; + bool matched = false; + for (const char* c : contractions) { + const std::size_t L = std::strlen(c); + if (i + L <= n && text.compare(i, L, c) == 0) { + out.emplace_back(text.substr(i, L)); + i += L; + matched = true; + break; + } + } + if (matched) + continue; + } + + // GPT-2 attaches an optional leading space to letters/digits/ + // punctuation pieces. Non-space whitespace (\t, \n, ...) is + // handled by the whitespace-run branch. + const bool leading_space = (text[i] == ' '); + const std::size_t start = i; + const std::size_t look = leading_space ? i + 1 : i; + + if (look < n && is_letter(static_cast(text[look]))) { + std::size_t j = look; + while (j < n && is_letter(static_cast(text[j]))) + ++j; + out.emplace_back(text.substr(start, j - start)); + i = j; + continue; + } + if (look < n && is_digit(static_cast(text[look]))) { + std::size_t j = look; + while (j < n && is_digit(static_cast(text[j]))) + ++j; + out.emplace_back(text.substr(start, j - start)); + i = j; + continue; + } + if (look < n + && !is_letter(static_cast(text[look])) + && !is_digit(static_cast(text[look])) + && !is_space(static_cast(text[look]))) { + std::size_t j = look; + while (j < n + && !is_letter(static_cast(text[j])) + && !is_digit(static_cast(text[j])) + && !is_space(static_cast(text[j]))) { + ++j; + } + out.emplace_back(text.substr(start, j - start)); + i = j; + continue; + } + // No content match — text[look] is whitespace or out of range. + // text[i] must be whitespace (since the only way ``look`` + // skips a character is leading_space). + if (is_space(static_cast(text[i]))) { + std::size_t j = i; + while (j < n && is_space(static_cast(text[j]))) + ++j; + const std::size_t run = j - i; + // Emit the whole run except for a single trailing space + // before non-space content; that space attaches to the + // next piece via the ``leading_space`` logic above. + if (j < n && text[j - 1] == ' ' && run >= 2) { + out.emplace_back(text.substr(i, run - 1)); + i = j - 1; + } else { + out.emplace_back(text.substr(i, run)); + i = j; + } + continue; + } + // Fallback: emit one byte at a time (handles non-ASCII). + out.emplace_back(text.substr(i, 1)); + ++i; + } + return out; + } + + // Apply BPE merges to a single byte-encoded piece, appending the + // resulting token ids to ``out``. Operates on a vector of UTF-8 + // codepoints (each codepoint is one BPE "symbol" initially). + void bpe_encode(const std::string& piece, std::vector& out) const + { + if (piece.empty()) + return; + // Split into codepoints. + std::vector symbols; + std::size_t i = 0; + while (i < piece.size()) { + const unsigned char c = static_cast(piece[i]); + std::size_t len = 1; + if ((c & 0x80) == 0) + len = 1; + else if ((c & 0xE0) == 0xC0) + len = 2; + else if ((c & 0xF0) == 0xE0) + len = 3; + else if ((c & 0xF8) == 0xF0) + len = 4; + if (i + len > piece.size()) + len = piece.size() - i; + symbols.emplace_back(piece.substr(i, len)); + i += len; + } + + // Iteratively merge the lowest-rank adjacent pair. + while (symbols.size() > 1) { + int best_rank = std::numeric_limits::max(); + std::size_t best_idx = symbols.size(); + for (std::size_t k = 0; k + 1 < symbols.size(); ++k) { + auto it = merge_ranks_.find({ symbols[k], symbols[k + 1] }); + if (it != merge_ranks_.end() && it->second < best_rank) { + best_rank = it->second; + best_idx = k; + } + } + if (best_idx == symbols.size()) + break; + symbols[best_idx] += symbols[best_idx + 1]; + symbols.erase(symbols.begin() + static_cast(best_idx) + 1); + } + + for (const auto& sym : symbols) { + auto it = vocab_.find(sym); + if (it != vocab_.end()) { + out.push_back(it->second); + } else { + // Unknown symbol: emit each byte as its own id (must exist + // since the byte-level BPE always covers all 256 bytes). + for (char ch : sym) { + const unsigned char ub = static_cast(ch); + const auto& enc = byte_to_unicode_[ub]; + auto it2 = vocab_.find(enc); + if (it2 != vocab_.end()) + out.push_back(it2->second); + } + } + } + } + + struct pair_hash { + std::size_t operator()(const std::pair& p) const noexcept + { + return std::hash()(p.first) ^ (std::hash()(p.second) << 1); + } + }; + + // Declaration order matters: ``unicode_to_byte_`` is filled in by + // ``make_byte_to_unicode`` (called from the initialiser of + // ``byte_to_unicode_``), so it must be constructed first. + std::unordered_map unicode_to_byte_; + std::array byte_to_unicode_; + std::unordered_map vocab_; + std::vector inv_vocab_; + std::unordered_map, int, pair_hash> merge_ranks_; + }; + +} // namespace llm +} // namespace fdeep diff --git a/include/fdeep/llm/gpt2_cached.hpp b/include/fdeep/llm/gpt2_cached.hpp new file mode 100644 index 00000000..57801c97 --- /dev/null +++ b/include/fdeep/llm/gpt2_cached.hpp @@ -0,0 +1,360 @@ +// Copyright 2026, Tobias Hermann. +// https://github.com/Dobiasd/frugally-deep +// Distributed under the MIT License. +// (See accompanying LICENSE file or at +// https://opensource.org/licenses/MIT) + +// Stateful GPT-2-style inference engine with a key/value cache, intended for +// the chat demo in examples/gpt2_chat. This bypasses the generic frugally-deep +// model runtime and computes the forward pass directly with Eigen so that +// successive ``step()`` calls reuse cached K/V tensors (cost ~constant per +// step rather than O(seq_len^2) for the full re-encode path). +// +// Reads the binary weights file produced by +// keras_export/save_gpt2_weights_bin.py. + +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace fdeep { +namespace llm { + + using ColMatrix = Eigen::Matrix; + using RowVector = Eigen::Matrix; + + namespace internal { + + constexpr std::int32_t GPT2_MAGIC = 0x47505432; + constexpr std::int32_t GPT2_VERSION = 1; + + inline void read_exact(std::ifstream& in, void* dst, std::size_t bytes, + const char* what) + { + in.read(reinterpret_cast(dst), static_cast(bytes)); + if (!in || static_cast(in.gcount()) != bytes) { + throw std::runtime_error(std::string("short read: ") + what); + } + } + + struct gpt2_header { + std::int32_t magic; + std::int32_t version; + std::int32_t num_layers; + std::int32_t hidden_dim; + std::int32_t num_heads; + std::int32_t head_dim; + std::int32_t intermediate_dim; + std::int32_t vocab_size; + std::int32_t max_position; + std::int32_t reserved; + float layer_norm_epsilon; + char preset[64]; + }; + + } // namespace internal + + struct gpt2_block_weights { + // Pre-attention LayerNorm. + std::vector attn_norm_gamma; + std::vector attn_norm_beta; + // Attention projections (Q/K/V flattened to (hidden, hidden)). + ColMatrix Wq, Wk, Wv; // shape (hidden, hidden) + std::vector bq, bk, bv; // length hidden + // Output projection (hidden, hidden). + ColMatrix Wo; + std::vector bo; + // Pre-FFN LayerNorm. + std::vector ffn_norm_gamma; + std::vector ffn_norm_beta; + // FFN. + ColMatrix W1; // (hidden, intermediate) + std::vector b1; + ColMatrix W2; // (intermediate, hidden) + std::vector b2; + }; + + class gpt2_cached_model { + public: + explicit gpt2_cached_model(const std::string& weights_path, + std::size_t max_seq_len = 256) + : max_seq_len_(max_seq_len) + { + std::ifstream in(weights_path, std::ios::binary); + if (!in) { + throw std::runtime_error("could not open " + weights_path); + } + internal::gpt2_header h {}; + internal::read_exact(in, &h, sizeof(h), "header"); + if (h.magic != internal::GPT2_MAGIC) { + throw std::runtime_error("not a gpt2 weights file"); + } + if (h.version != internal::GPT2_VERSION) { + throw std::runtime_error("unsupported weights version"); + } + num_layers_ = static_cast(h.num_layers); + hidden_dim_ = static_cast(h.hidden_dim); + num_heads_ = static_cast(h.num_heads); + head_dim_ = static_cast(h.head_dim); + intermediate_dim_ = static_cast(h.intermediate_dim); + vocab_size_ = static_cast(h.vocab_size); + max_position_ = static_cast(h.max_position); + layer_norm_epsilon_ = h.layer_norm_epsilon; + if (max_seq_len_ > max_position_) { + throw std::runtime_error("max_seq_len exceeds model's max_position"); + } + if (num_heads_ * head_dim_ != hidden_dim_) { + throw std::runtime_error("num_heads * head_dim != hidden_dim"); + } + + token_embedding_ = read_matrix(in, vocab_size_, hidden_dim_); + position_embedding_ = read_matrix(in, max_position_, hidden_dim_); + + blocks_.resize(num_layers_); + for (std::size_t i = 0; i < num_layers_; ++i) { + auto& b = blocks_[i]; + b.attn_norm_gamma = read_vec(in, hidden_dim_); + b.attn_norm_beta = read_vec(in, hidden_dim_); + b.Wq = read_matrix(in, hidden_dim_, hidden_dim_); + b.bq = read_vec(in, hidden_dim_); + b.Wk = read_matrix(in, hidden_dim_, hidden_dim_); + b.bk = read_vec(in, hidden_dim_); + b.Wv = read_matrix(in, hidden_dim_, hidden_dim_); + b.bv = read_vec(in, hidden_dim_); + b.Wo = read_matrix(in, hidden_dim_, hidden_dim_); + b.bo = read_vec(in, hidden_dim_); + b.ffn_norm_gamma = read_vec(in, hidden_dim_); + b.ffn_norm_beta = read_vec(in, hidden_dim_); + b.W1 = read_matrix(in, hidden_dim_, intermediate_dim_); + b.b1 = read_vec(in, intermediate_dim_); + b.W2 = read_matrix(in, intermediate_dim_, hidden_dim_); + b.b2 = read_vec(in, hidden_dim_); + } + + final_norm_gamma_ = read_vec(in, hidden_dim_); + final_norm_beta_ = read_vec(in, hidden_dim_); + lm_head_ = read_matrix(in, hidden_dim_, vocab_size_); + + // Allocate KV cache: (num_layers, max_seq_len, hidden_dim) each. + cache_k_.assign(num_layers_, + ColMatrix::Zero(max_seq_len_, hidden_dim_)); + cache_v_.assign(num_layers_, + ColMatrix::Zero(max_seq_len_, hidden_dim_)); + cur_len_ = 0; + } + + std::size_t vocab_size() const { return vocab_size_; } + std::size_t hidden_dim() const { return hidden_dim_; } + std::size_t num_layers() const { return num_layers_; } + std::size_t cur_len() const { return cur_len_; } + std::size_t max_seq_len() const { return max_seq_len_; } + + void reset() + { + cur_len_ = 0; + } + + // Process the entire prompt, building the cache. Returns logits for + // predicting the token after the last prompt token. + std::vector prefill(const std::vector& prompt_ids) + { + if (prompt_ids.empty()) { + throw std::runtime_error("prompt is empty"); + } + std::vector logits; + for (int id : prompt_ids) { + logits = step(id); + } + return logits; + } + + // Advance by one token. Returns logits for the next position. + std::vector step(int token_id) + { + if (cur_len_ >= max_seq_len_) { + throw std::runtime_error("KV cache full"); + } + if (token_id < 0 || static_cast(token_id) >= vocab_size_) { + throw std::runtime_error("token id out of range"); + } + const std::size_t pos = cur_len_; + // Embed: tok + pos. + RowVector x = token_embedding_.row(token_id) + position_embedding_.row(pos); + + for (std::size_t i = 0; i < num_layers_; ++i) { + forward_block(i, x, pos); + } + + // Final norm. + apply_layer_norm(x, final_norm_gamma_, final_norm_beta_); + + // LM head. + RowVector logits = x * lm_head_; + std::vector out(vocab_size_); + std::memcpy(out.data(), logits.data(), vocab_size_ * sizeof(float)); + + // Commit step. + ++cur_len_; + return out; + } + + private: + static std::vector read_vec(std::ifstream& in, std::size_t n) + { + std::vector v(n); + internal::read_exact(in, v.data(), n * sizeof(float), "vec"); + return v; + } + + static ColMatrix read_matrix(std::ifstream& in, std::size_t rows, std::size_t cols) + { + ColMatrix m(rows, cols); + internal::read_exact(in, m.data(), rows * cols * sizeof(float), "matrix"); + return m; + } + + void apply_layer_norm(RowVector& x, const std::vector& gamma, + const std::vector& beta) const + { + const float mean = x.mean(); + const float var = (x.array() - mean).square().mean(); + const float inv = 1.0f / std::sqrt(var + layer_norm_epsilon_); + for (Eigen::Index j = 0; j < x.size(); ++j) { + x(0, j) = ((x(0, j) - mean) * inv) * gamma[static_cast(j)] + + beta[static_cast(j)]; + } + } + + static float gelu_approx(float v) + { + constexpr float c0 = 0.7978845608028654f; // sqrt(2/pi) + constexpr float c1 = 0.044715f; + const float t = c0 * (v + c1 * v * v * v); + return 0.5f * v * (1.0f + std::tanh(t)); + } + + // One transformer block. Mutates ``x`` in place. ``pos`` is the + // sequence position of the new token (also where the new K/V is + // appended in the cache). + void forward_block(std::size_t layer, RowVector& x, std::size_t pos) + { + const auto& b = blocks_[layer]; + + // Pre-attention LN. + RowVector n = x; + apply_layer_norm(n, b.attn_norm_gamma, b.attn_norm_beta); + + // QKV projections (each (hidden,) = (hidden,) * (hidden,hidden)). + RowVector q = n * b.Wq; + RowVector k = n * b.Wk; + RowVector v = n * b.Wv; + for (std::size_t j = 0; j < hidden_dim_; ++j) { + q(0, static_cast(j)) += b.bq[j]; + k(0, static_cast(j)) += b.bk[j]; + v(0, static_cast(j)) += b.bv[j]; + } + + // Append K, V to cache at position `pos`. + cache_k_[layer].row(static_cast(pos)) = k; + cache_v_[layer].row(static_cast(pos)) = v; + + // Attention: for each head, compute scores against the cached + // keys, softmax, weighted sum of cached values. Causal masking is + // implicit — only positions [0..pos] have meaningful values. + const float scale = 1.0f / std::sqrt(static_cast(head_dim_)); + const std::size_t cur = pos + 1; + std::vector attn_out(hidden_dim_, 0.0f); + std::vector scores(cur); + for (std::size_t h = 0; h < num_heads_; ++h) { + const std::size_t q_off = h * head_dim_; + // Compute scores[t] = q_h . k_cache[t, h, :] * scale + float max_score = -std::numeric_limits::infinity(); + for (std::size_t t = 0; t < cur; ++t) { + float dot = 0.0f; + for (std::size_t d = 0; d < head_dim_; ++d) { + dot += q(0, static_cast(q_off + d)) + * cache_k_[layer](static_cast(t), + static_cast(q_off + d)); + } + const float s = dot * scale; + scores[t] = s; + if (s > max_score) + max_score = s; + } + float total = 0.0f; + for (std::size_t t = 0; t < cur; ++t) { + scores[t] = std::exp(scores[t] - max_score); + total += scores[t]; + } + const float inv_total = 1.0f / total; + for (std::size_t t = 0; t < cur; ++t) { + scores[t] *= inv_total; + } + // Weighted sum of values. + for (std::size_t d = 0; d < head_dim_; ++d) { + float s = 0.0f; + for (std::size_t t = 0; t < cur; ++t) { + s += scores[t] + * cache_v_[layer](static_cast(t), + static_cast(q_off + d)); + } + attn_out[q_off + d] = s; + } + } + // Output projection. + Eigen::Map attn_vec(attn_out.data(), 1, static_cast(hidden_dim_)); + RowVector attn_proj = attn_vec * b.Wo; + for (std::size_t j = 0; j < hidden_dim_; ++j) { + attn_proj(0, static_cast(j)) += b.bo[j]; + } + x += attn_proj; + + // Pre-FFN LN. + n = x; + apply_layer_norm(n, b.ffn_norm_gamma, b.ffn_norm_beta); + + // FFN: Dense(intermediate) -> GELU -> Dense(hidden). + RowVector h1 = n * b.W1; + for (std::size_t j = 0; j < intermediate_dim_; ++j) { + h1(0, static_cast(j)) + = gelu_approx(h1(0, static_cast(j)) + b.b1[j]); + } + RowVector h2 = h1 * b.W2; + for (std::size_t j = 0; j < hidden_dim_; ++j) { + h2(0, static_cast(j)) += b.b2[j]; + } + x += h2; + } + + std::size_t num_layers_ = 0; + std::size_t hidden_dim_ = 0; + std::size_t num_heads_ = 0; + std::size_t head_dim_ = 0; + std::size_t intermediate_dim_ = 0; + std::size_t vocab_size_ = 0; + std::size_t max_position_ = 0; + std::size_t max_seq_len_ = 0; + float layer_norm_epsilon_ = 1e-5f; + + ColMatrix token_embedding_; // (vocab, hidden) + ColMatrix position_embedding_; // (max_position, hidden) + std::vector blocks_; + std::vector final_norm_gamma_, final_norm_beta_; + ColMatrix lm_head_; // (hidden, vocab) + + std::vector cache_k_; // num_layers x (max_seq_len, hidden) + std::vector cache_v_; + std::size_t cur_len_ = 0; + }; + +} // namespace llm +} // namespace fdeep diff --git a/include/fdeep/llm/gpt2_generator.hpp b/include/fdeep/llm/gpt2_generator.hpp new file mode 100644 index 00000000..c9d0a951 --- /dev/null +++ b/include/fdeep/llm/gpt2_generator.hpp @@ -0,0 +1,191 @@ +// Copyright 2026, Tobias Hermann. +// https://github.com/Dobiasd/frugally-deep +// Distributed under the MIT License. +// (See accompanying LICENSE file or at +// https://opensource.org/licenses/MIT) + +// Autoregressive token generation on top of a frugally-deep model exported +// from a GPT-2-style backbone with a tied LM head. The model is expected to +// take two integer inputs (token_ids, position_ids) of shape (seq_len,) and +// produce vocabulary logits of shape (seq_len, vocab_size). +// +// This implementation does NOT use a key/value cache — every step re-runs the +// full forward pass over the (prompt + generated-so-far) tokens, padded out +// to the model's fixed sequence length. For long contexts this is slow. + +#pragma once + +#include + +#include +#include +#include +#include +#include +#include + +namespace fdeep { +namespace llm { + + struct gpt2_generation_params { + std::size_t max_new_tokens = 32; + // Temperature 0 means greedy/argmax. Otherwise scales logits. + float temperature = 0.0f; + // 0 disables top-k filtering. + std::size_t top_k = 0; + // Optional seed for the sampling RNG. + std::uint64_t seed = 0; + // Optional callback invoked with each newly generated token id, just + // after it is appended. Useful for streaming output. Returning ``false`` + // stops generation. + std::function on_token = {}; + }; + + class gpt2_generator { + public: + gpt2_generator(const fdeep::model& model_, std::size_t seq_len, + std::size_t vocab_size, int eos_token_id, std::size_t pad_token_id = 0) + : model_(model_) + , seq_len_(seq_len) + , vocab_size_(vocab_size) + , eos_token_id_(eos_token_id) + , pad_token_id_(pad_token_id) + { + if (seq_len_ == 0) { + throw std::runtime_error("seq_len must be positive"); + } + } + + std::vector generate(const std::vector& prompt_ids, + const gpt2_generation_params& params) const + { + if (prompt_ids.empty()) { + throw std::runtime_error("prompt_ids must not be empty"); + } + if (prompt_ids.size() >= seq_len_) { + throw std::runtime_error( + "prompt is at least as long as the model's seq_len; " + "no room to generate"); + } + + std::mt19937_64 rng(params.seed); + + // Pre-fill the fixed-length token / position buffers. + fdeep::float_vec tokens(seq_len_, static_cast(pad_token_id_)); + fdeep::float_vec positions(seq_len_); + for (std::size_t i = 0; i < seq_len_; ++i) { + positions[i] = static_cast(i); + } + for (std::size_t i = 0; i < prompt_ids.size(); ++i) { + tokens[i] = static_cast(prompt_ids[i]); + } + + std::vector generated; + generated.reserve(params.max_new_tokens); + + for (std::size_t step = 0; step < params.max_new_tokens; ++step) { + const std::size_t cur_len = prompt_ids.size() + step; + if (cur_len >= seq_len_) { + break; // out of room + } + + fdeep::float_vec tokens_copy = tokens; + fdeep::float_vec positions_copy = positions; + const fdeep::tensor token_t(fdeep::tensor_shape(seq_len_), + std::move(tokens_copy)); + const fdeep::tensor position_t(fdeep::tensor_shape(seq_len_), + std::move(positions_copy)); + const auto out = model_.predict({ token_t, position_t }); + if (out.empty()) { + throw std::runtime_error("model produced no output"); + } + const auto& logits = *out.front().as_vector(); + if (logits.size() != seq_len_ * vocab_size_) { + throw std::runtime_error( + "unexpected logit count: got " + + std::to_string(logits.size()) + ", expected " + + std::to_string(seq_len_ * vocab_size_)); + } + + // Read logits at position (cur_len - 1), the last filled slot. + const std::size_t row = cur_len - 1; + const float* row_ptr = logits.data() + row * vocab_size_; + + int next_id = sample_next(row_ptr, params, rng); + generated.push_back(next_id); + tokens[cur_len] = static_cast(next_id); + + if (params.on_token && !params.on_token(next_id)) { + break; + } + if (next_id == eos_token_id_) { + break; + } + } + return generated; + } + + private: + int sample_next(const float* logits, const gpt2_generation_params& p, + std::mt19937_64& rng) const + { + if (p.temperature <= 0.0f) { + // Greedy. + std::size_t best = 0; + float best_val = logits[0]; + for (std::size_t i = 1; i < vocab_size_; ++i) { + if (logits[i] > best_val) { + best_val = logits[i]; + best = i; + } + } + return static_cast(best); + } + + // Build a vector of (logit, idx). For top-k, keep only the top k. + std::vector> candidates; + candidates.reserve(vocab_size_); + for (std::size_t i = 0; i < vocab_size_; ++i) { + candidates.emplace_back(logits[i], static_cast(i)); + } + const float inv_temp = 1.0f / p.temperature; + for (auto& c : candidates) + c.first *= inv_temp; + + if (p.top_k > 0 && p.top_k < candidates.size()) { + std::partial_sort(candidates.begin(), + candidates.begin() + static_cast(p.top_k), + candidates.end(), + [](const auto& a, const auto& b) { return a.first > b.first; }); + candidates.resize(p.top_k); + } + + float m = candidates[0].first; + for (const auto& c : candidates) + m = std::max(m, c.first); + float total = 0.0f; + for (auto& c : candidates) { + c.first = std::exp(c.first - m); + total += c.first; + } + std::uniform_real_distribution dist(0.0f, total); + const float pick = dist(rng); + float acc = 0.0f; + for (const auto& c : candidates) { + acc += c.first; + if (acc >= pick) { + return c.second; + } + } + return candidates.back().second; + } + + const fdeep::model& model_; + std::size_t seq_len_; + std::size_t vocab_size_; + int eos_token_id_; + std::size_t pad_token_id_; + }; + +} // namespace llm +} // namespace fdeep diff --git a/keras_export/save_gpt2_backbone_for_fdeep.py b/keras_export/save_gpt2_backbone_for_fdeep.py new file mode 100644 index 00000000..0b36fc09 --- /dev/null +++ b/keras_export/save_gpt2_backbone_for_fdeep.py @@ -0,0 +1,199 @@ +#!/usr/bin/env python3 +"""Rebuild a keras_hub GPT-2 model as a plain Keras Functional model +using only layer types that frugally-deep supports, copy weights from +the keras_hub model, and save it in .keras format. + +The resulting model takes: + - token_ids: int32 of shape (batch, seq_len) + - position_ids: int32 of shape (batch, seq_len) + +and outputs either: + - the final hidden state of shape (batch, seq_len, hidden_dim) (default), or + - vocabulary logits of shape (batch, seq_len, vocab_size) when ``--with-lm-head`` + is set. The LM head is a tied projection: its kernel is the token-embedding + matrix transposed, matching keras_hub's GPT2CausalLM behaviour. + +Causal attention is enabled via MultiHeadAttention(use_causal_mask=True). +GELU uses the tanh-approximation form (matching keras_hub's GPT-2). The +``gelu_approximate`` helper from convert_model is registered as a +serializable Keras activation; the converter rewrites it to a plain +``gelu`` activation with ``approximate=True`` in the layer config, which +the C++ runtime understands. +""" + +import argparse +import os +import sys + +os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") +os.environ.setdefault("KERAS_BACKEND", "tensorflow") + +import keras +import keras_hub +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from convert_model import gelu_approximate # noqa: E402 (registers serializable) + + +def build_gpt2_like(seq_len: int, vocab_size: int, max_position: int, + hidden_dim: int, num_heads: int, intermediate_dim: int, + num_layers: int, layer_norm_epsilon: float = 1e-5, + with_lm_head: bool = False) -> keras.Model: + """Build a GPT-2-style decoder-only transformer with plain Keras layers.""" + head_dim = hidden_dim // num_heads + + token_ids = keras.Input(shape=(seq_len,), dtype="int32", name="token_ids") + position_ids = keras.Input(shape=(seq_len,), dtype="int32", name="position_ids") + + tok_emb = keras.layers.Embedding(vocab_size, hidden_dim, name="token_embedding")(token_ids) + pos_emb = keras.layers.Embedding(max_position, hidden_dim, name="position_embedding")(position_ids) + x = keras.layers.Add(name="embeddings_add")([tok_emb, pos_emb]) + + for i in range(num_layers): + # Pre-norm self-attention + n = keras.layers.LayerNormalization(epsilon=layer_norm_epsilon, + name=f"block_{i}_attn_norm")(x) + a = keras.layers.MultiHeadAttention( + num_heads=num_heads, key_dim=head_dim, value_dim=head_dim, + use_bias=True, name=f"block_{i}_attn")(n, n, n, use_causal_mask=True) + x = keras.layers.Add(name=f"block_{i}_attn_residual")([x, a]) + + # Pre-norm FFN + n = keras.layers.LayerNormalization(epsilon=layer_norm_epsilon, + name=f"block_{i}_ffn_norm")(x) + n = keras.layers.Dense(intermediate_dim, activation=gelu_approximate, + name=f"block_{i}_ffn_intermediate")(n) + n = keras.layers.Dense(hidden_dim, name=f"block_{i}_ffn_output")(n) + x = keras.layers.Add(name=f"block_{i}_ffn_residual")([x, n]) + + x = keras.layers.LayerNormalization(epsilon=layer_norm_epsilon, name="final_norm")(x) + + if with_lm_head: + # GPT-2 uses a tied LM head: logits = hidden @ E.T, where E is the + # token-embedding matrix. We materialise it as a plain bias-less Dense + # so frugally-deep can run it without extra layer support. + x = keras.layers.Dense(vocab_size, use_bias=False, name="lm_head")(x) + + return keras.Model(inputs=[token_ids, position_ids], outputs=x, name="gpt2_like") + + +def copy_weights_from_keras_hub(src_backbone, dst_model: keras.Model, num_layers: int, + with_lm_head: bool = False) -> None: + """Copy weights from a keras_hub GPT2Backbone into our plain-Keras model.""" + # Embeddings + src_token = src_backbone.get_layer("token_embedding") + embedding_matrix = src_token.embeddings.numpy() + dst_model.get_layer("token_embedding").set_weights([embedding_matrix]) + if with_lm_head: + # Tied weights: kernel is E.T (Dense maps hidden_dim -> vocab_size). + dst_model.get_layer("lm_head").set_weights([embedding_matrix.T]) + + src_pos = src_backbone.get_layer("position_embedding") + # PositionEmbedding stores its lookup table under .position_embeddings + pos_w = src_pos.position_embeddings.numpy() + dst_model.get_layer("position_embedding").set_weights([pos_w]) + + for i in range(num_layers): + src_block = src_backbone.get_layer(f"transformer_layer_{i}") + + # Attention: keras_hub uses Keras MultiHeadAttention internally with + # the same einsum-style weight layout. + src_attn = src_block._self_attention_layer + dst_attn = dst_model.get_layer(f"block_{i}_attn") + dst_attn.set_weights(src_attn.get_weights()) + + # Pre-attention LayerNorm + src_attn_ln = src_block._self_attention_layer_norm + dst_model.get_layer(f"block_{i}_attn_norm").set_weights(src_attn_ln.get_weights()) + + # FFN dense layers + src_ffn1 = src_block._feedforward_intermediate_dense + src_ffn2 = src_block._feedforward_output_dense + dst_model.get_layer(f"block_{i}_ffn_intermediate").set_weights(src_ffn1.get_weights()) + dst_model.get_layer(f"block_{i}_ffn_output").set_weights(src_ffn2.get_weights()) + + # Pre-FFN LayerNorm + src_ffn_ln = src_block._feedforward_layer_norm + dst_model.get_layer(f"block_{i}_ffn_norm").set_weights(src_ffn_ln.get_weights()) + + # Final LayerNorm + dst_model.get_layer("final_norm").set_weights(src_backbone.get_layer("layer_norm").get_weights()) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--preset", default="gpt2_base_en") + parser.add_argument("--seq-len", type=int, default=32) + parser.add_argument("--output", required=True) + parser.add_argument("--with-lm-head", action="store_true", + help="Append a tied LM head so the model outputs " + "vocabulary logits instead of hidden states.") + parser.add_argument("--verify", action="store_true", + help="Run a forward pass through both models and " + "report the max absolute difference.") + args = parser.parse_args() + + print(f"Loading keras_hub preset {args.preset!r}...") + gpt2_lm = keras_hub.models.GPT2CausalLM.from_preset(args.preset) + src = gpt2_lm.backbone + cfg = src.get_config() + + vocab_size = cfg["vocabulary_size"] + hidden_dim = cfg["hidden_dim"] + num_heads = cfg["num_heads"] + intermediate_dim = cfg["intermediate_dim"] + num_layers = cfg["num_layers"] + max_position = cfg["max_sequence_length"] + layer_norm_epsilon = cfg.get("layer_norm_epsilon", 1e-5) + + print(f" vocab_size={vocab_size}, hidden_dim={hidden_dim}, " + f"num_heads={num_heads}, intermediate_dim={intermediate_dim}, " + f"num_layers={num_layers}, max_position={max_position}, " + f"layer_norm_epsilon={layer_norm_epsilon}") + + print(f"Building plain-Keras GPT-2 with seq_len={args.seq_len}, " + f"with_lm_head={args.with_lm_head}...") + dst = build_gpt2_like(seq_len=args.seq_len, vocab_size=vocab_size, + max_position=max_position, hidden_dim=hidden_dim, + num_heads=num_heads, intermediate_dim=intermediate_dim, + num_layers=num_layers, + layer_norm_epsilon=layer_norm_epsilon, + with_lm_head=args.with_lm_head) + print("Copying weights...") + copy_weights_from_keras_hub(src, dst, num_layers, + with_lm_head=args.with_lm_head) + + if args.verify: + print("Verifying numerics against keras_hub reference...") + rng = np.random.default_rng(0) + token_ids = rng.integers(0, vocab_size, size=(1, args.seq_len), dtype=np.int32) + position_ids = np.arange(args.seq_len, dtype=np.int32)[None, :] + padding_mask = np.ones_like(token_ids, dtype=np.int32) + + if args.with_lm_head: + # Compare against full GPT2CausalLM logits. + ref = gpt2_lm({"token_ids": token_ids, "padding_mask": padding_mask}).numpy() + else: + ref = src({"token_ids": token_ids, "padding_mask": padding_mask}).numpy() + ours = dst([token_ids, position_ids]).numpy() + diff = np.max(np.abs(ref - ours)) + print(f" max |ref - ours| = {diff:.3e}") + if args.with_lm_head: + # Logits scale much larger than hidden states (often 50–200), so + # an absolute tolerance scaled by output magnitude is appropriate. + ref_max = float(np.max(np.abs(ref))) + print(f" ref output scale (max |ref|) = {ref_max:.3f}") + tol = max(1e-3 * ref_max, 1e-2) + else: + tol = 1e-3 + assert diff < tol, f"Numerics differ too much: {diff} (tol={tol})" + print(" OK") + + print(f"Saving to {args.output}...") + dst.save(args.output) + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/keras_export/save_gpt2_weights_bin.py b/keras_export/save_gpt2_weights_bin.py new file mode 100644 index 00000000..0da5b4e2 --- /dev/null +++ b/keras_export/save_gpt2_weights_bin.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +"""Dump GPT-2 weights to a flat binary file consumable by the C++ cached +inference engine in include/fdeep/llm/gpt2_cached.hpp. + +The binary layout is intentionally simple: + + Header (104 bytes, little-endian): + int32 magic = 0x47505432 ("GPT2") + int32 version = 1 + int32 num_layers + int32 hidden_dim + int32 num_heads + int32 head_dim + int32 intermediate_dim + int32 vocab_size + int32 max_position + int32 reserved = 0 + float layer_norm_epsilon + char[64] preset_name (NUL-padded) + + Body (all float32, little-endian, contiguous, native row-major layouts): + token_embedding: [vocab_size, hidden_dim] + position_embedding: [max_position, hidden_dim] + for layer in 0..num_layers-1: + attn_norm_gamma: [hidden_dim] + attn_norm_beta: [hidden_dim] + attn_q_kernel: [hidden_dim, num_heads, head_dim] + attn_q_bias: [num_heads, head_dim] + attn_k_kernel: [hidden_dim, num_heads, head_dim] + attn_k_bias: [num_heads, head_dim] + attn_v_kernel: [hidden_dim, num_heads, head_dim] + attn_v_bias: [num_heads, head_dim] + attn_o_kernel: [num_heads, head_dim, hidden_dim] + attn_o_bias: [hidden_dim] + ffn_norm_gamma: [hidden_dim] + ffn_norm_beta: [hidden_dim] + ffn1_kernel: [hidden_dim, intermediate_dim] + ffn1_bias: [intermediate_dim] + ffn2_kernel: [intermediate_dim, hidden_dim] + ffn2_bias: [hidden_dim] + final_norm_gamma: [hidden_dim] + final_norm_beta: [hidden_dim] + lm_head_kernel: [hidden_dim, vocab_size] (tied = token_embedding.T) +""" +import argparse +import os +import struct +import sys + +os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") +os.environ.setdefault("KERAS_BACKEND", "tensorflow") + +import numpy as np +import keras +import keras_hub # noqa: F401 (registers GPT2 layers) + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from convert_model import gelu_approximate # noqa: F401 (registers serializable) +from save_gpt2_backbone_for_fdeep import ( + build_gpt2_like, copy_weights_from_keras_hub, +) + + +MAGIC = 0x47505432 +VERSION = 1 + + +def write_f32(f, arr: np.ndarray) -> None: + a = np.ascontiguousarray(arr, dtype=np.float32) + f.write(a.tobytes()) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--preset", default="gpt2_base_en") + parser.add_argument("--output", required=True, + help="Path to write the binary weights file.") + args = parser.parse_args() + + print(f"Loading keras_hub preset {args.preset!r}...") + gpt2_lm = keras_hub.models.GPT2CausalLM.from_preset(args.preset) + src = gpt2_lm.backbone + cfg = src.get_config() + vocab_size = cfg["vocabulary_size"] + hidden_dim = cfg["hidden_dim"] + num_heads = cfg["num_heads"] + intermediate_dim = cfg["intermediate_dim"] + num_layers = cfg["num_layers"] + max_position = cfg["max_sequence_length"] + layer_norm_epsilon = float(cfg.get("layer_norm_epsilon", 1e-5)) + head_dim = hidden_dim // num_heads + + # Build a small (seq_len=2) plain-Keras copy purely as a vehicle for + # ``copy_weights_from_keras_hub`` -- only the weights are exported. + dst = build_gpt2_like(seq_len=2, vocab_size=vocab_size, + max_position=max_position, hidden_dim=hidden_dim, + num_heads=num_heads, intermediate_dim=intermediate_dim, + num_layers=num_layers, + layer_norm_epsilon=layer_norm_epsilon, + with_lm_head=True) + copy_weights_from_keras_hub(src, dst, num_layers, with_lm_head=True) + print("Weights copied. Writing binary...") + + with open(args.output, "wb") as f: + # Header + preset_bytes = args.preset.encode("utf-8")[:64].ljust(64, b"\0") + f.write(struct.pack( + "<10i f 64s", + MAGIC, VERSION, num_layers, hidden_dim, num_heads, head_dim, + intermediate_dim, vocab_size, max_position, 0, + layer_norm_epsilon, preset_bytes, + )) + # Embeddings + token_w = dst.get_layer("token_embedding").get_weights()[0] + write_f32(f, token_w) + pos_w = dst.get_layer("position_embedding").get_weights()[0] + write_f32(f, pos_w) + + for i in range(num_layers): + attn_norm = dst.get_layer(f"block_{i}_attn_norm").get_weights() + assert len(attn_norm) == 2 + write_f32(f, attn_norm[0]); write_f32(f, attn_norm[1]) + + attn = dst.get_layer(f"block_{i}_attn").get_weights() + # Keras MultiHeadAttention.get_weights returns (per Keras source, + # in the order they were created): + # query_kernel, query_bias, + # key_kernel, key_bias, + # value_kernel, value_bias, + # output_kernel, output_bias + # which matches the layout we declare in the binary. + assert len(attn) == 8 + for w in attn: + write_f32(f, w) + + ffn_norm = dst.get_layer(f"block_{i}_ffn_norm").get_weights() + assert len(ffn_norm) == 2 + write_f32(f, ffn_norm[0]); write_f32(f, ffn_norm[1]) + + ffn1 = dst.get_layer(f"block_{i}_ffn_intermediate").get_weights() + assert len(ffn1) == 2 + write_f32(f, ffn1[0]); write_f32(f, ffn1[1]) + ffn2 = dst.get_layer(f"block_{i}_ffn_output").get_weights() + assert len(ffn2) == 2 + write_f32(f, ffn2[0]); write_f32(f, ffn2[1]) + + fn = dst.get_layer("final_norm").get_weights() + assert len(fn) == 2 + write_f32(f, fn[0]); write_f32(f, fn[1]) + + lm = dst.get_layer("lm_head").get_weights() + assert len(lm) == 1 + write_f32(f, lm[0]) + + size = os.path.getsize(args.output) + print(f"Wrote {size:,} bytes to {args.output}") + + +if __name__ == "__main__": + main()