Commit Graph
641 Commits
Author SHA1 Message Date
ViperEkura 2e29ed45d3 perf: shrink decode tile to BC=16 for higher occupancy
- BC=32→16 halves smem (32KB→16KB for D=128), doubling blocks/SM (3→6)
- D=256 now fits STAGES=2 double-buffer in 32KB, eliminating 176-byte spill
- min_tiles_per_split=2 avoids excessive split overhead on small kv
- paged decode: require page_size multiple of BC so tiles stay page-aligned

Benchmark (L20 sm_89, D=128):
- B=1 kv=4096: 0.0134→0.0122ms (+9% BW)
- B=16 kv=2048: 0.0434→0.0352ms (+23% BW)
- B=32 kv=1024: 0.0343→0.0282ms (+22% BW)
2026-07-27 22:44:02 +08:00
ViperEkura 5ba21f4eb3 refactor: eliminate test duplication via shared helpers
- Add tests/helpers.py with shared config, dataset, tokenizer, executor, and assertion helpers
- Replace 15 copies of device one-liner with session-scoped fixture
- Collapse 5 near-identical Dataset subclasses into RandomTokenDataset
- Remove duplicate _make_config/_make_model/_make_frozen and FakeTokenizer/FakeExecutor definitions
- Make test_callbacks and test_early_stopping use existing train_config_factory
- Replace 6 duplicate meta.json read blocks with load_shard_meta
- Fix mkdtemp leaks in test_lora.py with TemporaryDirectory
2026-07-27 22:34:53 +08:00
ViperEkura c26a47b0df docs: sync docs with current code after refactor
- architecture: remove TaskManager.max_prompt_len (merged into max_seq_len in 53c804e)
- dataflow: fix DatasetFactory.load param name max_position_embeddings -> max_len
- params: add fsdp2 to parallel_mode, add --max_seq_len to server, add 4 missing generate options
- preprocessing: add missing batch_size config field
2026-07-27 21:43:29 +08:00
ViperEkura b1a87b22bb feat: add --device flag for GPU-accelerated SVD, default to cuda 2026-07-27 08:53:40 +08:00
ViperEkura 07625057f2 feat : add setup_logging with hierarchical astrai logger
- setup_logging(): attach handler only to astrai logger, not root
- all astrai.* sub-module loggers inherit automatically
- controlled by ASTR_LOG_LEVEL env var (default INFO)
- called in if __name__ == '__main__' of each CLI script
2026-07-27 08:13:48 +08:00
ViperEkura 53c804e233 refactor : merge max_prompt_len into max_seq_len, replace assert with raise
- Engine/Scheduler/TaskManager: merge max_prompt_len into max_seq_len
- train.py: replace bare assert with ValueError/FileNotFoundError
- server.py: add --max_seq_len CLI option
- engine.py: remove dead page_size param
2026-07-27 08:05:11 +08:00
ViperEkura 05c7432964 chore: remove AGENTS.md 2026-07-27 07:21:40 +08:00
ViperEkura 4de42d83c2 refactor: migrate scripts from argparse to click, add YAML config support
- Replace argparse with click in all scripts (train, server, generate,
  preprocess, benchmark)
