NLP第二阶段学习 Windows环境下纯CPU训练语言大模型

一、环境准备

1.0 安装Python和依赖

# 以管理员运行PowerShell

# 1. 安装Chocolatey包管理器
Set-ExecutionPolicy Bypass -Scope Process -Force; 
[System.Net.ServicePointManager]::SecurityProtocol = [System.Net.ServicePointManager]::SecurityProtocol -bor 3072; 
iex ((New-Object System.Net.WebClient).DownloadString('https://community.chocolatey.org/install.ps1'))

# 2. 安装基础工具
choco install -y git
choco install -y vscode
choco install -y notepadplusplus
choco install -y 7zip

# 3. 安装Python（推荐3.10，稳定且兼容性好）
choco install -y python310

# 4. 验证安装
python --version  # 应显示 3.10.x
pip --version
git --version

# 5. 创建虚拟环境
python -m venv llm_env
 
# 6. 激活虚拟环境 (Windows PowerShell)
 .\llm_env\Scripts\activate.bat

1.2 安装核心依赖

# 以管理员运行PowerShell

# 安装PyTorch CPU版本
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

# 安装其他依赖
pip install transformers datasets tokenizers accelerate
pip install sentencepiece tqdm numpy pandas
pip install psutil  # 用于监控系统资源

二、训练代码

2.1 项目结构

llm_project/
├── data/
│ └── train.txt # 训练数据
├── tokenizer/
│ └── tokenizer.json # 分词器
├── models/
│ └── checkpoint.pt # 模型检查点
├── src/
│ ├── config.py # 配置文件
│ ├── model.py # 模型定义
│ ├── tokenizer_train.py # 分词器训练
│ ├── dataset.py # 数据集
│ ├── train.py # 训练脚本
│ └── inference.py # 推理脚本
└── requirements.txt

2.2 配置文件 (src/config.py)

"""模型和训练配置"""

class ModelConfig:
    """模型架构配置"""
    vocab_size = 32000          # 词表大小
    hidden_size = 512          # 隐藏层维度 (CPU建议512-768)
    num_layers = 6               # Transformer层数
    num_heads = 8                # 注意力头数
    intermediate_size = 2048     # FFN中间层
    max_seq_length = 512         # 最大序列长度
    dropout = 0.1
    layer_norm_eps = 1e-12

class TrainingConfig:
    """训练配置"""
    # 训练参数
    learning_rate = 5e-4
    weight_decay = 0.01
    batch_size = 2               # CPU建议小batch
    num_epochs = 3
    gradient_accumulation_steps = 4  # 梯度累积步数
    
    # 保存和日志
    save_steps = 500
    log_steps = 10
    output_dir = "./models"
    
    # 数据
    data_file = "./data/train.txt"
    tokenizer_path = "./tokenizer/tokenizer.json"
    
    # 系统
    num_workers = 0  # Windows下用0避免多进程问题

# 根据CPU性能选择配置
CPU_CONFIGS = {
    "low": {        # 4核8GB内存
        "hidden_size": 256,
        "num_layers": 4,
        "num_heads": 4,
        "batch_size": 1,
    },
    "medium": {     # 8核16GB内存
        "hidden_size": 512,
        "num_layers": 6,
        "num_heads": 8,
        "batch_size": 2,
    },
    "high": {       # 16核32GB内存
        "hidden_size": 768,
        "num_layers": 8,
        "num_heads": 12,
        "batch_size": 4,
    }
}

2.3 模型定义 (src/model.py)

"""Transformer语言模型"""
import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    """多头注意力"""
    def __init__(self, config):
        super().__init__()
        self.num_heads = config.num_heads
        self.head_dim = config.hidden_size // config.num_heads
        
        self.q_proj = nn.Linear(config.hidden_size, config.hidden_size)
        self.k_proj = nn.Linear(config.hidden_size, config.hidden_size)
        self.v_proj = nn.Linear(config.hidden_size, config.hidden_size)
        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size)
        
        self.dropout = nn.Dropout(config.dropout)
        self.scale = math.sqrt(self.head_dim)
        
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.size()
        
        # 线性投影
        Q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.k_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.v_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        
        # 注意力计算
        scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        
        attn = torch.softmax(scores, dim=-1)
        attn = self.dropout(attn)
        
        out = torch.matmul(attn, V)
        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
        
        return self.o_proj(out)

