v16: fix critical audit defects (N6/N7/N8/N9 + S3/S4/M2/M5)

CRITICAL:
- N6: apply_rotary_pos_emb now slices cos/sin by position_ids, preventing
broadcast mismatch during incremental generation (Q shape (B,H,1,D) would
broadcast to (B,H,kv_seq_len,D) causing wrong output)
- N6: FusionAttention.forward() builds position_ids from KV cache offset
when not provided, ensuring correct RoPE at every step

SEVERE:
- N7: SBLA _cached_block_latents null safety - clear cache when use_cache=False,
validate batch_size matches to prevent cross-batch contamination
- S3: LoRA target_modules now includes SBLAttention latent projections
(latent_q/k/v_proj, latent_out_proj, v_to_hidden_proj)
- S4: Extract create_local_model() into shared train/model_utils.py,
both lora_finetune.py and full_finetune.py delegate to it

MEDIUM:
- N8: Remove dead depth_bias computation in ThinkingDialModel.forward()
- N9: apply_rotary_pos_emb accepts position_ids parameter (was unused)
- M2: THINK_END changed from fragile '|>' to proper '<|think_end|>' special
token; separated THINK_CLOSE for think_depth token closing bracket
- M5: generate_with_thinking() now passes thinking_depth to model.generate()

Tests: 12/12 passed

models/fusion_model.py

@@ -150,17 +150,23 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None):
 
     Args:
         q: (batch, num_heads, seq_len, head_dim)
