Commit Graph
70 Commits
Author SHA1 Message Date
ViperEkura c79d34eee1 refactor: simplify training and inference interfaces
- avoid constructing model_fn more than once when reading config
- keep inference package exports focused on public entry points
- rename extra strategy arguments to strategy_kwargs
2026-08-19 20:55:13 +08:00
ViperEkura 398e8a3ea3 refactor: deduplicate low-risk code paths 2026-08-19 16:17:40 +08:00
ViperEkura da6d94492d fix: parse yaml floats with yaml 1.2 schema
- register yaml 1.2 float resolver so scientific notation (2e-5) becomes float, not str
- replaces the decimal-point workaround in train configs
- add containerized training doc under docs/developer
2026-08-13 23:28:06 +08:00
ViperEkura cf4f5ab9f6 feat: add persistent DataLoader workers
- Keep training workers alive between epochs when enabled.
- Avoid invalid prefetch settings for single-process loading.
2026-08-09 11:38:50 +08:00
ViperEkura cbc584470d refactor: centralize logging in astrai.logging, replace ASTRAI_TIMED with log level
- move setup_logging to astrai/logging.py
- timed() now uses logger.isEnabledFor(DEBUG) instead of separate env var
- enable ASTR_LOG_LEVEL=DEBUG to see per-step timing logs
- call setup_logging() in stream_chat.py
2026-08-08 12:39:27 +08:00
ViperEkura 9b7e6c205f feat: add moe auxloss and metrics 2026-08-05 18:12:28 +08:00
ViperEkura d6bfb09863 feat: add grad_snr metric with EMA-based gradient SNR tracking
- add GradSNRTracker to metric_util.py computing SNR = E[g]^2 / Var(g) via per-parameter EMA moments
- add grad_snr_tracker field to TrainContext (instantiated by default)
- register grad_snr in MetricCallback, update tracker on each optimizer step before metrics are recorded
- add grad_snr to default --metrics in train.py CLI
2026-08-01 08:54:44 +08:00
ViperEkura 6db276f37a feat: add Mano manifold optimizer (mano_adamw)
- implement Mano (v2) with axis-rotating tangent projection and manifold normalization, replacing Newton-Schulz iteration
- composite ManoAdamW reuses partition_optimizer_parameters and composite helpers
- register mano_adamw in OptimizerFactory, export Mano and ManoAdamW
- add --mano_momentum and --mano_nesterov CLI options in Optimizer group
- add mano_adamw hyperparameters branch in train.py
- document mano_adamw in params.md
- add tests for single-step projection, axis alternation, factory registration, closure, and resume
2026-08-01 08:51:08 +08:00
ViperEkura 6c76c16480 feat: group train CLI options in --help output
- add GroupedOption/GroupedCommand (no third-party dep) that tags each option with a group label and renders help in labeled sections
- add opt() shorthand wrapping click.option with cls=GroupedOption
- tag all ~55 options into 10 groups aligned with params.md chapters
2026-08-01 08:40:22 +08:00
ViperEkura 25c9e81b2b refactor: keep muon_adamw as default optimizer and drop nora docs
- revert CLI/create_optimizer/display defaults to muon_adamw
- revert README, README-zh-CN, params.md to pre-merge state
2026-08-01 07:51:51 +08:00
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