class FeedForward(nn.Module):
    """前馈网络"""
    def __init__(self, config):
        super().__init__()
        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
        self.dropout = nn.Dropout(config.dropout)
        self.activation = nn.GELU()
        
    def forward(self, x):
        x = self.fc1(x)
        x = self.activation(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x

class TransformerBlock(nn.Module):
    """Transformer块"""
    def __init__(self, config):
        super().__init__()
        self.attention = MultiHeadAttention(config)
        self.feed_forward = FeedForward(config)
        self.ln1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.ln2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.dropout)
        
    def forward(self, x, mask=None):
        # 自注意力
        attn_out = self.attention(self.ln1(x), mask)
        x = x + self.dropout(attn_out)
        
        # 前馈网络
        ff_out = self.feed_forward(self.ln2(x))
        x = x + self.dropout(ff_out)
        
        return x

class SmallLLM(nn.Module):
    """小型语言模型"""
    def __init__(self, config):
        super().__init__()
        self.config = config
        
        # 嵌入层
        self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
        self.position_embedding = nn.Embedding(config.max_seq_length, config.hidden_size)
        
        # Transformer层
        self.layers = nn.ModuleList([
            TransformerBlock(config) for _ in range(config.num_layers)
        ])
        
        # 输出层
        self.ln_final = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        
        # 权重绑定
        self.lm_head.weight = self.token_embedding.weight
        
        self.dropout = nn.Dropout(config.dropout)
        self._init_weights()
        
    def _init_weights(self):
        """初始化"""
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, std=0.02)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, std=0.02)
                
    def forward(self, input_ids, attention_mask=None, labels=None):
        batch_size, seq_len = input_ids.shape
        
        # 位置编码
        positions = torch.arange(0, seq_len, device=input_ids.device).unsqueeze(0)
        
        # 嵌入
        x = self.token_embedding(input_ids)
        x = x + self.position_embedding(positions)
        x = self.dropout(x)
        
        # 创建因果掩码
        if attention_mask is not None:
            mask = attention_mask.unsqueeze(1).unsqueeze(2)
            causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=input_ids.device))
            mask = mask * causal_mask
        else:
            mask = torch.tril(torch.ones(seq_len, seq_len, device=input_ids.device))
        
        # Transformer层
        for layer in self.layers:
            x = layer(x, mask)
        
        # 输出
        x = self.ln_final(x)
        logits = self.lm_head(x)
        
        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
        
        return {"loss": loss, "logits": logits} if loss is not None else {"logits": logits}
    
    def generate(self, input_ids, max_length=100, temperature=1.0, top_k=50):
        """生成文本"""
        self.eval()
        device = input_ids.device
        
        with torch.no_grad():
            for _ in range(max_length):
                outputs = self.forward(input_ids)
                logits = outputs["logits"][:, -1, :] / temperature
                
                if top_k > 0:
                    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
                    logits[indices_to_remove] = -float('inf')
                
                probs = torch.softmax(logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
                
                input_ids = torch.cat([input_ids, next_token], dim=-1)
                
                if next_token.item() == 2:  # EOS
                    break
        
        return input_ids

2.4 分词器训练 (src/tokenizer_train.py)

"""训练BPE分词器"""
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import Whitespace
from tokenizers.processors import TemplateProcessing
import os

def train_tokenizer(input_file, output_dir="./tokenizer", vocab_size=32000):
    """
    训练BPE分词器
    
    Args:
        input_file: 训练数据文件路径
        output_dir: 输出目录
        vocab_size: 词表大小
    """
    os.makedirs(output_dir, exist_ok=True)
    
    # 初始化BPE
    tokenizer = Tokenizer(BPE(unk_token=""))
    tokenizer.pre_tokenizer = Whitespace()
    
    # 训练器配置
    trainer = BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=["", "", "", " ", " "],
        min_frequency=2
    )
    
    print(f"开始训练分词器，词表大小: {vocab_size}")
    print(f"训练数据: {input_file}")
    
    # 训练
    tokenizer.train([input_file], trainer)
    
    # 后处理
    tokenizer.post_processor = TemplateProcessing(
        single=" $A ",
        pair=" $A   $B:1 :2 ",
        special_tokens=[
            ("", 0),
            ("", 1),
        ],
    )
    
    # 保存
    output_path = os.path.join(output_dir, "tokenizer.json")
    tokenizer.save(output_path)
    print(f"分词器已保存到: {output_path}")
    
    # 测试
    test_text = "你好，这是一个测试。"
    output = tokenizer.encode(test_text)
    print(f"测试文本: {test_text}")
    print(f"编码结果: {output.tokens}")
    print(f"解码结果: {tokenizer.decode(output.ids)}")
    
    return tokenizer

