Commit Graph
25 Commits
Author SHA1 Message Date
0z5a 587b0ee046 fix: keep async rollouts version-consistent
- serialize shared-model optimizer updates with generation
- reject future or over-lagged rollout results after asynchronous scoring
- close cache publication races
- persist policy versions in online checkpoints
2026-09-03 12:01:24 +08:00
0z5a 4019ddac31 perf: reuse rollout behavior logprobs
Feed sampler-aligned behavior log-probabilities directly into online GRPO instead of allocating, synchronizing, and forwarding a duplicate old-policy model. Keep the old-model path as an offline compatibility fallback and validate supplied rollout tensors before loss computation.
2026-09-02 19:29:53 +08:00
0z5a e58a728b80 feat: version rollout weight updates
Track a monotonic policy version across optimizer steps, scheduler updates, and rollout results. Serialize synchronous generation with weight acknowledgements and invalidate reusable prefix KV entries so cached samples remain attributable to the behavior policy that generated them.
2026-09-02 19:01:41 +08:00
ViperEkuraand0z5a 1fad50d847 fix: publish checkpoints atomically
- Write checkpoint payloads to a hidden sibling staging directory, add a versioned checksum manifest, fsync the completed payload, and publish it with an atomic rename
- Republishing an existing step retires the old payload under a hidden sibling name before the atomic rename, so re-runs into the same output directory replace the previous checkpoint instead of raising FileExistsError
- Keep legacy checkpoints loadable, add optional checksum verification, and align metric flushing with checkpoint publication

Co-authored-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
2026-09-02 15:29:22 +08:00
0z5a 01bcd0d105 perf: batch ragged prefill requests
- Pack prompts with a shared prefix start and attention backend into one forward.

- Select per-request final logits from cumulative query lengths.

- Cover ragged tokens, logprobs, scheduling, and documentation.
2026-09-02 15:00:16 +08:00
ViperEkura 0e7dafad8e refactor: rename optimizer step callback hooks to before and after
- rename on_optimizer_step to before_optimizer_step across the callback protocol, built-in callbacks, and trainer call site
- rename on_after_optimizer_step to after_optimizer_step for the symmetric post-step hook
- document the hook pair and the checkpoint save location in developer and training guides
2026-09-01 14:22:54 +08:00
ViperEkura a29bdfae46 refactor: rework attention backend resolution
- explicit attn_backend() context wins over ASTR_BACKEND env
- polymorphic available()/supports_call() replace isinstance dispatch
- cache singleton backend instances to avoid hot-path allocation
- training (fwd=None) resolves cuda > flash > torch by capability
- flash dense supports mask-free calls only; masked training falls back to torch
2026-08-23 14:47:02 +08:00
ViperEkura 10fec8dca1 docs: fix stale docs and align with code
- update cuda_kernels layout, arch flags, and add FP8 section
- fix install docs: kernels auto-build when nvcc + CUDA detected
- mark ignored OpenAI request params and complete KVCache fields
- add docker docs to indexes and astrai.optim to module overview
- refresh document update timestamps
2026-08-23 14:23:28 +08:00
ViperEkura cb21af38ba feat: unify Docker serving configuration in YAML
- Add server.py --config serve.yaml; explicit CLI flags override YAML
- Add scripts/serve.sh and serve_runtime.py for the Compose lifecycle
- Template server/cpu ports and param mounts in docker-compose.yml
- Document schema in docs/developer/docker-serving.md and params guide
- Add tests for runtime parsing and server CLI merge logic
2026-08-21 23:16:45 +08:00
ViperEkura 0d0dc64884 docs: explain extension layer boundaries 2026-08-16 21:30:03 +08:00
ViperEkura 6ac3b51496 refactor: separate extension ops and backends 2026-08-16 21:15:52 +08:00
ViperEkura 056c1382ff 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
2026-08-07 23:52:27 +08:00
ViperEkura 6f09b1d2ee docs : clarify radix cache architecture
- document exact page-aligned radix prefix matching
- explain partial-page ownership and materialized KV boundaries
- remove bilingual wording from project overview
2026-08-06 11:50:45 +08:00
ViperEkura cdf9145ecf docs: align CUDA kernel and RoPE docs with code
- Fix rotary docs to describe cos/sin freqs_cis table, not complex buffer
- Replace attn_prefill with attn_paged_prefill for the CudaBackend path
- Register attn_paged_prefill in kernel overview, layout, and module list
- Add qo_indptr and InferenceWorkspace to architecture class diagram
- Add FrequencyPenaltyStrategy to sampling design patterns
2026-08-03 20:54:40 +08:00
ViperEkura 288ba20db1 docs: audit non-CUDA documentation
- Aligns CLI and strategy metric contracts
- Refreshes architecture, dataflow, preprocessing, distributed, and eval guides
- Corrects links, TOCs, defaults, and repository paths
2026-08-02 07:39:24 +08:00
ViperEkura 020e2eff4e refactor: emit strategy metrics as floats
- Converts detached strategy metrics before returning loss output
- Removes redundant item conversion from the trainer loop
- Updates the documented contract and regression tests
2026-08-02 06:38:28 +08:00
ViperEkura 1c7369f293 feat: add MoE auxiliary loss metrics
- Propagates MoE load-balancing loss through model outputs
- Logs task, auxiliary, and weighted losses across strategies
- Computes only explicitly requested callback metrics
- Preserves tensor compute_loss API and adds regression tests
2026-08-02 06:30:43 +08:00
ViperEkura 6db276f37a feat: add Mano manifold optimizer (mano_adamw)
- implement Mano (v2) with axis-rotating tangent projection and manifold normalization, replacing Newton-Schulz iteration
- composite ManoAdamW reuses partition_optimizer_parameters and composite helpers
- register mano_adamw in OptimizerFactory, export Mano and ManoAdamW
- add --mano_momentum and --mano_nesterov CLI options in Optimizer group
- add mano_adamw hyperparameters branch in train.py
- document mano_adamw in params.md
- add tests for single-step projection, axis alternation, factory registration, closure, and resume
2026-08-01 08:51:08 +08:00
ViperEkura 11073bd1d2 refactor: extract composite optimizer helpers and unify naming
- add astrai/optim/composite.py with shared step/zero_grad/state_dict/param_groups helpers and OptimizerFactory
- rename MuonMix to MuonAdamW (matches registered name muon_adamw) and file to muon_adamw.py
- use @OptimizerFactory.register decorator in each optimizer module instead of post-import registration in __init__
- fix closure being invoked once per sub-optimizer in MuonAdamW.step (now exactly once via composite_step)
- NoraNAdamW.step now forwards closure correctly
2026-08-01 08:07:45 +08:00
ViperEkura 25c9e81b2b refactor: keep muon_adamw as default optimizer and drop nora docs
- revert CLI/create_optimizer/display defaults to muon_adamw
- revert README, README-zh-CN, params.md to pre-merge state
2026-08-01 07:51:51 +08:00
QueenAmish 04899a2b15 Make Nora+NAdamW the default optimizer 2026-07-31 23:16:39 +08:00
ViperEkura 7aa5ed09d9 refactor: unify rotary embedding interface and update docs
- Merge cos/sin into single freqs_cis tensor [batch, seq, dim/2, 2] throughout the pipeline: RotaryEmbedding buffer, forward return type, apply_rotary_emb signature, CUDA kernel interface
- CUDA kernel now takes freqs_cis directly and reads cos/sin via stride offset internally, eliminating Python-side slice/copy overhead
- Kernel interface: rotary_emb(x, freqs_cis) replaces rotary_emb(x, cos, sin)
- All call sites pass rotary_emb as Tensor (was tuple), type annotations consistent
- Update build threads from 8 to 16
- Fix all docs: get-started, inference, training, cuda_kernels, architecture, internals — reflect new rotary interface, KVCache fields, rotary backend dispatch, .so path, kernel registry count, file layout
2026-07-31 16:52:25 +08:00
ViperEkura 97114b95a4 docs: update for attention backend and extension API
- Remove stale 'not yet wired' references
- Add AttentionBackend/CudaBackend sections to cuda_kernels.md, internals.md, inference.md
- Add astrai.extension to architecture.md module table and design patterns
- Update get-started.md: CUDA kernels activatable via attn_backend()
2026-07-30 18:50:16 +08:00
ViperEkura df979b4469 refactor: use single-index access and update docs for cache architecture
- Replace all buffer[layer_id][loc] double indexing with buffer[layer_id, loc] single advanced indexing in cache.py and attention.py
- Revert KVStorage buffers back to 4D [n_layers, size, n_kv_heads, head_dim], remove leftover 3D reshape/view in MLA path
- Update docs/guides/inference.md, docs/developer/internals.md, docs/developer/architecture.md to reflect new PagePool/KVStorage/ReqToTokenPool/KVCache classes
2026-07-30 17:47:04 +08:00
ViperEkura c8b1e40f71 docs: restructure to docs/, add guides and developer docs
- Rename assets/ to docs/, split into guides/ and developer/
- Add get-started.md: installation + 5-step quickstart
- Add guides/evaluation.md: 7 eval scripts with CLI args
- Add guides/distributed.md: DDP/FSDP, gradient accumulation, NCCL
- Add developer/internals.md: loss formulas, RoPE, KV cache math
- Add developer/cuda_kernels.md: build system, benchmarks, file layout
- Fix storage_format doc in preprocessing.md
- Update cross-references in README.md, README-zh-CN.md, Dockerfile
2026-07-30 00:49:04 +08:00