Commit Graph
69 Commits
Author SHA1 Message Date
ViperEkura 350e4a1849 feat: add online ppo with value-model critic and gae advantages
- register online_ppo train type backed by PPOStrategy: token-level clipped surrogate over GAE advantages plus masked value regression against rollout-pinned returns, with explained-variance metrics
- fold the reference-KL penalty (k3 estimator) into per-token rewards before GAE and pin advantages/returns on RolloutResult so replayed gradient steps optimize fixed targets
- add self-contained ValueModel critic with a zero-initialized value head and backbone warm-started from policy weights; AutoRegressiveLM stays untouched and trunk parity is pinned by tests
- step the critic's own optimizer outside the policy-version lock with the same max_grad_norm clipping as the policy
- persist critic state as value_model.pt/value_optimizer.pt checkpoint extras; resume restores it, fails loudly when missing, and the train.sh completeness check requires the extras for online_ppo configs
- extract shared rollout sequence/logprob helpers from GRPO (behavior unchanged) and add ppo_gamma/ppo_gae_lambda/ppo_vf_coef CLI options
2026-09-05 01:59:50 +08:00
ViperEkura ae7fc3059a refactor: harden inference cache state and attention dispatch
- split KVCache into phase-specific PrefillKVCache/DecodeKVCache types selected by start_pos
- unify steady-state detection in TaskCacheManager
- guard decode steady-state reuse with the cached task signature so recycled req slots cannot replay a prior generation's tokens and positions
- collapse attention backend fwd_decode/fwd_prefill into a single subclass-owned forward with a shared _check_fwd guard
- fix thread-safety gap in weight update and validate prefill inputs before KV allocation
- centralize magic constants in InferenceConfig and align docs with behavior
2026-09-04 14:28:04 +08:00
ViperEkura 76f1c10feb refactor: generate train and serve CLIs from config-backed option specs
- Add astrai/config/cli.py: OptSpec tables plus apply_specs infer click types and defaults from config fields, covering Optional[X], Union[X, None], PEP 604 X | None, stringified PEP 563 annotations, bool flag pairs, and repeatable list options
- Move GroupedCommand/GroupedOption and the three-layer YAML merge (option defaults < YAML < explicit CLI) into the config package, adding unknown-key warning and mapping validation
- Replace ~420 lines of hand-written @opt decorators in scripts/tools/train.py with a 66-entry spec table; option names, defaults, flag styles, and YAML semantics verified unchanged
- Migrate scripts/tools/server.py to the same mechanism with its section binding, integer coercion, and dtype validation preserved locally
- Add tests/config/test_cli.py covering type inference across annotation styles, default overrides, flag pairs, merge precedence, scientific notation, and help ordering
2026-09-03 20:56:11 +08:00
ViperEkura 7e98a419a7 fix: resolve audited dispatch, kernel, and rollout bugs
- re-register the linear family with the operator dispatcher (ASTR_OPS / op_backend / resolve)
- fix bf16 gemv misaligned-address faults and element mispairing for offset weights
- reject misaligned bf16_swiglu inputs with a clear error and fall back in the backend gate
- make the rollout reuse decision, validation, and return atomic under one policy snapshot
- add the documented post-scoring rollout version check
- derive live+1 under the scheduler lock in optimizer_step via apply_weight_update(None, ...)
- reject rollout_max_policy_lag below rollout_interval - 1 at config time
- sync gemv stream-test inputs before switching streams; drop dead loader imports
2026-09-03 16:58:14 +08:00
0z5a 587b0ee046 fix: keep async rollouts version-consistent
- serialize shared-model optimizer updates with generation
- reject future or over-lagged rollout results after asynchronous scoring
- close cache publication races
- persist policy versions in online checkpoints
2026-09-03 12:01:24 +08:00
ViperEkura 88c06db096 fix: resolve audited training and inference bugs
- reject prompts that encode to zero tokens in add_task instead of admitting a task whose prefill can never run, and surface empty-id run_batch calls as prompt_empty errors
- deliver the STOP stream callback when cancelling a live task so clients observe termination instead of hanging until socket timeout
- strip the torch.compile _orig_mod. prefix at every unwrap_model site and when loading checkpoints so FSDP state dicts and saved weights no longer leak the wrapper name into downstream keys
- reject online_* train strategies with nprocs > 1 at config validation time, explaining the NCCL all-gather deadlock they would otherwise hit mid-run
- apply the frequency penalty before temperature scaling (OpenAI semantics) so the penalty survives temperature=0 instead of being annihilated by the 1e8 logit blowup, and exclude penalty pipelines from the greedy fast path
- return logprobs from the raw pre-strategy distribution so they match training-side policy logprobs for PPO/GRPO importance ratios
2026-09-02 21:25:01 +08:00
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 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 9b7e6c205f feat: add moe auxloss and metrics 2026-08-05 18:12:28 +08:00
ViperEkura 1c7369f293 feat: add MoE auxiliary loss metrics
- Propagates MoE load-balancing loss through model outputs
- Logs task, auxiliary, and weighted losses across strategies
- Computes only explicitly requested callback metrics
- Preserves tensor compute_loss API and adds regression tests
2026-08-02 06:30:43 +08:00
ViperEkura d7db37a70f fix: preserve MoE routing defaults 2026-08-02 05:30:26 +08:00
Gaolingx 6d98bb4f9f 20260801-moe model impl
need to add aux loss for load balancing
2026-08-01 22:48:58 +08:00
QueenAmish 04899a2b15 Make Nora+NAdamW the default optimizer 2026-07-31 23:16:39 +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 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 99b5d2b2da perf: batch tokenizer preprocessing 2026-07-23 18:42:19 +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 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 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 4e508afa2d fix : SFT pipeline position_ids default & doc boundary preservation
- change position_ids_mode default from "none" to "doc_reset" so SFT preprocessing always generates position_ids (was causing dataset load KeyError)
- generate per-doc position_ids before packing (doc_reset mode), preserving document boundaries for BFD packing (cross-doc attention leak fix)
- change _align_bucket padding from [1] to [0] to avoid accidentally training on loss_mask padding
2026-07-04 15:59:11 +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 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 b1adc40cfb refactor : 将 config 对象直接传给 DecoderBlock,替代 16 个独立参数
- DecoderBlock.__init__ 改为 (config, layer_id),内部用 asdict
  展开字段给 AttnFactory/FFNFactory,factory 按 __init__ 签名自动过滤