if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1:
        train_tokenizer(sys.argv[1])
    else:
        # 使用示例
        train_tokenizer("./data/train.txt")

2.5 数据集 (src/dataset.py)

"""数据集定义"""
import torch
from torch.utils.data import Dataset
import os

class TextDataset(Dataset):
    def __init__(self, file_path, tokenizer, max_length=512):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.samples = []
        
        print(f"加载数据: {file_path}")
        
        # 读取数据
        with open(file_path, 'r', encoding='utf-8') as f:
            text = f.read()
        
        # 分块处理
        chunk_size = max_length
        for i in range(0, len(text), chunk_size):
            chunk = text[i:i + chunk_size]
            if len(chunk) > 50:  # 过滤太短的数据
                self.samples.append(chunk)
        
        print(f"加载了 {len(self.samples)} 个样本")
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        text = self.samples[idx]
        
        # 编码
        encoding = self.tokenizer.encode(text)
        input_ids = encoding.ids[:self.max_length]
        
        # 填充
        padding_length = self.max_length - len(input_ids)
        input_ids = input_ids + [0] * padding_length
        
        # Attention mask
        attention_mask = [1 if i < len(encoding.ids) else 0 for i in range(self.max_length)]
        
        return {
            "input_ids": torch.tensor(input_ids, dtype=torch.long),
            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
            "labels": torch.tensor(input_ids, dtype=torch.long)
        }

2.6 训练脚本 (src/train.py)

"""训练脚本"""
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm
import os
import time
import psutil
import json

from config import ModelConfig, TrainingConfig
from model import SmallLLM
from dataset import TextDataset

class CPUTrainer:
    """CPU训练器"""
    def __init__(self, model, train_dataset, config):
        self.model = model
        self.config = config
        
        # 设置CPU线程数
        torch.set_num_threads(os.cpu_count())
        print(f"使用 {os.cpu_count()} 个CPU线程")
        
        # 优化器
        self.optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=config.learning_rate,
            weight_decay=config.weight_decay,
            betas=(0.9, 0.95)
        )
        
        # 学习率调度
        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
            self.optimizer,
            T_0=1000,
            T_mult=2
        )
        
        # 数据加载器 (Windows下num_workers=0)
        self.train_loader = DataLoader(
            train_dataset,
            batch_size=config.batch_size,
            shuffle=True,
            num_workers=0,  # Windows必须设为0
            pin_memory=False
        )
        
        self.global_step = 0
        self.best_loss = float('inf')
        
    def train_epoch(self, epoch):
        """训练一个epoch"""
        self.model.train()
        total_loss = 0
        self.optimizer.zero_grad()
        
        progress_bar = tqdm(self.train_loader, desc=f"Epoch {epoch+1}")
        
        for batch_idx, batch in enumerate(progress_bar):
            # 前向传播
            outputs = self.model(
                input_ids=batch['input_ids'],
                attention_mask=batch['attention_mask'],
                labels=batch['labels']
            )
            
            loss = outputs['loss'] / self.config.gradient_accumulation_steps
            
            # 反向传播
            loss.backward()
            
            # 梯度累积
            if (batch_idx + 1) % self.config.gradient_accumulation_steps == 0:
                # 梯度裁剪
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
                
                self.optimizer.step()
                self.scheduler.step()
                self.optimizer.zero_grad()
                
                self.global_step += 1
            
            total_loss += loss.item() * self.config.gradient_accumulation_steps
            
            # 更新进度条
            if batch_idx % self.config.log_steps == 0:
                avg_loss = total_loss / (batch_idx + 1)
                progress_bar.set_postfix({
                    'loss': f"{loss.item() * self.config.gradient_accumulation_steps:.4f}",
                    'avg_loss': f"{avg_loss:.4f}",
                    'lr': f"{self.scheduler.get_last_lr()[0]:.6f}"
                })
            
            # 保存检查点
            if self.global_step > 0 and self.global_step % self.config.save_steps == 0:
                self.save_checkpoint(f"checkpoint_step{self.global_step}.pt")
        
        return total_loss / len(self.train_loader)
    
    def save_checkpoint(self, filename, is_best=False):
        """保存检查点"""
        os.makedirs(self.config.output_dir, exist_ok=True)
        
        checkpoint = {
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'scheduler_state_dict': self.scheduler.state_dict(),
            'global_step': self.global_step,
            'config': self.model.config
        }
        
        path = os.path.join(self.config.output_dir, filename)
        torch.save(checkpoint, path)
        print(f"\n保存检查点: {path}")
        
        # 保存最佳模型
        if is_best:
            best_path = os.path.join(self.config.output_dir, "best_model.pt")
            torch.save(checkpoint, best_path)
    
    def train(self, num_epochs):
        """训练"""
        print(f"\n开始训练 {num_epochs} 个epoch")
        print(f"总样本数: {len(self.train_loader.dataset)}")
        print(f"Batch size: {self.config.batch_size}")
        print(f"梯度累积步数: {self.config.gradient_accumulation_steps}")
        print(f"有效batch size: {self.config.batch_size * self.config.gradient_accumulation_steps}\n")
        
        for epoch in range(num_epochs):
            start_time = time.time()
            
            avg_loss = self.train_epoch(epoch)
            
            epoch_time = time.time() - start_time
            
            print(f"\nEpoch {epoch+1}/{num_epochs} 完成")
            print(f"平均损失: {avg_loss:.4f}")
            print(f"用时: {epoch_time/60:.2f} 分钟")
            
            # 保存每个epoch的模型
            self.save_checkpoint(f"model_epoch{epoch+1}.pt")
            
            # 保存最佳模型
            if avg_loss < self.best_loss:
                self.best_loss = avg_loss
                self.save_checkpoint("best_model.pt", is_best=True)
        
        print("\n训练完成！")

