ViperEkura
/
AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
503 Commits 1 Branch 9 Tags 3.9 MiB
Commit Graph

187 Commits

Author SHA1 Message Date
ViperEkura 2c3cef1c87 feat: wire up paged decode CUDA kernel to Python extension 
- Add attn_paged_decode wrapper in ops.py with gather fallback
- Register kernel in loader.py and export from __init__.py
- Extract test_utils.cuh shared by all attention unit tests
- Rename attn_paged_vs_contiguous.cu to attn_paged_decode_test.cu
- Refactor decode/prefill tests to use common bf16 helpers and cpu ref
- Fix k_cache dim check in attn_paged_decode.cu
 2026-07-11 18:40:49 +08:00
ViperEkura d923ebe38d refactor: rename gqa_* to attn_*, split-KV for all decode paths 
- Rename all csrc/kernels/gqa_*.cuh/cu to attn_*, with _split_q / _split_kv
  strategy suffix and optional _mma compute suffix
- Remove non-split MMA decode kernel, keep only split-KV path
- Convert scalar decode fallback to split-KV (o_part/ml_part + combine)
- Move combine kernel to attn_decode_split_kv.cuh (shared by both paths)
- Rename GQAParams to AttentionParams
- Update all C++ #include, PYBIND11, and Python extension references
 2026-07-10 23:35:14 +08:00
ViperEkura 841a582b28 refactor: split mask builder by single/multi output 
- Extract SingleOutputMaskBuilder for SFT and pretrain configs
- Extract MultiOutputMaskBuilder for DPO and GRPO configs
- Keep SectionedMaskBuilder as backward-compatible facade
- Register "single" and "multi" names in MaskBuilderFactory
- Add parity and rejection tests for concrete builders
 2026-07-08 21:18:34 +08:00
ViperEkura fd65b9bc23 feat: support HEAD_DIM=32 and split extension into loader/ops 
- add case 32 to decode/prefill dispatch switch
- fix swiz_col out-of-bounds for HEAD_DIM=32: XOR mask now limited to chunk count (3 for 32, 7 for >=64) instead of always 7, which produced column offsets >= LD=32 and corrupted shared memory
- restructure decode dispatch to #ifndef/#else/#endif matching prefill
- split astrai/extension/__init__.py into loader.py (kernel .so discovery) and ops.py (wrapper functions + torch SDPA fallback); __init__.py now re-exports the public API
 2026-07-08 14:14:11 +08:00
ViperEkura 536dbc0c9a fix: set tqdm postfix before update so first step shows metrics 2026-07-07 00:14:13 +08:00
ViperEkura 4af7acd449 fix: support single .h5 file loading in load_h5 2026-07-07 00:11:16 +08:00
ViperEkura 53ed52b4b8 refactor: extension dispatch layer with CUDA/torch fallback 
- Add gqa_decode_attn/gqa_prefill_attn dispatch functions
- Internal _available/__modules with underscore prefix
- CUDA kernel path with F.scaled_dot_product_attention fallback
- GQA head expansion in fallback path
 2026-07-06 21:07:16 +08:00
ViperEkura 579b8c3129 fix: correct gqa_decode_attn reduction + add gqa_prefill_attn 
- gqa_decode_attn: rewrite to per-KV-head, K in smem
- gqa_prefill_attn: new kernel for Q_len > 1 with GQA
 2026-07-06 13:45:18 +08:00
ViperEkura e8e228d035 feat: add optional CUDA kernel system (csrc/) + fused GQA decode attention 
Structure:
  csrc/               -- .cu sources + build.py registry
  astrai/extension/   -- compiled .so + __init__.py (import dispatcher)
  setup.py            -- CUDAExtension from csrc/build.py REGISTRY

Control: CSRC_KERNELS=true|false env var at install time.
Fallback: astrai.extension.available dict for runtime detection.
 2026-07-06 12:09:58 +08:00
ViperEkura bbe6ff2d8f release : v1.3.8 
- refactor: 重写 IFD 评估为三层架构，引入 BFD 装箱与自定义 attention mask 批处理打分
- refactor: 重写 HumanEval 评估为函数式流水线，修复测试超时与动态 pass@k
- perf: 替换 paged KV cache 为 ContiguousCache，解码所有 group
- feat: 新增 ROUGE 评估脚本、JSONL 数据集 store、stream_chat 参数
- fix: 修复 IFD token-set 不对称、SFT position_ids 默认值、文档边界保留
 2026-07-05 19:12:33 +08:00