- EncoderConfig 补充 attn_type 和 ffn_type 字段
- 314 个测试全部通过
2026-06-19 14:15:33 +08:00
ViperEkura fec376b0dd fix : 修复策略相关文件的类型注解与抽象方法体
- 修复 strategy.py 单元素 Union 与缺失的参数/返回类型注解
- 修复 train_context.py 8 个 default=None 字段缺 Optional 标记
- 修复 sample.py/packing.py/position_id.py 方法缺参数及返回类型注解
- 修复 factory.py _resolve_type/list_registered 缺类型注解
- 修复 train_config.py 裸 dict/list 缺泛型参数
- abstractmethod body 从 ... 改为 raise NotImplementedError
- feat : checkpoint meta.json 保存 TrainConfig 超参供人工查阅
2026-06-14 16:20:10 +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 31bc7f5c2a refactor : pipeline 策略化拆分,消除 _flush if/else
- PackingStrategy / PositionIdStrategy / StoreWriter 独立文件 + Factory
- Pipeline._flush 零 if/else,纯编排
- SectionRenderer 从 SectionedMaskBuilder 分离
- OutputConfig.position_ids_mode 默认改为 ""none""
2026-06-06 00:45:33 +08:00
ViperEkura 3057741de9 refactor : 合并 data config docstring 并实现 BFD 打包策略
- 将 ProcessingConfig/OutputConfig 参数描述合并到类级 docstring

