114_预训练：Masked LM优化与动态掩码效率深度解析

def create_static_mask(input_ids, mask_prob=0.15):
    # 随机生成掩码位置
    mask = torch.rand(input_ids.shape) < mask_prob
    
    # 确保至少有一个token被掩盖
    for i in range(input_ids.shape[0]):
        if not mask[i].any():
            mask[i, torch.randint(0, input_ids.shape[1], (1,))] = True
    
    # 80%用[MASK]替换
    mask_80 = mask & (torch.rand(input_ids.shape) < 0.8)
    # 10%用随机token替换
    mask_10 = mask & ~mask_80 & (torch.rand(input_ids.shape) < 0.5)
    # 10%保持不变
    
    masked_input = input_ids.clone()
    masked_input[mask_80] = mask_token_id
    
    # 生成随机token替换
    random_tokens = torch.randint(0, vocab_size, input_ids.shape, dtype=torch.long)
    masked_input[mask_10] = random_tokens[mask_10]
    
    # 构建标签（只有被掩盖的位置有标签）
    labels = input_ids.clone()
    labels[~mask] = -100  # -100在PyTorch中表示计算损失时忽略
    
    return masked_input, labels

def create_dynamic_mask(input_ids, mask_prob=0.15):
    # 在每次调用时动态生成新的掩码
    batch_size, seq_length = input_ids.shape
    
    # 随机生成掩码位置
    mask = torch.rand(batch_size, seq_length) < mask_prob
    
    # 确保每个序列至少有一个token被掩盖
    for i in range(batch_size):
        if not mask[i].any():
            mask[i, torch.randint(0, seq_length, (1,))] = True
    
    # 应用三种掩码策略（80%[MASK], 10%随机, 10%保持原token）
    # 实现代码与静态掩码类似，但每次都会生成新的掩码
    # ...
    
    return masked_input, labels

# 在数据加载器中使用动态掩码
class DynamicMaskDataLoader(DataLoader):
    def __iter__(self):
        for batch in super().__iter__():
            input_ids = batch['input_ids']
            masked_input, labels = create_dynamic_mask(input_ids)
            batch['input_ids'] = masked_input
            batch['labels'] = labels
            yield batch

def create_context_aware_mask(input_ids, attention_mask, model):
    # 首先获取token的上下文表示
    with torch.no_grad():
        outputs = model(input_ids, attention_mask=attention_mask, output_hidden_states=True)
        last_hidden_state = outputs.hidden_states[-1]
    
    # 计算token重要性分数（可以基于token的嵌入范数或其他指标）
    importance_scores = torch.norm(last_hidden_state, dim=-1)
    
    # 对每个序列，按重要性分数排序并选择前15%的token进行掩盖
    batch_size, seq_length = input_ids.shape
    mask = torch.zeros_like(input_ids, dtype=torch.bool)
    
    for i in range(batch_size):
        # 找到有效的token位置（非padding）
        valid_indices = attention_mask[i].nonzero().squeeze()
        
        # 获取有效token的重要性分数
        valid_scores = importance_scores[i, valid_indices]
        
        # 选择重要性最高的15%的token
        k = max(1, int(valid_indices.numel() * 0.15))
        top_indices = valid_indices[valid_scores.topk(k).indices]
        
        mask[i, top_indices] = True
    
    # 后续掩码处理（80%[MASK], 10%随机, 10%保持原token）
    # ...
    
    return masked_input, labels

def get_adaptive_mask_prob(epoch, model_size, data_complexity):
    # 基础掩码比例
    base_prob = 0.15
    
    # 根据模型规模调整（模型越大，掩码比例可能越高）
    size_factor = min(1.5, model_size / 1e9)  # 假设model_size以参数数量衡量
    
    # 根据数据复杂度调整
    complexity_factor = 1.0 + (data_complexity - 1.0) * 0.2
    
    # 根据训练阶段调整（预热阶段较低，稳定阶段较高）
    if epoch < 5:
        phase_factor = 0.7 + (epoch / 5) * 0.3  # 预热阶段从0.7线性增长到1.0
    elif epoch > 30:
        phase_factor = 0.9  # 后期稍微降低以稳定训练
    else:
        phase_factor = 1.0
    
    # 综合计算最终掩码比例，确保在合理范围内
    final_prob = base_prob * size_factor * complexity_factor * phase_factor
    return min(max(final_prob, 0.1), 0.3)  # 限制在0.1-0.3之间

def create_efficient_dynamic_mask(input_ids, mask_prob=0.15, device=None):
    if device is None:
        device = input_ids.device
    
    batch_size, seq_length = input_ids.shape
    
    # 直接在GPU上生成随机掩码（避免CPU-GPU传输）
    mask = torch.rand(batch_size, seq_length, device=device) < mask_prob
    
    # 确保每个序列至少有一个掩码
    # 批量处理而非循环（提高效率）
    has_mask = mask.any(dim=1)
    if not has_mask.all():
        # 找出没有掩码的样本
        no_mask_indices = (~has_mask).nonzero().squeeze()
        # 为这些样本随机选择一个位置进行掩码
        random_positions = torch.randint(0, seq_length, (no_mask_indices.numel(),), device=device)
        mask[no_mask_indices, random_positions] = True
    
    # 并行计算三种掩码策略
    # 使用向量化操作而非循环
    random_values = torch.rand(batch_size, seq_length, device=device)
    mask_80 = mask & (random_values < 0.8)
    mask_10_random = mask & ~mask_80 & (random_values < 0.9)
    
    # 生成masked_input
    masked_input = input_ids.clone()
    masked_input[mask_80] = mask_token_id
    
    # 并行生成随机token替换
    random_tokens = torch.randint(0, vocab_size, input_ids.shape, device=device)
    masked_input[mask_10_random] = random_tokens[mask_10_random]
    
    # 创建标签
    labels = input_ids.clone()
    labels[~mask] = -100
    
    return masked_input, labels

# 在分布式训练中使用动态掩码
class DistributedDynamicMask:
    def __init__(self, mask_prob=0.15, seed=None):
        self.mask_prob = mask_prob
        # 为每个rank设置不同的随机种子
        if seed is not None:
            local_rank = torch.distributed.get_rank()
            torch.manual_seed(seed + local_rank)
    
    def __call__(self, input_ids):
        # 每个worker独立生成掩码
        return create_efficient_dynamic_mask(input_ids, self.mask_prob)

# 在分布式训练中的使用示例
def train_step(input_ids, model, optimizer, distributed_masker):
    # 动态生成掩码
    masked_input, labels = distributed_masker(input_ids)
    
    # 前向传播
    outputs = model(input_ids=masked_input, labels=labels)
    loss = outputs.loss
    
    # 反向传播和优化
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    
    return loss.item()

def large_scale_dynamic_mask(input_ids, mask_prob=0.15, use_amp=True):
    # 内存优化：使用半精度生成随机数
    if use_amp:
        random_tensor = torch.rand(input_ids.shape, dtype=torch.half, device=input_ids.device)
    else:
        random_tensor = torch.rand(input_ids.shape, device=input_ids.device)
    
    # 生成掩码
    mask = random_tensor < mask_prob
    
    # 其余掩码处理逻辑...
    # ...
    
    return masked_input, labels