| #ifndef NOMINMAX
|
| #define NOMINMAX
|
| #endif
|
|
|
|
|
| #ifndef _WIN32
|
| #include <fcntl.h>
|
| #include <sys/mman.h>
|
| #include <sys/types.h>
|
| #include <unistd.h>
|
| #else
|
| #include <windows.h>
|
| #endif
|
|
|
|
|
| #include <algorithm>
|
| #include <chrono>
|
| #include <cmath>
|
| #include <cstdio>
|
| #include <cstdlib>
|
| #include <cstring>
|
| #include <filesystem>
|
| #include <fstream>
|
| #include <iostream>
|
| #include <numeric>
|
| #include <sstream>
|
| #include <string>
|
| #include <vector>
|
|
|
|
|
| #ifdef _OPENMP
|
| #include <omp.h>
|
| #endif
|
|
|
|
|
| #include "neuroflow/backprop.hpp"
|
| #include "neuroflow/generative.hpp"
|
| #include "neuroflow/model.hpp"
|
| #include "neuroflow/online_learning.hpp"
|
| #include "weight_io.hpp"
|
|
|
| using neuroflow::BPETokenizer;
|
| using neuroflow::CausalLMConfig;
|
| using neuroflow::CausalLMHead;
|
| using neuroflow::FullBackpropEngine;
|
| using neuroflow::InitStrategy;
|
| using neuroflow::NeuroFlowModel;
|
| using neuroflow::QuantType;
|
| using neuroflow::Tensor;
|
| using neuroflow::WeightInitializer;
|
|
|
| namespace {
|
|
|
| Tensor lm_linear_backward_input(const Tensor& output_grad, const Tensor& weight) {
|
| size_t batch = output_grad.shape_[0];
|
| size_t out_f = output_grad.shape_[1];
|
| size_t in_f = weight.shape_[1];
|
| Tensor input_grad({batch, in_f}, QuantType::FP32);
|
|
|
| #ifdef USE_CUDA
|
| if (CudaContext::instance().is_available() && output_grad.is_on_gpu()) {
|
| input_grad.to_gpu();
|
| CudaContext::instance().sgemm_rowmajor(false, false,
|
| static_cast<int>(batch), static_cast<int>(in_f), static_cast<int>(out_f),
|
| 1.0f, output_grad.as_gpu_fp32(), static_cast<int>(out_f),
|
| weight.as_gpu_fp32(), static_cast<int>(in_f),
|
| 0.0f, input_grad.as_gpu_fp32(), static_cast<int>(in_f));
|
| input_grad.gpu_dirty_ = true;
|
| return input_grad;
|
| }
|
| #endif
|
|
|
| memset(input_grad.as_fp32(), 0, input_grad.data_size_);
|
|
|
| if (out_f > 8192 && batch == 1) {
|
| const float* og = output_grad.as_fp32();
|
| const float* w = weight.as_fp32();
|
| float* ig = input_grad.as_fp32();
|
| float threshold = 1e-4f;
|
| size_t n_active = 0;
|
| for (size_t i = 0; i < out_f; ++i) {
|
| if (std::abs(og[i]) >= threshold) n_active++;
|
| }
|
| if (n_active < out_f / 4) {
|
| for (size_t i = 0; i < out_f; ++i) {
|
| float g = og[i];
|
| if (std::abs(g) < threshold) continue;
|
| const float* w_row = w + i * in_f;
|
| for (size_t j = 0; j < in_f; ++j) {
|
| ig[j] += g * w_row[j];
|
| }
|
| }
|
| return input_grad;
|
| }
|
| }
|
|
|
| #ifdef USE_CBLAS
|
| cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
|
| batch, in_f, out_f,
|
| 1.0f, output_grad.as_fp32(), out_f, weight.as_fp32(), in_f,
|
| 0.0f, input_grad.as_fp32(), in_f);
|
| #else
|
| const float* og2 = output_grad.as_fp32();
|
| const float* w2 = weight.as_fp32();
|
| float* ig2 = input_grad.as_fp32();
|
| for (size_t b = 0; b < batch; ++b) {
|
| for (size_t j = 0; j < in_f; ++j) {
|
| float sum = 0.0f;
|
| for (size_t i = 0; i < out_f; ++i) {
|
| sum += og2[b * out_f + i] * w2[i * in_f + j];
|
| }
|
| ig2[b * in_f + j] = sum;
|
| }
|
| }
|
| #endif
|
| return input_grad;
|
| }
|
|
|
| Tensor lm_linear_backward_weight(const Tensor& input, const Tensor& output_grad) {
|
| size_t batch = input.shape_[0];
|
| size_t in_f = input.shape_[1];
|
| size_t out_f = output_grad.shape_[1];
|
| Tensor weight_grad({out_f, in_f}, QuantType::FP32);
|
| #ifdef USE_CUDA
|
| if (CudaContext::instance().is_available() && output_grad.is_on_gpu()) {
|
| weight_grad.to_gpu();
|
| CudaContext::instance().sgemm_rowmajor(true, false,
|
| static_cast<int>(out_f), static_cast<int>(in_f), static_cast<int>(batch),
|
| 1.0f / batch, output_grad.as_gpu_fp32(), static_cast<int>(out_f),
|
| input.as_gpu_fp32(), static_cast<int>(in_f),
|
| 0.0f, weight_grad.as_gpu_fp32(), static_cast<int>(in_f));
|
| weight_grad.gpu_dirty_ = true;
|
| return weight_grad;
|
| }
|
| #endif
|
| #ifdef USE_CBLAS
|
| cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans,
|
| out_f, in_f, batch,
|
| 1.0f / batch, output_grad.as_fp32(), out_f, input.as_fp32(), in_f,
|
| 0.0f, weight_grad.as_fp32(), in_f);
|
| #else
|
| const float* inp = input.as_fp32();
|
| const float* og = output_grad.as_fp32();
|
| float* wg = weight_grad.as_fp32();
|
| for (size_t i = 0; i < out_f; ++i) {
|
| for (size_t j = 0; j < in_f; ++j) {
|
| float sum = 0.0f;
|
| for (size_t b = 0; b < batch; ++b) {
|
| sum += og[b * out_f + i] * inp[b * in_f + j];
|
| }
|
| wg[i * in_f + j] = sum / batch;
|
| }
|
| }
|
| #endif
|
| return weight_grad;
|
| }
|
|
|
| }
|
|
|
| struct TrainConfig {
|
| std::string config_path;
|
| std::string tokenizer_path;
|
| std::string data_path;
|
| std::string output_dir = "./output";
|
| int epochs = 10;
|
| float learning_rate = 1e-4f;
|
| uint32_t seed = 42;
|
| std::string resume_path = "";
|
| std::string init_strategy = "xavier";
|
| int batch_size = 32;
|
| int log_interval = 50;
|
| int save_interval = 5000;
|
| float grad_clip = 1.0f;
|
| bool use_adam = false;
|
| int grad_accum_steps = 4;
|
| std::string vocab_mask_path = "";
|
| int replay_buffer_size = 10000;
|
| float replay_ratio = 0.25f;
|
| bool use_cuda = false;
|
| bool verbose = false;
|
| };
|
|
|
| TrainConfig parse_args(int argc, char* argv[]) {
|
| TrainConfig cfg;
|
| for (int i = 1; i < argc; ++i) {
|
| std::string arg = argv[i];
|
| if (arg == "--config" && i + 1 < argc) cfg.config_path = argv[++i];
|
| else if (arg == "--tokenizer" && i + 1 < argc) cfg.tokenizer_path = argv[++i];
|
| else if (arg == "--data" && i + 1 < argc) cfg.data_path = argv[++i];
|
| else if (arg == "--output" && i + 1 < argc) cfg.output_dir = argv[++i];
|
| else if (arg == "--epochs" && i + 1 < argc) cfg.epochs = std::atoi(argv[++i]);
|
| else if (arg == "--lr" && i + 1 < argc) cfg.learning_rate = std::atof(argv[++i]);
|
| else if (arg == "--seed" && i + 1 < argc) cfg.seed = static_cast<uint32_t>(std::atoi(argv[++i]));
|
| else if (arg == "--resume" && i + 1 < argc) cfg.resume_path = argv[++i];
|
| else if (arg == "--init-weights" && i + 1 < argc) cfg.init_strategy = argv[++i];
|
| else if (arg == "--batch-size" && i + 1 < argc) cfg.batch_size = std::atoi(argv[++i]);
|
| else if (arg == "--log-interval" && i + 1 < argc) cfg.log_interval = std::atoi(argv[++i]);
|
| else if (arg == "--save-interval" && i + 1 < argc) cfg.save_interval = std::atoi(argv[++i]);
|
| else if (arg == "--grad-clip" && i + 1 < argc) cfg.grad_clip = std::atof(argv[++i]);
|
| else if (arg == "--adam") cfg.use_adam = true;
|
| else if (arg == "--grad-accum" && i + 1 < argc) cfg.grad_accum_steps = std::atoi(argv[++i]);
|
| else if (arg == "--vocab-mask" && i + 1 < argc) cfg.vocab_mask_path = argv[++i];
|
| else if (arg == "--replay-buffer" && i + 1 < argc) cfg.replay_buffer_size = std::atoi(argv[++i]);
|
| else if (arg == "--replay-ratio" && i + 1 < argc) cfg.replay_ratio = std::atof(argv[++i]);
|
| else if (arg == "--use-cuda") cfg.use_cuda = true;
|
| else if (arg == "--verbose") cfg.verbose = true;
|
| }
|
| return cfg;
|
| }
|
|
|
| namespace {
|
|
|
| std::string read_file(const std::string& path) {
|
| std::ifstream ifs(path);
|
| if (!ifs) return "";
|
| return std::string((std::istreambuf_iterator<char>(ifs)), std::istreambuf_iterator<char>());
|
| }
|
|
|
| std::string trim(const std::string& s) {
|
| size_t start = s.find_first_not_of(" \t\r\n");
|
| if (start == std::string::npos) return "";
|
| size_t end = s.find_last_not_of(" \t\r\n");
|
| return s.substr(start, end - start + 1);
|
| }
|
|
|
| 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 = json.find('"', pos + 1);
|
| if (end == std::string::npos) return "";
|
| return json.substr(pos + 1, end - pos - 1);
|
| }
|
|
|
| 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));
|
| }
|
|
|
| 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;
|
| }
|
|
|
| }
|
|
|
| NeuroFlowModel::Config load_model_config(const std::string& json_path) {
|
| NeuroFlowModel::Config cfg;
|
| std::string json = read_file(json_path);
|
| if (json.empty()) {
|
| std::cerr << "配置加载: " << json_path << " (使用默认配置)" << std::endl;
|
| return cfg;
|
| }
|
|
|
| cfg.input_dim = extract_json_number(json, "d_model", cfg.input_dim);
|
| cfg.hidden_dim = extract_json_number(json, "hidden_dim", cfg.hidden_dim);
|
| cfg.output_dim = extract_json_number(json, "output_dim", cfg.output_dim);
|
| cfg.memory_dim = extract_json_number(json, "memory_dim", cfg.memory_dim);
|
| cfg.memory_slots = extract_json_number(json, "memory_slots", cfg.memory_slots);
|
| cfg.num_layers = extract_json_number(json, "num_layers", cfg.num_layers);
|
| cfg.num_associations = extract_json_number(json, "num_associations", cfg.num_associations);
|
| cfg.use_quantization = extract_json_bool(json, "use_quantization", cfg.use_quantization);
|
| cfg.use_mla = extract_json_bool(json, "use_mla", cfg.use_mla);
|
| cfg.use_causal_lm = extract_json_bool(json, "use_causal_lm", cfg.use_causal_lm);
|
| cfg.mla_latent_dim = extract_json_number(json, "mla_latent_dim", cfg.mla_latent_dim);
|
| cfg.vocab_size = extract_json_number(json, "vocab_size", cfg.vocab_size);
|
| cfg.max_seq_len = extract_json_number(json, "max_seq_len", cfg.max_seq_len);
|
| cfg.causal_window_size = extract_json_number(json, "causal_window_size", cfg.causal_window_size);
|
| cfg.lm_num_attn_layers = extract_json_number(json, "lm_num_attn_layers", cfg.lm_num_attn_layers);
|
| cfg.lm_num_attn_heads = extract_json_number(json, "lm_num_attn_heads", cfg.lm_num_attn_heads);
|
| cfg.lm_n_kv_heads = extract_json_number(json, "lm_n_kv_heads", cfg.lm_n_kv_heads);
|
| cfg.lm_use_rope = extract_json_bool(json, "lm_use_rope", cfg.lm_use_rope);
|
| cfg.lm_use_qk_norm = extract_json_bool(json, "lm_use_qk_norm", cfg.lm_use_qk_norm);
|
| cfg.lm_use_swiglu = extract_json_bool(json, "lm_use_swiglu", cfg.lm_use_swiglu);
|
| cfg.lm_use_bridge = extract_json_bool(json, "lm_use_bridge", cfg.lm_use_bridge);
|
| cfg.lm_pooling = extract_json_string(json, "lm_pooling");
|
| if (cfg.lm_pooling.empty()) cfg.lm_pooling = "last";
|
|
|
| std::cerr << "配置加载: " << json_path << std::endl;
|
| std::cerr << " d_model=" << cfg.input_dim << " hidden_dim=" << cfg.hidden_dim
|
| << " output_dim=" << cfg.output_dim << " vocab_size=" << cfg.vocab_size << std::endl;
|
| return cfg;
|
| }
|
|
|
| struct ResumeState {
|
| bool found = false;
|
| size_t step = 0;
|
| int epoch = 0;
|
| float loss = 0.0f;
|
| float lr = 0.0f;
|
| size_t config_d_model = 0;
|
| size_t config_output_dim = 0;
|
| };
|
|
|
| ResumeState load_checkpoint(NeuroFlowModel& model, const std::string& ckpt_path) {
|
| ResumeState state;
|
|
|
| std::string model_path, state_path;
|
| if (ckpt_path.size() >= 5 && ckpt_path.substr(ckpt_path.size() - 5) == ".nfv1") {
|
| model_path = ckpt_path;
|
| state_path = ckpt_path.substr(0, ckpt_path.size() - 5) + "/../training_state.json";
|
| std::string dir = ckpt_path.substr(0, ckpt_path.find_last_of("/\\"));
|
| state_path = dir + "/training_state.json";
|
| } else {
|
| model_path = ckpt_path + "/model.nfv1";
|
| state_path = ckpt_path + "/training_state.json";
|
| }
|
|
|
| std::ifstream ifs(model_path, std::ios::binary);
|
| if (!ifs) return state;
|
| char magic[5] = {0};
|
| ifs.read(magic, 4);
|
| if (std::string(magic) != "NFv1") return state;
|
| ifs.close();
|
|
|
| std::ifstream state_ifs(state_path);
|
| if (state_ifs) {
|
| std::string content((std::istreambuf_iterator<char>(state_ifs)),
|
| std::istreambuf_iterator<char>());
|
| state.step = extract_json_number(content, "step", 0);
|
| state.epoch = static_cast<int>(extract_json_number(content, "epoch", 0));
|
| state.loss = std::atof(extract_json_string(content, "loss").c_str());
|
| state.lr = std::atof(extract_json_string(content, "learning_rate").c_str());
|
| state.config_d_model = extract_json_number(content, "config_d_model", 0);
|
| state.config_output_dim = extract_json_number(content, "config_output_dim", 0);
|
| }
|
|
|
| if (state.config_d_model > 0 && state.config_d_model != model.config.input_dim) {
|
| std::cerr << "警告: checkpoint的d_model=" << state.config_d_model
|
| << " 与当前config的d_model=" << model.config.input_dim
|
| << " 不匹配，可能导致维度错误！" << std::endl;
|
| }
|
| if (state.config_output_dim > 0 && state.config_output_dim != model.config.output_dim) {
|
| std::cerr << "警告: checkpoint的output_dim=" << state.config_output_dim
|
| << " 与当前config的output_dim=" << model.config.output_dim
|
| << " 不匹配！" << std::endl;
|
| }
|
|
|
| model.load(model_path);
|
| state.found = true;
|
| return state;
|
| }
|
|
|
| bool load_lm_head(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;
|
| }
|
|
|
| NeuroFlowModel build_model(const NeuroFlowModel::Config& cfg, const TrainConfig& train_cfg,
|
| ResumeState& resume_state) {
|
| NeuroFlowModel model(cfg);
|
| if (!train_cfg.resume_path.empty()) {
|
| resume_state = load_checkpoint(model, train_cfg.resume_path);
|
| if (resume_state.found) {
|
| std::cerr << "从checkpoint恢复: " << train_cfg.resume_path
|
| << " (step=" << resume_state.step
|
| << ", epoch=" << resume_state.epoch
|
| << ", loss=" << resume_state.loss << ")" << std::endl;
|
| } else {
|
| throw std::runtime_error("无法加载checkpoint: " + train_cfg.resume_path);
|
| }
|
| } else {
|
| InitStrategy strategy = InitStrategy::XAVIER_UNIFORM;
|
| if (train_cfg.init_strategy == "kaiming") strategy = InitStrategy::KAIMING_NORMAL;
|
| else if (train_cfg.init_strategy == "zeros") strategy = InitStrategy::ZEROS;
|
| WeightInitializer::init_model_weights(model, strategy, train_cfg.seed);
|
| std::cerr << "权重初始化: " << train_cfg.init_strategy << std::endl;
|
| }
|
| return model;
|
| }
|
|
|
| static 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++; }
|
| }
|
| }
|
|
|
| struct LMSample {
|
| std::vector<size_t> token_ids;
|
| };
|
|
|
| struct TextSpan {
|
| size_t file_index;
|
| size_t offset;
|
| size_t length;
|
| };
|
|
|
| struct TrainingSample {
|
| std::vector<float> input;
|
| std::vector<float> target;
|
| };
|
|
|
| class StreamingDataLoader {
|
| public:
|
| StreamingDataLoader(const std::string& path, int batch_size,
|
| BPETokenizer* tokenizer, size_t max_seq_len = 128,
|
| size_t max_samples = 500000)
|
| : batch_size_(batch_size), tokenizer_(tokenizer),
|
| max_seq_len_(max_seq_len), max_samples_(max_samples),
|
| current_idx_(0) {
|
| index_files(path);
|
| std::cerr << "StreamingDataLoader: " << spans_.size() << " 文本段已索引"
|
| << " (" << file_paths_.size() << " 文件)" << std::endl;
|
| }
|
|
|
| void index_files(const std::string& path) {
|
| namespace fs = std::filesystem;
|
| fs::path p(path);
|
|
|
| if (fs::is_directory(p)) {
|
| std::vector<fs::path> files;
|
| for (auto& entry : fs::recursive_directory_iterator(p, fs::directory_options::follow_directory_symlink)) {
|
| if (!entry.is_regular_file()) continue;
|
| auto ext = entry.path().extension().string();
|
| if (ext == ".txt" || ext == ".json" || ext == ".jsonl" || ext == ".csv" || ext == ".tsv") {
|
| files.push_back(entry.path());
|
| }
|
| }
|
| std::sort(files.begin(), files.end());
|
| for (auto& f : files) index_single_file(f.string());
|
| } else {
|
| index_single_file(p.string());
|
| }
|
| }
|
|
|
| void index_single_file(const std::string& file_path) {
|
| namespace fs = std::filesystem;
|
| std::ifstream ifs(file_path, std::ios::binary | std::ios::ate);
|
| if (!ifs) return;
|
|
|
| size_t file_size = static_cast<size_t>(ifs.tellg());
|
| ifs.seekg(0, std::ios::beg);
|
|
|
| std::string ext = fs::path(file_path).extension().string();
|
| size_t file_idx = file_paths_.size();
|
| file_paths_.push_back(file_path);
|
|
|
| if (ext == ".jsonl" || ext == ".json") {
|
| std::string line;
|
| size_t line_offset = 0;
|
| while (std::getline(ifs, line) && spans_.size() < max_samples_) {
|
| if (line.size() < 12) { line_offset += line.size() + 1; continue; }
|
| std::string text = extract_text_from_line(line);
|
| if (text.size() >= 10) {
|
| text_cache_.push_back(std::move(text));
|
| spans_.push_back(TextSpan{SIZE_MAX, text_cache_.size() - 1, text_cache_.back().size()});
|
| }
|
| line_offset += line.size() + 1;
|
| }
|
| } else if (ext == ".txt") {
|
| std::string line;
|
| while (std::getline(ifs, line) && spans_.size() < max_samples_) {
|
| if (line.size() < 10) continue;
|
| text_cache_.push_back(std::move(line));
|
| spans_.push_back(TextSpan{SIZE_MAX, text_cache_.size() - 1, text_cache_.back().size()});
|
| }
|
| } else if (ext == ".csv" || ext == ".tsv") {
|
| char delim = (ext == ".tsv") ? '\t' : ',';
|
| std::string line;
|
| while (std::getline(ifs, line) && spans_.size() < max_samples_) {
|
| size_t pos = line.find(delim);
|
| if (pos == std::string::npos) continue;
|
| std::string text = line.substr(pos + 1);
|
| if (text.size() >= 10) {
|
| text_cache_.push_back(std::move(text));
|
| spans_.push_back(TextSpan{SIZE_MAX, text_cache_.size() - 1, text_cache_.back().size()});
|
| }
|
| }
|
| }
|
| }
|
|
|
| std::string extract_text_from_line(const std::string& line) {
|
| const char* fields[] = {"\"text\"", "\"content\"", "\"title\"", "\"question\"", "\"answer\""};
|
| for (auto& field : fields) {
|
| size_t pos = line.find(field);
|
| if (pos == std::string::npos) continue;
|
| size_t colon = line.find(':', pos + strlen(field));
|
| if (colon == std::string::npos) continue;
|
| size_t q1 = line.find('"', colon);
|
| if (q1 == std::string::npos) continue;
|
| size_t q2 = line.find('"', q1 + 1);
|
| if (q2 == std::string::npos) continue;
|
| std::string text = line.substr(q1 + 1, q2 - q1 - 1);
|
| unescape_json(text);
|
| return text;
|
| }
|
| return "";
|
| }
|
|
|
| LMSample get(size_t idx) const {
|
| if (idx >= spans_.size()) return {};
|
| const auto& span = spans_[idx];
|
| const std::string& text = text_cache_[span.offset];
|
| auto ids = tokenizer_->encode(text, max_seq_len_);
|
| return LMSample{std::move(ids)};
|
| }
|
|
|
| bool has_next() const { return current_idx_ < spans_.size(); }
|
|
|
| std::vector<LMSample> next_lm_batch_raw() {
|
| std::vector<LMSample> batch;
|
| for (int i = 0; i < batch_size_ && current_idx_ < spans_.size(); ++i) {
|
| batch.push_back(get(current_idx_++));
|
| }
|
| return batch;
|
| }
|
|
|
| void reset() { current_idx_ = 0; }
|
|
|
| void shuffle(std::mt19937& rng) {
|
| std::shuffle(spans_.begin(), spans_.end(), rng);
|
| }
|
|
|
| size_t total_samples() const { return spans_.size(); }
|
|
|
| size_t memory_usage_bytes() const {
|
| size_t total = spans_.capacity() * sizeof(TextSpan);
|
| for (auto& t : text_cache_) total += t.capacity();
|
| return total;
|
| }
|
|
|
| private:
|
| int batch_size_;
|
| BPETokenizer* tokenizer_;
|
| size_t max_seq_len_;
|
| size_t max_samples_;
|
| size_t current_idx_;
|
| std::vector<std::string> file_paths_;
|
| std::vector<TextSpan> spans_;
|
| std::vector<std::string> text_cache_;
|
| };
|
|
|
| class DataLoader {
|
| public:
|
| enum class Mode { LM, CLASSIFICATION, REGRESSION };
|
|
|
| DataLoader(const std::string& path, int batch_size, size_t input_dim, size_t output_dim,
|
| Mode mode = Mode::LM, BPETokenizer* tokenizer = nullptr,
|
| size_t max_seq_len = 128, size_t max_samples = 500000,
|
| size_t max_memory_mb = 4096)
|
| : path_(path), batch_size_(batch_size), input_dim_(input_dim), output_dim_(output_dim),
|
| current_idx_(0), mode_(mode), tokenizer_(tokenizer), max_seq_len_(max_seq_len),
|
| max_samples_(max_samples), max_memory_bytes_(max_memory_mb * 1024ULL * 1024ULL) {
|
| load_data();
|
| }
|
|
|
| void load_data() {
|
| if (path_.empty()) {
|
| throw std::runtime_error("DataLoader: 数据路径为空，请指定--data参数");
|
| }
|
|
|
| namespace fs = std::filesystem;
|
| fs::path p(path_);
|
|
|
| if (!fs::exists(p)) {
|
| throw std::runtime_error("DataLoader: 路径不存在: " + path_);
|
| }
|
|
|
| if (fs::is_directory(p)) {
|
| load_directory(p.string());
|
| } else {
|
| load_file(p.string());
|
| }
|
|
|
| if (lm_samples_.empty() && samples_.empty()) {
|
| throw std::runtime_error(
|
| "DataLoader: 未能从 " + path_ + " 加载任何训练数据。"
|
| "请检查文件格式（支持txt/json/jsonl/csv/tsv）和内容是否有效。");
|
| }
|
|
|
| if (!lm_samples_.empty()) {
|
| std::cerr << "DataLoader: " << lm_samples_.size() << " LM样本已加载" << std::endl;
|
| } else {
|
| std::cerr << "DataLoader: " << samples_.size() << " 数值样本已加载" << std::endl;
|
| }
|
| }
|
|
|
| void load_directory(const std::string& dir_path) {
|
| namespace fs = std::filesystem;
|
| size_t file_count = 0;
|
| size_t skipped = 0;
|
|
|
| std::cerr << "DataLoader: 扫描目录 " << dir_path << " ..." << std::endl;
|
| std::vector<fs::path> files;
|
| for (auto& entry : fs::recursive_directory_iterator(dir_path, fs::directory_options::follow_directory_symlink)) {
|
| if (!entry.is_regular_file()) continue;
|
| files.push_back(entry.path());
|
| }
|
| std::cerr << "DataLoader: 发现 " << files.size() << " 个文件" << std::endl;
|
| std::sort(files.begin(), files.end());
|
|
|
| for (size_t fi = 0; fi < files.size(); ++fi) {
|
| auto& fpath = files[fi];
|
| std::string ext = fpath.extension().string();
|
| std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
|
|
|
| if (ext == ".json" || ext == ".jsonl" || ext == ".txt" ||
|
| ext == ".csv" || ext == ".tsv" || ext == ".tok1" || ext == ".bin") {
|
| size_t file_size_mb = 0;
|
| try { file_size_mb = fs::file_size(fpath) / 1024 / 1024; } catch (...) {}
|
| std::cerr << "[" << (fi + 1) << "/" << files.size() << "] 加载: "
|
| << fpath.filename().string()
|
| << " (" << file_size_mb << " MB)" << std::endl;
|
|
|
| auto load_start = std::chrono::steady_clock::now();
|
| try {
|
| load_file(fpath.string());
|
| file_count++;
|
| } catch (const std::exception& e) {
|
| std::cerr << " 跳过: " << e.what() << std::endl;
|
| skipped++;
|
| }
|
| auto load_end = std::chrono::steady_clock::now();
|
| double load_sec = std::chrono::duration<double>(load_end - load_start).count();
|
|
|
| size_t total = lm_samples_.size() + samples_.size();
|
| std::cerr << " 完成: " << load_sec << "s, 累计 "
|
| << total << " 样本" << std::endl;
|
|
|
| if (total >= max_samples_) {
|
| std::cerr << "DataLoader: 已达采样上限 " << max_samples_
|
| << "，停止遍历" << std::endl;
|
| break;
|
| }
|
| } else if (ext == ".parquet") {
|
| skipped++;
|
| }
|
| }
|
| std::cerr << "DataLoader: 遍历 " << file_count << " 个文件, 跳过 " << skipped << std::endl;
|
| }
|
|
|
| void load_file(const std::string& file_path) {
|
| namespace fs = std::filesystem;
|
| std::string ext = fs::path(file_path).extension().string();
|
| std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
|
|
|
| {
|
| std::ifstream ifs(file_path, std::ios::binary);
|
| if (ifs) {
|
| char magic[5] = {0};
|
| ifs.read(magic, 4);
|
| ifs.close();
|
| if (std::string(magic) == "TOK1") {
|
| load_tok1(file_path);
|
| return;
|
| }
|
| }
|
| }
|
|
|
| if (mode_ == Mode::LM && tokenizer_) {
|
| if (ext == ".json") {
|
| size_t file_size = fs::file_size(file_path);
|
| if (file_size > 512 * 1024 * 1024) {
|
| load_json_mmap(file_path);
|
| } else {
|
| load_json(file_path);
|
| }
|
| }
|
| else if (ext == ".jsonl") load_jsonl(file_path);
|
| else if (ext == ".txt") load_txt(file_path);
|
| else if (ext == ".csv" || ext == ".tsv") load_txt(file_path);
|
| else if (ext == ".tok1" || ext == ".bin") load_tok1(file_path);
|
| else if (ext == ".parquet") {
|
| throw std::runtime_error("Parquet格式需Python预处理: python3 scripts/preprocess_corpus.py split --input <path>");
|
| } else {
|
| throw std::runtime_error("不支持的文件格式: " + ext);
|
| }
|
| } else {
|
| if (ext == ".csv" || ext == ".tsv") load_csv(file_path);
|
| else if (ext == ".tok1" || ext == ".bin") load_tok1(file_path);
|
| else {
|
| throw std::runtime_error("数值模式仅支持csv/tsv/tok1格式，收到: " + ext);
|
| }
|
| }
|
| }
|
|
|
| void load_tok1(const std::string& file_path) {
|
| std::ifstream ifs(file_path, std::ios::binary);
|
| if (!ifs) throw std::runtime_error("无法打开TOK1文件: " + file_path);
|
|
|
| char magic[5] = {0};
|
| ifs.read(magic, 4);
|
| if (std::string(magic) != "TOK1") {
|
| throw std::runtime_error("TOK1格式Magic不匹配: " + std::string(magic));
|
| }
|
|
|
| uint16_t version = 0;
|
| ifs.read(reinterpret_cast<char*>(&version), 2);
|
|
|
| uint32_t vocab_size = 0, max_seq = 0, total_samples = 0;
|
| ifs.read(reinterpret_cast<char*>(&vocab_size), 4);
|
| ifs.read(reinterpret_cast<char*>(&max_seq), 4);
|
| ifs.read(reinterpret_cast<char*>(&total_samples), 4);
|
|
|
| std::cerr << "TOK1: version=" << version << " vocab=" << vocab_size
|
| << " max_seq=" << max_seq << " samples=" << total_samples << std::endl;
|
|
|
| size_t loaded = 0;
|
| for (size_t i = 0; i < total_samples; ++i) {
|
| uint16_t seq_len = 0;
|
| ifs.read(reinterpret_cast<char*>(&seq_len), 2);
|
| if (!ifs) break;
|
|
|
| LMSample sample;
|
| sample.token_ids.resize(seq_len);
|
| ifs.read(reinterpret_cast<char*>(sample.token_ids.data()), seq_len * 4);
|
| if (!ifs) break;
|
|
|
| lm_samples_.push_back(std::move(sample));
|
| loaded++;
|
|
|
| if (lm_samples_.size() >= max_samples_) {
|
| std::cerr << "TOK1: 达到采样上限 " << max_samples_ << std::endl;
|
| break;
|
| }
|
| }
|
|
|
| std::cerr << "TOK1: 加载 " << loaded << " 样本" << std::endl;
|
| }
|
|
|
| void load_json_mmap(const std::string& file_path) {
|
| #ifdef _WIN32
|
| load_json_chunked(file_path);
|
| #else
|
| int fd = open(file_path.c_str(), O_RDONLY);
|
| if (fd < 0) throw std::runtime_error("无法mmap: " + file_path);
|
|
|
| size_t file_size = lseek(fd, 0, SEEK_END);
|
| lseek(fd, 0, SEEK_SET);
|
|
|
| void* addr = mmap(nullptr, file_size, PROT_READ, MAP_PRIVATE, fd, 0);
|
| if (addr == MAP_FAILED) {
|
| close(fd);
|
| throw std::runtime_error("mmap失败: " + file_path);
|
| }
|
|
|
| const char* data = static_cast<const char*>(addr);
|
| std::cerr << "mmap: 映射 " << file_path << " (" << file_size / 1024 / 1024 << " MB)" << std::endl;
|
|
|
| size_t pos = 0;
|
| size_t brace_depth = 0;
|
| std::string record;
|
| const char* fields[] = {"\"text\"", "\"content\"", "\"title\"", "\"question\"", "\"answer\""};
|
|
|
| while (pos < file_size && lm_samples_.size() < max_samples_) {
|
| char ch = data[pos];
|
| if (ch == '{') {
|
| if (brace_depth == 0) record.clear();
|
| brace_depth++;
|
| } else if (ch == '}') {
|
| brace_depth--;
|
| if (brace_depth == 0 && !record.empty()) {
|
| record += '}';
|
| extract_text_from_json_record(record);
|
| record.clear();
|
| }
|
| }
|
| if (brace_depth > 0) record += ch;
|
| pos++;
|
|
|
| if (lm_samples_.size() % 100000 == 0 && lm_samples_.size() > 0 && pos % (100 * 1024 * 1024) == 0) {
|
| std::cerr << " mmap进度: " << pos * 100 / file_size << "%, "
|
| << lm_samples_.size() << " 样本" << std::endl;
|
| }
|
| }
|
|
|
| munmap(addr, file_size);
|
| close(fd);
|
| std::cerr << "mmap: 完成, " << lm_samples_.size() << " 样本" << std::endl;
|
| #endif
|
| }
|
|
|
| void load_json_chunked(const std::string& file_path) {
|
| const size_t chunk_size = 256 * 1024 * 1024;
|
| std::ifstream ifs(file_path, std::ios::binary);
|
| if (!ifs) throw std::runtime_error("无法打开: " + file_path);
|
|
|
| size_t file_size = std::filesystem::file_size(file_path);
|
| std::cerr << "分块读取: " << file_path << " (" << file_size / 1024 / 1024 << " MB)" << std::endl;
|
|
|
| size_t offset = 0;
|
| size_t brace_depth = 0;
|
| std::string record;
|
| std::string overlap;
|
|
|
| while (offset < file_size && lm_samples_.size() < max_samples_) {
|
| size_t read_size = std::min(chunk_size, file_size - offset);
|
| std::string buffer;
|
| buffer.resize(read_size);
|
| ifs.read(&buffer[0], read_size);
|
| if (!ifs) break;
|
|
|
| if (!overlap.empty()) {
|
| buffer = overlap + buffer;
|
| overlap.clear();
|
| }
|
|
|
| size_t pos = 0;
|
| while (pos < buffer.size() && lm_samples_.size() < max_samples_) {
|
| char ch = buffer[pos];
|
| if (ch == '{') {
|
| if (brace_depth == 0) record.clear();
|
| brace_depth++;
|
| } else if (ch == '}') {
|
| brace_depth--;
|
| if (brace_depth == 0 && !record.empty()) {
|
| record += '}';
|
| extract_text_from_json_record(record);
|
| record.clear();
|
| }
|
| }
|
| if (brace_depth > 0 && record.size() < 1048576) record += ch;
|
| pos++;
|
| }
|
|
|
| if (brace_depth > 0) {
|
| overlap = record;
|
| }
|
|
|
| offset += read_size;
|
| std::cerr << " 分块进度: " << offset * 100 / file_size << "%, "
|
| << lm_samples_.size() << " 样本" << std::endl;
|
| }
|
|
|
| if (!record.empty() && brace_depth == 0) {
|
| extract_text_from_json_record(record);
|
| }
|
| }
|
|
|
| void load_json(const std::string& file_path) {
|
| std::ifstream ifs(file_path);
|
| if (!ifs) throw std::runtime_error("无法打开: " + file_path);
|
|
|
| std::string line;
|
| size_t brace_depth = 0;
|
| std::string record;
|
| size_t records_found = 0;
|
|
|
| while (std::getline(ifs, line)) {
|
| for (size_t i = 0; i < line.size(); ++i) {
|
| if (line[i] == '{') {
|
| if (brace_depth == 0) record.clear();
|
| brace_depth++;
|
| } else if (line[i] == '}') {
|
| brace_depth--;
|
| if (brace_depth == 0 && !record.empty()) {
|
| record += '}';
|
| extract_text_from_json_record(record);
|
| record.clear();
|
| records_found++;
|
| if (lm_samples_.size() >= max_samples_) return;
|
| }
|
| }
|
| }
|
| if (brace_depth > 0) {
|
| if (record.size() < 1048576) {
|
| record += line;
|
| record += '\n';
|
| }
|
| }
|
| }
|
|
|
| if (records_found == 0) {
|
| std::cerr << " 警告: 流式record解析未找到数据，尝试逐行文本提取" << std::endl;
|
| ifs.clear();
|
| ifs.seekg(0);
|
| std::string line;
|
| while (std::getline(ifs, line)) {
|
| if (line.size() < 20) continue;
|
| size_t q1 = line.find('"');
|
| if (q1 == std::string::npos) continue;
|
| size_t q2 = line.rfind('"');
|
| if (q2 <= q1 + 1) continue;
|
| std::string text = line.substr(q1 + 1, q2 - q1 - 1);
|
| unescape_json(text);
|
| if (text.size() >= 10) add_lm_sample(text);
|
| if (lm_samples_.size() >= max_samples_) return;
|
| }
|
| }
|
| }
|
|
|
| void extract_text_from_json_record(const std::string& record) {
|
| const char* fields[] = {"\"text\"", "\"content\"", "\"title\"", "\"question\"", "\"answer\""};
|
| for (auto& field : fields) {
|
| size_t pos = record.find(field);
|
| if (pos == std::string::npos) continue;
|
|
|
| size_t colon = record.find(':', pos + strlen(field));
|
| if (colon == std::string::npos) continue;
|
| colon++;
|
| while (colon < record.size() && record[colon] != '"') colon++;
|
| if (colon >= record.size()) continue;
|
| colon++;
|
|
|
| size_t end = colon;
|
| while (end < record.size() && record[end] != '"') {
|
| if (record[end] == '\\') end++;
|
| end++;
|
| }
|
|
|
| std::string text = record.substr(colon, end - colon);
|
| unescape_json(text);
|
| if (text.size() >= 10) add_lm_sample(text);
|
| }
|
| }
|
|
|
| void extract_all_text_fields(const std::string& content) {
|
| size_t pos = 0;
|
| const char* fields[] = {"\"text\"", "\"content\"", "\"title\"", "\"question\"", "\"answer\""};
|
| while (pos < content.size()) {
|
| size_t best_pos = std::string::npos;
|
| for (auto& field : fields) {
|
| size_t p = content.find(field, pos);
|
| if (p != std::string::npos && (best_pos == std::string::npos || p < best_pos)) {
|
| best_pos = p;
|
| }
|
| }
|
| if (best_pos == std::string::npos) break;
|
|
|
| size_t colon = content.find(':', best_pos);
|
| if (colon == std::string::npos) break;
|
| colon++;
|
| while (colon < content.size() && content[colon] != '"') colon++;
|
| if (colon >= content.size()) break;
|
| colon++;
|
|
|
| size_t end = colon;
|
| while (end < content.size() && content[end] != '"') {
|
| if (content[end] == '\\') end++;
|
| end++;
|
| }
|
| if (end >= content.size()) break;
|
|
|
| std::string text = content.substr(colon, end - colon);
|
| unescape_json(text);
|
| add_lm_sample(text);
|
|
|
| pos = end + 1;
|
| if (lm_samples_.size() >= max_samples_) return;
|
| }
|
| }
|
|
|
| void load_jsonl(const std::string& file_path) {
|
| std::ifstream ifs(file_path);
|
| if (!ifs) throw std::runtime_error("无法打开: " + file_path);
|
| std::string line;
|
| size_t line_num = 0;
|
| while (std::getline(ifs, line)) {
|
| line_num++;
|
| if (line.empty() || line[0] == '#') continue;
|
|
|
| const char* fields[] = {"\"text\"", "\"content\"", "\"title\"", "\"question\"", "\"answer\""};
|
| bool found = false;
|
| for (auto& field : fields) {
|
| size_t text_pos = line.find(field);
|
| if (text_pos == std::string::npos) continue;
|
|
|
| size_t colon = line.find(':', text_pos);
|
| if (colon == std::string::npos) continue;
|
| colon++;
|
| while (colon < line.size() && line[colon] != '"') colon++;
|
| if (colon >= line.size()) continue;
|
| colon++;
|
|
|
| size_t end = colon;
|
| while (end < line.size() && line[end] != '"') {
|
| if (line[end] == '\\') end++;
|
| end++;
|
| }
|
|
|
| std::string text = line.substr(colon, end - colon);
|
| unescape_json(text);
|
| add_lm_sample(text);
|
| found = true;
|
| break;
|
| }
|
| if (!found && line_num <= 5) {
|
| std::cerr << " jsonl行" << line_num << ": 未找到text/content字段" << std::endl;
|
| }
|
| if (lm_samples_.size() >= max_samples_) return;
|
| }
|
| }
|
|
|
| void load_txt(const std::string& file_path) {
|
| std::ifstream ifs(file_path);
|
| if (!ifs) throw std::runtime_error("无法打开: " + file_path);
|
| std::string line;
|
| std::string paragraph;
|
| size_t line_count = 0;
|
|
|
| while (std::getline(ifs, line)) {
|
| line_count++;
|
| if (line.empty()) {
|
| if (paragraph.size() >= 10) {
|
| add_lm_sample(paragraph);
|
| paragraph.clear();
|
| if (lm_samples_.size() >= max_samples_) return;
|
| }
|
| } else {
|
| if (!paragraph.empty()) paragraph += '\n';
|
| paragraph += line;
|
| if (paragraph.size() > 10000) {
|
| add_lm_sample(paragraph);
|
| paragraph.clear();
|
| if (lm_samples_.size() >= max_samples_) return;
|
| }
|
| }
|
| if (line_count % 500000 == 0) {
|
| std::cerr << " load_txt: " << line_count << " 行, "
|
| << lm_samples_.size() << " 样本" << std::endl;
|
| }
|
| }
|
| if (paragraph.size() >= 10) add_lm_sample(paragraph);
|
| }
|
|
|
| void load_csv(const std::string& file_path) {
|
| std::ifstream ifs(file_path);
|
| if (!ifs) throw std::runtime_error("无法打开: " + file_path);
|
| std::string line;
|
| while (std::getline(ifs, line)) {
|
| if (line.empty() || line[0] == '#') continue;
|
| TrainingSample s;
|
| std::stringstream ss(line);
|
| std::string val;
|
| char delim = ',';
|
| if (line.find('\t') != std::string::npos && line.find(',') == std::string::npos) {
|
| delim = '\t';
|
| }
|
| bool in_target = false;
|
| while (std::getline(ss, val, delim)) {
|
| float v = std::atof(val.c_str());
|
| if (!in_target) {
|
| s.input.push_back(v);
|
| if (s.input.size() == input_dim_) in_target = true;
|
| } else {
|
| s.target.push_back(v);
|
| }
|
| }
|
| if (!s.input.empty() && !s.target.empty()) {
|
| samples_.push_back(std::move(s));
|
| }
|
| if (samples_.size() >= max_samples_) return;
|
| }
|
| }
|
|
|
| void add_lm_sample(const std::string& text) {
|
| if (text.size() < 10) return;
|
| if (!tokenizer_) return;
|
| if (lm_samples_.size() >= max_samples_) return;
|
| if (memory_usage_bytes() >= max_memory_bytes_) {
|
| if (lm_samples_.size() % 100000 == 1) {
|
| std::cerr << " 内存预算达到 " << max_memory_bytes_ / 1024 / 1024
|
| << " MB，停止加载" << std::endl;
|
| }
|
| return;
|
| }
|
|
|
|
|
|
|
|
|
| const size_t max_text_len = max_seq_len_ * 4;
|
| if (text.size() > max_text_len) {
|
|
|
| auto ids = tokenizer_->encode(text.substr(0, max_text_len), max_seq_len_);
|
| if (ids.size() >= 4) {
|
| lm_samples_.push_back(LMSample{std::move(ids)});
|
| }
|
| return;
|
| }
|
| auto ids = tokenizer_->encode(text, max_seq_len_);
|
| if (ids.size() < 4) return;
|
| lm_samples_.push_back(LMSample{std::move(ids)});
|
| }
|
|
|
| static void unescape_json_static(std::string& s) { unescape_json(s); }
|
|
|
| bool has_next() const {
|
| if (mode_ == Mode::LM) return current_idx_ < lm_samples_.size();
|
| return current_idx_ < samples_.size();
|
| }
|
|
|
| std::vector<TrainingSample> next_batch() {
|
| std::vector<TrainingSample> batch;
|
| if (mode_ == Mode::LM) {
|
| float vocab_scale = 1.0f;
|
| if (tokenizer_) vocab_scale = 1.0f / static_cast<float>(tokenizer_->vocab_size());
|
|
|
| for (int i = 0; i < batch_size_ && current_idx_ < lm_samples_.size(); ++i) {
|
| auto& sample = lm_samples_[current_idx_++];
|
| TrainingSample s;
|
| s.input.resize(input_dim_, 0.0f);
|
| s.target.resize(output_dim_, 0.0f);
|
| size_t n = std::min(sample.token_ids.size(), input_dim_);
|
|
|
| for (size_t j = 0; j < n; ++j) {
|
| s.input[j] = static_cast<float>(sample.token_ids[j]) * vocab_scale;
|
| }
|
|
|
| for (size_t j = 0; j + 1 < n; ++j) {
|
| size_t next_id = sample.token_ids[j + 1];
|
| if (next_id < output_dim_) {
|
| s.target[next_id] = 1.0f;
|
| }
|
| }
|
| batch.push_back(std::move(s));
|
| }
|
| } else {
|
| for (int i = 0; i < batch_size_ && current_idx_ < samples_.size(); ++i) {
|
| batch.push_back(samples_[current_idx_++]);
|
| }
|
| }
|
| return batch;
|
| }
|
|
|
| std::vector<LMSample> next_lm_batch_raw() {
|
| std::vector<LMSample> batch;
|
| for (int i = 0; i < batch_size_ && current_idx_ < lm_samples_.size(); ++i) {
|
| batch.push_back(lm_samples_[current_idx_++]);
|
| }
|
| return batch;
|
| }
|
|
|
| void reset() { current_idx_ = 0; }
|
| void shuffle(std::mt19937& rng) {
|
| if (mode_ == Mode::LM) {
|
| std::shuffle(lm_samples_.begin(), lm_samples_.end(), rng);
|
| } else {
|
| std::shuffle(samples_.begin(), samples_.end(), rng);
|
| }
|
| }
|
| size_t total_samples() const {
|
| if (mode_ == Mode::LM) return lm_samples_.size();
|
| return samples_.size();
|
| }
|
| Mode mode() const { return mode_; }
|
| size_t memory_usage_bytes() const {
|
| size_t total = 0;
|
| for (auto& s : lm_samples_) total += s.token_ids.size() * sizeof(size_t);
|
| for (auto& s : samples_) total += (s.input.size() + s.target.size()) * sizeof(float);
|
| return total;
|
| }
|
|
|
| private:
|
| std::string path_;
|
| int batch_size_;
|
| size_t input_dim_;
|
| size_t output_dim_;
|
| size_t current_idx_;
|
| Mode mode_;
|
| BPETokenizer* tokenizer_;
|
| size_t max_seq_len_;
|
| size_t max_samples_;
|
| size_t max_memory_bytes_;
|
| std::vector<TrainingSample> samples_;
|
| std::vector<LMSample> lm_samples_;
|
| };
|
|
|
| struct TrainMetrics {
|
| float loss = 0.0f;
|
| float lr = 0.0f;
|
| size_t step = 0;
|
| int epoch = 0;
|
| float elapsed_seconds = 0.0f;
|
| float grad_norm = 0.0f;
|
| };
|
|
|
| class MetricsLogger {
|
| public:
|
| MetricsLogger() : start_time_(std::chrono::steady_clock::now()) {}
|
|
|
| void log_step(const TrainMetrics& m, int interval) {
|
| if (m.step % interval != 0) return;
|
| std::cerr << "[Epoch " << m.epoch << "][Step " << m.step
|
| << "] loss=" << m.loss << " lr=" << m.lr
|
| << " grad_norm=" << m.grad_norm
|
| << " elapsed=" << m.elapsed_seconds << "s" << std::endl;
|
| }
|
|
|
| void log_epoch(int epoch, float avg_loss, float elapsed) {
|
| std::cerr << "=== Epoch " << epoch << " 完成, avg_loss=" << avg_loss
|
| << ", elapsed=" << elapsed << "s ===" << std::endl;
|
| }
|
|
|
| void save_training_log(const std::string& path, const std::vector<float>& losses) {
|
| std::ofstream ofs(path);
|
| if (!ofs) return;
|
| ofs << "{\n \"losses\": [";
|
| for (size_t i = 0; i < losses.size(); ++i) {
|
| if (i > 0) ofs << ", ";
|
| ofs << losses[i];
|
| }
|
| ofs << "],\n \"num_epochs\": " << losses.size() << "\n}";
|
| ofs.close();
|
| }
|
|
|
| double total_elapsed() const {
|
| auto now = std::chrono::steady_clock::now();
|
| return std::chrono::duration<double>(now - start_time_).count();
|
| }
|
|
|
| private:
|
| std::chrono::steady_clock::time_point start_time_;
|
| };
|
|
|
| void ensure_dir_exists(const std::string& dir) {
|
| std::filesystem::create_directories(dir);
|
| }
|
|
|
| void save_checkpoint(const NeuroFlowModel& model, const TrainMetrics& metrics,
|
| const std::string& output_dir, size_t step, int epoch,
|
| const TrainConfig& cfg) {
|
|
|
|
|
|
|
|
|
| std::string ckpt_name = "checkpoint_step" + std::to_string(step);
|
| std::string ckpt_dir = output_dir + "/" + ckpt_name;
|
| ensure_dir_exists(ckpt_dir);
|
|
|
|
|
| std::string shm_dir = "/dev/shm/nf_ckpt/" + ckpt_name;
|
| ensure_dir_exists(shm_dir);
|
| model.save(shm_dir + "/model.nfv1");
|
|
|
| std::ofstream shm_ofs(shm_dir + "/training_state.json");
|
| if (shm_ofs) {
|
| shm_ofs << "{\n"
|
| << " \"step\": " << step << ",\n"
|
| << " \"epoch\": " << epoch << ",\n"
|
| << " \"loss\": " << metrics.loss << ",\n"
|
| << " \"learning_rate\": " << metrics.lr << ",\n"
|
| << " \"grad_norm\": " << metrics.grad_norm << ",\n"
|
| << " \"config_d_model\": " << model.config.input_dim << ",\n"
|
| << " \"config_hidden_dim\": " << model.config.hidden_dim << ",\n"
|
| << " \"config_output_dim\": " << model.config.output_dim << ",\n"
|
| << " \"config_vocab_size\": " << model.config.vocab_size << ",\n"
|
| << " \"init_strategy\": \"" << cfg.init_strategy << "\",\n"
|
| << " \"seed\": " << cfg.seed << "\n"
|
| << "}\n";
|
| shm_ofs.close();
|
| }
|
|
|
|
|
| #ifndef _WIN32
|
| pid_t pid = fork();
|
| if (pid == 0) {
|
|
|
| execlp("cp", "cp", "-r", shm_dir.c_str(), ckpt_dir.c_str(), static_cast<char*>(nullptr));
|
| _exit(1);
|
| }
|
|
|
| std::cerr << " [ckpt] RAM→D盘后台拷贝: " << ckpt_name
|
| << " (pid=" << pid << ")" << std::endl;
|
| #else
|
|
|
| std::filesystem::copy(shm_dir, ckpt_dir,
|
| std::filesystem::copy_options::recursive | std::filesystem::copy_options::overwrite_existing);
|
| #endif
|
|
|
| std::cerr << "Checkpoint已保存: step=" << step << " epoch=" << epoch
|
| << " loss=" << metrics.loss << std::endl;
|
| }
|
|
|
|
|
| std::unique_ptr<BPETokenizer> load_tokenizer(const std::string& tokenizer_path) {
|
| auto tok = std::make_unique<BPETokenizer>(tokenizer_path);
|
| std::cerr << "词表加载成功: " << tokenizer_path
|
| << " (vocab_size=" << tok->vocab_size() << ")" << std::endl;
|
| return tok;
|
| }
|
|
|
| void train_loop(NeuroFlowModel& model, CausalLMHead& lm_head, DataLoader& loader, MetricsLogger& logger,
|
| const TrainConfig& cfg, const ResumeState& resume_state,
|
| const std::vector<uint8_t>& vocab_mask = {}) {
|
| #ifdef _OPENMP
|
| int max_threads = omp_get_max_threads();
|
| std::cerr << "OpenMP最大线程数: " << max_threads << std::endl;
|
| #else
|
| std::cerr << "警告: 未启用OpenMP，非GEMM运算为单线程" << std::endl;
|
| #endif
|
|
|
| size_t vocab_sz = lm_head.config_.vocab_size;
|
| size_t d_model = lm_head.config_.d_model;
|
| float learning_rate = cfg.learning_rate;
|
|
|
| std::cerr << "训练模式: CausalLMHead完整forward (embed→PE→Attn×"
|
| << lm_head.config_.num_attn_layers << "→Gate→SAE→NTM→LN→Pool→Proj→Out)" << std::endl;
|
| std::cerr << " d_model=" << d_model << " vocab=" << vocab_sz
|
| << " pooling=" << lm_head.config_.pooling << std::endl;
|
|
|
| std::vector<float> epoch_losses;
|
| int start_epoch = resume_state.found ? resume_state.epoch : 0;
|
| size_t start_step = resume_state.found ? resume_state.step : 0;
|
|
|
| if (start_epoch > 0) {
|
| std::cerr << "续训: 从epoch " << start_epoch + 1 << " 开始, 已完成 "
|
| << start_step << " 步" << std::endl;
|
| }
|
|
|
| auto save_lm_head = [&](const std::string& path) {
|
| #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();
|
| };
|
|
|
| lm_head.train();
|
|
|
| for (int epoch = start_epoch; epoch < cfg.epochs; ++epoch) {
|
| auto epoch_start = std::chrono::steady_clock::now();
|
| loader.reset();
|
| std::mt19937 shuffle_rng(cfg.seed + epoch);
|
| loader.shuffle(shuffle_rng);
|
| float epoch_loss = 0.0f;
|
| size_t step_count = 0;
|
| size_t global_step = start_step + epoch * ((loader.total_samples() + cfg.batch_size - 1) / cfg.batch_size);
|
|
|
| while (loader.has_next()) {
|
| auto batch = loader.next_lm_batch_raw();
|
| if (batch.empty()) break;
|
|
|
| float batch_loss = 0.0f;
|
| float batch_grad_norm = 0.0f;
|
| size_t batch_count = 0;
|
|
|
| for (auto& sample : batch) {
|
| if (sample.token_ids.size() < 2) continue;
|
|
|
| std::vector<size_t> input_ids(sample.token_ids.begin(), sample.token_ids.end() - 1);
|
| size_t target_id = sample.token_ids.back();
|
| if (target_id >= vocab_sz) target_id = 1;
|
|
|
| Tensor logits = lm_head.forward_for_training(input_ids);
|
|
|
| float loss = 0.0f;
|
| float grad_norm = 0.0f;
|
| Tensor logits_grad({1, vocab_sz}, QuantType::FP32);
|
|
|
| #ifdef USE_CUDA
|
| if (cfg.use_cuda && logits.is_on_gpu()) {
|
| logits_grad.to_gpu();
|
| launch_cross_entropy_backward(logits_grad.as_gpu_fp32(), logits.as_gpu_fp32(),
|
| static_cast<int>(target_id), static_cast<int>(vocab_sz), CudaContext::instance().stream());
|
| logits_grad.gpu_dirty_ = true;
|
|
|
| Tensor d_loss({1}, QuantType::FP32);
|
| d_loss.to_gpu();
|
| launch_cross_entropy(d_loss.as_gpu_fp32(), logits.as_gpu_fp32(),
|
| static_cast<int>(target_id), static_cast<int>(vocab_sz), CudaContext::instance().stream());
|
| CudaContext::instance().synchronize();
|
| float loss_arr[1];
|
| CudaContext::instance().copy_d2h(loss_arr, d_loss.as_gpu_fp32(), sizeof(float));
|
| CudaContext::instance().synchronize();
|
| loss = loss_arr[0];
|
|
|
| if (!std::isfinite(loss)) {
|
| std::cerr << "[WARN] NaN loss at step " << global_step << ", skipping" << std::endl;
|
| continue;
|
| }
|
|
|
| grad_norm = 0.0f;
|
| } else
|
| #endif
|
| {
|
| const float* pred = logits.as_fp32();
|
| float max_val = -1e30f;
|
| for (size_t j = 0; j < vocab_sz; ++j) {
|
| if (pred[j] > max_val) max_val = pred[j];
|
| }
|
| float sum_exp = 0.0f;
|
| for (size_t j = 0; j < vocab_sz; ++j) {
|
| sum_exp += std::exp(pred[j] - max_val);
|
| }
|
| float log_sum_exp = max_val + std::log(sum_exp);
|
| loss = -(pred[target_id] - log_sum_exp);
|
|
|
| if (!std::isfinite(loss)) {
|
| std::cerr << "[WARN] NaN loss at step " << global_step << ", skipping" << std::endl;
|
| continue;
|
| }
|
|
|
| float* lg = logits_grad.as_fp32();
|
| for (size_t j = 0; j < vocab_sz; ++j) {
|
| float softmax_val = std::exp(pred[j] - max_val) / sum_exp;
|
| lg[j] = softmax_val;
|
| if (j == target_id) lg[j] -= 1.0f;
|
| grad_norm += lg[j] * lg[j];
|
| }
|
| }
|
|
|
| auto lm_grads = lm_head.backward_from_logits(logits_grad);
|
|
|
| float clip_val = cfg.grad_clip;
|
| float gn = std::sqrt(grad_norm);
|
| float clip_scale = 1.0f;
|
| if (!std::isfinite(gn) || gn > clip_val && clip_val > 0.0f) {
|
| clip_scale = clip_val / gn;
|
| }
|
|
|
| lm_head.apply_lm_gradients(lm_grads, learning_rate * clip_scale);
|
|
|
|
|
| batch_loss += loss;
|
| batch_grad_norm += grad_norm;
|
| batch_count++;
|
| }
|
|
|
| if (batch_count > 0) {
|
| epoch_loss += batch_loss / batch_count;
|
| step_count++;
|
| global_step++;
|
|
|
| TrainMetrics metrics;
|
| metrics.loss = batch_loss / batch_count;
|
| metrics.lr = cfg.learning_rate;
|
| metrics.step = global_step;
|
| metrics.epoch = epoch + 1;
|
| metrics.grad_norm = std::sqrt(batch_grad_norm / batch_count);
|
| metrics.elapsed_seconds = static_cast<float>(logger.total_elapsed());
|
| logger.log_step(metrics, cfg.log_interval);
|
|
|
| if (cfg.save_interval > 0 && global_step % cfg.save_interval == 0) {
|
| std::string cdir = cfg.output_dir + "/checkpoint_step" + std::to_string(global_step);
|
| ensure_dir_exists(cdir);
|
| model.save(cdir + "/model.nfv1");
|
| save_lm_head(cdir + "/lm_head.nfv1");
|
| std::cerr << "Checkpoint: step=" << global_step << " loss=" << metrics.loss << std::endl;
|
| }
|
| }
|
| }
|
|
|
| float avg_loss = (step_count > 0) ? epoch_loss / step_count : 0.0f;
|
| epoch_losses.push_back(avg_loss);
|
|
|
| auto epoch_end = std::chrono::steady_clock::now();
|
| float epoch_elapsed = static_cast<float>(
|
| std::chrono::duration<double>(epoch_end - epoch_start).count());
|
| logger.log_epoch(epoch + 1, avg_loss, epoch_elapsed);
|
|
|
| std::string cdir = cfg.output_dir + "/checkpoint_epoch" + std::to_string(epoch + 1);
|
| ensure_dir_exists(cdir);
|
| model.save(cdir + "/model.nfv1");
|
| save_lm_head(cdir + "/lm_head.nfv1");
|
| }
|
|
|
| logger.save_training_log(cfg.output_dir + "/training_log.json", epoch_losses);
|
| save_lm_head(cfg.output_dir + "/lm_head_final.nfv1");
|
| std::cerr << "训练完成, LM Head已保存: " << cfg.output_dir << "/lm_head_final.nfv1" << std::endl;
|
| }
|
|
|
| int main(int argc, char* argv[]) {
|
| setvbuf(stderr, nullptr, _IONBF, 0);
|
| setvbuf(stdout, nullptr, _IONBF, 0);
|
| std::ios::sync_with_stdio(false);
|
|
|
| #ifdef USE_CBLAS
|
| int ncpu = 1;
|
| #ifdef _WIN32
|
| SYSTEM_INFO si; GetSystemInfo(&si); ncpu = si.dwNumberOfProcessors;
|
| #else
|
| ncpu = sysconf(_SC_NPROCESSORS_ONLN);
|
| #endif
|
| if (ncpu < 1) ncpu = 1;
|
| int blas_threads = std::max(1, ncpu / 2);
|
| std::string nt = std::to_string(blas_threads);
|
| #ifdef _WIN32
|
| _putenv_s("OPENBLAS_NUM_THREADS", nt.c_str());
|
| _putenv_s("GOTO_NUM_THREADS", nt.c_str());
|
| #else
|
| setenv("OPENBLAS_NUM_THREADS", nt.c_str(), 1);
|
| setenv("GOTO_NUM_THREADS", nt.c_str(), 1);
|
| #endif
|
| std::cerr << "GEMM后端: OpenBLAS (多线程), BLAS线程: " << blas_threads << std::endl;
|
| #else
|
| std::cerr << "GEMM后端: 内置AVX2 (单线程!) — 安装libopenblas-dev可大幅提升性能" << std::endl;
|
| #endif
|
|
|
| TrainConfig cfg = parse_args(argc, argv);
|
|
|
| if (cfg.config_path.empty() || cfg.tokenizer_path.empty()) {
|
| std::cerr << "用法: neuroflow_train_v2 --config <config.json> --tokenizer <tokenizer.json>"
|
| << " [--data <data>] [--output <dir>] [--epochs <N>] [--lr <lr>]"
|
| << " [--seed <N>] [--resume <path>] [--init-weights <strategy>]"
|
| << " [--batch-size <N>] [--log-interval <N>] [--save-interval <N>]"
|
| << " [--grad-clip <val>] [--adam] [--grad-accum <N>]"
|
| << " [--use-cuda] [--verbose]" << std::endl;
|
| return 1;
|
| }
|
|
|
| std::cerr << "======================================" << std::endl;
|
| std::cerr << "NeuroFlow Train v2 (Full Backprop)" << std::endl;
|
| std::cerr << "======================================" << std::endl;
|
| std::cerr << "配置: " << cfg.config_path << std::endl;
|
| std::cerr << "词表: " << cfg.tokenizer_path << std::endl;
|
| std::cerr << "数据: " << (cfg.data_path.empty() ? "(随机)" : cfg.data_path) << std::endl;
|
| std::cerr << "输出: " << cfg.output_dir << std::endl;
|
| std::cerr << "Epochs: " << cfg.epochs << " LR: " << cfg.learning_rate
|
| << " Batch: " << cfg.batch_size
|
| << " GradAccum: " << cfg.grad_accum_steps
|
| << " (等效batch=" << cfg.batch_size * cfg.grad_accum_steps << ")" << std::endl;
|
|
|
| #ifdef USE_CUDA
|
| bool cuda_active = false;
|
| if (cfg.use_cuda) {
|
| if (CudaContext::instance().initialize(0)) {
|
| cuda_active = true;
|
| std::cerr << "GPU后端: 已启用 (RTX 3090)" << std::endl;
|
| } else {
|
| std::cerr << "[CUDA WARNING] GPU初始化失败，回退CPU后端" << std::endl;
|
| cfg.use_cuda = false;
|
| }
|
| } else {
|
| std::cerr << "GPU后端: 未启用 (使用--use-cuda启用)" << std::endl;
|
| }
|
| #else
|
| if (cfg.use_cuda) {
|
| std::cerr << "[CUDA WARNING] 编译时未启用CUDA支持 (NEUROFLOW_USE_CUDA=OFF)，回退CPU后端" << std::endl;
|
| cfg.use_cuda = false;
|
| }
|
| #endif
|
|
|
| auto model_cfg = load_model_config(cfg.config_path);
|
| ensure_dir_exists(cfg.output_dir);
|
| ResumeState resume_state;
|
| auto model = build_model(model_cfg, cfg, resume_state);
|
| auto tokenizer = load_tokenizer(cfg.tokenizer_path);
|
|
|
| auto stats = model.get_stats();
|
| std::cerr << "模型参数: " << stats.total_params << " 内存: " << stats.memory_bytes << " bytes" << std::endl;
|
|
|
| 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 = model_cfg.lm_num_attn_heads;
|
| lm_cfg.n_kv_heads = model_cfg.lm_n_kv_heads;
|
| lm_cfg.use_rope = model_cfg.lm_use_rope;
|
| lm_cfg.use_qk_norm = model_cfg.lm_use_qk_norm;
|
| lm_cfg.use_swiglu = model_cfg.lm_use_swiglu;
|
| lm_cfg.use_bridge = model_cfg.lm_use_bridge;
|
| lm_cfg.pooling = model_cfg.lm_pooling;
|
| CausalLMHead lm_head(lm_cfg);
|
| if (lm_cfg.weight_tying) lm_head.tie_weights();
|
| if (!cfg.resume_path.empty()) {
|
| std::string lm_ckpt_path;
|
| if (cfg.resume_path.size() >= 5 && cfg.resume_path.substr(cfg.resume_path.size() - 5) == ".nfv1") {
|
| std::string dir = cfg.resume_path.substr(0, cfg.resume_path.find_last_of("/\\"));
|
| lm_ckpt_path = dir + "/lm_head.nfv1";
|
| } else {
|
| lm_ckpt_path = cfg.resume_path + "/lm_head.nfv1";
|
| }
|
| load_lm_head(lm_head, lm_ckpt_path);
|
| }
|
| std::cerr << "CausalLMHead: d_model=" << lm_cfg.d_model
|
| << " vocab=" << lm_cfg.vocab_size
|
| << " attn_layers=" << lm_cfg.num_attn_layers
|
| << " pooling=" << lm_cfg.pooling << std::endl;
|
|
|
| #ifdef USE_CUDA
|
| if (cfg.use_cuda && cuda_active) {
|
| 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;
|
| }
|
| #endif
|
|
|
| DataLoader loader(cfg.data_path, cfg.batch_size, model_cfg.input_dim, model_cfg.vocab_size,
|
| DataLoader::Mode::LM, tokenizer.get(), model_cfg.max_seq_len,
|
| 500000, 4096);
|
| std::cerr << "DataLoader内存预算: 4096 MB, 采样上限: 500000" << std::endl;
|
|
|
| std::vector<uint8_t> vocab_mask;
|
| if (!cfg.vocab_mask_path.empty()) {
|
| std::ifstream vmf(cfg.vocab_mask_path);
|
| if (vmf) {
|
| std::string vm_content((std::istreambuf_iterator<char>(vmf)), std::istreambuf_iterator<char>());
|
| size_t mask_start = vm_content.find("\"mask\"");
|
| if (mask_start != std::string::npos) {
|
| size_t arr_start = vm_content.find('[', mask_start);
|
| size_t arr_end = vm_content.find(']', arr_start);
|
| if (arr_start != std::string::npos && arr_end != std::string::npos) {
|
| std::string arr_str = vm_content.substr(arr_start + 1, arr_end - arr_start - 1);
|
| size_t pos = 0;
|
| while (pos < arr_str.size()) {
|
| while (pos < arr_str.size() && (arr_str[pos] == ' ' || arr_str[pos] == ',')) pos++;
|
| if (pos >= arr_str.size()) break;
|
| vocab_mask.push_back(static_cast<uint8_t>(arr_str[pos] - '0'));
|
| pos++;
|
| }
|
| }
|
| }
|
| size_t active = 0;
|
| for (auto m : vocab_mask) active += m;
|
| std::cerr << "Vocab mask: " << cfg.vocab_mask_path
|
| << " (" << active << "/" << vocab_mask.size() << " active)" << std::endl;
|
| } else {
|
| std::cerr << "Warning: vocab mask not found: " << cfg.vocab_mask_path << std::endl;
|
| }
|
| }
|
|
|
| MetricsLogger metrics_logger;
|
|
|
| train_loop(model, lm_head, loader, metrics_logger, cfg, resume_state, vocab_mask);
|
|
|
| model.save(cfg.output_dir + "/model_final.nfv1");
|
| std::cerr << "NF模型已保存: " << cfg.output_dir << "/model_final.nfv1" << std::endl;
|
|
|
| return 0;
|
| }
|
|
|
