ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
612 Commits 1 Branch 14 Tags
5b67d5865aea9ddc58fb72b27b5a7ea787d3d761
Commit Graph
87 Commits
Author SHA1 Message Date
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 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 d9a0c72149 feat: store metric logs inside each checkpoint dir, remove log_dir config 2026-07-28 00:22:29 +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 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 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 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 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 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
ViperEkura bf09a35c95 feat: optimizer 参数分组，bias/norm 不做 weight decay 2026-06-27 16:30:34 +08:00
ViperEkura 01d2da2893 feat : 训练支持 --schedule_type 及对应调度器参数 
- --schedule_type 可选 cosine/sgdr/wsd，默认 cosine
- --min_rate 统一控制最小 LR 比率
- --cycle_length / --t_mult 用于 sgdr
- --stable_steps / --decay_steps 用于 wsd，自动计算默认值
 2026-06-22 10:35:56 +08:00
ViperEkura 39985840c7 refactor : neftune_alpha 在 Embedding 构造时传入，由模型配置链路负责 
- BaseModelConfig 添加 neftune_alpha 字段 (默认 0.0)
- Embedding.__init__ 接受 neftune_alpha 参数，不再外部 set
- AutoRegressiveLM / EmbeddingEncoder 从 config 传入 neftune_alpha
- train.py 将 CLI 参数注入 config 后再创建模型
- TrainContextBuilder 移除 neftune 设置（不再是其职责）
 2026-06-19 14:23:27 +08:00
ViperEkura 7348bac6ab fix: 规范 generate.py 命令行接口 
- generate.py 清理描述文字，help 统一标注默认值
- max_tokens 默认改为 None，回退 model config max_len
- evaluate_ppl.py 同步清理描述文字
- params.md 同步 max_tokens 默认值
 2026-06-19 14:03:02 +08:00
ViperEkura 1818d06576 feat: 新增 IFD 数据质量评分工具, 移动 ppl 至 eval 
- 计算指令遵循难度分数用于数据筛选
- IFD = 条件交叉熵 / 无条件交叉熵
- perplexity 移至 scripts/eval/
 2026-06-16 22:03:45 +08:00
ViperEkura 4e8d1ee24e feat: 新增 IFEval 指令遵循评测 
- 实现 25 种正则约束 verifier
- 将评测脚本从 scripts/tools/ 移至 scripts/eval/
 2026-06-16 21:57:34 +08:00
ViperEkura 445378667f feat : NEFTune 噪声注入 + label_smoothing 默认值修正 
- Embedding.forward 训练时注入 randn 噪声，缩放系数 neftune_noise_alpha / sqrt(seq_len)
- TrainConfig.neftune_alpha 通过 config 传递（默认 0=关闭）
- TrainContextBuilder 将 config.neftune_alpha 写入 embed_tokens
- --neftune_alpha CLI 参数（典型值 5.0）
- label_smoothing 默认值 0.05 -> 0.0
 2026-06-11 15:32:43 +08:00
ViperEkura 6ae1828449 refactor : 清理工厂和配置系统中的死代码与冗余抽象 
- 删除 Registry 中未使用的 category/priority 字段，_entries 简化为直接存储类引用
- 修正 __init_subclass__ 避免叶子类（AutoRegressiveLM 等）创建空注册表
- 删除 5 个工厂的薄 create() 覆写，统一使用 BaseFactory.create(name, *args, **kwargs)
- 删除 3 处零调用的 available_types/available_strategies 别名死代码
- 删除零调用的 BaseModelConfig.to_file 死代码
- 将 BaseConfig.from_json/to_json 重命名为 from_file/to_file，消除与子类重复
- 移除两个 inference builder 中总是被覆写的 prompt_tokens=0
 2026-06-07 11:39:50 +08:00
ViperEkura 5e73ca20aa feat : train CLI 新增 val_split/val_step/metrics/log 参数 
- --val_split 从训练集按比例切分验证集
- --val_step 控制验证间隔 optimizer step 数
- --metrics 自定义日志指标列表，默认 loss lr
- --log_dir / --log_interval 控制日志输出目录和频率
 2026-06-03 14:31:22 +08:00
ViperEkura 438dc10391 fix : MMLU eval 使用 chat template 格式匹配 SFT 训练数据 
- 原 prompt 为纯文本格式，与 SFT chat template 不匹配导致模型输出随机
- 新增 apply_chat() 将 MMLU prompt 包装为 user/assistant 对话格式
- choice_text 改为单字母（去掉空格前缀）适配模板输出
- 5-shot 时 few-shot 示例作为独立 user/assistant 轮次插入
 2026-06-03 11:59:42 +08:00