-        k: (batch, num_heads, seq_len, head_dim) or (batch, num_kv_heads, seq_len, head_dim)
-        cos: (seq_len, head_dim) cosine part of rotary embedding
-        sin: (seq_len, head_dim) sine part of rotary embedding
-        position_ids: optional position ids (unused, for API compat)
+        k: (batch, num_kv_heads, seq_len, head_dim)
+        cos: (kv_seq_len, head_dim) cosine part of rotary embedding
+        sin: (kv_seq_len, head_dim) sine part of rotary embedding
+        position_ids: (batch, seq_len) position ids for slicing cos/sin
 
     Returns:
         Tuple of (q_embed, k_embed) with rotary position encoding applied.
     """
-    # cos/sin: (seq_len, head_dim) -> (1, 1, seq_len, head_dim)
-    cos = cos.unsqueeze(0).unsqueeze(0)
-    sin = sin.unsqueeze(0).unsqueeze(0)
+    if position_ids is not None:
+        # N6 FIX: Slice cos/sin by position_ids to match actual Q/K positions
+        # position_ids: (batch, seq_len), cos/sin: (kv_seq_len, head_dim)
+        cos = cos[position_ids].unsqueeze(1)  # (batch, 1, seq_len, head_dim)
+        sin = sin[position_ids].unsqueeze(1)  # (batch, 1, seq_len, head_dim)
+    else:
+        # Fallback: broadcast for full-sequence (prefill) when position_ids not provided
+        cos = cos.unsqueeze(0).unsqueeze(0)  # (1, 1, kv_seq_len, head_dim)
+        sin = sin.unsqueeze(0).unsqueeze(0)
     q_embed = (q * cos) + (rotate_half(q) * sin)
     k_embed = (k * cos) + (rotate_half(k) * sin)
     return q_embed, k_embed
@@ -250,8 +256,15 @@ class FusionAttention(nn.Module):
         emb = self.rotary_emb(kv_seq_len, device=hidden_states.device)
         cos = emb.cos()
         sin = emb.sin()
-        # Apply RoPE to Q (full position range) and K (full position range)
-        Q, K = apply_rotary_pos_emb(Q, K, cos, sin)
+        # N6 FIX: Build position_ids for proper RoPE slicing
+        if position_ids is None:
+            if past_key_value is not None:
+                offset = past_key_value[0].shape[2]
+                position_ids = torch.arange(offset, offset + seq_len, device=hidden_states.device).unsqueeze(0)
+            else:
+                position_ids = torch.arange(seq_len, device=hidden_states.device).unsqueeze(0)
+        # Apply RoPE with position_ids to prevent broadcast mismatch during incremental generation
+        Q, K = apply_rotary_pos_emb(Q, K, cos, sin, position_ids=position_ids)
         
         # Store RoPE'd K/V in SBLAttention's cache for incremental generation
         # S1 FIXED: KV Cache now works natively through SBLAttention.

models/sbla_attention.py

@@ -393,24 +393,26 @@ class SBLAttention(nn.Module):
             # Incremental step: use cached block latents if available
             if hasattr(self, '_cached_block_latents') and self._cached_block_latents is not None:
                 cached_q, cached_k, cached_v, cached_num_blocks = self._cached_block_latents
-                # Compute latent query for the single new token
-                # V_current is the single-step V before KV concat: (B, num_kv_heads, 1, kv_head_dim)
-                # After _repeat_kv on V_current it would be (B, num_heads, 1, head_dim)
-                V_current_expanded = self._repeat_kv(V_current, self.num_kv_groups)
-                V_reshaped_inc = V_current_expanded.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
-                hidden_approx_inc = self.v_to_hidden_proj(V_reshaped_inc)  # (batch, 1, hidden_size)
-                blk_q_inc = self.latent_q_proj(hidden_approx_inc)  # (batch, 1, latent_dim)
-                # Attend to cached block keys/values
-                latent_attn_scores = torch.matmul(
-                    blk_q_inc, cached_k.transpose(-1, -2)
-                ) / math.sqrt(self.latent_dim)
-                # Causal: new token can attend to all blocks
-                latent_attn_probs = F.softmax(latent_attn_scores, dim=-1)
-                latent_attn_probs = self.dropout(latent_attn_probs)
-                latent_context = torch.matmul(latent_attn_probs, cached_v)  # (batch, 1, latent_dim)
-                latent_output = self.latent_out_proj(latent_context)  # (batch, 1, hidden_size)
-                gate_value = torch.sigmoid(self.gate)
-                output = output_std + gate_value * latent_output
+                # N7 FIX: Validate batch size matches to prevent cross-batch contamination
+                if cached_q.size(0) != batch_size:
+                    # Batch size changed (e.g., different batch in concurrent usage)
+                    output = output_std
+                else:
+                    # Compute latent query for the single new token
+                    V_current_expanded = self._repeat_kv(V_current, self.num_kv_groups)
+                    V_reshaped_inc = V_current_expanded.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
+                    hidden_approx_inc = self.v_to_hidden_proj(V_reshaped_inc)
+                    blk_q_inc = self.latent_q_proj(hidden_approx_inc)
+                    # Attend to cached block keys/values
+                    latent_attn_scores = torch.matmul(
+                        blk_q_inc, cached_k.transpose(-1, -2)
+                    ) / math.sqrt(self.latent_dim)
+                    latent_attn_probs = F.softmax(latent_attn_scores, dim=-1)
+                    latent_attn_probs = self.dropout(latent_attn_probs)
+                    latent_context = torch.matmul(latent_attn_probs, cached_v)
+                    latent_output = self.latent_out_proj(latent_context)
+                    gate_value = torch.sigmoid(self.gate)
+                    output = output_std + gate_value * latent_output
             else:
                 # No cached latents: fall back to standard attention only
                 output = output_std
@@ -464,6 +466,9 @@ class SBLAttention(nn.Module):
         if use_cache and past_key_value is None:
             # Prefill step: cache block latents for subsequent incremental steps
             self._cached_block_latents = (blk_q, blk_k, blk_v, num_blocks)
+        elif past_key_value is None:
+            # N7 FIX: Ensure cache is cleared when not using cache, prevents stale data
+            self._cached_block_latents = None
 
         return output, present_key_value
 

models/thinking_dial.py

@@ -47,8 +47,9 @@ from transformers import PreTrainedModel, GenerationMixin
 # ============================================================
 
 THINK_START = "<|think_depth_"
-THINK_END = "|>"
-THINK_DEPTH_PATTERN = re.compile(r"<\|think_depth_(\d+)\|>")
+THINK_CLOSE = "|>"  # Closing bracket for think_depth token
+THINK_END = "<|think_end|>"  # End-of-thinking-block marker
+THINK_DEPTH_PATTERN = re.compile(r"<\|think_depth_(\d)\|>")
 
 # Depth 0-3 的描述
 THINK_DEPTH_DESCRIPTIONS = {
@@ -72,7 +73,7 @@ def build_think_token(depth: int) -> str:
     if not 0 <= depth <= 3:
         raise ValueError(f"depth 必须在 0-3 之间，当前值：{depth}")
     
-    return f"{THINK_START}{depth}{THINK_END}"
+    return f"{THINK_START}{depth}{THINK_CLOSE}"
 
 
 def parse_think_token(text: str) -> Optional[int]:
@@ -430,6 +431,8 @@ class GRPOTrainer:
         返回：
             生成的文本列表
         """
+        # M5 FIX: Pass thinking_depth to model.generate() so the ThinkingDialModel
+        # forward() can actually apply depth-dependent bias to logits
         outputs = self.model.generate(
             input_ids=input_ids,
             max_new_tokens=max_new_tokens,
@@ -438,6 +441,7 @@ class GRPOTrainer:
             do_sample=True,
             pad_token_id=self.model.config.pad_token_id or 0,
             eos_token_id=self.model.config.eos_token_id or 1,
+            thinking_depth=thinking_depth,
             **kwargs,
         )
         
