fix: F1 mask double-conversion NaN, F2 think token old format, S3 vocab sync, S1/S2 TODO, N1-N5

configs/ds_zero3.json

@@ -15,7 +15,7 @@
         },
         "overlap_comm": true,
         "contiguous_gradients": true,
-        "sub_group_size": 1e9,
+        "sub_group_size": 1e6,
         "reduce_bucket_size": "auto",
         "stage3_prefetch_bucket_size": "auto",
         "stage3_param_persistence_threshold": "auto",

configs/fusion-mini-config.json

@@ -26,7 +26,6 @@
   "tie_word_embeddings": false,
   "enable_thinking_dial": true,
   "num_thinking_depths": 4,
-  "think_rank": 0,
 
   "torch_dtype": "float32",
   "transformers_version": "4.36.0",

models/fusion_mini.py

@@ -357,11 +357,9 @@ class FusionMini(PreTrainedModel):
         # 1. Embeddings
         hidden_states = self.embeddings(input_ids)
         
-        # 2. 处理 attention_mask
-        if attention_mask is not None:
-            # 转换为 (batch, 1, 1, seq_len) 格式
-            attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
-            attention_mask = (1.0 - attention_mask) * -10000.0
+        # 2. 处理 attention_mask - pass raw HF format to layers
+        # Each FusionMiniLayer handles mask conversion internally
+        # DO NOT pre-convert here to avoid double-conversion NaN (F1)
         
         # 3. Transformer 层
         for layer in self.layers:

models/fusion_model.py

@@ -96,8 +96,8 @@ class FusionConfig(PretrainedConfig):
         # SBLA 参数
         self.block_size = block_size
         self.latent_dim = latent_dim
-        self.sbla_window_size = sbla_window_size or block_size
-        self.window_size = window_size or self.sbla_window_size  # Unified field
+        self.window_size = window_size or sbla_window_size or block_size
+        self.sbla_window_size = self.window_size  # Keep as alias for backward compat
         self.sbla_mode = sbla_mode
         
         # Thinking Dial 参数
@@ -134,6 +134,10 @@ class FusionAttention(nn.Module):
     def __init__(self, config: FusionConfig):
         super().__init__()
         
