File size: 8,016 Bytes
26d5b81 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | #ifndef NOMINMAX
#define NOMINMAX
#endif
#include <iostream>
#include <vector>
#include "neuroflow/generative.hpp"
#include "neuroflow/model.hpp"
#include "neuroflow/tensor.hpp"
using namespace neuroflow;
int main() {
std::cerr << "=== CausalLMHead 通路测试 ===" << std::endl;
CausalLMConfig cfg;
cfg.vocab_size = 100;
cfg.d_model = 32;
cfg.max_seq_len = 16;
cfg.causal_window_size = 4;
cfg.sae_k = 8;
cfg.ntm_memory_slots = 4;
cfg.use_mla = false;
cfg.weight_tying = true;
cfg.num_attn_layers = 1;
cfg.num_attn_heads = 2;
cfg.pooling = "mean";
std::cerr << "1. 构造 CausalLMHead..." << std::endl;
CausalLMHead lm_head(cfg);
lm_head.tie_weights();
std::cerr << " OK" << std::endl;
std::cerr << "2. forward (推理)..." << std::endl;
std::vector<size_t> token_ids = {1, 5, 10, 20, 30};
Tensor logits = lm_head.forward(token_ids);
std::cerr << " logits shape: [" << logits.shape_[0] << "," << logits.shape_[1] << "]" << std::endl;
std::cerr << " OK" << std::endl;
std::cerr << "3. forward_for_training..." << std::endl;
std::vector<size_t> train_ids = {1, 5, 10, 20};
Tensor train_logits = lm_head.forward_for_training(train_ids);
std::cerr << " train_logits shape: [" << train_logits.shape_[0] << "," << train_logits.shape_[1] << "]" << std::endl;
std::cerr << " OK" << std::endl;
std::cerr << "4. backward_from_logits..." << std::endl;
size_t target_id = 30;
const float* pred = train_logits.as_fp32();
float max_val = -1e30f;
for (size_t j = 0; j < cfg.vocab_size; ++j) {
if (pred[j] > max_val) max_val = pred[j];
}
float sum_exp = 0.0f;
for (size_t j = 0; j < cfg.vocab_size; ++j) {
sum_exp += std::exp(pred[j] - max_val);
}
float loss = -(pred[target_id] - max_val - std::log(sum_exp));
std::cerr << " loss = " << loss << std::endl;
Tensor logits_grad({1, cfg.vocab_size}, QuantType::FP32);
float* lg = logits_grad.as_fp32();
for (size_t j = 0; j < cfg.vocab_size; ++j) {
float softmax_val = std::exp(pred[j] - max_val) / sum_exp;
lg[j] = softmax_val;
if (j == target_id) lg[j] -= 1.0f;
}
auto grads = lm_head.backward_from_logits(logits_grad);
std::cerr << " attn_grads.size() = " << grads.attn_grads.size() << std::endl;
std::cerr << " w_proj_weight_grad shape: [" << grads.w_proj_weight_grad.shape_[0] << "," << grads.w_proj_weight_grad.shape_[1] << "]" << std::endl;
std::cerr << " embed_grad shape: [" << grads.embed_grad.shape_[0] << "," << grads.embed_grad.shape_[1] << "]" << std::endl;
std::cerr << " dw_kernel_grad shape: [" << grads.dw_kernel_grad.shape_[0] << "," << grads.dw_kernel_grad.shape_[1] << "]" << std::endl;
// NaN diagnostics
auto check_nan = [](const std::string& name, const Tensor& t) {
if (t.numel() == 0 || t.data_size_ == 0) return;
const float* d = t.as_fp32();
size_t nan_count = 0, inf_count = 0;
float max_abs = 0.0f;
for (size_t i = 0; i < t.numel(); ++i) {
if (std::isnan(d[i])) nan_count++;
else if (std::isinf(d[i])) inf_count++;
else max_abs = std::max(max_abs, std::abs(d[i]));
}
std::cerr << " [DIAG] " << name << ": nan=" << nan_count << " inf=" << inf_count << " max_abs=" << max_abs << std::endl;
};
check_nan("w_proj_weight_grad", grads.w_proj_weight_grad);
std::cerr << " [DEBUG] w_out_weight_grad numel=" << grads.w_out_weight_grad.numel() << std::endl;
check_nan("w_out_weight_grad", grads.w_out_weight_grad);
check_nan("ln_weight_grad", grads.ln_weight_grad);
check_nan("sae_encode_weight_grad", grads.sae_encode_weight_grad);
check_nan("sae_decode_weight_grad", grads.sae_decode_weight_grad);
check_nan("ntm_read_weight_grad", grads.ntm_read_weight_grad);
check_nan("dw_kernel_grad", grads.dw_kernel_grad);
check_nan("pw_conv_weight_grad", grads.pw_conv_weight_grad);
check_nan("embed_grad", grads.embed_grad);
if (!grads.attn_grads.empty()) {
check_nan("attn0.w_q_weight_grad", grads.attn_grads[0].w_q_weight_grad);
check_nan("attn0.w_k_weight_grad", grads.attn_grads[0].w_k_weight_grad);
check_nan("attn0.w_v_weight_grad", grads.attn_grads[0].w_v_weight_grad);
check_nan("attn0.w_out_weight_grad", grads.attn_grads[0].w_out_weight_grad);
check_nan("attn0.input_grad", grads.attn_grads[0].input_grad);
}
std::cerr << " OK" << std::endl;
std::cerr << "5. apply_lm_gradients..." << std::endl;
lm_head.apply_lm_gradients(grads, 1e-5f);
std::cerr << " OK" << std::endl;
std::cerr << "6. 第二次 forward_for_training (验证梯度更新有效)..." << std::endl;
Tensor train_logits2 = lm_head.forward_for_training(train_ids);
const float* pred2 = train_logits2.as_fp32();
float max_val2 = -1e30f;
for (size_t j = 0; j < cfg.vocab_size; ++j) {
if (pred2[j] > max_val2) max_val2 = pred2[j];
}
float sum_exp2 = 0.0f;
for (size_t j = 0; j < cfg.vocab_size; ++j) {
sum_exp2 += std::exp(pred2[j] - max_val2);
}
float loss2 = -(pred2[target_id] - max_val2 - std::log(sum_exp2));
std::cerr << " loss2 = " << loss2 << std::endl;
if (loss2 != loss) {
std::cerr << " 梯度更新有效 (loss变化)" << std::endl;
} else {
std::cerr << " 警告: loss未变化" << std::endl;
}
std::cerr << "7. 保存 LM Head..." << std::endl;
{
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("D:/neuroflow-C++/test_run/lm_head_test.nfv1", std::ios::binary);
o.write("LMH2", 4);
sl(o, "w_embed", lm_head.w_embed_);
sl(o, "w_proj.weight", lm_head.w_proj_->weight);
sl(o, "w_out.weight", lm_head.w_out_->weight);
sl(o, "dw_kernel", lm_head.dw_kernel_);
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, "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();
}
std::cerr << " OK" << std::endl;
std::cerr << "=== 所有通路测试通过! ===" << std::endl;
return 0;
} |