@@ -745,13 +749,8 @@ class ThinkingDialModel(nn.Module):
         if thinking_depth is not None:
             depth_idx = thinking_depth.long().clamp(0, self.thinking_config.num_thinking_depths - 1)
             depth_embedding = self.thinking_embedding(depth_idx)  # (batch, hidden_size)
-            # Project depth embedding to vocabulary space via tied weights
-            depth_bias = nn.functional.linear(
-                depth_embedding.unsqueeze(1),  # (batch, 1, hidden_size)
-                self.thinking_embedding.weight.t(),  # (hidden_size, num_depths) -> transpose gives (num_depths, hidden_size)
-            ).squeeze(1)  # (batch, num_depths)
-            # To properly project to vocab space, use base model's lm_head
-            # Simpler: add hidden_size bias via learned projection
+            # N8 FIX: Removed dead depth_bias computation (was computed but never used)
+            # Apply thinking_gate as scaling factor, then project to vocab space via lm_head
             depth_hidden = self.thinking_gate * depth_embedding  # (batch, hidden_size)
             # Add as residual to logits via lm_head projection
             if hasattr(self.base_model, 'lm_head'):

train/full_finetune.py

@@ -118,57 +118,21 @@ class FusionFullFinetuneDataset(Dataset):
         }
 
 
+from train.model_utils import create_local_model as _create_local_model_from_utils
+
+
 def create_local_model(
     model_size: str = "8B",
     torch_dtype: torch.dtype = torch.bfloat16,
     vocab_size_override: Optional[int] = None,
 ):
-    """
-    创建本地 FusionModel（无需预训练权重）
-    """
-    model_configs = {
-        "0.5B": dict(vocab_size=32000, hidden_size=2048, num_hidden_layers=16,
-                     num_attention_heads=16, num_key_value_heads=8, intermediate_size=5504),
-        "1.5B": dict(vocab_size=32000, hidden_size=3072, num_hidden_layers=24,
-                     num_attention_heads=24, num_key_value_heads=8, intermediate_size=8192),
-        "8B": dict(vocab_size=100000, hidden_size=4096, num_hidden_layers=32,
-                   num_attention_heads=32, num_key_value_heads=8, intermediate_size=11008),
-        "14B": dict(vocab_size=100000, hidden_size=5120, num_hidden_layers=40,
-                    num_attention_heads=40, num_key_value_heads=8, intermediate_size=13824),
-    }
-    
-    if model_size not in model_configs:
-        raise ValueError(f"不支持的模型大小：{model_size}")
-    
-    config_dict = model_configs[model_size]
-    
-    # S3 fix: override vocab_size to match actual tokenizer
-    if vocab_size_override is not None:
-        config_dict['vocab_size'] = vocab_size_override
-    
-    common_config = dict(
-        block_size=512,
-        latent_dim=64,
-        window_size=2048,
-        sbla_mode="hybrid",
-        rms_norm_eps=1e-6,
-        rope_theta=10000.0,
-        tie_word_embeddings=False,
-        enable_thinking_dial=True,
-        num_thinking_depths=4,
+    """S4 FIX: Delegate to shared model_utils.create_local_model, preserving API."""
+    model = _create_local_model_from_utils(
+        model_size=model_size,
+        torch_dtype=torch_dtype,
+        vocab_size_override=vocab_size_override,
     )
-    
-    config = FusionConfig(**config_dict, **common_config)
-    
-    logger.info(f"[create_local_model] 创建 Fusion-{model_size}（随机初始化）")
-    logger.info(f"  hidden_size={config.hidden_size}, layers={config.num_hidden_layers}, "
-                f"heads={config.num_attention_heads}")
-    
-    model = FusionModel(config)
-    
-    total_params = sum(p.numel() for p in model.parameters())
-    logger.info(f"[create_local_model] 参数总量：{total_params / 1e9:.2f}B")
-    
+    config = model.config
     return model, config
 
 

train/lora_finetune.py

@@ -142,6 +142,9 @@ class FusionDataset(Dataset):
         }
 
 
