feat: 新增LoRA微调模块

- LoRALinear基于register_parameter托管base weight，state_dict路径不变 - inject_lora/merge_lora/save_lora/load_lora完备封装 - 24个单元测试覆盖注入、合并、存取、边界场景
2026-05-25 20:11:25 +08:00
50 changed files with 1205 additions and 1914 deletions
@@ -22,8 +22,7 @@ classDiagram
            +int n_layers
            +float norm_eps
            +int dim_ffn
-            +Optional[bool] tie_weight
-            +Optional[dict] rope_scaling
+            +bool tie_weight
            +int max_len
            +float rope_theta
            +str attn_type
@@ -53,7 +52,6 @@ classDiagram
            +int n_kv_heads
            +bool use_qk_norm
            +bool use_gated_attention
-            +Optional[dict] rope_scaling
            +Optional[str] pooling_type
            +Optional[bool] normalize_embeddings
        }
@@ -65,7 +63,7 @@ classDiagram
        }

        class TrainConfig {
-            +Callable[[], nn.Module] model_fn
+            +nn.Module model
            +str strategy
            +Dataset dataset
            +Callable optimizer_fn
@@ -82,7 +80,6 @@ classDiagram
            +str log_dir
            +int log_interval
            +List[str] metrics
-            +Optional[LoRAConfig] lora
            +int random_seed
            +int num_workers
            +Optional[int] prefetch_factor
@@ -107,8 +104,8 @@ classDiagram
        class BaseDataset {
            +int window_size
            +int stride
-            +Optional[Store] storage
-            +load(load_path, storage_type)
+            +Optional[BaseStorage] storage
+            +load(load_path, storage_type, tokenizer)
            +__getitem__(index)
            +__len__()
        }
@@ -129,25 +126,38 @@ classDiagram
            +__getitem__(index) Dict
        }

-        class Store {
-            +Dict[str, List[Tensor]] _data
-            +Dict[str, List[int]] _cum
-            +int _length
+        class BaseSegmentFetcher {
+            +List[Tensor] segments
+            +List[int] cum_lengths
+            +int total_length
+            +fetch_data(begin_idx, end_idx) Tensor
+        }
+
+        class BaseStorage {
+            +MultiSegmentFetcher _fetcher
            +keys (property)
-            +load(path)
+            +load(load_path, tokenizer)
            +fetch(begin, end, keys)
            +__len__()
-            -_fetch_key(key, begin, end) Tensor
-            -_normalize(raw)
        }

-        class H5Store {
-            +load(path)
+        class H5Storage {
+            +load(load_path, tokenizer)
+            +fetch(begin, end, keys) Dict
+            +keys() List
        }

-        class MmapStore {
-            +List _mmap_refs
-            +load(path)
+        class JSONStorage {
+            +load(load_path, tokenizer)
+            +fetch(begin, end, keys) Dict
+            +keys() List
+        }
+
+        class MultiSegmentFetcher {
+            +Dict multi_fetchers
+            +List multi_keys
+            +key_fetch(begin_idx, end_idx, keys) Dict
+            +fetch_data(begin_idx, end_idx) Dict
        }

        class ResumableDistributedSampler {
@@ -155,17 +165,17 @@ classDiagram
            +int iter
        }

-        class StoreFactory {
+        class StorageFactory {
            +Registry _registry
            +register(name) decorator
-            +create(storage_type) Store
+            +create(storage_type) BaseStorage
        }

        class DatasetFactory {
            +Registry _registry
            +register(name) decorator
            +create(train_type, window_size, stride) BaseDataset
-            +load(train_type, load_path, window_size, stride, storage_type) BaseDataset
+            +load(train_type, load_path, window_size, stride, storage_type, tokenizer) BaseDataset
        }
    }

@@ -176,9 +186,8 @@ classDiagram
            +int iteration
            +dict extra
            +dict meta
-            +dict config
            +save(save_dir)
-            +load(save_dir, broadcast) Checkpoint
+            +load(save_dir) Checkpoint
        }
    }

@@ -186,8 +195,8 @@ classDiagram
        class AutoModel {
            +BaseModelConfig config
            +Registry _registry
-            +register(name) decorator
-            +get_component_class(name) Type
+            +register(model_type) decorator
+            +get_component_class(model_type) Type
            +from_pretrained(path, disable_random_init, strict) nn.Module
            +save_pretrained(save_directory)
            +to(*args, **kwargs) Self
@@ -201,7 +210,7 @@ classDiagram
            +RMSNorm norm
            +Linear lm_head
            +forward(input_ids, input_mask, paged_cache, position_ids) Dict[str, Tensor]
-            +load_state_dict(state_dict, strict, assign)
+            +load_state_dict(state_dict)
            +state_dict()
        }

@@ -226,7 +235,6 @@ classDiagram
        }

        class GQA {
-            +int dim
            +int n_heads
            +int n_kv_heads
            +int head_dim
@@ -241,7 +249,6 @@ classDiagram
        }

        class MLA {
-            +int dim
            +int n_heads
            +int n_kv_heads
            +int head_dim
@@ -302,7 +309,6 @@ classDiagram
            +int dim
            +int max_len
            +float base
-            +Optional[Dict] rope_scaling
            +forward(x, position_ids=None) Tensor
        }

@@ -315,10 +321,10 @@ classDiagram
    namespace tokenize {
        class AutoTokenizer {
            +vocab_size int
-            +encode(tokens, out_ids, is_pretokenized, add_special_tokens) List[int]
+            +encode(tokens, out_ids, add_special_tokens) List[int]
            +decode(tokens, skip_special_tokens) str
            +__getattr__(name) Any (bos_id, eos_id, pad_id, stop_ids)
-            +apply_chat_template(messages, system_prompt, tokenize, add_generation_prompt) Union[str, List[int]]
+            +apply_chat_template(messages, tokenize) Union[str, List[int]]
            +set_chat_template(template)
            +load(path)
            +from_pretrained(path) AutoTokenizer
@@ -326,7 +332,7 @@ classDiagram
        }

        class ChatTemplate {
-            +str template_str
+            +String template_str
            +render(messages, system_prompt, **extra_variables) str
            +from_string(template) ChatTemplate
        }
@@ -364,7 +370,6 @@ classDiagram
            +SchedulerProtocol scheduler
            +Checkpoint checkpoint
            +TrainConfig config
-            +dict model_config
            +BaseExecutor executor
            +int epoch
            +int iteration
@@ -378,7 +383,7 @@ classDiagram

        class TrainContextBuilder {
            +TrainConfig config
-            +with_resume_dir(resume_dir) TrainContextBuilder
+            +with_checkpoint(checkpoint) TrainContextBuilder
            +build() TrainContext
        }

@@ -452,15 +457,16 @@ classDiagram
            +on_train_end(context)
            +on_epoch_begin(context)
            +on_epoch_end(context)
+            +on_step_begin(context)
+            +on_step_end(context)
            +on_batch_begin(context)
            +on_batch_end(context)
-            +on_optimizer_step(context)
            +on_error(context)
        }

        class GradientClippingCallback {
            +float max_grad_norm
-            +on_optimizer_step(context)
+            +on_step_begin(context)
        }

        class GradientCheckpointingCallback {
@@ -473,12 +479,16 @@ classDiagram
            +str save_dir
            +int interval
            +bool weight_only
+            +Callable state_dict_fn
            +Callable save_extra_fn
+            +Callable load_extra_fn
            +_save_checkpoint(context)
+            +on_train_begin(context)
            +on_batch_end(context)
            +on_train_end(context)
            +on_error(context)
            +save_extra(context)$
+            +load_extra(extra, context)$
        }

        class ProgressBarCallback {
@@ -502,7 +512,7 @@ classDiagram

        class ValidationCallback {
            +_run_validation(context)
-            +on_optimizer_step(context)
+            +on_step_end(context)
        }

        class CallbackFactory {
@@ -515,12 +525,7 @@ classDiagram
            +float lr
            +float momentum
            +float weight_decay
-            +bool nesterov
            +int ns_steps
-            +float adamw_lr
-            +tuple adamw_betas
-            +float adamw_eps
-            +float adamw_wd
            +step(closure) Optional[float]
        }
    }
@@ -541,8 +546,6 @@ classDiagram
            +AutoModel model
            +AutoTokenizer tokenizer
            +KVCache page_cache
-            +Optional[str] device
-            +Optional[torch.dtype] dtype
            +execute_prefill(tasks, prompt_len, start_pos)
            +execute_decode(tasks) List[int]
        }
@@ -554,9 +557,7 @@ classDiagram
            +bool _running
            +Thread _loop_thread
            +int max_seq_len
-            +str device
-            +torch.dtype dtype
-            +add_task(prompt, **kwargs) str
+            +add_task(prompt, max_tokens, temperature, top_p, top_k, stream_callback) str
            +remove_task(task_id)
            +start()
            +stop()
@@ -659,19 +660,15 @@ classDiagram

        class TaskManager {
            +AutoTokenizer tokenizer
-            +int max_batch_size
-            +int max_seq_len
-            +int max_prompt_len
            +Deque waiting_queue
            +List active_tasks
-            +add_task(prompt, max_tokens, temperature, top_p, top_k, stream_callback) str
+            +add_task(prompt, **kwargs) str
            +remove_task(task_id) List[Task]
            +remove_finished_tasks(stop_ids) List[Task]
            +pull_candidates(n) List[Task]
            +activate(task)
            +return_to_waiting(tasks)
            +get_active_tasks() List[Task]
-            +get_stats() Dict
        }

        class GenerationRequest {
@@ -750,58 +747,56 @@ classDiagram
            +str model
            +List[AnthropicMessage] messages
            +Optional[str] system
-            +Optional[float] temperature
-            +Optional[float] top_p
-            +Optional[int] top_k
+            +float temperature
+            +float top_p
+            +int top_k
            +int max_tokens
-            +Optional[bool] stream
+            +bool stream
            +Optional[List[str]] stop_sequences
        }

-        class ResponseBuilder {
-            <<abstract>>
-            +prepare(request, engine) Tuple[str, GenContext, List[str]]
-            +format_stream_start(ctx) List[str]
-            +format_chunk(token) str
-            +format_stream_end(ctx, stop) List[str]
-            +format_response(ctx, content, stop) Dict
-        }
-
-        class OpenAIResponseBuilder {
-            +prepare(request, engine) Tuple
-            +format_stream_start(ctx) List[str]
-            +format_chunk(token) str
-            +format_stream_end(ctx, stop) List[str]
-            +format_response(ctx, content, stop) Dict
-        }
-
-        class AnthropicResponseBuilder {
-            +prepare(request, engine) Tuple
-            +format_stream_start(ctx) List[str]
-            +format_chunk(token) str
-            +format_stream_end(ctx, stop) List[str]
-            +format_response(ctx, content, stop) Dict
-        }
-
        class ProtocolHandler {
+            <<abstract>>
            +request
            +engine
-            +builder: ResponseBuilder
+            +build_prompt() str
+            +create_response_id() str
+            +get_stop_sequences() List[str]
+            +create_stop_checker() StopChecker
+            +on_token(ctx, token, stop_checker) Optional[str]
+            +format_stream_start(ctx) List[str]
+            +format_stream_token(ctx, token) str
+            +format_stream_end(ctx) List[str]
+            +format_non_stream_response(ctx, content) Dict
            +handle() Union[StreamingResponse, Dict]
-            -_handle_stream(agen, ctx, stops) StreamingResponse
-            -_handle_non_stream(agen, ctx, stops) Dict
+        }
+
+        class OpenAIHandler {
+            +build_prompt() str
+            +create_response_id() str
+        }
+
+        class AnthropicHandler {
+            +build_prompt() str
+            +create_response_id() str
+            +on_token(ctx, token, stop_checker) Optional[str]
        }

        class StopChecker {
+            +has_sequences (property) bool
            +check(text) Optional[str]
+            +trim(text, matched) str
        }

