cwenzi/neuroflow-cpp

1#!/usr/bin/env python3 -u2"""3NeuroFlow 蒸馏训练 — 用教师生成的数据训练学生模型4 5流程:6  1. 加载教师数据 (JSONL: {"prompt":..., "completion":...} 或纯文本)7  2. 加载 NeuroFlow 学生模型 (NF + LM head)8  3. 用 Cross-Entropy 训练学生预测教师文本 (滑动窗口多token预测)9  4. 保存 checkpoint (LMH1/LMH2 兼容格式)10 11用法:12  python3 -u scripts/train_distill.py \13    --teacher-data teacher_data.jsonl \14    --nf-model checkpoint/model.nfv1 \15    --lm-model checkpoint/lm_head.nfv1 \16    --tokenizer configs/tokenizer_128k.json \17    --output ./distill_output \18    --epochs 5 --lr 5e-6 --batch-size 32 \19    --train-nf   # 可选: 同时训练NF权重20 21注意: 必须使用 python3 -u 运行，或在后台运行时设置 PYTHONUNBUFFERED=122"""23 24import os25os.environ['PYTHONUNBUFFERED'] = '1'26 27import argparse, json, struct, sys, time, math28import numpy as np29 30sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))31from infer_full import load_nfv1, load_lmh1, load_tokenizer, encode, layernorm, gelu, softmax32 33 34def forward_with_cache(token_ids, nf_w, lm_w, vocab_size=128000):35    d_model = nf_w['input_proj.weight'].shape[1]36    hidden_dim = nf_w['input_proj.weight'].shape[0]37    cache = {}38 39    x = np.zeros(d_model, dtype=np.float32)40    copy_len = min(len(token_ids), d_model)41    for j in range(copy_len):42        x[j] = float(token_ids[j]) / float(vocab_size)43 44    h = nf_w['input_proj.weight'] @ x + nf_w['input_proj.bias']45    cache['input_proj.x'] = x.copy()46    cache['input_proj.pre_norm'] = h.copy()47    h = layernorm(h, nf_w['input_proj_norm.weight'], nf_w['input_proj_norm.bias'])48    cache['input_proj.post_norm'] = h.copy()49    h = gelu(h)50    cache['input_proj.post_gelu'] = h.copy()51 52    g1 = gelu(nf_w['sn.gate1.weight'] @ h + nf_w['sn.gate1.bias'])53    gates = softmax(nf_w['sn.gate2.weight'] @ g1 + nf_w['sn.gate2.bias'])54    cache['gates'] = gates.copy()55 56    h_ecn = h.copy()57    cache['ecn.h0'] = h_ecn.copy()58    for i in range(12):59        h_ecn = gelu(nf_w[f'ecn.dlpfc{i}.weight'] @ h_ecn + nf_w[f'ecn.dlpfc{i}.bias'])60    cache['ecn.last'] = h_ecn.copy()61 62    vmpfc = gelu(nf_w['ecn.vmpfc1.weight'] @ h_ecn + nf_w['ecn.vmpfc1.bias'])63    decision = nf_w['ecn.vmpfc2.weight'] @ vmpfc + nf_w['ecn.vmpfc2.bias']64    cache['decision'] = decision.copy()65 66    mem_encoded = nf_w['memory.encode.weight'] @ h + nf_w['memory.encode.bias']67    cache['mem_encoded'] = mem_encoded.copy()68 69    dmn_enc = gelu(nf_w['dmn.mem_encoder1.weight'] @ mem_encoded + nf_w['dmn.mem_encoder1.bias'])70    dmn_latent = nf_w['dmn.mem_encoder2.weight'] @ dmn_enc + nf_w['dmn.mem_encoder2.bias']71    cache['dmn_latent'] = dmn_latent.copy()72 73    assoc_outs = []74    for i in range(8):75        a1 = gelu(nf_w[f'dmn.head{i}.1.weight'] @ dmn_latent + nf_w[f'dmn.head{i}.1.bias'])76        a2 = nf_w[f'dmn.head{i}.2.weight'] @ a1 + nf_w[f'dmn.head{i}.2.bias']77        assoc_outs.append(a2)78    dmn_vision = gelu(nf_w['dmn.future_proj1.weight'] @ np.concatenate(assoc_outs) + nf_w['dmn.future_proj1.bias'])79    cache['dmn_vision'] = dmn_vision.copy()80 81    mem_bank = nf_w['memory.bank']82    att = softmax(mem_encoded @ mem_bank.T)83    retrieved = att @ mem_bank84    mem_retrieved = nf_w['memory.retrieve.weight'] @ retrieved + nf_w['memory.retrieve.bias']85    cache['mem_retrieved'] = mem_retrieved.copy()86 87    ecn_w = decision * gates[0]88    dmn_w = dmn_vision * gates[1]89    dmn_w_pad = np.zeros(hidden_dim, dtype=np.float32)90    dmn_w_pad[:dmn_w.shape[0]] = dmn_w91    mem_w = np.zeros(hidden_dim, dtype=np.float32)92    mem_w[:mem_retrieved.shape[0]] = mem_retrieved93    combined = np.concatenate([ecn_w, dmn_w_pad, mem_w])94    cache['combined'] = combined.copy()95 96    fused = nf_w['output_fusion.down.weight'] @ combined + nf_w['output_fusion.down.bias']97    fused_pre_relu = layernorm(fused, nf_w['output_fusion.bn_norm.weight'], nf_w['output_fusion.bn_norm.bias'])98    cache['fusion.pre_relu'] = fused_pre_relu.copy()99    fused_relu = np.maximum(0, fused_pre_relu)100    cache['fusion.post_relu'] = fused_relu.copy()101    nf_output = nf_w['output_fusion.up.weight'] @ fused_relu + nf_w['output_fusion.up.bias']102    nf_output = layernorm(nf_output, nf_w['output_fusion.norm.weight'], nf_w['output_fusion.norm.bias'])103    cache['nf_output'] = nf_output.copy()104 105    bridge_h = lm_w['bridge.weight'] @ nf_output + lm_w['bridge.bias']106    cache['bridge_h'] = bridge_h.copy()107 108    projected = lm_w['w_proj.weight'] @ bridge_h + lm_w['w_proj.bias']109    cache['projected'] = projected.copy()110 111    logits = lm_w['w_embed'] @ projected112    cache['logits'] = logits.copy()113 114    return logits, cache115 116 117 118def distill_step(token_ids, nf_w, lm_w, lr, vocab_size=128000, grad_clip=4.0,119                 train_nf=False, max_predictions=0):120    """单步蒸馏训练: 滑动窗口多token预测 + 反向传播121 122    对序列中每个位置 t，用 token_ids[:t+1] 前向预测 token_ids[t+1]。123    NF模型是"序列→单向量"架构，无法像Transformer那样单次前向获取所有位置hidden states，124    因此每个位置需要独立前向。max_predictions 限制每样本预测位置数以控制性能开销。125 126    --train-nf 反向传播路径止于 output_fusion.up/down，未穿过 ECN/DMN/Memory/SN gate。127    设计为渐进式解冻: 先训练 LM head + bridge，再解冻 output_fusion，最后解冻更深层。128    """129    seq_len = len(token_ids)130    if seq_len < 2:131        return 0.0132 133    total_loss = 0.0134    accum_lm_grads = {}135    accum_nf_grads = {}136    num_preds = 0137 138    positions = list(range(seq_len - 1))139    if max_predictions > 0 and len(positions) > max_predictions:140        step = max(1, len(positions) // max_predictions)141        positions = positions[::step][:max_predictions]142 143    for t in positions:144        prefix = token_ids[:t + 1]145        target_id = token_ids[t + 1]146        if target_id >= vocab_size:147            continue148 149        logits, cache = forward_with_cache(prefix, nf_w, lm_w, vocab_size)150 151        max_val = logits.max()152        exp_vals = np.exp(logits - max_val)153        sum_exp = exp_vals.sum()154        probs = exp_vals / sum_exp155 156        p_target = max(probs[target_id], 1e-10)157        total_loss += -math.log(p_target)158        num_preds += 1159 160        grad_logits = probs.copy()161        grad_logits[target_id] -= 1.0162 163        grad_w_embed = np.outer(grad_logits, cache['projected'])164        grad_projected = lm_w['w_embed'].T @ grad_logits165 166        grad_w_proj_weight = np.outer(grad_projected, cache['bridge_h'])167        grad_w_proj_bias = grad_projected.copy()168        grad_bridge_h = lm_w['w_proj.weight'].T @ grad_projected169 170        grad_bridge_weight = np.outer(grad_bridge_h, cache['nf_output'])171        grad_bridge_bias = grad_bridge_h.copy()172 173        step_lm_grads = {174            'w_embed': grad_w_embed,175            'w_proj.weight': grad_w_proj_weight,176            'w_proj.bias': grad_w_proj_bias,177            'bridge.weight': grad_bridge_weight,178            'bridge.bias': grad_bridge_bias,179        }180 181        for name, grad in step_lm_grads.items():182            if name not in accum_lm_grads:183                accum_lm_grads[name] = np.zeros_like(lm_w[name])184            accum_lm_grads[name] += grad185 186        if train_nf:187            grad_nf_output = lm_w['bridge.weight'].T @ grad_bridge_h188            grad_fused_relu = nf_w['output_fusion.up.weight'].T @ grad_nf_output189            grad_fused_pre_relu = grad_fused_relu * (cache['fusion.pre_relu'] > 0).astype(np.float32)190 191            step_nf_grads = {192                'output_fusion.up.weight': np.outer(grad_nf_output, cache['fusion.post_relu']),193                'output_fusion.up.bias': grad_nf_output.copy(),194                'output_fusion.down.weight': np.outer(grad_fused_pre_relu, cache['combined']),195                'output_fusion.down.bias': grad_fused_pre_relu.copy(),196            }197 198            for name, grad in step_nf_grads.items():199                if name in nf_w and nf_w[name].shape == grad.shape:200                    if name not in accum_nf_grads:201                        accum_nf_grads[name] = np.zeros_like(nf_w[name])202                    accum_nf_grads[name] += grad203 204    if num_preds == 0:205        return 0.0206 207    total_loss /= num_preds208 209    all_grads = {}210    for name, grad in accum_lm_grads.items():211        all_grads[f'lm.{name}'] = grad / num_preds212    for name, grad in accum_nf_grads.items():213        all_grads[f'nf.{name}'] = grad / num_preds214 215    total_norm = 0.0216    for g in all_grads.values():217        total_norm += np.sum(g ** 2)218    total_norm = math.sqrt(total_norm)219 220    clip_scale = 1.0221    if total_norm > grad_clip and grad_clip > 0:222        clip_scale = grad_clip / total_norm223 224    effective_lr = lr * clip_scale225    for name, grad in all_grads.items():226        if name.startswith('lm.'):227            key = name[3:]228            if key in lm_w and lm_w[key].shape == grad.shape:229                lm_w[key] -= effective_lr * grad230        elif name.startswith('nf.'):231            key = name[3:]232            if key in nf_w and nf_w[key].shape == grad.shape:233                nf_w[key] -= effective_lr * grad234 235    return total_loss236 237 238def save_lmh1(path, lm_w):239    with open(path, 'wb') as f:240        f.write(b'LMH1')241        for name, arr in lm_w.items():242            name_bytes = name.encode('utf-8')243            f.write(struct.pack('<I', len(name_bytes)))244            f.write(name_bytes)245            f.write(struct.pack('<I', len(arr.shape)))246            for d in arr.shape:247                f.write(struct.pack('<I', d))248            data = arr.astype(np.float32).tobytes()249            f.write(struct.pack('<I', len(data)))250            f.write(data)251        f.write(struct.pack('<I', 0))252 253 254def save_nfv1(path, nf_w):255    with open(path, 'wb') as f:256        f.write(b'NFv1')257        for name, arr in nf_w.items():258            name_bytes = name.encode('utf-8')259            f.write(struct.pack('<I', len(name_bytes)))260            f.write(name_bytes)261            f.write(struct.pack('<I', len(arr.shape)))262            for d in arr.shape:263                f.write(struct.pack('<I', d))264            data = arr.astype(np.float32).tobytes()265            f.write(struct.pack('<I', len(data)))266            f.write(data)267        f.write(struct.pack('<I', 0))268 269 270def main():271    parser = argparse.ArgumentParser(description='NeuroFlow Distillation Training')272    parser.add_argument('--teacher-data', required=True, help='教师数据 (JSONL/TXT)')273    parser.add_argument('--nf-model', required=True, help='学生 NF 模型路径')274    parser.add_argument('--lm-model', required=True, help='学生 LM head 路径')275    parser.add_argument('--tokenizer', required=True, help='分词器路径')276    parser.add_argument('--output', default='./distill_output', help='输出目录')277    parser.add_argument('--epochs', type=int, default=5)278    parser.add_argument('--lr', type=float, default=5e-6)279    parser.add_argument('--batch-size', type=int, default=32)280    parser.add_argument('--save-interval', type=int, default=500)281    parser.add_argument('--grad-clip', type=float, default=4.0)282    parser.add_argument('--train-nf', action='store_true', help='同时训练NF权重(默认只训练LM head)')283    parser.add_argument('--resume', default='', help='断点续训: 指定checkpoint目录')284    parser.add_argument('--max-predictions', type=int, default=8, help='每样本最大预测位置数(0=全部, 默认8)')285    args = parser.parse_args()286 287    os.makedirs(args.output, exist_ok=True)288 289    print("加载教师数据...")290    samples = []291    with open(args.teacher_data, 'r', encoding='utf-8') as f:292        for line in f:293            line = line.strip()294            if not line:295                continue296            try:297                rec = json.loads(line)298                text = rec.get('prompt', '') + rec.get('completion', '')299            except json.JSONDecodeError:300                text = line301            if len(text) >= 10:302                samples.append(text)303    print(f"   {len(samples)} 个样本")304 305    print("加载分词器...")306    vocab, id2token, merge_ranks = load_tokenizer(args.tokenizer)307    print(f"   词表: {len(vocab)} tokens")308 309    print("分词...")310    tokenized = []311    total_tokens = 0312    t0 = time.time()313    progress_interval = max(1, len(samples) // 20)314    for si, text in enumerate(samples):315        ids = encode(text, vocab, merge_ranks, max_len=128)316        if len(ids) >= 4:317            tokenized.append(ids)318            total_tokens += len(ids)319        if (si + 1) % progress_interval == 0 or si == len(samples) - 1:320            elapsed = time.time() - t0321            pct = (si + 1) * 100 // len(samples)322            rate = (si + 1) / max(elapsed, 0.01)323            eta = (len(samples) - si - 1) / max(rate, 0.01)324            print(f"   分词进度: {si+1}/{len(samples)} ({pct}%) | "325                  f"{rate:.0f} samples/s | ETA {eta:.0f}s | "326                  f"tokens={total_tokens:,}")327    avg_len = total_tokens / max(len(tokenized), 1)328    elapsed = time.time() - t0329    print(f"   完成: {total_tokens:,} tokens ({avg_len:.0f} avg/sample) | 耗时 {elapsed:.1f}s")330 331    print("加载学生模型...")332    nf_w = load_nfv1(args.nf_model)333    lm_w = load_lmh1(args.lm_model)334    d_model = nf_w['input_proj.weight'].shape[1]335    hidden_dim = nf_w['input_proj.weight'].shape[0]336    print(f"   NF: {len(nf_w)}层 | LM: {len(lm_w)}层 | d_model={d_model} hidden={hidden_dim}")337 338    start_step = 0339    start_epoch = 0340    if args.resume:341        print(f"断点续训: {args.resume}")342        nf_w = load_nfv1(f"{args.resume}/model.nfv1")343        lm_w = load_lmh1(f"{args.resume}/lm_head.nfv1")344        state_path = f"{args.resume}/training_state.json"345        if os.path.exists(state_path):346            with open(state_path) as sf:347                state = json.load(sf)348            start_step = state.get('step', 0)349            start_epoch = state.get('epoch', 1) - 1350            print(f"   恢复: step={start_step}, epoch={start_epoch+1}")351 352    mode_str = "LM+NF" if args.train_nf else "LM only"353    print(f"\n开始蒸馏训练 ({args.epochs} epochs, lr={args.lr}, batch={args.batch_size}, "354          f"grad_clip={args.grad_clip}, mode={mode_str})")355    print(f"   每样本预测位置数: {args.max_predictions if args.max_predictions > 0 else int(avg_len)} (max_predictions={args.max_predictions})")356    print("=" * 60)357 358    global_step = start_step359    for epoch in range(start_epoch, args.epochs):360        epoch_loss = 0.0361        steps = 0362 363        indices = list(range(len(tokenized)))364        np.random.shuffle(indices)365 366        for i in range(0, len(indices), args.batch_size):367            batch_indices = indices[i:i + args.batch_size]368 369            batch_loss = 0.0370            for idx in batch_indices:371                loss = distill_step(tokenized[idx], nf_w, lm_w, args.lr,372                                    vocab_size=len(vocab), grad_clip=args.grad_clip,373                                    train_nf=args.train_nf, max_predictions=args.max_predictions)374                batch_loss += loss375 376            batch_loss /= len(batch_indices)377            epoch_loss += batch_loss378            steps += 1379            global_step += 1380 381            if steps % 10 == 0:382                print(f"  [Epoch {epoch+1}][Step {global_step}] loss={batch_loss:.4f}")383 384            if args.save_interval > 0 and global_step % args.save_interval == 0:385                ckpt_dir = f"{args.output}/step_{global_step}"386                os.makedirs(ckpt_dir, exist_ok=True)387                save_nfv1(f"{ckpt_dir}/model.nfv1", nf_w)388                save_lmh1(f"{ckpt_dir}/lm_head.nfv1", lm_w)389                state = {"step": global_step, "epoch": epoch + 1, "loss": batch_loss, "lr": args.lr}390                with open(f"{ckpt_dir}/training_state.json", 'w') as sf:391                    json.dump(state, sf, indent=2)392                print(f"  Checkpoint: {ckpt_dir}")393 394        avg_loss = epoch_loss / max(steps, 1)395        print(f"=== Epoch {epoch+1} 完成, avg_loss={avg_loss:.4f} ===\n")396 397    final_dir = f"{args.output}/final"398    os.makedirs(final_dir, exist_ok=True)399    save_nfv1(f"{final_dir}/model.nfv1", nf_w)400    save_lmh1(f"{final_dir}/lm_head.nfv1", lm_w)401    print(f"最终模型已保存: {final_dir}")402    print("蒸馏训练完成")403 404 405if __name__ == '__main__':406    main()407