Update modeling_sdar.py

modeling_sdar.py

@@ -77,10 +77,8 @@ def modify_padded_position_ids_2d(position_ids: torch.LongTensor) -> torch.LongT
     使用完全向量化的 PyTorch 操作修改一个 batch 的 packed position_ids。
     这个函数假设输入是一个 2D Tensor，形状为 (batch_size, sequence_length)。
     它会独立地处理 batch 中的每一行。
-
     Args:
         position_ids: 二维 PyTorch Tensor, shape (batch_size, sequence_length).
-
     Returns:
         修改后的 position_ids Tensor, shape (batch_size, sequence_length).
     """
@@ -108,7 +106,6 @@ def modify_padded_position_ids_2d(position_ids: torch.LongTensor) -> torch.LongT
 def calculate_token_nums(position_ids: torch.Tensor):
     """
     使用 PyTorch 高效计算一个批次中每个打包序列的长度。
-
     Args:
         position_ids (torch.Tensor): 一个 2D Tensor，形状为 (batch_size, sequence_length)。
                                      例如：tensor([[0,1,2,3,4,0,1,2,3,4,5,0,1,2,3,0,0,0]])
@@ -162,11 +159,9 @@ def forward_add_noise_packed(
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """
     为一批打包（packed）序列的 token ID 添加噪声。
-
     此函数保留了为每个逻辑样本（在每个批次项内拼接）生成独立随机噪声率的逻辑。
     它会随机将一部分 token 的 ID 替换为 mask_id。
     这个过程会避开被 prompt_mask 标记的位置。
-
     Args:
         inputs_ids (torch.Tensor): 
             输入的 token ID 张量，形状为 (bsz, total_tokens)。
@@ -182,7 +177,6 @@ def forward_add_noise_packed(
             微小值，用于防止噪声率 t 恰好为 0，确保 p_mask > 0。
         max_tries (int): 
             为确保至少一个非 prompt token 被 mask，对每个批次项尝试的最大次数。
-
     Returns:
         Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         - noisy_input_ids (torch.Tensor): 
@@ -290,13 +284,11 @@ def block_diff_mask(b, h, q_idx, kv_idx, block_size=None, n=None):
     - **Block Diagonal Mask (M_BD)**: Self-attention within noised blocks
     - **Offset Block Causal Mask (M_OBC)**: Cross-attention for conditional context
     - **Block Causal Mask (M_BC)**: Attention to update x0
-
     Args:
         b, h: Batch and head indices (ignored for mask logic).
         q_idx, kv_idx: Query and Key indices.
         seq_len: Total sequence length.
         block_size: Defines the block structure.
-
     Returns:
         A boolean attention mask.
     """
@@ -410,7 +402,6 @@ def rotate_half(x):
 
 def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
     """Applies Rotary Position Embedding to the query and key tensors.
-
     Args:
         q (`torch.Tensor`): The query tensor.
         k (`torch.Tensor`): The key tensor.
@@ -970,7 +961,6 @@ class SDARModel(SDARPreTrainedModel):
         """
         Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
         `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
-
         Args:
             attention_mask (`torch.Tensor`):
                 A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`.
@@ -1160,7 +1150,6 @@ class SDARForCausalLM(SDARPreTrainedModel, GenerationMixin):
         compute_rl_loss: bool = False,
         p_mask: Optional[torch.Tensor] = None,
         adv: Optional[torch.Tensor] = None,
-        adv_optimization: bool = False,
         logp_old_tok: Optional[torch.Tensor] = None,
         logp_ref_tok: Optional[torch.Tensor] = None,
         is_real: Optional[torch.Tensor] = None,
@@ -1237,12 +1226,6 @@ class SDARForCausalLM(SDARPreTrainedModel, GenerationMixin):
 
                 # 选出 logits — 保持原样
                 logits_p = logits[p_to_keep_real]                       # (N, V)
-                N = p_to_keep_real.sum().item()
-                total_response_tokens = (labels != -100).sum().item()
-                total_p_mask = p_mask.sum().item()
-                total_masked_indices = masked_indices.sum().item()
-                total_is_real = is_real_tensor.sum().item() if is_real_tensor.dim() > 0 else (1 if is_real_tensor.item() else 0)
-
 
                 # log_softmax
                 log_probs_p = torch.nn.functional.log_softmax(logits_p, dim=-1)
@@ -1260,133 +1243,7 @@ class SDARForCausalLM(SDARPreTrainedModel, GenerationMixin):
 
                 # advantage 处理
                 adv_tensor = adv.to(device) if torch.is_tensor(adv) else torch.tensor(adv, dtype=torch.float, device=device)
-                adv_optimization=False
-                if adv_optimization:
-                    # token级别优化：对相同前缀取最大advantage（剪枝优化版本）
-                    response_mask = (labels != -100)  # (B, L)
-                    bsz, seq_len = input_ids.shape
-                    
-                    # 预计算每个样本的response起始位置
-                    response_starts = torch.full((bsz,), seq_len, dtype=torch.long, device=device)
-                    for b in range(bsz):
-                        if response_mask[b].any():
-                            response_starts[b] = response_mask[b].long().argmax()
-                    
-                    # 剪枝1: 找出已经是最大advantage的样本，直接填充不参与比较
-                    max_adv_value = adv_tensor.max()
-                    is_max_adv = (adv_tensor == max_adv_value)  # (B,) bool
-                    
-                    # 创建优化后的 advantage map (B, L)，确保dtype与adv_tensor一致
-                    optimized_adv = torch.zeros_like(labels, dtype=adv_tensor.dtype)
-                    
-                    # 对于已是最大advantage的样本，直接填充
-                    for b in range(bsz):
-                        if is_max_adv[b]:
-                            optimized_adv[b][response_mask[b]] = max_adv_value
-                    
-                    # 统计信息
-                    total_response_tokens = 0
-                    updated_tokens = 0
-                    skipped_tokens = 0
-                    original_adv_sum = 0.0
-                    optimized_adv_sum = 0.0
-                    
-                    # 按position处理，批量比较前缀
-                    for pos in range(seq_len):
-                        valid_samples = response_mask[:, pos]  # (B,)
-                        if not valid_samples.any():
-                            continue
-                        
-                        # 剪枝2: 排除已是最大advantage的样本
-                        valid_samples = valid_samples & ~is_max_adv
-                        if not valid_samples.any():
-                            # 所有样本都是最大值，统计后跳过
-                            max_count = (response_mask[:, pos] & is_max_adv).sum().item()
-                            total_response_tokens += max_count
-                            skipped_tokens += max_count
-                            original_adv_sum += max_adv_value.item() * max_count
-                            optimized_adv_sum += max_adv_value.item() * max_count
-                            continue
-                        
-                        # 获取所有需要处理的样本索引
-                        valid_indices = valid_samples.nonzero(as_tuple=True)[0]  # (N,)
-                        
-                        for b in valid_indices:
-                            b_item = b.item()
-                            response_start = response_starts[b_item].item()
-                            prefix_len = pos + 1 - response_start
-                            
-                            if prefix_len <= 0:
-                                optimized_adv[b_item, pos] = adv_tensor[b_item]
-                                continue
-                            
-                            # 找出所有response起始位置相同且在pos位置有效的样本（包括已是最大值的）
-                            same_start_mask = (response_starts == response_start) & response_mask[:, pos]
-                            same_start_indices = same_start_mask.nonzero(as_tuple=True)[0]
-                            
-                            if len(same_start_indices) == 1:
-                                # 只有自己，不需要比较
-                                optimized_adv[b_item, pos] = adv_tensor[b_item]
-                                total_response_tokens += 1
-                                original_adv_sum += adv_tensor[b_item].item()
-                                optimized_adv_sum += adv_tensor[b_item].item()
-                                continue
-                            
-                            # 剪枝3: 如果候选中有最大advantage样本，可以直接用最大值
-                            has_max_in_candidates = (same_start_mask & is_max_adv).any()
-                            
-                            prefix_end = pos + 1
-                            current_prefix = input_ids[b_item, response_start:prefix_end]
-                            
-                            # 批量比较：提取所有候选样本的前缀
-                            prefixes = input_ids[same_start_indices, response_start:prefix_end]  # (M, prefix_len)
-                            
-                            # 使用广播比较：(M, prefix_len) vs (prefix_len,)
-                            matches = (prefixes == current_prefix.unsqueeze(0)).all(dim=1)  # (M,)
-                            
-                            # 找到匹配的样本
-                            matching_indices = same_start_indices[matches]
-                            
-                            # 在相同前缀的样本中取最大 advantage
-                            original_adv_value = adv_tensor[b_item].item()
-                            if matching_indices.numel() > 0:
-                                # 剪枝4: 如果匹配中有最大值样本，直接用最大值
-                                if has_max_in_candidates and is_max_adv[matching_indices].any():
-                                    max_adv = max_adv_value
-                                else:
-                                    max_adv = adv_tensor[matching_indices].max()
-                                
-                                optimized_adv[b_item, pos] = max_adv
-                                # 统计
-                                if abs(max_adv.item() - original_adv_value) > 1e-6:
-                                    updated_tokens += 1
-                                original_adv_sum += original_adv_value
-                                optimized_adv_sum += max_adv.item()
-                            else:
-                                optimized_adv[b_item, pos] = adv_tensor[b_item]
-                                original_adv_sum += original_adv_value
-                                optimized_adv_sum += original_adv_value
-                            
-                            total_response_tokens += 1
-                    
-                    # 输出统计信息
-                    if total_response_tokens > 0:
-                        update_ratio = updated_tokens / total_response_tokens
-                        skip_ratio = skipped_tokens / total_response_tokens
-                        avg_original = original_adv_sum / total_response_tokens
-                        avg_optimized = optimized_adv_sum / total_response_tokens
-                        print(f"[Adv Optimization] Total: {total_response_tokens}, "
-                              f"Updated: {updated_tokens} ({update_ratio:.2%}), "
-                              f"Skipped: {skipped_tokens} ({skip_ratio:.2%}), "
-                              f"Avg adv: {avg_original:.4f} -> {avg_optimized:.4f} "
-                              f"(+{avg_optimized - avg_original:.4f})")
-                    
-                    # 使用优化后的 advantage
-                    adv_expanded = optimized_adv
-                else:
-                    # 不优化：直接使用原始 advantage
-                    adv_expanded = adv_tensor.unsqueeze(1).expand_as(p_mask)
-                
+                adv_expanded = adv_tensor.unsqueeze(1).expand_as(p_mask)
                 adv_p = adv_expanded[masked_indices][p_to_keep_real]
 
                 # old logp
@@ -1394,20 +1251,13 @@ class SDARForCausalLM(SDARPreTrainedModel, GenerationMixin):
                     logp_old_p = logp_old_tok.to(device)[masked_indices][p_to_keep_real]
                 else:
                     logp_old_p = logp_p.detach()
-
+                
                 # ratio/exp
                 ratio_p = (logp_p - logp_old_p).clamp(-10.0, 10.0).exp()
                 clipped = ratio_p.clamp(1 - ppo_eps, 1 + ppo_eps+0.08)
                 surrogate_p = torch.minimum(ratio_p * adv_p, clipped * adv_p)
-                # 输出离1最远的ratio值
-                # if not torch.allclose(ratio_p, torch.ones_like(ratio_p)):
-                furthest_value = ratio_p[torch.abs(ratio_p - 1).argmax()]
-                # print(f"Furthest ratio from 1: {furthest_value.item()}")
 
                 # Policy loss: use mean or sum based on loss_mean parameter
-                num_masked = masked_indices.sum().item()
-                num_loss_elements = surrogate_p.numel()
-                print(f"masked_indices.sum()={num_masked}, surrogate_p.numel()={num_loss_elements}")
                 if loss_mean:
                     policy_loss = -surrogate_p.mean()
                 else: