dspy-wsp 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (84) hide show
  1. dspy_wsp-0.1.0/LICENSE +25 -0
  2. dspy_wsp-0.1.0/PKG-INFO +345 -0
  3. dspy_wsp-0.1.0/README.md +313 -0
  4. dspy_wsp-0.1.0/pyproject.toml +44 -0
  5. dspy_wsp-0.1.0/setup.cfg +4 -0
  6. dspy_wsp-0.1.0/src/dspy_wsp/__init__.py +40 -0
  7. dspy_wsp-0.1.0/src/dspy_wsp/_kernel_main.py +205 -0
  8. dspy_wsp-0.1.0/src/dspy_wsp/backends/__init__.py +3 -0
  9. dspy_wsp-0.1.0/src/dspy_wsp/backends/_vllm_worker.py +48 -0
  10. dspy_wsp-0.1.0/src/dspy_wsp/backends/base.py +224 -0
  11. dspy_wsp-0.1.0/src/dspy_wsp/backends/decision.py +615 -0
  12. dspy_wsp-0.1.0/src/dspy_wsp/backends/hf.py +390 -0
  13. dspy_wsp-0.1.0/src/dspy_wsp/backends/mlx.py +668 -0
  14. dspy_wsp-0.1.0/src/dspy_wsp/backends/vllm.py +317 -0
  15. dspy_wsp-0.1.0/src/dspy_wsp/blocks.py +34 -0
  16. dspy_wsp-0.1.0/src/dspy_wsp/budget.py +31 -0
  17. dspy_wsp-0.1.0/src/dspy_wsp/cli.py +57 -0
  18. dspy_wsp-0.1.0/src/dspy_wsp/decision/__init__.py +6 -0
  19. dspy_wsp-0.1.0/src/dspy_wsp/decision/client.py +93 -0
  20. dspy_wsp-0.1.0/src/dspy_wsp/decision/families/__init__.py +56 -0
  21. dspy_wsp-0.1.0/src/dspy_wsp/decision/families/d1_lm/__init__.py +0 -0
  22. dspy_wsp-0.1.0/src/dspy_wsp/decision/families/d1_lm/encoding.py +297 -0
  23. dspy_wsp-0.1.0/src/dspy_wsp/decision/families/d1_lm/lfm3_mlx.py +167 -0
  24. dspy_wsp-0.1.0/src/dspy_wsp/decision/families/d1_lm/lfm3_torch.py +171 -0
  25. dspy_wsp-0.1.0/src/dspy_wsp/decision/families/d1_lm/mlx.py +84 -0
  26. dspy_wsp-0.1.0/src/dspy_wsp/decision/families/d1_lm/torch.py +92 -0
  27. dspy_wsp-0.1.0/src/dspy_wsp/decision/families/d1_omni/__init__.py +0 -0
  28. dspy_wsp-0.1.0/src/dspy_wsp/decision/families/d1_omni/encoding.py +141 -0
  29. dspy_wsp-0.1.0/src/dspy_wsp/decision/families/d1_omni/mlx.py +275 -0
  30. dspy_wsp-0.1.0/src/dspy_wsp/decision/families/d1_omni/torch.py +240 -0
  31. dspy_wsp-0.1.0/src/dspy_wsp/decision/families/nanojev/__init__.py +0 -0
  32. dspy_wsp-0.1.0/src/dspy_wsp/decision/families/nanojev/encoding.py +87 -0
  33. dspy_wsp-0.1.0/src/dspy_wsp/decision/families/nanojev/mlx.py +107 -0
  34. dspy_wsp-0.1.0/src/dspy_wsp/decision/families/nanojev/torch.py +107 -0
  35. dspy_wsp-0.1.0/src/dspy_wsp/decision/loading.py +211 -0
  36. dspy_wsp-0.1.0/src/dspy_wsp/decision/schema.py +106 -0
  37. dspy_wsp-0.1.0/src/dspy_wsp/decision/server.py +97 -0
  38. dspy_wsp-0.1.0/src/dspy_wsp/edit_gepa/__init__.py +11 -0
  39. dspy_wsp-0.1.0/src/dspy_wsp/edit_gepa/language.py +123 -0
  40. dspy_wsp-0.1.0/src/dspy_wsp/edit_gepa/search.py +188 -0
  41. dspy_wsp-0.1.0/src/dspy_wsp/edits.py +117 -0
  42. dspy_wsp-0.1.0/src/dspy_wsp/forward.py +551 -0
  43. dspy_wsp-0.1.0/src/dspy_wsp/kernel.py +211 -0
  44. dspy_wsp-0.1.0/src/dspy_wsp/optimizer.py +331 -0
  45. dspy_wsp-0.1.0/src/dspy_wsp/patch.py +226 -0
  46. dspy_wsp-0.1.0/src/dspy_wsp/prompts.py +309 -0
  47. dspy_wsp-0.1.0/src/dspy_wsp/search.py +328 -0
  48. dspy_wsp-0.1.0/src/dspy_wsp/session.py +235 -0
  49. dspy_wsp-0.1.0/src/dspy_wsp/stats.py +48 -0
  50. dspy_wsp-0.1.0/src/dspy_wsp/tensors.py +107 -0
  51. dspy_wsp-0.1.0/src/dspy_wsp.egg-info/PKG-INFO +345 -0
  52. dspy_wsp-0.1.0/src/dspy_wsp.egg-info/SOURCES.txt +82 -0
  53. dspy_wsp-0.1.0/src/dspy_wsp.egg-info/dependency_links.txt +1 -0
  54. dspy_wsp-0.1.0/src/dspy_wsp.egg-info/entry_points.txt +2 -0
  55. dspy_wsp-0.1.0/src/dspy_wsp.egg-info/requires.txt +23 -0
  56. dspy_wsp-0.1.0/src/dspy_wsp.egg-info/top_level.txt +1 -0
  57. dspy_wsp-0.1.0/tests/test_blocks_and_stats.py +30 -0
  58. dspy_wsp-0.1.0/tests/test_compile.py +115 -0
  59. dspy_wsp-0.1.0/tests/test_d1.py +112 -0
  60. dspy_wsp-0.1.0/tests/test_d1_lm_backend.py +146 -0
  61. dspy_wsp-0.1.0/tests/test_d1_lm_encoding.py +142 -0
  62. dspy_wsp-0.1.0/tests/test_d1_lm_network.py +184 -0
  63. dspy_wsp-0.1.0/tests/test_d1_lm_reference.py +100 -0
  64. dspy_wsp-0.1.0/tests/test_d1_reference.py +66 -0
  65. dspy_wsp-0.1.0/tests/test_decision_backend.py +191 -0
  66. dspy_wsp-0.1.0/tests/test_decision_dspy.py +117 -0
  67. dspy_wsp-0.1.0/tests/test_decision_loading.py +129 -0
  68. dspy_wsp-0.1.0/tests/test_decision_optimizer.py +105 -0
  69. dspy_wsp-0.1.0/tests/test_decision_registry.py +28 -0
  70. dspy_wsp-0.1.0/tests/test_decision_schema.py +77 -0
  71. dspy_wsp-0.1.0/tests/test_edit_gepa.py +67 -0
  72. dspy_wsp-0.1.0/tests/test_edits_and_patch.py +101 -0
  73. dspy_wsp-0.1.0/tests/test_engine.py +36 -0
  74. dspy_wsp-0.1.0/tests/test_forward.py +132 -0
  75. dspy_wsp-0.1.0/tests/test_forward_decision.py +126 -0
  76. dspy_wsp-0.1.0/tests/test_gepa_search.py +257 -0
  77. dspy_wsp-0.1.0/tests/test_hf_backend.py +139 -0
  78. dspy_wsp-0.1.0/tests/test_kernel.py +102 -0
  79. dspy_wsp-0.1.0/tests/test_mlx_backend.py +133 -0
  80. dspy_wsp-0.1.0/tests/test_nanojev.py +105 -0
  81. dspy_wsp-0.1.0/tests/test_nanojev_reference.py +74 -0
  82. dspy_wsp-0.1.0/tests/test_prompt_snapshot.py +26 -0
  83. dspy_wsp-0.1.0/tests/test_typed_decisions_example.py +106 -0
  84. dspy_wsp-0.1.0/tests/test_vllm_backend.py +179 -0
