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