diff --git a/CHANGELOG.md b/CHANGELOG.md index c8ef9e9..559cb84 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,19 @@ and this project aims to adhere to [Semantic Versioning](https://semver.org/). --- +## [0.4.0] — 2026-07-29 + +### Added + +- LLaMA-style RoPE (`model.OdysseyRoPE`, cache, math, visualizer) +- `configs/model.yaml` with rope hyperparameters +- Cross-implementation validator `scripts/validate_rope.py` (vs Phalanx) +- Benchmarks + assets under `assets/rope/` +- Papers: RoFormer, attention positional encodings, LLaMA RoPE +- Experiment `ODY-0004` (validation PASS, max abs error ~4.8e-7) + +--- + ## [0.3.0] — 2026-07-28 ### Added diff --git a/EXPERIMENTS.md b/EXPERIMENTS.md index 3ca2f8a..2fcdcc9 100644 --- a/EXPERIMENTS.md +++ b/EXPERIMENTS.md @@ -81,3 +81,21 @@ Details: [experiments/ODY-0002/README.md](experiments/ODY-0002/README.md) | Lessons | Embedding is a gather not a matmul; `math/` notes should ship with each neural phase; weight tying deferred | Details: [experiments/ODY-0003/README.md](experiments/ODY-0003/README.md) + +--- + +## ODY-0004 — Rotary Positional Embeddings + +| Field | Value | +| --- | --- | +| ID | ODY-0004 | +| Date | 2026-07-29 | +| Phase | 4 | +| Purpose | LLaMA-style RoPE + Phalanx numerical parity | +| Config | `configs/model.yaml` / experiment `config.yaml` | +| theta / rotary_dim | 10000 / 128 | +| Result | **Successful** | +| Validation | Max abs error ≈ 4.8e-7 vs Phalanx (`PASS` @ 1e-6) | +| Lessons | Keep inv_freq + adjacent-pair rotate bit-identical to Runtime; validate every layer | + +Details: [experiments/ODY-0004/README.md](experiments/ODY-0004/README.md) diff --git a/PAPERS.md b/PAPERS.md index 091ceee..4e61c70 100644 --- a/PAPERS.md +++ b/PAPERS.md @@ -39,12 +39,21 @@ Math companion: [math/embeddings.md](math/embeddings.md) --- +## Phase 4 — RoPE *(complete)* + +| Paper / Study | Summary | +| --- | --- | +| RoFormer (Su et al.) | [papers/roformer.md](papers/roformer.md) | +| Transformer positional encodings | [papers/attention_position.md](papers/attention_position.md) | +| LLaMA RoPE notes | [papers/llama_rope.md](papers/llama_rope.md) | + +--- + ## Planned Reading (later phases) | Topic | Canonical paper / resource | Phase | | --- | --- | --- | | Attention | *Attention Is All You Need* (Vaswani et al., 2017) | 6+ | -| RoPE | *RoFormer* (Su et al.) | 4 | | RMSNorm | *Root Mean Square Layer Normalization* | 5 | | SwiGLU | *GLU Variants Improve Transformer* (Shazeer) | 7 | | GPT-style LMs | GPT / Llama technical reports | 9–10 | diff --git a/README.md b/README.md index 9eeae9e..99e4ff7 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ --- -**Status:** Research Project — Phase 3 complete · **Spec v1.0.0** frozen +**Status:** Research Project — Phase 4 complete · **Spec v1.0.0** frozen **Language:** Python 3.12+ **Framework:** PyTorch **Target Runtime:** [Phalanx Runtime](https://github.com/404khai/phalanx) @@ -34,6 +34,7 @@ Odyssey is a research repository for building a small, carefully engineered deco | 1 | SentencePiece **reference** tokenizer | | 2 | **Owned** byte-level BPE library (`odyssey_tokenizer`) | | 3 | **Token embedding layer** (`OdysseyEmbedding`) | +| 4 | **RoPE** (`OdysseyRoPE`) + Phalanx numerical validation | ```mermaid flowchart TD @@ -42,12 +43,32 @@ flowchart TD TokenIDs --> EmbeddingLookup[Embedding Lookup] EmbeddingMatrix[Embedding Matrix] --> EmbeddingLookup EmbeddingLookup --> Vectors[Embedding Vectors] - Vectors --> RoPE[RoPE Phase 4] + Vectors --> RoPE[RoPE] RoPE --> TransformerBlock[Transformer Block] ``` --- +## Rotary Positional Embeddings (Phase 4) + +```python +from model import OdysseyRoPE, load_rope_config + +rope = OdysseyRoPE(load_rope_config()) +q_rot, k_rot = rope.apply_rotary(q, k, position_offset=0) +``` + +Cross-check against Phalanx Runtime (must PASS): + +```bash +python scripts/validate_rope.py +# Max Error ~5e-7 → PASS +``` + +Docs: [`docs/architecture/rope.md`](docs/architecture/rope.md) · Spec: [`spec/rope.md`](spec/rope.md) + +--- + ## Embedding Layer (Phase 3) ```python @@ -180,7 +201,8 @@ MYPYPATH=tokenizer mypy --explicit-package-bases -p odyssey_tokenizer | 1 | SentencePiece reference | **Complete** | | 2 | Odyssey BPE library | **Complete** | | 3 | Embedding layer | **Complete** | -| 4–20 | RoPE → Odyssey v1 | Planned | +| 4 | RoPE (+ Phalanx validation) | **Complete** | +| 5–20 | RMSNorm → Odyssey v1 | Planned | --- @@ -192,6 +214,7 @@ MYPYPATH=tokenizer mypy --explicit-package-bases -p odyssey_tokenizer | ODY-0001 | SentencePiece baseline | Successful | | ODY-0002 | Odyssey BPE tokenizer | Successful | | ODY-0003 | Token embedding layer | Successful | +| ODY-0004 | RoPE + Phalanx parity | Successful | --- diff --git a/ROADMAP.md b/ROADMAP.md index 699c68f..b4bb553 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -42,11 +42,12 @@ Development proceeds **one phase at a time**. Complete, document, commit, then w --- -## Phase 4 — Rotary Positional Embeddings +## Phase 4 — Rotary Positional Embeddings *(complete)* -- RoPE implementation -- Visualization -- Tests +- LLaMA-style RoPE (`OdysseyRoPE`) +- Cos/sin cache, partial rotary dims, linear scaling +- Cross-implementation validation vs Phalanx (`scripts/validate_rope.py`) +- ODY-0004 baseline --- diff --git a/assets/rope/cos_sin_curves.png b/assets/rope/cos_sin_curves.png new file mode 100644 index 0000000..dc95f61 Binary files /dev/null and b/assets/rope/cos_sin_curves.png differ diff --git a/assets/rope/inv_freq.png b/assets/rope/inv_freq.png new file mode 100644 index 0000000..9e8086d Binary files /dev/null and b/assets/rope/inv_freq.png differ diff --git a/assets/rope/rotation_demo.png b/assets/rope/rotation_demo.png new file mode 100644 index 0000000..feb5f5e Binary files /dev/null and b/assets/rope/rotation_demo.png differ diff --git a/configs/default.yaml b/configs/default.yaml index e3ff091..68374de 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -2,8 +2,8 @@ # All hyperparameters should eventually flow from config files. experiment: - id: ODY-0003 - name: odyssey-embedding-baseline + id: ODY-0004 + name: odyssey-rope-baseline seed: 42 model: @@ -20,6 +20,14 @@ model: init_std: 0.02 device: cpu dtype: float32 + rope: + theta: 10000.0 + rotary_dim: 64 + max_position_embeddings: 2048 + scaling: none + scaling_factor: 1.0 + device: cpu + dtype: float32 tokenizer: path: assets/tokenizer/bpe/odyssey.model diff --git a/configs/model.yaml b/configs/model.yaml new file mode 100644 index 0000000..a5b5075 --- /dev/null +++ b/configs/model.yaml @@ -0,0 +1,30 @@ +# Odyssey model configuration (Phase 4+) + +model: + name: odyssey-tiny + vocab_size: 32000 + hidden_size: 768 + intermediate_size: 2048 + num_layers: 12 + num_heads: 12 + num_kv_heads: 12 + context_length: 2048 + + embedding: + padding_idx: 0 + init_strategy: xavier_uniform + init_std: 0.02 + device: cpu + dtype: float32 + + # 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_factor: 1.0 + device: cpu + dtype: float32 diff --git a/docs/architecture/README.md b/docs/architecture/README.md index 042074a..0b81be5 100644 --- a/docs/architecture/README.md +++ b/docs/architecture/README.md @@ -7,7 +7,7 @@ Notes for Odyssey's decoder-only transformer. | Doc | Phase | Status | | --- | --- | --- | | [embeddings.md](embeddings.md) | 3 | Complete | -| RoPE | 4 | Planned | +| [rope.md](rope.md) | 4 | Complete | | RMSNorm | 5 | Planned | | Attention | 6 | Planned | diff --git a/docs/architecture/rope.md b/docs/architecture/rope.md new file mode 100644 index 0000000..35a4632 --- /dev/null +++ b/docs/architecture/rope.md @@ -0,0 +1,64 @@ +# Rotary Positional Embeddings (RoPE) + +## Motivation + +Token embeddings alone do not encode order. Absolute position tables add a vector per index but extrapolate poorly. **RoPE** rotates Q/K feature pairs by an angle that grows with position so attention sees relative offsets. + +## Mathematical Derivation + +See Spec [`spec/rope.md`](../../spec/rope.md) and pedagogy [`math/rope.md`](../../math/rope.md). + +\[ +\theta_i = 10000^{-2i/d_r} +\] + +\[ +\begin{pmatrix} x_0' \\ x_1' \end{pmatrix} += +\begin{pmatrix} \cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{pmatrix} +\begin{pmatrix} x_0 \\ x_1 \end{pmatrix} +\] + +Rotation preserves pairwise (and thus subspace) L2 norms. + +## Implementation + +| Module | Role | +| --- | --- | +| `model/rope_math.py` | inv_freq + adjacent rotate | +| `model/rope_cache.py` | lazy cos/sin cache | +| `model/rope.py` | `OdysseyRoPE` public API | +| `configs/model.yaml` | θ, rotary_dim, scaling | + +```python +from model import OdysseyRoPE, load_rope_config +rope = OdysseyRoPE(load_rope_config()) +q_rot, k_rot = rope.apply_rotary(q, k, position_offset=0) +``` + +## Tensor Shapes + +| Layout | Shape | +| --- | --- | +| Training | `(B, S, H, d)` | +| Phalanx | `(S, H, d)` or `(S, d)` | +| Output | Same as input | + +## Complexity + +- Time: \(O(B S H d)\) (same order as a cheap elementwise pass) +- Cache memory: \(O(S_{\max} \cdot d_r)\) for cos+sin + +## Phalanx Compatibility + +Cross-check with: + +```bash +python scripts/validate_rope.py +``` + +Tolerance default `1e-6` (float32). Report: `experiments/ODY-0004/rope_validation.json`. + +## Visualization + +`assets/rope/` — inv-freq spectrum, cos/sin curves, 2D rotation demo. diff --git a/docs/architecture/rope_math.md b/docs/architecture/rope_math.md new file mode 100644 index 0000000..563685d --- /dev/null +++ b/docs/architecture/rope_math.md @@ -0,0 +1,29 @@ +# RoPE Mathematics (Architecture Companion) + +Normative equations live in [`spec/rope.md`](../../spec/rope.md). + +## Variables + +| Symbol | Meaning | +| --- | --- | +| \(d\) | `head_dim` | +| \(d_r\) | `rotary_dim` (even, ≤ \(d\)) | +| \(\theta\) | base frequency (`theta`, default 10000) | +| \(m\) | absolute position | +| \(m'\) | \(m / \mathrm{factor}\) under linear scaling | +| \(\omega_i\) | \(\theta^{-2i/d_r}\) | + +## Why Adjacent Pairing + +LLaMA stores pairs as consecutive dims. The complex view \((x_{2i}+i x_{2i+1})\) maps directly onto that layout without a gather/transpose of even/odd channels across the full head. + +## Why Not Rotate V + +RoPE’s relative-phase argument applies to the \(QK^\top\) product. Rotating \(V\) is not part of the standard derivation and would mix content with position unnecessarily. + +## Cache Layout (shared with Phalanx) + +```text +cos[pos, pair], sin[pos, pair] +flat index = pos * (d_r/2) + pair +``` diff --git a/docs/architecture/rope_visualization.md b/docs/architecture/rope_visualization.md new file mode 100644 index 0000000..d8439a8 --- /dev/null +++ b/docs/architecture/rope_visualization.md @@ -0,0 +1,17 @@ +# RoPE Visualization Guide + +Artifacts are written to `assets/rope/` by `model.rope_visualizer.export_rope_assets`. + +| Figure | Meaning | +| --- | --- | +| `inv_freq.png` | Wavelength spectrum across pair indices | +| `cos_sin_curves.png` | Cos/sin vs position for one pair | +| `rotation_demo.png` | Unit vector `(1,0)` tracing a circle across positions | + +Regenerate: + +```bash +python -c "from model.rope_visualizer import export_rope_assets; export_rope_assets()" +``` + +Intuition: each pair is a 2D vector spun by an angle \(m\omega_i\); magnitude stays constant, only orientation encodes position. diff --git a/experiments/ODY-0004/README.md b/experiments/ODY-0004/README.md new file mode 100644 index 0000000..1e12935 --- /dev/null +++ b/experiments/ODY-0004/README.md @@ -0,0 +1,40 @@ +# ODY-0004 — Rotary Positional Embeddings + +| Field | Value | +| --- | --- | +| Phase | 4 | +| Date | 2026-07-29 | +| Purpose | Implement LLaMA-style RoPE + cross-validate vs Phalanx | +| Result | **Successful** | + +## Configuration + +| Knob | Value | +| --- | --- | +| theta | 10000 | +| rotary_dim | 128 (experiment) / 64 (Tiny default) | +| head_dim | 128 (experiment) / 64 (Tiny) | +| max_position_embeddings | 4096 (experiment) | +| scaling | none (linear supported) | + +## Cross-Implementation Validation + +```bash +python scripts/validate_rope.py --head-dim 128 --rotary-dim 128 --max-position 4096 +``` + +See `rope_validation.json` for max/mean absolute error vs Phalanx `layers::Rope`. + +| Metric | Value | +| --- | --- | +| Max abs error | ≈ 4.77e-7 | +| Mean abs error | ≈ 5.62e-9 | +| Tolerance | 1e-6 | +| Status | **PASS** | + +## Artifacts + +- Metrics: `metrics.json` +- Validation: `rope_validation.json` +- Plots: `assets/rope/` +- Config snapshot: `config.yaml` diff --git a/experiments/ODY-0004/config.yaml b/experiments/ODY-0004/config.yaml new file mode 100644 index 0000000..937fcde --- /dev/null +++ b/experiments/ODY-0004/config.yaml @@ -0,0 +1,11 @@ +rope: + theta: 10000.0 + head_dim: 128 + rotary_dim: 128 + max_position_embeddings: 4096 + scaling: none + scaling_factor: 1.0 + +experiment: + id: ODY-0004 + seed: 0 diff --git a/experiments/ODY-0004/metrics.json b/experiments/ODY-0004/metrics.json new file mode 100644 index 0000000..8b76bc7 --- /dev/null +++ b/experiments/ODY-0004/metrics.json @@ -0,0 +1,17 @@ +{ + "theta": 10000.0, + "rotary_dim": 128, + "head_dim": 128, + "max_position_embeddings": 4096, + "cache_build_seconds": 0.002301, + "cache_memory_bytes": 2097408, + "rotate_mean_seconds": 0.00114166, + "rotate_elems_per_second": 1377700103.75, + "device": "cpu", + "shape": [ + 2, + 512, + 12, + 128 + ] +} diff --git a/experiments/ODY-0004/rope_validation.json b/experiments/ODY-0004/rope_validation.json new file mode 100644 index 0000000..556c80d --- /dev/null +++ b/experiments/ODY-0004/rope_validation.json @@ -0,0 +1,38 @@ +{ + "component": "RoPE", + "odyssey_spec": "1.0.0", + "manifest": { + "shape": [ + 16, + 4, + 128 + ], + "head_dim": 128, + "rotary_dim": 128, + "theta": 10000.0, + "scale": 1.0, + "max_position": 4096, + "position_offset": 3, + "seed": 0 + }, + "q": { + "max_error": 4.76837158203125e-07, + "mean_error": 5.68458347061096e-09, + "tolerance": 1e-06, + "pass": true, + "status": "PASS" + }, + "k": { + "max_error": 4.76837158203125e-07, + "mean_error": 5.555648385779932e-09, + "tolerance": 1e-06, + "pass": true, + "status": "PASS" + }, + "max_error": 4.76837158203125e-07, + "mean_error": 5.620115928195446e-09, + "tolerance": 1e-06, + "status": "PASS", + "work_dir": "/var/folders/c0/32vfjvv9233c0nmghgqt_mp40000gn/T/rope_val__zmqwoye", + "message": "Odyssey and Phalanx are mathematically identical." +} diff --git a/math/README.md b/math/README.md index 93790b5..c0d58ff 100644 --- a/math/README.md +++ b/math/README.md @@ -17,7 +17,7 @@ Each note covers: | [linear_algebra.md](linear_algebra.md) | foundation | Written | | [tensor_shapes.md](tensor_shapes.md) | foundation | Written | | [embeddings.md](embeddings.md) | 3 | Written | -| [rope.md](rope.md) | 4 | Outline (pre-implementation) | +| [rope.md](rope.md) | 4 | Written + cross-validated | | [rmsnorm.md](rmsnorm.md) | 5 | Outline (pre-implementation) | | [attention.md](attention.md) | 6 | Outline (pre-implementation) | | [swiglu.md](swiglu.md) | 7 | Outline (pre-implementation) | diff --git a/math/rope.md b/math/rope.md index 274d2ca..47df837 100644 --- a/math/rope.md +++ b/math/rope.md @@ -1,47 +1,35 @@ -# RoPE — Mathematical Note (Phase 4 outline) +# RoPE — Mathematical Note (Phase 4) -> Implementation lands in Phase 4. Equations recorded here so Phase 3 readers see what comes next. +**Status:** Implemented in Odyssey + Phalanx (cross-validated) -## Equations (preview) - -For head dimension \(d\) (even), position \(m\), and pair index \(i\): - -\[ -\theta_i = 10000^{-2i/d} -\] +## Equations \[ -\begin{pmatrix} q'_{2i} \\ q'_{2i+1} \end{pmatrix} +\omega_i = \theta^{-2i/d_r}, \quad +\begin{pmatrix} x_0' \\ x_1' \end{pmatrix} = -\begin{pmatrix} -\cos m\theta_i & -\sin m\theta_i \\ -\sin m\theta_i & \cos m\theta_i -\end{pmatrix} -\begin{pmatrix} q_{2i} \\ q_{2i+1} \end{pmatrix} +R(m'\omega_i) +\begin{pmatrix} x_0 \\ x_1 \end{pmatrix} \] -Applied to query and key vectors; values typically untouched. +Linear scaling: \(m' = m / \mathrm{factor}\). -## Why it works +## Complexity -Relative position is encoded in the **relative rotation** between positions \(m\) and \(n\), improving length extrapolation vs absolute learned position embeddings. - -## Complexity (preview) - -- Time: \(O(B S D)\) to apply rotations -- Memory: \(O(S \cdot d/2)\) for cached \(\cos/\sin\) tables (or on-the-fly) +| Resource | Cost | +| --- | --- | +| Apply | \(O(BSH d)\) | +| Cache | \(O(S_{\max} d_r)\) | -## Numerical stability +## Numerical Stability -Use float32 for angle tables when possible; watch bf16 precision at long contexts. +- Build angles in float32 +- Default validation tolerance vs Phalanx: `1e-6` -## PyTorch vs Phalanx +## PyTorch -| Path | Status | -| --- | --- | -| Odyssey `model/` RoPE module | Phase 4 | -| Phalanx `layers::rope` | Already prototyping in runtime | +`OdysseyRoPE` / `rope_cache.RopeCacheManager` -## References +## Phalanx Runtime -- Su et al., *RoFormer* (RoPE) +`layers::Rope` — identical adjacent-pair kernel; validated via `scripts/validate_rope.py`. diff --git a/model/__init__.py b/model/__init__.py index 51c187a..eb9c029 100644 --- a/model/__init__.py +++ b/model/__init__.py @@ -1,18 +1,30 @@ """Model package — decoder-only transformer components. -Phase 3 delivers the token embedding layer. Later phases add RoPE, RMSNorm, -attention, SwiGLU, and the full decoder stack. +Phase 3: embeddings. Phase 4: LLaMA-style RoPE (Phalanx-parity). """ -from model.config import EmbeddingConfig, load_embedding_config +from model.config import ( + EmbeddingConfig, + ModelConfig, + RopeConfig, + load_embedding_config, + load_model_config, + load_rope_config, +) from model.embeddings import EmbeddingInspection, OdysseyEmbedding from model.initialization import describe_strategy, initialize_embedding +from model.rope import OdysseyRoPE __all__ = [ "EmbeddingConfig", "EmbeddingInspection", + "ModelConfig", "OdysseyEmbedding", + "OdysseyRoPE", + "RopeConfig", "describe_strategy", "initialize_embedding", "load_embedding_config", + "load_model_config", + "load_rope_config", ] diff --git a/model/config.py b/model/config.py index 014ffff..61d5218 100644 --- a/model/config.py +++ b/model/config.py @@ -1,14 +1,14 @@ -"""Configuration for Odyssey token embeddings. +"""Configuration for Odyssey model components (embeddings + RoPE). Why this exists: - Embedding hyperparameters (vocab size, hidden size, init, device, dtype) - must stay out of scattered script constants so experiments are reproducible - and Phalanx Runtime can mirror the same shapes later. + Embedding and RoPE hyperparameters must stay out of scattered script + constants so experiments are reproducible and Phalanx Runtime can mirror + the same shapes / θ / scaling. """ from __future__ import annotations -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Any, Literal @@ -18,6 +18,7 @@ from odyssey.config import REPO_ROOT DEFAULT_EMBEDDING_CONFIG_PATH = REPO_ROOT / "configs" / "embedding.yaml" +DEFAULT_MODEL_CONFIG_PATH = REPO_ROOT / "configs" / "model.yaml" InitStrategy = Literal[ "normal", @@ -35,6 +36,8 @@ "kaiming_normal", ) +RopeScalingType = Literal["none", "linear"] + DTYPE_MAP: dict[str, torch.dtype] = { "float32": torch.float32, "float16": torch.float16, @@ -87,11 +90,9 @@ def torch_device(self) -> torch.device: @property def parameter_count(self) -> int: - """Number of embedding parameters (vocab × hidden).""" return self.vocab_size * self.hidden_size def memory_bytes(self, dtype: torch.dtype | None = None) -> int: - """Bytes for the embedding matrix at the given (or config) dtype.""" resolved = dtype if dtype is not None else self.torch_dtype return self.parameter_count * resolved.itemsize @@ -114,6 +115,147 @@ def from_dict(cls, data: dict[str, Any]) -> EmbeddingConfig: ) +@dataclass(slots=True) +class RopeConfig: + """RoPE hyperparameters — must match Phalanx ``RopeConfig`` semantics.""" + + theta: float = 10000.0 + head_dim: int = 64 + rotary_dim: int = 64 + max_position_embeddings: int = 2048 + scaling: RopeScalingType = "none" + scaling_factor: float = 1.0 + device: str = "cpu" + dtype: str = "float32" + + def __post_init__(self) -> None: + if not (self.theta > 0): + raise ValueError("theta must be > 0") + if self.head_dim < 1: + raise ValueError("head_dim must be >= 1") + if self.rotary_dim < 2 or self.rotary_dim % 2 != 0: + raise ValueError("rotary_dim must be even and >= 2") + if self.rotary_dim > self.head_dim: + raise ValueError( + f"rotary_dim ({self.rotary_dim}) exceeds head_dim ({self.head_dim})" + ) + if self.max_position_embeddings < 1: + raise ValueError("max_position_embeddings must be >= 1") + if self.scaling not in ("none", "linear"): + raise ValueError( + f"scaling must be 'none' or 'linear' (NTK/YaRN deferred), got {self.scaling!r}" + ) + if self.scaling_factor <= 0: + raise ValueError("scaling_factor must be > 0") + 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 effective_scale(self) -> float: + return self.scaling_factor if self.scaling == "linear" else 1.0 + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> RopeConfig: + scaling = str(data.get("scaling", "none")).lower() + if scaling in ("", "null", "none"): + scaling = "none" + head_dim = int(data.get("head_dim", 64)) + return cls( + theta=float(data.get("theta", data.get("freq_base", 10000.0))), + head_dim=head_dim, + rotary_dim=int( + data.get("rotary_dim", data.get("dimension_count", head_dim)) + ), + max_position_embeddings=int( + data.get( + "max_position_embeddings", + data.get("max_position", data.get("context_length", 2048)), + ) + ), + scaling=scaling, # type: ignore[arg-type] + scaling_factor=float(data.get("scaling_factor", data.get("factor", 1.0))), + 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.""" + + name: str = "odyssey-tiny" + vocab_size: int = 32000 + hidden_size: int = 768 + intermediate_size: int = 2048 + num_layers: int = 12 + num_heads: int = 12 + num_kv_heads: int = 12 + context_length: int = 2048 + embedding: EmbeddingConfig = field(default_factory=EmbeddingConfig) + rope: RopeConfig = field(default_factory=RopeConfig) + + def __post_init__(self) -> None: + if self.hidden_size % self.num_heads != 0: + raise ValueError("hidden_size must be divisible by num_heads") + if self.num_heads % self.num_kv_heads != 0: + raise ValueError("num_heads must be divisible by num_kv_heads") + + @property + def head_dim(self) -> int: + return self.hidden_size // self.num_heads + + def to_dict(self) -> dict[str, Any]: + return { + "name": self.name, + "vocab_size": self.vocab_size, + "hidden_size": self.hidden_size, + "intermediate_size": self.intermediate_size, + "num_layers": self.num_layers, + "num_heads": self.num_heads, + "num_kv_heads": self.num_kv_heads, + "context_length": self.context_length, + "embedding": self.embedding.to_dict(), + "rope": self.rope.to_dict(), + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> ModelConfig: + emb_raw = data.get("embedding", {}) or {} + rope_raw = dict(data.get("rope", {}) or {}) + hidden = int(data.get("hidden_size", 768)) + heads = int(data.get("num_heads", 12)) + head_dim = hidden // heads + rope_raw.setdefault("head_dim", head_dim) + rope_raw.setdefault("rotary_dim", rope_raw.get("rotary_dim", head_dim)) + rope_raw.setdefault( + "max_position_embeddings", + int(data.get("context_length", 2048)), + ) + 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)), + num_layers=int(data.get("num_layers", 12)), + num_heads=heads, + num_kv_heads=int(data.get("num_kv_heads", heads)), + context_length=int(data.get("context_length", 2048)), + embedding=EmbeddingConfig.from_dict(emb_raw), + rope=RopeConfig.from_dict(rope_raw), + ) + + def load_embedding_config(path: Path | str | None = None) -> EmbeddingConfig: """Load embedding configuration from YAML.""" config_path = Path(path) if path is not None else DEFAULT_EMBEDDING_CONFIG_PATH @@ -125,8 +267,29 @@ def load_embedding_config(path: Path | str | None = None) -> EmbeddingConfig: if not isinstance(raw, dict): raise ValueError("Embedding config root must be a mapping") - # Support nested `embedding:` or flat keys. section = raw.get("embedding", raw) if not isinstance(section, dict): raise ValueError("embedding section must be a mapping") return EmbeddingConfig.from_dict(section) + + +def load_model_config(path: Path | str | None = None) -> ModelConfig: + """Load model configuration (including RoPE) from YAML.""" + config_path = Path(path) if path is not None else DEFAULT_MODEL_CONFIG_PATH + if not config_path.is_file(): + raise FileNotFoundError(f"Model config not found: {config_path}") + + with config_path.open("r", encoding="utf-8") as handle: + raw = yaml.safe_load(handle) or {} + if not isinstance(raw, dict): + raise ValueError("Model config root must be a mapping") + + section = raw.get("model", raw) + if not isinstance(section, dict): + raise ValueError("model section must be a mapping") + return ModelConfig.from_dict(section) + + +def load_rope_config(path: Path | str | None = None) -> RopeConfig: + """Load RoPE config from ``configs/model.yaml``.""" + return load_model_config(path).rope diff --git a/model/rope.py b/model/rope.py new file mode 100644 index 0000000..5a9fc45 --- /dev/null +++ b/model/rope.py @@ -0,0 +1,145 @@ +"""LLaMA-style Rotary Positional Embeddings for Odyssey. + +Canonical training-side RoPE. Must stay numerically identical to +``phalanx::layers::Rope`` (see ``scripts/validate_rope.py``). +""" + +from __future__ import annotations + +from typing import Any + +import torch +from torch import nn + +from model.config import RopeConfig +from model.rope_cache import RopeCacheManager +from model.rope_math import rotate_adjacent_pairs + + +class OdysseyRoPE(nn.Module): + """Apply rotary embeddings to query / key tensors. + + Accepted layouts (last dim = ``head_dim``): + + - ``(batch, seq, heads, head_dim)`` — Odyssey training default + - ``(seq, heads, head_dim)`` — Phalanx multi-head layout + - ``(seq, head_dim)`` — single-head / packed + + Shape is preserved. Values (V) are never rotated. + """ + + def __init__(self, config: RopeConfig) -> None: + super().__init__() + self.config = config + self._manager = RopeCacheManager( + rotary_dim=config.rotary_dim, + theta=config.theta, + scale=config.scaling_factor if config.scaling == "linear" else 1.0, + initial_max_position=config.max_position_embeddings, + device=config.device, + dtype=config.torch_dtype, + ) + + @classmethod + def from_config(cls, config: RopeConfig) -> OdysseyRoPE: + return cls(config) + + @property + def cache_max_position(self) -> int: + return self._manager.cache.max_position + + def cache_memory_bytes(self) -> int: + return self._manager.cache.memory_bytes() + + def forward( + self, + x: torch.Tensor, + *, + position_offset: int = 0, + ) -> torch.Tensor: + """Rotate ``x`` starting at absolute position ``position_offset``.""" + self._validate_input(x) + if position_offset < 0: + raise ValueError("position_offset must be >= 0") + + seq = self._seq_len(x) + if seq == 0: + return x + + required = position_offset + seq + cache = self._manager.get(required) + positions = torch.arange( + position_offset, + position_offset + seq, + device=x.device, + dtype=torch.long, + ) + cos = cache.cos[positions] # (seq, n_pairs) + sin = cache.sin[positions] + + # Align cos/sin to x layout: insert head axis when needed. + cos_b, sin_b = self._broadcast_cis(cos, sin, x.ndim) + # Compute in float32 for numerical parity with Phalanx, cast back. + x_f = x.float() + out = rotate_adjacent_pairs(x_f, cos_b.float(), sin_b.float()) + return out.to(dtype=x.dtype) + + def apply_rotary( + self, + q: torch.Tensor, + k: torch.Tensor, + *, + position_offset: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Rotate Q and K; leave V to the caller untouched.""" + return self.forward(q, position_offset=position_offset), self.forward( + k, position_offset=position_offset + ) + + def _validate_input(self, x: torch.Tensor) -> None: + if x.ndim not in (2, 3, 4): + raise ValueError( + "expected rank 2/3/4 with last dim = head_dim, " + f"got shape {tuple(x.shape)}" + ) + if x.shape[-1] != self.config.head_dim: + raise ValueError( + f"last dim {x.shape[-1]} != configured head_dim {self.config.head_dim}" + ) + if self.config.rotary_dim > self.config.head_dim: + raise ValueError("rotary_dim exceeds head_dim") + + @staticmethod + def _seq_len(x: torch.Tensor) -> int: + if x.ndim == 4: + return int(x.shape[1]) + return int(x.shape[0]) + + @staticmethod + def _broadcast_cis( + cos: torch.Tensor, sin: torch.Tensor, rank: int + ) -> tuple[torch.Tensor, torch.Tensor]: + """Expand ``(seq, n_pairs)`` for broadcasting onto ``x`` pair dims.""" + if rank == 2: + # x: (seq, head_dim) — cos already (seq, n_pairs) + return cos, sin + if rank == 3: + # x: (seq, heads, head_dim) → cos (seq, 1, n_pairs) + return cos.unsqueeze(1), sin.unsqueeze(1) + # rank 4: (batch, seq, heads, head_dim) → (1, seq, 1, n_pairs) + return cos.unsqueeze(0).unsqueeze(2), sin.unsqueeze(0).unsqueeze(2) + + def inspect(self) -> dict[str, Any]: + cache = self._manager.cache + return { + "theta": self.config.theta, + "rotary_dim": self.config.rotary_dim, + "head_dim": self.config.head_dim, + "scaling": self.config.scaling, + "scaling_factor": self.config.scaling_factor, + "max_position_embeddings": self.config.max_position_embeddings, + "cache_max_position": cache.max_position, + "cache_memory_bytes": cache.memory_bytes(), + "device": str(cache.device), + "dtype": str(cache.dtype).removeprefix("torch."), + } diff --git a/model/rope_cache.py b/model/rope_cache.py new file mode 100644 index 0000000..e42ca0b --- /dev/null +++ b/model/rope_cache.py @@ -0,0 +1,163 @@ +"""Cosine / sine cache for Rotary Positional Embeddings. + +Layout (matches Phalanx Runtime ``layers::Rope``): + + cos[pos * n_pairs + pair] + sin[pos * n_pairs + pair] + +Logical view exposed to PyTorch: ``(max_position, n_pairs)``. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from model.rope_math import apply_linear_position_scale, inverse_frequencies + + +@dataclass(slots=True) +class RopeCache: + """Precomputed cos/sin tables for absolute positions ``0 .. max_position-1``.""" + + cos: torch.Tensor # (max_position, n_pairs) + sin: torch.Tensor # (max_position, n_pairs) + inv_freq: torch.Tensor # (n_pairs,) + rotary_dim: int + theta: float + scale: float + max_position: int + + @property + def n_pairs(self) -> int: + return self.rotary_dim // 2 + + @property + def device(self) -> torch.device: + return self.cos.device + + @property + def dtype(self) -> torch.dtype: + return self.cos.dtype + + def memory_bytes(self) -> int: + return int(self.cos.nbytes + self.sin.nbytes + self.inv_freq.nbytes) + + def to( + self, device: torch.device | str | None = None, dtype: torch.dtype | None = None + ) -> RopeCache: + return RopeCache( + cos=self.cos.to(device=device, dtype=dtype), + sin=self.sin.to(device=device, dtype=dtype), + inv_freq=self.inv_freq.to(device=device, dtype=dtype), + rotary_dim=self.rotary_dim, + theta=self.theta, + scale=self.scale, + max_position=self.max_position, + ) + + +def build_rope_cache( + *, + rotary_dim: int, + max_position: int, + theta: float = 10000.0, + scale: float = 1.0, + device: torch.device | str | None = None, + dtype: torch.dtype = torch.float32, +) -> RopeCache: + """Build a dense cos/sin cache for ``max_position`` absolute indices.""" + if max_position < 1: + raise ValueError("max_position must be >= 1") + + inv_freq = inverse_frequencies( + rotary_dim, theta, dtype=torch.float32, device=device + ) + positions = torch.arange(max_position, device=device, dtype=torch.float32) + positions = apply_linear_position_scale(positions, scale) + # angles: (max_position, n_pairs) + angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0) + cos = torch.cos(angles).to(dtype=dtype) + sin = torch.sin(angles).to(dtype=dtype) + return RopeCache( + cos=cos, + sin=sin, + inv_freq=inv_freq.to(dtype=dtype), + rotary_dim=rotary_dim, + theta=theta, + scale=scale, + max_position=max_position, + ) + + +class RopeCacheManager: + """Lazy cos/sin cache with growth and device/dtype moves. + + The cache is never rebuilt when a request fits the existing table on the + same device/dtype/hyperparameters. + """ + + def __init__( + self, + *, + rotary_dim: int, + theta: float = 10000.0, + scale: float = 1.0, + initial_max_position: int = 2048, + device: torch.device | str = "cpu", + dtype: torch.dtype = torch.float32, + ) -> None: + self.rotary_dim = rotary_dim + self.theta = theta + self.scale = scale + self.device = torch.device(device) + self.dtype = dtype + self._cache: RopeCache | None = None + self._ensure(initial_max_position) + + @property + def cache(self) -> RopeCache: + if self._cache is None: + raise RuntimeError("RoPE cache not initialized") + return self._cache + + def _ensure(self, required_positions: int) -> RopeCache: + need = max(1, required_positions) + if ( + self._cache is not None + and self._cache.max_position >= need + and self._cache.device == self.device + and self._cache.dtype == self.dtype + and self._cache.rotary_dim == self.rotary_dim + and self._cache.theta == self.theta + and self._cache.scale == self.scale + ): + return self._cache + + # Grow geometrically to avoid frequent rebuilds. + current = 0 if self._cache is None else self._cache.max_position + target = max(need, current) + if self._cache is not None and need > current: + target = max(need, current * 2) + + self._cache = build_rope_cache( + rotary_dim=self.rotary_dim, + max_position=target, + theta=self.theta, + scale=self.scale, + device=self.device, + dtype=self.dtype, + ) + return self._cache + + def get(self, required_positions: int) -> RopeCache: + """Return a cache covering absolute positions ``0 .. required_positions-1``.""" + return self._ensure(required_positions) + + def to(self, device: torch.device | str, dtype: torch.dtype | None = None) -> None: + self.device = torch.device(device) + if dtype is not None: + self.dtype = dtype + if self._cache is not None: + self._cache = self._cache.to(device=self.device, dtype=self.dtype) diff --git a/model/rope_math.py b/model/rope_math.py new file mode 100644 index 0000000..b03024c --- /dev/null +++ b/model/rope_math.py @@ -0,0 +1,102 @@ +"""Rotary embedding mathematics (LLaMA / Phalanx-compatible). + +Frequency schedule +------------------ +For rotary width ``d_r`` (even) and base ``θ`` (default 10000): + + inv_freq[i] = θ ^ (-2i / d_r) for i = 0 .. d_r/2 - 1 + +Why θ = 10000: inherited from the original Transformer sinusoids and +RoFormer / LLaMA; it spaces wavelengths from ~2π to cover short and long +offsets without trainable parameters. + +Adjacent-pair rotation +---------------------- +For each pair ``(x0, x1)`` at absolute position ``m`` (after optional linear +scale ``m' = m / factor``): + + x0' = x0 * cos(m'·ω) - x1 * sin(m'·ω) + x1' = x0 * sin(m'·ω) + x1 * cos(m'·ω) + +This is the real form of multiplying ``(x0 + i x1)`` by ``e^{i m' ω}``. +Rotation preserves L2 norm of each pair (and thus the full rotary subspace). + +V is never rotated: values carry content, not relative-position phases in the +RoPE attention derivation. +""" + +from __future__ import annotations + +import math + +import torch + + +def inverse_frequencies( + rotary_dim: int, + theta: float = 10000.0, + *, + dtype: torch.dtype = torch.float32, + device: torch.device | str | None = None, +) -> torch.Tensor: + """Return ``inv_freq`` of shape ``(rotary_dim // 2,)``.""" + if rotary_dim <= 0 or rotary_dim % 2 != 0: + raise ValueError(f"rotary_dim must be even and > 0, got {rotary_dim}") + if not math.isfinite(theta) or theta <= 0: + raise ValueError(f"theta must be finite and > 0, got {theta}") + + n_pairs = rotary_dim // 2 + # Match Phalanx: exponent = (2 * i) / rotary_dim in f32-style arithmetic. + indices = torch.arange(n_pairs, device=device, dtype=torch.float32) + inv_freq = theta ** (-(2.0 * indices) / float(rotary_dim)) + return inv_freq.to(dtype=dtype) + + +def rotate_adjacent_pairs( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> torch.Tensor: + """Rotate the leading even-width slice of ``x`` with broadcastable cos/sin. + + Args: + x: ``(..., head_dim)`` — only the first ``2 * cos.shape[-1]`` dims rotate. + cos: ``(..., n_pairs)`` broadcastable to the pair axis. + sin: same shape as ``cos``. + + Returns: + Tensor with the same shape as ``x``. + """ + if x.shape[-1] < 2: + return x + + n_pairs = cos.shape[-1] + rotary_dim = n_pairs * 2 + if rotary_dim > x.shape[-1]: + raise ValueError(f"rotary width {rotary_dim} exceeds head_dim {x.shape[-1]}") + + x_rot = x[..., :rotary_dim] + x_pass = x[..., rotary_dim:] + + # Adjacent pairs: (x0, x1, x2, x3, ...) → rotate (x0,x1), (x2,x3), ... + x0 = x_rot[..., 0::2] + x1 = x_rot[..., 1::2] + # Broadcast cos/sin over batch/head dims as needed. + while cos.ndim < x0.ndim: + cos = cos.unsqueeze(0) + sin = sin.unsqueeze(0) + + out0 = x0 * cos - x1 * sin + out1 = x0 * sin + x1 * cos + rotated = torch.stack((out0, out1), dim=-1).flatten(-2) + + if x_pass.numel() == 0: + return rotated + return torch.cat((rotated, x_pass), dim=-1) + + +def apply_linear_position_scale(positions: torch.Tensor, factor: float) -> torch.Tensor: + """Return ``m' = m / factor`` for linear RoPE scaling (factor=1 → identity).""" + if not math.isfinite(factor) or factor <= 0: + raise ValueError(f"scaling factor must be finite and > 0, got {factor}") + return positions.to(dtype=torch.float32) / float(factor) diff --git a/model/rope_visualizer.py b/model/rope_visualizer.py new file mode 100644 index 0000000..26ff531 --- /dev/null +++ b/model/rope_visualizer.py @@ -0,0 +1,118 @@ +"""Educational visualizations for Odyssey RoPE.""" + +from __future__ import annotations + +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +import torch + +from model.config import RopeConfig +from model.rope import OdysseyRoPE +from model.rope_math import inverse_frequencies +from odyssey.config import REPO_ROOT + +DEFAULT_ASSET_DIR = REPO_ROOT / "assets" / "rope" + + +def plot_inverse_frequencies( + config: RopeConfig, + *, + output_path: Path | str | None = None, +) -> Path: + path = Path(output_path) if output_path else DEFAULT_ASSET_DIR / "inv_freq.png" + path.parent.mkdir(parents=True, exist_ok=True) + inv = inverse_frequencies(config.rotary_dim, config.theta).cpu().numpy() + fig, ax = plt.subplots(figsize=(7, 4)) + ax.plot(np.arange(len(inv)), inv, color="#2c5f6e", marker="o", markersize=3) + ax.set_xlabel("pair index i") + ax.set_ylabel("inv_freq[i]") + ax.set_title( + f"RoPE inverse frequencies (θ={config.theta}, d_r={config.rotary_dim})" + ) + ax.set_yscale("log") + fig.tight_layout() + fig.savefig(path, dpi=140) + plt.close(fig) + return path + + +def plot_cos_sin_curves( + config: RopeConfig, + *, + pair_index: int = 0, + output_path: Path | str | None = None, +) -> Path: + path = ( + Path(output_path) if output_path else DEFAULT_ASSET_DIR / "cos_sin_curves.png" + ) + path.parent.mkdir(parents=True, exist_ok=True) + rope = OdysseyRoPE(config) + cache = rope._manager.cache # noqa: SLF001 — visualizer introspection + pos = np.arange(min(256, cache.max_position)) + cos = cache.cos[: len(pos), pair_index].cpu().numpy() + sin = cache.sin[: len(pos), pair_index].cpu().numpy() + fig, ax = plt.subplots(figsize=(8, 4)) + ax.plot(pos, cos, label=f"cos pair={pair_index}", color="#2c5f6e") + ax.plot(pos, sin, label=f"sin pair={pair_index}", color="#b85c38") + ax.set_xlabel("position m") + ax.set_ylabel("value") + ax.set_title("RoPE cos/sin vs position") + ax.legend() + fig.tight_layout() + fig.savefig(path, dpi=140) + plt.close(fig) + return path + + +def plot_rotation_demo( + config: RopeConfig, + *, + output_path: Path | str | None = None, +) -> Path: + """2D scatter of a unit pair rotating across positions.""" + path = Path(output_path) if output_path else DEFAULT_ASSET_DIR / "rotation_demo.png" + path.parent.mkdir(parents=True, exist_ok=True) + # Minimal head: rotary_dim dims, demo on first pair. + demo = RopeConfig( + theta=config.theta, + head_dim=max(2, config.rotary_dim), + rotary_dim=max(2, min(config.rotary_dim, 8)), + max_position_embeddings=64, + scaling=config.scaling, + scaling_factor=config.scaling_factor, + ) + rope = OdysseyRoPE(demo) + x = torch.zeros(32, demo.head_dim) + x[:, 0] = 1.0 + x[:, 1] = 0.0 + y = rope(x, position_offset=0) + pts = y[:, :2].detach().cpu().numpy() + fig, ax = plt.subplots(figsize=(5, 5)) + ax.scatter(pts[:, 0], pts[:, 1], c=np.arange(len(pts)), cmap="viridis", s=28) + ax.axhline(0, color="#888", lw=0.5) + ax.axvline(0, color="#888", lw=0.5) + ax.set_aspect("equal") + ax.set_xlabel("x0'") + ax.set_ylabel("x1'") + ax.set_title("Unit pair (1,0) under RoPE positions 0..31") + fig.tight_layout() + fig.savefig(path, dpi=140) + plt.close(fig) + return path + + +def export_rope_assets(config: RopeConfig | None = None) -> dict[str, Path]: + cfg = config or RopeConfig( + theta=10000.0, + head_dim=128, + rotary_dim=128, + max_position_embeddings=4096, + ) + DEFAULT_ASSET_DIR.mkdir(parents=True, exist_ok=True) + return { + "inv_freq": plot_inverse_frequencies(cfg), + "cos_sin": plot_cos_sin_curves(cfg), + "rotation": plot_rotation_demo(cfg), + } diff --git a/odyssey/__init__.py b/odyssey/__init__.py index c534ecd..5ecfcd2 100644 --- a/odyssey/__init__.py +++ b/odyssey/__init__.py @@ -1,7 +1,6 @@ """Odyssey research package. -Phase 3 delivers the token embedding layer. Later phases add RoPE, -normalization, attention, and the full decoder stack. +Phase 4 delivers LLaMA-style RoPE with Phalanx numerical parity. """ -__version__ = "0.3.0" +__version__ = "0.4.0" diff --git a/papers/attention_position.md b/papers/attention_position.md new file mode 100644 index 0000000..5f929ac --- /dev/null +++ b/papers/attention_position.md @@ -0,0 +1,38 @@ +# Attention Is All You Need — Positional Encoding + +**Authors:** Vaswani et al. +**Link:** [https://arxiv.org/abs/1706.03762](https://arxiv.org/abs/1706.03762) +**Focus:** Positional Encoding section + +--- + +## Motivation + +Self-attention is permutation-invariant. Without position signals, token order is lost. + +--- + +## Absolute Sinusoidal Encoding + +\[ +\begin{aligned} +PE_{(pos, 2i)} &= \sin(pos / 10000^{2i/d}) \\ +PE_{(pos, 2i+1)} &= \cos(pos / 10000^{2i/d}) +\end{aligned} +\] + +Added to input embeddings. The \(10000\) base later reappears in RoPE’s θ schedule. + +--- + +## Limitations + +- Absolute, not relative +- Additive mixing with content embeddings +- Weaker long-context extrapolation than RoPE in practice + +--- + +## Notes for Odyssey + +Odyssey does **not** use additive sinusoids. Spec v1 mandates RoPE on Q/K (see `spec/rope.md`). This paper remains the historical root of the frequency schedule. diff --git a/papers/llama_rope.md b/papers/llama_rope.md new file mode 100644 index 0000000..1faa500 --- /dev/null +++ b/papers/llama_rope.md @@ -0,0 +1,38 @@ +# LLaMA RoPE Notes + +**References:** LLaMA / Llama 2 technical reports; llama.cpp / GGUF RoPE metadata + +--- + +## Formulation Used by Odyssey & Phalanx + +- Adjacent-pair rotation (interleaved even/odd dims) +- \(\theta_i = \texttt{rope.freq\_base}^{-2i/d_r}\) (default base 10000) +- Partial rotary dims: rotate first `rope.dimension_count` features +- Optional **linear** position scaling: \(m' = m / \mathrm{factor}\) +- Apply to **Q and K only** + +--- + +## Differences from RoFormer Presentation + +| Topic | RoFormer paper | LLaMA / Odyssey | +| --- | --- | --- | +| Complex vs real | Emphasizes complex view | Real 2×2 rotate (equivalent) | +| Pairing | Complex channel pairs | Adjacent dims `(0,1), (2,3), …` | +| Partial RoPE | Not always stressed | First `rotary_dim` only | +| Scaling | Later literature | Linear in Spec v1; YaRN/NTK later | + +--- + +## Phalanx Parity Contract + +Must match: + +- θ / inverse-frequency generation +- Adjacent-pair rotate +- Cos/sin cache indexing `pos * n_pairs + pair` +- Linear scaling +- Tensor layouts `[seq, heads, head_dim]` (Odyssey also accepts batched rank-4) + +Validated by `scripts/validate_rope.py` + `runtime` bin `validate_rope`. diff --git a/papers/roformer.md b/papers/roformer.md new file mode 100644 index 0000000..60c47e5 --- /dev/null +++ b/papers/roformer.md @@ -0,0 +1,50 @@ +# RoFormer: Enhanced Transformer with Rotary Position Embedding + +**Authors:** Jianlin Su et al. +**Link:** [https://arxiv.org/abs/2104.09864](https://arxiv.org/abs/2104.09864) + +--- + +## Motivation + +Absolute position embeddings (learned or sinusoidal) inject position as an additive vector. They do not naturally encode **relative** offsets inside attention scores and extrapolate poorly beyond the training context. + +--- + +## Mathematics + +Represent a 2D feature pair as a complex number and multiply by a position-dependent phase: + +\[ +(x_{2i} + i x_{2i+1})\, e^{i m \theta_i} +\] + +with \(\theta_i = 10000^{-2i/d}\). In real arithmetic this is the 2×2 rotation used by Odyssey / LLaMA / Phalanx. + +--- + +## Rotation Derivation + +The relative rotation between positions \(m\) and \(n\) depends on \(m-n\), so attention logits gain relative positional structure without extra bias parameters. + +--- + +## Advantages + +- No trainable position table +- Built-in relative geometry +- Efficient via cos/sin caches +- Strong length extrapolation vs absolute embeddings + +--- + +## Weaknesses + +- Pairing scheme and θ schedule are design choices +- Very long contexts still need scaling tricks (NTK, YaRN — deferred) + +--- + +## Notes for Odyssey + +Odyssey Phase 4 implements the LLaMA adjacent-pair form of RoPE and validates numerically against Phalanx Runtime (`scripts/validate_rope.py`). diff --git a/pyproject.toml b/pyproject.toml index 8e039d0..f222deb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "odyssey" -version = "0.3.0" +version = "0.4.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_rope.py b/scripts/benchmark_rope.py new file mode 100644 index 0000000..f5d5e1b --- /dev/null +++ b/scripts/benchmark_rope.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 +"""Benchmark Odyssey RoPE cache build + rotate throughput.""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import torch + +from model import OdysseyRoPE, RopeConfig +from model.rope_visualizer import export_rope_assets +from odyssey.config import REPO_ROOT + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--rotary-dim", type=int, default=128) + parser.add_argument("--max-position", type=int, default=4096) + parser.add_argument("--seq", type=int, default=512) + parser.add_argument("--heads", type=int, default=12) + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--repeats", type=int, default=50) + parser.add_argument("--device", default="cpu") + parser.add_argument( + "--output", + type=Path, + default=REPO_ROOT / "experiments" / "ODY-0004" / "metrics.json", + ) + parser.add_argument("--visualize", action="store_true") + args = parser.parse_args() + + cfg = RopeConfig( + theta=10000.0, + head_dim=args.head_dim, + rotary_dim=args.rotary_dim, + max_position_embeddings=args.max_position, + device=args.device, + dtype="float32", + ) + + t0 = time.perf_counter() + rope = OdysseyRoPE(cfg) + # Force full cache materialization + _ = rope(torch.zeros(1, 1, args.head_dim, device=args.device), position_offset=0) + cache_s = time.perf_counter() - t0 + + x = torch.randn(args.batch, args.seq, args.heads, args.head_dim, device=args.device) + for _ in range(args.warmup): + _ = rope(x, position_offset=0) + if args.device.startswith("cuda"): + torch.cuda.synchronize() + + times: list[float] = [] + for _ in range(args.repeats): + if args.device.startswith("cuda"): + torch.cuda.synchronize() + start = time.perf_counter() + _ = rope(x, position_offset=0) + if args.device.startswith("cuda"): + torch.cuda.synchronize() + times.append(time.perf_counter() - start) + + mean_s = sum(times) / len(times) + elems = args.batch * args.seq * args.heads * args.head_dim + metrics = { + "theta": 10000.0, + "rotary_dim": args.rotary_dim, + "head_dim": args.head_dim, + "max_position_embeddings": args.max_position, + "cache_build_seconds": round(cache_s, 6), + "cache_memory_bytes": rope.cache_memory_bytes(), + "rotate_mean_seconds": round(mean_s, 8), + "rotate_elems_per_second": round(elems / mean_s, 2), + "device": args.device, + "shape": list(x.shape), + } + 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 args.visualize: + paths = export_rope_assets(cfg) + print("visualizations:", {k: str(v) for k, v in paths.items()}) + + +if __name__ == "__main__": + main() diff --git a/scripts/validate_rope.py b/scripts/validate_rope.py new file mode 100644 index 0000000..e99fdc5 --- /dev/null +++ b/scripts/validate_rope.py @@ -0,0 +1,215 @@ +#!/usr/bin/env python3 +"""Cross-implementation RoPE validation: Odyssey (PyTorch) vs Phalanx (Rust). + +Workflow +-------- +Architecture Spec → Odyssey Implementation → this script → Phalanx Runtime +→ Numerical Comparison → PASS / FAIL + +Procedure +1. Generate identical random Q, K (float32). +2. Apply Odyssey ``OdysseyRoPE``. +3. Serialize inputs + manifest for the Phalanx ``validate_rope`` binary. +4. Run Phalanx RoPE on the same tensors. +5. Compare max / mean absolute error against a tolerance (default 1e-6). + +Example +------- + python scripts/validate_rope.py + python scripts/validate_rope.py --tolerance 1e-6 --seed 0 +""" + +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 OdysseyRoPE, RopeConfig +from odyssey.config import REPO_ROOT + +PHALANX_ROOT = REPO_ROOT.parent / "runtime" +DEFAULT_REPORT = REPO_ROOT / "experiments" / "ODY-0004" / "rope_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( + q: torch.Tensor, + k: torch.Tensor, + config: RopeConfig, + position_offset: int, +) -> tuple[torch.Tensor, torch.Tensor]: + rope = OdysseyRoPE(config) + return rope.apply_rotary(q, k, position_offset=position_offset) + + +def run_phalanx( + work_dir: Path, + *, + phalanx_root: Path, + release: bool, +) -> None: + """Invoke Phalanx via ``cargo run`` so ``CARGO_TARGET_DIR`` overrides work.""" + cmd = ["cargo", "run", "--quiet", "--bin", "validate_rope"] + 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 + passed = max_err <= tolerance + return { + "max_error": max_err, + "mean_error": mean_err, + "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("--seq", type=int, default=16) + parser.add_argument("--heads", type=int, default=4) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--rotary-dim", type=int, default=128) + parser.add_argument("--theta", type=float, default=10000.0) + parser.add_argument("--max-position", type=int, default=4096) + parser.add_argument("--position-offset", type=int, default=3) + parser.add_argument("--scale", type=float, default=1.0, help="linear scale factor") + parser.add_argument("--tolerance", type=float, default=1e-6) + 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) + + # Phalanx layout: [seq, heads, head_dim] + shape = (args.seq, args.heads, args.head_dim) + q = torch.randn(shape, dtype=torch.float32) + k = torch.randn(shape, dtype=torch.float32) + + scaling = "linear" if args.scale != 1.0 else "none" + config = RopeConfig( + theta=args.theta, + head_dim=args.head_dim, + rotary_dim=args.rotary_dim, + max_position_embeddings=args.max_position, + scaling=scaling, # type: ignore[arg-type] + scaling_factor=args.scale, + device="cpu", + dtype="float32", + ) + + q_ody, k_ody = run_odyssey(q, k, config, args.position_offset) + + work = ( + Path(tempfile.mkdtemp(prefix="rope_val_")) + if args.keep_work_dir is None + else args.keep_work_dir + ) + work.mkdir(parents=True, exist_ok=True) + + manifest = { + "shape": list(shape), + "head_dim": args.head_dim, + "rotary_dim": args.rotary_dim, + "theta": args.theta, + "scale": args.scale, + "max_position": args.max_position, + "position_offset": args.position_offset, + "seed": args.seed, + } + (work / "manifest.json").write_text( + json.dumps(manifest, indent=2) + "\n", encoding="utf-8" + ) + _write_f32(work / "q_in.bin", q.numpy()) + _write_f32(work / "k_in.bin", k.numpy()) + _write_f32(work / "q_ody.bin", q_ody.detach().numpy()) + _write_f32(work / "k_ody.bin", k_ody.detach().numpy()) + + run_phalanx(work, phalanx_root=args.phalanx_root, release=args.release) + + q_ph = _read_f32(work / "q_out.bin", shape) + k_ph = _read_f32(work / "k_out.bin", shape) + + q_cmp = compare(q_ody.detach().numpy(), q_ph, tolerance=args.tolerance) + k_cmp = compare(k_ody.detach().numpy(), k_ph, tolerance=args.tolerance) + max_err = max(float(q_cmp["max_error"]), float(k_cmp["max_error"])) + mean_err = (float(q_cmp["mean_error"]) + float(k_cmp["mean_error"])) / 2.0 + passed = max_err <= args.tolerance + + report = { + "component": "RoPE", + "odyssey_spec": "1.0.0", + "manifest": manifest, + "q": q_cmp, + "k": k_cmp, + "max_error": max_err, + "mean_error": mean_err, + "tolerance": args.tolerance, + "status": "PASS" if passed else "FAIL", + "work_dir": str(work), + "message": ( + "Odyssey and Phalanx are mathematically identical." + if passed + else "RoPE 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("RoPE Validation") + print() + print("Max Error") + print(f"{max_err:.8f}") + print() + print("Mean Error") + print(f"{mean_err:.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_cache.py b/tests/test_cache.py new file mode 100644 index 0000000..769dfbe --- /dev/null +++ b/tests/test_cache.py @@ -0,0 +1,40 @@ +"""RoPE cache tests.""" + +from __future__ import annotations + +import torch + +from model.rope_cache import RopeCacheManager, build_rope_cache + + +def test_build_cache_shapes() -> None: + cache = build_rope_cache(rotary_dim=8, max_position=16, theta=10000.0) + assert cache.cos.shape == (16, 4) + assert cache.sin.shape == (16, 4) + assert cache.inv_freq.shape == (4,) + + +def test_lazy_extension() -> None: + mgr = RopeCacheManager( + rotary_dim=8, theta=10000.0, initial_max_position=16, device="cpu" + ) + assert mgr.cache.max_position == 16 + c1 = mgr.get(16) + c2 = mgr.get(16) + assert c1 is c2 or torch.equal(c1.cos, c2.cos) + grown = mgr.get(40) + assert grown.max_position >= 40 + + +def test_dtype_device_consistency() -> None: + cache = build_rope_cache( + rotary_dim=4, max_position=8, dtype=torch.float32, device="cpu" + ) + assert cache.cos.dtype == torch.float32 + assert cache.cos.device.type == "cpu" + + +def test_position_zero_cos_one() -> None: + cache = build_rope_cache(rotary_dim=4, max_position=4) + assert torch.allclose(cache.cos[0], torch.ones(2), atol=1e-6) + assert torch.allclose(cache.sin[0], torch.zeros(2), atol=1e-6) diff --git a/tests/test_phase0.py b/tests/test_phase0.py index 2a008cb..dfa70ef 100644 --- a/tests/test_phase0.py +++ b/tests/test_phase0.py @@ -17,7 +17,7 @@ def test_version() -> None: - assert __version__ == "0.3.0" + assert __version__ == "0.4.0" def test_package_imports() -> None: @@ -43,8 +43,9 @@ 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-0003" + assert config["experiment"]["id"] == "ODY-0004" assert config["model"]["embedding"]["init_strategy"] == "xavier_uniform" + assert config["model"]["rope"]["theta"] == 10000.0 def test_config_yaml_parses_directly() -> None: @@ -115,6 +116,10 @@ def test_required_docs_exist() -> None: "spec/gguf_mapping.md", "spec/runtime_contract.md", "configs/embedding.yaml", + "configs/model.yaml", + "docs/architecture/rope.md", + "papers/roformer.md", + "scripts/validate_rope.py", "tokenizer/README.md", "tokenizer/docs/bpe.md", ] diff --git a/tests/test_rope.py b/tests/test_rope.py new file mode 100644 index 0000000..7f5c387 --- /dev/null +++ b/tests/test_rope.py @@ -0,0 +1,68 @@ +"""Unit tests for Odyssey RoPE.""" + +from __future__ import annotations + +import pytest +import torch + +from model import EmbeddingConfig, OdysseyEmbedding, OdysseyRoPE, RopeConfig + + +@pytest.fixture +def rope() -> OdysseyRoPE: + return OdysseyRoPE( + RopeConfig( + theta=10000.0, + head_dim=8, + rotary_dim=8, + max_position_embeddings=64, + ) + ) + + +def test_shape_preserved_rank4(rope: OdysseyRoPE) -> None: + q = torch.randn(2, 5, 3, 8) + out = rope(q) + assert out.shape == q.shape + + +def test_shape_preserved_phalanx_layout(rope: OdysseyRoPE) -> None: + x = torch.randn(5, 3, 8) + assert rope(x).shape == x.shape + + +def test_apply_rotary_returns_pair(rope: OdysseyRoPE) -> None: + q = torch.randn(4, 2, 8) + k = torch.randn(4, 2, 8) + q2, k2 = rope.apply_rotary(q, k, position_offset=1) + assert q2.shape == q.shape and k2.shape == k.shape + + +def test_position_zero_near_identity(rope: OdysseyRoPE) -> None: + x = torch.randn(1, 2, 8) + y = rope(x, position_offset=0) + assert torch.allclose(x, y, atol=1e-5) + + +def test_deterministic(rope: OdysseyRoPE) -> None: + x = torch.randn(3, 2, 8) + a = rope(x, position_offset=2) + b = rope(x, position_offset=2) + assert torch.equal(a, b) + + +def test_embedding_then_rope_shapes() -> None: + emb = OdysseyEmbedding( + EmbeddingConfig(vocab_size=50, hidden_size=32, padding_idx=0) + ) + # Reshape embedding to fake multi-head: D=32 → H=4, d=8 + ids = torch.randint(0, 50, (1, 6)) + h = emb(ids).view(1, 6, 4, 8) + rope = OdysseyRoPE(RopeConfig(head_dim=8, rotary_dim=8, max_position_embeddings=32)) + out = rope(h) + assert out.shape == (1, 6, 4, 8) + + +def test_rejects_bad_head_dim(rope: OdysseyRoPE) -> None: + with pytest.raises(ValueError, match="head_dim"): + rope(torch.randn(2, 2, 4)) diff --git a/tests/test_rope_shapes.py b/tests/test_rope_shapes.py new file mode 100644 index 0000000..456d123 --- /dev/null +++ b/tests/test_rope_shapes.py @@ -0,0 +1,32 @@ +"""Additional RoPE shape tests (batch layouts).""" + +from __future__ import annotations + +import pytest +import torch + +from model import OdysseyRoPE, RopeConfig, load_model_config + + +def test_load_model_yaml_rope() -> None: + cfg = load_model_config() + assert cfg.rope.theta == 10000.0 + assert cfg.rope.rotary_dim == 64 + assert cfg.head_dim == 64 + + +@pytest.mark.parametrize( + "shape", + [ + (8, 64), + (8, 4, 64), + (2, 8, 4, 64), + ], +) +def test_layouts(shape: tuple[int, ...]) -> None: + rope = OdysseyRoPE( + RopeConfig(head_dim=64, rotary_dim=64, max_position_embeddings=128) + ) + x = torch.randn(*shape) + y = rope(x, position_offset=0) + assert y.shape == shape diff --git a/tests/test_rotation.py b/tests/test_rotation.py new file mode 100644 index 0000000..d1ee5a1 --- /dev/null +++ b/tests/test_rotation.py @@ -0,0 +1,52 @@ +"""Property tests for RoPE rotations.""" + +from __future__ import annotations + +import torch + +from model import OdysseyRoPE, RopeConfig +from model.rope_math import inverse_frequencies, rotate_adjacent_pairs + + +def test_inverse_freq_length_and_positive() -> None: + inv = inverse_frequencies(16, theta=10000.0) + assert inv.shape == (8,) + assert torch.all(inv > 0) + # First freq is 1.0 when θ^{-0}=1 + assert torch.isclose(inv[0], torch.tensor(1.0)) + + +def test_rotation_preserves_norm() -> None: + rope = OdysseyRoPE( + RopeConfig(head_dim=16, rotary_dim=16, max_position_embeddings=128) + ) + x = torch.randn(7, 3, 16) + y = rope(x, position_offset=11) + nx = torch.linalg.vector_norm(x, dim=-1) + ny = torch.linalg.vector_norm(y, dim=-1) + assert torch.allclose(nx, ny, atol=1e-5) + + +def test_partial_rotary_leaves_tail() -> None: + rope = OdysseyRoPE(RopeConfig(head_dim=8, rotary_dim=4, max_position_embeddings=32)) + x = torch.randn(5, 2, 8) + y = rope(x, position_offset=4) + assert torch.equal(x[..., 4:], y[..., 4:]) + assert not torch.allclose(x[..., :4], y[..., :4]) + + +def test_adjacent_pair_formula() -> None: + x = torch.tensor([[1.0, 0.0, 0.0, 1.0]]) + cos = torch.tensor([[0.0, 1.0]]) # 90° on pair0, 0° on pair1 + sin = torch.tensor([[1.0, 0.0]]) + y = rotate_adjacent_pairs(x, cos, sin) + # pair0: (1,0) → (0,1); pair1: (0,1) → (0,1) + assert torch.allclose(y, torch.tensor([[0.0, 1.0, 0.0, 1.0]]), atol=1e-6) + + +def test_different_positions_change_output() -> None: + rope = OdysseyRoPE(RopeConfig(head_dim=4, rotary_dim=4, max_position_embeddings=64)) + x = torch.tensor([[1.0, 0.0, 0.0, 1.0]]) + y0 = rope(x, position_offset=0) + y7 = rope(x, position_offset=7) + assert (y0 - y7).abs().sum() > 1e-3 diff --git a/tests/test_scaling.py b/tests/test_scaling.py new file mode 100644 index 0000000..10b2864 --- /dev/null +++ b/tests/test_scaling.py @@ -0,0 +1,55 @@ +"""RoPE scaling tests.""" + +from __future__ import annotations + +import pytest +import torch + +from model import OdysseyRoPE, RopeConfig + + +def test_linear_scaling_changes_angles() -> None: + base = RopeConfig( + head_dim=8, + rotary_dim=8, + max_position_embeddings=64, + scaling="none", + scaling_factor=1.0, + ) + scaled = RopeConfig( + head_dim=8, + rotary_dim=8, + max_position_embeddings=64, + scaling="linear", + scaling_factor=2.0, + ) + x = torch.randn(1, 1, 8) + y1 = OdysseyRoPE(base)(x, position_offset=10) + y2 = OdysseyRoPE(scaled)(x, position_offset=10) + assert not torch.allclose(y1, y2) + + +def test_linear_scale_matches_half_position() -> None: + """m'=m/2 at offset 10 ≈ unscaled offset 5 for the same x.""" + cfg_s = RopeConfig( + head_dim=8, + rotary_dim=8, + max_position_embeddings=64, + scaling="linear", + scaling_factor=2.0, + ) + cfg = RopeConfig( + head_dim=8, + rotary_dim=8, + max_position_embeddings=64, + scaling="none", + ) + x = torch.randn(1, 2, 8) + y_scaled = OdysseyRoPE(cfg_s)(x, position_offset=10) + y_half = OdysseyRoPE(cfg)(x, position_offset=5) + assert torch.allclose(y_scaled, y_half, atol=1e-5) + + +def test_rejects_yarn() -> None: + with pytest.raises(ValueError, match="scaling"): + RopeConfig(scaling="yarn") # type: ignore[arg-type]