fix: 修复 to_dict list 类型丢失与 OpenAI stop 参数失效

- to_dict() 增加 list 类型序列化支持，metrics 等字段不再丢失 - OpenAIHandler 补充 get_stop_sequences/on_token，读取 request.stop 并检测停止序列 - 文档类图补充缺失字段、修正关系分类、ChatCompletionRequest 字段增加 Optional
feat: ProgressBarCallback 支持日志行输出到 stdout
2026-05-19 21:07:07 +08:00 · 2026-05-19 19:12:38 +08:00 · 2026-05-19 17:50:24 +08:00 · 2026-05-19 17:34:01 +08:00 · 2026-05-18 14:24:00 +08:00 · 2026-05-18 09:40:32 +08:00
22 changed files with 483 additions and 96 deletions
@@ -1,7 +1,7 @@
 # AstrAI Dockerfile - Multi-stage Build (Optimized)

 # Build stage - use base image with minimal build tools
-FROM nvidia/cuda:12.6.0-base-ubuntu24.04 AS builder
+FROM ubuntu:24.04 AS builder

 WORKDIR /app

@@ -18,7 +18,7 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-ins
 RUN python3.12 -m venv --copies /opt/venv
 ENV PATH="/opt/venv/bin:$PATH"

-# Copy source code and install dependencies
+# Copy source code and install (deps read from pyproject.toml)
 COPY astrai/ ./astrai/
 COPY pyproject.toml .
 RUN pip install --no-cache-dir --upgrade pip \
@@ -26,13 +26,14 @@ RUN pip install --no-cache-dir --upgrade pip \
    --extra-index-url https://download.pytorch.org/whl/cu126

 # Production stage
-FROM nvidia/cuda:12.6.0-base-ubuntu24.04 AS production
+FROM ubuntu:24.04 AS production

 WORKDIR /app

-# Install Python 3.12 runtime
+# Install Python 3.12 runtime and healthcheck dependency
 RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
    python3.12 \