- Pipeline 支持 packing_strategy/truncation_mode,新增 bfd 打包
2026-06-05 17:41:51 +08:00
ViperEkura 985d940db6 feat : 数据流水拼接策略支持 position_ids 预计算
- OutputConfig.position_ids_mode 三种模式控制边界策略
- pipeline._flush() 按配置生成扁平 position_ids 数组
- SFTDataset 在 __getitem__ 中返回 position_ids
- SFTStrategy 将 position_ids 传入 model.forward()
2026-06-04 13:56:19 +08:00
ViperEkura 02a7cb9fa0 feat : preprocessing 支持 DPO/GRPO 多输出格式
- InputConfig 新增 sources 字段驱动多输出映射
- SectionedMaskBuilder 提取 _process_sections/_build_multi 模板方法
- Pipeline 泛化 accumulate 逻辑处理多 key 结果
- 测试拆分为 config/builder/pipeline 三文件,纯函数风格
2026-06-03 10:32:10 +08:00
ViperEkura 9fe2121743 feat : TrainConfig 支持 val_split 从训练集自动切分验证集
- val_split 比例从 dataset 中划出验证集,用 random_seed 固定随机切分
- 若 val_dataset 已显式设置则跳过自动切分
2026-06-02 20:33:40 +08:00
ViperEkura 01ce1fb9e3 refactor : Pipeline 去除去重,ids 重命名为 sequence,泛型透传
- 移除 Pipeline 内置去重逻辑及 dedup_signature 工具函数
- 删除 ProcessingConfig.deduplicate 字段
- builder 返回 'sequence' 替代 'ids',与 dataset 层统一
- pipeline 纯透传,泛型处理任意 key 补齐默认值
2026-05-31 15:14:27 +08:00
ViperEkura dbe5891201 refactor : 统一 SectionedMaskBuilder,支持可配置 dtype
- 三合一 MaskBuilder,移除 chat/instruction/text,统一为 sections 配置
- OutputConfig 增加 dtype 字段 (per-key,默认 int32)
- 移除 from __future__ import annotations
- 测试适配新配置格式
2026-05-31 14:24:10 +08:00
ViperEkura 31ae2deeba refactor : BaseConfig 提供 from_json/to_json,嵌套 config 自动反序列化
- from_json/to_json 上提至 BaseConfig,所有子类自动继承
- _coerce 新增 dict 到 BaseConfig 子类的递归反序列化,消除子类 from_dict 重载
- PipelineConfig 等子类仅声明字段,零样板代码
- 测试 tokenizer 改为自包含 BPE(含 chat template),不依赖 params/ 目录
- 特殊 token 改用 ASCII 字符,兼容所有平台
2026-05-30 21:04:19 +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 c424dfc293 feat : checkpoint 支持保存 config.json
- Checkpoint.save 写入独立的 config.json(模型架构参数)
- Checkpoint.load 读取 config.json,恢复时覆盖 context.model_config
- TrainContext 新增 model_config 字段,builder 从 resume_dir/config.json 加载
- BaseConfig.to_dict 支持 tuple 和嵌套 dataclass(如 LoRAConfig)
- 删除 _get_meta/_get_config wrapper,直接使用 load_json
2026-05-28 20:21:51 +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 737585a32a feat: 新增NTK-Aware RoPE缩放支持
- RotaryEmbedding接受rope_scaling配置,自动计算scaled base
- AutoRegressiveLMConfig和EncoderConfig新增rope_scaling字段
2026-05-25 21:22:07 +08:00
ViperEkura a4688021bf feat: 新增LoRA微调模块
- LoRALinear基于register_parameter托管base weight,state_dict路径不变
- inject_lora/merge_lora/save_lora/load_lora完备封装
- 24个单元测试覆盖注入、合并、存取、边界场景
2026-05-25 20:15:31 +08:00
ViperEkura 7df6eb9211 feat: 新增FSDP并行后端
- FSDPExecutor通过**fsdp_kwargs直传FSDP参数
- unwrap_model同时支持DDP和FSDP
- parallel_mode新增fsdp选项
2026-05-25 19:43:14 +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 785d65436c fix: 修复 to_dict list 类型丢失与 OpenAI stop 参数失效
- to_dict() 增加 list 类型序列化支持,metrics 等字段不再丢失
- OpenAIHandler 补充 get_stop_sequences/on_token,读取 request.stop 并检测停止序列
- 文档类图补充缺失字段、修正关系分类、ChatCompletionRequest 字段增加 Optional
2026-05-19 21:07:07 +08:00