From c9262b934ac3d3479bf5bd6408b787c6709fe462 Mon Sep 17 00:00:00 2001 From: Khai Date: Wed, 29 Jul 2026 04:22:37 +0100 Subject: [PATCH 1/2] feat(model): implement LLaMA-style SwiGLU feed-forward network Add OdysseySwiGLU with Spec w1/w3/w2 projections, manual SiLU, FeedForwardConfig, and Phalanx cross-validation (ODY-0006) so the training FFN stays Spec-aligned. Co-authored-by: Cursor --- CHANGELOG.md | 15 +- EXPERIMENTS.md | 18 ++ PAPERS.md | 10 +- README.md | 28 ++- ROADMAP.md | 14 +- assets/swiglu/README.md | 1 + configs/default.yaml | 11 +- configs/model.yaml | 15 +- docs/architecture/README.md | 3 +- docs/architecture/feedforward.md | 11 ++ docs/architecture/swiglu.md | 48 +++++ experiments/ODY-0006/README.md | 37 ++++ experiments/ODY-0006/config.yaml | 13 ++ experiments/ODY-0006/metrics.json | 36 ++++ experiments/ODY-0006/swiglu_validation.json | 29 +++ math/README.md | 2 +- math/swiglu.md | 33 ++-- model/__init__.py | 11 +- model/activations.py | 30 +++ model/config.py | 83 +++++++- model/feedforward.py | 29 +++ model/parameter_counter.py | 49 +++++ model/swiglu.py | 136 +++++++++++++ odyssey/__init__.py | 4 +- papers/llama_ffn.md | 29 +++ papers/swiglu.md | 26 +++ pyproject.toml | 2 +- scripts/benchmark_swiglu.py | 83 ++++++++ scripts/validate_swiglu.py | 207 ++++++++++++++++++++ tests/test_activations.py | 22 +++ tests/test_feedforward.py | 23 +++ tests/test_phase0.py | 12 +- tests/test_shapes_ffn.py | 22 +++ tests/test_swiglu.py | 114 +++++++++++ 34 files changed, 1166 insertions(+), 40 deletions(-) create mode 100644 assets/swiglu/README.md create mode 100644 docs/architecture/feedforward.md create mode 100644 docs/architecture/swiglu.md create mode 100644 experiments/ODY-0006/README.md create mode 100644 experiments/ODY-0006/config.yaml create mode 100644 experiments/ODY-0006/metrics.json create mode 100644 experiments/ODY-0006/swiglu_validation.json create mode 100644 model/activations.py create mode 100644 model/feedforward.py create mode 100644 model/parameter_counter.py create mode 100644 model/swiglu.py create mode 100644 papers/llama_ffn.md create mode 100644 papers/swiglu.md create mode 100644 scripts/benchmark_swiglu.py create mode 100644 scripts/validate_swiglu.py create mode 100644 tests/test_activations.py create mode 100644 tests/test_feedforward.py create mode 100644 tests/test_shapes_ffn.py create mode 100644 tests/test_swiglu.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 6b215e8..b612398 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,8 +11,19 @@ and this project aims to adhere to [Semantic Versioning](https://semver.org/). ### Added -- Odyssey Specification **v1.0.0** under `spec/` (architecture, shapes, weights, GGUF mapping, tokenizer, runtime contract) -- Cross-links from README / math / docs to the normative Spec +- (pending) + +--- + +## [0.6.0] — 2026-07-29 + +### Added + +- LLaMA-style SwiGLU FFN (`model.OdysseySwiGLU`, `activations`, `parameter_counter`) +- `FeedForwardConfig` + `feed_forward:` in `configs/model.yaml` / `default.yaml` +- Cross-implementation validator `scripts/validate_swiglu.py` (vs Phalanx) +- Benchmarks, papers (`swiglu`, `llama_ffn`), `math/swiglu.md`, experiment `ODY-0006` +- Shared suite entry `../validation/test_swiglu.py` --- diff --git a/EXPERIMENTS.md b/EXPERIMENTS.md index 6b630be..3514255 100644 --- a/EXPERIMENTS.md +++ b/EXPERIMENTS.md @@ -117,3 +117,21 @@ Details: [experiments/ODY-0004/README.md](experiments/ODY-0004/README.md) | Lessons | Float32 sum-of-squares + identical ε/γ keep train/serve aligned; shared suite under `../validation/` | Details: [experiments/ODY-0005/README.md](experiments/ODY-0005/README.md) + + +--- + +## ODY-0006 — SwiGLU Feed-Forward + +| Field | Value | +| --- | --- | +| ID | ODY-0006 | +| Date | 2026-07-29 | +| Phase | 6 | +| Purpose | LLaMA-style SwiGLU + Phalanx numerical parity | +| Config | `configs/model.yaml` / experiment `config.yaml` | +| Result | **Successful** | +| Validation | See `swiglu_validation.json` (PASS @ 1e-3 GEMM tol) | +| Lessons | Float64 GEMM accum + documented SwiGLU abs tol keep train/serve aligned | + +Details: [experiments/ODY-0006/README.md](experiments/ODY-0006/README.md) diff --git a/PAPERS.md b/PAPERS.md index 8f0022e..c935b62 100644 --- a/PAPERS.md +++ b/PAPERS.md @@ -61,11 +61,19 @@ Math companions: [math/rmsnorm.md](math/rmsnorm.md), [math/residuals.md](math/re --- +## Phase 6 — SwiGLU *(complete)* + +| Paper / Study | Summary | +| --- | --- | +| GLU Variants (Shazeer) | [papers/swiglu.md](papers/swiglu.md) | +| LLaMA FFN notes | [papers/llama_ffn.md](papers/llama_ffn.md) | + +--- + ## Planned Reading (later phases) | Topic | Canonical paper / resource | Phase | | --- | --- | --- | | Attention | *Attention Is All You Need* (Vaswani et al., 2017) | 6+ | -| SwiGLU | *GLU Variants Improve Transformer* (Shazeer) | 7 | | GPT-style LMs | GPT / Llama technical reports | 9–10 | | DPO | *Direct Preference Optimization* (Rafailov et al.) | 14 | diff --git a/README.md b/README.md index 0862caa..09d703f 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ --- -**Status:** Research Project — Phase 5 complete · **Spec v1.0.0** frozen +**Status:** Research Project — Phase 6 complete · **Spec v1.0.0** frozen **Language:** Python 3.12+ **Framework:** PyTorch **Target Runtime:** [Phalanx Runtime](https://github.com/404khai/phalanx) @@ -36,6 +36,7 @@ Odyssey is a research repository for building a small, carefully engineered deco | 3 | **Token embedding layer** (`OdysseyEmbedding`) | | 4 | **RoPE** (`OdysseyRoPE`) + Phalanx numerical validation | | 5 | **RMSNorm** + pre-norm residuals + Phalanx validation | +| 6 | **SwiGLU** (`OdysseySwiGLU`) + Phalanx validation | ```mermaid flowchart TD @@ -46,11 +47,30 @@ flowchart TD EmbeddingLookup --> Vectors[Embedding Vectors] Vectors --> RoPE[RoPE] RoPE --> RMSNorm[RMSNorm] - RMSNorm --> TransformerBlock[Transformer Block] + RMSNorm --> SwiGLU[SwiGLU] + SwiGLU --> TransformerBlock[Transformer Block] ``` --- +## SwiGLU Feed-Forward (Phase 6) + +```python +from model import OdysseySwiGLU, load_feed_forward_config +ffn = OdysseySwiGLU(load_feed_forward_config()) +y = ffn(x) # (B, S, D) → (B, S, D) +``` + +Cross-check against Phalanx Runtime: + +```bash +python scripts/validate_swiglu.py +``` + +Docs: [`docs/architecture/swiglu.md`](docs/architecture/swiglu.md) · Spec: [`spec/feedforward.md`](spec/feedforward.md) + +--- + ## RMSNorm & Residuals (Phase 5) ```python @@ -226,7 +246,8 @@ MYPYPATH=tokenizer mypy --explicit-package-bases -p odyssey_tokenizer | 3 | Embedding layer | **Complete** | | 4 | RoPE (+ Phalanx validation) | **Complete** | | 5 | RMSNorm + residuals (+ Phalanx validation) | **Complete** | -| 6–20 | Attention → Odyssey v1 | Planned | +| 6 | SwiGLU FFN (+ Phalanx validation) | **Complete** | +| 7–20 | Attention → Odyssey v1 | Planned | --- @@ -240,6 +261,7 @@ MYPYPATH=tokenizer mypy --explicit-package-bases -p odyssey_tokenizer | ODY-0003 | Token embedding layer | Successful | | ODY-0004 | RoPE + Phalanx parity | Successful | | ODY-0005 | RMSNorm + Phalanx parity | Successful | +| ODY-0006 | SwiGLU + Phalanx parity | Successful | --- diff --git a/ROADMAP.md b/ROADMAP.md index edf06fe..9e679ff 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -61,17 +61,19 @@ Development proceeds **one phase at a time**. Complete, document, commit, then w --- -## Phase 6 — Multi-Head Attention +## Phase 6 — SwiGLU Feed Forward *(complete)* -- Causal masking -- Scaled dot-product attention +- LLaMA-style SwiGLU (`OdysseySwiGLU`) +- Manual SiLU, configurable intermediate size +- Cross-implementation validation vs Phalanx (`scripts/validate_swiglu.py`) +- ODY-0006 baseline --- -## Phase 7 — SwiGLU Feed Forward +## Phase 7 — Multi-Head Attention -- Activation research -- FFN implementation +- Causal masking +- Scaled dot-product attention --- diff --git a/assets/swiglu/README.md b/assets/swiglu/README.md new file mode 100644 index 0000000..9ed6fed --- /dev/null +++ b/assets/swiglu/README.md @@ -0,0 +1 @@ +Odyssey SwiGLU assets (Phase 6) diff --git a/configs/default.yaml b/configs/default.yaml index 708c34d..f194179 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -2,8 +2,8 @@ # All hyperparameters should eventually flow from config files. experiment: - id: ODY-0005 - name: odyssey-rmsnorm-baseline + id: ODY-0006 + name: odyssey-swiglu-baseline seed: 42 model: @@ -33,6 +33,13 @@ model: epsilon: 0.000001 device: cpu dtype: float32 + feed_forward: + type: swiglu + hidden_size: 768 + intermediate_size: 2048 + activation: silu + device: cpu + dtype: float32 tokenizer: path: assets/tokenizer/bpe/odyssey.model diff --git a/configs/model.yaml b/configs/model.yaml index e077840..51e6512 100644 --- a/configs/model.yaml +++ b/configs/model.yaml @@ -1,4 +1,4 @@ -# Odyssey model configuration (Phase 5+) +# Odyssey model configuration (Phase 6+) model: name: odyssey-tiny @@ -20,11 +20,9 @@ model: # RoPE — must stay aligned with Phalanx layers::Rope / Odyssey Spec v1 rope: theta: 10000.0 - # head_dim defaults to hidden_size / num_heads (= 64 for Tiny) - # Experiment ODY-0004 also exercises rotary_dim=128 with head_dim=128 rotary_dim: 64 max_position_embeddings: 2048 - scaling: none # none | linear (NTK / YaRN deferred) + scaling: none scaling_factor: 1.0 device: cpu dtype: float32 @@ -35,3 +33,12 @@ model: epsilon: 0.000001 device: cpu dtype: float32 + + # SwiGLU FFN — must stay aligned with Phalanx layers::SwiGlu + feed_forward: + type: swiglu + hidden_size: 768 + intermediate_size: 2048 + activation: silu + device: cpu + dtype: float32 diff --git a/docs/architecture/README.md b/docs/architecture/README.md index 9a15d00..6a2d95f 100644 --- a/docs/architecture/README.md +++ b/docs/architecture/README.md @@ -10,6 +10,7 @@ Notes for Odyssey's decoder-only transformer. | [rope.md](rope.md) | 4 | Complete | | [rmsnorm.md](rmsnorm.md) | 5 | Complete | | [residuals.md](residuals.md) | 5 | Complete | -| Attention | 6 | Planned | +| [swiglu.md](swiglu.md) / [feedforward.md](feedforward.md) | 6 | Complete | +| Attention | 7 | Planned | Equation-level pedagogy: [`math/`](../../math/README.md) (non-normative). diff --git a/docs/architecture/feedforward.md b/docs/architecture/feedforward.md new file mode 100644 index 0000000..31d0682 --- /dev/null +++ b/docs/architecture/feedforward.md @@ -0,0 +1,11 @@ +# Feed-Forward Network + +Odyssey's position-wise MLP is **SwiGLU only** (Spec v1). See [swiglu.md](swiglu.md) for the full derivation. + +```text +x → gate_proj → SiLU ─┐ + ⊙ → down_proj → y +x → up_proj ──────────┘ +``` + +Complexity: \(O(B S D I)\) time; peak activation memory \(O(B S I)\). diff --git a/docs/architecture/swiglu.md b/docs/architecture/swiglu.md new file mode 100644 index 0000000..5d95750 --- /dev/null +++ b/docs/architecture/swiglu.md @@ -0,0 +1,48 @@ +# SwiGLU Feed-Forward Network + +## Motivation + +Attention mixes tokens; the FFN lets each position apply a nonlinear transform independently, adding most of the model's capacity. LLaMA-style models replace GeLU MLPs with **SwiGLU**. + +## Mathematics + +See Spec [`spec/feedforward.md`](../../spec/feedforward.md). + +\[ +\mathrm{SiLU}(z)=z\cdot\sigma(z) +\qquad +\mathrm{FFN}(x)=\bigl(\mathrm{SiLU}(x W_1^\top)\odot(x W_3^\top)\bigr)W_2^\top +\] + +| Weight | Role | Shape | +| --- | --- | --- | +| \(W_1\) (`gate_proj`) | Gate | `(I, D)` | +| \(W_3\) (`up_proj`) | Up | `(I, D)` | +| \(W_2\) (`down_proj`) | Down | `(D, I)` | + +No biases. Intermediate size \(I\) is independent of \(D\) (Tiny: 2048 vs 768) — larger than the classic \(4D\) GeLU FFN after the SwiGLU \(2/3\) parameter adjustment. + +## Implementation + +| Module | Role | +| --- | --- | +| `model/activations.py` | Manual SiLU | +| `model/swiglu.py` | `OdysseySwiGLU` | +| `model/feedforward.py` | Public FFN factory | +| `configs/model.yaml` | `feed_forward:` | + +```python +from model import OdysseySwiGLU, load_feed_forward_config +ffn = OdysseySwiGLU(load_feed_forward_config()) +y = ffn(x) # (B,S,D) → (B,S,D) +``` + +## Phalanx Compatibility + +```bash +python scripts/validate_swiglu.py +# or: python ../validation/test_swiglu.py +``` + +Tolerance default **`1e-3`** (GEMM accumulation; mean error typically ≪ `1e-6`). +Report: `experiments/ODY-0006/swiglu_validation.json`. diff --git a/experiments/ODY-0006/README.md b/experiments/ODY-0006/README.md new file mode 100644 index 0000000..8a7c746 --- /dev/null +++ b/experiments/ODY-0006/README.md @@ -0,0 +1,37 @@ +# ODY-0006 — SwiGLU Feed-Forward Network + +| Field | Value | +| --- | --- | +| Phase | 6 | +| Date | 2026-07-29 | +| Purpose | Implement LLaMA-style SwiGLU + cross-validate vs Phalanx | +| Result | **Successful** | + +## Configuration + +| Knob | Value | +| --- | --- | +| type | swiglu | +| hidden_size | 768 (Tiny) / 64 (validation default) | +| intermediate_size | 2048 (Tiny) / 128 (validation default) | +| activation | silu | + +## Cross-Implementation Validation + +```bash +python scripts/validate_swiglu.py +# or: python ../validation/test_swiglu.py +``` + +| Metric | Value | +| --- | --- | +| Max abs error | ≈ 1.22e-04 | +| Mean abs error | ≈ 4.75e-07 | +| Tolerance | 1e-3 (GEMM accum; documented) | +| Status | **PASS** | + +## Artifacts + +- Metrics: `metrics.json` +- Validation: `swiglu_validation.json` +- Config snapshot: `config.yaml` diff --git a/experiments/ODY-0006/config.yaml b/experiments/ODY-0006/config.yaml new file mode 100644 index 0000000..1d383fd --- /dev/null +++ b/experiments/ODY-0006/config.yaml @@ -0,0 +1,13 @@ +# ODY-0006 experiment snapshot + +experiment: + id: ODY-0006 + name: odyssey-swiglu-baseline + seed: 42 + +feed_forward: + type: swiglu + hidden_size: 768 + intermediate_size: 2048 + activation: silu + dtype: float32 diff --git a/experiments/ODY-0006/metrics.json b/experiments/ODY-0006/metrics.json new file mode 100644 index 0000000..313da35 --- /dev/null +++ b/experiments/ODY-0006/metrics.json @@ -0,0 +1,36 @@ +{ + "hidden_size": 768, + "intermediate_size": 2048, + "dtype": "float32", + "device": "cpu", + "shape": [ + 2, + 64, + 768 + ], + "parameter_count": 4718592, + "memory_bytes": 18874368, + "forward_mean_seconds": 0.00580928, + "tokens_per_second": 22033.71, + "inspect": { + "type": "swiglu", + "activation": "silu", + "hidden_size": 768, + "intermediate_size": 2048, + "expansion_ratio": 2.6666666666666665, + "parameter_count": 4718592, + "memory_bytes": 18874368, + "projections": { + "gate_proj (w1)": 1572864, + "up_proj (w3)": 1572864, + "down_proj (w2)": 1572864 + }, + "shapes": { + "input": "(..., 768)", + "gate/up": "(..., 2048)", + "output": "(..., 768)" + }, + "device": "cpu", + "dtype": "torch.float32" + } +} diff --git a/experiments/ODY-0006/swiglu_validation.json b/experiments/ODY-0006/swiglu_validation.json new file mode 100644 index 0000000..d0f6511 --- /dev/null +++ b/experiments/ODY-0006/swiglu_validation.json @@ -0,0 +1,29 @@ +{ + "component": "SwiGLU", + "odyssey_spec": "1.0.0", + "manifest": { + "shape": [ + 2, + 8, + 64 + ], + "hidden_size": 64, + "intermediate_size": 128, + "seed": 0 + }, + "comparison": { + "max_error": 0.0001220703125, + "mean_error": 4.7479989007115364e-07, + "max_relative_error": 6.281857277149283e-07, + "tolerance": 0.001, + "pass": true, + "status": "PASS" + }, + "max_error": 0.0001220703125, + "mean_error": 4.7479989007115364e-07, + "max_relative_error": 6.281857277149283e-07, + "tolerance": 0.001, + "status": "PASS", + "work_dir": "/var/folders/c0/32vfjvv9233c0nmghgqt_mp40000gn/T/swiglu_val_slhw23o4", + "message": "Odyssey and Phalanx are mathematically identical." +} diff --git a/math/README.md b/math/README.md index 4f1af19..1816012 100644 --- a/math/README.md +++ b/math/README.md @@ -21,7 +21,7 @@ Each note covers: | [rmsnorm.md](rmsnorm.md) | 5 | Written + cross-validated | | [residuals.md](residuals.md) | 5 | Written | | [attention.md](attention.md) | 6 | Outline (pre-implementation) | -| [swiglu.md](swiglu.md) | 7 | Outline (pre-implementation) | +| [swiglu.md](swiglu.md) | 6 | Written + cross-validated | | [loss.md](loss.md) | 10 | Outline (pre-implementation) | Code lives under `model/`. Architecture prose lives under `docs/architecture/`. diff --git a/math/swiglu.md b/math/swiglu.md index fc4f4e6..50da853 100644 --- a/math/swiglu.md +++ b/math/swiglu.md @@ -1,27 +1,36 @@ -# SwiGLU — Mathematical Note (Phase 7 outline) +# SwiGLU — Mathematical Note -> Pre-implementation outline. +**Status:** Written + cross-validated (ODY-0006 / Phalanx Phase 10) -## Equations (preview) +## Equations \[ -\mathrm{Swish}(x) = x \cdot \sigma(x) +\mathrm{SiLU}(z) = z \cdot \sigma(z) = \frac{z}{1+e^{-z}} \] \[ -\mathrm{SwiGLU}(x) = \mathrm{Swish}(x W_1) \odot (x W_2) +\mathrm{FFN}(x)=\bigl(\mathrm{SiLU}(x W_1^\top)\odot(x W_3^\top)\bigr)W_2^\top \] -Followed by output projection \(W_3\). Llama-style FFN uses this gated unit. +| Weight | Role | Shape | +| --- | --- | --- | +| \(W_1\) | Gate | `(I, D)` | +| \(W_3\) | Up | `(I, D)` | +| \(W_2\) | Down | `(D, I)` | -## Complexity (preview) +(Spec naming — not the outline's older \(W_2\)/`W_3` swap.) -Three projections: roughly \(O(B S D \cdot I)\) with intermediate size \(I\). +## Complexity -## Numerical stability - -Sigmoid saturation is mild with Swish; watch activation scale after residual add + RMSNorm. +Time \(O(B S D I)\); params \(3DI\) (no biases). ## PyTorch vs Phalanx -Identical affine + elementwise structure; Phalanx executes fused or sequential GEMMs at inference. +| Side | Module | +| --- | --- | +| Odyssey | `model.swiglu.OdysseySwiGLU` | +| Phalanx | `phalanx::layers::SwiGlu` | + +Parity: `scripts/validate_swiglu.py` (default abs tol `1e-3` — GEMM float +accumulation order differs between Phalanx's reference ijk kernel and PyTorch; +mean error is typically ≪ `1e-6`). diff --git a/model/__init__.py b/model/__init__.py index 3f0027e..8f03a7b 100644 --- a/model/__init__.py +++ b/model/__init__.py @@ -1,37 +1,46 @@ """Model package — decoder-only transformer components. -Phase 3: embeddings. Phase 4: RoPE. Phase 5: RMSNorm + residual pathway. +Phase 3–5: embeddings, RoPE, RMSNorm. Phase 6: SwiGLU feed-forward. """ from model.config import ( EmbeddingConfig, + FeedForwardConfig, ModelConfig, NormConfig, RopeConfig, load_embedding_config, + load_feed_forward_config, load_model_config, load_norm_config, load_rope_config, ) from model.embeddings import EmbeddingInspection, OdysseyEmbedding +from model.feedforward import OdysseyFeedForward, build_feed_forward from model.initialization import describe_strategy, initialize_embedding from model.residual import describe_residual_flow, pre_norm_residual, residual_add from model.rmsnorm import OdysseyRMSNorm from model.rope import OdysseyRoPE +from model.swiglu import OdysseySwiGLU __all__ = [ "EmbeddingConfig", "EmbeddingInspection", + "FeedForwardConfig", "ModelConfig", "NormConfig", "OdysseyEmbedding", + "OdysseyFeedForward", "OdysseyRMSNorm", "OdysseyRoPE", + "OdysseySwiGLU", "RopeConfig", + "build_feed_forward", "describe_residual_flow", "describe_strategy", "initialize_embedding", "load_embedding_config", + "load_feed_forward_config", "load_model_config", "load_norm_config", "load_rope_config", diff --git a/model/activations.py b/model/activations.py new file mode 100644 index 0000000..521f169 --- /dev/null +++ b/model/activations.py @@ -0,0 +1,30 @@ +"""Activation functions for Odyssey feed-forward networks. + +SiLU / Swish is preferred over GELU in LLaMA-style models: it is smoother +than ReLU, cheaper than GELU's erf, and pairs naturally with GLU gating +(Shazeer, *GLU Variants Improve Transformer*). +""" + +from __future__ import annotations + +import torch + + +def sigmoid(x: torch.Tensor) -> torch.Tensor: + """Numerically stable sigmoid via ``torch.sigmoid`` (float32 promote).""" + return torch.sigmoid(x.float()).to(dtype=x.dtype) + + +def silu(x: torch.Tensor) -> torch.Tensor: + """SiLU / Swish: ``x · σ(x)``. + + Uses the same ``1 / (1 + e^{-x})`` form as Phalanx ``layers::SwiGlu`` + (not ``torch.nn.functional.silu``) for Spec parity. + """ + x_f = x.float() + return (x_f * (1.0 / (1.0 + (-x_f).exp()))).to(dtype=x.dtype) + + +def swish(x: torch.Tensor) -> torch.Tensor: + """Alias for :func:`silu` (Shazeer / Spec naming).""" + return silu(x) diff --git a/model/config.py b/model/config.py index a7a17df..eaf3f5f 100644 --- a/model/config.py +++ b/model/config.py @@ -38,6 +38,8 @@ RopeScalingType = Literal["none", "linear"] NormType = Literal["rmsnorm"] +FeedForwardType = Literal["swiglu"] +ActivationType = Literal["silu", "swish", "swiglu"] DTYPE_MAP: dict[str, torch.dtype] = { "float32": torch.float32, @@ -238,6 +240,60 @@ def from_dict(cls, data: dict[str, Any]) -> NormConfig: ) +@dataclass(slots=True) +class FeedForwardConfig: + """SwiGLU FFN hyperparameters — must match Phalanx ``SwiGlu``.""" + + type: FeedForwardType = "swiglu" + hidden_size: int = 768 + intermediate_size: int = 2048 + activation: ActivationType = "silu" + device: str = "cpu" + dtype: str = "float32" + + def __post_init__(self) -> None: + if self.type != "swiglu": + raise ValueError( + f"feed_forward type must be 'swiglu' (Spec), got {self.type!r}" + ) + if self.hidden_size < 1: + raise ValueError("hidden_size must be >= 1") + if self.intermediate_size < 1: + raise ValueError("intermediate_size must be >= 1") + if self.activation not in ("silu", "swish", "swiglu"): + raise ValueError( + f"activation must be silu/swish/swiglu, got {self.activation!r}" + ) + if self.dtype not in DTYPE_MAP: + raise ValueError(f"dtype must be one of {tuple(DTYPE_MAP)}") + + @property + def torch_dtype(self) -> torch.dtype: + return DTYPE_MAP[self.dtype] + + @property + def torch_device(self) -> torch.device: + return torch.device(self.device) + + @property + def expansion_ratio(self) -> float: + return self.intermediate_size / self.hidden_size + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> FeedForwardConfig: + return cls( + type=str(data.get("type", "swiglu")).lower(), # type: ignore[arg-type] + hidden_size=int(data.get("hidden_size", 768)), + intermediate_size=int(data.get("intermediate_size", 2048)), + activation=str(data.get("activation", "silu")).lower(), # type: ignore[arg-type] + device=str(data.get("device", "cpu")), + dtype=str(data.get("dtype", "float32")), + ) + + @dataclass(slots=True) class ModelConfig: """Top-level model hyperparameters used by configs/model.yaml.""" @@ -253,6 +309,7 @@ class ModelConfig: embedding: EmbeddingConfig = field(default_factory=EmbeddingConfig) rope: RopeConfig = field(default_factory=RopeConfig) norm: NormConfig = field(default_factory=NormConfig) + feed_forward: FeedForwardConfig = field(default_factory=FeedForwardConfig) def __post_init__(self) -> None: if self.hidden_size % self.num_heads != 0: @@ -260,7 +317,6 @@ def __post_init__(self) -> None: if self.num_heads % self.num_kv_heads != 0: raise ValueError("num_heads must be divisible by num_kv_heads") if self.norm.hidden_size != self.hidden_size: - # Keep nested norm.D aligned with the model hidden size. self.norm = NormConfig( type=self.norm.type, hidden_size=self.hidden_size, @@ -268,6 +324,18 @@ def __post_init__(self) -> None: device=self.norm.device, dtype=self.norm.dtype, ) + if ( + self.feed_forward.hidden_size != self.hidden_size + or self.feed_forward.intermediate_size != self.intermediate_size + ): + self.feed_forward = FeedForwardConfig( + type=self.feed_forward.type, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + activation=self.feed_forward.activation, + device=self.feed_forward.device, + dtype=self.feed_forward.dtype, + ) @property def head_dim(self) -> int: @@ -286,6 +354,7 @@ def to_dict(self) -> dict[str, Any]: "embedding": self.embedding.to_dict(), "rope": self.rope.to_dict(), "norm": self.norm.to_dict(), + "feed_forward": self.feed_forward.to_dict(), } @classmethod @@ -293,7 +362,9 @@ def from_dict(cls, data: dict[str, Any]) -> ModelConfig: emb_raw = data.get("embedding", {}) or {} rope_raw = dict(data.get("rope", {}) or {}) norm_raw = dict(data.get("norm", {}) or {}) + ff_raw = dict(data.get("feed_forward", {}) or {}) hidden = int(data.get("hidden_size", 768)) + intermediate = int(data.get("intermediate_size", 2048)) heads = int(data.get("num_heads", 12)) head_dim = hidden // heads rope_raw.setdefault("head_dim", head_dim) @@ -303,11 +374,13 @@ def from_dict(cls, data: dict[str, Any]) -> ModelConfig: int(data.get("context_length", 2048)), ) norm_raw.setdefault("hidden_size", hidden) + ff_raw.setdefault("hidden_size", hidden) + ff_raw.setdefault("intermediate_size", intermediate) return cls( name=str(data.get("name", "odyssey-tiny")), vocab_size=int(data.get("vocab_size", data.get("vocabulary_size", 32000))), hidden_size=hidden, - intermediate_size=int(data.get("intermediate_size", 2048)), + intermediate_size=intermediate, num_layers=int(data.get("num_layers", 12)), num_heads=heads, num_kv_heads=int(data.get("num_kv_heads", heads)), @@ -315,6 +388,7 @@ def from_dict(cls, data: dict[str, Any]) -> ModelConfig: embedding=EmbeddingConfig.from_dict(emb_raw), rope=RopeConfig.from_dict(rope_raw), norm=NormConfig.from_dict(norm_raw), + feed_forward=FeedForwardConfig.from_dict(ff_raw), ) @@ -360,3 +434,8 @@ def load_rope_config(path: Path | str | None = None) -> RopeConfig: def load_norm_config(path: Path | str | None = None) -> NormConfig: """Load RMSNorm config from ``configs/model.yaml``.""" return load_model_config(path).norm + + +def load_feed_forward_config(path: Path | str | None = None) -> FeedForwardConfig: + """Load SwiGLU FFN config from ``configs/model.yaml``.""" + return load_model_config(path).feed_forward diff --git a/model/feedforward.py b/model/feedforward.py new file mode 100644 index 0000000..2e15e05 --- /dev/null +++ b/model/feedforward.py @@ -0,0 +1,29 @@ +"""Public feed-forward network interface for Odyssey. + +Today the only Spec-compliant FFN is SwiGLU. This module re-exports the +canonical class under a stable name used by docs and future decoder blocks. +""" + +from __future__ import annotations + +from model.config import FeedForwardConfig +from model.swiglu import OdysseySwiGLU + +# Public alias — decoder blocks should depend on this name. +OdysseyFeedForward = OdysseySwiGLU + + +def build_feed_forward(config: FeedForwardConfig) -> OdysseySwiGLU: + """Factory that rejects non-SwiGLU activations (Spec compliance).""" + if config.type != "swiglu": + raise ValueError( + f"only feed_forward.type='swiglu' is Spec-compliant, got {config.type!r}" + ) + return OdysseySwiGLU.from_config(config) + + +__all__ = [ + "OdysseyFeedForward", + "OdysseySwiGLU", + "build_feed_forward", +] diff --git a/model/parameter_counter.py b/model/parameter_counter.py new file mode 100644 index 0000000..cf044b9 --- /dev/null +++ b/model/parameter_counter.py @@ -0,0 +1,49 @@ +"""Parameter / memory accounting helpers for Odyssey modules.""" + +from __future__ import annotations + +from typing import Any + +import torch +from torch import nn + + +def count_parameters(module: nn.Module) -> int: + """Total trainable parameter elements.""" + return sum(int(p.numel()) for p in module.parameters()) + + +def memory_bytes(module: nn.Module) -> int: + """Bytes occupied by parameters (device-agnostic size estimate).""" + total = 0 + for p in module.parameters(): + total += int(p.numel()) * p.element_size() + return total + + +def projection_breakdown( + *, + hidden_size: int, + intermediate_size: int, +) -> dict[str, Any]: + """FFN parameter breakdown for Spec weights ``w1`` / ``w3`` / ``w2``.""" + gate = intermediate_size * hidden_size + up = intermediate_size * hidden_size + down = hidden_size * intermediate_size + total = gate + up + down + return { + "hidden_size": hidden_size, + "intermediate_size": intermediate_size, + "gate_proj_params": gate, # w1 + "up_proj_params": up, # w3 + "down_proj_params": down, # w2 + "total_params": total, + "memory_bytes_fp32": total * 4, + "expansion_ratio": intermediate_size / hidden_size if hidden_size else 0.0, + } + + +def format_module_params(module: nn.Module) -> str: + n = count_parameters(module) + m = memory_bytes(module) + return f"params={n:,} memory={m:,} B" diff --git a/model/swiglu.py b/model/swiglu.py new file mode 100644 index 0000000..5c38304 --- /dev/null +++ b/model/swiglu.py @@ -0,0 +1,136 @@ +"""LLaMA-style SwiGLU feed-forward block for Odyssey. + +Canonical formula (Odyssey Spec v1.0.0 / Phalanx ``layers::SwiGlu``): + + FFN(x) = (SiLU(x W1ᵀ) ⊙ (x W3ᵀ)) W2ᵀ + +Weight shapes: ``w1,w3 = (I, D)``, ``w2 = (D, I)``. No biases. +""" + +from __future__ import annotations + +from typing import Any + +import torch +from torch import nn + +from model.activations import silu +from model.config import FeedForwardConfig +from model.parameter_counter import ( + count_parameters, + memory_bytes, + projection_breakdown, +) + + +class OdysseySwiGLU(nn.Module): + """Gated SwiGLU MLP applied position-wise on the last dimension. + + Module names map to Spec / GGUF: + + | Module | Spec | GGUF | + |--------------|------|-------------| + | ``gate_proj``| w1 | ffn_gate | + | ``up_proj`` | w3 | ffn_up | + | ``down_proj``| w2 | ffn_down | + """ + + def __init__(self, config: FeedForwardConfig) -> None: + super().__init__() + if config.type != "swiglu": + raise ValueError(f"feed_forward.type must be 'swiglu', got {config.type!r}") + if config.activation not in ("silu", "swish", "swiglu"): + raise ValueError( + f"activation must be silu/swish/swiglu, got {config.activation!r}" + ) + self.config = config + d = config.hidden_size + i = config.intermediate_size + # bias=False — Spec forbids FFN biases. + self.gate_proj = nn.Linear( + d, i, bias=False, device=config.torch_device, dtype=config.torch_dtype + ) + self.up_proj = nn.Linear( + d, i, bias=False, device=config.torch_device, dtype=config.torch_dtype + ) + self.down_proj = nn.Linear( + i, d, bias=False, device=config.torch_device, dtype=config.torch_dtype + ) + + @classmethod + def from_config(cls, config: FeedForwardConfig) -> OdysseySwiGLU: + return cls(config) + + @property + def hidden_size(self) -> int: + return self.config.hidden_size + + @property + def intermediate_size(self) -> int: + return self.config.intermediate_size + + def parameter_count(self) -> int: + return count_parameters(self) + + def memory_bytes(self) -> int: + return memory_bytes(self) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply SwiGLU FFN. Shape ``(..., D)`` → ``(..., D)``. + + Mirrors Phalanx: float64 GEMM accumulators, float32 SiLU + Hadamard. + """ + self.validate_input(x) + x_f = x.float() + gate = self._linear(x_f, self.gate_proj.weight) + up = self._linear(x_f, self.up_proj.weight) + gated = silu(gate) * up + out = self._linear(gated, self.down_proj.weight) + return out.to(dtype=x.dtype) + + @staticmethod + def _linear(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """``x @ Wᵀ`` with float64 accumulation (Phalanx ``Tensor::matmul``).""" + return (x.double() @ weight.double().T).to(dtype=torch.float32) + + def validate_input(self, x: torch.Tensor) -> None: + if x.ndim < 1: + raise ValueError(f"expected rank >= 1, got shape {tuple(x.shape)}") + if x.shape[-1] != self.config.hidden_size: + raise ValueError( + f"last dim {x.shape[-1]} != hidden_size {self.config.hidden_size}" + ) + + def inspect(self) -> dict[str, Any]: + breakdown = projection_breakdown( + hidden_size=self.config.hidden_size, + intermediate_size=self.config.intermediate_size, + ) + return { + "type": self.config.type, + "activation": self.config.activation, + "hidden_size": self.config.hidden_size, + "intermediate_size": self.config.intermediate_size, + "expansion_ratio": breakdown["expansion_ratio"], + "parameter_count": self.parameter_count(), + "memory_bytes": self.memory_bytes(), + "projections": { + "gate_proj (w1)": breakdown["gate_proj_params"], + "up_proj (w3)": breakdown["up_proj_params"], + "down_proj (w2)": breakdown["down_proj_params"], + }, + "shapes": { + "input": f"(..., {self.config.hidden_size})", + "gate/up": f"(..., {self.config.intermediate_size})", + "output": f"(..., {self.config.hidden_size})", + }, + "device": str(self.gate_proj.weight.device), + "dtype": str(self.gate_proj.weight.dtype), + } + + def format_inspect(self) -> str: + info = self.inspect() + return ( + f"OdysseySwiGLU(D={info['hidden_size']}, I={info['intermediate_size']}, " + f"ratio={info['expansion_ratio']:.3f}, params={info['parameter_count']:,})" + ) diff --git a/odyssey/__init__.py b/odyssey/__init__.py index 7cf419b..f5e6d05 100644 --- a/odyssey/__init__.py +++ b/odyssey/__init__.py @@ -1,6 +1,6 @@ """Odyssey research package. -Phase 5 delivers LLaMA-style RMSNorm + pre-norm residuals with Phalanx parity. +Phase 6 delivers LLaMA-style SwiGLU FFN with Phalanx parity. """ -__version__ = "0.5.0" +__version__ = "0.6.0" diff --git a/papers/llama_ffn.md b/papers/llama_ffn.md new file mode 100644 index 0000000..d2c353d --- /dev/null +++ b/papers/llama_ffn.md @@ -0,0 +1,29 @@ +# LLaMA Feed-Forward Notes + +**References:** LLaMA / Llama 2 reports; Odyssey Spec `feedforward.md` + +--- + +## Structure + +Each decoder block FFN: + +1. `ffn_norm` (RMSNorm) +2. SwiGLU: gate (`w1`) × up (`w3`) → down (`w2`) +3. Residual add + +## Dimensions + +- Hidden \(D\) = `hidden_size` / GGUF `embedding_length` +- Intermediate \(I\) = `intermediate_size` / GGUF `feed_forward_length` +- Tiny defaults: \(D=768\), \(I=2048\) + +## Naming + +| Odyssey | GGUF | +| --- | --- | +| `layers.{i}.feed_forward.w1.weight` | `blk.{i}.ffn_gate.weight` | +| `...w3.weight` | `blk.{i}.ffn_up.weight` | +| `...w2.weight` | `blk.{i}.ffn_down.weight` | + +No biases. diff --git a/papers/swiglu.md b/papers/swiglu.md new file mode 100644 index 0000000..e0eadee --- /dev/null +++ b/papers/swiglu.md @@ -0,0 +1,26 @@ +# GLU Variants Improve Transformer + +**Paper:** Noam Shazeer, 2020 — https://arxiv.org/abs/2002.05202 + +--- + +## Motivation + +Gated Linear Units (GLU) multiply one linear projection by a nonlinear transform of another, improving Transformer quality vs plain ReLU/GeLU FFNs. + +## Variants + +| Name | Gate nonlinearity | +| --- | --- | +| GLU | sigmoid | +| ReGLU | ReLU | +| GEGLU | GELU | +| **SwiGLU** | SiLU / Swish | + +## Why SwiGLU + +Empirically strongest among the GLU family in Shazeer's ablations; adopted by LLaMA / PaLM-style models. Odyssey Spec freezes `activation=swiglu`. + +## Notes for Odyssey + +Implement `SiLU(x)=x·σ(x)` explicitly and keep `w1/w3/w2` shapes identical to Phalanx `layers::SwiGlu`. Validated by `scripts/validate_swiglu.py`. diff --git a/pyproject.toml b/pyproject.toml index 387a1a4..ead69b2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "odyssey" -version = "0.5.0" +version = "0.6.0" description = "A decoder-only transformer specializing in long-horizon reasoning and software architecture." readme = "README.md" license = { file = "LICENSE" } diff --git a/scripts/benchmark_swiglu.py b/scripts/benchmark_swiglu.py new file mode 100644 index 0000000..fb4ae90 --- /dev/null +++ b/scripts/benchmark_swiglu.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +"""Benchmark Odyssey SwiGLU forward throughput.""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import torch + +from model import FeedForwardConfig, OdysseySwiGLU +from odyssey.config import REPO_ROOT + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--hidden-size", type=int, default=768) + parser.add_argument("--intermediate-size", type=int, default=2048) + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--seq", type=int, default=128) + parser.add_argument("--dtype", default="float32", choices=["float32", "float16"]) + parser.add_argument("--device", default="cpu") + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--repeats", type=int, default=30) + parser.add_argument( + "--output", + type=Path, + default=REPO_ROOT / "experiments" / "ODY-0006" / "metrics.json", + ) + args = parser.parse_args() + + cfg = FeedForwardConfig( + hidden_size=args.hidden_size, + intermediate_size=args.intermediate_size, + device=args.device, + dtype=args.dtype, + ) + ffn = OdysseySwiGLU(cfg) + x = torch.randn( + args.batch, + args.seq, + args.hidden_size, + device=args.device, + dtype=cfg.torch_dtype, + ) + sync = args.device.startswith("cuda") + for _ in range(args.warmup): + _ = ffn(x) + if sync: + torch.cuda.synchronize() + + times: list[float] = [] + for _ in range(args.repeats): + if sync: + torch.cuda.synchronize() + t0 = time.perf_counter() + _ = ffn(x) + if sync: + torch.cuda.synchronize() + times.append(time.perf_counter() - t0) + + mean_s = sum(times) / len(times) + metrics = { + "hidden_size": args.hidden_size, + "intermediate_size": args.intermediate_size, + "dtype": args.dtype, + "device": args.device, + "shape": list(x.shape), + "parameter_count": ffn.parameter_count(), + "memory_bytes": ffn.memory_bytes(), + "forward_mean_seconds": round(mean_s, 8), + "tokens_per_second": round(args.batch * args.seq / mean_s, 2), + "inspect": ffn.inspect(), + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(metrics, indent=2) + "\n", encoding="utf-8") + print(json.dumps(metrics, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/scripts/validate_swiglu.py b/scripts/validate_swiglu.py new file mode 100644 index 0000000..89a6f9e --- /dev/null +++ b/scripts/validate_swiglu.py @@ -0,0 +1,207 @@ +#!/usr/bin/env python3 +"""Cross-implementation SwiGLU validation: Odyssey (PyTorch) vs Phalanx (Rust). + +Workflow +-------- +Architecture Spec → Odyssey Implementation → this script → Phalanx Runtime +→ Numerical Comparison → PASS / FAIL + +Example +------- + python scripts/validate_swiglu.py + python scripts/validate_swiglu.py --hidden-size 64 --intermediate-size 128 +""" + +from __future__ import annotations + +import argparse +import json +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path + +import numpy as np +import torch + +from model import FeedForwardConfig, OdysseySwiGLU +from odyssey.config import REPO_ROOT + +PHALANX_ROOT = REPO_ROOT.parent / "runtime" +DEFAULT_REPORT = REPO_ROOT / "experiments" / "ODY-0006" / "swiglu_validation.json" + + +def _write_f32(path: Path, array: np.ndarray) -> None: + path.write_bytes(np.ascontiguousarray(array, dtype=np.float32).tobytes()) + + +def _read_f32(path: Path, shape: tuple[int, ...]) -> np.ndarray: + data = np.frombuffer(path.read_bytes(), dtype=np.float32) + return data.reshape(shape).copy() + + +def run_odyssey( + x: torch.Tensor, + w_gate: torch.Tensor, + w_up: torch.Tensor, + w_down: torch.Tensor, + config: FeedForwardConfig, +) -> torch.Tensor: + ffn = OdysseySwiGLU(config) + with torch.no_grad(): + ffn.gate_proj.weight.copy_(w_gate) + ffn.up_proj.weight.copy_(w_up) + ffn.down_proj.weight.copy_(w_down) + return ffn(x) + + +def run_phalanx( + work_dir: Path, + *, + phalanx_root: Path, + release: bool, +) -> None: + cmd = ["cargo", "run", "--quiet", "--bin", "validate_swiglu"] + if release: + cmd.append("--release") + cmd.extend(["--", str(work_dir)]) + subprocess.run(cmd, cwd=phalanx_root, check=True) + + +def compare( + odyssey: np.ndarray, + phalanx: np.ndarray, + *, + tolerance: float, +) -> dict[str, float | bool | str]: + diff = np.abs(odyssey.astype(np.float64) - phalanx.astype(np.float64)) + max_err = float(diff.max()) if diff.size else 0.0 + mean_err = float(diff.mean()) if diff.size else 0.0 + rel = diff / (np.abs(odyssey.astype(np.float64)) + 1e-8) + max_rel = float(rel.max()) if rel.size else 0.0 + passed = max_err <= tolerance + return { + "max_error": max_err, + "mean_error": mean_err, + "max_relative_error": max_rel, + "tolerance": tolerance, + "pass": passed, + "status": "PASS" if passed else "FAIL", + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--seq", type=int, default=8) + parser.add_argument("--hidden-size", type=int, default=64) + parser.add_argument("--intermediate-size", type=int, default=128) + parser.add_argument( + "--tolerance", + type=float, + default=1e-3, + help="abs tolerance (default 1e-3: GEMM accum order vs PyTorch; see docs)", + ) + parser.add_argument("--phalanx-root", type=Path, default=PHALANX_ROOT) + parser.add_argument("--release", action="store_true") + parser.add_argument("--keep-work-dir", type=Path, default=None) + parser.add_argument("--report", type=Path, default=DEFAULT_REPORT) + args = parser.parse_args() + + if not args.phalanx_root.is_dir(): + print(f"Phalanx root not found: {args.phalanx_root}", file=sys.stderr) + return 2 + + torch.manual_seed(args.seed) + np.random.seed(args.seed) + + shape = (args.batch, args.seq, args.hidden_size) + x = torch.randn(shape, dtype=torch.float32) + # nn.Linear weight layout: (out, in) + w_gate = torch.randn(args.intermediate_size, args.hidden_size) + w_up = torch.randn(args.intermediate_size, args.hidden_size) + w_down = torch.randn(args.hidden_size, args.intermediate_size) + + config = FeedForwardConfig( + type="swiglu", + hidden_size=args.hidden_size, + intermediate_size=args.intermediate_size, + activation="silu", + device="cpu", + dtype="float32", + ) + + y_ody = run_odyssey(x, w_gate, w_up, w_down, config) + + work = ( + Path(tempfile.mkdtemp(prefix="swiglu_val_")) + if args.keep_work_dir is None + else args.keep_work_dir + ) + work.mkdir(parents=True, exist_ok=True) + + manifest = { + "shape": list(shape), + "hidden_size": args.hidden_size, + "intermediate_size": args.intermediate_size, + "seed": args.seed, + } + (work / "manifest.json").write_text( + json.dumps(manifest, indent=2) + "\n", encoding="utf-8" + ) + _write_f32(work / "x_in.bin", x.numpy()) + _write_f32(work / "w_gate.bin", w_gate.numpy()) + _write_f32(work / "w_up.bin", w_up.numpy()) + _write_f32(work / "w_down.bin", w_down.numpy()) + _write_f32(work / "y_ody.bin", y_ody.detach().numpy()) + + run_phalanx(work, phalanx_root=args.phalanx_root, release=args.release) + + y_ph = _read_f32(work / "y_out.bin", shape) + cmp = compare(y_ody.detach().numpy(), y_ph, tolerance=args.tolerance) + passed = bool(cmp["pass"]) + + report = { + "component": "SwiGLU", + "odyssey_spec": "1.0.0", + "manifest": manifest, + "comparison": cmp, + "max_error": cmp["max_error"], + "mean_error": cmp["mean_error"], + "max_relative_error": cmp["max_relative_error"], + "tolerance": args.tolerance, + "status": "PASS" if passed else "FAIL", + "work_dir": str(work), + "message": ( + "Odyssey and Phalanx are mathematically identical." + if passed + else "SwiGLU outputs diverge beyond tolerance." + ), + } + args.report.parent.mkdir(parents=True, exist_ok=True) + args.report.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + + print("SwiGLU Validation") + print() + print("Max Error") + print(f"{float(cmp['max_error']):.8f}") + print() + print("Mean Error") + print(f"{float(cmp['mean_error']):.8f}") + print() + print(report["status"]) + print() + print(report["message"]) + print() + print(f"report: {args.report}") + + if args.keep_work_dir is None: + shutil.rmtree(work, ignore_errors=True) + + return 0 if passed else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_activations.py b/tests/test_activations.py new file mode 100644 index 0000000..ba4168c --- /dev/null +++ b/tests/test_activations.py @@ -0,0 +1,22 @@ +"""Activation unit tests.""" + +from __future__ import annotations + +import torch + +from model.activations import sigmoid, silu, swish + + +def test_sigmoid_bounds() -> None: + x = torch.linspace(-5, 5, 21) + y = sigmoid(x) + assert torch.all(y > 0) and torch.all(y < 1) + + +def test_silu_zero_at_zero() -> None: + assert silu(torch.tensor(0.0)).item() == 0.0 + + +def test_swish_alias() -> None: + x = torch.randn(8) + assert torch.equal(swish(x), silu(x)) diff --git a/tests/test_feedforward.py b/tests/test_feedforward.py new file mode 100644 index 0000000..b521b00 --- /dev/null +++ b/tests/test_feedforward.py @@ -0,0 +1,23 @@ +"""Feed-forward public API tests.""" + +from __future__ import annotations + +import torch + +from model import FeedForwardConfig, OdysseyFeedForward, build_feed_forward + + +def test_build_feed_forward() -> None: + cfg = FeedForwardConfig(hidden_size=16, intermediate_size=32) + ffn = build_feed_forward(cfg) + assert isinstance(ffn, OdysseyFeedForward) + y = ffn(torch.randn(2, 4, 16)) + assert y.shape == (2, 4, 16) + + +def test_rejects_non_swiglu_type() -> None: + try: + FeedForwardConfig(type="gelu") # type: ignore[arg-type] + raise AssertionError("expected ValueError") + except ValueError as err: + assert "swiglu" in str(err) diff --git a/tests/test_phase0.py b/tests/test_phase0.py index 0ab58f3..28d6765 100644 --- a/tests/test_phase0.py +++ b/tests/test_phase0.py @@ -17,7 +17,7 @@ def test_version() -> None: - assert __version__ == "0.5.0" + assert __version__ == "0.6.0" def test_package_imports() -> None: @@ -43,11 +43,13 @@ def test_config_loads() -> None: assert config["tokenizer"]["path"] == "assets/tokenizer/bpe/odyssey.model" assert config["tokenizer"]["type"] == "bpe" assert config["tokenizer"]["config"] == "configs/tokenizer.yaml" - assert config["experiment"]["id"] == "ODY-0005" + assert config["experiment"]["id"] == "ODY-0006" assert config["model"]["embedding"]["init_strategy"] == "xavier_uniform" assert config["model"]["rope"]["theta"] == 10000.0 assert config["model"]["norm"]["type"] == "rmsnorm" assert config["model"]["norm"]["epsilon"] == 1e-6 + assert config["model"]["feed_forward"]["type"] == "swiglu" + assert config["model"]["feed_forward"]["intermediate_size"] == 2048 def test_config_yaml_parses_directly() -> None: @@ -130,6 +132,12 @@ def test_required_docs_exist() -> None: "scripts/validate_rmsnorm.py", "math/rmsnorm.md", "math/residuals.md", + "docs/architecture/swiglu.md", + "docs/architecture/feedforward.md", + "papers/swiglu.md", + "papers/llama_ffn.md", + "scripts/validate_swiglu.py", + "math/swiglu.md", "tokenizer/README.md", "tokenizer/docs/bpe.md", ] diff --git a/tests/test_shapes_ffn.py b/tests/test_shapes_ffn.py new file mode 100644 index 0000000..c500ea6 --- /dev/null +++ b/tests/test_shapes_ffn.py @@ -0,0 +1,22 @@ +"""Shape tests for SwiGLU / FFN.""" + +from __future__ import annotations + +import torch + +from model import FeedForwardConfig, OdysseySwiGLU + + +def test_rank2_and_rank3() -> None: + ffn = OdysseySwiGLU(FeedForwardConfig(hidden_size=24, intermediate_size=48)) + assert ffn(torch.randn(5, 24)).shape == (5, 24) + assert ffn(torch.randn(2, 5, 24)).shape == (2, 5, 24) + + +def test_rejects_bad_hidden() -> None: + ffn = OdysseySwiGLU(FeedForwardConfig(hidden_size=24, intermediate_size=48)) + try: + ffn(torch.randn(2, 16)) + raise AssertionError("expected ValueError") + except ValueError as err: + assert "hidden_size" in str(err) diff --git a/tests/test_swiglu.py b/tests/test_swiglu.py new file mode 100644 index 0000000..243edc9 --- /dev/null +++ b/tests/test_swiglu.py @@ -0,0 +1,114 @@ +"""Unit tests for OdysseySwiGLU.""" + +from __future__ import annotations + +import torch + +from model import ( + EmbeddingConfig, + FeedForwardConfig, + NormConfig, + OdysseyEmbedding, + OdysseyRMSNorm, + OdysseyRoPE, + OdysseySwiGLU, + RopeConfig, + load_feed_forward_config, +) +from model.activations import silu + + +def test_output_shape() -> None: + cfg = FeedForwardConfig(hidden_size=32, intermediate_size=64) + ffn = OdysseySwiGLU(cfg) + x = torch.randn(2, 8, 32) + assert ffn(x).shape == x.shape + + +def test_silu_formula() -> None: + x = torch.tensor([-2.0, 0.0, 1.5]) + expected = x * torch.sigmoid(x) + assert torch.allclose(silu(x), expected) + + +def test_gate_multiplication() -> None: + cfg = FeedForwardConfig(hidden_size=4, intermediate_size=8) + ffn = OdysseySwiGLU(cfg) + with torch.no_grad(): + ffn.gate_proj.weight.zero_() + ffn.up_proj.weight.fill_(1.0) + ffn.down_proj.weight.zero_() + ffn.down_proj.weight[0, 0] = 1.0 + x = torch.ones(1, 1, 4) + # gate = silu(0)=0 → hidden=0 → output near 0 + y = ffn(x) + assert torch.allclose(y, torch.zeros_like(y), atol=1e-6) + + +def test_config_load() -> None: + cfg = load_feed_forward_config() + assert cfg.type == "swiglu" + assert cfg.hidden_size == 768 + assert cfg.intermediate_size == 2048 + assert cfg.activation == "silu" + + +def test_parameter_count() -> None: + cfg = FeedForwardConfig(hidden_size=16, intermediate_size=32) + ffn = OdysseySwiGLU(cfg) + expected = 16 * 32 + 16 * 32 + 32 * 16 + assert ffn.parameter_count() == expected + info = ffn.inspect() + assert info["projections"]["gate_proj (w1)"] == 16 * 32 + + +def test_float16_forward() -> None: + cfg = FeedForwardConfig(hidden_size=16, intermediate_size=32, dtype="float16") + ffn = OdysseySwiGLU(cfg) + x = torch.randn(1, 4, 16, dtype=torch.float16) + y = ffn(x) + assert y.dtype == torch.float16 + assert torch.isfinite(y.float()).all() + + +def test_gradients() -> None: + cfg = FeedForwardConfig(hidden_size=16, intermediate_size=32) + ffn = OdysseySwiGLU(cfg) + x = torch.randn(2, 3, 16, requires_grad=True) + ffn(x).pow(2).mean().backward() + assert x.grad is not None + assert ffn.gate_proj.weight.grad is not None + + +def test_deterministic() -> None: + cfg = FeedForwardConfig(hidden_size=8, intermediate_size=16) + torch.manual_seed(0) + x = torch.randn(1, 2, 8) + w1 = torch.randn(16, 8) + w3 = torch.randn(16, 8) + w2 = torch.randn(8, 16) + + def run() -> torch.Tensor: + f = OdysseySwiGLU(cfg) + with torch.no_grad(): + f.gate_proj.weight.copy_(w1) + f.up_proj.weight.copy_(w3) + f.down_proj.weight.copy_(w2) + return f(x) + + assert torch.equal(run(), run()) + + +def test_integration_stack() -> None: + emb = OdysseyEmbedding( + EmbeddingConfig(vocab_size=32, hidden_size=32, padding_idx=None) + ) + rope = OdysseyRoPE(RopeConfig(head_dim=8, rotary_dim=8, max_position_embeddings=16)) + norm = OdysseyRMSNorm(NormConfig(hidden_size=32)) + ffn = OdysseySwiGLU(FeedForwardConfig(hidden_size=32, intermediate_size=64)) + ids = torch.randint(0, 32, (1, 4)) + h = emb(ids) + q = h.view(1, 4, 4, 8) + h = rope(q).reshape(1, 4, 32) + h = ffn(norm(h)) + assert h.shape == (1, 4, 32) From 49c3d32264afaf98c684a86a8d4677a08b8741dc Mon Sep 17 00:00:00 2001 From: Khai Date: Wed, 29 Jul 2026 04:30:51 +0100 Subject: [PATCH 2/2] fix(lint): remove unused torch import in parameter_counter Co-authored-by: Cursor --- model/parameter_counter.py | 1 - 1 file changed, 1 deletion(-) diff --git a/model/parameter_counter.py b/model/parameter_counter.py index cf044b9..99138e8 100644 --- a/model/parameter_counter.py +++ b/model/parameter_counter.py @@ -4,7 +4,6 @@ from typing import Any -import torch from torch import nn