Audit LLM stack and harden training recovery

.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

.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
 

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

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
 

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,22 +106,22 @@ 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" }
 cubecl-zspace = { git = "https://github.com/ppodolsky/cubecl", rev = "bda6a68d86df41f468cac78d713f1563869daa8e" }
 
 # 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"]

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

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

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.

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`.

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