mirror of
https://github.com/hanzoai/gym.git
synced 2026-09-30 22:37:59 +00:00
gym: LoRA SFT and GRPO on hanzo-ml, end to end
model: one trainable Llama-family decoder (Qwen3 with q/k norm, Llama) built from differentiable hanzo-nn primitives, LoRA on the projections the config names, PEFT adapter save/load/merge; logits match hanzo-transformers to 1e-4. data: chat_template / alpaca / completion with Axolotl's label semantics (dummy-render turn spans, train_on_eos), jsonl/json/parquet, HF datasets and split grammar, seeded shuffle, validation tail, GRPO prompts, collate. train: chunked masked cross-entropy, warmup + cosine/linear/constant, accumulation, global-norm clipping, AdamW, checkpoints, eval. grpo: the hanzo-training core ported unchanged in its math, with LmPolicy over CausalLm (batched masked scoring), exact/numeric/format/length/http verifiers. 61 tests. gradient_checkpointing, QLoRA, full fine-tune, multi-GPU and the DPO family are rejected at config load rather than approximated. Co-authored-by: Hanzo Dev <[email protected]>
This commit is contained in:
Hanzo Dev
Generated
+247
@@ -21,6 +21,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"const-random",
|
||||
"getrandom 0.3.4",
|
||||
"once_cell",
|
||||
"serde",
|
||||
@@ -37,6 +38,21 @@ dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "alloc-no-stdlib"
|
||||
version = "2.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3"
|
||||
|
||||
[[package]]
|
||||
name = "alloc-stdlib"
|
||||
version = "0.2.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0e76a019e91224d279006ff972f1e984179a6e9feb050adba6ce8274aef23195"
|
||||
dependencies = [
|
||||
"alloc-no-stdlib",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "allocator-api2"
|
||||
version = "0.2.21"
|
||||
@@ -144,6 +160,21 @@ version = "0.23.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5"
|
||||
|
||||
[[package]]
|
||||
name = "bit-set"
|
||||
version = "0.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3"
|
||||
dependencies = [
|
||||
"bit-vec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bit-vec"
|
||||
version = "0.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7"
|
||||
|
||||
[[package]]
|
||||
name = "bitflags"
|
||||
version = "1.3.2"
|
||||
@@ -180,6 +211,27 @@ dependencies = [
|
||||
"objc2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "brotli"
|
||||
version = "8.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5cc91aac060a7a1e25823bdccbfb6af1875b88f17c6daac97894eed8207166b3"
|
||||
dependencies = [
|
||||
"alloc-no-stdlib",
|
||||
"alloc-stdlib",
|
||||
"brotli-decompressor",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "brotli-decompressor"
|
||||
version = "5.0.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3a32acac15fe1967bc3986b2a6347dffc965602354ea6f450ad07e8bfd253583"
|
||||
dependencies = [
|
||||
"alloc-no-stdlib",
|
||||
"alloc-stdlib",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bumpalo"
|
||||
version = "3.20.3"
|
||||
@@ -234,6 +286,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0ad534f4357a5264cce5019c989cf66a4f0dc4e0d1b1d15f8aacec0ff7360273"
|
||||
dependencies = [
|
||||
"find-msvc-tools",
|
||||
"jobserver",
|
||||
"libc",
|
||||
"shlex",
|
||||
]
|
||||
|
||||
@@ -359,6 +413,26 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "const-random"
|
||||
version = "0.1.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359"
|
||||
dependencies = [
|
||||
"const-random-macro",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "const-random-macro"
|
||||
version = "0.1.16"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e"
|
||||
dependencies = [
|
||||
"getrandom 0.2.17",
|
||||
"once_cell",
|
||||
"tiny-keccak",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cookie"
|
||||
version = "0.18.2"
|
||||
@@ -729,6 +803,17 @@ version = "0.1.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6"
|
||||
|
||||
[[package]]
|
||||
name = "fancy-regex"
|
||||
version = "0.18.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e1e1dacd0d2082dfcf1351c4bdd566bbe89a2b263235a2b50058f1e130a47277"
|
||||
dependencies = [
|
||||
"bit-set",
|
||||
"regex-automata",
|
||||
"regex-syntax",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "filetime"
|
||||
version = "0.2.29"
|
||||
@@ -1215,10 +1300,13 @@ dependencies = [
|
||||
"half",
|
||||
"hanzo-ml",
|
||||
"hanzo-nn",
|
||||
"hanzo-transformers",
|
||||
"hf-hub",
|
||||
"indicatif 0.18.6",
|
||||
"memmap2",
|
||||
"minijinja",
|
||||
"minijinja-contrib",
|
||||
"parquet",
|
||||
"rand 0.9.5",
|
||||
"rand_distr",
|
||||
"regex",
|
||||
@@ -1312,6 +1400,25 @@ dependencies = [
|
||||
"thiserror 2.0.20",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hanzo-transformers"
|
||||
version = "0.11.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b561ebf35167c463b2ce29ce318d44840f00d91a17dd9be021b0b2a1a9cca6a9"
|
||||
dependencies = [
|
||||
"byteorder",
|
||||
"fancy-regex",
|
||||
"hanzo-ml",
|
||||
"hanzo-nn",
|
||||
"num-traits",
|
||||
"rand 0.9.5",
|
||||
"rayon",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_plain",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hanzo-ug"
|
||||
version = "0.11.33"
|
||||
@@ -1642,6 +1749,12 @@ dependencies = [
|
||||
"web-time",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "integer-encoding"
|
||||
version = "3.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02"
|
||||
|
||||
[[package]]
|
||||
name = "intel-mkl-src"
|
||||
version = "0.8.1"
|
||||
@@ -1691,6 +1804,16 @@ version = "1.0.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682"
|
||||
|
||||
[[package]]
|
||||
name = "jobserver"
|
||||
version = "0.1.35"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3"
|
||||
dependencies = [
|
||||
"getrandom 0.4.3",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "js-sys"
|
||||
version = "0.3.104"
|
||||
@@ -1779,6 +1902,15 @@ version = "0.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154"
|
||||
|
||||
[[package]]
|
||||
name = "lz4_flex"
|
||||
version = "0.13.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7ef0d4ed8669f8f8826eb00dc878084aa8f253506c4fd5e8f58f5bce72ddb97e"
|
||||
dependencies = [
|
||||
"twox-hash",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "macro_rules_attribute"
|
||||
version = "0.2.3"
|
||||
@@ -1861,6 +1993,16 @@ dependencies = [
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "minijinja-contrib"
|
||||
version = "2.24.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bd3e5f077bc2379f0f7d911e7cfdd921114ed99fc884533dca502944cb355b11"
|
||||
dependencies = [
|
||||
"minijinja",
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "minimal-lexical"
|
||||
version = "0.2.1"
|
||||
@@ -2169,6 +2311,42 @@ version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d"
|
||||
|
||||
[[package]]
|
||||
name = "ordered-float"
|
||||
version = "2.10.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "68f19d67e5a2795c94e73e0bb1cc1a7edeb2e28efd39e2e1c9b7a40c1108b11c"
|
||||
dependencies = [
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "parquet"
|
||||
version = "58.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d298093b2dec60289dce0684c986d0f7679e9dd15771c2c65406e1aaf604a704"
|
||||
dependencies = [
|
||||
"ahash",
|
||||
"base64 0.22.1",
|
||||
"brotli",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"flate2",
|
||||
"half",
|
||||
"hashbrown 0.17.1",
|
||||
"lz4_flex",
|
||||
"num-bigint",
|
||||
"num-integer",
|
||||
"num-traits",
|
||||
"paste",
|
||||
"seq-macro",
|
||||
"serde_json",
|
||||
"snap",
|
||||
"thrift",
|
||||
"twox-hash",
|
||||
"zstd",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "paste"
|
||||
version = "1.0.15"
|
||||
@@ -2715,6 +2893,15 @@ dependencies = [
|
||||
"zmij",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_plain"
|
||||
version = "1.0.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9ce1fc6db65a611022b23a0dec6975d63fb80a302cb3388835ff02c097258d50"
|
||||
dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_spanned"
|
||||
version = "0.6.9"
|
||||
@@ -2793,6 +2980,12 @@ version = "1.15.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90"
|
||||
|
||||
[[package]]
|
||||
name = "snap"
|
||||
version = "1.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "199905e6153d6405f9728fe44daace35f8f837bbf830bb6e85fbd5828709a886"
|
||||
|
||||
[[package]]
|
||||
name = "socket2"
|
||||
version = "0.6.5"
|
||||
@@ -2985,6 +3178,17 @@ dependencies = [
|
||||
"cfg-if",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thrift"
|
||||
version = "0.17.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7e54bc85fc7faa8bc175c4bab5b92ba8d9a3ce893d0e9f42cc455c8ab16a9e09"
|
||||
dependencies = [
|
||||
"byteorder",
|
||||
"integer-encoding",
|
||||
"ordered-float",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "time"
|
||||
version = "0.3.55"
|
||||
@@ -3015,6 +3219,15 @@ dependencies = [
|
||||
"time-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tiny-keccak"
|
||||
version = "2.0.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237"
|
||||
dependencies = [
|
||||
"crunchy",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tinystr"
|
||||
version = "0.8.4"
|
||||
@@ -3296,6 +3509,12 @@ version = "0.2.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b"
|
||||
|
||||
[[package]]
|
||||
name = "twox-hash"
|
||||
version = "2.1.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5283634e518fe9e82c7b20520bb4bc209009fd16c82077c802f8111ecbb0117a"
|
||||
|
||||
[[package]]
|
||||
name = "typed-path"
|
||||
version = "0.12.3"
|
||||
@@ -4153,3 +4372,31 @@ name = "zmij"
|
||||
version = "1.0.23"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b"
|
||||
|
||||
[[package]]
|
||||
name = "zstd"
|
||||
version = "0.13.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a"
|
||||
dependencies = [
|
||||
"zstd-safe",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zstd-safe"
|
||||
version = "7.2.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d"
|
||||
dependencies = [
|
||||
"zstd-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zstd-sys"
|
||||
version = "2.0.16+zstd.1.5.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"pkg-config",
|
||||
]
|
||||
|
||||
@@ -39,6 +39,8 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||
ureq = { version = "3", features = ["json"] }
|
||||
memmap2 = "0.9"
|
||||
half = "2"
|
||||
parquet = { version = "58", default-features = false, features = ["json", "snap", "zstd", "lz4", "flate2-zlib-rs", "brotli"] }
|
||||
minijinja-contrib = { version = "2", features = ["pycompat"] }
|
||||
|
||||
[features]
|
||||
default = []
|
||||
@@ -53,3 +55,6 @@ lto = "thin"
|
||||
codegen-units = 1
|
||||
|
||||
[workspace]
|
||||
|
||||
[dev-dependencies]
|
||||
hanzo-transformers = "0.11.9"
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
# hanzo-gym
|
||||
|
||||
Rust trainer on hanzo-ml (candle fork, crates.io `hanzo-ml` / `hanzo-nn`).
|
||||
Binary `gym`, lib `gym`. Standalone Cargo workspace (`[workspace]` in Cargo.toml
|
||||
keeps it out of ~/work/hanzo's root workspace).
|
||||
|
||||
## Shape
|
||||
|
||||
- `src/lib.rs` — the contracts every module is written against: `CausalLm`
|
||||
(forward → logits, `trainable_vars`, `save_adapter`), `Example`, `Batch`
|
||||
(labels unshifted, `IGNORE_INDEX` = -100 at padding), `Prompt` (GRPO).
|
||||
- `src/config.rs` — Axolotl-shaped YAML, `deny_unknown_fields`: a key this
|
||||
crate does not implement fails at load rather than being ignored.
|
||||
- `src/hub.rs` — model files: local dir or Hugging Face repo via hf-hub.
|
||||
- `src/model/` — LoRA (`lora.rs`), one trainable transformer for Qwen3 and
|
||||
Llama built from differentiable primitives only, PEFT adapter save/load/merge.
|
||||
- `src/data/` — dataset sources (jsonl, json, parquet, HF datasets; HF split
|
||||
grammar), strategies `chat_template` / `alpaca` / `completion` with Axolotl's
|
||||
label semantics, `collate`, `load_prompts` for GRPO.
|
||||
- `src/train/` — SFT loop: chunked masked CE, schedules, accumulation,
|
||||
clipping, AdamW, checkpoints; `run` picks device/dtype and dispatches to GRPO.
|
||||
- `src/grpo/` — port of hanzo-ml's `hanzo-training/src/grpo` math with a real
|
||||
`LmPolicy` over `CausalLm`, verifiers (`exact_match`, `numeric_match`,
|
||||
`format:<re>`, `length:<n>`, `http:<url>`).
|
||||
|
||||
## Facts that shaped it
|
||||
|
||||
- hanzo-nn's fused `rms_norm`, `rope` and attention kernels have no backward;
|
||||
the model uses `rms_norm_slow`, `rope_slow`, `softmax_last_dim` and matmul.
|
||||
- No KV cache in GRPO sampling yet (forward re-run per token). No QLoRA, full
|
||||
fine-tune, multi-GPU, DPO family, sample packing — rejected at config load.
|
||||
- KL in GRPO is against the pre-update policy of the step (adapter cannot be
|
||||
disabled through `CausalLm`), not the frozen base.
|
||||
- Hosted `/v1/ai/finetune/jobs` (hanzoai/ai) runs a Python transformers+PEFT
|
||||
image today; this binary is meant to take over `hanzo-ft-lora` once it
|
||||
covers the presets.
|
||||
|
||||
## Build / test
|
||||
|
||||
`cargo test` (CPU), `--features metal` or `--features cuda` for GPU. CI in
|
||||
`.hanzo/workflows/ci.yml` (fmt, clippy -D warnings, tests; metal leg on
|
||||
`hanzo-build-macos-arm64`). Publish: `cargo publish` as `hanzo-gym`.
|
||||
+92
-28
@@ -198,7 +198,10 @@ impl Config {
|
||||
}
|
||||
|
||||
pub fn validate(&self) -> anyhow::Result<()> {
|
||||
anyhow::ensure!(!self.datasets.is_empty(), "at least one dataset is required");
|
||||
anyhow::ensure!(
|
||||
!self.datasets.is_empty(),
|
||||
"at least one dataset is required"
|
||||
);
|
||||
anyhow::ensure!(self.sequence_len > 1, "sequence_len must be > 1");
|
||||
anyhow::ensure!(self.micro_batch_size >= 1, "micro_batch_size must be >= 1");
|
||||
if self.adapter == Adapter::Lora {
|
||||
@@ -209,8 +212,14 @@ impl Config {
|
||||
);
|
||||
}
|
||||
if self.rl == Some(Rl::Grpo) {
|
||||
anyhow::ensure!(self.trl.num_generations >= 2, "trl.num_generations must be >= 2");
|
||||
anyhow::ensure!(!self.trl.reward_funcs.is_empty(), "grpo needs trl.reward_funcs");
|
||||
anyhow::ensure!(
|
||||
self.trl.num_generations >= 2,
|
||||
"trl.num_generations must be >= 2"
|
||||
);
|
||||
anyhow::ensure!(
|
||||
!self.trl.reward_funcs.is_empty(),
|
||||
"grpo needs trl.reward_funcs"
|
||||
);
|
||||
if !self.trl.reward_weights.is_empty() {
|
||||
anyhow::ensure!(
|
||||
self.trl.reward_weights.len() == self.trl.reward_funcs.len(),
|
||||
@@ -230,10 +239,18 @@ impl Config {
|
||||
/// projection in attention and MLP.
|
||||
pub fn lora_targets(&self) -> Vec<String> {
|
||||
if self.lora_target_linear {
|
||||
["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
|
||||
.iter()
|
||||
.map(|s| s.to_string())
|
||||
.collect()
|
||||
[
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
"o_proj",
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
"down_proj",
|
||||
]
|
||||
.iter()
|
||||
.map(|s| s.to_string())
|
||||
.collect()
|
||||
} else {
|
||||
self.lora_target_modules.clone()
|
||||
}
|
||||
@@ -249,26 +266,66 @@ impl Config {
|
||||
}
|
||||
}
|
||||
|
||||
fn d_seq_len() -> usize { 2048 }
|
||||
fn d_lora_r() -> usize { 8 }
|
||||
fn d_lora_alpha() -> f64 { 16.0 }
|
||||
fn d_one() -> usize { 1 }
|
||||
fn d_one_f() -> f64 { 1.0 }
|
||||
fn d_ten() -> usize { 10 }
|
||||
fn d_lr() -> f64 { 2e-4 }
|
||||
fn d_grad_norm() -> f64 { 1.0 }
|
||||
fn d_seed() -> u64 { 42 }
|
||||
fn d_true() -> bool { true }
|
||||
fn d_gens() -> usize { 4 }
|
||||
fn d_completion() -> usize { 256 }
|
||||
fn d_messages() -> String { "messages".into() }
|
||||
fn d_roles() -> Vec<String> { vec!["assistant".into()] }
|
||||
fn d_instruction() -> String { "instruction".into() }
|
||||
fn d_input() -> String { "input".into() }
|
||||
fn d_output() -> String { "output".into() }
|
||||
fn d_text() -> String { "text".into() }
|
||||
fn d_role() -> String { "role".into() }
|
||||
fn d_content() -> String { "content".into() }
|
||||
fn d_seq_len() -> usize {
|
||||
2048
|
||||
}
|
||||
fn d_lora_r() -> usize {
|
||||
8
|
||||
}
|
||||
fn d_lora_alpha() -> f64 {
|
||||
16.0
|
||||
}
|
||||
fn d_one() -> usize {
|
||||
1
|
||||
}
|
||||
fn d_one_f() -> f64 {
|
||||
1.0
|
||||
}
|
||||
fn d_ten() -> usize {
|
||||
10
|
||||
}
|
||||
fn d_lr() -> f64 {
|
||||
2e-4
|
||||
}
|
||||
fn d_grad_norm() -> f64 {
|
||||
1.0
|
||||
}
|
||||
fn d_seed() -> u64 {
|
||||
42
|
||||
}
|
||||
fn d_true() -> bool {
|
||||
true
|
||||
}
|
||||
fn d_gens() -> usize {
|
||||
4
|
||||
}
|
||||
fn d_completion() -> usize {
|
||||
256
|
||||
}
|
||||
fn d_messages() -> String {
|
||||
"messages".into()
|
||||
}
|
||||
fn d_roles() -> Vec<String> {
|
||||
vec!["assistant".into()]
|
||||
}
|
||||
fn d_instruction() -> String {
|
||||
"instruction".into()
|
||||
}
|
||||
fn d_input() -> String {
|
||||
"input".into()
|
||||
}
|
||||
fn d_output() -> String {
|
||||
"output".into()
|
||||
}
|
||||
fn d_text() -> String {
|
||||
"text".into()
|
||||
}
|
||||
fn d_role() -> String {
|
||||
"role".into()
|
||||
}
|
||||
fn d_content() -> String {
|
||||
"content".into()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
@@ -306,7 +363,14 @@ warmup_ratio: 0.03
|
||||
assert_eq!(c.lora_scale(), 2.0);
|
||||
assert_eq!(c.lora_targets().len(), 6);
|
||||
assert_eq!(c.warmup(100), 3);
|
||||
assert_eq!(c.datasets[0].message_property_mappings.as_ref().unwrap().role, "from");
|
||||
assert_eq!(
|
||||
c.datasets[0]
|
||||
.message_property_mappings
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.role,
|
||||
"from"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -0,0 +1,171 @@
|
||||
//! Alpaca instruct rows. The prompt and the response are tokenized separately
|
||||
//! and concatenated; the response ends with eos and is the only part labelled
|
||||
//! unless `train_on_inputs`.
|
||||
|
||||
use super::text::{render, Message, Text};
|
||||
use crate::config::Dataset;
|
||||
use crate::{Example, Prompt, IGNORE_INDEX};
|
||||
use anyhow::{Context, Result};
|
||||
use serde_json::Value;
|
||||
|
||||
pub const SYSTEM: &str = "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.";
|
||||
pub const SYSTEM_NO_INPUT: &str = "Below is an instruction that describes a task. Write a response that appropriately completes the request.";
|
||||
|
||||
pub fn prompt_text(instruction: &str, input: &str) -> String {
|
||||
if input.is_empty() {
|
||||
format!("{SYSTEM_NO_INPUT}\n\n### Instruction:\n{instruction}\n\n### Response:\n")
|
||||
} else {
|
||||
format!(
|
||||
"{SYSTEM}\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn fields<'a>(ds: &Dataset, row: &'a Value) -> Result<(&'a str, &'a str, &'a str)> {
|
||||
let s = |k: &str| {
|
||||
row[k]
|
||||
.as_str()
|
||||
.with_context(|| format!("row has no `{k}` string"))
|
||||
};
|
||||
Ok((
|
||||
s(&ds.field_instruction)?,
|
||||
row[&ds.field_input].as_str().unwrap_or(""),
|
||||
s(&ds.field_output)?,
|
||||
))
|
||||
}
|
||||
|
||||
pub fn example(ds: &Dataset, text: &Text, train_on_inputs: bool, row: &Value) -> Result<Example> {
|
||||
let (instruction, input, output) = fields(ds, row)?;
|
||||
let p = text.encode(&prompt_text(instruction, input), true)?;
|
||||
let mut r = text.encode(output, true)?;
|
||||
if text.bos_id.is_some() && r.first() == text.bos_id.as_ref() {
|
||||
r.remove(0);
|
||||
}
|
||||
if r.last() != Some(&text.eos_id) {
|
||||
r.push(text.eos_id);
|
||||
}
|
||||
let mut labels: Vec<i64> = p
|
||||
.iter()
|
||||
.map(|&i| {
|
||||
if train_on_inputs {
|
||||
i as i64
|
||||
} else {
|
||||
IGNORE_INDEX
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
labels.extend(r.iter().map(|&i| i as i64));
|
||||
let mut input_ids = p;
|
||||
input_ids.extend(r);
|
||||
Ok(Example { input_ids, labels })
|
||||
}
|
||||
|
||||
/// The instruction (and input) as one user turn through the chat template;
|
||||
/// the output is the answer.
|
||||
pub fn prompt(ds: &Dataset, text: &Text, row: &Value) -> Result<Prompt> {
|
||||
let (instruction, input, output) = fields(ds, row)?;
|
||||
let content = if input.is_empty() {
|
||||
instruction.to_string()
|
||||
} else {
|
||||
format!("{instruction}\n\n{input}")
|
||||
};
|
||||
let s = render(
|
||||
text,
|
||||
&[Message {
|
||||
role: "user".into(),
|
||||
content,
|
||||
}],
|
||||
true,
|
||||
)?;
|
||||
Ok(Prompt {
|
||||
input_ids: text.encode(&s, false)?,
|
||||
text: s,
|
||||
answer: Some(output.to_string()),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::config::DatasetKind;
|
||||
use crate::data::fixture;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn strings_are_byte_exact() {
|
||||
assert_eq!(
|
||||
prompt_text("Say hi", ""),
|
||||
"Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\nSay hi\n\n### Response:\n"
|
||||
);
|
||||
assert_eq!(
|
||||
prompt_text("Say hi", "to Bob"),
|
||||
"Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\nSay hi\n\n### Input:\nto Bob\n\n### Response:\n"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mask_and_eos() {
|
||||
let text = fixture::text();
|
||||
let ds = fixture::dataset(DatasetKind::Alpaca, "");
|
||||
let row = json!({ "instruction": "hi there", "input": "", "output": "fine thanks" });
|
||||
let n = text
|
||||
.encode(&prompt_text("hi there", ""), true)
|
||||
.unwrap()
|
||||
.len();
|
||||
let ex = example(&ds, &text, false, &row).unwrap();
|
||||
assert_eq!(ex.input_ids.len(), n + 3);
|
||||
assert!(ex.labels[..n].iter().all(|&l| l == IGNORE_INDEX));
|
||||
assert_eq!(
|
||||
&ex.labels[n..],
|
||||
&[
|
||||
ex.input_ids[n] as i64,
|
||||
ex.input_ids[n + 1] as i64,
|
||||
text.eos_id as i64
|
||||
]
|
||||
);
|
||||
assert_eq!(ex.input_ids[n], text.tokenizer.token_to_id("fine").unwrap());
|
||||
let ex = example(&ds, &text, true, &row).unwrap();
|
||||
assert!(ex
|
||||
.labels
|
||||
.iter()
|
||||
.zip(&ex.input_ids)
|
||||
.all(|(&l, &i)| l == i as i64));
|
||||
assert!(example(&ds, &text, false, &json!({ "instruction": "x" })).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_drops_leading_bos_and_keeps_one_eos() {
|
||||
let mut text = fixture::text();
|
||||
text.bos_id = text.tokenizer.token_to_id("<|endoftext|>");
|
||||
let ds = fixture::dataset(DatasetKind::Alpaca, "");
|
||||
let row = json!({ "instruction": "hi", "output": "<|endoftext|> fine <|im_end|>" });
|
||||
let ex = example(&ds, &text, false, &row).unwrap();
|
||||
let tail = &ex.input_ids[ex.input_ids.len() - 2..];
|
||||
assert_eq!(
|
||||
tail,
|
||||
&[text.tokenizer.token_to_id("fine").unwrap(), text.eos_id]
|
||||
);
|
||||
assert_eq!(
|
||||
ex.input_ids.iter().filter(|&&i| i == text.eos_id).count(),
|
||||
1
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompt_goes_through_chat_template() {
|
||||
let text = fixture::text();
|
||||
let ds = fixture::dataset(DatasetKind::Alpaca, "");
|
||||
let p = prompt(
|
||||
&ds,
|
||||
&text,
|
||||
&json!({ "instruction": "hi there", "input": "how are you", "output": "fine" }),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
p.text,
|
||||
"<|im_start|>user\nhi there\n\nhow are you<|im_end|>\n<|im_start|>assistant\n"
|
||||
);
|
||||
assert_eq!(p.answer.as_deref(), Some("fine"));
|
||||
assert_eq!(p.input_ids.len(), 10);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,263 @@
|
||||
//! Conversations rendered through the chat template. Labels cover the content
|
||||
//! of trainable turns, found by rendering the conversation up to each turn
|
||||
//! twice (once with the turn's content replaced by `[[dummy_message]]`) and
|
||||
//! diffing the token ids from both ends, plus end-of-sequence tokens as
|
||||
//! `train_on_eos` asks.
|
||||
|
||||
use super::text::{render, Message, Text};
|
||||
use crate::config::{Dataset, MessageMap, TrainOnEos};
|
||||
use crate::{Example, Prompt, IGNORE_INDEX};
|
||||
use anyhow::{Context, Result};
|
||||
use serde_json::Value;
|
||||
|
||||
const DUMMY: &str = "[[dummy_message]]";
|
||||
|
||||
/// Canonical role for a source role name.
|
||||
pub fn role(r: &str) -> &str {
|
||||
match r {
|
||||
"human" => "user",
|
||||
"gpt" => "assistant",
|
||||
_ => r,
|
||||
}
|
||||
}
|
||||
|
||||
/// The row's turns with fields mapped and roles canonical, with a system turn
|
||||
/// from the row's `system` field in front when the first turn is not system.
|
||||
pub fn messages(ds: &Dataset, row: &Value) -> Result<Vec<Message>> {
|
||||
let map = ds.message_property_mappings.clone().unwrap_or(MessageMap {
|
||||
role: "role".into(),
|
||||
content: "content".into(),
|
||||
});
|
||||
let list = row[&ds.field_messages]
|
||||
.as_array()
|
||||
.with_context(|| format!("row has no `{}` array", ds.field_messages))?;
|
||||
let mut out: Vec<Message> = list
|
||||
.iter()
|
||||
.map(|m| {
|
||||
let field = |k: &str| {
|
||||
m[k].as_str()
|
||||
.map(String::from)
|
||||
.with_context(|| format!("message has no `{k}` string"))
|
||||
};
|
||||
Ok(Message {
|
||||
role: role(&field(&map.role)?).to_string(),
|
||||
content: field(&map.content)?,
|
||||
})
|
||||
})
|
||||
.collect::<Result<_>>()?;
|
||||
if let Some(sys) = row["system"].as_str() {
|
||||
if out.first().is_none_or(|m| m.role != "system") {
|
||||
out.insert(
|
||||
0,
|
||||
Message {
|
||||
role: "system".into(),
|
||||
content: sys.into(),
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
pub fn example(ds: &Dataset, text: &Text, train_on_inputs: bool, row: &Value) -> Result<Example> {
|
||||
let turns = messages(ds, row)?;
|
||||
let ids = text.encode(&render(text, &turns, false)?, false)?;
|
||||
let mut labels = vec![IGNORE_INDEX; ids.len()];
|
||||
let roles: Vec<&str> = ds.roles_to_train.iter().map(|r| role(r)).collect();
|
||||
let eos = ds.train_on_eos;
|
||||
let mut last_eos = None;
|
||||
for i in 0..turns.len() {
|
||||
let train = train_on_inputs || roles.contains(&turns[i].role.as_str());
|
||||
if !train && eos != TrainOnEos::All {
|
||||
continue;
|
||||
}
|
||||
let Some((start, end)) = span(text, &turns, i)? else {
|
||||
continue;
|
||||
};
|
||||
let end = end.min(ids.len());
|
||||
if train {
|
||||
for j in start..end {
|
||||
labels[j] = ids[j] as i64;
|
||||
}
|
||||
}
|
||||
let Some(e) = ids[end..]
|
||||
.iter()
|
||||
.position(|&t| t == text.eos_id)
|
||||
.map(|p| p + end)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
if e - end > 3 {
|
||||
continue;
|
||||
}
|
||||
last_eos = Some(e);
|
||||
if eos == TrainOnEos::All || (eos == TrainOnEos::Turn && train) {
|
||||
labels[e] = ids[e] as i64;
|
||||
}
|
||||
}
|
||||
if let (TrainOnEos::Last, Some(e)) = (eos, last_eos) {
|
||||
labels[e] = ids[e] as i64;
|
||||
}
|
||||
Ok(Example {
|
||||
input_ids: ids,
|
||||
labels,
|
||||
})
|
||||
}
|
||||
|
||||
/// Token span `[start, end)` of turn `i`'s content within the render of
|
||||
/// `turns[..=i]`, or `None` when the template leaves no trace of it.
|
||||
fn span(text: &Text, turns: &[Message], i: usize) -> Result<Option<(usize, usize)>> {
|
||||
let mut dummy = turns[..=i].to_vec();
|
||||
dummy[i].content = DUMMY.into();
|
||||
let d = text.encode(&render(text, &dummy, false)?, false)?;
|
||||
let f = text.encode(&render(text, &turns[..=i], false)?, false)?;
|
||||
if d.is_empty() || f.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
let n = d.len().min(f.len());
|
||||
let Some(start) = (0..n).find(|&k| d[k] != f[k]) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(end) = (0..n)
|
||||
.find(|&k| d[d.len() - 1 - k] != f[f.len() - 1 - k])
|
||||
.map(|k| f.len() - k)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok((end > start).then_some((start, end)))
|
||||
}
|
||||
|
||||
/// Everything before the last assistant turn, rendered with a generation
|
||||
/// prompt; that turn's content is the answer.
|
||||
pub fn prompt(ds: &Dataset, text: &Text, row: &Value) -> Result<Prompt> {
|
||||
let turns = messages(ds, row)?;
|
||||
let (prefix, answer) = match turns.iter().rposition(|m| m.role == "assistant") {
|
||||
Some(i) => (&turns[..i], Some(turns[i].content.clone())),
|
||||
None => (&turns[..], super::answer(row)),
|
||||
};
|
||||
let s = render(text, prefix, true)?;
|
||||
Ok(Prompt {
|
||||
input_ids: text.encode(&s, false)?,
|
||||
text: s,
|
||||
answer,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::config::DatasetKind;
|
||||
use crate::data::fixture;
|
||||
use serde_json::json;
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
fn row() -> Value {
|
||||
json!({ "messages": [
|
||||
{ "role": "system", "content": "be good" },
|
||||
{ "role": "user", "content": "hi there" },
|
||||
{ "role": "assistant", "content": "fine thanks" },
|
||||
{ "role": "user", "content": "how are you" },
|
||||
{ "role": "assistant", "content": "you too" },
|
||||
] })
|
||||
}
|
||||
|
||||
fn trained(ex: &Example) -> BTreeSet<usize> {
|
||||
(0..ex.labels.len())
|
||||
.filter(|&i| ex.labels[i] != IGNORE_INDEX)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn set(v: &[usize]) -> BTreeSet<usize> {
|
||||
v.iter().copied().collect()
|
||||
}
|
||||
|
||||
/// Tokens per turn: `[<|im_start|>, role, words..., <|im_end|>]`, so the
|
||||
/// five turns occupy 0..5, 5..10, 10..15, 15..21, 21..26.
|
||||
#[test]
|
||||
fn labels_per_train_on_eos() {
|
||||
let text = fixture::text();
|
||||
let mut ds = fixture::dataset(DatasetKind::ChatTemplate, "");
|
||||
let ex = example(&ds, &text, false, &row()).unwrap();
|
||||
assert_eq!(ex.input_ids.len(), 26);
|
||||
assert_eq!(ex.input_ids[14], text.eos_id);
|
||||
assert_eq!(ex.input_ids[25], text.eos_id);
|
||||
assert_eq!(trained(&ex), set(&[12, 13, 14, 23, 24, 25]));
|
||||
for i in trained(&ex) {
|
||||
assert_eq!(ex.labels[i], ex.input_ids[i] as i64);
|
||||
}
|
||||
|
||||
ds.train_on_eos = TrainOnEos::All;
|
||||
assert_eq!(
|
||||
trained(&example(&ds, &text, false, &row()).unwrap()),
|
||||
set(&[4, 9, 12, 13, 14, 20, 23, 24, 25])
|
||||
);
|
||||
ds.train_on_eos = TrainOnEos::Last;
|
||||
assert_eq!(
|
||||
trained(&example(&ds, &text, false, &row()).unwrap()),
|
||||
set(&[12, 13, 23, 24, 25])
|
||||
);
|
||||
ds.train_on_eos = TrainOnEos::None;
|
||||
assert_eq!(
|
||||
trained(&example(&ds, &text, false, &row()).unwrap()),
|
||||
set(&[12, 13, 23, 24])
|
||||
);
|
||||
|
||||
ds.train_on_eos = TrainOnEos::Turn;
|
||||
let all = example(&ds, &text, true, &row()).unwrap();
|
||||
assert_eq!(
|
||||
trained(&all),
|
||||
set(&[2, 3, 4, 7, 8, 9, 12, 13, 14, 17, 18, 19, 20, 23, 24, 25])
|
||||
);
|
||||
|
||||
ds.roles_to_train = vec!["gpt".into(), "human".into()];
|
||||
assert_eq!(
|
||||
trained(&example(&ds, &text, false, &row()).unwrap()),
|
||||
set(&[7, 8, 9, 12, 13, 14, 17, 18, 19, 20, 23, 24, 25])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sharegpt_mapping_and_system_field() {
|
||||
let text = fixture::text();
|
||||
let mut ds = fixture::dataset(DatasetKind::ChatTemplate, "");
|
||||
ds.field_messages = "conversations".into();
|
||||
ds.message_property_mappings = Some(MessageMap {
|
||||
role: "from".into(),
|
||||
content: "value".into(),
|
||||
});
|
||||
let r = json!({ "system": "be good", "conversations": [
|
||||
{ "from": "human", "value": "hi there" },
|
||||
{ "from": "gpt", "value": "fine thanks" },
|
||||
] });
|
||||
let m = messages(&ds, &r).unwrap();
|
||||
assert_eq!(
|
||||
m.iter().map(|m| m.role.as_str()).collect::<Vec<_>>(),
|
||||
["system", "user", "assistant"]
|
||||
);
|
||||
assert_eq!(m[0].content, "be good");
|
||||
assert_eq!(m[2].content, "fine thanks");
|
||||
assert_eq!(
|
||||
trained(&example(&ds, &text, false, &r).unwrap()),
|
||||
set(&[12, 13, 14])
|
||||
);
|
||||
assert!(messages(&ds, &json!({ "conversations": [{ "from": "human" }] })).is_err());
|
||||
assert!(messages(&ds, &json!({ "text": "x" })).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompt_splits_final_assistant_turn() {
|
||||
let text = fixture::text();
|
||||
let ds = fixture::dataset(DatasetKind::ChatTemplate, "");
|
||||
let p = prompt(&ds, &text, &row()).unwrap();
|
||||
assert!(p
|
||||
.text
|
||||
.ends_with("how are you<|im_end|>\n<|im_start|>assistant\n"));
|
||||
assert_eq!(p.answer.as_deref(), Some("you too"));
|
||||
assert_eq!(p.input_ids.len(), 23);
|
||||
let r =
|
||||
json!({ "messages": [{ "role": "user", "content": "hi there" }], "answer": "fine" });
|
||||
let p = prompt(&ds, &text, &r).unwrap();
|
||||
assert_eq!(p.answer.as_deref(), Some("fine"));
|
||||
assert!(p.text.ends_with("<|im_start|>assistant\n"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
//! Raw text rows: every token is trained and eos is appended once.
|
||||
|
||||
use super::text::Text;
|
||||
use crate::config::Dataset;
|
||||
use crate::{Example, Prompt};
|
||||
use anyhow::{Context, Result};
|
||||
use serde_json::Value;
|
||||
|
||||
fn field<'a>(ds: &Dataset, row: &'a Value) -> Result<&'a str> {
|
||||
row[&ds.field]
|
||||
.as_str()
|
||||
.with_context(|| format!("row has no `{}` string", ds.field))
|
||||
}
|
||||
|
||||
pub fn example(ds: &Dataset, text: &Text, row: &Value) -> Result<Example> {
|
||||
let mut ids = text.encode(field(ds, row)?, true)?;
|
||||
if ids.last() != Some(&text.eos_id) {
|
||||
ids.push(text.eos_id);
|
||||
}
|
||||
let labels = ids.iter().map(|&i| i as i64).collect();
|
||||
Ok(Example {
|
||||
input_ids: ids,
|
||||
labels,
|
||||
})
|
||||
}
|
||||
|
||||
/// The text itself is the prompt; `answer` or `solution` is the answer.
|
||||
pub fn prompt(ds: &Dataset, text: &Text, row: &Value) -> Result<Prompt> {
|
||||
let s = field(ds, row)?.to_string();
|
||||
Ok(Prompt {
|
||||
input_ids: text.encode(&s, true)?,
|
||||
text: s,
|
||||
answer: super::answer(row),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::config::DatasetKind;
|
||||
use crate::data::fixture;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn all_tokens_trained_with_one_eos() {
|
||||
let text = fixture::text();
|
||||
let mut ds = fixture::dataset(DatasetKind::Completion, "");
|
||||
let tok = |s: &str| text.tokenizer.token_to_id(s).unwrap();
|
||||
let ex = example(&ds, &text, &json!({ "text": "hi there" })).unwrap();
|
||||
assert_eq!(ex.input_ids, [tok("hi"), tok("there"), text.eos_id]);
|
||||
assert_eq!(
|
||||
ex.labels,
|
||||
[tok("hi") as i64, tok("there") as i64, text.eos_id as i64]
|
||||
);
|
||||
let ex = example(&ds, &text, &json!({ "text": "hi <|im_end|>" })).unwrap();
|
||||
assert_eq!(ex.input_ids, [tok("hi"), text.eos_id]);
|
||||
ds.field = "content".into();
|
||||
assert!(example(&ds, &text, &json!({ "text": "hi" })).is_err());
|
||||
let p = prompt(
|
||||
&ds,
|
||||
&text,
|
||||
&json!({ "content": "how are you", "solution": "fine" }),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
(p.text.as_str(), p.answer.as_deref(), p.input_ids.len()),
|
||||
("how are you", Some("fine"), 3)
|
||||
);
|
||||
assert_eq!(
|
||||
prompt(&ds, &text, &json!({ "content": "hi" }))
|
||||
.unwrap()
|
||||
.answer,
|
||||
None
|
||||
);
|
||||
}
|
||||
}
|
||||
+387
-2
@@ -1,2 +1,387 @@
|
||||
//! data
|
||||
pub fn preprocess(_cfg: &crate::Config) -> anyhow::Result<()> { anyhow::bail!("preprocess: not implemented") }
|
||||
//! Datasets to tokens. Rows come from [`source`], each dataset `type` turns a
|
||||
//! row into an [`Example`] or a [`Prompt`], and this module truncates, drops
|
||||
//! rows with nothing to learn from, shuffles, splits off validation, and pads
|
||||
//! batches.
|
||||
|
||||
pub mod alpaca;
|
||||
pub mod chat_template;
|
||||
pub mod completion;
|
||||
pub mod source;
|
||||
pub mod text;
|
||||
|
||||
pub use text::{render, Message, Text};
|
||||
|
||||
use crate::config::{Dataset, DatasetKind};
|
||||
use crate::{Batch, Config, Example, Prompt, IGNORE_INDEX};
|
||||
use anyhow::Result;
|
||||
use hanzo_ml::{Device, Tensor};
|
||||
use rand::seq::SliceRandom;
|
||||
use rand::SeedableRng;
|
||||
use serde_json::Value;
|
||||
|
||||
pub struct Tokenized {
|
||||
pub train: Vec<Example>,
|
||||
pub val: Vec<Example>,
|
||||
}
|
||||
|
||||
/// One dataset tokenized, with what was cut.
|
||||
struct Set {
|
||||
examples: Vec<Example>,
|
||||
tokens: usize,
|
||||
trained: usize,
|
||||
truncated: usize,
|
||||
dropped: usize,
|
||||
}
|
||||
|
||||
fn set(cfg: &Config, ds: &Dataset, text: &Text) -> Result<Set> {
|
||||
let mut s = Set {
|
||||
examples: Vec::new(),
|
||||
tokens: 0,
|
||||
trained: 0,
|
||||
truncated: 0,
|
||||
dropped: 0,
|
||||
};
|
||||
for row in source::rows(ds)? {
|
||||
let mut ex = match ds.kind {
|
||||
DatasetKind::ChatTemplate => {
|
||||
chat_template::example(ds, text, cfg.train_on_inputs, &row)?
|
||||
}
|
||||
DatasetKind::Alpaca => alpaca::example(ds, text, cfg.train_on_inputs, &row)?,
|
||||
DatasetKind::Completion => completion::example(ds, text, &row)?,
|
||||
};
|
||||
if ex.input_ids.len() > cfg.sequence_len {
|
||||
ex.input_ids.truncate(cfg.sequence_len);
|
||||
ex.labels.truncate(cfg.sequence_len);
|
||||
s.truncated += 1;
|
||||
}
|
||||
let trained = ex.labels.iter().filter(|&&l| l != IGNORE_INDEX).count();
|
||||
if ex.input_ids.len() < 2 || trained == 0 {
|
||||
s.dropped += 1;
|
||||
continue;
|
||||
}
|
||||
s.tokens += ex.input_ids.len();
|
||||
s.trained += trained;
|
||||
s.examples.push(ex);
|
||||
}
|
||||
Ok(s)
|
||||
}
|
||||
|
||||
/// Every dataset tokenized and concatenated, shuffled with `cfg.seed`, with
|
||||
/// `val_set_size` (a fraction when <= 1, a count when > 1) taken from the end.
|
||||
pub fn load(cfg: &Config, text: &Text) -> Result<Tokenized> {
|
||||
let mut all = Vec::new();
|
||||
for ds in &cfg.datasets {
|
||||
all.extend(set(cfg, ds, text)?.examples);
|
||||
}
|
||||
all.shuffle(&mut rand::rngs::StdRng::seed_from_u64(cfg.seed));
|
||||
let n = all.len();
|
||||
let val = if cfg.val_set_size > 1.0 {
|
||||
cfg.val_set_size as usize
|
||||
} else {
|
||||
(cfg.val_set_size * n as f64).ceil() as usize
|
||||
}
|
||||
.min(n);
|
||||
let val = all.split_off(n - val);
|
||||
Ok(Tokenized { train: all, val })
|
||||
}
|
||||
|
||||
/// GRPO prompts from every dataset; rows longer than `sequence_len` are dropped.
|
||||
pub fn load_prompts(cfg: &Config, text: &Text) -> Result<Vec<Prompt>> {
|
||||
let mut out = Vec::new();
|
||||
let mut long = 0;
|
||||
for ds in &cfg.datasets {
|
||||
for row in source::rows(ds)? {
|
||||
let p = match ds.kind {
|
||||
DatasetKind::ChatTemplate => chat_template::prompt(ds, text, &row)?,
|
||||
DatasetKind::Alpaca => alpaca::prompt(ds, text, &row)?,
|
||||
DatasetKind::Completion => completion::prompt(ds, text, &row)?,
|
||||
};
|
||||
if p.input_ids.len() > cfg.sequence_len {
|
||||
long += 1;
|
||||
continue;
|
||||
}
|
||||
out.push(p);
|
||||
}
|
||||
}
|
||||
if long > 0 {
|
||||
tracing::info!(dropped = long, "prompts longer than sequence_len");
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// A row's reference answer for a verifier: `answer`, else `solution`.
|
||||
fn answer(row: &Value) -> Option<String> {
|
||||
row["answer"]
|
||||
.as_str()
|
||||
.or_else(|| row["solution"].as_str())
|
||||
.map(String::from)
|
||||
}
|
||||
|
||||
/// Right-pad to the longest row. Labels are not shifted; padding is
|
||||
/// [`IGNORE_INDEX`]. The mask is `None` when no row needed padding.
|
||||
pub fn collate(exs: &[Example], pad_id: u32, device: &Device) -> hanzo_ml::Result<Batch> {
|
||||
let b = exs.len();
|
||||
let t = exs.iter().map(|e| e.input_ids.len()).max().unwrap_or(0);
|
||||
let (mut ids, mut labels, mut mask) = (
|
||||
Vec::with_capacity(b * t),
|
||||
Vec::with_capacity(b * t),
|
||||
Vec::with_capacity(b * t),
|
||||
);
|
||||
let mut padded = false;
|
||||
for e in exs {
|
||||
let n = e.input_ids.len();
|
||||
ids.extend(&e.input_ids);
|
||||
ids.resize(ids.len() + t - n, pad_id);
|
||||
labels.extend(&e.labels);
|
||||
labels.resize(labels.len() + t - n, IGNORE_INDEX);
|
||||
mask.resize(mask.len() + n, 1u8);
|
||||
mask.resize(mask.len() + t - n, 0);
|
||||
padded |= n < t;
|
||||
}
|
||||
Ok(Batch {
|
||||
input_ids: Tensor::from_vec(ids, (b, t), device)?,
|
||||
labels: Tensor::from_vec(labels, (b, t), device)?,
|
||||
attention_mask: padded
|
||||
.then(|| Tensor::from_vec(mask, (b, t), device))
|
||||
.transpose()?,
|
||||
})
|
||||
}
|
||||
|
||||
/// Tokenize every dataset and report what training would see.
|
||||
pub fn preprocess(cfg: &Config) -> Result<()> {
|
||||
let text = Text::load(cfg)?;
|
||||
let mut total = Set {
|
||||
examples: Vec::new(),
|
||||
tokens: 0,
|
||||
trained: 0,
|
||||
truncated: 0,
|
||||
dropped: 0,
|
||||
};
|
||||
let line = |name: &str, s: &Set| {
|
||||
println!(
|
||||
"{name}: {} examples, {} tokens, {} trained, {} truncated, {} dropped",
|
||||
s.examples.len(),
|
||||
s.tokens,
|
||||
s.trained,
|
||||
s.truncated,
|
||||
s.dropped
|
||||
)
|
||||
};
|
||||
for ds in &cfg.datasets {
|
||||
let s = set(cfg, ds, &text)?;
|
||||
line(&ds.path, &s);
|
||||
total.tokens += s.tokens;
|
||||
total.trained += s.trained;
|
||||
total.truncated += s.truncated;
|
||||
total.dropped += s.dropped;
|
||||
total.examples.extend(s.examples);
|
||||
}
|
||||
line("total", &total);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) mod fixture {
|
||||
use super::Text;
|
||||
use crate::config::{Config, Dataset, DatasetKind};
|
||||
use std::path::PathBuf;
|
||||
|
||||
const WORDS: [&str; 24] = [
|
||||
"<unk>",
|
||||
"system",
|
||||
"user",
|
||||
"assistant",
|
||||
"be",
|
||||
"good",
|
||||
"hi",
|
||||
"there",
|
||||
"fine",
|
||||
"thanks",
|
||||
"how",
|
||||
"are",
|
||||
"you",
|
||||
"too",
|
||||
"one",
|
||||
"two",
|
||||
"three",
|
||||
"a",
|
||||
"b",
|
||||
"c",
|
||||
"d",
|
||||
"e",
|
||||
"f",
|
||||
"g",
|
||||
];
|
||||
|
||||
/// WordLevel over [`WORDS`] with chatml special tokens; eos is `<|im_end|>`.
|
||||
pub fn text() -> Text {
|
||||
let special = ["<|im_start|>", "<|im_end|>", "<|endoftext|>"];
|
||||
let vocab: serde_json::Map<String, serde_json::Value> = special
|
||||
.iter()
|
||||
.chain(WORDS.iter())
|
||||
.enumerate()
|
||||
.map(|(i, w)| (w.to_string(), i.into()))
|
||||
.collect();
|
||||
let added: Vec<serde_json::Value> = special
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, t)| {
|
||||
serde_json::json!({ "id": i, "content": t, "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": true })
|
||||
})
|
||||
.collect();
|
||||
let json = serde_json::json!({
|
||||
"version": "1.0",
|
||||
"added_tokens": added,
|
||||
"pre_tokenizer": { "type": "Whitespace" },
|
||||
"model": { "type": "WordLevel", "vocab": vocab, "unk_token": "<unk>" },
|
||||
});
|
||||
let tokenizer = tokenizers::Tokenizer::from_bytes(json.to_string().as_bytes()).unwrap();
|
||||
Text {
|
||||
template: Some(super::text::CHATML.into()),
|
||||
eos_id: 1,
|
||||
pad_id: 2,
|
||||
bos_id: None,
|
||||
tokenizer,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn dataset(kind: DatasetKind, path: &str) -> Dataset {
|
||||
serde_yaml::from_str(&format!(
|
||||
"{{path: '{path}', type: {}}}",
|
||||
serde_yaml::to_string(&kind).unwrap().trim()
|
||||
))
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
pub fn config(datasets: Vec<Dataset>) -> Config {
|
||||
let mut cfg: Config = serde_yaml::from_str(
|
||||
"base_model: x\noutput_dir: o\ndatasets: [{path: p, type: completion}]",
|
||||
)
|
||||
.unwrap();
|
||||
cfg.datasets = datasets;
|
||||
cfg
|
||||
}
|
||||
|
||||
/// A fresh empty directory under the system temp dir.
|
||||
pub fn dir(name: &str) -> PathBuf {
|
||||
let d = std::env::temp_dir().join(format!("gym-data-{}-{name}", std::process::id()));
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
std::fs::create_dir_all(&d).unwrap();
|
||||
d
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::config::DatasetKind;
|
||||
|
||||
fn ex(ids: &[u32]) -> Example {
|
||||
Example {
|
||||
input_ids: ids.to_vec(),
|
||||
labels: ids.iter().map(|&i| i as i64).collect(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn collate_pads_right() {
|
||||
let dev = Device::Cpu;
|
||||
let b = collate(&[ex(&[5, 6, 7]), ex(&[8])], 2, &dev).unwrap();
|
||||
assert_eq!(
|
||||
b.input_ids.to_vec2::<u32>().unwrap(),
|
||||
[[5, 6, 7], [8, 2, 2]]
|
||||
);
|
||||
assert_eq!(
|
||||
b.labels.to_vec2::<i64>().unwrap(),
|
||||
[[5, 6, 7], [8, IGNORE_INDEX, IGNORE_INDEX]]
|
||||
);
|
||||
assert_eq!(
|
||||
b.attention_mask.unwrap().to_vec2::<u8>().unwrap(),
|
||||
[[1, 1, 1], [1, 0, 0]]
|
||||
);
|
||||
let b = collate(&[ex(&[5, 6]), ex(&[7, 8])], 2, &dev).unwrap();
|
||||
assert!(b.attention_mask.is_none());
|
||||
assert_eq!(b.input_ids.dims(), [2, 2]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_truncates_drops_shuffles_and_splits() {
|
||||
let dir = fixture::dir("load");
|
||||
let mut rows = String::new();
|
||||
for w in WORDS10 {
|
||||
rows += &format!("{{\"text\": \"{w}\"}}\n");
|
||||
}
|
||||
rows += "{\"text\": \"\"}\n";
|
||||
std::fs::write(dir.join("train.jsonl"), rows).unwrap();
|
||||
let text = fixture::text();
|
||||
let mut cfg = fixture::config(vec![fixture::dataset(
|
||||
DatasetKind::Completion,
|
||||
dir.join("train.jsonl").to_str().unwrap(),
|
||||
)]);
|
||||
cfg.sequence_len = 3;
|
||||
cfg.val_set_size = 0.2;
|
||||
|
||||
let s = set(&cfg, &cfg.datasets[0], &text).unwrap();
|
||||
assert_eq!((s.examples.len(), s.truncated, s.dropped), (10, 2, 1));
|
||||
assert!(s.examples.iter().all(|e| e.input_ids.len() <= 3));
|
||||
assert_eq!(s.tokens, 7 * 2 + 3 * 3);
|
||||
assert_eq!(s.trained, s.tokens);
|
||||
|
||||
let t = load(&cfg, &text).unwrap();
|
||||
assert_eq!((t.train.len(), t.val.len()), (8, 2));
|
||||
let again = load(&cfg, &text).unwrap();
|
||||
assert_eq!(t.train, again.train);
|
||||
let order: Vec<u32> = t
|
||||
.train
|
||||
.iter()
|
||||
.chain(&t.val)
|
||||
.map(|e| e.input_ids[0])
|
||||
.collect();
|
||||
let mut sorted = order.clone();
|
||||
sorted.sort();
|
||||
assert_ne!(order, sorted);
|
||||
|
||||
cfg.val_set_size = 3.0;
|
||||
assert_eq!(load(&cfg, &text).unwrap().val.len(), 3);
|
||||
cfg.val_set_size = 0.0;
|
||||
assert_eq!(load(&cfg, &text).unwrap().val.len(), 0);
|
||||
|
||||
cfg.datasets.push(cfg.datasets[0].clone());
|
||||
assert_eq!(load(&cfg, &text).unwrap().train.len(), 20);
|
||||
}
|
||||
|
||||
const WORDS10: [&str; 10] = [
|
||||
"a",
|
||||
"b",
|
||||
"c",
|
||||
"d",
|
||||
"e",
|
||||
"f",
|
||||
"g",
|
||||
"hi there",
|
||||
"fine thanks you",
|
||||
"how are you too",
|
||||
];
|
||||
|
||||
#[test]
|
||||
fn load_prompts_drops_long_rows() {
|
||||
let dir = fixture::dir("prompts");
|
||||
std::fs::write(
|
||||
dir.join("train.jsonl"),
|
||||
"{\"text\": \"hi\", \"answer\": \"a\"}\n{\"text\": \"hi there fine thanks\"}\n",
|
||||
)
|
||||
.unwrap();
|
||||
let text = fixture::text();
|
||||
let mut cfg = fixture::config(vec![fixture::dataset(
|
||||
DatasetKind::Completion,
|
||||
dir.to_str().unwrap(),
|
||||
)]);
|
||||
cfg.sequence_len = 2;
|
||||
let p = load_prompts(&cfg, &text).unwrap();
|
||||
assert_eq!(p.len(), 1);
|
||||
assert_eq!(
|
||||
(p[0].text.as_str(), p[0].answer.as_deref()),
|
||||
("hi", Some("a"))
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,322 @@
|
||||
//! Rows of a dataset. `path` is a local file, a local directory, or a Hugging
|
||||
//! Face dataset id; files are `.jsonl`, `.json` (an array) or `.parquet`. Rows
|
||||
//! are read in file order and cut to the requested split slice.
|
||||
|
||||
use crate::config::Dataset;
|
||||
use anyhow::{bail, Context, Result};
|
||||
use hf_hub::api::sync::Api;
|
||||
use hf_hub::{Repo, RepoType};
|
||||
use parquet::file::reader::{FileReader, SerializedFileReader};
|
||||
use serde_json::Value;
|
||||
use std::fs::File;
|
||||
use std::io::{BufRead, BufReader};
|
||||
use std::ops::Range;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
pub fn rows(ds: &Dataset) -> Result<Vec<Value>> {
|
||||
let split = Split::parse(ds.split.as_deref().unwrap_or("train"))?;
|
||||
let mut all = Vec::new();
|
||||
for f in files(ds, &split.name)? {
|
||||
all.extend(read(&f).with_context(|| f.display().to_string())?);
|
||||
}
|
||||
let r = split.range(all.len());
|
||||
all.truncate(r.end);
|
||||
all.drain(..r.start);
|
||||
Ok(all)
|
||||
}
|
||||
|
||||
fn files(ds: &Dataset, split: &str) -> Result<Vec<PathBuf>> {
|
||||
let p = Path::new(&ds.path);
|
||||
if p.is_file() {
|
||||
return Ok(vec![p.to_path_buf()]);
|
||||
}
|
||||
if p.is_dir() {
|
||||
let names = match &ds.data_files {
|
||||
Some(d) => d.clone(),
|
||||
None => {
|
||||
let mut v: Vec<String> = std::fs::read_dir(p)?
|
||||
.filter_map(|e| e.ok()?.file_name().into_string().ok())
|
||||
.filter(|n| is_data(n))
|
||||
.collect();
|
||||
v.sort();
|
||||
v
|
||||
}
|
||||
};
|
||||
return Ok(choose(names, split, ds.data_files.is_some())?
|
||||
.iter()
|
||||
.map(|n| p.join(n))
|
||||
.collect());
|
||||
}
|
||||
let repo = Api::new()?.repo(Repo::new(ds.path.clone(), RepoType::Dataset));
|
||||
let names = match &ds.data_files {
|
||||
Some(d) => d.clone(),
|
||||
None => repo
|
||||
.info()?
|
||||
.siblings
|
||||
.into_iter()
|
||||
.map(|s| s.rfilename)
|
||||
.filter(|n| is_data(n))
|
||||
.collect(),
|
||||
};
|
||||
choose(names, split, ds.data_files.is_some())?
|
||||
.iter()
|
||||
.map(|n| Ok(repo.get(n)?))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn is_data(name: &str) -> bool {
|
||||
["jsonl", "json", "parquet"].contains(
|
||||
&Path::new(name)
|
||||
.extension()
|
||||
.and_then(|e| e.to_str())
|
||||
.unwrap_or(""),
|
||||
)
|
||||
}
|
||||
|
||||
/// Files for `split`: those in a directory named after it or whose name
|
||||
/// mentions it. Files that name no split at all belong to every split.
|
||||
fn choose(names: Vec<String>, split: &str, explicit: bool) -> Result<Vec<String>> {
|
||||
if explicit {
|
||||
return Ok(names);
|
||||
}
|
||||
let mentions = |n: &str, s: &str| {
|
||||
Path::new(n)
|
||||
.iter()
|
||||
.any(|c| c.to_str().is_some_and(|c| c.contains(s)))
|
||||
};
|
||||
let mut v: Vec<String> = names
|
||||
.iter()
|
||||
.filter(|n| mentions(n, split))
|
||||
.cloned()
|
||||
.collect();
|
||||
if v.is_empty() {
|
||||
v = names
|
||||
.into_iter()
|
||||
.filter(|n| {
|
||||
!["train", "test", "validation", "valid", "dev"]
|
||||
.iter()
|
||||
.any(|s| mentions(n, s))
|
||||
})
|
||||
.collect();
|
||||
}
|
||||
if v.is_empty() {
|
||||
bail!("no data files for split `{split}`");
|
||||
}
|
||||
Ok(v)
|
||||
}
|
||||
|
||||
fn read(p: &Path) -> Result<Vec<Value>> {
|
||||
match p.extension().and_then(|e| e.to_str()) {
|
||||
Some("jsonl") => BufReader::new(File::open(p)?)
|
||||
.lines()
|
||||
.filter(|l| l.as_ref().is_ok_and(|l| !l.trim().is_empty()))
|
||||
.map(|l| Ok(serde_json::from_str(&l?)?))
|
||||
.collect(),
|
||||
Some("json") => match serde_json::from_reader(BufReader::new(File::open(p)?))? {
|
||||
Value::Array(v) => Ok(v),
|
||||
_ => bail!("expected a JSON array"),
|
||||
},
|
||||
Some("parquet") => SerializedFileReader::new(File::open(p)?)?
|
||||
.get_row_iter(None)?
|
||||
.map(|r| Ok(r?.to_json_value()))
|
||||
.collect(),
|
||||
_ => bail!("unsupported file type"),
|
||||
}
|
||||
}
|
||||
|
||||
/// A Hugging Face split spec: `name`, `name[a:b]`, with bounds absolute
|
||||
/// (`1000`) or percent (`20%`, floored) and either side open.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct Split {
|
||||
pub name: String,
|
||||
lo: Bound,
|
||||
hi: Bound,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
enum Bound {
|
||||
Open,
|
||||
Abs(usize),
|
||||
Pct(f64),
|
||||
}
|
||||
|
||||
impl Split {
|
||||
pub fn parse(s: &str) -> Result<Split> {
|
||||
let (name, slice) = match s.split_once('[') {
|
||||
Some((n, rest)) => (
|
||||
n,
|
||||
Some(
|
||||
rest.strip_suffix(']')
|
||||
.with_context(|| format!("bad split `{s}`"))?,
|
||||
),
|
||||
),
|
||||
None => (s, None),
|
||||
};
|
||||
let (lo, hi) = match slice {
|
||||
Some(sl) => {
|
||||
let (a, b) = sl
|
||||
.split_once(':')
|
||||
.with_context(|| format!("bad split `{s}`"))?;
|
||||
(Bound::parse(a)?, Bound::parse(b)?)
|
||||
}
|
||||
None => (Bound::Open, Bound::Open),
|
||||
};
|
||||
Ok(Split {
|
||||
name: name.to_string(),
|
||||
lo,
|
||||
hi,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn range(&self, n: usize) -> Range<usize> {
|
||||
let lo = self.lo.at(n, 0);
|
||||
lo..self.hi.at(n, n).max(lo)
|
||||
}
|
||||
}
|
||||
|
||||
impl Bound {
|
||||
fn parse(s: &str) -> Result<Bound> {
|
||||
let s = s.trim();
|
||||
Ok(match s.strip_suffix('%') {
|
||||
_ if s.is_empty() => Bound::Open,
|
||||
Some(p) => Bound::Pct(
|
||||
p.parse::<f64>()
|
||||
.with_context(|| format!("bad percent `{s}`"))?,
|
||||
),
|
||||
None => Bound::Abs(s.parse().with_context(|| format!("bad index `{s}`"))?),
|
||||
})
|
||||
}
|
||||
|
||||
fn at(&self, n: usize, open: usize) -> usize {
|
||||
match *self {
|
||||
Bound::Open => open,
|
||||
Bound::Abs(i) => i.min(n),
|
||||
Bound::Pct(p) => ((p * n as f64 / 100.0).floor() as usize).min(n),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::data::fixture;
|
||||
use parquet::data_type::{ByteArray, ByteArrayType};
|
||||
use parquet::file::properties::WriterProperties;
|
||||
use parquet::file::writer::{SerializedFileWriter, SerializedRowGroupWriter};
|
||||
use parquet::schema::parser::parse_message_type;
|
||||
use std::sync::Arc;
|
||||
|
||||
fn r(spec: &str, n: usize) -> Range<usize> {
|
||||
Split::parse(spec).unwrap().range(n)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_grammar() {
|
||||
assert_eq!(Split::parse("train").unwrap().name, "train");
|
||||
assert_eq!(r("train", 10), 0..10);
|
||||
assert_eq!(r("test", 10), 0..10);
|
||||
assert_eq!(r("validation[:3]", 10), 0..3);
|
||||
assert_eq!(r("train[:20%]", 10), 0..2);
|
||||
assert_eq!(r("train[:1000]", 10), 0..10);
|
||||
assert_eq!(r("train[10%:]", 10), 1..10);
|
||||
assert_eq!(r("train[100:200]", 150), 100..150);
|
||||
assert_eq!(r("train[25%:75%]", 7), 1..5);
|
||||
assert_eq!(r("train[33%:]", 100), 33..100);
|
||||
assert!(Split::parse("train[1]").is_err());
|
||||
assert!(Split::parse("train[a:b]").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn jsonl_and_json_from_directory() {
|
||||
let dir = fixture::dir("source");
|
||||
std::fs::write(
|
||||
dir.join("train.jsonl"),
|
||||
"{\"text\":\"a\"}\n\n{\"text\":\"b\"}\n",
|
||||
)
|
||||
.unwrap();
|
||||
std::fs::write(dir.join("test.json"), "[{\"text\":\"t\"}]").unwrap();
|
||||
std::fs::write(dir.join("notes.txt"), "ignored").unwrap();
|
||||
let mut ds = fixture::dataset(
|
||||
crate::config::DatasetKind::Completion,
|
||||
dir.to_str().unwrap(),
|
||||
);
|
||||
let texts = |ds: &Dataset| {
|
||||
rows(ds)
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|v| v["text"].as_str().unwrap().to_string())
|
||||
.collect::<Vec<_>>()
|
||||
};
|
||||
assert_eq!(texts(&ds), ["a", "b"]);
|
||||
ds.split = Some("test".into());
|
||||
assert_eq!(texts(&ds), ["t"]);
|
||||
ds.split = Some("validation".into());
|
||||
assert!(rows(&ds).is_err());
|
||||
ds.split = Some("train[1:]".into());
|
||||
ds.data_files = Some(vec!["train.jsonl".into(), "test.json".into()]);
|
||||
assert_eq!(texts(&ds), ["b", "t"]);
|
||||
ds.split = None;
|
||||
ds.data_files = None;
|
||||
ds.path = dir.join("test.json").to_str().unwrap().into();
|
||||
assert_eq!(texts(&ds), ["t"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parquet_rows_with_nested_messages() {
|
||||
let dir = fixture::dir("parquet");
|
||||
let path = dir.join("train.parquet");
|
||||
let schema = Arc::new(
|
||||
parse_message_type(
|
||||
"message row {
|
||||
required binary text (UTF8);
|
||||
optional group messages (LIST) {
|
||||
repeated group list {
|
||||
required group element {
|
||||
required binary role (UTF8);
|
||||
required binary content (UTF8);
|
||||
}
|
||||
}
|
||||
}
|
||||
}",
|
||||
)
|
||||
.unwrap(),
|
||||
);
|
||||
let mut w = SerializedFileWriter::new(
|
||||
File::create(&path).unwrap(),
|
||||
schema,
|
||||
Arc::new(WriterProperties::builder().build()),
|
||||
)
|
||||
.unwrap();
|
||||
let mut g = w.next_row_group().unwrap();
|
||||
let col = |g: &mut SerializedRowGroupWriter<'_, File>,
|
||||
vals: &[&str],
|
||||
levels: Option<(&[i16], &[i16])>| {
|
||||
let vals: Vec<ByteArray> = vals
|
||||
.iter()
|
||||
.map(|s| ByteArray::from(s.as_bytes().to_vec()))
|
||||
.collect();
|
||||
let mut c = g.next_column().unwrap().unwrap();
|
||||
c.typed::<ByteArrayType>()
|
||||
.write_batch(&vals, levels.map(|l| l.0), levels.map(|l| l.1))
|
||||
.unwrap();
|
||||
c.close().unwrap();
|
||||
};
|
||||
let nested = Some((&[2i16, 2, 2, 2][..], &[0i16, 1, 0, 1][..]));
|
||||
col(&mut g, &["hi there", "fine thanks"], None);
|
||||
col(&mut g, &["user", "assistant", "user", "assistant"], nested);
|
||||
col(
|
||||
&mut g,
|
||||
&["hi there", "fine thanks", "how are you", "you too"],
|
||||
nested,
|
||||
);
|
||||
g.close().unwrap();
|
||||
w.close().unwrap();
|
||||
|
||||
let v = read(&path).unwrap();
|
||||
assert_eq!(v.len(), 2);
|
||||
assert_eq!(v[0]["text"], "hi there");
|
||||
assert_eq!(v[1]["messages"][1]["role"], "assistant");
|
||||
assert_eq!(v[1]["messages"][1]["content"], "you too");
|
||||
assert_eq!(v[0]["messages"].as_array().unwrap().len(), 2);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,344 @@
|
||||
//! The tokenizer, its special ids, and the chat template a run renders with.
|
||||
|
||||
use crate::hub;
|
||||
use crate::Config;
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
use hf_hub::api::sync::Api;
|
||||
use hf_hub::{Repo, RepoType};
|
||||
use minijinja::{context, Environment, Error, ErrorKind};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use std::path::{Path, PathBuf};
|
||||
use tokenizers::Tokenizer;
|
||||
|
||||
/// One chat turn as the template sees it.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct Message {
|
||||
pub role: String,
|
||||
pub content: String,
|
||||
}
|
||||
|
||||
pub struct Text {
|
||||
pub tokenizer: Tokenizer,
|
||||
pub template: Option<String>,
|
||||
pub eos_id: u32,
|
||||
pub pad_id: u32,
|
||||
pub bos_id: Option<u32>,
|
||||
}
|
||||
|
||||
impl Text {
|
||||
pub fn load(cfg: &Config) -> Result<Text> {
|
||||
let (tok, tc) = match &cfg.tokenizer_config {
|
||||
Some(src) => files(src)?,
|
||||
None => {
|
||||
let s = hub::snapshot(&cfg.base_model)?;
|
||||
(s.tokenizer, s.tokenizer_config)
|
||||
}
|
||||
};
|
||||
let tokenizer =
|
||||
Tokenizer::from_file(&tok).map_err(|e| anyhow!("{}: {e}", tok.display()))?;
|
||||
let tc: Value = match tc {
|
||||
Some(p) => serde_json::from_slice(&std::fs::read(&p)?)
|
||||
.with_context(|| p.display().to_string())?,
|
||||
None => Value::Null,
|
||||
};
|
||||
let id = |k: &str| token(&tc[k]).and_then(|t| tokenizer.token_to_id(t));
|
||||
let eos_id = id("eos_token").context("tokenizer_config.json: eos_token")?;
|
||||
Ok(Text {
|
||||
template: template(cfg.chat_template.as_deref(), &tc),
|
||||
eos_id,
|
||||
pad_id: id("pad_token").unwrap_or(eos_id),
|
||||
bos_id: id("bos_token"),
|
||||
tokenizer,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn encode(&self, s: &str, add_special_tokens: bool) -> Result<Vec<u32>> {
|
||||
Ok(self
|
||||
.tokenizer
|
||||
.encode(s, add_special_tokens)
|
||||
.map_err(|e| anyhow!("{e}"))?
|
||||
.get_ids()
|
||||
.to_vec())
|
||||
}
|
||||
}
|
||||
|
||||
/// `tokenizer.json` and `tokenizer_config.json` from a directory or a model repo.
|
||||
fn files(src: &str) -> Result<(PathBuf, Option<PathBuf>)> {
|
||||
let p = Path::new(src);
|
||||
if p.is_dir() {
|
||||
return Ok((
|
||||
p.join("tokenizer.json"),
|
||||
Some(p.join("tokenizer_config.json")).filter(|p| p.exists()),
|
||||
));
|
||||
}
|
||||
let repo = Api::new()?.repo(Repo::new(src.to_string(), RepoType::Model));
|
||||
Ok((
|
||||
repo.get("tokenizer.json").context("tokenizer.json")?,
|
||||
repo.get("tokenizer_config.json").ok(),
|
||||
))
|
||||
}
|
||||
|
||||
/// A special token entry: a string or `{"content": ...}`.
|
||||
fn token(v: &Value) -> Option<&str> {
|
||||
v.as_str().or_else(|| v["content"].as_str())
|
||||
}
|
||||
|
||||
/// `None` or `tokenizer_default` takes the tokenizer's own template; a known
|
||||
/// name takes the built-in; `tokenizer_default_fallback_<name>` prefers the
|
||||
/// tokenizer's and falls back to the name; anything else is inline Jinja.
|
||||
fn template(choice: Option<&str>, tc: &Value) -> Option<String> {
|
||||
let own = || match &tc["chat_template"] {
|
||||
Value::String(s) => Some(s.clone()),
|
||||
Value::Array(a) => a
|
||||
.iter()
|
||||
.find(|t| t["name"] == "default")
|
||||
.and_then(|t| t["template"].as_str().map(String::from)),
|
||||
_ => None,
|
||||
};
|
||||
let named = |n: &str| Some(builtin(n).unwrap_or(n).to_string());
|
||||
match choice {
|
||||
None | Some("tokenizer_default") => own(),
|
||||
Some(c) => match c.strip_prefix("tokenizer_default_fallback_") {
|
||||
Some(n) => own().or_else(|| named(n)),
|
||||
None => named(c),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn builtin(name: &str) -> Option<&'static str> {
|
||||
match name {
|
||||
"chatml" => Some(CHATML),
|
||||
"qwen3" => Some(QWEN3),
|
||||
"llama3" => Some(LLAMA3),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Render `messages` through the chat template with Hugging Face's variables:
|
||||
/// `messages`, `add_generation_prompt`, `bos_token`, `eos_token`,
|
||||
/// `enable_thinking` (false) and a `raise_exception` function.
|
||||
pub fn render(text: &Text, messages: &[Message], add_generation_prompt: bool) -> Result<String> {
|
||||
let src = text
|
||||
.template
|
||||
.as_deref()
|
||||
.context("no chat template: set `chat_template` or use a tokenizer that has one")?;
|
||||
let mut env = Environment::new();
|
||||
env.set_trim_blocks(true);
|
||||
env.set_lstrip_blocks(true);
|
||||
env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback);
|
||||
env.add_function("raise_exception", |msg: String| -> Result<(), Error> {
|
||||
Err(Error::new(ErrorKind::InvalidOperation, msg))
|
||||
});
|
||||
let tok = |id: Option<u32>| {
|
||||
id.and_then(|i| text.tokenizer.id_to_token(i))
|
||||
.unwrap_or_default()
|
||||
};
|
||||
Ok(env.template_from_str(src)?.render(context! {
|
||||
messages,
|
||||
add_generation_prompt,
|
||||
bos_token => tok(text.bos_id),
|
||||
eos_token => tok(Some(text.eos_id)),
|
||||
enable_thinking => false,
|
||||
})?)
|
||||
}
|
||||
|
||||
pub const CHATML: &str = "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}";
|
||||
|
||||
pub const LLAMA3: &str = "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}";
|
||||
|
||||
pub const QWEN3: &str = r##"{%- if tools %}
|
||||
{{- '<|im_start|>system\n' }}
|
||||
{%- if messages[0].role == 'system' %}
|
||||
{{- messages[0].content + '\n\n' }}
|
||||
{%- endif %}
|
||||
{{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
|
||||
{%- for tool in tools %}
|
||||
{{- "\n" }}
|
||||
{{- tool | tojson }}
|
||||
{%- endfor %}
|
||||
{{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
|
||||
{%- else %}
|
||||
{%- if messages[0].role == 'system' %}
|
||||
{{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
|
||||
{%- for message in messages[::-1] %}
|
||||
{%- set index = (messages|length - 1) - loop.index0 %}
|
||||
{%- if ns.multi_step_tool and message.role == "user" and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
|
||||
{%- set ns.multi_step_tool = false %}
|
||||
{%- set ns.last_query_index = index %}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- for message in messages %}
|
||||
{%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
|
||||
{{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
|
||||
{%- elif message.role == "assistant" %}
|
||||
{%- set content = message.content %}
|
||||
{%- set reasoning_content = '' %}
|
||||
{%- if message.reasoning_content is defined and message.reasoning_content is not none %}
|
||||
{%- set reasoning_content = message.reasoning_content %}
|
||||
{%- else %}
|
||||
{%- if '</think>' in message.content %}
|
||||
{%- set content = message.content.split('</think>')[-1].lstrip('\n') %}
|
||||
{%- set reasoning_content = message.content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{%- if loop.index0 > ns.last_query_index %}
|
||||
{%- if loop.last or (not loop.last and reasoning_content) %}
|
||||
{{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
|
||||
{%- else %}
|
||||
{{- '<|im_start|>' + message.role + '\n' + content }}
|
||||
{%- endif %}
|
||||
{%- else %}
|
||||
{{- '<|im_start|>' + message.role + '\n' + content }}
|
||||
{%- endif %}
|
||||
{%- if message.tool_calls %}
|
||||
{%- for tool_call in message.tool_calls %}
|
||||
{%- if (loop.first and content) or (not loop.first) %}
|
||||
{{- '\n' }}
|
||||
{%- endif %}
|
||||
{%- if tool_call.function %}
|
||||
{%- set tool_call = tool_call.function %}
|
||||
{%- endif %}
|
||||
{{- '<tool_call>\n{"name": "' }}
|
||||
{{- tool_call.name }}
|
||||
{{- '", "arguments": ' }}
|
||||
{%- if tool_call.arguments is string %}
|
||||
{{- tool_call.arguments }}
|
||||
{%- else %}
|
||||
{{- tool_call.arguments | tojson }}
|
||||
{%- endif %}
|
||||
{{- '}\n</tool_call>' }}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{{- '<|im_end|>\n' }}
|
||||
{%- elif message.role == "tool" %}
|
||||
{%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
|
||||
{{- '<|im_start|>user' }}
|
||||
{%- endif %}
|
||||
{{- '\n<tool_response>\n' }}
|
||||
{{- message.content }}
|
||||
{{- '\n</tool_response>' }}
|
||||
{%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
|
||||
{{- '<|im_end|>\n' }}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- if add_generation_prompt %}
|
||||
{{- '<|im_start|>assistant\n' }}
|
||||
{%- if enable_thinking is defined and enable_thinking is false %}
|
||||
{{- '<think>\n\n</think>\n\n' }}
|
||||
{%- else %}
|
||||
{{- '<think>\n\n' }}
|
||||
{%- endif %}
|
||||
{%- endif %}"##;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::data::fixture;
|
||||
|
||||
fn msgs(v: &[(&str, &str)]) -> Vec<Message> {
|
||||
v.iter()
|
||||
.map(|(r, c)| Message {
|
||||
role: r.to_string(),
|
||||
content: c.to_string(),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chatml_render() {
|
||||
let t = fixture::text();
|
||||
let m = msgs(&[
|
||||
("system", "be good"),
|
||||
("user", "hi there"),
|
||||
("assistant", "fine thanks"),
|
||||
]);
|
||||
assert_eq!(
|
||||
render(&t, &m, false).unwrap(),
|
||||
"<|im_start|>system\nbe good<|im_end|>\n<|im_start|>user\nhi there<|im_end|>\n<|im_start|>assistant\nfine thanks<|im_end|>\n"
|
||||
);
|
||||
assert!(render(&t, &m[..2], true)
|
||||
.unwrap()
|
||||
.ends_with("<|im_start|>assistant\n"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn qwen3_render_thinking_and_generation_prompt() {
|
||||
let mut t = fixture::text();
|
||||
t.template = Some(QWEN3.into());
|
||||
let m = msgs(&[
|
||||
("user", "hi there"),
|
||||
("assistant", "<think>\nponder\n</think>\n\nfine thanks"),
|
||||
]);
|
||||
assert_eq!(
|
||||
render(&t, &m, false).unwrap(),
|
||||
"<|im_start|>user\nhi there<|im_end|>\n<|im_start|>assistant\n<think>\nponder\n</think>\n\nfine thanks<|im_end|>\n"
|
||||
);
|
||||
assert_eq!(
|
||||
render(&t, &m[..1], true).unwrap(),
|
||||
"<|im_start|>user\nhi there<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
|
||||
);
|
||||
let m = msgs(&[
|
||||
("user", "hi"),
|
||||
("assistant", "one"),
|
||||
("user", "two"),
|
||||
("assistant", "three"),
|
||||
]);
|
||||
assert_eq!(
|
||||
render(&t, &m, false).unwrap(),
|
||||
"<|im_start|>user\nhi<|im_end|>\n<|im_start|>assistant\none<|im_end|>\n<|im_start|>user\ntwo<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\nthree<|im_end|>\n"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn llama3_render_uses_bos() {
|
||||
let mut t = fixture::text();
|
||||
t.template = Some(LLAMA3.into());
|
||||
t.bos_id = t.tokenizer.token_to_id("<|endoftext|>");
|
||||
let m = msgs(&[("user", " hi "), ("assistant", "fine")]);
|
||||
assert_eq!(
|
||||
render(&t, &m, true).unwrap(),
|
||||
"<|endoftext|><|start_header_id|>user<|end_header_id|>\n\nhi<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\nfine<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn raise_exception_is_an_error() {
|
||||
let mut t = fixture::text();
|
||||
t.template = Some("{{ raise_exception('nope') }}".into());
|
||||
assert!(render(&t, &[], false)
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("nope"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn template_resolution() {
|
||||
let tc = serde_json::json!({ "chat_template": "own" });
|
||||
assert_eq!(template(None, &tc).as_deref(), Some("own"));
|
||||
assert_eq!(
|
||||
template(Some("tokenizer_default"), &tc).as_deref(),
|
||||
Some("own")
|
||||
);
|
||||
assert_eq!(template(Some("chatml"), &tc).as_deref(), Some(CHATML));
|
||||
assert_eq!(template(Some("{{ x }}"), &tc).as_deref(), Some("{{ x }}"));
|
||||
assert_eq!(
|
||||
template(Some("tokenizer_default_fallback_chatml"), &tc).as_deref(),
|
||||
Some("own")
|
||||
);
|
||||
let list = serde_json::json!({ "chat_template": [{ "name": "rag", "template": "r" }, { "name": "default", "template": "d" }] });
|
||||
assert_eq!(template(None, &list).as_deref(), Some("d"));
|
||||
assert_eq!(template(None, &Value::Null), None);
|
||||
assert_eq!(
|
||||
template(Some("tokenizer_default_fallback_chatml"), &Value::Null).as_deref(),
|
||||
Some(CHATML)
|
||||
);
|
||||
assert_eq!(token(&serde_json::json!({ "content": "<s>" })), Some("<s>"));
|
||||
assert_eq!(token(&serde_json::json!("<s>")), Some("<s>"));
|
||||
assert_eq!(token(&Value::Null), None);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,105 @@
|
||||
//! GRPO settings, derived from `cfg.trl` and the shared optimizer keys.
|
||||
|
||||
use crate::Config;
|
||||
|
||||
/// How the next token is chosen during rollout.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum Sampling {
|
||||
/// Arg-max. Every completion in a group is identical, so advantages are zero;
|
||||
/// useful only for debugging.
|
||||
Greedy,
|
||||
/// Sample from the softmax of `logits / temperature`.
|
||||
#[default]
|
||||
Temperature,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct GrpoConfig {
|
||||
/// Completions sampled per prompt (`trl.num_generations`). Must be >= 2.
|
||||
pub group_size: usize,
|
||||
/// Maximum tokens generated per completion.
|
||||
pub max_completion_len: usize,
|
||||
pub sampling: Sampling,
|
||||
/// Only used when `sampling == Temperature`.
|
||||
pub temperature: f64,
|
||||
/// Divide the group-centred reward by `std + advantage_eps`.
|
||||
pub scale_rewards: bool,
|
||||
pub advantage_eps: f64,
|
||||
/// KL coefficient (`trl.beta`). The reference log-probabilities are the
|
||||
/// policy's own, computed at the start of the step and detached; the term
|
||||
/// is a penalty against the pre-update policy, not KL to the frozen base
|
||||
/// model, because a [`crate::CausalLm`] cannot run without its adapter.
|
||||
/// `None` disables the term.
|
||||
pub kl_beta: Option<f64>,
|
||||
pub learning_rate: f64,
|
||||
pub weight_decay: f64,
|
||||
/// Global gradient-norm clip; `<= 0` disables clipping.
|
||||
pub max_grad_norm: f64,
|
||||
pub seed: u64,
|
||||
}
|
||||
|
||||
impl Default for GrpoConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
group_size: 8,
|
||||
max_completion_len: 64,
|
||||
sampling: Sampling::Temperature,
|
||||
temperature: 1.0,
|
||||
scale_rewards: true,
|
||||
advantage_eps: 1e-4,
|
||||
kl_beta: None,
|
||||
learning_rate: 1e-3,
|
||||
weight_decay: 0.0,
|
||||
max_grad_norm: 1.0,
|
||||
seed: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&Config> for GrpoConfig {
|
||||
fn from(cfg: &Config) -> Self {
|
||||
let t = &cfg.trl;
|
||||
Self {
|
||||
group_size: t.num_generations,
|
||||
max_completion_len: t.max_completion_length,
|
||||
sampling: Sampling::Temperature,
|
||||
temperature: t.temperature,
|
||||
scale_rewards: t.scale_rewards,
|
||||
advantage_eps: 1e-4,
|
||||
kl_beta: (t.beta > 0.0).then_some(t.beta),
|
||||
learning_rate: cfg.learning_rate,
|
||||
weight_decay: cfg.weight_decay,
|
||||
max_grad_norm: cfg.max_grad_norm,
|
||||
seed: cfg.seed,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl GrpoConfig {
|
||||
pub fn validate(&self) -> anyhow::Result<()> {
|
||||
anyhow::ensure!(
|
||||
self.group_size >= 2,
|
||||
"group_size must be >= 2 (got {}); group-relative advantages need two completions",
|
||||
self.group_size
|
||||
);
|
||||
anyhow::ensure!(
|
||||
self.max_completion_len > 0,
|
||||
"max_completion_len must be >= 1"
|
||||
);
|
||||
if self.sampling == Sampling::Temperature {
|
||||
anyhow::ensure!(
|
||||
self.temperature > 0.0,
|
||||
"temperature must be > 0 (got {})",
|
||||
self.temperature
|
||||
);
|
||||
}
|
||||
if let Some(beta) = self.kl_beta {
|
||||
anyhow::ensure!(beta >= 0.0, "kl_beta must be >= 0 (got {beta})");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn kl_active(&self) -> bool {
|
||||
matches!(self.kl_beta, Some(b) if b > 0.0)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
//! Group-relative advantages, the policy-gradient loss and the KL penalty.
|
||||
//! Operates on reward scalars and log-probability tensors from any [`Policy`](super::Policy).
|
||||
//!
|
||||
//! Advantages follow TRL's GRPO trainer:
|
||||
//!
|
||||
//! ```text
|
||||
//! advantages = rewards - mean_group(rewards)
|
||||
//! if scale_rewards: advantages /= std_group(rewards) + 1e-4
|
||||
//! ```
|
||||
//!
|
||||
//! and the loss is the single-iteration, unclipped form:
|
||||
//!
|
||||
//! ```text
|
||||
//! loss = -mean_i(A_i * logprob_i / len_i) (+ beta * KL)
|
||||
//! ```
|
||||
|
||||
use hanzo_ml::{Result, Tensor};
|
||||
|
||||
/// Advantages for one group of `rewards` sharing a prompt: `r_i - mean(r)`,
|
||||
/// divided by `std(r) + eps` when `scale`. The std is unbiased (`n - 1`),
|
||||
/// matching `torch.std`.
|
||||
pub fn group_advantages(rewards: &[f32], scale: bool, eps: f64) -> Vec<f32> {
|
||||
let n = rewards.len();
|
||||
debug_assert!(n >= 2, "group_advantages needs a group of >= 2");
|
||||
let mean = rewards.iter().sum::<f32>() / n as f32;
|
||||
if !scale {
|
||||
return rewards.iter().map(|&r| r - mean).collect();
|
||||
}
|
||||
let var = if n > 1 {
|
||||
rewards.iter().map(|&r| (r - mean).powi(2)).sum::<f32>() / (n as f32 - 1.0)
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
let denom = var.sqrt() + eps as f32;
|
||||
rewards.iter().map(|&r| (r - mean) / denom).collect()
|
||||
}
|
||||
|
||||
/// [`group_advantages`] over consecutive groups of `group_size` in `rewards`.
|
||||
pub fn batch_advantages(rewards: &[f32], group_size: usize, scale: bool, eps: f64) -> Vec<f32> {
|
||||
debug_assert!(group_size >= 2);
|
||||
debug_assert_eq!(rewards.len() % group_size, 0);
|
||||
rewards
|
||||
.chunks(group_size)
|
||||
.flat_map(|g| group_advantages(g, scale, eps))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// `-mean_i(A_i * logprob_i / len_i)`. `advantages` are detached constants;
|
||||
/// `logprobs[i]` is the differentiable summed log-probability of completion `i`.
|
||||
pub fn policy_gradient_loss(
|
||||
advantages: &[f32],
|
||||
logprobs: &[Tensor],
|
||||
lengths: &[usize],
|
||||
) -> Result<Tensor> {
|
||||
assert_eq!(advantages.len(), logprobs.len());
|
||||
assert_eq!(advantages.len(), lengths.len());
|
||||
assert!(
|
||||
!advantages.is_empty(),
|
||||
"empty batch in policy_gradient_loss"
|
||||
);
|
||||
let mut total: Option<Tensor> = None;
|
||||
for ((&adv, lp), &len) in advantages.iter().zip(logprobs).zip(lengths) {
|
||||
let weighted = lp.affine(adv as f64 / len.max(1) as f64, 0.0)?;
|
||||
total = Some(match total {
|
||||
Some(acc) => (acc + weighted)?,
|
||||
None => weighted,
|
||||
});
|
||||
}
|
||||
total.unwrap().affine(-1.0 / advantages.len() as f64, 0.0)
|
||||
}
|
||||
|
||||
/// `beta * mean_i k3_i` with `k3 = exp(d) - d - 1 >= 0` for
|
||||
/// `d = ref_i/len_i - logprob_i/len_i`, the low-variance estimator of
|
||||
/// `KL(policy || ref)` used by GRPO. `ref_logprobs` are detached constants.
|
||||
pub fn kl_penalty(
|
||||
beta: f64,
|
||||
logprobs: &[Tensor],
|
||||
ref_logprobs: &[f32],
|
||||
lengths: &[usize],
|
||||
) -> Result<Tensor> {
|
||||
assert_eq!(logprobs.len(), ref_logprobs.len());
|
||||
assert_eq!(logprobs.len(), lengths.len());
|
||||
assert!(!logprobs.is_empty());
|
||||
let mut total: Option<Tensor> = None;
|
||||
for ((lp, &ref_lp), &len) in logprobs.iter().zip(ref_logprobs).zip(lengths) {
|
||||
let len = len.max(1) as f64;
|
||||
let delta = lp.affine(-1.0 / len, ref_lp as f64 / len)?;
|
||||
let k3 = ((delta.exp()? - &delta)? - 1.0)?;
|
||||
total = Some(match total {
|
||||
Some(acc) => (acc + k3)?,
|
||||
None => k3,
|
||||
});
|
||||
}
|
||||
total.unwrap().affine(beta / logprobs.len() as f64, 0.0)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use hanzo_ml::Device;
|
||||
|
||||
#[test]
|
||||
fn advantages_are_zero_mean_within_group() {
|
||||
let a = group_advantages(&[1.0, 2.0, 3.0, 4.0], false, 1e-4);
|
||||
let sum: f32 = a.iter().sum();
|
||||
assert!(sum.abs() < 1e-5, "advantages should sum to ~0, got {sum}");
|
||||
assert!(a[0] < a[1] && a[1] < a[2] && a[2] < a[3]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scaled_advantages_have_unit_scale() {
|
||||
let a = group_advantages(&[0.0, 10.0], true, 1e-4);
|
||||
assert!(
|
||||
(a[1] - std::f32::consts::FRAC_1_SQRT_2).abs() < 1e-2,
|
||||
"got {}",
|
||||
a[1]
|
||||
);
|
||||
assert!(
|
||||
(a[0] + std::f32::consts::FRAC_1_SQRT_2).abs() < 1e-2,
|
||||
"got {}",
|
||||
a[0]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn degenerate_group_gives_zero_advantage() {
|
||||
for v in group_advantages(&[5.0, 5.0, 5.0], true, 1e-4) {
|
||||
assert_eq!(v, 0.0);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_splits_groups() {
|
||||
let a = batch_advantages(&[0.0, 2.0, 5.0, 5.0], 2, false, 1e-4);
|
||||
assert_eq!(a, vec![-1.0, 1.0, 0.0, 0.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pg_loss_sign_and_finiteness() -> Result<()> {
|
||||
let dev = Device::Cpu;
|
||||
let lp0 = Tensor::new(-2.0f32, &dev)?;
|
||||
let lp1 = Tensor::new(-1.0f32, &dev)?;
|
||||
let v = policy_gradient_loss(&[-1.0, 1.0], &[lp0, lp1], &[1, 1])?.to_scalar::<f32>()?;
|
||||
assert!(v.is_finite());
|
||||
// -mean((-1)(-2) + (1)(-1)) = -1/2
|
||||
assert!((v + 0.5).abs() < 1e-5, "got {v}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kl_is_nonnegative_and_zero_at_match() -> Result<()> {
|
||||
let dev = Device::Cpu;
|
||||
let p = Tensor::new(-3.0f32, &dev)?;
|
||||
let v = kl_penalty(0.1, &[p], &[-3.0], &[1])?.to_scalar::<f32>()?;
|
||||
assert!(v.abs() < 1e-5, "KL at match should be 0, got {v}");
|
||||
let p2 = Tensor::new(-1.0f32, &dev)?;
|
||||
assert!(kl_penalty(0.1, &[p2], &[-3.0], &[1])?.to_scalar::<f32>()? > 0.0);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
//! What one GRPO step reports.
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Metrics {
|
||||
/// Policy-gradient loss plus KL penalty; the value optimized.
|
||||
pub loss: f32,
|
||||
pub pg_loss: f32,
|
||||
/// 0.0 when the KL term is disabled.
|
||||
pub kl_loss: f32,
|
||||
/// Mean of the group reward means.
|
||||
pub mean_reward: f32,
|
||||
/// Mean of the within-group reward stds.
|
||||
pub reward_std: f32,
|
||||
pub mean_completion_len: f32,
|
||||
/// Global gradient norm before clipping.
|
||||
pub grad_norm: f32,
|
||||
pub num_groups: usize,
|
||||
pub num_completions: usize,
|
||||
}
|
||||
|
||||
impl Metrics {
|
||||
pub fn is_finite(&self) -> bool {
|
||||
self.loss.is_finite() && self.pg_loss.is_finite() && self.kl_loss.is_finite()
|
||||
}
|
||||
}
|
||||
+81
-1
@@ -1 +1,81 @@
|
||||
//! grpo
|
||||
//! GRPO (group relative policy optimization) over a [`CausalLm`].
|
||||
//!
|
||||
//! One step: for every prompt, sample a group of `num_generations` completions
|
||||
//! from the current policy, score each with the verifiers in `trl.reward_funcs`,
|
||||
//! centre the rewards within the group (`math::group_advantages`), form
|
||||
//! `-mean(A_i * logprob_i / len_i)` (`math::policy_gradient_loss`), and take a
|
||||
//! norm-clipped AdamW step on `model.trainable_vars()`.
|
||||
//!
|
||||
//! When `trl.beta > 0` a KL term is added. A [`CausalLm`] cannot run without its
|
||||
//! adapter, so there is no frozen reference model: the reference
|
||||
//! log-probabilities are the policy's own, computed at the start of the step and
|
||||
//! detached. The term therefore penalises drift from the pre-update policy, not
|
||||
//! from the base model. With a single optimizer update per step the term and its
|
||||
//! gradient are exactly zero at the point they are evaluated.
|
||||
|
||||
pub mod config;
|
||||
pub mod math;
|
||||
pub mod metrics;
|
||||
pub mod policy;
|
||||
pub mod reward;
|
||||
pub mod sampler;
|
||||
pub mod trainer;
|
||||
|
||||
pub use config::{GrpoConfig, Sampling};
|
||||
pub use metrics::Metrics;
|
||||
pub use policy::{LmPolicy, Policy};
|
||||
pub use reward::{parse_rewards, Reward};
|
||||
pub use sampler::Sampler;
|
||||
pub use trainer::Trainer;
|
||||
|
||||
use crate::{CausalLm, Config, Prompt};
|
||||
use rand::seq::SliceRandom;
|
||||
use rand::{rngs::StdRng, SeedableRng};
|
||||
use std::path::Path;
|
||||
|
||||
/// Train `model` on `prompts` as described by `cfg`, writing the adapter to
|
||||
/// `cfg.output_dir` (and `checkpoint-{step}` under it every `save_steps`).
|
||||
pub fn run(
|
||||
cfg: &Config,
|
||||
model: &dyn CausalLm,
|
||||
text: &crate::data::Text,
|
||||
mut prompts: Vec<Prompt>,
|
||||
) -> anyhow::Result<()> {
|
||||
anyhow::ensure!(!prompts.is_empty(), "grpo: no prompts");
|
||||
let policy = LmPolicy {
|
||||
model,
|
||||
tokenizer: &text.tokenizer,
|
||||
eos: Some(text.eos_id),
|
||||
};
|
||||
let mut trainer = Trainer::new(GrpoConfig::from(cfg), policy, parse_rewards(cfg)?)?;
|
||||
let out = Path::new(&cfg.output_dir);
|
||||
let max_steps = cfg.max_steps.unwrap_or(usize::MAX);
|
||||
let log_every = cfg.logging_steps.max(1);
|
||||
let mut rng = StdRng::seed_from_u64(cfg.seed);
|
||||
let mut step = 0;
|
||||
'epochs: for _ in 0..cfg.num_epochs {
|
||||
prompts.shuffle(&mut rng);
|
||||
for batch in prompts.chunks(cfg.micro_batch_size) {
|
||||
if step >= max_steps {
|
||||
break 'epochs;
|
||||
}
|
||||
let m = trainer.step(batch)?;
|
||||
step += 1;
|
||||
if step % log_every == 0 {
|
||||
tracing::info!(
|
||||
step,
|
||||
reward = m.mean_reward,
|
||||
reward_std = m.reward_std,
|
||||
loss = m.loss,
|
||||
kl = m.kl_loss,
|
||||
completion_len = m.mean_completion_len,
|
||||
grad_norm = m.grad_norm,
|
||||
);
|
||||
}
|
||||
if cfg.save_steps.is_some_and(|s| step % s == 0) {
|
||||
model.save_adapter(&out.join(format!("checkpoint-{step}")))?;
|
||||
}
|
||||
}
|
||||
}
|
||||
model.save_adapter(out)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,355 @@
|
||||
//! What GRPO needs from a model: rollout, differentiable scoring, decoding.
|
||||
|
||||
use crate::grpo::sampler::Sampler;
|
||||
use crate::CausalLm;
|
||||
use hanzo_ml::{DType, Tensor, Var, D};
|
||||
use hanzo_nn::ops::log_softmax;
|
||||
use tokenizers::Tokenizer;
|
||||
|
||||
pub trait Policy {
|
||||
fn trainable_vars(&self) -> Vec<Var>;
|
||||
|
||||
/// `group_size` completions of at most `max_len` tokens for `prompt`, drawn
|
||||
/// with `sampler`. No autograd graph is kept.
|
||||
fn sample_group(
|
||||
&self,
|
||||
prompt: &[u32],
|
||||
group_size: usize,
|
||||
max_len: usize,
|
||||
sampler: &mut Sampler,
|
||||
) -> anyhow::Result<Vec<Vec<u32>>>;
|
||||
|
||||
/// `sum_t log p(c_t | prompt, c_<t)` as a scalar tensor differentiable
|
||||
/// w.r.t. [`Policy::trainable_vars`].
|
||||
fn sequence_logprob(&self, prompt: &[u32], completion: &[u32]) -> anyhow::Result<Tensor>;
|
||||
|
||||
/// [`Policy::sequence_logprob`] for every completion of one group.
|
||||
fn group_logprob(
|
||||
&self,
|
||||
prompt: &[u32],
|
||||
completions: &[Vec<u32>],
|
||||
) -> anyhow::Result<Vec<Tensor>> {
|
||||
completions
|
||||
.iter()
|
||||
.map(|c| self.sequence_logprob(prompt, c))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Text for `tokens`, used by text rewards. Empty when the policy has no tokenizer.
|
||||
fn decode(&self, _tokens: &[u32]) -> anyhow::Result<String> {
|
||||
Ok(String::new())
|
||||
}
|
||||
}
|
||||
|
||||
/// A [`CausalLm`] as a GRPO policy.
|
||||
///
|
||||
/// Rollout has no KV cache: every generated token re-runs `model.forward` over
|
||||
/// the whole group's growing sequences, so sampling `n` tokens costs `n`
|
||||
/// forwards of lengths `p+1 ..= p+n`, quadratic in `n`.
|
||||
pub struct LmPolicy<'a> {
|
||||
pub model: &'a dyn CausalLm,
|
||||
pub tokenizer: &'a Tokenizer,
|
||||
pub eos: Option<u32>,
|
||||
}
|
||||
|
||||
impl Policy for LmPolicy<'_> {
|
||||
fn trainable_vars(&self) -> Vec<Var> {
|
||||
self.model.trainable_vars()
|
||||
}
|
||||
|
||||
fn sample_group(
|
||||
&self,
|
||||
prompt: &[u32],
|
||||
group_size: usize,
|
||||
max_len: usize,
|
||||
sampler: &mut Sampler,
|
||||
) -> anyhow::Result<Vec<Vec<u32>>> {
|
||||
let dev = self.model.device();
|
||||
let mut live = vec![prompt.to_vec(); group_size];
|
||||
let mut done = Vec::with_capacity(group_size);
|
||||
for _ in 0..max_len {
|
||||
if live.is_empty() {
|
||||
break;
|
||||
}
|
||||
let t = live[0].len();
|
||||
let ids = Tensor::from_vec(live.concat(), (live.len(), t), dev)?;
|
||||
let last = self
|
||||
.model
|
||||
.forward(&ids, None)?
|
||||
.narrow(1, t - 1, 1)?
|
||||
.squeeze(1)?;
|
||||
let mut next = Vec::with_capacity(live.len());
|
||||
for i in 0..live.len() {
|
||||
next.push(sampler.sample(&last.get(i)?)?);
|
||||
}
|
||||
let mut still = Vec::with_capacity(live.len());
|
||||
for (mut seq, tok) in live.into_iter().zip(next) {
|
||||
seq.push(tok);
|
||||
if Some(tok) == self.eos {
|
||||
done.push(seq)
|
||||
} else {
|
||||
still.push(seq)
|
||||
}
|
||||
}
|
||||
live = still;
|
||||
}
|
||||
done.extend(live);
|
||||
Ok(done
|
||||
.into_iter()
|
||||
.map(|s| s[prompt.len()..].to_vec())
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn sequence_logprob(&self, prompt: &[u32], completion: &[u32]) -> anyhow::Result<Tensor> {
|
||||
Ok(self
|
||||
.group_logprob(prompt, &[completion.to_vec()])?
|
||||
.remove(0))
|
||||
}
|
||||
|
||||
/// One forward over the group, completions padded right with an attention mask.
|
||||
fn group_logprob(
|
||||
&self,
|
||||
prompt: &[u32],
|
||||
completions: &[Vec<u32>],
|
||||
) -> anyhow::Result<Vec<Tensor>> {
|
||||
let dev = self.model.device();
|
||||
let p = prompt.len();
|
||||
anyhow::ensure!(p > 0, "grpo: a prompt needs at least one token");
|
||||
let g = completions.len();
|
||||
let t = p + completions.iter().map(Vec::len).max().unwrap_or(0);
|
||||
let pad = self.eos.unwrap_or(0);
|
||||
let mut ids = Vec::with_capacity(g * t);
|
||||
let mut mask = Vec::with_capacity(g * t);
|
||||
for c in completions {
|
||||
ids.extend_from_slice(prompt);
|
||||
ids.extend_from_slice(c);
|
||||
ids.resize(ids.len() + t - p - c.len(), pad);
|
||||
mask.extend(std::iter::repeat_n(1u8, p + c.len()));
|
||||
mask.extend(std::iter::repeat_n(0u8, t - p - c.len()));
|
||||
}
|
||||
let padded = mask.contains(&0);
|
||||
let ids = Tensor::from_vec(ids, (g, t), dev)?;
|
||||
let mask = if padded {
|
||||
Some(Tensor::from_vec(mask, (g, t), dev)?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let logits = self
|
||||
.model
|
||||
.forward(&ids, mask.as_ref())?
|
||||
.to_dtype(DType::F32)?;
|
||||
let logp = log_softmax(&logits, D::Minus1)?;
|
||||
completions
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, c)| {
|
||||
if c.is_empty() {
|
||||
return Ok(Tensor::zeros((), DType::F32, dev)?);
|
||||
}
|
||||
let rows = logp.get(i)?.narrow(0, p - 1, c.len())?;
|
||||
let idx = Tensor::from_slice(c, (c.len(), 1), dev)?;
|
||||
Ok(rows.gather(&idx, 1)?.sum_all()?)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn decode(&self, tokens: &[u32]) -> anyhow::Result<String> {
|
||||
self.tokenizer
|
||||
.decode(tokens, true)
|
||||
.map_err(|e| anyhow::anyhow!("decode: {e}"))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) mod fake {
|
||||
//! A bigram table as a [`CausalLm`]: `logits[b, t, :] = table[input_ids[b, t], :]`.
|
||||
use super::*;
|
||||
use hanzo_ml::Device;
|
||||
|
||||
pub struct Bigram {
|
||||
pub table: Var,
|
||||
pub eos: Option<u32>,
|
||||
}
|
||||
|
||||
impl Bigram {
|
||||
pub fn new(rows: &[Vec<f32>], eos: Option<u32>) -> Self {
|
||||
let v = rows.len();
|
||||
let flat: Vec<f32> = rows.concat();
|
||||
Self {
|
||||
table: Var::from_vec(flat, (v, v), &Device::Cpu).unwrap(),
|
||||
eos,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn uniform(vocab: usize, eos: Option<u32>) -> Self {
|
||||
Self {
|
||||
table: Var::zeros((vocab, vocab), DType::F32, &Device::Cpu).unwrap(),
|
||||
eos,
|
||||
}
|
||||
}
|
||||
|
||||
/// Softmax row `from` of the table.
|
||||
pub fn probs(&self, from: u32) -> Vec<f32> {
|
||||
let row = self.table.get(from as usize).unwrap();
|
||||
hanzo_nn::ops::softmax(&row, 0).unwrap().to_vec1().unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
impl CausalLm for Bigram {
|
||||
fn forward(&self, input_ids: &Tensor, _mask: Option<&Tensor>) -> hanzo_ml::Result<Tensor> {
|
||||
let (b, t) = input_ids.dims2()?;
|
||||
let v = self.vocab_size();
|
||||
self.table
|
||||
.index_select(&input_ids.flatten_all()?, 0)?
|
||||
.reshape((b, t, v))
|
||||
}
|
||||
fn trainable_vars(&self) -> Vec<Var> {
|
||||
vec![self.table.clone()]
|
||||
}
|
||||
fn device(&self) -> &Device {
|
||||
self.table.device()
|
||||
}
|
||||
fn dtype(&self) -> DType {
|
||||
DType::F32
|
||||
}
|
||||
fn eos_token_id(&self) -> Option<u32> {
|
||||
self.eos
|
||||
}
|
||||
fn vocab_size(&self) -> usize {
|
||||
self.table.dim(0).unwrap()
|
||||
}
|
||||
fn save_adapter(&self, _dir: &std::path::Path) -> anyhow::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Word-level tokenizer over `words`; id = index, `<eos>` is special.
|
||||
pub fn tokenizer(words: &[&str]) -> Tokenizer {
|
||||
let vocab: Vec<String> = words
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, w)| format!("{w:?}: {i}"))
|
||||
.collect();
|
||||
let eos = words
|
||||
.iter()
|
||||
.position(|&w| w == "<eos>")
|
||||
.expect("<eos> in vocab");
|
||||
let json = format!(
|
||||
r#"{{"version":"1.0","added_tokens":[{{"id":{eos},"content":"<eos>","special":true,"single_word":false,"lstrip":false,"rstrip":false,"normalized":false}}],
|
||||
"pre_tokenizer":{{"type":"Whitespace"}},
|
||||
"model":{{"type":"WordLevel","vocab":{{{}}},"unk_token":"{}"}}}}"#,
|
||||
vocab.join(","),
|
||||
words[0]
|
||||
);
|
||||
json.parse().expect("tokenizer json")
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::fake::{tokenizer, Bigram};
|
||||
use super::*;
|
||||
use crate::grpo::config::{GrpoConfig, Sampling};
|
||||
|
||||
const BIG: f32 = 30.0;
|
||||
|
||||
#[test]
|
||||
fn sequence_logprob_matches_hand_computation() -> anyhow::Result<()> {
|
||||
let rows = vec![
|
||||
vec![1.0, 2.0, 0.5],
|
||||
vec![0.0, -1.0, 3.0],
|
||||
vec![0.0, 0.0, 0.0],
|
||||
];
|
||||
let model = Bigram::new(&rows, None);
|
||||
let tk = tokenizer(&["a", "b", "<eos>"]);
|
||||
let policy = LmPolicy {
|
||||
model: &model,
|
||||
tokenizer: &tk,
|
||||
eos: None,
|
||||
};
|
||||
let got = policy.sequence_logprob(&[0], &[1, 2])?.to_scalar::<f32>()?;
|
||||
let lsm = |r: &[f32], i: usize| r[i] - r.iter().map(|x| x.exp()).sum::<f32>().ln();
|
||||
let want = lsm(&rows[0], 1) + lsm(&rows[1], 2);
|
||||
assert!((got - want).abs() < 1e-5, "got {got} want {want}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batched_group_equals_single() -> anyhow::Result<()> {
|
||||
let rows = vec![
|
||||
vec![1.0, 2.0, 0.5, 0.1],
|
||||
vec![0.0, -1.0, 3.0, 0.2],
|
||||
vec![2.0, 0.0, 0.0, 1.0],
|
||||
vec![0.0; 4],
|
||||
];
|
||||
let model = Bigram::new(&rows, Some(3));
|
||||
let tk = tokenizer(&["a", "b", "c", "<eos>"]);
|
||||
let policy = LmPolicy {
|
||||
model: &model,
|
||||
tokenizer: &tk,
|
||||
eos: Some(3),
|
||||
};
|
||||
let group = vec![vec![1, 2, 3], vec![2], vec![0, 0, 1, 3]];
|
||||
let batched = policy.group_logprob(&[0, 1], &group)?;
|
||||
for (c, b) in group.iter().zip(&batched) {
|
||||
let single = policy.sequence_logprob(&[0, 1], c)?.to_scalar::<f32>()?;
|
||||
assert!((b.to_scalar::<f32>()? - single).abs() < 1e-5);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn logprob_is_differentiable() -> anyhow::Result<()> {
|
||||
let model = Bigram::uniform(3, None);
|
||||
let tk = tokenizer(&["a", "b", "<eos>"]);
|
||||
let policy = LmPolicy {
|
||||
model: &model,
|
||||
tokenizer: &tk,
|
||||
eos: None,
|
||||
};
|
||||
let lp = policy.sequence_logprob(&[0], &[1])?;
|
||||
let grads = lp.backward()?;
|
||||
let g = grads
|
||||
.get(&model.table)
|
||||
.expect("gradient on table")
|
||||
.to_vec2::<f32>()?;
|
||||
// d/dlogit log softmax(row 0)[1] = onehot(1) - softmax = [-1/3, 2/3, -1/3]
|
||||
assert!((g[0][1] - 2.0 / 3.0).abs() < 1e-5 && (g[0][0] + 1.0 / 3.0).abs() < 1e-5);
|
||||
assert!(g[1].iter().all(|&x| x == 0.0));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sample_group_stops_at_eos_and_max_len() -> anyhow::Result<()> {
|
||||
// 0 -> 1 -> eos(2), near-deterministically.
|
||||
let rows = vec![vec![0.0, BIG, 0.0], vec![0.0, 0.0, BIG], vec![0.0; 3]];
|
||||
let model = Bigram::new(&rows, Some(2));
|
||||
let tk = tokenizer(&["a", "b", "<eos>"]);
|
||||
let cfg = GrpoConfig {
|
||||
sampling: Sampling::Temperature,
|
||||
temperature: 1.0,
|
||||
..GrpoConfig::default()
|
||||
};
|
||||
let mut sampler = Sampler::new(&cfg, 0);
|
||||
|
||||
let policy = LmPolicy {
|
||||
model: &model,
|
||||
tokenizer: &tk,
|
||||
eos: Some(2),
|
||||
};
|
||||
let group = policy.sample_group(&[0], 4, 5, &mut sampler)?;
|
||||
assert_eq!(group.len(), 4);
|
||||
assert!(group.iter().all(|c| c == &[1, 2]), "{group:?}");
|
||||
|
||||
let none = LmPolicy {
|
||||
model: &model,
|
||||
tokenizer: &tk,
|
||||
eos: None,
|
||||
};
|
||||
let group = none.sample_group(&[0], 3, 5, &mut sampler)?;
|
||||
assert_eq!(group.len(), 3);
|
||||
assert!(group.iter().all(|c| c.len() == 5), "{group:?}");
|
||||
assert_eq!(policy.decode(&[1, 2])?, "b");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,254 @@
|
||||
//! Verifiers named in `trl.reward_funcs`, weighted by `trl.reward_weights`.
|
||||
//!
|
||||
//! - `exact_match`: completion equals `Prompt.answer` after whitespace and case normalisation.
|
||||
//! - `numeric_match`: the last number in the completion (inside `\boxed{…}` when
|
||||
//! present) equals the answer's number within 1e-6.
|
||||
//! - `format:<regex>`: 1.0 when the regex matches the completion.
|
||||
//! - `length:<max_tokens>`: 0 within the budget, -1 over it.
|
||||
//! - `http:<url>`: POST `{"prompt","completion","answer"}`, read `{"reward": number}`.
|
||||
|
||||
use crate::{Config, Prompt};
|
||||
use anyhow::Context;
|
||||
use regex::Regex;
|
||||
use std::sync::LazyLock;
|
||||
|
||||
static NUMBER: LazyLock<Regex> =
|
||||
LazyLock::new(|| Regex::new(r"-?\d+(?:,\d{3})*(?:\.\d+)?").unwrap());
|
||||
static BOXED: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\\boxed\{([^{}]*)\}").unwrap());
|
||||
|
||||
pub struct Reward {
|
||||
pub name: String,
|
||||
pub weight: f32,
|
||||
kind: Kind,
|
||||
}
|
||||
|
||||
enum Kind {
|
||||
Exact,
|
||||
Numeric,
|
||||
Format(Regex),
|
||||
Length(usize),
|
||||
Http(String),
|
||||
/// Occurrences of a token id; a test fixture for policies without a tokenizer.
|
||||
#[cfg(test)]
|
||||
Token(u32),
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct Score {
|
||||
reward: f64,
|
||||
}
|
||||
|
||||
impl Reward {
|
||||
pub fn parse(spec: &str, weight: f32) -> anyhow::Result<Self> {
|
||||
let kind = match spec {
|
||||
"exact_match" => Kind::Exact,
|
||||
"numeric_match" => Kind::Numeric,
|
||||
s if s.starts_with("http://") || s.starts_with("https://") => Kind::Http(s.to_string()),
|
||||
s => match s.split_once(':') {
|
||||
Some(("format", re)) => {
|
||||
Kind::Format(Regex::new(re).with_context(|| format!("reward {spec}"))?)
|
||||
}
|
||||
Some(("length", n)) => {
|
||||
Kind::Length(n.parse().with_context(|| format!("reward {spec}"))?)
|
||||
}
|
||||
Some(("http", url)) => Kind::Http(url.to_string()),
|
||||
_ => anyhow::bail!("unknown reward function {spec:?}"),
|
||||
},
|
||||
};
|
||||
Ok(Self {
|
||||
name: spec.to_string(),
|
||||
weight,
|
||||
kind,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn token(id: u32) -> Self {
|
||||
Self {
|
||||
name: format!("token:{id}"),
|
||||
weight: 1.0,
|
||||
kind: Kind::Token(id),
|
||||
}
|
||||
}
|
||||
|
||||
/// Unweighted score of one completion.
|
||||
pub fn score(&self, prompt: &Prompt, tokens: &[u32], text: &str) -> anyhow::Result<f32> {
|
||||
let answer = prompt.answer.as_deref();
|
||||
Ok(match &self.kind {
|
||||
Kind::Exact => answer.is_some_and(|a| normal(a) == normal(text)) as u8 as f32,
|
||||
Kind::Numeric => match (answer.and_then(number), number(text)) {
|
||||
(Some(a), Some(c)) => ((a - c).abs() < 1e-6) as u8 as f32,
|
||||
_ => 0.0,
|
||||
},
|
||||
Kind::Format(re) => re.is_match(text) as u8 as f32,
|
||||
Kind::Length(max) => {
|
||||
if tokens.len() <= *max {
|
||||
0.0
|
||||
} else {
|
||||
-1.0
|
||||
}
|
||||
}
|
||||
Kind::Http(url) => {
|
||||
let body = serde_json::json!({ "prompt": prompt.text, "completion": text, "answer": answer });
|
||||
let score: Score = ureq::post(url)
|
||||
.send_json(body)
|
||||
.with_context(|| format!("reward {url}"))?
|
||||
.body_mut()
|
||||
.read_json()
|
||||
.with_context(|| format!("reward {url}: expected {{\"reward\": number}}"))?;
|
||||
score.reward as f32
|
||||
}
|
||||
#[cfg(test)]
|
||||
Kind::Token(id) => tokens.iter().filter(|&t| t == id).count() as f32,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn normal(s: &str) -> String {
|
||||
s.split_whitespace()
|
||||
.collect::<Vec<_>>()
|
||||
.join(" ")
|
||||
.to_lowercase()
|
||||
}
|
||||
|
||||
/// The last number in `s`, preferring the contents of the last `\boxed{…}`.
|
||||
fn number(s: &str) -> Option<f64> {
|
||||
let field = BOXED
|
||||
.captures_iter(s)
|
||||
.last()
|
||||
.map(|c| c.get(1).unwrap().as_str())
|
||||
.unwrap_or(s);
|
||||
NUMBER
|
||||
.find_iter(field)
|
||||
.last()?
|
||||
.as_str()
|
||||
.replace(',', "")
|
||||
.parse()
|
||||
.ok()
|
||||
}
|
||||
|
||||
pub fn parse_rewards(cfg: &Config) -> anyhow::Result<Vec<Reward>> {
|
||||
let t = &cfg.trl;
|
||||
t.reward_funcs
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, spec)| {
|
||||
Reward::parse(spec, t.reward_weights.get(i).copied().unwrap_or(1.0) as f32)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::{Read, Write};
|
||||
use std::net::TcpListener;
|
||||
|
||||
fn prompt(answer: Option<&str>) -> Prompt {
|
||||
Prompt {
|
||||
input_ids: vec![1],
|
||||
text: "q?".into(),
|
||||
answer: answer.map(String::from),
|
||||
}
|
||||
}
|
||||
|
||||
fn score(spec: &str, answer: Option<&str>, text: &str) -> f32 {
|
||||
Reward::parse(spec, 1.0)
|
||||
.unwrap()
|
||||
.score(&prompt(answer), &[0; 3], text)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exact_normalises_whitespace_and_case() {
|
||||
assert_eq!(
|
||||
score("exact_match", Some("Paris France"), " paris\nfrance "),
|
||||
1.0
|
||||
);
|
||||
assert_eq!(score("exact_match", Some("Paris"), "London"), 0.0);
|
||||
assert_eq!(score("exact_match", None, "Paris"), 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn numeric_takes_last_number_and_boxed() {
|
||||
assert_eq!(score("numeric_match", Some("42"), "First 7, then 42."), 1.0);
|
||||
assert_eq!(score("numeric_match", Some("42"), "First 42, then 7."), 0.0);
|
||||
assert_eq!(
|
||||
score("numeric_match", Some("42"), r"\boxed{42} and then 7"),
|
||||
1.0
|
||||
);
|
||||
assert_eq!(
|
||||
score("numeric_match", Some("#### 1,234"), r"so \boxed{1234.0}"),
|
||||
1.0
|
||||
);
|
||||
assert_eq!(score("numeric_match", Some("-3.5"), "answer is -3.5"), 1.0);
|
||||
assert_eq!(score("numeric_match", Some("3"), "no digits"), 0.0);
|
||||
assert_eq!(score("numeric_match", None, "3"), 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_and_length() {
|
||||
assert_eq!(
|
||||
score(r"format:^<think>.*</think>", None, "<think>x</think> y"),
|
||||
1.0
|
||||
);
|
||||
assert_eq!(score(r"format:^<think>.*</think>", None, "y"), 0.0);
|
||||
assert_eq!(score("length:3", None, ""), 0.0);
|
||||
assert_eq!(score("length:2", None, ""), -1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_and_malformed_are_errors() {
|
||||
assert!(Reward::parse("bleu", 1.0).is_err());
|
||||
assert!(Reward::parse("length:many", 1.0).is_err());
|
||||
assert!(Reward::parse("format:(", 1.0).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn weights_default_to_one() {
|
||||
let y = "base_model: x\noutput_dir: o\ndatasets: [{path: p, type: completion}]\nrl: grpo\n\
|
||||
trl: {reward_funcs: [exact_match, 'length:8'], reward_weights: [2.0]}\n";
|
||||
let cfg: Config = serde_yaml::from_str(y).unwrap();
|
||||
let r = parse_rewards(&cfg).unwrap();
|
||||
assert_eq!((r[0].weight, r[1].weight), (2.0, 1.0));
|
||||
assert_eq!(r[1].name, "length:8");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_posts_json_and_reads_reward() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let url = format!("http://{}/score", listener.local_addr().unwrap());
|
||||
let server = std::thread::spawn(move || {
|
||||
let (mut s, _) = listener.accept().unwrap();
|
||||
let mut buf = vec![0u8; 4096];
|
||||
let mut n = 0;
|
||||
while !String::from_utf8_lossy(&buf[..n]).contains("\r\n\r\n") {
|
||||
n += s.read(&mut buf[n..]).unwrap();
|
||||
}
|
||||
let head = String::from_utf8_lossy(&buf[..n]).to_string();
|
||||
let len: usize = head
|
||||
.lines()
|
||||
.find_map(|l| {
|
||||
l.to_ascii_lowercase()
|
||||
.strip_prefix("content-length:")
|
||||
.map(|v| v.trim().parse().unwrap())
|
||||
})
|
||||
.unwrap();
|
||||
let start = head.find("\r\n\r\n").unwrap() + 4;
|
||||
while n < start + len {
|
||||
n += s.read(&mut buf[n..]).unwrap();
|
||||
}
|
||||
let body: serde_json::Value = serde_json::from_slice(&buf[start..start + len]).unwrap();
|
||||
let reply = r#"{"reward":0.75}"#;
|
||||
write!(s, "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{reply}", reply.len()).unwrap();
|
||||
body
|
||||
});
|
||||
let r = Reward::parse(&format!("http:{url}"), 1.0).unwrap();
|
||||
assert_eq!(r.score(&prompt(Some("4")), &[], "2+2=4").unwrap(), 0.75);
|
||||
let body = server.join().unwrap();
|
||||
assert_eq!(
|
||||
body,
|
||||
serde_json::json!({"prompt": "q?", "completion": "2+2=4", "answer": "4"})
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
//! Next-token sampling during rollout: arg-max or temperature sampling from a
|
||||
//! seeded generator, so a step replays exactly for a given seed.
|
||||
|
||||
use crate::grpo::config::{GrpoConfig, Sampling};
|
||||
use hanzo_ml::{DType, Result, Tensor};
|
||||
use rand::{rngs::StdRng, Rng, SeedableRng};
|
||||
|
||||
pub struct Sampler {
|
||||
rng: StdRng,
|
||||
mode: Sampling,
|
||||
temperature: f64,
|
||||
}
|
||||
|
||||
impl Sampler {
|
||||
pub fn new(cfg: &GrpoConfig, seed: u64) -> Self {
|
||||
Self {
|
||||
rng: StdRng::seed_from_u64(seed),
|
||||
mode: cfg.sampling,
|
||||
temperature: cfg.temperature,
|
||||
}
|
||||
}
|
||||
|
||||
/// Next token id for `logits` of shape `[vocab]`.
|
||||
pub fn sample(&mut self, logits: &Tensor) -> Result<u32> {
|
||||
let logits = logits.to_dtype(DType::F32)?.to_vec1::<f32>()?;
|
||||
let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
|
||||
if self.mode == Sampling::Greedy {
|
||||
return Ok(logits.iter().position(|&l| l == max).unwrap_or(0) as u32);
|
||||
}
|
||||
let t = self.temperature as f32;
|
||||
let weights: Vec<f32> = logits.iter().map(|&l| ((l - max) / t).exp()).collect();
|
||||
let mut u = self.rng.random::<f32>() * weights.iter().sum::<f32>();
|
||||
for (i, w) in weights.iter().enumerate() {
|
||||
u -= w;
|
||||
if u < 0.0 {
|
||||
return Ok(i as u32);
|
||||
}
|
||||
}
|
||||
Ok((weights.len() - 1) as u32)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use hanzo_ml::Device;
|
||||
|
||||
fn cfg(mode: Sampling, temperature: f64) -> GrpoConfig {
|
||||
GrpoConfig {
|
||||
sampling: mode,
|
||||
temperature,
|
||||
..GrpoConfig::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn greedy_picks_argmax() -> Result<()> {
|
||||
let logits = Tensor::new(&[0.1f32, 3.0, -1.0, 2.9], &Device::Cpu)?;
|
||||
let mut s = Sampler::new(&cfg(Sampling::Greedy, 1.0), 0);
|
||||
assert_eq!(s.sample(&logits)?, 1);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn temperature_matches_softmax_frequencies() -> Result<()> {
|
||||
// softmax([0, ln 3]) = [0.25, 0.75]
|
||||
let logits = Tensor::new(&[0.0f32, 3f32.ln()], &Device::Cpu)?;
|
||||
let mut s = Sampler::new(&cfg(Sampling::Temperature, 1.0), 1);
|
||||
let n = 20_000;
|
||||
let ones = (0..n)
|
||||
.map(|_| s.sample(&logits))
|
||||
.filter(|t| matches!(t, Ok(1)))
|
||||
.count();
|
||||
let p = ones as f64 / n as f64;
|
||||
assert!((p - 0.75).abs() < 0.02, "P(1) = {p}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_seed_replays() -> Result<()> {
|
||||
let logits = Tensor::new(&[0.0f32, 0.0, 0.0, 0.0], &Device::Cpu)?;
|
||||
let draw = |seed| -> Result<Vec<u32>> {
|
||||
let mut s = Sampler::new(&cfg(Sampling::Temperature, 0.7), seed);
|
||||
(0..16).map(|_| s.sample(&logits)).collect()
|
||||
};
|
||||
assert_eq!(draw(3)?, draw(3)?);
|
||||
assert_ne!(draw(3)?, draw(4)?);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,372 @@
|
||||
//! One GRPO step: rollout, reward, group-relative advantage, loss, clipped AdamW update.
|
||||
|
||||
use crate::grpo::config::GrpoConfig;
|
||||
use crate::grpo::math;
|
||||
use crate::grpo::metrics::Metrics;
|
||||
use crate::grpo::policy::Policy;
|
||||
use crate::grpo::reward::Reward;
|
||||
use crate::grpo::sampler::Sampler;
|
||||
use crate::Prompt;
|
||||
use hanzo_ml::backprop::GradStore;
|
||||
use hanzo_ml::{DType, Var};
|
||||
use hanzo_nn::{AdamW, Optimizer, ParamsAdamW};
|
||||
|
||||
pub struct Trainer<P: Policy> {
|
||||
cfg: GrpoConfig,
|
||||
policy: P,
|
||||
rewards: Vec<Reward>,
|
||||
vars: Vec<Var>,
|
||||
optimizer: AdamW,
|
||||
steps: usize,
|
||||
/// Advances every step so rollouts do not replay.
|
||||
seed: u64,
|
||||
}
|
||||
|
||||
impl<P: Policy> Trainer<P> {
|
||||
pub fn new(cfg: GrpoConfig, policy: P, rewards: Vec<Reward>) -> anyhow::Result<Self> {
|
||||
cfg.validate()?;
|
||||
anyhow::ensure!(
|
||||
!rewards.is_empty(),
|
||||
"grpo needs at least one reward function"
|
||||
);
|
||||
let vars = policy.trainable_vars();
|
||||
anyhow::ensure!(
|
||||
!vars.is_empty(),
|
||||
"grpo: the policy has no trainable variables"
|
||||
);
|
||||
let params = ParamsAdamW {
|
||||
lr: cfg.learning_rate,
|
||||
weight_decay: cfg.weight_decay,
|
||||
..Default::default()
|
||||
};
|
||||
let optimizer = AdamW::new(vars.clone(), params)?;
|
||||
Ok(Self {
|
||||
seed: cfg.seed,
|
||||
cfg,
|
||||
policy,
|
||||
rewards,
|
||||
vars,
|
||||
optimizer,
|
||||
steps: 0,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn policy(&self) -> &P {
|
||||
&self.policy
|
||||
}
|
||||
|
||||
pub fn steps(&self) -> usize {
|
||||
self.steps
|
||||
}
|
||||
|
||||
/// `sum_f w_f * r_f(prompt, completion)`.
|
||||
fn score(&self, prompt: &Prompt, tokens: &[u32], text: &str) -> anyhow::Result<f32> {
|
||||
self.rewards
|
||||
.iter()
|
||||
.map(|r| Ok(r.weight * r.score(prompt, tokens, text)?))
|
||||
.sum()
|
||||
}
|
||||
|
||||
/// One update over `prompts`; each prompt yields a group of `group_size` completions.
|
||||
pub fn step(&mut self, prompts: &[Prompt]) -> anyhow::Result<Metrics> {
|
||||
anyhow::ensure!(!prompts.is_empty(), "grpo: step with no prompts");
|
||||
let g = self.cfg.group_size;
|
||||
let n = prompts.len() * g;
|
||||
|
||||
let mut rewards = Vec::with_capacity(n);
|
||||
let mut lengths = Vec::with_capacity(n);
|
||||
let mut logprobs = Vec::with_capacity(n);
|
||||
let mut group_mean = Vec::with_capacity(prompts.len());
|
||||
let mut group_std = Vec::with_capacity(prompts.len());
|
||||
for (pi, prompt) in prompts.iter().enumerate() {
|
||||
let seed = self
|
||||
.seed
|
||||
.wrapping_add((pi as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15));
|
||||
let mut sampler = Sampler::new(&self.cfg, seed);
|
||||
let group = self.policy.sample_group(
|
||||
&prompt.input_ids,
|
||||
g,
|
||||
self.cfg.max_completion_len,
|
||||
&mut sampler,
|
||||
)?;
|
||||
anyhow::ensure!(
|
||||
group.len() == g,
|
||||
"grpo: policy returned {} completions, expected {g}",
|
||||
group.len()
|
||||
);
|
||||
|
||||
let mut scores = Vec::with_capacity(g);
|
||||
for tokens in &group {
|
||||
let text = self.policy.decode(tokens)?;
|
||||
scores.push(self.score(prompt, tokens, &text)?);
|
||||
lengths.push(tokens.len());
|
||||
}
|
||||
let mean = scores.iter().sum::<f32>() / g as f32;
|
||||
let var = scores.iter().map(|&r| (r - mean).powi(2)).sum::<f32>() / (g as f32 - 1.0);
|
||||
group_mean.push(mean);
|
||||
group_std.push(var.sqrt());
|
||||
rewards.extend(scores);
|
||||
logprobs.extend(self.policy.group_logprob(&prompt.input_ids, &group)?);
|
||||
}
|
||||
|
||||
let advantages =
|
||||
math::batch_advantages(&rewards, g, self.cfg.scale_rewards, self.cfg.advantage_eps);
|
||||
let pg_loss = math::policy_gradient_loss(&advantages, &logprobs, &lengths)?;
|
||||
let pg_val = pg_loss.to_scalar::<f32>()?;
|
||||
let (loss, kl_val) = match self.cfg.kl_beta.filter(|&b| b > 0.0) {
|
||||
Some(beta) => {
|
||||
let reference: Vec<f32> = logprobs
|
||||
.iter()
|
||||
.map(|lp| lp.to_scalar::<f32>())
|
||||
.collect::<hanzo_ml::Result<_>>()?;
|
||||
let kl = math::kl_penalty(beta, &logprobs, &reference, &lengths)?;
|
||||
let kl_val = kl.to_scalar::<f32>()?;
|
||||
((&pg_loss + &kl)?, kl_val)
|
||||
}
|
||||
None => (pg_loss, 0.0),
|
||||
};
|
||||
let loss_val = loss.to_scalar::<f32>()?;
|
||||
|
||||
let mut grad_norm = 0.0;
|
||||
if loss_val.is_finite() {
|
||||
let mut grads = loss.backward()?;
|
||||
grad_norm = clip(&mut grads, &self.vars, self.cfg.max_grad_norm)? as f32;
|
||||
self.optimizer.step(&grads)?;
|
||||
}
|
||||
self.steps += 1;
|
||||
self.seed = self.seed.wrapping_add(prompts.len() as u64 + 1);
|
||||
|
||||
let groups = prompts.len() as f32;
|
||||
Ok(Metrics {
|
||||
loss: loss_val,
|
||||
pg_loss: pg_val,
|
||||
kl_loss: kl_val,
|
||||
mean_reward: group_mean.iter().sum::<f32>() / groups,
|
||||
reward_std: group_std.iter().sum::<f32>() / groups,
|
||||
mean_completion_len: lengths.iter().sum::<usize>() as f32 / n as f32,
|
||||
grad_norm,
|
||||
num_groups: prompts.len(),
|
||||
num_completions: n,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Scale the gradients of `vars` so their global L2 norm is at most `max`
|
||||
/// (`max <= 0` never scales). Returns the norm before scaling.
|
||||
pub fn clip(grads: &mut GradStore, vars: &[Var], max: f64) -> hanzo_ml::Result<f64> {
|
||||
let mut sq = 0f64;
|
||||
for v in vars {
|
||||
if let Some(g) = grads.get(v) {
|
||||
sq += g
|
||||
.sqr()?
|
||||
.sum_all()?
|
||||
.to_dtype(DType::F64)?
|
||||
.to_scalar::<f64>()?;
|
||||
}
|
||||
}
|
||||
let norm = sq.sqrt();
|
||||
if max > 0.0 && norm > max {
|
||||
let scale = max / (norm + 1e-6);
|
||||
for v in vars {
|
||||
if let Some(g) = grads.get(v) {
|
||||
let g = (g * scale)?;
|
||||
grads.insert(v, g);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(norm)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
//! The toy policy is a contextual bandit: one learnable `[max_len, vocab]`
|
||||
//! logits table, one categorical per output position, prompt-independent.
|
||||
//! It exercises the whole step: rollout through [`Sampler`], a
|
||||
//! differentiable `log_softmax` + gather log-probability, the advantage and
|
||||
//! loss math, clipping and AdamW.
|
||||
use super::*;
|
||||
use crate::grpo::config::Sampling;
|
||||
use crate::grpo::policy::fake::{tokenizer, Bigram};
|
||||
use crate::grpo::policy::LmPolicy;
|
||||
use hanzo_ml::{Device, Tensor};
|
||||
use hanzo_nn::ops::log_softmax;
|
||||
|
||||
struct Toy {
|
||||
logits: Var,
|
||||
max_len: usize,
|
||||
}
|
||||
|
||||
impl Toy {
|
||||
fn new(vocab: usize, max_len: usize) -> Self {
|
||||
Self {
|
||||
logits: Var::zeros((max_len, vocab), DType::F32, &Device::Cpu).unwrap(),
|
||||
max_len,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Policy for Toy {
|
||||
fn trainable_vars(&self) -> Vec<Var> {
|
||||
vec![self.logits.clone()]
|
||||
}
|
||||
|
||||
fn sample_group(
|
||||
&self,
|
||||
_: &[u32],
|
||||
group_size: usize,
|
||||
max_len: usize,
|
||||
sampler: &mut Sampler,
|
||||
) -> anyhow::Result<Vec<Vec<u32>>> {
|
||||
let steps = max_len.min(self.max_len);
|
||||
(0..group_size)
|
||||
.map(|_| {
|
||||
(0..steps)
|
||||
.map(|pos| Ok(sampler.sample(&self.logits.get(pos)?)?))
|
||||
.collect()
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn sequence_logprob(&self, _: &[u32], completion: &[u32]) -> anyhow::Result<Tensor> {
|
||||
let logp = log_softmax(self.logits.as_tensor(), 1)?;
|
||||
let idx = Tensor::from_slice(completion, (completion.len(), 1), &Device::Cpu)?;
|
||||
Ok(logp
|
||||
.narrow(0, 0, completion.len())?
|
||||
.gather(&idx, 1)?
|
||||
.sum_all()?)
|
||||
}
|
||||
}
|
||||
|
||||
fn prompt(ids: &[u32]) -> Prompt {
|
||||
Prompt {
|
||||
input_ids: ids.to_vec(),
|
||||
text: String::new(),
|
||||
answer: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn toy_reward_increases() -> anyhow::Result<()> {
|
||||
let (vocab, max_len, target) = (5, 4, 3);
|
||||
let cfg = GrpoConfig {
|
||||
group_size: 16,
|
||||
max_completion_len: max_len,
|
||||
sampling: Sampling::Temperature,
|
||||
temperature: 1.0,
|
||||
scale_rewards: true,
|
||||
advantage_eps: 1e-4,
|
||||
kl_beta: None,
|
||||
learning_rate: 0.5,
|
||||
weight_decay: 0.0,
|
||||
max_grad_norm: 0.0,
|
||||
seed: 42,
|
||||
};
|
||||
let mut trainer = Trainer::new(cfg, Toy::new(vocab, max_len), vec![Reward::token(target)])?;
|
||||
let prompts = [prompt(&[1, 2]), prompt(&[0]), prompt(&[4, 4, 4])];
|
||||
let first = trainer.step(&prompts)?;
|
||||
assert!(first.is_finite(), "{first:?}");
|
||||
assert_eq!((first.num_groups, first.num_completions), (3, 48));
|
||||
let mut last = first.clone();
|
||||
for _ in 0..40 {
|
||||
last = trainer.step(&prompts)?;
|
||||
assert!(last.is_finite(), "{last:?}");
|
||||
}
|
||||
println!(
|
||||
"[toy] reward {:.3} -> {:.3} of {max_len}",
|
||||
first.mean_reward, last.mean_reward
|
||||
);
|
||||
assert!(
|
||||
last.mean_reward > first.mean_reward + 0.5,
|
||||
"{} -> {}",
|
||||
first.mean_reward,
|
||||
last.mean_reward
|
||||
);
|
||||
assert!(
|
||||
last.mean_reward > max_len as f32 * 0.6,
|
||||
"{} vs {max_len}",
|
||||
last.mean_reward
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn toy_with_kl_stays_finite() -> anyhow::Result<()> {
|
||||
let cfg = GrpoConfig {
|
||||
group_size: 8,
|
||||
max_completion_len: 3,
|
||||
kl_beta: Some(0.05),
|
||||
learning_rate: 0.2,
|
||||
seed: 7,
|
||||
..GrpoConfig::default()
|
||||
};
|
||||
let mut trainer = Trainer::new(cfg, Toy::new(5, 3), vec![Reward::token(2)])?;
|
||||
let prompts = [prompt(&[1]), prompt(&[2])];
|
||||
let mut last = trainer.step(&prompts)?;
|
||||
for _ in 0..10 {
|
||||
last = trainer.step(&prompts)?;
|
||||
assert!(last.is_finite(), "{last:?}");
|
||||
}
|
||||
assert!(last.kl_loss >= 0.0);
|
||||
assert_eq!(trainer.steps(), 11);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clip_bounds_global_norm() -> hanzo_ml::Result<()> {
|
||||
let a = Var::from_slice(&[3.0f32, 0.0], 2, &Device::Cpu)?;
|
||||
let b = Var::from_slice(&[0.0f32, 4.0], 2, &Device::Cpu)?;
|
||||
let loss = (a.sqr()?.sum_all()? + b.sqr()?.sum_all()?)?.affine(0.5, 0.0)?;
|
||||
let mut grads = loss.backward()?;
|
||||
let vars = vec![a.clone(), b.clone()];
|
||||
assert!((clip(&mut grads, &vars, 0.0)? - 5.0).abs() < 1e-5);
|
||||
assert_eq!(grads.get(&a).unwrap().to_vec1::<f32>()?, vec![3.0, 0.0]);
|
||||
let before = clip(&mut grads, &vars, 1.0)?;
|
||||
assert!((before - 5.0).abs() < 1e-5);
|
||||
let after = clip(&mut grads, &vars, 1.0)?;
|
||||
assert!((after - 1.0).abs() < 1e-5, "{after}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lm_step_raises_rewarded_token() -> anyhow::Result<()> {
|
||||
// Vocab: unk q 3 7 <eos>. Prompt "q", answer "3", one-token completions.
|
||||
let model = Bigram::uniform(5, Some(4));
|
||||
let tk = tokenizer(&["<unk>", "q", "3", "7", "<eos>"]);
|
||||
let policy = LmPolicy {
|
||||
model: &model,
|
||||
tokenizer: &tk,
|
||||
eos: Some(4),
|
||||
};
|
||||
let cfg = GrpoConfig {
|
||||
group_size: 8,
|
||||
max_completion_len: 1,
|
||||
learning_rate: 0.3,
|
||||
max_grad_norm: 1.0,
|
||||
seed: 1,
|
||||
..GrpoConfig::default()
|
||||
};
|
||||
let mut trainer = Trainer::new(cfg, policy, vec![Reward::parse("exact_match", 1.0)?])?;
|
||||
let prompts = [Prompt {
|
||||
input_ids: vec![1],
|
||||
text: "q".into(),
|
||||
answer: Some("3".into()),
|
||||
}];
|
||||
let before = model.probs(1)[2];
|
||||
let mut m = trainer.step(&prompts)?;
|
||||
assert!(m.grad_norm > 0.0, "{m:?}");
|
||||
for _ in 0..30 {
|
||||
m = trainer.step(&prompts)?;
|
||||
}
|
||||
let after = model.probs(1)[2];
|
||||
println!(
|
||||
"[lm] P(3 | q) {before:.3} -> {after:.3}, reward {:.3}",
|
||||
m.mean_reward
|
||||
);
|
||||
assert!(after > 0.6 && after > before, "{before} -> {after}");
|
||||
assert!(
|
||||
model.probs(0).iter().all(|&p| (p - 0.2).abs() < 1e-6),
|
||||
"rows never seen stay uniform"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
+11
-2
@@ -21,7 +21,11 @@ pub fn snapshot(id_or_path: &str) -> Result<Snapshot> {
|
||||
return local(p);
|
||||
}
|
||||
let api = Api::new()?;
|
||||
let repo = api.repo(Repo::with_revision(id_or_path.to_string(), RepoType::Model, "main".into()));
|
||||
let repo = api.repo(Repo::with_revision(
|
||||
id_or_path.to_string(),
|
||||
RepoType::Model,
|
||||
"main".into(),
|
||||
));
|
||||
let config = repo.get("config.json").context("config.json")?;
|
||||
let tokenizer = repo.get("tokenizer.json").context("tokenizer.json")?;
|
||||
let tokenizer_config = repo.get("tokenizer_config.json").ok();
|
||||
@@ -29,7 +33,12 @@ pub fn snapshot(id_or_path: &str) -> Result<Snapshot> {
|
||||
Ok(index) => shards(&repo, &index)?,
|
||||
Err(_) => vec![repo.get("model.safetensors").context("model.safetensors")?],
|
||||
};
|
||||
Ok(Snapshot { config, tokenizer, tokenizer_config, weights })
|
||||
Ok(Snapshot {
|
||||
config,
|
||||
tokenizer,
|
||||
tokenizer_config,
|
||||
weights,
|
||||
})
|
||||
}
|
||||
|
||||
fn shards(repo: &ApiRepo, index: &Path) -> Result<Vec<PathBuf>> {
|
||||
|
||||
+11
-3
@@ -1,7 +1,11 @@
|
||||
use clap::{Parser, Subcommand};
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(name = "gym", version, about = "Hanzo Gym — LoRA fine-tuning and GRPO on hanzo-ml")]
|
||||
#[command(
|
||||
name = "gym",
|
||||
version,
|
||||
about = "Hanzo Gym — LoRA fine-tuning and GRPO on hanzo-ml"
|
||||
)]
|
||||
struct Cli {
|
||||
#[command(subcommand)]
|
||||
cmd: Cmd,
|
||||
@@ -24,11 +28,15 @@ enum Cmd {
|
||||
|
||||
fn main() -> anyhow::Result<()> {
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env().add_directive("gym=info".parse()?))
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::from_default_env().add_directive("gym=info".parse()?),
|
||||
)
|
||||
.init();
|
||||
match Cli::parse().cmd {
|
||||
Cmd::Train { config } => gym::train::run(&gym::Config::load(config)?),
|
||||
Cmd::Preprocess { config } => gym::data::preprocess(&gym::Config::load(config)?),
|
||||
Cmd::Merge { config, out } => gym::model::merge(&gym::Config::load(config)?, out.as_deref()),
|
||||
Cmd::Merge { config, out } => {
|
||||
gym::model::merge(&gym::Config::load(config)?, out.as_deref())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
//! Llama: `head_dim = hidden_size / num_attention_heads`, default rope only.
|
||||
|
||||
use super::transformer::Arch;
|
||||
|
||||
pub fn arch(v: &serde_json::Value) -> anyhow::Result<Arch> {
|
||||
anyhow::ensure!(
|
||||
v.get("rope_scaling").is_none_or(serde_json::Value::is_null),
|
||||
"rope_scaling is not supported; only default rope"
|
||||
);
|
||||
Arch::parse(v, false)
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
//! LoRA on a frozen linear layer, as PEFT defines it.
|
||||
|
||||
use hanzo_ml::{DType, Result, Tensor, Var};
|
||||
use hanzo_nn::Dropout;
|
||||
|
||||
/// `x·Wᵀ (+ b)` for `x` of shape `[b, t, in]` and `w` of shape `[out, in]`.
|
||||
pub fn linear(x: &Tensor, w: &Tensor, b: Option<&Tensor>) -> Result<Tensor> {
|
||||
let (bs, t, i) = x.dims3()?;
|
||||
let y = x
|
||||
.reshape((bs * t, i))?
|
||||
.matmul(&w.t()?)?
|
||||
.reshape((bs, t, ()))?;
|
||||
match b {
|
||||
Some(b) => y.broadcast_add(b),
|
||||
None => Ok(y),
|
||||
}
|
||||
}
|
||||
|
||||
/// Frozen projection.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Linear {
|
||||
pub weight: Tensor,
|
||||
pub bias: Option<Tensor>,
|
||||
}
|
||||
|
||||
/// Frozen projection plus a trainable low-rank update `scale·B·A`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LoraLinear {
|
||||
pub weight: Tensor,
|
||||
pub bias: Option<Tensor>,
|
||||
pub a: Var,
|
||||
pub b: Var,
|
||||
pub scale: f64,
|
||||
pub dropout: Option<Dropout>,
|
||||
}
|
||||
|
||||
impl LoraLinear {
|
||||
/// `A` is uniform in `±1/√in` (PEFT's kaiming init), `B` is zero, so the
|
||||
/// layer starts identical to the frozen one.
|
||||
pub fn new(
|
||||
weight: Tensor,
|
||||
bias: Option<Tensor>,
|
||||
r: usize,
|
||||
scale: f64,
|
||||
dropout: f64,
|
||||
) -> Result<Self> {
|
||||
let (out, inp) = weight.dims2()?;
|
||||
let bound = 1.0 / (inp as f32).sqrt();
|
||||
let a = Tensor::rand(-bound, bound, (r, inp), weight.device())?.to_dtype(weight.dtype())?;
|
||||
Ok(Self {
|
||||
a: Var::from_tensor(&a)?,
|
||||
b: Var::zeros((out, r), weight.dtype(), weight.device())?,
|
||||
scale,
|
||||
dropout: (dropout > 0.0).then(|| Dropout::new(dropout as f32)),
|
||||
weight,
|
||||
bias,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Tensor, train: bool) -> Result<Tensor> {
|
||||
let y = linear(x, &self.weight, self.bias.as_ref())?;
|
||||
let x = match &self.dropout {
|
||||
Some(d) => d.forward(x, train)?,
|
||||
None => x.clone(),
|
||||
};
|
||||
let d = linear(&linear(&x, &self.a, None)?, &self.b, None)?;
|
||||
y + (d * self.scale)?
|
||||
}
|
||||
|
||||
/// `W + scale·B·A`, computed in f32 and returned in the weight's dtype.
|
||||
pub fn merged_weight(&self) -> Result<Tensor> {
|
||||
let d = (self
|
||||
.b
|
||||
.to_dtype(DType::F32)?
|
||||
.matmul(&self.a.to_dtype(DType::F32)?)?
|
||||
* self.scale)?;
|
||||
(self.weight.to_dtype(DType::F32)? + d)?.to_dtype(self.weight.dtype())
|
||||
}
|
||||
}
|
||||
|
||||
/// A projection in the transformer: frozen, or frozen with a LoRA update.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum Proj {
|
||||
Base(Linear),
|
||||
Lora(LoraLinear),
|
||||
}
|
||||
|
||||
impl Proj {
|
||||
pub fn forward(&self, x: &Tensor, train: bool) -> Result<Tensor> {
|
||||
match self {
|
||||
Proj::Base(l) => linear(x, &l.weight, l.bias.as_ref()),
|
||||
Proj::Lora(l) => l.forward(x, train),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn lora(&self) -> Option<&LoraLinear> {
|
||||
match self {
|
||||
Proj::Lora(l) => Some(l),
|
||||
Proj::Base(_) => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
+178
-2
@@ -1,2 +1,178 @@
|
||||
//! model
|
||||
pub fn merge(_cfg: &crate::Config, _out: Option<&str>) -> anyhow::Result<()> { anyhow::bail!("merge: not implemented") }
|
||||
//! Trainable models. `load` builds a [`Transformer`] with LoRA from a Hugging
|
||||
//! Face checkpoint; adapters are written and read in PEFT layout; `merge`
|
||||
//! folds an adapter into the base weights.
|
||||
|
||||
pub mod llama;
|
||||
pub mod lora;
|
||||
pub mod qwen3;
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
pub mod transformer;
|
||||
|
||||
use crate::config::{Adapter, Config};
|
||||
use crate::{hub, CausalLm};
|
||||
use anyhow::Context;
|
||||
use hanzo_ml::{DType, Device, Result, Tensor, Var};
|
||||
use hanzo_nn::VarBuilder;
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
pub use transformer::{Arch, Lora, Transformer};
|
||||
|
||||
/// The transformer described by `cfg.base_model`, with LoRA on `cfg.lora_targets()`.
|
||||
pub fn build(cfg: &Config, device: &Device, dtype: DType) -> anyhow::Result<Transformer> {
|
||||
anyhow::ensure!(
|
||||
cfg.adapter == Adapter::Lora,
|
||||
"full fine-tuning is not implemented; set adapter: lora"
|
||||
);
|
||||
anyhow::ensure!(
|
||||
!cfg.gradient_checkpointing,
|
||||
"gradient_checkpointing is not implemented: hanzo-ml's autograd cannot recompute a layer in the backward pass; set gradient_checkpointing: false"
|
||||
);
|
||||
let snap = hub::snapshot(&cfg.base_model)?;
|
||||
let (arch, json) = arch(&snap.config)?;
|
||||
let lora = Lora {
|
||||
r: cfg.lora_r,
|
||||
alpha: cfg.lora_alpha,
|
||||
dropout: cfg.lora_dropout,
|
||||
targets: cfg.lora_targets(),
|
||||
};
|
||||
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&snap.weights, dtype, device)? };
|
||||
let mut model = Transformer::new(arch, lora, &vb)?;
|
||||
let eos = &json["eos_token_id"];
|
||||
model.eos = eos
|
||||
.as_u64()
|
||||
.or_else(|| eos.get(0)?.as_u64())
|
||||
.map(|e| e as u32);
|
||||
model.base = cfg.base_model.clone();
|
||||
Ok(model)
|
||||
}
|
||||
|
||||
/// `config.json` parsed by its `model_type`.
|
||||
fn arch(config: &Path) -> anyhow::Result<(Arch, serde_json::Value)> {
|
||||
let json: serde_json::Value = serde_json::from_slice(&std::fs::read(config)?)?;
|
||||
let arch = match json["model_type"].as_str() {
|
||||
Some("qwen3") => qwen3::arch(&json)?,
|
||||
Some("llama") => llama::arch(&json)?,
|
||||
other => anyhow::bail!("model_type {other:?} is not supported (qwen3, llama)"),
|
||||
};
|
||||
Ok((arch, json))
|
||||
}
|
||||
|
||||
pub fn load(cfg: &Config, device: &Device, dtype: DType) -> anyhow::Result<Box<dyn CausalLm>> {
|
||||
Ok(Box::new(build(cfg, device, dtype)?))
|
||||
}
|
||||
|
||||
/// Reads `adapter_model.safetensors` from `dir` into the model's LoRA vars.
|
||||
pub fn load_adapter(model: &Transformer, dir: &Path) -> anyhow::Result<()> {
|
||||
let file = dir.join("adapter_model.safetensors");
|
||||
let tensors = hanzo_ml::safetensors::load(&file, model.device())?;
|
||||
for (name, l) in model.adapters() {
|
||||
for (suffix, var) in [("lora_A", &l.a), ("lora_B", &l.b)] {
|
||||
let key = format!("{name}.{suffix}.weight");
|
||||
let t = tensors
|
||||
.get(&key)
|
||||
.with_context(|| format!("{}: missing {key}", file.display()))?;
|
||||
var.set(&t.to_dtype(model.dtype())?)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Writes the base weights with the adapter from `cfg.output_dir` folded in,
|
||||
/// plus config and tokenizer files, to `out` or `<output_dir>/merged`.
|
||||
pub fn merge(cfg: &Config, out: Option<&str>) -> anyhow::Result<()> {
|
||||
let snap = hub::snapshot(&cfg.base_model)?;
|
||||
let device = Device::Cpu;
|
||||
let mut base = HashMap::new();
|
||||
for w in &snap.weights {
|
||||
base.extend(hanzo_ml::safetensors::load(w, &device)?);
|
||||
}
|
||||
let dtype = base
|
||||
.get("model.embed_tokens.weight")
|
||||
.context("model.embed_tokens.weight")?
|
||||
.dtype();
|
||||
let (arch, _) = arch(&snap.config)?;
|
||||
let lora = Lora {
|
||||
r: cfg.lora_r,
|
||||
alpha: cfg.lora_alpha,
|
||||
dropout: 0.0,
|
||||
targets: cfg.lora_targets(),
|
||||
};
|
||||
let model = Transformer::new(
|
||||
arch,
|
||||
lora,
|
||||
&VarBuilder::from_tensors(base.clone(), dtype, &device),
|
||||
)?;
|
||||
load_adapter(&model, Path::new(&cfg.output_dir))?;
|
||||
for (name, l) in model.adapters() {
|
||||
let key = format!("{}.weight", name.trim_start_matches("base_model.model."));
|
||||
base.insert(key, l.merged_weight()?);
|
||||
}
|
||||
let out = out
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|| Path::new(&cfg.output_dir).join("merged"));
|
||||
std::fs::create_dir_all(&out)?;
|
||||
hanzo_ml::safetensors::save(&base, out.join("model.safetensors"))?;
|
||||
for f in [
|
||||
Some(&snap.config),
|
||||
Some(&snap.tokenizer),
|
||||
snap.tokenizer_config.as_ref(),
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
{
|
||||
std::fs::copy(f, out.join(f.file_name().context("file name")?))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
impl CausalLm for Transformer {
|
||||
fn forward(&self, input_ids: &Tensor, attention_mask: Option<&Tensor>) -> Result<Tensor> {
|
||||
Transformer::forward(self, input_ids, attention_mask)
|
||||
}
|
||||
|
||||
fn trainable_vars(&self) -> Vec<Var> {
|
||||
Transformer::trainable_vars(self)
|
||||
}
|
||||
|
||||
fn device(&self) -> &Device {
|
||||
Transformer::device(self)
|
||||
}
|
||||
|
||||
fn dtype(&self) -> DType {
|
||||
Transformer::dtype(self)
|
||||
}
|
||||
|
||||
fn eos_token_id(&self) -> Option<u32> {
|
||||
self.eos
|
||||
}
|
||||
|
||||
fn vocab_size(&self) -> usize {
|
||||
self.arch.vocab
|
||||
}
|
||||
|
||||
fn save_adapter(&self, dir: &Path) -> anyhow::Result<()> {
|
||||
std::fs::create_dir_all(dir)?;
|
||||
let mut tensors = HashMap::new();
|
||||
for (name, l) in self.adapters() {
|
||||
tensors.insert(format!("{name}.lora_A.weight"), l.a.as_tensor().clone());
|
||||
tensors.insert(format!("{name}.lora_B.weight"), l.b.as_tensor().clone());
|
||||
}
|
||||
hanzo_ml::safetensors::save(&tensors, dir.join("adapter_model.safetensors"))?;
|
||||
let config = serde_json::json!({
|
||||
"peft_type": "LORA",
|
||||
"task_type": "CAUSAL_LM",
|
||||
"r": self.lora.r,
|
||||
"lora_alpha": self.lora.alpha,
|
||||
"lora_dropout": self.lora.dropout,
|
||||
"target_modules": self.lora.targets,
|
||||
"base_model_name_or_path": self.base,
|
||||
"bias": "none",
|
||||
});
|
||||
std::fs::write(
|
||||
dir.join("adapter_config.json"),
|
||||
serde_json::to_string_pretty(&config)?,
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
//! Qwen3: the Llama layout with per-head RMSNorm on q and k and an explicit `head_dim`.
|
||||
|
||||
use super::transformer::Arch;
|
||||
|
||||
pub fn arch(v: &serde_json::Value) -> anyhow::Result<Arch> {
|
||||
anyhow::ensure!(
|
||||
v.get("head_dim").is_some_and(|d| d.is_u64()),
|
||||
"qwen3 config.json needs head_dim"
|
||||
);
|
||||
Arch::parse(v, true)
|
||||
}
|
||||
@@ -0,0 +1,377 @@
|
||||
use super::*;
|
||||
use hanzo_ml::{DType, Device, Tensor};
|
||||
use hanzo_nn::{loss::cross_entropy, AdamW, Optimizer, ParamsAdamW, VarBuilder};
|
||||
use std::collections::HashMap;
|
||||
|
||||
const V: usize = 64;
|
||||
const H: usize = 32;
|
||||
const I: usize = 64;
|
||||
const L: usize = 2;
|
||||
const NH: usize = 4;
|
||||
const KV: usize = 2;
|
||||
const HD: usize = 8;
|
||||
const T: usize = 6;
|
||||
const ALL: [&str; 7] = [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
"o_proj",
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
"down_proj",
|
||||
];
|
||||
|
||||
fn arch(qk_norm: bool) -> Arch {
|
||||
Arch {
|
||||
vocab: V,
|
||||
hidden: H,
|
||||
inter: I,
|
||||
layers: L,
|
||||
heads: NH,
|
||||
kv_heads: KV,
|
||||
head_dim: HD,
|
||||
eps: 1e-6,
|
||||
theta: 10_000.0,
|
||||
tied: false,
|
||||
attn_bias: false,
|
||||
qk_norm,
|
||||
}
|
||||
}
|
||||
|
||||
/// Random base weights under Hugging Face names.
|
||||
fn weights(dev: &Device, qk_norm: bool) -> HashMap<String, Tensor> {
|
||||
let r = |s: &[usize]| Tensor::randn(0f32, 0.1, s, dev).unwrap();
|
||||
let g = |n: usize| Tensor::rand(0.5f32, 1.5, n, dev).unwrap();
|
||||
let mut m = HashMap::from([
|
||||
("model.embed_tokens.weight".to_string(), r(&[V, H])),
|
||||
("lm_head.weight".to_string(), r(&[V, H])),
|
||||
("model.norm.weight".to_string(), g(H)),
|
||||
]);
|
||||
let shapes = [
|
||||
("self_attn.q_proj", [NH * HD, H]),
|
||||
("self_attn.k_proj", [KV * HD, H]),
|
||||
("self_attn.v_proj", [KV * HD, H]),
|
||||
("self_attn.o_proj", [H, NH * HD]),
|
||||
("mlp.gate_proj", [I, H]),
|
||||
("mlp.up_proj", [I, H]),
|
||||
("mlp.down_proj", [H, I]),
|
||||
];
|
||||
for i in 0..L {
|
||||
let p = format!("model.layers.{i}");
|
||||
for (n, s) in shapes {
|
||||
m.insert(format!("{p}.{n}.weight"), r(&s));
|
||||
}
|
||||
m.insert(format!("{p}.input_layernorm.weight"), g(H));
|
||||
m.insert(format!("{p}.post_attention_layernorm.weight"), g(H));
|
||||
if qk_norm {
|
||||
m.insert(format!("{p}.self_attn.q_norm.weight"), g(HD));
|
||||
m.insert(format!("{p}.self_attn.k_norm.weight"), g(HD));
|
||||
}
|
||||
}
|
||||
m
|
||||
}
|
||||
|
||||
fn model(w: &HashMap<String, Tensor>, qk_norm: bool, targets: &[&str]) -> Transformer {
|
||||
let lora = Lora {
|
||||
r: 4,
|
||||
alpha: 8.0,
|
||||
dropout: 0.0,
|
||||
targets: targets.iter().map(|s| s.to_string()).collect(),
|
||||
};
|
||||
Transformer::new(
|
||||
arch(qk_norm),
|
||||
lora,
|
||||
&VarBuilder::from_tensors(w.clone(), DType::F32, &Device::Cpu),
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn ids(dev: &Device) -> Tensor {
|
||||
let v: Vec<u32> = (0..2 * T).map(|i| ((i * 7 + 3) % V) as u32).collect();
|
||||
Tensor::from_vec(v, (2, T), dev).unwrap()
|
||||
}
|
||||
|
||||
/// Random `B` so the LoRA path contributes.
|
||||
fn perturb(m: &Transformer) {
|
||||
for (_, l) in m.adapters() {
|
||||
l.b.set(&Tensor::randn(0f32, 0.1, l.b.dims(), l.b.device()).unwrap())
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
fn max_diff(a: &Tensor, b: &Tensor) -> f32 {
|
||||
(a - b)
|
||||
.unwrap()
|
||||
.abs()
|
||||
.unwrap()
|
||||
.max_all()
|
||||
.unwrap()
|
||||
.to_scalar()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
/// Next-token cross-entropy over every position.
|
||||
fn loss(m: &Transformer, ids: &Tensor) -> Tensor {
|
||||
let logits = m.forward(ids, None).unwrap();
|
||||
let t = ids.dim(1).unwrap();
|
||||
let inp = logits.narrow(1, 0, t - 1).unwrap().flatten(0, 1).unwrap();
|
||||
let target = ids.narrow(1, 1, t - 1).unwrap().flatten_all().unwrap();
|
||||
cross_entropy(&inp, &target).unwrap()
|
||||
}
|
||||
|
||||
fn tmp(name: &str) -> std::path::PathBuf {
|
||||
let d = std::env::temp_dir().join(format!("gym-{name}-{}", std::process::id()));
|
||||
std::fs::create_dir_all(&d).unwrap();
|
||||
d
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn qwen3_matches_reference() {
|
||||
use hanzo_transformers::models::qwen3::{Config, Model};
|
||||
let dev = Device::Cpu;
|
||||
let w = weights(&dev, true);
|
||||
let ids = ids(&dev);
|
||||
let ours = model(&w, true, &ALL).forward(&ids, None).unwrap();
|
||||
let cfg = Config {
|
||||
vocab_size: V,
|
||||
hidden_size: H,
|
||||
intermediate_size: I,
|
||||
num_hidden_layers: L,
|
||||
num_attention_heads: NH,
|
||||
head_dim: HD,
|
||||
attention_bias: false,
|
||||
num_key_value_heads: KV,
|
||||
max_position_embeddings: 64,
|
||||
sliding_window: None,
|
||||
max_window_layers: 0,
|
||||
tie_word_embeddings: false,
|
||||
rope_theta: 10_000.0,
|
||||
rms_norm_eps: 1e-6,
|
||||
use_sliding_window: false,
|
||||
hidden_act: hanzo_nn::Activation::Silu,
|
||||
};
|
||||
let mut theirs =
|
||||
Model::new(&cfg, VarBuilder::from_tensors(w.clone(), DType::F32, &dev)).unwrap();
|
||||
let hidden = theirs.forward(&ids, 0).unwrap();
|
||||
let logits = hidden
|
||||
.broadcast_matmul(&w["lm_head.weight"].t().unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(ours.dims(), &[2, T, V]);
|
||||
assert!(
|
||||
max_diff(&ours, &logits) < 1e-4,
|
||||
"qwen3 diff {}",
|
||||
max_diff(&ours, &logits)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn llama_matches_reference() {
|
||||
use hanzo_transformers::models::llama::{Cache, Config, Llama};
|
||||
let dev = Device::Cpu;
|
||||
let w = weights(&dev, false);
|
||||
let ids = ids(&dev);
|
||||
let ours = model(&w, false, &ALL).forward(&ids, None).unwrap();
|
||||
let cfg = Config {
|
||||
hidden_size: H,
|
||||
intermediate_size: I,
|
||||
vocab_size: V,
|
||||
num_hidden_layers: L,
|
||||
num_attention_heads: NH,
|
||||
num_key_value_heads: KV,
|
||||
use_flash_attn: false,
|
||||
rms_norm_eps: 1e-6,
|
||||
rope_theta: 10_000.0,
|
||||
bos_token_id: None,
|
||||
eos_token_id: None,
|
||||
rope_scaling: None,
|
||||
max_position_embeddings: 64,
|
||||
tie_word_embeddings: false,
|
||||
};
|
||||
let theirs = Llama::load(VarBuilder::from_tensors(w.clone(), DType::F32, &dev), &cfg).unwrap();
|
||||
let mut cache = Cache::new(false, DType::F32, &cfg, &dev).unwrap();
|
||||
// Their forward returns the last position only; every prefix covers every position.
|
||||
for n in 1..=T {
|
||||
let last = theirs
|
||||
.forward(&ids.narrow(1, 0, n).unwrap(), 0, &mut cache)
|
||||
.unwrap();
|
||||
let ours_n = ours.narrow(1, n - 1, 1).unwrap().squeeze(1).unwrap();
|
||||
assert!(
|
||||
max_diff(&ours_n, &last) < 1e-4,
|
||||
"llama position {} diff {}",
|
||||
n - 1,
|
||||
max_diff(&ours_n, &last)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gradients_reach_only_lora() {
|
||||
let dev = Device::Cpu;
|
||||
let w = weights(&dev, true);
|
||||
let m = model(&w, true, &ALL);
|
||||
let grads = loss(&m, &ids(&dev)).backward().unwrap();
|
||||
let vars = m.trainable_vars();
|
||||
assert_eq!(vars.len(), 2 * ALL.len() * L);
|
||||
for v in &vars {
|
||||
assert!(grads.get(v).is_some());
|
||||
}
|
||||
assert!(grads.get(&w["lm_head.weight"]).is_none());
|
||||
assert!(grads
|
||||
.get(&w["model.layers.0.self_attn.q_proj.weight"])
|
||||
.is_none());
|
||||
assert_eq!(
|
||||
model(&w, true, &["q_proj", "v_proj"])
|
||||
.trainable_vars()
|
||||
.len(),
|
||||
2 * 2 * L
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn adamw_reduces_loss() {
|
||||
let dev = Device::Cpu;
|
||||
let w = weights(&dev, true);
|
||||
let m = model(&w, true, &ALL);
|
||||
let ids = ids(&dev);
|
||||
let mut opt = AdamW::new(
|
||||
m.trainable_vars(),
|
||||
ParamsAdamW {
|
||||
lr: 1e-2,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
let before: f32 = loss(&m, &ids).to_scalar().unwrap();
|
||||
for _ in 0..30 {
|
||||
opt.backward_step(&loss(&m, &ids)).unwrap();
|
||||
}
|
||||
let after: f32 = loss(&m, &ids).to_scalar().unwrap();
|
||||
println!("loss {before} -> {after}");
|
||||
assert!(after < 0.7 * before, "loss {before} -> {after}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn padding_leaves_real_positions_unchanged() {
|
||||
let dev = Device::Cpu;
|
||||
let w = weights(&dev, true);
|
||||
let m = model(&w, true, &ALL);
|
||||
perturb(&m);
|
||||
let ids = ids(&dev);
|
||||
let mask = Tensor::from_vec(vec![1u8, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], (2, T), &dev).unwrap();
|
||||
let padded = m.forward(&ids, Some(&mask)).unwrap();
|
||||
for (row, len) in [(0, T), (1, T - 2)] {
|
||||
let alone = m
|
||||
.forward(
|
||||
&ids.narrow(0, row, 1).unwrap().narrow(1, 0, len).unwrap(),
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
let real = padded.narrow(0, row, 1).unwrap().narrow(1, 0, len).unwrap();
|
||||
assert!(
|
||||
max_diff(&alone, &real) < 1e-4,
|
||||
"row {row} diff {}",
|
||||
max_diff(&alone, &real)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn adapter_saves_loads_and_merges() {
|
||||
let dev = Device::Cpu;
|
||||
let w = weights(&dev, true);
|
||||
let m = model(&w, true, &ALL);
|
||||
perturb(&m);
|
||||
let ids = ids(&dev);
|
||||
let want = m.forward(&ids, None).unwrap();
|
||||
|
||||
let dir = tmp("adapter");
|
||||
m.save_adapter(&dir).unwrap();
|
||||
let saved = hanzo_ml::safetensors::load(dir.join("adapter_model.safetensors"), &dev).unwrap();
|
||||
assert_eq!(saved.len(), 2 * ALL.len() * L);
|
||||
assert_eq!(
|
||||
saved["base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight"].dims(),
|
||||
&[4, H]
|
||||
);
|
||||
assert_eq!(
|
||||
saved["base_model.model.model.layers.1.mlp.down_proj.lora_B.weight"].dims(),
|
||||
&[H, 4]
|
||||
);
|
||||
let cfg: serde_json::Value =
|
||||
serde_json::from_slice(&std::fs::read(dir.join("adapter_config.json")).unwrap()).unwrap();
|
||||
assert_eq!(cfg["peft_type"], "LORA");
|
||||
assert_eq!(cfg["r"], 4);
|
||||
assert_eq!(cfg["target_modules"].as_array().unwrap().len(), ALL.len());
|
||||
|
||||
let fresh = model(&w, true, &ALL);
|
||||
load_adapter(&fresh, &dir).unwrap();
|
||||
assert_eq!(max_diff(&fresh.forward(&ids, None).unwrap(), &want), 0.0);
|
||||
|
||||
let mut merged = w.clone();
|
||||
for (name, l) in m.adapters() {
|
||||
merged.insert(
|
||||
format!("{}.weight", name.trim_start_matches("base_model.model.")),
|
||||
l.merged_weight().unwrap(),
|
||||
);
|
||||
}
|
||||
let plain = model(&merged, true, &[]);
|
||||
assert!(plain.trainable_vars().is_empty());
|
||||
let got = plain.forward(&ids, None).unwrap();
|
||||
assert!(
|
||||
max_diff(&got, &want) < 1e-4,
|
||||
"merged diff {}",
|
||||
max_diff(&got, &want)
|
||||
);
|
||||
std::fs::remove_dir_all(dir).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_and_merge_from_local_checkpoint() {
|
||||
let dev = Device::Cpu;
|
||||
let w = weights(&dev, true);
|
||||
let dir = tmp("checkpoint");
|
||||
hanzo_ml::safetensors::save(&w, dir.join("model.safetensors")).unwrap();
|
||||
let config = serde_json::json!({
|
||||
"model_type": "qwen3", "vocab_size": V, "hidden_size": H, "intermediate_size": I,
|
||||
"num_hidden_layers": L, "num_attention_heads": NH, "num_key_value_heads": KV, "head_dim": HD,
|
||||
"rms_norm_eps": 1e-6, "rope_theta": 10000.0, "tie_word_embeddings": false,
|
||||
"attention_bias": false, "eos_token_id": [151645, 151643],
|
||||
});
|
||||
std::fs::write(dir.join("config.json"), config.to_string()).unwrap();
|
||||
std::fs::write(dir.join("tokenizer.json"), "{}").unwrap();
|
||||
let yaml = |extra: &str| -> Config {
|
||||
serde_yaml::from_str(&format!(
|
||||
"base_model: {d}\noutput_dir: {d}/out\nlora_r: 4\nlora_alpha: 8\nlora_target_linear: true\n{extra}",
|
||||
d = dir.display()
|
||||
))
|
||||
.unwrap()
|
||||
};
|
||||
let cfg = yaml("adapter: lora");
|
||||
assert!(build(
|
||||
&yaml("adapter: lora\ngradient_checkpointing: true"),
|
||||
&dev,
|
||||
DType::F32
|
||||
)
|
||||
.is_err());
|
||||
assert!(build(&yaml("adapter: none"), &dev, DType::F32).is_err());
|
||||
|
||||
let m = build(&cfg, &dev, DType::F32).unwrap();
|
||||
assert_eq!(m.eos, Some(151645));
|
||||
assert_eq!(m.base, dir.display().to_string());
|
||||
perturb(&m);
|
||||
let ids = ids(&dev);
|
||||
let want = m.forward(&ids, None).unwrap();
|
||||
m.save_adapter(Path::new(&cfg.output_dir)).unwrap();
|
||||
|
||||
merge(&cfg, None).unwrap();
|
||||
let out = dir.join("out/merged");
|
||||
assert!(out.join("config.json").exists() && out.join("tokenizer.json").exists());
|
||||
let merged = hanzo_ml::safetensors::load(out.join("model.safetensors"), &dev).unwrap();
|
||||
assert_eq!(merged.len(), w.len());
|
||||
let got = model(&merged, true, &[]).forward(&ids, None).unwrap();
|
||||
assert!(
|
||||
max_diff(&got, &want) < 1e-4,
|
||||
"merged diff {}",
|
||||
max_diff(&got, &want)
|
||||
);
|
||||
std::fs::remove_dir_all(dir).unwrap();
|
||||
}
|
||||
@@ -0,0 +1,335 @@
|
||||
//! The Llama-family decoder (Llama, Qwen3), built only from ops hanzo-ml can
|
||||
//! differentiate: `rms_norm_slow`, `rope_slow`, `softmax`, `silu`, matmul.
|
||||
|
||||
use crate::model::lora::{linear, Linear, LoraLinear, Proj};
|
||||
use hanzo_ml::{DType, Device, Result, Tensor, Var, D};
|
||||
use hanzo_nn::ops::{rms_norm_slow, silu, softmax};
|
||||
use hanzo_nn::rotary_emb::rope_slow;
|
||||
use hanzo_nn::VarBuilder;
|
||||
use serde::Deserialize;
|
||||
|
||||
/// Hyperparameters of one checkpoint.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Arch {
|
||||
pub vocab: usize,
|
||||
pub hidden: usize,
|
||||
pub inter: usize,
|
||||
pub layers: usize,
|
||||
pub heads: usize,
|
||||
pub kv_heads: usize,
|
||||
pub head_dim: usize,
|
||||
pub eps: f64,
|
||||
pub theta: f64,
|
||||
pub tied: bool,
|
||||
pub attn_bias: bool,
|
||||
/// Per-head RMSNorm on q and k (Qwen3).
|
||||
pub qk_norm: bool,
|
||||
}
|
||||
|
||||
/// The `config.json` keys shared by the family.
|
||||
#[derive(Deserialize)]
|
||||
struct Json {
|
||||
vocab_size: usize,
|
||||
hidden_size: usize,
|
||||
intermediate_size: usize,
|
||||
num_hidden_layers: usize,
|
||||
num_attention_heads: usize,
|
||||
num_key_value_heads: Option<usize>,
|
||||
head_dim: Option<usize>,
|
||||
rms_norm_eps: f64,
|
||||
#[serde(default = "ten_thousand")]
|
||||
rope_theta: f64,
|
||||
#[serde(default)]
|
||||
tie_word_embeddings: bool,
|
||||
#[serde(default)]
|
||||
attention_bias: bool,
|
||||
}
|
||||
|
||||
fn ten_thousand() -> f64 {
|
||||
10_000.0
|
||||
}
|
||||
|
||||
impl Arch {
|
||||
pub fn parse(v: &serde_json::Value, qk_norm: bool) -> anyhow::Result<Self> {
|
||||
let j: Json = serde_json::from_value(v.clone())?;
|
||||
Ok(Self {
|
||||
vocab: j.vocab_size,
|
||||
hidden: j.hidden_size,
|
||||
inter: j.intermediate_size,
|
||||
layers: j.num_hidden_layers,
|
||||
heads: j.num_attention_heads,
|
||||
kv_heads: j.num_key_value_heads.unwrap_or(j.num_attention_heads),
|
||||
head_dim: j.head_dim.unwrap_or(j.hidden_size / j.num_attention_heads),
|
||||
eps: j.rms_norm_eps,
|
||||
theta: j.rope_theta,
|
||||
tied: j.tie_word_embeddings,
|
||||
attn_bias: j.attention_bias,
|
||||
qk_norm,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// LoRA settings applied while building.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Lora {
|
||||
pub r: usize,
|
||||
pub alpha: f64,
|
||||
pub dropout: f64,
|
||||
/// Projection names (`q_proj`, `down_proj`, ...).
|
||||
pub targets: Vec<String>,
|
||||
}
|
||||
|
||||
/// Projection paths under `model.layers.{i}`, in HF order.
|
||||
const PROJS: [&str; 7] = [
|
||||
"self_attn.q_proj",
|
||||
"self_attn.k_proj",
|
||||
"self_attn.v_proj",
|
||||
"self_attn.o_proj",
|
||||
"mlp.gate_proj",
|
||||
"mlp.up_proj",
|
||||
"mlp.down_proj",
|
||||
];
|
||||
|
||||
struct Layer {
|
||||
ln1: Tensor,
|
||||
ln2: Tensor,
|
||||
projs: [Proj; 7],
|
||||
qk_norm: Option<(Tensor, Tensor)>,
|
||||
}
|
||||
|
||||
pub struct Transformer {
|
||||
pub arch: Arch,
|
||||
pub lora: Lora,
|
||||
/// Applies LoRA dropout in `forward`.
|
||||
pub train: bool,
|
||||
pub eos: Option<u32>,
|
||||
/// Checkpoint id, written to `adapter_config.json`.
|
||||
pub base: String,
|
||||
embed: Tensor,
|
||||
layers: Vec<Layer>,
|
||||
norm: Tensor,
|
||||
head: Tensor,
|
||||
inv_freq: Tensor,
|
||||
device: Device,
|
||||
dtype: DType,
|
||||
}
|
||||
|
||||
fn proj(
|
||||
vb: &VarBuilder,
|
||||
name: &str,
|
||||
out: usize,
|
||||
inp: usize,
|
||||
bias: bool,
|
||||
lora: &Lora,
|
||||
) -> Result<Proj> {
|
||||
let weight = vb.get((out, inp), &format!("{name}.weight"))?;
|
||||
let bias = bias
|
||||
.then(|| vb.get(out, &format!("{name}.bias")))
|
||||
.transpose()?;
|
||||
let short = name.rsplit('.').next().unwrap_or(name);
|
||||
if lora.targets.iter().any(|t| t == short) {
|
||||
let scale = lora.alpha / lora.r as f64;
|
||||
Ok(Proj::Lora(LoraLinear::new(
|
||||
weight,
|
||||
bias,
|
||||
lora.r,
|
||||
scale,
|
||||
lora.dropout,
|
||||
)?))
|
||||
} else {
|
||||
Ok(Proj::Base(Linear { weight, bias }))
|
||||
}
|
||||
}
|
||||
|
||||
impl Layer {
|
||||
fn new(a: &Arch, vb: &VarBuilder, lora: &Lora) -> Result<Self> {
|
||||
let (h, kv, d) = (a.heads * a.head_dim, a.kv_heads * a.head_dim, a.hidden);
|
||||
let shapes = [
|
||||
(h, d),
|
||||
(kv, d),
|
||||
(kv, d),
|
||||
(d, h),
|
||||
(a.inter, d),
|
||||
(a.inter, d),
|
||||
(d, a.inter),
|
||||
];
|
||||
let mut projs = Vec::with_capacity(7);
|
||||
for (i, name) in PROJS.iter().enumerate() {
|
||||
let bias = a.attn_bias && i < 4;
|
||||
projs.push(proj(vb, name, shapes[i].0, shapes[i].1, bias, lora)?);
|
||||
}
|
||||
let qk_norm = if a.qk_norm {
|
||||
Some((
|
||||
vb.get(a.head_dim, "self_attn.q_norm.weight")?,
|
||||
vb.get(a.head_dim, "self_attn.k_norm.weight")?,
|
||||
))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
Ok(Self {
|
||||
ln1: vb.get(a.hidden, "input_layernorm.weight")?,
|
||||
ln2: vb.get(a.hidden, "post_attention_layernorm.weight")?,
|
||||
projs: projs
|
||||
.try_into()
|
||||
.map_err(|_| hanzo_ml::Error::Msg("seven projections".into()))?,
|
||||
qk_norm,
|
||||
})
|
||||
}
|
||||
|
||||
fn forward(
|
||||
&self,
|
||||
x: &Tensor,
|
||||
a: &Arch,
|
||||
cos: &Tensor,
|
||||
sin: &Tensor,
|
||||
mask: &Tensor,
|
||||
train: bool,
|
||||
) -> Result<Tensor> {
|
||||
let eps = a.eps as f32;
|
||||
let h = rms_norm_slow(x, &self.ln1, eps)?;
|
||||
let x = (x + self.attn(&h, a, cos, sin, mask, train)?)?;
|
||||
let h = rms_norm_slow(&x, &self.ln2, eps)?;
|
||||
let [_, _, _, _, gate, up, down] = &self.projs;
|
||||
let h = (silu(&gate.forward(&h, train)?)? * up.forward(&h, train)?)?;
|
||||
x + down.forward(&h, train)?
|
||||
}
|
||||
|
||||
fn attn(
|
||||
&self,
|
||||
x: &Tensor,
|
||||
a: &Arch,
|
||||
cos: &Tensor,
|
||||
sin: &Tensor,
|
||||
mask: &Tensor,
|
||||
train: bool,
|
||||
) -> Result<Tensor> {
|
||||
let (b, t, _) = x.dims3()?;
|
||||
let (h, kv, d) = (a.heads, a.kv_heads, a.head_dim);
|
||||
let [q, k, v, o, ..] = &self.projs;
|
||||
let split = |x: Tensor, n: usize| x.reshape((b, t, n, d))?.transpose(1, 2);
|
||||
let mut q = split(q.forward(x, train)?, h)?;
|
||||
let mut k = split(k.forward(x, train)?, kv)?;
|
||||
let v = split(v.forward(x, train)?, kv)?;
|
||||
if let Some((qn, kn)) = &self.qk_norm {
|
||||
q = rms_norm_slow(&q, qn, a.eps as f32)?;
|
||||
k = rms_norm_slow(&k, kn, a.eps as f32)?;
|
||||
}
|
||||
let q = rope_slow(&q, cos, sin)?.to_dtype(DType::F32)?;
|
||||
let k = repeat(&rope_slow(&k, cos, sin)?, h / kv)?.to_dtype(DType::F32)?;
|
||||
let v = repeat(&v, h / kv)?.to_dtype(DType::F32)?;
|
||||
let scores = (q.matmul(&k.transpose(2, 3)?.contiguous()?)? / (d as f64).sqrt())?;
|
||||
let p = softmax(&scores.broadcast_add(mask)?, D::Minus1)?;
|
||||
let ctx = p.matmul(&v.contiguous()?)?.to_dtype(x.dtype())?;
|
||||
o.forward(&ctx.transpose(1, 2)?.reshape((b, t, h * d))?, train)
|
||||
}
|
||||
}
|
||||
|
||||
/// Repeats each kv head `n` times along the head axis: `[b, kv, t, d]` to `[b, kv·n, t, d]`.
|
||||
fn repeat(x: &Tensor, n: usize) -> Result<Tensor> {
|
||||
if n == 1 {
|
||||
return Ok(x.clone());
|
||||
}
|
||||
let (b, kv, t, d) = x.dims4()?;
|
||||
Tensor::cat(&vec![x; n], 2)?.reshape((b, kv * n, t, d))
|
||||
}
|
||||
|
||||
impl Transformer {
|
||||
pub fn new(arch: Arch, lora: Lora, vb: &VarBuilder) -> Result<Self> {
|
||||
let embed = vb.get((arch.vocab, arch.hidden), "model.embed_tokens.weight")?;
|
||||
let head = if arch.tied {
|
||||
embed.clone()
|
||||
} else {
|
||||
vb.get((arch.vocab, arch.hidden), "lm_head.weight")?
|
||||
};
|
||||
let layers = (0..arch.layers)
|
||||
.map(|i| Layer::new(&arch, &vb.pp(format!("model.layers.{i}")), &lora))
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
let d = arch.head_dim;
|
||||
let inv_freq: Vec<f32> = (0..d)
|
||||
.step_by(2)
|
||||
.map(|i| 1.0 / arch.theta.powf(i as f64 / d as f64) as f32)
|
||||
.collect();
|
||||
Ok(Self {
|
||||
norm: vb.get(arch.hidden, "model.norm.weight")?,
|
||||
inv_freq: Tensor::from_vec(inv_freq, (1, d / 2), vb.device())?,
|
||||
device: vb.device().clone(),
|
||||
dtype: vb.dtype(),
|
||||
embed,
|
||||
layers,
|
||||
head,
|
||||
arch,
|
||||
lora,
|
||||
train: true,
|
||||
eos: None,
|
||||
base: String::new(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn device(&self) -> &Device {
|
||||
&self.device
|
||||
}
|
||||
|
||||
pub fn dtype(&self) -> DType {
|
||||
self.dtype
|
||||
}
|
||||
|
||||
/// Logits `[b, t, vocab]` over the full sequence. `attention_mask` is
|
||||
/// `[b, t]` u8 with 1 at real tokens.
|
||||
pub fn forward(&self, ids: &Tensor, attention_mask: Option<&Tensor>) -> Result<Tensor> {
|
||||
let (b, t) = ids.dims2()?;
|
||||
let a = &self.arch;
|
||||
let mut x = self
|
||||
.embed
|
||||
.index_select(&ids.flatten_all()?, 0)?
|
||||
.reshape((b, t, a.hidden))?;
|
||||
let pos = Tensor::arange(0u32, t as u32, &self.device)?
|
||||
.to_dtype(DType::F32)?
|
||||
.reshape((t, 1))?;
|
||||
let freqs = pos.matmul(&self.inv_freq)?;
|
||||
let (cos, sin) = (
|
||||
freqs.cos()?.to_dtype(self.dtype)?,
|
||||
freqs.sin()?.to_dtype(self.dtype)?,
|
||||
);
|
||||
let mask = self.mask(t, attention_mask)?;
|
||||
for l in &self.layers {
|
||||
x = l.forward(&x, a, &cos, &sin, &mask, self.train)?;
|
||||
}
|
||||
linear(
|
||||
&rms_norm_slow(&x, &self.norm, a.eps as f32)?,
|
||||
&self.head,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
/// Additive f32 mask: causal `[1, 1, t, t]`, or `[b, 1, t, t]` with `-inf`
|
||||
/// on padded keys.
|
||||
fn mask(&self, t: usize, attention_mask: Option<&Tensor>) -> Result<Tensor> {
|
||||
let causal: Vec<f32> = (0..t)
|
||||
.flat_map(|i| (0..t).map(move |j| if j <= i { 0.0 } else { f32::NEG_INFINITY }))
|
||||
.collect();
|
||||
let causal = Tensor::from_vec(causal, (1, 1, t, t), &self.device)?;
|
||||
let Some(m) = attention_mask else {
|
||||
return Ok(causal);
|
||||
};
|
||||
let m = m.unsqueeze(1)?.unsqueeze(1)?;
|
||||
let keep = Tensor::zeros(m.dims(), DType::F32, &self.device)?;
|
||||
let drop = Tensor::full(f32::NEG_INFINITY, m.dims(), &self.device)?;
|
||||
causal.broadcast_add(&m.where_cond(&keep, &drop)?)
|
||||
}
|
||||
|
||||
/// LoRA layers with their PEFT names (`base_model.model.model.layers.{i}.{path}`).
|
||||
pub fn adapters(&self) -> impl Iterator<Item = (String, &LoraLinear)> {
|
||||
self.layers.iter().enumerate().flat_map(|(i, l)| {
|
||||
PROJS.iter().zip(&l.projs).filter_map(move |(p, proj)| {
|
||||
proj.lora()
|
||||
.map(|l| (format!("base_model.model.model.layers.{i}.{p}"), l))
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
pub fn trainable_vars(&self) -> Vec<Var> {
|
||||
self.adapters()
|
||||
.flat_map(|(_, l)| [l.a.clone(), l.b.clone()])
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
//! Gradient norm clipping and accumulation across micro-batches.
|
||||
|
||||
use hanzo_ml::backprop::GradStore;
|
||||
use hanzo_ml::{DType, Result, Var};
|
||||
|
||||
/// Scale every gradient in `grads` so their global L2 norm is at most
|
||||
/// `max_norm`. Returns the norm before clipping.
|
||||
pub fn clip_grad_norm(grads: &mut GradStore, vars: &[Var], max_norm: f64) -> Result<f64> {
|
||||
let mut sq = 0f64;
|
||||
for v in vars {
|
||||
if let Some(g) = grads.get(v) {
|
||||
sq += g
|
||||
.sqr()?
|
||||
.sum_all()?
|
||||
.to_dtype(DType::F32)?
|
||||
.to_scalar::<f32>()? as f64;
|
||||
}
|
||||
}
|
||||
let norm = sq.sqrt();
|
||||
if norm > max_norm {
|
||||
let scale = max_norm / (norm + 1e-6);
|
||||
for v in vars {
|
||||
let scaled = match grads.get(v) {
|
||||
Some(g) => (g * scale)?,
|
||||
None => continue,
|
||||
};
|
||||
grads.insert(v, scaled);
|
||||
}
|
||||
}
|
||||
Ok(norm)
|
||||
}
|
||||
|
||||
/// Add the gradients in `from` to those in `into`, per variable.
|
||||
pub fn accumulate(into: &mut GradStore, from: GradStore, vars: &[Var]) -> Result<()> {
|
||||
for v in vars {
|
||||
let Some(g) = from.get(v) else { continue };
|
||||
let sum = match into.get(v) {
|
||||
Some(acc) => (acc + g)?,
|
||||
None => g.clone(),
|
||||
};
|
||||
into.insert(v, sum);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use hanzo_ml::{Device, Tensor};
|
||||
|
||||
fn grads_for(vars: &[Var], k: f64) -> GradStore {
|
||||
// d/dv sum(k * v^2 / 2) = k * v
|
||||
let loss = vars
|
||||
.iter()
|
||||
.map(|v| (v.sqr().unwrap().sum_all().unwrap() * (k / 2.0)).unwrap())
|
||||
.reduce(|a, b| (a + b).unwrap())
|
||||
.unwrap();
|
||||
loss.backward().unwrap()
|
||||
}
|
||||
|
||||
fn vars() -> Vec<Var> {
|
||||
let dev = Device::Cpu;
|
||||
vec![
|
||||
Var::new(&[3f32, 0.0], &dev).unwrap(),
|
||||
Var::new(&[[0f32, 4.0]], &dev).unwrap(),
|
||||
]
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clips_to_max_norm_and_reports_pre_clip() {
|
||||
let vars = vars();
|
||||
let mut g = grads_for(&vars, 1.0);
|
||||
let before = clip_grad_norm(&mut g, &vars, 1.0).unwrap();
|
||||
assert!((before - 5.0).abs() < 1e-5);
|
||||
let after = clip_grad_norm(&mut g, &vars, 10.0).unwrap();
|
||||
assert!((after - 1.0).abs() < 1e-5, "{after}");
|
||||
let x: Vec<f32> = g.get(&vars[0]).unwrap().to_vec1().unwrap();
|
||||
assert!((x[0] - 0.6).abs() < 1e-5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn under_max_norm_is_untouched() {
|
||||
let vars = vars();
|
||||
let mut g = grads_for(&vars, 1.0);
|
||||
assert!((clip_grad_norm(&mut g, &vars, 6.0).unwrap() - 5.0).abs() < 1e-5);
|
||||
let x: Vec<f32> = g.get(&vars[0]).unwrap().to_vec1().unwrap();
|
||||
assert_eq!(x, vec![3.0, 0.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accumulate_sums_per_var() {
|
||||
let vars = vars();
|
||||
let mut acc = grads_for(&vars, 1.0);
|
||||
accumulate(&mut acc, grads_for(&vars, 2.0), &vars).unwrap();
|
||||
let x: Vec<f32> = acc.get(&vars[0]).unwrap().to_vec1().unwrap();
|
||||
assert_eq!(x, vec![9.0, 0.0]);
|
||||
let y: Vec<Vec<f32>> = acc.get(&vars[1]).unwrap().to_vec2().unwrap();
|
||||
assert_eq!(y, vec![vec![0.0, 12.0]]);
|
||||
let t = Tensor::zeros(2, DType::F32, &Device::Cpu).unwrap();
|
||||
assert!(acc.get(&t).is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
//! Next-token cross entropy.
|
||||
|
||||
use crate::IGNORE_INDEX;
|
||||
use hanzo_ml::{DType, Result, Tensor, D};
|
||||
use hanzo_nn::ops::log_softmax;
|
||||
|
||||
/// Positions scored per log_softmax call, so the f32 `[b, chunk, vocab]`
|
||||
/// temporary stays small whatever the sequence length.
|
||||
const CHUNK: usize = 512;
|
||||
|
||||
/// Mean negative log-likelihood of `labels[:, 1:]` under `logits[:, :-1]`,
|
||||
/// skipping positions labelled [`IGNORE_INDEX`], and the number of positions
|
||||
/// scored. Softmax runs in f32 whatever the logits dtype. With nothing to
|
||||
/// score the loss is zero.
|
||||
pub fn causal_lm_loss(logits: &Tensor, labels: &Tensor) -> Result<(Tensor, usize)> {
|
||||
let (_, t, _) = logits.dims3()?;
|
||||
let logits = logits.narrow(1, 0, t - 1)?;
|
||||
let labels = labels.narrow(1, 1, t - 1)?;
|
||||
let scored = labels.ne(IGNORE_INDEX)?;
|
||||
let count = scored.to_dtype(DType::F32)?.sum_all()?.to_scalar::<f32>()? as usize;
|
||||
if count == 0 {
|
||||
return Ok((logits.zeros_like()?.sum_all()?.to_dtype(DType::F32)?, 0));
|
||||
}
|
||||
let targets = scored
|
||||
.where_cond(&labels, &labels.zeros_like()?)?
|
||||
.unsqueeze(2)?;
|
||||
let weights = scored.to_dtype(DType::F32)?;
|
||||
let mut nll = None;
|
||||
for start in (0..t - 1).step_by(CHUNK) {
|
||||
let len = CHUNK.min(t - 1 - start);
|
||||
let logp = log_softmax(
|
||||
&logits.narrow(1, start, len)?.to_dtype(DType::F32)?,
|
||||
D::Minus1,
|
||||
)?;
|
||||
let picked = logp
|
||||
.gather(&targets.narrow(1, start, len)?, 2)?
|
||||
.squeeze(2)?;
|
||||
let part = (picked * weights.narrow(1, start, len)?)?.sum_all()?;
|
||||
nll = Some(match nll {
|
||||
None => part,
|
||||
Some(acc) => (acc + part)?,
|
||||
});
|
||||
}
|
||||
let nll = nll.expect("t > 1 with a scored position");
|
||||
Ok(((nll / -(count as f64))?, count))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use hanzo_ml::Device;
|
||||
|
||||
/// Straight loop over the same definition, in f64.
|
||||
fn reference(logits: &[Vec<Vec<f32>>], labels: &[Vec<i64>]) -> (f64, usize) {
|
||||
let (mut sum, mut n) = (0.0, 0);
|
||||
for (b, rows) in logits.iter().enumerate() {
|
||||
for t in 0..rows.len() - 1 {
|
||||
let y = labels[b][t + 1];
|
||||
if y == IGNORE_INDEX {
|
||||
continue;
|
||||
}
|
||||
let row: Vec<f64> = rows[t].iter().map(|&x| x as f64).collect();
|
||||
let max = row.iter().cloned().fold(f64::MIN, f64::max);
|
||||
let lse = max + row.iter().map(|x| (x - max).exp()).sum::<f64>().ln();
|
||||
sum += lse - row[y as usize];
|
||||
n += 1;
|
||||
}
|
||||
}
|
||||
(if n == 0 { 0.0 } else { sum / n as f64 }, n)
|
||||
}
|
||||
|
||||
fn example() -> (Vec<Vec<Vec<f32>>>, Vec<Vec<i64>>) {
|
||||
let mut v = 0.0f32;
|
||||
let logits = (0..2)
|
||||
.map(|_| {
|
||||
(0..4)
|
||||
.map(|_| {
|
||||
(0..5)
|
||||
.map(|_| {
|
||||
v = (v * 1.7 + 0.3) % 2.9;
|
||||
v - 1.4
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.collect();
|
||||
let labels = vec![
|
||||
vec![IGNORE_INDEX, 3, IGNORE_INDEX, 1],
|
||||
vec![2, IGNORE_INDEX, 4, 0],
|
||||
];
|
||||
(logits, labels)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matches_reference_with_ignored_positions() {
|
||||
let (logits, labels) = example();
|
||||
let (want, n) = reference(&logits, &labels);
|
||||
let dev = Device::Cpu;
|
||||
let l = Tensor::new(logits, &dev).unwrap();
|
||||
let y = Tensor::new(labels, &dev).unwrap();
|
||||
let (loss, count) = causal_lm_loss(&l, &y).unwrap();
|
||||
assert_eq!(count, n);
|
||||
assert_eq!(n, 4);
|
||||
let got = loss.to_scalar::<f32>().unwrap() as f64;
|
||||
assert!((got - want).abs() < 1e-5, "got {got} want {want}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chunks_agree_with_one_pass() {
|
||||
let dev = Device::Cpu;
|
||||
let t = 2 * CHUNK + 7;
|
||||
let l = Tensor::randn(0f32, 1.0, (1, t, 11), &dev).unwrap();
|
||||
let ids: Vec<i64> = (0..t).map(|i| ((i * 5 + 3) % 11) as i64).collect();
|
||||
let y = Tensor::from_vec(ids, (1, t), &dev).unwrap();
|
||||
let (loss, count) = causal_lm_loss(&l, &y).unwrap();
|
||||
assert_eq!(count, t - 1);
|
||||
let logits: Vec<Vec<Vec<f32>>> = l.to_vec3().unwrap();
|
||||
let labels: Vec<Vec<i64>> = y.to_vec2().unwrap();
|
||||
let (want, _) = reference(&logits, &labels);
|
||||
let got = loss.to_scalar::<f32>().unwrap() as f64;
|
||||
assert!((got - want).abs() < 1e-4, "got {got} want {want}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn nothing_scored_is_zero() {
|
||||
let dev = Device::Cpu;
|
||||
let l = Tensor::zeros((2, 4, 5), DType::F32, &dev).unwrap();
|
||||
let y = Tensor::full(IGNORE_INDEX, (2, 4), &dev).unwrap();
|
||||
let (loss, count) = causal_lm_loss(&l, &y).unwrap();
|
||||
assert_eq!(count, 0);
|
||||
assert_eq!(loss.to_scalar::<f32>().unwrap(), 0.0);
|
||||
}
|
||||
}
|
||||
+50
-2
@@ -1,2 +1,50 @@
|
||||
//! train
|
||||
pub fn run(_cfg: &crate::Config) -> anyhow::Result<()> { anyhow::bail!("train: not implemented") }
|
||||
//! Training entry point: pick device and dtype, load text, data and model,
|
||||
//! then hand off to SFT or GRPO.
|
||||
|
||||
mod clip;
|
||||
mod loss;
|
||||
mod sched;
|
||||
mod trainer;
|
||||
|
||||
pub use clip::{accumulate, clip_grad_norm};
|
||||
pub use loss::causal_lm_loss;
|
||||
pub use sched::lr_at;
|
||||
pub use trainer::sft;
|
||||
|
||||
use crate::config::{Precision, Rl};
|
||||
use crate::Config;
|
||||
use hanzo_ml::{DType, Device};
|
||||
|
||||
pub fn run(cfg: &Config) -> anyhow::Result<()> {
|
||||
let device = if cfg!(feature = "cuda") {
|
||||
Device::new_cuda(0)?
|
||||
} else if cfg!(feature = "metal") {
|
||||
Device::new_metal(0)?
|
||||
} else {
|
||||
Device::Cpu
|
||||
};
|
||||
let dtype = match cfg.bf16 {
|
||||
Precision::Auto if device.is_cpu() => DType::F32,
|
||||
Precision::Auto | Precision::On => DType::BF16,
|
||||
Precision::Off => DType::F32,
|
||||
};
|
||||
std::fs::create_dir_all(&cfg.output_dir)?;
|
||||
let text = crate::data::Text::load(cfg)?;
|
||||
let data = crate::data::load(cfg, &text)?;
|
||||
let model = crate::model::load(cfg, &device, dtype)?;
|
||||
let params: usize = model.trainable_vars().iter().map(|v| v.elem_count()).sum();
|
||||
tracing::info!(
|
||||
train = data.train.len(),
|
||||
val = data.val.len(),
|
||||
params,
|
||||
device = ?device,
|
||||
dtype = ?dtype,
|
||||
"run"
|
||||
);
|
||||
if cfg.rl == Some(Rl::Grpo) {
|
||||
let prompts = crate::data::load_prompts(cfg, &text)?;
|
||||
crate::grpo::run(cfg, model.as_ref(), &text, prompts)
|
||||
} else {
|
||||
sft(cfg, model.as_ref(), &data.train, &data.val, text.pad_id)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
//! Learning-rate schedule.
|
||||
|
||||
use crate::config::{Config, Scheduler};
|
||||
|
||||
/// Learning rate for optimizer step `step` of `total`: linear warmup over
|
||||
/// `cfg.warmup(total)` steps to `learning_rate`, then decay per `lr_scheduler`.
|
||||
pub fn lr_at(cfg: &Config, step: usize, total: usize) -> f64 {
|
||||
let peak = cfg.learning_rate;
|
||||
let warmup = cfg.warmup(total);
|
||||
if step < warmup {
|
||||
return peak * (step + 1) as f64 / warmup as f64;
|
||||
}
|
||||
let span = total.saturating_sub(warmup).max(1) as f64;
|
||||
let progress = ((step - warmup) as f64 / span).min(1.0);
|
||||
match cfg.lr_scheduler {
|
||||
Scheduler::Cosine => peak * 0.5 * (1.0 + (std::f64::consts::PI * progress).cos()),
|
||||
Scheduler::Linear => peak * (1.0 - progress),
|
||||
Scheduler::Constant => peak,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn cfg(scheduler: &str, warmup: usize) -> Config {
|
||||
serde_yaml::from_str(&format!(
|
||||
"base_model: x\noutput_dir: o\ndatasets: [{{path: p, type: completion}}]\n\
|
||||
learning_rate: 1.0\nlr_scheduler: {scheduler}\nwarmup_steps: {warmup}\n"
|
||||
))
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn close(a: f64, b: f64) -> bool {
|
||||
(a - b).abs() < 1e-9
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn warmup_then_cosine() {
|
||||
let c = cfg("cosine", 10);
|
||||
assert!(close(lr_at(&c, 0, 110), 0.1));
|
||||
assert!(close(lr_at(&c, 9, 110), 1.0));
|
||||
assert!(close(lr_at(&c, 10, 110), 1.0));
|
||||
assert!(close(lr_at(&c, 60, 110), 0.5));
|
||||
assert!(close(lr_at(&c, 110, 110), 0.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linear_and_constant() {
|
||||
let c = cfg("linear", 0);
|
||||
assert!(close(lr_at(&c, 0, 100), 1.0));
|
||||
assert!(close(lr_at(&c, 25, 100), 0.75));
|
||||
assert!(close(lr_at(&c, 100, 100), 0.0));
|
||||
let c = cfg("constant", 4);
|
||||
assert!(close(lr_at(&c, 1, 100), 0.5));
|
||||
assert!(close(lr_at(&c, 99, 100), 1.0));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,269 @@
|
||||
//! Supervised fine-tuning loop.
|
||||
|
||||
use super::clip::{accumulate, clip_grad_norm};
|
||||
use super::loss::causal_lm_loss;
|
||||
use super::sched::lr_at;
|
||||
use crate::data::collate;
|
||||
use crate::{CausalLm, Config, Example};
|
||||
use hanzo_ml::backprop::GradStore;
|
||||
use hanzo_nn::{AdamW, Optimizer, ParamsAdamW};
|
||||
use rand::rngs::StdRng;
|
||||
use rand::seq::SliceRandom;
|
||||
use rand::SeedableRng;
|
||||
use std::path::Path;
|
||||
use std::time::Instant;
|
||||
|
||||
/// Train `model` on `train` per `cfg`, evaluating on `val` at the end of each
|
||||
/// epoch and writing the adapter to `cfg.output_dir`.
|
||||
pub fn sft(
|
||||
cfg: &Config,
|
||||
model: &dyn CausalLm,
|
||||
train: &[Example],
|
||||
val: &[Example],
|
||||
pad_id: u32,
|
||||
) -> anyhow::Result<()> {
|
||||
fit(cfg, model, train, val, pad_id).map(|_| ())
|
||||
}
|
||||
|
||||
/// Micro-batch losses, tokens and last gradient norm since the previous log line.
|
||||
struct Window {
|
||||
losses: Vec<f32>,
|
||||
tokens: usize,
|
||||
norm: f64,
|
||||
since: Instant,
|
||||
}
|
||||
|
||||
impl Window {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
losses: Vec::new(),
|
||||
tokens: 0,
|
||||
norm: 0.0,
|
||||
since: Instant::now(),
|
||||
}
|
||||
}
|
||||
|
||||
fn mean(&self) -> f32 {
|
||||
self.losses.iter().sum::<f32>() / self.losses.len() as f32
|
||||
}
|
||||
|
||||
fn log(&mut self, step: usize, total: usize, lr: f64) -> f32 {
|
||||
let loss = self.mean();
|
||||
let tps = self.tokens as f64 / self.since.elapsed().as_secs_f64();
|
||||
tracing::info!(
|
||||
step,
|
||||
total,
|
||||
loss,
|
||||
lr,
|
||||
grad_norm = self.norm,
|
||||
tokens_per_s = tps,
|
||||
"train"
|
||||
);
|
||||
*self = Self::new();
|
||||
loss
|
||||
}
|
||||
}
|
||||
|
||||
/// The loop behind [`sft`]; returns the window-mean loss at every log line.
|
||||
fn fit(
|
||||
cfg: &Config,
|
||||
model: &dyn CausalLm,
|
||||
train: &[Example],
|
||||
val: &[Example],
|
||||
pad_id: u32,
|
||||
) -> anyhow::Result<Vec<f32>> {
|
||||
let vars = model.trainable_vars();
|
||||
let device = model.device();
|
||||
let out = Path::new(&cfg.output_dir);
|
||||
let micro = cfg.micro_batch_size;
|
||||
let accum = cfg.gradient_accumulation_steps;
|
||||
let per_epoch = train.len().div_ceil(micro * accum);
|
||||
let total = cfg
|
||||
.max_steps
|
||||
.unwrap_or(usize::MAX)
|
||||
.min(cfg.num_epochs * per_epoch);
|
||||
let mut opt = AdamW::new(
|
||||
vars.clone(),
|
||||
ParamsAdamW {
|
||||
lr: 0.0,
|
||||
beta1: 0.9,
|
||||
beta2: 0.999,
|
||||
eps: 1e-8,
|
||||
weight_decay: cfg.weight_decay,
|
||||
},
|
||||
)?;
|
||||
let mut step = 0;
|
||||
let mut lr = 0.0;
|
||||
let mut window = Window::new();
|
||||
let mut logged = Vec::new();
|
||||
let mut done = total == 0;
|
||||
for epoch in 0..cfg.num_epochs {
|
||||
if done {
|
||||
break;
|
||||
}
|
||||
let mut order: Vec<usize> = (0..train.len()).collect();
|
||||
order.shuffle(&mut StdRng::seed_from_u64(cfg.seed + epoch as u64));
|
||||
let chunks: Vec<&[usize]> = order.chunks(micro).collect();
|
||||
let mut grads: Option<GradStore> = None;
|
||||
let mut pending = 0;
|
||||
for (i, ids) in chunks.iter().enumerate() {
|
||||
let examples: Vec<Example> = ids.iter().map(|&j| train[j].clone()).collect();
|
||||
let batch = collate(&examples, pad_id, device)?;
|
||||
let logits = model.forward(&batch.input_ids, batch.attention_mask.as_ref())?;
|
||||
let (loss, scored) = causal_lm_loss(&logits, &batch.labels)?;
|
||||
window.losses.push(loss.to_scalar::<f32>()?);
|
||||
window.tokens += batch.input_ids.elem_count();
|
||||
if scored > 0 {
|
||||
let g = (loss / accum as f64)?.backward()?;
|
||||
match grads.as_mut() {
|
||||
Some(acc) => accumulate(acc, g, &vars)?,
|
||||
None => grads = Some(g),
|
||||
}
|
||||
}
|
||||
pending += 1;
|
||||
if pending < accum && i + 1 < chunks.len() {
|
||||
continue;
|
||||
}
|
||||
if let Some(mut g) = grads.take() {
|
||||
window.norm = clip_grad_norm(&mut g, &vars, cfg.max_grad_norm)?;
|
||||
lr = lr_at(cfg, step, total);
|
||||
opt.set_learning_rate(lr);
|
||||
opt.step(&g)?;
|
||||
}
|
||||
pending = 0;
|
||||
step += 1;
|
||||
if step % cfg.logging_steps == 0 {
|
||||
logged.push(window.log(step, total, lr));
|
||||
}
|
||||
if cfg.save_steps.is_some_and(|s| step % s == 0) {
|
||||
model.save_adapter(&out.join(format!("checkpoint-{step}")))?;
|
||||
}
|
||||
if step >= total {
|
||||
done = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if !val.is_empty() {
|
||||
let loss = evaluate(model, val, pad_id, micro)?;
|
||||
tracing::info!(epoch, step, loss, "eval");
|
||||
}
|
||||
}
|
||||
if !window.losses.is_empty() {
|
||||
logged.push(window.log(step, total, lr));
|
||||
}
|
||||
model.save_adapter(out)?;
|
||||
Ok(logged)
|
||||
}
|
||||
|
||||
/// Token-weighted mean loss over `val`, no parameter update.
|
||||
fn evaluate(
|
||||
model: &dyn CausalLm,
|
||||
val: &[Example],
|
||||
pad_id: u32,
|
||||
micro: usize,
|
||||
) -> anyhow::Result<f32> {
|
||||
let (mut nll, mut n) = (0f64, 0usize);
|
||||
for examples in val.chunks(micro) {
|
||||
let batch = collate(examples, pad_id, model.device())?;
|
||||
let logits = model.forward(&batch.input_ids, batch.attention_mask.as_ref())?;
|
||||
let (loss, scored) = causal_lm_loss(&logits, &batch.labels)?;
|
||||
nll += loss.to_scalar::<f32>()? as f64 * scored as f64;
|
||||
n += scored;
|
||||
}
|
||||
Ok(if n == 0 { 0.0 } else { (nll / n as f64) as f32 })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use hanzo_ml::{DType, Device, Result, Tensor, Var};
|
||||
|
||||
/// Embedding followed by one linear map: exactly a bigram model.
|
||||
struct Bigram {
|
||||
emb: Var,
|
||||
out: Var,
|
||||
dev: Device,
|
||||
}
|
||||
|
||||
const VOCAB: usize = 16;
|
||||
|
||||
impl Bigram {
|
||||
fn new(dev: &Device) -> Self {
|
||||
Self {
|
||||
emb: Var::randn(0f32, 0.5, (VOCAB, 8), dev).unwrap(),
|
||||
out: Var::randn(0f32, 0.5, (8, VOCAB), dev).unwrap(),
|
||||
dev: dev.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CausalLm for Bigram {
|
||||
fn forward(&self, ids: &Tensor, _mask: Option<&Tensor>) -> Result<Tensor> {
|
||||
let (b, t) = ids.dims2()?;
|
||||
let h = self.emb.index_select(&ids.flatten_all()?, 0)?;
|
||||
h.matmul(self.out.as_tensor())?.reshape((b, t, VOCAB))
|
||||
}
|
||||
fn trainable_vars(&self) -> Vec<Var> {
|
||||
vec![self.emb.clone(), self.out.clone()]
|
||||
}
|
||||
fn device(&self) -> &Device {
|
||||
&self.dev
|
||||
}
|
||||
fn dtype(&self) -> DType {
|
||||
DType::F32
|
||||
}
|
||||
fn eos_token_id(&self) -> Option<u32> {
|
||||
None
|
||||
}
|
||||
fn vocab_size(&self) -> usize {
|
||||
VOCAB
|
||||
}
|
||||
fn save_adapter(&self, dir: &Path) -> anyhow::Result<()> {
|
||||
std::fs::create_dir_all(dir)?;
|
||||
std::fs::write(dir.join("adapter_model.safetensors"), b"")?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Sequences that count upwards mod `VOCAB` from a varying start.
|
||||
fn counting(n: usize, len: usize) -> Vec<Example> {
|
||||
(0..n)
|
||||
.map(|i| {
|
||||
let input_ids: Vec<u32> = (0..len).map(|k| ((i * 7 + k) % VOCAB) as u32).collect();
|
||||
let labels = input_ids.iter().map(|&x| x as i64).collect();
|
||||
Example { input_ids, labels }
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn learns_the_counting_rule_and_saves() {
|
||||
let dir = std::env::temp_dir().join(format!("gym-sft-{}", std::process::id()));
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
let cfg: Config = serde_yaml::from_str(&format!(
|
||||
"base_model: x\noutput_dir: {}\ndatasets: [{{path: p, type: completion}}]\n\
|
||||
micro_batch_size: 4\ngradient_accumulation_steps: 2\nnum_epochs: 10\nmax_steps: 40\n\
|
||||
learning_rate: 0.05\nlr_scheduler: constant\nlogging_steps: 5\nsave_steps: 10\n\
|
||||
max_grad_norm: 1.0\nseed: 7\n",
|
||||
dir.display()
|
||||
))
|
||||
.unwrap();
|
||||
let dev = Device::Cpu;
|
||||
let model = Bigram::new(&dev);
|
||||
let train = counting(64, 8);
|
||||
let val = counting(8, 8);
|
||||
let windows = fit(&cfg, &model, &train, &val, 0).unwrap();
|
||||
assert_eq!(windows.len(), 8, "{windows:?}");
|
||||
let (first, last) = (windows[0], *windows.last().unwrap());
|
||||
assert!(last < 0.5 * first, "first {first} last {last}");
|
||||
for step in [10, 20, 30, 40] {
|
||||
let marker = dir
|
||||
.join(format!("checkpoint-{step}"))
|
||||
.join("adapter_model.safetensors");
|
||||
assert!(marker.is_file(), "{}", marker.display());
|
||||
}
|
||||
assert!(!dir.join("checkpoint-50").exists());
|
||||
assert!(dir.join("adapter_model.safetensors").is_file());
|
||||
std::fs::remove_dir_all(&dir).unwrap();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user