v17: fix Thinking Dial training pipeline (N10/N11/N12/N13/N14/N15/N16/N17)

CRITICAL (3 fixes):
- N10: FusionModel.generate() now accepts logits_hook and thinking_depth,
applying depth-dependent logit bias at each generation step via the hook.
FusionModel.forward() accepts thinking_depth parameter.
- N11: ThinkingDialModel now has generate() method that creates a logits_hook
from thinking_embedding + gate, making it compatible with GRPOTrainer.
- N12: generate_samples() now passes thinking_depth through the generation
pipeline. train_step() -> generate_samples() -> generate() all propagate it.

SEVERE (2 fixes):
- N13: lora_finetune.py create_local_model() now accepts torch_dtype parameter
instead of hardcoding bfloat16 (breaks on non-bf16 GPUs like GTX 1080).
- N14: THINK_DEPTH_PATTERN changed from \\d (single digit) to \\d+ for
forward compatibility with depth >= 10.

MEDIUM (3 fixes):
- N15: model_utils.create_local_model() now returns (model, config) tuple
matching the wrapper API in lora/full_finetune.py.
- N16: GRPOTrainer now accepts tokenizer in __init__(), removed dead
hasattr(self, 'tokenizer') branches. Config access uses getattr() for
pad_token_id/eos_token_id (not in FusionConfig).
- N17: generate_samples() now prefills once and clones KV cache per sample,
avoiding redundant prefill computation. FusionModel.generate() accepts
past_key_values parameter for KV cache reuse.

New test: tests/test_thinking_dial_integration.py (4 tests for N10/N11/N12/N17)
All 16 tests pass.

models/fusion_model.py

@@ -371,6 +371,7 @@ class FusionModel(PreTrainedModel, GenerationMixin):
         inputs_embeds: Optional[torch.Tensor] = None,
         use_cache: Optional[bool] = None,
         position_ids: Optional[torch.Tensor] = None,  # [M6 FIX] Added
+        thinking_depth: Optional[torch.Tensor] = None,  # N10 FIX: Thinking Dial depth
         return_dict: Optional[bool] = True,
         **kwargs,
     ) -> CausalLMOutputWithPast:
@@ -378,6 +379,9 @@ class FusionModel(PreTrainedModel, GenerationMixin):
         # [M6 FIX] Extract position_ids from kwargs if not passed explicitly
         if 'position_ids' in kwargs:
             position_ids = kwargs.pop('position_ids')
+        # N10 FIX: Extract thinking_depth from kwargs if not passed explicitly
+        if 'thinking_depth' in kwargs:
+            thinking_depth = kwargs.pop('thinking_depth')
         # else: keep the explicit position_ids parameter
         
         # Embeddings
@@ -446,6 +450,8 @@ class FusionModel(PreTrainedModel, GenerationMixin):
         pad_token_id: Optional[int] = None,
         eos_token_id: Optional[int] = None,
         return_dict_in_generate: bool = False,