def get_system_info():
    """获取系统信息"""
    print("=" * 50)
    print("系统信息")
    print("=" * 50)
    print(f"CPU核心数: {os.cpu_count()}")
    print(f"内存总量: {psutil.virtual_memory().total / (1024**3):.2f} GB")
    print(f"可用内存: {psutil.virtual_memory().available / (1024**3):.2f} GB")
    print(f"PyTorch版本: {torch.__version__}")
    print("=" * 50)

def main():
    """主函数"""
    get_system_info()
    
    # 配置
    model_config = ModelConfig()
    train_config = TrainingConfig()
    
    # 创建模型
    print("\n创建模型...")
    model = SmallLLM(model_config)
    
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    print(f"总参数量: {total_params:,} ({total_params/1e6:.2f}M)")
    print(f"可训练参数量: {trainable_params:,} ({trainable_params/1e6:.2f}M)")
    
    # 加载分词器
    print("\n加载分词器...")
    from tokenizers import Tokenizer
    tokenizer = Tokenizer.from_file(train_config.tokenizer_path)
    
    # 创建数据集
    print("\n创建数据集...")
    dataset = TextDataset(train_config.data_file, tokenizer, model_config.max_seq_length)
    
    # 创建训练器
    trainer = CPUTrainer(model, dataset, train_config)
    
    # 开始训练
    trainer.train(train_config.num_epochs)

if __name__ == "__main__":
    main()

2.7 推理脚本 (src/inference.py)

"""推理脚本"""
import torch
from tokenizers import Tokenizer
import sys
import os

# 添加src到路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from model import SmallLLM

def load_model(checkpoint_path):
    """加载模型"""
    print(f"加载模型: {checkpoint_path}")
    
    checkpoint = torch.load(checkpoint_path, map_location='cpu')
    config = checkpoint['config']
    
    model = SmallLLM(config)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    
    print("模型加载完成")
    return model, config

def generate_text(model, tokenizer, prompt, max_length=100, temperature=0.8, top_k=50):
    """生成文本"""
    # 编码
    encoding = tokenizer.encode(prompt)
    input_ids = torch.tensor([encoding.ids], dtype=torch.long)
    
    print(f"\n输入: {prompt}")
    print(f"输入token数: {len(encoding.ids)}")
    
    # 生成
    with torch.no_grad():
        output_ids = model.generate(input_ids, max_length=max_length, 
                                   temperature=temperature, top_k=top_k)
    
    # 解码
    output_text = tokenizer.decode(output_ids[0].tolist())
    
    return output_text

