Commit Graph
235 Commits
Author SHA1 Message Date
ViperEkura ef8783b7e3 fix: separate attn_mask and loss_mask in get_logprobs, compose causal masking in strategy
- add loss_mask parameter to get_logprobs to decouple attention from loss masking
- DPO/GRPO strategies compose key-padding + causal mask before model forward
- prevents prompt tokens from being masked out of attention and missing causal masking
2026-07-21 23:47:28 +08:00
ViperEkura ccf728a1b7 perf: eliminate GPU syncs in contiguous cache write/gather hot paths
- Replace .tolist() calls with _total_len in gather(); move _slot_len updates from per-layer write to once-per-step bind_tasks

- Use torch.as_tensor instead of torch.tensor in decode penalty history construction
2026-07-21 16:40:56 +08:00
ViperEkura f1b4b05d08 feat: add attention dimension dispatch 2026-07-21 12:52:45 +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 d7ac66fb73 refactor: simplify attention mask handling 2026-07-20 20:36:16 +08:00
ViperEkura a6e920fdb0 Merge pull request #20 from ccx1324/lora-device-fix
fix: LoRA device mismatch and checkpoint resume
2026-07-20 19:38:10 +08:00
ViperEkura 958df58f9d refactor: unify tokenizer encode and apply_chat_template for batch support
- encode(str) single-thread, encode(List[str]) Rust parallel encode_batch
- apply_chat_template accepts single Messages or List[Messages] for batch
- add Message/Messages type aliases at module level
2026-07-20 17:49:26 +08:00
ViperEkura e0f102c4d9 feat: support SFT directly from JSONL without dataset_config.json
- JsonlStore falls back to built-in messages config when no config file found and tokenizer_path is provided
- DatasetFactory.load forwards tokenizer_path to store for SFT/SEQ+jsonl
- assistant turns train, other roles masked, position_ids doc_reset
2026-07-20 17:25:09 +08:00
ccx1324andccx 5a942527b2 fix: inject LoRA before loading checkpoint state_dict
move inject_lora() before load_state_dict in _before_wrap so that
  LoRA adapter weights from a checkpoint are properly restored on
  training resume. Previously, inject happened after load, causing
  lora_A/lora_B keys to be silently ignored (strict=False).

  Co-Authored-By: ccx1324 <2424441089@qq.com>
2026-07-20 17:00:24 +08:00
ViperEkura 37a3036934 refactor: split LoRA param init into local vars 2026-07-20 16:23:47 +08:00
ccx1324andClaude Opus 4.7 a5678c9185 fix: create LoRA parameters on base weight device instead of CPU
When `inject_lora()` replaces Linear layers with LoRALinear after the model
has been moved to CUDA, the new lora_A and lora_B parameters were always
created on CPU, causing a device mismatch error during the forward pass.