+from train.model_utils import create_local_model as _create_local_model_from_utils
+
+
 def create_local_model(
     model_size: str = "8B",
     quantize: bool = False,
@@ -149,75 +152,23 @@ def create_local_model(
     load_in_8bit: bool = False,
     vocab_size_override: int | None = None,
 ):
-    """
-    创建本地 FusionModel（无需预训练权重）
-    
-    参数：
-        model_size: "0.5B", "1.5B", "8B", "14B"
-        quantize: 是否量化
-        load_in_4bit: 4-bit 量化（NF4）
-        load_in_8bit: 8-bit 量化
-        vocab_size_override: S3 fix - sync vocab to actual tokenizer size
-    """
-    # 模型配置（基于尺寸）
-    model_configs = {
-        "0.5B": dict(vocab_size=32000, hidden_size=2048, num_hidden_layers=16,
-                     num_attention_heads=16, num_key_value_heads=8, intermediate_size=5504),
-        "1.5B": dict(vocab_size=32000, hidden_size=3072, num_hidden_layers=24,
-                     num_attention_heads=24, num_key_value_heads=8, intermediate_size=8192),
-        "8B": dict(vocab_size=100000, hidden_size=4096, num_hidden_layers=32,
-                   num_attention_heads=32, num_key_value_heads=8, intermediate_size=11008),
-        "14B": dict(vocab_size=100000, hidden_size=5120, num_hidden_layers=40,
-                    num_attention_heads=40, num_key_value_heads=8, intermediate_size=13824),
-    }
-    
-    if model_size not in model_configs:
-        raise ValueError(f"不支持的模型大小：{model_size}，可选：{list(model_configs.keys())}")
-    
-    config_dict = model_configs[model_size]
-    
-    # S3 fix: override vocab_size to match actual tokenizer
-    if vocab_size_override is not None:
-        config_dict['vocab_size'] = vocab_size_override
-    
-    # 通用配置
-    common_config = dict(
-        block_size=512,
-        latent_dim=64,
-        window_size=2048,
-        sbla_mode="hybrid",
-        rms_norm_eps=1e-6,
-        rope_theta=10000.0,
-        tie_word_embeddings=False,
-        enable_thinking_dial=True,
-        num_thinking_depths=4,
+    """S4 FIX: Delegate to shared model_utils.create_local_model, then apply quantization."""
+    model = _create_local_model_from_utils(
+        model_size=model_size,
+        torch_dtype=torch.bfloat16,
+        vocab_size_override=vocab_size_override,
     )
-    
-    config = FusionConfig(**config_dict, **common_config)
-    
-    logger.info(f"[create_local_model] 创建 Fusion-{model_size} 模型")
-    logger.info(f"  vocab_size={config.vocab_size}, hidden_size={config.hidden_size}, "
-                f"layers={config.num_hidden_layers}, heads={config.num_attention_heads}")
-    
-    # 创建模型（随机初始化）
-    model = FusionModel(config)
-    
-    total_params = sum(p.numel() for p in model.parameters())
-    logger.info(f"[create_local_model] 模型参数总量：{total_params / 1e9:.2f}B")
-    
+    config = model.config
+
     # S-NEW-9 FIX: QLoRA requires proper bitsandbytes integration
-    # For local models created from scratch, we can't use HF's AutoModel quantization.
-    # Instead, we quantize the model directly with bitsandbytes if available.
     if quantize:
         if load_in_4bit:
             logger.info("[create_local_model] Using 4-bit quantization (QLoRA)")
             try:
                 import bitsandbytes as bnb
-                # S-NEW-12 FIX: Cache module dict to avoid O(n^2) traversal
                 name_to_module = dict(model.named_modules())
                 for name, module in model.named_modules():
                     if isinstance(module, nn.Linear) and not any(x in name for x in ['lora', 'head', 'embed']):
-                        # Create 4-bit quantized linear (using bitsandbytes nf4)
                         quantized = bnb.nn.Linear4bit(
                             module.in_features,
                             module.out_features,
@@ -225,7 +176,6 @@ def create_local_model(
                             quant_type='nf4',
                             compute_dtype=torch.float16
                         )
-                        # Replace in model
                         parent_name = '.'.join(name.split('.')[:-1])
                         child_name = name.split('.')[-1]
                         if parent_name:
@@ -237,14 +187,12 @@ def create_local_model(
             except ImportError:
                 logger.warning("bitsandbytes not installed, 4-bit quantization DISABLED")
                 logger.warning("Model will train in FP32 - install bitsandbytes for true QLoRA")
-                # M-NEW-16 FIX: Skip prepare_model_for_kbit_training when bnb unavailable
                 return model, config
             model = prepare_model_for_kbit_training(model)
         elif load_in_8bit:
             logger.info("[create_local_model] Using 8-bit quantization")
             try:
                 import bitsandbytes as bnb
-                # S-NEW-12 FIX: Cache module dict to avoid O(n^2) traversal
                 name_to_module = dict(model.named_modules())
                 for name, module in model.named_modules():
                     if isinstance(module, nn.Linear) and not any(x in name for x in ['lora', 'head', 'embed']):
@@ -264,10 +212,9 @@ def create_local_model(
                 logger.info("[create_local_model] 8-bit quantization applied")
             except ImportError:
                 logger.warning("bitsandbytes not installed, 8-bit quantization DISABLED")
-                # M-NEW-16 FIX: Skip prepare_model_for_kbit_training when bnb unavailable
                 return model, config
             model = prepare_model_for_kbit_training(model)
-    
+
     return model, config
 
 
@@ -305,7 +252,13 @@ def apply_lora(
         # L-NEW-2 FIX: Remove "out_proj" (doesn't match model);
         # FusionModel follows LLaMA naming: q_proj/k_proj/v_proj/o_proj for attention,
         # gate_proj/up_proj/down_proj for MLP
-        target_modules = ["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+        # S3 FIX: Include SBLAttention latent projections for proper LoRA coverage
+        target_modules = [
+            "q_proj", "v_proj", "k_proj", "o_proj",
+            "gate_proj", "up_proj", "down_proj",
+            "latent_q_proj", "latent_k_proj", "latent_v_proj", "latent_out_proj",
+            "v_to_hidden_proj",
+        ]
     
     logger.info(f"[apply_lora] 应用 LoRA（rank={lora_rank}, alpha={lora_alpha}）")
     logger.info(f"[apply_lora] 目标模块：{target_modules}")

train/model_utils.py

@@ -0,0 +1,76 @@
+"""Shared model creation utilities for Fusion-LLM training scripts.
+
+S4 FIX: Extract duplicated create_local_model() from lora_finetune.py and
+full_finetune.py into a single source of truth.
+"""
+
+import torch
+import logging
+
+logger = logging.getLogger(__name__)
+
+# Model size presets
+MODEL_CONFIGS = {
+    "0.5B": dict(vocab_size=32000, hidden_size=2048, num_hidden_layers=16,
+                 num_attention_heads=16, num_key_value_heads=8, intermediate_size=5504),
+    "1.5B": dict(vocab_size=32000, hidden_size=3072, num_hidden_layers=24,
+                 num_attention_heads=24, num_key_value_heads=8, intermediate_size=8192),
+    "8B": dict(vocab_size=100000, hidden_size=4096, num_hidden_layers=32,
+               num_attention_heads=32, num_key_value_heads=8, intermediate_size=11008),
+    "14B": dict(vocab_size=100000, hidden_size=5120, num_hidden_layers=40,
+                num_attention_heads=40, num_key_value_heads=8, intermediate_size=13824),
+}
+
+COMMON_CONFIG = dict(
+    block_size=512,
+    latent_dim=64,
+    window_size=2048,
+    sbla_mode="hybrid",
+    rms_norm_eps=1e-6,
+    rope_theta=10000.0,
+    tie_word_embeddings=False,
+    enable_thinking_dial=True,
+    num_thinking_depths=4,
+)
+
+
+def create_local_model(
+    model_size: str = "8B",
+    torch_dtype: torch.dtype = torch.bfloat16,
+    vocab_size_override: int | None = None,
+):
+    """Create a locally-initialized FusionModel (no pretrained weights required).
+
+    Args:
+        model_size: One of "0.5B", "1.5B", "8B", "14B"
+        torch_dtype: Model dtype (default bfloat16)
+        vocab_size_override: Override vocab_size to match actual tokenizer
+
+    Returns:
+        FusionModel instance with random initialization
+    """
+    from models.fusion_model import FusionConfig, FusionModel
+
+    if model_size not in MODEL_CONFIGS:
+        raise ValueError(f"Unsupported model size: {model_size}, options: {list(MODEL_CONFIGS.keys())}")
+
+    config_dict = MODEL_CONFIGS[model_size].copy()
+
+    if vocab_size_override is not None:
+        config_dict['vocab_size'] = vocab_size_override
+
+    config = FusionConfig(**config_dict, **COMMON_CONFIG)
+
+    logger.info(f"[create_local_model] Creating Fusion-{model_size} model")
+    logger.info(f"  vocab_size={config.vocab_size}, hidden_size={config.hidden_size}, "
+                f"layers={config.num_hidden_layers}, heads={config.num_attention_heads}")
+
+    model = FusionModel(config)
+
+    if torch_dtype is not None:
+        model = model.to(torch_dtype)
+
+    param_count = sum(p.numel() for p in model.parameters())
+    logger.info(f"  Parameters: {param_count / 1e9:.2f}B")
+
+    return model