Commit Graph
123 Commits
Author SHA1 Message Date
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 d9a0c72149 feat: store metric logs inside each checkpoint dir, remove log_dir config 2026-07-28 00:22: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 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 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 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 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 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 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 e5f9b1a3a9 fix: default max_grad_norm to 1.0 and drop None branch 2026-07-20 01:08:13 +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 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 d7695b40e3 feat: make max_grad_norm optional (None disables clipping)
- TrainConfig.max_grad_norm defaults to None
- executor.clip_grad_norm returns grad norm without clipping when None
- train.py --max_grad_norm defaults to None
2026-07-19 00:08:18 +08:00
ViperEkura b133fc9c07 refactor: split Store into StreamStore and RecordStore
- StreamStore: fetch(begin, end, key) for stream access (SEQ/SFT)
- RecordStore: mixin with fetch_record(i, key) for record access
- H5Store/MmapStore/JsonlStore now dual-inherit both (C3 MRO)
- JsonlStore supports lazy mode via processor= (no TokenizeTransform)
- RecordDataset base class holds processor, DPO/GRPO simplified
- dpo_tokenize pure function for on-the-fly JSONL tokenisation
- DatasetFactory builds processor for jsonl+record datasets
- train.py passes tokenizer_path=param_path uniformly
- progress: len(dataset) returns sample count (stream=windows, record=records)
- json no longer auto-detected as jsonl format
2026-07-18 23:04:31 +08:00
ViperEkura a74e5b91a3 feat: add record-mode to Store for DPO/GRPO
- Store gains fetch_record/num_records alongside stream fetch/__len__
- save_bin/load_bin support per-record offsets via record_keys param
- H5Store/MmapStore/JsonlStore all support dual stream+record access
- DPODataset/GRPODataset use fetch_record, no cross-record concat
- dpo_collate_fn + collate_fn wired through TrainConfig
- fixes attention context leakage in DPO from windowed concatenation
2026-07-18 21:02:29 +08:00
ViperEkura 28886e4241 fix: make system prompt optional across scripts
- stream_chat: default empty system_prompt, single-turn mode
- generate_batch: drop hardcoded system role
- generate.py: preserve original fields in messages branch
  and use response_key for the output column name
2026-07-18 14:10:37 +08:00
ViperEkura a24a7b4da5 perf: merge decode batch for 10x throughput
- merge all active decode tasks into single forward pass (was grouped by next_pos)
- add per-task write_positions to ContiguousCacheView for correct KV writes
- override ContiguousCache.task_cached (base returned 0, caused prefill loops)
- add --cache_len/--frequency_penalty/--rep_window to generate.py
- chunked batch processing with tqdm progress

bench (1.2B model, 128 prompts, 64 tok, batch=128):
  before: 77.2s, ~111 tok/s
  after:   7.1s, ~1210 tok/s (10.9x)
2026-07-18 08:50:46 +08:00
ViperEkura f7df02f9a3 feat: add --num_samples to batch generation script 2026-07-18 01:13:38 +08:00
ViperEkura ee450686f3 fix: add option permutation to MMLU eval
- Few-shot examples now include subject preamble (consistent format)
- Add --seed flag for option permutation (default 0, -1 to disable)
- Shuffles A/B/C/D positions per-question to neutralise positional bias
2026-07-18 00:14:34 +08:00
ViperEkura 2565755e45 refactor: switch eval datasets to HuggingFace source
- Replace GitHub/berkeley direct downloads with HF datasets API
- MMLU: cais/mmlu (all config), map val->validation split, write per-subject CSV
- HumanEval: openai/openai_humaneval
- IFEval: google/IFEval
- Enables HF_ENDPOINT mirror for faster downloads in CN
2026-07-18 00:09:31 +08:00
ViperEkura d08a92c7bd feat: add frequency penalty to inference sampling pipeline
- Add FrequencyPenaltyStrategy (logit -= penalty * count)
- Per-task rep_window for penalty history lookup
- Wire through engine, task, executor, API layer
- Add --frequency_penalty and --rep_window to stream_chat.py
- 9 unit tests for frequency penalty strategy
2026-07-17 21:28:31 +08:00
ViperEkura c17aa0dc54 fix: eval script bugs and add missing features
- evaluate_mmlu: fix double few-shot injection (build_prompt no longer
  adds few-shot, apply_chat handles it once)