-        class GenContext {
+        class StreamContext {
            +str resp_id
            +int created
            +str model
            +int prompt_tokens
            +int completion_tokens
+            +str accumulated
+            +Optional[str] stop_matched
+            +str last_yield_trimmed
        }

        class app {
@@ -881,11 +876,6 @@ classDiagram
            +unwrap_model(model) nn.Module
        }

-        class FSDPExecutor {
-            +_prepare_model(model) nn.Module
-            +unwrap_model(model) nn.Module
-        }
-
        class ExecutorFactory {
            +Registry _registry
            +register(name) decorator
@@ -921,13 +911,12 @@ classDiagram
    TrainCallback <|-- CheckpointCallback
    TrainCallback <|-- ProgressBarCallback
    TrainCallback <|-- MetricLoggerCallback
-    TrainCallback <|-- ValidationCallback
    BaseDataset <|-- SEQDataset
    BaseDataset <|-- SFTDataset
    BaseDataset <|-- DPODataset
    BaseDataset <|-- GRPODataset
-    Store <|-- H5Store
-    Store <|-- MmapStore
+    BaseStorage <|-- H5Storage
+    BaseStorage <|-- JSONStorage
    BaseSamplingStrategy <|-- TemperatureStrategy
    BaseSamplingStrategy <|-- TopKStrategy
    BaseSamplingStrategy <|-- TopPStrategy
@@ -947,19 +936,20 @@ classDiagram
    BaseFactory <|-- StrategyFactory
    BaseFactory <|-- SchedulerFactory
    BaseFactory <|-- CallbackFactory
-    BaseFactory <|-- StoreFactory
+    BaseFactory <|-- StorageFactory
    BaseFactory <|-- ExecutorFactory
    BaseFactory <|-- ConfigFactory
    BaseExecutor <|-- NoneExecutor
    BaseExecutor <|-- DDPExecutor
-    BaseExecutor <|-- FSDPExecutor
-    ResponseBuilder <|-- OpenAIResponseBuilder
-    ResponseBuilder <|-- AnthropicResponseBuilder
+    ProtocolHandler <|-- OpenAIHandler
+    ProtocolHandler <|-- AnthropicHandler

    %% --- Composition (strong ownership, part destroyed with whole) ---
    KVCache *-- PagePool
    KVCache *-- Storage
    KVCache *-- TaskTable
+    PagePool *-- Allocator
+    PagePool *-- PrefixCache
    InferenceEngine *-- InferenceScheduler
    InferenceScheduler *-- KVCache
    InferenceScheduler *-- Executor
@@ -973,6 +963,7 @@ classDiagram
    DecoderBlock *-- RMSNorm
    ChatCompletionRequest *-- ChatMessage
    MessagesRequest *-- AnthropicMessage
+    AutoTokenizer *-- ChatTemplate
    BaseFactory *-- Registry
    BaseExecutor *-- GradientState
    AccumOptimizer o-- GradientState
@@ -980,9 +971,6 @@ classDiagram

    %% --- Aggregation (weak ownership) ---
    AutoModel o-- BaseModelConfig
-    AutoTokenizer o-- ChatTemplate
-    PagePool o-- Allocator
-    PagePool o-- PrefixCache
    Trainer o-- TrainCallback
    TrainContext o-- BaseStrategy
    TrainContext o-- BaseScheduler
@@ -990,7 +978,7 @@ classDiagram
    TrainContext o-- BaseExecutor
    KvcacheView o-- Storage
    SamplingPipeline o-- BaseSamplingStrategy
-    BaseDataset o-- Store
+    BaseDataset o-- BaseStorage

    %% --- Dependency (uses temporarily) ---
    TrainConfig ..> BaseStrategy : selects
@@ -1004,13 +992,12 @@ classDiagram
    FFNFactory ..> DeepSeekMoE : creates
    DecoderBlock ..> AttnFactory : uses
    DecoderBlock ..> FFNFactory : uses
-    StoreFactory ..> H5Store : creates
-    StoreFactory ..> MmapStore : creates
+    StorageFactory ..> H5Storage : creates
+    StorageFactory ..> JSONStorage : creates
    ConfigFactory ..> AutoRegressiveLMConfig : creates
    ConfigFactory ..> EncoderConfig : creates
    ExecutorFactory ..> NoneExecutor : creates
    ExecutorFactory ..> DDPExecutor : creates
-    ExecutorFactory ..> FSDPExecutor : creates
    TrainContextBuilder ..> ExecutorFactory : creates
    Trainer ..> TrainContextBuilder : uses
    TrainContextBuilder ..> TrainContext : creates
@@ -1022,10 +1009,10 @@ classDiagram
    KVCache ..> KvcacheView : binds
    InferenceEngine ..> GenerationRequest : uses
    InferenceEngine ..> GenerateResult : creates
-    OpenAIResponseBuilder ..> ChatCompletionRequest : receives
-    AnthropicResponseBuilder ..> MessagesRequest : receives
+    OpenAIHandler ..> ChatCompletionRequest : receives
+    AnthropicHandler ..> MessagesRequest : receives
    ProtocolHandler ..> StopChecker : creates
-    ProtocolHandler ..> GenContext : creates
+    ProtocolHandler ..> StreamContext : creates

    %% --- Association (general usage) ---
    Trainer --> TrainConfig
@@ -1038,6 +1025,8 @@ classDiagram
    Executor --> AutoModel
    Executor --> AutoTokenizer
    TaskManager --> AutoTokenizer
+    MultiSegmentFetcher --> BaseSegmentFetcher
+    ResumableDistributedSampler --> BaseDataset

 ```

@@ -1047,13 +1036,13 @@ classDiagram
 | Module | Components | Description |
 |--------|------------|-------------|
 | **astrai.config** | BaseConfig, BaseModelConfig, AutoRegressiveLMConfig, EncoderConfig, ConfigFactory, TrainConfig | Configuration management (to_dict/from_dict, to_file/from_file) |
-| **astrai.dataset** | BaseDataset–GRPODataset, Store–MmapStore, StoreFactory, ResumableDistributedSampler, DatasetFactory | Dataset loading and management |
+| **astrai.dataset** | BaseDataset–GRPODataset, BaseStorage–JSONStorage, StorageFactory, BaseSegmentFetcher, MultiSegmentFetcher, ResumableDistributedSampler, DatasetFactory | Dataset loading and management |
 | **astrai.serialization** | Checkpoint | Model serialization |
 | **astrai.model** | AutoModel, AutoRegressiveLM, EmbeddingEncoder, DecoderBlock, GQA, MLA, MLP, DeepSeekMoE, AttnFactory, FFNFactory, RMSNorm, Linear, RotaryEmbedding, Embedding | Neural network model |
 | **astrai.tokenize** | AutoTokenizer, ChatTemplate | Tokenizer and chat template |
 | **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy–GRPOStrategy, StrategyFactory, BaseScheduler–SGDRScheduler, SchedulerFactory, TrainCallback(Protocol)–ValidationCallback, CallbackFactory, Muon | Training workflow |
-| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, KVCache–KvcacheView, Allocator–Storage, Task, TaskManager, TaskStatus, GenerationRequest, GenerateResult, BaseSamplingStrategy–SamplingPipeline, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, ChatMessage–MessagesRequest, app | Inference service |
-| **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank/get_world_size/get_current_device, only_on_rank, BaseExecutor, ExecutorFactory, NoneExecutor, DDPExecutor, FSDPExecutor, GradientState, AccumOptimizer, AccumScheduler, ParallelModel, RowParallelLinear, ColumnParallelLinear | Distributed parallel & gradient accumulation |
+| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, KVCache–KvcacheView, Allocator–Storage, Task, TaskManager, TaskStatus, GenerationRequest, GenerateResult, BaseSamplingStrategy–SamplingPipeline, ProtocolHandler–AnthropicHandler, StopChecker, StreamContext, ChatMessage–MessagesRequest, app | Inference service |
+| **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank/get_world_size/get_current_device, only_on_rank, BaseExecutor, ExecutorFactory, NoneExecutor, DDPExecutor, GradientState, AccumOptimizer, AccumScheduler, ParallelModel, RowParallelLinear, ColumnParallelLinear | Distributed parallel & gradient accumulation |
 | **astrai.factory** | Registry, BaseFactory[T] | Component registration |
 | **astrai.protocols** | OptimizerProtocol, SchedulerProtocol | Structural subtyping for optimizer/scheduler wrappers |

@@ -1061,17 +1050,17 @@ classDiagram

 | Pattern | Classes | Purpose |
 |---------|---------|---------|
-| **Factory** | `AttnFactory`, `FFNFactory`, `StrategyFactory`, `DatasetFactory`, `SchedulerFactory`, `CallbackFactory`, `StoreFactory`, `ConfigFactory`, `ExecutorFactory` | Decorator-based component creation |
+| **Factory** | `AttnFactory`, `FFNFactory`, `StrategyFactory`, `DatasetFactory`, `SchedulerFactory`, `CallbackFactory`, `StorageFactory`, `ConfigFactory`, `ExecutorFactory` | Decorator-based component creation |
 | **Registry** | `BaseFactory`, `Registry` | Component registration with category/priority |
 | **Strategy** | `SEQStrategy`, `SFTStrategy`, `DPOStrategy`, `GRPOStrategy` | Training strategy switching |
 | **Strategy (Sampling)** | `TemperatureStrategy`, `TopKStrategy`, `TopPStrategy`, `SamplingPipeline` | Composable logit transformations |
-| **Strategy (API)** | `ResponseBuilder`, `OpenAIResponseBuilder`, `AnthropicResponseBuilder` | HTTP API handler with format hooks |
+| **Template Method** | `ProtocolHandler`, `OpenAIHandler`, `AnthropicHandler` | HTTP API handler with format hooks |
 | **Builder** | `TrainContextBuilder` | Chain-building training context |
 | **Observer** | `TrainCallback`, callback implementations | Training process monitoring |
 | **Context** | `TrainContext` | Unified training state bag |
 | **Object Pool** | `Allocator`, `PagePool` | Page-based KV cache with LRU eviction |
 | **Executor** | `BaseExecutor`, `NoneExecutor`, `DDPExecutor` | Gradient accumulation & model distribution |
-| **Storage** | `Store`, `H5Store`, `MmapStore` | Format-agnostic data access with multi-segment support |
+| **Storage** | `BaseStorage`, `H5Storage`, `JSONStorage` | Format-agnostic data access |
 | **Producer-Consumer** | `InferenceScheduler`, `Task`, queues | Continuous batching |
 | **AutoModel Registry** | `AutoModel`, `AutoRegressiveLM`, `EmbeddingEncoder` | Model-type dynamic loading |

@@ -1080,13 +1069,13 @@ classDiagram
 1. **Config → Training**: `TrainConfig` holds model, dataset, optimizer_fn, scheduler_fn, `parallel_mode`, `executor_kwargs`
 2. **Training Flow**: `Trainer` → `TrainContextBuilder` → `TrainContext`, uses `BaseStrategy` for loss, `BaseExecutor` for gradient accumulation + model distribution
 3. **Strategy Selection**: `StrategyFactory` creates strategy by `train_type`
-4. **Executor Selection**: `ExecutorFactory.create(cfg.parallel_mode, grad_accum_steps=cfg.grad_accum_steps, **cfg.executor_kwargs)` → `NoneExecutor` / `DDPExecutor` / `FSDPExecutor`
+4. **Executor Selection**: `ExecutorFactory.create(parallel_mode, **executor_kwargs)` → `NoneExecutor` (single) / `DDPExecutor` (distributed)
 5. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `KVCache` + `SamplingPipeline`
 6. **Distributed**: `spawn_parallel_fn` + `setup_parallel` for multi-process DDP
-7. **Dataset Loading**: `DatasetFactory` creates datasets, `Store` (H5Store/MmapStore) loads data with explicit `_length` and multi-segment `_data`
+7. **Dataset Loading**: `DatasetFactory` creates datasets, `BaseStorage` (H5Storage/JSONStorage) loads via `BaseSegmentFetcher` + `MultiSegmentFetcher`
 8. **Checkpoint**: `Checkpoint` saves/loads safetensors + metadata (rank-0 only), extra state saved as `{key}.pt`
 9. **Scheduler**: `SchedulerFactory` creates `CosineScheduler`/`SGDRScheduler`
 10. **AutoModel**: `from_pretrained()` loads `config.json` + `model.safetensors`, `_disable_random_init` replaces `nn.init.*` with no-ops
 11. **Protocols**: `OptimizerProtocol` / `SchedulerProtocol` — structural subtyping for `AccumOptimizer` / `AccumScheduler` wrappers

-> Document Update Time: 2026-05-28
+> Document Update Time: 2026-05-24
@@ -5,21 +5,21 @@ This document describes the data pipeline: from raw text to model input tensors.
 ## Overview

 ```
-Raw Text → AutoTokenizer → Token IDs → .h5/.bin → Dataset → Sampler → DataLoader → Training/Inference
+Raw Text → AutoTokenizer → Token IDs → .h5/.json → Dataset → Sampler → DataLoader → Training/Inference
 ```

 ## Data Preparation

-Raw text is tokenized via `AutoTokenizer.encode()` and saved as HDF5 (`.h5`) or binary (`.bin` + `meta.json`) files with keyed tensor groups.
+Raw text is tokenized via `AutoTokenizer.encode()` and saved as HDF5 (`.h5`) or JSON (`.json`/`.jsonl`) files with keyed tensor groups.

 Storage format is auto-detected by `detect_format()`; backends are dispatched via registry:

 ```
-StoreFactory.create("h5")  → H5Store
-StoreFactory.create("bin") → MmapStore
+StorageFactory.create("h5")   → H5Storage
+StorageFactory.create("json") → JSONStorage
 ```

-H5 backend supports shared memory via `.share_memory_()`. Bin (mmap) uses OS page-cache sharing natively.
+Both support shared memory via `.share_memory_()`.

 ## Data Keys by Training Type

@@ -33,14 +33,14 @@ H5 backend supports shared memory via `.share_memory_()`. Bin (mmap) uses OS pag
 ## Dataset Architecture

 ```
-DatasetFactory.load(train_type, load_path, window_size, stride, storage_type)
-  → StoreFactory.create(detect_format(path))
-    → Store._data[Dict[str, List[Tensor]]] + _cum[Dict[str, List[int]]]
+DatasetFactory.load(train_type, path, window_size, stride)
+  → StorageFactory.create(detect_format(path))
+    → MultiSegmentFetcher(BaseSegmentFetcher per key)
      → BaseDataset.__getitem__(idx)
        → sliding window [begin, end) via get_index(idx)
 ```

-`window_size` = max input length, `stride` = step between consecutive samples (defaults to `window_size`).
+`window_size` = max input length, `stride` = step between consecutive samples.

 ## Sampler

@@ -54,4 +54,4 @@ DatasetFactory.load(train_type, load_path, window_size, stride, storage_type)

 Standard PyTorch `DataLoader` with configurable `batch_size`, `num_workers`, `pin_memory`, `prefetch_factor`. Sampler produces indices; dataloader fetches tensor batches via `__getitem__`.

-> Document Update Time: 2026-05-28
+> Document Update Time: 2026-05-17
@@ -16,12 +16,12 @@ Six classes working together:

 ```
 KVCache (facade)
-  ├── PagePool         orchestrates page allocation + prefix matching
-  │     ├── Allocator   bitmask-based page allocator + ref-count + LRU eviction (inside PagePool)
-  │     └── PrefixCache hash-based prefix matching (page_hash via polynomial hash) (inside PagePool)
+  ├── Allocator        bitmask-based page allocator + ref-count + LRU eviction
+  ├── PrefixCache      hash-based prefix matching (page_hash via rolling hash)
+  ├── PagePool         orchestrates Allocator + PrefixCache
  ├── TaskTable        maps task_id → page_table + cached token count
  ├── Storage          k_cache / v_cache tensors (n_layers × n_pages × page_size × n_kv_heads × head_dim)
-  └── KvcacheView      bundles Storage + page_table + total_len for attention layers (returned by bind())
+  └── KvcacheView      bundles Storage + page_table + total_len for attention layers
 ```

 `KVCache.bind(page_table, total_len)` returns a `KvcacheView` used by attention layers via `write()` / `gather()`.
@@ -40,32 +40,26 @@ KVCache (facade)
 ## Sampling (Strategy Pattern)

 ```
-BaseSamplingStrategy (ABC)
-  ├── TemperatureStrategy
-  ├── TopKStrategy
-  └── TopPStrategy
+BaseSamplingStrategy → TemperatureStrategy → TopKStrategy → TopPStrategy
 ```

 `SamplingPipeline` composes them: Temperature → Top-K → Top-P → softmax → multinomial.  
 `sample()` is a convenience shortcut for one-shot usage.

-## Protocol Handlers (Strategy Pattern)
+## Protocol Handlers (Template Method)

 ```python
-class ProtocolHandler:  # concrete orchestrator
-    def __init__(self, request, engine, builder): ...
-    async def handle(self):
-        prompt, ctx, stops = builder.prepare(request, engine)
+class ProtocolHandler(ABC):
+    def handle(self):
+        ctx = StreamContext(...)
        agen = engine.generate_async(prompt, ...)
-        if stream: self._handle_stream(agen, ctx, stops)
-        else:      return await self._handle_non_stream(agen, ctx, stops)
+        if stream: self._handle_stream(agen, ctx)
+        else:      self._handle_non_stream(agen, ctx)
 ```

-`ResponseBuilder` (ABC): `prepare()`, `format_stream_start()`, `format_chunk()`, `format_stream_end()`, `format_response()`.
+Subclass hooks: `build_prompt()`, `create_response_id()`, `format_stream_start/token/end()`, `format_non_stream_response()`.

-`OpenAIResponseBuilder` → `/v1/chat/completions`, `AnthropicResponseBuilder` → `/v1/messages`.
-
-Adding a protocol = one builder file, no handler subclassing needed.
+`OpenAIHandler` → `/v1/chat/completions`, `AnthropicHandler` → `/v1/messages`.

 ## Engine & GenerateResult

@@ -100,14 +94,12 @@ Response:
 {
  "id": "chatcmpl-abc123",
  "object": "chat.completion",
-  "created": 1717000000,
-  "model": "astrai",
-  "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello!"}, "finish_reason": "stop"}],
+  "choices": [{"message": {"role": "assistant", "content": "Hello!"}, "finish_reason": "stop"}],
  "usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15}
 }
 ```

-Streaming SSE: `object: "chat.completion.chunk"` — starts with role delta, then token chunks, ends with finish chunk + usage stats, then `data: [DONE]`.
+Streaming SSE: `data: {"choices":[{"delta":{"role":"assistant"}}]}` → token chunks → `data: [DONE]`

 ### Anthropic

@@ -124,10 +116,10 @@ Supports `stop_sequences` and streaming via `event: content_block_delta`.
 | Param | Type | Default | Description |
 |-------|------|---------|-------------|
 | `messages` | List[dict] | required | Chat messages (role, content) |
-| `temperature` | float | 1.0 | Sampling temperature (>= 0.0) |
+| `temperature` | float | 1.0 | Sampling temperature (0.0–2.0) |
 | `top_p` | float | 1.0 | Nucleus threshold |
 | `top_k` | int | 50 | Top-k count |
-| `max_tokens` | Optional[int] | None | Max generation length |
+| `max_tokens` | int | None | Max generation length |
 | `stream` | bool | False | Stream output |

 ## Engine API
@@ -145,4 +137,4 @@ engine.generate(["A", "B"], stream=True)        # -> Generator[Tuple[int, str]]
 await engine.generate_async("Hello", ...)       # -> AsyncGenerator[str]
 ```

-> Document Update Time: 2026-05-28
+> Document Update Time: 2026-05-17
@@ -53,7 +53,7 @@
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--nprocs` | Number of GPUs / processes | 1 |
-| `--parallel_mode` | Parallel strategy (`none`, `ddp`, or `fsdp`) | none |
+| `--parallel_mode` | Parallel strategy (`none` or `ddp`) | none |
 | `--device_type` | Device type | cuda |
 | `--start_method` | Multiprocessing start method (`spawn`, `fork`, `forkserver`) | spawn |

@@ -74,17 +74,16 @@ on_train_begin
      on_batch_begin
      with executor.accumulate(model):
        loss = strategy(batch)
-        stand_loss = loss / executor.grad_accum_steps
-        executor.backward(stand_loss)
+        (loss / grad_accum_steps).backward()
        iteration += 1
-        on_batch_end
+      on_batch_end

-        if executor.sync_gradients:
-          on_optimizer_step
-          optimizer.step()
-          optimizer.zero_grad()
-          if scheduler:
-            scheduler.step()
+      if executor.sync_gradients:
+        on_optimizer_step
+        optimizer.step()
+        optimizer.zero_grad()
+
+      scheduler.step()  # called every iteration
    on_epoch_end
 on_train_end
 ```
@@ -171,27 +170,27 @@ Callback wraps each `DecoderBlock.forward` with `torch.utils.checkpoint.checkpoi
 ## Checkpoint

 ```
-Checkpoint(state_dict, epoch, iteration, extra, meta, config)
-  ├── save(save_dir)    rank-0 only: meta.json (epoch/iteration/timestamp) + config.json (model config) + state_dict.safetensors + optional {key}.pt (optimizer.pt, scheduler.pt)
+Checkpoint(state_dict, epoch, iteration, extra, meta)
+  ├── save(save_dir)    rank-0 only: meta.json (includes training config) + state_dict.safetensors + optional optimizer.pt / scheduler.pt
  └── load(save_dir)    broadcasts metadata from rank-0
 ```

 Optimizer/scheduler state persisted by default via `Checkpoint.extra`.  
-Model config (`context.model_config`) saved into `config.json` during training via `CheckpointCallback`.
+Training config (`TrainConfig.to_dict()`) saved into `meta.json` during training via `CheckpointCallback`.

 ## TrainContextBuilder (Builder Pattern)

 ```python
 context = (
    TrainContextBuilder(config)
-        .with_resume_dir(resume_dir)
+        .with_checkpoint(checkpoint)
        .build()
 )
 # Returns TrainContext with model, strategy, optimizer, scheduler, dataloader, checkpoint
 ```

 - Loads checkpoint weights if provided
- Creates executor via `ExecutorFactory.create(cfg.parallel_mode, grad_accum_steps=cfg.grad_accum_steps, **cfg.executor_kwargs)`
+- Creates executor via `ExecutorFactory.create(parallel_mode, **executor_kwargs)`
 - Calls `executor.prepare(model, optimizer, dataloader, scheduler)` for model distribution (e.g. DDP) + gradient accumulation wrappers
 - Creates `ResumableDistributedSampler` for shuffle+resume
 - Builds strategy via `StrategyFactory.create(train_type, ...)`
@@ -224,4 +223,4 @@ nohup python scripts/tools/train.py \

 Full parameter reference at [params.md](params.md).

-> Document Update Time: 2026-05-28
+> Document Update Time: 2026-05-24
@@ -1,4 +1,4 @@
-__version__ = "1.3.7"
+__version__ = "1.3.6"
 __author__ = "ViperEkura"

 from astrai.config import (
@@ -13,21 +13,12 @@ class BaseConfig:
                d[fld.name] = v
            elif v is None:
                d[fld.name] = None
-            elif isinstance(v, (dict, list, tuple)):
+            elif isinstance(v, (dict, list)):
                try:
-                    val = list(v) if isinstance(v, tuple) else v
-                    json.dumps(val)
-                    d[fld.name] = val
+                    json.dumps(v)
+                    d[fld.name] = v
                except (TypeError, ValueError):
                    pass
-            elif isinstance(v, BaseConfig):
-                d[fld.name] = v.to_dict()
-            elif hasattr(v, "__dataclass_fields__"):
-                sub = {}
-                for f in fields(v):
-                    a = getattr(v, f.name)
-                    sub[f.name] = list(a) if isinstance(a, tuple) else a
-                d[fld.name] = sub
        return d

    @classmethod
@@ -49,7 +49,6 @@ class AutoRegressiveLMConfig(BaseModelConfig):

    max_len: Optional[int] = None
    rope_theta: Optional[float] = None
-    rope_scaling: Optional[dict] = None

    attn_type: str = "gqa"
    n_heads: Optional[int] = None
@@ -81,7 +80,6 @@ class EncoderConfig(BaseModelConfig):

    max_len: Optional[int] = None
    rope_theta: Optional[float] = None
-    rope_scaling: Optional[dict] = None

    n_heads: Optional[int] = None
    n_kv_heads: Optional[int] = None
@@ -7,7 +7,6 @@ from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import Dataset

 from astrai.config.base import BaseConfig
-from astrai.model.components.lora import LoRAConfig


 def required(**kw):
@@ -17,8 +16,8 @@ def required(**kw):
@dataclass
 class TrainConfig(BaseConfig):
    # basic setting
-    model_fn: Callable[[], nn.Module] = field(
-        default=None, metadata=required(help="Model factory for training.")
+    model: nn.Module = field(
+        default=None, metadata=required(help="Model for training.")
    )
    strategy: str = field(default=None, metadata=required(help="Training strategy."))
    dataset: Dataset = field(
@@ -57,12 +56,6 @@ class TrainConfig(BaseConfig):
        default=5000, metadata={"help": "Number of iterations between checkpoints."}
    )

-    # lora setting
-    lora: Optional[LoRAConfig] = field(
-        default=None,
-        metadata={"help": "LoRA config. None means full fine-tuning."},
-    )
-
    # metric setting
    log_dir: str = field(
        default="./checkpoint/logs", metadata={"help": "Directory for metric logs."}
@@ -4,28 +4,32 @@ from astrai.dataset.dataset import (
 )
 from astrai.dataset.sampler import ResumableDistributedSampler
 from astrai.dataset.storage import (
-    H5Store,
-    MmapStore,
-    Store,
-    StoreFactory,
+    BaseSegmentFetcher,
+    BaseStorage,
+    H5Storage,
+    JSONStorage,
+    MultiSegmentFetcher,
+    StorageFactory,
    detect_format,
-    load_bin,
    load_h5,
-    save_bin,
+    load_json,
    save_h5,
+    save_json,
 )

 __all__ = [
    "BaseDataset",
    "DatasetFactory",
-    "Store",
-    "StoreFactory",
-    "H5Store",
-    "MmapStore",
+    "BaseSegmentFetcher",
+    "MultiSegmentFetcher",
+    "BaseStorage",
+    "H5Storage",
+    "JSONStorage",
+    "StorageFactory",
    "detect_format",
    "save_h5",
    "load_h5",
-    "save_bin",
-    "load_bin",
+    "save_json",
+    "load_json",
    "ResumableDistributedSampler",
 ]
@@ -8,8 +8,8 @@ from torch import Tensor
 from torch.utils.data import Dataset

 from astrai.dataset.storage import (
-    Store,
-    StoreFactory,
+    BaseStorage,
+    StorageFactory,
    detect_format,
 )
 from astrai.factory import BaseFactory
@@ -26,7 +26,7 @@ class BaseDataset(Dataset, ABC):
        super().__init__()
        self.window_size = window_size
        self.stride = stride
-        self.storage: Optional[Store] = None
+        self.storage: Optional[BaseStorage] = None

    @property
    def required_keys(self) -> List[str]:
@@ -48,26 +48,37 @@ class BaseDataset(Dataset, ABC):
                f"Missing: {missing}"
            )

-    def load(self, load_path: str, storage_type: Optional[str] = None):
+    def load(self, load_path: str, storage_type: Optional[str] = None, tokenizer=None):
        """Load dataset from the given path.

        Auto-detects the storage format if not specified.

        Args:
            load_path: Path to the data directory or file
-            storage_type: Force a specific storage type ("h5", "bin"),
+            storage_type: Force a specific storage type ("h5", "json"),
                          or None for auto-detection
+            tokenizer: Callable str -> List[int], used to tokenize raw text
+                       in JSON files. Ignored for HDF5.

        Raises:
            KeyError: If the loaded storage is missing required keys.
        """
        if storage_type is None:
            storage_type = detect_format(load_path)
-        self.storage = StoreFactory.create(storage_type)
+        self.storage = StorageFactory.create(storage_type)
        self._load_path = load_path
-        self.storage.load(load_path)
+        self.storage.load(load_path, tokenizer=tokenizer)
        self._validate_keys()

+    def load_json(self, load_path: str, tokenizer=None):
+        """Load dataset from JSON files explicitly.
+
+        Args:
+            load_path: Path to the JSON data file or directory
+            tokenizer: Optional tokenizer callable for raw text JSON.
+        """
+        self.load(load_path, storage_type="json", tokenizer=tokenizer)
+
    @property
    def count(self) -> int:
        """Return the total number of raw elements (tokens) in the dataset."""
@@ -137,7 +148,7 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
    """

    @classmethod
-    def _validate_component(cls, dataset_cls: type):
+    def _validate_component(cls, dataset_cls: type) -> None:
        """Validate that the dataset class inherits from BaseDataset."""
        if not issubclass(dataset_cls, BaseDataset):
            raise TypeError(f"{dataset_cls.__name__} must inherit from BaseDataset")
@@ -164,6 +175,7 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
        window_size: int,
        stride: Optional[int] = None,
        storage_type: Optional[str] = None,
+        tokenizer=None,
    ) -> "BaseDataset":
        """Create and load a dataset in one step.

@@ -172,7 +184,8 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
            load_path: Path to the data file
            window_size: Window size for data sampling
            stride: Stride between consecutive samples (default: same as window_size)
-            storage_type: Storage type ("h5", "bin") or None for auto-detection
+            storage_type: Storage type ("h5", "json") or None for auto-detection
+            tokenizer: Callable str -> List[int] for raw text JSON tokenization

        Returns:
            Loaded dataset instance
@@ -181,7 +194,7 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
            stride = window_size

        dataset = cls.create(train_type, window_size, stride)
-        dataset.load(load_path, storage_type=storage_type)
+        dataset.load(load_path, storage_type=storage_type, tokenizer=tokenizer)

        return dataset

@@ -293,11 +306,9 @@ class GRPODataset(BaseDataset):
    def __getitem__(self, index: int) -> Dict[str, Tensor]:
        begin_idx, end_idx = self.get_index(index)

-        prompts = self._fetch_data(begin_idx, end_idx, "prompts").to(dtype=torch.long)
-        responses = self._fetch_data(begin_idx, end_idx, "responses").to(
-            dtype=torch.long
-        )
-        masks = self._fetch_data(begin_idx, end_idx, "masks").to(dtype=torch.bool)
+        prompts = self._fetch_data(begin_idx, end_idx, "prompts")
+        responses = self._fetch_data(begin_idx, end_idx, "responses")
+        masks = self._fetch_data(begin_idx, end_idx, "masks")
        rewards = self._fetch_data(begin_idx, end_idx, "rewards")

        return {
@@ -43,7 +43,6 @@ class ResumableDistributedSampler(Sampler[int]):
        offset = 0 if drop_last else self.num_replicas - 1
        self.num_samples_per_replica = (self.num_samples + offset) // self.num_replicas
        self.total_size = self.num_samples_per_replica * self.num_replicas
-        self.iter = self.iter % self.num_samples_per_replica

        self._indices = None

@@ -75,10 +74,5 @@ class ResumableDistributedSampler(Sampler[int]):
        self.epoch += 1
        self._indices = None

-    @property
-    def _remaining(self):
-        remaining = self.num_samples_per_replica - self.iter
-        return max(remaining, 0)
-
    def __len__(self):
-        return self._remaining
+        return self.num_samples_per_replica
@@ -1,20 +1,7 @@
 """Storage backends for different data formats.

-Layers:
-  - I/O layer:       save_* / load_* functions, read/write raw files (HDF5/bin)
-                      return Dict[str, List[Tensor]] — format-specific, no state
-  - Store (ABC):     central abstraction, normalizes multi-segment into
-                      Dict[str, List[Tensor]] per key via _normalize(),
-                      fetch() uses bisect across segments — no forced concat
-  - Dataset layer:   BaseDataset owns a Store, only calls store.fetch(begin, end, key)
-
-Key properties:
-  - Multi-segment:   segments kept as-is, no forced concatenation — safe for
-                      datasets larger than RAM
-  - Explicit length: _length = min(total elements across keys), set at load,
-                      __len__ returns O(1)
-  - Zero-copy mmap:  MmapStore wraps np.memmap(mode="r"), all DataLoader
-                      workers share OS page-cache pages
+Each storage handles format-specific loading (HDF5, JSON, etc.) and provides
+a uniform interface for data access and length observation via fetchers.
 """

 import bisect
@@ -22,10 +9,9 @@ import json
 import os
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Dict, List, Union
+from typing import Callable, Dict, List, Optional, Union

 import h5py
-import numpy as np
 import torch
 from torch import Tensor

@@ -68,30 +54,54 @@ def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
    return tensor_group


-def save_bin(file_path: str, tensor_group: Dict[str, List[Tensor]]):
+def save_json(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
    os.makedirs(file_path, exist_ok=True)
-    meta = {}
+    full_file_path = os.path.join(file_path, f"{file_name}.json")
+    json_data = {}
    for key, tensors in tensor_group.items():
-        cat = torch.cat(tensors, dim=0)
-        meta[key] = {"shape": list(cat.shape), "dtype": str(cat.dtype).split(".")[-1]}
-        np.asarray(cat.cpu().numpy()).tofile(os.path.join(file_path, f"{key}.bin"))
-    with open(os.path.join(file_path, "meta.json"), "w") as f:
-        json.dump(meta, f)
+        json_data[key] = [tensor.tolist() for tensor in tensors]
+    with open(full_file_path, "w", encoding="utf-8") as f:
+        json.dump(json_data, f, ensure_ascii=False)


-def load_bin(file_path: str) -> Dict[str, List[Tensor]]:
-    with open(os.path.join(file_path, "meta.json"), "r") as f:
-        meta = json.load(f)
-    segments: Dict[str, List[Tensor]] = {}
-    for key, info in meta.items():
-        arr = np.memmap(
-            os.path.join(file_path, f"{key}.bin"),
-            dtype=info["dtype"],
-            mode="r+",
-            shape=tuple(info["shape"]),
-        )
-        segments[key] = [torch.from_numpy(arr)]
-    return segments
+def load_json(
+    file_path: str,
+    share_memory: bool = True,
+    tokenizer: Optional[Callable[[str], List[int]]] = None,
+) -> Dict[str, List[Tensor]]:
+    """Load tensor data from JSON files.
+
+    Supports two modes:
+    - Pre-tokenized: JSON values are List[List[int]] (token IDs), loaded as-is.
+    - Raw text: JSON values are List[str], tokenized via ``tokenizer`` callable
+      at load time. A ``tokenizer`` receives a str and returns List[int].
+
+    Non-data JSON files (e.g. config.json) with scalar/object values are
+    silently skipped.
+    """
+    tensor_group: Dict[str, List[Tensor]] = {}
+    root_path = Path(file_path)
+    json_files = list(root_path.rglob("*.json")) + list(root_path.rglob("*.jsonl"))
+    for json_file in json_files:
+        with open(json_file, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        if not isinstance(data, dict):
+            continue
+        for key, sequences in data.items():
+            if not isinstance(sequences, list):
+                continue
+            tensors = []
+            for seq in sequences:
+                if tokenizer is not None and isinstance(seq, str):
+                    seq = tokenizer(seq)
+                tensor = torch.tensor(seq, dtype=torch.long)
+                if share_memory:
+                    tensor = tensor.share_memory_()
+                tensors.append(tensor)
+            if tensor_group.get(key) is None:
+                tensor_group[key] = []
+            tensor_group[key].extend(tensors)
+    return tensor_group


 def detect_format(load_path: str) -> str:
@@ -101,7 +111,7 @@ def detect_format(load_path: str) -> str:
        load_path: Directory or file path

    Returns:
-        Format string ("h5" or "bin")
+        Format string ("h5" or "json")

    Raises:
        FileNotFoundError: If no supported data files are found
@@ -111,140 +121,181 @@ def detect_format(load_path: str) -> str:
        suffix = root.suffix.lower()
        if suffix in (".h5", ".hdf5"):
            return "h5"
+        if suffix in (".json", ".jsonl"):
+            return "json"
        raise ValueError(f"Unsupported file format: {suffix}")

    h5_files = list(root.rglob("*.h5")) + list(root.rglob("*.hdf5"))
    if h5_files:
        return "h5"
-    bin_files = list(root.rglob("*.bin"))
-    if bin_files and (root / "meta.json").exists():
-        return "bin"
+    json_files = list(root.rglob("*.json")) + list(root.rglob("*.jsonl"))
+    if json_files:
+        return "json"
    raise FileNotFoundError(f"No supported data files found at {load_path}")


-class Store(ABC):
-    """String keys -> segmented tensors with ``fetch(begin, end, keys)``.
+class BaseSegmentFetcher:
+    """Fetches data segments across multiple tensor segments.

-    Each key maps to one or more tensor segments (no forced concatenation).
-    ``len(store)`` returns ``self._length`` (explicit, O(1)), the minimum
-    total element count across all keys.
+    Maintains cumulative lengths for efficient range queries across
+    multiple discontinuous segments.
+    """

-    Subclasses fill ``self._data`` and ``self._cum`` during ``load()``
-    via ``_normalize()``.
+    def __init__(self, segments: List[Tensor]):
+        self.segments = segments
+        self.cum_lengths = []
+
+        total = 0
+        for seg in segments:
+            total += torch.numel(seg)
+            self.cum_lengths.append(total)
+
+        self.total_length = total
+
+    def __len__(self) -> int:
+        return self.total_length
+
+    def fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
+        """Fetch data in the range [begin_idx, end_idx)."""
+        if not (
+            0 <= begin_idx < self.total_length and 0 <= end_idx <= self.total_length
+        ):
+            raise ValueError("begin_idx or end_idx out of bounds")
+        if begin_idx >= end_idx:
+            return torch.tensor([], dtype=torch.long)
+
+        seg_start_idx = bisect.bisect_right(self.cum_lengths, begin_idx)
+        seg_end_idx = bisect.bisect_left(self.cum_lengths, end_idx)
+
+        result_segments = []
+
+        for i in range(seg_start_idx, seg_end_idx + 1):
+            prev_cum = self.cum_lengths[i - 1] if i > 0 else 0
+            start = max(begin_idx - prev_cum, 0)
+            end = min(end_idx - prev_cum, len(self.segments[i]))
+            result_segments.append(self.segments[i][start:end])
+
+        return torch.cat(result_segments, dim=0)
+
+
+class MultiSegmentFetcher:
+    """Manages multiple segment fetchers for different data keys."""
+
+    def __init__(self, multi_segments: Dict):
+        self.multi_keys = list(multi_segments.keys())
+        self.multi_fetchers = {
+            key: BaseSegmentFetcher(segments)
+            for key, segments in multi_segments.items()
+        }
+
+    def __len__(self) -> int:
+        """Returns the minimum length across all fetchers."""
+        if not self.multi_fetchers:
+            return 0
+        len_list = [len(seg) for seg in self.multi_fetchers.values()]
+        return min(len_list)
+
+    def key_fetch(
+        self, begin_idx: int, end_idx: int, keys: Union[str, List[str]]
+    ) -> Dict:
+        """Fetch data for specific keys."""
+        fetch_dict = {}
+        keys = [keys] if isinstance(keys, str) else keys
+
+        for key in keys:
+            fetcher = self.multi_fetchers[key]
+            fetch_tensor = fetcher.fetch_data(begin_idx, end_idx)
+            fetch_dict[key] = fetch_tensor
+
+        return fetch_dict if len(keys) > 1 else fetch_dict[keys[0]]
+
+    def fetch_data(self, begin_idx: int, end_idx: int) -> Dict:
+        """Fetch all keys."""
+        return self.key_fetch(begin_idx, end_idx, self.multi_keys)
+
+
+class BaseStorage(ABC):
+    """Abstract storage backend for loading and dispatching data.
+
+    Storage encapsulates format-specific loading and provides a uniform
+    interface for data access and length observation. Subclasses handle
+    different data formats (HDF5, JSON, etc.) while exposing the same
+    fetch interface.
    """

    def __init__(self):
-        self._data: Dict[str, List[Tensor]] = {}
-        self._cum: Dict[str, List[int]] = {}
-        self._length: int = 0
+        self._fetcher: Optional[MultiSegmentFetcher] = None

    @abstractmethod
-    def load(self, path: str) -> None:
+    def load(self, load_path: str, tokenizer=None) -> None:
+        """Load data from the given path into internal fetcher."""
        raise NotImplementedError

+    def __len__(self) -> int:
+        """Total number of raw elements (tokens) in storage."""
+        if self._fetcher is None:
+            return 0
+        return len(self._fetcher)
+
+    def fetch(self, begin_idx: int, end_idx: int, keys: Union[str, List[str]]):
+        """Fetch data for the given keys and index range.
+
+        Args:
+            begin_idx: Starting index (inclusive)
+            end_idx: Ending index (exclusive)
+            keys: Single key or list of keys to fetch
+
+        Returns:
+            Tensor if single key, Dict[str, Tensor] if multiple keys
+        """
+        if self._fetcher is None:
+            raise RuntimeError("Storage not loaded")
+        return self._fetcher.key_fetch(begin_idx, end_idx, keys)
+
    @property
    def keys(self) -> List[str]:
-        return list(self._data.keys())
-
-    def __len__(self) -> int:
-        return self._length
-
-    def fetch(
-        self,
-        begin: int,
-        end: int,
-        keys: Union[str, List[str]],
-    ):
-        if not self._data:
-            raise RuntimeError("Store not loaded")
-        if not (0 <= begin < self._length and 0 <= end <= self._length):
-            raise ValueError(
-                f"Index out of bounds: begin={begin}, end={end}, length={self._length}"
-            )
-        if isinstance(keys, str):
-            return self._fetch_key(keys, begin, end)
-        return {k: self._fetch_key(k, begin, end) for k in keys}
-
-    def _fetch_key(self, key: str, begin: int, end: int) -> Tensor:
-        """Fetch slice [begin, end) across potentially multiple segments."""
-        segments = self._data[key]
-        cum = self._cum[key]
-        seg_start = bisect.bisect_right(cum, begin)
-        seg_end = bisect.bisect_left(cum, end)
-
-        results = []
-        for i in range(seg_start, seg_end + 1):
-            prev = cum[i - 1] if i > 0 else 0
-            s = max(begin - prev, 0)
-            e = min(end - prev, segments[i].shape[0])
-            results.append(segments[i][s:e])
-
-        return results[0] if len(results) == 1 else torch.cat(results, dim=0)
-
-    def _normalize(self, raw: Dict[str, List[Tensor]]):
-        """Register segments and pre-compute cumulative lengths.
-
-        Does NOT concatenate — segments are kept as-is to avoid OOM on
-        large datasets.  Sets ``self._length`` to the minimum total
-        element count across all keys.
-        """
-        for key, tensors in raw.items():
-            self._data[key] = tensors
-            cum = []
-            total = 0
-            for t in tensors:
-                total += t.shape[0]
-                cum.append(total)
-            self._cum[key] = cum
-        self._length = (
-            min((cum[-1] if cum else 0) for cum in self._cum.values())
-            if self._cum
-            else 0
-        )
+        """Return the data keys available in this storage."""
+        if self._fetcher is None:
+            return []
+        return self._fetcher.multi_keys


-class StoreFactory(BaseFactory["Store"]):
-    """Factory for creating Store instances by type name.
+class StorageFactory(BaseFactory["BaseStorage"]):
+    """Factory for creating storage backends by type name.

-    Example::
-
-        @StoreFactory.register("custom")
-        class CustomStore(Store):
+    Example:
+        @StorageFactory.register("custom")
+        class CustomStorage(BaseStorage):
            ...
+
+        storage = StorageFactory.create("custom")
    """

    @classmethod
-    def _validate_component(cls, store_cls: type):
-        if not issubclass(store_cls, Store):
-            raise TypeError(f"{store_cls.__name__} must inherit from Store")
+    def _validate_component(cls, storage_cls: type) -> None:
+        if not issubclass(storage_cls, BaseStorage):
+            raise TypeError(f"{storage_cls.__name__} must inherit from BaseStorage")


-@StoreFactory.register("h5")
-class H5Store(Store):
+@StorageFactory.register("h5")
+class H5Storage(BaseStorage):
    """HDF5-based storage backend (pre-tokenized data)."""

-    def load(self, path: str):
-        self._normalize(load_h5(path))
+    def load(self, load_path: str, tokenizer=None) -> None:
+        segments = load_h5(load_path)
+        self._fetcher = MultiSegmentFetcher(segments)


-@StoreFactory.register("bin")
-class MmapStore(Store):
-    """Memory-mapped binary storage backend.
+@StorageFactory.register("json")
+class JSONStorage(BaseStorage):
+    """JSON-based storage backend.

-    Each key is a single .bin file backed by ``np.memmap(mode="r")``.
-    No per-process memory duplication — all DataLoader workers share the
-    same OS page-cache pages.
-
-    Format on disk::
-
-        data_root/
-          meta.json          # {key: {shape, dtype}, ...}
-          <key>.bin          # raw numpy array, one per key
+    Supports two modes:
+    - Pre-tokenized: JSON values are List[List[int]], loaded as-is.
+    - Raw text: JSON values are List[str], tokenized via ``tokenizer``
+      callable (str -> List[int]) at load time.
    """

-    def load(self, path: str):
-        self._mmap_refs = []
-        raw = load_bin(path)
-        self._normalize(raw)
-        for tensors in self._data.values():
-            self._mmap_refs.extend(tensors)
+    def load(self, load_path: str, tokenizer=None) -> None:
+        segments = load_json(load_path, tokenizer=tokenizer)
+        self._fetcher = MultiSegmentFetcher(segments)
@@ -23,7 +23,7 @@ class Registry:
        component_cls: Type,
        category: Optional[str] = None,
        priority: int = 0,
-    ):
+    ) -> None:
        """Register a component class with optional category and priority."""
        if name in self._entries:
            raise ValueError(f"Component '{name}' is already registered")
@@ -158,7 +158,7 @@ class BaseFactory(ABC, Generic[T]):
        return component_cls(*args, **kwargs)

    @classmethod
-    def _validate_component(cls, component_cls: Type[T]):
+    def _validate_component(cls, component_cls: Type[T]) -> None:
        """Validate that the component class is valid for this factory.

        Override this method in subclasses to add custom validation.
@@ -1,27 +1,25 @@
 """Inference module for continuous batching.

 Layers:
-  - core/:        Core inference loop (cache, executor, scheduler, task)
-  - api/:         HTTP orchestration (ProtocolHandler, server)
-  - protocols/:   Response builders (OpenAI, Anthropic)
-  - transport/:   SSE transport utilities
-  - engine.py:    Facade (InferenceEngine), Value Object (GenerationRequest)
-  - sample.py:    Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy)
+  - core/:           Core inference loop (cache, executor, scheduler, task)
+  - api/:            HTTP protocol handlers (OpenAI, Anthropic)
+  - engine.py:       Facade (InferenceEngine), Value Object (GenerationRequest)
+  - sample.py:       Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy)
 """

 from astrai.inference.api import (
+    AnthropicHandler,
    AnthropicMessage,
    ChatCompletionRequest,
    ChatMessage,
-    GenContext,
    MessagesRequest,
+    OpenAIHandler,
    ProtocolHandler,
    StopChecker,
+    StreamContext,
    app,
    run_server,
 )
-from astrai.inference.api.anthropic import AnthropicResponseBuilder
-from astrai.inference.api.openai import OpenAIResponseBuilder
 from astrai.inference.core import (
    STOP,
    Allocator,
@@ -38,7 +36,10 @@ from astrai.inference.core import (
    TaskTable,
    page_hash,
 )
-from astrai.inference.engine import GenerationRequest, InferenceEngine
+from astrai.inference.engine import (
+    GenerationRequest,
+    InferenceEngine,
+)
 from astrai.inference.sample import (
    BaseSamplingStrategy,
    SamplingPipeline,
@@ -49,14 +50,17 @@ from astrai.inference.sample import (
 )

 __all__ = [
+    # Engine / Requests
    "InferenceEngine",
    "GenerationRequest",
+    # Core scheduler
    "InferenceScheduler",
    "Executor",
    "STOP",
    "Task",
    "TaskManager",
    "TaskStatus",
+    # Core cache
    "Allocator",
    "KVCache",
    "KvcacheView",
@@ -65,17 +69,20 @@ __all__ = [
    "Storage",
    "TaskTable",
    "page_hash",
+    # Sampling (Strategy pattern)
    "sample",
    "BaseSamplingStrategy",
    "TemperatureStrategy",
    "TopKStrategy",
    "TopPStrategy",
    "SamplingPipeline",
+    # Protocol
    "ProtocolHandler",
    "StopChecker",
-    "GenContext",
-    "OpenAIResponseBuilder",
-    "AnthropicResponseBuilder",
+    "StreamContext",
+    "AnthropicHandler",
+    "OpenAIHandler",
+    # Server
    "ChatMessage",
    "ChatCompletionRequest",
    "AnthropicMessage",
@@ -1,6 +1,12 @@
-"""Inference API: protocol handler, stop checker, and FastAPI server."""
+"""Inference API: protocol handlers and FastAPI server."""

-from astrai.inference.api.protocol import GenContext, ProtocolHandler, StopChecker
+from astrai.inference.api.protocol import (
+    AnthropicHandler,
+    OpenAIHandler,
+    ProtocolHandler,
+    StopChecker,
+    StreamContext,
+)
 from astrai.inference.api.server import (
    AnthropicMessage,
    ChatCompletionRequest,
@@ -11,9 +17,11 @@ from astrai.inference.api.server import (
 )

 __all__ = [
+    "AnthropicHandler",
+    "OpenAIHandler",
    "ProtocolHandler",
    "StopChecker",
-    "GenContext",
+    "StreamContext",
    "AnthropicMessage",
    "ChatCompletionRequest",
    "ChatMessage",
@@ -1,140 +0,0 @@
-"""Anthropic message completion response builder."""
-
-import uuid
-from typing import Any, Dict, List, Tuple, Union
-
-from pydantic import BaseModel
-
-from astrai.inference.api.protocol import (
-    GenContext,
-    ResponseBuilder,
-    StopInfo,
-    sse_event,
-)
-from astrai.inference.engine import InferenceEngine
-
-
-def _extract_text(content: Union[str, List[Dict[str, Any]]]) -> str:
-    if isinstance(content, str):
-        return content
-    if isinstance(content, list):
-        for block in content:
-            if isinstance(block, dict) and block.get("type") == "text":
-                return block.get("text", "")
-    return ""
-
-
-class AnthropicResponseBuilder(ResponseBuilder):
-    def prepare(
-        self, request: BaseModel, engine: InferenceEngine
-    ) -> Tuple[str, GenContext, List[str]]:
-        messages: List[Dict[str, str]] = []
-        system = getattr(request, "system", None)
-        if system:
-            messages.append({"role": "system", "content": system})
-        for m in request.messages:
-            text = _extract_text(m.content)
-            if text:
-                messages.append({"role": m.role, "content": text})
-        prompt = engine.tokenizer.apply_chat_template(messages, tokenize=False)
-        ctx = GenContext(
-            resp_id=f"msg_{uuid.uuid4().hex[:24]}",
-            created=0,
-            model=request.model,
-            prompt_tokens=0,
-        )
-        stop_sequences = getattr(request, "stop_sequences", None) or []
-        return prompt, ctx, stop_sequences
-
-    def format_stream_start(self, ctx: GenContext) -> List[str]:
-        return [
-            sse_event(
-                {
-                    "type": "message_start",
-                    "message": {
-                        "id": ctx.resp_id,
-                        "type": "message",
-                        "role": "assistant",
-                        "model": ctx.model,
-                        "content": [],
-                        "usage": {"input_tokens": ctx.prompt_tokens},
-                    },
-                },
-                event="message_start",
-            ),
-            sse_event(
-                {
-                    "type": "content_block_start",
-                    "index": 0,
-                    "content_block": {"type": "text", "text": ""},
-                },
-                event="content_block_start",
-            ),
-        ]
-
-    def format_chunk(self, token: str) -> str:
-        return sse_event(
-            {
-                "type": "content_block_delta",
-                "index": 0,
-                "delta": {"type": "text_delta", "text": token},
-            },
-            event="content_block_delta",
-        )
-
-    def format_stream_end(self, ctx: GenContext, stop: StopInfo) -> List[str]:
-        events: List[str] = []
-        if stop.matched:
-            trimmed = stop.body[: stop.body.rfind(stop.matched)]
-            unyielded = trimmed[len(stop.yielded) :]
-            if unyielded:
-                events.append(
-                    sse_event(
-                        {
-                            "type": "content_block_delta",
-                            "index": 0,
-                            "delta": {"type": "text_delta", "text": unyielded},
-                        },
-                        event="content_block_delta",
-                    )
-                )
-        events.append(
-            sse_event(
-                {"type": "content_block_stop", "index": 0},
-                event="content_block_stop",
-            )
-        )
-        events.append(
-            sse_event(
-                {
-                    "type": "message_delta",
-                    "delta": {
-                        "stop_reason": "stop_sequence" if stop.matched else "end_turn",
-                        "stop_sequence": stop.matched,
-                    },
-                    "usage": {"output_tokens": ctx.completion_tokens},
-                },
-                event="message_delta",
-            )
-        )
-        events.append(sse_event({"type": "message_stop"}, event="message_stop"))
-        return events
-
-    def format_response(
-        self, ctx: GenContext, content: str, stop: StopInfo
-    ) -> Dict[str, Any]:
-        if stop.matched:
-            content = content[: content.rfind(stop.matched)]
-        return {
-            "id": ctx.resp_id,
-            "type": "message",
-            "role": "assistant",
-            "model": ctx.model,
-            "content": [{"type": "text", "text": content}],
-            "stop_reason": "stop_sequence" if stop.matched else "end_turn",
-            "stop_sequence": stop.matched,
-            "usage": {
-                "input_tokens": ctx.prompt_tokens,
-                "output_tokens": ctx.completion_tokens,
-            },
-        }
@@ -1,111 +0,0 @@
-"""OpenAI chat completion response builder."""
-
-import uuid
-from typing import Any, Dict, List, Tuple
-
-from pydantic import BaseModel
-
-from astrai.inference.api.protocol import (
-    GenContext,
-    ResponseBuilder,
-    StopInfo,
-    sse_event,
-)
-from astrai.inference.engine import InferenceEngine
-
-
-class OpenAIResponseBuilder(ResponseBuilder):
-    def prepare(
-        self, request: BaseModel, engine: InferenceEngine
-    ) -> Tuple[str, GenContext, List[str]]:
-        messages = [{"role": m.role, "content": m.content} for m in request.messages]
-        prompt = engine.tokenizer.apply_chat_template(messages, tokenize=False)
-
-        self._resp_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
-        self._model = request.model
-
-        ctx = GenContext(
-            resp_id=self._resp_id,
-            created=0,
-            model=self._model,
-            prompt_tokens=0,
-        )
-        stop = request.stop
-        stop_sequences = (
-            [] if stop is None else [stop] if isinstance(stop, str) else stop
-        )
-        return prompt, ctx, stop_sequences
-
-    def format_stream_start(self, ctx: GenContext) -> List[str]:
-        return [
-            sse_event(
-                {
-                    "id": self._resp_id,
-                    "object": "chat.completion.chunk",
-                    "created": ctx.created,
-                    "model": self._model,
-                    "choices": [
-                        {
-                            "index": 0,
-                            "delta": {"role": "assistant"},
-                            "finish_reason": None,
-                        }
-                    ],
-                }
-            )
-        ]
-
-    def format_chunk(self, token: str) -> str:
-        return sse_event(
-            {
-                "id": self._resp_id,
-                "object": "chat.completion.chunk",
-                "created": 0,
-                "model": self._model,
-                "choices": [
-                    {"index": 0, "delta": {"content": token}, "finish_reason": None}
-                ],
-            }
-        )
-
-    def format_stream_end(self, ctx: GenContext, stop: StopInfo) -> List[str]:
-        return [
-            sse_event(
-                {
-                    "id": self._resp_id,
-                    "object": "chat.completion.chunk",
-                    "created": ctx.created,
-                    "model": self._model,
-                    "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
-                }
-            ),
-            sse_event(
-                {
-                    "prompt_tokens": ctx.prompt_tokens,
-                    "completion_tokens": ctx.completion_tokens,
-                    "total_tokens": ctx.prompt_tokens + ctx.completion_tokens,
-                }
-            ),
-        ]
-
-    def format_response(
-        self, ctx: GenContext, content: str, stop: StopInfo
-    ) -> Dict[str, Any]:
-        return {
-            "id": self._resp_id,
-            "object": "chat.completion",
-            "created": ctx.created,
-            "model": self._model,
-            "choices": [
-                {
-                    "index": 0,
-                    "message": {"role": "assistant", "content": content},
-                    "finish_reason": "stop",
-                }
-            ],
-            "usage": {
-                "prompt_tokens": ctx.prompt_tokens,
-                "completion_tokens": ctx.completion_tokens,
-                "total_tokens": ctx.prompt_tokens + ctx.completion_tokens,
-            },
-        }
@@ -1,13 +1,15 @@
-"""Orchestration layer: ProtocolHandler, StopChecker, GenContext, StopInfo, ResponseBuilder, SSE utils.
+"""Protocol handlers for OpenAI and Anthropic chat completion APIs.

-ProtocolHandler orchestrates the async generation loop and delegates
-protocol-specific formatting to a ResponseBuilder.
+Template Method + Builder patterns eliminate the 45% code duplication between
+stream/non-stream branches and across protocol adapters.
 """

 import json
+import time
+import uuid
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Union

 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
@@ -15,7 +17,7 @@ from pydantic import BaseModel
 from astrai.inference.engine import InferenceEngine


-def sse_event(data: Dict[str, Any], event: Optional[str] = None) -> str:
+def _sse_event(data: Dict[str, Any], event: Optional[str] = None) -> str:
    lines: List[str] = []
    if event:
        lines.append(f"event: {event}")
@@ -24,28 +26,22 @@ def sse_event(data: Dict[str, Any], event: Optional[str] = None) -> str:
    return "\n".join(lines)


-def sse_done() -> str:
+def _sse_done() -> str:
    return "data: [DONE]\n\n"


@dataclass
-class GenContext:
-    """Per-generation metadata passed to builder format methods."""
+class StreamContext:
+    """Shared state across the streaming generation lifecycle."""

    resp_id: str
    created: int
    model: str
    prompt_tokens: int
    completion_tokens: int = 0
-
-
-@dataclass
-class StopInfo:
-    """Stop-check result passed to format_stream_end / format_response."""
-
-    matched: Optional[str] = None
-    body: str = ""
-    yielded: str = ""
+    accumulated: str = ""
+    stop_matched: Optional[str] = None
+    last_yield_trimmed: str = ""


 class StopChecker:
@@ -60,60 +56,95 @@ class StopChecker:
                return seq
        return None

+    def trim(self, text: str, matched: str) -> str:
+        idx = text.rfind(matched)
+        return text[:idx] if idx != -1 else text

-class ResponseBuilder(ABC):
-    """Interface for protocol-specific response formatting.
+    @property
+    def has_sequences(self) -> bool:
+        return len(self._sequences) > 0

-    A new protocol requires one concrete builder implementing 6 methods.
+
+class ProtocolHandler(ABC):
+    """Template-method base for API protocol handlers.
+
+    Subclasses implement format hooks; the base class orchestrates the
+    generate-async loop and SSE/JSON response construction.
+
+    Lifecycle::
+
+        handle()
+          ├─ build_prompt()          # protocol-specific prompt assembly
+          ├─ create_response_id()    # unique response identifier
+          ├─ [stream]
+          │   ├─ format_stream_start()
+          │   ├─ format_stream_token()  × N
+          │   │   └─ on_token() hook for stop-sequence interception
+          │   └─ format_stream_end()
+          └─ [non-stream]
+              ├─ (accumulate tokens)
+              └─ format_non_stream_response()
    """

-    @abstractmethod
-    def prepare(
-        self, request: BaseModel, engine: InferenceEngine
-    ) -> Tuple[str, GenContext, List[str]]:
-        """Return (prompt, ctx, stop_sequences) for a generation request."""
+    request_model: type[BaseModel]

-    @abstractmethod
-    def format_stream_start(self, ctx: GenContext) -> List[str]:
-        """SSE events that open the stream."""
-
-    @abstractmethod
-    def format_chunk(self, token: str) -> str:
-        """SSE event for a single generated token."""
-
-    @abstractmethod
-    def format_stream_end(self, ctx: GenContext, stop: StopInfo) -> List[str]:
-        """SSE events that close the stream."""
-
-    @abstractmethod
-    def format_response(
-        self, ctx: GenContext, content: str, stop: StopInfo
-    ) -> Dict[str, Any]:
-        """JSON response body for non-streaming mode."""
-
-
-class ProtocolHandler:
-    """Orchestrates the generation loop, delegates formatting to a builder.
-
-    Usage::
-
-        handler = ProtocolHandler(request, engine, OpenAIResponseBuilder())
-        response = await handler.handle()
-    """
-
-    def __init__(
-        self, request: BaseModel, engine: InferenceEngine, builder: ResponseBuilder
-    ):
+    def __init__(self, request: BaseModel, engine: InferenceEngine):
        self.request = request
        self.engine = engine
-        self.builder = builder
+
+    @abstractmethod
+    def build_prompt(self) -> str:
+        """Build the full prompt string from the request messages."""
+
+    @abstractmethod
+    def create_response_id(self) -> str:
+        """Generate a unique response ID following the protocol convention."""
+
+    @abstractmethod
+    def format_stream_start(self, ctx: StreamContext) -> List[str]:
+        """Yield SSE events that open the stream (role marker, metadata)."""
+
+    @abstractmethod
+    def format_stream_token(self, ctx: StreamContext, token: str) -> str:
+        """Yield an SSE event for a single generated token."""
+
+    @abstractmethod
+    def format_stream_end(self, ctx: StreamContext) -> List[str]:
+        """Yield SSE events that close the stream (finish reason, usage stats)."""
+
+    @abstractmethod
+    def format_non_stream_response(
+        self, ctx: StreamContext, content: str
+    ) -> Dict[str, Any]:
+        """Build the JSON response body for non-streaming mode."""
+
+    def get_stop_sequences(self) -> List[str]:
+        return []
+
+    def create_stop_checker(self) -> StopChecker:
+        return StopChecker(self.get_stop_sequences())
+
+    def on_token(
+        self, ctx: StreamContext, token: str, stop_checker: StopChecker
+    ) -> Optional[str]:
+        """Hook after each token is appended to accumulated.
+
+        Return a matched stop-sequence string to break the loop,
+        or None to continue.
+
+        """
+        return None

    async def handle(self) -> Union[StreamingResponse, Dict[str, Any]]:
-        prompt, ctx, stop_sequences = self.builder.prepare(self.request, self.engine)
-        ctx.prompt_tokens = len(self.engine.tokenizer.encode(prompt))
+        ctx = StreamContext(
+            resp_id=self.create_response_id(),
+            created=int(time.time()),
+            model=self.request.model,
+            prompt_tokens=self._count_prompt_tokens(),
+        )

        agen = self.engine.generate_async(
-            prompt=prompt,
+            prompt=self.build_prompt(),
            max_tokens=self.request.max_tokens,
            temperature=self.request.temperature,
            top_p=self.request.top_p,
@@ -121,37 +152,33 @@ class ProtocolHandler:
        )

        if self.request.stream:
-            return self._handle_stream(agen, ctx, stop_sequences)
+            return self._handle_stream(agen, ctx)
        else:
-            return await self._handle_non_stream(agen, ctx, stop_sequences)
+            return await self._handle_non_stream(agen, ctx)

-    def _handle_stream(
-        self, agen: AsyncGenerator, ctx: GenContext, stop_sequences: List[str]
-    ) -> StreamingResponse:
-        checker = StopChecker(stop_sequences)
+    def _count_prompt_tokens(self) -> int:
+        return len(self.engine.tokenizer.encode(self.build_prompt()))
+
+    def _handle_stream(self, agen, ctx: StreamContext) -> StreamingResponse:
+        stop_checker = self.create_stop_checker()

        async def event_stream():
-            for event in self.builder.format_stream_start(ctx):
+            for event in self.format_stream_start(ctx):
                yield event

-            body = ""
-            yielded = ""
-            matched = None
            async for token in agen:
                ctx.completion_tokens += 1
-                body += token
+                ctx.accumulated += token

-                matched = checker.check(body)
+                matched = self.on_token(ctx, token, stop_checker)
                if matched:
                    break

-                yield self.builder.format_chunk(token)
-                yielded += token
+                yield self.format_stream_token(ctx, token)

-            stop = StopInfo(matched=matched, body=body, yielded=yielded)
-            for event in self.builder.format_stream_end(ctx, stop):
+            for event in self.format_stream_end(ctx):
                yield event
-            yield sse_done()
+            yield _sse_done()

        return StreamingResponse(
            event_stream(),
@@ -159,23 +186,260 @@ class ProtocolHandler:
            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
        )

-    async def _handle_non_stream(
-        self, agen: AsyncGenerator, ctx: GenContext, stop_sequences: List[str]
-    ) -> Dict[str, Any]:
-        checker = StopChecker(stop_sequences)
+    async def _handle_non_stream(self, agen, ctx: StreamContext) -> Dict[str, Any]:
+        stop_checker = self.create_stop_checker()
        chunks: List[str] = []
-        body = ""
-        matched = None

        async for token in agen:
            ctx.completion_tokens += 1
+            ctx.accumulated += token
            chunks.append(token)
-            body += token

-            matched = checker.check(body)
+            matched = self.on_token(ctx, token, stop_checker)
            if matched:
                break

        content = "".join(chunks)
-        stop = StopInfo(matched=matched, body=body)
-        return self.builder.format_response(ctx, content, stop)
+        return self.format_non_stream_response(ctx, content)
+
+
+def _extract_text_content(content: Union[str, List[Dict[str, Any]]]) -> str:
+    """Extract plain text from an Anthropic content block (string or list)."""
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        for block in content:
+            if isinstance(block, dict) and block.get("type") == "text":
+                return block.get("text", "")
+    return ""
+
+
+class OpenAIHandler(ProtocolHandler):
+    """OpenAI-compatible /v1/chat/completions handler."""
+
+    def build_prompt(self) -> str:
+        messages = [
+            {"role": m.role, "content": m.content} for m in self.request.messages
+        ]
+        return self.engine.tokenizer.apply_chat_template(messages, tokenize=False)
+
+    def create_response_id(self) -> str:
+        return f"chatcmpl-{uuid.uuid4().hex[:12]}"
+
+    def get_stop_sequences(self) -> List[str]:
+        stop = self.request.stop
+        if stop is None:
+            return []
+        return [stop] if isinstance(stop, str) else stop
+
+    def on_token(
+        self, ctx: StreamContext, token: str, stop_checker: StopChecker
+    ) -> Optional[str]:
+        return stop_checker.check(ctx.accumulated)
+
+    def format_stream_start(self, ctx: StreamContext) -> List[str]:
+        return [
+            _sse_event(
+                {
+                    "id": ctx.resp_id,
+                    "object": "chat.completion.chunk",
+                    "created": ctx.created,
+                    "model": ctx.model,
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"role": "assistant"},
+                            "finish_reason": None,
+                        }
+                    ],
+                }
+            )
+        ]
+
+    def format_stream_token(self, ctx: StreamContext, token: str) -> str:
+        return _sse_event(
+            {
+                "id": ctx.resp_id,
+                "object": "chat.completion.chunk",
+                "created": ctx.created,
+                "model": ctx.model,
+                "choices": [
+                    {"index": 0, "delta": {"content": token}, "finish_reason": None}
+                ],
+            }
+        )
+
+    def format_stream_end(self, ctx: StreamContext) -> List[str]:
+        return [
+            _sse_event(
+                {
+                    "id": ctx.resp_id,
+                    "object": "chat.completion.chunk",
+                    "created": ctx.created,
+                    "model": ctx.model,
+                    "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+                }
+            ),
+            _sse_event(
+                {
+                    "prompt_tokens": ctx.prompt_tokens,
+                    "completion_tokens": ctx.completion_tokens,
+                    "total_tokens": ctx.prompt_tokens + ctx.completion_tokens,
+                }
+            ),
+        ]
+
+    def format_non_stream_response(
+        self, ctx: StreamContext, content: str
+    ) -> Dict[str, Any]:
+        return {
+            "id": ctx.resp_id,
+            "object": "chat.completion",
+            "created": ctx.created,
+            "model": ctx.model,
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": content},
+                    "finish_reason": "stop",
+                }
+            ],
+            "usage": {
+                "prompt_tokens": ctx.prompt_tokens,
+                "completion_tokens": ctx.completion_tokens,
+                "total_tokens": ctx.prompt_tokens + ctx.completion_tokens,
+            },
+        }
+
+
+class AnthropicHandler(ProtocolHandler):
+    """Anthropic-compatible /v1/messages handler."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._yielded = ""
+
+    def build_prompt(self) -> str:
+        messages: List[Dict[str, str]] = []
+        system = getattr(self.request, "system", None)
+        if system:
+            messages.append({"role": "system", "content": system})
+        for m in self.request.messages:
+            content = _extract_text_content(m.content)
+            if content:
+                messages.append({"role": m.role, "content": content})
+        return self.engine.tokenizer.apply_chat_template(messages, tokenize=False)
+
+    def create_response_id(self) -> str:
+        return f"msg_{uuid.uuid4().hex[:24]}"
+
+    def get_stop_sequences(self) -> List[str]:
+        return getattr(self.request, "stop_sequences", None) or []
+
+    def on_token(
+        self, ctx: StreamContext, token: str, stop_checker: StopChecker
+    ) -> Optional[str]:
+        matched = stop_checker.check(ctx.accumulated)
+        if not matched:
+            return None
+
+        ctx.stop_matched = matched
+        trimmed = ctx.accumulated[: ctx.accumulated.rfind(matched)]
+        unyielded = trimmed[len(self._yielded) :]
+        if unyielded:
+            ctx.last_yield_trimmed = unyielded
+        return matched
+
+    def format_stream_start(self, ctx: StreamContext) -> List[str]:
+        return [
+            _sse_event(
+                {
+                    "type": "message_start",
+                    "message": {
+                        "id": ctx.resp_id,
+                        "type": "message",
+                        "role": "assistant",
+                        "model": ctx.model,
+                        "content": [],
+                        "usage": {"input_tokens": ctx.prompt_tokens},
+                    },
+                },
+                event="message_start",
+            ),
+            _sse_event(
+                {
+                    "type": "content_block_start",
+                    "index": 0,
+                    "content_block": {"type": "text", "text": ""},
+                },
+                event="content_block_start",
+            ),
+        ]
+
+    def format_stream_token(self, ctx: StreamContext, token: str) -> str:
+        self._yielded += token
+        return _sse_event(
+            {
+                "type": "content_block_delta",
+                "index": 0,
+                "delta": {"type": "text_delta", "text": token},
+            },
+            event="content_block_delta",
+        )
+
+    def format_stream_end(self, ctx: StreamContext) -> List[str]:
+        matched = ctx.stop_matched
+        events: List[str] = []
+        last_yielded = ctx.last_yield_trimmed
+        if last_yielded:
+            events.append(
+                _sse_event(
+                    {
+                        "type": "content_block_delta",
+                        "index": 0,
+                        "delta": {"type": "text_delta", "text": last_yielded},
+                    },
+                    event="content_block_delta",
+                )
+            )
+        events.append(
+            _sse_event(
+                {"type": "content_block_stop", "index": 0},
+                event="content_block_stop",
+            )
+        )
+        events.append(
+            _sse_event(
+                {
+                    "type": "message_delta",
+                    "delta": {
+                        "stop_reason": "stop_sequence" if matched else "end_turn",
+                        "stop_sequence": matched,
+                    },
+                    "usage": {"output_tokens": ctx.completion_tokens},
+                },
+                event="message_delta",
+            )
+        )
+        events.append(_sse_event({"type": "message_stop"}, event="message_stop"))
+        return events
+
+    def format_non_stream_response(
+        self, ctx: StreamContext, content: str
+    ) -> Dict[str, Any]:
+        matched = ctx.stop_matched
+        if matched:
+            content = content[: content.rfind(matched)]
+        return {
+            "id": ctx.resp_id,
+            "type": "message",
+            "role": "assistant",
+            "model": ctx.model,
+            "content": [{"type": "text", "text": content}],
+            "stop_reason": "stop_sequence" if matched else "end_turn",
+            "stop_sequence": matched,
+            "usage": {
+                "input_tokens": ctx.prompt_tokens,
+                "output_tokens": ctx.completion_tokens,
+            },
+        }
@@ -15,9 +15,7 @@ import uvicorn
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel, Field

-from astrai.inference.api.anthropic import AnthropicResponseBuilder
-from astrai.inference.api.openai import OpenAIResponseBuilder
-from astrai.inference.api.protocol import ProtocolHandler
+from astrai.inference.api.protocol import AnthropicHandler, OpenAIHandler
 from astrai.inference.engine import InferenceEngine
 from astrai.model import AutoModel
 from astrai.tokenize import AutoTokenizer
@@ -135,14 +133,14 @@ async def get_stats():
@app.post("/v1/chat/completions")
 async def chat_completion(request: ChatCompletionRequest):
    engine = _get_engine()
-    handler = ProtocolHandler(request, engine, OpenAIResponseBuilder())
+    handler = OpenAIHandler(request, engine)
    return await handler.handle()


@app.post("/v1/messages")
 async def create_message(request: MessagesRequest):
    engine = _get_engine()
-    handler = ProtocolHandler(request, engine, AnthropicResponseBuilder())
+    handler = AnthropicHandler(request, engine)
    return await handler.handle()


@@ -42,7 +42,7 @@ class Allocator:
                return idx
            return -1

-    def free(self, idx: int, keep_cached: bool = False):
+    def free(self, idx: int, keep_cached: bool = False) -> None:
        with self._lock:
            self._refs[idx] -= 1
            if self._refs[idx] == 0:
@@ -51,7 +51,7 @@ class Allocator:
                else:
                    self._free_mask |= 1 << idx

-    def inc_ref(self, idx: int):
+    def inc_ref(self, idx: int) -> None:
        with self._lock:
            self._refs[idx] += 1
            self._lru.pop(idx, None)
@@ -60,7 +60,7 @@ class Allocator:
        with self._lock:
            return self._refs[idx]

-    def touch(self, idx: int):
+    def touch(self, idx: int) -> None:
        with self._lock:
            self._lru.move_to_end(idx)

@@ -74,7 +74,7 @@ class PrefixCache:
        self._hash_to_page: Dict[int, int] = {}
        self._lock = threading.Lock()

-    def evict(self, idx: int):
+    def evict(self, idx: int) -> None:
        with self._lock:
            h = self._page_to_hash.pop(idx, None)
            if h is not None:
@@ -96,7 +96,9 @@ class PrefixCache:
                hits.append(p)
            return hits

-    def record(self, page_idx: int, token_ids: List[int], logical_page_idx: int):
+    def record(
+        self, page_idx: int, token_ids: List[int], logical_page_idx: int
+    ) -> None:
        with self._lock:
            h = page_hash(token_ids, logical_page_idx, self._page_size)
            old_h = self._page_to_hash.pop(page_idx, None)
@@ -125,13 +127,13 @@ class PagePool:
    def alloc(self) -> int:
        return self._alloc.alloc()

-    def free(self, idx: int):
+    def free(self, idx: int) -> None:
        keep = self._prefix.has_page(idx)
        self._alloc.free(idx, keep_cached=keep)
        if not keep:
            self._prefix.evict(idx)

-    def inc_ref(self, idx: int):
+    def inc_ref(self, idx: int) -> None:
        self._alloc.inc_ref(idx)

    def lookup(self, token_ids: List[int]) -> List[int]:
@@ -140,7 +142,9 @@ class PagePool:
            self._alloc.touch(p)
        return hits

-    def record(self, page_idx: int, token_ids: List[int], logical_page_idx: int):
+    def record(
+        self, page_idx: int, token_ids: List[int], logical_page_idx: int
+    ) -> None:
        self._prefix.record(page_idx, token_ids, logical_page_idx)


@@ -153,7 +157,7 @@ class TaskTable:
        self._cached: Dict[str, int] = {}
        self._lock = threading.Lock()

-    def set(self, task_id: str, page_table: List[int], cached: int):
+    def set(self, task_id: str, page_table: List[int], cached: int) -> None:
        with self._lock:
            self._pages[task_id] = page_table
            self._cached[task_id] = cached
@@ -216,7 +220,7 @@ class Storage:
        start_pos: int,
        k: Tensor,
        v: Tensor,
-    ):
+    ) -> None:
        seq_len = k.size(1)
        if seq_len == 0:
            return
@@ -282,7 +286,7 @@ class KvcacheView:
        self._page_table = page_table
        self._total_len = total_len

-    def write(self, layer_id: int, k: Tensor, v: Tensor):
+    def write(self, layer_id: int, k: Tensor, v: Tensor) -> None:
        start_pos = self._total_len - k.size(1)
        self._storage.write(layer_id, self._page_table, start_pos, k, v)

@@ -335,7 +339,7 @@ class KVCache:
        self._table.set(task_id, hits + new_pages, cached)
        return True

-    def task_free(self, task_id: str):
+    def task_free(self, task_id: str) -> None:
        page_table, _ = self._table.pop(task_id)
        for idx in page_table:
            self._pool.free(idx)
@@ -355,7 +359,7 @@ class KVCache:

    def task_record_hashes(
        self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0
-    ):
+    ) -> None:
        page_table = self._table.get(task_id)
        full_pages = len(prompt_ids) // self.page_size
        for i in range(start_logical_page, full_pages):
@@ -29,7 +29,9 @@ class Executor:
        self.device = device or next(model.parameters()).device
        self.dtype = dtype or next(model.parameters()).dtype

-    def execute_prefill(self, tasks: List[Task], prompt_len: int, start_pos: int = 0):
+    def execute_prefill(
+        self, tasks: List[Task], prompt_len: int, start_pos: int = 0
+    ) -> None:
        if start_pos >= prompt_len:
            return

@@ -75,14 +75,14 @@ class InferenceScheduler:
    def add_task(self, prompt: str, **kwargs) -> str:
        return self._task_mgr.add_task(prompt, **kwargs)

-    def remove_task(self, task_id: str):
+    def remove_task(self, task_id: str) -> None:
        for task in self._task_mgr.remove_task(task_id):
            self._page_cache.task_free(task.task_id)

    def get_stats(self) -> Dict[str, Any]:
        return self._task_mgr.get_stats()

-    def _run_generation_loop(self):
+    def _run_generation_loop(self) -> None:
        stop_ids = self._task_mgr.tokenizer.stop_ids
        try:
            while self._running:
@@ -108,10 +108,7 @@ class InferenceScheduler:
                    continue

                to_prefill = [
-                    t
-                    for t in self._task_mgr.get_active_tasks()
-                    if t.output_tokens == 0
-                    and self._page_cache.task_cached(t.task_id) < len(t.prompt_ids)
+                    t for t in self._task_mgr.get_active_tasks() if t.output_tokens == 0
                ]
                if to_prefill:
                    for t in to_prefill:
@@ -186,14 +183,14 @@ class InferenceScheduler:
            self._task_mgr.clear_queues()
            raise

-    def start(self):
+    def start(self) -> None:
        if not self._running:
            self._running = True
            t = threading.Thread(target=self._run_generation_loop, daemon=True)
            t.start()
            self._loop_thread = t

-    def stop(self):
+    def stop(self) -> None:
        self._running = False
        self._task_mgr.wake()
        if hasattr(self, "_loop_thread"):
@@ -172,12 +172,12 @@ class TaskManager:
                to_add.append(self.waiting_queue.popleft())
        return to_add

-    def activate(self, task: Task):
+    def activate(self, task: Task) -> None:
        task.status = TaskStatus.RUNNING
        with self._lock:
            self.active_tasks.append(task)

-    def return_to_waiting(self, tasks: List[Task]):
+    def return_to_waiting(self, tasks: List[Task]) -> None:
        with self._lock:
            for task in reversed(tasks):
                self.waiting_queue.appendleft(task)
@@ -185,7 +185,7 @@ class TaskManager:
    def has_work(self) -> bool:
        return bool(self.active_tasks or self.waiting_queue)

-    def wait_for_tasks(self, timeout: float = 1.0):
+    def wait_for_tasks(self, timeout: float = 1.0) -> None:
        self._task_event.clear()
        self._task_event.wait(timeout=timeout)

@@ -197,10 +197,10 @@ class TaskManager:
        with self._lock:
            return list(self.waiting_queue)

-    def clear_queues(self):
+    def clear_queues(self) -> None:
        with self._lock:
            self.waiting_queue.clear()
            self.active_tasks.clear()

-    def wake(self):
+    def wake(self) -> None:
        self._task_event.set()
@@ -13,6 +13,17 @@ from astrai.inference.core.task import STOP
 from astrai.tokenize import AutoTokenizer


+def _validate_sampling_params(
+    top_k: int, top_p: float, temperature: float, max_tokens: Optional[int] = None
+):
+    if not (isinstance(top_k, int) and top_k >= 0):
+        raise ValueError("top_k must be a non-negative integer")
+    if not (0.0 <= top_p <= 1.0):
+        raise ValueError("top_p must be a float between 0.0 and 1.0")
+    if not (isinstance(temperature, (int, float)) and temperature >= 0):
+        raise ValueError("temperature must be a non-negative number")
+
+
 class GenerateResult:
    """Thread-safe token accumulator for streaming and non-streaming modes."""

@@ -48,7 +59,7 @@ class GenerateResult:
    def wait(self, timeout: Optional[float] = None) -> bool:
        return self._event.wait(timeout=timeout)

-    def wait_completion(self, timeout: float = 300.0):
+    def wait_completion(self, timeout: float = 300.0) -> None:
        with self._cond:
            if not self._cond.wait_for(
                lambda: self._completed >= self._total, timeout=timeout
@@ -75,12 +86,7 @@ class GenerationRequest:
        max_tokens: Optional[int] = None,
        stream: bool = False,
    ):
-        if not (isinstance(top_k, int) and top_k >= 0):
-            raise ValueError("top_k must be a non-negative integer")
-        if not (0.0 <= top_p <= 1.0):
-            raise ValueError("top_p must be a float between 0.0 and 1.0")
-        if not (isinstance(temperature, (int, float)) and temperature >= 0):
-            raise ValueError("temperature must be a non-negative number")
+        _validate_sampling_params(top_k, top_p, temperature, max_tokens)

        self.messages = messages
        self.top_k = top_k
@@ -131,6 +137,7 @@ class InferenceEngine:
        top_p: float = 1.0,
        top_k: int = 50,
    ) -> Union[Generator, str, List[str]]:
+        _validate_sampling_params(top_k, top_p, temperature, max_tokens)
        is_batch = isinstance(prompt, list)
        prompts = prompt if is_batch else [prompt]

@@ -151,6 +158,7 @@ class InferenceEngine:
        top_p: float = 1.0,
        top_k: int = 50,
    ) -> AsyncGenerator[str, None]:
+        _validate_sampling_params(top_k, top_p, temperature, max_tokens)
        sync_gen = self._generate_streaming(
            [prompt], False, max_tokens, temperature, top_p, top_k
        )
@@ -281,7 +289,7 @@ class InferenceEngine:
    def get_stats(self) -> Dict[str, Any]:
        return self.scheduler.get_stats()

-    def shutdown(self):
+    def shutdown(self) -> None:
        self.scheduler.stop()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
@@ -2,24 +2,21 @@
 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
 def _disable_random_init(enable: bool = True):
-    if not enable:
-        yield
-        return
-
-    names = (
+    init_functions = [
        "xavier_normal_",
        "xavier_uniform_",
        "kaiming_normal_",
@@ -29,15 +26,18 @@ def _disable_random_init(enable: bool = True):
        "constant_",
        "normal_",
        "uniform_",
-    )
-    orig = {n: getattr(nn.init, n) for n in names if hasattr(nn.init, n)}
-    for n in orig:
-        setattr(nn.init, n, lambda *a, **kw: None)
+    ]
+    original_funcs = {}
+    for name in init_functions:
+        if enable and hasattr(nn.init, name):
+            original_funcs[name] = getattr(nn.init, name)
+            setattr(nn.init, name, lambda *args, **kwargs: None)
    try:
        yield
    finally:
-        for n, fn in orig.items():
-            setattr(nn.init, n, fn)
+        if enable:
+            for name, orig_func in original_funcs.items():
+                setattr(nn.init, name, orig_func)


 class AutoModel(BaseFactory["AutoModel"], nn.Module):
@@ -60,22 +60,25 @@ class AutoModel(BaseFactory["AutoModel"], nn.Module):

        model_path = Path(path)

+        # Load config
        config_path = model_path / "config.json"
-        if not config_path.exists():
+        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:
            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 = load_model_weights(str(model_path))
+            state_dict = st.load_file(str(weights_path))
            model.load_state_dict(state_dict, strict=strict)

        return model
@@ -83,12 +86,15 @@ class AutoModel(BaseFactory["AutoModel"], nn.Module):
    def save_pretrained(
        self,
        save_directory: Union[str, Path],
-    ):
-        save_model(
-            config=self.config.to_dict(),
-            state_dict=self.state_dict(),
-            save_directory=str(save_directory),
-        )
+    ) -> 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"))

    def to(self, *args, **kwargs) -> Self:
        """Move model to device/dtype."""
@@ -1,19 +1,15 @@
+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__)

@@ -132,14 +128,16 @@ def save_lora(model: nn.Module, save_dir: str, config: LoRAConfig):

    path = Path(save_dir)
    path.mkdir(parents=True, exist_ok=True)
-    save_safetensors(lora_sd, path / "adapter_model.safetensors")
-    save_json(asdict(config), path / "adapter_config.json")
+    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)
    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)
-    raw = load_json(path / "adapter_config.json")
+    with open(path / "adapter_config.json") as f:
+        raw = json.load(f)
    config = LoRAConfig(
        r=raw["r"], alpha=raw["alpha"], target_modules=tuple(raw["target_modules"])
    )
@@ -159,7 +157,7 @@ def load_lora(model: nn.Module, load_dir: str) -> LoRAConfig:
            target_modules=set(config.target_modules),
        )

-    weights = load_safetensors(path / "adapter_model.safetensors")
+    weights = st.load_file(str(path / "adapter_model.safetensors"))
    try:
        missing, unexpected = model.load_state_dict(weights, strict=False)
    except RuntimeError as e:
@@ -1,4 +1,4 @@
-from typing import Dict, Optional
+from typing import Optional

 import torch
 import torch.nn as nn
@@ -19,10 +19,6 @@ def get_rotary_emb(
    return torch.complex(cos, sin)


-def ntk_base(base: float, dim: int, factor: float) -> float:
-    return base * (factor ** (dim / (dim - 2)))
-
-
 def apply_rotary_emb(x: torch.Tensor, freqs_cis: Tensor) -> Tensor:
    dtype = x.dtype
    x_ = x.float().reshape(*x.shape[:-1], -1, 2)
@@ -34,25 +30,11 @@ def apply_rotary_emb(x: torch.Tensor, freqs_cis: Tensor) -> Tensor:


 class RotaryEmbedding(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        max_len: int,
-        base: float = 10000,
-        rope_scaling: Optional[Dict] = None,
-    ):
+    def __init__(self, dim: int, max_len: int, base: float = 10000):
        super().__init__()
        self.dim = dim
        self.max_len = max_len
        self.base = base
-        self.rope_scaling = rope_scaling
-
-        if rope_scaling is not None:
-            scaling_type = rope_scaling.get("type", "ntk")
-            factor = rope_scaling.get("factor", 1.0)
-            if scaling_type == "ntk":
-                self.base = ntk_base(base, dim, factor)
-
        self._set_rotary_buffer(self.max_len)

    def _set_rotary_buffer(self, max_len: int):
@@ -20,9 +20,7 @@ class EmbeddingEncoder(AutoModel):
        self.config = config
        rope_dim = config.dim // config.n_heads
        rope_base = config.rope_theta if config.rope_theta is not None else 10000
-        self.rotary_embedding = RotaryEmbedding(
-            rope_dim, config.max_len, rope_base, rope_scaling=config.rope_scaling
-        )
+        self.rotary_embedding = RotaryEmbedding(rope_dim, config.max_len, rope_base)
        self.embed_tokens = Embedding(config.vocab_size, config.dim)

        self.layers = nn.ModuleList(
@@ -68,6 +66,9 @@ class EmbeddingEncoder(AutoModel):

        x = self.embed_tokens(input_ids)

+        if position_ids is None:
+            position_ids = torch.arange(S, device=x.device).unsqueeze(0).expand(B, -1)
+
        rotary_emb = self.rotary_embedding(x, position_ids)
        attn_mask = process_attention_mask(x, position_ids, input_mask, is_causal=False)

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Mapping, Optional
+from typing import Any, Mapping, Optional

 import torch
 import torch.nn as nn
@@ -59,9 +59,7 @@ class AutoRegressiveLM(AutoModel):
            else config.dim // config.n_heads
        )
        rope_base = config.rope_theta if config.rope_theta is not None else 10000
-        self.rotary_embedding = RotaryEmbedding(
-            rope_dim, config.max_len, rope_base, rope_scaling=config.rope_scaling
-        )
+        self.rotary_embedding = RotaryEmbedding(rope_dim, config.max_len, rope_base)
        self.embed_tokens = Embedding(config.vocab_size, config.dim)

        self.layers = nn.ModuleList(
@@ -136,7 +134,7 @@ class AutoRegressiveLM(AutoModel):
        input_mask: Optional[Tensor] = None,
        paged_cache: Optional[KvcacheView] = None,
        position_ids: Optional[Tensor] = None,
-    ) -> Dict[str, Tensor]:
+    ) -> Tensor:
        assert input_ids.ndim == 2

        x = self.embed_tokens(input_ids)
@@ -203,45 +203,9 @@ class DDPExecutor(BaseExecutor):

@ExecutorFactory.register("fsdp")
 class FSDPExecutor(BaseExecutor):
-    def __init__(
-        self,
-        grad_accum_steps: int = 1,
-        process_group=None,
-        sharding_strategy=None,
-        cpu_offload=None,
-        auto_wrap_policy=None,
-        backward_prefetch=None,
-        mixed_precision=None,
-        ignored_modules=None,
-        param_init_fn=None,
-        sync_module_states: bool = False,
-        forward_prefetch: bool = False,
-        limit_all_gathers: bool = True,
-        use_orig_params: bool = False,
-        ignored_states=None,
-        device_mesh=None,
-    ):
+    def __init__(self, grad_accum_steps: int = 1, **fsdp_kwargs):
        super().__init__(grad_accum_steps=grad_accum_steps)
-        self._fsdp_kwargs = {
-            k: v
-            for k, v in dict(
-                process_group=process_group,
-                sharding_strategy=sharding_strategy,
-                cpu_offload=cpu_offload,
-                auto_wrap_policy=auto_wrap_policy,
-                backward_prefetch=backward_prefetch,
-                mixed_precision=mixed_precision,
-                ignored_modules=ignored_modules,
-                param_init_fn=param_init_fn,
-                sync_module_states=sync_module_states,
-                forward_prefetch=forward_prefetch,
-                limit_all_gathers=limit_all_gathers,
-                use_orig_params=use_orig_params,
-                ignored_states=ignored_states,
-                device_mesh=device_mesh,
-            ).items()
-            if v is not None
-        }
+        self._fsdp_kwargs = fsdp_kwargs
        self._original_model: Optional[nn.Module] = None

    def _prepare_model(self, model: nn.Module) -> nn.Module:
@@ -1,9 +1,7 @@
-import io
 import json
 import time
-from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Dict, Union
+from typing import Any, Dict, Optional

 import safetensors.torch as st
 import torch
@@ -11,172 +9,75 @@ import torch.distributed as dist

 from astrai.parallel.setup import get_rank

-_META_FILE = "meta.json"
-_CONFIG_FILE = "config.json"
-_WEIGHTS_FILE = "model.safetensors"

-
-def save_safetensors(state_dict: dict, path: Union[str, Path]):
-    st.save_file(state_dict, str(path))
-
-
-def load_safetensors(path: Union[str, Path], broadcast: bool = False) -> dict:
-    if not broadcast or not dist.is_initialized():
-        return st.load_file(str(path))
-
-    rank = get_rank()
-    if rank == 0:
-        state_dict = st.load_file(str(path))
-    else:
-        state_dict = {}
-    tmp = [state_dict]
-    dist.broadcast_object_list(tmp, src=0)
-    return tmp[0]
-
-
-def save_json(data: dict, path: Union[str, Path]):
-    with open(str(path), "w") as f:
-        json.dump(data, f, indent=2)
-
-
-def load_json(path: Union[str, Path], broadcast: bool = False) -> dict:
-    if not broadcast or not dist.is_initialized():
-        with open(str(path), "r") as f:
-            return json.load(f)
-
-    rank = get_rank()
-    if rank == 0:
-        with open(str(path), "r") as f:
-            data = json.load(f)
-    else:
-        data = {}
-    tmp = [data]
-    dist.broadcast_object_list(tmp, src=0)
-    return tmp[0]
-
-
-def save_torch(obj: Any, path: Union[str, Path]):
-    torch.save(obj, str(path))
-
-
-def load_torch(path: Union[str, Path], broadcast: bool = False) -> Any:
-    if not broadcast or not dist.is_initialized():
-        return torch.load(str(path), map_location="cpu", weights_only=False)
-
-    path = Path(path)
-    rank = get_rank()
-
-    if rank == 0:
-        with open(path, "rb") as f:
-            raw = f.read()
-        data_tensor = torch.frombuffer(bytearray(raw), dtype=torch.uint8)
-        num_bytes = torch.tensor([len(raw)], dtype=torch.long)
-    else:
-        num_bytes = torch.tensor([0], dtype=torch.long)
-
-    dist.broadcast(num_bytes, src=0)
-
-    if rank != 0:
-        data_tensor = torch.empty(num_bytes.item(), dtype=torch.uint8)
-
-    dist.broadcast(data_tensor, src=0)
-
-    buf = io.BytesIO(data_tensor.numpy().tobytes())
-    return torch.load(buf, map_location="cpu", weights_only=False)
-
-
-def save_model(config: dict, state_dict: dict, save_directory: str):
-    save_path = Path(save_directory)
-    save_path.mkdir(parents=True, exist_ok=True)
-    save_json(config, save_path / _CONFIG_FILE)
-    save_safetensors(state_dict, save_path / _WEIGHTS_FILE)
-
-
-def load_model_config(save_directory: str) -> dict:
-    return load_json(Path(save_directory) / _CONFIG_FILE)
-
-
-def load_model_weights(save_directory: str) -> dict:
-    return load_state_dict(Path(save_directory) / _WEIGHTS_FILE)
-
-
-def load_state_dict(path: Union[str, Path], broadcast: bool = False) -> dict:
-    path = Path(path)
-    if not broadcast or not dist.is_initialized():
-        return load_safetensors(path)
-
-    rank = get_rank()
-    if rank == 0:
-        state_dict = load_safetensors(path)
-        specs = [
-            (k, list(state_dict[k].shape), str(state_dict[k].dtype).split(".")[-1])
-            for k in sorted(state_dict)
-        ]
-    else:
-        state_dict = {}
-        specs = []
-
-    specs_list = [specs]
-    dist.broadcast_object_list(specs_list, src=0)
-    specs = specs_list[0]
-
-    for key, shape, dtype_name in specs:
-        dtype = getattr(torch, dtype_name)
-        if rank != 0:
-            tensor = torch.empty(shape, dtype=dtype, device="cpu")
-        else:
-            tensor = state_dict[key].contiguous().cpu()
-        dist.broadcast(tensor, src=0)
-        if rank != 0:
-            state_dict[key] = tensor
-    return state_dict
-
-
-@dataclass
 class Checkpoint:
-    state_dict: Dict[str, Any] = field(default_factory=dict)
-    epoch: int = 0
-    iteration: int = 0
-    extra: Dict[str, Any] = field(default_factory=dict)
-    meta: Dict[str, Any] = field(default_factory=dict)
-    config: Dict[str, Any] = field(default_factory=dict)
+    def __init__(
+        self,
+        state_dict: Dict[str, Any],
+        epoch: int = 0,
+        iteration: int = 0,
+        extra: Optional[Dict[str, Any]] = None,
+        meta: Optional[Dict[str, Any]] = None,
+    ):
+        self.state_dict = state_dict
+        self.epoch = epoch
+        self.iteration = iteration
+        self.extra = extra or {}
+        self.meta = meta or {}
+
+    def save(
+        self,
+        save_dir: str,
+    ) -> None:

-    def save(self, save_dir: str):
        save_path = Path(save_dir)
        save_path.mkdir(parents=True, exist_ok=True)

-        if get_rank() != 0:
-            return
+        rank = get_rank()
+        if rank == 0:
+            meta = {
+                "epoch": self.epoch,
+                "iteration": self.iteration,
+                "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
+            }
+            meta.update(self.meta)
+            with open(save_path / "meta.json", "w") as f:
+                json.dump(meta, f, indent=2)

-        meta = {
-            "epoch": self.epoch,
-            "iteration": self.iteration,
-            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
-            **self.meta,
-        }
-        save_json(meta, save_path / _META_FILE)
-        save_json(self.config, save_path / _CONFIG_FILE)
-        save_safetensors(self.state_dict, save_path / _WEIGHTS_FILE)
-        for key, value in self.extra.items():
-            save_torch(value, save_path / f"{key}.pt")
+            st.save_file(self.state_dict, save_path / "state_dict.safetensors")
+            if self.extra:
+                for key, value in self.extra.items():
+                    torch.save(value, save_path / f"{key}.pt")

    @classmethod
-    def load(cls, save_dir: str, broadcast: bool = False) -> "Checkpoint":
+    def load(
+        cls,
+        save_dir: str,
+    ) -> "Checkpoint":
+
+        rank = get_rank()
        save_path = Path(save_dir)

-        meta = load_json(save_path / _META_FILE, broadcast)
-        config = load_json(save_path / _CONFIG_FILE, broadcast)
-        state_dict = load_state_dict(save_path / _WEIGHTS_FILE, broadcast=broadcast)
+        meta = {}
+        if rank == 0:
+            with open(Path(save_dir) / "meta.json", "r") as f:
+                meta = json.load(f)
+
+        if dist.is_initialized():
+            meta_list = [meta]
+            dist.broadcast_object_list(meta_list, src=0)
+            meta = meta_list[0]
+
+        state_dict = st.load_file(save_path / "state_dict.safetensors")

        extra = {}
-        for f in sorted(save_path.iterdir()):
-            if f.suffix == ".pt":
-                extra[f.stem] = load_torch(f, broadcast=broadcast)
+        for f in save_path.iterdir():
+            if f.suffix == ".pt" and f.stem not in ("meta",):
+                extra[f.stem] = torch.load(f, map_location="cpu", weights_only=False)

        return cls(
            state_dict=state_dict,
-            epoch=meta.get("epoch", 0),
-            iteration=meta.get("iteration", 0),
-            extra=extra,
-            config=config,
+            epoch=meta["epoch"],
+            iteration=meta["iteration"],
+            extra=extra or None,
        )
@@ -42,7 +42,7 @@ class SchedulerFactory(BaseFactory["BaseScheduler"]):
    """

    @classmethod
-    def _validate_component(cls, scheduler_cls: Type[BaseScheduler]):
+    def _validate_component(cls, scheduler_cls: Type[BaseScheduler]) -> None:
        """Validate that the scheduler class inherits from BaseScheduler."""
        if not issubclass(scheduler_cls, BaseScheduler):
            raise TypeError(f"{scheduler_cls.__name__} must inherit from BaseScheduler")
@@ -125,7 +125,7 @@ class StrategyFactory(BaseFactory["BaseStrategy"]):
    """

    @classmethod
-    def _validate_component(cls, strategy_cls: type):
+    def _validate_component(cls, strategy_cls: type) -> None:
        """Validate that the strategy class inherits from BaseStrategy."""
        if not issubclass(strategy_cls, BaseStrategy):
            raise TypeError(f"{strategy_cls.__name__} must inherit from BaseStrategy")
@@ -15,7 +15,7 @@ from tqdm import tqdm

 from astrai.factory import BaseFactory
 from astrai.parallel import only_on_rank
-from astrai.parallel.setup import get_current_device, get_rank
+from astrai.parallel.setup import get_current_device
 from astrai.serialization import Checkpoint
 from astrai.trainer.metric_util import (
    ctx_get_grad_max,
@@ -137,32 +137,44 @@ class CheckpointCallback(TrainCallback):
        save_dir: str,
        interval: int,
        weight_only: bool = False,
+        state_dict_fn: Optional[Callable[[nn.Module], dict]] = None,
        save_extra_fn: Optional[Callable[["TrainContext"], dict]] = None,
+        load_extra_fn: Optional[Callable[[dict, "TrainContext"], None]] = None,
    ):
        self.save_dir = save_dir
        self.interval = interval
        self.weight_only = weight_only
+        self.state_dict_fn = state_dict_fn
        self.save_extra_fn = save_extra_fn or CheckpointCallback.save_extra
+        self.load_extra_fn = load_extra_fn or CheckpointCallback.load_extra
        self.last_ckpt_iter = 0

+    @only_on_rank(0)
    def _save_checkpoint(self, context: TrainContext):
-        unwrapped = context.executor.unwrap_model(context.model)
-        state_dict = unwrapped.state_dict()
+        save_path = os.path.join(
+            self.save_dir, f"epoch_{context.epoch}_iter_{context.iteration}"
+        )
+        state_dict = (
+            self.state_dict_fn(context.model)
+            if self.state_dict_fn
+            else context.model.state_dict()
+        )
+
+        extra = self.save_extra_fn(context)
+        context.checkpoint = Checkpoint(
+            state_dict=state_dict,
+            epoch=context.epoch,
+            iteration=context.iteration,
+            extra=extra,
+            meta=context.config.to_dict(),
+        )
+
+        context.checkpoint.save(save_path)
        self.last_ckpt_iter = context.iteration

-        if get_rank() == 0:
-            save_path = os.path.join(
-                self.save_dir, f"epoch_{context.epoch}_iter_{context.iteration}"
-            )
-            extra = self.save_extra_fn(context)
-            context.checkpoint = Checkpoint(
-                state_dict=state_dict,
-                epoch=context.epoch,
-                iteration=context.iteration,
-                extra=extra,
-                config=context.model_config,
-            )
-            context.checkpoint.save(save_path)
+    def on_train_begin(self, context: TrainContext):
+        if context.checkpoint and context.checkpoint.extra:
+            self.load_extra_fn(context.checkpoint.extra, context)

    def on_batch_end(self, context: TrainContext):
        if context.iteration - self.last_ckpt_iter >= self.interval:
@@ -184,6 +196,12 @@ class CheckpointCallback(TrainCallback):
                extra[name] = obj.state_dict()
        return extra

+    @staticmethod
+    def load_extra(extra: dict, context: TrainContext):
+        for name in CheckpointCallback.extra_keys:
+            if name in extra:
+                getattr(context, name).load_state_dict(extra[name])
+

@CallbackFactory.register("progress_bar")
 class ProgressBarCallback(TrainCallback):
@@ -192,7 +210,7 @@ class ProgressBarCallback(TrainCallback):
    """

    def __init__(
-        self, num_epoch: int, log_interval: int = 100, file: Optional[IO[str]] = None
+        self, num_epoch: int, log_interval: int = 100, file: IO[str] = sys.stdout
    ):
        self.num_epoch = num_epoch
        self.log_interval = log_interval
@@ -205,7 +223,7 @@ class ProgressBarCallback(TrainCallback):
            context.dataloader,
            desc=f"Epoch {context.epoch + 1}/{self.num_epoch}",
            dynamic_ncols=True,
-            file=self.file or sys.stdout,
+            file=self.file,
        )

    @only_on_rank(0)
@@ -1,5 +1,4 @@
 from dataclasses import dataclass, field
-from pathlib import Path
 from typing import Optional, Self

 import torch.nn as nn
@@ -7,11 +6,10 @@ from torch.utils.data import DataLoader

 from astrai.config.train_config import TrainConfig
 from astrai.dataset import ResumableDistributedSampler
-from astrai.model.components.lora import inject_lora
 from astrai.parallel.executor import BaseExecutor, ExecutorFactory
 from astrai.parallel.setup import get_current_device, get_rank, get_world_size
 from astrai.protocols import OptimizerProtocol, SchedulerProtocol
-from astrai.serialization import Checkpoint, load_json, load_model_weights
+from astrai.serialization import Checkpoint
 from astrai.trainer.strategy import BaseStrategy, StrategyFactory


@@ -24,7 +22,6 @@ class TrainContext:
    scheduler: SchedulerProtocol = field(default=None)
    checkpoint: Checkpoint = field(default=None)
    config: TrainConfig = field(default=None)
-    model_config: dict = field(default_factory=dict)
    executor: BaseExecutor = field(default=None)

    epoch: int = field(default=0)
@@ -44,10 +41,10 @@ class TrainContextBuilder:
        config: TrainConfig,
    ):
        self.config = config
-        self._resume_dir: Optional[str] = None
+        self._checkpoint: Optional[Checkpoint] = None

-    def with_resume_dir(self, resume_dir: Optional[str]) -> Self:
-        self._resume_dir = resume_dir
+    def with_checkpoint(self, checkpoint: Optional[Checkpoint]) -> Self:
+        self._checkpoint = checkpoint
        return self

    def build(self) -> TrainContext:
@@ -60,52 +57,27 @@ class TrainContextBuilder:
            **cfg.executor_kwargs,
        )

-        model = cfg.model_fn()
-        model = model.to(device=device)
-
-        model_config = {}
-        if self._resume_dir:
-            config_path = Path(self._resume_dir) / "config.json"
-            if config_path.exists():
-                model_config = load_json(config_path)
-
-        if not model_config and hasattr(model, "config"):
-            model_config = model.config.to_dict()
-
        context = TrainContext(
-            model=model,
+            model=cfg.model,
            world_size=get_world_size(),
            rank=get_rank(),
            config=cfg,
-            model_config=model_config,
            executor=executor,
        )

-        if self._resume_dir is not None:
-            resume_path = Path(self._resume_dir)
-            if (resume_path / "meta.json").exists():
-                checkpoint = Checkpoint.load(self._resume_dir)
-                state_dict = checkpoint.state_dict
-                if checkpoint.config:
-                    context.model_config = checkpoint.config
-            else:
-                checkpoint = None
-                state_dict = load_model_weights(self._resume_dir)
-            model.load_state_dict(state_dict, strict=False)
-            if checkpoint is not None:
-                context.epoch = cfg.start_epoch
-                context.iteration = cfg.start_batch
-            context.checkpoint = checkpoint
+        context.model = context.model.to(device=device)

-        if cfg.lora is not None:
-            inject_lora(
-                model,
-                r=cfg.lora.r,
-                alpha=cfg.lora.alpha,
-                target_modules=set(cfg.lora.target_modules),
+        if self._checkpoint is not None:
+            context.epoch = max(self._checkpoint.epoch, cfg.start_epoch)
+            context.iteration = max(self._checkpoint.iteration, cfg.start_batch)
+            context.model.load_state_dict(self._checkpoint.state_dict)
+            context.checkpoint = self._checkpoint
+        else:
+            context.checkpoint = Checkpoint(
+                state_dict=context.model.state_dict(),
            )

-        context.optimizer = cfg.optimizer_fn(model)
+        context.optimizer = cfg.optimizer_fn(context.model)
        context.scheduler = cfg.scheduler_fn(context.optimizer)

        sampler_offset = context.iteration * cfg.batch_per_device
@@ -143,21 +115,13 @@ class TrainContextBuilder:

        context.model, context.optimizer, context.dataloader, context.scheduler = (
            executor.prepare(
-                model,
+                context.model,
                context.optimizer,
                context.dataloader,
                context.scheduler,
            )
        )

-        if context.checkpoint and context.checkpoint.extra:
-            extra = context.checkpoint.extra
-            for name in ("optimizer", "scheduler"):
-                if name in extra:
-                    obj = getattr(context, name, None)
-                    if obj is not None:
-                        obj.load_state_dict(extra[name])
-
        context.strategy = StrategyFactory.create(
            model=context.model,
            train_type=cfg.strategy,
@@ -3,6 +3,7 @@ from typing import List, Optional

 from astrai.config import TrainConfig
 from astrai.parallel.setup import spawn_parallel_fn
+from astrai.serialization import Checkpoint
 from astrai.trainer.train_callback import (
    CallbackFactory,
    TrainCallback,
@@ -53,9 +54,9 @@ class Trainer:
            if method:
                method(context)

-    def _trainer_loop(self, resume_dir: Optional[str] = None):
+    def _trainer_loop(self, checkpoint: Optional[Checkpoint] = None):
        context = (
-            TrainContextBuilder(self.train_config).with_resume_dir(resume_dir).build()
+            TrainContextBuilder(self.train_config).with_checkpoint(checkpoint).build()
        )
        executor = context.executor
        self._call_callbacks("on_train_begin", context)
@@ -83,19 +84,19 @@ class Trainer:
                            context.optimizer.step()
                            context.optimizer.zero_grad()

-                            if context.scheduler:
-                                context.scheduler.step()
+                        if context.scheduler:
+                            context.scheduler.step()

                self._call_callbacks("on_epoch_end", context)

        except Exception as e:
-            logger.error("Training failed: %s", str(e), exc_info=True)
+            logger.error(f"Training failed: {str(e)}", exc_info=True)
            self._call_callbacks("on_error", context)
            raise
        finally:
            self._call_callbacks("on_train_end", context)

-    def train(self, resume_dir: Optional[str] = None):
+    def train(self, checkpoint: Optional[Checkpoint] = None):
        cfg = self.train_config
        spawn_parallel_fn(
            self._trainer_loop,
@@ -105,5 +106,5 @@ class Trainer:
            master_port=cfg.master_port,
            device_type=cfg.device_type,
            start_method=cfg.start_method,
-            resume_dir=resume_dir,
+            checkpoint=checkpoint,
        )
@@ -1,279 +0,0 @@
-"""MMLU evaluation via log-likelihood ranking."""
-
-import argparse
-import csv
-import json
-import os
-import shutil
-import urllib.request
-import zipfile
-
-import torch
-import torch.nn.functional as F
-import tqdm
-
-from astrai.model import AutoModel
-from astrai.tokenize import AutoTokenizer
-
-MMLU_URL = "https://github.com/hendrycks/test/archive/refs/heads/master.zip"
-MMLU_SUBJECTS = [
-    "abstract_algebra",
-    "anatomy",
-    "astronomy",
-    "business_ethics",
-    "clinical_knowledge",
-    "college_biology",
-    "college_chemistry",
-    "college_computer_science",
-    "college_mathematics",
-    "college_medicine",
-    "college_physics",
-    "computer_security",
-    "conceptual_physics",
-    "econometrics",
-    "electrical_engineering",
-    "elementary_mathematics",
-    "formal_logic",
-    "global_facts",
-    "high_school_biology",
-    "high_school_chemistry",
-    "high_school_computer_science",
-    "high_school_european_history",
-    "high_school_geography",
-    "high_school_government_and_politics",
-    "high_school_macroeconomics",
-    "high_school_mathematics",
-    "high_school_microeconomics",
-    "high_school_physics",
-    "high_school_psychology",
-    "high_school_statistics",
-    "high_school_us_history",
-    "high_school_world_history",
-    "human_aging",
-    "human_sexuality",
-    "international_law",
-    "jurisprudence",
-    "logical_fallacies",
-    "machine_learning",
-    "management",
-    "marketing",
-    "medical_genetics",
-    "miscellaneous",
-    "moral_disputes",
-    "moral_scenarios",
-    "nutrition",
-    "philosophy",
-    "prehistory",
-    "professional_accounting",
-    "professional_law",
-    "professional_medicine",
-    "professional_psychology",
-    "public_relations",
-    "security_studies",
-    "sociology",
-    "us_foreign_policy",
-    "virology",
-    "world_religions",
-]
-
-
-def _download_and_extract(url: str, data_dir: str):
-    zip_path = os.path.join(data_dir, "mmlu.zip")
-    os.makedirs(data_dir, exist_ok=True)
-    print(f"Downloading MMLU data from {url}...")
-    urllib.request.urlretrieve(url, zip_path)
-    print("Extracting...")
-    with zipfile.ZipFile(zip_path, "r") as zf:
-        zf.extractall(data_dir)
-    os.remove(zip_path)
-
-
-def download_mmlu(data_dir: str):
-    _download_and_extract(MMLU_URL, data_dir)
-    src = os.path.join(data_dir, "test-master", "data")
-    if os.path.exists(src):
-        for item in os.listdir(src):
-            os.rename(os.path.join(src, item), os.path.join(data_dir, item))
-        shutil.rmtree(os.path.join(data_dir, "test-master"))
-    print(f"MMLU data saved to {data_dir}")
-
-
-def _strip_prefix(text: str, prefix: str) -> str:
-    if text.startswith(prefix):
-        return text[len(prefix) :].strip()
-    return text
-
-
-def load_csv(path: str) -> list[dict]:
-    data = []
-    with open(path, "r", encoding="utf-8") as f:
-        for row in csv.reader(f):
-            if len(row) < 6:
-                continue
-            if row[0].strip().lower() == "question":
-                continue
-            data.append(
-                {
-                    "question": row[0].strip(),
-                    "A": _strip_prefix(row[1].strip(), "A)"),
-                    "B": _strip_prefix(row[2].strip(), "B)"),
-                    "C": _strip_prefix(row[3].strip(), "C)"),
-                    "D": _strip_prefix(row[4].strip(), "D)"),
-                    "answer": row[5].strip(),
-                }
-            )
-    return data
-
-
-def build_prompt(
-    question: str, choices: dict, subject: str, n_shot: int, dev_data: list[dict]
-) -> str:
-    prompt = ""
-    if n_shot > 0 and dev_data:
-        prompt = f"The following are multiple choice questions (with answers) about {subject}.\n\n"
-        for item in dev_data[:n_shot]:
-            prompt += f"Question: {item['question']}\n"
-            for k in ("A", "B", "C", "D"):
-                prompt += f"{k}. {item[k]}\n"
-            prompt += f"Answer: {item['answer']}\n\n"
-    prompt += f"Question: {question}\n"
-    for k in ("A", "B", "C", "D"):
-        prompt += f"{k}. {choices[k]}\n"
-    prompt += "Answer:"
-    return prompt
-
-
-def choice_logprob(
-    model, tokenizer, context_ids: list[int], choice_letter: str, device: str
-) -> float:
-    choice_text = f" {choice_letter}"
-    choice_ids = tokenizer.encode(choice_text, add_special_tokens=False)
-    input_ids = context_ids + choice_ids
-    max_len = model.config.max_len
-    if len(input_ids) > max_len:
-        overflow = len(input_ids) - max_len
-        input_ids = input_ids[overflow:]
-        ctx_len = len(input_ids) - len(choice_ids)
-    else:
-        ctx_len = len(context_ids)
-
-    input_tensor = torch.tensor([input_ids], device=device, dtype=torch.long)
-    with torch.inference_mode():
-        logits = model(input_tensor)["logits"][0]
-
-    score = 0.0
-    for i, tid in enumerate(choice_ids):
-        pos = ctx_len - 1 + i
-        if pos >= len(logits):
-            break
-        score += F.log_softmax(logits[pos], dim=-1)[tid].item()
-    return score
-
-
-def evaluate_subject(
-    model,
-    tokenizer,
-    subject: str,
-    test_data: list[dict],
-    dev_data: list[dict] | None,
-    device: str,
-    n_shot: int,
-) -> tuple[float, int, int]:
-    correct = 0
-    total = 0
-    for item in tqdm.tqdm(test_data, desc=f"{subject:40s}", leave=False):
-        prompt = build_prompt(item["question"], item, subject, n_shot, dev_data or [])
-        context_ids = tokenizer.encode(prompt)
-        scores = {
-            c: choice_logprob(model, tokenizer, context_ids, c, device)
-            for c in ("A", "B", "C", "D")
-        }
-        if max(scores, key=scores.get) == item["answer"]:
-            correct += 1
-        total += 1
-    return correct / total, correct, total
-
-
-def main():
-    parser = argparse.ArgumentParser(description="MMLU evaluation")
-    parser.add_argument(
-        "--param_path", type=str, default="./params", help="Model directory"
-    )
-    parser.add_argument(
-        "--data_dir", type=str, default="./mmlu_data", help="MMLU data directory"
-    )
-    parser.add_argument("--download", action="store_true", help="Download MMLU data")
-    parser.add_argument(
-        "--n_shot", type=int, default=5, help="Few-shot examples (0 for zero-shot)"
-    )
-    parser.add_argument(
-        "--subjects", type=str, nargs="+", help="Specific subjects (default: all)"
-    )
-    parser.add_argument("--output", type=str, help="Output JSON path")
-    parser.add_argument("--split", type=str, default="test", choices=["test", "val"])
-    parser.add_argument(
-        "--device",
-        type=str,
-        default="cuda" if torch.cuda.is_available() else "cpu",
-        help="Device",
-    )
-    parser.add_argument(
-        "--dtype",
-        type=str,
-        default="bfloat16" if torch.cuda.is_available() else "float32",
-        help="Torch dtype",
-    )
-    args = parser.parse_args()
-
-    if args.download or not os.path.exists(args.data_dir):
-        download_mmlu(args.data_dir)
-
-    model = AutoModel.from_pretrained(args.param_path)
-    tokenizer = AutoTokenizer.from_pretrained(args.param_path)
-    device = args.device
-    dtype = getattr(torch, args.dtype)
-    model.to(device=device, dtype=dtype)
-
-    subjects = args.subjects or MMLU_SUBJECTS
-    results = {}
-    total_correct = 0
-    total_questions = 0
-
-    for subject in subjects:
-        dev_path = os.path.join(args.data_dir, "dev", f"{subject}_dev.csv")
-        test_path = os.path.join(
-            args.data_dir, args.split, f"{subject}_{args.split}.csv"
-        )
-
-        if not os.path.exists(test_path):
-            print(f"  Skipping {subject}: test file not found")
-            continue
-
-        dev_data = load_csv(dev_path) if os.path.exists(dev_path) else None
-        test_data = load_csv(test_path)
-
-        acc, corr, tot = evaluate_subject(
-            model, tokenizer, subject, test_data, dev_data, device, args.n_shot
-        )
-        results[subject] = {"accuracy": round(acc, 4), "correct": corr, "total": tot}
-        total_correct += corr
-        total_questions += tot
-        print(f"  {subject:40s}  {acc:.2%}  ({corr}/{tot})")
-
-    overall = total_correct / total_questions if total_questions else 0
-    print(f"\n{'=' * 70}")
-    print(f"  Overall: {overall:.2%}  ({total_correct}/{total_questions})")
-    results["_overall"] = {
-        "accuracy": round(overall, 4),
-        "correct": total_correct,
-        "total": total_questions,
-    }
-
-    if args.output:
-        with open(args.output, "w", encoding="utf-8") as f:
-            json.dump(results, f, indent=2)
-        print(f"Results saved to {args.output}")
-
-
-if __name__ == "__main__":
-    main()
@@ -10,11 +10,11 @@ from astrai.tokenize import AutoTokenizer


 def process_file(
-    param_path: str, input_file: str, output_file: str, batch_size: int, text_key: str
+    model_dir: str, input_file: str, output_file: str, batch_size: int, text_key: str
 ):
    # Load model and tokenizer
-    model = AutoModel.from_pretrained(param_path)
-    tokenizer = AutoTokenizer.from_pretrained(param_path)
+    model = AutoModel.from_pretrained(model_dir)
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model.to(device="cuda", dtype=torch.bfloat16)

    with open(input_file, "r", encoding="utf-8") as f:
@@ -44,8 +44,8 @@ def process_file(

        for seq in batch_encoded:
            pad_len = max_len - len(seq)
-            padded_seq = seq + [tokenizer.pad_id] * pad_len
-            mask = [True] * len(seq) + [False] * pad_len
+            padded_seq = [tokenizer.pad_id] * pad_len + seq
+            mask = [False] * pad_len + [True] * len(seq)
            padded_ids.append(padded_seq)
            masks.append(mask)

@@ -88,7 +88,7 @@ def process_file(
 if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run perplexity with a Khaosz model.")
    parser.add_argument(
-        "--param_path", type=str, required=True, help="Path to the model directory."
+        "--model_dir", type=str, required=True, help="Path to the model directory."
    )
    parser.add_argument(
        "--input_file", type=str, required=True, help="Path to the input file."
@@ -18,7 +18,7 @@ def main():
        "--reload", action="store_true", help="Enable auto-reload for development"
    )
    parser.add_argument(
-        "--param_path",
+        "--param-path",
        type=Path,
        default=None,
        help="Path to model parameters (default: project_root/params)",
@@ -2,6 +2,7 @@ import argparse
 import os
 from functools import partial

+import safetensors.torch as st
 import torch
 import torch.optim as optim

@@ -146,8 +147,8 @@ def parse_args() -> argparse.Namespace:
        "--parallel_mode",
        type=str,
        default="none",
-        choices=["none", "ddp", "fsdp"],
-        help="Parallel training strategy (none, ddp, fsdp).",
+        choices=["none", "ddp"],
+        help="Parallel training strategy.",
    )
    parser.add_argument(
        "--device_type", type=str, default="cuda", help="Device type to use."
@@ -165,10 +166,6 @@ def parse_args() -> argparse.Namespace:
    return args


-def create_model(config):
-    return AutoRegressiveLM(config).to(dtype=torch.bfloat16)
-
-
 def create_optimizer(model, **kwargs) -> optim.Optimizer:
    return optim.AdamW(model.parameters(), fused=True, **kwargs)

@@ -231,8 +228,6 @@ def train(
 ):
    assert train_type in ["seq", "sft", "dpo", "grpo"]
    assert os.path.exists(param_path)
-    if nprocs > 1 and parallel_mode == "none":
-        raise ValueError("--nprocs > 1 requires --parallel_mode to be 'ddp' or 'fsdp'")

    # Load config
    config_path = os.path.join(param_path, "config.json")
@@ -241,6 +236,17 @@ def train(
    if window_size is None:
        window_size = config.max_len

+    # Create bare AutoRegressiveLM (for training, no tokenizer needed)
+    model = AutoRegressiveLM(config)
+
+    # Load weights if available
+    weights_path = os.path.join(param_path, "model.safetensors")
+    if os.path.exists(weights_path):
+        state_dict = st.load_file(weights_path)
+        model.load_state_dict(state_dict, strict=False)
+
+    model = model.to(dtype=torch.bfloat16)
+
    strategy_kwargs = {
        "beta": dpo_beta,
        "label_smoothing": label_smoothing,
@@ -251,11 +257,12 @@ def train(
    }

    executor_kwargs = {
+        "static_graph": True,
+        "find_unused_parameters": False,
        "gradient_as_bucket_view": True,
        "broadcast_buffers": False,
    }

-    model_fn = partial(create_model, config)
    dataset = DatasetFactory.load(
        train_type=train_type,
        load_path=data_root_path,
@@ -287,7 +294,7 @@ def train(
    )

    train_config = TrainConfig(
-        model_fn=model_fn,
+        model=model,
        strategy=train_type,
        dataset=dataset,
        optimizer_fn=optimizer_fn,
@@ -312,7 +319,7 @@ def train(
    )

    trainer = Trainer(train_config)
-    trainer.train(resume_dir=param_path)
+    trainer.train()


 if __name__ == "__main__":
@@ -1,4 +1,3 @@
-import os
 import tempfile

 import torch
@@ -37,6 +36,7 @@ def test_single_process():


 def test_checkpoint_with_extra():
+    """Verify extra keys are saved as individual .pt files and loaded back."""
    model = torch.nn.Linear(10, 5)
    optimizer = AdamW(model.parameters(), lr=1e-3)
    optimizer.step()
@@ -52,6 +52,8 @@ def test_checkpoint_with_extra():
    with tempfile.TemporaryDirectory() as tmpdir:
        checkpoint.save(tmpdir)

+        import os
+
        assert os.path.exists(os.path.join(tmpdir, "optimizer.pt"))
        assert os.path.exists(os.path.join(tmpdir, "scheduler.pt"))

@@ -7,12 +7,12 @@ import torch

 from astrai.dataset.dataset import DatasetFactory, SEQDataset
 from astrai.dataset.storage import (
-    H5Store,
-    MmapStore,
-    StoreFactory,
+    BaseSegmentFetcher,
+    H5Storage,
+    MultiSegmentFetcher,
+    StorageFactory,
    detect_format,
-    load_bin,
-    save_bin,
+    load_json,
    save_h5,
 )

@@ -157,6 +157,111 @@ def test_dataset_with_custom_stride(base_test_env):
    assert len(dataset) > len(default_stride_dataset)


+# ============== JSON Storage Tests (raw text + tokenizer) ==============
+
+
+def _make_tokenizer_fn(tokenizer):
+    """Wrap tokenizer.encode() as a str -> List[int] callable."""
+    return lambda text: tokenizer.encode(text, add_special_tokens=False)
+
+
+def test_seq_dataset_from_json_text(base_test_env):
+    """Test loading SEQ dataset from raw-text JSON with tokenizer"""
+    tokenizer = base_test_env["tokenizer"]
+    tokenizer_fn = _make_tokenizer_fn(tokenizer)
+    test_dir = base_test_env["test_dir"]
+    data_dir = os.path.join(test_dir, "json_text")
+    os.makedirs(data_dir, exist_ok=True)
+
+    texts = [
+        "hello world this is a test sentence for tokenizer",
+        "another sentence with different words and tokens",
+        "machine learning is fascinating and powerful",
+    ]
+
+    json_path = os.path.join(data_dir, "seq_data.json")
+    with open(json_path, "w", encoding="utf-8") as f:
+        json.dump({"sequence": texts}, f, ensure_ascii=False)
+
+    dataset = DatasetFactory.load(
+        train_type="seq",
+        load_path=data_dir,
+        window_size=16,
+        tokenizer=tokenizer_fn,
+    )
+    assert dataset is not None
+    assert len(dataset) > 0
+    assert dataset.count > 0
+    assert "sequence" in dataset.keys
+
+    item = dataset[0]
+    assert "input_ids" in item
+    assert "target_ids" in item
+    assert item["input_ids"].shape[0] == 16
+
+
+def test_sft_dataset_from_json_text(base_test_env):
+    """Test loading SFT dataset from raw-text JSON with tokenizer"""
+    tokenizer = base_test_env["tokenizer"]
+    tokenizer_fn = _make_tokenizer_fn(tokenizer)
+    test_dir = base_test_env["test_dir"]
+    data_dir = os.path.join(test_dir, "json_sft")
+    os.makedirs(data_dir, exist_ok=True)
+
+    texts = [
+        "user asks a question about the weather",
+        "assistant provides a helpful response to the user",
+    ]
+
+    json_path = os.path.join(data_dir, "sft_data.json")
+    with open(json_path, "w", encoding="utf-8") as f:
+        json.dump(
+            {"sequence": texts, "loss_mask": texts},
+            f,
+            ensure_ascii=False,
+        )
+
+    dataset = DatasetFactory.load(
+        train_type="sft",
+        load_path=data_dir,
+        window_size=16,
+        tokenizer=tokenizer_fn,
+    )
+    assert dataset is not None
+    assert len(dataset) > 0
+
+    item = dataset[0]
+    assert "loss_mask" in item
+
+
+def test_json_storage_explicit_tokenizer(base_test_env):
+    """Test explicit JSON storage with tokenizer"""
+    tokenizer = base_test_env["tokenizer"]
+    tokenizer_fn = _make_tokenizer_fn(tokenizer)
+    test_dir = base_test_env["test_dir"]
+    data_dir = os.path.join(test_dir, "json_explicit")
+    os.makedirs(data_dir, exist_ok=True)
+
+    texts = ["abcdefghijklmnopqrstuvwxyz" * 10]
+
+    json_path = os.path.join(data_dir, "data.json")
+    with open(json_path, "w", encoding="utf-8") as f:
+        json.dump({"sequence": texts}, f, ensure_ascii=False)
+
+    token_count = len(tokenizer_fn(texts[0]))
+
+    dataset = DatasetFactory.load(
+        train_type="seq",
+        load_path=data_dir,
+        window_size=32,
+        storage_type="json",
+        tokenizer=tokenizer_fn,
+    )
+    assert dataset is not None
+    assert len(dataset) > 0
+    assert dataset.count == token_count
+
+
 def test_dataset_count_property(base_test_env):
    """Test the count property returns correct raw token count"""
    test_dir = base_test_env["test_dir"]
@@ -213,29 +318,37 @@ def test_unloaded_dataset_len():
    assert len(dataset) == 0


-def test_store_unloaded_len():
-    """Unloaded Store has __len__ == 0"""
-    store = H5Store()
-    assert len(store) == 0
-    assert store.keys == []
+def test_base_segment_fetcher_empty():
+    """BaseSegmentFetcher with empty segments list"""
+    fetcher = BaseSegmentFetcher([])
+    assert len(fetcher) == 0
+    with pytest.raises(ValueError, match="out of bounds"):
+        fetcher.fetch_data(0, 1)


-def test_store_fetch_begin_equals_end(base_test_env):
-    """Store.fetch with begin == end returns empty tensor"""
+def test_base_segment_fetcher_begin_equals_end(base_test_env):
+    """fetch_data with begin == end returns empty tensor"""
    test_dir = base_test_env["test_dir"]
    dummy = {"sequence": [torch.randint(0, 1000, (100,), dtype=torch.int64)]}
    save_h5(test_dir, "empty_fetch", dummy)

    dataset = DatasetFactory.load("seq", test_dir, window_size=32)
-    result = dataset.storage.fetch(10, 10, "sequence")
+    fetcher = dataset.storage._fetcher.multi_fetchers["sequence"]
+    result = fetcher.fetch_data(10, 10)
    assert result.numel() == 0


-def test_store_fetch_before_load():
-    """Store.fetch before load raises RuntimeError"""
-    store = H5Store()
+def test_multi_segment_fetcher_empty_dict():
+    """MultiSegmentFetcher with empty dict has __len__ == 0"""
+    fetcher = MultiSegmentFetcher({})
+    assert len(fetcher) == 0
+
+
+def test_storage_fetch_before_load():
+    """BaseStorage.fetch before load raises RuntimeError"""
+    storage = H5Storage()
    with pytest.raises(RuntimeError, match="not loaded"):
-        store.fetch(0, 10, "sequence")
+        storage.fetch(0, 10, "sequence")


 def test_detect_format_nonexistent_path():
@@ -254,192 +367,54 @@ def test_detect_format_unsupported_file(base_test_env):
        detect_format(path)


-def test_create_store_invalid_type():
-    """StoreFactory.create raises ValueError for unknown type"""
+def test_create_storage_invalid_type():
+    """StorageFactory.create raises ValueError for unknown type"""
    with pytest.raises(ValueError, match="Unknown component"):
-        StoreFactory.create("parquet")
+        StorageFactory.create("parquet")


-def test_store_multi_segment_concat(base_test_env):
-    """Multi-segment H5 data is concatenated into single tensor at load time"""
-    import os
-
+def test_json_pretokenized_without_tokenizer(base_test_env):
+    """Pre-tokenized JSON (List[List[int]]) loads without tokenizer"""
    test_dir = base_test_env["test_dir"]
-    data_dir = os.path.join(test_dir, "multi_seg")
+    data_dir = os.path.join(test_dir, "json_pretok")
    os.makedirs(data_dir, exist_ok=True)

+    json_path = os.path.join(data_dir, "data.json")
+    with open(json_path, "w", encoding="utf-8") as f:
+        json.dump({"sequence": [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]}, f)
+
+    dataset = DatasetFactory.load("seq", data_dir, window_size=4, storage_type="json")
+    assert len(dataset) > 0
+    assert dataset.count == 10
+
+    item = dataset[0]
+    assert item["input_ids"].tolist() == [1, 2, 3, 4]
+    assert item["target_ids"].tolist() == [2, 3, 4, 5]
+
+
+def test_load_json_skips_config_file(base_test_env):
+    """load_json skips scalar-value config files"""
+    test_dir = base_test_env["test_dir"]
+    with open(os.path.join(test_dir, "config.json"), "w") as f:
+        json.dump({"vocab_size": 1000, "dim": 16}, f)
+
+    with open(os.path.join(test_dir, "data.json"), "w") as f:
+        json.dump({"sequence": [[1, 2, 3, 4, 5]]}, f)
+
+    result = load_json(test_dir)
+    assert "sequence" in result
+    assert "vocab_size" not in result
+    assert len(result["sequence"]) == 1
+
+
+def test_base_segment_fetcher_multi_segment():
+    """fetch_data across multiple segment boundaries"""
    segs = [
        torch.tensor([1, 2, 3]),
        torch.tensor([4, 5, 6, 7]),
        torch.tensor([8, 9]),
    ]
-    save_h5(data_dir, "data", {"sequence": segs})
-
-    store = StoreFactory.create("h5")
-    store.load(data_dir)
-    assert len(store) == 9
-    result = store.fetch(2, 7, "sequence")
+    fetcher = BaseSegmentFetcher(segs)
+    assert len(fetcher) == 9
+    result = fetcher.fetch_data(2, 7)
    assert result.tolist() == [3, 4, 5, 6, 7]
-
-
-def test_save_load_bin_roundtrip(base_test_env):
-    """save_bin + load_bin roundtrip preserves data"""
-    test_dir = base_test_env["test_dir"]
-
-    data = {
-        "sequence": [torch.tensor([1, 2, 3, 4, 5], dtype=torch.int64)],
-        "loss_mask": [torch.tensor([0, 1, 1, 0, 1], dtype=torch.int64)],
-    }
-    save_bin(test_dir, data)
-    result = load_bin(test_dir)
-
-    assert "sequence" in result
-    assert "loss_mask" in result
-    assert result["sequence"][0].tolist() == [1, 2, 3, 4, 5]
-    assert result["loss_mask"][0].tolist() == [0, 1, 1, 0, 1]
-
-
-def test_mmap_store_load_and_fetch(base_test_env):
-    """MmapStore loads bin data and fetches correctly"""
-    test_dir = base_test_env["test_dir"]
-
-    data = {
-        "sequence": [torch.randint(0, 1000, (200,), dtype=torch.int64)],
-    }
-    save_bin(test_dir, data)
-
-    store = StoreFactory.create("bin")
-    store.load(test_dir)
-    assert len(store) == 200
-    assert "sequence" in store.keys
-
-    result = store.fetch(10, 20, "sequence")
-    assert result.tolist() == data["sequence"][0][10:20].tolist()
-
-
-def test_mmap_dataset_load(base_test_env):
-    """DatasetFactory.load auto-detects bin format"""
-    test_dir = base_test_env["test_dir"]
-
-    data = {
-        "sequence": [torch.randint(0, 1000, (200,), dtype=torch.int64)],
-    }
-    save_bin(test_dir, data)
-
-    dataset = DatasetFactory.load("seq", test_dir, window_size=64)
-    assert len(dataset) > 0
-    assert dataset.count == 200
-    assert dataset[0]["input_ids"].shape[0] == 64
-
-
-def test_normalize_empty_key():
-    """_normalize with empty tensor list does not crash"""
-    store = H5Store()
-    store._normalize({"sequence": []})
-    assert len(store) == 0
-    assert store.keys == ["sequence"]
-
-
-def test_normalize_mixed_empty_key():
-    """_normalize with empty + non-empty keys returns min=0"""
-    store = H5Store()
-    store._normalize({"sequence": [torch.tensor([1, 2, 3])], "loss_mask": []})
-    assert len(store) == 0
-    assert set(store.keys) == {"sequence", "loss_mask"}
-
-
-def test_grpo_dataset_dtype(base_test_env):
-    """GRPODataset returns correct dtypes"""
-    test_dir = base_test_env["test_dir"]
-
-    seq_len = 100
-    data = {
-        "prompts": [torch.randint(0, 100, (seq_len,), dtype=torch.int32)],
-        "responses": [torch.randint(0, 100, (seq_len,), dtype=torch.int32)],
-        "masks": [torch.ones(seq_len, dtype=torch.int32)],
-        "rewards": [torch.ones(seq_len, dtype=torch.float32)],
-    }
-    save_h5(test_dir, "grpo_dtype", data)
-
-    dataset = DatasetFactory.load("grpo", test_dir, window_size=32)
-    item = dataset[0]
-
-    assert item["prompts"].dtype == torch.long
-    assert item["responses"].dtype == torch.long
-    assert item["masks"].dtype == torch.bool
-    assert item["rewards"].dtype == torch.float32
-
-
-def test_grpo_dataset_load(base_test_env):
-    """GRPODataset loads and returns correct keys"""
-    test_dir = base_test_env["test_dir"]
-    seq_len = 200
-    data = {
-        "prompts": [torch.randint(0, 1000, (seq_len,), dtype=torch.int64)],
-        "responses": [torch.randint(0, 1000, (seq_len,), dtype=torch.int64)],
-        "masks": [torch.ones(seq_len, dtype=torch.int64)],
-        "rewards": [torch.rand(seq_len, dtype=torch.float32)],
-    }
-    save_h5(test_dir, "grpo_test", data)
-
-    dataset = DatasetFactory.load("grpo", test_dir, window_size=64)
-    assert len(dataset) > 0
-    item = dataset[0]
-    assert "prompts" in item
-    assert "responses" in item
-    assert "masks" in item
-    assert "rewards" in item
-    assert item["prompts"].shape[0] == 64
-    assert item["responses"].shape[0] == 64
-
-
-def test_detect_format_bin_dir(base_test_env):
-    """detect_format returns 'bin' for directory with .bin + meta.json"""
-    test_dir = base_test_env["test_dir"]
-    save_bin(test_dir, {"sequence": [torch.randint(0, 100, (10,))]})
-    assert detect_format(test_dir) == "bin"
-
-
-def test_store_fetch_multi_key(base_test_env):
-    """Store.fetch with List[str] returns Dict[str, Tensor]"""
-    test_dir = base_test_env["test_dir"]
-    save_h5(
-        test_dir,
-        "multi_key",
-        {
-            "sequence": [torch.randint(0, 100, (100,), dtype=torch.int64)],
-            "loss_mask": [torch.ones(100, dtype=torch.int64)],
-        },
-    )
-
-    store = StoreFactory.create("h5")
-    store.load(test_dir)
-    result = store.fetch(10, 20, ["sequence", "loss_mask"])
-    assert isinstance(result, dict)
-    assert result["sequence"].shape[0] == 10
-    assert result["loss_mask"].shape[0] == 10
-
-
-def test_store_fetch_out_of_bounds(base_test_env):
-    """Store.fetch raises ValueError for out-of-bounds indices"""
-    test_dir = base_test_env["test_dir"]
-    save_h5(test_dir, "bounds", {"sequence": [torch.randint(0, 100, (50,))]})
-
-    store = StoreFactory.create("h5")
-    store.load(test_dir)
-    with pytest.raises(ValueError, match="out of bounds"):
-        store.fetch(-1, 10, "sequence")
-    with pytest.raises(ValueError, match="out of bounds"):
-        store.fetch(0, 51, "sequence")
-    with pytest.raises(ValueError, match="out of bounds"):
-        store.fetch(50, 50, "sequence")
-
-
-def test_dataset_load_explicit_storage_type(base_test_env):
-    """DatasetFactory.load with explicit storage_type bypasses auto-detect"""
-    test_dir = base_test_env["test_dir"]
-    save_h5(test_dir, "explicit", {"sequence": [torch.randint(0, 100, (200,))]})
-
-    dataset = DatasetFactory.load("seq", test_dir, window_size=64, storage_type="h5")
-    assert len(dataset) > 0
-    assert dataset.count == 200
@@ -1,286 +0,0 @@
-"""Unit tests for protocol builders, StopChecker, GenContext, StopInfo."""
-
-import json
-from unittest.mock import MagicMock
-
-import pytest
-
-from astrai.inference.api.anthropic import AnthropicResponseBuilder
-from astrai.inference.api.openai import OpenAIResponseBuilder
-from astrai.inference.api.protocol import GenContext, StopChecker, StopInfo
-from astrai.inference.engine import GenerationRequest
-
-
-def _make_ctx(**kwargs):
-    defaults = {
-        "resp_id": "test-123",
-        "created": 1000,
-        "model": "test-model",
-        "prompt_tokens": 10,
-        "completion_tokens": 5,
-    }
-    defaults.update(kwargs)
-    return GenContext(**defaults)
-
-
-def _sse_payloads(events):
-    payloads = []
-    for chunk in events:
-        for line in chunk.strip().split("\n"):
-            if line.startswith("data: "):
-                try:
-                    payloads.append(json.loads(line[6:]))
-                except json.JSONDecodeError:
-                    pass
-    return payloads
-
-
-class TestStopChecker:
-    def test_check_finds_match(self):
-        sc = StopChecker(["stop", "end"])
-        assert sc.check("hello stop world") == "stop"
-
-    def test_check_returns_none_when_no_match(self):
-        sc = StopChecker(["stop"])
-        assert sc.check("hello world") is None
-
-    def test_check_empty_sequences(self):
-        sc = StopChecker([])
-        assert sc.check("hello") is None
-
-
-class TestGenContext:
-    def test_defaults(self):
-        ctx = GenContext(resp_id="a", created=1, model="m", prompt_tokens=10)
-        assert ctx.completion_tokens == 0
-
-    def test_fields_mutable(self):
-        ctx = GenContext(resp_id="a", created=1, model="m", prompt_tokens=10)
-        ctx.completion_tokens = 42
-        assert ctx.completion_tokens == 42
-
-
-class TestStopInfo:
-    def test_defaults(self):
-        s = StopInfo()
-        assert s.matched is None
-        assert s.body == ""
-        assert s.yielded == ""
-
-    def test_with_values(self):
-        s = StopInfo(matched="stop", body="hello stop", yielded="hello ")
-        assert s.matched == "stop"
-        assert s.body == "hello stop"
-        assert s.yielded == "hello "
-
-
-class TestOpenAIResponseBuilder:
-    @pytest.fixture
-    def builder(self):
-        builder = OpenAIResponseBuilder()
-        req = MagicMock()
-        req.messages = [MagicMock(role="user", content="Hello")]
-        req.stop = None
-        req.model = "astrai"
-        engine = MagicMock()
-        engine.tokenizer.apply_chat_template.return_value = "Hello"
-        builder.prepare(req, engine)
-        return builder
-
-    def test_prepare_returns_prompt_ctx_stops(self, builder):
-        req = MagicMock()
-        req.messages = [MagicMock(role="user", content="Hi")]
-        req.stop = ["END"]
-        req.model = "gpt"
-        engine = MagicMock()
-        engine.tokenizer.apply_chat_template.return_value = "Hi"
-        prompt, ctx, stops = builder.prepare(req, engine)
-        assert prompt == "Hi"
-        assert ctx.model == "gpt"
-        assert ctx.prompt_tokens == 0
-        assert stops == ["END"]
-
-    def test_prepare_no_stop_returns_empty_list(self, builder):
-        req = MagicMock()
-        req.messages = []
-        req.stop = None
-        req.model = "x"
-        engine = MagicMock()
-        engine.tokenizer.apply_chat_template.return_value = ""
-        _, _, stops = builder.prepare(req, engine)
-        assert stops == []
-
-    def test_format_stream_start(self, builder):
-        ctx = _make_ctx()
-        events = builder.format_stream_start(ctx)
-        payloads = _sse_payloads(events)
-        assert len(payloads) == 1
-        p = payloads[0]
-        assert p["object"] == "chat.completion.chunk"
-        assert p["choices"][0]["delta"]["role"] == "assistant"
-        assert p["choices"][0]["finish_reason"] is None
-
-    def test_format_chunk(self, builder):
-        event = builder.format_chunk("hello")
-        payload = json.loads(event.split("data: ", 1)[1])
-        assert payload["choices"][0]["delta"]["content"] == "hello"
-        assert payload["choices"][0]["finish_reason"] is None
-
-    def test_format_stream_end(self, builder):
-        ctx = _make_ctx(completion_tokens=5)
-        stop = StopInfo(matched="stop")
-        events = builder.format_stream_end(ctx, stop)
-        payloads = _sse_payloads(events)
-        finish = payloads[0]
-        assert finish["choices"][0]["finish_reason"] == "stop"
-        usage = payloads[1]
-        assert usage["completion_tokens"] == 5
-        assert usage["total_tokens"] == 15
-
-    def test_format_response(self, builder):
-        ctx = _make_ctx()
-        stop = StopInfo()
-        resp = builder.format_response(ctx, "hello", stop)
-        assert resp["object"] == "chat.completion"
-        assert resp["choices"][0]["message"]["content"] == "hello"
-        assert resp["usage"]["prompt_tokens"] == 10
-
-
-class TestAnthropicResponseBuilder:
-    @pytest.fixture
-    def builder(self):
-        builder = AnthropicResponseBuilder()
-        req = MagicMock()
-        req.messages = [MagicMock(role="user", content="Hello")]
-        req.model = "claude"
-        engine = MagicMock()
-        engine.tokenizer.apply_chat_template.return_value = "Hello"
-        req.system = None
-        builder.prepare(req, engine)
-        return builder
-
-    def test_prepare_messages(self, builder):
-        req = MagicMock()
-        req.messages = [MagicMock(role="user", content="Hi")]
-        req.model = "claude"
-        req.system = None
-        req.stop_sequences = None
-        engine = MagicMock()
-        engine.tokenizer.apply_chat_template.return_value = "Hi"
-        prompt, ctx, stops = builder.prepare(req, engine)
-        assert prompt == "Hi"
-        assert stops == []
-
-    def test_prepare_with_stop_sequences(self, builder):
-        req = MagicMock()
-        req.messages = []
-        req.model = "x"
-        req.stop_sequences = ["stop", "end"]
-        req.system = None
-        engine = MagicMock()
-        engine.tokenizer.apply_chat_template.return_value = ""
-        _, _, stops = builder.prepare(req, engine)
-        assert stops == ["stop", "end"]
-
-    def test_format_stream_start(self, builder):
-        ctx = _make_ctx(prompt_tokens=3)
-        events = builder.format_stream_start(ctx)
-        payloads = _sse_payloads(events)
-        assert len(payloads) == 2
-        assert payloads[0]["type"] == "message_start"
-        assert payloads[0]["message"]["usage"]["input_tokens"] == 3
-        assert payloads[1]["type"] == "content_block_start"
-
-    def test_format_chunk(self, builder):
-        event = builder.format_chunk("tok")
-        payload = json.loads(event.split("data: ", 1)[1])
-        assert payload["type"] == "content_block_delta"
-        assert payload["delta"]["text"] == "tok"
-
-    def test_format_stream_end_no_stop(self, builder):
-        ctx = _make_ctx(completion_tokens=3)
-        stop = StopInfo()
-        events = builder.format_stream_end(ctx, stop)
-        payloads = _sse_payloads(events)
-        # content_block_stop, message_delta, message_stop
-        types = [p["type"] for p in payloads]
-        assert types == ["content_block_stop", "message_delta", "message_stop"]
-        assert payloads[1]["delta"]["stop_reason"] == "end_turn"
-
-    def test_format_stream_end_with_stop_trims_and_emits_remaining(self, builder):
-        ctx = _make_ctx(completion_tokens=7)
-        stop = StopInfo(
-            matched="END",
-            body="Hello world END extra",
-            yielded="Hello ",
-        )
-        events = builder.format_stream_end(ctx, stop)
-        payloads = _sse_payloads(events)
-        # unyielded delta, content_block_stop, message_delta, message_stop
-        types = [p["type"] for p in payloads]
-        assert types == [
-            "content_block_delta",
-            "content_block_stop",
-            "message_delta",
-            "message_stop",
-        ]
-        assert payloads[0]["delta"]["text"] == "world "
-        assert payloads[2]["delta"]["stop_reason"] == "stop_sequence"
-        assert payloads[2]["delta"]["stop_sequence"] == "END"
-
-    def test_format_stream_end_stop_trimmed_already_yielded(self, builder):
-        ctx = _make_ctx()
-        stop = StopInfo(
-            matched="END",
-            body="Hello END",
-            yielded="Hello ",
-        )
-        events = builder.format_stream_end(ctx, stop)
-        payloads = _sse_payloads(events)
-        # No unyielded delta (everything already sent)
-        types = [p["type"] for p in payloads]
-        assert types == ["content_block_stop", "message_delta", "message_stop"]
-
-    def test_format_response_with_stop_trims_content(self, builder):
-        ctx = _make_ctx()
-        stop = StopInfo(matched="STOP", body="text STOP extra", yielded="text ")
-        resp = builder.format_response(ctx, "text STOP extra", stop)
-        assert resp["content"][0]["text"] == "text "
-        assert resp["stop_reason"] == "stop_sequence"
-        assert resp["stop_sequence"] == "STOP"
-
-    def test_format_response_no_stop(self, builder):
-        ctx = _make_ctx()
-        stop = StopInfo()
-        resp = builder.format_response(ctx, "full text", stop)
-        assert resp["content"][0]["text"] == "full text"
-        assert resp["stop_reason"] == "end_turn"
-
-
-class TestGenerationRequestValidation:
-    def test_valid_params(self):
-        gr = GenerationRequest(
-            messages=[{"role": "user", "content": "hi"}],
-            top_k=50,
-            top_p=0.9,
-            temperature=0.7,
-        )
-        assert gr.top_k == 50
-
-    def test_invalid_top_p_raises(self):
-        with pytest.raises(ValueError, match="top_p"):
-            GenerationRequest(messages=[{"role": "user", "content": "hi"}], top_p=1.5)
-
-    def test_invalid_top_k_raises(self):
-        with pytest.raises(ValueError, match="top_k"):
-            GenerationRequest(messages=[{"role": "user", "content": "hi"}], top_k=-1)
-
-    def test_invalid_temperature_raises(self):
-        with pytest.raises(ValueError, match="temperature"):
-            GenerationRequest(
-                messages=[{"role": "user", "content": "hi"}], temperature=-0.1
-            )
-
-    def test_top_k_zero_valid(self):
-        gr = GenerationRequest(messages=[{"role": "user", "content": "hi"}], top_k=0)
-        assert gr.top_k == 0
@@ -173,21 +173,3 @@ def test_scheduler_concurrent_get_stats(mock_model_and_tokenizer):
    for stats in results["stats"]:
        assert "total_tasks" in stats
        assert stats["total_tasks"] >= 0
-
-
-def test_prefill_skips_fully_cached_tasks(mock_model_and_tokenizer):
-    """Tasks whose entire prompt is cached skip the prefill phase."""
-    mock_model, mock_tokenizer = mock_model_and_tokenizer
-
-    with patch("astrai.inference.core.scheduler.AutoModel"):
-        with patch("astrai.inference.core.scheduler.AutoTokenizer"):
-            scheduler = InferenceScheduler(
-                model=mock_model,
-                tokenizer=mock_tokenizer,
-                max_batch_size=4,
-                device="cpu",
-            )
-
-    task_id = scheduler.add_task("short prompt", stream_callback=lambda t: None)
-    scheduler.stop()
-    assert task_id.startswith("task_")
@@ -5,17 +5,17 @@ import torch

 from astrai.config.model_config import AutoRegressiveLMConfig
 from astrai.model import AutoRegressiveLM
-from astrai.model.components.linear import Linear
 from astrai.model.components.lora import (
    LoRAConfig,
    LoRALinear,
-    _collect_lora_info,
    _get_lora_count,
+    _collect_lora_info,
    inject_lora,
    load_lora,
    merge_lora,
    save_lora,
 )
+from astrai.model.components.linear import Linear

 MODEL_KWARGS = dict(
    vocab_size=1000,
@@ -27,7 +27,7 @@ class TrainerDataset(Dataset):


 def create_train_config(
-    model_fn,
+    model: torch.nn.Module,
    dataset: Dataset,
    test_dir: str,
    device: str,
@@ -43,7 +43,7 @@ def create_train_config(
    """Factory function to create common TrainConfig for tests.

    Args:
-        model_fn: Model factory (callable returning nn.Module)
+        model: The model to train
        dataset: Training dataset
        test_dir: Checkpoint directory
        device: Device type ("cuda" or "cpu")
@@ -70,7 +70,7 @@ def create_train_config(

    return TrainConfig(
        strategy=strategy,
-        model_fn=model_fn,
+        model=model,
        dataset=dataset,
        optimizer_fn=optimizer_fn,
        scheduler_fn=scheduler_fn,
@@ -106,7 +106,7 @@ def test_gradient_checkpointing_trainer_integration(base_test_env, random_datase
        )

    train_config = TrainConfig(
-        model_fn=lambda: base_test_env["model"],
+        model=base_test_env["model"],
        strategy="seq",
        dataset=random_dataset,
        optimizer_fn=optimizer_fn,
@@ -140,7 +140,7 @@ def test_callback_integration(base_test_env, random_dataset):
        )

    train_config = TrainConfig(
-        model_fn=lambda: base_test_env["model"],
+        model=base_test_env["model"],
        strategy="seq",
        dataset=random_dataset,
        optimizer_fn=optimizer_fn,
@@ -4,6 +4,7 @@ import numpy as np
 import torch

 from astrai.config.train_config import TrainConfig
+from astrai.serialization import Checkpoint
 from astrai.trainer.schedule import SchedulerFactory
 from astrai.trainer.trainer import Trainer

@@ -23,7 +24,7 @@ def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
        strategy="seq",
        optimizer_fn=optimizer_fn,
        scheduler_fn=scheduler_fn,
-        model_fn=lambda: base_test_env["model"],
+        model=base_test_env["model"],
        dataset=early_stopping_dataset,
        ckpt_dir=base_test_env["test_dir"],
        log_dir=os.path.join(base_test_env["test_dir"], "logs"),
@@ -38,20 +39,17 @@ def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
    trainer = Trainer(train_config)

    # Should handle early stopping gracefully
+    checkpoint = None
    try:
-        trainer.train()
+        checkpoint = trainer.train()
    except Exception:
+        # Handle any exceptions
        pass

-    # Resume from latest checkpoint
    load_dir = os.path.join(base_test_env["test_dir"], "epoch_0_iter_2")
-    trainer = Trainer(train_config)
-    trainer.train(resume_dir=load_dir)
+    checkpoint = Checkpoint.load(load_dir)
+    trainer.train(checkpoint)

-    # Verify checkpoint was saved at expected iteration
    load_dir = os.path.join(base_test_env["test_dir"], "epoch_1_iter_10")
-    import json
-
-    with open(os.path.join(load_dir, "meta.json")) as f:
-        meta = json.load(f)
-    assert meta["iteration"] == 10
+    checkpoint = Checkpoint.load(load_dir)
+    assert checkpoint.iteration == 10
@@ -9,7 +9,7 @@ def test_different_batch_sizes(base_test_env, random_dataset, train_config_facto

    for batch_per_device in batch_sizes:
        train_config = train_config_factory(
-            model_fn=lambda: base_test_env["model"],
+            model=base_test_env["model"],
            dataset=random_dataset,
            test_dir=base_test_env["test_dir"],
            device=base_test_env["device"],
@@ -25,7 +25,7 @@ def test_gradient_accumulation(base_test_env, random_dataset, train_config_facto

    for grad_accum_steps in grad_accum_steps_list:
        train_config = train_config_factory(
-            model_fn=lambda: base_test_env["model"],
+            model=base_test_env["model"],
            dataset=random_dataset,
            test_dir=base_test_env["test_dir"],
            device=base_test_env["device"],
@@ -50,7 +50,7 @@ def test_memory_efficient_training(base_test_env, random_dataset, train_config_f

    for config in small_batch_configs:
        train_config = train_config_factory(
-            model_fn=lambda: base_test_env["model"],
+            model=base_test_env["model"],
            dataset=random_dataset,
            test_dir=base_test_env["test_dir"],
            device=base_test_env["device"],