Now lora_A and lora_B are created on the same device and dtype as the
parent weight, matching the model's current device.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-07-20 16:11:48 +08:00
ViperEkura 06eeeead79 refactor: map instruction/input/output to chat roles
- RolloutGenerator._instruction_to_messages builds system/user/assistant list (instruction->system, input->user, output->assistant), replacing single-user-turn concatenation
- Remove _iter_samples helper; _prepare_prompts zips parallel list-of-strings fields directly per the collate_fn contract
- Tests adopt a system-aware chat template and pin the three-field role mapping
- Drop unused imports caught by ruff F401 (torch.Tensor in scheduler.py, iter_raw_records in pipeline.py, Tuple in evaluate_rouge.py)
2026-07-20 13:55:25 +08:00
ViperEkura e8ff7f5321 fix: use batch_per_device for rollout scheduler batch sizing
- train_context.py referenced non-existent cfg.batch_size, replaced with cfg.batch_per_device
- default group_size lowered from 8 to 1: without a group concept (DPO), scheduler batch equals batch_per_device; rollout-based DPO can opt in via extra_kwargs['group_size']>=2
- inline expressions (rollout_batch_size, max_seq_len) extracted for readability
- add tests/trainer/test_online_e2e.py: end-to-end online_dpo via Trainer.train, exercising KV-cache-backed rollout path
2026-07-20 13:32:04 +08:00
ViperEkura a6e1f26cd4 refactor: simplify sample return_logprobs path
- SamplingPipeline.sample gains return_logprobs; both greedy and multinomial paths now share a single log_softmax+gather instead of duplicating the sampling logic
- module-level sample() becomes a thin forwarder instead of re-implementing the three-branch logic
- eliminates ~10 lines of duplicated softmax/gather code; no caller-facing API change
2026-07-20 13:16:18 +08:00
ViperEkura 95c43368ae refactor: unify rollout onto inference engine KV-cache path
- RolloutGenerator now delegates prefill/decode to InferenceScheduler.run_batch (sync API, no background thread), sharing one KV-cache code path with the inference server and eliminating O(n^2) recompute in rollout
- Add sample(return_logprobs=) and Executor.execute_decode(return_logprobs=) to expose behaviour-policy log-probs through the engine; Task gains output_logprobs
- RolloutResult now subclasses RawRollout (adds rewards only), removing duplicated fields
- RolloutRunner.__call__ returns (result, is_fresh) instead of relying on object identity, removing the fragile refresh-detection contract
- Remove O(n^2) generate_responses helper and dead code (_tokenize_prompts, unused old_model arg)
- train_context.py wires InferenceScheduler directly instead of hand-rolling SamplingPipeline
- Tests: +11 covering return_logprobs, run_batch, and KV-cache-backed rollout semantics; 404 pass
2026-07-20 12:52:20 +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 74b9308883 refactor: pass model_fn/optimizer_fn to executor.prepare
- BaseExecutor.prepare now takes factories and instantiates model via model_fn(), runs before_wrap hook, wraps DDP/FSDP, then builds optimizer/scheduler on the wrapped model
- optimizer/scheduler creation moved into executor.prepare, eliminating the old 'create-then-wrap' hack reliance on use_orig_params=True
- FSDPExecutor/BaseExecutor accept **_extra kwargs to tolerate DDP-only keys (broadcast_buffers, gradient_as_bucket_view) being forwarded via executor_kwargs
- dataloader builds stay external; executor only handles model/optimizer/scheduler
- train_context.py rewritten to load checkpoint state_dict before prepare via a before_wrap closure
2026-07-20 01:32:05 +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 31d33ccdf0 chore: bump to 1.3.10 2026-07-19 16:40:27 +08:00
ViperEkura 88ec786e39 fix: memmap mode=r, tool parser json.loads, greedy decode 2026-07-19 16:38:28 +08:00
ViperEkura 663ef900fc refactor: move sample-id indexing from dataset to store
- Store owns window_size/stride and __getitem__/__len__/sample_window
- Dataset classes become thin delegators binding a Store to a train-type key mapping
- Drop BaseDataset.get_index and the RecordDataset中间类 (window死代码)
- DatasetFactory forces window_size=0 for record datasets so record semantics never get window-tainted
- token_count/num_records split the legacy len() semantics (raw stream length vs record count)
- Update tests to the new .store/.token_count API and window/record mode switching
2026-07-19 16:02:50 +08:00
ViperEkura f3eaaef842 refactor: remove redundant strategy/executor code
- Drop BaseStrategy.model_fn (stored but never read)
- Drop model_fn= passed to StrategyFactory.create in train_context
- Simplify FSDPExecutor.clip_grad_norm None branch to delegate to super()
- Remove DDPExecutor._gather_state_dict override (identical to base)
2026-07-19 12:45:58 +08:00
ViperEkura 31c22dc043 refactor: deduplicate preprocessing kernel and BFD packing
- Extract shared core (mask building, primary-id extraction, tensorisation, position-id generation) to astrai/preprocessing/core.py; Pipeline and TokenizeTransform both consume it, eliminating ~60% duplicated logic
- Promote BFD _plan to module-level plan_bfd(lengths, max_len) returning pure index bins; BFDPacking.apply and evaluate_ifd._pack_bins both call it, removing the second BFD implementation
- Split Pipeline._flush (49 lines) into _inject_doc_reset_position_ids + _inject_continuous_position_ids + _to_tensors; split Pipeline.run by delegating record iteration to core.iter_raw_records
- Remove dead no-op pop/塞回 in Pipeline.run (L110-111)
2026-07-19 12:27:56 +08:00
ViperEkura 17127f8b3c fix: make tokenizer picklable for spawn multiprocessing
- ChatTemplate: defer Jinja2 compilation to cached_property, exclude compiled template from __getstate__ (its dynamic root function has __module__=None and falls back to __main__, breaking pickle)
- AutoTokenizer: bypass __getattr__ for underscore-prefixed attrs to prevent infinite recursion during unpickle when __dict__ is empty
2026-07-19 11:59:55 +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 fc62890e70 fix: apply chat template in DPO tokenization
- dpo_tokenize now uses tokenizer.apply_chat_template to match SFT format
- Prompt rendered with add_generation_prompt=True
- Chosen/rejected appended as assistant turn
- Remove leftover dead code from _extract_text
- Update tests to mock apply_chat_template
2026-07-19 00:00:51 +08:00
ViperEkura f433672140 fix: use sum reduction for DPO sequence logprob
- DPO requires sequence-level sum of token logprobs, not per-token mean
- mean reduction made beta*ratio_diff ~0.03 (near-zero gradient)
- loss stalled at 0.6931 because logsigmoid(0.03) has vanishing grad
- sum gives beta*ratio_diff ~10 with meaningful gradients
2026-07-18 23:48:35 +08:00
ViperEkura 7e1e5b6e6a refactor: DatasetFactory.load accepts pre-built store instance
- load(store=...) binds directly, skipping format detection/processor
- load_path now optional when store is given
- Remove redundant from_store (merged into load)
- Caller can fully control Store construction + processor setup
2026-07-18 23:23:51 +08:00
ViperEkura 553a42702d refactor: replace diamond inheritance with mixin composition
- StreamStore/RecordStore → Streamable/Recordable (stateless mixins)
- Store is sole base class, no MRO ambiguity
- H5Store/MmapStore/JsonlStore mix in both traits explicitly
- segments_are_records declared per-subclass (H5/Jsonl=True, bin=False)
- Add tests for dpo_tokenize, lazy jsonl, dual-mode H5, stream-only bin
- Remove unused _to_tensor helper
2026-07-18 23:20:41 +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 b33250dc28 refactor: decouple tokenizer from Store into Transform layer
- Extract tokenization/mask/position logic from JsonlStore into TokenizeTransform
- JsonlStore now pure reader: reads JSON records, delegates to transform
- Store no longer imports tokenizer or preprocessing components
- Replace per_record param with segments_are_records class attribute
- Store subclasses declare segment semantics as format-level property
2026-07-18 21:37: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 9d3ccfdffc fix: incremental decode to avoid U+FFFD in streaming
- StreamDecoder buffers incomplete multi-byte sequences
- Task.decode_new_token replaces per-token decode in scheduler
- flush_remaining emits final buffered text on task finish
2026-07-18 13:05:36 +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 d08a92c7bd feat: add frequency penalty to inference sampling pipeline
- Add FrequencyPenaltyStrategy (logit -= penalty * count)
- Per-task rep_window for penalty history lookup
- Wire through engine, task, executor, API layer
- Add --frequency_penalty and --rep_window to stream_chat.py
- 9 unit tests for frequency penalty strategy
2026-07-17 21:28:31 +08:00
ViperEkura a1ea26d367 fix: rewrite GRPO data pipeline for offline record-level access
- process_list_field returns List[List[int]] preserving per-response boundaries
- GRPODataset rewritten to record-level __getitem__ (no windowing/stride)
- grpo_collate_fn pads variable-length responses into [B, G, R] tensors
- JsonlStore detects nested List[List[int]] and stores List[Tensor] per record
- Store._normalize skips nested-list keys from cumsum bookkeeping
- Pipeline._flush handles nested lists without cross-record flattening
- Export grpo_collate_fn from astrai.dataset
- 6 new GRPO tests + 2 updated builder tests, 114 total pass
2026-07-17 14:34:41 +08:00
ViperEkura cd14d53707 feat: implement bfd_split packing strategy
- BFDSplitPacking splits over-length sequences into chunks before BFD
- All keys (loss_mask, position_ids, ...) split in lockstep for alignment
- No tokens lost vs bfd which truncates over-length sequences
- Tests: token preservation, chunk alignment, short unchanged, vs bfd
2026-07-17 12:38:56 +08:00
ViperEkura e220413035 feat: support raw JSON files in dataset pipeline and JsonlStore
- detect_format now recognizes .json directories as jsonl store
- JsonlStore loads .json arrays and dicts alongside .jsonl
- tokenizer_path defaults to dataset dir when omitted
- Pipeline._iter_items handles .json files (arrays/single dict)
- Tests: detect_format, seq load, self-contained dataset dir
2026-07-17 12:20:03 +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 b14f301730 fix: init last_ckpt_step and last_log_flush_step from context.optimizer_step 2026-07-15 22:15:55 +08:00
ViperEkura bb175fda91 fix: resume optimizer LR, step display, and consumed_samples alignment 2026-07-15 08:59:52 +08:00
ViperEkura 57729fd92d refactor: stride-based attn interface with layout and causal mask
- Replace is_causal + causal_offset with unified causal_offset (-1 = off, >=0 = first Q pos)
- Causal and mask can now coexist (was mutually exclusive)
- Add stride-based addressing for Q/KV/O (layout-agnostic, zero-copy)
- Add layout param ("bhld"/"blhd") parsed in Python, passed as int to C++
- Support 2D [batch, kv_len] and 3D [batch, q_len, kv_len] mask
- Vectorize paged KV gather in Python fallback (was per-token Python loop)
- Extract shared helpers: compute_num_splits, alloc_split_partials, dispatch_head_dim
- Unify paged_decode entry via attn_pack_paged_params
- Update mma_softmax_tile for 3D mask with per-row qrow indexing
2026-07-14 21:34:42 +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 b092316385 feat : add distributed checkpoint via executor checkpoint_context
- Add checkpoint_context context manager to BaseExecutor with entry/exit barrier
- Add _gather_state_dict hook overridden per executor (template method)
- DDPExecutor skips unwrap on non-rank-0 to avoid redundant state_dict gather
- FSDPExecutor uses rank0_only=True to reduce memory on non-writers
- Remove redundant rank-0 guard from Checkpoint.save and manual barrier from Callback
2026-07-13 12:27:09 +08:00
ViperEkura 9bcd696580 fix: token-level ratio and prompt masking in GRPO strategy
- Mask prompt tokens to 0 so their logprobs excluded from ratio/KL
- Switch to token-level ratio + PPO clipping via reduction='none'
- Slice response token logprobs from full sequence output
- Replace k3 KL estimator with non-negative k1 estimator
- Fix epsilon from finifo.eps (~1e-38) to 1e-8
- Remove unused 'reduction' param from GRPOStrategy.__init__
- Clarify offline batch semantics in docstring
- Add 11 unit tests for masking, advantage, KL, sync, clipping
- Sync training.md and architecture.md docs
2026-07-12 21:24:09 +08:00
ViperEkura 8f89c82d55 chore: bump version to 1.3.9 2026-07-12 21:04:20 +08:00
ViperEkura 4c35d36146 fix: auto-assign free port in spawn_parallel_fn to avoid EADDRINUSE 2026-07-12 19:21:56 +08:00
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