refactor: 统一序列化层，消除分散的 I/O 路径

- Checkpoint 改为 @dataclass，内聚 save/load 方法
- 提取 save_safetensors/load_safetensors/save_json/load_json 共享工具
- 新增 save_model/load_model_config/load_model_weights 模块函数
- automodel 和 lora 统一委托到 serialization 模块

astrai/model/automodel.py

@@ -2,16 +2,15 @@
 AutoModel base class for model loading and saving.
 """
 
-import json
 from contextlib import contextmanager
 from pathlib import Path
 from typing import Self, Union
 
-import safetensors.torch as st
 import torch.nn as nn
 
 from astrai.config.model_config import BaseModelConfig, ConfigFactory
 from astrai.factory import BaseFactory
+from astrai.serialization import load_model_config, load_model_weights, save_model
 
 
 @contextmanager
@@ -60,25 +59,22 @@ class AutoModel(BaseFactory["AutoModel"], nn.Module):
 
         model_path = Path(path)
 
-        # Load config
         config_path = model_path / "config.json"
-        if config_path.exists():
-            with open(config_path, "r") as f:
-                raw = json.load(f)
-            config = ConfigFactory.load(raw)
-            model_type = config.model_type or "autoregressive_lm"
-        else:
+        if not config_path.exists():
             raise FileNotFoundError(f"Config file not found: {config_path}")
 
+        raw = load_model_config(str(model_path))
+        config = ConfigFactory.load(raw)
+        model_type = config.model_type or "autoregressive_lm"
+
         actual_cls = AutoModel.get_component_class(model_type)
 
         with _disable_random_init(enable=disable_random_init):
             model = actual_cls(config)
 
-        # Load weights
         weights_path = model_path / "model.safetensors"
         if weights_path.exists():
-            state_dict = st.load_file(str(weights_path))
+            state_dict = load_model_weights(str(model_path))
             model.load_state_dict(state_dict, strict=strict)
 
         return model
@@ -87,14 +83,11 @@ class AutoModel(BaseFactory["AutoModel"], nn.Module):
         self,
         save_directory: Union[str, Path],
     ) -> None:
-        save_path = Path(save_directory)
-        save_path.mkdir(parents=True, exist_ok=True)
-
-        # Save config
-        self.config.to_file(str(save_path / "config.json"))
-
-        # Save weights
-        st.save_file(self.state_dict(), str(save_path / "model.safetensors"))
+        save_model(
+            config=self.config.to_dict(),
+            state_dict=self.state_dict(),
+            save_directory=str(save_directory),
+        )
 
     def to(self, *args, **kwargs) -> Self:
         """Move model to device/dtype."""

astrai/model/components/lora.py

@@ -1,15 +1,19 @@
-import json
 import logging
 from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Optional, Set
 
-import safetensors.torch as st
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
 from astrai.model.components.linear import Linear
+from astrai.serialization import (
+    load_json,
+    load_safetensors,
+    save_json,
+    save_safetensors,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -128,16 +132,14 @@ def save_lora(model: nn.Module, save_dir: str, config: LoRAConfig):
 
     path = Path(save_dir)
     path.mkdir(parents=True, exist_ok=True)
-    st.save_file(lora_sd, str(path / "adapter_model.safetensors"))
-    with open(path / "adapter_config.json", "w") as f:
-        json.dump(asdict(config), f, indent=2)
+    save_safetensors(lora_sd, path / "adapter_model.safetensors")
+    save_json(asdict(config), path / "adapter_config.json")
     logger.info("LoRA adapter saved to %s (%d keys)", save_dir, len(lora_sd))
 
 
 def load_lora(model: nn.Module, load_dir: str) -> LoRAConfig:
     path = Path(load_dir)
-    with open(path / "adapter_config.json") as f:
-        raw = json.load(f)
+    raw = load_json(path / "adapter_config.json")
     config = LoRAConfig(
         r=raw["r"], alpha=raw["alpha"], target_modules=tuple(raw["target_modules"])
     )
@@ -157,7 +159,7 @@ def load_lora(model: nn.Module, load_dir: str) -> LoRAConfig:
             target_modules=set(config.target_modules),
         )
 
-    weights = st.load_file(str(path / "adapter_model.safetensors"))
+    weights = load_safetensors(path / "adapter_model.safetensors")
     try:
         missing, unexpected = model.load_state_dict(weights, strict=False)
     except RuntimeError as e: