refactor : 合并 data config docstring 并实现 BFD 打包策略

- 将 ProcessingConfig/OutputConfig 参数描述合并到类级 docstring

- Pipeline 支持 packing_strategy/truncation_mode，新增 bfd 打包

astrai/config/preprocess_config.py

@@ -33,25 +33,70 @@ class InputConfig(BaseConfig):
 
 @dataclass
 class ProcessingConfig(BaseConfig):
+    """Processing configuration.
+
+    Parameters
+    ----------
+    max_seq_len : int
+        Maximum sequence length (default: 2048).
+    min_chars : int
+        Minimum number of characters to keep (default: 50).
+    max_chars : int
+        Maximum number of characters to keep (default: 2_000_000).
+    max_items : Optional[int]
+        Maximum number of items to process (default: None, unlimited).
+    packing_strategy : str
+        How to pack sequences into a contiguous stream.
+
+        - ``"simple"``: sequential concatenation (default, backward compatible).
+        - ``"bfd"``: best-fit decreasing bin packing, minimises wasted tokens.
+        - ``"bfd_split"``: BFD with over-length sequences split into chunks.
+    max_packed_len : int
+        Maximum length of a packed bin. Sequences longer than this are
+        truncated or split depending on ``packing_strategy`` (default: 8192).
+    truncation_mode : str
+        How to truncate sequences longer than ``max_packed_len``.
+
+        - ``"keep_start"``: keep the first ``max_packed_len`` tokens (default).
+        - ``"keep_end"``: keep the last ``max_packed_len`` tokens.
+    """
+
     max_seq_len: int = 2048
     min_chars: int = 50
     max_chars: int = 2_000_000
     max_items: Optional[int] = None
+    packing_strategy: str = "simple"
+    max_packed_len: int = 8192
+    truncation_mode: str = "keep_start"
 
 
 @dataclass
 class OutputConfig(BaseConfig):
+    """Output configuration.
+
+    Parameters
+    ----------
+    domain_key : Optional[str]
+        Domain key for the output store (default: None).
+    storage_format : str
+        Storage format, one of ``"bin"``, ``"jsonl"`` (default: ``"bin"``).
+    max_tokens_per_shard : int
+        Maximum tokens per shard before splitting (default: 100_000_000).
+    dtype : Dict[str, str]
+        Per-key dtype overrides, e.g. ``{"input_ids": "int32"}`` (default: {}).
+    position_ids_mode : Optional[str]
+        How to compute position_ids in packed sequences.
+
+        - ``None`` / ``"none"``: do not generate (backward compatible).
+        - ``"doc_reset"``: reset to 0 at each document boundary.
+        - ``"continuous"``: sequential 0, 1, 2, ... (pretrain, single doc).
+    """
+
     domain_key: Optional[str] = None
     storage_format: str = "bin"
     max_tokens_per_shard: int = 100_000_000
     dtype: Dict[str, str] = field(default_factory=dict)
     position_ids_mode: Optional[str] = None
-    """How to compute position_ids in packed sequences.
-
-    - ``None`` / ``"none"``: do not generate (backward compatible).
-    - ``"doc_reset"``: reset to 0 at each document boundary.
-    - ``"continuous"``: sequential 0, 1, 2, ... (pretrain, single doc).
-    """
 
 
 @dataclass

astrai/inference/__init__.py

@@ -17,7 +17,7 @@ from astrai.inference.api import (
     MessagesRequest,
     ProtocolHandler,
     StopChecker,
-    app,
+    get_app,
     run_server,
 )
 from astrai.inference.api.anthropic import AnthropicResponseBuilder
@@ -80,6 +80,6 @@ __all__ = [
     "ChatCompletionRequest",
     "AnthropicMessage",
     "MessagesRequest",
-    "app",
+    "get_app",
     "run_server",
 ]

astrai/inference/api/__init__.py

@@ -1,4 +1,8 @@
-"""Inference API: protocol handler, stop checker, and FastAPI server."""
+"""Inference API: protocol handler, stop checker, and FastAPI server.
+
+``app`` is no longer a module-level global. Use :func:`get_app` to access the
+lazy singleton FastAPI instance.
+"""
 
 from astrai.inference.api.protocol import GenContext, ProtocolHandler, StopChecker
 from astrai.inference.api.server import (
@@ -6,7 +10,7 @@ from astrai.inference.api.server import (
     ChatCompletionRequest,
     ChatMessage,
     MessagesRequest,
-    app,
+    get_app,
     run_server,
 )
 
@@ -18,6 +22,6 @@ __all__ = [
     "ChatCompletionRequest",
     "ChatMessage",
     "MessagesRequest",
-    "app",
+    "get_app",
     "run_server",
 ]

astrai/inference/api/server.py

@@ -3,6 +3,9 @@ OpenAI / Anthropic-compatible chat completion server backed by continuous-batchi
 
 Protocol-specific formatting is delegated to ``astrai.inference.protocol``.
 This module owns the FastAPI app, request/response schemas, and dependency wiring.
+
+``app`` is lazily constructed — importing this module does NOT create a FastAPI instance.
+Use :func:`get_app` to access the singleton.
 """
 
 import logging
@@ -12,7 +15,7 @@ from typing import Any, Dict, List, Optional, Union
 
 import torch
 import uvicorn
-from fastapi import FastAPI, HTTPException
+from fastapi import APIRouter, FastAPI, HTTPException
 from pydantic import BaseModel, Field
 
 from astrai.inference.api.anthropic import AnthropicResponseBuilder
@@ -24,7 +27,7 @@ from astrai.tokenize import AutoTokenizer
 
 logger = logging.getLogger(__name__)
 
-_project_root = Path(__file__).parent.parent.parent
+_app_instance: Optional[FastAPI] = None
 
 
 class ChatMessage(BaseModel):
@@ -84,17 +87,15 @@ async def lifespan(app: FastAPI):
         logger.info("Inference engine shutdown complete")
 
 
-app = FastAPI(title="AstrAI Inference Server", version="0.2.0", lifespan=lifespan)
+router = APIRouter()
 
 
 def _create_engine(
-    param_path: Optional[Path] = None,
+    param_path: Path,
     device: str = "cuda",
     dtype: torch.dtype = torch.bfloat16,
     max_batch_size: int = 16,
 ) -> InferenceEngine:
-    if param_path is None:
-        param_path = _project_root / "params"
     if not param_path.exists():
         raise FileNotFoundError(f"Parameter directory not found: {param_path}")
 
@@ -112,34 +113,50 @@ def _create_engine(
     return engine
 
 
+def get_app() -> FastAPI:
+    """Return the singleton FastAPI instance (lazily created on first call)."""
+    global _app_instance
+    if _app_instance is None:
+        _app_instance = FastAPI(
+            title="AstrAI Inference Server",
+            version="0.2.0",
+            lifespan=lifespan,
+        )
+        _app_instance.include_router(router)
+        _app_instance.state.server_config = {}
+        _app_instance.state.engine = None
+    return _app_instance
+
+
 def _get_engine() -> InferenceEngine:
-    engine = app.state.engine
+    engine = get_app().state.engine
     if engine is None:
         raise HTTPException(status_code=503, detail="Engine not initialized")
     return engine
 
 
-@app.get("/health")
+@router.get("/health")
 async def health():
+    app = get_app()
     return {
         "status": "ok",
         "model_loaded": app.state.engine is not None,
     }
 
 
-@app.get("/stats")
+@router.get("/stats")
 async def get_stats():
     return _get_engine().get_stats()
 
 
-@app.post("/v1/chat/completions")
+@router.post("/v1/chat/completions")
 async def chat_completion(request: ChatCompletionRequest):
     engine = _get_engine()
     handler = ProtocolHandler(request, engine, OpenAIResponseBuilder())
     return await handler.handle()
 
 
-@app.post("/v1/messages")
+@router.post("/v1/messages")
 async def create_message(request: MessagesRequest):
     engine = _get_engine()
     handler = ProtocolHandler(request, engine, AnthropicResponseBuilder())
@@ -147,14 +164,15 @@ async def create_message(request: MessagesRequest):
 
 
 def run_server(
+    param_path: Path,
     host: str = "0.0.0.0",
     port: int = 8000,
     reload: bool = False,
     device: str = "cuda",
     dtype: torch.dtype = torch.bfloat16,
-    param_path: Optional[Path] = None,
     max_batch_size: int = 16,
 ):
+    app = get_app()
     app.state.server_config = {
         "device": device,
         "dtype": dtype,

astrai/model/transformer.py

@@ -26,24 +26,21 @@ def process_attention_mask(
         return input_mask
 
     device = input_tensor.device
-    dtype = input_tensor.dtype
-    B, S = input_tensor.size()[:2]
+    B = input_tensor.size(0)
     T = position_ids.max().item() + 1
 
     if input_mask is None:
         if position_ids.min().item() == 0 and is_causal:
             return None
-        pad = torch.ones(B, T, dtype=torch.bool, device=device)
+        attend = torch.ones(B, 1, T, dtype=torch.bool, device=device)
     else:
-        pad = input_mask[:, :T].to(device=device, dtype=torch.bool)
+        attend = input_mask[:, :T].to(device=device, dtype=torch.bool).unsqueeze(1)
 
-    attend = pad.view(B, 1, T).expand(B, S, T).clone()
     if is_causal:
-        attend &= position_ids.unsqueeze(-1) >= torch.arange(T, device=device)
+        causal = position_ids.unsqueeze(-1) >= torch.arange(T, device=device)
+        attend = attend & causal
 
-    return torch.full(
-        (B, 1, S, T), -torch.finfo(dtype).max / 2, dtype=dtype, device=device
-    ).masked_fill_(attend.unsqueeze(1), 0.0)
+    return attend.unsqueeze(1)
 
 
 @AutoModel.register("autoregressive_lm")

astrai/preprocessing/pipeline.py

@@ -8,7 +8,7 @@ import json
 import os
 from collections import defaultdict
 from itertools import chain
-from typing import Optional
+from typing import List, Optional, Tuple
 
 import torch
 import tqdm
@@ -35,6 +35,65 @@ def filter_by_length(text: str, min_len: int = 50, max_len: int = 2_000_000) ->
     return min_len <= len(text) <= max_len
 
 
+def _truncate(seq: list, max_len: int, mode: str) -> list:
+    if len(seq) <= max_len:
+        return seq
+    if mode == "keep_end":
+        return seq[-max_len:]
+    return seq[:max_len]
+
+
+def pack_sequences(
+    sequences: List[list],
+    max_packed_len: int,
+    strategy: str,
+    truncation_mode: str,
+) -> List[Tuple[int, int]]:
+    """Pack *sequences* into bins and return a reorder plan.
+
+    Returns a list of ``(orig_idx, truncated_length)`` in flush order.
+    All keys (sequence, loss_mask, …) must be reordered and truncated
+    identically according to this plan.
+
+    Supported *strategy* values:
+
+    - ``"simple"``: sequential, no reordering.
+    - ``"bfd"``: best-fit decreasing bin packing.
+    """
+    n = len(sequences)
+    if strategy == "simple":
+        return [(i, min(len(sequences[i]), max_packed_len)) for i in range(n)]
+
+    order = sorted(range(n), key=lambda i: len(sequences[i]), reverse=True)
+    bins: List[List[int]] = []
+    bin_lengths: List[int] = []
+
+    for orig_idx in order:
+        seq_len = min(len(sequences[orig_idx]), max_packed_len)
+
+        best_bin = None
+        best_remain = max_packed_len + 1
+        for i, bl in enumerate(bin_lengths):
+            remain = max_packed_len - bl
+            if seq_len <= remain < best_remain:
+                best_remain = remain
+                best_bin = i
+
+        if best_bin is not None:
+            bins[best_bin].append(orig_idx)
+            bin_lengths[best_bin] += seq_len
+        else:
+            bins.append([orig_idx])
+            bin_lengths.append(seq_len)
+
+    plan: List[Tuple[int, int]] = []
+    for bin_indices in bins:
+        for orig_idx in bin_indices:
+            plan.append((orig_idx, min(len(sequences[orig_idx]), max_packed_len)))
+
+    return plan
+
+
 class Pipeline:
     """Tokenization pipeline driven by a declarative :class:`PipelineConfig`.
 
@@ -145,6 +204,25 @@ class Pipeline:
         for domain, keys in domains.items():
             idx = shard_idx[domain]
             chunk_dir = os.path.join(self.output_dir, domain)
+
+            pp = self.config.preprocessing
+            if pp.packing_strategy != "simple" and "sequence" in keys:
+                plan = pack_sequences(
+                    keys["sequence"],
+                    pp.max_packed_len,
+                    pp.packing_strategy,
+                    pp.truncation_mode,
+                )
+                reordered = defaultdict(list)
+                for orig_idx, truncated_len in plan:
+                    for k, vals in keys.items():
+                        reordered[k].append(
+                            _truncate(
+                                vals[orig_idx], pp.max_packed_len, pp.truncation_mode
+                            )
+                        )
+                keys = reordered
+
             tensors = {}
             for key, ids_list in keys.items():
                 dt = _STR_TO_DTYPE.get(

astrai/trainer/strategy.py

@@ -68,6 +68,22 @@ def get_logprobs(
         return token_logprobs * shifted_mask
 
 
+def make_doc_boundary_mask(position_ids: Tensor) -> Tensor:
+    S = position_ids.size(1)
+    device = position_ids.device
+    boundaries = position_ids[:, 1:] <= position_ids[:, :-1]
+    doc_ids = torch.cat(
+        [
+            torch.zeros(position_ids.size(0), 1, dtype=torch.long, device=device),
+            boundaries.long().cumsum(dim=1),
+        ],
+        dim=1,
+    )
+    same_doc = doc_ids.unsqueeze(-1) == doc_ids.unsqueeze(-2)
+    causal = torch.tril(torch.ones(S, S, dtype=torch.bool, device=device))
+    return (same_doc & causal).unsqueeze(1)
+
+
 class BaseStrategy(ABC):
     """Abstract base class for training strategies."""
 
@@ -188,8 +204,11 @@ class SFTStrategy(BaseStrategy):
         )
 
         ignore_index = -100
-        logits = self.model(input_ids=input_ids, position_ids=position_ids)["logits"]
+        input_mask = make_doc_boundary_mask(position_ids)
         target_ids = target_ids.masked_fill(loss_mask == 0, ignore_index)
+        logits = self.model(
+            input_ids=input_ids, position_ids=position_ids, input_mask=input_mask
+        )["logits"]
 
         loss = F.cross_entropy(
             input=logits.flatten(0, 1).float(),

tests/inference/conftest.py

@@ -5,21 +5,22 @@ from unittest.mock import MagicMock
 import pytest
 from fastapi.testclient import TestClient
 
-from astrai.inference import app
+from astrai.inference import get_app
 
 
 @pytest.fixture
 def client():
     """Provide a test client for the FastAPI app."""
-    app.state.server_config = {
+    _app = get_app()
+    _app.state.server_config = {
         "device": "cpu",
         "dtype": "bfloat16",
         "param_path": None,
         "max_batch_size": 1,
         "_test": True,
     }
-    app.state.engine = None
-    return TestClient(app)
+    _app.state.engine = None
+    return TestClient(_app)
 
 
 @pytest.fixture
@@ -49,5 +50,5 @@ def mock_engine():
 @pytest.fixture
 def loaded_model(client, mock_engine):
     """Simulate that the engine is loaded."""
-    app.state.engine = mock_engine
+    get_app().state.engine = mock_engine
     return mock_engine

tests/inference/test_server.py

@@ -2,12 +2,12 @@
 
 import pytest
 
-from astrai.inference import app
+from astrai.inference import get_app
 
 
 def test_health_no_model(client):
     """GET /health should return 200 even when engine not loaded."""
-    app.state.engine = None
+    get_app().state.engine = None
     response = client.get("/health")
     assert response.status_code == 200
     data = response.json()
@@ -30,7 +30,7 @@ def test_chat_completions_non_stream(client, loaded_model):
     async def async_gen():
         yield "Assistant reply"
 
-    app.state.engine = loaded_model
+    get_app().state.engine = loaded_model
     loaded_model.generate_async.return_value = async_gen()
     response = client.post(
         "/v1/chat/completions",
@@ -56,7 +56,7 @@ def test_chat_completions_stream(client, loaded_model):
         yield "cumulative1"
         yield "cumulative2"
 
-    app.state.engine = loaded_model
+    get_app().state.engine = loaded_model
     loaded_model.generate_async.return_value = async_gen()
     response = client.post(
         "/v1/chat/completions",
@@ -83,7 +83,7 @@ def test_messages_non_stream(client, loaded_model):
     async def async_gen():
         yield "Assistant reply"
 
-    app.state.engine = loaded_model
+    get_app().state.engine = loaded_model
     loaded_model.generate_async.return_value = async_gen()
     response = client.post(
         "/v1/messages",
@@ -111,7 +111,7 @@ def test_messages_stream(client, loaded_model):
         yield "cumulative1"
         yield "cumulative2"
 
-    app.state.engine = loaded_model
+    get_app().state.engine = loaded_model
     loaded_model.generate_async.return_value = async_gen()
     response = client.post(
         "/v1/messages",
@@ -141,7 +141,7 @@ def test_messages_with_system(client, loaded_model):
     async def async_gen():
         yield "Reply"
 
-    app.state.engine = loaded_model
+    get_app().state.engine = loaded_model
     loaded_model.generate_async.return_value = async_gen()
     response = client.post(
         "/v1/messages",
@@ -165,7 +165,7 @@ def test_chat_completions_stop_sequence(client, loaded_model):
         yield "X"
         yield "world"
 
-    app.state.engine = loaded_model
+    get_app().state.engine = loaded_model
     loaded_model.generate_async.return_value = async_gen()
     response = client.post(
         "/v1/chat/completions",
@@ -191,7 +191,7 @@ def test_chat_completions_stop_sequence_stream(client, loaded_model):
         yield "X"
         yield "world"
 
-    app.state.engine = loaded_model
+    get_app().state.engine = loaded_model
     loaded_model.generate_async.return_value = async_gen()
     response = client.post(
         "/v1/chat/completions",