ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
628 Commits 1 Branch 14 Tags
04899a2b159ff293e9eacebf2f2608fb78d28fd0
Commit Graph
60 Commits
Author SHA1 Message Date
QueenAmish 04899a2b15 Make Nora+NAdamW the default optimizer 2026-07-31 23:16:39 +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 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 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 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 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 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 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 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
ViperEkura 836e02a166 docs: 同步 architecture/inference/training 文档至实际代码，CLI 补充 fsdp 选项 
- 修正 ProtocolHandler 架构：concrete + ResponseBuilder(ABC) 策略模式
- 修正训练循环 scheduler.step() 在 sync_gradients 块内
- 修正组合/聚合关系：注入组件改为 o--，删除不持有引用的关联
- --parallel_mode CLI choices 加入 fsdp
- nprocs > 1 且 parallel_mode=none 时 raise error
 2026-05-26 19:37:00 +08:00
ViperEkura 1d26aa2e93 fix: 禁用DDP static_graph避免PyTorch 2.7.1下no_sync与backward冲突 
- static_graph=True时DDP.no_sync() + loss.backward()触发expect_autograd_hooks_内部断言
- PyTorch 2.7.1中no_sync上下文切换与静态图hook状态管理存在兼容性bug
- 将static_graph设为False恢复梯度累积正常执行
- find_unused_parameters保持False（模型无不参与计算的参数）
 2026-05-26 15:08:01 +08:00
ViperEkura a548d4553e fix: 断点续训恢复优化器/调度器状态及采样器剩余长度 
- 使用Checkpoint.load()替代手动加载model.safetensors，恢复optimizer/scheduler状态
- TrainContextBuilder从checkpoint.extra恢复优化器和调度器state_dict
- ResumableDistributedSampler.__len__返回剩余样本数而非总数
- 训练前对state_dict置空避免mp.spawn pickle 7GB大对象
 2026-05-26 13:50:25 +08:00
ViperEkura 3ab4f237e5 refactor: 重构训练后端为 Executor 模式 
- backend.py → executor.py，BaseTrainingBackend → BaseExecutor
- 新增 NoneExecutor（单卡）和 DDPExecutor（DDP，world_size=1 自动降级）
- 新增 GradientState 分离梯度同步状态，AccumOptimizer/AccumScheduler 包裹拦截
- 新增 astrai/protocols.py：OptimizerProtocol/SchedulerProtocol 结构子类型
- TrainContext.backend → executor，TrainConfig 移除 parallel_wrapper/state_dict_fn，新增 parallel_mode/executor_kwargs
- 训练循环用 accumulate() 包裹，on_optimizer_step 命名约定=gate
- scripts/tools/train.py 移除 ddp_wrap/prepare_checkpoint，新增 --parallel_mode
 2026-05-24 20:35:44 +08:00
ViperEkura 7621f05d3f docs: AdamW beta 默认值改为 (0.9, 0.95) 
- 与 Muon 优化器的 AdamW 子优化器保持一致
- 同步更新 train.py/training.md/params.md/README
 2026-05-17 17:08:31 +08:00
ViperEkura 97c7ac0f4f refactor: Transformer更名为AutoRegressiveLM并新增EmbeddingEncoder 
- AutoRegressiveLM 注册名改为 autoregressive_lm
- 新增 EmbeddingEncoder 支持 mean/cls/last pooling
- ModelConfig 增加 pooling_type / normalize_embeddings 字段
- 导入、注释、测试全部同步更新
 2026-05-17 15:29:20 +08:00
ViperEkura c241a5dcef refactor: 优化并行训练配置与启动管理 
- 配置新增 start_method 支持 spawn/fork/forkserver 选择
- 启动方式 mp.spawn 改为 mp.start_processes，支持 daemon=True
- validate() 改为基于 metadata 的反射式校验，不再硬编码字段列表
- CLI 新增 --start_method 参数
 2026-05-17 12:33:10 +08:00
ViperEkura 8a11a7d444 fix: 修复训练脚本两处参数传递问题 
- prepare_checkpoint 增加 DDP 判断，单卡时不访问 .module
- dpo_beta 改为 beta，对齐 DPOStrategy 参数名
 2026-05-17 11:04:40 +08:00
ViperEkura e1638a7ade fix: 修正AdamW超参数默认值与文档示例 
- 交换adamw_beta1/adamw_beta2默认值：beta1=0.95, beta2=0.99
- label_smoothing默认值改为0.05
- 文档示例统一更新：train_type=pt, weight_decay=0.01
- 移除文档中过时的strategy default标注
 2026-05-16 22:46:17 +08:00
ViperEkura f91bfee33e refactor: Config序列化统一BaseConfig基类 
- 新增astrai/config/base.py，提供to_dict/from_dict基类
- 统一命名：load/save → from_file/to_file
- Checkpoint.meta合并训练配置到meta.json
- sys.stderr.warn → warnings.warn
- from_file改为classmethod
 2026-05-16 22:06:39 +08:00
ViperEkura d7a7f570ed refactor: 训练循环改为两重迭代并统一参数命名 
- 训练循环从三重(epoch→batched→batch)改为二重(epoch→batch)
- batch_size → batch_per_device, accumulation_steps → grad_accum_steps
- scheduler 移入 step block 对齐 optimizer 更新步
- GradientClippingCallback 改用 on_step_begin 避免零梯度裁剪
- 移除 _train_impl 误导性的 -> Checkpoint 标注
- total_steps 修除为向下取整并精简为一行
- warmup_steps 改为 warmup_ratio (默认0.05)
 2026-05-16 21:27:35 +08:00
ViperEkura 026d1fc33d fix: total_steps 改用 ceiling 匹配实际步数 
原公式全用 floor 少算 optimizer step，改用逐层 ceiling
（ceil_div via (a+b-1)//b）对齐 DDP sampler padding +
DataLoader drop_last=False 尾批 + batched 尾组截断。
 2026-05-16 17:53:18 +08:00
ViperEkura 7242eedbf4 fix: 学习率调度按 optimizer step 计数并防止 warmup 越界 
- total_steps 除以 accumulation_steps，匹配 optimizer.step() 频率
- warmup_steps 用 min 截断，避免 lr_decay_steps 为负
 2026-05-16 17:07:36 +08:00
ViperEkura 9d5e9fa6c4 perf: DDP 加 gradient_as_bucket_view/static_graph/broadcast_buffers，AdamW fused 
- gradient_as_bucket_view=True 零拷贝梯度归并
- static_graph=True 跳过每轮 bucket 重建
- broadcast_buffers=False 省 buffer 广播
- AdamW fused=True 融合优化器 kernel
 2026-05-15 15:30:24 +08:00
ViperEkura b98c9cefdc refactor: 移除 device_ids 参数设计，统一通过 CUDA_VISIBLE_DEVICES 控制 GPU 分配；更新 README 训练示例 
- setup.py: 移除 device_ids 参数，setup_parallel 直接用 rank 作为设备索引
- train_config.py: 移除 device_ids 字段
- trainer.py: 不再传递 device_ids
- train.py: ddp_wrap 用 get_rank() 直接取值
- README.md, README-zh-CN.md: 训练示例改为多行命令风格，去掉参数表格
 2026-05-09 14:55:43 +08:00