StandardOne-3B-SH-GGUF / runner /sh_states.cpp
MyeongHoJeong's picture
Standard One 3B SH v1
eb6c5ec verified
Raw History Blame Contribute Delete
6.76 kB
// sh-states: dump the hidden states a Standard One schema head reads (final norm output + one decoder layer output)
// from a llama.cpp GGUF model, for already-tokenized inputs.
//
// sh-states -m model.gguf --tap 25 [-t threads] [-c n_ctx] [-b chunk]
//
// Protocol (line based, stdin -> stdout): each request line is
// <ids.bin> <out_prefix>
// where ids.bin holds little-endian int32 token ids (BOS included, exactly as the renderer produced them). For every
// request the program writes <out_prefix>.final.f32 and <out_prefix>.tap.f32 (row-major float32 [n_tokens, n_embd])
// and answers "ok <n_tokens> <n_embd>" (or "err <message>"). The KV cache is cleared between requests.
//
// The final state is llama.cpp's per-token embedding output (embeddings=true, pooling NONE: the final RMSNorm output =
// HF last_hidden_state). The tap is taken with the scheduler eval callback on "l_out-<tap>" (= the output of decoder
// layer <tap>, as a forward hook on layers[tap] sees it). The prompt is decoded in chunks with every token marked as
// an output, so both cover all tokens.
#include "llama.h"
#include "ggml.h"
#include "ggml-backend.h"
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>
#include <vector>
struct capture {
std::string tap_name;
std::vector<float> final_rows, tap_rows;
int64_t n_embd = 0;
};
static bool eval_cb(struct ggml_tensor * t, bool ask, void * user_data) {
auto * cap = (capture *) user_data;
const bool want = cap->tap_name == t->name;
if (ask) {
return want;
}
if (!want) {
return true;
}
if (t->type != GGML_TYPE_F32 || ggml_n_dims(t) > 2) {
fprintf(stderr, "unexpected tensor %s type %d dims %d\n", t->name, (int) t->type, ggml_n_dims(t));
return false;
}
const int64_t n_embd = t->ne[0], n_tok = t->ne[1];
cap->n_embd = n_embd;
std::vector<float> buf(n_embd * n_tok);
if (t->nb[1] == (size_t) n_embd * sizeof(float)) {
ggml_backend_tensor_get(t, buf.data(), 0, buf.size() * sizeof(float));
} else { // non-contiguous rows
for (int64_t r = 0; r < n_tok; r++) {
ggml_backend_tensor_get(t, buf.data() + r * n_embd, r * t->nb[1], n_embd * sizeof(float));
}
}
cap->tap_rows.insert(cap->tap_rows.end(), buf.begin(), buf.end());
return true;
}
static bool write_f32(const std::string & path, const std::vector<float> & v) {
std::ofstream f(path, std::ios::binary);
f.write((const char *) v.data(), v.size() * sizeof(float));
return (bool) f;
}
int main(int argc, char ** argv) {
std::string model_path;
int tap = 25, threads = 8, n_ctx = 34000, chunk = 512;
for (int i = 1; i < argc; i++) {
std::string a = argv[i];
auto next = [&]() { return std::string(i + 1 < argc ? argv[++i] : ""); };
if (a == "-m") model_path = next();
else if (a == "--tap") tap = std::stoi(next());
else if (a == "-t") threads = std::stoi(next());
else if (a == "-c") n_ctx = std::stoi(next());
else if (a == "-b") chunk = std::stoi(next());
else { fprintf(stderr, "unknown argument %s\n", a.c_str()); return 2; }
}
if (model_path.empty()) { fprintf(stderr, "usage: sh-states -m model.gguf --tap 25\n"); return 2; }
llama_backend_init();
llama_model_params mp = llama_model_default_params();
llama_model * model = llama_model_load_from_file(model_path.c_str(), mp);
if (!model) { fprintf(stderr, "failed to load %s\n", model_path.c_str()); return 1; }
if (tap < 0 || tap >= llama_model_n_layer(model)) { fprintf(stderr, "tap out of range\n"); return 2; }
capture cap;
cap.tap_name = "l_out-" + std::to_string(tap);
llama_context_params cp = llama_context_default_params();
cp.n_ctx = n_ctx;
cp.n_batch = chunk;
cp.n_ubatch = chunk;
cp.n_threads = threads;
cp.n_threads_batch = threads;
cp.embeddings = true;
cp.pooling_type = LLAMA_POOLING_TYPE_NONE;
cp.cb_eval = eval_cb;
cp.cb_eval_user_data = &cap;
llama_context * ctx = llama_init_from_model(model, cp);
if (!ctx) { fprintf(stderr, "failed to create context\n"); return 1; }
printf("ready n_layer=%d n_embd=%d tap=%d\n", llama_model_n_layer(model), llama_model_n_embd(model), tap);
fflush(stdout);
std::string line;
while (std::getline(std::cin, line)) {
std::istringstream ss(line);
std::string ids_path, out_prefix;
if (!(ss >> ids_path >> out_prefix)) { printf("err bad request line\n"); fflush(stdout); continue; }
std::ifstream f(ids_path, std::ios::binary | std::ios::ate);
if (!f) { printf("err cannot read %s\n", ids_path.c_str()); fflush(stdout); continue; }
const size_t bytes = f.tellg();
f.seekg(0);
std::vector<llama_token> ids(bytes / sizeof(int32_t));
f.read((char *) ids.data(), bytes);
const int n = (int) ids.size();
if (n == 0 || n > n_ctx) { printf("err %d tokens (context %d)\n", n, n_ctx); fflush(stdout); continue; }
llama_memory_clear(llama_get_memory(ctx), true);
cap.final_rows.clear();
cap.tap_rows.clear();
bool ok = true;
llama_batch batch = llama_batch_init(chunk, 0, 1);
for (int s = 0; s < n && ok; s += chunk) {
const int m = std::min(chunk, n - s);
batch.n_tokens = m;
for (int j = 0; j < m; j++) {
batch.token[j] = ids[s + j];
batch.pos[j] = s + j;
batch.n_seq_id[j] = 1;
batch.seq_id[j][0] = 0;
batch.logits[j] = 1;
}
if (llama_decode(ctx, batch) != 0) { ok = false; break; }
const float * emb = llama_get_embeddings(ctx);
if (!emb) { ok = false; break; }
const int64_t h = llama_model_n_embd(model);
cap.final_rows.insert(cap.final_rows.end(), emb, emb + (size_t) m * h);
}
llama_batch_free(batch);
const size_t want = (size_t) n * cap.n_embd;
if (!ok || cap.final_rows.size() != want || cap.tap_rows.size() != want) {
printf("err decode=%d final=%zu tap=%zu want=%zu\n", ok, cap.final_rows.size(), cap.tap_rows.size(), want);
fflush(stdout);
continue;
}
if (!write_f32(out_prefix + ".final.f32", cap.final_rows) || !write_f32(out_prefix + ".tap.f32", cap.tap_rows)) {
printf("err write failed\n");
} else {
printf("ok %d %lld\n", n, (long long) cap.n_embd);
}
fflush(stdout);
}
llama_free(ctx);
llama_model_free(model);
llama_backend_free();
return 0;
}