- Add --config YAML support to train.py with CLI flag override
- Add --dry-run mode to validate config before training
- Add type annotations throughout benchmark.py
- Unify docstring format across all commands
- Remove redundant deps httpx, requests, pyyaml, rich from pyproject.toml
- Net -346 lines while adding YAML config support
2026-07-27 06:55:46 +08:00
ViperEkura b99485f462 chore: bump version to 1.3.11 v1.3.11 2026-07-27 01:23:13 +08:00
ViperEkura 20041d7aa9 perf: extend MMA decode to arbitrary GQA ratio, add launch bounds, vectorize combine
- Multi-pass MMA: encode pass in grid blockIdx.x, compute q_head0/G in-kernel
- Fixes crash for G>32 (previously block(32,G) exceeded 1024 threads)
- Fixes alloc_split_partials using uninitialized num_splits (MAX_SPLITS=32)
- __launch_bounds__ on all MMA and prefill kernels for better register allocation
- 4x vectorized combine kernel (4 head_dim per thread)
- uint4 vectorized K loads in scalar decode kernels
- cp.async .L2::128B cache hint for K/V tile streaming
- Extract warp_reduce_sum, bf16, MAX_SPLITS to attn_warp_utils.cuh
2026-07-27 00:35:34 +08:00
ViperEkura 59248032dc chore: fix ruff lint warnings and signal handling edge cases
- Fix pre-existing ruff lint warnings (F401, F541, F841, E741)
- Exclude .md/.json/.yml from ruff format check
- Unblock SIGTERM/SIGINT via pthread_sigmask in early signal handler
- Do not restore SIG_DFL on unregister to prevent pending signal kills
2026-07-25 21:08:30 +08:00
ViperEkura ceadc34ea9 feat: auto-checkpoint on SIGTERM/SIGINT with DDP support
- Register SIGTERM/SIGINT handlers in training loop, set stop flag on signal
- Check stop_requested at each epoch/batch boundary, break and call on_error to save checkpoint
- LocalStrategy parent forwards signal to child processes via terminate(), waits up to 600s for graceful exit
- TrainContext gains threading.Event-based stop_requested/request_stop
- Tests verify SIGTERM/SIGINT trigger checkpoint save with exit code 0, works on both CPU and GPU
2026-07-25 20:40:54 +08:00
ViperEkura 8ab5631446 fix: correct online rollout lifecycle 2026-07-23 19:01:37 +08:00
ViperEkura 99b5d2b2da perf: batch tokenizer preprocessing 2026-07-23 18:42:19 +08:00
ViperEkura 021e6f3788 style: apply ruff formatting to FSDP2 changes 2026-07-23 16:30:10 +08:00
ViperEkura 4e38183e86 fix: make FSDP2 executor work with ABC+Generic model hierarchy
- Wrap each child module individually, skip root (CPython layout
  conflict between ABC+Generic and FSDP2 __class__ assignment)
- Remove manual unshard in clip_grad_norm (DTensor compatible)
- Fix _no_sync to iterate modules() instead of checking root
- Add reshard after unwrap_model
- Guard __init_subclass__ type resolution against dynamic subclasses
- Add fsdp2 to --parallel_mode CLI choices
2026-07-23 16:11:02 +08:00
ViperEkura 4eeb23e2b3 fix: use copy-on-write mmap mode to silence non-writable tensor warning 2026-07-22 17:37:41 +08:00
ViperEkura ef8783b7e3 fix: separate attn_mask and loss_mask in get_logprobs, compose causal masking in strategy
- add loss_mask parameter to get_logprobs to decouple attention from loss masking
- DPO/GRPO strategies compose key-padding + causal mask before model forward
- prevents prompt tokens from being masked out of attention and missing causal masking
2026-07-21 23:47:28 +08:00
ViperEkura 60d7ee614a fix: improve attention kernel numerical stability and test precision checks
- use fmaf() for V-accumulation in scalar decode paths to reduce rounding
- delay scale multiplication to after dot-product in scalar prefill
- unify __expf/expf across MMA and scalar paths for consistent numerics
- harmonize divide-by-zero guards to 1e-20f
- add both absolute and relative error checks in standalone tests (atol=0.01, rtol=0.01)
2026-07-21 23:05:17 +08:00
ViperEkura f7a16efc9d refactor: extract shared dispatcher header, unify MMA/scalar dispatch format
- Merge 3 duplicated dispatch blocks into single attn_dispatchers.cuh
- Merge compute_num_splits from attn_utils.cuh into dispatcher header
- All dim3 grid/block declarations and <<<>>> launches are single-line
- Production .cu files (35-42 loc) only handle torch wrapping + pybind11
- Test files include dispatcher header directly, removing all #ifndef ASTRAI_NO_MMA duplication
2026-07-21 22:21:39 +08:00
ViperEkura a01e8bbe98 refactor: adopt FA2-style KernelTraits + compile-time causal/mask dispatch
- Introduce KernelTraits<HEAD_DIM, BC, WARPS, STAGES> compile-time config bundle, replacing scattered <KD, NC8, KT2, ...> template params
- Template all MMA and scalar kernels on IsCausal/HasMask bools to eliminate inner-loop runtime branches
- Dispatch to 4-path IsCausal/HasMask kernel variants at entry points based on p.causal_offset and p.use_mask
- Update standalone test files with new kernel signatures, add causal test cases
- Fix duplicate using bf16 in MMA kernels that include attn_mma_utils.cuh
2026-07-21 21:52:46 +08:00
ViperEkura ccf728a1b7 perf: eliminate GPU syncs in contiguous cache write/gather hot paths
- Replace .tolist() calls with _total_len in gather(); move _slot_len updates from per-layer write to once-per-step bind_tasks

