Commit Graph
810 Commits
Author SHA1 Message Date
ViperEkura deb2d7e127 refactor: rebuild KV cache with three-layer separation architecture
- Replace CacheView/ContiguousCache/PageCache with SGLang-inspired design: KVStorage (flat token-level NHD buffers [n_layers, size, H, D]), ReqToTokenPool (index table [req_idx, pos] -> token_slot), Allocator + PrefixCache (slot allocation with LRU and prefix sharing)
- Add KVCache as pure dataclass passed to model: k_buffer, v_buffer, req_to_token, req_pool_indices, seq_lens, out_cache_loc
- PagePool orchestrates all three layers, supports contiguous mode (pre-allocated per-request blocks, default) and paged mode (page_size=1 or >1 with dynamic allocation and prefix caching)
- Attention layers now do raw buffer indexing instead of opaque write/gather method calls on CacheView objects
- Update executor.bind_tasks signature: seq_lens list + start_pos
- Rename paged_cache -> kv_cache throughout model/ and inference/
2026-07-30 17:19:06 +08:00
ViperEkura fc47319240 refactor: simplify BaseFactory and separate ModelFactory from AutoModel
- Extract _resolve_base_type and _validate_component as module-level helpers
- Replace ForwardRef._evaluate private API with eval in module namespace
- Remove broad except Exception in __init_subclass__, _component_base always set
- Replace direct _entries mutation in strategy.py with register() call form
- Remove dead TOKENIZER_CLASSES registry from AutoTokenizer
- Extract ModelFactory(BaseFactory[nn.Module]) as pure factory
- AutoModel now inherits only nn.Module, no factory state
- Move @AutoModel.register to @ModelFactory.register in transformer.py and encoder.py
2026-07-30 09:38:20 +08:00
ViperEkura 22cf798d81 feat: add field and model validators to config classes
- TrainConfig: enum validators (strategy, parallel_mode, backend, start_method, compile_mode), positive/non-negative/range validators, model_validator requiring reward_model_fn for online RL strategies
- AutoRegressiveLMConfig/EncoderConfig: attn_type, ffn_type enum validators
- ProcessingConfig: packing_strategy, truncation_mode enums, positive int validators
- OutputConfig: storage_format, position_ids_mode enum validators
2026-07-30 08:41:14 +08:00
ViperEkura 164be9708b refactor: migrate config system to Pydantic dataclasses
- Replace hand-rolled BaseConfig (from_dict/to_dict/_coerce/_unwrap_optional) with pydantic.dataclasses
- from_dict now uses cls(**d), to_dict uses dataclasses.asdict + json.dumps filter
- TrainConfig: required fields are now truly required (no default=None), delete manual validate()/__post_init__
- Remove dead required() helper and metadata={'help': ...} annotations
- Fix gradient_checkpointing_modules type from List[str] to List[type]
- Add pydantic>=2.0 as direct dependency in pyproject.toml
- Add numpy-style Parameters docstrings to all config classes
- Enable use_attribute_docstrings in BaseConfig for schema generation
- LoRAConfig also migrated to pydantic dataclass
2026-07-30 08:25:32 +08:00
ViperEkura 6a97524db4 refactor: inline parallel utils into executor module
- Move create_ref_model from astrai/parallel/utils.py into executor.py
- Remove unused ColumnParallelLinear/RowParallelLinear (module.py)
- Update imports in strategy.py and train_context.py
- Drop unused astrai.parallel.utils and astrai.parallel.module
2026-07-30 07:54:54 +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
ViperEkura bcaa2d1ae0 fix: FSDP unwrap_model collective op and None guard
- unshard() and full_tensor() are collective ops, all ranks must participate
- Old code returned None on non-rank-0 before calling unshard, causing deadlock
- Fix: all ranks unshard/full_tensor, only rank-0 keeps the result
- Move create_ref_model to parallel/utils.py, accept executor+model directly
- Guard create_ref_model and sync_old_model against None on non-rank-0
2026-07-29 23:41:10 +08:00
ViperEkura 8206afefd9 fix: FSDP clip_grad_norm and default reshard_after_forward=False
- FSDP params are DTensors sharded across ranks
- torch.nn.utils.clip_grad_norm_ computes LOCAL norm only
- Each rank would clip by a different factor, causing gradient divergence
- Fix: compute local norm, all-reduce squared sum, sqrt for global norm
- Default reshard_after_forward=False (forward then backward makes reshard redundant)
- Reduces per-step time by ~19% (1033ms to 839ms on 2xL20)
2026-07-29 23:27:10 +08:00
ViperEkura 646b1b0f46 refactor: replace FSDP with FSDP2 as default parallel backend
- Remove FSDPExecutor (FullyShardedDataParallel wrapper)
- Rename FSDP2Executor to FSDPExecutor, register as 'fsdp'
- Remove 'fsdp2' from CLI choices, make 'fsdp' the default parallel_mode
- Pass after_wrap to executor.prepare for compile-after-wrap ordering
- Update architecture.md, params.md, AGENTS.md references
- FSDP2 uses per-module fully_shard: no FlatParameter, better compile compat
2026-07-29 23:09:37 +08:00
ViperEkura 8150ab6c32 feat: add torch.compile CLI option for training
- Add --compile flag (default/reduce-overhead/max-autotune)
- Apply torch.compile in _before_wrap before DDP/FSDP wrapping
- Profiling shows MFU 85.5% -> 88.5% (+3%), time -3.2%, memory -7.9%
2026-07-29 22:06:51 +08:00
ViperEkura 0b0693a0a2 fix: make ChatTemplate picklable for spawn multiprocessing
- Add __getstate__/__setstate__ to drop cached _compiled Jinja2 template
- Jinja2 Template.root_render_func is a dynamic closure unpicklable by reference
- cached_property rebuilds the template lazily on first render after unpickle
2026-07-29 13:24:13 +08:00
ViperEkura 115192c67c refactor: remove H5 storage backend in favor of mmap bin
- Remove H5Store, H5Writer, save_h5/load_h5 and h5py dependency
- MmapStore (bin) is the sole pre-tokenized storage backend
- Move setup_logging after imports to fix E402 in __init__.py
- Clean up unused imports across test files
- Move inline test imports to file top
2026-07-29 12:50:27 +08:00
ViperEkura c2b04d8458 refactor: align generate.py params with engine API
- Remove --max_tokens, let scheduler use max_seq_len - prompt_len
- Rename --cache_len to --max_seq_len to match engine naming
- Unify sampling defaults to 0.8/50/0.95
2026-07-29 09:47:53 +08:00
ViperEkura db487ab48b feat: append EOS to response in IFD evaluation
- Add EOS token at end of response in both conditional and unconditional passes so model also predicts when response should end
- New --append_eos/--no-append_eos CLI flag (default: enabled) with graceful fallback when tokenizer has no EOS
2026-07-28 22:22:59 +08:00
ViperEkura a95794d3db perf: use Rust-native DecodeStream for O(n) streaming decode
- Replace hand-rolled StreamDecoder (O(n^2) full-history re-decode per token) with tokenizers.decoders.DecodeStream
- Keep O(1) bounded token buffer internally via prefix drain instead of accumulating all token IDs
- Simplify flush_remaining to no-op since stream always emits completed text per step
- Benchmark on 8000 tokens: 2305ms -> 3.9ms (~592x speedup)
2026-07-28 14:32:10 +08:00
ViperEkura 39f84f3b4c refactor: move signal_handler from parallel/ to top-level for broader reuse 2026-07-28 10:36:17 +08:00
ViperEkura 9f7cf50c56 fix: keep metric logs cumulative instead of segmental in each checkpoint 2026-07-28 09:18:48 +08:00
ViperEkura d9a0c72149 feat: store metric logs inside each checkpoint dir, remove log_dir config 2026-07-28 00:22:29 +08:00
ViperEkura 5ab18bec48 fix: correct epoch computation on resume to avoid redoing whole epoch 2026-07-28 00:01:29 +08:00
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