cwenzi/neuroflow-cpp

1/**2 * 最小 NaN 诊断测试 — 绕过 DataLoader，直接测试 forward + CE loss3 */4#include <cmath>5#include <cstdio>6#include <cstdlib>7#include <random>8#include <vector>9 10#include "neuroflow/backprop.hpp"11#include "neuroflow/generative.hpp"12#include "neuroflow/model.hpp"13 14using namespace neuroflow;15 16int main() {17    // Load model config18    NeuroFlowModel::Config cfg;19    // Use defaults: d_model=512, hidden_dim=2048, output_dim=2048, vocab_size=12800020 21    NeuroFlowModel model(cfg);22    FullBackpropEngine backprop(model);23 24    // Load tokenizer25    BPETokenizer tok("configs/tokenizer_128k.json");26    printf("Tokenizer loaded: vocab=%zu\n", tok.vocab_size());27 28    // Create CausalLMHead29    CausalLMConfig lm_cfg;30    lm_cfg.vocab_size = cfg.vocab_size;31    lm_cfg.d_model = cfg.hidden_dim;  // 204832    lm_cfg.max_seq_len = cfg.max_seq_len;33    lm_cfg.causal_window_size = cfg.causal_window_size;34    lm_cfg.sae_k = cfg.sae_k;35    lm_cfg.ntm_memory_slots = cfg.ntm_memory_slots;36    lm_cfg.weight_tying = true;37    lm_cfg.num_attn_layers = cfg.lm_num_attn_layers;38    lm_cfg.pooling = cfg.lm_pooling;39    CausalLMHead lm_head(lm_cfg);40    lm_head.tie_weights();41 42    printf("CausalLMHead: d_model=%zu vocab=%zu\n", lm_cfg.d_model, lm_cfg.vocab_size);43 44    // Create bridge projection (matching train_v2.cpp)45    size_t d_model = lm_cfg.d_model;46    size_t hidden_dim = cfg.hidden_dim;47    Tensor lm_bridge_weight({d_model, hidden_dim}, QuantType::FP32);48    Tensor lm_bridge_bias({d_model}, QuantType::FP32);49    {50        float* bw = lm_bridge_weight.as_fp32();51        float scale = std::sqrt(2.0f / (hidden_dim + d_model));52        std::mt19937 br_rng(hidden_dim * 31 + d_model);53        std::uniform_real_distribution<float> br_dist(-scale, scale);54        for (size_t i = 0; i < lm_bridge_weight.numel(); ++i) bw[i] = br_dist(br_rng);55        memset(lm_bridge_bias.as_fp32(), 0, lm_bridge_bias.data_size_);56    }57 58    // Create random test data59    size_t batch_sz = 4;60    size_t seq_len = 16;61    std::mt19937 rng(42);62    std::uniform_int_distribution<size_t> id_dist(4, tok.vocab_size() - 1);63 64    printf("\n=== Testing %zu batches ===\n", batch_sz);65 66    for (int step = 0; step < 200; ++step) {67        // Create random token sequence68        std::vector<size_t> all_ids;69        for (size_t b = 0; b < batch_sz; ++b) {70            for (size_t s = 0; s < seq_len; ++s) {71                all_ids.push_back(id_dist(rng));72            }73        }74 75        // Create input tensor [batch_sz, input_dim=512]76        Tensor input({batch_sz, cfg.input_dim}, QuantType::FP32);77        float* inp = input.as_fp32();78        for (size_t b = 0; b < batch_sz; ++b) {79            for (size_t j = 0; j < cfg.input_dim; ++j) {80                inp[b * cfg.input_dim + j] = static_cast<float>(all_ids[b * seq_len + (j % seq_len)]) / 128000.0f;81            }82        }83 84        // Create target tensor [batch_sz, vocab_size]85        size_t vocab_sz = lm_cfg.vocab_size;86        Tensor target({batch_sz, vocab_sz}, QuantType::FP32);87        float* tgt = target.as_fp32();88        memset(tgt, 0, target.data_size_);89        for (size_t b = 0; b < batch_sz; ++b) {90            size_t next_id = all_ids[b * seq_len + 1]; // next token as target91            if (next_id < vocab_sz) tgt[b * vocab_sz + next_id] = 1.0f;92        }93 94        // Forward pass95        auto nf_output = backprop.forward_with_cache(input);96        const Tensor& nf_hidden = nf_output.output;97 98        // Bridge projection99        Tensor hidden_proj({batch_sz, d_model}, QuantType::FP32);100#ifdef USE_CBLAS101        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans,102                    batch_sz, d_model, hidden_dim,103                    1.0f, nf_hidden.as_fp32(), hidden_dim,104                    lm_bridge_weight.as_fp32(), hidden_dim,105                    0.0f, hidden_proj.as_fp32(), d_model);106#else107        {108            float* hp = hidden_proj.as_fp32();109            const float* nh = nf_hidden.as_fp32();110            const float* bw = lm_bridge_weight.as_fp32();111            for (size_t b = 0; b < batch_sz; ++b) {112                for (size_t j = 0; j < d_model; ++j) {113                    float sum = 0.0f;114                    for (size_t k = 0; k < hidden_dim; ++k) {115                        sum += nh[b * hidden_dim + k] * bw[j * hidden_dim + k];116                    }117                    hp[b * d_model + j] = sum;118                }119            }120        }121#endif122        // Add bias123        {124            float* hp = hidden_proj.as_fp32();125            const float* bb = lm_bridge_bias.as_fp32();126            for (size_t b = 0; b < batch_sz; ++b)127                for (size_t j = 0; j < d_model; ++j)128                    hp[b * d_model + j] += bb[j];129        }130 131        Tensor projected = lm_head.w_proj_->forward(hidden_proj);132        Tensor logits = lm_head.w_out_->forward(projected);133 134        // Check logits for NaN135        const float* pred = logits.as_fp32();136        size_t n = logits.numel();137        bool pred_nan = false;138        for (size_t i = 0; i < n; ++i) {139            if (!std::isfinite(pred[i])) { pred_nan = true; break; }140        }141        if (pred_nan) {142            printf("Step %d: logits NaN BEFORE loss!\n", step);143            return 1;144        }145 146        // Compute CE loss147        size_t dim = vocab_sz;148        float loss = 0.0f;149        for (size_t b = 0; b < batch_sz; ++b) {150            float max_val = -1e30f;151            for (size_t j = 0; j < dim; ++j) {152                if (pred[b * dim + j] > max_val) max_val = pred[b * dim + j];153            }154            float sum_exp = 0.0f;155            for (size_t j = 0; j < dim; ++j) {156                sum_exp += std::exp(pred[b * dim + j] - max_val);157            }158            if (!std::isfinite(sum_exp) || sum_exp <= 0) {159                printf("Step %d batch %zu: sum_exp=%f max_val=%f\n", step, b, sum_exp, max_val);160                return 1;161            }162            float log_sum_exp = max_val + std::log(sum_exp);163            if (!std::isfinite(log_sum_exp)) {164                printf("Step %d batch %zu: log_sum_exp=%f\n", step, b, log_sum_exp);165                return 1;166            }167 168            for (size_t j = 0; j < dim; ++j) {169                float softmax_val = std::exp(pred[b * dim + j] - max_val) / sum_exp;170                float t = tgt[b * dim + j];171                if (t > 0.5f) {172                    float contrib = pred[b * dim + j] - log_sum_exp;173                    loss -= contrib;174                    if (!std::isfinite(loss)) {175                        printf("Step %d batch %zu j=%zu: LOSS NaN! pred=%f log_sum_exp=%f contrib=%f\n",176                               step, b, j, pred[b*dim+j], log_sum_exp, contrib);177                        return 1;178                    }179                }180            }181        }182        loss /= batch_sz;183 184        // Backward185        Tensor output_grad({batch_sz, dim}, QuantType::FP32);186        float* og = output_grad.as_fp32();187        for (size_t b = 0; b < batch_sz; ++b) {188            float max_val = -1e30f;189            for (size_t j = 0; j < dim; ++j)190                if (pred[b*dim+j] > max_val) max_val = pred[b*dim+j];191            float sum_exp = 0.0f;192            for (size_t j = 0; j < dim; ++j)193                sum_exp += std::exp(pred[b*dim+j] - max_val);194            for (size_t j = 0; j < dim; ++j) {195                float softmax_val = std::exp(pred[b*dim+j] - max_val) / sum_exp;196                float t = tgt[b*dim+j];197                og[b*dim+j] = (softmax_val - t) / batch_sz;198            }199        }200 201        // Check output_grad202        for (size_t i = 0; i < output_grad.numel(); ++i) {203            if (!std::isfinite(og[i])) {204                printf("Step %d: output_grad NaN at i=%zu\n", step, i);205                return 1;206            }207        }208 209        // Backward through LM head210        // (simplified for test — just check that backward doesn't crash)211        auto nf_grads = backprop.backward(output_grad);212 213        // Check nf_grads for NaN214        auto check_nan = [](const Tensor& t, const char* name) {215            if (t.numel() == 0) return;216            const float* d = t.as_fp32();217            for (size_t i = 0; i < t.numel(); ++i) {218                if (!std::isfinite(d[i])) {219                    printf("  GRAD NaN in %s at i=%zu\n", name, i);220                    return;221                }222            }223        };224        check_nan(nf_grads.input_proj_weight_grad, "input_proj_weight");225        check_nan(nf_grads.output_fusion_up_weight_grad, "output_fusion_up_weight");226        check_nan(nf_grads.ecn_vmpfc2_weight_grad, "ecn_vmpfc2_weight");227        check_nan(nf_grads.sn_gate1_weight_grad, "sn_gate1_weight");228        check_nan(nf_grads.dmn_mem_encoder1_weight_grad, "dmn_mem_encoder1_weight");229        check_nan(nf_grads.mem_encode_proj_weight_grad, "mem_encode_proj_weight");230 231        if (step % 50 == 0) {232            float lmin = 1e30f, lmax = -1e30f;233            for (size_t i = 0; i < std::min(n, size_t(1000)); ++i) {234                if (pred[i] < lmin) lmin = pred[i];235                if (pred[i] > lmax) lmax = pred[i];236            }237            printf("Step %d: loss=%.4f logits=[%.4f, %.4f] OK\n", step, loss, lmin, lmax);238        }239    }240 241    printf("\nAll %d steps passed! No NaN detected.\n", 200);242    return 0;243}244