| #ifndef NOMINMAX
|
| #define NOMINMAX
|
| #endif
|
| #include <iostream>
|
| #include <string>
|
| #include <vector>
|
| #include <fstream>
|
| #include "neuroflow/generative.hpp"
|
| #include "neuroflow/model.hpp"
|
|
|
| using namespace neuroflow;
|
|
|
| int main(int argc, char* argv[]) {
|
| std::string config_path = "configs/config_distill_small.json";
|
| std::string tokenizer_path = "configs/tokenizer_cn_013.json";
|
| std::string lm_head_path = "";
|
| int max_new_tokens = 64;
|
| float temperature = 0.8f;
|
| int top_k = 40;
|
| float top_p = 0.9f;
|
| float repetition_penalty = 1.0f;
|
| bool use_cuda = false;
|
| std::string strategy_name = "top_k";
|
| int yarn_max_seq_len = -1;
|
|
|
| for (int i = 1; i < argc; ++i) {
|
| std::string arg = argv[i];
|
| if (arg == "--config" && i + 1 < argc) config_path = argv[++i];
|
| else if (arg == "--tokenizer" && i + 1 < argc) tokenizer_path = argv[++i];
|
| else if (arg == "--lm-head" && i + 1 < argc) lm_head_path = argv[++i];
|
| else if (arg == "--max-tokens" && i + 1 < argc) max_new_tokens = std::atoi(argv[++i]);
|
| else if (arg == "--temperature" && i + 1 < argc) temperature = std::atof(argv[++i]);
|
| else if (arg == "--top-k" && i + 1 < argc) top_k = std::atoi(argv[++i]);
|
| else if (arg == "--top-p" && i + 1 < argc) top_p = std::atof(argv[++i]);
|
| else if (arg == "--repetition-penalty" && i + 1 < argc) repetition_penalty = std::atof(argv[++i]);
|
| else if (arg == "--strategy" && i + 1 < argc) strategy_name = argv[++i];
|
| else if (arg == "--max-seq-len" && i + 1 < argc) yarn_max_seq_len = std::atoi(argv[++i]);
|
| else if (arg == "--use-cuda") use_cuda = true;
|
| }
|
|
|
| if (lm_head_path.empty()) {
|
| std::cerr << "用法: neuroflow_infer --lm-head <path> [--config <path>] [--tokenizer <path>]"
|
| << " [--max-tokens <N>] [--temperature <f>] [--top-k <N>] [--use-cuda]" << std::endl;
|
| return 1;
|
| }
|
|
|
| NeuroFlowModel::Config model_cfg;
|
| {
|
| std::ifstream jf(config_path);
|
| if (jf) {
|
| std::string json_str((std::istreambuf_iterator<char>(jf)), std::istreambuf_iterator<char>());
|
| auto ex_num = [&](const std::string& key, size_t def) -> size_t {
|
| std::string pat = "\"" + key + "\"";
|
| size_t pos = json_str.find(pat);
|
| if (pos == std::string::npos) return def;
|
| pos = json_str.find(':', pos + pat.size());
|
| if (pos == std::string::npos) return def;
|
| pos++;
|
| while (pos < json_str.size() && (json_str[pos] == ' ' || json_str[pos] == '\t')) pos++;
|
| return std::stoull(json_str.substr(pos));
|
| };
|
| auto ex_str = [&](const std::string& key, const std::string& def) -> std::string {
|
| std::string pat = "\"" + key + "\"";
|
| size_t pos = json_str.find(pat);
|
| if (pos == std::string::npos) return def;
|
| pos = json_str.find(':', pos + pat.size());
|
| if (pos == std::string::npos) return def;
|
| size_t q1 = json_str.find('"', pos);
|
| if (q1 == std::string::npos) return def;
|
| size_t q2 = json_str.find('"', q1 + 1);
|
| if (q2 == std::string::npos) return def;
|
| return json_str.substr(q1 + 1, q2 - q1 - 1);
|
| };
|
| model_cfg.input_dim = ex_num("d_model", model_cfg.input_dim);
|
| model_cfg.hidden_dim = ex_num("hidden_dim", model_cfg.hidden_dim);
|
| model_cfg.output_dim = ex_num("output_dim", model_cfg.output_dim);
|
| model_cfg.vocab_size = ex_num("vocab_size", model_cfg.vocab_size);
|
| model_cfg.max_seq_len = ex_num("max_seq_len", model_cfg.max_seq_len);
|
| model_cfg.causal_window_size = ex_num("causal_window_size", model_cfg.causal_window_size);
|
| model_cfg.sae_k = ex_num("sae_k", model_cfg.sae_k);
|
| model_cfg.ntm_memory_slots = ex_num("ntm_memory_slots", model_cfg.ntm_memory_slots);
|
| model_cfg.lm_num_attn_layers = ex_num("lm_num_attn_layers", model_cfg.lm_num_attn_layers);
|
| model_cfg.lm_pooling = ex_str("lm_pooling", "mean");
|
| }
|
| }
|
|
|
| CausalLMConfig lm_cfg;
|
| lm_cfg.vocab_size = model_cfg.vocab_size;
|
| lm_cfg.d_model = model_cfg.hidden_dim;
|
| lm_cfg.max_seq_len = model_cfg.max_seq_len;
|
| lm_cfg.causal_window_size = model_cfg.causal_window_size;
|
| lm_cfg.sae_k = model_cfg.sae_k;
|
| lm_cfg.ntm_memory_slots = model_cfg.ntm_memory_slots;
|
| lm_cfg.use_mla = model_cfg.use_mla;
|
| lm_cfg.mla_latent_dim = model_cfg.mla_latent_dim;
|
| lm_cfg.mla_n_heads = model_cfg.mla_n_heads;
|
| lm_cfg.mla_max_cache_len = 4096;
|
| lm_cfg.weight_tying = true;
|
| lm_cfg.num_attn_layers = model_cfg.lm_num_attn_layers;
|
| lm_cfg.num_attn_heads = 4;
|
| lm_cfg.pooling = model_cfg.lm_pooling;
|
|
|
| CausalLMHead lm_head(lm_cfg);
|
| if (lm_cfg.weight_tying) lm_head.tie_weights();
|
|
|
| std::cerr << "加载 LM Head: " << lm_head_path << std::endl;
|
| {
|
| std::ifstream ifs(lm_head_path, std::ios::binary);
|
| if (!ifs) {
|
| std::cerr << "错误: 无法打开 " << lm_head_path << std::endl;
|
| return 1;
|
| }
|
| char magic[5] = {0};
|
| ifs.read(magic, 4);
|
| if (std::string(magic) != "LMH2" && std::string(magic) != "LMH1") {
|
| std::cerr << "错误: 格式不匹配 " << std::string(magic) << std::endl;
|
| return 1;
|
| }
|
|
|
| std::unordered_map<std::string, Tensor*> tensor_map;
|
| tensor_map["w_embed"] = &lm_head.w_embed_;
|
| tensor_map["w_pos"] = &lm_head.w_pos_;
|
| tensor_map["dw_kernel"] = &lm_head.dw_kernel_;
|
| tensor_map["pw_conv.weight"] = &lm_head.pw_conv_->weight;
|
| tensor_map["sae_encode.weight"] = &lm_head.sae_w_encode_->weight;
|
| tensor_map["sae_decode.weight"] = &lm_head.sae_w_decode_->weight;
|
| tensor_map["ntm_read.weight"] = &lm_head.ntm_w_read_->weight;
|
| tensor_map["ntm_write.weight"] = &lm_head.ntm_w_write_->weight;
|
| tensor_map["ntm_erase.weight"] = &lm_head.ntm_w_erase_->weight;
|
| tensor_map["ntm_memory"] = &lm_head.ntm_memory_;
|
| tensor_map["w_proj.weight"] = &lm_head.w_proj_->weight;
|
| tensor_map["w_proj.bias"] = &lm_head.w_proj_->bias;
|
| tensor_map["w_out.weight"] = &lm_head.w_out_->weight;
|
| tensor_map["w_out.bias"] = &lm_head.w_out_->bias;
|
| tensor_map["ln.weight"] = &lm_head.ln_->weight;
|
| tensor_map["ln.bias"] = &lm_head.ln_->bias;
|
| for (size_t i = 0; i < lm_head.attn_layers_.size(); ++i) {
|
| std::string p = "attn" + std::to_string(i) + ".";
|
| tensor_map[p + "w_q.weight"] = &lm_head.attn_layers_[i]->w_q->weight;
|
| tensor_map[p + "w_q.bias"] = &lm_head.attn_layers_[i]->w_q->bias;
|
| tensor_map[p + "w_k.weight"] = &lm_head.attn_layers_[i]->w_k->weight;
|
| tensor_map[p + "w_k.bias"] = &lm_head.attn_layers_[i]->w_k->bias;
|
| tensor_map[p + "w_v.weight"] = &lm_head.attn_layers_[i]->w_v->weight;
|
| tensor_map[p + "w_v.bias"] = &lm_head.attn_layers_[i]->w_v->bias;
|
| tensor_map[p + "w_out.weight"] = &lm_head.attn_layers_[i]->w_out->weight;
|
| tensor_map[p + "w_out.bias"] = &lm_head.attn_layers_[i]->w_out->bias;
|
| tensor_map[p + "norm.weight"] = &lm_head.attn_layers_[i]->norm->weight;
|
| tensor_map[p + "norm.bias"] = &lm_head.attn_layers_[i]->norm->bias;
|
| }
|
|
|
| size_t loaded = 0;
|
| while (ifs) {
|
| uint32_t nl = 0;
|
| ifs.read((char*)&nl, 4);
|
| if (nl == 0 || ifs.eof()) break;
|
| if (nl > 256) { std::cerr << "异常名称长度: " << nl << ", 可能文件损坏" << std::endl; break; }
|
| std::string name(nl, '\0'); ifs.read(&name[0], nl);
|
| uint32_t nd = 0; ifs.read((char*)&nd, 4);
|
| if (nd > 10) { std::cerr << "异常维度: " << nd << ", name=" << name << std::endl; break; }
|
| std::vector<size_t> shape(nd);
|
| for (size_t i = 0; i < nd; ++i) { uint32_t d = 0; ifs.read((char*)&d, 4); shape[i] = d; }
|
| uint32_t ds = 0; ifs.read((char*)&ds, 4);
|
| Tensor t(shape, QuantType::FP32);
|
| if (t.data_size_ == ds) {
|
| ifs.read((char*)t.data_.get(), ds);
|
| } else {
|
| ifs.seekg(ds, std::ios::cur);
|
| continue;
|
| }
|
| auto it = tensor_map.find(name);
|
| if (it != tensor_map.end() && it->second->shape_ == t.shape_) {
|
| memcpy(it->second->data_.get(), t.data_.get(), t.data_size_);
|
| loaded++;
|
| } else if (name.find("w_qkv.") != std::string::npos) {
|
| std::string base = name.substr(0, name.find("w_qkv."));
|
| std::string suffix = name.substr(name.find("w_qkv.") + 6);
|
| for (size_t i = 0; i < lm_head.attn_layers_.size(); ++i) {
|
| std::string p = "attn" + std::to_string(i) + ".";
|
| if (base != p) continue;
|
| size_t d_model = lm_head.config_.d_model;
|
| size_t head_dim = d_model / lm_head.config_.num_attn_heads;
|
| size_t n_q = lm_head.attn_layers_[i]->n_q_heads_;
|
| size_t n_kv = lm_head.attn_layers_[i]->n_kv_heads_;
|
| if (suffix == "weight" && t.shape_.size() == 2 && t.shape_[0] == 3 * d_model) {
|
| const float* src = t.as_fp32();
|
| float* dq = lm_head.attn_layers_[i]->w_q->weight.as_fp32();
|
| float* dk = lm_head.attn_layers_[i]->w_k->weight.as_fp32();
|
| float* dv = lm_head.attn_layers_[i]->w_v->weight.as_fp32();
|
| for (size_t r = 0; r < d_model; ++r) {
|
| memcpy(dq + r * n_q * head_dim, src + r * 3 * d_model, n_q * head_dim * sizeof(float));
|
| memcpy(dk + r * n_kv * head_dim, src + r * 3 * d_model + d_model, n_kv * head_dim * sizeof(float));
|
| memcpy(dv + r * n_kv * head_dim, src + r * 3 * d_model + 2 * d_model, n_kv * head_dim * sizeof(float));
|
| }
|
| loaded++;
|
| } else if (suffix == "bias" && t.shape_[0] == 3 * d_model) {
|
| const float* src = t.as_fp32();
|
| float* bq = lm_head.attn_layers_[i]->w_q->bias.as_fp32();
|
| float* bk = lm_head.attn_layers_[i]->w_k->bias.as_fp32();
|
| float* bv = lm_head.attn_layers_[i]->w_v->bias.as_fp32();
|
| memcpy(bq, src, n_q * head_dim * sizeof(float));
|
| memcpy(bk, src + d_model, n_kv * head_dim * sizeof(float));
|
| memcpy(bv, src + 2 * d_model, n_kv * head_dim * sizeof(float));
|
| loaded++;
|
| }
|
| }
|
| }
|
| }
|
| ifs.close();
|
| if (lm_cfg.weight_tying) lm_head.tie_weights();
|
| std::cerr << "已加载 " << loaded << " 个张量" << std::endl;
|
| }
|
|
|
| BPETokenizer tokenizer(tokenizer_path);
|
| std::cerr << "词表: " << tokenizer.vocab_size() << " tokens" << std::endl;
|
|
|
| #ifdef USE_CUDA
|
| bool cuda_active = false;
|
| if (use_cuda) {
|
| if (CudaContext::instance().initialize(0)) {
|
| cuda_active = true;
|
| std::cerr << "GPU后端: 已启用" << std::endl;
|
|
|
| lm_head.w_embed_.to_gpu();
|
| lm_head.w_pos_.to_gpu();
|
| lm_head.dw_kernel_.to_gpu();
|
| lm_head.pw_conv_->weight.to_gpu();
|
| lm_head.sae_w_encode_->weight.to_gpu();
|
| lm_head.sae_w_decode_->weight.to_gpu();
|
| lm_head.ntm_w_read_->weight.to_gpu();
|
| lm_head.ntm_w_write_->weight.to_gpu();
|
| lm_head.ntm_w_erase_->weight.to_gpu();
|
| lm_head.ntm_memory_.to_gpu();
|
| lm_head.w_proj_->weight.to_gpu();
|
| lm_head.w_proj_->bias.to_gpu();
|
| lm_head.w_out_->weight.to_gpu();
|
| if (lm_head.w_out_->bias.data_) lm_head.w_out_->bias.to_gpu();
|
| lm_head.ln_->weight.to_gpu();
|
| lm_head.ln_->bias.to_gpu();
|
| for (auto& attn : lm_head.attn_layers_) {
|
| attn->w_q->weight.to_gpu();
|
| attn->w_q->bias.to_gpu();
|
| attn->w_k->weight.to_gpu();
|
| attn->w_k->bias.to_gpu();
|
| attn->w_v->weight.to_gpu();
|
| attn->w_v->bias.to_gpu();
|
| attn->w_out->weight.to_gpu();
|
| attn->w_out->bias.to_gpu();
|
| attn->norm->weight.to_gpu();
|
| attn->norm->bias.to_gpu();
|
| }
|
| CudaContext::instance().synchronize();
|
| size_t free_mem = CudaContext::instance().free_memory();
|
| size_t total_mem = CudaContext::instance().total_memory();
|
| std::cerr << "GPU显存: " << (total_mem - free_mem) / (1024*1024)
|
| << " MB 已用 / " << total_mem / (1024*1024) << " MB 总计" << std::endl;
|
| } else {
|
| std::cerr << "[CUDA WARNING] GPU初始化失败，回退CPU后端" << std::endl;
|
| use_cuda = false;
|
| }
|
| } else {
|
| std::cerr << "GPU后端: 未启用 (使用--use-cuda启用)" << std::endl;
|
| }
|
| #else
|
| if (use_cuda) {
|
| std::cerr << "[CUDA WARNING] 编译时未启用CUDA支持 (NEUROFLOW_USE_CUDA=OFF)，回退CPU后端" << std::endl;
|
| use_cuda = false;
|
| }
|
| #endif
|
|
|
| lm_head.eval();
|
|
|
| if (yarn_max_seq_len > 0 && static_cast<size_t>(yarn_max_seq_len) > lm_cfg.max_seq_len) {
|
| float scale_factor = static_cast<float>(yarn_max_seq_len) / static_cast<float>(lm_cfg.max_seq_len);
|
| lm_head.set_yarn_scale(scale_factor);
|
| }
|
|
|
| std::cerr << "\n=== NeuroFlow 推理模式 ===" << std::endl;
|
| std::cerr << "输入提示（Ctrl+C退出）:" << std::endl;
|
|
|
| std::string line;
|
| while (std::getline(std::cin, line)) {
|
| if (line.empty()) continue;
|
| if (line == "quit" || line == "exit") break;
|
|
|
| auto token_ids = tokenizer.encode(line, lm_cfg.max_seq_len);
|
| if (token_ids.empty()) {
|
| std::cerr << "(空token序列)" << std::endl;
|
| continue;
|
| }
|
|
|
| lm_head.clear_cache();
|
|
|
| std::vector<size_t> prefix(token_ids.begin(), token_ids.end() - 1);
|
| Tensor logits;
|
| if (prefix.size() > 0) {
|
| logits = lm_head.forward(prefix);
|
| #ifdef USE_CUDA
|
| if (cuda_active && logits.is_on_gpu()) {
|
| logits.to_cpu();
|
| }
|
| #endif
|
| }
|
|
|
| std::mt19937 rng(42);
|
| std::vector<size_t> generated;
|
| size_t last_id = token_ids.back();
|
|
|
| GenerateConfig gen_cfg;
|
| gen_cfg.max_new_tokens = static_cast<size_t>(max_new_tokens);
|
| gen_cfg.temperature = temperature;
|
| gen_cfg.top_k = static_cast<size_t>(top_k);
|
| gen_cfg.top_p = top_p;
|
| gen_cfg.repetition_penalty = repetition_penalty;
|
| gen_cfg.eos_id = 0;
|
|
|
| std::unique_ptr<SamplingStrategy> sampler;
|
| if (strategy_name == "greedy") {
|
| sampler = std::make_unique<GreedyDecoding>();
|
| } else if (strategy_name == "top_p") {
|
| sampler = std::make_unique<TopPSampling>();
|
| } else if (strategy_name == "top_k_top_p") {
|
| sampler = std::make_unique<TopKTopPSampling>();
|
| } else {
|
| sampler = std::make_unique<TopKSampling>();
|
| }
|
|
|
| for (int step = 0; step < max_new_tokens; ++step) {
|
| size_t pos = token_ids.size() - 1 + step;
|
| logits = lm_head.forward_step(last_id, pos);
|
|
|
| #ifdef USE_CUDA
|
| if (cuda_active && logits.is_on_gpu()) {
|
| int* d_sampled_token = nullptr;
|
| int h_sampled_token = 0;
|
| cudaError_t alloc_err = cudaMalloc(reinterpret_cast<void**>(&d_sampled_token), sizeof(int));
|
|
|
| if (alloc_err == cudaSuccess) {
|
| unsigned int seed = static_cast<unsigned int>(step * 7919 + 42);
|
| bool ok = launch_topk_topp_sampling(
|
| logits.as_gpu_fp32(), d_sampled_token,
|
| static_cast<int>(lm_cfg.vocab_size),
|
| top_k, top_p, temperature, seed,
|
| CudaContext::instance().stream());
|
|
|
| if (ok) {
|
| CudaContext::instance().synchronize();
|
| cudaMemcpy(&h_sampled_token, d_sampled_token, sizeof(int), cudaMemcpyDeviceToHost);
|
| cudaFree(d_sampled_token);
|
|
|
| size_t chosen = static_cast<size_t>(h_sampled_token);
|
| generated.push_back(chosen);
|
| last_id = chosen;
|
|
|
| if (chosen == 0 || chosen == 1) break;
|
| continue;
|
| } else {
|
| cudaFree(d_sampled_token);
|
| std::cerr << "[INFER WARNING] GPU sampling failed, falling back to CPU sampling" << std::endl;
|
| }
|
| } else {
|
| std::cerr << "[INFER WARNING] cudaMalloc failed: "
|
| << cudaGetErrorString(alloc_err)
|
| << ", falling back to CPU sampling" << std::endl;
|
| }
|
|
|
| logits.to_cpu();
|
| }
|
| #endif
|
|
|
| Tensor probs = sampler->apply(std::move(logits), gen_cfg, generated);
|
| size_t chosen = sampler->sample(probs, rng);
|
|
|
| generated.push_back(chosen);
|
| last_id = chosen;
|
|
|
| if (chosen == 0 || chosen == 1) break;
|
| }
|
|
|
| std::string output = tokenizer.decode(generated);
|
| std::cout << output << std::endl;
|
| }
|
|
|
| return 0;
|
| } |
