| #include "neuroflow/train_lm.hpp"
|
|
|
| #include <algorithm>
|
| #include <cmath>
|
| #include <cstring>
|
| #include <fstream>
|
| #include <iostream>
|
| #include <numeric>
|
| #include <random>
|
|
|
| #ifdef USE_CUDA
|
| #include "cuda_context.hpp"
|
| #include "cuda_kernels.hpp"
|
| #endif
|
|
|
| namespace neuroflow {
|
|
|
| TrainLM::TrainLM(const TrainLMConfig& config) : cfg_(config) {
|
| CausalLMConfig lm_cfg;
|
| lm_cfg.vocab_size = cfg_.vocab_size;
|
| lm_cfg.d_model = cfg_.d_model;
|
| lm_cfg.max_seq_len = cfg_.max_seq_len;
|
| lm_cfg.num_attn_layers = cfg_.num_attn_layers;
|
| lm_cfg.num_attn_heads = cfg_.num_attn_heads;
|
| lm_cfg.causal_window_size = cfg_.causal_window_size;
|
| lm_cfg.sae_k = cfg_.sae_k;
|
| lm_cfg.ntm_memory_slots = cfg_.ntm_memory_slots;
|
| lm_cfg.weight_tying = cfg_.weight_tying;
|
| lm_cfg.use_rope = cfg_.use_rope;
|
| lm_cfg.use_bridge = cfg_.use_bridge;
|
| lm_cfg.use_swiglu = cfg_.use_swiglu;
|
| lm_cfg.use_qk_norm = cfg_.use_qk_norm;
|
| lm_cfg.swiglu_intermediate_size = cfg_.swiglu_intermediate_size;
|
| lm_cfg.pooling = cfg_.pooling;
|
|
|
| lm_head_ = std::make_unique<CausalLMHead>(lm_cfg);
|
| setup_optimizer();
|
|
|
| size_t total = cfg_.total_steps > 0 ? cfg_.total_steps : cfg_.epochs * 1000;
|
| scheduler_ = std::make_unique<CosineScheduler>(
|
| cfg_.learning_rate, total, cfg_.lr_min_ratio, cfg_.warmup_ratio);
|
| }
|
|
|
| void TrainLM::setup_optimizer() {
|
| optimizer_ = std::make_unique<AdamW>(
|
| cfg_.learning_rate, cfg_.adam_beta1, cfg_.adam_beta2,
|
| cfg_.adam_eps, cfg_.adam_weight_decay);
|
|
|
| lm_head_->register_trainable_params(*optimizer_, cfg_.learning_rate, cfg_.adam_weight_decay);
|
| }
|
|
|
| float TrainLM::compute_loss_and_grad(const std::vector<size_t>& input_ids,
|
| const std::vector<size_t>& target_ids,
|
| Tensor& logits_grad) {
|
| Tensor logits = lm_head_->forward_for_training(input_ids);
|
|
|
| size_t seq_len = logits.shape_[0];
|
| size_t vocab_size = logits.shape_[1];
|
| size_t n_targets = std::min(seq_len, target_ids.size());
|
|
|
| float total_loss = 0.0f;
|
| logits_grad = Tensor(logits.shape_, QuantType::FP32);
|
| float* lg = logits_grad.as_fp32();
|
| const float* lp = logits.as_fp32();
|
| memset(lg, 0, logits_grad.data_size_);
|
|
|
| float grad_norm = 0.0f;
|
|
|
| for (size_t t = 0; t < n_targets; ++t) {
|
| size_t target_id = target_ids[t];
|
| if (target_id >= vocab_size) continue;
|
|
|
| const float* row = &lp[t * vocab_size];
|
| float* grad_row = &lg[t * vocab_size];
|
|
|
| float max_val = -1e30f;
|
| for (size_t j = 0; j < vocab_size; ++j) {
|
| max_val = std::max(max_val, row[j]);
|
| }
|
|
|
| float sum_exp = 0.0f;
|
| for (size_t j = 0; j < vocab_size; ++j) {
|
| sum_exp += std::exp(row[j] - max_val);
|
| }
|
|
|
| float log_sum_exp = max_val + std::log(sum_exp);
|
| float loss = -(row[target_id] - log_sum_exp);
|
| total_loss += loss;
|
|
|
| for (size_t j = 0; j < vocab_size; ++j) {
|
| float softmax_val = std::exp(row[j] - max_val) / sum_exp;
|
| grad_row[j] = softmax_val;
|
| if (j == target_id) grad_row[j] -= 1.0f;
|
| grad_norm += grad_row[j] * grad_row[j];
|
| }
|
| }
|
|
|
| if (n_targets > 0) {
|
| float inv_n = 1.0f / static_cast<float>(n_targets);
|
| total_loss *= inv_n;
|
| for (size_t i = 0; i < logits_grad.numel(); ++i) {
|
| lg[i] *= inv_n;
|
| }
|
| }
|
|
|
| float gn = std::sqrt(grad_norm);
|
| if (cfg_.grad_clip > 0.0f && gn > cfg_.grad_clip && std::isfinite(gn)) {
|
| float scale = cfg_.grad_clip / gn;
|
| for (size_t i = 0; i < logits_grad.numel(); ++i) {
|
| lg[i] *= scale;
|
| }
|
| }
|
|
|
| return total_loss;
|
| }
|
|
|
| void TrainLM::train(const std::vector<std::vector<size_t>>& dataset) {
|
| if (dataset.empty()) {
|
| std::cerr << "[ERROR] Empty dataset" << std::endl;
|
| return;
|
| }
|
|
|
| std::cerr << "=== TrainLM: Standard Causal LM Training ===" << std::endl;
|
| std::cerr << "Dataset: " << dataset.size() << " samples" << std::endl;
|
| std::cerr << "Optimizer: AdamW (lr=" << cfg_.learning_rate
|
| << ", wd=" << cfg_.adam_weight_decay << ")" << std::endl;
|
| std::cerr << "Scheduler: Cosine with " << cfg_.warmup_ratio * 100 << "% warmup" << std::endl;
|
|
|
| size_t global_step = 0;
|
| std::mt19937 rng(42);
|
|
|
| for (size_t epoch = 0; epoch < cfg_.epochs; ++epoch) {
|
| std::vector<size_t> indices(dataset.size());
|
| std::iota(indices.begin(), indices.end(), 0);
|
| std::shuffle(indices.begin(), indices.end(), rng);
|
|
|
| float epoch_loss = 0.0f;
|
| size_t step_count = 0;
|
|
|
| for (size_t idx : indices) {
|
| const auto& sample = dataset[idx];
|
| if (sample.size() < 2) continue;
|
|
|
| std::vector<size_t> input_ids(sample.begin(), sample.end() - 1);
|
| std::vector<size_t> target_ids(sample.begin() + 1, sample.end());
|
|
|
| float lr = scheduler_->get_lr(global_step);
|
| optimizer_->set_lr(lr);
|
|
|
| Tensor logits_grad;
|
| float loss = compute_loss_and_grad(input_ids, target_ids, logits_grad);
|
|
|
| auto lm_grads = lm_head_->backward_from_logits(logits_grad);
|
|
|
| lm_head_->assign_grads_to_optimizer(*optimizer_, lm_grads);
|
|
|
| optimizer_->step();
|
| global_step++;
|
|
|
| epoch_loss += loss;
|
| step_count++;
|
|
|
| if (step_count % cfg_.log_interval == 0) {
|
| std::cerr << "Epoch " << epoch + 1 << " Step " << step_count
|
| << " loss=" << loss << " lr=" << lr << std::endl;
|
| }
|
|
|
| if (cfg_.save_interval > 0 && global_step % cfg_.save_interval == 0) {
|
| save_checkpoint(global_step, epoch_loss / step_count);
|
| }
|
| }
|
|
|
| if (step_count > 0) {
|
| std::cerr << "=== Epoch " << epoch + 1 << " avg_loss="
|
| << epoch_loss / step_count << " ===" << std::endl;
|
| }
|
| }
|
|
|
| save_checkpoint(global_step, 0.0f);
|
| std::cerr << "Training complete. Total steps: " << global_step << std::endl;
|
| }
|
|
|
| void TrainLM::save_checkpoint(size_t step, float loss) {
|
| std::string path = cfg_.output_dir + "/lm_head_step" + std::to_string(step) + ".nfv1";
|
|
|
| #ifdef _WIN32
|
| std::string mkdir_cmd = "if not exist \"" + cfg_.output_dir + "\" mkdir \"" + cfg_.output_dir + "\"";
|
| system(mkdir_cmd.c_str());
|
| #else
|
| std::string mkdir_cmd = "mkdir -p " + cfg_.output_dir;
|
| system(mkdir_cmd.c_str());
|
| #endif
|
|
|
| std::ofstream ofs(path, std::ios::binary);
|
| if (!ofs) {
|
| std::cerr << "[ERROR] Cannot save checkpoint: " << path << std::endl;
|
| return;
|
| }
|
|
|
| ofs.write("LMH2", 4);
|
|
|
| auto sl = [&ofs](const std::string& n, const Tensor& t) {
|
| uint32_t nl = static_cast<uint32_t>(n.size());
|
| ofs.write(reinterpret_cast<const char*>(&nl), 4);
|
| ofs.write(n.data(), nl);
|
| uint32_t nd = static_cast<uint32_t>(t.shape_.size());
|
| ofs.write(reinterpret_cast<const char*>(&nd), 4);
|
| for (auto d : t.shape_) {
|
| uint32_t dd = static_cast<uint32_t>(d);
|
| ofs.write(reinterpret_cast<const char*>(&dd), 4);
|
| }
|
| uint32_t ds = static_cast<uint32_t>(t.data_size_);
|
| ofs.write(reinterpret_cast<const char*>(&ds), 4);
|
| ofs.write(reinterpret_cast<const char*>(t.data_.get()), ds);
|
| };
|
|
|
| sl("w_embed", lm_head_->w_embed_);
|
| sl("w_pos", lm_head_->w_pos_);
|
| sl("dw_kernel", lm_head_->dw_kernel_);
|
| sl("pw_conv.weight", lm_head_->pw_conv_->weight);
|
| sl("sae_encode.weight", lm_head_->sae_w_encode_->weight);
|
| sl("sae_decode.weight", lm_head_->sae_w_decode_->weight);
|
| sl("ntm_read.weight", lm_head_->ntm_w_read_->weight);
|
| sl("ntm_write.weight", lm_head_->ntm_w_write_->weight);
|
| sl("ntm_erase.weight", lm_head_->ntm_w_erase_->weight);
|
| sl("ntm_memory", lm_head_->ntm_memory_);
|
| sl("w_proj.weight", lm_head_->w_proj_->weight);
|
| sl("w_proj.bias", lm_head_->w_proj_->bias);
|
| if (lm_head_->bridge_) {
|
| sl("bridge.weight", lm_head_->bridge_->weight);
|
| sl("bridge.bias", lm_head_->bridge_->bias);
|
| }
|
| sl("w_out.weight", lm_head_->w_out_->weight);
|
| if (lm_head_->w_out_->bias.data_) sl("w_out.bias", lm_head_->w_out_->bias);
|
| sl("ln.weight", lm_head_->ln_->weight);
|
| sl("ln.bias", lm_head_->ln_->bias);
|
| for (size_t i = 0; i < lm_head_->attn_layers_.size(); ++i) {
|
| std::string prefix = "attn" + std::to_string(i) + ".";
|
| sl(prefix + "w_q.weight", lm_head_->attn_layers_[i]->w_q->weight);
|
| sl(prefix + "w_q.bias", lm_head_->attn_layers_[i]->w_q->bias);
|
| sl(prefix + "w_k.weight", lm_head_->attn_layers_[i]->w_k->weight);
|
| sl(prefix + "w_k.bias", lm_head_->attn_layers_[i]->w_k->bias);
|
| sl(prefix + "w_v.weight", lm_head_->attn_layers_[i]->w_v->weight);
|
| sl(prefix + "w_v.bias", lm_head_->attn_layers_[i]->w_v->bias);
|
| sl(prefix + "w_out.weight", lm_head_->attn_layers_[i]->w_out->weight);
|
| sl(prefix + "w_out.bias", lm_head_->attn_layers_[i]->w_out->bias);
|
| sl(prefix + "norm.weight", lm_head_->attn_layers_[i]->norm->weight);
|
| sl(prefix + "norm.bias", lm_head_->attn_layers_[i]->norm->bias);
|
| }
|
|
|
| ofs.close();
|
| std::cerr << "[CHECKPOINT] Saved: " << path << " (loss=" << loss << ")" << std::endl;
|
| }
|
|
|
| } |