+        # TODO(S2): GQA not yet implemented - num_key_value_heads is stored in config
+        # but SBLAttention uses num_attention_heads for all Q/K/V projections.
+        # To implement GQA: split Q/K/V projections, use num_key_value_heads for K/V,
+        # and add repeat_kv() expansion for attention computation.
         mode = getattr(config, 'sbla_mode', 'hybrid')
         self.sbla = SBLAttention(
             hidden_size=config.hidden_size,
@@ -152,12 +156,15 @@ class FusionAttention(nn.Module):
         past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         use_cache: bool = False,
     ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
-        # Delegate to SBLAttention
+        # TODO(S1): KV Cache is currently decorative - SBLA recomputes full Q/K/V internally.
+        # For true incremental generation, SBLA.forward needs to accept past_key_value
+        # and only compute Q for new tokens, reusing cached K/V.
+        # Present implementation: always full recomputation, cache returned but unused.
         output = self.sbla(hidden_states, attention_mask)
         
-        # KV Cache: extract from SBLA's Q/K projections for cache reuse
         present_key_value = None
         if use_cache:
+            # Cache K/V for potential future use when SBLA supports incremental mode
             batch_size, seq_len, _ = hidden_states.shape
             K = self.sbla.k_proj(hidden_states)
             V = self.sbla.v_proj(hidden_states)
@@ -277,12 +284,14 @@ class FusionModel(PreTrainedModel, GenerationMixin):
         else:
             raise ValueError("Either input_ids or inputs_embeds must be provided")
         
-        # 处理 attention_mask
-        if attention_mask is not None:
-            if attention_mask.dim() == 2:
-                attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
-            float_mask = attention_mask.to(dtype=hidden_states.dtype)
-            attention_mask = (1.0 - float_mask) * torch.finfo(hidden_states.dtype).min
+        # 处理 attention_mask - pass raw HF format (1=valid, 0=padding) to SBLA
+        # SBLAttention handles the conversion internally
+        # DO NOT convert here - it would cause double-conversion NaN (F1)
+        # if attention_mask is not None:
+        #     if attention_mask.dim() == 2:
+        #         attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+        #     float_mask = attention_mask.to(dtype=hidden_states.dtype)
+        #     attention_mask = (1.0 - float_mask) * torch.finfo(hidden_states.dtype).min
         
         # Transformer 层（支持 KV Cache）
         past_key_values = kwargs.get("past_key_values", None)

models/sbla_attention.py

@@ -275,12 +275,24 @@ class SBLAttention(nn.Module):
             combined_mask = causal_mask
         
         # 应用外部 attention_mask（padding mask）
+        # Supports both raw HF format (batch, seq_len) with 1=valid/0=padding,
+        # and pre-expanded format (batch, 1, 1, seq_len).
         if attention_mask is not None:
-            # attention_mask: (batch, 1, 1, seq_len) -> 扩展为 (batch, 1, seq_len, seq_len)
-            ext_mask = attention_mask.squeeze(1)  # (batch, 1, seq_len)
-            # 将 padding 位置设为 -inf
-            padding_mask = (1.0 - ext_mask) * float('-inf')  # (batch, 1, seq_len)
-            combined_mask = combined_mask.unsqueeze(0) + padding_mask.unsqueeze(1)  # (batch, 1, seq_len, seq_len)
+            if attention_mask.dim() == 2:
+                # Raw HF format: (batch, seq_len), 1=valid, 0=padding
+                padding_mask = (1.0 - attention_mask.float()).unsqueeze(1).unsqueeze(2)  # (batch, 1, 1, seq_len)
+                padding_mask = padding_mask * torch.finfo(hidden_states.dtype).min
+                combined_mask = combined_mask.unsqueeze(0).unsqueeze(0) + padding_mask  # (1, 1, seq, seq) + (batch, 1, 1, seq)
+            elif attention_mask.dim() == 4:
+                # Pre-expanded format (already converted, e.g. from FusionMini)
+                ext_mask = attention_mask.squeeze(1)  # (batch, 1, seq_len)
+                padding_mask = (1.0 - ext_mask) * float('-inf')  # (batch, 1, seq_len)
+                combined_mask = combined_mask.unsqueeze(0) + padding_mask.unsqueeze(1)
+            else:
+                # 3D fallback
+                padding_mask = (1.0 - attention_mask.float()).unsqueeze(1)  # (batch, 1, 1, seq_len)
+                padding_mask = padding_mask * torch.finfo(hidden_states.dtype).min
+                combined_mask = combined_mask.unsqueeze(0).unsqueeze(0) + padding_mask
         else:
             combined_mask = combined_mask.unsqueeze(0)  # (1, 1, seq_len, seq_len)
         

models/thinking_dial.py

@@ -36,7 +36,7 @@ Thinking Dial（动态推理强度控制）- 真实实现
 import torch
 import torch.nn as nn
 import re
-from typing import List, Dict, Optional, Any
+from typing import List, Dict, Optional, Any, Tuple
 from dataclasses import dataclass
 from transformers import PreTrainedModel, GenerationMixin
 

scripts/fix_think_tokens.py

@@ -1,57 +1,65 @@
 #!/usr/bin/env python3
-"""Fix think token naming consistency across the project."""
+"""Fix think token format across the project.
+
+Ensures all files use the unified <|think_depth_N|> format,
+not the old <|think| depth=N|> format.
+"""
 import re
-import glob
+import sys
+from pathlib import Path
+
+# Old format patterns (to replace)
+OLD_INLINE = re.compile(r'<\|think\|\s*depth=(\d+)\|>')
+OLD_FSTRING = re.compile(r'f"<\|think\|\s*depth=\{(\w+)\}\|>"')
+OLD_TEMPLATE = '<|think| depth='
 
-# Target format: <|think_depth_0|>, <|think_depth_1|>, etc.
+# New format
+NEW_INLINE = r'<|think_depth_\1|>'
+NEW_FSTRING = r'f"<|think_depth_{\1}|>"'
 
-def fix_file(filepath):
-    with open(filepath, 'r', encoding='utf-8') as f:
-        content = f.read()
+
+def fix_file(filepath: Path) -> bool:
+    """Fix think tokens in a single file. Returns True if changes were made."""
+    try:
+        content = filepath.read_text(encoding='utf-8')
+    except Exception:
+        return False
     
     original = content
     
-    # Replace THINK_START/THINK_END constants
-    content = content.replace('THINK_START = "<|think_depth_"', 'THINK_START = "<|think_depth_"')
-    content = content.replace('THINK_END = "|>"', 'THINK_END = "|>"')
-    
-    # Replace build_think_token return
-    content = content.replace(
-        'return f"{THINK_START}{depth}{THINK_END}"',
-        'return f"{THINK_START}{depth}{THINK_END}"'
-    )
+    # Fix inline occurrences: <|think| depth=N|> -> <|think_depth_N|>
+    content = OLD_INLINE.sub(NEW_INLINE, content)
     
-    # Replace THINK_DEPTH_PATTERN regex
-    content = content.replace(
-        'THINK_DEPTH_PATTERN = re.compile(r"<\\|think\\| depth=(\\d+)\\|>")',
-        'THINK_DEPTH_PATTERN = re.compile(r"<\\|think_depth_(\\d+)\\|>")'
-    )
-    
-    # Replace any inline <|think| depth=N|> with <|think_depth_N|>
-    content = re.sub(r'<\|think\|\s*depth=(\d+)\|>', r'<|think_depth_\1|>', content)
+    # Fix f-string templates: f"<|think| depth={var}|>" -> f"<|think_depth_{var}|>"
+    content = OLD_FSTRING.sub(NEW_FSTRING, content)
     
     if content != original:
-        with open(filepath, 'w', encoding='utf-8') as f:
-            f.write(content)
-        print(f"  Fixed: {filepath}")
+        filepath.write_text(content, encoding='utf-8')
+        count = sum(1 for a, b in zip(original.split('\n'), content.split('\n')) if a != b)
+        print(f"  Fixed {filepath} ({count} lines)")
         return True
-    else:
-        print(f"  No change: {filepath}")
-        return False
+    return False
 
-files = glob.glob("**/*.py", recursive=True) + glob.glob("**/*.json", recursive=True)
-fixed = 0
-for f in sorted(files):
-    # Skip data files and output
-    if any(skip in f for skip in ['node_modules', '.git', 'output/']):
-        continue
-    try:
-        with open(f, 'r', encoding='utf-8') as fh:
-            text = fh.read()
-        if '<|think|' in text or 'think| depth=' in text:
-            if fix_file(f):
-                fixed += 1
-    except:
-        pass
-
-print(f"\nTotal files fixed: {fixed}")
+
+def main():
+    project_root = Path(__file__).parent.parent
+    
+    # Scan all Python files
+    changed = 0
+    for py_file in project_root.rglob('*.py'):
+        if fix_file(py_file):
+            changed += 1
+    
+    # Also check JSON data files
+    for json_file in project_root.rglob('*.json'):
+        if 'node_modules' in str(json_file):
+            continue
+        if fix_file(json_file):
+            changed += 1
+    
+    print(f"\nTotal files fixed: {changed}")
+    return 0
+
+
+if __name__ == '__main__':
+    sys.exit(main())

train/full_finetune.py

@@ -89,7 +89,7 @@ class FusionFullFinetuneDataset(Dataset):
         think_rank = item.get("think_rank", 0)
         
         if think_rank > 0:
-            thinking_token = f"<|think| depth={think_rank}|>"
+            thinking_token = f"<|think_depth_{think_rank}|>"
             full_text = f"{thinking_token}\n{prompt}\n{response}"
         else:
             full_text = f"{prompt}\n{response}"
@@ -112,6 +112,7 @@ class FusionFullFinetuneDataset(Dataset):
 def create_local_model(
     model_size: str = "8B",
     torch_dtype: torch.dtype = torch.bfloat16,
+    vocab_size_override: Optional[int] = None,
 ):
     """
     创建本地 FusionModel（无需预训练权重）
@@ -132,6 +133,12 @@ def create_local_model(
     
     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
+    if vocab_size_override is not None:
+        config_dict['vocab_size'] = vocab_size_override
+    
     common_config = dict(
         block_size=512,
         latent_dim=64,
@@ -146,6 +153,18 @@ def create_local_model(
     
     config = FusionConfig(**config_dict, **common_config)
     
+    # Override vocab_size if tokenizer has different size (S3 fix)
+    if vocab_size_override is not None and vocab_size_override != config.vocab_size:
+        logger.warning(f"[S3-fix] Overriding model vocab_size: {config.vocab_size} -> {vocab_size_override}")
+        config.vocab_size = vocab_size_override
+    
+    # S3-fix: sync vocab_size to actual tokenizer if provided
+    if vocab_size_override is not None:
+        if vocab_size_override != config.vocab_size:
+            logger.warning(f"[S3-fix] Overriding vocab_size: {config.vocab_size} -> {vocab_size_override}")
+            model.resize_token_embeddings(vocab_size_override)
+        config.vocab_size = vocab_size_override
+    
     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}")
@@ -162,8 +181,7 @@ def create_tokenizer(tokenizer_type: str = "fusion", vocab_size: int = 32000):
     """
     Create tokenizer using the unified tokenizer module.
     """
-    effective_vocab = get_effective_vocab_size(tokenizer_type, vocab_size)
-    logger.info(f"[create_tokenizer] Creating tokenizer: type={tokenizer_type}, effective_vocab={effective_vocab}")
+    logger.info(f"[create_tokenizer] Creating tokenizer: type={tokenizer_type}, vocab_size={vocab_size}")
     tokenizer = get_tokenizer(tokenizer_type, vocab_size=vocab_size)
     return tokenizer
 
@@ -192,8 +210,14 @@ def train(args):
     vocab_size_map = {"0.5B": 32000, "1.5B": 32000, "8B": 100000, "14B": 100000}
     tokenizer = create_tokenizer(vocab_size=vocab_size_map.get(args.model_size, 32000))
     
+    # Sync vocab_size to actual tokenizer size to prevent index-out-of-range (S3)
+    actual_vocab_size = len(tokenizer)
+    if actual_vocab_size != vocab_size_map.get(args.model_size, 32000):
+        logger.warning(f"[S3-fix] Vocab size mismatch: config={vocab_size_map.get(args.model_size, 32000)}, tokenizer={actual_vocab_size}. Syncing to tokenizer.")
+        vocab_size_map[args.model_size] = actual_vocab_size
+    
     # 3. 创建模型（本地随机初始化）
-    model, config = create_local_model(args.model_size, torch_dtype=args.torch_dtype)
+    model, config = create_local_model(args.model_size, torch_dtype=args.torch_dtype, vocab_size_override=actual_vocab_size)
     
     # 4. 加载数据集
     train_dataset = FusionFullFinetuneDataset(

train/lora_finetune.py

@@ -104,7 +104,7 @@ class FusionDataset(Dataset):
         
         # 注入 Thinking Dial 控制 token
         if self.add_thinking_token and think_rank > 0:
-            thinking_token = f"<|think| depth={think_rank}|>"
+            thinking_token = f"<|think_depth_{think_rank}|>"
             full_text = f"{thinking_token}\n{prompt}\n{response}"
         else:
             full_text = f"{prompt}\n{response}"