[AI学习笔记]预训练范式演进：DeepSeek多阶段预训练技巧解密

# 传统单阶段训练的问题演示
def train(model, dataset):
    optimizer = AdamW(model.parameters(), lr=5e-4)
    for epoch in range(100):
        for batch in dataset:
            loss = model(batch)
            loss.backward()
            optimizer.step()  # 持续的高学习率导致后期不稳定
            optimizer.zero_grad()

class CurriculumSampler:
    def __init__(self, dataset, alpha=0.5):
        self.scores = compute_difficulty(dataset)  # 预计算样本难度
        self.alpha = alpha  # 课程进度系数
        
    def sample(self, current_step):
        # 动态调整采样概率
        temp = 1.0 + self.alpha * current_step
        probs = F.softmax(self.scores / temp, dim=0)
        return torch.multinomial(probs, batch_size)

def hybrid_loss(batch):
    text_loss = model(batch['text'])  # 语言建模损失
    kg_loss = model(batch['triples'])  # 知识图谱三元组损失
    math_loss = model(batch['equations'])  # 数学公式推导损失
    return 0.6*text_loss + 0.3*kg_loss + 0.1*math_loss

class DeepSeekScheduler:
    def get_lr(self, step):
        if step < warmup_steps:  # 初始阶段
            return base_lr * step / warmup_steps
        elif step < transition_step:  # 平稳阶段
            return base_lr
        else:  # 衰减阶段
            return base_lr * 0.5 ** ((step - transition_step) // decay_interval)

def freeze_layers(model, current_stage):
    # 筑基阶段：仅训练顶层
    if current_stage == 1:
        for param in model.base_layers.parameters():
            param.requires_grad = False
    # 增强阶段：解冻中间层
    elif current_stage == 2:
        for param in model.mid_layers.parameters():
            param.requires_grad = True
    # 精调阶段：全参数训练
    else:
        for param in model.parameters():
            param.requires_grad = True

# 基于Megatron-LM的层内并行
from megatron.core import tensor_parallel

class ParallelMLP(nn.Module):
    def __init__(self, hidden_size):
        self.dense_h_to_4h = tensor_parallel.ColumnParallelLinear(
            hidden_size, 4*hidden_size
        )
        self.dense_4h_to_h = tensor_parallel.RowParallelLinear(
            4*hidden_size, hidden_size
        )
    
    def forward(self, x):
        x = self.dense_h_to_4h(x)
        x = gelu(x)
        return self.dense_4h_to_h(x)

def activation_checkpoint(module):
    # 使用选择性激活检查点
    from torch.utils.checkpoint import checkpoint_sequential
    
    def forward(*inputs):
        def custom_forward(*tensors):
            return module(*tensors)
        
        return checkpoint_sequential(
            [custom_forward], 
            chunks=4,
            input=inputs
        )
    return forward

# 使用Torch Quantization
model_fp32 = load_pretrained_model()
model_fp32.eval()

# 量化配置
qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
model_fp32.qconfig = qconfig

# 准备量化
model_int8 = torch.quantization.prepare_qat(model_fp32)

# 校准
for data in calibration_data:
    model_int8(data)

# 转换
model_final = torch.quantization.convert(model_int8)

# 使用Triton推理服务器配置
name: "deepseek_67b"
platform: "pytorch_libtorch"
max_batch_size: 16

input [
  { name: "input_ids", data_type: TYPE_INT64, dims: [ 2048 ] }
]

output [
  { name: "logits", data_type: TYPE_FP32, dims: [ 50257 ] }
]

instance_group [
  { count: 4, kind: KIND_GPU }
]

optimization {
  cuda {
    graphs: true
    busy_wait_events: true
  }
}

import torch
import torch.nn as nn
from transformers import BertConfig, BertModel
from torchvision.models import resnet50

# 图像特征提取模块
class ImageFeatureExtractor(nn.Module):
    def __init__(self):
        super(ImageFeatureExtractor, self).__init__()
        self.resnet = resnet50(pretrained=True)
        self.features = nn.Sequential(*list(self.resnet.children())[:-1])  # 去掉最后一层全连接层

    def forward(self, x):
        x = self.features(x)
        x = x.squeeze(-1).squeeze(-1)  # 将二维特征图转换为一维特征向量
        return x

# 文本特征提取模块
class TextFeatureExtractor(nn.Module):
    def __init__(self, bert_model_name='bert-base-uncased'):
        super(TextFeatureExtractor, self).__init__()
        self.config = BertConfig.from_pretrained(bert_model_name)
        self.bert = BertModel.from_pretrained(bert_model_name, config=self.config)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        return outputs.last_hidden_state

# 多模态融合模块
class MultimodalFusion(nn.Module):
    def __init__(self, image_feature_dim, text_feature_dim, fusion_dim):
        super(MultimodalFusion, self).__init__()
        self.image_fc = nn.Linear(image_feature_dim, fusion_dim)
        self.text_fc = nn.Linear(text_feature_dim, fusion_dim)
        self.activation = nn.ReLU()

    def forward(self, image_features, text_features):
        image_features = self.image_fc(image_features)
        text_features = self.text_fc(text_features)
        fused_features = image_features + text_features  # 简单的加法融合，也可以尝试其他融合方式
        fused_features = self.activation(fused_features)
        return fused_features

# 完整的多模态预训练模型
class DeepSeekMultimodalModel(nn.Module):
    def __init__(self, bert_model_name='bert-base-uncased'):
        super(DeepSeekMultimodalModel, self).__init__()
        self.image_extractor = ImageFeatureExtractor()
        self.text_extractor = TextFeatureExtractor(bert_model_name)
        self.fusion = MultimodalFusion(image_feature_dim=2048, text_feature_dim=768, fusion_dim=768)
        self.classifier = nn.Linear(768, 2)  # 假设为二分类任务

    def forward(self, images, input_ids, attention_mask):
        image_features = self.image_extractor(images)
        text_features = self.text_extractor(input_ids, attention_mask)
        text_features = text_features[:, 0, :]  # 取[CLS] token的特征作为文本的整体特征
        fused_features = self.fusion(image_features, text_features)
        logits = self.classifier(fused_features)
        return logits

# 图像特征预训练（以图像分类为例）
model = ImageFeatureExtractor()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 假设已经加载好图像数据和对应的分类标签
for epoch in range(num_epochs):
    for images, labels in image_dataloader:
        images = images.to(device)
        labels = labels.to(device)
        
        features = model(images)
        logits = classifier(features)  # classifier是一个简单的全连接分类层
        loss = criterion(logits, labels)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# 文本特征预训练（以Masked Language Model为例）
model = TextFeatureExtractor()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 假设已经加载好文本数据，并进行了Mask处理
for epoch in range(num_epochs):
    for batch in text_dataloader:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)
        
        outputs = model(input_ids, attention_mask)
        loss = criterion(outputs, labels)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# 跨模态关联预训练（以图像-文本匹配为例）