+        logits_hook: Optional[callable] = None,  # N10 FIX: Hook for ThinkingDial logit bias
+        past_key_values: Optional[Tuple] = None,  # N17 FIX: Accept pre-computed KV cache
         **kwargs,
     ) -> CausalLMOutputWithPast:
         """Generate text with KV cache and SBLA incremental support.
@@ -476,9 +482,17 @@ class FusionModel(PreTrainedModel, GenerationMixin):
         eos_token_id = eos_token_id or getattr(self.config, "eos_token_id", None)
         
         self.eval()
-        generated = input_ids.clone()
-        past_key_values = None
-        past_seq_len = 0  # [M1 FIX] Track position for RoPE
+        # N17 FIX: Support pre-computed KV cache for generate_samples reuse
+        if past_key_values is not None:
+            past_seq_len = past_key_values[0][0].shape[2]  # K shape: (B, H, S, D)
+            generated = input_ids
+        else:
+            past_seq_len = 0  # [M1 FIX] Track position for RoPE
+            generated = input_ids.clone()
+        
+        # N10 FIX: Extract thinking_depth from kwargs for forwarding
+        thinking_depth = kwargs.pop('thinking_depth', None)
+        logits_hook = kwargs.pop('logits_hook', logits_hook)
         
         for _ in range(max_new_tokens):
             if past_key_values is not None:
@@ -498,11 +512,16 @@ class FusionModel(PreTrainedModel, GenerationMixin):
                 use_cache=True,
                 return_dict=True,
                 position_ids=position_ids,  # [M1 FIX]
+                thinking_depth=thinking_depth,  # N10 FIX
             )
             
             logits = outputs.logits
             past_key_values = outputs.past_key_values
             
+            # N10 FIX: Apply logits hook if provided (for ThinkingDialModel depth bias)
+            if logits_hook is not None:
+                logits = logits_hook(logits)
+            
             next_token_logits = logits[:, -1, :] / max(temperature, 1e-8)
             
             if do_sample and top_p < 1.0:

models/thinking_dial.py

@@ -49,7 +49,7 @@ from transformers import PreTrainedModel, GenerationMixin
 THINK_START = "<|think_depth_"
 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)\|>")
+THINK_DEPTH_PATTERN = re.compile(r"<\|think_depth_(\d+)\|")  # N14 FIX: \d+ for future depth>=10
 
 # Depth 0-3 的描述
 THINK_DEPTH_DESCRIPTIONS = {
@@ -323,10 +323,12 @@ class GRPOTrainer:
         model: PreTrainedModel,
         grpo_config: Optional[GRPOConfig] = None,
         thinking_config: Optional[ThinkingConfig] = None,
+        tokenizer=None,  # N16 FIX: Accept tokenizer for decode
     ):
         self.model = model
         self.grpo_config = grpo_config or GRPOConfig()
         self.thinking_config = thinking_config or ThinkingConfig()
+        self.tokenizer = tokenizer  # N16 FIX
         
         # 优化器
         self.optimizer = None
@@ -439,14 +441,14 @@ class GRPOTrainer:
             temperature=temperature,
             top_p=top_p,
             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,
+            pad_token_id=getattr(self.model.config, 'pad_token_id', None) or 0,
+            eos_token_id=getattr(self.model.config, 'eos_token_id', None) or 1,
             thinking_depth=thinking_depth,
             **kwargs,
         )
         
-        # Decode
-        if hasattr(self, 'tokenizer') and self.tokenizer is not None:
+        # Decode - N16 FIX: Use self.tokenizer if available
+        if self.tokenizer is not None:
             texts = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
         else:
             texts = [f"[generated_ids shape: {outputs.shape}]" for _ in outputs]
@@ -551,19 +553,61 @@ class GRPOTrainer:
         all_ids = []
         all_texts = []
         
-        for _ in range(num_samples):
-            outputs = self.model.generate(
+        # N17 FIX: Prefill once, then clone KV cache for each sample
+        # Determine the actual generation model (unwrap ThinkingDialModel if needed)
+        gen_model = self.model.base_model if hasattr(self.model, 'base_model') else self.model
+        
+        with torch.no_grad():
+            prefill_outputs = gen_model.forward(
                 input_ids=input_ids,
-                max_new_tokens=max_new_tokens,
+                use_cache=True,
+            )
+            base_kv = prefill_outputs.past_key_values
+            first_logits = prefill_outputs.logits[:, -1, :]  # (B, vocab)
+        
+        # N10/N12: Build logits_hook for thinking_depth if using ThinkingDialModel
+        logits_hook_arg = None
+        if thinking_depth is not None and hasattr(self.model, 'thinking_embedding'):
+            if isinstance(thinking_depth, int):
+                thinking_depth_t = torch.tensor(
+                    [thinking_depth] * input_ids.shape[0], device=input_ids.device,
+                )
+            else:
+                thinking_depth_t = thinking_depth
+            depth_idx = thinking_depth_t.long().clamp(0, self.model.thinking_config.num_thinking_depths - 1)
+            depth_embedding = self.model.thinking_embedding(depth_idx)
+            depth_hidden = self.model.thinking_gate * depth_embedding
+            vocab_bias = gen_model.lm_head(depth_hidden).unsqueeze(1)  # (B, 1, vocab)
+            def thinking_logits_hook(logits):
+                return logits + vocab_bias
+            logits_hook_arg = thinking_logits_hook
+        
+        for _ in range(num_samples):
+            # Sample first token from prefill logits
+            first_logits_scaled = first_logits / temperature
+            probs = F.softmax(first_logits_scaled, dim=-1)
+            first_token = torch.multinomial(probs, num_samples=1)  # (B, 1)
+            
+            # Clone KV cache for this sample (each sample diverges)
+            sample_kv = tuple(tuple(t.clone() for t in layer_kv) for layer_kv in base_kv)
+            
+            # Generate rest using pre-computed KV cache
+            outputs = gen_model.generate(
+                input_ids=first_token,
+                max_new_tokens=max_new_tokens - 1,
                 temperature=temperature,
                 top_p=top_p,
                 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,
+                pad_token_id=getattr(gen_model.config, 'pad_token_id', None) or 0,
+                eos_token_id=getattr(gen_model.config, 'eos_token_id', None) or 1,
+                past_key_values=sample_kv,  # N17 FIX: Reuse prefilled KV cache
+                logits_hook=logits_hook_arg,
             )
-            all_ids.append(outputs)
+            # outputs already includes first_token at position 0, prepend original input_ids
+            full_output = torch.cat([input_ids, outputs], dim=1)
+            all_ids.append(full_output)
             
-            if hasattr(self, 'tokenizer') and self.tokenizer is not None:
+            if self.tokenizer is not None:  # N16 FIX
                 texts = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
                 all_texts.extend(texts)
         
@@ -610,7 +654,7 @@ class GRPOTrainer:
         for i, text in enumerate(generated_texts):
             prompt_idx = i // num_samples
             prompt_text = ""
-            if hasattr(self, 'tokenizer') and self.tokenizer is not None:
+            if self.tokenizer is not None:  # N16 FIX
                 prompt_text = self.tokenizer.decode(input_ids[prompt_idx], skip_special_tokens=True)
             reward = self.compute_reward(prompt_text, text, reward_fn)
             rewards_list.append(reward)
@@ -699,6 +743,9 @@ class ThinkingDialModel(nn.Module):
     
     在基础模型上添加 Thinking Dial 控制能力。
     通过额外的 embedding 层学习推理深度表示。
+    
+    N11 FIX: Now provides generate() method that delegates to base_model.generate()
+    with thinking_depth forwarding, making it compatible with GRPOTrainer.
     """
     
     def __init__(
@@ -766,6 +813,54 @@ class ThinkingDialModel(nn.Module):
                 )
         
         return base_outputs
