Tiny-R2: A Hybrid Architecture Combining SWA, CSA, HCA, mHC and DSMoE Under the DeepSeek V4 Design Paradigm

Tiny-R2 模型架构与训练流程文档

📋 目录

项目概述

模型架构总览

快速启动

2.1 安装依赖

pip install tiktoken datasets transformers huggingface_hub
pip install git+https://github.com/KellerJordan/Muon
pip install --upgrade transformers
hf auth login --force

2.2 启动训练

2.2.1支持采用Agent进行自主观察和训练调整lr、clip超参,以实现更加稳定的智能化训练;默认不开启;开启后需要使用gemini的api key

python train.py --n_layer 6 --n_embd 1536 --hc 'True' --mhc 'True' --n_experts 8 --max_iters 10000 --attention_types 'Sparse' --batch_size 8 --ctx_len 2048 --hf_dataset 'karpathy/climbmix-400b-shuffle' --resume True --save_best_only True

2.2.2 设置 --use_agent_observe开启Agent智能化训练供能,需要填入你的geimini的api key

python train.py --n_layer 6 --n_embd 1536 --hc 'True' --mhc 'True' --n_experts 8 --max_iters 10000 --attention_types 'Sparse' --batch_size 8 --ctx_len 2048 --hf_dataset 'karpathy/climbmix-400b-shuffle' --resume True --save_best_only True --use_agent_observe True --gemini_api_key "your gemini apikey"

2.3 验证模型训练效果PPL

python evaluate.py --checkpoint checkpoints/best_model_step_xxx.pt 

2.4 启动OPD在线蒸馏(以下三选一即可)

2.4.1使用Qwen3.5-9B模型作为教师模型、Qwen3.5-0.8B作为学生模型进行OPD训练,用RAG增加教师模型,RAG数据集来自问答数据集集medquad,可以复现Readme中的结果;(目前实测在pubmed数据集上Qwen3.5-0.8B可以提升12%以上准确率)

python opd_train.py --dataset pubmed_qa --hf_teacher_model Qwen/Qwen3.5-9B --student_model_name Qwen/Qwen3.5-0.8B --batch_size 2 --grad_accum_steps 4

2.4.2 使用Qwen3.5-9B模型作为教师模型、Qwen3.5-0.8B作为学生模型进行OPD训练,用RAG增加教师模型;--rag_corpus_path外挂RAG数据集,--custom_qa_path 自定义问题集

python opd_train.py --hf_teacher_model Qwen/Qwen3.5-9B --student_model_name Qwen/Qwen3.5-0.8B --batch_size 2 --grad_accum_steps 4 --custom_qa_path baoxianqa.jsonl --rag_corpus_path baoxianqa.txt
 

2.4.3 使用Qwen3.5-9B模型作为教师模型、Tiny-R2作为学生模型进行OPD训练,用RAG增加教师模型;--rag_corpus_path外挂RAG数据集,--custom_qa_path 自定义问题集

python opd_train.py --hf_teacher_model Qwen/Qwen3.5-9B --student_model_name tiny-r2 --student_ckpt checkpoints/best_model_step_40.pt --enable_rag_teacher --batch_size 2 --grad_accum_steps 4 --custom_qa_path baoxianqa.jsonl --rag_corpus_path baoxianqa.txt 

核心组件详解

3.1 注意力机制

3.1.1 CausalSelfAttention (Full Attention)

class CausalSelfAttention(nn.Module):
 def __init__(self, config):
 # Projections: Q, K, V from single linear
 self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
 self.c_proj = nn.Linear(config.n_embd, config.n_embd)
 
 # Value residual connections
 self.v_residual = config.v_residual
 self.lamb1 = nn.Parameter(torch.tensor(0.5))
 self.lamb2 = nn.Parameter(torch.tensor(0.5))
 
 # Flash Attention support
 self.flash = hasattr(F, "scaled_dot_product_attention")

3.1.2 HCA-CSA Hybrid Attention

3.2 HCA-CSA 混合注意力

架构流程

 Input x
 ↓
┌──────────────────────────────────────────────────────────────┐
│ Query Preparation (HCA style) │
│ compress_q → q_norm → decompress_q → RoPE → Query │
└──────────────────────────────────────────────────────────────┘
 ↓
┌──────────────────┬──────────────────┬──────────────────────┐
│ Branch 1 │ Branch 2 │ Branch 3 │
│ Compression │ Selection │ Sliding Window │
│ (HCA) │ (CSA) │ (SWA) │
├──────────────────┼──────────────────┼──────────────────────┤
│ compress_kv │ importance_score │ window_k/v │
│ kv_norm │ topk selection │ sliding_window │
│ decompress_k/v │ selection_k/v │ RoPE │
│ k_rope │ RoPE │ │
│ K/V Recombine │ K/V Selected │ K/V Window │
└──────────────────┴──────────────────┴──────────────────────┘
 ↓ ↓ ↓
┌──────────────────────────────────────────────────────────────┐
│ Attention Computation │
│ Attention 1: (Q @ K1.T) @ V1 │
│ Attention 2: (Q @ K2.T) @ V2 │
│ Attention 3: (Q @ K3.T) @ V3 │
└──────────────────────────────────────────────────────────────┘
 ↓