- evaluate_humaneval: fix pass@k k-filtering to be per-problem instead
  of using first problem's n globally; reuse ProcessPoolExecutor across
  problems; fix closure UnboundLocalError in test_one; handle None in
  report when k > n
- evaluate_ifd: remove dead code (score_plain/score_messages); add
  multi-file/directory input support with --input_path/--output_dir;
  add summary.json aggregation and --max_samples; add --dtype flag
- evaluate_ppl: add --device and --dtype flags (was hardcoded to cuda)
- evaluate_ifeval: fix docstring path (scripts/tools -> scripts/eval)
- analyze_weights: add --output JSON export; fix dead code filter
  ("_norm" not in r was always True)
2026-07-17 14:02:58 +08:00
ViperEkura b12b24eadc feat: rewrite evaluate_ppl with token-level loss and multi-file support
- Support multiple input files, glob patterns, and directory input
- Add --token_level flag: per-record token_ids + log_probs JSONL output
- Add --max_samples for random subsampling per file
- LossAccumulator: streaming mode (histogram-based percentiles, low memory) vs exact mode (full token list)
- Token type analysis (ascii/cjk/non_ascii/special) when token_level=True
- Fix token_ids/log_probs alignment (shift offset)
- Cache frozenset(stop_ids) outside loop for performance
- Aggregate stats: mean/median/ppl/p50/p90/p95/p99
- Summary JSON with all datasets in one file
2026-07-17 13:14:15 +08:00
ViperEkura 84ed2327f5 feat: add --resume flag to decouple weight loading from training resumption
- Add --resume bool flag to train.py CLI
- --param_path always loads weights only by default
- --resume restores epoch, consumed_samples, optimizer & scheduler
- Checkpoint.load() now preserves full meta dict
- Update test_early_stopping to use new param_path/resume API
2026-07-16 14:23:23 +08:00
ViperEkura bb175fda91 fix: resume optimizer LR, step display, and consumed_samples alignment 2026-07-15 08:59:52 +08:00
ViperEkura 2c7a71a9c0 refactor: separate old policy and ref model in GRPO strategy
- Split single ref_model into old_model (importance sampling ratio) and ref_model (frozen KL regularizer)
- Move ref_model/old_model creation from strategy __init__ to TrainContextBuilder, pass as explicit parameters
- Remove periodic sync_ref_model + sync_interval; add sync_old_model for external rollout loop to call
- DPOStrategy also receives ref_model from builder
- Fix std to use unbiased=False (population std per GRPO paper)
- Remove redundant tests (test_grpo_kl_zero_at_init, test_grpo_no_sync_interval_param)
- Remove --grpo_sync_interval CLI arg
2026-07-14 20:03:45 +08:00
ViperEkura c8567a6f65 fix: exclude embedding, lm_head, bias, and norm params from Muon optimizer, use AdamW 2026-07-08 19:42:20 +08:00
ViperEkura 8035be9b1f fix: make MuonMix inherit from torch.optim.Optimizer 2026-07-08 17:03:05 +08:00
ViperEkura c50adbaac0 feat : replace AdamW with MuonMix (Muon + AdamW) optimizer
- Muon for 2D matrix params, AdamW for 1D (norm/bias/embed)
- MuonMix wrapper handles combined step/zero_grad/state_dict
- New CLI args: weight_decay, muon_momentum, muon_nesterov, muon_ns_steps, muon_adjust_lr
- Removed adamw_beta1/adamw_beta2/adamw_weight_decay
- Moved optimizer/strategy params from signature to **kwargs
2026-07-07 14:10:36 +08:00
ViperEkura f0cd0134c6 fix : update benchmark for v1.3.8 cache API, add argparse and cache type switch
- Replaced old KVCage API with PageCache/ContiguousCache
- Added --cache contiguous|paged switch for decoding comparison
- Added argparse for all params (batch/prompt/gen/device/dtype)
- Fixed PageCache decode crash by extending pages for full sequence
2026-07-05 20:30:26 +08:00
ViperEkura db9b39b084 fix: resolve IFD token-set asymmetry and support single-token answers
- Sentinel-anchored unconditional pass: both branches now predict the same N response tokens
- Single-token responses (rl=1) fully supported
- ctx_len tracked per sample; skip_reason replaces silent None
- --per_token flag for per-token IFD breakdown
2026-07-05 17:48:26 +08:00
ViperEkura 599a51f4f7 fix: reliable test timeout, separate generate/test phases, dynamic pass@k
- Replace SIGALRM+exec() with subprocess.run(timeout=) for test execution
- Add --test_only flag to skip generation and test existing completions
- Add --generate_only flag for generation-only runs
- Derive pass@k values from num_samples (filter k > n)
- Support loading completions from array JSON (not just JSONL)
2026-07-05 08:47:30 +08:00
ViperEkura 17d6eaa2f2 refactor: rewrite humaneval evaluation with functional pipeline design
- fix KeyError race condition in inference cache touch()
- EvalConfig dataclass for centralized configuration
- load->generate->extract->test->score->report pipeline
- two-phase generation+testing for max GPU utilization
- signal-based SIGALRM timeout protection for code exec
- suppress subprocess stdout/stderr pollution
2026-07-05 07:58:28 +08:00
ViperEkura 2d908639e9 feat : add ROUGE evaluation script (manual impl, no deps)
- ROUGE-1/2 via n-gram overlap (Counter)
- ROUGE-L via LCS (DP)
- CLI: python scripts/eval/evaluate_rouge.py --data_path ... --output ...
- Library: compute_rouge(ref, cand) -> dict of precision/recall/f1
2026-07-05 01:15:01 +08:00
ViperEkura c7158418dd perf: add BFD bin-packing and custom attention mask to IFD batch scoring 2026-07-04 18:58:13 +08:00
ViperEkura 4d3c9341c1 refactor: rewrite IFD evaluation with clean three-layer architecture 2026-07-04 18:33:51 +08:00
ViperEkura 1adca39cd8 fix: handle long sequences and optimize IFD computation 2026-07-04 08:35:45 +08:00
ViperEkura 204873fa2f fix: handle long sequences and optimize IFD computation 2026-07-04 08:23:32 +08:00
ViperEkura a5c1de6b1b feat: add model_path temperature top_p top_k max_tokens system_prompt args to stream_chat 2026-07-04 07:33:32 +08:00
ViperEkura 70c0e5de90 refactor: merge validation into MetricCallback, simplify progress bar to optimizer steps
- Remove separate ValidationCallback, merge into MetricCallback
- Progress bar now tracks optimizer steps instead of micro-steps
- Remove unused log_interval config field and CLI flag
- Fix validation all_reduce: use SUM(loss, count) instead of AVG
- Simplify metric logging: always log every optimizer step
- Add grad_norm display to progress bar
2026-07-03 21:43:08 +08:00
ViperEkura aabb0d83e9 refactor : replace iteration with consumed_samples
- Replace context.iteration with consumed_samples (global sample count)
- Add optimizer_step property derived from consumed_samples
- Checkpoint meta.json stores consumed_samples, drops iteration
- CLI --start_batch renamed to --start_samples (per-rank samples)
- Checkpoint dir naming: epoch_X_step_Y instead of epoch_X_iter_Y
- Metric log entries use step and consumed_samples fields
- Backward compat removed (old iteration checkpoints unsupported)
2026-06-30 18:42:42 +08:00
ViperEkura 44579ea6dc refactor : metric 日志改为以 optimizer step 为单位,默认每步记录
- log_interval 默认 100 -> 1,语义从 batch iteration 改为 optimizer step
- step 指标从 on_batch_end 移到 on_optimizer_step,不受梯度累积影响
- JSONL 条目新增 step 字段,保留 iter
- flush 落盘仍在 on_batch_end
2026-06-30 15:12:31 +08:00
ViperEkura 0f1fcb079f refactor : grad_norm 指标简化,clip_grad_norm 移至 executor
- metrics 默认加入 grad_norm,移除 grad_std/max/min/mean/nan_num
- grad_norm 默认返回总 L2 范数,per_param=True 返回各参数范数
- clip_grad_norm 从 callback 移至 BaseExecutor/FSDPExecutor
- FSDPExecutor 覆盖为 model.clip_grad_norm_() 保证分布式正确
- ctx_get_grad_norm 改为读取 context.grad_norm
2026-06-30 14:59:43 +08:00