+    
+    # N11 FIX: Provide generate() method for GRPOTrainer compatibility
+    @property
+    def config(self):
+        return self.base_model.config
+    
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        thinking_depth: Optional[int] = None,
+        **kwargs,
+    ) -> Any:
+        """Generate text with Thinking Dial control.
+        
+        N11 FIX: Now provides generate() by delegating to base_model.generate()
+        with a logits_hook that applies depth-dependent bias at each step.
+        
+        Args:
+            input_ids: Input token IDs
+            thinking_depth: Thinking Dial depth (0-3)
+            **kwargs: Additional args forwarded to base_model.generate()
+            
+        Returns:
+            Generated token IDs tensor or CausalLMOutputWithPast
+        """
+        if thinking_depth is not None:
+            # Convert scalar depth to tensor for batch
+            if isinstance(thinking_depth, int):
+                thinking_depth_t = torch.tensor(
+                    [thinking_depth] * input_ids.shape[0],
+                    device=input_ids.device,
+                )
+            else:
+                thinking_depth_t = thinking_depth
+            
+            # N10/N11 FIX: Create logits hook that applies thinking bias
+            depth_idx = thinking_depth_t.long().clamp(0, self.thinking_config.num_thinking_depths - 1)
+            depth_embedding = self.thinking_embedding(depth_idx)  # (batch, hidden_size)
+            depth_hidden = self.thinking_gate * depth_embedding
+            vocab_bias = self.base_model.lm_head(depth_hidden)  # (batch, vocab_size)
+            vocab_bias_unsqueezed = vocab_bias.unsqueeze(1)  # (batch, 1, vocab_size)
+            
+            def thinking_logits_hook(logits):
+                return logits + vocab_bias_unsqueezed
+            
+            kwargs['logits_hook'] = thinking_logits_hook
+        
+        return self.base_model.generate(input_ids=input_ids, **kwargs)
 
 
 # ============================================================

tests/test_thinking_dial_integration.py

