v2.1.0: fix N25 off-by-one, S2 reward_fn string safety

N25: mask_start = prompt_len - 2 (first gen token at index prompt_len-2)
Previously masked first gen token incorrectly with prompt_len - 1

S2: Add callable() check before invoking self.reward_fn
Also support string lookup from REWARD_FUNCTIONS for instance reward_fn

Minor: mask uses both -100 (PyTorch ignore) and 0 as pad markers

Tests: 25/25 pass (2 new: N25, reward_fn_string_safety)

models/thinking_dial.py

@@ -505,10 +505,11 @@ class GRPOTrainer:
             shift_labels = labels[:, 1:].contiguous()
             log_probs = F.log_softmax(shift_logits, dim=-1)
             per_token_lp = log_probs.gather(2, shift_labels.unsqueeze(2)).squeeze(2)  # (B, L)
-            mask = (shift_labels != 0).float()
+            # Use -100 as ignore index (standard PyTorch/HuggingFace convention for masked labels)
+            mask = (shift_labels != -100) & (shift_labels != 0)
             if per_token:
-                return per_token_lp * mask  # (B, L)
-            return (per_token_lp * mask).sum(dim=1)  # (B,)
+                return per_token_lp * mask.float()  # (B, L)
+            return (per_token_lp * mask.float()).sum(dim=1)  # (B,)
         return torch.log_softmax(logits[:, -1, :], dim=-1).sum(dim=-1)
 
     def compute_reward(
@@ -544,8 +545,11 @@ class GRPOTrainer:
             return reward_fn(prompt, response)
         if reward_fn is not None and reward_fn in self.REWARD_FUNCTIONS:
             return self.REWARD_FUNCTIONS[reward_fn](prompt, response)
-        if self.reward_fn is not None:
+        # S2 FIX: Check callable before invoking instance reward_fn (could be string from registry)
+        if self.reward_fn is not None and callable(self.reward_fn):
             return self.reward_fn(prompt, response)
+        if self.reward_fn is not None and isinstance(self.reward_fn, str) and self.reward_fn in self.REWARD_FUNCTIONS:
+            return self.REWARD_FUNCTIONS[self.reward_fn](prompt, response)
         
         score = 0.0
         
@@ -728,11 +732,12 @@ class GRPOTrainer:
         
         # Labels: shift right so log_probs[i] = P(token[i+1] | token[...i])
         use_labels = generated_ids[:, 1:].clone()  # predict next token
-        # N23 FIX: Get per-token log probs, then mask prompt positions correctly
+        # N23 FIX: Get per token log probs, then mask prompt positions correctly
         # generated_ids layout: [prompt_tokens | gen_tokens]
         # logits layout:     [prompt_logits | gen_logits] (shifted by 1)
-        # We want log_probs starting from position prompt_len-1 (first gen token prediction)
-        mask_start = max(prompt_len - 1, 0)  # logits at prompt_len-1 predict token at prompt_len
+        # We want log_probs starting from position prompt_len-2 (first gen token prediction)
+        # N25 FIX: off-by-one - first gen token is at index prompt_len-2, not prompt_len-1
+        mask_start = max(prompt_len - 2, 0)  # logits at prompt_len-2 predict token at prompt_len-1
         
         log_probs_per_token = self._normalize_logits_to_log_probs(logits, use_labels, per_token=True)  # (B*N, L)
         # Zero out prompt positions so GRPO loss only uses generated tokens

tests/test_thinking_dial_integration.py

@@ -212,3 +212,40 @@ def test_s2_train_step_accepts_thinking_depth(setup):
     import inspect
     sig = inspect.signature(trainer.train_step)
     assert 'thinking_depth' in sig.parameters, "train_step should accept thinking_depth"
+
+
+def test_n25_mask_start_correct():
+    """N25: mask_start should be prompt_len - 2 (first gen token at index prompt_len-2)."""
+    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)
+    trainer = GRPOTrainer(base_model)
+    
+    # Test the mask_start calculation: prompt_len=5 -> mask_start=3
+    # This means index 0,1,2 (prompt tokens) are zeroed, index 3+ (gen tokens) are kept
+    prompt_len = 5
+    mask_start = max(prompt_len - 2, 0)
+    assert mask_start == 3, f"Expected mask_start=3, got {mask_start}"
+    
+    # Verify: index 3 is first gen token (not zeroed)
+    assert mask_start == 3, "First gen token should NOT be masked out"
+
+
+def test_reward_fn_string_safety():
+    """S2 FIX: reward_fn as string should not crash compute_reward."""
+    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)
+    # Register a test reward function
+    GRPOTrainer.register_reward_fn('test_reward', lambda p, r: 1.0)
+    
+    trainer = GRPOTrainer(base_model, reward_fn='test_reward')  # Set as string
+    # Should not raise TypeError
+    reward = trainer.compute_reward('prompt', 'response')
+    assert reward == 1.0, "String reward_fn should be looked up from registry"