dspy_wsp-0.1.0/LICENSE ADDED
@@ -0,0 +1,25 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Drew Breunig
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
22
+
23
+ The writer's task prompt in src/dspy_wsp/prompts.py is adapted from the
24
+ Experiment 8 prompt by Rikiya Takehi, Jayden Teoh, Alex Zhang and Omar Khattab,
25
+ released under CC BY 4.0 at https://github.com/RikiyaT/lm_understand-logs.
@@ -0,0 +1,345 @@
1
+ Metadata-Version: 2.4
2
+ Name: dspy-wsp
3
+ Version: 0.1.0
4
+ Summary: Weight-space programming as a DSPy optimizer: a frontier LM reads and edits a small model's weights, with no gradients.
5
+ Author-email: Drew Breunig <dbreunig@gmail.com>
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/cmpnd-ai/dspy-wsp
8
+ Project-URL: Issues, https://github.com/cmpnd-ai/dspy-wsp/issues
9
+ Keywords: dspy,llm,optimizer,interpretability,weight editing
10
+ Requires-Python: >=3.10
11
+ Description-Content-Type: text/markdown
12
+ License-File: LICENSE
13
+ Requires-Dist: dspy[typesafe]<3.5,>=3.4
14
+ Requires-Dist: typesafe-sdk==0.7.5
15
+ Requires-Dist: numpy>=1.26
16
+ Requires-Dist: safetensors>=0.4
17
+ Requires-Dist: tokenizers>=0.20
18
+ Requires-Dist: huggingface_hub>=0.24
19
+ Requires-Dist: torch>=2.2
20
+ Requires-Dist: transformers>=4.45
21
+ Requires-Dist: mlx>=0.30; sys_platform == "darwin" and platform_machine == "arm64"
22
+ Requires-Dist: mlx-lm>=0.32; sys_platform == "darwin" and platform_machine == "arm64"
23
+ Provides-Extra: vllm
24
+ Requires-Dist: vllm>=0.32; extra == "vllm"
25
+ Provides-Extra: examples
26
+ Requires-Dist: datasets>=2.14; extra == "examples"
27
+ Requires-Dist: math-verify>=0.7; extra == "examples"
28
+ Provides-Extra: dev
29
+ Requires-Dist: pytest>=8; extra == "dev"
30
+ Requires-Dist: requests>=2.31; extra == "dev"
31
+ Dynamic: license-file
32
+
33
+ # dspy-wsp
34
+
35
+ Weight-space programming as a DSPy optimizer.
36
+
37
+ `WeightSpace` improves a DSPy program by having a frontier LM, the *writer*, edit the student model's weights directly. The writer reads the weights in a sandboxed Python kernel. It runs your program on practice examples through a budgeted forward-pass tool, works out what goes wrong inside the network, and saves targeted edits. It uses no gradients and no training.
38
+
39
+ The method comes from [*Sparks of Weight-Space Programming*](https://rikiyat.github.io/blog/weight-space-programming/) by Rikiya Takehi, Jayden Teoh, Alex Zhang and Omar Khattab (MIT CSAIL, 2026). In their Experiment 8, a frontier model edited 0.016% of Qwen3.5-0.8B's weights. The edit raised MATH-500 from 7.9% to 48.1% by removing a direction that drove the model's looping. This package rebuilds that harness around DSPy programs and metrics.
40
+
41
+ This is an early release. We are releasing it outside the core DSPy library so that more people can try it and tell us what breaks. Please report problems and results, and tell us which models you would like us to support, in [GitHub issues](https://github.com/cmpnd-ai/dspy-wsp/issues) or on the DSPy Discord.
42
+
43
+ ## Install
44
+
45
+ ```bash
46
+ pip install dspy-wsp
47
+ apt install bubblewrap # sandboxes the writer's code (Linux)
48
+ ```
49
+
50
+ The base install runs every supported decision model, with MLX on Apple silicon and PyTorch elsewhere. Two extras add more:
51
+
52
+ - `pip install "dspy-wsp[vllm]"` adds the vLLM backend for text-generating LLMs on NVIDIA GPUs.
53
+ - `pip install "dspy-wsp[examples]"` adds the datasets for the benchmark scripts in `examples/`.
54
+
55
+ ## Use
56
+
57
+ `WeightSpace` works like the other DSPy optimizers. It takes your metric and your module, and `compile` returns the improved module. Here the student is d1-omni-600M, a decision model that runs on a laptop:
58
+
59
+ ```python
60
+ from typing import Literal
61
+
62
+ import dspy
63
+ from dspy_wsp import WeightSpace, load_decision_backend
64
+
65
+
66
+ class Triage(dspy.Signature):
67
+ """Triage a support ticket."""
68
+
69
+ ticket: str = dspy.InputField()
70
+ urgent: bool = dspy.OutputField(desc="Does this ticket need a reply today?")
71
+ team: Literal["billing", "technical", "account"] = dspy.OutputField(desc="Which team should handle it?")
72
+
73
+
74
+ def metric(example, pred, trace=None):
75
+ return ((pred.urgent == example.urgent) + (pred.team == example.team)) / 2
76
+
77
+
78
+ backend = load_decision_backend("d1") # d1-omni-600M at the tested revision
79
+ optimizer = WeightSpace(
80
+ metric=metric,
81
+ writer_lm=dspy.LM("openai/gpt-6-astra", model_type="responses", reasoning_effort="max"),
82
+ backend=backend,
83
+ max_passes=25,
84
+ holdout=300,
85
+ run_dir="runs/triage",
86
+ )
87
+ compiled = optimizer.compile(dspy.Predict(Triage), trainset=train, valset=test)
88
+
89
+ compiled(ticket="I was charged twice this month.") # runs on the edited weights
90
+ print(compiled.wsp_patch.summary())
91
+ print(compiled.wsp_report["validation"]["paired"])
92
+ ```
93
+
94
+ `trainset` becomes the writer's practice set. `holdout` moves 300 of those examples to a validation set that the writer can score edits on but never see. The writer never sees `valset`; the optimizer uses it once at the end, comparing original and edited weights on the same examples with the same seeds. Metrics can return a number, a bool, or `dspy.Prediction(score=..., feedback=...)` as GEPA metrics do. The writer reads any feedback text in the records.
95
+
96
+ A decision model reads `bool` outputs as yes or no questions and `Literal` outputs as choices. Each output needs a question in its `desc`. The `Choice`, `Noul` and `Score` types in `dspy.experimental` add criteria for each option and ordered score levels, and they return each option's probability.
97
+
98
+ ## Results
99
+
100
+ We edited three decision models on the Typed Decisions benchmark with `gpt-6-astra` at maximum reasoning effort as the writer. The writer had one conversation with 25 forward passes for each model. The table scores the answers each program returns on 400 held-out test cases. Each gain is the paired gain over the plain model in points, with its standard error.
101
+
102
+ | Version | d1-omni-600M | NanoJev (Qwen3-0.6B) | d1-2B |
103
+ |---|---|---|---|
104
+ | Plain | 50.9% | 39.2% | 60.6% |
105
+ | ReAnchor | 60.2% (+9.3 ± 1.1) | 50.7% (+11.6 ± 1.3) | 70.1% (+9.5 ± 1.2) |
106
+ | WeightSpace | 66.9% (+16.1 ± 1.2) | 52.4% (+13.3 ± 1.1) | 74.0% (+13.4 ± 1.2) |
107
+ | WeightSpace + ReAnchor | 69.1% (+18.3 ± 1.3) | 56.4% (+17.3 ± 1.3) | 77.9% (+17.3 ± 1.3) |
108
+
109
+ ReAnchor fits how the program turns probabilities into answers and changes no weights. Each run took between 78 and 167 minutes and cost between $5.42 and $7.63 in writer calls. With `gpt-6.1-sol` as the writer, the WeightSpace edits alone gained 5.8, 12.1 and −0.6 points on the same three models.
110
+
111
+ [The results folder](results/2026-11-07-typed-decisions-astra/) has every run, edit and comparison. [Decision models](#decision-models) covers the earlier runs and the Procedural Typed Decisions benchmark.
112
+
113
+ ## Decision models
114
+
115
+ `WeightSpace` also edits System One decision models: models that read a state and typed questions (Choice, Noul and Score) and return a probability for every option, with no text generation. Three families are supported, each with our own loader in MLX and in PyTorch. Each family's code is in its own folder under `src/dspy_wsp/decision/families/`:
116
+
117
+ - `D1Backend` runs Liquid AI's d1-omni models (text path), e.g. d1-omni-600M: an encoder with a decision head.
118
+ - `NanoJevBackend` runs NanoJev models: a Qwen3 backbone with a decision head.
119
+ - `D1LMBackend` runs Liquid AI's d1 models built on LFM3, e.g. d1-2B. It reads each answer from the language model's output at the question's answer position, so it has no separate head. Its output layer is the token embedding, which stays protected, so the writer's edits land in the backbone. It reads text states and answers Choice, Noul and Score questions. It refuses d1-2B's other question types (text, number, date, box and multiple choice).
120
+
121
+ We have tested `WeightSpace` on three decision models so far: d1-omni-600M, NanoJev and d1-2B. If you run it on another open decision model, please tell us how it went in [GitHub issues](https://github.com/cmpnd-ai/dspy-wsp/issues) or on the DSPy Discord. Please also tell us which models you would like us to support, especially decision models.
122
+
123
+ `load_decision_backend` opens a model by alias, Hugging Face repo id or local folder. It reads the checkpoint's config to pick the family, and it refuses other architectures with a message that names the supported ones. The aliases `"d1"`, `"nanojev"` and `"d1-2b"` open the checkpoints we tested, at pinned revisions.
124
+
125
+ ```python
126
+ from dspy_wsp import load_decision_backend
127
+
128
+ backend = load_decision_backend("nanojev") # C-Tianyu/NanoJev at the tested revision
129
+ backend = load_decision_backend("someone/nanojev-finetune") # any checkpoint of a supported family
130
+ backend = load_decision_backend("/path/to/checkpoint", framework="torch")
131
+ ```
132
+
133
+ Before you run a new checkpoint, check it:
134
+
135
+ ```bash
136
+ dspy-wsp check someone/nanojev-finetune
137
+ ```
138
+
139
+ The check loads the model, confirms that its tensors map to our network, and answers sample Choice, Noul and Score questions. It confirms that the answers are valid distributions and that a request gives the same answer alone and in a batch. Finally, it edits one tensor, confirms that the answers move, and confirms that restoring the tensor restores them. `check_decision_model` runs the same check from Python.
140
+
141
+ The backend runs a local `/v1/systemone` server, and a DSPy program calls it through DSPy's TypeSafe client, as it would call TypeSafe's hosted Jev. Give the program's signatures `Choice`, `Noul` or `Score` outputs from `dspy.experimental`.
142
+
143
+ ```python
144
+ from dspy_wsp import WeightSpace, load_decision_backend
145
+
146
+ backend = load_decision_backend("d1", framework="mlx")
147
+ optimizer = WeightSpace(metric=metric, writer_lm=dspy.LM("openai/gpt-6.1-sol"), backend=backend,
148
+ max_passes=25, holdout=300, holdout_key=lambda ex: ex.workflow)
149
+ compiled = optimizer.compile(program, trainset=train, valset=test)
150
+ compiled.wsp_patch.write_checkpoint(backend.weights_dir, "edited-d1/")
151
+ ```
152
+
153
+ `write_checkpoint` copies the model's folder, including its subfolders and the model's own code files, with the edited tensors substituted and checked against their hashes. The result has the same layout as the base folder, so the model's own code and `load_decision_backend` both open it. The example scripts take `--model` with an alias, repo id or folder, and they save the edited checkpoint in the run folder's `checkpoint/`.
154
+
155
+ For a decision model, the writer gets `decide` passes, which run the program and score it, and `read` passes, which return each option's logit and the decision head's inputs. It may recompute the head from those reads and derive edits in closed form. `holdout` moves training examples to a validation set that the writer can score edits on but never see. `examples/typed_decisions.py` runs the Typed Decisions benchmark. Two runs on a MacBook Pro with MLX, with `openai/gpt-6.1-sol` as the writer and 25 passes each, gave these results on the 400 held-out test cases:
156
+
157
+ | Model | Tensor edited | Original | Edited | Paired gain | Improved / worse | Run |
158
+ |---|---|---|---|---|---|---|
159
+ | d1-omni-600M | `head.scorer.3.weight` (1,024 weights) | 53.4% | 60.6% | +7.25 ± 1.05 points | 150 / 63 | [results](results/2026-11-04-typed-decisions-d1/) |
160
+ | NanoJev (Qwen3-0.6B) | `scalar.weight` (1,024 weights) | 38.6% | 55.7% | +17.10 ± 1.35 points | 243 / 63 | [results](results/2026-11-04-typed-decisions-nanojev/) |
161
+
162
+ In both runs the writer derived the edit in closed form from the decision head's inputs, confirmed it on 300 held-out validation cases, and never saw the test cases. Each run took between 75 and 110 minutes.
163
+
164
+ WeightSpace also combines with `dspy.experimental.ReAnchor`, which fits how a program turns probabilities into answers (Noul thresholds, Score cut points and Choice weights) without changing weights. Scored on the answers NanoJev's program returns, on the same 400 test cases:
165
+
166
+ | NanoJev | Decided accuracy | Paired gain over plain |
167
+ |---|---|---|
168
+ | Plain | 39.2% | |
169
+ | ReAnchor | 50.7% | +11.6 ± 1.3 points |
170
+ | WeightSpace | 51.3% | +12.1 ± 1.2 points |
171
+ | WeightSpace + ReAnchor | 55.7% | +16.6 ± 1.4 points |
172
+
173
+ See [the comparison](results/2026-11-04-typed-decisions-nanojev/reanchor/) for how the two methods divide the gain.
174
+
175
+ On three balanced, rule-based tasks from Procedural Typed Decisions, overnight on the same laptop, the same kind of edit gave larger gains on 400 test cases:
176
+
177
+ | Model | Tensor edited | Original | Edited | Paired gain | Improved / worse | Run |
178
+ |---|---|---|---|---|---|---|
179
+ | d1-omni-600M | `head.scorer.3.weight` (1,024 weights) | 52.3% | 72.3% | +20.0 ± 2.3 points | 179 / 54 | [results](results/2026-11-04-procedural-decisions/) |
180
+ | NanoJev (Qwen3-0.6B) | `scalar.weight` (1,024 weights) | 49.3% | 72.3% | +23.1 ± 2.3 points | 193 / 54 | [results](results/2026-11-04-procedural-decisions/) |
181
+
182
+ These tasks have balanced labels. Edits that only shifted option preferences gained nothing, so the gains come from how the head reads each record.
183
+
184
+ Scored on the answers each program returns, on the same 400 test cases:
185
+
186
+ | Version | d1-omni-600M | NanoJev |
187
+ |---|---|---|
188
+ | Plain | 52.3% | 47.8% |
189
+ | ReAnchor | 58.7% (+6.4 ± 2.2) | 52.8% (+5.0 ± 1.7) |
190
+ | WeightSpace | 64.2% (+11.9 ± 2.3) | 68.4% (+20.6 ± 2.3) |
191
+ | WeightSpace + ReAnchor | 69.9% (+17.7 ± 2.4) | 71.3% (+23.4 ± 2.3) |
192
+ | WeightSpace (GEPA) | 66.9% (+14.7 ± 2.2) | 65.6% (+17.8 ± 2.4) |
193
+ | WeightSpace (GEPA) + ReAnchor | 71.2% (+18.9 ± 2.3) | 67.9% (+20.1 ± 2.4) |
194
+
195
+ The WeightSpace rows come from one 25-pass conversation. The GEPA rows come from a [population search](#population-search) with 51 passes and at least 10 rounds. The search beat the conversation on d1-omni-600M and trailed it on NanoJev. Every difference between the two is smaller than two standard errors.
196
+
197
+ ## Generating LLMs
198
+
199
+ `WeightSpace` also edits text-generating LLMs, as the paper did. A pass on an LLM generates many long answers, so we recommend `VLLMBackend` on a GPU machine. `HFBackend` and `MLXBackend` suit small experiments.
200
+
201
+ A backend holds the student model, swaps edits in and out, and runs every forward pass.
202
+
203
+ - **`VLLMBackend`** runs vLLM on NVIDIA GPUs for runs at the paper's scale. vLLM schedules every concurrent LM call, so the forward-pass tool keeps a whole pass in flight (`concurrency=512`). Teacher-forced reads go to an `HFBackend` on the same GPU, because vLLM cannot return hidden states or MLP activations. Both models hold the same edits. To swap edits, the backend waits for running requests to finish, reloads the checkpoint in place with the edited tensors substituted, and clears vLLM's prefix cache. `gpu_memory_utilization` (default 0.6) sets vLLM's share of GPU memory; the Hugging Face model uses the rest. `batch_invariant=True` makes each answer depend only on its prompt and seed, at some cost in speed.
204
+
205
+ ```python
206
+ from dspy_wsp import VLLMBackend
207
+
208
+ backend = VLLMBackend("Qwen/Qwen3.5-0.8B", revision="<commit>", max_tokens=32_768, presence_penalty=1.5,
209
+ chat_template_kwargs={"enable_thinking": True})
210
+ ```
211
+
212
+ - **`HFBackend`** runs transformers in-process on CUDA, MPS or CPU. It generates one answer at a time, so it suits small models and short answers.
213
+ - **`MLXBackend`** runs mlx-lm on Apple silicon. It decodes up to `batch_size` answers together, and the forward-pass tool runs that many programs at once to fill the batch. Each answer samples from its own seed. It also supports the presence penalty in Qwen's recommended sampling settings. On an M2 Max it generates about 650 tokens per second across a batch of 32 answers from Qwen3.5-0.8B, which is several times faster than `HFBackend` on the same machine.
214
+
215
+ ```python
216
+ from dspy_wsp import MLXBackend
217
+
218
+ backend = MLXBackend("Qwen/Qwen3.5-0.8B", revision="<commit>", batch_size=64, max_tokens=32_768,
219
+ presence_penalty=1.5, chat_template_kwargs={"enable_thinking": True})
220
+ ```
221
+
222
+ `MLXBackend` reads attention probabilities only from softmax-attention blocks. In Qwen3.5 those are blocks 3, 7, 11, 15, 19 and 23; the other blocks use linear attention.
223
+
224
+ ## Replicate the paper
225
+
226
+ The paper's Experiment 8 edited Qwen3.5-0.8B to stop its looping on MATH-500. Two scripts rebuild that setup:
227
+
228
+ - `examples/math500_qwen_vllm.py` runs it on one GPU: 50 passes of up to 512 answers, and a final check on all 500 MATH-500 problems with 4 samples each.
229
+ - `examples/math500_qwen_mlx.py` runs a smaller version on a Mac. Its `--baseline N` flag scores the unedited model without a writer.
230
+
231
+ This path is a bigger step than the decision models, and we have tested it less. A full run at the paper's scale takes many GPU hours and a long writer session. `results/2026-11-04-math500-qwen-vllm-1h/` documents a one-hour proof-of-concept run on eight A100 GPUs, with the checks we ran before it. We have not yet repeated the paper's full result. If you run this setup, please share what you find in [GitHub issues](https://github.com/cmpnd-ai/dspy-wsp/issues) or on the DSPy Discord.
232
+
233
+ ## How a run works
234
+
235
+ The writer works in turns. Each turn it explains its next step and sends exactly one block:
236
+
237
+ - **` ```python `** runs in a persistent kernel with no GPU and, under the sandbox, no network. The kernel predefines `NAMES`, `SHAPES`, `load(name)`, `load_original(name)`, `save_edit(name, tensor)`, `drop_edit(name)`, `edited()` and `tok`. Gradients are disabled.
238
+ - **` ```forward `** sends a JSON request to the forward-pass tool, the only way to run the model:
239
+ - `generate` runs your program on up to 512 practice items with the original weights, the edited weights, or `"both"`. Every answer is scored with your metric. `"both"` uses the same seed for each item and reports the paired gain with its standard error. The records keep every LM call's full prompt and completion. Pooled activations are optional.
240
+ - `teacher_forced` reads up to 1,024 texts and returns log-probs, top-k predictions, hidden states, MLP activations or attention probabilities as `.npy` files.
241
+
242
+ Each forward pass costs one unit of budget, whatever its size. A refused request costs nothing. The session continues until the writer has used every pass. Then it leaves its chosen edit in `work/edited/` and writes a final answer.
243
+
244
+ The harness follows the paper's rules:
245
+
246
+ - **Guardrails:** `save_edit` refuses edits to the embeddings, the output layer and the final norm. It also refuses edits that only rescale a whole tensor. The host checks every edit again before each pass.
247
+ - **Limits:** the writer gets 300 Python runs, 15 minutes per run, and 10,000 characters of output per run.
248
+ - **Explanations:** a block sent without an explanation is returned unrun, at most twice in a row.
249
+ - **Handoffs:** every 40 turns the conversation restarts with the task, the writer's `notes.md` and its current edit.
250
+
251
+ Every turn is written to disk. Pass `resume=True` with the same `run_dir` to continue an interrupted run.
252
+
253
+ ### Population search
254
+
255
+ By default, a run is one long conversation. For decision models, `search=GEPASearch(...)` runs GEPA's population search instead. The writer still works in the Python kernel with the same tools. The search keeps a population of candidate edits and runs in rounds:
256
+
257
+ 1. GEPA picks a parent from its Pareto frontier over the validation examples.
258
+ 2. GEPA scores the parent on a minibatch of practice examples.
259
+ 3. The writer gets a short session that starts from the parent's edit and notes, with the minibatch's feedback and a small budget of forward passes. The edit it leaves becomes the child.
260
+ 4. GEPA scores the child on the same minibatch. If the child scores higher, GEPA scores it on the validation set and adds it to the population.
261
+
262
+ ```python
263
+ from dspy_wsp import GEPASearch, WeightSpace
264
+
265
+ optimizer = WeightSpace(metric=metric, writer_lm=writer, backend=backend, max_passes=51, holdout=300,
266
+ search=GEPASearch(min_rounds=10, minibatch_size=8, passes_per_proposal=2))
267
+ ```
268
+
269
+ - `max_passes` counts every pass: the writer's passes and each of GEPA's scorings. A round costs at most `3 + passes_per_proposal` passes.
270
+ - `WeightSpace` refuses a `max_passes` that can't pay for `min_rounds` rounds, and a `max_code_runs` below `min_rounds * code_runs_per_proposal` (30 by default).
271
+ - A parent that scores perfectly on its minibatch can't be beaten there, so GEPA skips that round's writer session. The parent's scoring still costs one pass.
272
+ - A writer call that times out, drops its connection or runs past the model's output limit is retried twice. If it still fails, or every example fails to run, the search stops. The run then finishes with the best candidate so far, and the report's `search.error` says why.
273
+ - Writers can't run validation passes in a round, because GEPA uses the validation set to select candidates.
274
+ - Each round is written to `rounds.jsonl` and `progress.md` when it ends. Every candidate's edit and notes are kept in `candidates/`.
275
+ - Pass `resume=True` with the same `run_dir` to continue a stopped search. It picks up GEPA's saved state, the round log and the pass count, and it reuses the seed's saved validation scores.
276
+
277
+ The example scripts take `--search gepa --rounds 10`. To compare the two searches, give the conversation the same `--passes`.
278
+
279
+ ## What you get back
280
+
281
+ ```
282
+ runs/qwen-math/
283
+ task_prompt.txt what the writer was told
284
+ turns/tNNN/ each turn's prompt and reply
285
+ transcript.jsonl every Python run and forward pass
286
+ work/ the writer's folder: notes.md, edited/, gen_k.jsonl, fwd_k/, practice/
287
+ final_answer.md the writer's account of its edit and the evidence
288
+ final_check/ the valset comparison (never shown to the writer)
289
+ patch/ manifest.json + edit.npz
290
+ report.json
291
+ ```
292
+
293
+ `compiled.save("program.json")` records the base model, its settings and the patch location. `program.load("program.json", allow_unsafe_lm_state=True)` rebuilds the edited model; DSPy asks for that flag before it imports a custom engine.
294
+
295
+ The patch stores only the changed entries, plus the sha256 of every base and edited tensor. Applying it checks both hashes, so anyone with the public base model can rebuild the edited weights bit for bit.
296
+
297
+ To deploy, write a full checkpoint and serve it like any other model:
298
+
299
+ ```bash
300
+ dspy-wsp apply ~/.cache/huggingface/.../Qwen3.5-0.8B runs/qwen-math/patch qwen-edited/
301
+ vllm serve qwen-edited/
302
+ ```
303
+
304
+ ```python
305
+ program.set_lm(dspy.LM("openai/qwen-edited", api_base="http://localhost:8000/v1"))
306
+ ```
307
+
308
+ ## Safety
309
+
310
+ The writer's code runs on your machine. With `sandbox="auto"` (the default), the kernel runs under bubblewrap: the filesystem is read-only except the work folder, and there is no network. If `bwrap` is missing, the optimizer refuses to start. `sandbox="none"` runs the writer's code with your permissions and network; use it only on a machine you're willing to hand to the writer.
311
+
312
+ ## Differences from the paper
313
+
314
+ - The task is your DSPy program and metric, not fixed math benchmarks. Generation passes run the program; the records show every LM call.
315
+ - The paper's writer was `gpt-6-astra` at maximum reasoning effort, driven through the Codex CLI with every built-in tool switched off. Here the writer is any `dspy.LM`. The writer matters: on our decision model runs, `gpt-6-astra` wrote better edits than `gpt-6.1-sol` on two of three models.
316
+ - `activations` takes separate `layers` (hidden-state indices) and `mlp_layers` (block indices) lists, because one list served both in the paper.
317
+ - `fallback_if_worse=True` returns the original weights when the final check shows a loss. This is an addition to the paper's protocol, and it is off by default.
318
+ - Patches store changed entries or whole tensors. The paper's release also has a low-rank form, which this version doesn't produce yet.
319
+
320
+ ## Limits and roadmap
321
+
322
+ - `HFBackend` generates one sequence at a time in-process. That's fine for small models and short answers, but too slow at the paper's scale. There, one pass produced 512 answers of up to 32k tokens, and a pass took up to 30 minutes on a server. `MLXBackend` batches on a Mac, but a 64-problem paired pass at 32k tokens still takes one to two hours.
323
+ - A run is long and costly. The paper's 50-pass run took 158 turns and 20 hours. Budget for the writer's tokens accordingly.
324
+ - Editing removes obstacles; it doesn't add abilities. The paper's Gemma 3 270M follow-up barely moved. That model didn't loop, so the writer had no failure mechanism to repair.
325
+ - Neither the paper nor this package compares against prompt optimization yet. Run GEPA on the same student and budget before you trust a gain.
326
+
327
+ ## Credits and license
328
+
329
+ Code: MIT.
330
+
331
+ - The method and the harness come from [*Sparks of Weight-Space Programming*](https://rikiyat.github.io/blog/weight-space-programming/) by Rikiya Takehi, Jayden Teoh, Alex Zhang and Omar Khattab. Please cite it when you use this package.
332
+ - The writer's task prompt is adapted from the Experiment 8 prompt released by Takehi et al. in [lm_understand-logs](https://github.com/RikiyaT/lm_understand-logs) under CC BY 4.0.
333
+ - Haoze "Hector" He's [weight editing design](https://github.com/HectorHHZ/llm-honing/blob/main/system-design/weight-edit.md) for `llm-honing` is earlier work on editing decision models, including NanoJev, and releasing the edits as checked checkpoints. See [his post](https://x.com/HaozeHe_Hector/status/2118048048148623453).
334
+ - Michael Isaac contributed the GEPA population search. It uses Lakshya Agrawal's [GEPA](https://github.com/gepa-ai/gepa) library.
335
+
336
+ ```bibtex
337
+ @misc{takehi2026weightspace,
338
+ title = {{Sparks of Weight-Space Programming: LLMs can natively read \& write neural network weights}},
339
+ author = {Takehi, Rikiya and Teoh, Jayden and Zhang, Alex and Khattab, Omar},
340
+ year = {2026},
341
+ month = nov,
342
+ howpublished = {\url{https://rikiyat.github.io/blog/weight-space-programming/}},
343
+ note = {Blog post}
344
+ }
345
+ ```