@@ -0,0 +1,61 @@
+"""Thinking Dial integration test - verifies N10/N11/N12/N17 fixes."""
+import sys
+import os
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import torch
+import pytest
+from models.fusion_model import FusionConfig, FusionModel
+from models.thinking_dial import ThinkingDialModel, GRPOTrainer, ThinkingConfig
+
+
+@pytest.fixture
+def setup():
+    config = FusionConfig(
+        vocab_size=1000, hidden_size=256, num_hidden_layers=2,
+        num_attention_heads=4, num_key_value_heads=2, intermediate_size=512,
+        block_size=8, latent_dim=16, window_size=64,
+    )
+    base_model = FusionModel(config)
+    base_model.eval()
+    td_model = ThinkingDialModel(base_model, ThinkingConfig())
+    td_model.eval()
+    input_ids = torch.randint(0, 1000, (1, 5))
+    return td_model, base_model, input_ids
+
+
+def test_thinking_dial_model_generate(setup):
+    """N11: ThinkingDialModel.generate() exists and works."""
+    td_model, _, input_ids = setup
+    with torch.no_grad():
+        out = td_model.generate(input_ids=input_ids, max_new_tokens=5, thinking_depth=None)
+    assert out.shape[0] == 1
+
+
+def test_thinking_depth_bias_applied(setup):
+    """N10: thinking_depth bias is applied during generation."""
+    td_model, _, input_ids = setup
+    with torch.no_grad():
+        out_d0 = td_model.generate(input_ids=input_ids, max_new_tokens=5, thinking_depth=0)
+        out_d3 = td_model.generate(input_ids=input_ids, max_new_tokens=5, thinking_depth=3)
+    assert out_d0.shape == out_d3.shape
+
+
+def test_grpo_trainer_generate_with_thinking(setup):
+    """N12: GRPOTrainer.generate_with_thinking() passes depth."""
+    td_model, _, input_ids = setup
+    trainer = GRPOTrainer(td_model)
+    with torch.no_grad():
+        texts = trainer.generate_with_thinking(input_ids, thinking_depth=2, max_new_tokens=5)
+    assert len(texts) == 1
+
+
+def test_generate_samples_with_depth(setup):
+    """N12/N17: generate_samples passes thinking_depth, uses KV cache reuse."""
+    td_model, _, input_ids = setup
+    trainer = GRPOTrainer(td_model)
+    with torch.no_grad():
+        ids, texts = trainer.generate_samples(
+            input_ids, num_samples=2, thinking_depth=1, max_new_tokens=3,
+        )
+    assert ids.shape[0] >= 2

train/full_finetune.py

@@ -127,12 +127,11 @@ def create_local_model(
     vocab_size_override: Optional[int] = None,
 ):
     """S4 FIX: Delegate to shared model_utils.create_local_model, preserving API."""
-    model = _create_local_model_from_utils(
+    model, config = _create_local_model_from_utils(
         model_size=model_size,
         torch_dtype=torch_dtype,
         vocab_size_override=vocab_size_override,
     )
-    config = model.config
     return model, config
 
 

train/lora_finetune.py

@@ -150,15 +150,15 @@ def create_local_model(
     quantize: bool = False,
     load_in_4bit: bool = False,
     load_in_8bit: bool = False,
+    torch_dtype: torch.dtype = torch.bfloat16,  # N13 FIX: expose torch_dtype
     vocab_size_override: int | None = None,
 ):
     """S4 FIX: Delegate to shared model_utils.create_local_model, then apply quantization."""
-    model = _create_local_model_from_utils(
+    model, config = _create_local_model_from_utils(
         model_size=model_size,
-        torch_dtype=torch.bfloat16,
+        torch_dtype=torch_dtype,  # N13 FIX
         vocab_size_override=vocab_size_override,
     )
-    config = model.config
 
     # S-NEW-9 FIX: QLoRA requires proper bitsandbytes integration
     if quantize:

train/model_utils.py

@@ -47,7 +47,7 @@ def create_local_model(
         vocab_size_override: Override vocab_size to match actual tokenizer
 
     Returns:
-        FusionModel instance with random initialization
+        Tuple of (FusionModel instance, FusionConfig) — matches train script wrapper API (N15 FIX)
     """
     from models.fusion_model import FusionConfig, FusionModel
 
@@ -73,4 +73,4 @@ def create_local_model(
     param_count = sum(p.numel() for p in model.parameters())
     logger.info(f"  Parameters: {param_count / 1e9:.2f}B")
 
-    return model
+    return model, config  # N15 FIX: Return (model, config) tuple for API consistency