branch_gate (Linear + Softmax) → Weighted Sum
 ↓
proj (Linear) → res_dropout → Output

关键参数

# HCA 参数
self.v_head_dim = 32
self.kv_lora_rank = 32
self.q_lora_rank = 3 * self.kv_lora_rank
self.rope_head_dim = 64
self.nope_head_dim = 32
# CSA 参数
self.block_size = config.block_size # Token压缩块大小
self.window_size = config.window_size # 滑动窗口大小
self.num_tokens_to_keep = config.num_tokens_to_keep # 选择保留的token数

3.3 前馈网络与 MoE

3.3.1 MLP

class MLP(nn.Module):
 def __init__(self):
 self.c_fc = nn.Linear(n_embd, 4 * n_embd)
 self.c_proj = nn.Linear(4 * n_embd, n_embd)
 
 def forward(self, x):
 x = self.c_fc(x)
 x = F.relu(x).square() # ReLU squared
 x = self.c_proj(x)
 return x

3.3.2 DSMoE (DeepSeek Mixture of Experts)

Input x [B, T, C]
 ↓
Gate Network (Linear + UnitCenteredNoise)
 ↓
Softmax → Top-k Selection
 ↓
┌─────────────────────────────────────────────────────────────┐
│ Expert Networks │
│ ┌──────────────┐ ┌──────────┐ ┌──────────┐ ┌────────┐ │
│ │ Shared Exp 0 │ │ Expert 1 │ │ Expert 2 │ │ ... │ │
│ │ (Always On) │ │ (Top-k) │ │ (Top-k) │ │ (Top-k)│ │
│ └──────────────┘ └──────────┘ └──────────┘ └────────┘ │
└─────────────────────────────────────────────────────────────┘
 ↓
Weighted Sum of Expert Outputs
 ↓
Output [B, T, C]

def moe_load_balance_loss(router_weights, num_experts):
 load = router_weights.sum(dim=0)
 load = load / load.sum()
 ideal = torch.full_like(load, 1.0 / num_experts)
 loss = num_experts * torch.sum((load - ideal) ** 2)
 return loss

3.4 Hyper-Connections

# 初始化 Hyper-Connections
self.init_hc, self.expand_stream, self.reduce_stream = \
 get_init_and_expand_reduce_stream_functions(
 config.hc_num_streams,
 num_fracs=config.hc_num_fracs,
 disable=config.hc_disable,
 )
# 在每个 Block 中使用
self.hc_attn = init_hc(
 dim=config.n_embd,
 branch=self.attn_branch,
 layer_index=index * 2,
 mhc=config.mhc,
 sinkhorn_iters=config.sinkhorn_iters,
 sinkhorn_tau=config.sinkhorn_tau,
)

训练流程

4.1 初始化阶段

Parse Arguments → Update Config → Init WandB → Setup Distributed → Setup AMP

4.2 数据准备

Load HF Dataset (flytech/python-codes-25k)
 ↓
Init GPT2 Tokenizer
 ↓
Create TokenBuffer

4.3 模型初始化

Create Transformer
 ↓
Configure Optimizers (Muon + AdamW)
 ↓
Create LR Scheduler (Warmup + Cosine)

4.4 训练循环

For iter in range(max_iters):
 │
 ├── For step in grad_accum_steps:
 │ ├── Get Batch (TokenBuffer)
 │ ├── Forward Pass (model)
 │ ├── Backward Pass (scaler.scale)
 │ └── Collect Router Weights
 │
 ├── Gradient Clipping (clip_grad_norm_)
 ├── Optimizer Steps (Muon + AdamW)
 ├── Update Scaler (scaler.update)
 ├── LR Scheduler Step
 ├── Update Expert Biases (load balancing)
 └── Log Metrics (WandB)

4.5 评估与保存

If iter % eval_interval == 0:
 ├── Estimate Loss (eval mode)
 ├── Save Checkpoint (if val_loss < 5.27)
 └── Log to WandB

4.6 优化器配置

def configure_optimizers(self, weight_decay, learning_rate, device):
 muon_params = [] # ≥2D parameters in blocks
 adamw_params = [] # Other parameters
 
 for name, param in self.named_parameters():
 if 'blocks' in name and param.ndim >= 2:
 muon_params.append(param)
 else:
 adamw_params.append(param)
 
 return [
 Muon(muon_params, lr=0.02, momentum=0.95),
 torch.optim.AdamW(adamw_params, lr=learning_rate, 
 betas=(0.90, 0.95), weight_decay=weight_decay)
 ]

关键技术特性

5.1 注意力机制对比

5.2 FFN 类型对比

5.3 核心配置参数

# 模型架构
n_embd = 512 # 嵌入维度
n_head = 8 # 注意力头数
n_layer = 8 # 层数
n_experts = 8 # 专家数量
num_exp = 2 # 每token激活的专家数
# 注意力配置
attention_types = ["FULL", "Spares", ...] # 每层注意力类型
attention_mode = ["FULL", "SWA", "CSA"] # 稀疏注意力模式
# Hyper-Connections
hc = True # 启用超连接
hc_num_streams = 4 # 流数量
# 训练
batch_size = 32
ctx_len = 512 # 上下文长度
lr = 1e-3
warmup_iters = 1000
max_iters = 100000

附录:图表索引

参考资料