ViperEkura 849e1e00a3 refactor: clean up inference design patterns 
1. KVCache base: add default task_cached/task_record_hashes, remove getattr from scheduler
2. Remove page_size param from scheduler constructor (ContiguousCache-only)
3. InferenceEngine expose cache param for KVCache injection
4. Rename page_cache -> kv_cache in Executor
5. Move stream_callback from Task to TaskManager._callbacks dict
6. TaskManager.clear_queues clears callbacks
 2026-07-05 11:41:54 +08:00
ViperEkura 5416c2e8fb perf: replace paged KV cache with contiguous ContiguousCache, decode all groups 
- Add KVCache/CacheView abstract base classes in cache.py
- Add ContiguousCache (contiguous per-slot buffer, default) alongside PageCache (paged, renamed from old KVCache)
- Merge make_table_tensor + bind into bind_tasks on KVCache interface
- Remove task_cached/task_record_hashes from base class (PageCache-only)
- Scheduler: decode all position groups instead of just the largest (eliminates 63% group skip rate)
- Scheduler: accept optional cache param for swapping implementations
- Model layer type hints use CacheView base class
- Batch 1-32: 1-7% speedup from eliminating Storage.gather overhead
- All 183 inference tests pass
 2026-07-05 11:34:36 +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 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 8999ca89b8 feat: add JSONL dataset store with on-the-fly tokenization 
- Add JsonlStore registered under "jsonl" in astrai/dataset/storage.py
- Reuse PipelineConfig schema for JSONL dataset configuration
- Update detect_format to recognize JSONL directories and files
- Move save_h5/load_h5/save_bin/load_bin to astrai/serialization
- Split astrai/serialization.py into checkpoint/dataset submodules
- Add tests for JSONL detection, seq/SFT stores, and config roundtrip
 2026-07-04 15:42:33 +08:00
ViperEkura 27524ad085 fix: reset sampler iter at epoch end so progress bar shows total after first epoch 2026-07-04 06:35:55 +08:00
ViperEkura 27d1921d9c fix: scheduler division-by-zero, loss_mask bool 
- schedule.py: guard warmup_steps/lr_decay_steps against zero
- strategy.py: use ~loss_mask instead of loss_mask==0 on bool tensor
 2026-07-03 22:04:55 +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 dfb151537b fix: ForwardRef._evaluate Python 3.12 compatibility 2026-07-03 18:41:19 +08:00
ViperEkura 500c605fad fix: unify scheduler min_rate default to 0.01, clamp WSD warmup 2026-07-03 17:52:23 +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 6715461a36 chore : 升级 torch 2.11.0+cu128，移除自定义 Muon，修复 gloo device_id 
- torch 2.7.1-cu126 升级至 2.11.0-cu128，numpy 2.3.2 升级至 2.4.4
- 移除 astrai/trainer/optim.py，改用 torch.optim.Muon
- parallel setup: gloo 后端不再传递 device_id，单卡多进程不再报错
 2026-06-27 16:10:37 +08:00
ViperEkura b4587c5d08 refactor : metric_logger 改用事件类型 (type=step/validation/epoch) 
- 每种事件独立 schema，不再混入 null 字段
- 回调顺序 validation 移到 metric_logger 之前，确保 on_optimizer_step 先跑
- 用内部 _last_val_loss 代替 TrainContext.last_val_iter 判断新验证
- 修复 factory.py 未使用导入、evaluate_ifeval.py 多余 f 前缀
 2026-06-25 17:18:20 +08:00
ViperEkura 88ec63121d feat : GPT-2 residual scaling weight init 
- Linear: normal(0, init_std) replaces kaiming_uniform_(a=sqrt(5))
- o_proj / mlp.down: init_std = 0.02 / sqrt(2 * n_layers)
- MoE: expert down scaled by 1/sqrt(1/n_shared + 1/K)
- Embedding: normal(0, 0.02), unchanged
 2026-06-25 15:08:31 +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 d88a41f8f1 fix: 修复预处理流水线 4 个致命问题 
- pipeline: 单条数据异常不再崩溃整条流水线, 改 log warning 后跳过
- pipeline: _align_bucket 统一用 len(ids) 填充, 修复多输出模式下长度错配
- writer: BinWriter/H5Writer 写入失败自动清理残留文件并记录详细错误
- packing: BFDPacking 真正将序列打包进 bin 而非仅重排, 减少碎片
 2026-06-18 17:38:01 +08:00
