docs: sync all documentation with current codebase
- remove GenerationRequest and generate_with_request references (class deleted) - document cuda>flash>torch default priority and FlashAttnBackend - add ASTR_BACKEND env var to backend docs, TorchNativeBackend (default) → (fallback) - fix JsonlStore transform routing → DatasetFactory ownership - fix CudaBackend fallback chain description (FlashAttn → TorchNative) - add FlashAttnBackend to architecture strategy table - add router_stats to DecoderOutput/FFNOutput TypedDict diagrams - add decode_o_part/ml_part/decode_out to KVCache diagram - add --append_eos/--no-append_eos to IFD evaluation parameter table - update get-started CUDA kernel note (no longer requires explicit attn_backend activation) - fix python -m scripts.tools.server (no __init__.py) → direct script call
This commit is contained in:
@@ -315,6 +315,7 @@ classDiagram
|
||||
<<TypedDict>>
|
||||
+Tensor hidden_states
|
||||
+Optional[Tensor] aux_loss
|
||||
+Optional[RouterStats] router_stats
|
||||
}
|
||||
|
||||
class GQA {
|
||||
@@ -361,6 +362,7 @@ classDiagram
|
||||
<<TypedDict>>
|
||||
+Tensor hidden_states
|
||||
+Optional[Tensor] aux_loss
|
||||
+Optional[RouterStats] router_stats
|
||||
}
|
||||
|
||||
class DeepSeekMoE {
|
||||
@@ -807,7 +809,6 @@ classDiagram
|
||||
+AutoTokenizer tokenizer
|
||||
+InferenceScheduler scheduler
|
||||
+generate(prompt, stream, max_tokens, temperature, top_p, top_k, frequency_penalty, rep_window) Union[Generator, str, List[str]]
|
||||
+generate_with_request(request) Union[Generator, str, List[str]]
|
||||
+generate_async(prompt, max_tokens, temperature, top_p, top_k, frequency_penalty, rep_window) AsyncGenerator
|
||||
+get_stats() Dict
|
||||
+shutdown()
|
||||
@@ -915,6 +916,9 @@ classDiagram
|
||||
+int max_len
|
||||
+Optional[Tensor] kv_indptr
|
||||
+Optional[Tensor] qo_indptr
|
||||
+Optional[Tensor] decode_o_part
|
||||
+Optional[Tensor] decode_ml_part
|
||||
+Optional[Tensor] decode_out
|
||||
}
|
||||
|
||||
class PagePool {
|
||||
@@ -980,17 +984,6 @@ classDiagram
|
||||
+get_stats() Dict
|
||||
}
|
||||
|
||||
class GenerationRequest {
|
||||
+List[Dict] messages
|
||||
+int top_k
|
||||
+float top_p
|
||||
+float temperature
|
||||
+Optional[int] max_tokens
|
||||
+float frequency_penalty
|
||||
+int rep_window
|
||||
+bool stream
|
||||
}
|
||||
|
||||
class BaseSamplingStrategy {
|
||||
<<abstract>>
|
||||
+apply(logits, filter_value, input_ids, input_mask) Tensor
|
||||
@@ -1407,7 +1400,7 @@ classDiagram
|
||||
CheckpointCallback ..> Checkpoint : creates
|
||||
PagePool ..> KVCache : binds
|
||||
PagePool ..> InferenceWorkspace : fills
|
||||
InferenceEngine ..> GenerationRequest : uses
|
||||
InferenceEngine ..> GenerateResult : uses
|
||||
InferenceEngine ..> GenerateResult : creates
|
||||
OpenAIResponseBuilder ..> ChatCompletionRequest : receives
|
||||
AnthropicResponseBuilder ..> MessagesRequest : receives
|
||||
@@ -1443,7 +1436,7 @@ classDiagram
|
||||
| **astrai.model** | ModelFactory, AutoModel, AutoRegressiveLM, EmbeddingEncoder, DecoderBlock, GQA, MLA, MLP, DeepSeekMoE, AttnFactory, FFNFactory, RMSNorm, Linear, LoRAConfig, LoRALinear, RotaryEmbedding, Embedding | Neural network model |
|
||||
| **astrai.tokenize** | AutoTokenizer, ChatTemplate | Tokenizer and chat template |
|
||||
| **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy–GRPOStrategy, StrategyFactory, BaseScheduler–WSDScheduler, SchedulerFactory, TrainCallback(Protocol)–MetricCallback, CallbackFactory, RawRollout, RolloutResult, BaseRewardModel, RolloutGenerator, RolloutRunner | Training workflow |
|
||||
| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, InferenceWorkspace, PagePool, KVStorage, ReqToTokenPool, KVCache, Allocator, RadixCache, Task, TaskManager, TaskStatus, StreamDecoder, GenerationRequest, GenerateResult, BaseSamplingStrategy–SamplingPipeline, FrequencyPenaltyStrategy, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, StopInfo, ChatMessage, FunctionDef, ToolDef, ChatCompletionRequest, AnthropicMessage, MessagesRequest, BaseToolParser, ToolParserFactory, SimpleJsonToolParser | Inference service |
|
||||
| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, InferenceWorkspace, PagePool, KVStorage, ReqToTokenPool, KVCache, Allocator, RadixCache, Task, TaskManager, TaskStatus, StreamDecoder, GenerateResult, BaseSamplingStrategy–SamplingPipeline, FrequencyPenaltyStrategy, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, StopInfo, ChatMessage, FunctionDef, ToolDef, ChatCompletionRequest, AnthropicMessage, MessagesRequest, BaseToolParser, ToolParserFactory, SimpleJsonToolParser | Inference service |
|
||||
| **astrai.extension** | AttentionBackend, TorchNativeBackend, CudaBackend, attn_backend, ATTN_BACKEND, attn_decode, attn_prefill, attn_paged_decode, attn_paged_prefill, rotary_emb, apply_rotary_emb, rotary_backend, is_available | CUDA attention + rotary kernels, backend abstraction, auto-dispatch |
|
||||
| **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank/get_world_size/get_current_device, only_on_rank, LaunchStrategy, TorchrunStrategy, LocalStrategy, BaseExecutor, ExecutorFactory, NoneExecutor, DDPExecutor, FSDPExecutor, GradientState, AccumOptimizer, AccumScheduler | Distributed parallel & gradient accumulation |
|
||||
| **astrai.factory** | BaseFactory | Component registration |
|
||||
@@ -1462,7 +1455,7 @@ classDiagram
|
||||
| **Observer** | `TrainCallback`, callback implementations | Training process monitoring |
|
||||
| **Context** | `TrainContext` | Unified training state bag |
|
||||
| **Object Pool** | `Allocator`, `PagePool` | Page-based KV cache with LRU eviction |
|
||||
| **Strategy (Attention)** | `AttentionBackend`, `TorchNativeBackend`, `CudaBackend` | Attention computation backend switching via context manager |
|
||||
| **Strategy (Attention)** | `AttentionBackend`, `CudaBackend`, `FlashAttnBackend`, `TorchNativeBackend` | Attention computation backend switching via context manager |
|
||||
| **Auto-dispatch (Rotary)** | `apply_rotary_emb`, `rotary_backend.py`, `rotary_ops.py` | Rotary embedding CUDA kernel auto-dispatch with torch fallback |
|
||||
| **Executor** | `BaseExecutor`, `NoneExecutor`, `DDPExecutor`, `FSDPExecutor` | Gradient accumulation & model distribution |
|
||||
| **Storage** | `Store`, `MmapStore`, `JsonlStore` | Format-agnostic data access with multi-segment support |
|
||||
@@ -1475,7 +1468,7 @@ classDiagram
|
||||
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`
|
||||
5. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `PagePool` + `KVCache` + `SamplingPipeline`. Attention backend selected via `attn_backend()` context manager (`TorchNativeBackend` default, `CudaBackend` for CUDA kernels). Rotary embedding auto-dispatches to CUDA kernel when available (inference mode), else torch complex multiply (training).
|
||||
5. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `PagePool` + `KVCache` + `SamplingPipeline`. Attention backend selected via `attn_backend()` context manager (cuda > flash > torch priority; `ASTR_BACKEND` env var overrides default; `TorchNativeBackend` fallback). Rotary embedding auto-dispatches to CUDA kernel when available, else torch complex multiply.
|
||||
6. **Distributed**: `spawn_parallel_fn` + `setup_parallel` for multi-process DDP
|
||||
7. **Dataset Loading**: `DatasetFactory` creates datasets, `Store` (`MmapStore`/`JsonlStore`) loads data with explicit `_length` and multi-segment `_data`
|
||||
8. **Checkpoint**: `Checkpoint` saves/loads safetensors + metadata; `CheckpointCallback` performs rank-0 training saves, with extra state saved as `{key}.pt`
|
||||
|
||||
Reference in New Issue
Block a user