+    curl \
    && rm -rf /var/lib/apt/lists/*

 # Copy virtual environment from builder
@@ -82,7 +82,7 @@ export CUDA_VISIBLE_DEVICES=0,1,2,3

 nohup python scripts/tools/train.py \
    --nprocs=4 \
-    --train_type=pt \
+    --train_type=seq \
    --data_root_path=/path/to/dataset \
    --param_path=/path/to/model \
    --batch_per_device=4 \
@@ -90,8 +90,8 @@ nohup python scripts/tools/train.py \
    --warmup_ratio=0.05 \
    --max_lr=1e-4 \
    --max_grad_norm=1.0 \
-    --adamw_beta1=0.95 \
-    --adamw_beta2=0.99 \
+    --adamw_beta1=0.9 \
+    --adamw_beta2=0.95 \
    --adamw_weight_decay=0.01 \
    --window_size=2048 \
    --ckpt_interval=10000 \
@@ -213,7 +213,7 @@ python scripts/demo/generate_batch.py
 python scripts/demo/generate_ar.py
 ```

-Watch a video walkthrough on [bilibili](https://www.bilibili.com/video/BV1z5RPYHEkd).
+Watch a video walkthrough on [bilibili](https://www.bilibili.com/video/BV1fuLB6yEj6).

 ### Documentation

@@ -88,7 +88,7 @@ export CUDA_VISIBLE_DEVICES=0,1,2,3

 nohup python scripts/tools/train.py \
    --nprocs=4 \
-    --train_type=pt \
+    --train_type=seq \
    --data_root_path=/path/to/dataset \
    --param_path=/path/to/model \
    --batch_per_device=4 \
@@ -96,8 +96,8 @@ nohup python scripts/tools/train.py \
    --warmup_ratio=0.05 \
    --max_lr=1e-4 \
    --max_grad_norm=1.0 \
-    --adamw_beta1=0.95 \
-    --adamw_beta2=0.99 \
+    --adamw_beta1=0.9 \
+    --adamw_beta2=0.95 \
    --adamw_weight_decay=0.01 \
    --window_size=2048 \
    --ckpt_interval=10000 \
@@ -219,7 +219,7 @@ python scripts/demo/generate_batch.py
 python scripts/demo/generate_ar.py
 ```

-观看 [bilibili](https://www.bilibili.com/video/BV1z5RPYHEkd) 上的视频演示。
+观看 [bilibili](https://www.bilibili.com/video/BV1fuLB6yEj6) 上的视频演示。

 ### 文档

@@ -16,7 +16,7 @@ classDiagram
            +to_file(config_path)
        }

-        class ModelConfig {
+        class AutoRegressiveLMConfig {
            +int vocab_size
            +int dim
            +int n_layers
@@ -25,21 +25,41 @@ classDiagram
            +bool tie_weight
            +int max_len
            +float rope_theta
+            +str attn_type
            +int n_heads
            +int n_kv_heads
            +bool use_qk_norm
            +bool use_gated_attention
-            +str attn_type
+            +Optional[int] kv_lora_rank
+            +Optional[int] qk_nope_head_dim
+            +Optional[int] qk_rope_head_dim
            +str ffn_type
            +int n_routed_experts
            +int n_shared_experts
            +int n_activated_experts
-            +str moe_topk_method
-            +Optional[int] kv_lora_rank
-            +Optional[int] qk_nope_head_dim
-            +Optional[int] qk_rope_head_dim
-            +load(config_path) ModelConfig
-            +save(config_path)
+            +Optional[str] topk_method
+        }
+
+        class EncoderConfig {
+            +int vocab_size
+            +int dim
+            +int n_layers
+            +float norm_eps
+            +int dim_ffn
+            +int max_len
+            +float rope_theta
+            +int n_heads
+            +int n_kv_heads
+            +bool use_qk_norm
+            +bool use_gated_attention
+            +Optional[str] pooling_type
+            +Optional[bool] normalize_embeddings
+        }
+
+        class ConfigFactory {
+            +Registry _registry
+            +register(name) decorator
+            +load(raw) BaseConfig
        }

        class TrainConfig {
@@ -52,10 +72,14 @@ classDiagram
            +int batch_per_device
            +int grad_accum_steps
            +float max_grad_norm
+            +list gradient_checkpointing_modules
            +int start_epoch
            +int start_batch
            +str ckpt_dir
            +int ckpt_interval
+            +str log_dir
+            +int log_interval
+            +List[str] metrics
            +int random_seed
            +int num_workers
            +Optional[int] prefetch_factor
@@ -66,7 +90,10 @@ classDiagram
            +str master_port
            +Callable parallel_wrapper
            +Callable state_dict_fn
+            +str start_method
            +str device_type
+            +Optional[Dataset] val_dataset
+            +int val_step
            +dict extra_kwargs
            +validate()
        }
@@ -138,11 +165,17 @@ classDiagram
            +int iter
        }

+        class StorageFactory {
+            +Registry _registry
+            +register(name) decorator
+            +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) BaseDataset
+            +load(train_type, load_path, window_size, stride, storage_type, tokenizer) BaseDataset
        }
    }

@@ -160,7 +193,7 @@ classDiagram

    namespace model {
        class AutoModel {
-            +ModelConfig config
+            +BaseModelConfig config
            +Registry _registry
            +register(model_type) decorator
            +get_component_class(model_type) Type
@@ -169,8 +202,8 @@ classDiagram
            +to(*args, **kwargs) Self
        }

-        class Transformer {
-            +ModelConfig config
+        class AutoRegressiveLM {
+            +AutoRegressiveLMConfig config
            +RotaryEmbedding rotary_embedding
            +Embedding embed_tokens
            +ModuleList layers
@@ -181,6 +214,18 @@ classDiagram
            +state_dict()
        }

+        class EmbeddingEncoder {
+            +EncoderConfig config
+            +RotaryEmbedding rotary_embedding
+            +Embedding embed_tokens
+            +ModuleList layers
+            +RMSNorm norm
+            +str pooling_type
+            +bool normalize_embeddings
+            +forward(input_ids, input_mask, position_ids) Tensor
+            +load_state_dict(state_dict)
+        }
+
        class DecoderBlock {
            +nn.Module attention  # GQA or MLA via AttnFactory
            +RMSNorm input_norm
@@ -322,11 +367,15 @@ classDiagram
            +Optimizer optimizer
            +LRScheduler scheduler
            +Checkpoint checkpoint
+            +TrainConfig config
            +int epoch
            +int iteration
            +float loss
+            +DataLoader val_dataloader
+            +float val_loss
            +int world_size
            +int rank
+            +dict kwargs
        }

        class TrainContextBuilder {
@@ -372,6 +421,7 @@ classDiagram
            +str reduction
            +int sync_interval
            +compute_loss(batch) Tensor
+            +sync_ref_model()
        }

        class BaseScheduler {
@@ -399,6 +449,7 @@ classDiagram
        }

        class TrainCallback {
+            <<protocol>>
            +on_train_begin(context)
            +on_train_end(context)
            +on_epoch_begin(context)
@@ -415,17 +466,32 @@ classDiagram
            +on_step_begin(context)
        }

+        class GradientCheckpointingCallback {
+            +tuple modules
+            +on_train_begin(context)
+            +on_train_end(context)
+        }
+
        class CheckpointCallback {
            +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 {
            +int num_epoch
+            +int log_interval
+            +IO file
            +on_epoch_begin(context)
            +on_batch_end(context)
            +on_epoch_end(context)
@@ -434,8 +500,16 @@ classDiagram
        class MetricLoggerCallback {
            +str log_dir
            +int save_interval
+            +int log_interval
+            +List[str] metrics
            +on_batch_end(context)
            +on_train_end(context)
+            +on_error(context)
+        }
+
+        class ValidationCallback {
+            +_run_validation(context)
+            +on_step_end(context)
        }

        class CallbackFactory {
@@ -443,6 +517,14 @@ classDiagram
            +register(name) decorator
            +create(name, **kwargs) TrainCallback
        }
+
+        class Muon {
+            +float lr
+            +float momentum
+            +float weight_decay
+            +int ns_steps
+            +step(closure) Optional[float]
+        }
    }

    namespace inference {
@@ -616,7 +698,7 @@ classDiagram
        }

        class SamplingPipeline {
-            +List strategies
+            +List[BaseSamplingStrategy] strategies
            +apply(logits, filter_value) Tensor
            +sample(logits, filter_value) Tensor
        }
@@ -638,14 +720,19 @@ classDiagram
        }

        class ChatCompletionRequest {
+            +str model
            +List[ChatMessage] messages
-            +float temperature
-            +float top_p
-            +int top_k
-            +int max_tokens
-            +bool stream
-            +Optional[str] stop
+            +Optional[float] temperature
+            +Optional[float] top_p
+            +Optional[int] top_k
+            +Optional[int] max_tokens
+            +Optional[bool] stream
+            +Optional[Union[str, List[str]]] stop
            +Optional[int] n
+            +Optional[float] presence_penalty
+            +Optional[float] frequency_penalty
+            +Optional[Dict[int, float]] logit_bias
+            +Optional[str] user
        }

        class AnthropicMessage {
@@ -654,6 +741,7 @@ classDiagram
        }

        class MessagesRequest {
+            +str model
            +List[AnthropicMessage] messages
            +Optional[str] system
            +float temperature
@@ -666,8 +754,13 @@ classDiagram

        class ProtocolHandler {
            <<abstract>>
+            +request
+            +engine
            +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]
@@ -687,6 +780,7 @@ classDiagram
        }

        class StopChecker {
+            +has_sequences (property) bool
            +check(text) Optional[str]
            +trim(text, matched) str
        }
@@ -699,6 +793,7 @@ classDiagram
            +int completion_tokens
            +str accumulated
            +Optional[str] stop_matched
+            +str last_yield_trimmed
        }

        class app {
@@ -709,11 +804,13 @@ classDiagram

    namespace parallel {
        class Functions {
-            +spawn_parallel_fn(func, world_size, backend, master_addr, master_port, device_type, **kwargs)
+            <<module>>
+            +spawn_parallel_fn(func, world_size, backend, master_addr, master_port, device_type, start_method, **kwargs)
            +setup_parallel(rank, world_size, backend, master_addr, master_port, device_type)
            +get_current_device() str
            +get_world_size() int
            +get_rank() int
+            +only_on_rank(rank, sync) decorator
        }

        class ParallelModel {
@@ -741,6 +838,7 @@ classDiagram
    BaseScheduler <|-- CosineScheduler
    BaseScheduler <|-- SGDRScheduler
    TrainCallback <|-- GradientClippingCallback
+    TrainCallback <|-- GradientCheckpointingCallback
    TrainCallback <|-- CheckpointCallback
    TrainCallback <|-- ProgressBarCallback
    TrainCallback <|-- MetricLoggerCallback
@@ -753,12 +851,15 @@ classDiagram
    BaseSamplingStrategy <|-- TemperatureStrategy
    BaseSamplingStrategy <|-- TopKStrategy
    BaseSamplingStrategy <|-- TopPStrategy
+    BaseSamplingStrategy <|-- SamplingPipeline
    ParallelModel <|-- RowParallelLinear
    ParallelModel <|-- ColumnParallelLinear
-    AutoModel <|-- Transformer
+    AutoModel <|-- AutoRegressiveLM
+    AutoModel <|-- EmbeddingEncoder
    BaseConfig <|-- BaseModelConfig
    BaseConfig <|-- TrainConfig
-    BaseModelConfig <|-- ModelConfig
+    BaseModelConfig <|-- AutoRegressiveLMConfig
+    BaseModelConfig <|-- EncoderConfig
    BaseFactory <|-- AutoModel
    BaseFactory <|-- AttnFactory
    BaseFactory <|-- FFNFactory
@@ -766,6 +867,9 @@ classDiagram
    BaseFactory <|-- StrategyFactory
    BaseFactory <|-- SchedulerFactory
    BaseFactory <|-- CallbackFactory
+    BaseFactory <|-- StorageFactory
+    BaseFactory <|-- ConfigFactory
+    TrainCallback <|-- ValidationCallback
    ProtocolHandler <|-- OpenAIHandler
    ProtocolHandler <|-- AnthropicHandler

@@ -773,31 +877,33 @@ classDiagram
    KVCache *-- PagePool
    KVCache *-- Storage
    KVCache *-- TaskTable
-    KVCache *-- Allocator
-    KVCache *-- PrefixCache
+    PagePool *-- Allocator
+    PagePool *-- PrefixCache
    InferenceEngine *-- InferenceScheduler
    InferenceScheduler *-- KVCache
    InferenceScheduler *-- Executor
    InferenceScheduler *-- TaskManager
-    SamplingPipeline *-- BaseSamplingStrategy
-    TrainContextBuilder *-- TrainContext
-    Transformer *-- DecoderBlock
-    Transformer *-- RotaryEmbedding
-    Transformer *-- Embedding
+    AutoRegressiveLM *-- DecoderBlock
+    AutoRegressiveLM *-- RotaryEmbedding
+    AutoRegressiveLM *-- Embedding
+    EmbeddingEncoder *-- DecoderBlock
+    EmbeddingEncoder *-- RotaryEmbedding
+    EmbeddingEncoder *-- Embedding
    DecoderBlock *-- RMSNorm
-    BaseDataset *-- BaseStorage
    ChatCompletionRequest *-- ChatMessage
    MessagesRequest *-- AnthropicMessage
+    AutoTokenizer *-- ChatTemplate
+    BaseFactory *-- Registry

    %% --- Aggregation (weak ownership) ---
-    AutoModel o-- ModelConfig
+    AutoModel o-- BaseModelConfig
    Trainer o-- TrainCallback
    TrainContext o-- BaseStrategy
    TrainContext o-- BaseScheduler
    TrainContext o-- Checkpoint
-    AutoTokenizer o-- ChatTemplate
    KvcacheView o-- Storage
-    BaseFactory o-- Registry
+    SamplingPipeline o-- BaseSamplingStrategy
+    BaseDataset o-- BaseStorage

    %% --- Dependency (uses temporarily) ---
    TrainConfig ..> BaseStrategy : selects
@@ -811,7 +917,12 @@ classDiagram
    FFNFactory ..> DeepSeekMoE : creates
    DecoderBlock ..> AttnFactory : uses
    DecoderBlock ..> FFNFactory : uses
+    StorageFactory ..> H5Storage : creates
+    StorageFactory ..> JSONStorage : creates
+    ConfigFactory ..> AutoRegressiveLMConfig : creates
+    ConfigFactory ..> EncoderConfig : creates
    Trainer ..> TrainContextBuilder : uses
+    TrainContextBuilder ..> TrainContext : creates
    Trainer ..> Functions : spawns
    TrainContextBuilder ..> StrategyFactory : uses
    TrainContextBuilder ..> ResumableDistributedSampler : creates
@@ -827,13 +938,13 @@ classDiagram

    %% --- Association (general usage) ---
    Trainer --> TrainConfig
-    DPOStrategy --> Transformer
-    GRPOStrategy --> Transformer
+    DPOStrategy --> AutoModel
+    GRPOStrategy --> AutoModel
    InferenceScheduler --> Task
    InferenceScheduler --> TaskStatus
    Task --> TaskStatus
-    InferenceEngine --> Transformer
-    Executor --> Transformer
+    InferenceEngine --> AutoModel
+    Executor --> AutoModel
    Executor --> AutoTokenizer
    TaskManager --> AutoTokenizer
    MultiSegmentFetcher --> BaseSegmentFetcher
@@ -846,12 +957,12 @@ classDiagram

 | Module | Components | Description |
 |--------|------------|-------------|
-| **astrai.config** | BaseConfig, BaseModelConfig, ModelConfig, TrainConfig | Configuration management (to_dict/from_dict, to_file/from_file) |
-| **astrai.dataset** | BaseDataset–GRPODataset, BaseStorage–JSONStorage, BaseSegmentFetcher, MultiSegmentFetcher, ResumableDistributedSampler, DatasetFactory | Dataset loading and management |
+| **astrai.config** | BaseConfig, BaseModelConfig, AutoRegressiveLMConfig, EncoderConfig, ConfigFactory, TrainConfig | Configuration management (to_dict/from_dict, to_file/from_file) |
+| **astrai.dataset** | BaseDataset–GRPODataset, BaseStorage–JSONStorage, StorageFactory, BaseSegmentFetcher, MultiSegmentFetcher, ResumableDistributedSampler, DatasetFactory | Dataset loading and management |
 | **astrai.serialization** | Checkpoint | Model serialization |
-| **astrai.model** | AutoModel, Transformer, DecoderBlock, GQA, MLA, MLP, DeepSeekMoE, AttnFactory, FFNFactory, RMSNorm, Linear, RotaryEmbedding, Embedding | Neural network model |
+| **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–MetricLoggerCallback, CallbackFactory | Training workflow |
+| **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, BaseSamplingStrategy–SamplingPipeline, ProtocolHandler–AnthropicHandler, ChatMessage–MessagesRequest, app | Inference service |
 | **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank/get_world_size/get_current_device, only_on_rank, ParallelModel, RowParallelLinear, ColumnParallelLinear | Distributed parallel |
 | **astrai.factory** | Registry, BaseFactory[T] | Component registration |
@@ -860,7 +971,7 @@ classDiagram

 | Pattern | Classes | Purpose |
 |---------|---------|---------|
-| **Factory** | `AttnFactory`, `FFNFactory`, `StrategyFactory`, `DatasetFactory`, `SchedulerFactory`, `CallbackFactory` | Decorator-based component creation |
+| **Factory** | `AttnFactory`, `FFNFactory`, `StrategyFactory`, `DatasetFactory`, `SchedulerFactory`, `CallbackFactory`, `StorageFactory`, `ConfigFactory` | 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 |
@@ -871,18 +982,18 @@ classDiagram
 | **Object Pool** | `Allocator`, `PagePool` | Page-based KV cache with LRU eviction |
 | **Storage** | `BaseStorage`, `H5Storage`, `JSONStorage` | Format-agnostic data access |
 | **Producer-Consumer** | `InferenceScheduler`, `Task`, queues | Continuous batching |
-| **AutoModel Registry** | `AutoModel`, `Transformer` | Model-type dynamic loading |
+| **AutoModel Registry** | `AutoModel`, `AutoRegressiveLM`, `EmbeddingEncoder` | Model-type dynamic loading |

 ## Core Relationships

 1. **Config → Training**: `TrainConfig` holds model, dataset, optimizer_fn, scheduler_fn
 2. **Training Flow**: `Trainer` → `TrainContextBuilder` → `TrainContext`, uses `BaseStrategy` for loss
 3. **Strategy Selection**: `StrategyFactory` creates strategy by `train_type`
-4. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `Transformer`, backed by `KVCache` + `SamplingPipeline`
+4. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `KVCache` + `SamplingPipeline`
 5. **Distributed**: `spawn_parallel_fn` + `setup_parallel` for multi-process DDP
 6. **Dataset Loading**: `DatasetFactory` creates datasets, `BaseStorage` (H5Storage/JSONStorage) loads via `BaseSegmentFetcher` + `MultiSegmentFetcher`
 7. **Checkpoint**: `Checkpoint` saves/loads safetensors + metadata (rank-0 only)
 8. **Scheduler**: `SchedulerFactory` creates `CosineScheduler`/`SGDRScheduler`
 9. **AutoModel**: `from_pretrained()` loads `config.json` + `model.safetensors`, `_disable_random_init` replaces `nn.init.*` with no-ops

-> Document Update Time: 2026-05-16
+> Document Update Time: 2026-05-17
@@ -15,8 +15,8 @@ Raw text is tokenized via `AutoTokenizer.encode()` and saved as HDF5 (`.h5`) or
 Storage format is auto-detected by `detect_format()`; backends are dispatched via registry:

 ```
-create_storage("h5")   → H5Storage
-create_storage("json") → JSONStorage
+StorageFactory.create("h5")   → H5Storage
+StorageFactory.create("json") → JSONStorage
 ```

 Both support shared memory via `.share_memory_()`.
@@ -34,7 +34,7 @@ Both support shared memory via `.share_memory_()`.

 ```
 DatasetFactory.load(train_type, path, window_size, stride)
-  → create_storage(detect_format(path))
+  → StorageFactory.create(detect_format(path))
    → MultiSegmentFetcher(BaseSegmentFetcher per key)
      → BaseDataset.__getitem__(idx)
        → sliding window [begin, end) via get_index(idx)
@@ -54,4 +54,4 @@ DatasetFactory.load(train_type, path, window_size, stride)

 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-15
+> Document Update Time: 2026-05-17
@@ -137,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-15
+> Document Update Time: 2026-05-17
@@ -25,8 +25,8 @@

 | Parameter | Description | Default |
 |-----------|-------------|---------|
-| `--adamw_beta1` | AdamW beta1 | 0.95 |
-| `--adamw_beta2` | AdamW beta2 | 0.99 |
+| `--adamw_beta1` | AdamW beta1 | 0.9 |
+| `--adamw_beta2` | AdamW beta2 | 0.95 |
 | `--adamw_weight_decay` | AdamW weight decay | 0.01 |

 ### Data Loading
@@ -73,7 +73,7 @@ export CUDA_VISIBLE_DEVICES=0,1,2,3

 nohup python scripts/tools/train.py \
    --nprocs=4 \
-    --train_type=pt \
+    --train_type=seq \
    --data_root_path=/path/to/dataset \
    --param_path=/path/to/model \
    --batch_per_device=4 \
@@ -81,8 +81,8 @@ nohup python scripts/tools/train.py \
    --warmup_ratio=0.05 \
    --max_lr=1e-4 \
    --max_grad_norm=1.0 \
-    --adamw_beta1=0.95 \
-    --adamw_beta2=0.99 \
+    --adamw_beta1=0.9 \
+    --adamw_beta2=0.95 \
    --adamw_weight_decay=0.01 \
    --window_size=2048 \
    --ckpt_interval=10000 \
@@ -94,4 +94,4 @@ nohup python scripts/tools/train.py \

 ---

-> Document Update Time: 2026-05-16
+> Document Update Time: 2026-05-17
@@ -91,11 +91,13 @@ on_train_end

 | Hook | Fires | Default callback |
 |------|-------|-----------------|
+| `on_train_begin` | Before training starts | `GradientCheckpointingCallback` |
 | `on_step_begin` | Every accumulation window | `GradientClippingCallback` |
 | `on_batch_end` | Every batch | `CheckpointCallback`, `MetricLoggerCallback`, `ProgressBarCallback` |
+| `on_step_end` | Every accumulation window | `ValidationCallback` |
 | `on_train_end` | Training ends | `CheckpointCallback`, `MetricLoggerCallback` (final save) |

-Default callbacks: `progress_bar` (tqdm), `checkpoint` (safetensors, rank-0), `metric_logger` (JSONL, rank-0), `gradient_clipping`.
+Default callbacks: `gradient_checkpointing` (activation checkpointing, optional), `progress_bar` (tqdm), `checkpoint` (safetensors, rank-0), `metric_logger` (JSONL, rank-0), `gradient_clipping`, `validation` (periodic validation on val_dataset).

 ## Strategies

@@ -154,6 +156,17 @@ Keys: `prompts`, `responses`, `masks`, `rewards`.

 Created by `SchedulerFactory.create(optimizer, schedule_type, **kwargs)`.

+## Gradient Checkpointing
+
+Trades compute for memory by recomputing activations during backward pass. Specify module types via `gradient_checkpointing_modules`:
+
+```python
+from astrai.model.components.decoder_block import DecoderBlock
+config = TrainConfig(..., gradient_checkpointing_modules=[DecoderBlock])
+```
+
+Callback wraps each `DecoderBlock.forward` with `torch.utils.checkpoint.checkpoint(use_reentrant=False)`, compatible with `torch.compile`. Uses `nn.Module.apply()` for traversal — works through DDP wrappers without manual unwrap. Empty list (default) means no-op.
+
 ## Checkpoint

 ```
@@ -188,7 +201,7 @@ export CUDA_VISIBLE_DEVICES=0,1,2,3

 nohup python scripts/tools/train.py \
    --nprocs=4 \
-    --train_type=pt \
+    --train_type=seq \
    --data_root_path=/path/to/dataset \
    --param_path=/path/to/model \
    --batch_per_device=4 \
@@ -196,8 +209,8 @@ nohup python scripts/tools/train.py \
    --warmup_ratio=0.05 \
    --max_lr=1e-4 \
    --max_grad_norm=1.0 \
-    --adamw_beta1=0.95 \
-    --adamw_beta2=0.99 \
+    --adamw_beta1=0.9 \
+    --adamw_beta2=0.95 \
    --adamw_weight_decay=0.01 \
    --window_size=2048 \
    --ckpt_interval=10000 \
@@ -209,4 +222,4 @@ nohup python scripts/tools/train.py \

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

-> Document Update Time: 2026-05-16
+> Document Update Time: 2026-05-17
@@ -1,4 +1,4 @@
-__version__ = "1.3.5"
+__version__ = "1.3.6"
 __author__ = "ViperEkura"

 from astrai.config import (
@@ -11,7 +11,6 @@ __all__ = [
    "BaseModelConfig",
    "AutoRegressiveLMConfig",
    "EncoderConfig",
-    "ModelConfig",
    "ConfigFactory",
    "TrainConfig",
 ]
@@ -13,7 +13,7 @@ class BaseConfig:
                d[fld.name] = v
            elif v is None:
                d[fld.name] = None
-            elif isinstance(v, dict):
+            elif isinstance(v, (dict, list)):
                try:
                    json.dumps(v)
                    d[fld.name] = v
@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field, fields
-from typing import Callable, Optional
+from typing import Callable, List, Optional

 import torch.nn as nn
 from torch.optim import Optimizer
@@ -39,6 +39,10 @@ class TrainConfig(BaseConfig):
    max_grad_norm: float = field(
        default=1.0, metadata={"help": "Maximum gradient norm."}
    )
+    gradient_checkpointing_modules: list = field(
+        default_factory=list,
+        metadata={"help": "Module types to enable activation checkpointing for."},
+    )

    # checkpoint setting
    start_epoch: int = field(default=0, metadata={"help": "Start epoch for training."})
@@ -52,6 +56,19 @@ class TrainConfig(BaseConfig):
        default=5000, metadata={"help": "Number of iterations between checkpoints."}
    )

+    # metric setting
+    log_dir: str = field(
+        default="./checkpoint/logs", metadata={"help": "Directory for metric logs."}
+    )
+    log_interval: int = field(
+        default=100,
+        metadata={"help": "Number of batch iterations between metric logs."},
+    )
+    metrics: List[str] = field(
+        default_factory=lambda: ["loss", "lr"],
+        metadata={"help": "Metrics to record during training."},
+    )
+
    # dataloader setting
    random_seed: int = field(default=3407, metadata={"help": "Random seed."})
    num_workers: int = field(
@@ -226,6 +226,17 @@ class OpenAIHandler(ProtocolHandler):
    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(
@@ -163,5 +163,4 @@ def spawn_parallel_fn(
        nprocs=world_size,
        start_method=start_method,
        join=True,
-        daemon=True,
    )
@@ -38,7 +38,7 @@ class Checkpoint:
            meta = {
                "epoch": self.epoch,
                "iteration": self.iteration,
-                "timestamp": time.time(),
+                "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
            }
            meta.update(self.meta)
            with open(save_path / "meta.json", "w") as f:
@@ -1,14 +1,16 @@
 import json
 import logging
 import os
+import sys
 import time
 from pathlib import Path
-from typing import Callable, List, Optional, Protocol, runtime_checkable
+from typing import IO, Callable, List, Optional, Protocol, runtime_checkable

 import torch
 import torch.distributed as dist
 import torch.nn as nn
 from torch.nn.utils import clip_grad_norm_
+from torch.utils.checkpoint import checkpoint as torch_checkpoint
 from tqdm import tqdm

 from astrai.factory import BaseFactory
@@ -90,6 +92,41 @@ class GradientClippingCallback(TrainCallback):
        clip_grad_norm_(context.model.parameters(), self.max_grad_norm)


+@CallbackFactory.register("gradient_checkpointing")
+class GradientCheckpointingCallback(TrainCallback):
+    """
+    Activation checkpointing callback — trades compute for memory
+    by recomputing specified module activations during the backward pass.
+
+    Args:
+        modules: Module types to apply checkpointing to.
+    """
+
+    def __init__(self, modules: Optional[List[type]] = None):
+        self.modules = tuple(modules) if modules else ()
+
+    def _enable(self, module: nn.Module):
+        if self.modules and isinstance(module, self.modules):
+            fn = module.forward
+            module._original_forward = fn
+            module.forward = lambda *a, **kw: torch_checkpoint(
+                fn, *a, use_reentrant=False, **kw
+            )
+
+    @staticmethod
+    def _disable(module: nn.Module):
+        if hasattr(module, "_original_forward"):
+            module.forward = module._original_forward
+            del module._original_forward
+
+    def on_train_begin(self, context: TrainContext):
+        context.model.apply(self._enable)
+        logger.info("Gradient checkpointing enabled")
+
+    def on_train_end(self, context: TrainContext):
+        context.model.apply(self._disable)
+
+
@CallbackFactory.register("checkpoint")
 class CheckpointCallback(TrainCallback):
    """
@@ -175,8 +212,12 @@ class ProgressBarCallback(TrainCallback):
    Progress bar callback for trainer.
    """

-    def __init__(self, num_epoch: int):
+    def __init__(
+        self, num_epoch: int, log_interval: int = 100, file: IO[str] = sys.stdout
+    ):
        self.num_epoch = num_epoch
+        self.log_interval = log_interval
+        self.file = file
        self.progress_bar: tqdm = None

    @only_on_rank(0)
@@ -185,6 +226,7 @@ class ProgressBarCallback(TrainCallback):
            context.dataloader,
            desc=f"Epoch {context.epoch + 1}/{self.num_epoch}",
            dynamic_ncols=True,
+            file=self.file,
        )

    @only_on_rank(0)
@@ -238,7 +280,7 @@ class MetricLoggerCallback(TrainCallback):

    def _get_log_data(self, context: TrainContext):
        return {
-            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
            "epoch": context.epoch,
            "iter": context.iteration,
            **{m: self._metric_funcs[m](context) for m in self.metrics},
@@ -25,18 +25,29 @@ class Trainer:

    def _get_default_callbacks(self) -> List[TrainCallback]:
        cfg = self.train_config
-        return [
+        callbacks = [
+            CallbackFactory.create(
+                "gradient_checkpointing",
+                modules=cfg.gradient_checkpointing_modules,
+            ),
            CallbackFactory.create(
                "checkpoint",
                cfg.ckpt_dir,
                cfg.ckpt_interval,
                state_dict_fn=cfg.state_dict_fn,
            ),
+            CallbackFactory.create(
+                "metric_logger",
+                log_dir=cfg.log_dir,
+                save_interval=cfg.ckpt_interval,
+                log_interval=cfg.log_interval,
+                metrics=cfg.metrics,
+            ),
            CallbackFactory.create("progress_bar", cfg.n_epoch),
-            CallbackFactory.create("metric_logger", cfg.ckpt_dir, cfg.ckpt_interval),
            CallbackFactory.create("gradient_clipping", cfg.max_grad_norm),
            CallbackFactory.create("validation"),
        ]
+        return callbacks

    def _call_callbacks(self, method_name: str, context: TrainContext):
        for callback in self.callbacks:
@@ -1,12 +1,13 @@
 services:
  server:
-    build: .
-    image: astrai:latest
+    build:
+      context: .
+      dockerfile: Dockerfile
+    user: "${UID:-1000}:${GID:-1000}"
    ports:
      - "8000:8000"
    volumes:
      - ./params:/app/params:ro
-      - ./checkpoints:/app/checkpoints
    command: python -m scripts.tools.server --port 8000 --device cuda
    deploy:
      resources:
@@ -25,13 +26,14 @@ services:

  server-cpu:
    profiles: [cpu]
-    build: .
-    image: astrai:latest
+    build:
+      context: .
+      dockerfile: Dockerfile
+    user: "${UID:-1000}:${GID:-1000}"
    ports:
      - "8000:8000"
    volumes:
      - ./params:/app/params:ro
-      - ./checkpoints:/app/checkpoints
    command: python -m scripts.tools.server --port 8000 --device cpu
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
@@ -16,6 +16,7 @@ NC='\033[0m' # No Color
 IMAGE_NAME="astrai"
 IMAGE_TAG="latest"
 REGISTRY=""
+CONTAINER_ID=""

 # Print colored messages
 print_info() {
@@ -175,6 +176,10 @@ main() {
                PORT="$2"
                shift 2
                ;;
+            --container)
+                CONTAINER_ID="$2"
+                shift 2
+                ;;
            --gpu)
                GPU=true
                shift
@@ -197,6 +202,7 @@ main() {
                echo "  --dockerfile FILE  Dockerfile path (default: Dockerfile)"
                echo "  --context PATH     Build context (default: .)"
                echo "  --port PORT        Port for run (default: 8000)"
+                echo "  --container ID     Container ID for logs"
                echo "  --gpu              Enable GPU support"
                echo "  --help             Show this help message"
                echo ""
@@ -205,6 +211,7 @@ main() {
                echo "  $0 build --tag v1.0.0"
                echo "  $0 run --port 8080"
                echo "  $0 run --gpu"
+                echo "  $0 logs --container abc123"
                echo "  $0 push --registry ghcr.io/username"
                exit 0
                ;;
@@ -237,7 +244,7 @@ main() {
            show_info
            ;;
        logs)
-            show_logs "$2"
+            show_logs "$CONTAINER_ID"
            ;;
        "")
            print_error "No command specified. Use --help for usage"
@@ -69,14 +69,14 @@ def parse_args() -> argparse.Namespace:
    parser.add_argument(
        "--adamw_beta1",
        type=float,
-        default=0.95,
-        help="Beta values for AdamW optimizer.",
+        default=0.9,
+        help="Beta1 for AdamW optimizer.",
    )
    parser.add_argument(
        "--adamw_beta2",
        type=float,
-        default=0.99,
-        help="Beta values for AdamW optimizer.",
+        default=0.95,
+        help="Beta2 for AdamW optimizer.",
    )
    parser.add_argument(
        "--adamw_weight_decay",
@@ -157,5 +157,60 @@ def test_messages_with_system(client, loaded_model):
    assert data["type"] == "message"


+def test_chat_completions_stop_sequence(client, loaded_model):
+    """POST /v1/chat/completions with stop parameter truncates at stop sequence."""
+
+    async def async_gen():
+        yield "Hello"
+        yield "X"
+        yield "world"
+
+    app.state.engine = loaded_model
+    loaded_model.generate_async.return_value = async_gen()
+    response = client.post(
+        "/v1/chat/completions",
+        json={
+            "messages": [{"role": "user", "content": "Hello"}],
+            "max_tokens": 100,
+            "stream": False,
+            "stop": ["X"],
+        },
+    )
+    assert response.status_code == 200
+    data = response.json()
+    content = data["choices"][0]["message"]["content"]
+    assert "X" in content
+    assert "world" not in content
+
+
+def test_chat_completions_stop_sequence_stream(client, loaded_model):
+    """POST /v1/chat/completions with stop parameter truncates SSE stream."""
+
+    async def async_gen():
+        yield "Hello"
+        yield "X"
+        yield "world"
+
+    app.state.engine = loaded_model
+    loaded_model.generate_async.return_value = async_gen()
+    response = client.post(
+        "/v1/chat/completions",
+        json={
+            "messages": [{"role": "user", "content": "Hello"}],
+            "max_tokens": 100,
+            "stream": True,
+            "stop": ["X"],
+        },
+        headers={"Accept": "text/event-stream"},
+    )
+    assert response.status_code == 200
+    content = response.content.decode("utf-8")
+    assert "Hello" in content
+    assert "world" not in content
+    assert any(
+        "finish_reason" in line for line in content.split("\n") if "stop" in line
+    )
+
+
 if __name__ == "__main__":
    pytest.main([__file__, "-v"])
@@ -1,11 +1,130 @@
 import torch

 from astrai.config.train_config import TrainConfig
+from astrai.model.components.decoder_block import DecoderBlock
 from astrai.trainer.schedule import SchedulerFactory
-from astrai.trainer.train_callback import TrainCallback
+from astrai.trainer.train_callback import GradientCheckpointingCallback, TrainCallback
 from astrai.trainer.trainer import Trainer


+def test_gradient_checkpointing_enable_disable(test_model):
+    """Enable wraps forward, _disable restores it."""
+    model = test_model["model"]
+    callback = GradientCheckpointingCallback(modules=[DecoderBlock])
+
+    originals = [layer.forward for layer in model.layers]
+
+    for layer in model.layers:
+        callback._enable(layer)
+
+    for layer in model.layers:
+        assert hasattr(layer, "_original_forward")
+        assert layer.forward is not originals[0]
+
+    for layer in model.layers:
+        callback._disable(layer)
+
+    for layer in model.layers:
+        assert not hasattr(layer, "_original_forward")
+
+
+def test_gradient_checkpointing_empty_modules_noop(test_model):
+    """modules=None should leave forwards untouched."""
+    model = test_model["model"]
+    callback = GradientCheckpointingCallback()
+
+    originals = [layer.forward for layer in model.layers]
+
+    for layer in model.layers:
+        callback._enable(layer)
+
+    for layer, orig in zip(model.layers, originals):
+        assert layer.forward is orig
+
+
+def test_gradient_checkpointing_forward_unchanged(test_model):
+    """Forward output unchanged after patching (no_grad)."""
+    model = test_model["model"]
+    device = test_model["device"]
+    callback = GradientCheckpointingCallback(modules=[DecoderBlock])
+
+    input_ids = torch.randint(0, 1000, (2, 32)).to(device)
+
+    with torch.no_grad():
+        ref = model(input_ids)["logits"].clone()
+
+    for layer in model.layers:
+        callback._enable(layer)
+
+    with torch.no_grad():
+        out = model(input_ids)["logits"]
+
+    assert torch.equal(ref, out)
+
+
+def test_gradient_checkpointing_backward(test_model):
+    """backward passes gradients through checkpointed layers."""
+    model = test_model["model"]
+    device = test_model["device"]
+    callback = GradientCheckpointingCallback(modules=[DecoderBlock])
+
+    for layer in model.layers:
+        callback._enable(layer)
+
+    input_ids = torch.randint(0, 1000, (2, 32)).to(device)
+    target_ids = torch.randint(0, 1000, (2, 32)).to(device)
+
+    logits = model(input_ids)["logits"]
+    loss = torch.nn.functional.cross_entropy(
+        logits.flatten(0, 1).float(), target_ids.flatten()
+    )
+    loss.backward()
+
+    for name, param in model.named_parameters():
+        if param.requires_grad:
+            assert param.grad is not None, f"{name} gradient is None"
+
+    for layer in model.layers:
+        callback._disable(layer)
+
+    model.zero_grad()
+    for name, p in model.named_parameters():
+        assert p.grad is None or p.grad.sum().item() == 0, f"{name} grad not zeroed"
+
+
+def test_gradient_checkpointing_trainer_integration(base_test_env, random_dataset):
+    """Gradient checkpointing runs end-to-end via Trainer."""
+
+    def optimizer_fn(model):
+        return torch.optim.AdamW(model.parameters())
+
+    def scheduler_fn(optim):
+        return SchedulerFactory.create(
+            optim, "cosine", warmup_steps=10, lr_decay_steps=10, min_rate=0.05
+        )
+
+    train_config = TrainConfig(
+        model=base_test_env["model"],
+        strategy="seq",
+        dataset=random_dataset,
+        optimizer_fn=optimizer_fn,
+        scheduler_fn=scheduler_fn,
+        ckpt_dir=base_test_env["test_dir"],
+        n_epoch=1,
+        batch_per_device=2,
+        ckpt_interval=3,
+        grad_accum_steps=1,
+        max_grad_norm=1.0,
+        random_seed=42,
+        device_type=base_test_env["device"],
+        gradient_checkpointing_modules=[DecoderBlock],
+    )
+
+    trainer = Trainer(train_config)
+    trainer.train()
+    # no crash = callback correctly enabled/disabled
+
+
 def test_callback_integration(base_test_env, random_dataset):
    """Test that all callbacks are properly integrated"""
Author	SHA1	Message	Date
ViperEkura	785d65436c	fix: 修复 to_dict list 类型丢失与 OpenAI stop 参数失效 - to_dict() 增加 list 类型序列化支持，metrics 等字段不再丢失 - OpenAIHandler 补充 get_stop_sequences/on_token，读取 request.stop 并检测停止序列 - 文档类图补充缺失字段、修正关系分类、ChatCompletionRequest 字段增加 Optional	2026-05-19 21:07:07 +08:00
ViperEkura	64be81b7b3	feat: ProgressBarCallback 支持日志行输出到 stdout - serialization 和 metric_logger 的 timestamp 统一使用 ISO 8601 格式 - ProgressBarCallback 新增 log_interval/file 参数，默认输出到 sys.stdout	2026-05-19 19:12:38 +08:00
ViperEkura	45479b5731	feat: metric 参数通过 TrainConfig 传递 - TrainConfig 新增 log_dir/log_interval/metrics 配置字段 - metric_logger 调用改用 **kwargs 传递，BaseFactory.create 自动过滤	2026-05-19 17:50:24 +08:00
ViperEkura	e0a3337c22	docs: 更新视频链接	2026-05-19 17:34:01 +08:00
ViperEkura	812238060b	fix: docker-compose UID/GID 添加默认值，修复 docker.sh logs 命令	2026-05-18 14:24:00 +08:00
ViperEkura	14b0d56197	fix: 修复无法创建子进程的问题 - mp.start_processes daemon=False	2026-05-18 09:40:32 +08:00
ViperEkura	6c8533f1d2	docs: 修正文档中类名/字段名与代码不一致之处 - ModelConfig → AutoRegressiveLMConfig, Transformer → AutoRegressiveLM - 新增缺失类: EncoderConfig, EmbeddingEncoder, ConfigFactory, StorageFactory, ValidationCallback - TrainConfig/TrainContext/ChatCompletionRequest 补充缺失字段 - dataflow.md 中 create_storage → StorageFactory.create - 示例 --train_type=pt → seq 与代码一致	2026-05-17 21:02:21 +08:00
ViperEkura	2c2697390d	feat: 新增 GradientCheckpointingCallback - TrainConfig.gradient_checkpointing_modules 指定模块类型 - apply 递归遍历，兼容 DDP，不硬编码模型结构 - modules=None 时静默跳过，零开销	2026-05-17 18:21:05 +08:00
ViperEkura	7621f05d3f	docs: AdamW beta 默认值改为 (0.9, 0.95) - 与 Muon 优化器的 AdamW 子优化器保持一致 - 同步更新 train.py/training.md/params.md/README	2026-05-17 17:08:31 +08:00