ViperEkura a4e5a8c81c feat: 新增 WSD 学习率调度器 
- 支持 Warmup-Stable-Decay 三段式调度
- stable 阶段保持最高 lr，decay 阶段 sqrt 衰减
- 适用于持续预训练、SFT、RLHF 场景
 2026-06-18 15:55:15 +08:00
ViperEkura 3e234c46f6 fix: 使用 threading.Event 替代裸 bool，补全公共 API 
- scheduler 停止信号改用 threading.Event，跨解释器安全
- 移除 _fatal_error 和 check_health，异常仅用 logger.error 记录
- 补全 astrai/__init__.py，暴露所有主要模块
 2026-06-18 15:38:35 +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 a2512f8a5a fix : resume_dir 无权重文件时不强制加载，支持仅配置训练 
- Checkpoint.load_any 统一处理 meta.json / model.safetensors / 无文件三种情况
- train_context.py 调用简化为单一路径，移除 load_model_weights 直接依赖
 2026-06-13 15:40:14 +08:00
ViperEkura 457e16ea3c fix : val_loss 默认改为 None，日志跳过空值；val_dataloader 补 Optional 注解 2026-06-13 14:24:13 +08:00
ViperEkura daf627a6de fix : _save_log 前确保日志目录存在，防止跨进程反序列化后目录丢失 2026-06-12 15:39:54 +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 e7b18b7c03 refactor : BaseFactory 基类类型自动推导 + 移除冗余代码 
- _validate_component 从 BaseFactory[T] 泛型参数自动解析基类类型，9 个子类覆写移除
- Registry 类内联到 BaseFactory._entries，移除未用的 list_by_category/list_by_priority
- _component_base 在 __init_subclass__ 时立即解析
- 数据集 4 个子类冗余 __init__ 移除
 2026-06-06 21:23:41 +08:00
ViperEkura 9e31d4ef2b feat : BaseToolParser.feed 增加可选 token_ids 参数 
- format_chunk ABC 改为 (token, **kwargs)，body/token_ids 通过 kw 传入
- ProtocolHandler._handle_stream 逐 token encode 并透传
- Anthropic builder 用 **kwargs 吸收不使用的参数，零变更
- 新增 3 个 token_ids 参数测试
 2026-06-06 11:19:30 +08:00
ViperEkura 52aa4d01d5 feat : 推理层增加 vLLM 风格工具调用解析 
- 新增 BaseToolParser 抽象基类，定义 feed/parse_complete 流式接口
- 新增 SimpleJsonToolParser，解析 {"name":"...","arguments":{...}} 格式
- 新增 ToolParserFactory，基于 BaseFactory 实现可插拔注册
- 集成 parser 到 OpenAIResponseBuilder，支持流式/非流式工具调用
- 扩展 ChatMessage 和 ChatCompletionRequest，增加 tools/tool_choice 字段
- 重构 format_chunk 接口，传入累积文本支持全量重新解析
- 新增 74 个单元测试，覆盖扫描/查找/流式解析/完整解析/工厂
 2026-06-06 08:54:10 +08:00
ViperEkura 986be957ec refactor : on_batch_begin 移入 accumulate 上下文 2026-06-06 01:19:21 +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 acd1103bd0 fix : 使用 bool 注意力掩码并支持打包 SFT 文档边界阻断 
- 简化 process_attention_mask，通过广播返回 bool 掩码
- 新增 make_doc_boundary_mask 生成块对角因果掩码
- SFT strategy 传入文档边界掩码
 2026-06-05 17:02:28 +08:00
ViperEkura dc7d2cfbca refactor : FastAPI 懒加载单例，消除模块级副作用 
- import astrai.inference 不再在模块加载时创建 FastAPI 实例
- 路由移至 APIRouter；get_app() 首次调用时懒构造单例
- _create_engine 和 run_server 的 param_path 改为必填
- 更新测试改用 get_app() 替代模块级 app
 2026-06-04 15:52:27 +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 0422d6d38e refactor : 移除 LocalStrategy._clear_env 冗余清理 
- setup_parallel 已覆盖所有环境变量写入，无需前置清空
 2026-06-02 11:40:45 +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