model = DeepSeekMultimodalModel()
criterion = nn.CosineEmbeddingLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 假设已经加载好图像和对应的文本数据，并生成了匹配和不匹配的样本对
for epoch in range(num_epochs):
    for images, pos_texts, neg_texts in multimodal_dataloader:
        images = images.to(device)
        pos_input_ids = pos_texts['input_ids'].to(device)
        pos_attention_mask = pos_texts['attention_mask'].to(device)
        neg_input_ids = neg_texts['input_ids'].to(device)
        neg_attention_mask = neg_texts['attention_mask'].to(device)
        
        # 计算匹配对的特征
        pos_text_features = model.text_extractor(pos_input_ids, pos_attention_mask)
        pos_text_features = pos_text_features[:, 0, :]
        image_features = model.image_extractor(images)
        pos_similarity = torch.cosine_similarity(image_features, pos_text_features)
        
        # 计算不匹配对的特征
        neg_text_features = model.text_extractor(neg_input_ids, neg_attention_mask)
        neg_text_features = neg_text_features[:, 0, :]
        neg_similarity = torch.cosine_similarity(image_features, neg_text_features)
        
        # 定义目标，匹配对的相似度应大于不匹配对
        targets = torch.ones(images.size(0), dtype=torch.long).to(device)
        loss = criterion(pos_similarity, neg_similarity, targets)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# 多模态情感分析微调
model = DeepSeekMultimodalModel()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)

# 假设已经加载好包含图像、文本和情感标签的数据
for epoch in range(num_epochs):
    for images, input_ids, attention_mask, labels in sentiment_dataloader:
        images = images.to(device)
        input_ids = input_ids.to(device)
        attention_mask = attention_mask.to(device)
        labels = labels.to(device)
        
        logits = model(images, input_ids, attention_mask)
        loss = criterion(logits, labels)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# 在测试集上评估模型性能
model.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for images, input_ids, attention_mask, labels in test_dataloader:
        images = images.to(device)
        input_ids = input_ids.to(device)
        attention_mask = attention_mask.to(device)
        labels = labels.to(device)
        
        logits = model(images, input_ids, attention_mask)
        _, predicted = torch.max(logits, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
    
    accuracy = correct / total
    print(f'Test Accuracy: {accuracy:.4f}')

# 保存模型
torch.save(model.state_dict(), 'deepseek_multimodal_sentiment_analysis.pth')

# 加载模型
model = DeepSeekMultimodalModel()
model.load_state_dict(torch.load('deepseek_multimodal_sentiment_analysis.pth'))
model.eval()

# 应用示例
def analyze_multimodal_sentiment(image_path, text):
    # 加载和预处理图像
    image = load_and_preprocess_image(image_path)
    image_tensor = torch.tensor(image, dtype=torch.float32).unsqueeze(0).to(device)
    
    # 预处理文本
    text_inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True)
    input_ids = text_inputs['input_ids'].to(device)
    attention_mask = text_inputs['attention_mask'].to(device)
    
    # 情感分析
    logits = model(image_tensor, input_ids, attention_mask)
    probabilities = torch.softmax(logits, dim=1)
    predicted_class = torch.argmax(probabilities).item()
    
    return predicted_class, probabilities.cpu().numpy()[0]