- Use torch.as_tensor instead of torch.tensor in decode penalty history construction
2026-07-21 16:40:56 +08:00
ViperEkura f1b4b05d08 feat: add attention dimension dispatch 2026-07-21 12:52:45 +08:00
ViperEkura 0c86c89af4 refactor : align config field names with Hugging Face
- dim -> hidden_size, n_layers -> num_hidden_layers
- dim_ffn -> intermediate_size, n_heads -> num_attention_heads
- n_kv_heads -> num_key_value_heads, max_len -> max_position_embeddings
- norm_eps -> rms_norm_eps, tie_weight -> tie_word_embeddings
- update model, inference, training, scripts, tests, docs
2026-07-20 22:05:31 +08:00
ViperEkura d7ac66fb73 refactor: simplify attention mask handling 2026-07-20 20:36:16 +08:00
ViperEkura a6e920fdb0 Merge pull request #20 from ccx1324/lora-device-fix
fix: LoRA device mismatch and checkpoint resume
2026-07-20 19:38:10 +08:00
ViperEkura 958df58f9d refactor: unify tokenizer encode and apply_chat_template for batch support
- encode(str) single-thread, encode(List[str]) Rust parallel encode_batch
- apply_chat_template accepts single Messages or List[Messages] for batch
- add Message/Messages type aliases at module level
2026-07-20 17:49:26 +08:00
ViperEkura e0f102c4d9 feat: support SFT directly from JSONL without dataset_config.json
- JsonlStore falls back to built-in messages config when no config file found and tokenizer_path is provided
- DatasetFactory.load forwards tokenizer_path to store for SFT/SEQ+jsonl
- assistant turns train, other roles masked, position_ids doc_reset
2026-07-20 17:25:09 +08:00
ccx1324andccx 5a942527b2 fix: inject LoRA before loading checkpoint state_dict
move inject_lora() before load_state_dict in _before_wrap so that
  LoRA adapter weights from a checkpoint are properly restored on
  training resume. Previously, inject happened after load, causing
  lora_A/lora_B keys to be silently ignored (strict=False).

  Co-Authored-By: ccx1324 <2424441089@qq.com>
2026-07-20 17:00:24 +08:00
ViperEkura 37a3036934 refactor: split LoRA param init into local vars 2026-07-20 16:23:47 +08:00
ViperEkura 121a7bf8b4 Merge pull request #19 from ccx1324/lora-device-fix
fix: create LoRA parameters on base weight device instead of CPU
2026-07-20 16:14:30 +08:00
ccx1324andClaude Opus 4.7 a5678c9185 fix: create LoRA parameters on base weight device instead of CPU
When `inject_lora()` replaces Linear layers with LoRALinear after the model
has been moved to CUDA, the new lora_A and lora_B parameters were always
created on CPU, causing a device mismatch error during the forward pass.

