| #ifndef NEUROFLOW_ALIGNMENT_COMMON_HPP
|
| #define NEUROFLOW_ALIGNMENT_COMMON_HPP
|
|
|
| #include <algorithm>
|
| #include <cmath>
|
| #include <cstring>
|
| #include <fstream>
|
| #include <iostream>
|
| #include <string>
|
| #include <unordered_map>
|
| #include <vector>
|
|
|
| #include "causal_lm.hpp"
|
| #include "tensor.hpp"
|
|
|
| namespace neuroflow {
|
|
|
| struct SFTSample {
|
| std::string instruction;
|
| std::string response;
|
| };
|
|
|
| struct DPOSample {
|
| std::string instruction;
|
| std::string chosen;
|
| std::string rejected;
|
| };
|
|
|
| struct SFTTrainingTensors {
|
| std::vector<size_t> input_ids;
|
| std::vector<size_t> target_ids;
|
| std::vector<float> loss_mask;
|
| size_t instruction_len;
|
| };
|
|
|
| struct DPOTrainingTensors {
|
| std::vector<size_t> chosen_ids;
|
| std::vector<size_t> rejected_ids;
|
| size_t chosen_prompt_len;
|
| size_t rejected_prompt_len;
|
| };
|
|
|
| inline std::string extract_json_string(const std::string& json, const std::string& key) {
|
| std::string search = "\"" + key + "\"";
|
| size_t pos = json.find(search);
|
| if (pos == std::string::npos) return "";
|
| pos = json.find(':', pos + search.size());
|
| if (pos == std::string::npos) return "";
|
| pos++;
|
| while (pos < json.size() && (json[pos] == ' ' || json[pos] == '\t')) pos++;
|
| if (pos >= json.size() || json[pos] != '"') return "";
|
| size_t end = pos + 1;
|
| while (end < json.size()) {
|
| if (json[end] == '"' && json[end - 1] != '\\') break;
|
| if (json[end - 1] == '\\') { end++; continue; }
|
| end++;
|
| }
|
| if (end >= json.size()) return "";
|
| return json.substr(pos + 1, end - pos - 1);
|
| }
|
|
|
| inline size_t extract_json_number(const std::string& json, const std::string& key, size_t default_val = 0) {
|
| std::string search = "\"" + key + "\"";
|
| size_t pos = json.find(search);
|
| if (pos == std::string::npos) return default_val;
|
| pos = json.find(':', pos + search.size());
|
| if (pos == std::string::npos) return default_val;
|
| pos++;
|
| while (pos < json.size() && (json[pos] == ' ' || json[pos] == '\t')) pos++;
|
| size_t end = pos;
|
| while (end < json.size() && (json[end] >= '0' && json[end] <= '9')) end++;
|
| if (end == pos) return default_val;
|
| return std::stoul(json.substr(pos, end - pos));
|
| }
|
|
|
| inline bool extract_json_bool(const std::string& json, const std::string& key, bool default_val = false) {
|
| std::string search = "\"" + key + "\"";
|
| size_t pos = json.find(search);
|
| if (pos == std::string::npos) return default_val;
|
| pos = json.find(':', pos + search.size());
|
| if (pos == std::string::npos) return default_val;
|
| pos++;
|
| while (pos < json.size() && (json[pos] == ' ' || json[pos] == '\t')) pos++;
|
| if (pos + 3 < json.size() && json.substr(pos, 4) == "true") return true;
|
| if (pos + 4 < json.size() && json.substr(pos, 5) == "false") return false;
|
| return default_val;
|
| }
|
|
|
| inline void unescape_json(std::string& s) {
|
| size_t pos = 0;
|
| while ((pos = s.find('\\', pos)) != std::string::npos) {
|
| if (pos + 1 < s.size()) {
|
| char c = s[pos + 1];
|
| if (c == 'n') { s.replace(pos, 2, "\n"); pos++; }
|
| else if (c == 't') { s.replace(pos, 2, "\t"); pos++; }
|
| else if (c == 'r') { s.replace(pos, 2, "\r"); pos++; }
|
| else if (c == '"') { s.replace(pos, 2, "\""); pos++; }
|
| else if (c == '\\') { s.replace(pos, 2, "\\"); pos++; }
|
| else if (c == '/') { s.replace(pos, 2, "/"); pos++; }
|
| else if (c == 'u' && pos + 5 < s.size()) { pos += 6; }
|
| else { pos += 2; }
|
| } else { pos++; }
|
| }
|
| }
|
|
|
| inline bool validate_path(const std::string& path) {
|
| if (path.find("..") != std::string::npos) return false;
|
| return true;
|
| }
|
|
|
| inline void save_lm_checkpoint(const CausalLMHead& model, const std::string& path) {
|
| CausalLMHead& lm_head = const_cast<CausalLMHead&>(model);
|
| #ifdef USE_CUDA
|
| auto sync_to_cpu = [](const Tensor& t) {
|
| if (t.is_on_gpu()) { const_cast<Tensor&>(t).to_cpu(); }
|
| };
|
| sync_to_cpu(lm_head.w_embed_);
|
| sync_to_cpu(lm_head.w_pos_);
|
| sync_to_cpu(lm_head.dw_kernel_);
|
| sync_to_cpu(lm_head.pw_conv_->weight);
|
| sync_to_cpu(lm_head.sae_w_encode_->weight);
|
| sync_to_cpu(lm_head.sae_w_decode_->weight);
|
| sync_to_cpu(lm_head.ntm_w_read_->weight);
|
| sync_to_cpu(lm_head.ntm_w_write_->weight);
|
| sync_to_cpu(lm_head.ntm_w_erase_->weight);
|
| sync_to_cpu(lm_head.ntm_memory_);
|
| sync_to_cpu(lm_head.w_proj_->weight);
|
| sync_to_cpu(lm_head.w_proj_->bias);
|
| if (lm_head.bridge_) {
|
| sync_to_cpu(lm_head.bridge_->weight);
|
| sync_to_cpu(lm_head.bridge_->bias);
|
| }
|
| sync_to_cpu(lm_head.w_out_->weight);
|
| sync_to_cpu(lm_head.w_out_->bias);
|
| sync_to_cpu(lm_head.ln_->weight);
|
| sync_to_cpu(lm_head.ln_->bias);
|
| for (auto& attn : lm_head.attn_layers_) {
|
| sync_to_cpu(attn->w_q->weight);
|
| sync_to_cpu(attn->w_q->bias);
|
| sync_to_cpu(attn->w_k->weight);
|
| sync_to_cpu(attn->w_k->bias);
|
| sync_to_cpu(attn->w_v->weight);
|
| sync_to_cpu(attn->w_v->bias);
|
| sync_to_cpu(attn->w_out->weight);
|
| sync_to_cpu(attn->w_out->bias);
|
| sync_to_cpu(attn->norm->weight);
|
| sync_to_cpu(attn->norm->bias);
|
| }
|
| #endif
|
| auto sl = [](std::ofstream& o, const std::string& n, const Tensor& t) {
|
| uint32_t nl = n.size(); o.write((char*)&nl, 4); o.write(n.data(), nl);
|
| uint32_t nd = t.shape_.size(); o.write((char*)&nd, 4);
|
| for (auto d : t.shape_) { uint32_t dd = d; o.write((char*)&dd, 4); }
|
| uint32_t ds = t.data_size_; o.write((char*)&ds, 4);
|
| o.write((char*)t.data_.get(), ds);
|
| };
|
| std::ofstream o(path, std::ios::binary);
|
| o.write("LMH2", 4);
|
| sl(o, "w_embed", lm_head.w_embed_);
|
| sl(o, "w_pos", lm_head.w_pos_);
|
| sl(o, "dw_kernel", lm_head.dw_kernel_);
|
| sl(o, "pw_conv.weight", lm_head.pw_conv_->weight);
|
| sl(o, "sae_encode.weight", lm_head.sae_w_encode_->weight);
|
| sl(o, "sae_decode.weight", lm_head.sae_w_decode_->weight);
|
| sl(o, "ntm_read.weight", lm_head.ntm_w_read_->weight);
|
| sl(o, "ntm_write.weight", lm_head.ntm_w_write_->weight);
|
| sl(o, "ntm_erase.weight", lm_head.ntm_w_erase_->weight);
|
| sl(o, "ntm_memory", lm_head.ntm_memory_);
|
| sl(o, "w_proj.weight", lm_head.w_proj_->weight);
|
| sl(o, "w_proj.bias", lm_head.w_proj_->bias);
|
| if (lm_head.bridge_) {
|
| sl(o, "bridge.weight", lm_head.bridge_->weight);
|
| sl(o, "bridge.bias", lm_head.bridge_->bias);
|
| }
|
| sl(o, "w_out.weight", lm_head.w_out_->weight);
|
| if (lm_head.w_out_->bias.data_) sl(o, "w_out.bias", lm_head.w_out_->bias);
|
| sl(o, "ln.weight", lm_head.ln_->weight);
|
| sl(o, "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) + ".";
|
| sl(o, p + "w_q.weight", lm_head.attn_layers_[i]->w_q->weight);
|
| sl(o, p + "w_q.bias", lm_head.attn_layers_[i]->w_q->bias);
|
| sl(o, p + "w_k.weight", lm_head.attn_layers_[i]->w_k->weight);
|
| sl(o, p + "w_k.bias", lm_head.attn_layers_[i]->w_k->bias);
|
| sl(o, p + "w_v.weight", lm_head.attn_layers_[i]->w_v->weight);
|
| sl(o, p + "w_v.bias", lm_head.attn_layers_[i]->w_v->bias);
|
| sl(o, p + "w_out.weight", lm_head.attn_layers_[i]->w_out->weight);
|
| sl(o, p + "w_out.bias", lm_head.attn_layers_[i]->w_out->bias);
|
| sl(o, p + "norm.weight", lm_head.attn_layers_[i]->norm->weight);
|
| sl(o, p + "norm.bias", lm_head.attn_layers_[i]->norm->bias);
|
| }
|
| uint32_t z = 0; o.write((char*)&z, 4); o.close();
|
| }
|
|
|
| inline bool load_lm_checkpoint(CausalLMHead& lm_head, const std::string& lm_path) {
|
| std::ifstream ifs(lm_path, std::ios::binary);
|
| if (!ifs) {
|
| std::cerr << "LM Head checkpoint未找到: " << lm_path << std::endl;
|
| return false;
|
| }
|
|
|
| char magic[5] = {0};
|
| ifs.read(magic, 4);
|
| if (std::string(magic) != "LMH2" && std::string(magic) != "LMH1") {
|
| std::cerr << "LM Head格式不匹配: " << std::string(magic) << " (期望LMH2/LMH1)" << std::endl;
|
| ifs.close();
|
| return false;
|
| }
|
|
|
| auto read_named_tensor = [&]() -> std::pair<std::string, Tensor> {
|
| uint32_t nl = 0; ifs.read((char*)&nl, 4);
|
| if (nl == 0 || ifs.eof()) return {"", Tensor()};
|
| std::string name(nl, '\0'); ifs.read(&name[0], nl);
|
| uint32_t nd = 0; ifs.read((char*)&nd, 4);
|
| 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);
|
| t = Tensor();
|
| }
|
| return {name, std::move(t)};
|
| };
|
|
|
| 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;
|
| if (lm_head.bridge_) {
|
| tensor_map["bridge.weight"] = &lm_head.bridge_->weight;
|
| tensor_map["bridge.bias"] = &lm_head.bridge_->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) {
|
| auto [name, tensor] = read_named_tensor();
|
| if (name.empty()) break;
|
| auto it = tensor_map.find(name);
|
| if (it != tensor_map.end() && tensor.numel() > 0) {
|
| if (it->second->shape_ == tensor.shape_) {
|
| memcpy(it->second->data_.get(), tensor.data_.get(), tensor.data_size_);
|
| loaded++;
|
| } else {
|
| std::cerr << " 跳过 '" << name << "': 形状不匹配" << std::endl;
|
| }
|
| } 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" && tensor.shape_.size() == 2 && tensor.shape_[0] == 3 * d_model) {
|
| const float* src = tensor.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++;
|
| std::cerr << " 拆分旧格式 '" << name << "' -> w_q/w_k/w_v" << std::endl;
|
| } else if (suffix == "bias" && tensor.shape_[0] == 3 * d_model) {
|
| const float* src = tensor.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++;
|
| std::cerr << " 拆分旧格式 '" << name << "' -> w_q/w_k/w_v bias" << std::endl;
|
| }
|
| }
|
| }
|
| }
|
|
|
| ifs.close();
|
| if (lm_head.config_.weight_tying) lm_head.tie_weights();
|
| std::cerr << "LM Head已恢复: " << loaded << " 个张量从 " << lm_path << std::endl;
|
| return loaded > 0;
|
| }
|
|
|
| }
|
|
|
| #endif |