def interactive_mode(model, tokenizer):
    """交互模式"""
    print("\n" + "=" * 50)
    print("交互模式 - 输入 'quit' 退出")
    print("=" * 50)
    
    while True:
        prompt = input("\n请输入提示词: ").strip()
        if prompt.lower() == 'quit':
            break
        
        if not prompt:
            continue
        
        try:
            result = generate_text(model, tokenizer, prompt, max_length=100)
            print(f"\n生成结果:\n{result}")
        except Exception as e:
            print(f"错误: {e}")

def main():
    import argparse
    parser = argparse.ArgumentParser(description='模型推理')
    parser.add_argument('--model', type=str, default='./models/best_model.pt',
                       help='模型路径')
    parser.add_argument('--tokenizer', type=str, default='./tokenizer/tokenizer.json',
                       help='分词器路径')
    parser.add_argument('--prompt', type=str, default=None,
                       help='提示词（不指定则进入交互模式）')
    
    args = parser.parse_args()
    
    # 加载
    model, config = load_model(args.model)
    tokenizer = Tokenizer.from_file(args.tokenizer)
    
    if args.prompt:
        result = generate_text(model, tokenizer, args.prompt)
        print(f"\n生成结果:\n{result}")
    else:
        interactive_mode(model, tokenizer)

if __name__ == "__main__":
    main()

三、Windows批处理脚本

3.1 一键训练脚本 (train.bat)

@echo off
chcp 65001
echo ==========================================
echo  语言模型训练脚本 (Windows CPU)
echo ==========================================

:: 激活虚拟环境
call llm_env\Scripts\activate.bat

:: 检查数据
if not exist "data\train.txt" (
    echo 错误: 未找到训练数据 data\train.txt
    echo 请准备训练数据并放在 data/train.txt
    pause
    exit /b 1
)

:: 训练分词器
echo.
echo [1/3] 训练分词器...
if not exist "tokenizer\tokenizer.json" (
    python src\tokenizer_train.py data\train.txt
) else (
    echo 分词器已存在，跳过训练
)

:: 开始训练
echo.
echo [2/3] 开始训练模型...
python src\train.py

:: 完成
echo.
echo [3/3] 训练完成！
echo 模型保存在 models/ 目录
pause

3.2 推理脚本 (inference.bat)

@echo off
chcp 65001
echo ==========================================
echo  模型推理脚本
echo ==========================================

call llm_env\Scripts\activate.bat

:: 检查模型
if not exist "models\best_model.pt" (
    echo 错误: 未找到模型文件
    pause
    exit /b 1
)

:: 运行推理
python src\inference.py --model models\best_model.pt

pause

四、快速开始

4.1 准备训练数据

# 创建数据目录
mkdir data

# 准备训练文本（可以是任何纯文本文件）
# 例如：小说、百科、对话数据等
# 保存为 data/train.txt

4.2 执行训练

# 方法1: 使用批处理脚本
.\train.bat

# 方法2: 手动执行
.\llm_env\Scripts\Activate.ps1
python src\tokenizer_train.py data\train.txt
python src\train.py

4.3 运行推理

# 交互模式
.\inference.bat

# 或直接运行
python src\inference.py --prompt "你好"

五、Windows优化

5.1 内存优化

# 在训练脚本中添加内存监控

import psutil

def print_memory_usage():
    """打印内存使用情况"""
    mem = psutil.virtual_memory()
    print(f"内存使用: {mem.percent}% ({mem.used/1024**3:.2f}GB / {mem.total/1024**3:.2f}GB)")

5.2 性能优化

# 设置PyTorch CPU优化
import torch

# 启用MKL-DNN加速
torch.backends.mkldnn.enabled = True

# 设置线程数
torch.set_num_threads(os.cpu_count())
torch.set_num_interop_threads(2)

5.3 数据加载优化

# 预分词并保存，避免重复处理
def preprocess_and_save(input_file, output_file, tokenizer):
    """预处理后保存"""
    with open(input_file, 'r', encoding='utf-8') as f:
        text = f.read()
    
    # 分词
    tokens = tokenizer.encode(text).ids
    
    # 保存为二进制
    import numpy as np
    np.save(output_file, np.array(tokens))
    
    return tokens