ViperEkura
45cc048fe9
- shard the Muon Newton-Schulz orthogonalization over the FSDP mesh instead of partial local slices - import HF checkpoints faithfully: per-head RoPE permutation for q/k projections and qk-norm, qwen3, shared experts, and qk-norm before RoPE (changes numerics for existing use_qk_norm checkpoints) - make preprocessing and resume self-contained: backfill realigned bucket keys by semantics (masks ones, rest zeros) and snapshot tokenizer files into every checkpoint - keep RL consistent: sync the offline GRPO old_model each optimizer step and validate online strategies through a public one-off-rollout hook that leaves the replay cache untouched - fix streaming serving: withhold partial tool-call prefixes with a stream-end flush, stream tool-call arguments from the raw source span, and terminate SSE frames with a blank line - fix sampling semantics: capture logprobs before top-k/top-p mutate logits in place and detect greedy pipelines polymorphically instead of isinstance bookkeeping
27 lines
900 B
Python
27 lines
900 B
Python
"""Tests for preprocessing pipeline bucket alignment."""
|
|
|
|
from astrai.preprocessing.pipeline import Pipeline
|
|
|
|
|
|
def test_align_bucket_backfills_missing_mask_with_ones():
|
|
bucket = {
|
|
"sequence": [[1, 2], [3, 4]],
|
|
"loss_mask": [[0, 1]],
|
|
"chosen_mask": [[1]],
|
|
"position_ids": [[0, 1]],
|
|
}
|
|
result = {"sequence": [5, 6, 7]}
|
|
Pipeline._align_bucket(bucket, result, [5, 6, 7])
|
|
assert bucket["loss_mask"][-1] == [1, 1, 1]
|
|
assert bucket["chosen_mask"][-1] == [1, 1, 1]
|
|
assert bucket["position_ids"][-1] == [0, 0, 0]
|
|
assert bucket["sequence"] == [[1, 2], [3, 4]]
|
|
|
|
|
|
def test_align_bucket_keeps_present_keys():
|
|
bucket = {"sequence": [[1, 2]], "loss_mask": [[0, 1]]}
|
|
result = {"sequence": [9], "loss_mask": [1]}
|
|
Pipeline._align_bucket(bucket, result, [9])
|
|
assert bucket["loss_mask"] == [[0, 1]]
|
|
assert bucket["sequence"] == [[1, 2]]
|