Upload folder using huggingface_hub

export_to_hf.py

@@ -4,6 +4,13 @@ import json
 import shutil
 from pathlib import Path
 
+import torch
+
+
+IGNORED_CHECKPOINT_SUFFIXES = (
+    ".rotary.inv_freq",
+)
+
 
 def normalize_checkpoint(checkpoint):
     if isinstance(checkpoint, dict):
@@ -26,6 +33,18 @@ def infer_special_token_paths(repo_dir):
     return repo_dir / "tokenizer" / "tokenizer.model", subdir_metadata, repo_dir / "tokenizer" / "tokenizer.vocab"
 
 
+def sanitize_model_state(model_state):
+    """Drop deterministic non-persistent buffers that should not participate in HF weight export."""
+    sanitized = {}
+    ignored = []
+    for key, value in model_state.items():
+        if key.endswith(IGNORED_CHECKPOINT_SUFFIXES):
+            ignored.append(key)
+            continue
+        sanitized[key] = value
+    return sanitized, ignored
+
+
 def write_clean_tokenizer_metadata(repo_dir, special_tokens):
     tokenizer_config = {
         "backend": "custom",
@@ -67,8 +86,6 @@ def write_clean_tokenizer_metadata(repo_dir, special_tokens):
 
 
 def main():
-    import torch
-
     from configuration_taonet import TaoNetConfig
     from modeling_taonet import TaoNetForCausalLM
     from tokenization_taonet import TaoNetTokenizer
@@ -78,6 +95,7 @@ def main():
 
     checkpoint = torch.load(checkpoint_path, map_location="cpu")
     model_state, train_config = normalize_checkpoint(checkpoint)
+    model_state, ignored_keys = sanitize_model_state(model_state)
     model_config = dict(train_config.get("model", {}))
 
     metadata_model_path, metadata_path, vocab_path = infer_special_token_paths(repo_dir)
@@ -100,8 +118,43 @@ def main():
     }
 
     model = TaoNetForCausalLM(hf_config)
-    missing, unexpected = model.model.load_state_dict(model_state, strict=False)
+    current_state = model.model.state_dict()
+    missing = sorted(set(current_state) - set(model_state))
+    unexpected = sorted(set(model_state) - set(current_state))
+    if missing or unexpected:
+        if missing:
+            print("Missing keys while loading model state:")
+            for key in missing:
+                print(f"  - {key}")
+        if unexpected:
+            print("Unexpected keys while loading model state:")
+            for key in unexpected:
+                print(f"  - {key}")
+        raise ValueError("Checkpoint/model key mismatch detected. Refusing to export a partial model.")
+
+    model.model.load_state_dict(model_state, strict=True)
     model.tie_weights()
+
+    exported_state = model.model.state_dict()
+    mismatched_tensors = []
+    for key, value in model_state.items():
+        exported_value = exported_state[key]
+        if value.shape != exported_value.shape:
+            mismatched_tensors.append((key, "shape"))
+            continue
+        if value.dtype.is_floating_point:
+            if not torch.equal(value, exported_value.to(dtype=value.dtype)):
+                mismatched_tensors.append((key, "value"))
+        else:
+            if not torch.equal(value, exported_value):
+                mismatched_tensors.append((key, "value"))
+
+    if mismatched_tensors:
+        print("Tensor mismatches detected after strict load:")
+        for key, mismatch_type in mismatched_tensors[:20]:
+            print(f"  - {key} ({mismatch_type})")
+        raise ValueError("Checkpoint tensors do not match the HF wrapper after load.")
+
     model.save_pretrained(repo_dir, safe_serialization=False)
 
     tokenizer = TaoNetTokenizer(
@@ -115,13 +168,9 @@ def main():
     shutil.copyfile(metadata_path, repo_dir / "tokenizer.special_tokens.json")
     shutil.copyfile(vocab_path, repo_dir / "tokenizer.vocab")
 
-    if missing:
-        print("Missing keys while loading model state:")
-        for key in missing:
-            print(f"  - {key}")
-    if unexpected:
-        print("Unexpected keys while loading model state:")
-        for key in unexpected:
+    if ignored_keys:
+        print("Ignored non-persistent checkpoint buffers:")
+        for key in ignored_keys:
             print(f"  - {key}")
     print(f"Saved Hugging Face package to: {repo_dir}")
 

verify_export_weights.py

@@ -0,0 +1,82 @@
+"""Verify that TaoTrain checkpoint weights match the exported HF package exactly."""
+
+from pathlib import Path
+
+import torch
+
+from export_to_hf import normalize_checkpoint, sanitize_model_state
+from modeling_taonet import TaoNetForCausalLM
+
+
+def summarize_tensor_diff(left: torch.Tensor, right: torch.Tensor) -> str:
+    if left.shape != right.shape:
+        return f"shape mismatch: {tuple(left.shape)} != {tuple(right.shape)}"
+    if left.dtype != right.dtype:
+        right = right.to(dtype=left.dtype)
+    if left.dtype.is_floating_point:
+        diff = (left - right).abs()
+        return f"max_abs_diff={diff.max().item():.8g}, mean_abs_diff={diff.mean().item():.8g}"
+    unequal = (left != right).sum().item()
+    return f"unequal_values={unequal}"
+
+
+def main():
+    repo_dir = Path(__file__).resolve().parent
+    checkpoint_path = repo_dir / "checkpoints" / "sft" / "final_model.pt"
+
+    checkpoint = torch.load(checkpoint_path, map_location="cpu")
+    checkpoint_state, _ = normalize_checkpoint(checkpoint)
+    checkpoint_state, ignored_keys = sanitize_model_state(checkpoint_state)
+
+    model = TaoNetForCausalLM.from_pretrained(str(repo_dir))
+    exported_state = model.model.state_dict()
+
+    checkpoint_keys = set(checkpoint_state)
+    exported_keys = set(exported_state)
+
+    missing = sorted(exported_keys - checkpoint_keys)
+    unexpected = sorted(checkpoint_keys - exported_keys)
+
+    print(f"checkpoint tensors: {len(checkpoint_keys)}")
+    print(f"exported tensors:   {len(exported_keys)}")
+    if ignored_keys:
+        print(f"ignored checkpoint buffers: {len(ignored_keys)}")
+
+    if missing:
+        print("\nKeys present in exported model but missing from checkpoint:")
+        for key in missing[:50]:
+            print(f"  - {key}")
+
+    if unexpected:
+        print("\nKeys present in checkpoint but missing from exported model:")
+        for key in unexpected[:50]:
+            print(f"  - {key}")
+
+    common_keys = sorted(checkpoint_keys & exported_keys)
+    mismatches = []
+    exact_matches = 0
+    for key in common_keys:
+        left = checkpoint_state[key].cpu()
+        right = exported_state[key].cpu()
+        if left.shape == right.shape and torch.equal(left, right.to(dtype=left.dtype) if right.dtype != left.dtype else right):
+            exact_matches += 1
+            continue
+        mismatches.append((key, summarize_tensor_diff(left, right)))
+
+    print(f"\nexact tensor matches: {exact_matches}/{len(common_keys)}")
+    print(f"tensor mismatches:    {len(mismatches)}")
+
+    if mismatches:
+        print("\nFirst mismatches:")
+        for key, summary in mismatches[:50]:
+            print(f"  - {key}: {summary}")
+        raise SystemExit(1)
+
+    if missing or unexpected:
+        raise SystemExit(1)
+
+    print("\nWeight verification passed.")
+
+
+if __name__ == "__main__":
+    main()