Now lora_A and lora_B are created on the same device and dtype as the
parent weight, matching the model's current device.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-07-20 16:11:48 +08:00
ViperEkura 2c50b3cf37 ci: preserve both release wheel artifacts 2026-07-20 15:33:43 +08:00
ViperEkura eee7f54789 docs: sync training and architecture guides 2026-07-20 15:23:30 +08:00
ViperEkura 06eeeead79 refactor: map instruction/input/output to chat roles
- RolloutGenerator._instruction_to_messages builds system/user/assistant list (instruction->system, input->user, output->assistant), replacing single-user-turn concatenation
- Remove _iter_samples helper; _prepare_prompts zips parallel list-of-strings fields directly per the collate_fn contract
- Tests adopt a system-aware chat template and pin the three-field role mapping
- Drop unused imports caught by ruff F401 (torch.Tensor in scheduler.py, iter_raw_records in pipeline.py, Tuple in evaluate_rouge.py)
2026-07-20 13:55:25 +08:00
ViperEkura e8ff7f5321 fix: use batch_per_device for rollout scheduler batch sizing
- train_context.py referenced non-existent cfg.batch_size, replaced with cfg.batch_per_device
- default group_size lowered from 8 to 1: without a group concept (DPO), scheduler batch equals batch_per_device; rollout-based DPO can opt in via extra_kwargs['group_size']>=2
- inline expressions (rollout_batch_size, max_seq_len) extracted for readability
- add tests/trainer/test_online_e2e.py: end-to-end online_dpo via Trainer.train, exercising KV-cache-backed rollout path
2026-07-20 13:32:04 +08:00
ViperEkura a6e1f26cd4 refactor: simplify sample return_logprobs path
- SamplingPipeline.sample gains return_logprobs; both greedy and multinomial paths now share a single log_softmax+gather instead of duplicating the sampling logic
- module-level sample() becomes a thin forwarder instead of re-implementing the three-branch logic
- eliminates ~10 lines of duplicated softmax/gather code; no caller-facing API change
2026-07-20 13:16:18 +08:00
ViperEkura 95c43368ae refactor: unify rollout onto inference engine KV-cache path
- RolloutGenerator now delegates prefill/decode to InferenceScheduler.run_batch (sync API, no background thread), sharing one KV-cache code path with the inference server and eliminating O(n^2) recompute in rollout
- Add sample(return_logprobs=) and Executor.execute_decode(return_logprobs=) to expose behaviour-policy log-probs through the engine; Task gains output_logprobs
- RolloutResult now subclasses RawRollout (adds rewards only), removing duplicated fields
- RolloutRunner.__call__ returns (result, is_fresh) instead of relying on object identity, removing the fragile refresh-detection contract
- Remove O(n^2) generate_responses helper and dead code (_tokenize_prompts, unused old_model arg)
- train_context.py wires InferenceScheduler directly instead of hand-rolling SamplingPipeline
- Tests: +11 covering return_logprobs, run_batch, and KV-cache-backed rollout semantics; 404 pass
2026-07-20 12:52:20 +08:00
ViperEkura 754624acf0 feat: add online rollout framework for RL strategies
- RolloutRunner: generate + score responses with cached re-rollout trigger
- BaseStrategy.__call__ switches online/offline via runner injection
- GRPO/DPO implement prepare_from_rollout; aliases online_grpo/online_dpo
- TrainConfig + train.py add rollout params and CLI flags
- Tests cover generate_responses, RolloutRunner cache, shared __call__
2026-07-20 03:49:56 +08:00
ViperEkura 0b6a17330f feat: add FSDP2Executor using torch.distributed.fsdp.fully_shard API
- New FSDP2Executor registers as 'fsdp2' in ExecutorFactory, using per-module fully_shard() instead of FSDP1 FlatParameter wrapper
- FSDP2 preserves original Parameter objects as DTensors, eliminating use_orig_params=True hack
- FSDP2Executor implements _no_sync via set_requires_gradient_sync, clip_grad_norm via unshard, unwrap_model via DTensor.full_tensor
- Drop **_extra/**_ddp_only_kwargs fallbacks in BaseExecutor/FSDPExecutor/FSDP2Executor, replaced by parallel_mode-aware executor_kwargs dispatch in train.py (ddp-only kwargs only passed when parallel_mode=ddp)
- Export FSDP2Executor in astrai.parallel.__init__
2026-07-20 01:46:25 +08:00
ViperEkura 74b9308883 refactor: pass model_fn/optimizer_fn to executor.prepare
- BaseExecutor.prepare now takes factories and instantiates model via model_fn(), runs before_wrap hook, wraps DDP/FSDP, then builds optimizer/scheduler on the wrapped model
- optimizer/scheduler creation moved into executor.prepare, eliminating the old 'create-then-wrap' hack reliance on use_orig_params=True
- FSDPExecutor/BaseExecutor accept **_extra kwargs to tolerate DDP-only keys (broadcast_buffers, gradient_as_bucket_view) being forwarded via executor_kwargs
- dataloader builds stay external; executor only handles model/optimizer/scheduler
- train_context.py rewritten to load checkpoint state_dict before prepare via a before_wrap closure
2026-07-20 01:32:05 +08:00
ViperEkura e5f9b1a3a9 fix: default max_grad_norm to 1.0 and drop None branch 2026-07-20 01:08:13 +08:00
ViperEkura 31d33ccdf0 chore: bump to 1.3.10 v1.3.10 2026-07-19 16:40:27 +08:00
ViperEkura 88ec786e39 fix: memmap mode=r, tool parser json.loads, greedy decode 2026-07-19 16:38:28 +08:00
ViperEkura 663ef900fc refactor: move sample-id indexing from dataset to store
- Store owns window_size/stride and __getitem__/__len__/sample_window
- Dataset classes become thin delegators binding a Store to a train-type key mapping
- Drop BaseDataset.get_index and the RecordDataset中间类 (window死代码)
- DatasetFactory forces window_size=0 for record datasets so record semantics never get window-tainted
- token_count/num_records split the legacy len() semantics (raw stream length vs record count)
- Update tests to the new .store/.token_count API and window/record mode switching
2026-07-19 16:02:50 +08:00
ViperEkura 7d478a54db docs: update HF org from ViperEk to ViperEkura
- Replace 4 HF links in README.md and README-zh-CN.md to point to ViperEkura
- Update download.py default repo to AstrAI-V1-instruct under ViperEkura
2026-07-19 14:49:58 +08:00
ViperEkura f3eaaef842 refactor: remove redundant strategy/executor code
- Drop BaseStrategy.model_fn (stored but never read)
- Drop model_fn= passed to StrategyFactory.create in train_context
- Simplify FSDPExecutor.clip_grad_norm None branch to delegate to super()
- Remove DDPExecutor._gather_state_dict override (identical to base)
2026-07-19 12:45:58 +08:00
ViperEkura d655b65027 docs: sync architecture/dataflow/training/params with code
- dataflow.md: update DatasetFactory.load signature, stream vs record access, Store._offsets
- architecture.md: add tokenizer to Pipeline, TokenizeTransform class, RecordDataset, Streamable/Recordable mixins, fix GRPOStrategy (old_model/sync_old_model)
- training.md: DPO reduction="sum", GRPO rho_t uses pi_old, gradient_clipping always registered
- params.md: --max_grad_norm default None
2026-07-19 12:33:35 +08:00
ViperEkura 31c22dc043 refactor: deduplicate preprocessing kernel and BFD packing
- Extract shared core (mask building, primary-id extraction, tensorisation, position-id generation) to astrai/preprocessing/core.py; Pipeline and TokenizeTransform both consume it, eliminating ~60% duplicated logic
- Promote BFD _plan to module-level plan_bfd(lengths, max_len) returning pure index bins; BFDPacking.apply and evaluate_ifd._pack_bins both call it, removing the second BFD implementation
- Split Pipeline._flush (49 lines) into _inject_doc_reset_position_ids + _inject_continuous_position_ids + _to_tensors; split Pipeline.run by delegating record iteration to core.iter_raw_records
- Remove dead no-op pop/塞回 in Pipeline.run (L110-111)
2026-07-19 12:27:56 +08:00
ViperEkura 17127f8b3c fix: make tokenizer picklable for spawn multiprocessing
- ChatTemplate: defer Jinja2 compilation to cached_property, exclude compiled template from __getstate__ (its dynamic root function has __module__=None and falls back to __main__, breaking pickle)
- AutoTokenizer: bypass __getattr__ for underscore-prefixed attrs to prevent infinite recursion during unpickle when __dict__ is empty
2026-07-19 11:59:55 +08:00