ViperEkura 615ba5d8ef feat : 新增 HumanEval pass@k 代码生成评测 
- InferenceEngine.generate() 批量生成 n 个补全
- 正则提取函数体 + 停止符截断
- multiprocessing sandbox 执行 + timeout 保护
- 标准无偏 pass@k 公式 (1, 10, 100)
 2026-06-03 10:52:32 +08:00
ViperEkura 9b416c1bbb refactor : 并行启动 Strategy 模式重构，local_rank 解耦 
- setup_parallel 接收 local_rank 参数，不再读环境变量推导
- TorchrunStrategy 从 env 读取 LOCAL_RANK，LocalStrategy 用 rank
- _detect_launcher() 分级检测替代内联 RANK 检查
- _run_single_rank 统一入口，消除 _run_single/_run_multi 重复
- 优雅退出：except BaseException 终止子进程并 re-join
- gradient_checkpointing_modules 判定提取到外部变量
 2026-06-02 11:22:24 +08:00
yegroup001 0deee48602 feat : 训练脚本新增 gradient_checkpointing 与多机 DDP 参数 2026-06-02 01:01:00 +08:00
ViperEkura 69207e2c57 refactor : 基于声明式 JSON 配置的预处理管线重构 
- 用工厂注册的 MaskBuilder（chat/instruction/text）替换硬编码的 _transform_* 方法
- mask 规则以 role-to-action 映射声明在配置中，与 chat_template 完全解耦
- 单次编码 + role-span 追踪替代两次编码 + 长度差计算 mask 的方式
- 支持多轮对话训练：所有 assistant 轮次参与训练，而非仅最后一轮
- 新建 astrai.preprocessing 包（builder.py + pipeline.py），删除 astrai/preprocess.py
- CLI 精简为 --config 参数，所有参数通过 PipelineConfig JSON 配置
- 新增 PipelineConfig、InputConfig、ProcessingConfig、OutputConfig dataclass
- 文档：assets/docs/preprocessing.md
- 27 个测试覆盖 mask builder、pipeline、配置序列化、工厂注册
 2026-05-30 20:45:09 +08:00
ViperEkura 138c5bcc08 feat : 添加 JSONL 预处理管线 
- Pipeline 模板, Reader 加 transform 加 Writer 可组合
- 自动检测 JSONL 格式, 支持 messages 文本 prompt 加 response 三种
- chat 数据通过 apply_chat_template 适配, 自动生成 loss_mask
- 输出对齐 Store 和 DatasetFactory, 直接用于训练
- 默认 bin 格式, CLI 入口 scripts/tools/preprocess.py
 2026-05-30 17:12:42 +08:00
ViperEkura a923e0a23a fix : 修复 MMLU 评测脚本数据源和依赖 
- 数据源改为 Berkeley data.tar（GitHub zip 不含数据文件）
- urllib 替换为 requests，支持代理下载
- zip 解压替换为 tar，增加目录 flatten 逻辑
- 添加 model.eval() 确保推理模式正确
 2026-05-30 16:51:24 +08:00
ViperEkura 2d5dc93b3d fix : 修正类型标注与统一 CLI 参数命名 
- AutoRegressiveLM.forward 返回类型标注 -> Dict[str, Tensor]
- EmbeddingEncoder 移除冗余 position_ids 自动创建
- CLI 脚本模型目录参数统一为 --param_path
 2026-05-27 20:49:44 +08:00
ViperEkura 4145d35e3c refactor: 检查点加载重构，路径替代对象传递 
- model: nn.Module -> model_fn 工厂函数，spawn 边界只传字符串
- Trainer.train(resume_dir=path) — Checkpoint 不再通过 pickle 传递
- TrainContextBuilder.with_resume_dir(path) — 自动检测 meta.json 分流 resume/from-scratch
- CheckpointCallback: 拆分 state_dict 收集（全 rank）与磁盘写入（rank-0），修复 FSDP 死锁
- serialization: load_torch 支持 broadcast，消除 _load_extra/_load_torch_broadcast
- optimizer/scheduler 恢复逻辑内联到 build()，在 executor.prepare() 之后执行
- pyproject.toml: ruff exclude build/ 避免 CI 扫描构建产物
 2026-05-27 20:15:29 +08:00