diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7f859e48..a24c0d50 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -49,6 +49,21 @@ jobs: - name: Build run: cargo build --release + training-supervisor: + name: Training Relaunch + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Install shellcheck + run: sudo apt-get update && sudo apt-get install -y shellcheck + + - name: Check and test relaunch supervisor + run: | + shellcheck -x hermes-train/scripts/relaunch.sh hermes-train/scripts/relaunch_test.sh + python3 -m py_compile hermes-train/scripts/wandb_tail.py + hermes-train/scripts/relaunch_test.sh + # Security audit - check for known vulnerabilities audit: name: Cargo Audit diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 9e9a7e28..b96881ef 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -99,7 +99,8 @@ jobs: - name: Update pyproject.toml versions run: | sed -i 's/^version = ".*"/version = "${{ steps.bump.outputs.new_version }}"/' hermes-client-python/pyproject.toml - echo "Updated pyproject.toml to version ${{ steps.bump.outputs.new_version }}" + sed -i 's/^version = ".*"/version = "${{ steps.bump.outputs.new_version }}"/' hermes-mal-python/pyproject.toml + echo "Updated Python packages to version ${{ steps.bump.outputs.new_version }}" - name: Update TypeScript client version run: | @@ -112,7 +113,7 @@ jobs: - name: Commit and push run: | - git add Cargo.toml Cargo.lock hermes-client-python/pyproject.toml hermes-client-typescript/package.json + git add Cargo.toml Cargo.lock hermes-client-python/pyproject.toml hermes-mal-python/pyproject.toml hermes-client-typescript/package.json git commit -m "chore: bump version to ${{ steps.bump.outputs.new_version }}" git push diff --git a/CLAUDE.md b/CLAUDE.md index ddb80f2e..be7e86cd 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -70,8 +70,9 @@ cargo clippy --all-targets --all-features -- -D warnings # Build WASM (requires Homebrew LLVM for zstd cross-compilation) cd hermes-wasm && bash build.sh -# Build Python wheel -cd hermes-core-python && maturin build --release +# Build Python packages +cd hermes-client-python && uv build +cd hermes-mal-python && maturin build --release # hermes-train (LLM training) tests cargo test -p hermes-train diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 602d2a71..eb557007 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -42,7 +42,8 @@ cargo test --all-features | `cargo fmt --all` | Format all Rust code | | `cargo clippy --all-targets --all-features -- -D warnings` | Run lints | | `cd hermes-wasm && wasm-pack build --release --target web` | Build WASM package | -| `cd hermes-core-python && maturin build --release` | Build Python wheel | +| `cd hermes-client-python && uv build` | Build Python gRPC client | +| `cd hermes-mal-python && maturin build --release` | Build MAL Python binding | | `pre-commit run --all-files` | Run all pre-commit hooks | ## Project Structure @@ -54,11 +55,16 @@ cargo test --all-features | **hermes-tool** | CLI for index management and data processing pipelines | | **hermes-wasm** | WebAssembly bindings for browser-based search | | **hermes-web** | Vue.js web UI | -| **hermes-llm** | LLM training framework built on Candle ML | +| **hermes-mal** | Model Architecture Language parser and well-known definitions | +| **hermes-mal-python** | Thin PyO3 binding around the shared `hermes-mal` parser | +| **hermes-llm** | Burn-based shared model, inference, generation, and accelerator kernels | +| **hermes-train** | Autodiff training for the same `hermes-llm` model and safetensors checkpoints | | **hermes-proto** | Protocol Buffer definitions for gRPC services | | **hermes-client-python** | Python gRPC client library | -For a deeper look at the core architecture, see `CLAUDE.md`. +For a deeper look at the core architecture, see `CLAUDE.md`. The shared LLM +stack is mapped in `docs/llm-code-map.md`; temporary GPU forks and their +upstream exit criteria live in `docs/forked-dependencies.md`. ## Submitting Pull Requests diff --git a/Cargo.lock b/Cargo.lock index a8ff62cf..20d1306f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3240,7 +3240,7 @@ dependencies = [ ] [[package]] -name = "hermes-mal-py" +name = "hermes-mal-python" version = "1.8.68" dependencies = [ "hermes-mal", diff --git a/Cargo.toml b/Cargo.toml index 15aed06c..3dbfce99 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,7 +3,7 @@ members = [ "hermes-core", "hermes-llm", "hermes-mal", - "hermes-mal-py", + "hermes-mal-python", "hermes-train", "hermes-server", "hermes-tool", @@ -106,14 +106,13 @@ hashbrown = "0.17" memmap2 = "0.9" uuid = { version = "1.19", features = ["v4"] } -# Keep Burn's CubeK dependency on the same strided-attention fix used directly -# by hermes-llm. +# Forward softmax-LSE emission used by Hermes attention training. Upstream: +# https://github.com/tracel-ai/cubek/pull/428 [patch."https://github.com/tracel-ai/cubek"] cubek = { git = "https://github.com/ppodolsky/cubek", rev = "b4fe9788a774c0f5e5af26122776b5bf4c8ee94d" } -# Burn main still points at upstream CubeCL. Use our fork (allocation-retry -# fix + asynchronous cuBLAS BF16 GEMM server dispatch) for Burn's direct -# CubeCL dependencies as well as hermes-llm's. +# Allocation retry and asynchronous cuBLASLt BF16 GEMM dispatch. Consolidated upstream: +# https://github.com/tracel-ai/cubecl/pull/1440 [patch."https://github.com/tracel-ai/cubecl"] cubecl = { git = "https://github.com/ppodolsky/cubecl", rev = "bda6a68d86df41f468cac78d713f1563869daa8e" } cubecl-common = { git = "https://github.com/ppodolsky/cubecl", rev = "bda6a68d86df41f468cac78d713f1563869daa8e" } @@ -121,7 +120,8 @@ cubecl-zspace = { git = "https://github.com/ppodolsky/cubecl", rev = "bda6a68d86 # Route every Burn crate to the pinned fork: upstream `bd6e8fa2f` plus the # BF16 native-GEMM matmul dispatch (autotune-arbitrated against fused CubeK) -# and the upstream `f31e7513a` foreign-stream drop-ordering backport. +# and the now-upstream foreign-stream drop-ordering fix. Native GEMM upstream: +# https://github.com/tracel-ai/burn/pull/5190 # Measured on A100: +15.6%/+14.8% training tokens/s at batch 16/20 with # identical losses (docs/cublas-gemm-dispatch.md). [patch."https://github.com/tracel-ai/burn"] diff --git a/README.md b/README.md index 6cf6a0c8..48c788c2 100644 --- a/README.md +++ b/README.md @@ -261,8 +261,11 @@ cargo build --release # Build WASM (requires Homebrew LLVM on macOS for zstd cross-compilation) cd hermes-wasm && bash build.sh -# Build Python wheel -cd hermes-core-python && maturin build --release +# Build the Python gRPC client +cd hermes-client-python && uv build + +# Build the MAL Python binding +cd hermes-mal-python && maturin build --release ``` Alternatively you may build everything in docker via `docker compose`. @@ -278,6 +281,10 @@ Examples: cargo test --all-features ``` +LLM contributors should start with the [inference and training code map](docs/llm-code-map.md). +Temporary GPU dependency forks and their upstream removal criteria are listed +in [the fork register](docs/forked-dependencies.md). + ### Linting ```bash diff --git a/docs/bf16-residual-stream.md b/docs/bf16-residual-stream.md index 73410302..8bebe6c4 100644 --- a/docs/bf16-residual-stream.md +++ b/docs/bf16-residual-stream.md @@ -34,13 +34,11 @@ Boundary contracts this forced: fails at runtime, which the plain-CUDA suite never exercises. The end-to-end gate for that class is `training_fusion_hybrid_bf16_stream_loss_and_gradients_are_finite`. -- The selective scan's non-segmented paths (decode, prefill without saved - states, and the small-problem `checkpoint_interval == 1` full-state path) - have no BF16 kernels: the dispatcher normalizes those to FP32 at the - boundary and returns BF16 to the caller. The training hot path - (segment-parallel forward/backward) is BF16-native and unaffected. The - e2e gate caught exactly this on a small hybrid model before any benchmark - ran. +- The selective scan's inference paths (decode and prefill without saved + states) have no BF16 kernels: the dispatcher normalizes those to FP32 at the + boundary and returns BF16 to the caller. The single checkpointed training + path is BF16-native. The e2e gate caught this boundary on a small hybrid + model before any benchmark ran. - The same gate then surfaced two attention-backward issues at small-model shapes: the fused probabilities kernel read a BF16 `correction` tensor as raw FP32 (fixed by pinning the correction product to FP32 + loud dtype diff --git a/docs/forked-dependencies.md b/docs/forked-dependencies.md new file mode 100644 index 00000000..a05919e5 --- /dev/null +++ b/docs/forked-dependencies.md @@ -0,0 +1,25 @@ +# Forked dependency register + +Hermes vendors no source repositories. It temporarily pins three Git forks +for GPU changes that are not yet available from upstream. Every pin below has +an upstream submission and an explicit removal condition. + +| Dependency | Hermes pin | Upstream status | Remove the fork when | +| ---------- | --------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| CubeK | `ppodolsky/cubek@b4fe9788` | [consolidated forward softmax-LSE and CubeCL compatibility PR #428](https://github.com/tracel-ai/cubek/pull/428) | The LSE API is merged and Burn/Hermes can pin an upstream revision containing it. Do not include fork commit `ee578923`: its tensor-core backward was measured and rejected. | +| CubeCL | `ppodolsky/cubecl@bda6a68d` | [consolidated allocation-retry and cuBLASLt PR #1440](https://github.com/tracel-ai/cubecl/pull/1440) | Both runtime changes are merged, or Hermes drops the corresponding feature. | +| Burn | `ppodolsky/burn@e2fe6651` | [integration PR #5190](https://github.com/tracel-ai/burn/pull/5190), with the [pre-PR discussion #5189](https://github.com/tracel-ai/burn/issues/5189) required by Burn's contribution guide. The foreign-stream ordering prerequisite already merged upstream as Burn PR #5166. | The integration is merged against upstream CubeCL/CubeK revisions and Hermes passes its CUDA parity and throughput gates on that upstream stack. | + +The Apache Arrow `object_store` Git revision is an upstream security-fix pin, +not a fork. It can return to crates.io after a release containing the pinned +quick-xml update is available. + +## Update procedure + +1. Check the linked submissions before advancing any fork revision. +2. Rebase only the still-required commits onto the dependency's current main; + exclude changes already merged or recorded as rejected experiments. +3. Run CPU tests and clippy for `hermes-mal`, `hermes-llm`, and + `hermes-train`, then the CUDA parity and end-to-end training gates. +4. Replace fork URLs with upstream URLs as soon as the last required commit is + upstream. Keep this register until the lockfile no longer contains the fork. diff --git a/docs/llm-code-map.md b/docs/llm-code-map.md new file mode 100644 index 00000000..2f71eb90 --- /dev/null +++ b/docs/llm-code-map.md @@ -0,0 +1,32 @@ +# LLM inference and training code map + +Hermes uses one MAL model definition and one Burn `Transformer` implementation +for training, generation, and retrieval. There is no alternate PyTorch or +Candle model stack. + +| Area | Entry points | Responsibility | +| ------------------------ | ------------------------------------------------------------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Architecture | `hermes-mal/src/lib.rs`, `hermes-mal/src/mal.pest`, `hermes-mal/well-known/` | Parse composable MAL definitions, resolve references, and expose the serializable `ModelDef`. `parse_mal` requires exactly one model; tools that select among several use `parse_mal_full`. | +| Model assembly | `hermes-llm/src/model/transformer.rs`, `block.rs` | Validate dimensions and numeric settings, construct homogeneous or patterned attention/Mamba layers, and own the shared forward/loss/stateful paths. | +| Attention | `hermes-llm/src/model/attention.rs`, `fused_attention.rs`, `cube_attention.rs` | Grouped-query projection, RoPE, causal/window masks, KV caching, backend selection, and the CUDA training backward. | +| Mamba | `hermes-llm/src/model/mamba.rs`, `model/scan/`, `model/conv.rs` | Stateful selective-SSM mixing, CPU correctness references, autodiff nodes, and checkpointed CUDA/Metal kernels. | +| Loss and numerics | `hermes-llm/src/model/linear_cross_entropy.rs`, `norm.rs`, `matmul.rs` | Chunked vocabulary loss, normalization, precision policy, and native/fused matmul entry points. | +| Generation and artifacts | `hermes-llm/src/generate.rs`, `tokenizer.rs`, `remote.rs`, `model/weights.rs` | Tokenization/sampling, local or cached remote artifact resolution, and safetensors loading/saving. | +| Corpus pipeline | `hermes-train/src/data.rs` | Stream text/JSONL/Zstandard data, batch-tokenize, EOS-pack without padding, and bounded deterministic shuffle. | +| Checkpoints | `hermes-train/src/checkpoint.rs` | Atomically publish model/optimizer/training state and restore Burn parameter IDs for exact resume. | +| Optimization loop | `hermes-train/src/main.rs`, `muon.rs` | CLI validation, curriculum/epoch position, gradient accumulation/clipping, schedule, Muon + AdamW steps, and metrics. | + +## Validation layers + +- `cargo test -p hermes-mal -p hermes-llm -p hermes-train` covers the parser, + CPU model paths, streaming corpus logic, optimizer behavior, and checkpoint + resume. +- `cargo clippy -p hermes-mal -p hermes-llm -p hermes-train --all-targets -- -D warnings` + is the required host lint gate. +- CUDA and Metal kernel parity tests compare accelerator results with the + tensor-operation references. Performance changes additionally require the + end-to-end loss/gradient and steady-state throughput gates documented in the + relevant file under `docs/`. + +Temporary dependency forks and their upstream exit criteria are tracked in +`docs/forked-dependencies.md`. diff --git a/docs/mal-single-parser.md b/docs/mal-single-parser.md index 8731c619..5119b39f 100644 --- a/docs/mal-single-parser.md +++ b/docs/mal-single-parser.md @@ -3,17 +3,17 @@ The Model Architecture Language has one parser and one schema: the Rust `hermes-mal` crate. -| Consumer | Integration | -| -------------- | ----------------------------------------------------------- | -| `hermes-llm` | Re-exports `hermes-mal` as `hermes_llm::mal` | -| `hermes-train` | Uses the `hermes-llm` re-export and shared `ModelDef` | -| Python tools | Optional `hermes-mal-py` PyO3 wrapper around the same crate | +| Consumer | Integration | +| -------------- | --------------------------------------------------------------- | +| `hermes-llm` | Re-exports `hermes-mal` as `hermes_llm::mal` | +| `hermes-train` | Uses the `hermes-llm` re-export and shared `ModelDef` | +| Python tools | Optional `hermes-mal-python` PyO3 wrapper around the same crate | The grammar, AST, serde representation, embedded well-known models, and computed properties live under `hermes-mal/`. Neither training nor inference has a parallel parser or copied configuration type. -`hermes-mal-py` is a general binding for external Python tools. It exposes +`hermes-mal-python` is a general binding for external Python tools. It exposes `parse_mal(source) -> JSON`; it is not part of the training path. When changing MAL, update the grammar/schema and Rust tests in `hermes-mal`, then diff --git a/docs/segmented-selective-scan.md b/docs/segmented-selective-scan.md index 5170ae6f..e7662f1e 100644 --- a/docs/segmented-selective-scan.md +++ b/docs/segmented-selective-scan.md @@ -1,44 +1,36 @@ # Segment-parallel selective scan -The Mamba selective scan is a linear recurrence `h_t = α_t·h_{t-1} + β_t` with -a diagonal state transition, so transitions compose associatively: a run of -timesteps collapses to one `(decay, partial)` pair. The CUDA/Metal training -kernels exploit this by cutting each sequence at the existing checkpoint -boundaries (`CHECKPOINTED_SCAN_INTERVAL = 32`) and running every segment -concurrently instead of walking the whole sequence serially per scan. - -## Forward (training, `save_states` path) - -1. `selective_scan_forward_segment_partials` — one thread per - `(batch, segment, channel)`: scans its segment from a zero state, emitting - the segment's `partial` state and `decay = exp(A·Σdt)` (the exact product - of per-step decays for a diagonal `A`). -2. `selective_scan_forward_segment_carry` — one thread per - `(batch, channel, n)`: serially folds the per-segment transitions into the - state _entering_ each segment. Touches `segments × state_dim` values per - scan — microseconds. -3. `selective_scan_forward_segment_apply` — same grid as (1): re-runs each - segment from its stitched entering state, writing outputs, per-segment - checkpoints (the same tensor backward always consumed), and the final - state. - -## Backward - -The adjoint recurrence `adj_{t-1} = (adj_t + dy_t·C_t)·α_t` is linear in -`adj`, so the same trick applies right-to-left: -`selective_scan_backward_segment_partials` + `_carry` stitch the adjoint -across segments, then `selective_scan_backward_segmented` launches one block -per `(batch, channel-tile, segment)` — the fused gradient kernel with the -sequence walk replaced by a single segment and the incoming adjoint read from -the carry. Gradients for `A`/`D` accumulate through the pre-existing atomics. - -Inference decode, prefill without saved states, and the small-batch -full-state path (`checkpoint_interval == 1`) keep the serial/parallel/step -kernels. - -The segment kernels run with CubeCL fast-math (`ReducedPrecision | NotNaN | -NotInf`), which lowers `exp` to the hardware `__expf`; state loops are -unrolled so the 16-wide recurrent state stays in registers. +The Mamba selective scan is a linear recurrence `h_t = α_t·h_{t-1} + β_t` +with a diagonal state transition. CUDA and Metal use one checkpointed sweep +per direction for training; the earlier partial/carry/apply chain described in +the benchmark ledger below has been removed. + +## Current GPU training path + +- `selective_scan_forward_swept` launches one block per `(batch, channel +tile)` and walks checkpoint segments left-to-right with the recurrent state + in registers. It reads each sequence input once and writes a state checkpoint + every `CHECKPOINTED_SCAN_INTERVAL = 32` tokens. +- `selective_scan_backward_segmented` launches the same block ownership in + reverse. It carries the adjoint in registers, reconstructs each segment from + its entering checkpoint, and pre-reduces shared `B`/`C` gradients across the + channel tile before global atomics. +- Training sequence I/O may be BF16, while recurrent state, checkpoints, and + parameter-gradient accumulation stay FP32. State widths are validated as + powers of two in `4..=16` at the GPU dispatch boundary. + +Inference decode and prefill without saved states retain their dedicated FP32 +paths. There is no interval-1 training implementation. + +The sweep kernels use CubeCL fast-math (`ReducedPrecision | NotNaN | NotInf`), +which lowers `exp` to the hardware `__expf`; state loops are unrolled so the +16-wide recurrent state stays in registers. + +## Historical optimization ledger + +The sections below preserve the measured progression and rejected experiments. +Names such as partials/carry/apply describe superseded implementations unless +a later section explicitly says the change landed. ## Measured (A100 40GB, retriever-100m, B16/T1024/ga8, steps 5–8, 2026-07-16) diff --git a/hermes-core/src/index/helpers.rs b/hermes-core/src/index/helpers.rs index 874a5aa8..f197a567 100644 --- a/hermes-core/src/index/helpers.rs +++ b/hermes-core/src/index/helpers.rs @@ -1,7 +1,7 @@ //! Indexing helper functions //! //! This module provides high-level helper functions for creating indexes -//! and indexing documents, used by hermes-tool, hermes-server, and hermes-core-python. +//! and indexing documents, used by `hermes-tool` and `hermes-server`. use std::io::BufRead; use std::path::Path; diff --git a/hermes-llm/src/model/attention.rs b/hermes-llm/src/model/attention.rs index bc2ed4c5..5387bf19 100644 --- a/hermes-llm/src/model/attention.rs +++ b/hermes-llm/src/model/attention.rs @@ -153,9 +153,16 @@ impl MultiHeadAttention { (q, k, v) } - /// Scaled-dot-product attention over K/V ([B, H, T, hd]) for queries at - /// global positions `q_start..q_start+S`. Applies causal + window masking. - fn sdpa(&self, q: Tensor<4>, k: Tensor<4>, v: Tensor<4>, q_start: usize) -> Tensor<4> { + /// Scaled-dot-product attention over K/V ([B, H, T, hd]) for queries and + /// keys at their respective global offsets. Applies causal + window masking. + fn sdpa( + &self, + q: Tensor<4>, + k: Tensor<4>, + v: Tensor<4>, + q_start: usize, + key_start: usize, + ) -> Tensor<4> { let device = q.device(); let [_, _, seq_q, _] = q.dims(); let total = k.dims()[2]; @@ -165,7 +172,7 @@ impl MultiHeadAttention { let mut scores = q .matmul(k.transpose()) .div_scalar((self.head_dim as f32).sqrt()); - if let Some(mask) = self.build_mask(seq_q, total, q_start, &device) { + if let Some(mask) = self.build_mask(seq_q, total, q_start, key_start, &device) { scores = scores.mask_fill(mask, f32::NEG_INFINITY); } let weights = self.attention_dropout.forward(softmax(scores, 3)); @@ -174,14 +181,14 @@ impl MultiHeadAttention { // Full-sequence attention has a custom fused backward on CUDA. Cached, // offset, and sliding-window attention retain Burn's mask-capable path. - let fused_causal = - self.causal && self.window_size.is_none() && q_start == 0 && seq_q == total; + let aligned_full_sequence = q_start == key_start && seq_q == total; + let fused_causal = self.causal && self.window_size.is_none() && aligned_full_sequence; let mask = if fused_causal { None } else { - self.build_mask(seq_q, total, q_start, &device) + self.build_mask(seq_q, total, q_start, key_start, &device) }; - if mask.is_none() && q_start == 0 && seq_q == total { + if mask.is_none() && aligned_full_sequence { fused_attention(q, k, v, fused_causal) } else { attention( @@ -200,16 +207,18 @@ impl MultiHeadAttention { } /// Bool mask [1, 1, S, T] (true = blocked), or None when nothing is masked. - /// Position `q_start + i` (query row i) attends to key j in `0..T`. + /// Position `q_start + i` (query row i) attends to position + /// `key_start + j` (key column j). fn build_mask( &self, seq_q: usize, total: usize, q_start: usize, + key_start: usize, device: &Device, ) -> Option> { - let needs = - (self.causal && total > q_start + 1) || (self.window_size.is_some() && total > 1); + let last_key = key_start + total.saturating_sub(1); + let needs = (self.causal && last_key > q_start) || self.window_size.is_some(); if !needs { return None; } @@ -217,10 +226,11 @@ impl MultiHeadAttention { for i in 0..seq_q { let gi = q_start + i; for j in 0..total { - let causal_block = self.causal && j > gi; + let gj = key_start + j; + let causal_block = self.causal && gj > gi; let window_block = self .window_size - .map(|w| gi.saturating_sub(j) > w || j.saturating_sub(gi) > w) + .map(|window| gi.abs_diff(gj) > window) .unwrap_or(false); mask[i * total + j] = causal_block || window_block; } @@ -238,7 +248,7 @@ impl MultiHeadAttention { /// Full (stateless) attention over the given sequence. pub fn forward(&self, x: Tensor<3>, rope: &RotaryEncoding, start_pos: usize) -> Tensor<3> { let (q, k, v) = self.project_qkv(x, rope, start_pos); - let out = self.sdpa(q, k, v, start_pos); + let out = self.sdpa(q, k, v, start_pos, start_pos); // The output projection joins the training residual stream directly // in the matmul compute dtype; decode (`forward_cached`) keeps the // FP32 promotion for its FP32 stream. @@ -294,7 +304,7 @@ impl MultiHeadAttention { cache.v = Some(v_cache); cache.len = end; - let out = self.sdpa(q, k, v, start_pos); + let out = self.sdpa(q, k, v, start_pos, 0); linear(&self.o_proj, self.merge_heads(out)) } } diff --git a/hermes-llm/src/model/mod.rs b/hermes-llm/src/model/mod.rs index a25acf72..9cbe9721 100644 --- a/hermes-llm/src/model/mod.rs +++ b/hermes-llm/src/model/mod.rs @@ -186,6 +186,31 @@ mod tests { assert!(max_abs_diff(at_zero, at_offset) < 1e-6); } + #[test] + fn test_attention_stateless_offset_preserves_relative_mask() { + let mut config = get_builtin_model("tiny").unwrap(); + config.hidden_size = 16; + config.max_seq_len = 16; + config.block.attention.num_heads = Some(4); + config.block.attention.num_kv_heads = Some(2); + config.block.attention.head_dim = Some(4); + config.block.attention.position_encoding = PositionEncoding::None; + config.block.attention.causal = true; + config.block.attention.window_size = Some(2); + let device = Device::ndarray(); + device.seed(10); + let attention = MultiHeadAttention::new(&config, &config.block, &device); + let rope = RotaryEncodingConfig::new(16, 4).init(&device); + let data: Vec = (0..4 * config.hidden_size) + .map(|i| (i as f32 * 0.089).sin()) + .collect(); + let x = Tensor::from_data(TensorData::new(data, [1, 4, config.hidden_size]), &device); + + let at_zero = attention.forward(x.clone(), &rope, 0); + let at_offset = attention.forward(x, &rope, 5); + assert!(max_abs_diff(at_zero, at_offset) < 1e-6); + } + #[test] fn test_mamba_stateful_matches_stateless() { let mut config = get_builtin_model("hybrid-tiny").unwrap(); @@ -247,6 +272,51 @@ mod tests { assert_eq!(state.pos(), 6); } + #[test] + fn test_transformer_rejects_zero_sized_ssm_dimensions() { + let device = Device::ndarray(); + let mut cases = Vec::new(); + + let mut config = hybrid_test_config(); + config.pattern.as_mut().unwrap()[0] + .ssm + .as_mut() + .unwrap() + .expand = 0; + cases.push(("expand", config)); + + let mut config = hybrid_test_config(); + config.pattern.as_mut().unwrap()[0] + .ssm + .as_mut() + .unwrap() + .state_dim = 0; + cases.push(("state_dim", config)); + + let mut config = hybrid_test_config(); + config.pattern.as_mut().unwrap()[0] + .ssm + .as_mut() + .unwrap() + .conv_kernel = 0; + cases.push(("conv_kernel", config)); + + let mut config = hybrid_test_config(); + config.pattern.as_mut().unwrap()[0] + .ssm + .as_mut() + .unwrap() + .dt_rank = Some(0); + cases.push(("dt_rank", config)); + + for (field, config) in cases { + let err = Transformer::new(&config, &device) + .err() + .unwrap_or_else(|| panic!("zero {field} was accepted")); + assert!(err.to_string().contains(field), "{field}: {err}"); + } + } + #[cfg(all(feature = "cuda", target_os = "linux"))] #[test] fn test_cuda_prepared_inference_matches_ephemeral_weight_casts() { diff --git a/hermes-llm/src/model/transformer.rs b/hermes-llm/src/model/transformer.rs index 4592a77d..8cad6ea5 100644 --- a/hermes-llm/src/model/transformer.rs +++ b/hermes-llm/src/model/transformer.rs @@ -7,7 +7,7 @@ use burn::tensor::Int; use burn_nn::{Dropout, DropoutConfig, Embedding, EmbeddingConfig, Linear, LinearConfig}; use burn_nn::{RotaryEncoding, RotaryEncodingConfig}; -use crate::mal::{BlockDef, ModelDef, PositionEncoding}; +use crate::mal::{BlockDef, ModelDef, NormConfig, PositionEncoding}; use super::linear_cross_entropy::linear_cross_entropy; use super::matmul::{matmul_2, matmul_input, prepare_linear_for_inference, stream_cast}; @@ -16,6 +16,125 @@ use super::{InferenceState, Norm, TransformerBlock}; const EMBEDDING_STD: f64 = 0.02; const LOSS_CHUNKS: usize = 4; +fn validate_norm(name: &str, norm: &NormConfig) -> Result<()> { + if norm.eps != 0.0 && (!norm.eps.is_finite() || norm.eps <= 0.0) { + bail!( + "{name} epsilon must be finite and positive, got {}", + norm.eps + ); + } + Ok(()) +} + +fn validate_config(config: &ModelDef) -> Result<()> { + if config.num_layers == 0 { + bail!("model must contain at least one layer"); + } + if config.hidden_size == 0 || config.vocab_size == 0 || config.max_seq_len == 0 { + bail!("vocab_size, hidden_size, and max_seq_len must all be positive"); + } + if !(0.0..1.0).contains(&config.embeddings.dropout) { + bail!( + "embedding dropout must be in [0, 1), got {}", + config.embeddings.dropout + ); + } + if config + .embeddings + .scale + .is_some_and(|scale| !scale.is_finite() || scale <= 0.0) + { + bail!("embedding scale must be finite and positive"); + } + if let Some(norm) = &config.output.norm { + validate_norm("output norm", norm)?; + } + + for i in 0..config.num_layers { + let block = config.block_for_layer(i); + for (name, dropout) in [ + ("block", block.dropout), + ("attention", block.attention.dropout), + ("ffn", block.ffn.dropout), + ] { + if !(0.0..1.0).contains(&dropout) { + bail!("layer {i} {name} dropout must be in [0, 1), got {dropout}"); + } + } + validate_norm(&format!("layer {i} norm"), &block.norm)?; + let intermediate = match block.ffn.hidden_dim { + Some(size) => size, + None => config + .hidden_size + .checked_mul(4) + .ok_or_else(|| anyhow::anyhow!("layer {i} default FFN size overflows usize"))?, + }; + if intermediate == 0 { + bail!("layer {i} FFN hidden_dim must be positive"); + } + + if let Some(ssm) = &block.ssm { + for (name, size) in [ + ("expand", ssm.expand), + ("state_dim", ssm.state_dim), + ("conv_kernel", ssm.conv_kernel), + ("dt_rank", config.dt_rank(ssm)), + ] { + if size == 0 { + bail!("layer {i} Mamba {name} must be positive"); + } + } + if ssm.expand.checked_mul(config.hidden_size).is_none() { + bail!("layer {i} Mamba expand * hidden_size overflows usize"); + } + continue; + } + + let heads = block.attention.num_heads.unwrap_or(12); + if heads == 0 { + bail!("layer {i} attention num_heads must be positive"); + } + let kv_heads = block.attention.num_kv_heads.unwrap_or(heads); + let head_dim = block + .attention + .head_dim + .unwrap_or(config.hidden_size / heads); + if kv_heads == 0 || head_dim == 0 { + bail!("layer {i} attention num_kv_heads and head_dim must be positive"); + } + if !heads.is_multiple_of(kv_heads) { + bail!("layer {i} num_heads ({heads}) must be divisible by num_kv_heads ({kv_heads})"); + } + if heads.checked_mul(head_dim) != Some(config.hidden_size) { + bail!( + "layer {i} num_heads ({heads}) * head_dim ({head_dim}) must equal hidden_size ({})", + config.hidden_size + ); + } + if block.attention.window_size == Some(0) { + bail!("layer {i} attention window_size must be positive"); + } + match &block.attention.position_encoding { + PositionEncoding::Rope { theta, scaling } => { + if !head_dim.is_multiple_of(2) { + bail!("layer {i} RoPE head_dim must be even, got {head_dim}"); + } + if !theta.is_finite() + || *theta <= 0.0 + || scaling.is_some_and(|scale| !scale.is_finite() || scale <= 0.0) + { + bail!("layer {i} RoPE theta and scaling must be finite and positive"); + } + } + PositionEncoding::None => {} + other => { + bail!("position_encoding {other:?} is not implemented; use rope or none") + } + } + } + Ok(()) +} + fn pad_embedding(mut embedding: Embedding, stored_vocab_size: usize) -> Embedding { let [vocab_size, hidden_size] = embedding.weight.shape().dims(); if vocab_size < stored_vocab_size { @@ -84,67 +203,12 @@ pub struct Transformer { impl Transformer { pub fn new(config: &ModelDef, device: &Device) -> Result { - if config.num_layers == 0 { - bail!("model must contain at least one layer"); - } - if config.hidden_size == 0 || config.vocab_size == 0 || config.max_seq_len == 0 { - bail!("vocab_size, hidden_size, and max_seq_len must all be positive"); - } + validate_config(config)?; let attn_blocks: Vec<&BlockDef> = (0..config.num_layers) .map(|i| config.block_for_layer(i)) .filter(|block| !block.is_ssm()) .collect(); - for block in &attn_blocks { - let heads = block.num_heads(); - let kv_heads = block.num_kv_heads(); - let head_dim = block.head_dim(config.hidden_size); - if heads == 0 || kv_heads == 0 || head_dim == 0 { - bail!("attention head counts and head_dim must be positive"); - } - if heads % kv_heads != 0 { - bail!("num_heads ({heads}) must be divisible by num_kv_heads ({kv_heads})"); - } - if heads * head_dim != config.hidden_size { - bail!( - "num_heads ({heads}) * head_dim ({head_dim}) must equal hidden_size ({})", - config.hidden_size - ); - } - match &block.attention.position_encoding { - PositionEncoding::Rope { theta, scaling } => { - if head_dim % 2 != 0 { - bail!("RoPE head_dim must be even, got {head_dim}"); - } - if *theta <= 0.0 || scaling.is_some_and(|scale| scale <= 0.0) { - bail!("RoPE theta and scaling must be positive"); - } - } - PositionEncoding::None => {} - other => bail!("position_encoding {other:?} is not implemented; use rope or none"), - } - } - for i in 0..config.num_layers { - let block = config.block_for_layer(i); - for (name, dropout) in [ - ("block", block.dropout), - ("attention", block.attention.dropout), - ("ffn", block.ffn.dropout), - ] { - if !(0.0..1.0).contains(&dropout) { - bail!("{name} dropout must be in [0, 1), got {dropout}"); - } - } - if block.intermediate_size(config.hidden_size) == 0 { - bail!("FFN hidden_dim must be positive"); - } - } - if !(0.0..1.0).contains(&config.embeddings.dropout) { - bail!( - "embedding dropout must be in [0, 1), got {}", - config.embeddings.dropout - ); - } let rope_blocks: Vec<_> = attn_blocks .iter() diff --git a/hermes-llm/src/remote.rs b/hermes-llm/src/remote.rs index ef3c39cb..be4e810a 100644 --- a/hermes-llm/src/remote.rs +++ b/hermes-llm/src/remote.rs @@ -64,24 +64,49 @@ mod imp { Ok(base) } - /// Cache path: `-` so the extension (e.g. - /// `.json`, `.safetensors`) is preserved for downstream loaders while the - /// hash keeps distinct URIs (and buckets) from colliding. - fn cache_path(uri: &str) -> Result { + /// Return a URI suitable for diagnostics. Authentication material, signed + /// query parameters, and fragments are deliberately omitted. + pub(super) fn redacted_uri(uri: &str) -> String { + let Ok(mut url) = url::Url::parse(uri) else { + return "".to_owned(); + }; + let _ = url.set_username(""); + let _ = url.set_password(None); + url.set_query(None); + url.set_fragment(None); + url.to_string() + } + + /// Cache filename: `-`. Hashing the full + /// URI keeps distinct signed URLs from colliding, while deriving the human + /// readable suffix from the URL path prevents credentials and query + /// parameters from becoming filenames. + pub(super) fn cache_file_name(uri: &str) -> Result { let mut h = DefaultHasher::new(); uri.hash(&mut h); - let name = uri.rsplit(['/', '\\']).next().unwrap_or("artifact"); - let name = if name.is_empty() { "artifact" } else { name }; - Ok(cache_dir()?.join(format!("{:016x}-{name}", h.finish()))) + + let safe_uri = redacted_uri(uri); + let url = url::Url::parse(uri).with_context(|| format!("parsing URI {safe_uri}"))?; + let name = url + .path_segments() + .and_then(|mut segments| segments.rfind(|segment| !segment.is_empty())) + .filter(|name| name.len() <= 160) + .unwrap_or("artifact"); + Ok(format!("{:016x}-{name}", h.finish())) + } + + fn cache_path(uri: &str) -> Result { + Ok(cache_dir()?.join(cache_file_name(uri)?)) } fn download_cached(uri: &str) -> Result { let dest = cache_path(uri)?; + let safe_uri = redacted_uri(uri); if dest.exists() { - tracing::info!("using cached {} ({})", uri, dest.display()); + tracing::info!("using cached {} ({})", safe_uri, dest.display()); return Ok(dest); } - tracing::info!("downloading {} …", uri); + tracing::info!("downloading {} …", safe_uri); // Publish only after the complete object has been written and synced, so // an interrupted run never exposes a partial cache file. let tmp = dest.with_extension("part"); @@ -90,7 +115,7 @@ mod imp { .with_context(|| format!("publishing cache file {}", dest.display()))?; tracing::info!( "downloaded {} ({:.1} MiB) → {}", - uri, + safe_uri, size as f64 / (1024.0 * 1024.0), dest.display() ); @@ -106,10 +131,11 @@ mod imp { fn download(uri: &str, dest: &std::path::Path) -> Result { use object_store::{ObjectStore, ObjectStoreExt}; - let url = url::Url::parse(uri).with_context(|| format!("parsing URI {uri}"))?; + let safe_uri = redacted_uri(uri); + let url = url::Url::parse(uri).with_context(|| format!("parsing URI {safe_uri}"))?; let (store, path): (Box, object_store::path::Path) = object_store::parse_url(&url) - .with_context(|| format!("no object-store backend for {uri}"))?; + .with_context(|| format!("no object-store backend for {safe_uri}"))?; // object_store is async; run a single-threaded runtime for this blocking // CLI call rather than making the whole inference path async. @@ -212,6 +238,23 @@ mod tests { ); } + #[cfg(feature = "remote")] + #[test] + fn signed_http_uri_is_redacted_and_not_used_as_cache_filename() { + use super::imp::{cache_file_name, redacted_uri}; + + let uri = "https://alice:secret@example.com/models/weights.safetensors?X-Amz-Signature=private#fragment"; + let display = redacted_uri(uri); + assert_eq!(display, "https://example.com/models/weights.safetensors"); + + let name = cache_file_name(uri).unwrap(); + assert!(name.ends_with("-weights.safetensors"), "{name}"); + assert!( + !name.contains("secret") && !name.contains("X-Amz"), + "{name}" + ); + } + #[cfg(not(feature = "remote"))] #[test] fn remote_without_feature_errors_clearly() { diff --git a/hermes-mal-py/Cargo.toml b/hermes-mal-python/Cargo.toml similarity index 95% rename from hermes-mal-py/Cargo.toml rename to hermes-mal-python/Cargo.toml index 36158ad0..23e80318 100644 --- a/hermes-mal-py/Cargo.toml +++ b/hermes-mal-python/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "hermes-mal-py" +name = "hermes-mal-python" version.workspace = true edition.workspace = true rust-version.workspace = true diff --git a/hermes-mal-py/README.md b/hermes-mal-python/README.md similarity index 72% rename from hermes-mal-py/README.md rename to hermes-mal-python/README.md index df5cbe25..42f83990 100644 --- a/hermes-mal-py/README.md +++ b/hermes-mal-python/README.md @@ -1,7 +1,11 @@ -# hermes-mal (Python) +# hermes-mal Python bindings PyO3 bindings for the Hermes Model Architecture Language (MAL) parser. +The repository directory and Rust extension crate are named +`hermes-mal-python`; the published Python distribution remains `hermes-mal` +and its import module remains `hermes_mal`. + This wheel is a thin wrapper around the Rust `hermes-mal` crate — the single source of truth for parsing `.mal` model definitions. It exposes one function: diff --git a/hermes-mal-py/build.rs b/hermes-mal-python/build.rs similarity index 100% rename from hermes-mal-py/build.rs rename to hermes-mal-python/build.rs diff --git a/hermes-mal-py/pyproject.toml b/hermes-mal-python/pyproject.toml similarity index 97% rename from hermes-mal-py/pyproject.toml rename to hermes-mal-python/pyproject.toml index 37138a0b..6b206e27 100644 --- a/hermes-mal-py/pyproject.toml +++ b/hermes-mal-python/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "hermes-mal" -version = "1.8.54" +version = "1.8.68" description = "Python bindings for the Hermes MAL parser (single source of truth, backed by Rust)" readme = "README.md" license = "MIT" diff --git a/hermes-mal-py/src/lib.rs b/hermes-mal-python/src/lib.rs similarity index 100% rename from hermes-mal-py/src/lib.rs rename to hermes-mal-python/src/lib.rs diff --git a/hermes-mal/src/lib.rs b/hermes-mal/src/lib.rs index 03824020..7a9bbd8b 100644 --- a/hermes-mal/src/lib.rs +++ b/hermes-mal/src/lib.rs @@ -2,15 +2,18 @@ //! //! A composable DSL for defining LLM model architectures using pest parser. //! -//! # Example MAL - Simple (flat) style +//! # Example MAL - Minimal inline style //! //! ```text //! model tiny { //! vocab_size: 32000 +//! max_seq_len: 2048 //! hidden_size: 128 //! num_layers: 4 -//! num_heads: 4 -//! intermediate_size: 512 +//! block: { +//! attention: { num_heads: 4 } +//! ffn: { hidden_dim: 512 } +//! } //! } //! ``` //! @@ -714,13 +717,40 @@ fn parse_model_def(pair: pest::iterators::Pair, file: &MalFile) -> Result< Ok(def) } -/// Parse MAL from a string (returns first model found) +/// Parse a MAL string containing exactly one model definition. +/// +/// Use [`parse_mal_full`] when a source intentionally defines multiple models. pub fn parse_mal(input: &str) -> Result { let file = parse_mal_full(input)?; - file.models - .into_values() - .next() - .ok_or_else(|| anyhow!("No model definition found")) + match file.models.len() { + 0 => Err(anyhow!("no model definition found")), + 1 => Ok(file.models.into_values().next().expect("length checked")), + _ => { + let mut names = file.models.keys().cloned().collect::>(); + names.sort(); + Err(anyhow!( + "multiple model definitions found ({}); use parse_mal_full to select one", + names.join(", ") + )) + } + } +} + +fn insert_unique( + definitions: &mut HashMap, + kind: &str, + name: String, + value: T, +) -> Result<()> { + match definitions.entry(name) { + std::collections::hash_map::Entry::Vacant(entry) => { + entry.insert(value); + Ok(()) + } + std::collections::hash_map::Entry::Occupied(entry) => { + Err(anyhow!("duplicate {kind} '{}'", entry.key())) + } + } } /// Parse complete MAL file with all definitions @@ -737,23 +767,38 @@ pub fn parse_mal_full(input: &str) -> Result { match def.as_rule() { Rule::model_def => { let model = parse_model_def(def, &file)?; - file.models.insert(model.name.clone(), model); + insert_unique( + &mut file.models, + "model", + model.name.clone(), + model, + )?; } Rule::attention_def => { let attn = parse_attention_def(def)?; - file.attentions.insert(attn.name.clone(), attn); + insert_unique( + &mut file.attentions, + "attention", + attn.name.clone(), + attn, + )?; } Rule::ssm_def => { let ssm = parse_ssm_def(def)?; - file.ssms.insert(ssm.name.clone(), ssm); + insert_unique(&mut file.ssms, "ssm", ssm.name.clone(), ssm)?; } Rule::ffn_def => { let ffn = parse_ffn_def(def)?; - file.ffns.insert(ffn.name.clone(), ffn); + insert_unique(&mut file.ffns, "ffn", ffn.name.clone(), ffn)?; } Rule::block_def => { let block = parse_block_def(def, &file)?; - file.blocks.insert(block.name.clone(), block); + insert_unique( + &mut file.blocks, + "block", + block.name.clone(), + block, + )?; } _ => {} } @@ -955,12 +1000,10 @@ fn parse_norm_config(pair: pest::iterators::Pair) -> Result { }; for param in cfg.into_inner() { // norm_param -> norm_eps_prop -> number - if let Some(eps) = param - .into_inner() - .next() - .and_then(|p| p.into_inner().next().and_then(|n| n.as_str().parse().ok())) + if let Some(prop) = param.into_inner().next() + && let Some(number) = prop.into_inner().next() { - norm.eps = eps; + norm.eps = number.as_str().parse()?; } } } @@ -1352,6 +1395,28 @@ mod tests { } } + #[test] + fn test_parse_mal_rejects_multiple_models() { + let err = parse_mal( + r#" + model alpha { vocab_size: 10 hidden_size: 8 num_layers: 1 block: { attention: { num_heads: 1 } ffn: { hidden_dim: 16 } } } + model beta { vocab_size: 10 hidden_size: 8 num_layers: 1 block: { attention: { num_heads: 1 } ffn: { hidden_dim: 16 } } } + "#, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("multiple model definitions"), "{err}"); + assert!(err.contains("alpha") && err.contains("beta"), "{err}"); + } + + #[test] + fn test_parse_mal_full_rejects_duplicate_definitions() { + let err = parse_mal_full("attention repeated {} attention repeated {}") + .unwrap_err() + .to_string(); + assert!(err.contains("duplicate attention 'repeated'"), "{err}"); + } + #[test] fn test_wellknown_models() { for name in list_wellknown_models() { diff --git a/hermes-server/Dockerfile b/hermes-server/Dockerfile index fa6fa46d..93e75ef6 100644 --- a/hermes-server/Dockerfile +++ b/hermes-server/Dockerfile @@ -21,7 +21,7 @@ COPY Cargo.toml Cargo.lock ./ COPY hermes-core ./hermes-core COPY hermes-llm ./hermes-llm COPY hermes-mal ./hermes-mal -COPY hermes-mal-py ./hermes-mal-py +COPY hermes-mal-python ./hermes-mal-python COPY hermes-train ./hermes-train COPY hermes-server ./hermes-server COPY hermes-tool ./hermes-tool diff --git a/hermes-train/README.md b/hermes-train/README.md index e046a3b9..d77d7e9a 100644 --- a/hermes-train/README.md +++ b/hermes-train/README.md @@ -67,3 +67,41 @@ Outputs are deliberately minimal: The checkpoint loads directly in `hermes-llm` with strict tensor and shape validation. Experiment services such as W&B can tail `metrics.jsonl` without being linked into the training process. + +## Reliable relaunch and W&B + +[`scripts/relaunch.sh`](scripts/relaunch.sh) is the boot-safe supervisor for +long-running or spot-instance jobs. It owns `--output` and automatically adds +`--resume` only when all model, AdamW, Muon, and training-state files form a +complete checkpoint. A lock makes repeated boot hooks idempotent, failed +trainer processes are relaunched after a configurable delay, and termination +attempts one final remote sync. + +Remote backups use either `gs://` (through `gcloud storage`) or `file://`. +Checkpoints are uploaded to an immutable `checkpoints//` directory and +`latest.json` is published last. On boot, a newer complete remote checkpoint +is restored, while a newer persistent-disk checkpoint is never overwritten by +an older backup. The first sync migrates the earlier flat `gcloud rsync` +layout automatically. An interrupted local checkpoint is not resumed unless a +complete remote copy can replace it. + +Copy and edit the example configuration, then run the supervisor as the same +user that owns the training files: + +```bash +cp hermes-train/scripts/relaunch.conf.example /opt/hermes-run/relaunch.conf +hermes-train/scripts/relaunch.sh /opt/hermes-run/relaunch.conf +``` + +For boot and process supervision, use that command as `ExecStart` in a systemd +service with `Restart=on-failure`, or from an `@reboot` cron entry. The script +itself keeps the trainer alive after ordinary process failures, so systemd is +mainly protection for the supervisor and machine lifecycle. + +Set `HERMES_TRAIN_WANDB_ENV` and `HERMES_TRAIN_WANDB_PYTHON` in the run +configuration to supervise [`scripts/wandb_tail.py`](scripts/wandb_tail.py) +with the trainer. The environment file should be mode 600 and contain the API +key plus a stable `WANDB_RUN_ID`; the reporter then backfills `metrics.jsonl`, +survives file replacement during restore, and reconnects to the same run after +every restart. W&B configuration is validated before training starts so a +requested reporter cannot silently disappear. diff --git a/hermes-train/scripts/relaunch.conf.example b/hermes-train/scripts/relaunch.conf.example new file mode 100644 index 00000000..b6cf63cb --- /dev/null +++ b/hermes-train/scripts/relaunch.conf.example @@ -0,0 +1,43 @@ +# Trusted Bash configuration for scripts/relaunch.sh. +# Keep credentials out of this file; point HERMES_TRAIN_WANDB_ENV at a mode-600 +# file containing `WANDB_API_KEY=...` and the non-secret W&B run settings. + +HERMES_TRAIN_OUTPUT=/opt/hermes-run/checkpoint +HERMES_TRAIN_STATE_DIR=/opt/hermes-run/relaunch-state +HERMES_TRAIN_REMOTE_URL=gs://example-training-bucket/retriever-100m-stage1 + +# `--output` and `--resume` are owned by relaunch.sh and must not appear here. +HERMES_TRAIN_COMMAND=( + /opt/hermes-run/bin/hermes-train + train + --config /opt/hermes-run/retriever-100m.json + --tokenizer /opt/hermes-run/tokenizer.json + --data /opt/hermes-run/corpus/stage1.jsonl.zst + --batch-size 6 + --grad-accum 22 + --seq-len 4096 + --epochs 2 + --max-steps 19000 + --checkpoint-every 100 +) + +# A failed trainer is resumed indefinitely by default. Set a positive value to +# stop after that many retries so a boot-time service can apply its own policy. +HERMES_TRAIN_RESTART_DELAY=30 +HERMES_TRAIN_MAX_RESTARTS=0 + +# A complete checkpoint is published under checkpoints// and latest.json +# is updated last. The interval also controls retries after a cloud error. +HERMES_TRAIN_SYNC_INTERVAL=900 + +# Install W&B once in this virtual environment; relaunch.sh deliberately never +# installs packages or prints credentials during a boot. +HERMES_TRAIN_WANDB_ENV=/opt/hermes-run/wandb.env +HERMES_TRAIN_WANDB_PYTHON=/opt/hermes-run/wandb-venv/bin/python + +# Example /opt/hermes-run/wandb.env (chmod 600): +# WANDB_API_KEY=replace-me +# WANDB_ENTITY=space-frontiers +# WANDB_PROJECT=hermes-retriever +# WANDB_NAME=retriever-100m +# WANDB_RUN_ID=retriever-100m-stage1 diff --git a/hermes-train/scripts/relaunch.sh b/hermes-train/scripts/relaunch.sh new file mode 100755 index 00000000..431e0dc9 --- /dev/null +++ b/hermes-train/scripts/relaunch.sh @@ -0,0 +1,511 @@ +#!/usr/bin/env bash +# Supervise a hermes-train run across process failures and machine reboots. +# +# The single argument is a trusted Bash configuration file. See +# relaunch.conf.example for the supported settings. + +set -Eeuo pipefail +umask 077 + +RELAUNCH_SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +readonly RELAUNCH_SCRIPT_DIR +readonly -a CHECKPOINT_FILES=( + weights.safetensors + adamw-state.bpk + muon-state.bpk + training-state.json +) + +log() { + printf '%s hermes-train-relaunch: %s\n' "$(date -u +'%Y-%m-%dT%H:%M:%SZ')" "$*" >&2 +} + +die() { + log "error: $*" + exit 1 +} + +usage() { + cat >&2 <<'EOF' +Usage: relaunch.sh + +The configuration file must define: + HERMES_TRAIN_OUTPUT=/path/to/checkpoint + HERMES_TRAIN_COMMAND=(/path/to/hermes-train train ...) + +The supervisor appends --output and, when a complete checkpoint exists, +--resume. See relaunch.conf.example for cloud sync, W&B, and retry settings. +EOF +} + +[[ $# -eq 1 ]] || { + usage + exit 2 +} + +readonly RELAUNCH_CONFIG=$1 +[[ -r "$RELAUNCH_CONFIG" ]] || die "configuration is not readable: $RELAUNCH_CONFIG" + +# The configuration is trusted shell syntax so the training command can be a +# real Bash array without lossy string splitting or eval. +# shellcheck source=/dev/null +source "$RELAUNCH_CONFIG" + +: "${HERMES_TRAIN_OUTPUT:?set HERMES_TRAIN_OUTPUT in the configuration}" +declare -p HERMES_TRAIN_COMMAND >/dev/null 2>&1 \ + || die "set HERMES_TRAIN_COMMAND as a Bash array in the configuration" +[[ $(declare -p HERMES_TRAIN_COMMAND) == "declare -a"* ]] \ + || die "HERMES_TRAIN_COMMAND must be a Bash array" +(( ${#HERMES_TRAIN_COMMAND[@]} > 0 )) || die "HERMES_TRAIN_COMMAND is empty" + +readonly OUTPUT=${HERMES_TRAIN_OUTPUT%/} +readonly REMOTE=${HERMES_TRAIN_REMOTE_URL:-} +readonly STATE_DIR=${HERMES_TRAIN_STATE_DIR:-"$OUTPUT/.relaunch"} +readonly TRAIN_LOG=${HERMES_TRAIN_LOG:-"$STATE_DIR/train.log"} +readonly SYNC_LOG=${HERMES_TRAIN_SYNC_LOG:-"$STATE_DIR/sync.log"} +readonly WANDB_LOG=${HERMES_TRAIN_WANDB_LOG:-"$STATE_DIR/wandb.log"} +readonly LOCK_FILE=${HERMES_TRAIN_LOCK_FILE:-"$STATE_DIR/lock"} +readonly PYTHON_BIN=${HERMES_TRAIN_PYTHON:-python3} +readonly GCLOUD_BIN=${HERMES_TRAIN_GCLOUD:-gcloud} +readonly SYNC_INTERVAL=${HERMES_TRAIN_SYNC_INTERVAL:-900} +readonly RESTART_DELAY=${HERMES_TRAIN_RESTART_DELAY:-30} +readonly MAX_RESTARTS=${HERMES_TRAIN_MAX_RESTARTS:-0} +readonly WANDB_ENV=${HERMES_TRAIN_WANDB_ENV:-} +readonly WANDB_PYTHON=${HERMES_TRAIN_WANDB_PYTHON:-python3} +readonly WANDB_SCRIPT=${HERMES_TRAIN_WANDB_SCRIPT:-"$RELAUNCH_SCRIPT_DIR/wandb_tail.py"} +readonly WANDB_RESTART_DELAY=${HERMES_TRAIN_WANDB_RESTART_DELAY:-15} +readonly WANDB_FLUSH_DELAY=${HERMES_TRAIN_WANDB_FLUSH_DELAY:-6} + +is_nonnegative_integer() { + [[ $1 =~ ^[0-9]+$ ]] +} + +is_positive_integer() { + [[ $1 =~ ^[1-9][0-9]*$ ]] +} + +is_positive_integer "$SYNC_INTERVAL" || die "HERMES_TRAIN_SYNC_INTERVAL must be positive" +is_nonnegative_integer "$RESTART_DELAY" || die "HERMES_TRAIN_RESTART_DELAY must be non-negative" +is_nonnegative_integer "$MAX_RESTARTS" || die "HERMES_TRAIN_MAX_RESTARTS must be non-negative" +is_nonnegative_integer "$WANDB_RESTART_DELAY" \ + || die "HERMES_TRAIN_WANDB_RESTART_DELAY must be non-negative" +is_nonnegative_integer "$WANDB_FLUSH_DELAY" \ + || die "HERMES_TRAIN_WANDB_FLUSH_DELAY must be non-negative" + +for argument in "${HERMES_TRAIN_COMMAND[@]}"; do + case "$argument" in + --resume | --output | --output=* | -o) + die "leave $argument out of HERMES_TRAIN_COMMAND; the supervisor owns resume and output" + ;; + esac +done + +if command -v flock >/dev/null 2>&1; then + readonly LOCK_TOOL=flock +elif command -v shlock >/dev/null 2>&1; then + readonly LOCK_TOOL=shlock +else + die "flock (Linux) or shlock (macOS/BSD) is required" +fi +command -v "$PYTHON_BIN" >/dev/null 2>&1 || die "Python is required: $PYTHON_BIN" +command -v "${HERMES_TRAIN_COMMAND[0]}" >/dev/null 2>&1 \ + || die "trainer is unavailable: ${HERMES_TRAIN_COMMAND[0]}" +if [[ -n "$REMOTE" && $REMOTE != file://* ]]; then + [[ $REMOTE == gs://* ]] || die "remote URL must use gs:// or file://" + command -v "$GCLOUD_BIN" >/dev/null 2>&1 || die "gcloud is required for $REMOTE" +fi + +mkdir -p -- "$OUTPUT" "$STATE_DIR" "$(dirname -- "$TRAIN_LOG")" \ + "$(dirname -- "$SYNC_LOG")" "$(dirname -- "$WANDB_LOG")" + +# Children explicitly close fd 9 on Linux, so only this supervisor owns the +# flock. macOS/BSD shlock records the supervisor PID and rejects a live owner. +if [[ $LOCK_TOOL == flock ]]; then + exec 9>"$LOCK_FILE" + if ! flock -n 9; then + log "another supervisor already owns $LOCK_FILE; nothing to do" + exit 0 + fi +elif ! shlock -f "$LOCK_FILE" -p "$$"; then + log "another supervisor already owns $LOCK_FILE; nothing to do" + exit 0 +fi +printf '%s\n' "$$" >"$STATE_DIR/supervisor.pid" + +read_checkpoint_step() { + "$PYTHON_BIN" - "$1" <<'PY' +import json +import sys + +with open(sys.argv[1], encoding="utf-8") as handle: + state = json.load(handle) +step = state.get("step") +if not isinstance(step, int) or isinstance(step, bool) or step < 0: + raise SystemExit("training-state.json has an invalid step") +print(step) +PY +} + +checkpoint_step() { + local directory=$1 + local file + [[ ! -e "$directory/.checkpoint-in-progress" ]] || return 1 + for file in "${CHECKPOINT_FILES[@]}"; do + [[ -s "$directory/$file" ]] || return 1 + done + read_checkpoint_step "$directory/training-state.json" +} + +checkpoint_artifacts_exist() { + local file + [[ -e "$OUTPUT/.checkpoint-in-progress" ]] && return 0 + for file in "${CHECKPOINT_FILES[@]}"; do + [[ -e "$OUTPUT/$file" || -e "$OUTPUT/$file.tmp" ]] && return 0 + done + return 1 +} + +remote_path() { + printf '%s/%s' "${REMOTE%/}" "${1#/}" +} + +local_remote_root() { + printf '%s' "${REMOTE#file://}" +} + +remote_download() { + local relative=$1 + local destination=$2 + if [[ $REMOTE == file://* ]]; then + cp -- "$(local_remote_root)/$relative" "$destination" + else + "$GCLOUD_BIN" storage cp "$(remote_path "$relative")" "$destination" + fi +} + +remote_upload_file() { + local source=$1 + local relative=$2 + if [[ $REMOTE == file://* ]]; then + mkdir -p -- "$(dirname -- "$(local_remote_root)/$relative")" + cp -- "$source" "$(local_remote_root)/$relative" + else + "$GCLOUD_BIN" storage cp "$source" "$(remote_path "$relative")" + fi +} + +remote_upload_directory() { + local relative=$1 + shift + if [[ $REMOTE == file://* ]]; then + local destination + destination="$(local_remote_root)/$relative" + mkdir -p -- "$destination" + cp -- "$@" "$destination/" + else + "$GCLOUD_BIN" storage cp "$@" "$(remote_path "$relative")/" + fi +} + +remote_promote_checkpoint() { + local step=$1 + if [[ $REMOTE == file://* ]]; then + cp -- "$(local_remote_root)/checkpoints/$step/training-state.json" \ + "$(local_remote_root)/latest.json" + else + "$GCLOUD_BIN" storage cp \ + "$(remote_path "checkpoints/$step/training-state.json")" \ + "$(remote_path latest.json)" + fi +} + +REMOTE_STEP= +REMOTE_LAYOUT= +refresh_remote_checkpoint() { + local temporary + REMOTE_STEP= + REMOTE_LAYOUT= + [[ -n "$REMOTE" ]] || return 1 + temporary=$(mktemp "$STATE_DIR/remote-state.XXXXXX") + if remote_download latest.json "$temporary" >/dev/null 2>&1; then + if REMOTE_STEP=$(read_checkpoint_step "$temporary" 2>/dev/null); then + REMOTE_LAYOUT=versioned + rm -f -- "$temporary" + return 0 + fi + elif remote_download training-state.json "$temporary" >/dev/null 2>&1; then + # Compatibility with the original flat `gcloud storage rsync` layout. + if REMOTE_STEP=$(read_checkpoint_step "$temporary" 2>/dev/null); then + REMOTE_LAYOUT=legacy + rm -f -- "$temporary" + return 0 + fi + fi + rm -f -- "$temporary" + return 1 +} + +restore_remote_checkpoint() { + local expected_step=$1 + local layout=$2 + local restore_dir prefix file downloaded_step + restore_dir=$(mktemp -d "$STATE_DIR/restore.XXXXXX") + prefix= + [[ $layout == versioned ]] && prefix="checkpoints/$expected_step/" + + for file in "${CHECKPOINT_FILES[@]}"; do + if ! remote_download "$prefix$file" "$restore_dir/$file" >>"$SYNC_LOG" 2>&1; then + log "remote checkpoint $expected_step is incomplete ($file is unavailable)" + rm -rf -- "$restore_dir" + return 1 + fi + done + if ! downloaded_step=$(checkpoint_step "$restore_dir") \ + || [[ $downloaded_step != "$expected_step" ]]; then + log "remote checkpoint manifest says $expected_step but its state is invalid" + rm -rf -- "$restore_dir" + return 1 + fi + + # Keep the marker present until all data files and training-state.json have + # been atomically published. A reboot during restore is therefore detected. + printf '%s\n' "$expected_step" >"$OUTPUT/.checkpoint-in-progress" + for file in weights.safetensors adamw-state.bpk muon-state.bpk; do + mv -- "$restore_dir/$file" "$OUTPUT/$file.restore" + mv -f -- "$OUTPUT/$file.restore" "$OUTPUT/$file" + done + mv -- "$restore_dir/training-state.json" "$OUTPUT/training-state.json.restore" + mv -f -- "$OUTPUT/training-state.json.restore" "$OUTPUT/training-state.json" + rm -f -- "$OUTPUT/.checkpoint-in-progress" + rmdir -- "$restore_dir" + log "restored remote checkpoint at step $expected_step" +} + +RESUME_STEP= +prepare_checkpoint() { + local local_step= + local remote_available=false + RESUME_STEP= + + if local_step=$(checkpoint_step "$OUTPUT" 2>/dev/null); then + log "found complete local checkpoint at step $local_step" + fi + if [[ -n "$REMOTE" ]] && refresh_remote_checkpoint; then + remote_available=true + log "found remote checkpoint at step $REMOTE_STEP ($REMOTE_LAYOUT layout)" + fi + + if [[ $remote_available == true \ + && ( -z "$local_step" || $REMOTE_STEP -gt $local_step ) ]]; then + restore_remote_checkpoint "$REMOTE_STEP" "$REMOTE_LAYOUT" \ + || die "cannot restore the newest remote checkpoint" + local_step=$REMOTE_STEP + fi + + if [[ -z "$local_step" ]]; then + if checkpoint_artifacts_exist; then + die "local checkpoint is incomplete and no usable remote checkpoint is available" + fi + log "no checkpoint found; starting a new run" + return 1 + fi + RESUME_STEP=$local_step +} + +sync_checkpoint_once() ( + local step after_step remote_step=-1 + local -a sources=() + local file sync_owner + exec 9>&- + [[ -n "$REMOTE" ]] || return 0 + + if [[ $LOCK_TOOL == flock ]]; then + exec 8>"$STATE_DIR/sync.lock" + flock -n 8 || return 0 + else + # Bash 3.2 has no BASHPID. A short child observes this subshell as PPID. + sync_owner=$(sh -c 'printf "%s" "$PPID"') + shlock -f "$STATE_DIR/sync.lock" -p "$sync_owner" || return 0 + trap 'rm -f -- "$STATE_DIR/sync.lock"' EXIT + fi + + if [[ -s "$OUTPUT/metrics.jsonl" ]]; then + remote_upload_file "$OUTPUT/metrics.jsonl" metrics.jsonl || return 1 + fi + step=$(checkpoint_step "$OUTPUT" 2>/dev/null) || { + log "checkpoint sync skipped: no complete local checkpoint" + return 0 + } + if refresh_remote_checkpoint; then + remote_step=$REMOTE_STEP + fi + if (( remote_step > step )) \ + || { (( remote_step == step )) && [[ $REMOTE_LAYOUT == versioned ]]; }; then + return 0 + fi + + for file in "${CHECKPOINT_FILES[@]}"; do + sources+=("$OUTPUT/$file") + done + [[ -s "$OUTPUT/config.json" ]] && sources+=("$OUTPUT/config.json") + remote_upload_directory "checkpoints/$step" "${sources[@]}" || return 1 + + # The trainer may have started publishing another checkpoint while the + # upload was in flight. Only publish `latest.json` for an unchanged, + # complete local snapshot; incomplete version directories are never read. + after_step=$(checkpoint_step "$OUTPUT" 2>/dev/null) || { + log "checkpoint changed during upload; leaving remote latest unchanged" + return 1 + } + if [[ $after_step != "$step" ]]; then + log "checkpoint advanced from $step to $after_step during upload; retrying later" + return 1 + fi + # Promote the immutable state object that was just uploaded. Copying the + # live local file here would race with the next checkpoint publication. + remote_promote_checkpoint "$step" || return 1 + log "published checkpoint step $step to $REMOTE" +) + +validate_wandb() { + [[ -n "$WANDB_ENV" ]] || return 0 + [[ -r "$WANDB_ENV" ]] || die "W&B environment is not readable: $WANDB_ENV" + [[ -r "$WANDB_SCRIPT" ]] || die "W&B reporter is not readable: $WANDB_SCRIPT" + command -v "$WANDB_PYTHON" >/dev/null 2>&1 \ + || die "W&B Python is unavailable: $WANDB_PYTHON" + if ! ( + set -a + # shellcheck source=/dev/null + source "$WANDB_ENV" + set +a + [[ -n ${WANDB_API_KEY:-} ]] && "$WANDB_PYTHON" -c 'import wandb' + ); then + die "W&B is configured but WANDB_API_KEY or the wandb package is unavailable" + fi +} + +wandb_supervisor() { + exec 9>&- + local reporter_pid='' reporter_status + trap '[[ -z $reporter_pid ]] || kill "$reporter_pid" 2>/dev/null; wait "$reporter_pid" 2>/dev/null || true; exit 0' TERM INT + set -a + # shellcheck source=/dev/null + source "$WANDB_ENV" + set +a + export PYTHONUNBUFFERED=1 + while true; do + "$WANDB_PYTHON" "$WANDB_SCRIPT" "$OUTPUT/metrics.jsonl" & + reporter_pid=$! + set +e + wait "$reporter_pid" + reporter_status=$? + set -e + reporter_pid= + log "W&B reporter exited with status $reporter_status; restarting in ${WANDB_RESTART_DELAY}s" + sleep "$WANDB_RESTART_DELAY" + done +} + +sync_supervisor() { + exec 9>&- + local child_pid='' sync_status + trap '[[ -z $child_pid ]] || kill "$child_pid" 2>/dev/null; wait "$child_pid" 2>/dev/null || true; exit 0' TERM INT + while true; do + sync_checkpoint_once & + child_pid=$! + set +e + wait "$child_pid" + sync_status=$? + set -e + child_pid= + if (( sync_status != 0 )); then + log "checkpoint sync failed; retrying in ${SYNC_INTERVAL}s" + fi + sleep "$SYNC_INTERVAL" & + child_pid=$! + wait "$child_pid" || true + child_pid= + done +} + +TRAIN_PID= +SYNC_PID= +WANDB_PID= + +cleanup() { + local status=$? + trap - EXIT TERM INT + set +e + if [[ -n "$TRAIN_PID" ]]; then + kill "$TRAIN_PID" 2>/dev/null + wait "$TRAIN_PID" 2>/dev/null + fi + if [[ -n "$WANDB_PID" ]]; then + kill "$WANDB_PID" 2>/dev/null + wait "$WANDB_PID" 2>/dev/null + fi + if [[ -n "$SYNC_PID" ]]; then + kill "$SYNC_PID" 2>/dev/null + wait "$SYNC_PID" 2>/dev/null + fi + if [[ -n "$REMOTE" ]]; then + sync_checkpoint_once >>"$SYNC_LOG" 2>&1 + fi + rm -f -- "$STATE_DIR/supervisor.pid" + [[ $LOCK_TOOL != shlock ]] || rm -f -- "$LOCK_FILE" + exit "$status" +} + +trap cleanup EXIT +trap 'exit 143' TERM INT + +validate_wandb +if [[ -n "$REMOTE" ]]; then + sync_supervisor >>"$SYNC_LOG" 2>&1 & + SYNC_PID=$! +fi +if [[ -n "$WANDB_ENV" ]]; then + wandb_supervisor >>"$WANDB_LOG" 2>&1 & + WANDB_PID=$! + log "W&B reporter is supervised (log: $WANDB_LOG)" +else + log "W&B reporting is disabled; set HERMES_TRAIN_WANDB_ENV to enable it" +fi + +restart_count=0 +while true; do + if prepare_checkpoint; then + trainer=("${HERMES_TRAIN_COMMAND[@]}" --output "$OUTPUT" --resume) + log "launching training from checkpoint step $RESUME_STEP" + else + trainer=("${HERMES_TRAIN_COMMAND[@]}" --output "$OUTPUT") + log "launching training from scratch" + fi + + ( + exec 9>&- + exec "${trainer[@]}" + ) >>"$TRAIN_LOG" 2>&1 & + TRAIN_PID=$! + set +e + wait "$TRAIN_PID" + trainer_status=$? + set -e + TRAIN_PID= + + if [[ -n "$REMOTE" ]]; then + sync_checkpoint_once >>"$SYNC_LOG" 2>&1 || true + fi + if (( trainer_status == 0 )); then + log "training completed successfully" + [[ -z "$WANDB_PID" || $WANDB_FLUSH_DELAY -eq 0 ]] || sleep "$WANDB_FLUSH_DELAY" + exit 0 + fi + + (( restart_count += 1 )) + log "trainer exited with status $trainer_status (restart $restart_count)" + if (( MAX_RESTARTS > 0 && restart_count > MAX_RESTARTS )); then + die "trainer exceeded HERMES_TRAIN_MAX_RESTARTS=$MAX_RESTARTS" + fi + sleep "$RESTART_DELAY" +done diff --git a/hermes-train/scripts/relaunch_test.sh b/hermes-train/scripts/relaunch_test.sh new file mode 100755 index 00000000..5fea1cba --- /dev/null +++ b/hermes-train/scripts/relaunch_test.sh @@ -0,0 +1,244 @@ +#!/usr/bin/env bash + +set -Eeuo pipefail + +TEST_SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +readonly TEST_SCRIPT_DIR +TEST_ROOT=$(mktemp -d) +readonly TEST_ROOT +trap 'rm -rf -- "$TEST_ROOT"' EXIT + +fail() { + printf 'relaunch_test: %s\n' "$*" >&2 + exit 1 +} + +write_checkpoint() { + local directory=$1 + local step=$2 + mkdir -p -- "$directory" + printf 'weights-%s\n' "$step" >"$directory/weights.safetensors" + printf 'adamw-%s\n' "$step" >"$directory/adamw-state.bpk" + printf 'muon-%s\n' "$step" >"$directory/muon-state.bpk" + printf '{"step":%s}\n' "$step" >"$directory/training-state.json" +} + +fake_trainer=$TEST_ROOT/fake-trainer +fake_wandb_python=$TEST_ROOT/fake-wandb-python + +cat >"$fake_trainer" <<'EOF' +#!/usr/bin/env bash +set -euo pipefail +output= +resume=false +while (( $# > 0 )); do + case "$1" in + --output) + output=$2 + shift 2 + ;; + --resume) + resume=true + shift + ;; + *) + shift + ;; + esac +done +printf '%s\n' "$resume" >>"$TEST_CALLS" +if [[ ${TEST_BLOCK:-false} == true ]]; then + : >"$TEST_READY" + while [[ ! -e $TEST_RELEASE ]]; do + sleep 0.05 + done + exit 0 +fi +if [[ ${TEST_FAIL_ONCE:-false} == true && ! -e $TEST_FAILURE_MARKER ]]; then + mkdir -p -- "$output" + printf 'weights-3\n' >"$output/weights.safetensors" + printf 'adamw-3\n' >"$output/adamw-state.bpk" + printf 'muon-3\n' >"$output/muon-state.bpk" + printf '{"step":3}\n' >"$output/training-state.json" + printf '{"step":3,"loss":1.0}\n' >"$output/metrics.jsonl" + : >"$TEST_FAILURE_MARKER" + exit 17 +fi +if [[ -n ${TEST_EXPECT_STEP:-} ]]; then + [[ $resume == true ]] || exit 91 + actual=$(python3 -c 'import json,sys; print(json.load(open(sys.argv[1]))["step"])' \ + "$output/training-state.json") + [[ $actual == "$TEST_EXPECT_STEP" ]] || exit 92 +fi +EOF + +cat >"$fake_wandb_python" <<'EOF' +#!/usr/bin/env bash +set -euo pipefail +if [[ ${1:-} == -c ]]; then + exit 0 +fi +printf 'started\n' >>"$TEST_WANDB_CALLS" +trap 'exit 0' TERM INT +while true; do + sleep 1 +done +EOF +chmod +x "$fake_trainer" "$fake_wandb_python" + +run_restart_and_reporting_test() { + local case_root=$TEST_ROOT/restart + local config=$case_root/relaunch.conf + mkdir -p -- "$case_root/remote" + printf 'WANDB_API_KEY=test-only\n' >"$case_root/wandb.env" + chmod 600 "$case_root/wandb.env" + cat >"$config" <&2 || true + fail "checkpoint payload was not synced" + fi + [[ $(python3 -c 'import json,sys; print(json.load(open(sys.argv[1]))["step"])' \ + "$case_root/remote/latest.json") == 3 ]] || fail "latest manifest was not published last" +} + +run_remote_restore_test() { + local case_root=$TEST_ROOT/restore + local config=$case_root/relaunch.conf + mkdir -p -- "$case_root/remote/checkpoints/7" "$case_root/output" + write_checkpoint "$case_root/remote/checkpoints/7" 7 + cp -- "$case_root/remote/checkpoints/7/training-state.json" "$case_root/remote/latest.json" + printf 'interrupted\n' >"$case_root/output/.checkpoint-in-progress" + printf 'stale\n' >"$case_root/output/weights.safetensors" + cat >"$config" <"$config" <"$config" <"$config" <"$case_root/first.log" 2>&1 & + supervisor_pid=$! + for _attempt in {1..100}; do + if [[ -e $TEST_READY ]]; then + ready=true + break + fi + sleep 0.05 + done + "$TEST_SCRIPT_DIR/relaunch.sh" "$config" >"$case_root/second.log" 2>&1 + : >"$TEST_RELEASE" + wait "$supervisor_pid" + + [[ $ready == true ]] || fail "first supervisor did not launch its trainer" + [[ $(wc -l <"$TEST_CALLS") -eq 1 ]] || fail "duplicate supervisor launched a trainer" + grep -q 'another supervisor already owns' "$case_root/second.log" \ + || fail "duplicate supervisor did not report the held lock" +} + +run_restart_and_reporting_test +run_remote_restore_test +run_newer_local_wins_test +run_legacy_remote_migration_test +run_idempotent_lock_test +printf 'relaunch_test: ok\n' diff --git a/hermes-train/scripts/wandb_tail.py b/hermes-train/scripts/wandb_tail.py old mode 100644 new mode 100755 index 39e413cb..45b12749 --- a/hermes-train/scripts/wandb_tail.py +++ b/hermes-train/scripts/wandb_tail.py @@ -14,8 +14,9 @@ import json import os +import signal import sys -import time +import threading def main() -> int: @@ -29,41 +30,71 @@ def main() -> int: import wandb # deferred so a missing package never blocks training setup + project = os.environ.get("WANDB_PROJECT", "hermes-retriever") name = os.environ.get("WANDB_NAME", "retriever-100m") + run_id = os.environ.get("WANDB_RUN_ID", f"{name}-stage1") run = wandb.init( - project=os.environ.get("WANDB_PROJECT", "hermes-retriever"), + project=project, name=name, - id=os.environ.get("WANDB_RUN_ID", f"{name}-stage1"), + id=run_id, resume="allow", ) - last_step = run.step or 0 + # `run.step` starts from zero in some W&B SDK versions even when attaching + # to an existing run. The public run record is authoritative and prevents + # a resumed reporter from attempting to emit thousands of duplicate steps. + remote_run = wandb.Api().run(f"{run.entity}/{project}/{run_id}") + last_step = max(run.step or 0, remote_run.lastHistoryStep or 0) position = 0 - while True: - if not os.path.exists(path): - time.sleep(5) - continue - with open(path, encoding="utf-8") as handle: - handle.seek(position) - while True: - line = handle.readline() - if not line: + identity = None + stop = threading.Event() + + def request_stop(_signum, _frame): + stop.set() + + signal.signal(signal.SIGINT, request_stop) + signal.signal(signal.SIGTERM, request_stop) + try: + while not stop.is_set(): + try: + stat = os.stat(path) + except FileNotFoundError: + stop.wait(5) + continue + current_identity = (stat.st_dev, stat.st_ino) + if current_identity != identity or stat.st_size < position: + # A restore may atomically replace or truncate metrics.jsonl. + # Re-read it; last_step filters the overlapping history. + identity = current_identity + position = 0 + with open(path, encoding="utf-8") as handle: + handle.seek(position) + while True: + line_position = handle.tell() + line = handle.readline() + if not line: + position = handle.tell() + break + if not line.endswith("\n"): + # Partial write: re-read this line on the next pass. + position = line_position + break position = handle.tell() - break - if not line.endswith("\n"): - # Partial write: re-read this line on the next pass. - break - position = handle.tell() - try: - record = json.loads(line) - except json.JSONDecodeError: - continue - step = int(record.get("step", 0)) - if step <= last_step: - continue # already logged before a resume/backfill overlap - wandb.log(record, step=step) - last_step = step - time.sleep(5) + try: + record = json.loads(line) + except json.JSONDecodeError: + continue + raw_step = record.get("step") + if not isinstance(raw_step, int) or isinstance(raw_step, bool): + continue + step = raw_step + if step <= last_step: + continue # already logged before a resume/backfill overlap + wandb.log(record, step=step) + last_step = step + stop.wait(5) + finally: + run.finish() if __name__ == "__main__": diff --git a/hermes-train/src/checkpoint.rs b/hermes-train/src/checkpoint.rs new file mode 100644 index 00000000..fe2a9e7b --- /dev/null +++ b/hermes-train/src/checkpoint.rs @@ -0,0 +1,118 @@ +//! Atomic resumable training checkpoints. +//! +//! Parameter IDs are persisted alongside model and optimizer state because +//! Burn optimizers key their state by those IDs. A marker prevents resuming a +//! checkpoint whose multi-file publication was interrupted. + +use std::fs; +use std::path::Path; + +use anyhow::{Context, Result, ensure}; +use burn::module::{AutodiffModule, Module, ModuleMapper, Param, ParamId}; +use burn::tensor::{Device, Tensor}; +use burn_optim::ModuleOptimizer; +use hermes_llm::{Transformer, load_safetensors, save_safetensors}; +use serde::{Deserialize, Serialize}; + +use crate::muon::BatchedMuon; + +pub(crate) type AdamWOptimizer = ModuleOptimizer; + +#[derive(Clone, Deserialize, Serialize)] +pub(crate) struct TrainingState { + pub(crate) step: usize, + pub(crate) stage: usize, + pub(crate) epoch: usize, + pub(crate) samples_in_stage: usize, + pub(crate) parameter_ids: Vec, +} + +pub(crate) fn parameter_ids(model: &Transformer) -> Vec { + burn::module::list_param_ids(model) + .into_iter() + .map(|id| id.val()) + .collect() +} + +struct ParameterIdMapper<'a> { + ids: std::slice::Iter<'a, u64>, +} + +impl ModuleMapper for ParameterIdMapper<'_> { + fn map_float(&mut self, param: Param>) -> Param> { + let (_, tensor, mapper) = param.consume(); + let id = self + .ids + .next() + .copied() + .expect("checkpoint contains too few parameter IDs"); + Param::from_mapped_value(ParamId::from(id), tensor, mapper) + } +} + +fn restore_parameter_ids(model: &mut Transformer, ids: &[u64]) -> Result<()> { + ensure!( + ids.len() == burn::module::list_param_ids(model).len(), + "checkpoint has {} parameter IDs, model has {}", + ids.len(), + burn::module::list_param_ids(model).len() + ); + let mut mapper = ParameterIdMapper { ids: ids.iter() }; + *model = model.clone().map(&mut mapper); + ensure!( + mapper.ids.next().is_none(), + "checkpoint contains too many parameter IDs" + ); + Ok(()) +} + +pub(crate) fn save_training_checkpoint( + model: &Transformer, + adamw: &AdamWOptimizer, + muon: &BatchedMuon, + state: &TrainingState, + output: &Path, +) -> Result<()> { + let marker = output.join(".checkpoint-in-progress"); + let weights_temporary = output.join("weights.safetensors.tmp"); + let adamw_temporary = output.join("adamw-state.bpk.tmp"); + let muon_temporary = output.join("muon-state.bpk.tmp"); + let state_temporary = output.join("training-state.json.tmp"); + + fs::write(&marker, state.step.to_string())?; + save_safetensors(&model.clone().valid(), &weights_temporary)?; + adamw + .save(&adamw_temporary) + .context("failed to save AdamW state")?; + muon.save(&muon_temporary)?; + fs::write(&state_temporary, serde_json::to_vec_pretty(&state)?)?; + fs::rename(weights_temporary, output.join("weights.safetensors"))?; + fs::rename(adamw_temporary, output.join("adamw-state.bpk"))?; + fs::rename(muon_temporary, output.join("muon-state.bpk"))?; + fs::rename(state_temporary, output.join("training-state.json"))?; + fs::remove_file(marker)?; + Ok(()) +} + +pub(crate) fn load_training_state( + model: &mut Transformer, + adamw: AdamWOptimizer, + muon: &mut BatchedMuon, + output: &Path, + device: &Device, +) -> Result<(AdamWOptimizer, TrainingState)> { + ensure!( + !output.join(".checkpoint-in-progress").exists(), + "checkpoint was interrupted while being saved" + ); + let state: TrainingState = + serde_json::from_slice(&fs::read(output.join("training-state.json"))?)?; + restore_parameter_ids(model, &state.parameter_ids)?; + load_safetensors(model, output.join("weights.safetensors"))?; + muon.set_parameter_ids(model.muon_parameter_ids()); + muon.load(output.join("muon-state.bpk"), &device.clone().inner())?; + let adamw = adamw + .load(output.join("adamw-state.bpk")) + .context("failed to load AdamW state")?; + Ok((adamw, state)) +} diff --git a/hermes-train/src/data.rs b/hermes-train/src/data.rs new file mode 100644 index 00000000..28bd7860 --- /dev/null +++ b/hermes-train/src/data.rs @@ -0,0 +1,315 @@ +//! Streaming corpus ingestion, tokenization, shuffling, and batch packing. +//! +//! Documents are joined with EOS and packed without padding. JSONL and plain +//! text inputs may be Zstandard-compressed; only a bounded shuffle buffer and +//! one tokenizer batch are retained in memory. + +use std::fs::File; +use std::io::{BufRead, BufReader, Read}; +use std::path::Path; + +use anyhow::{Context, Result, ensure}; +use burn::tensor::{Device, Int, Tensor, TensorData}; +use hermes_llm::Tokenizer; +use rand::rngs::StdRng; +use rand::seq::SliceRandom; +use rand::{Rng, SeedableRng}; + +const TOKENIZE_BATCH: usize = 1_000; + +fn open_data(path: &Path) -> Result> { + let file = File::open(path) + .with_context(|| format!("failed to open training data {}", path.display()))?; + if path.extension().is_some_and(|ext| ext == "zst") { + let decoder = zstd::stream::read::Decoder::new(file) + .with_context(|| format!("failed to open zstd stream {}", path.display()))?; + Ok(Box::new(BufReader::new(decoder))) + } else { + Ok(Box::new(BufReader::new(file))) + } +} + +struct ShuffleBuffer { + samples: Vec>, + rng: StdRng, + capacity: usize, +} + +impl ShuffleBuffer { + fn new(capacity: usize, seed: u64) -> Self { + assert!(capacity > 0); + Self { + samples: Vec::with_capacity(capacity), + rng: StdRng::seed_from_u64(seed), + capacity, + } + } + + fn push(&mut self, sample: Vec) -> Option> { + if self.samples.len() < self.capacity { + self.samples.push(sample); + return None; + } + let index = self.rng.random_range(0..self.samples.len()); + Some(std::mem::replace(&mut self.samples[index], sample)) + } + + fn finish(mut self) -> Vec> { + self.samples.shuffle(&mut self.rng); + self.samples + } +} + +struct SamplePacker { + pending: Vec, + consumed: usize, + seq_len: usize, +} + +impl SamplePacker { + fn new(seq_len: usize) -> Self { + Self { + pending: Vec::new(), + consumed: 0, + seq_len, + } + } + + fn push( + &mut self, + tokens: impl IntoIterator, + count: &mut usize, + visit: &mut impl FnMut(Vec) -> Result, + ) -> Result { + if self.consumed > 0 { + self.pending.drain(..self.consumed); + self.consumed = 0; + } + self.pending.extend(tokens); + while self.pending.len() - self.consumed > self.seq_len { + let end = self.consumed + self.seq_len + 1; + let sample = self.pending[self.consumed..end].to_vec(); + self.consumed += self.seq_len; + *count += 1; + if !visit(sample)? { + return Ok(false); + } + } + Ok(true) + } +} + +fn push_documents( + documents: &mut Vec, + tokenizer: &Tokenizer, + packer: &mut SamplePacker, + count: &mut usize, + visit: &mut impl FnMut(Vec) -> Result, +) -> Result { + if documents.is_empty() { + return Ok(true); + } + let encodings = tokenizer.encode_batch(std::mem::take(documents), false)?; + for tokens in encodings { + let tokens = tokens + .into_iter() + .map(i64::from) + .chain(std::iter::once(i64::from(tokenizer.eos_token_id()))); + if !packer.push(tokens, count, visit)? { + return Ok(false); + } + } + Ok(true) +} + +/// Visit fixed-length next-token samples in their source order. +fn visit_samples_in_order( + path: &Path, + tokenizer: &Tokenizer, + seq_len: usize, + mut visit: impl FnMut(Vec) -> Result, +) -> Result { + ensure!(seq_len > 0, "seq_len must be positive"); + let mut count = 0; + let mut packer = SamplePacker::new(seq_len); + let is_jsonl = path + .file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.ends_with(".jsonl") || name.ends_with(".jsonl.zst")); + let mut reader = open_data(path)?; + if is_jsonl { + let mut documents = Vec::with_capacity(TOKENIZE_BATCH); + let mut line = String::new(); + let mut line_number = 0; + loop { + line.clear(); + if reader.read_line(&mut line)? == 0 { + break; + } + line_number += 1; + if line.trim().is_empty() { + continue; + } + let value: serde_json::Value = serde_json::from_str(&line) + .with_context(|| format!("invalid JSONL at {}:{line_number}", path.display()))?; + let document = value + .get("text") + .and_then(|value| value.as_str()) + .with_context(|| { + format!( + "JSONL row at {}:{line_number} must contain a string `text` field", + path.display() + ) + })?; + documents.push(document.to_owned()); + if documents.len() == TOKENIZE_BATCH + && !push_documents( + &mut documents, + tokenizer, + &mut packer, + &mut count, + &mut visit, + )? + { + return Ok(count); + } + } + if !push_documents( + &mut documents, + tokenizer, + &mut packer, + &mut count, + &mut visit, + )? { + return Ok(count); + } + } else { + let mut document = String::new(); + reader.read_to_string(&mut document)?; + if !push_documents( + &mut vec![document], + tokenizer, + &mut packer, + &mut count, + &mut visit, + )? { + return Ok(count); + } + } + Ok(count) +} + +pub(crate) fn visit_samples( + path: &Path, + tokenizer: &Tokenizer, + seq_len: usize, + shuffle_buffer: usize, + seed: u64, + mut visit: impl FnMut(Vec) -> Result, +) -> Result { + if shuffle_buffer == 0 { + return visit_samples_in_order(path, tokenizer, seq_len, visit); + } + + let mut shuffler = ShuffleBuffer::new(shuffle_buffer, seed); + let mut keep_going = true; + let count = visit_samples_in_order(path, tokenizer, seq_len, |sample| { + if let Some(sample) = shuffler.push(sample) { + keep_going = visit(sample)?; + } + Ok(keep_going) + })?; + + if keep_going { + for sample in shuffler.finish() { + if !visit(sample)? { + break; + } + } + } + Ok(count) +} + +pub(crate) fn count_samples(path: &Path, tokenizer: &Tokenizer, seq_len: usize) -> Result { + visit_samples_in_order(path, tokenizer, seq_len, |_| Ok(true)) +} + +pub(crate) fn make_batch( + samples: &[Vec], + seq_len: usize, + device: &Device, +) -> (Tensor<2, Int>, Tensor<2, Int>) { + let mut inputs = Vec::with_capacity(samples.len() * seq_len); + let mut targets = Vec::with_capacity(samples.len() * seq_len); + for sample in samples { + inputs.extend_from_slice(&sample[..seq_len]); + targets.extend_from_slice(&sample[1..]); + } + ( + Tensor::from_data(TensorData::new(inputs, [samples.len(), seq_len]), device), + Tensor::from_data(TensorData::new(targets, [samples.len(), seq_len]), device), + ) +} + +#[cfg(test)] +mod tests { + use std::fs; + use std::io::{Cursor, Read}; + + use super::*; + + #[test] + fn zstd_data_reader_streams_decompressed_text() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("data.jsonl.zst"); + let source = b"{\"text\":\"one\"}\n{\"text\":\"two\"}\n"; + let compressed = zstd::stream::encode_all(Cursor::new(source), 1).unwrap(); + fs::write(&path, compressed).unwrap(); + + let mut reader = open_data(&path).unwrap(); + let mut decoded = String::new(); + reader.read_to_string(&mut decoded).unwrap(); + assert_eq!(decoded.as_bytes(), source); + } + + #[test] + fn streaming_shuffle_is_bounded_and_deterministic() { + let shuffle = |seed| { + let mut buffer = ShuffleBuffer::new(4, seed); + let mut output = Vec::new(); + for value in 0..32_i64 { + if let Some(sample) = buffer.push(vec![value]) { + output.push(sample[0]); + } + assert!(buffer.samples.len() <= 4); + } + output.extend(buffer.finish().into_iter().map(|sample| sample[0])); + output + }; + + let first = shuffle(7); + assert_eq!(first, shuffle(7)); + assert_ne!(first, (0..32_i64).collect::>()); + let mut sorted = first; + sorted.sort_unstable(); + assert_eq!(sorted, (0..32_i64).collect::>()); + } + + #[test] + fn sample_packer_joins_documents_without_dropping_tokens() { + let mut packer = SamplePacker::new(3); + let mut samples = Vec::new(); + let mut count = 0; + let mut collect = |sample| { + samples.push(sample); + Ok(true) + }; + + for document in [vec![1, 2, 0], vec![3, 4, 0], vec![5, 6, 0]] { + assert!(packer.push(document, &mut count, &mut collect).unwrap()); + } + + assert_eq!(count, 2); + assert_eq!(samples, [vec![1, 2, 0, 3], vec![3, 4, 0, 5]]); + } +} diff --git a/hermes-train/src/main.rs b/hermes-train/src/main.rs index 5279fb13..3598724d 100644 --- a/hermes-train/src/main.rs +++ b/hermes-train/src/main.rs @@ -1,27 +1,26 @@ -use std::fs::{self, File, OpenOptions}; -use std::io::{BufRead, BufReader, BufWriter, Read, Write}; +use std::fs::{self, OpenOptions}; +use std::io::{BufWriter, Write}; use std::path::{Path, PathBuf}; use std::time::Instant; -use anyhow::{Context, Result, bail, ensure}; -use burn::module::{AutodiffModule, Module, ModuleMapper, ModuleVisitor, Param, ParamId}; -use burn::tensor::{Device, Int, Tensor, TensorData}; -use burn_optim::{AdamWConfig, GradientsAccumulator, GradientsParams, ModuleOptimizer}; +use anyhow::{Result, bail, ensure}; +use burn::module::{Module, ModuleVisitor, Param}; +use burn::tensor::Tensor; +use burn_optim::{AdamWConfig, GradientsAccumulator, GradientsParams}; use clap::{Parser, Subcommand, ValueEnum}; -use hermes_llm::{ModelDef, Tokenizer, Transformer, load_safetensors, save_safetensors}; -use rand::rngs::StdRng; -use rand::seq::SliceRandom; -use rand::{Rng, SeedableRng}; -use serde::{Deserialize, Serialize}; +use hermes_llm::{ModelDef, Tokenizer, Transformer, load_safetensors}; +mod checkpoint; +mod data; mod muon; +use checkpoint::{ + AdamWOptimizer, TrainingState, load_training_state, parameter_ids, save_training_checkpoint, +}; +use data::{count_samples, make_batch, visit_samples}; use muon::BatchedMuon; -type AdamWOptimizer = ModuleOptimizer; - const MUON_LR_SCALE: f64 = 20.0; -const TOKENIZE_BATCH: usize = 1_000; #[derive(Parser)] #[command(name = "hermes-train", about = "Hermes model training")] @@ -91,15 +90,6 @@ struct TrainArgs { seed: u64, } -#[derive(Clone, Deserialize, Serialize)] -struct TrainingState { - step: usize, - stage: usize, - epoch: usize, - samples_in_stage: usize, - parameter_ids: Vec, -} - fn load_config(path: &Path) -> Result { if path.extension().is_some_and(|ext| ext == "mal") { return hermes_llm::parse_mal_file(path); @@ -107,240 +97,6 @@ fn load_config(path: &Path) -> Result { ModelDef::from_json(path) } -fn open_data(path: &Path) -> Result> { - let file = File::open(path) - .with_context(|| format!("failed to open training data {}", path.display()))?; - if path.extension().is_some_and(|ext| ext == "zst") { - let decoder = zstd::stream::read::Decoder::new(file) - .with_context(|| format!("failed to open zstd stream {}", path.display()))?; - Ok(Box::new(BufReader::new(decoder))) - } else { - Ok(Box::new(BufReader::new(file))) - } -} - -struct ShuffleBuffer { - samples: Vec>, - rng: StdRng, - capacity: usize, -} - -impl ShuffleBuffer { - fn new(capacity: usize, seed: u64) -> Self { - assert!(capacity > 0); - Self { - samples: Vec::with_capacity(capacity), - rng: StdRng::seed_from_u64(seed), - capacity, - } - } - - fn push(&mut self, sample: Vec) -> Option> { - if self.samples.len() < self.capacity { - self.samples.push(sample); - return None; - } - let index = self.rng.random_range(0..self.samples.len()); - Some(std::mem::replace(&mut self.samples[index], sample)) - } - - fn finish(mut self) -> Vec> { - self.samples.shuffle(&mut self.rng); - self.samples - } -} - -struct SamplePacker { - pending: Vec, - consumed: usize, - seq_len: usize, -} - -impl SamplePacker { - fn new(seq_len: usize) -> Self { - Self { - pending: Vec::new(), - consumed: 0, - seq_len, - } - } - - fn push( - &mut self, - tokens: impl IntoIterator, - count: &mut usize, - visit: &mut impl FnMut(Vec) -> Result, - ) -> Result { - if self.consumed > 0 { - self.pending.drain(..self.consumed); - self.consumed = 0; - } - self.pending.extend(tokens); - while self.pending.len() - self.consumed > self.seq_len { - let end = self.consumed + self.seq_len + 1; - let sample = self.pending[self.consumed..end].to_vec(); - self.consumed += self.seq_len; - *count += 1; - if !visit(sample)? { - return Ok(false); - } - } - Ok(true) - } -} - -fn push_documents( - documents: &mut Vec, - tokenizer: &Tokenizer, - packer: &mut SamplePacker, - count: &mut usize, - visit: &mut impl FnMut(Vec) -> Result, -) -> Result { - if documents.is_empty() { - return Ok(true); - } - let encodings = tokenizer.encode_batch(std::mem::take(documents), false)?; - for tokens in encodings { - let tokens = tokens - .into_iter() - .map(i64::from) - .chain(std::iter::once(i64::from(tokenizer.eos_token_id()))); - if !packer.push(tokens, count, visit)? { - return Ok(false); - } - } - Ok(true) -} - -/// Pack the EOS-joined token stream into fixed-length next-token samples. -fn visit_samples_in_order( - path: &Path, - tokenizer: &Tokenizer, - seq_len: usize, - mut visit: impl FnMut(Vec) -> Result, -) -> Result { - ensure!(seq_len > 0, "seq_len must be positive"); - let mut count = 0; - let mut packer = SamplePacker::new(seq_len); - let is_jsonl = path - .file_name() - .and_then(|name| name.to_str()) - .is_some_and(|name| name.ends_with(".jsonl") || name.ends_with(".jsonl.zst")); - let mut reader = open_data(path)?; - if is_jsonl { - let mut documents = Vec::with_capacity(TOKENIZE_BATCH); - let mut line = String::new(); - let mut line_number = 0; - loop { - line.clear(); - if reader.read_line(&mut line)? == 0 { - break; - } - line_number += 1; - if line.trim().is_empty() { - continue; - } - let value: serde_json::Value = serde_json::from_str(&line) - .with_context(|| format!("invalid JSONL at {}:{line_number}", path.display()))?; - let document = value - .get("text") - .and_then(|value| value.as_str()) - .with_context(|| { - format!( - "JSONL row at {}:{line_number} must contain a string `text` field", - path.display() - ) - })?; - documents.push(document.to_owned()); - if documents.len() == TOKENIZE_BATCH - && !push_documents( - &mut documents, - tokenizer, - &mut packer, - &mut count, - &mut visit, - )? - { - return Ok(count); - } - } - if !push_documents( - &mut documents, - tokenizer, - &mut packer, - &mut count, - &mut visit, - )? { - return Ok(count); - } - } else { - let mut document = String::new(); - reader.read_to_string(&mut document)?; - if !push_documents( - &mut vec![document], - tokenizer, - &mut packer, - &mut count, - &mut visit, - )? { - return Ok(count); - } - } - Ok(count) -} - -fn visit_samples( - path: &Path, - tokenizer: &Tokenizer, - seq_len: usize, - shuffle_buffer: usize, - seed: u64, - mut visit: impl FnMut(Vec) -> Result, -) -> Result { - if shuffle_buffer == 0 { - return visit_samples_in_order(path, tokenizer, seq_len, visit); - } - - let mut shuffler = ShuffleBuffer::new(shuffle_buffer, seed); - let mut keep_going = true; - let count = visit_samples_in_order(path, tokenizer, seq_len, |sample| { - if let Some(sample) = shuffler.push(sample) { - keep_going = visit(sample)?; - } - Ok(keep_going) - })?; - - if keep_going { - for sample in shuffler.finish() { - if !visit(sample)? { - break; - } - } - } - Ok(count) -} - -fn count_samples(path: &Path, tokenizer: &Tokenizer, seq_len: usize) -> Result { - visit_samples_in_order(path, tokenizer, seq_len, |_| Ok(true)) -} - -fn make_batch( - samples: &[Vec], - seq_len: usize, - device: &Device, -) -> (Tensor<2, Int>, Tensor<2, Int>) { - let mut inputs = Vec::with_capacity(samples.len() * seq_len); - let mut targets = Vec::with_capacity(samples.len() * seq_len); - for sample in samples { - inputs.extend_from_slice(&sample[..seq_len]); - targets.extend_from_slice(&sample[1..]); - } - ( - Tensor::from_data(TensorData::new(inputs, [samples.len(), seq_len]), device), - Tensor::from_data(TensorData::new(targets, [samples.len(), seq_len]), device), - ) -} - struct SquaredGradientNorm<'a> { grads: &'a GradientsParams, sum: Option>, @@ -428,100 +184,36 @@ fn learning_rate(args: &TrainArgs, step: usize, total_steps: usize) -> f64 { min_lr + cosine * (args.lr - min_lr) } -fn parameter_ids(model: &Transformer) -> Vec { - burn::module::list_param_ids(model) - .into_iter() - .map(|id| id.val()) - .collect() -} - -struct ParameterIdMapper<'a> { - ids: std::slice::Iter<'a, u64>, -} - -impl ModuleMapper for ParameterIdMapper<'_> { - fn map_float(&mut self, param: Param>) -> Param> { - let (_, tensor, mapper) = param.consume(); - let id = self - .ids - .next() - .copied() - .expect("checkpoint contains too few parameter IDs"); - Param::from_mapped_value(ParamId::from(id), tensor, mapper) - } -} - -fn restore_parameter_ids(model: &mut Transformer, ids: &[u64]) -> Result<()> { +fn validate_train_args(args: &TrainArgs) -> Result<()> { ensure!( - ids.len() == burn::module::list_param_ids(model).len(), - "checkpoint has {} parameter IDs, model has {}", - ids.len(), - burn::module::list_param_ids(model).len() + !args.data.is_empty(), + "at least one data source is required" ); - let mut mapper = ParameterIdMapper { ids: ids.iter() }; - *model = model.clone().map(&mut mapper); + ensure!(args.batch_size > 0, "batch_size must be positive"); + ensure!(args.grad_accum > 0, "grad_accum must be positive"); + ensure!(args.epochs > 0, "epochs must be positive"); + ensure!(args.seq_len > 0, "seq_len must be positive"); ensure!( - mapper.ids.next().is_none(), - "checkpoint contains too many parameter IDs" + args.lr.is_finite() && args.lr > 0.0, + "lr must be finite and positive" + ); + ensure!( + args.weight_decay.is_finite() && args.weight_decay >= 0.0, + "weight_decay must be finite and non-negative" + ); + ensure!( + args.grad_clip.is_finite() && args.grad_clip >= 0.0, + "grad_clip must be finite and non-negative" ); - Ok(()) -} - -fn save_training_checkpoint( - model: &Transformer, - adamw: &AdamWOptimizer, - muon: &BatchedMuon, - state: &TrainingState, - output: &Path, -) -> Result<()> { - let marker = output.join(".checkpoint-in-progress"); - let weights_temporary = output.join("weights.safetensors.tmp"); - let adamw_temporary = output.join("adamw-state.bpk.tmp"); - let muon_temporary = output.join("muon-state.bpk.tmp"); - let state_temporary = output.join("training-state.json.tmp"); - - fs::write(&marker, state.step.to_string())?; - save_safetensors(&model.clone().valid(), &weights_temporary)?; - adamw - .save(&adamw_temporary) - .context("failed to save AdamW state")?; - muon.save(&muon_temporary)?; - fs::write(&state_temporary, serde_json::to_vec_pretty(&state)?)?; - fs::rename(weights_temporary, output.join("weights.safetensors"))?; - fs::rename(adamw_temporary, output.join("adamw-state.bpk"))?; - fs::rename(muon_temporary, output.join("muon-state.bpk"))?; - fs::rename(state_temporary, output.join("training-state.json"))?; - fs::remove_file(marker)?; - Ok(()) -} - -fn load_training_state( - model: &mut Transformer, - adamw: AdamWOptimizer, - muon: &mut BatchedMuon, - output: &Path, - device: &Device, -) -> Result<(AdamWOptimizer, TrainingState)> { ensure!( - !output.join(".checkpoint-in-progress").exists(), - "checkpoint was interrupted while being saved" + args.max_steps.is_none_or(|steps| steps > 0), + "max_steps must be positive when set" ); - let state: TrainingState = - serde_json::from_slice(&fs::read(output.join("training-state.json"))?)?; - restore_parameter_ids(model, &state.parameter_ids)?; - load_safetensors(model, output.join("weights.safetensors"))?; - muon.set_parameter_ids(model.muon_parameter_ids()); - muon.load(output.join("muon-state.bpk"), &device.clone().inner())?; - let adamw = adamw - .load(output.join("adamw-state.bpk")) - .context("failed to load AdamW state")?; - Ok((adamw, state)) + Ok(()) } fn train(args: TrainArgs) -> Result<()> { - ensure!(args.batch_size > 0, "batch_size must be positive"); - ensure!(args.grad_accum > 0, "grad_accum must be positive"); - ensure!(args.epochs > 0, "epochs must be positive"); + validate_train_args(&args)?; let tokenizer = Tokenizer::from_file(&args.tokenizer)?; let mut config = load_config(&args.config)?; @@ -842,8 +534,9 @@ fn main() -> Result<()> { #[cfg(test)] mod tests { + use burn::module::{AutodiffModule, ParamId}; + use burn::tensor::{Int, TensorData}; use hermes_llm::get_builtin_model; - use std::io::Cursor; use super::*; @@ -865,6 +558,51 @@ mod tests { config } + fn valid_train_args() -> TrainArgs { + TrainArgs { + config: "config.mal".into(), + tokenizer: "tokenizer.json".into(), + data: vec!["corpus.jsonl".into()], + shuffle_buffer: 8, + output: "checkpoint".into(), + batch_size: 2, + grad_accum: 1, + epochs: 1, + seq_len: 8, + lr: 3e-4, + weight_decay: 0.1, + grad_clip: 1.0, + warmup_steps: 10, + schedule: Schedule::Wsd, + max_steps: Some(1), + checkpoint_every: 0, + checkpoint: None, + resume: false, + seed: 0, + } + } + + #[test] + fn invalid_numeric_training_arguments_fail_before_loading_files() { + type Invalidate = fn(&mut TrainArgs); + let cases: [(&str, Invalidate); 8] = [ + ("batch_size", |args| args.batch_size = 0), + ("grad_accum", |args| args.grad_accum = 0), + ("epochs", |args| args.epochs = 0), + ("seq_len", |args| args.seq_len = 0), + ("lr", |args| args.lr = f64::NAN), + ("weight_decay", |args| args.weight_decay = -0.1), + ("grad_clip", |args| args.grad_clip = f32::INFINITY), + ("max_steps", |args| args.max_steps = Some(0)), + ]; + for (field, invalidate) in cases { + let mut args = valid_train_args(); + invalidate(&mut args); + let err = validate_train_args(&args).unwrap_err().to_string(); + assert!(err.contains(field), "{field}: {err}"); + } + } + #[test] fn training_decreases_loss_and_checkpoint_roundtrips() { let config = small_hybrid(); @@ -987,59 +725,4 @@ mod tests { .fold(0.0, f32::max); assert!(max_diff < 1e-6, "checkpoint max diff: {max_diff}"); } - - #[test] - fn zstd_data_reader_streams_decompressed_text() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("data.jsonl.zst"); - let source = b"{\"text\":\"one\"}\n{\"text\":\"two\"}\n"; - let compressed = zstd::stream::encode_all(Cursor::new(source), 1).unwrap(); - fs::write(&path, compressed).unwrap(); - - let mut reader = open_data(&path).unwrap(); - let mut decoded = String::new(); - reader.read_to_string(&mut decoded).unwrap(); - assert_eq!(decoded.as_bytes(), source); - } - - #[test] - fn streaming_shuffle_is_bounded_and_deterministic() { - let shuffle = |seed| { - let mut buffer = ShuffleBuffer::new(4, seed); - let mut output = Vec::new(); - for value in 0..32_i64 { - if let Some(sample) = buffer.push(vec![value]) { - output.push(sample[0]); - } - assert!(buffer.samples.len() <= 4); - } - output.extend(buffer.finish().into_iter().map(|sample| sample[0])); - output - }; - - let first = shuffle(7); - assert_eq!(first, shuffle(7)); - assert_ne!(first, (0..32_i64).collect::>()); - let mut sorted = first; - sorted.sort_unstable(); - assert_eq!(sorted, (0..32_i64).collect::>()); - } - - #[test] - fn sample_packer_joins_documents_without_dropping_tokens() { - let mut packer = SamplePacker::new(3); - let mut samples = Vec::new(); - let mut count = 0; - let mut collect = |sample| { - samples.push(sample); - Ok(true) - }; - - for document in [vec![1, 2, 0], vec![3, 4, 0], vec![5, 6, 0]] { - assert!(packer.push(document, &mut count, &mut collect).unwrap()); - } - - assert_eq!(count, 2); - assert_eq!(samples, [vec![1, 2, 0, 3], vec![3, 4, 0, 5]]); - } }