MLX
Safetensors
qwen3_5
kev
decision-model
ternary
mixed-precision
post-training-quantization
experimental
Instructions to use Jakevin/kev-4b-ternary-mlx with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- MLX
How to use Jakevin/kev-4b-ternary-mlx with MLX:
# Download the model from the Hub pip install huggingface_hub[hf_xet] hf download Jakevin/kev-4b-ternary-mlx --local-dir kev-4b-ternary-mlx
- Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- LM Studio
- Atomic Chat
Publish standalone Kev-4B ternary MLX v0.1.0 experimental
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- .gitattributes +1 -0
- LICENSE +202 -0
- NOTICE +7 -0
- README.md +176 -0
- added_tokens.json +28 -0
- config.json +104 -0
- evaluation/development-summary.json +91 -0
- evaluation/memory-long.json +31 -0
- evaluation/memory-short.json +31 -0
- evaluation/native-comparison.json +9 -0
- head.pt +3 -0
- merges.txt +0 -0
- model.safetensors +3 -0
- packed.json +1016 -0
- release-manifest.json +252 -0
- runtime/LICENSE +202 -0
- runtime/README.md +466 -0
- runtime/docs/packed-mlx.md +132 -0
- runtime/kev/__init__.py +0 -0
- runtime/kev/anchors.py +71 -0
- runtime/kev/api.py +165 -0
- runtime/kev/autoresearch.py +181 -0
- runtime/kev/benchmark.py +214 -0
- runtime/kev/budget.py +107 -0
- runtime/kev/calibrate.py +73 -0
- runtime/kev/checkpoint.py +392 -0
- runtime/kev/compare.py +62 -0
- runtime/kev/composition.py +321 -0
- runtime/kev/contrastive.py +306 -0
- runtime/kev/cuda_graphs.py +375 -0
- runtime/kev/data.py +410 -0
- runtime/kev/device.py +30 -0
- runtime/kev/evaluate.py +219 -0
- runtime/kev/experiment.py +492 -0
- runtime/kev/full_ft.py +454 -0
- runtime/kev/fused_qwen35.py +155 -0
- runtime/kev/jev.py +63 -0
- runtime/kev/metrics.py +426 -0
- runtime/kev/mirror.py +84 -0
- runtime/kev/mlx_model.py +229 -0
- runtime/kev/model.py +508 -0
- runtime/kev/packed_model.py +102 -0
- runtime/kev/plot.py +59 -0
- runtime/kev/predictors.py +341 -0
- runtime/kev/publish.py +106 -0
- runtime/kev/rounds.py +1364 -0
- runtime/kev/serve.py +343 -0
- runtime/kev/shared_prefix.py +141 -0
- runtime/kev/study_v3.py +232 -0
- runtime/kev/suite.py +450 -0
.gitattributes
CHANGED
|
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
tokenizer.json filter=lfs diff=lfs merge=lfs -text
|
LICENSE
ADDED
|
@@ -0,0 +1,202 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
|
| 2 |
+
Apache License
|
| 3 |
+
Version 2.0, January 2004
|
| 4 |
+
http://www.apache.org/licenses/
|
| 5 |
+
|
| 6 |
+
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
| 7 |
+
|
| 8 |
+
1. Definitions.
|
| 9 |
+
|
| 10 |
+
"License" shall mean the terms and conditions for use, reproduction,
|
| 11 |
+
and distribution as defined by Sections 1 through 9 of this document.
|
| 12 |
+
|
| 13 |
+
"Licensor" shall mean the copyright owner or entity authorized by
|
| 14 |
+
the copyright owner that is granting the License.
|
| 15 |
+
|
| 16 |
+
"Legal Entity" shall mean the union of the acting entity and all
|
| 17 |
+
other entities that control, are controlled by, or are under common
|
| 18 |
+
control with that entity. For the purposes of this definition,
|
| 19 |
+
"control" means (i) the power, direct or indirect, to cause the
|
| 20 |
+
direction or management of such entity, whether by contract or
|
| 21 |
+
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
| 22 |
+
outstanding shares, or (iii) beneficial ownership of such entity.
|
| 23 |
+
|
| 24 |
+
"You" (or "Your") shall mean an individual or Legal Entity
|
| 25 |
+
exercising permissions granted by this License.
|
| 26 |
+
|
| 27 |
+
"Source" form shall mean the preferred form for making modifications,
|
| 28 |
+
including but not limited to software source code, documentation
|
| 29 |
+
source, and configuration files.
|
| 30 |
+
|
| 31 |
+
"Object" form shall mean any form resulting from mechanical
|
| 32 |
+
transformation or translation of a Source form, including but
|
| 33 |
+
not limited to compiled object code, generated documentation,
|
| 34 |
+
and conversions to other media types.
|
| 35 |
+
|
| 36 |
+
"Work" shall mean the work of authorship, whether in Source or
|
| 37 |
+
Object form, made available under the License, as indicated by a
|
| 38 |
+
copyright notice that is included in or attached to the work
|
| 39 |
+
(an example is provided in the Appendix below).
|
| 40 |
+
|
| 41 |
+
"Derivative Works" shall mean any work, whether in Source or Object
|
| 42 |
+
form, that is based on (or derived from) the Work and for which the
|
| 43 |
+
editorial revisions, annotations, elaborations, or other modifications
|
| 44 |
+
represent, as a whole, an original work of authorship. For the purposes
|
| 45 |
+
of this License, Derivative Works shall not include works that remain
|
| 46 |
+
separable from, or merely link (or bind by name) to the interfaces of,
|
| 47 |
+
the Work and Derivative Works thereof.
|
| 48 |
+
|
| 49 |
+
"Contribution" shall mean any work of authorship, including
|
| 50 |
+
the original version of the Work and any modifications or additions
|
| 51 |
+
to that Work or Derivative Works thereof, that is intentionally
|
| 52 |
+
submitted to Licensor for inclusion in the Work by the copyright owner
|
| 53 |
+
or by an individual or Legal Entity authorized to submit on behalf of
|
| 54 |
+
the copyright owner. For the purposes of this definition, "submitted"
|
| 55 |
+
means any form of electronic, verbal, or written communication sent
|
| 56 |
+
to the Licensor or its representatives, including but not limited to
|
| 57 |
+
communication on electronic mailing lists, source code control systems,
|
| 58 |
+
and issue tracking systems that are managed by, or on behalf of, the
|
| 59 |
+
Licensor for the purpose of discussing and improving the Work, but
|
| 60 |
+
excluding communication that is conspicuously marked or otherwise
|
| 61 |
+
designated in writing by the copyright owner as "Not a Contribution."
|
| 62 |
+
|
| 63 |
+
"Contributor" shall mean Licensor and any individual or Legal Entity
|
| 64 |
+
on behalf of whom a Contribution has been received by Licensor and
|
| 65 |
+
subsequently incorporated within the Work.
|
| 66 |
+
|
| 67 |
+
2. Grant of Copyright License. Subject to the terms and conditions of
|
| 68 |
+
this License, each Contributor hereby grants to You a perpetual,
|
| 69 |
+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
| 70 |
+
copyright license to reproduce, prepare Derivative Works of,
|
| 71 |
+
publicly display, publicly perform, sublicense, and distribute the
|
| 72 |
+
Work and such Derivative Works in Source or Object form.
|
| 73 |
+
|
| 74 |
+
3. Grant of Patent License. Subject to the terms and conditions of
|
| 75 |
+
this License, each Contributor hereby grants to You a perpetual,
|
| 76 |
+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
| 77 |
+
(except as stated in this section) patent license to make, have made,
|
| 78 |
+
use, offer to sell, sell, import, and otherwise transfer the Work,
|
| 79 |
+
where such license applies only to those patent claims licensable
|
| 80 |
+
by such Contributor that are necessarily infringed by their
|
| 81 |
+
Contribution(s) alone or by combination of their Contribution(s)
|
| 82 |
+
with the Work to which such Contribution(s) was submitted. If You
|
| 83 |
+
institute patent litigation against any entity (including a
|
| 84 |
+
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
| 85 |
+
or a Contribution incorporated within the Work constitutes direct
|
| 86 |
+
or contributory patent infringement, then any patent licenses
|
| 87 |
+
granted to You under this License for that Work shall terminate
|
| 88 |
+
as of the date such litigation is filed.
|
| 89 |
+
|
| 90 |
+
4. Redistribution. You may reproduce and distribute copies of the
|
| 91 |
+
Work or Derivative Works thereof in any medium, with or without
|
| 92 |
+
modifications, and in Source or Object form, provided that You
|
| 93 |
+
meet the following conditions:
|
| 94 |
+
|
| 95 |
+
(a) You must give any other recipients of the Work or
|
| 96 |
+
Derivative Works a copy of this License; and
|
| 97 |
+
|
| 98 |
+
(b) You must cause any modified files to carry prominent notices
|
| 99 |
+
stating that You changed the files; and
|
| 100 |
+
|
| 101 |
+
(c) You must retain, in the Source form of any Derivative Works
|
| 102 |
+
that You distribute, all copyright, patent, trademark, and
|
| 103 |
+
attribution notices from the Source form of the Work,
|
| 104 |
+
excluding those notices that do not pertain to any part of
|
| 105 |
+
the Derivative Works; and
|
| 106 |
+
|
| 107 |
+
(d) If the Work includes a "NOTICE" text file as part of its
|
| 108 |
+
distribution, then any Derivative Works that You distribute must
|
| 109 |
+
include a readable copy of the attribution notices contained
|
| 110 |
+
within such NOTICE file, excluding those notices that do not
|
| 111 |
+
pertain to any part of the Derivative Works, in at least one
|
| 112 |
+
of the following places: within a NOTICE text file distributed
|
| 113 |
+
as part of the Derivative Works; within the Source form or
|
| 114 |
+
documentation, if provided along with the Derivative Works; or,
|
| 115 |
+
within a display generated by the Derivative Works, if and
|
| 116 |
+
wherever such third-party notices normally appear. The contents
|
| 117 |
+
of the NOTICE file are for informational purposes only and
|
| 118 |
+
do not modify the License. You may add Your own attribution
|
| 119 |
+
notices within Derivative Works that You distribute, alongside
|
| 120 |
+
or as an addendum to the NOTICE text from the Work, provided
|
| 121 |
+
that such additional attribution notices cannot be construed
|
| 122 |
+
as modifying the License.
|
| 123 |
+
|
| 124 |
+
You may add Your own copyright statement to Your modifications and
|
| 125 |
+
may provide additional or different license terms and conditions
|
| 126 |
+
for use, reproduction, or distribution of Your modifications, or
|
| 127 |
+
for any such Derivative Works as a whole, provided Your use,
|
| 128 |
+
reproduction, and distribution of the Work otherwise complies with
|
| 129 |
+
the conditions stated in this License.
|
| 130 |
+
|
| 131 |
+
5. Submission of Contributions. Unless You explicitly state otherwise,
|
| 132 |
+
any Contribution intentionally submitted for inclusion in the Work
|
| 133 |
+
by You to the Licensor shall be under the terms and conditions of
|
| 134 |
+
this License, without any additional terms or conditions.
|
| 135 |
+
Notwithstanding the above, nothing herein shall supersede or modify
|
| 136 |
+
the terms of any separate license agreement you may have executed
|
| 137 |
+
with Licensor regarding such Contributions.
|
| 138 |
+
|
| 139 |
+
6. Trademarks. This License does not grant permission to use the trade
|
| 140 |
+
names, trademarks, service marks, or product names of the Licensor,
|
| 141 |
+
except as required for reasonable and customary use in describing the
|
| 142 |
+
origin of the Work and reproducing the content of the NOTICE file.
|
| 143 |
+
|
| 144 |
+
7. Disclaimer of Warranty. Unless required by applicable law or
|
| 145 |
+
agreed to in writing, Licensor provides the Work (and each
|
| 146 |
+
Contributor provides its Contributions) on an "AS IS" BASIS,
|
| 147 |
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
| 148 |
+
implied, including, without limitation, any warranties or conditions
|
| 149 |
+
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
| 150 |
+
PARTICULAR PURPOSE. You are solely responsible for determining the
|
| 151 |
+
appropriateness of using or redistributing the Work and assume any
|
| 152 |
+
risks associated with Your exercise of permissions under this License.
|
| 153 |
+
|
| 154 |
+
8. Limitation of Liability. In no event and under no legal theory,
|
| 155 |
+
whether in tort (including negligence), contract, or otherwise,
|
| 156 |
+
unless required by applicable law (such as deliberate and grossly
|
| 157 |
+
negligent acts) or agreed to in writing, shall any Contributor be
|
| 158 |
+
liable to You for damages, including any direct, indirect, special,
|
| 159 |
+
incidental, or consequential damages of any character arising as a
|
| 160 |
+
result of this License or out of the use or inability to use the
|
| 161 |
+
Work (including but not limited to damages for loss of goodwill,
|
| 162 |
+
work stoppage, computer failure or malfunction, or any and all
|
| 163 |
+
other commercial damages or losses), even if such Contributor
|
| 164 |
+
has been advised of the possibility of such damages.
|
| 165 |
+
|
| 166 |
+
9. Accepting Warranty or Additional Liability. While redistributing
|
| 167 |
+
the Work or Derivative Works thereof, You may choose to offer,
|
| 168 |
+
and charge a fee for, acceptance of support, warranty, indemnity,
|
| 169 |
+
or other liability obligations and/or rights consistent with this
|
| 170 |
+
License. However, in accepting such obligations, You may act only
|
| 171 |
+
on Your own behalf and on Your sole responsibility, not on behalf
|
| 172 |
+
of any other Contributor, and only if You agree to indemnify,
|
| 173 |
+
defend, and hold each Contributor harmless for any liability
|
| 174 |
+
incurred by, or claims asserted against, such Contributor by reason
|
| 175 |
+
of your accepting any such warranty or additional liability.
|
| 176 |
+
|
| 177 |
+
END OF TERMS AND CONDITIONS
|
| 178 |
+
|
| 179 |
+
APPENDIX: How to apply the Apache License to your work.
|
| 180 |
+
|
| 181 |
+
To apply the Apache License to your work, attach the following
|
| 182 |
+
boilerplate notice, with the fields enclosed by brackets "[]"
|
| 183 |
+
replaced with your own identifying information. (Don't include
|
| 184 |
+
the brackets!) The text should be enclosed in the appropriate
|
| 185 |
+
comment syntax for the file format. We also recommend that a
|
| 186 |
+
file or class name and description of purpose be included on the
|
| 187 |
+
same "printed page" as the copyright notice for easier
|
| 188 |
+
identification within third-party archives.
|
| 189 |
+
|
| 190 |
+
Copyright 2026 Jared Palmer
|
| 191 |
+
|
| 192 |
+
Licensed under the Apache License, Version 2.0 (the "License");
|
| 193 |
+
you may not use this file except in compliance with the License.
|
| 194 |
+
You may obtain a copy of the License at
|
| 195 |
+
|
| 196 |
+
http://www.apache.org/licenses/LICENSE-2.0
|
| 197 |
+
|
| 198 |
+
Unless required by applicable law or agreed to in writing, software
|
| 199 |
+
distributed under the License is distributed on an "AS IS" BASIS,
|
| 200 |
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 201 |
+
See the License for the specific language governing permissions and
|
| 202 |
+
limitations under the License.
|
NOTICE
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
Kev source and original decision model: Copyright 2026 Jared Palmer. Apache-2.0.
|
| 2 |
+
Base backbone: Qwen/Qwen3.5-4B-Base, Apache-2.0.
|
| 3 |
+
Standalone mixed ternary/4-bit compression export and packaging: Jakevin, 2026-10-06.
|
| 4 |
+
This distribution modifies kev/checkpoint.py and kev/mlx_model.py and adds kev/packed_model.py.
|
| 5 |
+
Upstream source: https://github.com/jaredpalmer/kev
|
| 6 |
+
Upstream Kev checkpoint: jaredpalmer/kev-4b@6cfce5c2fa4b4bd64026336ab649c5ca78857d52
|
| 7 |
+
Base checkpoint: Qwen/Qwen3.5-4B-Base@1001bb4d826a52d1f399e183466143f4da7b741b
|
README.md
ADDED
|
@@ -0,0 +1,176 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: apache-2.0
|
| 3 |
+
base_model: jaredpalmer/kev-4b
|
| 4 |
+
base_model_relation: quantized
|
| 5 |
+
library_name: mlx
|
| 6 |
+
tags:
|
| 7 |
+
- kev
|
| 8 |
+
- decision-model
|
| 9 |
+
- ternary
|
| 10 |
+
- mixed-precision
|
| 11 |
+
- post-training-quantization
|
| 12 |
+
- mlx
|
| 13 |
+
- experimental
|
| 14 |
+
---
|
| 15 |
+
|
| 16 |
+
# Kev-4B Ternary + CLoQ16 for MLX
|
| 17 |
+
|
| 18 |
+
An experimental, standalone compression of **Jared Palmer's Kev-4B decision
|
| 19 |
+
model**, published by Jakevin. The entire distribution is approximately
|
| 20 |
+
**1.71 GB**, including the compressed backbone, compensation adapter, pointer
|
| 21 |
+
head, configuration, tokenizer, and the required Kev runtime source.
|
| 22 |
+
|
| 23 |
+
On **decision-v7 development**, 1,264 scored clean questions, native packed
|
| 24 |
+
inference reaches **84.41% accuracy**, retaining **96.82%** of the archived
|
| 25 |
+
bf16-MLX accuracy. These are development diagnostics for a suite related to
|
| 26 |
+
Kev's training recipe, not a fresh blind test or a general retention guarantee.
|
| 27 |
+
|
| 28 |
+
This is a **decision scorer**, not a text-generating chat model. It implements
|
| 29 |
+
Kev's existing `state` / `questions` -> `answers` / `usage` System One API.
|
| 30 |
+
**Apple Silicon and the included custom MLX runtime are required.** Stock
|
| 31 |
+
Transformers, stock mlx-lm text generation, and llama.cpp are not supported
|
| 32 |
+
loading paths for this artifact.
|
| 33 |
+
|
| 34 |
+
## Download and Run
|
| 35 |
+
|
| 36 |
+
Download this repository, including `runtime/`. It contains the packed-aware
|
| 37 |
+
Kev loader and serving source; the upstream Kev release alone does not include
|
| 38 |
+
this loader. Python 3.12 or 3.13 is required.
|
| 39 |
+
|
| 40 |
+
```sh
|
| 41 |
+
python3 -m pip install huggingface_hub
|
| 42 |
+
python3 -c "from huggingface_hub import snapshot_download; snapshot_download('Jakevin/kev-4b-ternary-mlx', revision='v0.1.0-experimental', local_dir='kev-4b-ternary-mlx')"
|
| 43 |
+
cd kev-4b-ternary-mlx
|
| 44 |
+
python3.13 -m venv .venv
|
| 45 |
+
.venv/bin/python -m pip install './runtime[serve]'
|
| 46 |
+
|
| 47 |
+
HF_HUB_OFFLINE=1 HF_HOME=/tmp/kev-empty-hf-cache KEV_BACKEND=mlx \
|
| 48 |
+
KEV_PREFIX_CACHE=1 KEV_PREFIX_MAX_TOKENS=4096 \
|
| 49 |
+
.venv/bin/python -m kev.serve --run . --host 127.0.0.1 --port 8010
|
| 50 |
+
```
|
| 51 |
+
|
| 52 |
+
```sh
|
| 53 |
+
curl http://127.0.0.1:8010/v1/systemone \
|
| 54 |
+
-H 'Content-Type: application/json' \
|
| 55 |
+
-d '{"state":"I was charged twice for order 8812.","questions":{"team":{"type":"choice","instructions":"Which team should handle this ticket?","criteria":{"billing":"Charges and payments","shipping":"Delivery and delays"}}}}'
|
| 56 |
+
```
|
| 57 |
+
|
| 58 |
+
The library interface is unchanged:
|
| 59 |
+
|
| 60 |
+
```python
|
| 61 |
+
from kev.checkpoint import Checkpoint
|
| 62 |
+
|
| 63 |
+
tokenizer, model = Checkpoint("path/to/downloaded-repository").load("mps")
|
| 64 |
+
```
|
| 65 |
+
|
| 66 |
+
After dependencies are installed, loading and inference work offline with an
|
| 67 |
+
empty Hugging Face cache. They do not load the original backbone or original
|
| 68 |
+
PEFT adapter, reconstruct a dense bf16 backbone, or construct the dense
|
| 69 |
+
compensation matrix. Architecture initialization is lazy and its random dense
|
| 70 |
+
parameters are replaced before evaluation.
|
| 71 |
+
|
| 72 |
+
## Compression Recipe
|
| 73 |
+
|
| 74 |
+
Source checkpoint: [`jaredpalmer/kev-4b`](https://huggingface.co/jaredpalmer/kev-4b),
|
| 75 |
+
revision `6cfce5c2fa4b4bd64026336ab649c5ca78857d52`, based on
|
| 76 |
+
[`Qwen/Qwen3.5-4B-Base`](https://huggingface.co/Qwen/Qwen3.5-4B-Base), revision
|
| 77 |
+
`1001bb4d826a52d1f399e183466143f4da7b741b`.
|
| 78 |
+
|
| 79 |
+
- The source Kev LoRA was merged before the original compression experiment.
|
| 80 |
+
- GPTQ-style error compensation uses 128 records from the decision-v7
|
| 81 |
+
**calibration** partition, RNG 1234, damping 0.01 and sequential 2-layer chunks.
|
| 82 |
+
- Most backbone linear weights are ternary, with an fp16 scale per 64 input columns.
|
| 83 |
+
- Gated DeltaNet `out_proj` and full-attention `q/k/v/o_proj` use 4-bit affine
|
| 84 |
+
weights with group size 64. The embedding is also 4-bit.
|
| 85 |
+
- Small `in_proj_a/b`, norms, conv kernels and DeltaNet parameters retain their
|
| 86 |
+
saved source precision. The pointer head and source temperature are unchanged.
|
| 87 |
+
- CLoQ-style rank-16 compensation uses a closed-form calibration-based solution;
|
| 88 |
+
no gradient retraining or new temperature fitting was performed.
|
| 89 |
+
- No online rotation is required by this recipe.
|
| 90 |
+
|
| 91 |
+
At inference, ternary codes use native **2-bit** MLX affine storage. Native
|
| 92 |
+
storage is larger than the original five-trits-per-byte archival format. The
|
| 93 |
+
CLoQ factors remain bf16 and execute as two separate small matrix products
|
| 94 |
+
with fp32 accumulation. Activations are bf16. Only requested embedding rows
|
| 95 |
+
are dequantized. This is a weight compression method; activations and the
|
| 96 |
+
attention/prefix cache are not claimed to be 2-bit or 4-bit.
|
| 97 |
+
|
| 98 |
+
The method combines existing quantization/compensation ideas in a layer-specific
|
| 99 |
+
recipe; it does not claim a new GPTQ or CLoQ algorithm. Relevant work:
|
| 100 |
+
[GPTQ](https://arxiv.org/abs/2210.17323),
|
| 101 |
+
[CLoQ](https://arxiv.org/abs/2501.18475).
|
| 102 |
+
|
| 103 |
+
## Development Evaluation
|
| 104 |
+
|
| 105 |
+
All 1,204 decision-v7 development records were evaluated: 1,468 total rows,
|
| 106 |
+
1,264 scored clean questions, zero rejected records and zero truncations.
|
| 107 |
+
The source checkpoint's temperature was retained.
|
| 108 |
+
|
| 109 |
+
| Path | Backbone + compensation | Accuracy | Relative accuracy retention | Brier |
|
| 110 |
+
|---|---:|---:|---:|---:|
|
| 111 |
+
| Archived bf16-MLX baseline | 8.412 GB | 0.87184 | 100% | 0.18248 |
|
| 112 |
+
| Archived five-trits-per-byte artifact, decoded/merged to bf16 for scoring | 1.434 GB | 0.84335 | 96.73% | 0.23260 |
|
| 113 |
+
| This native MLX packed inference path | 1.679 GB | **0.84415** | **96.82%** | **0.23251** |
|
| 114 |
+
|
| 115 |
+
The complete published folder is approximately 1.71 GB; sizes above exclude
|
| 116 |
+
tokenizer and pointer head. GB means decimal GB.
|
| 117 |
+
|
| 118 |
+
Native quantized matmul plus separate compensation has different arithmetic
|
| 119 |
+
from the archived merged-bf16 simulation. Across 1,468 rows, 1 argmax changed;
|
| 120 |
+
mean per-row maximum probability difference was 0.00268, maximum 0.06209.
|
| 121 |
+
This is not bit-identical serving. The original bf16 model was not loaded again
|
| 122 |
+
for this comparison; its metrics came from the frozen local experiment.
|
| 123 |
+
|
| 124 |
+
Aggregate evaluation and memory records are in `evaluation/`. No evaluation
|
| 125 |
+
sentences, labels or per-example prediction dumps are uploaded. No NAER /
|
| 126 |
+
Chinese Linguipedia data, derivatives or predictions were used in this export
|
| 127 |
+
or included in this repository.
|
| 128 |
+
|
| 129 |
+
## Memory Measurements
|
| 130 |
+
|
| 131 |
+
Fresh-process probes on an **Apple M4 Pro, 64 GiB RAM**, macOS 26.6.2:
|
| 132 |
+
MLX 0.32.2, mlx-lm 0.31.3, torch 2.8.0, transformers 5.17.0.
|
| 133 |
+
Original loaders and downloads were explicitly blocked during the probes.
|
| 134 |
+
|
| 135 |
+
| Measurement | Short request | Longer request |
|
| 136 |
+
|---|---:|---:|
|
| 137 |
+
| State / total input tokens | 19 / 42 | 3,519 / 3,542 |
|
| 138 |
+
| Load time | 0.57 s | 0.57 s |
|
| 139 |
+
| MLX active weights after load | 1.679 GB | 1.679 GB |
|
| 140 |
+
| Process physical-footprint peak while loading | 2.396 GB | 2.398 GB |
|
| 141 |
+
| MLX inference peak | 2.013 GB | 3.066 GB |
|
| 142 |
+
| Process physical-footprint inference peak | **2.846 GB** | **3.862 GB** |
|
| 143 |
+
| First inference time | 0.155 s | 6.599 s |
|
| 144 |
+
|
| 145 |
+
Physical footprint includes Metal buffers; RSS alone undercounts them. These
|
| 146 |
+
are single-request probes, not minimum-RAM device certification. **No 8-GB
|
| 147 |
+
hardware was tested.** More state tokens, questions, or retained prefixes can
|
| 148 |
+
increase memory. The synthetic longer probe measures memory/latency, not
|
| 149 |
+
long-context accuracy. `KEV_PREFIX_MAX_TOKENS` limits retained prefix cache,
|
| 150 |
+
not input admission. `mx.set_memory_limit(3 GB)` is a guideline rather than a
|
| 151 |
+
hard RAM cap; the longer probe exceeded it.
|
| 152 |
+
|
| 153 |
+
## Reproducibility and Limitations
|
| 154 |
+
|
| 155 |
+
- The runtime and the native exporter are in `runtime/`. Source and model
|
| 156 |
+
hashes are recorded in `packed.json`, `head.pt["compression"]` and
|
| 157 |
+
`release-manifest.json`.
|
| 158 |
+
- The test source is included. The local release check passed 7 relevant
|
| 159 |
+
tests, including native code/adapter arithmetic, loading with base loaders
|
| 160 |
+
blocked, missing-weight rejection, and the System One API.
|
| 161 |
+
- Decision-v7 development was used to select compression recipes. No new
|
| 162 |
+
blind benchmark, transfer-v9, hard-v1 or breadth-v1 evaluation is claimed.
|
| 163 |
+
- Full 3-bit and 4-bit same-condition baselines were not evaluated here.
|
| 164 |
+
- The source temperature is retained; Brier/ECE worsened relative to the
|
| 165 |
+
bf16 baseline. Confidence-sensitive users should evaluate their own task.
|
| 166 |
+
- Native inference has not been validated on non-Apple devices or llama.cpp.
|
| 167 |
+
- This is an experimental derivative of Kev-4B, not a new independently
|
| 168 |
+
pretrained foundation model and not the Taiwan/mainland wording model.
|
| 169 |
+
|
| 170 |
+
## Attribution and License
|
| 171 |
+
|
| 172 |
+
Apache-2.0. The original Kev decision model and source are by **Jared Palmer**;
|
| 173 |
+
the base backbone is from **Qwen**. This compression experiment, standalone
|
| 174 |
+
export and publication are by **Jakevin**. The runtime distribution modifies
|
| 175 |
+
Kev's checkpoint and MLX loaders and adds the packed model implementation.
|
| 176 |
+
See `LICENSE` and `NOTICE` for attribution.
|
added_tokens.json
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"</think>": 151668,
|
| 3 |
+
"</tool_call>": 151658,
|
| 4 |
+
"</tool_response>": 151666,
|
| 5 |
+
"<think>": 151667,
|
| 6 |
+
"<tool_call>": 151657,
|
| 7 |
+
"<tool_response>": 151665,
|
| 8 |
+
"<|box_end|>": 151649,
|
| 9 |
+
"<|box_start|>": 151648,
|
| 10 |
+
"<|endoftext|>": 151643,
|
| 11 |
+
"<|file_sep|>": 151664,
|
| 12 |
+
"<|fim_middle|>": 151660,
|
| 13 |
+
"<|fim_pad|>": 151662,
|
| 14 |
+
"<|fim_prefix|>": 151659,
|
| 15 |
+
"<|fim_suffix|>": 151661,
|
| 16 |
+
"<|im_end|>": 151645,
|
| 17 |
+
"<|im_start|>": 151644,
|
| 18 |
+
"<|image_pad|>": 151655,
|
| 19 |
+
"<|object_ref_end|>": 151647,
|
| 20 |
+
"<|object_ref_start|>": 151646,
|
| 21 |
+
"<|quad_end|>": 151651,
|
| 22 |
+
"<|quad_start|>": 151650,
|
| 23 |
+
"<|repo_name|>": 151663,
|
| 24 |
+
"<|video_pad|>": 151656,
|
| 25 |
+
"<|vision_end|>": 151653,
|
| 26 |
+
"<|vision_pad|>": 151654,
|
| 27 |
+
"<|vision_start|>": 151652
|
| 28 |
+
}
|
config.json
ADDED
|
@@ -0,0 +1,104 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"Qwen3_5ForConditionalGeneration"
|
| 4 |
+
],
|
| 5 |
+
"image_token_id": 248056,
|
| 6 |
+
"model_type": "qwen3_5",
|
| 7 |
+
"text_config": {
|
| 8 |
+
"attention_bias": false,
|
| 9 |
+
"attention_dropout": 0.0,
|
| 10 |
+
"attn_output_gate": true,
|
| 11 |
+
"dtype": "bfloat16",
|
| 12 |
+
"eos_token_id": 248044,
|
| 13 |
+
"full_attention_interval": 4,
|
| 14 |
+
"head_dim": 256,
|
| 15 |
+
"hidden_act": "silu",
|
| 16 |
+
"hidden_size": 2560,
|
| 17 |
+
"initializer_range": 0.02,
|
| 18 |
+
"intermediate_size": 9216,
|
| 19 |
+
"layer_types": [
|
| 20 |
+
"linear_attention",
|
| 21 |
+
"linear_attention",
|
| 22 |
+
"linear_attention",
|
| 23 |
+
"full_attention",
|
| 24 |
+
"linear_attention",
|
| 25 |
+
"linear_attention",
|
| 26 |
+
"linear_attention",
|
| 27 |
+
"full_attention",
|
| 28 |
+
"linear_attention",
|
| 29 |
+
"linear_attention",
|
| 30 |
+
"linear_attention",
|
| 31 |
+
"full_attention",
|
| 32 |
+
"linear_attention",
|
| 33 |
+
"linear_attention",
|
| 34 |
+
"linear_attention",
|
| 35 |
+
"full_attention",
|
| 36 |
+
"linear_attention",
|
| 37 |
+
"linear_attention",
|
| 38 |
+
"linear_attention",
|
| 39 |
+
"full_attention",
|
| 40 |
+
"linear_attention",
|
| 41 |
+
"linear_attention",
|
| 42 |
+
"linear_attention",
|
| 43 |
+
"full_attention",
|
| 44 |
+
"linear_attention",
|
| 45 |
+
"linear_attention",
|
| 46 |
+
"linear_attention",
|
| 47 |
+
"full_attention",
|
| 48 |
+
"linear_attention",
|
| 49 |
+
"linear_attention",
|
| 50 |
+
"linear_attention",
|
| 51 |
+
"full_attention"
|
| 52 |
+
],
|
| 53 |
+
"linear_conv_kernel_dim": 4,
|
| 54 |
+
"linear_key_head_dim": 128,
|
| 55 |
+
"linear_num_key_heads": 16,
|
| 56 |
+
"linear_num_value_heads": 32,
|
| 57 |
+
"linear_value_head_dim": 128,
|
| 58 |
+
"max_position_embeddings": 262144,
|
| 59 |
+
"mlp_only_layers": [],
|
| 60 |
+
"model_type": "qwen3_5_text",
|
| 61 |
+
"mtp_num_hidden_layers": 1,
|
| 62 |
+
"mtp_use_dedicated_embeddings": false,
|
| 63 |
+
"num_attention_heads": 16,
|
| 64 |
+
"num_hidden_layers": 32,
|
| 65 |
+
"num_key_value_heads": 4,
|
| 66 |
+
"rms_norm_eps": 1e-06,
|
| 67 |
+
"tie_word_embeddings": true,
|
| 68 |
+
"use_cache": true,
|
| 69 |
+
"vocab_size": 248320,
|
| 70 |
+
"mamba_ssm_dtype": "float32",
|
| 71 |
+
"rope_parameters": {
|
| 72 |
+
"mrope_interleaved": true,
|
| 73 |
+
"mrope_section": [
|
| 74 |
+
11,
|
| 75 |
+
11,
|
| 76 |
+
10
|
| 77 |
+
],
|
| 78 |
+
"rope_type": "default",
|
| 79 |
+
"rope_theta": 10000000,
|
| 80 |
+
"partial_rotary_factor": 0.25
|
| 81 |
+
}
|
| 82 |
+
},
|
| 83 |
+
"tie_word_embeddings": true,
|
| 84 |
+
"transformers_version": "4.57.0.dev0",
|
| 85 |
+
"video_token_id": 248057,
|
| 86 |
+
"vision_config": {
|
| 87 |
+
"deepstack_visual_indexes": [],
|
| 88 |
+
"depth": 24,
|
| 89 |
+
"hidden_act": "gelu_pytorch_tanh",
|
| 90 |
+
"hidden_size": 1024,
|
| 91 |
+
"in_channels": 3,
|
| 92 |
+
"initializer_range": 0.02,
|
| 93 |
+
"intermediate_size": 4096,
|
| 94 |
+
"model_type": "qwen3_5",
|
| 95 |
+
"num_heads": 16,
|
| 96 |
+
"num_position_embeddings": 2304,
|
| 97 |
+
"out_hidden_size": 2560,
|
| 98 |
+
"patch_size": 16,
|
| 99 |
+
"spatial_merge_size": 2,
|
| 100 |
+
"temporal_patch_size": 2
|
| 101 |
+
},
|
| 102 |
+
"vision_end_token_id": 248054,
|
| 103 |
+
"vision_start_token_id": 248053
|
| 104 |
+
}
|
evaluation/development-summary.json
ADDED
|
@@ -0,0 +1,91 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"clean": {
|
| 3 |
+
"n": 1264,
|
| 4 |
+
"nll": 0.43551092937105246,
|
| 5 |
+
"acc": 0.8441455696202531,
|
| 6 |
+
"ece": 0.043164737775791336,
|
| 7 |
+
"brier": 0.23251480498478275,
|
| 8 |
+
"mean_conf": 0.8088378144123675,
|
| 9 |
+
"confident_error_rate": 0.010284810126582278,
|
| 10 |
+
"coverage_at_0_9": 0.4849683544303797,
|
| 11 |
+
"accuracy_at_0_9": 0.9787928221859706,
|
| 12 |
+
"coverage_at_5pct_error": 0.6724683544303798,
|
| 13 |
+
"coverage_at_1pct_error": 0.16772151898734178,
|
| 14 |
+
"aurc": 0.046630371957041034,
|
| 15 |
+
"error_rate_at_0_9": 0.021207177814029365,
|
| 16 |
+
"confidence_bias": -0.03530775520788565,
|
| 17 |
+
"top_bins": {
|
| 18 |
+
"0.9": {
|
| 19 |
+
"n": 613,
|
| 20 |
+
"errors": 13,
|
| 21 |
+
"error_rate": 0.021207177814029365
|
| 22 |
+
},
|
| 23 |
+
"0.95": {
|
| 24 |
+
"n": 443,
|
| 25 |
+
"errors": 6,
|
| 26 |
+
"error_rate": 0.013544018058690745
|
| 27 |
+
},
|
| 28 |
+
"0.99": {
|
| 29 |
+
"n": 152,
|
| 30 |
+
"errors": 2,
|
| 31 |
+
"error_rate": 0.013157894736842105
|
| 32 |
+
}
|
| 33 |
+
},
|
| 34 |
+
"selective": {
|
| 35 |
+
"0.5": {
|
| 36 |
+
"coverage": 0.5,
|
| 37 |
+
"accuracy": 0.9778481012658228,
|
| 38 |
+
"confidence_cutoff": 0.8935324549674988
|
| 39 |
+
},
|
| 40 |
+
"0.8": {
|
| 41 |
+
"coverage": 0.8006329113924051,
|
| 42 |
+
"accuracy": 0.9061264822134387,
|
| 43 |
+
"confidence_cutoff": 0.5975180090213416
|
| 44 |
+
}
|
| 45 |
+
},
|
| 46 |
+
"score_mae": 0.5421033783798573,
|
| 47 |
+
"ranked_probability_score": 0.0842686503245251
|
| 48 |
+
},
|
| 49 |
+
"coverage": {
|
| 50 |
+
"requested_records": 1204,
|
| 51 |
+
"requested_questions": 1468,
|
| 52 |
+
"evaluated_records": 1204,
|
| 53 |
+
"evaluated_questions": 1468,
|
| 54 |
+
"rejected_records": 0,
|
| 55 |
+
"truncated_records": 0
|
| 56 |
+
},
|
| 57 |
+
"latency_ms": {
|
| 58 |
+
"median": 304.7052710317075,
|
| 59 |
+
"p95": 1490.851858188398
|
| 60 |
+
},
|
| 61 |
+
"calibration": {
|
| 62 |
+
"inference_temperature": 2.406050072164233,
|
| 63 |
+
"additional_temperature": 1.0,
|
| 64 |
+
"logits_recorded": true
|
| 65 |
+
},
|
| 66 |
+
"environment": {
|
| 67 |
+
"packages": {
|
| 68 |
+
"torch": "2.8.0",
|
| 69 |
+
"transformers": "5.17.0",
|
| 70 |
+
"peft": "0.21.0",
|
| 71 |
+
"flash-linear-attention": null,
|
| 72 |
+
"triton": null,
|
| 73 |
+
"causal-conv1d": null,
|
| 74 |
+
"mlx": "0.32.2",
|
| 75 |
+
"mlx-lm": "0.31.3"
|
| 76 |
+
},
|
| 77 |
+
"device": "mps",
|
| 78 |
+
"gpu": null,
|
| 79 |
+
"backend": "mlx",
|
| 80 |
+
"dtype": "bfloat16",
|
| 81 |
+
"triton_f32_default": null,
|
| 82 |
+
"attention": null,
|
| 83 |
+
"deltanet": null
|
| 84 |
+
},
|
| 85 |
+
"suite": "decision-v7",
|
| 86 |
+
"partition": "development",
|
| 87 |
+
"role": "in-distribution development diagnostic, not a blind test",
|
| 88 |
+
"source_baseline": "archived bf16-MLX metrics; original weights not loaded for this comparison",
|
| 89 |
+
"baseline_accuracy": 0.8718354430379747,
|
| 90 |
+
"accuracy_retention": 0.9682395644283122
|
| 91 |
+
}
|
evaluation/memory-long.json
ADDED
|
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"backend": "mlx",
|
| 3 |
+
"dtype": "bfloat16",
|
| 4 |
+
"platform": "macOS-26.6.2-arm64-arm-64bit-Mach-O",
|
| 5 |
+
"original_loaders_blocked": true,
|
| 6 |
+
"mlx_limit_bytes": 3000000000,
|
| 7 |
+
"load_seconds": 0.5740630407817662,
|
| 8 |
+
"mlx_load_peak_bytes": 1678767640,
|
| 9 |
+
"mlx_after_load_bytes": 1678767624,
|
| 10 |
+
"load_footprint_peak_bytes": 2397899416,
|
| 11 |
+
"load_rss_peak_bytes": 2381381632,
|
| 12 |
+
"state_tokens": 3519,
|
| 13 |
+
"input_tokens": 3542,
|
| 14 |
+
"inference_seconds": 6.599218292161822,
|
| 15 |
+
"probabilities": {
|
| 16 |
+
"billing": 0.941597580909729,
|
| 17 |
+
"shipping": 0.05840245634317398
|
| 18 |
+
},
|
| 19 |
+
"mlx_inference_peak_bytes": 3065702282,
|
| 20 |
+
"inference_footprint_peak_bytes": 3862170744,
|
| 21 |
+
"inference_rss_peak_bytes": 2520678400,
|
| 22 |
+
"mlx_limit_is_hard_cap": false,
|
| 23 |
+
"machine": "Apple M4 Pro",
|
| 24 |
+
"machine_ram_bytes": 68719476736,
|
| 25 |
+
"versions": {
|
| 26 |
+
"mlx": "0.32.2",
|
| 27 |
+
"mlx-lm": "0.31.3",
|
| 28 |
+
"torch": "2.8.0",
|
| 29 |
+
"transformers": "5.17.0"
|
| 30 |
+
}
|
| 31 |
+
}
|
evaluation/memory-short.json
ADDED
|
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"backend": "mlx",
|
| 3 |
+
"dtype": "bfloat16",
|
| 4 |
+
"platform": "macOS-26.6.2-arm64-arm-64bit-Mach-O",
|
| 5 |
+
"original_loaders_blocked": true,
|
| 6 |
+
"mlx_limit_bytes": 3000000000,
|
| 7 |
+
"load_seconds": 0.5707792090252042,
|
| 8 |
+
"mlx_load_peak_bytes": 1678767640,
|
| 9 |
+
"mlx_after_load_bytes": 1678767624,
|
| 10 |
+
"load_footprint_peak_bytes": 2396031616,
|
| 11 |
+
"load_rss_peak_bytes": 2307145728,
|
| 12 |
+
"state_tokens": 19,
|
| 13 |
+
"input_tokens": 42,
|
| 14 |
+
"inference_seconds": 0.15510379197075963,
|
| 15 |
+
"probabilities": {
|
| 16 |
+
"billing": 0.9229481220245361,
|
| 17 |
+
"shipping": 0.07705190777778625
|
| 18 |
+
},
|
| 19 |
+
"mlx_inference_peak_bytes": 2012569374,
|
| 20 |
+
"inference_footprint_peak_bytes": 2846362528,
|
| 21 |
+
"inference_rss_peak_bytes": 2516697088,
|
| 22 |
+
"mlx_limit_is_hard_cap": false,
|
| 23 |
+
"machine": "Apple M4 Pro",
|
| 24 |
+
"machine_ram_bytes": 68719476736,
|
| 25 |
+
"versions": {
|
| 26 |
+
"mlx": "0.32.2",
|
| 27 |
+
"mlx-lm": "0.31.3",
|
| 28 |
+
"torch": "2.8.0",
|
| 29 |
+
"transformers": "5.17.0"
|
| 30 |
+
}
|
| 31 |
+
}
|
evaluation/native-comparison.json
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"rows": 1468,
|
| 3 |
+
"max_abs_dp": 0.06208749928657609,
|
| 4 |
+
"mean_max_abs_dp": 0.0026799194979893045,
|
| 5 |
+
"argmax_flips": 1,
|
| 6 |
+
"accuracy": 0.8441455696202531,
|
| 7 |
+
"retention": 0.9682395644283122,
|
| 8 |
+
"reference": "archival packed merged-bf16 evaluation; no original weights loaded"
|
| 9 |
+
}
|
head.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:6f6740d94a49187770b68c0f058d85bedc0cf7efc20dff7ca9a418c8b461dc0e
|
| 3 |
+
size 5247807
|
merges.txt
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
model.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:c4899eaa9fb206e77260f832d00188c9b90750724422ea0eea43ed69faeff43d
|
| 3 |
+
size 1678909327
|
packed.json
ADDED
|
@@ -0,0 +1,1016 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"format": "kev-packed-mlx-v1",
|
| 3 |
+
"linears": {
|
| 4 |
+
"language_model.model.layers.0.mlp.up_proj": {
|
| 5 |
+
"bits": 2,
|
| 6 |
+
"group_size": 64,
|
| 7 |
+
"rank": 16
|
| 8 |
+
},
|
| 9 |
+
"language_model.model.layers.0.mlp.down_proj": {
|
| 10 |
+
"bits": 2,
|
| 11 |
+
"group_size": 64,
|
| 12 |
+
"rank": 16
|
| 13 |
+
},
|
| 14 |
+
"language_model.model.layers.0.mlp.gate_proj": {
|
| 15 |
+
"bits": 2,
|
| 16 |
+
"group_size": 64,
|
| 17 |
+
"rank": 16
|
| 18 |
+
},
|
| 19 |
+
"language_model.model.layers.0.linear_attn.out_proj": {
|
| 20 |
+
"bits": 4,
|
| 21 |
+
"group_size": 64,
|
| 22 |
+
"rank": 16
|
| 23 |
+
},
|
| 24 |
+
"language_model.model.layers.0.linear_attn.in_proj_z": {
|
| 25 |
+
"bits": 2,
|
| 26 |
+
"group_size": 64,
|
| 27 |
+
"rank": 16
|
| 28 |
+
},
|
| 29 |
+
"language_model.model.layers.0.linear_attn.in_proj_qkv": {
|
| 30 |
+
"bits": 2,
|
| 31 |
+
"group_size": 64,
|
| 32 |
+
"rank": 16
|
| 33 |
+
},
|
| 34 |
+
"language_model.model.layers.1.mlp.up_proj": {
|
| 35 |
+
"bits": 2,
|
| 36 |
+
"group_size": 64,
|
| 37 |
+
"rank": 16
|
| 38 |
+
},
|
| 39 |
+
"language_model.model.layers.1.mlp.down_proj": {
|
| 40 |
+
"bits": 2,
|
| 41 |
+
"group_size": 64,
|
| 42 |
+
"rank": 16
|
| 43 |
+
},
|
| 44 |
+
"language_model.model.layers.1.mlp.gate_proj": {
|
| 45 |
+
"bits": 2,
|
| 46 |
+
"group_size": 64,
|
| 47 |
+
"rank": 16
|
| 48 |
+
},
|
| 49 |
+
"language_model.model.layers.1.linear_attn.out_proj": {
|
| 50 |
+
"bits": 4,
|
| 51 |
+
"group_size": 64,
|
| 52 |
+
"rank": 16
|
| 53 |
+
},
|
| 54 |
+
"language_model.model.layers.1.linear_attn.in_proj_z": {
|
| 55 |
+
"bits": 2,
|
| 56 |
+
"group_size": 64,
|
| 57 |
+
"rank": 16
|
| 58 |
+
},
|
| 59 |
+
"language_model.model.layers.1.linear_attn.in_proj_qkv": {
|
| 60 |
+
"bits": 2,
|
| 61 |
+
"group_size": 64,
|
| 62 |
+
"rank": 16
|
| 63 |
+
},
|
| 64 |
+
"language_model.model.layers.2.mlp.up_proj": {
|
| 65 |
+
"bits": 2,
|
| 66 |
+
"group_size": 64,
|
| 67 |
+
"rank": 16
|
| 68 |
+
},
|
| 69 |
+
"language_model.model.layers.2.mlp.down_proj": {
|
| 70 |
+
"bits": 2,
|
| 71 |
+
"group_size": 64,
|
| 72 |
+
"rank": 16
|
| 73 |
+
},
|
| 74 |
+
"language_model.model.layers.2.mlp.gate_proj": {
|
| 75 |
+
"bits": 2,
|
| 76 |
+
"group_size": 64,
|
| 77 |
+
"rank": 16
|
| 78 |
+
},
|
| 79 |
+
"language_model.model.layers.2.linear_attn.out_proj": {
|
| 80 |
+
"bits": 4,
|
| 81 |
+
"group_size": 64,
|
| 82 |
+
"rank": 16
|
| 83 |
+
},
|
| 84 |
+
"language_model.model.layers.2.linear_attn.in_proj_z": {
|
| 85 |
+
"bits": 2,
|
| 86 |
+
"group_size": 64,
|
| 87 |
+
"rank": 16
|
| 88 |
+
},
|
| 89 |
+
"language_model.model.layers.2.linear_attn.in_proj_qkv": {
|
| 90 |
+
"bits": 2,
|
| 91 |
+
"group_size": 64,
|
| 92 |
+
"rank": 16
|
| 93 |
+
},
|
| 94 |
+
"language_model.model.layers.3.mlp.up_proj": {
|
| 95 |
+
"bits": 2,
|
| 96 |
+
"group_size": 64,
|
| 97 |
+
"rank": 16
|
| 98 |
+
},
|
| 99 |
+
"language_model.model.layers.3.mlp.down_proj": {
|
| 100 |
+
"bits": 2,
|
| 101 |
+
"group_size": 64,
|
| 102 |
+
"rank": 16
|
| 103 |
+
},
|
| 104 |
+
"language_model.model.layers.3.mlp.gate_proj": {
|
| 105 |
+
"bits": 2,
|
| 106 |
+
"group_size": 64,
|
| 107 |
+
"rank": 16
|
| 108 |
+
},
|
| 109 |
+
"language_model.model.layers.3.self_attn.o_proj": {
|
| 110 |
+
"bits": 4,
|
| 111 |
+
"group_size": 64,
|
| 112 |
+
"rank": 16
|
| 113 |
+
},
|
| 114 |
+
"language_model.model.layers.3.self_attn.v_proj": {
|
| 115 |
+
"bits": 4,
|
| 116 |
+
"group_size": 64,
|
| 117 |
+
"rank": 16
|
| 118 |
+
},
|
| 119 |
+
"language_model.model.layers.3.self_attn.k_proj": {
|
| 120 |
+
"bits": 4,
|
| 121 |
+
"group_size": 64,
|
| 122 |
+
"rank": 16
|
| 123 |
+
},
|
| 124 |
+
"language_model.model.layers.3.self_attn.q_proj": {
|
| 125 |
+
"bits": 4,
|
| 126 |
+
"group_size": 64,
|
| 127 |
+
"rank": 16
|
| 128 |
+
},
|
| 129 |
+
"language_model.model.layers.4.mlp.up_proj": {
|
| 130 |
+
"bits": 2,
|
| 131 |
+
"group_size": 64,
|
| 132 |
+
"rank": 16
|
| 133 |
+
},
|
| 134 |
+
"language_model.model.layers.4.mlp.down_proj": {
|
| 135 |
+
"bits": 2,
|
| 136 |
+
"group_size": 64,
|
| 137 |
+
"rank": 16
|
| 138 |
+
},
|
| 139 |
+
"language_model.model.layers.4.mlp.gate_proj": {
|
| 140 |
+
"bits": 2,
|
| 141 |
+
"group_size": 64,
|
| 142 |
+
"rank": 16
|
| 143 |
+
},
|
| 144 |
+
"language_model.model.layers.4.linear_attn.out_proj": {
|
| 145 |
+
"bits": 4,
|
| 146 |
+
"group_size": 64,
|
| 147 |
+
"rank": 16
|
| 148 |
+
},
|
| 149 |
+
"language_model.model.layers.4.linear_attn.in_proj_z": {
|
| 150 |
+
"bits": 2,
|
| 151 |
+
"group_size": 64,
|
| 152 |
+
"rank": 16
|
| 153 |
+
},
|
| 154 |
+
"language_model.model.layers.4.linear_attn.in_proj_qkv": {
|
| 155 |
+
"bits": 2,
|
| 156 |
+
"group_size": 64,
|
| 157 |
+
"rank": 16
|
| 158 |
+
},
|
| 159 |
+
"language_model.model.layers.5.mlp.up_proj": {
|
| 160 |
+
"bits": 2,
|
| 161 |
+
"group_size": 64,
|
| 162 |
+
"rank": 16
|
| 163 |
+
},
|
| 164 |
+
"language_model.model.layers.5.mlp.down_proj": {
|
| 165 |
+
"bits": 2,
|
| 166 |
+
"group_size": 64,
|
| 167 |
+
"rank": 16
|
| 168 |
+
},
|
| 169 |
+
"language_model.model.layers.5.mlp.gate_proj": {
|
| 170 |
+
"bits": 2,
|
| 171 |
+
"group_size": 64,
|
| 172 |
+
"rank": 16
|
| 173 |
+
},
|
| 174 |
+
"language_model.model.layers.5.linear_attn.out_proj": {
|
| 175 |
+
"bits": 4,
|
| 176 |
+
"group_size": 64,
|
| 177 |
+
"rank": 16
|
| 178 |
+
},
|
| 179 |
+
"language_model.model.layers.5.linear_attn.in_proj_z": {
|
| 180 |
+
"bits": 2,
|
| 181 |
+
"group_size": 64,
|
| 182 |
+
"rank": 16
|
| 183 |
+
},
|
| 184 |
+
"language_model.model.layers.5.linear_attn.in_proj_qkv": {
|
| 185 |
+
"bits": 2,
|
| 186 |
+
"group_size": 64,
|
| 187 |
+
"rank": 16
|
| 188 |
+
},
|
| 189 |
+
"language_model.model.layers.6.mlp.up_proj": {
|
| 190 |
+
"bits": 2,
|
| 191 |
+
"group_size": 64,
|
| 192 |
+
"rank": 16
|
| 193 |
+
},
|
| 194 |
+
"language_model.model.layers.6.mlp.down_proj": {
|
| 195 |
+
"bits": 2,
|
| 196 |
+
"group_size": 64,
|
| 197 |
+
"rank": 16
|
| 198 |
+
},
|
| 199 |
+
"language_model.model.layers.6.mlp.gate_proj": {
|
| 200 |
+
"bits": 2,
|
| 201 |
+
"group_size": 64,
|
| 202 |
+
"rank": 16
|
| 203 |
+
},
|
| 204 |
+
"language_model.model.layers.6.linear_attn.out_proj": {
|
| 205 |
+
"bits": 4,
|
| 206 |
+
"group_size": 64,
|
| 207 |
+
"rank": 16
|
| 208 |
+
},
|
| 209 |
+
"language_model.model.layers.6.linear_attn.in_proj_z": {
|
| 210 |
+
"bits": 2,
|
| 211 |
+
"group_size": 64,
|
| 212 |
+
"rank": 16
|
| 213 |
+
},
|
| 214 |
+
"language_model.model.layers.6.linear_attn.in_proj_qkv": {
|
| 215 |
+
"bits": 2,
|
| 216 |
+
"group_size": 64,
|
| 217 |
+
"rank": 16
|
| 218 |
+
},
|
| 219 |
+
"language_model.model.layers.7.mlp.up_proj": {
|
| 220 |
+
"bits": 2,
|
| 221 |
+
"group_size": 64,
|
| 222 |
+
"rank": 16
|
| 223 |
+
},
|
| 224 |
+
"language_model.model.layers.7.mlp.down_proj": {
|
| 225 |
+
"bits": 2,
|
| 226 |
+
"group_size": 64,
|
| 227 |
+
"rank": 16
|
| 228 |
+
},
|
| 229 |
+
"language_model.model.layers.7.mlp.gate_proj": {
|
| 230 |
+
"bits": 2,
|
| 231 |
+
"group_size": 64,
|
| 232 |
+
"rank": 16
|
| 233 |
+
},
|
| 234 |
+
"language_model.model.layers.7.self_attn.o_proj": {
|
| 235 |
+
"bits": 4,
|
| 236 |
+
"group_size": 64,
|
| 237 |
+
"rank": 16
|
| 238 |
+
},
|
| 239 |
+
"language_model.model.layers.7.self_attn.v_proj": {
|
| 240 |
+
"bits": 4,
|
| 241 |
+
"group_size": 64,
|
| 242 |
+
"rank": 16
|
| 243 |
+
},
|
| 244 |
+
"language_model.model.layers.7.self_attn.k_proj": {
|
| 245 |
+
"bits": 4,
|
| 246 |
+
"group_size": 64,
|
| 247 |
+
"rank": 16
|
| 248 |
+
},
|
| 249 |
+
"language_model.model.layers.7.self_attn.q_proj": {
|
| 250 |
+
"bits": 4,
|
| 251 |
+
"group_size": 64,
|
| 252 |
+
"rank": 16
|
| 253 |
+
},
|
| 254 |
+
"language_model.model.layers.8.mlp.up_proj": {
|
| 255 |
+
"bits": 2,
|
| 256 |
+
"group_size": 64,
|
| 257 |
+
"rank": 16
|
| 258 |
+
},
|
| 259 |
+
"language_model.model.layers.8.mlp.down_proj": {
|
| 260 |
+
"bits": 2,
|
| 261 |
+
"group_size": 64,
|
| 262 |
+
"rank": 16
|
| 263 |
+
},
|
| 264 |
+
"language_model.model.layers.8.mlp.gate_proj": {
|
| 265 |
+
"bits": 2,
|
| 266 |
+
"group_size": 64,
|
| 267 |
+
"rank": 16
|
| 268 |
+
},
|
| 269 |
+
"language_model.model.layers.8.linear_attn.out_proj": {
|
| 270 |
+
"bits": 4,
|
| 271 |
+
"group_size": 64,
|
| 272 |
+
"rank": 16
|
| 273 |
+
},
|
| 274 |
+
"language_model.model.layers.8.linear_attn.in_proj_z": {
|
| 275 |
+
"bits": 2,
|
| 276 |
+
"group_size": 64,
|
| 277 |
+
"rank": 16
|
| 278 |
+
},
|
| 279 |
+
"language_model.model.layers.8.linear_attn.in_proj_qkv": {
|
| 280 |
+
"bits": 2,
|
| 281 |
+
"group_size": 64,
|
| 282 |
+
"rank": 16
|
| 283 |
+
},
|
| 284 |
+
"language_model.model.layers.9.mlp.up_proj": {
|
| 285 |
+
"bits": 2,
|
| 286 |
+
"group_size": 64,
|
| 287 |
+
"rank": 16
|
| 288 |
+
},
|
| 289 |
+
"language_model.model.layers.9.mlp.down_proj": {
|
| 290 |
+
"bits": 2,
|
| 291 |
+
"group_size": 64,
|
| 292 |
+
"rank": 16
|
| 293 |
+
},
|
| 294 |
+
"language_model.model.layers.9.mlp.gate_proj": {
|
| 295 |
+
"bits": 2,
|
| 296 |
+
"group_size": 64,
|
| 297 |
+
"rank": 16
|
| 298 |
+
},
|
| 299 |
+
"language_model.model.layers.9.linear_attn.out_proj": {
|
| 300 |
+
"bits": 4,
|
| 301 |
+
"group_size": 64,
|
| 302 |
+
"rank": 16
|
| 303 |
+
},
|
| 304 |
+
"language_model.model.layers.9.linear_attn.in_proj_z": {
|
| 305 |
+
"bits": 2,
|
| 306 |
+
"group_size": 64,
|
| 307 |
+
"rank": 16
|
| 308 |
+
},
|
| 309 |
+
"language_model.model.layers.9.linear_attn.in_proj_qkv": {
|
| 310 |
+
"bits": 2,
|
| 311 |
+
"group_size": 64,
|
| 312 |
+
"rank": 16
|
| 313 |
+
},
|
| 314 |
+
"language_model.model.layers.10.mlp.up_proj": {
|
| 315 |
+
"bits": 2,
|
| 316 |
+
"group_size": 64,
|
| 317 |
+
"rank": 16
|
| 318 |
+
},
|
| 319 |
+
"language_model.model.layers.10.mlp.down_proj": {
|
| 320 |
+
"bits": 2,
|
| 321 |
+
"group_size": 64,
|
| 322 |
+
"rank": 16
|
| 323 |
+
},
|
| 324 |
+
"language_model.model.layers.10.mlp.gate_proj": {
|
| 325 |
+
"bits": 2,
|
| 326 |
+
"group_size": 64,
|
| 327 |
+
"rank": 16
|
| 328 |
+
},
|
| 329 |
+
"language_model.model.layers.10.linear_attn.out_proj": {
|
| 330 |
+
"bits": 4,
|
| 331 |
+
"group_size": 64,
|
| 332 |
+
"rank": 16
|
| 333 |
+
},
|
| 334 |
+
"language_model.model.layers.10.linear_attn.in_proj_z": {
|
| 335 |
+
"bits": 2,
|
| 336 |
+
"group_size": 64,
|
| 337 |
+
"rank": 16
|
| 338 |
+
},
|
| 339 |
+
"language_model.model.layers.10.linear_attn.in_proj_qkv": {
|
| 340 |
+
"bits": 2,
|
| 341 |
+
"group_size": 64,
|
| 342 |
+
"rank": 16
|
| 343 |
+
},
|
| 344 |
+
"language_model.model.layers.11.mlp.up_proj": {
|
| 345 |
+
"bits": 2,
|
| 346 |
+
"group_size": 64,
|
| 347 |
+
"rank": 16
|
| 348 |
+
},
|
| 349 |
+
"language_model.model.layers.11.mlp.down_proj": {
|
| 350 |
+
"bits": 2,
|
| 351 |
+
"group_size": 64,
|
| 352 |
+
"rank": 16
|
| 353 |
+
},
|
| 354 |
+
"language_model.model.layers.11.mlp.gate_proj": {
|
| 355 |
+
"bits": 2,
|
| 356 |
+
"group_size": 64,
|
| 357 |
+
"rank": 16
|
| 358 |
+
},
|
| 359 |
+
"language_model.model.layers.11.self_attn.o_proj": {
|
| 360 |
+
"bits": 4,
|
| 361 |
+
"group_size": 64,
|
| 362 |
+
"rank": 16
|
| 363 |
+
},
|
| 364 |
+
"language_model.model.layers.11.self_attn.v_proj": {
|
| 365 |
+
"bits": 4,
|
| 366 |
+
"group_size": 64,
|
| 367 |
+
"rank": 16
|
| 368 |
+
},
|
| 369 |
+
"language_model.model.layers.11.self_attn.k_proj": {
|
| 370 |
+
"bits": 4,
|
| 371 |
+
"group_size": 64,
|
| 372 |
+
"rank": 16
|
| 373 |
+
},
|
| 374 |
+
"language_model.model.layers.11.self_attn.q_proj": {
|
| 375 |
+
"bits": 4,
|
| 376 |
+
"group_size": 64,
|
| 377 |
+
"rank": 16
|
| 378 |
+
},
|
| 379 |
+
"language_model.model.layers.12.mlp.up_proj": {
|
| 380 |
+
"bits": 2,
|
| 381 |
+
"group_size": 64,
|
| 382 |
+
"rank": 16
|
| 383 |
+
},
|
| 384 |
+
"language_model.model.layers.12.mlp.down_proj": {
|
| 385 |
+
"bits": 2,
|
| 386 |
+
"group_size": 64,
|
| 387 |
+
"rank": 16
|
| 388 |
+
},
|
| 389 |
+
"language_model.model.layers.12.mlp.gate_proj": {
|
| 390 |
+
"bits": 2,
|
| 391 |
+
"group_size": 64,
|
| 392 |
+
"rank": 16
|
| 393 |
+
},
|
| 394 |
+
"language_model.model.layers.12.linear_attn.out_proj": {
|
| 395 |
+
"bits": 4,
|
| 396 |
+
"group_size": 64,
|
| 397 |
+
"rank": 16
|
| 398 |
+
},
|
| 399 |
+
"language_model.model.layers.12.linear_attn.in_proj_z": {
|
| 400 |
+
"bits": 2,
|
| 401 |
+
"group_size": 64,
|
| 402 |
+
"rank": 16
|
| 403 |
+
},
|
| 404 |
+
"language_model.model.layers.12.linear_attn.in_proj_qkv": {
|
| 405 |
+
"bits": 2,
|
| 406 |
+
"group_size": 64,
|
| 407 |
+
"rank": 16
|
| 408 |
+
},
|
| 409 |
+
"language_model.model.layers.13.mlp.up_proj": {
|
| 410 |
+
"bits": 2,
|
| 411 |
+
"group_size": 64,
|
| 412 |
+
"rank": 16
|
| 413 |
+
},
|
| 414 |
+
"language_model.model.layers.13.mlp.down_proj": {
|
| 415 |
+
"bits": 2,
|
| 416 |
+
"group_size": 64,
|
| 417 |
+
"rank": 16
|
| 418 |
+
},
|
| 419 |
+
"language_model.model.layers.13.mlp.gate_proj": {
|
| 420 |
+
"bits": 2,
|
| 421 |
+
"group_size": 64,
|
| 422 |
+
"rank": 16
|
| 423 |
+
},
|
| 424 |
+
"language_model.model.layers.13.linear_attn.out_proj": {
|
| 425 |
+
"bits": 4,
|
| 426 |
+
"group_size": 64,
|
| 427 |
+
"rank": 16
|
| 428 |
+
},
|
| 429 |
+
"language_model.model.layers.13.linear_attn.in_proj_z": {
|
| 430 |
+
"bits": 2,
|
| 431 |
+
"group_size": 64,
|
| 432 |
+
"rank": 16
|
| 433 |
+
},
|
| 434 |
+
"language_model.model.layers.13.linear_attn.in_proj_qkv": {
|
| 435 |
+
"bits": 2,
|
| 436 |
+
"group_size": 64,
|
| 437 |
+
"rank": 16
|
| 438 |
+
},
|
| 439 |
+
"language_model.model.layers.14.mlp.up_proj": {
|
| 440 |
+
"bits": 2,
|
| 441 |
+
"group_size": 64,
|
| 442 |
+
"rank": 16
|
| 443 |
+
},
|
| 444 |
+
"language_model.model.layers.14.mlp.down_proj": {
|
| 445 |
+
"bits": 2,
|
| 446 |
+
"group_size": 64,
|
| 447 |
+
"rank": 16
|
| 448 |
+
},
|
| 449 |
+
"language_model.model.layers.14.mlp.gate_proj": {
|
| 450 |
+
"bits": 2,
|
| 451 |
+
"group_size": 64,
|
| 452 |
+
"rank": 16
|
| 453 |
+
},
|
| 454 |
+
"language_model.model.layers.14.linear_attn.out_proj": {
|
| 455 |
+
"bits": 4,
|
| 456 |
+
"group_size": 64,
|
| 457 |
+
"rank": 16
|
| 458 |
+
},
|
| 459 |
+
"language_model.model.layers.14.linear_attn.in_proj_z": {
|
| 460 |
+
"bits": 2,
|
| 461 |
+
"group_size": 64,
|
| 462 |
+
"rank": 16
|
| 463 |
+
},
|
| 464 |
+
"language_model.model.layers.14.linear_attn.in_proj_qkv": {
|
| 465 |
+
"bits": 2,
|
| 466 |
+
"group_size": 64,
|
| 467 |
+
"rank": 16
|
| 468 |
+
},
|
| 469 |
+
"language_model.model.layers.15.mlp.up_proj": {
|
| 470 |
+
"bits": 2,
|
| 471 |
+
"group_size": 64,
|
| 472 |
+
"rank": 16
|
| 473 |
+
},
|
| 474 |
+
"language_model.model.layers.15.mlp.down_proj": {
|
| 475 |
+
"bits": 2,
|
| 476 |
+
"group_size": 64,
|
| 477 |
+
"rank": 16
|
| 478 |
+
},
|
| 479 |
+
"language_model.model.layers.15.mlp.gate_proj": {
|
| 480 |
+
"bits": 2,
|
| 481 |
+
"group_size": 64,
|
| 482 |
+
"rank": 16
|
| 483 |
+
},
|
| 484 |
+
"language_model.model.layers.15.self_attn.o_proj": {
|
| 485 |
+
"bits": 4,
|
| 486 |
+
"group_size": 64,
|
| 487 |
+
"rank": 16
|
| 488 |
+
},
|
| 489 |
+
"language_model.model.layers.15.self_attn.v_proj": {
|
| 490 |
+
"bits": 4,
|
| 491 |
+
"group_size": 64,
|
| 492 |
+
"rank": 16
|
| 493 |
+
},
|
| 494 |
+
"language_model.model.layers.15.self_attn.k_proj": {
|
| 495 |
+
"bits": 4,
|
| 496 |
+
"group_size": 64,
|
| 497 |
+
"rank": 16
|
| 498 |
+
},
|
| 499 |
+
"language_model.model.layers.15.self_attn.q_proj": {
|
| 500 |
+
"bits": 4,
|
| 501 |
+
"group_size": 64,
|
| 502 |
+
"rank": 16
|
| 503 |
+
},
|
| 504 |
+
"language_model.model.layers.16.mlp.up_proj": {
|
| 505 |
+
"bits": 2,
|
| 506 |
+
"group_size": 64,
|
| 507 |
+
"rank": 16
|
| 508 |
+
},
|
| 509 |
+
"language_model.model.layers.16.mlp.down_proj": {
|
| 510 |
+
"bits": 2,
|
| 511 |
+
"group_size": 64,
|
| 512 |
+
"rank": 16
|
| 513 |
+
},
|
| 514 |
+
"language_model.model.layers.16.mlp.gate_proj": {
|
| 515 |
+
"bits": 2,
|
| 516 |
+
"group_size": 64,
|
| 517 |
+
"rank": 16
|
| 518 |
+
},
|
| 519 |
+
"language_model.model.layers.16.linear_attn.out_proj": {
|
| 520 |
+
"bits": 4,
|
| 521 |
+
"group_size": 64,
|
| 522 |
+
"rank": 16
|
| 523 |
+
},
|
| 524 |
+
"language_model.model.layers.16.linear_attn.in_proj_z": {
|
| 525 |
+
"bits": 2,
|
| 526 |
+
"group_size": 64,
|
| 527 |
+
"rank": 16
|
| 528 |
+
},
|
| 529 |
+
"language_model.model.layers.16.linear_attn.in_proj_qkv": {
|
| 530 |
+
"bits": 2,
|
| 531 |
+
"group_size": 64,
|
| 532 |
+
"rank": 16
|
| 533 |
+
},
|
| 534 |
+
"language_model.model.layers.17.mlp.up_proj": {
|
| 535 |
+
"bits": 2,
|
| 536 |
+
"group_size": 64,
|
| 537 |
+
"rank": 16
|
| 538 |
+
},
|
| 539 |
+
"language_model.model.layers.17.mlp.down_proj": {
|
| 540 |
+
"bits": 2,
|
| 541 |
+
"group_size": 64,
|
| 542 |
+
"rank": 16
|
| 543 |
+
},
|
| 544 |
+
"language_model.model.layers.17.mlp.gate_proj": {
|
| 545 |
+
"bits": 2,
|
| 546 |
+
"group_size": 64,
|
| 547 |
+
"rank": 16
|
| 548 |
+
},
|
| 549 |
+
"language_model.model.layers.17.linear_attn.out_proj": {
|
| 550 |
+
"bits": 4,
|
| 551 |
+
"group_size": 64,
|
| 552 |
+
"rank": 16
|
| 553 |
+
},
|
| 554 |
+
"language_model.model.layers.17.linear_attn.in_proj_z": {
|
| 555 |
+
"bits": 2,
|
| 556 |
+
"group_size": 64,
|
| 557 |
+
"rank": 16
|
| 558 |
+
},
|
| 559 |
+
"language_model.model.layers.17.linear_attn.in_proj_qkv": {
|
| 560 |
+
"bits": 2,
|
| 561 |
+
"group_size": 64,
|
| 562 |
+
"rank": 16
|
| 563 |
+
},
|
| 564 |
+
"language_model.model.layers.18.mlp.up_proj": {
|
| 565 |
+
"bits": 2,
|
| 566 |
+
"group_size": 64,
|
| 567 |
+
"rank": 16
|
| 568 |
+
},
|
| 569 |
+
"language_model.model.layers.18.mlp.down_proj": {
|
| 570 |
+
"bits": 2,
|
| 571 |
+
"group_size": 64,
|
| 572 |
+
"rank": 16
|
| 573 |
+
},
|
| 574 |
+
"language_model.model.layers.18.mlp.gate_proj": {
|
| 575 |
+
"bits": 2,
|
| 576 |
+
"group_size": 64,
|
| 577 |
+
"rank": 16
|
| 578 |
+
},
|
| 579 |
+
"language_model.model.layers.18.linear_attn.out_proj": {
|
| 580 |
+
"bits": 4,
|
| 581 |
+
"group_size": 64,
|
| 582 |
+
"rank": 16
|
| 583 |
+
},
|
| 584 |
+
"language_model.model.layers.18.linear_attn.in_proj_z": {
|
| 585 |
+
"bits": 2,
|
| 586 |
+
"group_size": 64,
|
| 587 |
+
"rank": 16
|
| 588 |
+
},
|
| 589 |
+
"language_model.model.layers.18.linear_attn.in_proj_qkv": {
|
| 590 |
+
"bits": 2,
|
| 591 |
+
"group_size": 64,
|
| 592 |
+
"rank": 16
|
| 593 |
+
},
|
| 594 |
+
"language_model.model.layers.19.mlp.up_proj": {
|
| 595 |
+
"bits": 2,
|
| 596 |
+
"group_size": 64,
|
| 597 |
+
"rank": 16
|
| 598 |
+
},
|
| 599 |
+
"language_model.model.layers.19.mlp.down_proj": {
|
| 600 |
+
"bits": 2,
|
| 601 |
+
"group_size": 64,
|
| 602 |
+
"rank": 16
|
| 603 |
+
},
|
| 604 |
+
"language_model.model.layers.19.mlp.gate_proj": {
|
| 605 |
+
"bits": 2,
|
| 606 |
+
"group_size": 64,
|
| 607 |
+
"rank": 16
|
| 608 |
+
},
|
| 609 |
+
"language_model.model.layers.19.self_attn.o_proj": {
|
| 610 |
+
"bits": 4,
|
| 611 |
+
"group_size": 64,
|
| 612 |
+
"rank": 16
|
| 613 |
+
},
|
| 614 |
+
"language_model.model.layers.19.self_attn.v_proj": {
|
| 615 |
+
"bits": 4,
|
| 616 |
+
"group_size": 64,
|
| 617 |
+
"rank": 16
|
| 618 |
+
},
|
| 619 |
+
"language_model.model.layers.19.self_attn.k_proj": {
|
| 620 |
+
"bits": 4,
|
| 621 |
+
"group_size": 64,
|
| 622 |
+
"rank": 16
|
| 623 |
+
},
|
| 624 |
+
"language_model.model.layers.19.self_attn.q_proj": {
|
| 625 |
+
"bits": 4,
|
| 626 |
+
"group_size": 64,
|
| 627 |
+
"rank": 16
|
| 628 |
+
},
|
| 629 |
+
"language_model.model.layers.20.mlp.up_proj": {
|
| 630 |
+
"bits": 2,
|
| 631 |
+
"group_size": 64,
|
| 632 |
+
"rank": 16
|
| 633 |
+
},
|
| 634 |
+
"language_model.model.layers.20.mlp.down_proj": {
|
| 635 |
+
"bits": 2,
|
| 636 |
+
"group_size": 64,
|
| 637 |
+
"rank": 16
|
| 638 |
+
},
|
| 639 |
+
"language_model.model.layers.20.mlp.gate_proj": {
|
| 640 |
+
"bits": 2,
|
| 641 |
+
"group_size": 64,
|
| 642 |
+
"rank": 16
|
| 643 |
+
},
|
| 644 |
+
"language_model.model.layers.20.linear_attn.out_proj": {
|
| 645 |
+
"bits": 4,
|
| 646 |
+
"group_size": 64,
|
| 647 |
+
"rank": 16
|
| 648 |
+
},
|
| 649 |
+
"language_model.model.layers.20.linear_attn.in_proj_z": {
|
| 650 |
+
"bits": 2,
|
| 651 |
+
"group_size": 64,
|
| 652 |
+
"rank": 16
|
| 653 |
+
},
|
| 654 |
+
"language_model.model.layers.20.linear_attn.in_proj_qkv": {
|
| 655 |
+
"bits": 2,
|
| 656 |
+
"group_size": 64,
|
| 657 |
+
"rank": 16
|
| 658 |
+
},
|
| 659 |
+
"language_model.model.layers.21.mlp.up_proj": {
|
| 660 |
+
"bits": 2,
|
| 661 |
+
"group_size": 64,
|
| 662 |
+
"rank": 16
|
| 663 |
+
},
|
| 664 |
+
"language_model.model.layers.21.mlp.down_proj": {
|
| 665 |
+
"bits": 2,
|
| 666 |
+
"group_size": 64,
|
| 667 |
+
"rank": 16
|
| 668 |
+
},
|
| 669 |
+
"language_model.model.layers.21.mlp.gate_proj": {
|
| 670 |
+
"bits": 2,
|
| 671 |
+
"group_size": 64,
|
| 672 |
+
"rank": 16
|
| 673 |
+
},
|
| 674 |
+
"language_model.model.layers.21.linear_attn.out_proj": {
|
| 675 |
+
"bits": 4,
|
| 676 |
+
"group_size": 64,
|
| 677 |
+
"rank": 16
|
| 678 |
+
},
|
| 679 |
+
"language_model.model.layers.21.linear_attn.in_proj_z": {
|
| 680 |
+
"bits": 2,
|
| 681 |
+
"group_size": 64,
|
| 682 |
+
"rank": 16
|
| 683 |
+
},
|
| 684 |
+
"language_model.model.layers.21.linear_attn.in_proj_qkv": {
|
| 685 |
+
"bits": 2,
|
| 686 |
+
"group_size": 64,
|
| 687 |
+
"rank": 16
|
| 688 |
+
},
|
| 689 |
+
"language_model.model.layers.22.mlp.up_proj": {
|
| 690 |
+
"bits": 2,
|
| 691 |
+
"group_size": 64,
|
| 692 |
+
"rank": 16
|
| 693 |
+
},
|
| 694 |
+
"language_model.model.layers.22.mlp.down_proj": {
|
| 695 |
+
"bits": 2,
|
| 696 |
+
"group_size": 64,
|
| 697 |
+
"rank": 16
|
| 698 |
+
},
|
| 699 |
+
"language_model.model.layers.22.mlp.gate_proj": {
|
| 700 |
+
"bits": 2,
|
| 701 |
+
"group_size": 64,
|
| 702 |
+
"rank": 16
|
| 703 |
+
},
|
| 704 |
+
"language_model.model.layers.22.linear_attn.out_proj": {
|
| 705 |
+
"bits": 4,
|
| 706 |
+
"group_size": 64,
|
| 707 |
+
"rank": 16
|
| 708 |
+
},
|
| 709 |
+
"language_model.model.layers.22.linear_attn.in_proj_z": {
|
| 710 |
+
"bits": 2,
|
| 711 |
+
"group_size": 64,
|
| 712 |
+
"rank": 16
|
| 713 |
+
},
|
| 714 |
+
"language_model.model.layers.22.linear_attn.in_proj_qkv": {
|
| 715 |
+
"bits": 2,
|
| 716 |
+
"group_size": 64,
|
| 717 |
+
"rank": 16
|
| 718 |
+
},
|
| 719 |
+
"language_model.model.layers.23.mlp.up_proj": {
|
| 720 |
+
"bits": 2,
|
| 721 |
+
"group_size": 64,
|
| 722 |
+
"rank": 16
|
| 723 |
+
},
|
| 724 |
+
"language_model.model.layers.23.mlp.down_proj": {
|
| 725 |
+
"bits": 2,
|
| 726 |
+
"group_size": 64,
|
| 727 |
+
"rank": 16
|
| 728 |
+
},
|
| 729 |
+
"language_model.model.layers.23.mlp.gate_proj": {
|
| 730 |
+
"bits": 2,
|
| 731 |
+
"group_size": 64,
|
| 732 |
+
"rank": 16
|
| 733 |
+
},
|
| 734 |
+
"language_model.model.layers.23.self_attn.o_proj": {
|
| 735 |
+
"bits": 4,
|
| 736 |
+
"group_size": 64,
|
| 737 |
+
"rank": 16
|
| 738 |
+
},
|
| 739 |
+
"language_model.model.layers.23.self_attn.v_proj": {
|
| 740 |
+
"bits": 4,
|
| 741 |
+
"group_size": 64,
|
| 742 |
+
"rank": 16
|
| 743 |
+
},
|
| 744 |
+
"language_model.model.layers.23.self_attn.k_proj": {
|
| 745 |
+
"bits": 4,
|
| 746 |
+
"group_size": 64,
|
| 747 |
+
"rank": 16
|
| 748 |
+
},
|
| 749 |
+
"language_model.model.layers.23.self_attn.q_proj": {
|
| 750 |
+
"bits": 4,
|
| 751 |
+
"group_size": 64,
|
| 752 |
+
"rank": 16
|
| 753 |
+
},
|
| 754 |
+
"language_model.model.layers.24.mlp.up_proj": {
|
| 755 |
+
"bits": 2,
|
| 756 |
+
"group_size": 64,
|
| 757 |
+
"rank": 16
|
| 758 |
+
},
|
| 759 |
+
"language_model.model.layers.24.mlp.down_proj": {
|
| 760 |
+
"bits": 2,
|
| 761 |
+
"group_size": 64,
|
| 762 |
+
"rank": 16
|
| 763 |
+
},
|
| 764 |
+
"language_model.model.layers.24.mlp.gate_proj": {
|
| 765 |
+
"bits": 2,
|
| 766 |
+
"group_size": 64,
|
| 767 |
+
"rank": 16
|
| 768 |
+
},
|
| 769 |
+
"language_model.model.layers.24.linear_attn.out_proj": {
|
| 770 |
+
"bits": 4,
|
| 771 |
+
"group_size": 64,
|
| 772 |
+
"rank": 16
|
| 773 |
+
},
|
| 774 |
+
"language_model.model.layers.24.linear_attn.in_proj_z": {
|
| 775 |
+
"bits": 2,
|
| 776 |
+
"group_size": 64,
|
| 777 |
+
"rank": 16
|
| 778 |
+
},
|
| 779 |
+
"language_model.model.layers.24.linear_attn.in_proj_qkv": {
|
| 780 |
+
"bits": 2,
|
| 781 |
+
"group_size": 64,
|
| 782 |
+
"rank": 16
|
| 783 |
+
},
|
| 784 |
+
"language_model.model.layers.25.mlp.up_proj": {
|
| 785 |
+
"bits": 2,
|
| 786 |
+
"group_size": 64,
|
| 787 |
+
"rank": 16
|
| 788 |
+
},
|
| 789 |
+
"language_model.model.layers.25.mlp.down_proj": {
|
| 790 |
+
"bits": 2,
|
| 791 |
+
"group_size": 64,
|
| 792 |
+
"rank": 16
|
| 793 |
+
},
|
| 794 |
+
"language_model.model.layers.25.mlp.gate_proj": {
|
| 795 |
+
"bits": 2,
|
| 796 |
+
"group_size": 64,
|
| 797 |
+
"rank": 16
|
| 798 |
+
},
|
| 799 |
+
"language_model.model.layers.25.linear_attn.out_proj": {
|
| 800 |
+
"bits": 4,
|
| 801 |
+
"group_size": 64,
|
| 802 |
+
"rank": 16
|
| 803 |
+
},
|
| 804 |
+
"language_model.model.layers.25.linear_attn.in_proj_z": {
|
| 805 |
+
"bits": 2,
|
| 806 |
+
"group_size": 64,
|
| 807 |
+
"rank": 16
|
| 808 |
+
},
|
| 809 |
+
"language_model.model.layers.25.linear_attn.in_proj_qkv": {
|
| 810 |
+
"bits": 2,
|
| 811 |
+
"group_size": 64,
|
| 812 |
+
"rank": 16
|
| 813 |
+
},
|
| 814 |
+
"language_model.model.layers.26.mlp.up_proj": {
|
| 815 |
+
"bits": 2,
|
| 816 |
+
"group_size": 64,
|
| 817 |
+
"rank": 16
|
| 818 |
+
},
|
| 819 |
+
"language_model.model.layers.26.mlp.down_proj": {
|
| 820 |
+
"bits": 2,
|
| 821 |
+
"group_size": 64,
|
| 822 |
+
"rank": 16
|
| 823 |
+
},
|
| 824 |
+
"language_model.model.layers.26.mlp.gate_proj": {
|
| 825 |
+
"bits": 2,
|
| 826 |
+
"group_size": 64,
|
| 827 |
+
"rank": 16
|
| 828 |
+
},
|
| 829 |
+
"language_model.model.layers.26.linear_attn.out_proj": {
|
| 830 |
+
"bits": 4,
|
| 831 |
+
"group_size": 64,
|
| 832 |
+
"rank": 16
|
| 833 |
+
},
|
| 834 |
+
"language_model.model.layers.26.linear_attn.in_proj_z": {
|
| 835 |
+
"bits": 2,
|
| 836 |
+
"group_size": 64,
|
| 837 |
+
"rank": 16
|
| 838 |
+
},
|
| 839 |
+
"language_model.model.layers.26.linear_attn.in_proj_qkv": {
|
| 840 |
+
"bits": 2,
|
| 841 |
+
"group_size": 64,
|
| 842 |
+
"rank": 16
|
| 843 |
+
},
|
| 844 |
+
"language_model.model.layers.27.mlp.up_proj": {
|
| 845 |
+
"bits": 2,
|
| 846 |
+
"group_size": 64,
|
| 847 |
+
"rank": 16
|
| 848 |
+
},
|
| 849 |
+
"language_model.model.layers.27.mlp.down_proj": {
|
| 850 |
+
"bits": 2,
|
| 851 |
+
"group_size": 64,
|
| 852 |
+
"rank": 16
|
| 853 |
+
},
|
| 854 |
+
"language_model.model.layers.27.mlp.gate_proj": {
|
| 855 |
+
"bits": 2,
|
| 856 |
+
"group_size": 64,
|
| 857 |
+
"rank": 16
|
| 858 |
+
},
|
| 859 |
+
"language_model.model.layers.27.self_attn.o_proj": {
|
| 860 |
+
"bits": 4,
|
| 861 |
+
"group_size": 64,
|
| 862 |
+
"rank": 16
|
| 863 |
+
},
|
| 864 |
+
"language_model.model.layers.27.self_attn.v_proj": {
|
| 865 |
+
"bits": 4,
|
| 866 |
+
"group_size": 64,
|
| 867 |
+
"rank": 16
|
| 868 |
+
},
|
| 869 |
+
"language_model.model.layers.27.self_attn.k_proj": {
|
| 870 |
+
"bits": 4,
|
| 871 |
+
"group_size": 64,
|
| 872 |
+
"rank": 16
|
| 873 |
+
},
|
| 874 |
+
"language_model.model.layers.27.self_attn.q_proj": {
|
| 875 |
+
"bits": 4,
|
| 876 |
+
"group_size": 64,
|
| 877 |
+
"rank": 16
|
| 878 |
+
},
|
| 879 |
+
"language_model.model.layers.28.mlp.up_proj": {
|
| 880 |
+
"bits": 2,
|
| 881 |
+
"group_size": 64,
|
| 882 |
+
"rank": 16
|
| 883 |
+
},
|
| 884 |
+
"language_model.model.layers.28.mlp.down_proj": {
|
| 885 |
+
"bits": 2,
|
| 886 |
+
"group_size": 64,
|
| 887 |
+
"rank": 16
|
| 888 |
+
},
|
| 889 |
+
"language_model.model.layers.28.mlp.gate_proj": {
|
| 890 |
+
"bits": 2,
|
| 891 |
+
"group_size": 64,
|
| 892 |
+
"rank": 16
|
| 893 |
+
},
|
| 894 |
+
"language_model.model.layers.28.linear_attn.out_proj": {
|
| 895 |
+
"bits": 4,
|
| 896 |
+
"group_size": 64,
|
| 897 |
+
"rank": 16
|
| 898 |
+
},
|
| 899 |
+
"language_model.model.layers.28.linear_attn.in_proj_z": {
|
| 900 |
+
"bits": 2,
|
| 901 |
+
"group_size": 64,
|
| 902 |
+
"rank": 16
|
| 903 |
+
},
|
| 904 |
+
"language_model.model.layers.28.linear_attn.in_proj_qkv": {
|
| 905 |
+
"bits": 2,
|
| 906 |
+
"group_size": 64,
|
| 907 |
+
"rank": 16
|
| 908 |
+
},
|
| 909 |
+
"language_model.model.layers.29.mlp.up_proj": {
|
| 910 |
+
"bits": 2,
|
| 911 |
+
"group_size": 64,
|
| 912 |
+
"rank": 16
|
| 913 |
+
},
|
| 914 |
+
"language_model.model.layers.29.mlp.down_proj": {
|
| 915 |
+
"bits": 2,
|
| 916 |
+
"group_size": 64,
|
| 917 |
+
"rank": 16
|
| 918 |
+
},
|
| 919 |
+
"language_model.model.layers.29.mlp.gate_proj": {
|
| 920 |
+
"bits": 2,
|
| 921 |
+
"group_size": 64,
|
| 922 |
+
"rank": 16
|
| 923 |
+
},
|
| 924 |
+
"language_model.model.layers.29.linear_attn.out_proj": {
|
| 925 |
+
"bits": 4,
|
| 926 |
+
"group_size": 64,
|
| 927 |
+
"rank": 16
|
| 928 |
+
},
|
| 929 |
+
"language_model.model.layers.29.linear_attn.in_proj_z": {
|
| 930 |
+
"bits": 2,
|
| 931 |
+
"group_size": 64,
|
| 932 |
+
"rank": 16
|
| 933 |
+
},
|
| 934 |
+
"language_model.model.layers.29.linear_attn.in_proj_qkv": {
|
| 935 |
+
"bits": 2,
|
| 936 |
+
"group_size": 64,
|
| 937 |
+
"rank": 16
|
| 938 |
+
},
|
| 939 |
+
"language_model.model.layers.30.mlp.up_proj": {
|
| 940 |
+
"bits": 2,
|
| 941 |
+
"group_size": 64,
|
| 942 |
+
"rank": 16
|
| 943 |
+
},
|
| 944 |
+
"language_model.model.layers.30.mlp.down_proj": {
|
| 945 |
+
"bits": 2,
|
| 946 |
+
"group_size": 64,
|
| 947 |
+
"rank": 16
|
| 948 |
+
},
|
| 949 |
+
"language_model.model.layers.30.mlp.gate_proj": {
|
| 950 |
+
"bits": 2,
|
| 951 |
+
"group_size": 64,
|
| 952 |
+
"rank": 16
|
| 953 |
+
},
|
| 954 |
+
"language_model.model.layers.30.linear_attn.out_proj": {
|
| 955 |
+
"bits": 4,
|
| 956 |
+
"group_size": 64,
|
| 957 |
+
"rank": 16
|
| 958 |
+
},
|
| 959 |
+
"language_model.model.layers.30.linear_attn.in_proj_z": {
|
| 960 |
+
"bits": 2,
|
| 961 |
+
"group_size": 64,
|
| 962 |
+
"rank": 16
|
| 963 |
+
},
|
| 964 |
+
"language_model.model.layers.30.linear_attn.in_proj_qkv": {
|
| 965 |
+
"bits": 2,
|
| 966 |
+
"group_size": 64,
|
| 967 |
+
"rank": 16
|
| 968 |
+
},
|
| 969 |
+
"language_model.model.layers.31.mlp.up_proj": {
|
| 970 |
+
"bits": 2,
|
| 971 |
+
"group_size": 64,
|
| 972 |
+
"rank": 16
|
| 973 |
+
},
|
| 974 |
+
"language_model.model.layers.31.mlp.down_proj": {
|
| 975 |
+
"bits": 2,
|
| 976 |
+
"group_size": 64,
|
| 977 |
+
"rank": 16
|
| 978 |
+
},
|
| 979 |
+
"language_model.model.layers.31.mlp.gate_proj": {
|
| 980 |
+
"bits": 2,
|
| 981 |
+
"group_size": 64,
|
| 982 |
+
"rank": 16
|
| 983 |
+
},
|
| 984 |
+
"language_model.model.layers.31.self_attn.o_proj": {
|
| 985 |
+
"bits": 4,
|
| 986 |
+
"group_size": 64,
|
| 987 |
+
"rank": 16
|
| 988 |
+
},
|
| 989 |
+
"language_model.model.layers.31.self_attn.v_proj": {
|
| 990 |
+
"bits": 4,
|
| 991 |
+
"group_size": 64,
|
| 992 |
+
"rank": 16
|
| 993 |
+
},
|
| 994 |
+
"language_model.model.layers.31.self_attn.k_proj": {
|
| 995 |
+
"bits": 4,
|
| 996 |
+
"group_size": 64,
|
| 997 |
+
"rank": 16
|
| 998 |
+
},
|
| 999 |
+
"language_model.model.layers.31.self_attn.q_proj": {
|
| 1000 |
+
"bits": 4,
|
| 1001 |
+
"group_size": 64,
|
| 1002 |
+
"rank": 16
|
| 1003 |
+
}
|
| 1004 |
+
},
|
| 1005 |
+
"embedding": {
|
| 1006 |
+
"bits": 4,
|
| 1007 |
+
"group_size": 64
|
| 1008 |
+
},
|
| 1009 |
+
"source": {
|
| 1010 |
+
"format": "kev-packed-mlx-v1",
|
| 1011 |
+
"source_packed_sha256": "af8703c3c8f848c18f49a40e91978363fae1e0e3b6c4678a9436e3ebc5da73e0",
|
| 1012 |
+
"source_cloq_sha256": "52b8699b25a819542d8073b0de52c79bfb77e46b3bc969e6f553b6a047e0afcc",
|
| 1013 |
+
"arithmetic": "native affine quantized matmul plus separate fp32 low-rank compensation; not merged-bf16 parity"
|
| 1014 |
+
},
|
| 1015 |
+
"model_sha256": "c4899eaa9fb206e77260f832d00188c9b90750724422ea0eea43ed69faeff43d"
|
| 1016 |
+
}
|
release-manifest.json
ADDED
|
@@ -0,0 +1,252 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"repo_id": "Jakevin/kev-4b-ternary-mlx",
|
| 3 |
+
"tag": "v0.1.0-experimental",
|
| 4 |
+
"date": "2026-10-06",
|
| 5 |
+
"format": "kev-packed-mlx-v1",
|
| 6 |
+
"source_kev_revision": "6cfce5c2fa4b4bd64026336ab649c5ca78857d52",
|
| 7 |
+
"files": {
|
| 8 |
+
"LICENSE": {
|
| 9 |
+
"bytes": 11343,
|
| 10 |
+
"sha256": "6b08bb37982c233aa12bcbdf19106da12f3f4fcf800773ea42e28ebeddd34fb8"
|
| 11 |
+
},
|
| 12 |
+
"NOTICE": {
|
| 13 |
+
"bytes": 528,
|
| 14 |
+
"sha256": "9021541d1f4af146262ef38449dd517c13d059404132e3f7e017ab550cef62bc"
|
| 15 |
+
},
|
| 16 |
+
"README.md": {
|
| 17 |
+
"bytes": 8520,
|
| 18 |
+
"sha256": "c11cf88162f1739d10f5e68f1a1d2bdd05d794381ebb99dbaa65a06f0f0c4dea"
|
| 19 |
+
},
|
| 20 |
+
"added_tokens.json": {
|
| 21 |
+
"bytes": 707,
|
| 22 |
+
"sha256": "c0284b582e14987fbd3d5a2cb2bd139084371ed9acbae488829a1c900833c680"
|
| 23 |
+
},
|
| 24 |
+
"config.json": {
|
| 25 |
+
"bytes": 2692,
|
| 26 |
+
"sha256": "fc1557ea241bfc2c8abd87c1213729a44562fe5da9edf0fd5dc9f95b6cdbdaed"
|
| 27 |
+
},
|
| 28 |
+
"evaluation/development-summary.json": {
|
| 29 |
+
"bytes": 2475,
|
| 30 |
+
"sha256": "d17ad38950313c7513bd67a73cb2b15358c0a1f42dbec2db02c6c9388a392f5f"
|
| 31 |
+
},
|
| 32 |
+
"evaluation/memory-long.json": {
|
| 33 |
+
"bytes": 895,
|
| 34 |
+
"sha256": "b551f722360e3420a30aed1fada9d44557bef26730938df1a2a73b961a54354d"
|
| 35 |
+
},
|
| 36 |
+
"evaluation/memory-short.json": {
|
| 37 |
+
"bytes": 894,
|
| 38 |
+
"sha256": "d3111a799b8760323eba78cae3bf00fd61c61d24073a232cff33a9c7baf3d41d"
|
| 39 |
+
},
|
| 40 |
+
"evaluation/native-comparison.json": {
|
| 41 |
+
"bytes": 275,
|
| 42 |
+
"sha256": "a1355c9afa8b8f7f986c355941233164457d05503810d49fd698bd224adf8fc5"
|
| 43 |
+
},
|
| 44 |
+
"head.pt": {
|
| 45 |
+
"bytes": 5247807,
|
| 46 |
+
"sha256": "6f6740d94a49187770b68c0f058d85bedc0cf7efc20dff7ca9a418c8b461dc0e"
|
| 47 |
+
},
|
| 48 |
+
"merges.txt": {
|
| 49 |
+
"bytes": 1671853,
|
| 50 |
+
"sha256": "8831e4f1a044471340f7c0a83d7bd71306a5b867e95fd870f74d0c5308a904d5"
|
| 51 |
+
},
|
| 52 |
+
"model.safetensors": {
|
| 53 |
+
"bytes": 1678909327,
|
| 54 |
+
"sha256": "c4899eaa9fb206e77260f832d00188c9b90750724422ea0eea43ed69faeff43d"
|
| 55 |
+
},
|
| 56 |
+
"packed.json": {
|
| 57 |
+
"bytes": 24929,
|
| 58 |
+
"sha256": "1ae5d2d98e967e49ebdf29842da0847faf87fa03ace9370f8e05715ef583dd29"
|
| 59 |
+
},
|
| 60 |
+
"runtime/LICENSE": {
|
| 61 |
+
"bytes": 11343,
|
| 62 |
+
"sha256": "6b08bb37982c233aa12bcbdf19106da12f3f4fcf800773ea42e28ebeddd34fb8"
|
| 63 |
+
},
|
| 64 |
+
"runtime/README.md": {
|
| 65 |
+
"bytes": 42304,
|
| 66 |
+
"sha256": "19c9d7d30cd723b9f112f1f68d0b415a28c5ee616dbc21b358fc29bd07a66311"
|
| 67 |
+
},
|
| 68 |
+
"runtime/docs/packed-mlx.md": {
|
| 69 |
+
"bytes": 6233,
|
| 70 |
+
"sha256": "c70c572479cb7e8e8094e7cdd5982ade28084c9468e38190ef0d13ffa675631b"
|
| 71 |
+
},
|
| 72 |
+
"runtime/kev/__init__.py": {
|
| 73 |
+
"bytes": 0,
|
| 74 |
+
"sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
|
| 75 |
+
},
|
| 76 |
+
"runtime/kev/anchors.py": {
|
| 77 |
+
"bytes": 3930,
|
| 78 |
+
"sha256": "6963eafdb276db5a6c94d939eae675c761555553b8d0f963ffd803448beaabbb"
|
| 79 |
+
},
|
| 80 |
+
"runtime/kev/api.py": {
|
| 81 |
+
"bytes": 7526,
|
| 82 |
+
"sha256": "a6f53af5354f8ea17915abc78360fdf291a5108ce39a53a2d9135d48c603ae58"
|
| 83 |
+
},
|
| 84 |
+
"runtime/kev/autoresearch.py": {
|
| 85 |
+
"bytes": 11866,
|
| 86 |
+
"sha256": "7c62762b99c1b081bef06aa1ce83f8dfccce8ac7e9d2c5c0ab4d058055c7eab1"
|
| 87 |
+
},
|
| 88 |
+
"runtime/kev/benchmark.py": {
|
| 89 |
+
"bytes": 14203,
|
| 90 |
+
"sha256": "82e6a935472f170da48ec37eee8986c519595c4943f6f3040ec67b5da3afb278"
|
| 91 |
+
},
|
| 92 |
+
"runtime/kev/budget.py": {
|
| 93 |
+
"bytes": 8457,
|
| 94 |
+
"sha256": "246b5ca6450567a61b881308c9482cf1af6ecba599ca161af1febb23d6dc28a7"
|
| 95 |
+
},
|
| 96 |
+
"runtime/kev/calibrate.py": {
|
| 97 |
+
"bytes": 4846,
|
| 98 |
+
"sha256": "435ac98696fbd329c3a07f0c12633ba3c09a846f0e9c92e17540bbf8ba2986f8"
|
| 99 |
+
},
|
| 100 |
+
"runtime/kev/checkpoint.py": {
|
| 101 |
+
"bytes": 24924,
|
| 102 |
+
"sha256": "efe96cde9ef989c4fea3b6a0b631338bb1870bfe4b0f1eb99d71bb0761af34b7"
|
| 103 |
+
},
|
| 104 |
+
"runtime/kev/compare.py": {
|
| 105 |
+
"bytes": 3198,
|
| 106 |
+
"sha256": "3606dbf98bb305d66420158498cb837304d2edb8dab01cdaf2ec42964bef3435"
|
| 107 |
+
},
|
| 108 |
+
"runtime/kev/composition.py": {
|
| 109 |
+
"bytes": 15847,
|
| 110 |
+
"sha256": "f335ed17e18e0a544893db5e22b9a059e6ce1b2e14dbb863ac7d7bbf8f3e0536"
|
| 111 |
+
},
|
| 112 |
+
"runtime/kev/contrastive.py": {
|
| 113 |
+
"bytes": 20624,
|
| 114 |
+
"sha256": "cbb979aa5d40265ad0e64695f94b281d91751fa405811ddfa8212fede111edcf"
|
| 115 |
+
},
|
| 116 |
+
"runtime/kev/cuda_graphs.py": {
|
| 117 |
+
"bytes": 25356,
|
| 118 |
+
"sha256": "5ce84f0e4e5ad2177314ed3ecc9564d45d14ad9eb422e0fb079c9e1cbbb011b6"
|
| 119 |
+
},
|
| 120 |
+
"runtime/kev/data.py": {
|
| 121 |
+
"bytes": 25861,
|
| 122 |
+
"sha256": "e77a2fffeb8ee7e05b118893be8e38b8358aef97c3ca2e143cdce208201b92f9"
|
| 123 |
+
},
|
| 124 |
+
"runtime/kev/device.py": {
|
| 125 |
+
"bytes": 1164,
|
| 126 |
+
"sha256": "bda81dff5be33c59eee8de76e534976cc137a2661da6afd2a67abb85d8dad4a2"
|
| 127 |
+
},
|
| 128 |
+
"runtime/kev/evaluate.py": {
|
| 129 |
+
"bytes": 14408,
|
| 130 |
+
"sha256": "b2a6fbad7a1494728ea8d9ed0ad54ae23ff2ee1d543444fd0c3c110da4ab07cc"
|
| 131 |
+
},
|
| 132 |
+
"runtime/kev/experiment.py": {
|
| 133 |
+
"bytes": 34376,
|
| 134 |
+
"sha256": "3fd13a86f111dfd3292af9448df0ee28bd9596e3a77c1391f318cd654be14e8a"
|
| 135 |
+
},
|
| 136 |
+
"runtime/kev/full_ft.py": {
|
| 137 |
+
"bytes": 28171,
|
| 138 |
+
"sha256": "c3abef39cd9d486b3c075a37286afbcd1aa42a04c78a1b7b9b2e77b5a3963929"
|
| 139 |
+
},
|
| 140 |
+
"runtime/kev/fused_qwen35.py": {
|
| 141 |
+
"bytes": 10294,
|
| 142 |
+
"sha256": "5d2e587808822ce539ff52f63024357630b67a65f88cd734b82fd488655aebed"
|
| 143 |
+
},
|
| 144 |
+
"runtime/kev/jev.py": {
|
| 145 |
+
"bytes": 3275,
|
| 146 |
+
"sha256": "5488645c4fc2a53ca430e556ebfb744736e1a156bd121bd10ba0a9935f423e1c"
|
| 147 |
+
},
|
| 148 |
+
"runtime/kev/metrics.py": {
|
| 149 |
+
"bytes": 24401,
|
| 150 |
+
"sha256": "167a586334cfcced4adcc197758535c3bcda832b855c151095ef2d1968cf8983"
|
| 151 |
+
},
|
| 152 |
+
"runtime/kev/mirror.py": {
|
| 153 |
+
"bytes": 4982,
|
| 154 |
+
"sha256": "ca2284c670ce2cf1db57231968b9e6cbc01955f8f3c9a3207fd82e73d7b476b6"
|
| 155 |
+
},
|
| 156 |
+
"runtime/kev/mlx_model.py": {
|
| 157 |
+
"bytes": 15144,
|
| 158 |
+
"sha256": "8aef8aef93f50663395d73b5ae6cc964a01c4f4bf7f92a3843c1f32681cc9c09"
|
| 159 |
+
},
|
| 160 |
+
"runtime/kev/model.py": {
|
| 161 |
+
"bytes": 35153,
|
| 162 |
+
"sha256": "2ff3b2e67c048e3bc5a0d1a32aafe50804ed315a5564fce06bb4b95697c5271c"
|
| 163 |
+
},
|
| 164 |
+
"runtime/kev/packed_model.py": {
|
| 165 |
+
"bytes": 4892,
|
| 166 |
+
"sha256": "72cf553bd8d6f6e6d1e9f903eab4eba44d33d11571803752c0a89404f39e3499"
|
| 167 |
+
},
|
| 168 |
+
"runtime/kev/plot.py": {
|
| 169 |
+
"bytes": 3368,
|
| 170 |
+
"sha256": "0874bcff2885d8155a1cceade0de8a275a7163d3c3fd6aa0c294e8ccf5ff7e02"
|
| 171 |
+
},
|
| 172 |
+
"runtime/kev/predictors.py": {
|
| 173 |
+
"bytes": 21963,
|
| 174 |
+
"sha256": "29263ae3297e5cde6cf03da5c5aa3ed377c9681c6a5a639ea4aa259acec7575b"
|
| 175 |
+
},
|
| 176 |
+
"runtime/kev/publish.py": {
|
| 177 |
+
"bytes": 7512,
|
| 178 |
+
"sha256": "fcb490a374fc9ef330424ad4a75ce281ab08d2cb04dff2ad912c5f4dae6ab8a7"
|
| 179 |
+
},
|
| 180 |
+
"runtime/kev/rounds.py": {
|
| 181 |
+
"bytes": 96716,
|
| 182 |
+
"sha256": "2975bede7250812aecea36b17c0aefda6b389dc3a50e782e82e7566a7168c270"
|
| 183 |
+
},
|
| 184 |
+
"runtime/kev/serve.py": {
|
| 185 |
+
"bytes": 23057,
|
| 186 |
+
"sha256": "4f9885ce179dcb7de743f59be4b561758f47fac09373525e2994860b3345b162"
|
| 187 |
+
},
|
| 188 |
+
"runtime/kev/shared_prefix.py": {
|
| 189 |
+
"bytes": 9671,
|
| 190 |
+
"sha256": "bc6dedbad734cca227f3ec6746223b3b03363b0c045aa9f2493d886b1828e71e"
|
| 191 |
+
},
|
| 192 |
+
"runtime/kev/study_v3.py": {
|
| 193 |
+
"bytes": 14919,
|
| 194 |
+
"sha256": "7891150aa623e4479c7c789d4dc4662f186132b37b7bac1e1f072b9e08ee13e3"
|
| 195 |
+
},
|
| 196 |
+
"runtime/kev/suite.py": {
|
| 197 |
+
"bytes": 26388,
|
| 198 |
+
"sha256": "8e7d257a10b4704d7bf4f8ec2c74dddf26aa5a2a610ea78610360e6baf43eff5"
|
| 199 |
+
},
|
| 200 |
+
"runtime/kev/train.py": {
|
| 201 |
+
"bytes": 53543,
|
| 202 |
+
"sha256": "de0b0971defd79c053d56f8e8b17b0501b2e264e93a66a8d5493f1b1ba839416"
|
| 203 |
+
},
|
| 204 |
+
"runtime/kev/transfer_v9.py": {
|
| 205 |
+
"bytes": 10969,
|
| 206 |
+
"sha256": "0902409742151250a28af1fd8f42258b70fcc161f3deed3ee35fe3f473f3c763"
|
| 207 |
+
},
|
| 208 |
+
"runtime/pyproject.toml": {
|
| 209 |
+
"bytes": 1721,
|
| 210 |
+
"sha256": "7c17fbe9efcadda3eb488b59adf1b6db5938cd33af4d3bd389ce4359b73fcc7d"
|
| 211 |
+
},
|
| 212 |
+
"runtime/scripts/build_packed_model.py": {
|
| 213 |
+
"bytes": 7130,
|
| 214 |
+
"sha256": "1f4aba3f4c754753c3c654b934cd410acc7fb9c0f27fa35366e9d4a1899e5d5c"
|
| 215 |
+
},
|
| 216 |
+
"runtime/scripts/mlx_long_states.py": {
|
| 217 |
+
"bytes": 29651,
|
| 218 |
+
"sha256": "ec1ae5ff040da1060143dcbcd551bdd5982e583951af58677cdb0b6624ed005d"
|
| 219 |
+
},
|
| 220 |
+
"runtime/scripts/profile_packed_model.py": {
|
| 221 |
+
"bytes": 3543,
|
| 222 |
+
"sha256": "14b8e346771f978fb61b2d167cbb5f56011600f3e59ae4a16b84518fd99d2a18"
|
| 223 |
+
},
|
| 224 |
+
"runtime/tests/test_packed_model.py": {
|
| 225 |
+
"bytes": 7193,
|
| 226 |
+
"sha256": "9ff620541b964886f1e13117ffb4a6e143b3c843a144fd2fcc1a087baa3c82de"
|
| 227 |
+
},
|
| 228 |
+
"special_tokens_map.json": {
|
| 229 |
+
"bytes": 616,
|
| 230 |
+
"sha256": "6676f091c8bc4d1b50146427cfde92073402866b87b6e39223227931b70083e9"
|
| 231 |
+
},
|
| 232 |
+
"tokenizer.json": {
|
| 233 |
+
"bytes": 19989325,
|
| 234 |
+
"sha256": "06b9509352d2af50381ab2247e083b80d32d5c0aba91c272ca9ff729b6a0e523"
|
| 235 |
+
},
|
| 236 |
+
"tokenizer_config.json": {
|
| 237 |
+
"bytes": 1128,
|
| 238 |
+
"sha256": "8671bed7c852ce9e661be94f179a7b4ffd091c2a65aea0363e5501c20318ee45"
|
| 239 |
+
},
|
| 240 |
+
"vocab.json": {
|
| 241 |
+
"bytes": 2776833,
|
| 242 |
+
"sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910"
|
| 243 |
+
}
|
| 244 |
+
},
|
| 245 |
+
"excluded": [
|
| 246 |
+
"original backbone and PEFT weights",
|
| 247 |
+
"NAER/Chinese Linguipedia sources, derivatives, labels and predictions",
|
| 248 |
+
"training and validation JSONL",
|
| 249 |
+
"per-example evaluation prediction dumps"
|
| 250 |
+
],
|
| 251 |
+
"note": "Runtime code included; model SHA verified against packed.json. Manifest excludes its own hash."
|
| 252 |
+
}
|
runtime/LICENSE
ADDED
|
@@ -0,0 +1,202 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
|
| 2 |
+
Apache License
|
| 3 |
+
Version 2.0, January 2004
|
| 4 |
+
http://www.apache.org/licenses/
|
| 5 |
+
|
| 6 |
+
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
| 7 |
+
|
| 8 |
+
1. Definitions.
|
| 9 |
+
|
| 10 |
+
"License" shall mean the terms and conditions for use, reproduction,
|
| 11 |
+
and distribution as defined by Sections 1 through 9 of this document.
|
| 12 |
+
|
| 13 |
+
"Licensor" shall mean the copyright owner or entity authorized by
|
| 14 |
+
the copyright owner that is granting the License.
|
| 15 |
+
|
| 16 |
+
"Legal Entity" shall mean the union of the acting entity and all
|
| 17 |
+
other entities that control, are controlled by, or are under common
|
| 18 |
+
control with that entity. For the purposes of this definition,
|
| 19 |
+
"control" means (i) the power, direct or indirect, to cause the
|
| 20 |
+
direction or management of such entity, whether by contract or
|
| 21 |
+
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
| 22 |
+
outstanding shares, or (iii) beneficial ownership of such entity.
|
| 23 |
+
|
| 24 |
+
"You" (or "Your") shall mean an individual or Legal Entity
|
| 25 |
+
exercising permissions granted by this License.
|
| 26 |
+
|
| 27 |
+
"Source" form shall mean the preferred form for making modifications,
|
| 28 |
+
including but not limited to software source code, documentation
|
| 29 |
+
source, and configuration files.
|
| 30 |
+
|
| 31 |
+
"Object" form shall mean any form resulting from mechanical
|
| 32 |
+
transformation or translation of a Source form, including but
|
| 33 |
+
not limited to compiled object code, generated documentation,
|
| 34 |
+
and conversions to other media types.
|
| 35 |
+
|
| 36 |
+
"Work" shall mean the work of authorship, whether in Source or
|
| 37 |
+
Object form, made available under the License, as indicated by a
|
| 38 |
+
copyright notice that is included in or attached to the work
|
| 39 |
+
(an example is provided in the Appendix below).
|
| 40 |
+
|
| 41 |
+
"Derivative Works" shall mean any work, whether in Source or Object
|
| 42 |
+
form, that is based on (or derived from) the Work and for which the
|
| 43 |
+
editorial revisions, annotations, elaborations, or other modifications
|
| 44 |
+
represent, as a whole, an original work of authorship. For the purposes
|
| 45 |
+
of this License, Derivative Works shall not include works that remain
|
| 46 |
+
separable from, or merely link (or bind by name) to the interfaces of,
|
| 47 |
+
the Work and Derivative Works thereof.
|
| 48 |
+
|
| 49 |
+
"Contribution" shall mean any work of authorship, including
|
| 50 |
+
the original version of the Work and any modifications or additions
|
| 51 |
+
to that Work or Derivative Works thereof, that is intentionally
|
| 52 |
+
submitted to Licensor for inclusion in the Work by the copyright owner
|
| 53 |
+
or by an individual or Legal Entity authorized to submit on behalf of
|
| 54 |
+
the copyright owner. For the purposes of this definition, "submitted"
|
| 55 |
+
means any form of electronic, verbal, or written communication sent
|
| 56 |
+
to the Licensor or its representatives, including but not limited to
|
| 57 |
+
communication on electronic mailing lists, source code control systems,
|
| 58 |
+
and issue tracking systems that are managed by, or on behalf of, the
|
| 59 |
+
Licensor for the purpose of discussing and improving the Work, but
|
| 60 |
+
excluding communication that is conspicuously marked or otherwise
|
| 61 |
+
designated in writing by the copyright owner as "Not a Contribution."
|
| 62 |
+
|
| 63 |
+
"Contributor" shall mean Licensor and any individual or Legal Entity
|
| 64 |
+
on behalf of whom a Contribution has been received by Licensor and
|
| 65 |
+
subsequently incorporated within the Work.
|
| 66 |
+
|
| 67 |
+
2. Grant of Copyright License. Subject to the terms and conditions of
|
| 68 |
+
this License, each Contributor hereby grants to You a perpetual,
|
| 69 |
+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
| 70 |
+
copyright license to reproduce, prepare Derivative Works of,
|
| 71 |
+
publicly display, publicly perform, sublicense, and distribute the
|
| 72 |
+
Work and such Derivative Works in Source or Object form.
|
| 73 |
+
|
| 74 |
+
3. Grant of Patent License. Subject to the terms and conditions of
|
| 75 |
+
this License, each Contributor hereby grants to You a perpetual,
|
| 76 |
+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
| 77 |
+
(except as stated in this section) patent license to make, have made,
|
| 78 |
+
use, offer to sell, sell, import, and otherwise transfer the Work,
|
| 79 |
+
where such license applies only to those patent claims licensable
|
| 80 |
+
by such Contributor that are necessarily infringed by their
|
| 81 |
+
Contribution(s) alone or by combination of their Contribution(s)
|
| 82 |
+
with the Work to which such Contribution(s) was submitted. If You
|
| 83 |
+
institute patent litigation against any entity (including a
|
| 84 |
+
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
| 85 |
+
or a Contribution incorporated within the Work constitutes direct
|
| 86 |
+
or contributory patent infringement, then any patent licenses
|
| 87 |
+
granted to You under this License for that Work shall terminate
|
| 88 |
+
as of the date such litigation is filed.
|
| 89 |
+
|
| 90 |
+
4. Redistribution. You may reproduce and distribute copies of the
|
| 91 |
+
Work or Derivative Works thereof in any medium, with or without
|
| 92 |
+
modifications, and in Source or Object form, provided that You
|
| 93 |
+
meet the following conditions:
|
| 94 |
+
|
| 95 |
+
(a) You must give any other recipients of the Work or
|
| 96 |
+
Derivative Works a copy of this License; and
|
| 97 |
+
|
| 98 |
+
(b) You must cause any modified files to carry prominent notices
|
| 99 |
+
stating that You changed the files; and
|
| 100 |
+
|
| 101 |
+
(c) You must retain, in the Source form of any Derivative Works
|
| 102 |
+
that You distribute, all copyright, patent, trademark, and
|
| 103 |
+
attribution notices from the Source form of the Work,
|
| 104 |
+
excluding those notices that do not pertain to any part of
|
| 105 |
+
the Derivative Works; and
|
| 106 |
+
|
| 107 |
+
(d) If the Work includes a "NOTICE" text file as part of its
|
| 108 |
+
distribution, then any Derivative Works that You distribute must
|
| 109 |
+
include a readable copy of the attribution notices contained
|
| 110 |
+
within such NOTICE file, excluding those notices that do not
|
| 111 |
+
pertain to any part of the Derivative Works, in at least one
|
| 112 |
+
of the following places: within a NOTICE text file distributed
|
| 113 |
+
as part of the Derivative Works; within the Source form or
|
| 114 |
+
documentation, if provided along with the Derivative Works; or,
|
| 115 |
+
within a display generated by the Derivative Works, if and
|
| 116 |
+
wherever such third-party notices normally appear. The contents
|
| 117 |
+
of the NOTICE file are for informational purposes only and
|
| 118 |
+
do not modify the License. You may add Your own attribution
|
| 119 |
+
notices within Derivative Works that You distribute, alongside
|
| 120 |
+
or as an addendum to the NOTICE text from the Work, provided
|
| 121 |
+
that such additional attribution notices cannot be construed
|
| 122 |
+
as modifying the License.
|
| 123 |
+
|
| 124 |
+
You may add Your own copyright statement to Your modifications and
|
| 125 |
+
may provide additional or different license terms and conditions
|
| 126 |
+
for use, reproduction, or distribution of Your modifications, or
|
| 127 |
+
for any such Derivative Works as a whole, provided Your use,
|
| 128 |
+
reproduction, and distribution of the Work otherwise complies with
|
| 129 |
+
the conditions stated in this License.
|
| 130 |
+
|
| 131 |
+
5. Submission of Contributions. Unless You explicitly state otherwise,
|
| 132 |
+
any Contribution intentionally submitted for inclusion in the Work
|
| 133 |
+
by You to the Licensor shall be under the terms and conditions of
|
| 134 |
+
this License, without any additional terms or conditions.
|
| 135 |
+
Notwithstanding the above, nothing herein shall supersede or modify
|
| 136 |
+
the terms of any separate license agreement you may have executed
|
| 137 |
+
with Licensor regarding such Contributions.
|
| 138 |
+
|
| 139 |
+
6. Trademarks. This License does not grant permission to use the trade
|
| 140 |
+
names, trademarks, service marks, or product names of the Licensor,
|
| 141 |
+
except as required for reasonable and customary use in describing the
|
| 142 |
+
origin of the Work and reproducing the content of the NOTICE file.
|
| 143 |
+
|
| 144 |
+
7. Disclaimer of Warranty. Unless required by applicable law or
|
| 145 |
+
agreed to in writing, Licensor provides the Work (and each
|
| 146 |
+
Contributor provides its Contributions) on an "AS IS" BASIS,
|
| 147 |
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
| 148 |
+
implied, including, without limitation, any warranties or conditions
|
| 149 |
+
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
| 150 |
+
PARTICULAR PURPOSE. You are solely responsible for determining the
|
| 151 |
+
appropriateness of using or redistributing the Work and assume any
|
| 152 |
+
risks associated with Your exercise of permissions under this License.
|
| 153 |
+
|
| 154 |
+
8. Limitation of Liability. In no event and under no legal theory,
|
| 155 |
+
whether in tort (including negligence), contract, or otherwise,
|
| 156 |
+
unless required by applicable law (such as deliberate and grossly
|
| 157 |
+
negligent acts) or agreed to in writing, shall any Contributor be
|
| 158 |
+
liable to You for damages, including any direct, indirect, special,
|
| 159 |
+
incidental, or consequential damages of any character arising as a
|
| 160 |
+
result of this License or out of the use or inability to use the
|
| 161 |
+
Work (including but not limited to damages for loss of goodwill,
|
| 162 |
+
work stoppage, computer failure or malfunction, or any and all
|
| 163 |
+
other commercial damages or losses), even if such Contributor
|
| 164 |
+
has been advised of the possibility of such damages.
|
| 165 |
+
|
| 166 |
+
9. Accepting Warranty or Additional Liability. While redistributing
|
| 167 |
+
the Work or Derivative Works thereof, You may choose to offer,
|
| 168 |
+
and charge a fee for, acceptance of support, warranty, indemnity,
|
| 169 |
+
or other liability obligations and/or rights consistent with this
|
| 170 |
+
License. However, in accepting such obligations, You may act only
|
| 171 |
+
on Your own behalf and on Your sole responsibility, not on behalf
|
| 172 |
+
of any other Contributor, and only if You agree to indemnify,
|
| 173 |
+
defend, and hold each Contributor harmless for any liability
|
| 174 |
+
incurred by, or claims asserted against, such Contributor by reason
|
| 175 |
+
of your accepting any such warranty or additional liability.
|
| 176 |
+
|
| 177 |
+
END OF TERMS AND CONDITIONS
|
| 178 |
+
|
| 179 |
+
APPENDIX: How to apply the Apache License to your work.
|
| 180 |
+
|
| 181 |
+
To apply the Apache License to your work, attach the following
|
| 182 |
+
boilerplate notice, with the fields enclosed by brackets "[]"
|
| 183 |
+
replaced with your own identifying information. (Don't include
|
| 184 |
+
the brackets!) The text should be enclosed in the appropriate
|
| 185 |
+
comment syntax for the file format. We also recommend that a
|
| 186 |
+
file or class name and description of purpose be included on the
|
| 187 |
+
same "printed page" as the copyright notice for easier
|
| 188 |
+
identification within third-party archives.
|
| 189 |
+
|
| 190 |
+
Copyright 2026 Jared Palmer
|
| 191 |
+
|
| 192 |
+
Licensed under the Apache License, Version 2.0 (the "License");
|
| 193 |
+
you may not use this file except in compliance with the License.
|
| 194 |
+
You may obtain a copy of the License at
|
| 195 |
+
|
| 196 |
+
http://www.apache.org/licenses/LICENSE-2.0
|
| 197 |
+
|
| 198 |
+
Unless required by applicable law or agreed to in writing, software
|
| 199 |
+
distributed under the License is distributed on an "AS IS" BASIS,
|
| 200 |
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 201 |
+
See the License for the specific language governing permissions and
|
| 202 |
+
limitations under the License.
|
runtime/README.md
ADDED
|
@@ -0,0 +1,466 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Kev
|
| 2 |
+
|
| 3 |
+
Small Jev-like decision models you can train and run yourself.
|
| 4 |
+
|
| 5 |
+
<p>
|
| 6 |
+
<a href="https://github.com/jaredpalmer/kev/actions/workflows/ci.yml"><img alt="CI" src="https://img.shields.io/github/actions/workflow/status/jaredpalmer/kev/ci.yml?style=for-the-badge&labelColor=000000" height="28"></a>
|
| 7 |
+
<a href="https://huggingface.co/collections/jaredpalmer/kev-6aad9d0ea49f2589665e07cd"><img alt="Weights: Kev-0.8B · 4B · 9B · 27B" src="https://img.shields.io/badge/WEIGHTS-0.8B%20%C2%B7%204B%20%C2%B7%209B%20%C2%B7%2027B-0a0a0a.svg?style=for-the-badge&labelColor=000000" height="28"></a>
|
| 8 |
+
<a href="https://huggingface.co/spaces/jaredpalmer/kev"><img alt="Demo on Hugging Face Spaces" src="https://img.shields.io/badge/DEMO-HF%20Spaces-0a0a0a.svg?style=for-the-badge&labelColor=000000" height="28"></a>
|
| 9 |
+
<a href="https://huggingface.co/datasets/jaredpalmer/kev-suites"><img alt="Frozen eval suites" src="https://img.shields.io/badge/EVAL%20SUITES-frozen-0a0a0a.svg?style=for-the-badge&labelColor=000000" height="28"></a>
|
| 10 |
+
<a href="LICENSE"><img alt="License: Apache-2.0" src="https://img.shields.io/badge/license-Apache--2.0-0a0a0a.svg?style=for-the-badge&labelColor=000000" height="28"></a>
|
| 11 |
+
</p>
|
| 12 |
+
|
| 13 |
+
Kev is a family of small decision models built on Qwen3.5 and Qwen3.8 and based on the architecture described in [Jev's Architecture Unmasked](https://archerhume.com/posts/jevs-architecture-unmasked). You can use the pretrained weights or train your own. The API matches TypeSafe's [System One](https://docs.typesafe.ai/api), so you can point their Python SDK at your local server.
|
| 14 |
+
|
| 15 |
+
## Highlights
|
| 16 |
+
|
| 17 |
+
- Yes/no (`noul`), multiple-choice (`choice`) and rating (`score`) questions in one request. The questions share the text but can't read each other.
|
| 18 |
+
- Calibrated probabilities by default: each checkpoint ships with a fitted temperature.
|
| 19 |
+
- Drop-in for Jev: the TypeSafe Python SDK works against a Kev server unchanged.
|
| 20 |
+
- Four sizes, versioned together as Kev 1.0: from a 0.8B that runs on a laptop to a 27B for a single data-centre GPU.
|
| 21 |
+
- Documents of up to 65,536 tokens, on CUDA and on Apple Silicon through MLX. Each model card says how long a document can get before accuracy drops.
|
| 22 |
+
- Fine-tune on your own labelled examples. A coding-agent skill runs the whole loop on Modal, from finding your questions to serving the result.
|
| 23 |
+
- Deploy your own HTTPS endpoint with one command. It scales to zero when idle.
|
| 24 |
+
- Try it in the browser first: [huggingface.co/spaces/jaredpalmer/kev](https://huggingface.co/spaces/jaredpalmer/kev).
|
| 25 |
+
|
| 26 |
+
## Models
|
| 27 |
+
|
| 28 |
+
Start with Kev-4B. Move to Kev-9B if you have a bigger GPU, or to Kev-27B if you have an 80 GB GPU and want the most accurate Kev. Use Kev-0.8B when size matters more than accuracy.
|
| 29 |
+
|
| 30 |
+
| Model | Base (license) | Runs on: CUDA | Runs on: Mac (MLX) | Validated context | Held-out datasets: index | Card |
|
| 31 |
+
|---|---|---|---|---|---|---|
|
| 32 |
+
| [Kev-0.8B](https://huggingface.co/jaredpalmer/kev-0.8b) | Qwen3.5-0.8B-Base (Apache-2.0) | L4, any 4 GB GPU | Any Apple Silicon Mac; measured to 65k tokens | 8,192 | 23.3 | [Details](docs/model-cards/kev-0.8b.md) |
|
| 33 |
+
| [Kev-4B](https://huggingface.co/jaredpalmer/kev-4b) | Qwen3.5-4B-Base (Apache-2.0) | L40S, H100 | 32 GB Mac; measured to 65k tokens | 8,192 | 38.0 | [Details](docs/model-cards/kev-4b.md) |
|
| 34 |
+
| [Kev-9B](https://huggingface.co/jaredpalmer/kev-9b) | Qwen3.5-9B-Base (Apache-2.0) | L40S, H100 | 32 GB Mac or larger (expected, not measured) | 8,192 | 41.0 | [Details](docs/model-cards/kev-9b.md) |
|
| 35 |
+
| [Kev-27B](https://huggingface.co/jaredpalmer/kev-27b) | Qwen3.8-27B, post-trained (Apache-2.0) | B200, H200, H100 80 GB | 96–128 GB Mac (expected, not measured) | 65,536 | **52.3** | [Details](docs/model-cards/kev-27b.md) |
|
| 36 |
+
| Jev | Hosted | TypeSafe's API | – | – | 54.0 | – |
|
| 37 |
+
|
| 38 |
+
"Held-out datasets" is the chance-corrected index of the community Decision Index, scored on the test split of `breadth-v1`: 14 public datasets in five areas that no Kev trained on. "Validated context" is the longest document, in tokens, for which accuracy on real contracts (CUAD) stays within 3 points of the same model's accuracy at 8k tokens, at the 95 % lower bound; each model card has the measurement by length.
|
| 39 |
+
|
| 40 |
+
| Model | Accuracy: New Sources | Accuracy: Trained Sources | Brier: New Sources |
|
| 41 |
+
|---|---|---|---|
|
| 42 |
+
| Kev-0.8B | 0.648 / 0.697 | 0.827 / 0.838 | 0.481 / 0.416 |
|
| 43 |
+
| Kev-4B | 0.817 / 0.838 | 0.873 / 0.865 | 0.269 / 0.242 |
|
| 44 |
+
| Kev-9B | 0.820 / 0.852 | 0.874 / 0.873 | 0.289 / 0.217 |
|
| 45 |
+
| Kev-27B | **0.851 / 0.889** | 0.865 / 0.866 | **0.225 / 0.156** |
|
| 46 |
+
| Jev | 0.857 / – | 0.845 / – | 0.211 / – |
|
| 47 |
+
|
| 48 |
+
Each cell is **development / test**. "New sources" means datasets and policy rules Kev never saw during training. It is the closest thing here to your own questions. "Trained sources" means held-out examples from the datasets Kev was trained on. We pick checkpoints using the development sets and read each test set only once per released model. Jev has only been run on the development sets of these two suites. Brier scores the whole probability distribution, not just the top answer; lower is better.
|
| 49 |
+
|
| 50 |
+
On new sources Kev-27B is within a point of Jev (0.851 vs 0.857), and Kev-4B and Kev-9B are within four points. We don't know what Jev was trained on, so this isn't a controlled comparison of the two architectures. [What to Expect](#what-to-expect) says where Kev is as good as Jev and where it isn't.
|
| 51 |
+
|
| 52 |
+
Kev-0.8B, 4B and 9B start from Qwen base models and share one training recipe: a small adapter on a frozen base. Kev-27B starts from Qwen's post-trained release, and we don't know what that was trained on; every one of its weights is fine-tuned, so it ships as 51 GB of full weights rather than an adapter. Each model card has the full recipe, all results, and the earlier versions kept as Hub tags.
|
| 53 |
+
|
| 54 |
+
## Kev 1.0
|
| 55 |
+
|
| 56 |
+
The four models above are released together as Kev 1.0. Each Hub repo has a `v1.0` tag, so `--run jaredpalmer/kev-4b@v1.0` always loads the same weights, and the GitHub release [`kev-1.0`](https://github.com/jaredpalmer/kev/releases/tag/kev-1.0) has the 0.8B, 4B and 9B checkpoints with SHA-256 checksums. Kev-27B's 51 GB of weights are too large for a release asset and are on the Hub only.
|
| 57 |
+
|
| 58 |
+
| Model | Hub revision of the weights | Temperature | Trained on states up to |
|
| 59 |
+
|---|---|---|---|
|
| 60 |
+
| Kev-0.8B | `9a45d25e` | 2.35 | 7,552 tokens |
|
| 61 |
+
| Kev-4B | `139fdd94` | 2.41 | 7,552 tokens |
|
| 62 |
+
| Kev-9B | `b5d8c18e` (v2) | 2.19 | 7,552 tokens |
|
| 63 |
+
| Kev-27B | `28be62e9` (v2, full weights) | 1.32 | 32,768 tokens |
|
| 64 |
+
|
| 65 |
+
Kev 1.0 trains nothing new. It fixes the checkpoints, cards, evaluation suites and serving code that the next generation of Kev will be compared against. The [release notes](docs/releases/kev-1.0.md) list what changed since the previous family release and what is known not to work well.
|
| 66 |
+
|
| 67 |
+
## Quick Start
|
| 68 |
+
|
| 69 |
+
### Try It in the Browser
|
| 70 |
+
|
| 71 |
+
The [Hugging Face Space](https://huggingface.co/spaces/jaredpalmer/kev) runs Kev-4B and Kev-0.8B, with nothing to install.
|
| 72 |
+
|
| 73 |
+
### Run It Locally
|
| 74 |
+
|
| 75 |
+
You'll need Python 3.12 or 3.13 and [uv](https://docs.astral.sh/uv/). The repo's `.python-version` makes `uv sync` use 3.13; torch has no wheels for 3.14 yet.
|
| 76 |
+
|
| 77 |
+
```bash
|
| 78 |
+
git clone https://github.com/jaredpalmer/kev.git && cd kev
|
| 79 |
+
uv sync --extra serve
|
| 80 |
+
uv run --extra serve python -m kev.serve --run jaredpalmer/kev-4b --port 8009
|
| 81 |
+
```
|
| 82 |
+
|
| 83 |
+
This starts Kev-4B on your machine: CUDA or ROCm if you have a GPU, MLX on Apple Silicon. The first run downloads the adapter and the base model. `--run` also accepts a local checkpoint directory or a Hub revision like `jaredpalmer/kev-4b@qwen3`.
|
| 84 |
+
|
| 85 |
+
In another terminal, send it a ticket:
|
| 86 |
+
|
| 87 |
+
```bash
|
| 88 |
+
curl -s localhost:8009/v1/systemone -H 'content-type: application/json' -d '{
|
| 89 |
+
"state": "Shoes arrived two weeks late and in the wrong size. Also I see two charges on my card.",
|
| 90 |
+
"model": "kev-latest",
|
| 91 |
+
"questions": {
|
| 92 |
+
"department": {"type": "choice", "instructions": "Which team should handle this?",
|
| 93 |
+
"criteria": {"returns": "Exchanges, refunds, wrong or damaged items",
|
| 94 |
+
"shipping": "Delivery status, delays, lost packages",
|
| 95 |
+
"billing": "Charges, invoices, payment problems"}},
|
| 96 |
+
"escalate": {"type": "noul", "instructions": "Does this need urgent human attention?"},
|
| 97 |
+
"frustration": {"type": "score", "instructions": "How frustrated is the customer?",
|
| 98 |
+
"criteria": ["Calm", "Frustrated", "Very angry"]}
|
| 99 |
+
}}'
|
| 100 |
+
```
|
| 101 |
+
|
| 102 |
+
Example response from Kev-4B, running in bf16 on an Apple M5:
|
| 103 |
+
|
| 104 |
+
```json
|
| 105 |
+
{
|
| 106 |
+
"model": "kev-latest",
|
| 107 |
+
"answers": {
|
| 108 |
+
"department": { "type": "choice", "choice": "returns", "confidence": 0.21,
|
| 109 |
+
"probabilities": { "returns": 0.47, "shipping": 0.28, "billing": 0.25 } },
|
| 110 |
+
"escalate": { "type": "noul", "noul": 0.93 },
|
| 111 |
+
"frustration": { "type": "score", "score": 1.44, "confidence": 0.34,
|
| 112 |
+
"legend": { "0": "Calm", "1": "Frustrated", "2": "Very angry" },
|
| 113 |
+
"probabilities": { "0": 0.00, "1": 0.56, "2": 0.44 } }
|
| 114 |
+
},
|
| 115 |
+
"usage": { "input_tokens": 101, "output_tokens": 161 },
|
| 116 |
+
"latency_ms": 495
|
| 117 |
+
}
|
| 118 |
+
```
|
| 119 |
+
|
| 120 |
+
The ticket mentions a return, a late delivery and a billing problem, and the department probabilities say so. That's why Kev returns probabilities instead of a single label: your code can route the confident cases and send the rest to a person.
|
| 121 |
+
|
| 122 |
+
### Use It From Python
|
| 123 |
+
|
| 124 |
+
If you already call Jev, point your client at Kev and keep the rest of your code. The TypeSafe SDK is included in `uv sync --extra serve`:
|
| 125 |
+
|
| 126 |
+
```python
|
| 127 |
+
from typesafe_sdk import Choice, Noul, Score, TypeSafeClient
|
| 128 |
+
|
| 129 |
+
client = TypeSafeClient(
|
| 130 |
+
api_key="local",
|
| 131 |
+
base_url="http://127.0.0.1:8009",
|
| 132 |
+
model="kev-latest",
|
| 133 |
+
)
|
| 134 |
+
response = client.system_one(
|
| 135 |
+
state="I was charged twice. Please fix this ASAP.",
|
| 136 |
+
questions={
|
| 137 |
+
"billing": Noul(instructions="Is this ticket about billing?"),
|
| 138 |
+
"tone": Choice(
|
| 139 |
+
instructions="What is the customer's tone?",
|
| 140 |
+
criteria={"calm": None, "frustrated": None, "angry": None},
|
| 141 |
+
),
|
| 142 |
+
"urgency": Score(
|
| 143 |
+
instructions="How urgent is this ticket?",
|
| 144 |
+
criteria=["can wait", "this week", "today"],
|
| 145 |
+
),
|
| 146 |
+
},
|
| 147 |
+
)
|
| 148 |
+
print(response.nouls["billing"].noul)
|
| 149 |
+
print(response.choices["tone"].choice)
|
| 150 |
+
print(response.scores["urgency"].score)
|
| 151 |
+
```
|
| 152 |
+
|
| 153 |
+
## Fine-Tune on Your Own Data
|
| 154 |
+
|
| 155 |
+
The released models were trained on public datasets and generated policy examples. If your questions look different, like your own routing categories, your own escalation rules or another language, a short fine-tune usually helps more than any prompt change. It also fits the temperature to your data, so the confidence you set thresholds on is measured on your own labels.
|
| 156 |
+
|
| 157 |
+
What to expect: on an example support workload (three questions, 1,050 generated records, 15 minutes on an H100), fine-tuning took Kev-4B from 67.7% to 73.6% accuracy, and from automating 34% of decisions at a 5% error budget to 48% ([details](skills/kev-finetune/README.md#why-fine-tune-at-all)). On real data, one epoch on 5,219 labelled consumer-finance complaints took Kev-4B from 0.804 to 0.904 accuracy on complaints it had never seen. Gains like these are in distribution: they tell you how well Kev learns your task, not how it does on everything else. Size your dataset first. With 400 records, the gain on the example workload was inside the noise.
|
| 158 |
+
|
| 159 |
+
### With a Coding Agent
|
| 160 |
+
|
| 161 |
+
```bash
|
| 162 |
+
npx skills add jaredpalmer/kev@kev-finetune
|
| 163 |
+
```
|
| 164 |
+
|
| 165 |
+
Then ask your agent to "fine-tune Kev on my support tickets". The [`kev-finetune` skill](skills/kev-finetune/) interviews you, finds the questions your code already asks Jev or TypeSafe, converts the labels you have or generates enough with any LLM to measure a gain, fine-tunes from a released checkpoint on Modal, fits the temperature on a held-out slice, scores the result against the untouched model, deploys an endpoint, and tears everything down at the end. You don't need a local GPU or a clone of this repo. A Kev-4B training run costs about $1 on an H100.
|
| 166 |
+
|
| 167 |
+
### By Hand
|
| 168 |
+
|
| 169 |
+
The skill's [README](skills/kev-finetune/README.md) is the same recipe for people: six short standard-library scripts and one Modal app. To train from this repo instead, put your examples in a JSONL file, one request per line. It's the same shape as an API request, plus a `label` on every question:
|
| 170 |
+
|
| 171 |
+
```jsonl
|
| 172 |
+
{"state": {"subject": "Charged twice", "body": "I see two charges for order #4411. Please refund one."},
|
| 173 |
+
"questions": {
|
| 174 |
+
"team": {"type": "choice", "instructions": "Which team should handle this ticket?",
|
| 175 |
+
"criteria": {"billing": "Payments and refunds", "shipping": "Delivery problems", "access": "Login and account access"}, "label": "billing"},
|
| 176 |
+
"angry": {"type": "noul", "instructions": "Is the customer angry?", "label": false},
|
| 177 |
+
"priority": {"type": "score", "instructions": "How urgent is this ticket?", "criteria": ["low", "normal", "high"], "label": 1}}}
|
| 178 |
+
```
|
| 179 |
+
|
| 180 |
+
For `choice` the label is the option name, for `noul` it's `true` or `false`, and for `score` it's the level's position starting at 0. Keep 10–20% of the file aside for evaluation.
|
| 181 |
+
|
| 182 |
+
Then start from a released checkpoint with `--init_from`:
|
| 183 |
+
|
| 184 |
+
```bash
|
| 185 |
+
uv run python -m kev.train --data train.jsonl --base Qwen/Qwen3.5-4B-Base --init_from jaredpalmer/kev-4b \
|
| 186 |
+
--epochs 2 --lr 2e-5 --batch 1 --accum 8 --dtype bf16 --checkpointing 1 --device cuda --out runs/mine
|
| 187 |
+
|
| 188 |
+
uv run python -m kev.benchmark --run runs/mine --data heldout.jsonl --out runs/mine-eval
|
| 189 |
+
uv run --extra serve python -m kev.serve --run runs/mine --port 8009
|
| 190 |
+
```
|
| 191 |
+
|
| 192 |
+
`--init_from` loads the adapter and pointer head from the released model before training, so you keep what Kev already knows and add your domain on top. Starting from the base model instead throws that away: in one user's test on 836 support-tool decisions, a fine-tune from the base scored 0.33 on Kev's own evaluation set, against 0.84 for the released model; the same data with `--init_from` kept 0.83 there and reached 0.88 on the new domain. Use a smaller learning rate than the from-scratch recipe (`2e-5` is a good start), and pick `--base` to match the checkpoint you start from; the trainer checks that the base, revision, LoRA rank and head size agree before it loads anything.
|
| 193 |
+
|
| 194 |
+
`--batch 1 --accum 8` in bf16 fits the 0.8B model on a 4 GB GPU. The benchmark reports accuracy, Brier score and calibration per question type, so you can see which of your questions the fine-tune helped. The checkpoint you started from is recorded in `runs/mine/training_config.json`. On a Mac, run one training job at a time; two jobs on the same Apple GPU are much slower.
|
| 195 |
+
|
| 196 |
+
## Deploy Your Own Endpoint
|
| 197 |
+
|
| 198 |
+
To get an HTTPS endpoint instead of a local server, you don't need this repo, just a [Modal](https://modal.com) account:
|
| 199 |
+
|
| 200 |
+
```bash
|
| 201 |
+
pip install modal && modal setup
|
| 202 |
+
curl -LO https://raw.githubusercontent.com/jaredpalmer/kev/main/skills/kev-deploy/scripts/kev_serve.py
|
| 203 |
+
KEV_API_KEY=$(openssl rand -hex 24) modal deploy kev_serve.py
|
| 204 |
+
```
|
| 205 |
+
|
| 206 |
+
That serves Kev-4B on an L40S at `https://<your-workspace>--kev-api.modal.run`, with the same API as above behind `Authorization: Bearer <key>`. It scales to zero when idle, so an unused endpoint costs nothing. The first request after idle waits about 35 seconds for a container to start. `KEV_MODEL=jaredpalmer/kev-9b` serves another model on the GPU that suits it; Kev-27B goes to a B200, falling back to an H200 or H100. If you use a coding agent, `npx skills add jaredpalmer/kev@kev-deploy` does the same and wires the URL into your code. [skills/kev-deploy](skills/kev-deploy/) has the GPU and cost table.
|
| 207 |
+
|
| 208 |
+
A model you fine-tuned with the `kev-finetune` skill deploys the same way from its own Modal app (`KEV_SERVE_SECRET=kev-serve-key KEV_SERVE_RUN=<run> modal deploy scripts/kev_modal.py`; see [its deploy guide](skills/kev-finetune/references/deploy.md)). To host Kev on your own machines instead, run `kev.serve` from [Run It Locally](#run-it-locally) on a GPU box with `--host 0.0.0.0` and put it behind your own proxy; [Serving Performance](#serving-performance) says which GPU to pick.
|
| 209 |
+
|
| 210 |
+
## What to Expect
|
| 211 |
+
|
| 212 |
+
**Accuracy.** Kev-27B is within three points of Jev, or ahead of it, on 9 of the 11 new-source categories in the chart below. Kev-4B and Kev-9B are about as close on classification-shaped sources like routing, entailment and science questions. Knowledge questions depend mostly on the base model: on MMLU Kev-9B scores 0.73 and Kev-27B matches Jev at 0.90, but on the harder MMLU-Pro Kev-27B scores 0.675 against Jev's 0.840. The smaller models also trail on day-precision date arithmetic.
|
| 213 |
+
|
| 214 |
+

|
| 215 |
+
|
| 216 |
+
**Confidence.** Each checkpoint ships with a fitted temperature, so its probabilities are calibrated by default. As served, Kev-9B puts at least 0.9 probability on a wrong answer for 2.4% of new-source questions, against Jev's 3.7%. Jev still ranks its answers better: at a 5% error budget, Kev-4B, 9B and 27B can automate 0.52–0.69 of new-source decisions, Kev-0.8B 0.14 and Jev 0.70. Check a threshold on your own data before you rely on it.
|
| 217 |
+
|
| 218 |
+
**Speed.** Kev-4B answers six questions about a new short text in 18.1 ms of model time on an H100 and 41.5 ms on an L40S, and a container serves around 101 requests per second on an H100. On an Apple M5, Kev-4B takes 721 ms for five questions, or 136 ms when the text repeats and comes from the cache. [Serving Performance](#serving-performance) has every GPU and batch size.
|
| 219 |
+
|
| 220 |
+
**Length.** Kev-0.8B, 4B and 9B trained mostly on states of up to 384 tokens, with longer ones in their document and skill fine-tunes (up to 7,552 tokens), Kev-27B on states of up to 32,768. The server accepts states of up to 65,536 tokens, and 8,192 more for each question, and refuses a longer one with a 422 instead of cutting it. How far past its training length each model stays accurate is the "Validated context" column in [Models](#models). For Kev-0.8B, 4B and 9B that is 8,192 tokens: at 16k the measurement on real contracts can no longer rule out a drop of more than 3 points, and at 32k all three are measurably less accurate than at 8k. Kev-27B holds to the 65,536-token limit. On real contracts of up to 64k tokens (CUAD) Kev-27B scores 0.874, and its confidence there is less reliable than on short text; its [model card](docs/model-cards/kev-27b.md) has the numbers by length.
|
| 221 |
+
|
| 222 |
+
## Playground
|
| 223 |
+
|
| 224 |
+
With the server running, open another terminal. You'll need Node 20.9+:
|
| 225 |
+
|
| 226 |
+
```bash
|
| 227 |
+
cd playground
|
| 228 |
+
npm install
|
| 229 |
+
npm run dev -- -p 3001
|
| 230 |
+
```
|
| 231 |
+
|
| 232 |
+
Open [localhost:3001](http://localhost:3001), load a preset, and edit the text and questions. Press `⌘↵` to run it. "Packed vs separate" compares asking all questions at once with asking them one at a time. "Permute" runs a Choice question with six option orders. There are also presets for testing question isolation and fake delimiter tokens.
|
| 233 |
+
|
| 234 |
+

|
| 235 |
+
|
| 236 |
+
There's a [chess demo](http://localhost:3001/chess), too. The board is the input, legal moves are Choice options, and a Score question rates the position. You can play against Kev or let it play itself. Games are saved in `localStorage`.
|
| 237 |
+
|
| 238 |
+
## API
|
| 239 |
+
|
| 240 |
+
### `POST /v1/systemone`
|
| 241 |
+
|
| 242 |
+
`state` is the text to evaluate. Each question has instructions and, where needed, a set of answers to choose from.
|
| 243 |
+
|
| 244 |
+
```jsonc
|
| 245 |
+
{
|
| 246 |
+
"state": "…", // string | object | array — the content to evaluate
|
| 247 |
+
"model": "kev-latest",
|
| 248 |
+
"questions": {
|
| 249 |
+
"<id>": { // you choose the id; the model never sees it
|
| 250 |
+
"type": "noul" | "choice" | "score",
|
| 251 |
+
"instructions": "…", // string | object | array, optional
|
| 252 |
+
"criteria": … // noul: {true?, false?} choice: {option: description|null} score: [level, …]
|
| 253 |
+
}
|
| 254 |
+
}
|
| 255 |
+
}
|
| 256 |
+
```
|
| 257 |
+
|
| 258 |
+
| Type | Criteria | Answer |
|
| 259 |
+
|---|---|---|
|
| 260 |
+
| `noul` | Optional descriptions for `true` and `false` | `noul`: probability of yes |
|
| 261 |
+
| `choice` | 1–255 option names, each with a description or `null` | `choice`: most likely option; `probabilities` and `confidence` |
|
| 262 |
+
| `score` | 1–255 descriptions, ordered from lowest to highest | `score`: mean level index, starting at 0; `legend`, `probabilities`, and `confidence` |
|
| 263 |
+
|
| 264 |
+
For Choice with `K > 1` options, confidence is `(p_max − 1/K) / (1 − 1/K)`. A single option has confidence 1. Score confidence is `max(0, 1 − E|level − mode| / D)`: `mode` is the most likely level and `D` is the mean distance of a uniform distribution over the levels from its middle (2/3 for three levels), so all probability on one level gives 1 and a uniform or wider spread gives 0. Both formulas are the ones in TypeSafe's reference adapter ([`system-one-adapter`](https://github.com/typesafe-ai/system-one-adapter-python) 0.2.1). Neither field is a measured accuracy rate.
|
| 265 |
+
|
| 266 |
+
Objects and arrays are converted to labeled text. Delimiter-like strings in user input are escaped before tokenization. Invalid requests return `422`, and so does a state longer than 65,536 tokens: the server never drops part of a document silently, and the error gives the state's token count and the limit. `usage.output_tokens` counts tokens in the serialized answers, not generated tokens.
|
| 267 |
+
|
| 268 |
+
| Method | Path | Purpose |
|
| 269 |
+
|---|---|---|
|
| 270 |
+
| `GET` | `/v1/models` | Model cards (`name`, `description`, `release_date`) plus the loaded checkpoint's details |
|
| 271 |
+
| `POST` | `/v1/systemone/permute` | Run one Choice question with different option orders (`n_perm` 1 to 64, default 6) |
|
| 272 |
+
| `POST` | `/v1/systemone/separate` | Run each question in its own forward pass |
|
| 273 |
+
|
| 274 |
+
A request may carry any number of questions. The server runs them a token budget at a time (one maximal row of 16,384 tokens per forward pass, counting the cached document once per question in that pass), so memory does not grow with the question count and the answers do not depend on the split. Every response carries an `x-typesafe-request-id` header. The server binds to `127.0.0.1` (`--host 0.0.0.0` to accept other machines) and is open by default; set `KEV_API_KEY` to require `Authorization: Bearer <key>` on `/v1/*`, as the TypeSafe clients always send it.
|
| 275 |
+
|
| 276 |
+
| Variable | Effect |
|
| 277 |
+
|---|---|
|
| 278 |
+
| `KEV_TEMPERATURE=1.0` | Return raw probabilities instead of the calibrated ones |
|
| 279 |
+
| `KEV_DATE_FACTS=1` | Append the number of days between any two dates in the state (see [Benchmarks](#benchmarks)) |
|
| 280 |
+
| `KEV_TRUNCATE_STATES=1` | Read the first 65,536 tokens of a longer state instead of refusing it; every response then has `truncated` and `usage.state_tokens` / `state_tokens_used` |
|
| 281 |
+
| `KEV_DTYPE=fp32` | Serve the exact fp32 path the evaluations use (bf16 is the default on GPUs) |
|
| 282 |
+
| `KEV_API_KEY` | Require a bearer key |
|
| 283 |
+
|
| 284 |
+
## How It Works
|
| 285 |
+
|
| 286 |
+
Each checkpoint is a rank-16 LoRA adapter and a small pointer head on a Qwen base model. On an attention-only base (Qwen3), the state and questions go into one token sequence:
|
| 287 |
+
|
| 288 |
+
```text
|
| 289 |
+
<state> …state…
|
| 290 |
+
<q> instructions <opt> option 1 </opt> <opt> option 2 </opt> … <decide>
|
| 291 |
+
<q> instructions <opt> option 1 </opt> <opt> option 2 </opt> … <decide>
|
| 292 |
+
```
|
| 293 |
+
|
| 294 |
+
The attention mask lets a token read the state and its own question, but not other questions or future tokens. Each question's position IDs restart just after the state. This lets the model process the state once and answer each question independently.
|
| 295 |
+
|
| 296 |
+
Qwen3.5 and Qwen3.8 mix attention layers with Gated DeltaNet layers, which are recurrent and ignore attention masks. For those models, which is every current Kev, each question runs as its own row: the state followed by that question, with the same positions as above. The rows are independent, so isolation is exact, and the server and `DecisionModel.probs()` compute the state once and reuse its cache for every row. `forward()`, which `kev.benchmark` scores and every published number comes from, keeps the plain rows and runs the state once per question; the two agree to fp32 rounding. On attention-only models the rows and the mask above give identical probabilities (`tests/test_model.py`).
|
| 297 |
+
|
| 298 |
+
Kev-27B uses the same design on `Qwen/Qwen3.8-27B`, with two differences. Its base is Qwen's post-trained release rather than a `-Base` checkpoint, and we don't know what it was post-trained on. And every backbone weight is trained, not just an adapter, and kept in bf16, so the checkpoint is the whole model: 51 GB of bf16 weights plus the pointer head. It serves in bf16 only (about 66 GB resident with the serving buffers), which is why it needs an 80 GB card. On Apple Silicon the MLX backend loads those weights as they are, with no merge (see [Serving Performance](#serving-performance)); we expect that to fit a 96–128 GB Mac but haven't measured it. Its served probabilities stay within 0.022 of the evaluation path on an H200 (`runs/serving-27b-r23`).
|
| 299 |
+
|
| 300 |
+
The pointer head scores each option's `</opt>` hidden state against the question's `<decide>` hidden state. A softmax turns those scores into probabilities. Because `<decide>` comes last, it can attend to the full option list.
|
| 301 |
+
|
| 302 |
+
Training uses cross-entropy on the correct answer. The adapter and head are trained together; the rest of the base weights stay fixed (Kev-27B trains them all). Training examples and API requests use the same text format. No Jev outputs were used for training.
|
| 303 |
+
|
| 304 |
+
Asking questions together or separately produces probabilities within 4e-6 in the fp32 tests. This does **not** mean option order is irrelevant: options within a question can still affect one another. See [the model code](kev/model.py) and [parity tests](tests/test_model.py).
|
| 305 |
+
|
| 306 |
+
## Training
|
| 307 |
+
|
| 308 |
+
The released models share one base training set, `decision-v7`: 10,000 examples from ten public datasets, 896 generated policy examples, and 1,680 examples from 60 generated rule structures. Kev-0.8B, 4B and 9B train on it for two epochs with LoRA rank 16 and cross-entropy. The learning rate is `1e-4` for 0.8B and `5e-5` for 4B and 9B. On these hybrid bases the adapter covers the attention, MLP and DeltaNet projections; `kev.train` picks the right targets from the model config.
|
| 309 |
+
|
| 310 |
+
Kev-0.8B, 4B and 9B then get short follow-up fine-tunes from their released checkpoints, through the same `--init_from` path you'd use for your own data: generated cases that state day counts or have the deciding evidence removed (all three), then real documents and generated skill data (all three; Kev-9B since v2, 2026-09-30). Kev-27B is trained differently. Every weight of the base is fine-tuned for one epoch on eight H200s (`--full_ft 1`, learning rate `2e-6`) on a 145,840-record corpus: Kev's own data, the document, skill and developer-tooling suites, public datasets, licensed task families and generated long-document, tool-routing, agent-log and guardrail records, with states of up to 32,768 tokens. The result is then averaged with the earlier adapter-trained Kev-27B, 0.85 to 0.15. The model cards list every stage with its data and cost.
|
| 311 |
+
|
| 312 |
+
```bash
|
| 313 |
+
# sanity run, ~1 minute
|
| 314 |
+
uv run python -m kev.train --n_per_source 40 --accum 4 --out runs/smoke
|
| 315 |
+
|
| 316 |
+
# the first stage of Kev-0.8B (~20 min on one H100; the Mac path works but is slow for Qwen3.5 bases)
|
| 317 |
+
uv run python -m kev.train --suite evals/v7/decision-v7 --base Qwen/Qwen3.5-0.8B-Base --base_revision dc7cdfe2ee4154fa7e30f5b51ca41bfa40174e68 \
|
| 318 |
+
--epochs 2 --lr 1e-4 --batch 8 --dtype bf16 --p_none_pair 0.25 --device cuda --out runs/kev-0.8b
|
| 319 |
+
|
| 320 |
+
# the first stage of Kev-4B (one H100 via Modal, ~1 h; see below). Swap in Qwen/Qwen3-4B-Base for the previous generation.
|
| 321 |
+
uv run python -m kev.train --suite evals/v7/decision-v7 --base Qwen/Qwen3.5-4B-Base --base_revision 1001bb4d826a52d1f399e183466143f4da7b741b \
|
| 322 |
+
--epochs 2 --lr 5e-5 --batch 4 --accum 2 --dtype bf16 --checkpointing 1 --p_none_pair 0.25 --device cuda --out runs/kev-4b
|
| 323 |
+
```
|
| 324 |
+
|
| 325 |
+
Use `uv run python -m kev.train --help` for all training options. The released models don't use the optional `--perm_kl` or `--ord_w` losses. [PLAN.md](PLAN.md) records what was tried, what helped, and what didn't.
|
| 326 |
+
|
| 327 |
+
### Modal
|
| 328 |
+
|
| 329 |
+
Each trial gets its own H100. The study keeps running if you disconnect, and you can download the results when it finishes:
|
| 330 |
+
|
| 331 |
+
```bash
|
| 332 |
+
uv run modal token new # once; opens the browser
|
| 333 |
+
KEV_GPU=T4 uv run modal run modal_app.py::smoke # end-to-end check, ~1 minute of GPU
|
| 334 |
+
|
| 335 |
+
uv run modal deploy modal_app.py # once; studies run on the deployed app and survive disconnects
|
| 336 |
+
uv run modal run modal_app.py::study \
|
| 337 |
+
--suite evals/v7/decision-v7 --plan experiments/v7-final.json \
|
| 338 |
+
--name my-study --transfer evals/v4/transfer-v4 --budget 30 --timeout 7200
|
| 339 |
+
uv run modal run modal_app.py::pull --name my-study # results -> runs/my-study, ranked
|
| 340 |
+
```
|
| 341 |
+
|
| 342 |
+
[Study plans](experiments/v7-final.json) list training settings. Each trial saves the settings, code hashes, dataset hashes, and results. Choose models using the development results, not the locked test. After choosing a final candidate, you can read its test results once:
|
| 343 |
+
|
| 344 |
+
```bash
|
| 345 |
+
uv run modal run modal_app.py::locked_test --trial my-study/00-trial-0 --name my-candidate # one read, ever
|
| 346 |
+
```
|
| 347 |
+
|
| 348 |
+
## Benchmarks
|
| 349 |
+
|
| 350 |
+
The evaluation data under `evals/` is frozen: dataset versions and file checksums are recorded in each manifest. Large files are downloaded from [the Hub mirror](https://huggingface.co/datasets/jaredpalmer/kev-suites) and checked against those hashes. Every model in the tables above is scored on the same items. The numbers in this README and the model cards are checked in CI against the committed reports they come from (`docs/claims.json`, `uv run python scripts/verify_claims.py`).
|
| 351 |
+
|
| 352 |
+
| Suite | What it measures |
|
| 353 |
+
|---|---|
|
| 354 |
+
| `decision-v7` | Held-out examples from the ten training datasets, the generated policies and the rule structures ("trained sources") |
|
| 355 |
+
| `transfer-v4` | 764 records from datasets and policy and rule types Kev never trained on: QNLI, SciQ, PAWS, MMLU, Emotion, TweetEval, held-out policies and rules ("new sources") |
|
| 356 |
+
| `transfer-v9` | `transfer-v4` plus 10-way MMLU-Pro, records buried in unrelated text, and "unknowable" records whose deciding evidence was removed |
|
| 357 |
+
|
| 358 |
+
```bash
|
| 359 |
+
uv run python -m kev.benchmark --run jaredpalmer/kev-4b --suite evals/v4/transfer-v4 --out runs/my-eval # new sources
|
| 360 |
+
uv run python -m kev.benchmark --run jaredpalmer/kev-4b --suite evals/v9/transfer-v9 --out runs/my-eval-v9 # + MMLU-Pro, buried states, unknowable items
|
| 361 |
+
uv run python -m kev.benchmark --run jaredpalmer/kev-4b --suite evals/v7/decision-v7 --out runs/my-eval-id # trained sources
|
| 362 |
+
uv run python -m kev.benchmark --remote http://127.0.0.1:8009 --suite evals/v4/transfer-v4 --out runs/my-remote # any System One endpoint, Jev included
|
| 363 |
+
```
|
| 364 |
+
|
| 365 |
+
These commands use development data. Test data requires `--allow-test`. The benchmark reports accuracy, Brier score, calibration error, the share of decisions you could automate at a 5% error budget, option-order changes, and question isolation. On the unknowable records it reports how often the model still answers with at least 0.9 confidence (Kev-9B 0%, Jev 9%). Published accuracy numbers use fp32 evaluation, not the bf16 serving path. `kev.jev` runs the same questions against Jev through Vercel AI Gateway, and `kev.compare` compares two saved runs with paired bootstrap confidence intervals.
|
| 366 |
+
|
| 367 |
+
**Calibration.** Each checkpoint stores a temperature, and the pointer head applies it when the model is loaded. Kev-4B (2.41) and Kev-0.8B (2.35) fitted theirs on their in-distribution development sets; Kev-27B (1.32) and Kev-9B (2.19) fitted theirs on held-out datasets they never trained on. A refit of the two smaller models on those held-out datasets was tested and kept neither: it did not improve Kev-4B and made Kev-0.8B worse calibrated on its document and skill suites (the model cards have the numbers). A temperature never changes which answer wins. On new sources it takes Kev-9B's calibration error from 0.103 to 0.041 and its confident errors (wrong answers with probability ≥ 0.9) from 8.2% to 2.4%, below Jev's 3.7%. The accuracy numbers above are the same either way; the Brier numbers are for the raw probabilities. `scripts/calibrate_checkpoint.py` also reports an out-of-fold estimate, so the in-sample fit can be checked against records it didn't see.
|
| 368 |
+
|
| 369 |
+
**Dates.** Kev can't subtract dates reliably, but it can use a day count it's given. `KEV_DATE_FACTS=1` appends one sentence per pair of dates in the state ("June 26, 2026 is 8 days before July 4, 2026"). On the deadline policy questions this takes Kev-9B from 0.80 to 0.90 (Jev 0.93). None of the tables use it.
|
| 370 |
+
|
| 371 |
+
**Other people's test sets.** `evals/external/` holds test sets from other projects, converted to this format, with their published live Jev results. Some were scored on earlier versions of the Kev weights, which the Kev column names. Three were removed because they cannot serve as a gate, and the model cards keep the numbers their releases were decided on: scienthoon's synthetic support tickets on 2026-09-27 (templated text; one of its three questions depends on a rule the text does not state), and on 2026-09-30 WANLI (`wanli-v1`, `wanli-v2`: a quarter of the pairs are ones WANLI's two annotators labelled differently, with the gold set to one of them) and TypeSafe's public evals (`typesafe-v1`: the gold is the averaged answer of two closed frontier models, and on 89 questions it cannot tell checkpoints apart).
|
| 372 |
+
|
| 373 |
+
| Suite | What it is | Jev | Kev |
|
| 374 |
+
|---|---|---|---|
|
| 375 |
+
| [SemIf](https://github.com/TheoLeeCJ/SemIf) | 144 authored decisions | 0.965 | 0.917 (Kev-9B at `v7-base`) |
|
| 376 |
+
|
| 377 |
+
SemIf's labels hold up, but it is close to saturated: every Kev-27B checkpoint answers 130 of the 144 correctly, so it is a sanity check, not a way to rank models.
|
| 378 |
+
|
| 379 |
+
## Serving Performance
|
| 380 |
+
|
| 381 |
+
Pick the GPU by the model:
|
| 382 |
+
|
| 383 |
+
| Model | GPU ($/h) | 6 questions, short text | 5 questions, 2,200-token text | Requests/s, 64 clients |
|
| 384 |
+
|---|---|---|---|---|
|
| 385 |
+
| Kev-0.8B | L4 (0.80) | 22.7 / 16.1 ms | 108.6 / 32.3 ms | 62.8 |
|
| 386 |
+
| Kev-4B | L40S (1.95) | 41.5 / 27.7 ms | 145.2 / 43.0 ms | 51.4 |
|
| 387 |
+
| Kev-4B | H100 (3.95) | 18.1 / 12.9 ms | 89.4 / 22.5 ms | 100.8 |
|
| 388 |
+
| Kev-9B | L40S (1.95) | 66.4 / 42.7 ms | 235.6 / 57.5 ms | 32.7 |
|
| 389 |
+
| Kev-9B | H100 (3.95) | 24.0 / 16.6 ms | 88.5 / 26.4 ms | 79.5 |
|
| 390 |
+
| Kev-27B | B200 (6.25) | 46.5 / 32.2 ms | 178.0 / 52.1 ms | 44.2 |
|
| 391 |
+
| Kev-27B | H200 (4.54) | 67.2 / 50.0 ms | 274.8 / 73.8 ms | 28.6 |
|
| 392 |
+
| Kev-27B | H100 (3.95) | 75.0 / 52.0 ms | 277.5 / 79.3 ms | 28.9 |
|
| 393 |
+
|
| 394 |
+
Times are model time per request (the `latency_ms` the API returns), median of 20, for a new text / the same text again. The server caches the text, so asking more questions about a document you've already sent only pays for the questions. Requests per second are for 64 concurrent clients sending six questions about a new short text each; the server batches them. Kev-27B's B200 and H100 rows were measured on its previous version, the same architecture served in bf16 (`runs/fused-27b-*`); the H200 row is the current checkpoint (`runs/serving-27b-r23`). Network time is extra: about 65 ms per round trip through a Modal web endpoint in the same region.
|
| 395 |
+
|
| 396 |
+
An L4 is enough for Kev-0.8B but too slow for Kev-4B. The A100 is slower than the L40S here and costs more. Kev-9B needs about 17 GB of GPU memory and Kev-27B 51 GB of weights (about 66 GB with the batching buffers); under load Kev-27B is compute-bound, and a B200, H200 or H100 costs about the same per request. On CUDA, install `flash-linear-attention` for the Qwen3.5 models (`kev_serve.py` and the Modal images already do).
|
| 397 |
+
|
| 398 |
+
On Apple Silicon, `uv sync --extra serve` installs [MLX](https://github.com/ml-explore/mlx-lm) and the server uses it automatically. Five questions about a ~270-token text on an M5 (32 GB):
|
| 399 |
+
|
| 400 |
+
| Model | New text | Same text again |
|
| 401 |
+
|---|---|---|
|
| 402 |
+
| Kev-0.8B | 149 ms | 28 ms |
|
| 403 |
+
| Kev-4B | 721 ms | 136 ms |
|
| 404 |
+
|
| 405 |
+
Long documents are read into the cache 1,024 tokens at a time, so memory stays close to the weights. With a 65,000-token document, Kev-0.8B takes 21.2 s the first time and 202 ms after that, at a 3.8 GB peak, and Kev-4B 84.5 s and 716 ms at 13.0 GB (`runs/mlx-long-states`; the model cards have every length). Kev-9B has not been measured this way yet.
|
| 406 |
+
|
| 407 |
+
Adapter checkpoints are folded into the base as they load, which briefly holds a second copy of the weights. Full-weight checkpoints like Kev-27B load as saved, with nothing merged, so loading needs only the weights. We checked this on Kev-4B written out as full bf16 weights: loading peaked at 8.4 GB for 8.4 GB of weights, against 15.9 GB for the adapter path. Its answers matched the adapter path exactly once both hold the same bf16 values, and stayed within 0.015 of the fp32 path on 60 questions (`runs/mlx-full-4b`). Kev-27B's weights are 51 GB. By the same measurements it needs about 51 GB plus working memory, so a 64 GB Mac is borderline and a 96–128 GB Mac should fit. We haven't run it on a Mac that large yet. Kev-27B's first version, an adapter, did run this way on a 128 GB M5 Max, matching the published accuracy (thanks to Sean Connelly, [#175](https://github.com/jaredpalmer/kev/pull/175)).
|
| 408 |
+
|
| 409 |
+
The server runs in bf16 on GPUs and Macs. Its probabilities differ from the fp32 path the published evaluations use by at most about 0.03 on a GPU and 0.05 on a Mac, and the top answer changes on about one question in 300. Set `KEV_DTYPE=fp32` for the exact path. `/v1/models` reports the backend and precision in use. `uv run modal run modal_app.py::serving --run jaredpalmer/kev-4b --gpu L40S --name <name>` measures a row of the table on your own account (the rows above: `runs/serve-*`, `runs/grouping-4b-h100`, `runs/fused-27b-*`, `runs/serving-27b-r23`).
|
| 410 |
+
|
| 411 |
+
## Limitations
|
| 412 |
+
|
| 413 |
+
- Calibration is a single temperature. It can't reorder confidences, so the share of new-source decisions you can automate at a 5% error budget (0.52–0.69 for Kev-4B, 9B and 27B) is still below Jev's 0.70. Test a probability threshold on your own data before you rely on it.
|
| 414 |
+
- Knowledge questions are set by the base model. MMLU is 0.73 for Kev-9B against Jev's 0.90, and MMLU-Pro 0.59 against 0.84.
|
| 415 |
+
- Fine-tuning can make the base model worse at individual tasks. Date arithmetic was the clearest case ([issue #8](https://github.com/jaredpalmer/kev/issues/8)); training on stated day counts plus `KEV_DATE_FACTS=1` recovers it.
|
| 416 |
+
- Changing option order can change an answer. Question isolation doesn't prevent this.
|
| 417 |
+
- Kev-0.8B, 4B and 9B trained mostly on at most 384 state tokens and 1,024 tokens for the state plus one question (their document and skill fine-tunes on states of up to 7,552 tokens), Kev-27B on states of up to 32,768 tokens. Serving allows a 65,536-token state; the validated context length of each model is in [Models](#models).
|
| 418 |
+
- On a Mac, answers take hundreds of milliseconds, not tens. Kev-27B needs an 80 GB GPU. On a Mac it needs about 51 GB plus working memory; we expect a 96–128 GB Mac to fit it but haven't measured one.
|
| 419 |
+
- Kev-27B starts from a post-trained model whose training data we don't know.
|
| 420 |
+
|
| 421 |
+
## Development
|
| 422 |
+
|
| 423 |
+
```bash
|
| 424 |
+
uv run --extra serve python -m pytest tests/test_unit.py tests/test_research.py tests/test_generators.py tests/test_conventions.py \
|
| 425 |
+
tests/test_documents_tools.py tests/test_hard_v1.py tests/test_devtools_v1.py tests/test_breadth_v1.py tests/test_rounds.py tests/test_skill_scripts.py -q # no weights, no server; what CI runs
|
| 426 |
+
KEV_BASE_URL=http://127.0.0.1:8009 uv run --extra serve python -m pytest tests/test_api.py -q # against a running server
|
| 427 |
+
cd playground && npm run lint && npx next typegen && npx tsc --noEmit -p .
|
| 428 |
+
```
|
| 429 |
+
|
| 430 |
+
The API tests run TypeSafe's example requests and the official SDK against your local server. [PLAN.md](PLAN.md) is the research plan: what we have learned, the rules every experiment follows, and one line per round. The full log (every experiment, the criteria set before it ran, and how it came out) is at the git tag `research-archive-2026-09-24`.
|
| 431 |
+
|
| 432 |
+
<details>
|
| 433 |
+
<summary>Previous generation (Qwen3) and the prototype</summary>
|
| 434 |
+
|
| 435 |
+
The first Kev family used Qwen3 bases with the same data and settings. Those weights stay published and run on plain PyTorch on a Mac, but they are no longer developed.
|
| 436 |
+
|
| 437 |
+
| Model | Base | Accuracy: Trained Sources | Accuracy: New Sources | Brier: New Sources | Model Card |
|
| 438 |
+
|---|---|---|---|---|---|
|
| 439 |
+
| Kev-0.6B (Qwen3) — `jaredpalmer/kev-0.6b` | Qwen3-0.6B-Base | 0.801 / 0.808 | 0.620 / 0.642 | 0.536 / 0.483 | [Details](docs/model-cards/kev-0.6b-qwen3.md) |
|
| 440 |
+
| Kev-4B (Qwen3) — `jaredpalmer/kev-4b@qwen3` | Qwen3-4B-Base | 0.854 / 0.856 | 0.790 / 0.806 | 0.328 / 0.294 | [Details](docs/model-cards/kev-4b-qwen3.md) |
|
| 441 |
+
| Kev-8B (Qwen3) — `jaredpalmer/kev-8b` | Qwen3-8B-Base | 0.863 / 0.870 | 0.796 / 0.780 | 0.337 / 0.327 | [Details](docs/model-cards/kev-8b-qwen3.md) |
|
| 442 |
+
|
| 443 |
+
The original [Kev-0.5B](https://huggingface.co/jaredpalmer/kev-0.5b) used Qwen2.5-0.5B and is kept for reference; see its [model card](docs/model-cards/kev-0.5b.md).
|
| 444 |
+
|
| 445 |
+
</details>
|
| 446 |
+
|
| 447 |
+
<details>
|
| 448 |
+
<summary>Troubleshooting</summary>
|
| 449 |
+
|
| 450 |
+
- If MPS runs out of memory during training, check that you're running only one job. Don't enable `output_hidden_states` or add tokens with peft's `trainable_token_indices`; both have caused memory problems here.
|
| 451 |
+
- If the playground loads but buttons don't work, use `localhost:3001`. Next.js checks development hostnames. Other hosts need an entry in `allowedDevOrigins` in `playground/next.config.ts`.
|
| 452 |
+
- If dataset loading reports `Dataset scripts are no longer supported`, use `legacy-datasets/banking77`. This repo already uses it.
|
| 453 |
+
|
| 454 |
+
</details>
|
| 455 |
+
|
| 456 |
+
## Authors
|
| 457 |
+
|
| 458 |
+
- Jared Palmer ([@jaredpalmer](https://github.com/jaredpalmer))
|
| 459 |
+
|
| 460 |
+
Built with [Devin](https://devin.ai). Thanks to [Archer Hume](https://archerhume.com/posts/jevs-architecture-unmasked) for the architecture write-up, [TypeSafe](https://docs.typesafe.ai/api) for the API design, and [Qwen](https://huggingface.co/Qwen/Qwen3.5-9B-Base) for the base models.
|
| 461 |
+
|
| 462 |
+
Related work: [Hydragen](https://arxiv.org/abs/2402.05099), [DeFT](https://arxiv.org/abs/2404.00242), [FIRST](https://arxiv.org/abs/2406.15657).
|
| 463 |
+
|
| 464 |
+
## License
|
| 465 |
+
|
| 466 |
+
[Apache-2.0](LICENSE). The Qwen3, Qwen3.5 and Qwen3.8 base models are also Apache-2.0. Training datasets have their own licenses; see the [model cards](docs/model-cards/).
|
runtime/docs/packed-mlx.md
ADDED
|
@@ -0,0 +1,132 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Standalone Packed MLX Models
|
| 2 |
+
|
| 3 |
+
The packed checkpoint runs on Apple Silicon using MLX. It carries the entire
|
| 4 |
+
compressed backbone, compensation adapter, pointer head, temperature, model
|
| 5 |
+
configuration and tokenizer. Loading and inference never read the original
|
| 6 |
+
backbone or the original PEFT adapter, and work with an empty Hugging Face cache
|
| 7 |
+
and networking disabled.
|
| 8 |
+
|
| 9 |
+
Ternary weights stay in MLX's native 2-bit affine format; selected linear layers
|
| 10 |
+
and the embedding stay 4-bit. CLoQ is applied as two small matrix products,
|
| 11 |
+
without building its dense delta. The embedding dequantizes only requested
|
| 12 |
+
token rows. Activations are bf16. The constructor discards lazy random dense
|
| 13 |
+
parameters before evaluating any model parameters.
|
| 14 |
+
|
| 15 |
+
## Run
|
| 16 |
+
|
| 17 |
+
From the Kev repository root, with the existing MLX environment:
|
| 18 |
+
|
| 19 |
+
```sh
|
| 20 |
+
HF_HUB_OFFLINE=1 HF_HOME=/tmp/kev-empty-cache \
|
| 21 |
+
KEV_BACKEND=mlx KEV_PREFIX_CACHE=1 KEV_PREFIX_MAX_TOKENS=4096 \
|
| 22 |
+
.venv/bin/python -m kev.serve \
|
| 23 |
+
--run runs/ternary-rotation-20261005/standalone/kev-4b-T-cloq16-mlx \
|
| 24 |
+
--host 127.0.0.1 --port 8010
|
| 25 |
+
```
|
| 26 |
+
|
| 27 |
+
The server accepts the existing `POST /v1/systemone` API:
|
| 28 |
+
|
| 29 |
+
```sh
|
| 30 |
+
curl http://127.0.0.1:8010/v1/systemone \
|
| 31 |
+
-H 'Content-Type: application/json' \
|
| 32 |
+
-d '{"state":"I was charged twice for order 8812.","questions":{"team":{"type":"choice","instructions":"Which team should handle this ticket?","criteria":{"billing":"Charges and payments","shipping":"Delivery and delays"}}}}'
|
| 33 |
+
```
|
| 34 |
+
|
| 35 |
+
Library callers use the usual checkpoint interface:
|
| 36 |
+
|
| 37 |
+
```python
|
| 38 |
+
from kev.checkpoint import Checkpoint
|
| 39 |
+
|
| 40 |
+
tokenizer, model = Checkpoint("path/to/standalone-model").load("mps")
|
| 41 |
+
```
|
| 42 |
+
|
| 43 |
+
Copy the **whole standalone model directory** to the target computer. Keep the
|
| 44 |
+
Kev code and its MLX/serving dependencies available. No original model files or
|
| 45 |
+
Hugging Face cache are required on the target computer. This is an MLX format;
|
| 46 |
+
llama.cpp and non-Apple inference have not been implemented or validated.
|
| 47 |
+
|
| 48 |
+
## Build from the Compressed Experiment
|
| 49 |
+
|
| 50 |
+
The export reads the compressed packed file and its optional `.cloq.safetensors`
|
| 51 |
+
sidecar, plus only the source's small head/tokenizer files and the base's
|
| 52 |
+
`config.json`. The base/PEFT weight files are never read. Rotated and
|
| 53 |
+
activation-order recipes are rejected rather than loaded incorrectly.
|
| 54 |
+
|
| 55 |
+
```sh
|
| 56 |
+
.venv/bin/python scripts/build_packed_model.py \
|
| 57 |
+
--packed path/to/experimental.safetensors \
|
| 58 |
+
--checkpoint path/to/head-and-tokenizer-directory \
|
| 59 |
+
--config path/to/base-config.json \
|
| 60 |
+
--out path/to/new-standalone-model
|
| 61 |
+
```
|
| 62 |
+
|
| 63 |
+
The export copies the tokenizer files and retains source hashes in `packed.json`
|
| 64 |
+
and `head.pt`. It does not quantize again, train, or fit a temperature. Native
|
| 65 |
+
2-bit storage uses more bits than the experiment's five-trits-per-byte storage,
|
| 66 |
+
so the standalone file is larger than the archival packed file. Original
|
| 67 |
+
experimental files and results remain unchanged.
|
| 68 |
+
|
| 69 |
+
## Verify
|
| 70 |
+
|
| 71 |
+
```sh
|
| 72 |
+
.venv/bin/python -m pytest tests/test_packed_model.py -q
|
| 73 |
+
|
| 74 |
+
HF_HUB_OFFLINE=1 HF_HOME=/tmp/kev-empty-cache \
|
| 75 |
+
.venv/bin/python scripts/profile_packed_model.py \
|
| 76 |
+
--run path/to/standalone-model --mlx-limit-gb 3 --out path/to/memory.json
|
| 77 |
+
|
| 78 |
+
HF_HUB_OFFLINE=1 HF_HOME=/tmp/kev-empty-cache KEV_BACKEND=mlx \
|
| 79 |
+
.venv/bin/python -m kev.benchmark --run path/to/standalone-model \
|
| 80 |
+
--suite evals/v7/decision-v7 --device mps --out path/to/new-evaluation
|
| 81 |
+
```
|
| 82 |
+
|
| 83 |
+
The profiler blocks the original loaders and downloads, and reports both MLX
|
| 84 |
+
allocator memory and sampled process physical footprint. Physical footprint
|
| 85 |
+
includes Metal buffers; RSS alone does not. The MLX allocator cap is not a cap
|
| 86 |
+
on the entire process or the operating system. A test on a larger-memory Mac
|
| 87 |
+
does not establish compatibility with every low-memory device.
|
| 88 |
+
|
| 89 |
+
Native quantized matmul plus a separate compensation adapter uses different
|
| 90 |
+
arithmetic from the original experiment's merged, dequantized bf16 weights.
|
| 91 |
+
Report its newly measured accuracy and Brier score; do not inherit the archival
|
| 92 |
+
96.7% retention or claim bit-identical probabilities. Context length, option
|
| 93 |
+
count and cached states also consume memory, so report those with memory results.
|
| 94 |
+
`KEV_PREFIX_MAX_TOKENS` controls retained cache size, not input admission.
|
| 95 |
+
|
| 96 |
+
The development partition is already used for recipe selection. Results remain
|
| 97 |
+
development diagnostics, not an untouched final test or a general guarantee.
|
| 98 |
+
|
| 99 |
+
## Measured 4B Result (2026-10-06)
|
| 100 |
+
|
| 101 |
+
`standalone/kev-4b-T-cloq16-mlx` was evaluated with an empty HF cache in offline
|
| 102 |
+
mode on decision-v7 development: all 1,204 records / 1,468 rows, 1,264 scored
|
| 103 |
+
clean questions, zero rejections and zero truncations. Accuracy is **0.84415**
|
| 104 |
+
(**96.82%** of the archived bf16-MLX accuracy 0.87184), Brier **0.23251**.
|
| 105 |
+
Against the archived packed merged-bf16 probabilities, 1/1,468 argmaxes changed;
|
| 106 |
+
mean per-row max probability difference is 0.00268, maximum 0.06209. The
|
| 107 |
+
original bf16 model was not loaded or re-evaluated for this comparison.
|
| 108 |
+
|
| 109 |
+
On an M4 Pro with 64 GiB RAM, fresh-process probes explicitly blocked the base
|
| 110 |
+
loaders and downloads. MLX 0.32.2, mlx-lm 0.31.3, torch 2.8.0, transformers 5.17.0:
|
| 111 |
+
|
| 112 |
+
| Measurement | Short request | Longer request |
|
| 113 |
+
|---|---:|---:|
|
| 114 |
+
| State tokens / total tokens | 19 / 42 | 3,519 / 3,542 |
|
| 115 |
+
| Load time | 0.57 s | 0.57 s |
|
| 116 |
+
| MLX active weights after load | 1.679 GB | 1.679 GB |
|
| 117 |
+
| Process physical footprint peak during load | 2.396 GB | 2.398 GB |
|
| 118 |
+
| MLX inference peak | 2.013 GB | 3.066 GB |
|
| 119 |
+
| Process physical footprint inference peak | 2.846 GB | 3.862 GB |
|
| 120 |
+
| First inference time | 0.155 s | 6.599 s |
|
| 121 |
+
|
| 122 |
+
These are single-request probes, not a minimum-RAM device certification. No
|
| 123 |
+
8-GB machine was tested. The 3-GB `mx.set_memory_limit` setting is explicitly a
|
| 124 |
+
**guideline**, not a hard cap; the longer request exceeded it. Use the observed
|
| 125 |
+
physical footprint when judging fit, and leave room for the OS and other apps.
|
| 126 |
+
GB above means decimal GB. The full model folder is approximately 1.709 GB;
|
| 127 |
+
native weights including compensation are 1.679 GB.
|
| 128 |
+
|
| 129 |
+
Artifacts under `runs/ternary-rotation-20261005/standalone-checks/`:
|
| 130 |
+
`kev-4b-T-cloq16-native/report.json`, `native-comparison.json`,
|
| 131 |
+
`memory-short.json`, and `memory-long.json`. The longer synthetic filler probe
|
| 132 |
+
measures memory/latency only and does not establish long-context accuracy.
|
runtime/kev/__init__.py
ADDED
|
File without changes
|
runtime/kev/anchors.py
ADDED
|
@@ -0,0 +1,71 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Anchor targets: the frozen base model's zero-shot distribution over each training question's options, read from
|
| 2 |
+
next-token letter logits (the same readout as scripts/base_mmlu_probe.py). Used by `kev.train --anchor --anchor_w` as a
|
| 3 |
+
KL target so the fine-tune keeps what the base already knows (PLAN.md: fine-tuning erodes base capability).
|
| 4 |
+
|
| 5 |
+
uv run python -m kev.anchors --base Qwen/Qwen3-4B-Base --suite evals/v6/decision-v6 --out runs/anchors/v6-4b.json
|
| 6 |
+
|
| 7 |
+
Targets are keyed by record id and question id, values keyed by option key, so they align with any option permutation
|
| 8 |
+
the trainer applies; questions whose option set changed (none-of-the-above inserted) are skipped at training time.
|
| 9 |
+
No Jev or other external model is involved: the teacher is the student's own base.
|
| 10 |
+
"""
|
| 11 |
+
import argparse
|
| 12 |
+
import json
|
| 13 |
+
from pathlib import Path
|
| 14 |
+
|
| 15 |
+
import torch
|
| 16 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
| 17 |
+
|
| 18 |
+
from kev.api import question_keys
|
| 19 |
+
from kev.device import default_device
|
| 20 |
+
from kev.suite import load_split, write_json
|
| 21 |
+
|
| 22 |
+
LETTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def question_prompt(state, q):
|
| 26 |
+
text = state if isinstance(state, str) else " ".join(f"{k}: {v}" for k, v in state.items()) if isinstance(state, dict) else json.dumps(state)
|
| 27 |
+
keys = question_keys(q["type"], q.get("criteria"))
|
| 28 |
+
if q["type"] == "noul": texts = ["No", "Yes"]
|
| 29 |
+
elif q["type"] == "score": texts = list(q["criteria"])
|
| 30 |
+
else: texts = [v if v is not None else k for k, v in q["criteria"].items()]
|
| 31 |
+
prompt = f"{text}\n{q['instructions']}\n" + "\n".join(f"{LETTERS[i]}. {t}" for i, t in enumerate(texts)) + "\nAnswer:"
|
| 32 |
+
return prompt, keys
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
def build(base, suite, out, split="train", device=None, revision=None, max_options=26, batch=16):
|
| 36 |
+
device = device or default_device()
|
| 37 |
+
tok = AutoTokenizer.from_pretrained(base, revision=revision); tok.padding_side = "left"
|
| 38 |
+
if tok.pad_token is None: tok.pad_token = tok.eos_token
|
| 39 |
+
model = AutoModelForCausalLM.from_pretrained(base, revision=revision, dtype=torch.bfloat16 if device != "cpu" else torch.float32).to(device).eval()
|
| 40 |
+
letter_ids = [tok.encode(" " + L, add_special_tokens=False)[0] for L in LETTERS]
|
| 41 |
+
records = load_split(suite, split)
|
| 42 |
+
jobs = [(r["_meta"]["id"], qid, *question_prompt(r["state"], q)) for r in records for qid, q in r["questions"].items() if len(q.get("criteria", [0, 1])) <= max_options or q["type"] == "noul"]
|
| 43 |
+
targets, skipped = {}, 0
|
| 44 |
+
with torch.no_grad():
|
| 45 |
+
for i in range(0, len(jobs), batch):
|
| 46 |
+
chunk = jobs[i : i + batch]
|
| 47 |
+
enc = tok([p for _, _, p, _ in chunk], return_tensors="pt", padding=True, truncation=True, max_length=3072).to(device)
|
| 48 |
+
logits = model(**enc).logits[:, -1].float()
|
| 49 |
+
for (rid, qid, _, keys), row in zip(chunk, logits):
|
| 50 |
+
if len(keys) > len(letter_ids): skipped += 1; continue
|
| 51 |
+
p = torch.softmax(row[letter_ids[: len(keys)]], -1).tolist()
|
| 52 |
+
targets.setdefault(rid, {})[qid] = dict(zip(keys, p))
|
| 53 |
+
if (i // batch) % 50 == 0: print(f"anchors {i}/{len(jobs)}", flush=True)
|
| 54 |
+
meta = {"base": base, "revision": revision, "suite": str(suite), "split": split, "readout": "zero-shot next-token letter logits",
|
| 55 |
+
"records": len(targets), "questions": sum(len(v) for v in targets.values()), "skipped_over_26_options": skipped}
|
| 56 |
+
write_json(out, {"_meta": meta, "targets": targets})
|
| 57 |
+
print(json.dumps(meta))
|
| 58 |
+
return meta
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def main():
|
| 62 |
+
ap = argparse.ArgumentParser()
|
| 63 |
+
ap.add_argument("--base", required=True); ap.add_argument("--suite", required=True); ap.add_argument("--out", required=True)
|
| 64 |
+
ap.add_argument("--split", default="train"); ap.add_argument("--device"); ap.add_argument("--revision")
|
| 65 |
+
a = ap.parse_args()
|
| 66 |
+
if Path(a.out).exists(): raise FileExistsError(a.out)
|
| 67 |
+
build(a.base, a.suite, a.out, a.split, a.device, a.revision)
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
if __name__ == "__main__":
|
| 71 |
+
main()
|
runtime/kev/api.py
ADDED
|
@@ -0,0 +1,165 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""TypeSafe-compatible request/response shapes (POST /v1/systemone) mapped onto the single pointer primitive.
|
| 2 |
+
|
| 3 |
+
Noul -> 2 options [false, true]; answer = p(true)
|
| 4 |
+
Choice -> options 'name' or 'name: desc'; answer = argmax, probabilities by name, confidence
|
| 5 |
+
Score -> options = ordered level descriptions; answer = expected level, legend, probabilities by index
|
| 6 |
+
"""
|
| 7 |
+
import json
|
| 8 |
+
import re
|
| 9 |
+
from datetime import datetime
|
| 10 |
+
from typing import Any, Literal, Union
|
| 11 |
+
from pydantic import BaseModel, Field, model_validator
|
| 12 |
+
|
| 13 |
+
JSONContent = Union[str, dict, list, int, float, bool, None]
|
| 14 |
+
MAX_OPTIONS = 255
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
class Noul(BaseModel):
|
| 18 |
+
type: Literal["noul"]
|
| 19 |
+
instructions: JSONContent = None
|
| 20 |
+
criteria: dict[str, JSONContent] | None = None
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
class Choice(BaseModel):
|
| 24 |
+
type: Literal["choice"]
|
| 25 |
+
instructions: JSONContent = None
|
| 26 |
+
criteria: dict[str, JSONContent]
|
| 27 |
+
|
| 28 |
+
@model_validator(mode="after")
|
| 29 |
+
def _check(self):
|
| 30 |
+
if not 1 <= len(self.criteria) <= MAX_OPTIONS: raise ValueError(f"criteria must have 1..{MAX_OPTIONS} options")
|
| 31 |
+
return self
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
class Score(BaseModel):
|
| 35 |
+
type: Literal["score"]
|
| 36 |
+
instructions: JSONContent = None
|
| 37 |
+
criteria: list[JSONContent] = Field(min_length=1, max_length=MAX_OPTIONS)
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
Question = Union[Noul, Choice, Score]
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
class SystemOneRequest(BaseModel):
|
| 44 |
+
state: JSONContent
|
| 45 |
+
model: str = "kev-latest"
|
| 46 |
+
questions: dict[str, Question] = Field(min_length=1)
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def render(v: JSONContent, indent: int = 0) -> str:
|
| 50 |
+
"""Flatten str | object | array into text the model sees. Field names are kept as labels."""
|
| 51 |
+
pad = " " * indent
|
| 52 |
+
if v is None: return ""
|
| 53 |
+
if isinstance(v, (str, int, float, bool)): return str(v)
|
| 54 |
+
if isinstance(v, list): return "\n".join(f"{pad}- {render(x, indent + 1).lstrip()}" for x in v)
|
| 55 |
+
return "\n".join(f"{pad}{k}:\n{render(x, indent + 1)}" if isinstance(x, (dict, list)) else f"{pad}{k}: {render(x)}" for k, x in v.items())
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
def option_text(name: str, desc: JSONContent) -> str:
|
| 59 |
+
return name if desc is None or desc == "" else f"{name}: {render(desc)}"
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
MONTHS = "January|February|March|April|May|June|July|August|September|October|November|December"
|
| 63 |
+
_DATE = re.compile(rf"\b(?:{MONTHS}) \d{{1,2}}, \d{{4}}\b|\b\d{{4}}-\d{{2}}-\d{{2}}\b")
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def date_facts(text: str) -> str:
|
| 67 |
+
"""Deterministic date arithmetic for the model: every pair of absolute dates found in `text`, as one sentence each
|
| 68 |
+
("August 3, 2026 is 12 days after July 22, 2026."). The model cannot subtract dates reliably (issue #8); it can use a
|
| 69 |
+
stated day count. Returns "" when fewer than two dates are found. Dates are listed in order of first appearance."""
|
| 70 |
+
found = []
|
| 71 |
+
for m in _DATE.finditer(text):
|
| 72 |
+
raw = m.group(0)
|
| 73 |
+
try: d = datetime.strptime(raw, "%B %d, %Y") if "," in raw else datetime.strptime(raw, "%Y-%m-%d")
|
| 74 |
+
except ValueError: continue
|
| 75 |
+
if raw not in [r for r, _ in found]: found.append((raw, d))
|
| 76 |
+
facts = []
|
| 77 |
+
for i in range(len(found)):
|
| 78 |
+
for j in range(i + 1, len(found)):
|
| 79 |
+
n = (found[j][1] - found[i][1]).days
|
| 80 |
+
facts.append(f"{found[j][0]} is {abs(n)} day{'s' if abs(n) != 1 else ''} {'after' if n > 0 else 'before'} {found[i][0]}." if n else f"{found[j][0]} is the same day as {found[i][0]}.")
|
| 81 |
+
return " ".join(facts)
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
def with_date_facts(state):
|
| 85 |
+
"""State with a `date_facts` field (object states) or an appended paragraph (string states) when two or more absolute
|
| 86 |
+
dates appear. Opt-in preprocessing (KEV_DATE_FACTS=1 in kev.serve, --date_facts in kev.benchmark)."""
|
| 87 |
+
facts = date_facts(render(state))
|
| 88 |
+
if not facts: return state
|
| 89 |
+
if isinstance(state, dict): return {**state, "date_facts": facts}
|
| 90 |
+
if isinstance(state, list): return state + [{"date_facts": facts}]
|
| 91 |
+
return f"{state}\n\ndate_facts: {facts}"
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
def question_keys(qtype: str, criteria) -> list[str]:
|
| 95 |
+
"""The keys a question's probabilities are reported under, in option order: the criteria names (choice),
|
| 96 |
+
["false", "true"] (noul), the level indices as strings (score). Labels, targets and anchors use the same keys."""
|
| 97 |
+
if qtype == "choice": return list(criteria)
|
| 98 |
+
if qtype == "noul": return ["false", "true"]
|
| 99 |
+
return [str(i) for i in range(len(criteria))]
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
def to_record(req: SystemOneRequest):
|
| 103 |
+
"""-> internal record for encode(), plus per-question metadata ({"id", "type", "keys", "legend" for score}) to map
|
| 104 |
+
probabilities back."""
|
| 105 |
+
qs, meta = [], []
|
| 106 |
+
for qid, q in req.questions.items():
|
| 107 |
+
m = {"id": qid, "type": q.type, "keys": question_keys(q.type, q.criteria)}
|
| 108 |
+
if q.type == "noul":
|
| 109 |
+
c = q.criteria or {}
|
| 110 |
+
opts = [option_text("no", c.get("false")), option_text("yes", c.get("true"))]
|
| 111 |
+
elif q.type == "choice":
|
| 112 |
+
opts = [option_text(k, v) for k, v in q.criteria.items()]
|
| 113 |
+
else:
|
| 114 |
+
opts = [render(x) for x in q.criteria]
|
| 115 |
+
m["legend"] = dict(zip(m["keys"], opts))
|
| 116 |
+
qs.append({"instr": render(q.instructions), "options": opts, "label": 0}); meta.append(m)
|
| 117 |
+
return {"state": render(req.state), "questions": qs}, meta
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
# Both confidence formulas mirror TypeSafe's reference adapter, system-one-adapter 0.2.1
|
| 121 |
+
# (src/system_one_adapter/_utils/confidence_metrics.py): p is normalised to sum 1 (all zeros -> uniform), one option -> 1.
|
| 122 |
+
def _normalize(p: list[float]) -> list[float]:
|
| 123 |
+
t = sum(p)
|
| 124 |
+
return [1 / len(p)] * len(p) if t == 0 else [x / t for x in p]
|
| 125 |
+
|
| 126 |
+
|
| 127 |
+
def choice_confidence(p: list[float]) -> float:
|
| 128 |
+
"""(p_max - 1/K) / (1 - 1/K): 0 at uniform, 1 at certainty."""
|
| 129 |
+
K = len(p)
|
| 130 |
+
return 1.0 if K == 1 else (max(_normalize(p)) - 1 / K) / (1 - 1 / K)
|
| 131 |
+
|
| 132 |
+
|
| 133 |
+
def score_confidence(p: list[float]) -> float:
|
| 134 |
+
"""max(0, 1 - E|level - mode| / D), D = mean absolute deviation of a uniform distribution over the L levels around its
|
| 135 |
+
mean (L-1)/2; mode = first most likely level. 1 when all mass is on one level, 0 at uniform or anything as spread."""
|
| 136 |
+
L = len(p)
|
| 137 |
+
if L == 1: return 1.0
|
| 138 |
+
p = _normalize(p); mode = max(range(L), key=p.__getitem__)
|
| 139 |
+
D = sum(abs(i - (L - 1) / 2) for i in range(L)) / L
|
| 140 |
+
return max(0.0, 1.0 - sum(pi * abs(i - mode) for i, pi in enumerate(p)) / D)
|
| 141 |
+
|
| 142 |
+
|
| 143 |
+
def round_prob(x: float) -> float:
|
| 144 |
+
"""Serialization precision for probabilities and derived scalars. 4 decimals keeps the sum of a rounded distribution
|
| 145 |
+
within TypeSafe's tolerance (|sum - 1| < 0.02) at the 255-option maximum: 255 * 0.00005 < 0.02."""
|
| 146 |
+
return round(float(x), 4)
|
| 147 |
+
|
| 148 |
+
|
| 149 |
+
def to_answers(probs: list[list[float]], meta: list[dict]) -> dict[str, Any]:
|
| 150 |
+
out = {}
|
| 151 |
+
for p, m in zip(probs, meta):
|
| 152 |
+
if m["type"] == "noul":
|
| 153 |
+
out[m["id"]] = {"type": "noul", "noul": round_prob(p[1])}
|
| 154 |
+
elif m["type"] == "choice":
|
| 155 |
+
dist = {k: round_prob(v) for k, v in zip(m["keys"], p)}
|
| 156 |
+
out[m["id"]] = {"type": "choice", "choice": m["keys"][max(range(len(p)), key=lambda i: p[i])], "confidence": round_prob(choice_confidence(p)), "probabilities": dist}
|
| 157 |
+
else:
|
| 158 |
+
score = sum(i * pi for i, pi in enumerate(p))
|
| 159 |
+
out[m["id"]] = {"type": "score", "score": round_prob(score), "legend": m["legend"], "probabilities": {str(i): round_prob(v) for i, v in enumerate(p)}, "confidence": round_prob(score_confidence(p))}
|
| 160 |
+
return out
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
def output_tokens(tok, answers: dict) -> int:
|
| 164 |
+
"""Billing-style figure: tokens of the serialised answers. Not a measure of generation (there is none)."""
|
| 165 |
+
return len(tok(json.dumps(answers), add_special_tokens=False).input_ids)
|
runtime/kev/autoresearch.py
ADDED
|
@@ -0,0 +1,181 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Unattended research sessions on the round harness, plus the study leaderboard.
|
| 2 |
+
|
| 3 |
+
uv run python -m kev.autoresearch session experiments/rounds/r19.json [...] --spend-start 1377 --spend-cap 300
|
| 4 |
+
uv run python -m kev.autoresearch leaderboard # runs/leaderboard.{jsonl,md} from every trial's result.json
|
| 5 |
+
uv run python -m kev.autoresearch release-check --study X # every seed of a config must pass its gates
|
| 6 |
+
uv run python -m kev.autoresearch compare --studies a,b --reference runs/<study>/<trial>
|
| 7 |
+
|
| 8 |
+
A session runs registered rounds in order through kev.rounds (validate, launch the studies, watch: pull and read each
|
| 9 |
+
finished trial, then the read-out). What it may run is what the specs register: config-only trials through
|
| 10 |
+
kev.experiment.execute_trial with full provenance, and the reads the rule names. It never changes the evaluator, the
|
| 11 |
+
frozen suites or the gates, and it never confirms: test partitions and the locked read are read once, by a deliberate
|
| 12 |
+
`python -m kev.rounds confirm` after a read-out names a candidate (the session prints those commands).
|
| 13 |
+
|
| 14 |
+
Budget: before each round, `modal billing summary` is read; the session stops when metered spend since --spend-start
|
| 15 |
+
plus the round's study budgets (each at least its admission bound, checked by kev.rounds.validate) would pass
|
| 16 |
+
--spend-cap. Each round appends a line to runs/autoresearch-sessions.jsonl.
|
| 17 |
+
"""
|
| 18 |
+
import argparse
|
| 19 |
+
import hashlib
|
| 20 |
+
import json
|
| 21 |
+
import subprocess
|
| 22 |
+
import sys
|
| 23 |
+
from collections import defaultdict
|
| 24 |
+
from datetime import datetime, timezone
|
| 25 |
+
from pathlib import Path
|
| 26 |
+
|
| 27 |
+
from kev import rounds
|
| 28 |
+
from kev.experiment import CHOICE_DEFAULTS, DEFAULTS
|
| 29 |
+
from kev.metrics import paired_bootstrap
|
| 30 |
+
from kev.suite import ENCODING, read_json, write_jsonl
|
| 31 |
+
|
| 32 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 33 |
+
H100_RATE = 3.95
|
| 34 |
+
SNAPSHOT_KEYS = ("snapshot_fractions", "snapshot_every_steps", "snapshot_hub_repo") # when and where a full-weight trial saves intermediate checkpoints: not what it trains
|
| 35 |
+
INFRA_KEYS = ("base", "seed", "base_revision", "dtype", "checkpointing", "batch", "accum", "perm_frac", "shared_prefix", *SNAPSHOT_KEYS) # execution shape, not recipe
|
| 36 |
+
ALL_DEFAULTS = {**DEFAULTS, **CHOICE_DEFAULTS} # what kev.train does when a knob is not given
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def config_digest(value):
|
| 40 |
+
"""Canonical (key-sorted) digest for config identity; kev.suite.record_digest keeps insertion order for provenance."""
|
| 41 |
+
return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
def knobs(cfg):
|
| 45 |
+
"""The recipe a config expresses: every non-default, non-infrastructure parameter."""
|
| 46 |
+
return {k: v for k, v in sorted(cfg.items()) if k not in INFRA_KEYS and v != ALL_DEFAULTS.get(k)}
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def recipe(row):
|
| 50 |
+
"""A trial's recipe identity: its config without the seed (config_sha256 hashes the seed too) and snapshot schedule."""
|
| 51 |
+
return config_digest({k: v for k, v in row["config"].items() if k != "seed" and k not in SNAPSHOT_KEYS})
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
# --- sessions --------------------------------------------------------------------------------------------------------
|
| 55 |
+
|
| 56 |
+
def metered_spend():
|
| 57 |
+
out = subprocess.run([sys.executable, "-m", "modal", "billing", "summary", "--json"], capture_output=True, text=True, cwd=ROOT)
|
| 58 |
+
return float(json.loads(out.stdout)["metered_cost"]) if out.returncode == 0 else None
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def session(paths, spend_start, spend_cap, spend=metered_spend, log=print):
|
| 62 |
+
"""Run each registered round to its read-out while the spend cap allows; returns the ledger entries written."""
|
| 63 |
+
entries = []
|
| 64 |
+
for path in paths:
|
| 65 |
+
spec = rounds.load(path)
|
| 66 |
+
problems = rounds.launchable(spec)
|
| 67 |
+
if problems: log(f"{path}: does not validate, session stops:\n" + "\n".join(problems)); break
|
| 68 |
+
spent, bound = spend(), sum(s["budget"] for s in spec["studies"].values())
|
| 69 |
+
if spent is None: log("cannot read `modal billing summary`; session stops"); break
|
| 70 |
+
if spent - spend_start + bound > spend_cap:
|
| 71 |
+
log(f"round {spec['round']}: spent ${spent - spend_start:.2f} + budgets ${bound:.2f} > cap ${spend_cap:.2f}; session stops"); break
|
| 72 |
+
rounds.launch_studies(spec)
|
| 73 |
+
report = rounds.watch(spec)
|
| 74 |
+
entry = {"round": spec["round"], "spec": str(path), "spend_before": round(spent - spend_start, 2), "budgets": bound, "candidates": report["candidates"],
|
| 75 |
+
"at": datetime.now(timezone.utc).isoformat(timespec="minutes")}
|
| 76 |
+
with (ROOT / "runs/autoresearch-sessions.jsonl").open("a", encoding=ENCODING) as f: f.write(json.dumps(entry) + "\n")
|
| 77 |
+
entries.append(entry)
|
| 78 |
+
for arm in filter(None, report["candidates"].values()):
|
| 79 |
+
for stage in spec.get("confirm", {}):
|
| 80 |
+
log(f"candidate {arm}: uv run python -m kev.rounds launch-reads {path} --stage {stage} --arm {arm} then confirm {path} --stage {stage} --arm {arm}")
|
| 81 |
+
return entries
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
# --- leaderboard -----------------------------------------------------------------------------------------------------
|
| 85 |
+
|
| 86 |
+
def collect():
|
| 87 |
+
"""Every completed trial in runs/*/ with its provenance, as flat leaderboard rows."""
|
| 88 |
+
rows = []
|
| 89 |
+
for result in sorted(ROOT.glob("runs/*/*/result.json")):
|
| 90 |
+
r = read_json(result)
|
| 91 |
+
prov = r.get("provenance", {}); cfg = prov.get("config") or {}
|
| 92 |
+
tr = r.get("transfer") or {}
|
| 93 |
+
rows.append({"study": result.parent.parent.name, "trial": result.parent.name, "legacy": prov.get("legacy_checkpoint", False),
|
| 94 |
+
"base": cfg.get("base"), "seed": cfg.get("seed"), "config": cfg, "config_sha256": prov.get("config_sha256"),
|
| 95 |
+
"suite_sha256": prov.get("suite_sha256"), "transfer_suite_sha256": tr.get("suite_sha256"), "git": prov.get("git_commit", "")[:8],
|
| 96 |
+
"dev_acc": r["clean"]["acc"], "dev_nll": r["clean"]["nll"], "dev_ece": r["clean"]["ece"],
|
| 97 |
+
"transfer_acc": tr.get("clean", {}).get("acc"), "transfer_brier": tr.get("clean", {}).get("brier"),
|
| 98 |
+
"transfer_conf_err": tr.get("clean", {}).get("confident_error_rate"),
|
| 99 |
+
"heldout_pairs": (tr.get("paired_flip") or {}).get("both_correct_rate"),
|
| 100 |
+
"none_present": r["variants"].get("none_present", {}).get("acc"), "perm_flip": r["permutation"]["flip_rate"],
|
| 101 |
+
"gates": r["gates"]["checks"], "gates_passed": r["gates"]["passed"],
|
| 102 |
+
"train_s": (r.get("training_resources") or {}).get("wall_seconds"), "wall_s": r.get("wall_seconds"),
|
| 103 |
+
"est_usd": round(H100_RATE * r.get("wall_seconds", 0) / 3600, 2) if prov.get("device") == "cuda" else None})
|
| 104 |
+
return rows
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
def leaderboard_md(rows):
|
| 108 |
+
lines = ["# Leaderboard", "", f"Generated {datetime.now(timezone.utc).isoformat(timespec='minutes')} from runs/*/result.json. "
|
| 109 |
+
"Selection on development partitions only; the locked test is never read here. Rounds are decided by their "
|
| 110 |
+
"registered rules (`python -m kev.rounds readout`), not by this table.", "",
|
| 111 |
+
"| study/trial | base | seed | dev acc | transfer acc | Brier | conf-err | held-out pairs | none_present | perm flip | gates | $ | knobs |", "|---|---|---|---|---|---|---|---|---|---|---|---|---|"]
|
| 112 |
+
for r in sorted(rows, key=lambda r: (r["transfer_acc"] is None, -(r["transfer_acc"] or 0))):
|
| 113 |
+
b = (r["base"] or "legacy").split("/")[-1]
|
| 114 |
+
f = lambda x, p=3: "" if x is None else f"{x:.{p}f}"
|
| 115 |
+
lines.append(f"| {r['study']}/{r['trial']} | {b} | {r['seed'] if r['seed'] is not None else ''} | {f(r['dev_acc'])} | {f(r['transfer_acc'])} | {f(r['transfer_brier'])} | {f(r['transfer_conf_err'])} | {f(r['heldout_pairs'],2)} | {f(r['none_present'],2)} | {f(r['perm_flip'],2)} | {'pass' if r['gates_passed'] else 'fail'} | {r['est_usd'] or ''} | {', '.join(f'{k}={v}' for k, v in knobs(r['config']).items()) or 'defaults'} |")
|
| 116 |
+
return "\n".join(lines) + "\n"
|
| 117 |
+
|
| 118 |
+
|
| 119 |
+
def refresh_leaderboard():
|
| 120 |
+
rows = collect()
|
| 121 |
+
write_jsonl(ROOT / "runs/leaderboard.jsonl", rows)
|
| 122 |
+
(ROOT / "runs/leaderboard.md").write_text(leaderboard_md(rows), encoding=ENCODING)
|
| 123 |
+
return rows
|
| 124 |
+
|
| 125 |
+
|
| 126 |
+
def compare(studies, reference, tasks=("mmlu", "paws", "qnli", "emotion", "tweet_offensive", "contrastive_deadline")):
|
| 127 |
+
"""Print every trial of the given studies with a record-clustered paired bootstrap on transfer accuracy vs `reference`
|
| 128 |
+
(a runs/<study>/<trial> path). Development-set selection only."""
|
| 129 |
+
def rows(p): return read_json(ROOT / p / "transfer/rows.json")
|
| 130 |
+
ref_rows = rows(reference)
|
| 131 |
+
print(f"reference: {reference}")
|
| 132 |
+
print(f"{'trial':34} {'dev':>6} {'trf':>6} {'brier':>6} {'cerr':>6} {'pairs':>6} {'d trf':>7} {'ci95':>18} knobs / tasks")
|
| 133 |
+
for study in studies:
|
| 134 |
+
for d in sorted((ROOT / "runs" / study).iterdir()):
|
| 135 |
+
if not (d / "result.json").exists(): continue
|
| 136 |
+
r = read_json(d / "result.json"); tr = r.get("transfer")
|
| 137 |
+
if not tr: continue
|
| 138 |
+
cfg = r["provenance"]["config"]
|
| 139 |
+
try:
|
| 140 |
+
b = paired_bootstrap(rows(f"runs/{study}/{d.name}"), ref_rows, metric="acc"); delta, ci = b["macro_acc_delta"], [round(x, 3) for x in b["ci95"]]
|
| 141 |
+
except ValueError:
|
| 142 |
+
delta, ci = float("nan"), "n/a (different suite)"
|
| 143 |
+
print(f"{study + '/' + d.name:34} {r['clean']['acc']:6.3f} {tr['clean']['acc']:6.3f} {tr['clean']['brier']:6.3f} {tr['clean']['confident_error_rate']:6.3f} {tr['paired_flip']['both_correct_rate']:6.2f} {delta:+7.3f} {str(ci):>18} {knobs(cfg)} {({k: round(tr['tasks'][k]['acc'], 2) for k in tasks if k in tr['tasks']})}")
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
def release_check(study):
|
| 147 |
+
"""Release screen across seeds: every trial of the same config in the study must pass its gates (including the 70%
|
| 148 |
+
held-out-pair screen) - a single seed clearing the bar is not enough. Prints the verdict per config."""
|
| 149 |
+
rows = [r for r in collect() if r["study"] == study]
|
| 150 |
+
by_cfg = defaultdict(list)
|
| 151 |
+
for r in rows: by_cfg[recipe(r)].append(r)
|
| 152 |
+
verdicts = {}
|
| 153 |
+
for group in by_cfg.values():
|
| 154 |
+
seeds = sorted(r["seed"] for r in group)
|
| 155 |
+
passed = all(r["gates_passed"] for r in group)
|
| 156 |
+
pairs = [round(r["heldout_pairs"], 2) for r in group]
|
| 157 |
+
failing = sorted({k for r in group for k, v in r["gates"].items() if not v})
|
| 158 |
+
base = group[0]["base"].split("/")[-1]
|
| 159 |
+
verdicts[base] = {"seeds": seeds, "all_gates_passed": passed and len(seeds) >= 2, "heldout_pairs": pairs,
|
| 160 |
+
"transfer_acc": [round(r["transfer_acc"], 3) for r in group], "failing_gates": failing,
|
| 161 |
+
"candidate": passed and len(seeds) >= 2, "trials": [f"{r['study']}/{r['trial']}" for r in group]}
|
| 162 |
+
print(f"{base}: seeds {seeds} pairs {pairs} transfer {verdicts[base]['transfer_acc']} -> {'RELEASE CANDIDATE (gated locked read allowed)' if verdicts[base]['candidate'] else 'not a candidate: ' + ', '.join(failing or ['fewer than two seeds'])}")
|
| 163 |
+
return verdicts
|
| 164 |
+
|
| 165 |
+
|
| 166 |
+
def main():
|
| 167 |
+
ap = argparse.ArgumentParser()
|
| 168 |
+
sub = ap.add_subparsers(dest="cmd", required=True)
|
| 169 |
+
p = sub.add_parser("session"); p.add_argument("specs", nargs="+"); p.add_argument("--spend-start", type=float, required=True); p.add_argument("--spend-cap", type=float, required=True)
|
| 170 |
+
sub.add_parser("leaderboard")
|
| 171 |
+
p = sub.add_parser("release-check"); p.add_argument("--study", required=True)
|
| 172 |
+
p = sub.add_parser("compare"); p.add_argument("--studies", required=True); p.add_argument("--reference", required=True, help="runs/<study>/<trial>")
|
| 173 |
+
a = ap.parse_args()
|
| 174 |
+
if a.cmd == "session": session(a.specs, a.spend_start, a.spend_cap)
|
| 175 |
+
elif a.cmd == "leaderboard": print(f"{len(refresh_leaderboard())} trials -> runs/leaderboard.md")
|
| 176 |
+
elif a.cmd == "release-check": release_check(a.study)
|
| 177 |
+
else: compare(a.studies.split(","), a.reference)
|
| 178 |
+
|
| 179 |
+
|
| 180 |
+
if __name__ == "__main__":
|
| 181 |
+
main()
|
runtime/kev/benchmark.py
ADDED
|
@@ -0,0 +1,214 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Score a predictor on a frozen suite partition (or your own labelled JSONL).
|
| 2 |
+
|
| 3 |
+
uv run python -m kev.benchmark --run runs/<run>/checkpoint --suite evals/<v>/decision-<v> --out runs/<name>
|
| 4 |
+
uv run python -m kev.benchmark --remote http://127.0.0.1:8008 --suite ... --out ... # any System One endpoint
|
| 5 |
+
|
| 6 |
+
Every prediction becomes one row per question (prediction_rows); kev.metrics scores rows; evaluate_records writes
|
| 7 |
+
predictions.jsonl, rows.json and report.json. Predictors live in kev.predictors.
|
| 8 |
+
"""
|
| 9 |
+
import argparse
|
| 10 |
+
import functools
|
| 11 |
+
import json
|
| 12 |
+
import math
|
| 13 |
+
import os
|
| 14 |
+
from concurrent.futures import ThreadPoolExecutor
|
| 15 |
+
from pathlib import Path
|
| 16 |
+
|
| 17 |
+
import numpy as np
|
| 18 |
+
|
| 19 |
+
from kev.api import question_keys, with_date_facts
|
| 20 |
+
from kev.checkpoint import LoadOptions
|
| 21 |
+
from kev.contrastive import paired_flip
|
| 22 |
+
from kev.data import api_request, load_records
|
| 23 |
+
from kev.device import default_device
|
| 24 |
+
from kev.metrics import EPSILON, grouped_metrics, metrics, unknowable_report
|
| 25 |
+
from kev.model import ROW_PASS_TOKENS, ContextOverflow
|
| 26 |
+
from kev.predictors import LocalPredictor, RemotePredictor, RotationAveraged
|
| 27 |
+
from kev.suite import CONTEXT, ENCODING, digest, load_split, read_manifest, record_digest, write_json
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
def labels(q):
|
| 31 |
+
"""(option keys, label index) of a labelled request question."""
|
| 32 |
+
keys = question_keys(q["type"], q.get("criteria"))
|
| 33 |
+
return keys, keys.index(q["label"]) if q["type"] == "choice" else int(q["label"])
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
def validate_distribution(raw, keys):
|
| 37 |
+
if set(raw) != set(keys):
|
| 38 |
+
raise ValueError("probability keys do not match requested options")
|
| 39 |
+
p = np.array([raw[k] for k in keys], dtype=float)
|
| 40 |
+
if not np.isfinite(p).all() or (p < 0).any() or (p > 1).any():
|
| 41 |
+
raise ValueError("non-finite or out-of-range probabilities")
|
| 42 |
+
total = float(p.sum())
|
| 43 |
+
if total <= 0 or abs(total - 1) > max(1e-5, len(keys) * 0.005 + 1e-8):
|
| 44 |
+
raise ValueError(f"invalid probability sum: {total}")
|
| 45 |
+
return p / total, total
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def prediction_rows(record, prediction):
|
| 49 |
+
if set(prediction["probabilities"]) != set(record["questions"]):
|
| 50 |
+
raise ValueError("answer IDs differ from request IDs")
|
| 51 |
+
meta = record["_meta"]
|
| 52 |
+
rows = []
|
| 53 |
+
for qid, q in record["questions"].items():
|
| 54 |
+
keys, y = labels(q)
|
| 55 |
+
p, total = validate_distribution(prediction["probabilities"][qid], keys)
|
| 56 |
+
row = {"id": meta["id"], "group": meta["group_id"], "question": qid,
|
| 57 |
+
"source": meta["source"], "task": q["src"], "type": q["type"],
|
| 58 |
+
"variant": meta["variant"], "keys": keys, "label": y, "control_id": meta.get("control_id"),
|
| 59 |
+
"pair_id": meta.get("pair_id"), "sibling": meta.get("sibling"), # suites frozen before parent_id existed stored the parent's id in group_id for variants
|
| 60 |
+
"parent": meta.get("parent_id") or (meta["id"] if meta["variant"] == "clean" else meta["group_id"]),
|
| 61 |
+
"p": p.tolist(), "raw_probability_sum": total, "zero_count": int((p == 0).sum())}
|
| 62 |
+
if "logits" in prediction:
|
| 63 |
+
raw_logits = prediction["logits"][qid]
|
| 64 |
+
if set(raw_logits) != set(keys) or not all(math.isfinite(raw_logits[k]) for k in keys):
|
| 65 |
+
raise ValueError("logit keys or values do not match the requested options")
|
| 66 |
+
row["logits"] = [float(raw_logits[k]) for k in keys]
|
| 67 |
+
row["inference_temperature"] = prediction["inference_temperature"]
|
| 68 |
+
if "kernels" in prediction: row["kernels"] = prediction["kernels"] # a long row off the fp32-exact kernels (LocalPredictor)
|
| 69 |
+
rows.append(row)
|
| 70 |
+
return rows
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
def summarize(rows, temperature=1.0, heldout_sources=()):
|
| 74 |
+
"""Report over benchmark rows. heldout_sources: sources the scored model never trained on; their tasks are also
|
| 75 |
+
reported as a separate block."""
|
| 76 |
+
clean = [r for r in rows if r["variant"] == "clean"]
|
| 77 |
+
tasks = grouped_metrics(clean, "task")
|
| 78 |
+
variants = grouped_metrics(rows, "variant")
|
| 79 |
+
lookup = {(r["id"], r["question"]): r for r in clean}
|
| 80 |
+
diffs, flips = [], []
|
| 81 |
+
for row in rows:
|
| 82 |
+
if row["variant"] == "permuted" and row["type"] == "choice":
|
| 83 |
+
original = lookup[(row["parent"], row["question"])]
|
| 84 |
+
aligned = [row["p"][row["keys"].index(k)] for k in original["keys"]]
|
| 85 |
+
diffs.append(float(np.max(np.abs(np.array(aligned) - original["p"]))))
|
| 86 |
+
flips.append(int(np.argmax(aligned) != np.argmax(original["p"])))
|
| 87 |
+
knowable = [r for r in clean if r["source"] != "unknowable"] # unknowable records are scored on confidence, never on accuracy
|
| 88 |
+
return {"objective": -float(np.mean([v["nll"] for k, v in tasks.items() if not k.startswith("unknowable_") or k.startswith("unknowable_control")])),
|
| 89 |
+
"paired_flip": paired_flip(clean), "unknowable": unknowable_report(clean),
|
| 90 |
+
"clean": metrics(knowable), "tasks": tasks, "variants": variants,
|
| 91 |
+
"heldout_tasks": grouped_metrics([r for r in clean if r["source"] in heldout_sources], "task") if any(r["source"] in heldout_sources for r in clean) else {},
|
| 92 |
+
"permutation": {"n": len(diffs), "mean_max_delta": float(np.mean(diffs)) if diffs else None,
|
| 93 |
+
"flip_rate": float(np.mean(flips)) if flips else None},
|
| 94 |
+
"temperature": temperature, "calibrated_clean": metrics(knowable, temperature),
|
| 95 |
+
"metric_policy": {"version": 2, "selective_ties": "whole_confidence_groups",
|
| 96 |
+
"coverage_at_error": "in-sample maximum over confidence thresholds; not a deployed error guarantee",
|
| 97 |
+
"aurc": "right-step integral over whole confidence groups",
|
| 98 |
+
"confident_error_rate": "high-confidence errors divided by all questions",
|
| 99 |
+
"error_rate_at_0_9": "errors divided by questions accepted at p_max >= 0.9",
|
| 100 |
+
"nll": "exact from logits when recorded; otherwise from floored probabilities",
|
| 101 |
+
"nll_floor": EPSILON, "renormalize_returned_probabilities": True,
|
| 102 |
+
"raw_sums_outside_1e_5": sum(abs(r["raw_probability_sum"] - 1) > 1e-5 for r in rows),
|
| 103 |
+
"returned_zeros": sum(r["zero_count"] for r in rows)}}
|
| 104 |
+
|
| 105 |
+
|
| 106 |
+
def predictions(records, predictor):
|
| 107 |
+
"""Yield, in record order, a zero-argument callable that returns predictor(record) or raises what it raised. A
|
| 108 |
+
predictor with `concurrency` > 1 (RemotePredictor: one independent HTTP request per call) keeps that many calls in
|
| 109 |
+
flight on a thread pool; in-process predictors (one GPU) have no `concurrency` and run one record at a time. Either
|
| 110 |
+
way the caller sees each outcome in the same order as the plain sequential loop."""
|
| 111 |
+
workers = getattr(predictor, "concurrency", 1)
|
| 112 |
+
if workers <= 1 or len(records) <= 1:
|
| 113 |
+
for record in records:
|
| 114 |
+
yield functools.partial(predictor, record)
|
| 115 |
+
return
|
| 116 |
+
executor = ThreadPoolExecutor(max_workers=min(workers, len(records)))
|
| 117 |
+
try:
|
| 118 |
+
for future in [executor.submit(predictor, record) for record in records]:
|
| 119 |
+
yield future.result
|
| 120 |
+
finally:
|
| 121 |
+
executor.shutdown(wait=False, cancel_futures=True)
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
def evaluate_records(records, predictor, directory, temperature=1.0, heldout_sources=(), skip_overlong=False):
|
| 125 |
+
"""skip_overlong: for external data that was not admitted to a context (--data, or an eval-only suite frozen as published),
|
| 126 |
+
records the predictor cannot encode are counted in coverage["rejected_records"] and listed in rejected.json instead of
|
| 127 |
+
aborting. Admitted suites never trigger it; reports must state how rejected records enter any headline number."""
|
| 128 |
+
directory = Path(directory)
|
| 129 |
+
directory.mkdir(parents=True, exist_ok=False)
|
| 130 |
+
coverage = {"requested_records": len(records), "requested_questions": sum(len(r["questions"]) for r in records),
|
| 131 |
+
"evaluated_records": 0, "evaluated_questions": 0, "rejected_records": 0, "truncated_records": 0}
|
| 132 |
+
rows, latencies, rejected = [], [], []
|
| 133 |
+
with (directory / "predictions.jsonl").open("w", encoding=ENCODING) as output:
|
| 134 |
+
preds = predictions(records, predictor)
|
| 135 |
+
for record in records:
|
| 136 |
+
try:
|
| 137 |
+
pred = next(preds)()
|
| 138 |
+
new_rows = prediction_rows(record, pred)
|
| 139 |
+
except ContextOverflow as error:
|
| 140 |
+
if skip_overlong:
|
| 141 |
+
coverage["rejected_records"] += 1; rejected.append({"id": record["_meta"]["id"], "error": str(error)}); continue
|
| 142 |
+
coverage["rejected_records"] += 1
|
| 143 |
+
write_json(directory / "failure.json", {"coverage": coverage, "record_id": record["_meta"]["id"], "error_type": type(error).__name__})
|
| 144 |
+
raise
|
| 145 |
+
except Exception as error:
|
| 146 |
+
coverage["rejected_records"] += 1
|
| 147 |
+
write_json(directory / "failure.json", {"coverage": coverage, "record_id": record["_meta"]["id"], "error_type": type(error).__name__})
|
| 148 |
+
raise
|
| 149 |
+
output.write(json.dumps({"request_sha256": record_digest(api_request(record)), "id": record["_meta"]["id"],
|
| 150 |
+
"prediction": pred, "rows": new_rows}, allow_nan=False) + "\n")
|
| 151 |
+
output.flush()
|
| 152 |
+
rows.extend(new_rows)
|
| 153 |
+
latencies.append(pred["latency_ms"])
|
| 154 |
+
coverage["evaluated_records"] += 1
|
| 155 |
+
coverage["evaluated_questions"] += len(new_rows)
|
| 156 |
+
if coverage["evaluated_records"] % 50 == 0:
|
| 157 |
+
print(f"evaluated {coverage['evaluated_records']}/{len(records)}", flush=True)
|
| 158 |
+
write_json(directory / "rows.json", rows)
|
| 159 |
+
if rejected: write_json(directory / "rejected.json", rejected)
|
| 160 |
+
report = summarize(rows, temperature, heldout_sources)
|
| 161 |
+
report.update(coverage=coverage, latency_ms={"median": float(np.median(latencies)), "p95": float(np.quantile(latencies, .95))},
|
| 162 |
+
calibration={"inference_temperature": getattr(predictor, "temperature", None),
|
| 163 |
+
"additional_temperature": temperature, "logits_recorded": all("logits" in r for r in rows)})
|
| 164 |
+
long = [r for r in rows if "kernels" in r]
|
| 165 |
+
if long: # absent when every row ran the exact kernels, so such reports are unchanged
|
| 166 |
+
report["long_rows"] = {"count": len(long), "records": len({r["id"] for r in long}), "kernels": sorted({r["kernels"] for r in long}), "threshold": ROW_PASS_TOKENS}
|
| 167 |
+
write_json(directory / "report.json", report)
|
| 168 |
+
return report, rows
|
| 169 |
+
|
| 170 |
+
|
| 171 |
+
def main():
|
| 172 |
+
ap = argparse.ArgumentParser()
|
| 173 |
+
ap.add_argument("--run", help="checkpoint dir or Hub id (local scoring)")
|
| 174 |
+
ap.add_argument("--remote", help="base URL of a System One-compatible endpoint to score instead of a local checkpoint")
|
| 175 |
+
ap.add_argument("--remote-model", default="kev-latest")
|
| 176 |
+
ap.add_argument("--remote-concurrency", type=int, default=1, help="requests kept in flight against --remote (1 = sequential); rows and their order do not depend on it")
|
| 177 |
+
ap.add_argument("--suite", help="frozen suite directory (scores its development partition)")
|
| 178 |
+
ap.add_argument("--data", help="your own labelled requests, one JSON object per line (kev.data.load_records); an alternative to --suite")
|
| 179 |
+
ap.add_argument("--out", required=True)
|
| 180 |
+
ap.add_argument("--device", choices=["cpu", "mps", "cuda"], default=default_device())
|
| 181 |
+
ap.add_argument("--allow-test", action="store_true")
|
| 182 |
+
ap.add_argument("--split", choices=["development", "calibration", "train"], default="development",
|
| 183 |
+
help="suite partition to score (train: teacher predictions for distillation; --allow-test reads the locked test instead)")
|
| 184 |
+
ap.add_argument("--date_facts", action="store_true", help="apply kev.api.with_date_facts to every state before scoring (the opt-in serving preprocessor); reported in report.json")
|
| 185 |
+
ap.add_argument("--rotations", type=int, default=1, help="average every Choice question over this many cyclic option rotations (kev.predictors.RotationAveraged); 1 = one order")
|
| 186 |
+
a = ap.parse_args()
|
| 187 |
+
if bool(a.run) == bool(a.remote): ap.error("give exactly one of --run or --remote")
|
| 188 |
+
if a.rotations < 1: ap.error("--rotations must be >= 1")
|
| 189 |
+
if a.remote_concurrency < 1: ap.error("--remote-concurrency must be >= 1")
|
| 190 |
+
if bool(a.suite) == bool(a.data): ap.error("give exactly one of --suite or --data")
|
| 191 |
+
if a.data:
|
| 192 |
+
records, heldout, split, source_hash = load_records(a.data), [], "custom", digest(Path(a.data))
|
| 193 |
+
context, skip_overlong = CONTEXT, True
|
| 194 |
+
else:
|
| 195 |
+
split = "test" if a.allow_test else a.split
|
| 196 |
+
records = load_split(a.suite, split, allow_test=a.allow_test)
|
| 197 |
+
manifest = read_manifest(a.suite)
|
| 198 |
+
heldout = manifest["holdout_sources"]; source_hash = digest(Path(a.suite) / "manifest.json")
|
| 199 |
+
context, skip_overlong = manifest.get("context", CONTEXT), bool(manifest.get("eval_only"))
|
| 200 |
+
if a.date_facts:
|
| 201 |
+
records = [{**r, "state": with_date_facts(r["state"])} for r in records]
|
| 202 |
+
predictor = RemotePredictor(a.remote, a.remote_model, os.environ.get("KEV_REMOTE_API_KEY", "local"), concurrency=a.remote_concurrency) if a.remote else LocalPredictor(a.run, a.device, LoadOptions.from_env(), context=context)
|
| 203 |
+
scorer = RotationAveraged(predictor, a.rotations) if a.rotations > 1 else predictor
|
| 204 |
+
report, _ = evaluate_records(records, scorer, a.out, heldout_sources=tuple(heldout), skip_overlong=skip_overlong)
|
| 205 |
+
report.update(suite_sha256=source_hash, data=a.data, date_facts=a.date_facts, rotations=a.rotations, run=a.run or a.remote, split=split,
|
| 206 |
+
calibration_applied=predictor.temperature != 1.0 if not a.remote else None,
|
| 207 |
+
environment=predictor.environment if not a.remote else None, # the kernel set (kev.predictors.kernel_environment)
|
| 208 |
+
remote={"base_url": a.remote, "requested_model": a.remote_model, "served_model": predictor.served_model, "concurrency": a.remote_concurrency} if a.remote else None)
|
| 209 |
+
write_json(Path(a.out) / "report.json", report)
|
| 210 |
+
print(json.dumps({"objective": report["objective"], "clean": report["clean"], "coverage": report["coverage"]}, indent=2))
|
| 211 |
+
|
| 212 |
+
|
| 213 |
+
if __name__ == "__main__":
|
| 214 |
+
main()
|
runtime/kev/budget.py
ADDED
|
@@ -0,0 +1,107 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""The resources a Modal trial container asks for and the admission bound on a study's cost: one home for
|
| 2 |
+
modal_app.run_trial/admit_study and kev.rounds.validate. Pure Python (no torch, no modal)."""
|
| 3 |
+
TRIAL_CPU, TRIAL_MEMORY = 4, (65536, 196608) # cores, (request, limit) MiB
|
| 4 |
+
GPU_HOURLY = {"H100": 3.95, "H200": 4.54, "B200": 6.25, "T4": 0.59} # USD per GPU hour
|
| 5 |
+
# Full-weight trials (kev.train --full_ft, kev.full_ft). One GPU keeps the fp32 masters and AdamW moments in host memory:
|
| 6 |
+
# 12 bytes per parameter, ~310 GB for a 27B, plus the staging buffers and the loader; the CPU runs the AdamW step. Several
|
| 7 |
+
# GPUs (FSDP2) keep that state on the GPUs; host memory holds the copy of it a resume point is written from in the
|
| 8 |
+
# background (~307 GB over the ranks) and the gathered checkpoint rank 0 writes (~51 GB; a snapshot being written holds
|
| 9 |
+
# the same, one at a time). Disk: the runs volume stages what a container writes on its local disk, up to two resume
|
| 10 |
+
# points (the new one is complete before the old one goes), the final checkpoint and every snapshot (kept, never deleted),
|
| 11 |
+
# beyond Modal's default 512 GiB quota. A run may plan at most MAX_SNAPSHOTS snapshots (kev.train and
|
| 12 |
+
# kev.experiment.validated_trial refuse more): with a 27B, 2 x 307 + (1 + 8) x 51.2 ≈ 1,075 GB inside the 1 TiB
|
| 13 |
+
# (1,099.5 GB), even if the volume keeps committed files staged. A full disk would fail a snapshot, withhold its resume
|
| 14 |
+
# point and fail the trial (failed.json, not retried), so the cap is a hard limit, not a guideline.
|
| 15 |
+
FULL_FT_SINGLE = 24, (368640, 409600) # 360 / 400 GiB
|
| 16 |
+
FULL_FT_SHARDED = 16, (409600, 471040) # 400 / 460 GiB
|
| 17 |
+
FULL_FT_DISK = 1048576 # MiB of ephemeral disk (1 TiB); Modal bills disk as memory at 20:1
|
| 18 |
+
MAX_SNAPSHOTS = 8 # snapshots one full-weight run may plan (fractions + every-N steps)
|
| 19 |
+
RESUME_POINT_GB, CHECKPOINT_GB = 307, 51.2 # a 27B's resume point (fp32 masters + moments) and bf16 checkpoint, the disk math above
|
| 20 |
+
# A study's limits. Full-weight trials may run to Modal's 24 h cap per attempt and are continued after a timeout (each
|
| 21 |
+
# continuation is a new call that resumes from the trial's last resume point, kev.experiment.continue_trial), so their
|
| 22 |
+
# bound counts every attempt and their budget cap covers one 8 x H200 day (a LoRA study keeps the $250 cap).
|
| 23 |
+
MAX_TIMEOUT = {False: 28800, True: 86400} # a 24 h attempt fits the $1,000 cap on one H200 (3 x 24 h x ~$11/h); on
|
| 24 |
+
# H200:8 (~$41/h) the cap allows ~8 h attempts, a day with the continuations
|
| 25 |
+
MAX_BUDGET = {False: 250, True: 1000}
|
| 26 |
+
# Continuations per full-weight trial: at most 1 + FULL_FT_RETRIES calls, each one attempt. They are spawned by
|
| 27 |
+
# `kev.rounds watch` (or modal_app.py::resume --trial), not by Modal: calls go out with Modal's retries off, because Modal
|
| 28 |
+
# charges a timed-out attempt twice when the container then outlives the 30 s cancellation grace (a timeout and a killed
|
| 29 |
+
# task; a full-weight trial always does, its trainer and volume commits run on), and the retry the kill triggers can start
|
| 30 |
+
# while another attempt of the same input is running. Round 22's trial got two 8 h attempts of the three its bound
|
| 31 |
+
# counted (scripts/modal_retry_probe.py: Retries(2) and killed attempts -> 2 attempts; retries=0 -> exactly 1).
|
| 32 |
+
FULL_FT_RETRIES = 2
|
| 33 |
+
# The attempt lease: a running full-weight attempt holds <study>/<trial>/attempt.json on the kev-leases volume
|
| 34 |
+
# (modal_app.TrialLease), refreshed every LEASE_HEARTBEAT seconds from a thread of its own and marked ended when the
|
| 35 |
+
# attempt exits cleanly. A lease that is not ended and whose heartbeat is younger than LEASE_STALE belongs to a container
|
| 36 |
+
# that may still be writing the trial: a new attempt refuses to start and continue_full_trial waits. LEASE_STALE = 900 s:
|
| 37 |
+
# 30 s Modal's cancellation grace after a timeout, during which the old container keeps running (and committing)
|
| 38 |
+
# + 523 s the longest resume-point commit of the runs volume seen (round 22: 206-523 s); the heartbeat does not wait for it,
|
| 39 |
+
# but a commit the old container started just before its kill is covered even if a heartbeat had waited behind it
|
| 40 |
+
# + 60 s one heartbeat interval (the last heartbeat can be that old when the container dies)
|
| 41 |
+
# = 613 s, rounded up to 15 min for clock skew between the containers and the launcher (seconds, NTP) and a few failed
|
| 42 |
+
# heartbeat commits in a row: a live attempt is mistaken for a dead one only after 15 heartbeats fail in a row.
|
| 43 |
+
# The cost is a continuation that starts up to ~15 min after a killed attempt, against 8 h attempts.
|
| 44 |
+
LEASE_HEARTBEAT, LEASE_STALE = 60, 900
|
| 45 |
+
CPU_HOURLY, MEMORY_GIB_HOURLY = 0.04730, 0.008 # USD per core hour and per GiB hour (memory billed at the limit here)
|
| 46 |
+
# Checkpoint interpolation (scripts/interpolate_checkpoint.py, modal_app.py::interpolate): CPU only. The base backbone stays
|
| 47 |
+
# resident in bf16 (51 GB for Qwen3.8-27B, 25.6B parameters; loading it through transformers also stages the 2.5 GB LM
|
| 48 |
+
# head); the SFT side streams tensor by tensor, one output shard (<= 5 GB) is held until written, and the fp32 arithmetic
|
| 49 |
+
# runs in 16M-element chunks (~200 MB). Peak ~60 GB; the limit leaves room for the loader's transients.
|
| 50 |
+
INTERPOLATE_CPU, INTERPOLATE_MEMORY, INTERPOLATE_TIMEOUT = 8, (81920, 131072), 10800 # cores, 80 / 128 GiB, 3 h
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
def gpu_count(gpu):
|
| 54 |
+
"""Modal's "H200:8" -> 8; a bare type is one GPU."""
|
| 55 |
+
return int(gpu.partition(":")[2] or 1)
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
def trial_resources(gpu, full_ft=False):
|
| 59 |
+
"""(cpu cores, (memory request, limit) MiB) for one trial container."""
|
| 60 |
+
if not full_ft: return TRIAL_CPU, TRIAL_MEMORY
|
| 61 |
+
return FULL_FT_SHARDED if gpu_count(gpu) > 1 else FULL_FT_SINGLE
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def trial_disk(full_ft=False):
|
| 65 |
+
"""ephemeral_disk (MiB) for one trial container; None keeps Modal's default."""
|
| 66 |
+
return FULL_FT_DISK if full_ft else None
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
def hourly_rate(gpu, full_ft=False):
|
| 70 |
+
"""USD per hour of one trial container of `gpu` ("H200", or "H200:8" for 8 in one container): GPU + CPU + memory at
|
| 71 |
+
the resources trial_resources gives it."""
|
| 72 |
+
kind = gpu.partition(":")[0]
|
| 73 |
+
if kind not in GPU_HOURLY or gpu_count(gpu) < 1: raise ValueError(f"invalid GPU {gpu!r}")
|
| 74 |
+
cpu, memory = trial_resources(gpu, full_ft)
|
| 75 |
+
return GPU_HOURLY[kind] * gpu_count(gpu) + cpu * CPU_HOURLY + (memory[1] + (trial_disk(full_ft) or 0) / 20) / 1024 * MEMORY_GIB_HOURLY
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
def interpolation_bound(timeout=INTERPOLATE_TIMEOUT, jobs=1):
|
| 79 |
+
"""The most `jobs` interpolation containers (INTERPOLATE_CPU, INTERPOLATE_MEMORY; no GPU, no retries) can cost at
|
| 80 |
+
`timeout` seconds each."""
|
| 81 |
+
if timeout <= 0 or jobs < 1: raise ValueError("invalid timeout or job count")
|
| 82 |
+
return (INTERPOLATE_CPU * CPU_HOURLY + INTERPOLATE_MEMORY[1] / 1024 * MEMORY_GIB_HOURLY) * timeout / 3600 * jobs
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def compute_bound(gpu, timeout, trials, full_ft=False):
|
| 86 |
+
"""The most `trials` containers can cost if each runs to `timeout` seconds, counting a full-weight trial's retries."""
|
| 87 |
+
if timeout <= 0 or trials < 1: raise ValueError("invalid GPU, timeout, or trial count")
|
| 88 |
+
return hourly_rate(gpu, full_ft) * timeout / 3600 * trials * (1 + FULL_FT_RETRIES if full_ft else 1)
|
| 89 |
+
|
| 90 |
+
|
| 91 |
+
def trial_attempts(spawn, label):
|
| 92 |
+
"""(attempts used, attempts allowed) for one trial of a spawn record (runs/<study>.spawn.json, modal_app.launch_detached).
|
| 93 |
+
A record written with Modal's retries off ("modal_retries": 0) keeps each trial's attempts in `attempts`, one per entry
|
| 94 |
+
({"nonce", "call", ...}; a pending entry, written before its spawn, and an abandoned one count too, since their call
|
| 95 |
+
may run), and allows compute_bound's count: 1 + FULL_FT_RETRIES for a full-weight trial, 1 otherwise. An older record
|
| 96 |
+
(Modal retried each call, and nothing local says how many attempts ran) is taken to have used them all."""
|
| 97 |
+
allowed = 1 + FULL_FT_RETRIES if spawn.get("full_ft") else 1
|
| 98 |
+
if spawn.get("modal_retries") != 0: return allowed, allowed
|
| 99 |
+
return max(1, len(spawn.get("attempts", {}).get(label) or [])), allowed
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
def lease_state(lease, now):
|
| 103 |
+
"""'none' (no lease), 'ended' (its attempt exited cleanly), 'fresh' (a heartbeat younger than LEASE_STALE: its container
|
| 104 |
+
may still be writing the trial) or 'stale' (no heartbeat for LEASE_STALE: the container is gone)."""
|
| 105 |
+
if not lease: return "none"
|
| 106 |
+
if lease.get("ended"): return "ended"
|
| 107 |
+
return "fresh" if now - lease["heartbeat"] < LEASE_STALE else "stale"
|
runtime/kev/calibrate.py
ADDED
|
@@ -0,0 +1,73 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Per-workload calibration report from a saved rows.json: what one temperature fitted on *these* rows would do.
|
| 2 |
+
|
| 3 |
+
uv run python -m kev.calibrate --rows runs/kev-9b-wanli-v1/rows.json [--out runs/kev-9b-wanli-v1/calibration.json]
|
| 4 |
+
|
| 5 |
+
A released checkpoint ships one temperature fitted on its in-distribution development rows (scripts/calibrate_checkpoint.py,
|
| 6 |
+
which now refuses such rows unless --allow-in-distribution: new temperatures come from held-out datasets). It
|
| 7 |
+
transfers to our out-of-domain suites, not necessarily to a deployer's workload: Kev-9B on WANLI-256 is served at
|
| 8 |
+
mean confidence 0.82 against accuracy 0.70. This module reports, on one rows file, four arms of the same predictions:
|
| 9 |
+
|
| 10 |
+
raw the checkpoint's logits at T=1 (restored from the recorded inference_temperature)
|
| 11 |
+
shipped as served (the checkpoint's temperature)
|
| 12 |
+
workload one temperature fitted on all these rows (in-sample; the number a deployer would write down)
|
| 13 |
+
workload_oof the same fit, group-disjoint out-of-fold: what fitting on your labelled rows buys on rows you did not fit
|
| 14 |
+
|
| 15 |
+
and a record-clustered paired bootstrap of workload_oof against shipped on ECE, Brier, coverage at <= 5% error and
|
| 16 |
+
AURC. Argmax and accuracy are identical in every arm by construction. Report only: nothing is written to a checkpoint.
|
| 17 |
+
"""
|
| 18 |
+
import argparse
|
| 19 |
+
from pathlib import Path
|
| 20 |
+
|
| 21 |
+
from kev.metrics import TEMPERATURE_FIT, fit_temperature, metrics, out_of_fold_rows, paired_bootstrap, raw_row, scored_rows, tempered_row
|
| 22 |
+
from kev.suite import read_json, write_json
|
| 23 |
+
|
| 24 |
+
KEYS = ("n", "acc", "ece", "brier", "nll", "mean_conf", "confident_error_rate", "coverage_at_5pct_error", "aurc")
|
| 25 |
+
COMPARED = ("ece", "brier", "coverage_at_5pct_error", "aurc")
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
def workload_report(rows, folds=5, seed=0, samples=1000):
|
| 29 |
+
served = scored_rows(rows)
|
| 30 |
+
if not served:
|
| 31 |
+
raise ValueError("no clean, knowable rows to calibrate")
|
| 32 |
+
if any("logits" not in row or "inference_temperature" not in row for row in served):
|
| 33 |
+
raise ValueError("rows without recorded logits and inference_temperature (remote or Jev predictions): a fit would be "
|
| 34 |
+
"relative to the served probabilities, not the checkpoint's raw logits")
|
| 35 |
+
raw = [raw_row(row) for row in served]
|
| 36 |
+
temperature = fit_temperature(raw, **TEMPERATURE_FIT)
|
| 37 |
+
oof, fold_temperatures = out_of_fold_rows(raw, folds, seed, **TEMPERATURE_FIT)
|
| 38 |
+
arms = {"raw": raw, "shipped": served, "workload": [tempered_row(row, temperature) for row in raw], "workload_oof": oof}
|
| 39 |
+
return {"n": len(served), "shipped_temperature": sorted({row.get("inference_temperature", 1.0) for row in served}),
|
| 40 |
+
"workload_temperature": temperature, "fold_temperatures": fold_temperatures,
|
| 41 |
+
"arms": {name: {k: metrics(arm)[k] for k in KEYS} for name, arm in arms.items()},
|
| 42 |
+
"oof_vs_shipped": {metric: paired_bootstrap(oof, served, samples=samples, seed=seed, metric=metric, aggregation="micro") for metric in COMPARED},
|
| 43 |
+
"fit": TEMPERATURE_FIT, "folds": folds, "seed": seed,
|
| 44 |
+
"note": "argmax and accuracy are temperature-invariant; workload is in-sample, workload_oof is the held-out estimate"}
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def format_report(report):
|
| 48 |
+
lines = [f"n={report['n']} shipped T={', '.join(f'{t:.2f}' for t in report['shipped_temperature'])} workload T={report['workload_temperature']:.2f}"
|
| 49 |
+
f" folds T=[{', '.join(f'{t:.2f}' for t in report['fold_temperatures'])}]",
|
| 50 |
+
f"{'arm':13}{'acc':>7}{'ece':>7}{'brier':>7}{'nll':>7}{'conf':>7}{'c-err':>7}{'cov@5%':>8}{'aurc':>7}"]
|
| 51 |
+
for name, m in report["arms"].items():
|
| 52 |
+
lines.append(f"{name:13}{m['acc']:7.3f}{m['ece']:7.3f}{m['brier']:7.3f}{m['nll']:7.3f}{m['mean_conf']:7.3f}{m['confident_error_rate']:7.3f}{m['coverage_at_5pct_error']:8.2f}{m['aurc']:7.3f}")
|
| 53 |
+
for metric, b in report["oof_vs_shipped"].items():
|
| 54 |
+
lo, hi = b["ci95"]
|
| 55 |
+
lines.append(f"workload_oof - shipped {metric:24}{b[f'micro_{metric}_delta']:+.3f} 95% CI [{lo:+.3f}, {hi:+.3f}]")
|
| 56 |
+
return "\n".join(lines)
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def main():
|
| 60 |
+
ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
|
| 61 |
+
ap.add_argument("--rows", required=True, help="rows.json written by kev.benchmark")
|
| 62 |
+
ap.add_argument("--out", help="where to write the report (default: calibration.json next to the rows)")
|
| 63 |
+
ap.add_argument("--folds", type=int, default=5); ap.add_argument("--seed", type=int, default=0); ap.add_argument("--samples", type=int, default=1000)
|
| 64 |
+
a = ap.parse_args()
|
| 65 |
+
report = workload_report(read_json(a.rows), folds=a.folds, seed=a.seed, samples=a.samples)
|
| 66 |
+
report["rows"] = a.rows
|
| 67 |
+
out = Path(a.out) if a.out else Path(a.rows).with_name("calibration.json")
|
| 68 |
+
write_json(out, report)
|
| 69 |
+
print(format_report(report)); print(f"wrote {out}")
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
if __name__ == "__main__":
|
| 73 |
+
main()
|
runtime/kev/checkpoint.py
ADDED
|
@@ -0,0 +1,392 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Trained checkpoints: a run directory or a Hub repo holding a LoRA adapter (or, for a full-weight run, the whole bf16
|
| 2 |
+
backbone), `head.pt` and the tokenizer.
|
| 3 |
+
|
| 4 |
+
Loader rule: `adapter_config.json` present -> a LoRA adapter on `meta.base` at `meta.base_revision`; no adapter and
|
| 5 |
+
`config.json` + `model*.safetensors` (save_pretrained of the backbone, `meta.weights == "full"`) -> the backbone is loaded from
|
| 6 |
+
the checkpoint directory itself, nothing is merged, in the dtype head.pt's `weights_dtype` names (it must match the `dtype`
|
| 7 |
+
save_pretrained wrote to config.json). The tokenizer always comes from the base (both layouts carry a copy).
|
| 8 |
+
|
| 9 |
+
This is the one place that knows the layout of `head.pt` and how a checkpoint becomes a `DecisionModel`:
|
| 10 |
+
`kev.serve`, `kev.benchmark`, `kev.train --init_from`, `kev.publish`, the scripts and the Hugging Face Space all go
|
| 11 |
+
through it. The Space vendors this file next to `model.py` and `api.py` (scripts/publish_space.sh), so it must not
|
| 12 |
+
import the data or suite modules at import time.
|
| 13 |
+
|
| 14 |
+
ck = Checkpoint("jaredpalmer/kev-4b") # or a local run directory; `@tag` pins a Hub revision
|
| 15 |
+
tok, model = ck.load("mps", LoadOptions.from_env())
|
| 16 |
+
ck.meta.temperature # the calibration the checkpoint carries
|
| 17 |
+
"""
|
| 18 |
+
import datetime
|
| 19 |
+
import importlib.util
|
| 20 |
+
import json
|
| 21 |
+
import os
|
| 22 |
+
import re
|
| 23 |
+
from dataclasses import dataclass, field
|
| 24 |
+
from pathlib import Path
|
| 25 |
+
|
| 26 |
+
import torch
|
| 27 |
+
|
| 28 |
+
from .model import DecisionModel, is_hybrid, load_tokenizer, pad_id
|
| 29 |
+
|
| 30 |
+
HUB_ID = re.compile(r"[\w.-]+/[\w.-]+(@[\w.-]+)?")
|
| 31 |
+
MTIME_FLOOR = datetime.datetime(2000, 1, 1, tzinfo=datetime.timezone.utc).timestamp() # older file times are zeroed, not dates (Checkpoint.release_date)
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def is_hub_id(run):
|
| 35 |
+
return not os.path.isdir(run) and HUB_ID.fullmatch(str(run)) is not None
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def resolve_run(run):
|
| 39 |
+
"""Local run directory as given, or a Hub repo id like jaredpalmer/kev-4b, optionally pinned to a revision or tag
|
| 40 |
+
with `@` (jaredpalmer/kev-4b@qwen3), downloaded to the HF cache. Returns a str path."""
|
| 41 |
+
if os.path.isdir(run):
|
| 42 |
+
return str(run)
|
| 43 |
+
from huggingface_hub import snapshot_download
|
| 44 |
+
repo, _, revision = str(run).partition("@")
|
| 45 |
+
return snapshot_download(repo, revision=revision or None, allow_patterns=["*.json", "*.safetensors", "*.pt", "*.txt", "*.jinja"])
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
@dataclass
|
| 49 |
+
class Meta:
|
| 50 |
+
"""Contents of `head.pt`. Every reader gets the same defaults for fields older checkpoints did not write.
|
| 51 |
+
`extra` keeps the rest of the file (training args, suite hash, init provenance, temperature fit) so a
|
| 52 |
+
read-modify-write round trip loses nothing."""
|
| 53 |
+
base: str
|
| 54 |
+
head: dict | None = None
|
| 55 |
+
base_revision: str | None = None
|
| 56 |
+
lora: int = 0
|
| 57 |
+
head_dim: int = 256
|
| 58 |
+
option_isolation: bool = False
|
| 59 |
+
special_embeddings: bool = False
|
| 60 |
+
weights_dtype: str = "fp32"
|
| 61 |
+
temperature: float = 1.0
|
| 62 |
+
holdout: list = field(default_factory=list)
|
| 63 |
+
weights: str = "lora" # lora on a base, full backbone, or packed MLX backbone (no base weights)
|
| 64 |
+
extra: dict = field(default_factory=dict)
|
| 65 |
+
|
| 66 |
+
KNOWN = ("base", "head", "base_revision", "lora", "head_dim", "option_isolation", "special_embeddings", "weights_dtype", "temperature", "holdout", "weights")
|
| 67 |
+
|
| 68 |
+
@classmethod
|
| 69 |
+
def from_dict(cls, d):
|
| 70 |
+
return cls(**{k: d[k] for k in cls.KNOWN if k in d}, extra={k: v for k, v in d.items() if k not in cls.KNOWN})
|
| 71 |
+
|
| 72 |
+
def to_dict(self):
|
| 73 |
+
return {**self.extra, **{k: getattr(self, k) for k in self.KNOWN}} # known fields win over a stray key in extra
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
def read_meta(run):
|
| 77 |
+
return Meta.from_dict(torch.load(f"{run}/head.pt", map_location="cpu"))
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
def write_meta(run, meta):
|
| 81 |
+
torch.save(meta.to_dict(), f"{run}/head.pt")
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
@dataclass(frozen=True)
|
| 85 |
+
class LoadOptions:
|
| 86 |
+
"""How a checkpoint is turned into a model. Defaults are the exact path every reported number uses; the fields
|
| 87 |
+
are the same knobs the KEV_* environment variables expose to the command-line tools (see from_env).
|
| 88 |
+
|
| 89 |
+
dtype None = fp32, the exact path every reported number uses (bf16 when the checkpoint was trained with a bf16
|
| 90 |
+
backbone). kev.serve defaults to bf16 on CUDA and MPS instead: half the memory, 2-4.5x lower latency on an
|
| 91 |
+
L4 (Kev-4B: 209 -> 118 ms at 101 tokens, 850 -> 189 ms at 330 tokens), probabilities within ~0.01 and
|
| 92 |
+
the same argmax on the checks run so far. KEV_DTYPE=fp32 restores the exact path when serving.
|
| 93 |
+
merge fold the LoRA into the base weights: the delta is computed from the fp32 adapter and added in fp32 with
|
| 94 |
+
one rounding to the load dtype, so a bf16 model holds exactly round(W + delta), the same bits as merging
|
| 95 |
+
an fp32 copy and casting, without the fp32 copy (Kev-9B needed 36 GB of GPU memory to load for that).
|
| 96 |
+
Identical because the Qwen bases are stored in bf16; a base stored in fp32 would be rounded twice.
|
| 97 |
+
Exact in fp32; in bf16 it is faster (~15%) and closer to the fp32 numbers than the unmerged adapter
|
| 98 |
+
(kev-4b, 24 dev records: max |dp| 0.017 vs 0.029, 0 vs 1 argmax flips). Ignored for adapters that carry
|
| 99 |
+
trained token embeddings. A checkpoint trained on a bf16 backbone (--weights_dtype bf16: Kev-27B) keeps
|
| 100 |
+
its adapter unmerged, as it was trained, unless `fused` is asked for (serving); then it is folded the same
|
| 101 |
+
way (Kev-27B served against unmerged: max |dp| 0.009, 0 flips in 280 questions, runs/fused-27b-h200).
|
| 102 |
+
attn attention backend; None = the model default (SDPA on CUDA, eager elsewhere). "sdpa" on MPS measured
|
| 103 |
+
parity with eager and is a few percent faster.
|
| 104 |
+
lora_scale WiSE-FT-style interpolation between base (0) and fine-tuned weights (1), at inference.
|
| 105 |
+
temperature None = the temperature the checkpoint carries (fitted by scripts/calibrate_checkpoint.py); 1.0 = raw logits.
|
| 106 |
+
backend None = torch, the path every reported number uses. "mlx" = kev.mlx_model (Metal kernels for the hybrid
|
| 107 |
+
Qwen3.5 backbones through mlx-lm; the pointer head and encoder are shared; refused for attention-only
|
| 108 |
+
bases, which MPS already runs well). "auto" = mlx when the device is mps, the base is hybrid, mlx-lm is
|
| 109 |
+
installed and fp32 was not asked for (an explicit dtype=float32 means "the exact path"), else torch;
|
| 110 |
+
kev.serve uses auto. The MLX path always merges a LoRA adapter, loads a full-weight checkpoint as
|
| 111 |
+
saved (no merge, no copy: kev.mlx_model.load_full), and ignores `attn` and `dtype` (the backbone runs
|
| 112 |
+
as stored: bf16 for every released checkpoint).
|
| 113 |
+
cuda_graphs replay the serving passes of a hybrid backbone on CUDA (state prefix, question rows on a cached state) as
|
| 114 |
+
CUDA graphs, batched across requests (kev.cuda_graphs, DecisionModel.probs_batch). None = off, the eager
|
| 115 |
+
path every reported number uses; kev.serve turns it on for CUDA. Exact up to floating-point
|
| 116 |
+
reassociation, not bit for bit (the passes are padded to buckets).
|
| 117 |
+
fused rewrite a merged hybrid backbone on CUDA with fused Triton kernels (kev.fused_qwen35; needs
|
| 118 |
+
flash-linear-attention fused_qwen35.FLA_VERSION and refuses any other). None = off; kev.serve turns it on
|
| 119 |
+
for CUDA when fused_available() (KEV_FUSED=0 to decline). Equal to the reference layers up to bf16 rounding.
|
| 120 |
+
"""
|
| 121 |
+
dtype: torch.dtype | None = None
|
| 122 |
+
merge: bool = True
|
| 123 |
+
attn: str | None = None
|
| 124 |
+
lora_scale: float = 1.0
|
| 125 |
+
temperature: float | None = None
|
| 126 |
+
backend: str | None = None
|
| 127 |
+
cuda_graphs: bool | None = None
|
| 128 |
+
fused: bool | None = None
|
| 129 |
+
|
| 130 |
+
BACKENDS = (None, "torch", "mlx", "auto")
|
| 131 |
+
|
| 132 |
+
@classmethod
|
| 133 |
+
def from_env(cls, env=os.environ):
|
| 134 |
+
"""KEV_DTYPE=bf16|fp16|fp32, KEV_MERGE=0, KEV_ATTN=sdpa|eager, KEV_LORA_SCALE, KEV_TEMPERATURE, KEV_BACKEND=torch|mlx|auto,
|
| 135 |
+
KEV_CUDA_GRAPHS=0|1, KEV_FUSED=0|1.
|
| 136 |
+
For command-line entry points only; library code passes an explicit LoadOptions. Explicit values that equal a
|
| 137 |
+
library default are kept (fp32 as torch.float32, "torch" as a string) so a caller with its own default, like
|
| 138 |
+
kev.serve, can tell "asked for it" from "did not say"."""
|
| 139 |
+
backend = env.get("KEV_BACKEND") or None
|
| 140 |
+
if backend not in cls.BACKENDS: raise ValueError(f"KEV_BACKEND must be one of torch, mlx, auto; got {backend!r}")
|
| 141 |
+
return cls(dtype={"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}.get(env.get("KEV_DTYPE", "")),
|
| 142 |
+
merge=env.get("KEV_MERGE", "1") != "0", attn=env.get("KEV_ATTN") or None,
|
| 143 |
+
lora_scale=float(env.get("KEV_LORA_SCALE", "1")),
|
| 144 |
+
temperature=float(env["KEV_TEMPERATURE"]) if env.get("KEV_TEMPERATURE") else None, backend=backend,
|
| 145 |
+
cuda_graphs={"0": False, "1": True}.get(env.get("KEV_CUDA_GRAPHS", "")),
|
| 146 |
+
fused={"0": False, "1": True}.get(env.get("KEV_FUSED", "")))
|
| 147 |
+
|
| 148 |
+
|
| 149 |
+
def mlx_available():
|
| 150 |
+
try:
|
| 151 |
+
import mlx_lm # noqa: F401
|
| 152 |
+
return True
|
| 153 |
+
except ImportError:
|
| 154 |
+
return False
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
def fused_available():
|
| 158 |
+
"""Whether kev.fused_qwen35 can run: flash-linear-attention importable at its FLA_VERSION (kev.serve's CUDA default;
|
| 159 |
+
not in the serve extra, pinned in the Modal images). An explicit fused=True skips this and fails loudly instead."""
|
| 160 |
+
if importlib.util.find_spec("fla") is None: return False
|
| 161 |
+
try:
|
| 162 |
+
import fla
|
| 163 |
+
from .fused_qwen35 import FLA_VERSION
|
| 164 |
+
except ImportError:
|
| 165 |
+
return False
|
| 166 |
+
return fla.__version__ == FLA_VERSION
|
| 167 |
+
|
| 168 |
+
|
| 169 |
+
class Checkpoint:
|
| 170 |
+
def __init__(self, run):
|
| 171 |
+
self.requested = str(run) # what the caller asked for (a Hub id stays a Hub id in labels)
|
| 172 |
+
self.path = resolve_run(run)
|
| 173 |
+
self.meta = read_meta(self.path)
|
| 174 |
+
|
| 175 |
+
def file(self, name):
|
| 176 |
+
return Path(self.path) / name
|
| 177 |
+
|
| 178 |
+
def adapter_config(self):
|
| 179 |
+
return json.loads(self.file("adapter_config.json").read_text(encoding="utf-8"))
|
| 180 |
+
|
| 181 |
+
@property
|
| 182 |
+
def packed(self):
|
| 183 |
+
found = self.file("packed.json").exists()
|
| 184 |
+
if found != (self.meta.weights == "packed"):
|
| 185 |
+
raise ValueError("packed.json and head.pt weights must agree")
|
| 186 |
+
return found
|
| 187 |
+
|
| 188 |
+
@property
|
| 189 |
+
def full(self):
|
| 190 |
+
"""The loader rule (module docstring): True for a full-weight checkpoint, False for a LoRA adapter; head.pt's
|
| 191 |
+
`weights` must agree with the files."""
|
| 192 |
+
if self.packed:
|
| 193 |
+
if self.file("adapter_config.json").exists() or not self.file("config.json").exists() or not self.file("model.safetensors").exists():
|
| 194 |
+
raise ValueError("packed checkpoint requires config.json and model.safetensors, without a PEFT adapter")
|
| 195 |
+
return False
|
| 196 |
+
found = "lora" if self.file("adapter_config.json").exists() else "full" if self.file("config.json").exists() and self.shards() else None
|
| 197 |
+
if found != self.meta.weights:
|
| 198 |
+
raise ValueError(f"{self.path}: head.pt says weights={self.meta.weights!r} but the directory holds "
|
| 199 |
+
f"{ {'lora': 'an adapter', 'full': 'backbone weights'}.get(found, 'neither an adapter nor backbone weights') }")
|
| 200 |
+
return found == "full"
|
| 201 |
+
|
| 202 |
+
def shards(self):
|
| 203 |
+
"""The backbone's safetensors files of a full-weight checkpoint (model.safetensors or model-*-of-*.safetensors)."""
|
| 204 |
+
return sorted(Path(self.path).glob("model*.safetensors"))
|
| 205 |
+
|
| 206 |
+
def weights_sha256(self):
|
| 207 |
+
"""What a run's provenance pins: the adapter file's sha256, or for full weights the sha256 over every shard's."""
|
| 208 |
+
from .suite import digest # lazy: the Space vendors this module without kev/suite.py
|
| 209 |
+
if not self.full and not self.packed: return digest(self.file("adapter_model.safetensors"))
|
| 210 |
+
import hashlib
|
| 211 |
+
return hashlib.sha256("".join(f"{p.name}:{digest(p)}\n" for p in self.shards()).encode()).hexdigest()
|
| 212 |
+
|
| 213 |
+
def release_date(self):
|
| 214 |
+
"""ISO date for the TypeSafe model card: the Hub commit date for a Hub checkpoint. For a local run, or a Hub
|
| 215 |
+
checkpoint offline (its cached files), the UTC date head.pt was written, else the newest of the checkpoint's
|
| 216 |
+
files, else "unknown". An mtime before 2000 is not a date: archives zero it for reproducibility (Kev 1.0's first
|
| 217 |
+
release tarballs did), and it reads as 1970-01-01 (1969-12-31 west of UTC). scripts/build_release_assets.py stamps
|
| 218 |
+
the release date instead."""
|
| 219 |
+
if is_hub_id(self.requested):
|
| 220 |
+
from huggingface_hub import HfApi
|
| 221 |
+
repo, _, revision = self.requested.partition("@")
|
| 222 |
+
try:
|
| 223 |
+
return HfApi().model_info(repo, revision=revision or None).last_modified.date().isoformat()
|
| 224 |
+
except Exception:
|
| 225 |
+
pass
|
| 226 |
+
files = sorted((p.stat().st_mtime for p in Path(self.path).iterdir() if p.is_file()), reverse=True)
|
| 227 |
+
stamp = next((t for t in [self.file("head.pt").stat().st_mtime, *files] if t >= MTIME_FLOOR), None)
|
| 228 |
+
return "unknown" if stamp is None else datetime.datetime.fromtimestamp(stamp, datetime.timezone.utc).date().isoformat()
|
| 229 |
+
|
| 230 |
+
def hybrid_base(self):
|
| 231 |
+
"""Whether the backbone has Gated DeltaNet layers (Qwen3.5), read from a config without loading weights: the base's
|
| 232 |
+
for a LoRA adapter, the checkpoint's own config.json for full weights (the backbone that is actually loaded)."""
|
| 233 |
+
from transformers import AutoConfig
|
| 234 |
+
if self.packed or self.full:
|
| 235 |
+
config = AutoConfig.from_pretrained(self.path, local_files_only=self.packed)
|
| 236 |
+
else:
|
| 237 |
+
config = AutoConfig.from_pretrained(self.meta.base, revision=self.meta.base_revision)
|
| 238 |
+
return is_hybrid(config.get_text_config())
|
| 239 |
+
|
| 240 |
+
def backend(self, device, opts=LoadOptions()):
|
| 241 |
+
"""The backend `load` will use: LoadOptions.backend resolved ("auto" -> mlx only where it pays and is installed)."""
|
| 242 |
+
if opts.backend not in LoadOptions.BACKENDS: raise ValueError(f"unknown backend {opts.backend!r}")
|
| 243 |
+
if self.packed:
|
| 244 |
+
if opts.backend not in (None, "auto", "mlx") or opts.dtype not in (None, torch.bfloat16):
|
| 245 |
+
raise ValueError("packed checkpoints require the MLX backend with bf16 activations")
|
| 246 |
+
return "mlx"
|
| 247 |
+
if opts.backend != "auto": return opts.backend or "torch"
|
| 248 |
+
exact = opts.dtype is torch.float32 # KEV_DTYPE=fp32: the caller wants the reported-numbers path, not a faster one
|
| 249 |
+
return "mlx" if str(device) == "mps" and not exact and mlx_available() and self.hybrid_base() else "torch"
|
| 250 |
+
|
| 251 |
+
def load(self, device, opts=LoadOptions()):
|
| 252 |
+
"""-> (tokenizer, model) in eval mode with the LoRA applied (or the full backbone loaded) and the pointer head loaded. The model is a
|
| 253 |
+
DecisionModel (torch) or an MLXDecisionModel (backend mlx); both expose the same scoring interface."""
|
| 254 |
+
meta = self.meta
|
| 255 |
+
if self.packed:
|
| 256 |
+
self.full # validate the standalone layout before any tensor loading
|
| 257 |
+
self.backend(device, opts)
|
| 258 |
+
if opts.lora_scale != 1 or meta.option_isolation or meta.special_embeddings:
|
| 259 |
+
raise ValueError("packed checkpoints do not support interpolation, isolated options, or trained token embeddings")
|
| 260 |
+
from transformers import AutoTokenizer
|
| 261 |
+
from .packed_model import load_packed_model
|
| 262 |
+
tok = AutoTokenizer.from_pretrained(self.path, local_files_only=True)
|
| 263 |
+
m = load_packed_model(self.path, pad_id(tok), meta.head_dim)
|
| 264 |
+
m.head.load_state_dict(meta.head); m.eval()
|
| 265 |
+
m.head.temperature = meta.temperature if opts.temperature is None else opts.temperature
|
| 266 |
+
return tok, m
|
| 267 |
+
if self.full and opts.lora_scale != 1: raise ValueError("lora_scale interpolates an adapter; a full-weight checkpoint has none")
|
| 268 |
+
tok = load_tokenizer(meta.base, revision=meta.base_revision)
|
| 269 |
+
m = self._load_mlx(tok, opts) if self.backend(device, opts) == "mlx" else self._load_torch(tok, device, opts)
|
| 270 |
+
m.head.load_state_dict(meta.head); m.eval()
|
| 271 |
+
m.head.temperature = meta.temperature if opts.temperature is None else opts.temperature
|
| 272 |
+
return tok, m
|
| 273 |
+
|
| 274 |
+
def _load_mlx(self, tok, opts):
|
| 275 |
+
"""-> MLXDecisionModel. A LoRA checkpoint: the base's mlx-lm model with the adapter merged (fp32, one rounding). A
|
| 276 |
+
full-weight checkpoint: mlx-lm's model built from the checkpoint's own config and shards in the dtype head.pt names
|
| 277 |
+
(kev.mlx_model.load_full; strict names/shapes/dtype), nothing merged. Refusals come before the mlx-lm import, which
|
| 278 |
+
would otherwise hide them on a machine without it."""
|
| 279 |
+
full = self.full
|
| 280 |
+
if not full and not opts.merge: raise ValueError("the MLX backend always merges the adapter (KEV_MERGE=0 needs backend=torch)")
|
| 281 |
+
if self.meta.option_isolation: raise ValueError("option_isolation needs the packed mask; not available on the MLX backend")
|
| 282 |
+
if not self.hybrid_base(): raise ValueError(f"the MLX backend is for the hybrid (Qwen3.5) bases; {self.path if full else self.meta.base} is attention-only and runs on MPS with backend=torch")
|
| 283 |
+
dtype = self.saved_dtype() if full else None
|
| 284 |
+
from .mlx_model import MLXDecisionModel, load_base, load_full, merge_lora
|
| 285 |
+
if full:
|
| 286 |
+
lm = load_full(self.path, self.shards(), dtype)
|
| 287 |
+
else:
|
| 288 |
+
lm = load_base(resolve_run(f"{self.meta.base}@{self.meta.base_revision or ''}")) # the base snapshot the torch path already cached
|
| 289 |
+
merge_lora(lm, self.path, opts.lora_scale)
|
| 290 |
+
return MLXDecisionModel(lm, pad_id(tok), head_dim=self.meta.head_dim)
|
| 291 |
+
|
| 292 |
+
def _load_torch(self, tok, device, opts):
|
| 293 |
+
m, merged = self._full_torch(tok, device, opts) if self.full else self._adapted_torch(tok, device, opts)
|
| 294 |
+
serving = str(device).startswith("cuda") and m.hybrid
|
| 295 |
+
if opts.fused and serving and merged: # fused projections need plain (merged or full) weights
|
| 296 |
+
from .fused_qwen35 import fuse
|
| 297 |
+
fuse(m.lm)
|
| 298 |
+
if opts.cuda_graphs and serving:
|
| 299 |
+
from .cuda_graphs import CudaGraphs
|
| 300 |
+
m.graphs = CudaGraphs(m.lm, m.pad_id)
|
| 301 |
+
return m
|
| 302 |
+
|
| 303 |
+
SAVED_DTYPES = {"bf16": "bfloat16", "fp32": "float32"} # head.pt weights_dtype -> the dtype save_pretrained writes to config.json
|
| 304 |
+
|
| 305 |
+
def saved_dtype(self):
|
| 306 |
+
"""A full-weight checkpoint's dtype name ("bfloat16" | "float32"): the one head.pt's `weights_dtype` names (bf16 for
|
| 307 |
+
every kev.train --full_ft run: the dtype they were trained in), which must be the dtype save_pretrained recorded in
|
| 308 |
+
config.json; otherwise a mislabelled export would be silently cast (fp32 weights rounded to bf16, or bf16 upcast to
|
| 309 |
+
twice the memory). Both backends load full weights in it."""
|
| 310 |
+
cfg = json.loads(self.file("config.json").read_text(encoding="utf-8"))
|
| 311 |
+
expected, saved = self.SAVED_DTYPES.get(self.meta.weights_dtype), cfg.get("dtype") or cfg.get("torch_dtype")
|
| 312 |
+
if expected is None or saved not in (None, expected):
|
| 313 |
+
raise ValueError(f"{self.path}: config.json records the weights as {saved} but head.pt says weights_dtype={self.meta.weights_dtype!r}")
|
| 314 |
+
return expected
|
| 315 |
+
|
| 316 |
+
def _full_torch(self, tok, device, opts):
|
| 317 |
+
"""-> (model, True). Full weights load in `saved_dtype`; an explicit dtype still casts on purpose (fp32: the same
|
| 318 |
+
values computed in fp32). Nothing to merge (`load` refuses lora_scale for full weights on either backend)."""
|
| 319 |
+
meta = self.meta
|
| 320 |
+
return DecisionModel(meta.base, tok, device, head_dim=meta.head_dim, option_isolation=meta.option_isolation,
|
| 321 |
+
dtype=opts.dtype or getattr(torch, self.saved_dtype()), attn=opts.attn, weights=self.path), True
|
| 322 |
+
|
| 323 |
+
def _adapted_torch(self, tok, device, opts):
|
| 324 |
+
"""-> (model, whether the adapter was merged): the base with this checkpoint's LoRA."""
|
| 325 |
+
from peft import PeftModel
|
| 326 |
+
meta = self.meta
|
| 327 |
+
dtype, merge = opts.dtype or torch.float32, opts.merge
|
| 328 |
+
if meta.weights_dtype == "bf16":
|
| 329 |
+
# trained with a bf16 backbone (--weights_dtype bf16: Kev-27B, the 35B-A3B MoE whose fused experts need bf16):
|
| 330 |
+
# load it the same way. The exact path keeps the fp32 adapter unmerged; the fused serving path folds it in
|
| 331 |
+
# (one rounding of W + delta, as for every served Kev; parity in runs/serving-27b-*).
|
| 332 |
+
dtype, merge = torch.bfloat16, merge and bool(opts.fused)
|
| 333 |
+
merge = merge and not self.adapter_config().get("trainable_token_indices") # token-trained adapters stay unmerged
|
| 334 |
+
m = DecisionModel(meta.base, tok, device, lora=None, revision=meta.base_revision, head_dim=meta.head_dim,
|
| 335 |
+
option_isolation=meta.option_isolation, dtype=dtype, attn=opts.attn)
|
| 336 |
+
m.lm = PeftModel.from_pretrained(m.lm, self.path, torch_device=str(device)).to(device) # trainable token embeddings, if any, live in the adapter
|
| 337 |
+
if opts.lora_scale != 1:
|
| 338 |
+
for module in m.lm.modules():
|
| 339 |
+
if isinstance(getattr(module, "scaling", None), dict):
|
| 340 |
+
for k in module.scaling: module.scaling[k] *= opts.lora_scale
|
| 341 |
+
m.lora_scale = opts.lora_scale
|
| 342 |
+
if merge: m.lm = m.lm.merge_and_unload() # W += delta: fp32 math, one rounding (see LoadOptions.merge)
|
| 343 |
+
if dtype != torch.float32: m.lm = m.lm.to(dtype)
|
| 344 |
+
return m, merge
|
| 345 |
+
|
| 346 |
+
COMPAT_FIELDS = ("base", "base_revision", "lora", "head_dim", "option_isolation", "special_embeddings", "weights")
|
| 347 |
+
|
| 348 |
+
def warm_start(self, model, ours):
|
| 349 |
+
"""Delta training: load this checkpoint's weights (adapter, or full backbone) and pointer head into `model` (a fresh
|
| 350 |
+
DecisionModel built the way `ours` says: with LoRA, or for full-weight training without). `ours` is the Meta the new
|
| 351 |
+
run will save; every architecture field is compared BEFORE loading, because peft and load_state_dict(strict=False)
|
| 352 |
+
load matching keys silently and a half-loaded model still trains and still reports a loss. Returns provenance."""
|
| 353 |
+
from .suite import digest # lazy: the Space vendors this module without kev/suite.py
|
| 354 |
+
for name in self.COMPAT_FIELDS:
|
| 355 |
+
theirs, mine = getattr(self.meta, name), getattr(ours, name)
|
| 356 |
+
if theirs != mine and not (name == "base_revision" and None in (theirs, mine)):
|
| 357 |
+
raise ValueError(f"--init_from {self.path}: {name} is {theirs!r} there and {mine!r} here")
|
| 358 |
+
tensors = self._load_backbone_into(model.lm) if self.full else self._load_adapter_into(model.lm)
|
| 359 |
+
model.head.load_state_dict(self.meta.head)
|
| 360 |
+
return {"init_from": self.requested, "resolved": self.path, "weights_sha256": self.weights_sha256(),
|
| 361 |
+
"head_sha256": digest(self.file("head.pt")), "tensors": tensors}
|
| 362 |
+
|
| 363 |
+
def _load_backbone_into(self, lm):
|
| 364 |
+
"""Copy the saved backbone over `lm` shard by shard (the base's copy in memory is replaced, never merged); every
|
| 365 |
+
tensor of `lm` must be covered exactly once. -> tensor count."""
|
| 366 |
+
from safetensors.torch import load_file
|
| 367 |
+
have, seen = set(lm.state_dict()), set()
|
| 368 |
+
for shard in self.shards():
|
| 369 |
+
part = load_file(shard)
|
| 370 |
+
unexpected = sorted(set(part) - have)
|
| 371 |
+
if unexpected: raise ValueError(f"--init_from {self.path}: {shard.name} carries tensors this backbone does not have (e.g. {unexpected[:2]})")
|
| 372 |
+
lm.load_state_dict(part, strict=False); seen |= set(part)
|
| 373 |
+
missing = sorted(have - seen)
|
| 374 |
+
if missing: raise ValueError(f"--init_from {self.path} does not cover {len(missing)} of this backbone's tensors (e.g. {missing[:2]})")
|
| 375 |
+
return len(seen)
|
| 376 |
+
|
| 377 |
+
def _load_adapter_into(self, lm):
|
| 378 |
+
from peft import get_peft_model_state_dict, load_peft_weights, set_peft_model_state_dict
|
| 379 |
+
weights = load_peft_weights(self.path, device="cpu")
|
| 380 |
+
have = set(get_peft_model_state_dict(lm))
|
| 381 |
+
unexpected, missing = sorted(set(weights) - have), sorted(have - set(weights))
|
| 382 |
+
if unexpected:
|
| 383 |
+
raise ValueError(f"--init_from {self.path} carries {len(unexpected)} adapter tensors this model does not have (e.g. {unexpected[:2]}); check --lora_targets / --lora against its adapter_config.json")
|
| 384 |
+
if missing:
|
| 385 |
+
raise ValueError(f"--init_from {self.path} does not cover {len(missing)} of this model's adapter tensors (e.g. {missing[:2]}); check --lora_targets")
|
| 386 |
+
set_peft_model_state_dict(lm, weights)
|
| 387 |
+
return len(weights)
|
| 388 |
+
|
| 389 |
+
|
| 390 |
+
def load(run, device, opts=LoadOptions()):
|
| 391 |
+
"""Convenience: Checkpoint(run).load(device, opts)."""
|
| 392 |
+
return Checkpoint(run).load(device, opts)
|
runtime/kev/compare.py
ADDED
|
@@ -0,0 +1,62 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import argparse
|
| 2 |
+
import json
|
| 3 |
+
import math
|
| 4 |
+
from collections import defaultdict
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
|
| 7 |
+
import numpy as np
|
| 8 |
+
|
| 9 |
+
from kev.metrics import paired_bootstrap
|
| 10 |
+
from kev.suite import read_json, write_json
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
def nll_sensitivity(rows, floor):
|
| 14 |
+
tasks = defaultdict(list)
|
| 15 |
+
for row in rows:
|
| 16 |
+
if row["variant"] == "clean":
|
| 17 |
+
tasks[row["task"]].append(-math.log(max(row["p"][row["label"]], floor)))
|
| 18 |
+
return float(np.mean([np.mean(v) for v in tasks.values()]))
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
def none_diagnostics(rows):
|
| 22 |
+
out = {}
|
| 23 |
+
for variant in ("none_present", "none_absent"):
|
| 24 |
+
subset = [r for r in rows if r["variant"] == variant]
|
| 25 |
+
values = [r["p"][r["keys"].index("none_of_these")] for r in subset]
|
| 26 |
+
out[variant] = {"n": len(values), "mean_p_none": float(np.mean(values)),
|
| 27 |
+
"p_none_above_half": float(np.mean(np.array(values) > .5))}
|
| 28 |
+
return out
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
def main():
|
| 32 |
+
ap = argparse.ArgumentParser()
|
| 33 |
+
ap.add_argument("--candidate", required=True)
|
| 34 |
+
ap.add_argument("--reference", required=True)
|
| 35 |
+
ap.add_argument("--out", required=True)
|
| 36 |
+
a = ap.parse_args()
|
| 37 |
+
candidate, reference = Path(a.candidate), Path(a.reference)
|
| 38 |
+
left, right = [read_json(p / "report.json") for p in (candidate, reference)]
|
| 39 |
+
if not left.get("suite_sha256") or left["suite_sha256"] != right.get("suite_sha256"):
|
| 40 |
+
raise ValueError("comparison requires matching frozen-suite hashes")
|
| 41 |
+
lr, rr = [read_json(p / "rows.json") for p in (candidate, reference)]
|
| 42 |
+
result = {"candidate": str(candidate), "reference": str(reference), "suite_sha256": left["suite_sha256"],
|
| 43 |
+
"clean": {"candidate": left["clean"], "reference": right["clean"]},
|
| 44 |
+
"paired": {metric: paired_bootstrap(lr, rr, metric=metric) for metric in ("nll", "acc", "brier")},
|
| 45 |
+
"nll_floor_sensitivity": {str(floor): {"candidate_macro_nll": nll_sensitivity(lr, floor),
|
| 46 |
+
"reference_macro_nll": nll_sensitivity(rr, floor)} for floor in (1e-3, 1e-6, 1e-9)},
|
| 47 |
+
"none_of_the_above": {"candidate": none_diagnostics(lr), "reference": none_diagnostics(rr)},
|
| 48 |
+
"caveats": ["Uncalibrated returned distributions; probabilities normalized for scoring; raw responses preserved.",
|
| 49 |
+
"Historical kev checkpoint trained on all six sources; its heldout-source status does not transfer to this study.",
|
| 50 |
+
"Jev uses a hosted alias and rounds some probabilities to zero; NLL depends on the declared floor.",
|
| 51 |
+
"CPU local latency and hosted network latency are not comparable hardware benchmarks.",
|
| 52 |
+
"This is a development-suite comparison, not the locked final test."]}
|
| 53 |
+
write_json(a.out, result)
|
| 54 |
+
print("task kev acc Jev acc kev NLL Jev NLL")
|
| 55 |
+
for task in left["tasks"]:
|
| 56 |
+
l, r = left["tasks"][task], right["tasks"][task]
|
| 57 |
+
print(f"{task:20} {l['acc']:8.3f} {r['acc']:10.3f} {l['nll']:10.3f} {r['nll']:10.3f}")
|
| 58 |
+
print(json.dumps({"paired": result["paired"], "none_of_the_above": result["none_of_the_above"]}, indent=2))
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
if __name__ == "__main__":
|
| 62 |
+
main()
|
runtime/kev/composition.py
ADDED
|
@@ -0,0 +1,321 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import copy
|
| 2 |
+
import hashlib
|
| 3 |
+
import json
|
| 4 |
+
import random
|
| 5 |
+
from datetime import date, timedelta
|
| 6 |
+
|
| 7 |
+
UNKNOWN = None
|
| 8 |
+
SHAPES = {
|
| 9 |
+
"atom": 0,
|
| 10 |
+
"negation": ("not", 0),
|
| 11 |
+
"conjunction": ("and", 0, 1),
|
| 12 |
+
"disjunction": ("or", 0, 1),
|
| 13 |
+
"exception": ("unless", 0, 1),
|
| 14 |
+
"conditional": ("if", 0, 1, 2),
|
| 15 |
+
"nested_and": ("and", ("and", 0, 1), 2),
|
| 16 |
+
"nested_or": ("or", 0, ("or", 1, 2)),
|
| 17 |
+
"held_and_or": ("and", ("or", 0, 1), 2),
|
| 18 |
+
"held_or_not": ("or", ("and", 0, 1), ("not", 2)),
|
| 19 |
+
"held_conditional": ("if", 0, ("not", 1), 2),
|
| 20 |
+
"final_combination": ("and", ("if", 0, 1, 2), 3),
|
| 21 |
+
"final_negation": ("not", ("or", ("and", 0, 1), 2)),
|
| 22 |
+
"final_exception": ("or", ("unless", 0, 1), 2),
|
| 23 |
+
}
|
| 24 |
+
TRAIN_SHAPES = tuple(list(SHAPES)[:8])
|
| 25 |
+
DEV_SHAPES = tuple(list(SHAPES)[8:11])
|
| 26 |
+
TEST_SHAPES = tuple(list(SHAPES)[11:])
|
| 27 |
+
KINDS = ("lt", "le", "gt", "ge", "eq", "range", "match", "elapsed", "flag")
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
def atom_value(atom, facts):
|
| 31 |
+
if any(key not in facts for key in atom["fields"]):
|
| 32 |
+
return UNKNOWN
|
| 33 |
+
values = [facts[k] for k in atom["fields"]]
|
| 34 |
+
kind, threshold = atom["kind"], atom["threshold"]
|
| 35 |
+
x = values[0]
|
| 36 |
+
if kind == "lt": return x < threshold
|
| 37 |
+
if kind == "le": return x <= threshold
|
| 38 |
+
if kind == "gt": return x > threshold
|
| 39 |
+
if kind == "ge": return x >= threshold
|
| 40 |
+
if kind == "eq": return x == threshold
|
| 41 |
+
if kind == "range": return threshold <= x <= threshold + 10
|
| 42 |
+
if kind == "match": return x == values[1]
|
| 43 |
+
if kind == "elapsed": return (date.fromisoformat(values[1]) - date.fromisoformat(x)).days <= threshold
|
| 44 |
+
if kind == "flag": return x
|
| 45 |
+
raise ValueError(f"unknown atom {kind}")
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def evaluate_rule(tree, atoms, facts):
|
| 49 |
+
if isinstance(tree, int):
|
| 50 |
+
return atom_value(atoms[tree], facts)
|
| 51 |
+
op, *children = tree
|
| 52 |
+
values = [evaluate_rule(t, atoms, facts) for t in children]
|
| 53 |
+
if op == "not": return None if values[0] is None else not values[0]
|
| 54 |
+
if op == "unless":
|
| 55 |
+
a, exception = values
|
| 56 |
+
values = [a, None if exception is None else not exception]
|
| 57 |
+
op = "and"
|
| 58 |
+
if op == "and":
|
| 59 |
+
return False if False in values else None if None in values else True
|
| 60 |
+
if op == "or":
|
| 61 |
+
return True if True in values else None if None in values else False
|
| 62 |
+
if op == "if":
|
| 63 |
+
condition, yes, no = values
|
| 64 |
+
return yes if condition is True else no if condition is False else yes if yes == no else None
|
| 65 |
+
raise ValueError(f"unknown operation {op}")
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def atom_text(atom):
|
| 69 |
+
kind, fields, t = atom["kind"], atom["fields"], atom["threshold"]
|
| 70 |
+
key = fields[0]
|
| 71 |
+
words = {"lt": "less than", "le": "at most", "gt": "greater than", "ge": "at least", "eq": "equal to"}
|
| 72 |
+
if kind in words: return f"{key} is {words[kind]} {t}"
|
| 73 |
+
if kind == "range": return f"{key} is between {t} and {t + 10}, including both endpoints"
|
| 74 |
+
if kind == "match": return f"{fields[0]} is the same person as {fields[1]}"
|
| 75 |
+
if kind == "elapsed": return f"the elapsed days from {fields[0]} to {fields[1]} are at most {t}"
|
| 76 |
+
return f"{key} is yes"
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
def render_rule(tree, atoms, style):
|
| 80 |
+
"""Five surface styles for the same rule tree. 0 = logic-like, 1 = 'both/at least one of' lists, 2 = reserved for the
|
| 81 |
+
locked test, 3 = plain prose, 4 = clause per line."""
|
| 82 |
+
if isinstance(tree, int): return atom_text(atoms[tree])
|
| 83 |
+
op, *children = tree
|
| 84 |
+
parts = [render_rule(c, atoms, style) for c in children]
|
| 85 |
+
if style == 3:
|
| 86 |
+
if op == "not": return f"the condition \"{parts[0]}\" fails"
|
| 87 |
+
if op == "and": return f"{parts[0]}, and also {parts[1]}"
|
| 88 |
+
if op == "or": return f"either {parts[0]}, or else {parts[1]}"
|
| 89 |
+
if op == "unless": return f"{parts[0]}, except when {parts[1]}"
|
| 90 |
+
return f"when {parts[0]} the requirement is that {parts[1]}, and when it is not the requirement is that {parts[2]}"
|
| 91 |
+
if style == 4:
|
| 92 |
+
if op == "not": return f"[NOT: {parts[0]}]"
|
| 93 |
+
if op in ("and", "or"): return f"[{'ALL' if op == 'and' else 'ANY'} of: {parts[0]} | {parts[1]}]"
|
| 94 |
+
if op == "unless": return f"[{parts[0]} UNLESS {parts[1]}]"
|
| 95 |
+
return f"[IF {parts[0]} THEN {parts[1]} ELSE {parts[2]}]"
|
| 96 |
+
if op == "not": return f"NOT ({parts[0]})" if style == 0 else f"it is not the case that ({parts[0]})"
|
| 97 |
+
if op in ("and", "or"):
|
| 98 |
+
if style == 0: return f"({parts[0]}) {op.upper()} ({parts[1]})"
|
| 99 |
+
connector = "both" if op == "and" else "at least one of"
|
| 100 |
+
return f"{connector} these conditions hold: [({parts[0]}); ({parts[1]})]"
|
| 101 |
+
if op == "unless": return f"({parts[0]}) holds and the exception ({parts[1]}) does not hold"
|
| 102 |
+
return f"if ({parts[0]}), use ({parts[1]}); otherwise use ({parts[2]})"
|
| 103 |
+
|
| 104 |
+
|
| 105 |
+
# --- random rule trees (PLAN.md: structural diversity, negation anywhere) ---------------------------------------------
|
| 106 |
+
|
| 107 |
+
def skeleton(tree):
|
| 108 |
+
"""Structure with leaf identities erased: used to order commutative children independently of atom numbering."""
|
| 109 |
+
if isinstance(tree, int): return "_"
|
| 110 |
+
op, *children = tree
|
| 111 |
+
parts = [skeleton(c) for c in children]
|
| 112 |
+
if op in ("and", "or"): parts = sorted(parts)
|
| 113 |
+
return f"{op}({','.join(parts)})"
|
| 114 |
+
|
| 115 |
+
|
| 116 |
+
def sort_commutative(tree):
|
| 117 |
+
if isinstance(tree, int): return tree
|
| 118 |
+
op, *children = tree
|
| 119 |
+
children = [sort_commutative(c) for c in children]
|
| 120 |
+
if op in ("and", "or"): children = sorted(children, key=skeleton)
|
| 121 |
+
return (op, *children)
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
def canonical(tree):
|
| 125 |
+
"""Structure key: commutative children ordered by skeleton, then leaves renumbered in traversal order, so
|
| 126 |
+
(A and B) or not C and not C or (B and A) share one key."""
|
| 127 |
+
t = relabel(sort_commutative(tree))
|
| 128 |
+
def render(t):
|
| 129 |
+
if isinstance(t, int): return str(t)
|
| 130 |
+
return f"{t[0]}({','.join(render(c) for c in t[1:])})"
|
| 131 |
+
return render(t)
|
| 132 |
+
|
| 133 |
+
|
| 134 |
+
def push_negation(tree):
|
| 135 |
+
"""De Morgan normal form, so a random tree equivalent to a held-out shape under negation pushing is also excluded."""
|
| 136 |
+
if isinstance(tree, int): return tree
|
| 137 |
+
op, *children = tree
|
| 138 |
+
if op == "not":
|
| 139 |
+
inner = children[0]
|
| 140 |
+
if isinstance(inner, int): return tree
|
| 141 |
+
iop, *ic = inner
|
| 142 |
+
if iop == "not": return push_negation(ic[0])
|
| 143 |
+
if iop in ("and", "or"): return ("or" if iop == "and" else "and", *[push_negation(("not", c)) for c in ic])
|
| 144 |
+
if iop == "unless": return push_negation(("or", ("not", ic[0]), ic[1]))
|
| 145 |
+
return ("not", push_negation(inner))
|
| 146 |
+
if op == "unless": return ("and", push_negation(children[0]), push_negation(("not", children[1])))
|
| 147 |
+
return (op, *[push_negation(c) for c in children])
|
| 148 |
+
|
| 149 |
+
|
| 150 |
+
def relabel(tree):
|
| 151 |
+
"""Renumber leaves in first-appearance order so structure keys do not depend on which atom index was drawn."""
|
| 152 |
+
mapping = {}
|
| 153 |
+
def walk(t):
|
| 154 |
+
if isinstance(t, int):
|
| 155 |
+
mapping.setdefault(t, len(mapping)); return mapping[t]
|
| 156 |
+
return (t[0], *[walk(c) for c in t[1:]])
|
| 157 |
+
return walk(tree)
|
| 158 |
+
|
| 159 |
+
|
| 160 |
+
def structure_keys(tree):
|
| 161 |
+
return {canonical(tree), canonical(push_negation(tree))}
|
| 162 |
+
|
| 163 |
+
|
| 164 |
+
HELD_OUT_KEYS = set().union(*(structure_keys(SHAPES[s]) for s in DEV_SHAPES + TEST_SHAPES))
|
| 165 |
+
|
| 166 |
+
|
| 167 |
+
def random_tree(rng, depth, next_leaf):
|
| 168 |
+
"""Random rule over {and, or, not, unless, if} with negation allowed at any position; leaves are fresh atom indices."""
|
| 169 |
+
if depth == 0 or (depth < 3 and rng.random() < 0.25):
|
| 170 |
+
return next_leaf()
|
| 171 |
+
op = rng.choice(["and", "or", "not", "unless", "if", "and", "or"])
|
| 172 |
+
if op == "not": return ("not", random_tree(rng, depth - 1, next_leaf))
|
| 173 |
+
if op == "if": return ("if", random_tree(rng, depth - 1, next_leaf), random_tree(rng, depth - 1, next_leaf), random_tree(rng, depth - 1, next_leaf))
|
| 174 |
+
return (op, random_tree(rng, depth - 1, next_leaf), random_tree(rng, depth - 1, next_leaf))
|
| 175 |
+
|
| 176 |
+
|
| 177 |
+
def sample_trees(n, seed, min_leaves=2, max_leaves=4, exclude=HELD_OUT_KEYS):
|
| 178 |
+
"""n distinct rule structures (by canonical key), none structurally equal to a held-out or locked shape, none a
|
| 179 |
+
bare leaf, each with negation somewhere in ~60% of cases. Deterministic in `seed`."""
|
| 180 |
+
rng = random.Random(f"{seed}:trees"); trees, keys = [], set()
|
| 181 |
+
for _ in range(20000):
|
| 182 |
+
if len(trees) >= n: break
|
| 183 |
+
counter = [0]
|
| 184 |
+
def next_leaf():
|
| 185 |
+
counter[0] += 1; return counter[0] - 1
|
| 186 |
+
t = random_tree(rng, rng.choice([2, 2, 3]), next_leaf)
|
| 187 |
+
if isinstance(t, int) or not min_leaves <= counter[0] <= max_leaves: continue
|
| 188 |
+
k = structure_keys(t)
|
| 189 |
+
if k & exclude or k & keys: continue
|
| 190 |
+
keys |= k; trees.append(relabel(t))
|
| 191 |
+
if len(trees) < n: raise ValueError(f"only {len(trees)} distinct structures found")
|
| 192 |
+
return trees
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
def leaf_indices(tree):
|
| 196 |
+
if isinstance(tree, int): return {tree}
|
| 197 |
+
return set().union(*(leaf_indices(t) for t in tree[1:]))
|
| 198 |
+
|
| 199 |
+
|
| 200 |
+
def make_atoms(tree, rng):
|
| 201 |
+
nouns = rng.sample(["request", "account", "package", "review", "member", "shipment", "entry", "case"], 4)
|
| 202 |
+
atoms = []
|
| 203 |
+
for i in range(max(leaf_indices(tree)) + 1):
|
| 204 |
+
kind = rng.choice(KINDS)
|
| 205 |
+
prefix = nouns[i]
|
| 206 |
+
fields = [f"{prefix} value"]
|
| 207 |
+
if kind == "match": fields = [f"{prefix} signer", f"{prefix} designated approver"]
|
| 208 |
+
elif kind == "elapsed": fields = [f"{prefix} start date", f"{prefix} end date"]
|
| 209 |
+
elif kind == "flag": fields = [f"{prefix} verified"]
|
| 210 |
+
atoms.append({"kind": kind, "fields": fields, "threshold": rng.randint(5, 60)})
|
| 211 |
+
return atoms
|
| 212 |
+
|
| 213 |
+
|
| 214 |
+
def fact_domains(atoms, rng):
|
| 215 |
+
domains = {}
|
| 216 |
+
for a in atoms:
|
| 217 |
+
kind, fs, t = a["kind"], a["fields"], a["threshold"]
|
| 218 |
+
if kind == "match":
|
| 219 |
+
people = rng.sample(["Mira", "Noah", "Aiko", "Ravi", "Sana", "Elin", "Tomas", "Kofi"], 3)
|
| 220 |
+
domains.update({k: people for k in fs})
|
| 221 |
+
elif kind == "elapsed":
|
| 222 |
+
day = date(2027, rng.randint(1, 8), rng.randint(1, 28))
|
| 223 |
+
domains[fs[0]] = [day.isoformat()]
|
| 224 |
+
domains[fs[1]] = [(day + timedelta(days=n)).isoformat() for n in (max(0, t - 1), t, t + 1, t + 10)]
|
| 225 |
+
elif kind == "flag": domains[fs[0]] = [False, True]
|
| 226 |
+
else: domains[fs[0]] = [t - 1, t, t + 1, t + 10, t + 11]
|
| 227 |
+
domains["routing reference"] = [rng.randint(100, 500), rng.randint(501, 999)]
|
| 228 |
+
return domains
|
| 229 |
+
|
| 230 |
+
|
| 231 |
+
def rendered_facts(facts, order):
|
| 232 |
+
def value(v):
|
| 233 |
+
return "yes" if v is True else "no" if v is False else str(v)
|
| 234 |
+
return [f"The {k} is {value(facts[k])}." for k in order]
|
| 235 |
+
|
| 236 |
+
|
| 237 |
+
POLICY_WRAPPERS = {
|
| 238 |
+
0: "Approve exactly when {rule}. Otherwise deny. The routing reference does not affect eligibility.",
|
| 239 |
+
1: "Approve exactly when {rule}. Otherwise deny. The routing reference does not affect eligibility.",
|
| 240 |
+
2: "Approval requires the following rule to be true: {rule}. A false rule means denial. Routing references are irrelevant.",
|
| 241 |
+
3: "A case is approved when {rule}; any other case is denied. Routing references play no part in the decision.",
|
| 242 |
+
4: "DECISION RULE {rule} -> approve; otherwise reject. Ignore the routing reference.",
|
| 243 |
+
}
|
| 244 |
+
|
| 245 |
+
|
| 246 |
+
def generate(groups_per_shape, seed, shapes=TRAIN_SHAPES, styles=(0, 1), source="compositional", trees=None):
|
| 247 |
+
"""trees: optional {name: tree} to generate from random structures instead of SHAPES (names must not collide)."""
|
| 248 |
+
records = []
|
| 249 |
+
shape_trees = trees if trees is not None else {s: SHAPES[s] for s in shapes}
|
| 250 |
+
for shape, tree in shape_trees.items():
|
| 251 |
+
rng = random.Random(f"{seed}:{shape}")
|
| 252 |
+
for i in range(groups_per_shape):
|
| 253 |
+
atoms = make_atoms(tree, rng)
|
| 254 |
+
domains = fact_domains(atoms, rng)
|
| 255 |
+
for attempt in range(1000):
|
| 256 |
+
facts = {k: rng.choice(v) for k, v in domains.items()}
|
| 257 |
+
label = evaluate_rule(tree, atoms, facts)
|
| 258 |
+
changed = None
|
| 259 |
+
keys = list(domains); rng.shuffle(keys)
|
| 260 |
+
for key in keys:
|
| 261 |
+
for value in domains[key]:
|
| 262 |
+
edited = {**facts, key: value}
|
| 263 |
+
if evaluate_rule(tree, atoms, edited) != label:
|
| 264 |
+
missing = {k: v for k, v in facts.items() if k != key}
|
| 265 |
+
if evaluate_rule(tree, atoms, missing) is None:
|
| 266 |
+
changed = (edited, key)
|
| 267 |
+
break
|
| 268 |
+
if changed: break
|
| 269 |
+
if changed: break
|
| 270 |
+
if changed is None:
|
| 271 |
+
raise ValueError(f"cannot create a decisive edit for {shape}")
|
| 272 |
+
edited, deciding = changed
|
| 273 |
+
nuisance = {**facts, "routing reference": next(v for v in domains["routing reference"] if v != facts["routing reference"])}
|
| 274 |
+
order = list(facts); rng.shuffle(order)
|
| 275 |
+
style = rng.choice(styles)
|
| 276 |
+
rule = render_rule(tree, atoms, style)
|
| 277 |
+
policy = POLICY_WRAPPERS[style].format(rule=rule)
|
| 278 |
+
keys = ["accept", "reject"]; rng.shuffle(keys)
|
| 279 |
+
criteria = {k: "The policy permits this case" if k == "accept" else "The policy does not permit this case" for k in keys}
|
| 280 |
+
group = f"composition/{seed}/{shape}/{i}"
|
| 281 |
+
for kind, a, b in (("relevant", facts, edited), ("irrelevant", facts, nuisance)):
|
| 282 |
+
for sibling, values in (("a", a), ("b", b)):
|
| 283 |
+
result = evaluate_rule(tree, atoms, values)
|
| 284 |
+
state = {"policy": policy, "case": " ".join(rendered_facts(values, order))}
|
| 285 |
+
identifier = f"{group}/{kind}/{sibling}"
|
| 286 |
+
records.append({"state": state, "questions": {"decision": {"type": "choice",
|
| 287 |
+
"instructions": "Apply the policy to this case.", "criteria": dict(criteria),
|
| 288 |
+
"label": "accept" if result else "reject", "src": f"composition_{shape}"}},
|
| 289 |
+
"_meta": {"id": identifier, "group_id": group, "source": source, "variant": "clean",
|
| 290 |
+
"pair_id": f"{group}/{kind}", "sibling": sibling, "pair_kind": kind,
|
| 291 |
+
"family": shape, "family_id": shape, "render_style": style, "structure": canonical(tree),
|
| 292 |
+
"text_sha256": hashlib.sha256(json.dumps(state, sort_keys=True).encode()).hexdigest(),
|
| 293 |
+
"certificate": {"tree": tree, "atoms": atoms, "facts": values, "order": order,
|
| 294 |
+
"deciding_field": deciding, "label": result}}})
|
| 295 |
+
return records
|
| 296 |
+
|
| 297 |
+
|
| 298 |
+
def check_group(records):
|
| 299 |
+
if len(records) != 4: raise ValueError("a composition group requires four records")
|
| 300 |
+
for r in records:
|
| 301 |
+
c = r["_meta"]["certificate"]
|
| 302 |
+
if evaluate_rule(c["tree"], c["atoms"], c["facts"]) != c["label"]:
|
| 303 |
+
raise ValueError("incorrect certificate label")
|
| 304 |
+
if r["state"]["case"] != " ".join(rendered_facts(c["facts"], c["order"])):
|
| 305 |
+
raise ValueError("rendered facts differ from certificate")
|
| 306 |
+
q = r["questions"]["decision"]
|
| 307 |
+
if q["label"] != ("accept" if c["label"] else "reject"):
|
| 308 |
+
raise ValueError("answer differs from certificate")
|
| 309 |
+
for a, b in (records[:2], records[2:]):
|
| 310 |
+
ca, cb = a["_meta"]["certificate"], b["_meta"]["certificate"]
|
| 311 |
+
if ca["order"] != cb["order"] or a["state"]["policy"] != b["state"]["policy"]:
|
| 312 |
+
raise ValueError("pair changes more than one fact")
|
| 313 |
+
if sum(ca["facts"][k] != cb["facts"][k] for k in ca["facts"]) != 1:
|
| 314 |
+
raise ValueError("pair must change exactly one fact")
|
| 315 |
+
expected_flip = a["_meta"]["pair_kind"] == "relevant"
|
| 316 |
+
if (ca["label"] != cb["label"]) != expected_flip:
|
| 317 |
+
raise ValueError("incorrect intervention label")
|
| 318 |
+
missing = {k: v for k, v in ca["facts"].items() if k != ca["deciding_field"]}
|
| 319 |
+
if expected_flip and evaluate_rule(ca["tree"], ca["atoms"], missing) is not None:
|
| 320 |
+
raise ValueError("deciding evidence ablation did not remove the answer")
|
| 321 |
+
return True
|
runtime/kev/contrastive.py
ADDED
|
@@ -0,0 +1,306 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Programmatic contrastive pairs with certain labels (PLAN.md step 2).
|
| 2 |
+
|
| 3 |
+
Each item is a policy plus a case described by a few sentences. Every sentence carries zero or more explicit facts;
|
| 4 |
+
the label is a pure function `evaluate(facts)` of the facts present, returning UNDETERMINED when a required fact is
|
| 5 |
+
missing. A pair is the same item with exactly one sentence changed so that the label flips.
|
| 6 |
+
|
| 7 |
+
Two code-level checks run on every pair before it is kept (no model anywhere):
|
| 8 |
+
ablation : removing any evidence sentence (one that carries a required fact) makes the label UNDETERMINED,
|
| 9 |
+
so the answer provably lives in the state, not in the instruction, the options, or the policy alone;
|
| 10 |
+
invariance : removing any filler sentence leaves the label unchanged.
|
| 11 |
+
Evidence status is derived from the evaluator, not declared. Both siblings share `pair_id`/`family_id`, and
|
| 12 |
+
splits are assigned by family so siblings never straddle a split.
|
| 13 |
+
"""
|
| 14 |
+
import hashlib
|
| 15 |
+
import json
|
| 16 |
+
import random
|
| 17 |
+
from datetime import date, timedelta
|
| 18 |
+
|
| 19 |
+
UNDETERMINED = "UNDETERMINED"
|
| 20 |
+
|
| 21 |
+
NAMES = ["Mira", "Noah", "Priya", "Tomas", "Aiko", "Lena", "Omar", "Sana", "Jonas", "Ravi", "Elin", "Kofi"]
|
| 22 |
+
ROLES = ["account owner", "billing manager", "support agent", "warehouse lead"]
|
| 23 |
+
ITEMS = ["a pair of running shoes", "a desk lamp", "a wireless keyboard", "a rain jacket", "a coffee grinder", "a backpack"]
|
| 24 |
+
PROGRAMS = ["the volunteer driver program", "the apprenticeship", "the rental agreement", "the night-shift roster"]
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def _day(d):
|
| 28 |
+
return d.strftime("%B %-d, %Y")
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
def _need(facts, *keys):
|
| 32 |
+
return all(k in facts for k in keys)
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
# A family returns (item_a, item_b). An item is {policy, sentences: [(text, facts)], evaluate, question}.
|
| 36 |
+
# The two items must differ in exactly one sentence, and evaluate() must give different labels.
|
| 37 |
+
|
| 38 |
+
def family_return_window(rng):
|
| 39 |
+
window = rng.choice([14, 30, 45, 60]); item = rng.choice(ITEMS); name = rng.choice(NAMES)
|
| 40 |
+
bought = date(2026, rng.randint(1, 9), rng.randint(1, 28))
|
| 41 |
+
def evaluate(f):
|
| 42 |
+
return (f["request"] - f["purchase"]).days <= window if _need(f, "request", "purchase") else UNDETERMINED
|
| 43 |
+
def build(days):
|
| 44 |
+
request = bought + timedelta(days=days)
|
| 45 |
+
return {"policy": f"Returns are accepted only if the return request is submitted within {window} days of the purchase date.",
|
| 46 |
+
"sentences": [(f"{name} bought {item} on {_day(bought)}.", {"purchase": bought}),
|
| 47 |
+
(f"The return request was submitted on {_day(request)}.", {"request": request}),
|
| 48 |
+
(f"The order was paid by card and shipped to {name}'s home address.", {})],
|
| 49 |
+
"evaluate": evaluate, "question": {"type": "noul", "instructions": "Is this return request within the policy window?"}}
|
| 50 |
+
return build(rng.randint(1, window - 1)), build(window + rng.randint(1, 30))
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
def family_spend_threshold(rng):
|
| 54 |
+
limit = rng.choice([250, 500, 1000, 2500]); name = rng.choice(NAMES); role = rng.choice(ROLES)
|
| 55 |
+
def evaluate(f):
|
| 56 |
+
return ("auto_approved" if f["amount"] <= limit else "director_signoff") if _need(f, "amount") else UNDETERMINED
|
| 57 |
+
def build(amount):
|
| 58 |
+
return {"policy": f"Expense claims of ${limit:,} or less are approved automatically. Claims above ${limit:,} require director sign-off.",
|
| 59 |
+
"sentences": [(f"{name}, the {role}, submitted an expense claim.", {}),
|
| 60 |
+
(f"The claim total is ${amount:,}.", {"amount": amount}),
|
| 61 |
+
("Receipts were attached for every line item.", {})],
|
| 62 |
+
"evaluate": evaluate, "question": {"type": "choice", "instructions": "How is this claim handled under the policy?",
|
| 63 |
+
"criteria": {"auto_approved": "Approved without further review", "director_signoff": "Requires director sign-off", "rejected": "Rejected outright"}}}
|
| 64 |
+
return build(limit - rng.randint(1, limit // 2)), build(limit + rng.randint(1, limit))
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
def family_authorization(rng):
|
| 68 |
+
approver, other = rng.sample(NAMES, 2); account = rng.randint(10, 99); amount = rng.choice([40, 120, 350, 900])
|
| 69 |
+
def evaluate(f):
|
| 70 |
+
return f["signer"] == f["approver"] if _need(f, "signer", "approver") else UNDETERMINED
|
| 71 |
+
def build(signer):
|
| 72 |
+
return {"policy": "A refund is authorized only when its sole authorization was signed by someone who may authorize refunds for that account.",
|
| 73 |
+
"sentences": [(f"Only {approver} may authorize refunds for account {account}.", {"approver": approver}),
|
| 74 |
+
(f"The sole authorization for this refund on account {account} was signed by {signer}.", {"signer": signer}),
|
| 75 |
+
(f"The refund amount is ${amount}.", {})],
|
| 76 |
+
"evaluate": evaluate, "question": {"type": "noul", "instructions": "Is the refund authorized?"}}
|
| 77 |
+
return build(approver), build(other)
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
def family_age_eligibility(rng):
|
| 81 |
+
minimum = rng.choice([16, 18, 21, 25]); name = rng.choice(NAMES); program = rng.choice(PROGRAMS)
|
| 82 |
+
def evaluate(f):
|
| 83 |
+
return f["age"] >= minimum if _need(f, "age") else UNDETERMINED
|
| 84 |
+
def build(age):
|
| 85 |
+
return {"policy": f"Applicants must be at least {minimum} years old to be eligible for {program}.",
|
| 86 |
+
"sentences": [(f"{name} applied to join {program}.", {}),
|
| 87 |
+
(f"{name} is {age} years old.", {"age": age}),
|
| 88 |
+
("The application form was complete and signed.", {})],
|
| 89 |
+
"evaluate": evaluate, "question": {"type": "noul", "instructions": "Is the applicant eligible?"}}
|
| 90 |
+
return build(minimum + rng.randint(0, 20)), build(minimum - rng.randint(1, 5))
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
def family_quantity_limit(rng):
|
| 94 |
+
limit = rng.choice([2, 3, 5, 10]); item = rng.choice(ITEMS); name = rng.choice(NAMES)
|
| 95 |
+
def evaluate(f):
|
| 96 |
+
if not _need(f, "qty"): return UNDETERMINED
|
| 97 |
+
return "within_limit" if f["qty"] <= limit else "slightly_over" if f["qty"] <= 2 * limit else "far_over"
|
| 98 |
+
def build(qty):
|
| 99 |
+
return {"policy": f"Customers may order at most {limit} units of any single item per order. Orders up to double the limit are held for review; larger orders are cancelled.",
|
| 100 |
+
"sentences": [(f"{name} placed an order for {item}.", {}),
|
| 101 |
+
(f"The order quantity is {qty}.", {"qty": qty}),
|
| 102 |
+
("Delivery was requested to a residential address.", {})],
|
| 103 |
+
"evaluate": evaluate, "question": {"type": "choice", "instructions": "What happens to this order?",
|
| 104 |
+
"criteria": {"within_limit": "Processed normally", "slightly_over": "Held for review", "far_over": "Cancelled"}}}
|
| 105 |
+
return build(rng.randint(1, limit)), build(rng.choice([rng.randint(limit + 1, 2 * limit), rng.randint(2 * limit + 1, 4 * limit)]))
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
def family_deadline(rng):
|
| 109 |
+
name = rng.choice(NAMES); due = date(2026, rng.randint(2, 11), rng.randint(1, 28)); grace = rng.choice([3, 7, 14])
|
| 110 |
+
def evaluate(f):
|
| 111 |
+
if not _need(f, "received", "due"): return UNDETERMINED
|
| 112 |
+
late = (f["received"] - f["due"]).days
|
| 113 |
+
return 0 if late <= 0 else 1 if late <= grace else 2
|
| 114 |
+
def build(offset):
|
| 115 |
+
received = due + timedelta(days=offset)
|
| 116 |
+
return {"policy": f"Reports received by the deadline are on time. Reports received within {grace} days after the deadline are late but accepted. Later reports are refused.",
|
| 117 |
+
"sentences": [(f"The filing deadline for {name}'s report was {_day(due)}.", {"due": due}),
|
| 118 |
+
(f"The report was received on {_day(received)}.", {"received": received}),
|
| 119 |
+
("The report was submitted through the online portal.", {})],
|
| 120 |
+
"evaluate": evaluate, "question": {"type": "score", "instructions": "How late is this report?", "criteria": ["On time", "Late but accepted", "Refused"]}}
|
| 121 |
+
return build(-rng.randint(0, 10)), build(rng.choice([rng.randint(1, grace), grace + rng.randint(1, 20)]))
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
# --- ordinal threshold families as Score questions (PLAN.md: the deadline failure is "no idea how ordinal thresholds work").
|
| 125 |
+
# Each has three ordered levels split by two thresholds; one is date-based like deadline but with a different template.
|
| 126 |
+
|
| 127 |
+
def family_warranty_claim(rng):
|
| 128 |
+
name = rng.choice(NAMES); item = rng.choice(ITEMS); bought = date(2026, rng.randint(1, 6), rng.randint(1, 28))
|
| 129 |
+
standard, extended = rng.choice([(90, 365), (180, 730), (365, 1095)]); part = rng.choice(['stitching', 'battery', 'housing', 'zipper', 'switch'])
|
| 130 |
+
def evaluate(f):
|
| 131 |
+
if not _need(f, "claim", "purchase"): return UNDETERMINED
|
| 132 |
+
age = (f["claim"] - f["purchase"]).days
|
| 133 |
+
return 0 if age <= standard else 1 if age <= extended else 2
|
| 134 |
+
def build(age):
|
| 135 |
+
claim = bought + timedelta(days=age)
|
| 136 |
+
return {"policy": f"Warranty claims made within {standard} days of purchase are covered in full. Claims made after that but within {extended} days are covered at half cost. Later claims are not covered.",
|
| 137 |
+
"sentences": [(f"{name} purchased {item} on {_day(bought)}.", {"purchase": bought}),
|
| 138 |
+
(f"A warranty claim for it was filed on {_day(claim)}.", {"claim": claim}),
|
| 139 |
+
(f"The claim describes a defect in the {part}.", {})],
|
| 140 |
+
"evaluate": evaluate, "question": {"type": "score", "instructions": "How is this claim covered?", "criteria": ["Covered in full", "Covered at half cost", "Not covered"]}}
|
| 141 |
+
a = rng.choice([rng.randint(1, standard), rng.randint(standard + 1, extended), extended + rng.randint(1, 200)])
|
| 142 |
+
b = rng.choice([x for x in [rng.randint(1, standard), rng.randint(standard + 1, extended), extended + rng.randint(1, 200)] if evaluate({"claim": bought + timedelta(days=x), "purchase": bought}) != evaluate({"claim": bought + timedelta(days=a), "purchase": bought})] or [a])
|
| 143 |
+
return build(a), build(b)
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
def family_sla_response(rng):
|
| 147 |
+
name = rng.choice(NAMES); target, breach = rng.choice([(4, 24), (8, 48), (24, 72), (1, 8)])
|
| 148 |
+
unit = "hours"; topic = rng.choice(['a login failure', 'a duplicate charge', 'a missing invoice', 'an export error']); queue = rng.choice(['email', 'chat', 'phone'])
|
| 149 |
+
def evaluate(f):
|
| 150 |
+
if not _need(f, "hours"): return UNDETERMINED
|
| 151 |
+
return 0 if f["hours"] <= target else 1 if f["hours"] <= breach else 2
|
| 152 |
+
def build(h):
|
| 153 |
+
return {"policy": f"Support responses within {target} {unit} meet the service level. Responses after {target} but within {breach} {unit} are a minor breach. Anything slower is a major breach.",
|
| 154 |
+
"sentences": [(f"{name} opened a priority ticket about {topic}.", {}),
|
| 155 |
+
(f"The first response arrived {h} {unit} after the ticket was opened.", {"hours": h}),
|
| 156 |
+
(f"The ticket was routed through the {queue} queue.", {})],
|
| 157 |
+
"evaluate": evaluate, "question": {"type": "score", "instructions": "How does this response time rate against the service level?", "criteria": ["Met", "Minor breach", "Major breach"]}}
|
| 158 |
+
levels = [rng.randint(1, target), rng.randint(target + 1, breach), breach + rng.randint(1, 100)]
|
| 159 |
+
a, b = rng.sample(levels, 2)
|
| 160 |
+
return build(a), build(b)
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
def family_late_fee(rng):
|
| 164 |
+
name = rng.choice(NAMES); grace, cap = rng.choice([(5, 30), (10, 60), (15, 45)]); amount = rng.choice([120, 450, 980, 2300]); method = rng.choice(['bank transfer', 'card', 'cheque'])
|
| 165 |
+
def evaluate(f):
|
| 166 |
+
if not _need(f, "days_late"): return UNDETERMINED
|
| 167 |
+
return 0 if f["days_late"] <= grace else 1 if f["days_late"] <= cap else 2
|
| 168 |
+
def build(d):
|
| 169 |
+
return {"policy": f"Invoices paid within {grace} days after the due date incur no fee. Payments between {grace + 1} and {cap} days late incur a 2% fee. Payments later than {cap} days incur a 10% fee and a hold on the account.",
|
| 170 |
+
"sentences": [(f"{name}'s invoice for ${amount:,} fell due last quarter.", {}),
|
| 171 |
+
(f"Payment was received {d} days after the due date.", {"days_late": d}),
|
| 172 |
+
(f"The payment was made by {method}.", {})],
|
| 173 |
+
"evaluate": evaluate, "question": {"type": "score", "instructions": "Which fee tier applies?", "criteria": ["No fee", "2% fee", "10% fee and account hold"]}}
|
| 174 |
+
levels = [rng.randint(0, grace), rng.randint(grace + 1, cap), cap + rng.randint(1, 60)]
|
| 175 |
+
a, b = rng.sample(levels, 2)
|
| 176 |
+
return build(a), build(b)
|
| 177 |
+
|
| 178 |
+
|
| 179 |
+
def family_volume_discount(rng):
|
| 180 |
+
name = rng.choice(NAMES); item = rng.choice(ITEMS); t1, t2 = rng.choice([(10, 50), (25, 100), (5, 20), (100, 500)]); dest = rng.choice(['warehouse', 'storefront', 'branch office'])
|
| 181 |
+
def evaluate(f):
|
| 182 |
+
if not _need(f, "units"): return UNDETERMINED
|
| 183 |
+
return 0 if f["units"] < t1 else 1 if f["units"] < t2 else 2
|
| 184 |
+
def build(u):
|
| 185 |
+
return {"policy": f"Orders of fewer than {t1} units are charged the list price. Orders of {t1} to {t2 - 1} units receive the volume discount. Orders of {t2} units or more receive the wholesale rate.",
|
| 186 |
+
"sentences": [(f"{name} placed a business order for {item}.", {}),
|
| 187 |
+
(f"The order is for {u} units.", {"units": u}),
|
| 188 |
+
(f"Delivery is to a {dest}.", {})],
|
| 189 |
+
"evaluate": evaluate, "question": {"type": "score", "instructions": "Which pricing tier applies?", "criteria": ["List price", "Volume discount", "Wholesale rate"]}}
|
| 190 |
+
levels = [rng.randint(1, t1 - 1), rng.randint(t1, t2 - 1), t2 + rng.randint(0, t2)]
|
| 191 |
+
a, b = rng.sample(levels, 2)
|
| 192 |
+
return build(a), build(b)
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
def family_shipping_delay(rng):
|
| 196 |
+
"""Day-precision date arithmetic with small thresholds (the skill `deadline` tests): promised vs delivered dates."""
|
| 197 |
+
name = rng.choice(NAMES); item = rng.choice(ITEMS); promised = date(2026, rng.randint(1, 11), rng.randint(1, 28))
|
| 198 |
+
minor, major = rng.choice([(1, 4), (2, 7), (3, 10)]); carrier = rng.choice(["the courier", "the postal service", "a freight partner"])
|
| 199 |
+
def evaluate(f):
|
| 200 |
+
if not _need(f, "delivered", "promised"): return UNDETERMINED
|
| 201 |
+
late = (f["delivered"] - f["promised"]).days
|
| 202 |
+
return 0 if late <= minor else 1 if late <= major else 2
|
| 203 |
+
def build(offset):
|
| 204 |
+
delivered = promised + timedelta(days=offset)
|
| 205 |
+
return {"policy": f"Deliveries up to {minor} day{'s' if minor > 1 else ''} after the promised date count as on time. Deliveries {minor + 1} to {major} days after it are a minor delay and earn a shipping refund. Later deliveries are a major delay and earn a full refund.",
|
| 206 |
+
"sentences": [(f"{name} ordered {item} with delivery promised for {_day(promised)}.", {"promised": promised}),
|
| 207 |
+
(f"The parcel was delivered on {_day(delivered)}.", {"delivered": delivered}),
|
| 208 |
+
(f"It was shipped by {carrier}.", {})],
|
| 209 |
+
"evaluate": evaluate, "question": {"type": "score", "instructions": "How is this delivery classified?", "criteria": ["On time", "Minor delay: shipping refund", "Major delay: full refund"]}}
|
| 210 |
+
levels = [rng.randint(-3, minor), rng.randint(minor + 1, major), major + rng.randint(1, 30)]
|
| 211 |
+
a, b = rng.sample(levels, 2)
|
| 212 |
+
return build(a), build(b)
|
| 213 |
+
|
| 214 |
+
|
| 215 |
+
FAMILIES = {"return_window": family_return_window, "spend_threshold": family_spend_threshold, "authorization": family_authorization,
|
| 216 |
+
"age_eligibility": family_age_eligibility, "quantity_limit": family_quantity_limit, "deadline": family_deadline,
|
| 217 |
+
"warranty_claim": family_warranty_claim, "sla_response": family_sla_response, "late_fee": family_late_fee, "volume_discount": family_volume_discount,
|
| 218 |
+
"shipping_delay": family_shipping_delay}
|
| 219 |
+
ORDINAL_FAMILIES = ("warranty_claim", "sla_response", "late_fee", "volume_discount", "shipping_delay") # trainable Score-threshold families (deadline stays held out)
|
| 220 |
+
|
| 221 |
+
|
| 222 |
+
def label_of(item, drop=None):
|
| 223 |
+
facts = {}
|
| 224 |
+
for i, (_, f) in enumerate(item["sentences"]):
|
| 225 |
+
if i != drop: facts.update(f)
|
| 226 |
+
return item["evaluate"](facts)
|
| 227 |
+
|
| 228 |
+
|
| 229 |
+
def check_pair(a, b):
|
| 230 |
+
"""None if the pair is valid, else the reason it is rejected."""
|
| 231 |
+
la, lb = label_of(a), label_of(b)
|
| 232 |
+
if UNDETERMINED in (la, lb): return "label_undetermined"
|
| 233 |
+
if la == lb: return "labels_equal"
|
| 234 |
+
if sum(x[0] != y[0] for x, y in zip(a["sentences"], b["sentences"])) != 1 or len(a["sentences"]) != len(b["sentences"]):
|
| 235 |
+
return "not_exactly_one_sentence_differs"
|
| 236 |
+
for item, label in ((a, la), (b, lb)):
|
| 237 |
+
evidence = 0
|
| 238 |
+
for i, (_, facts) in enumerate(item["sentences"]):
|
| 239 |
+
got = label_of(item, drop=i)
|
| 240 |
+
if facts and got != UNDETERMINED: return "ablation_failed" # evidence removed must make it undeterminable
|
| 241 |
+
if not facts and got != label: return "invariance_failed" # filler removed must not change the label
|
| 242 |
+
evidence += bool(facts)
|
| 243 |
+
if evidence < 1: return "no_evidence_sentence"
|
| 244 |
+
return None
|
| 245 |
+
|
| 246 |
+
|
| 247 |
+
def to_request(item, family, pair_id, sibling, rng):
|
| 248 |
+
order = list(range(len(item["sentences"]))); rng.shuffle(order)
|
| 249 |
+
sentences = [item["sentences"][i][0] for i in order]
|
| 250 |
+
q = {**item["question"], "label": label_of(item), "src": f"contrastive_{family}"}
|
| 251 |
+
return {"state": {"policy": item["policy"], "case": " ".join(sentences)}, "questions": {"decision": q},
|
| 252 |
+
"_meta": {"source": "contrastive", "family": family, "family_id": f"{family}/{pair_id}", "pair_id": pair_id, "sibling": sibling,
|
| 253 |
+
"repo": None, "revision": None, "split": "generated", "row": pair_id, "id": f"contrastive/{family}/{pair_id}/{sibling}",
|
| 254 |
+
"text_sha256": hashlib.sha256(" ".join(sentences).casefold().encode()).hexdigest(),
|
| 255 |
+
"row_sha256": hashlib.sha256(json.dumps([t for t, _ in item["sentences"]]).encode()).hexdigest()}}
|
| 256 |
+
|
| 257 |
+
|
| 258 |
+
def generate(n_pairs_per_family, seed, families=None):
|
| 259 |
+
"""Sibling-adjacent records plus a per-family report of rejections."""
|
| 260 |
+
records, report = [], {}
|
| 261 |
+
for family in (families or FAMILIES):
|
| 262 |
+
rng = random.Random(f"{seed}:{family}")
|
| 263 |
+
kept, rejected, reasons, attempts = 0, 0, {}, 0
|
| 264 |
+
while kept < n_pairs_per_family and attempts < 50 * n_pairs_per_family:
|
| 265 |
+
attempts += 1
|
| 266 |
+
a, b = FAMILIES[family](rng)
|
| 267 |
+
why = check_pair(a, b)
|
| 268 |
+
if why:
|
| 269 |
+
rejected += 1; reasons[why] = reasons.get(why, 0) + 1; continue
|
| 270 |
+
pair_id = f"{seed}-{family}-{kept:04d}"
|
| 271 |
+
order_seed = rng.getrandbits(64)
|
| 272 |
+
records += [to_request(a, family, pair_id, "a", random.Random(order_seed)),
|
| 273 |
+
to_request(b, family, pair_id, "b", random.Random(order_seed))]
|
| 274 |
+
kept += 1
|
| 275 |
+
if kept < n_pairs_per_family:
|
| 276 |
+
raise ValueError(f"{family}: only {kept}/{n_pairs_per_family} pairs passed checks ({reasons})")
|
| 277 |
+
report[family] = {"pairs": kept, "rejected": rejected, "reasons": reasons}
|
| 278 |
+
return records, report
|
| 279 |
+
|
| 280 |
+
|
| 281 |
+
def paired_flip(rows):
|
| 282 |
+
"""Pair-level metrics from benchmark rows carrying pair_id/sibling. A model that ignores the state cannot flip."""
|
| 283 |
+
by_pair = {}
|
| 284 |
+
for r in rows:
|
| 285 |
+
if r.get("pair_id"):
|
| 286 |
+
key = (r["pair_id"], r.get("question", "decision"))
|
| 287 |
+
pair = by_pair.setdefault(key, {})
|
| 288 |
+
if r["sibling"] in pair:
|
| 289 |
+
raise ValueError("duplicate contrastive sibling")
|
| 290 |
+
pair[r["sibling"]] = r
|
| 291 |
+
if not by_pair:
|
| 292 |
+
return None
|
| 293 |
+
if any(set(p) != {"a", "b"} for p in by_pair.values()):
|
| 294 |
+
raise ValueError("incomplete contrastive pair")
|
| 295 |
+
prediction = lambda r: r["keys"][max(range(len(r["p"])), key=r["p"].__getitem__)]
|
| 296 |
+
truth = lambda r: r["keys"][r["label"]]
|
| 297 |
+
relevant = [p for p in by_pair.values() if truth(p["a"]) != truth(p["b"])]
|
| 298 |
+
invariant = [p for p in by_pair.values() if truth(p["a"]) == truth(p["b"])]
|
| 299 |
+
both = lambda ps: sum(all(prediction(r) == truth(r) for r in p.values()) for p in ps) / len(ps) if ps else None
|
| 300 |
+
result = {"pairs": len(relevant),
|
| 301 |
+
"flip_rate": sum(prediction(p["a"]) != prediction(p["b"]) for p in relevant) / len(relevant) if relevant else None,
|
| 302 |
+
"both_correct_rate": both(relevant)}
|
| 303 |
+
if invariant:
|
| 304 |
+
result.update(invariant_pairs=len(invariant), invariance_rate=sum(prediction(p["a"]) == prediction(p["b"]) for p in invariant) / len(invariant),
|
| 305 |
+
invariant_both_correct_rate=both(invariant))
|
| 306 |
+
return result
|
runtime/kev/cuda_graphs.py
ADDED
|
@@ -0,0 +1,375 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""CUDA graphs for serving the hybrid (Qwen3.5) backbones on CUDA, batched across requests.
|
| 2 |
+
|
| 3 |
+
A served request is two kinds of forward pass: the state pass (the document, on a prefix-cache miss) and the row pass (one
|
| 4 |
+
causal row per question, continuing the cached state). At a few hundred tokens the GPU finishes a pass long before Python
|
| 5 |
+
finishes issuing it: a pass launches ~2,000 kernels, and the flash-linear-attention wrapper around every DeltaNet layer
|
| 6 |
+
costs ~1 ms of CPU. Replaying a captured graph issues the same kernels in one call. And because one request leaves most of
|
| 7 |
+
the GPU idle, kev.serve hands this module every request that is waiting when the GPU frees up: one state pass computes all
|
| 8 |
+
their new states, one row pass all their questions.
|
| 9 |
+
|
| 10 |
+
Two sets of buffers, shared by every graph:
|
| 11 |
+
- the state bank: per layer the attention keys/values [GRAPH_STATES, heads, BANK_WIDTH, dim], each state right-aligned at
|
| 12 |
+
the end, and the DeltaNet conv/recurrent states. The state pass writes new states into it; cached states are copied in.
|
| 13 |
+
One fixed layout, so a row pass reads any state the same way whatever state pass wrote it.
|
| 14 |
+
- the row buffers: [rows, heads, Sr + Lb, dim] per attention layer and the DeltaNet states per question row. The row pass
|
| 15 |
+
starts by gathering each row's state from the bank (one index_select per tensor, inside the graph), then writes the
|
| 16 |
+
row's own keys after it.
|
| 17 |
+
Both are flat buffers viewed at the graph's shape, so memory is fixed however many shapes appear. Graph keys are kept
|
| 18 |
+
coarse because a batch's shape changes with the traffic: state passes by (states, state bucket), row passes by (rows,
|
| 19 |
+
row bucket, Sr = the longest state rounded up to a power of two, which only lengthens attention).
|
| 20 |
+
|
| 21 |
+
A graph has fixed shapes, so passes are padded to buckets: token counts to bucket() (under a quarter more), row and state counts
|
| 22 |
+
to count_bucket() (empty entries are masked). The padding is masked exactly:
|
| 23 |
+
- the state pass is LEFT-padded. The DeltaNet layers see zeroed inputs at the pads (their padding mask), so keys, values
|
| 24 |
+
and queries are zero and the recurrent state is still zero when the real tokens start; the causal convolution sees
|
| 25 |
+
the same zeros its own padding supplies. The attention layers mask the pads as keys, and a pad query attends to itself
|
| 26 |
+
so no row is fully masked (a NaN there would survive the zeroing: NaN * 0 = NaN). Real tokens keep their position ids.
|
| 27 |
+
- question rows are right-padded (after every real token, as the eager path already does); a state's keys are
|
| 28 |
+
right-aligned in the bucket and the slots before them are masked.
|
| 29 |
+
So the results equal the eager passes up to floating-point reassociation (another chunking of the DeltaNet scan, other
|
| 30 |
+
GEMM shapes), not bit for bit, and a request's answers do not depend on what else shares its batch.
|
| 31 |
+
|
| 32 |
+
Capturing a graph costs ~0.4 s (an H100, Kev-4B), so the first pass of a new bucket runs the same code eagerly from the
|
| 33 |
+
same buffers and the bucket joins `pending`; capture_pending() captures them later (kev.serve does it when no request is
|
| 34 |
+
waiting). Replays run one at a time (kev.serve has one model thread), each refills what it reads, and its results are
|
| 35 |
+
copied out before the next replay; that is what makes the shared buffers and the shared memory pool safe.
|
| 36 |
+
"""
|
| 37 |
+
from collections import OrderedDict
|
| 38 |
+
from typing import NamedTuple
|
| 39 |
+
|
| 40 |
+
import torch
|
| 41 |
+
from transformers import DynamicCache
|
| 42 |
+
from transformers.cache_utils import DynamicLayer, LinearAttentionLayer
|
| 43 |
+
|
| 44 |
+
# Limits of the graphed passes (not the model's context: that is kev.model.MAX_STATE / SERVE_MAX_STATE)
|
| 45 |
+
GRAPH_TOKENS = 32768 # rows x positions one row pass may hold in the attention buffers (about 1 GB on Kev-4B and 9B)
|
| 46 |
+
GRAPH_STATE = 1024 # longest state bucket the state pass graphs. A longer state pass is compute-bound, and the padding plus
|
| 47 |
+
# the explicit mask (no causal flash attention) made a 2,200-token one slower as a graph (L40S, Kev-4B:
|
| 48 |
+
# 223 vs 208 ms per request), so it runs eagerly
|
| 49 |
+
GRAPH_ROW = 1024 # longest question-row bucket the row pass graphs
|
| 50 |
+
GRAPH_ROWS = 32 # question rows per graphed pass; more run as several replays
|
| 51 |
+
GRAPH_STATES = 16 # states per state pass: the bank's entries
|
| 52 |
+
BANK_WIDTH = 4096 # positions per bank entry (states are right-aligned in it; about 2 GB for 16 entries on Kev-4B and 9B)
|
| 53 |
+
GRAPHS_KEPT = 256 # captured graphs kept, least recently used evicted (a busy server met ~160 on mixed traffic)
|
| 54 |
+
PASS_TOKENS = 256 # tokens a pass costs however few it holds: below about this many, a pass is bound by reading the
|
| 55 |
+
# weights, not by compute (bf16 GEMMs: peak FLOP/s over memory bandwidth is ~200-400 on H100, H200, B200, L40S)
|
| 56 |
+
HOT_BUCKET = 3 # eager passes after which a busy server captures a bucket's graph anyway (capture_due)
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def bucket(n, steps=8, floor=16):
|
| 60 |
+
"""n rounded up to one of `steps` steps per power of two, in steps of at least `floor`: the padding adds less than
|
| 61 |
+
2/steps of n (just past a power of two, one step is almost that), or less than `floor` tokens for short n."""
|
| 62 |
+
step = max(floor, (1 << (n - 1).bit_length()) // steps)
|
| 63 |
+
return -(-n // step) * step
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def pow2(n):
|
| 67 |
+
return 1 << (n - 1).bit_length()
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
def count_bucket(n):
|
| 71 |
+
"""A row or state count rounded up to 1, 2, 3, 4, 6, 8, 12, 16, 24, 32: the empty entries padding a pass add less than
|
| 72 |
+
half, and a traffic mix meets few distinct counts."""
|
| 73 |
+
return bucket(n, steps=4, floor=1)
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
def length_groups(lengths, cap):
|
| 77 |
+
"""Indices grouped into padded passes of at most `cap` items, in order of length, computing the fewest tokens: a pass
|
| 78 |
+
costs its padded size, count_bucket(items) x bucket(longest), but at least PASS_TOKENS (below that, reading the
|
| 79 |
+
weights dominates). Optimal among groupings of consecutive lengths (dynamic programming; ties keep larger passes)."""
|
| 80 |
+
order = sorted(range(len(lengths)), key=lengths.__getitem__)
|
| 81 |
+
padded = [bucket(lengths[i]) for i in order]
|
| 82 |
+
counts = [0] + [count_bucket(n) for n in range(1, cap + 1)]
|
| 83 |
+
best, cut = [0] + [float("inf")] * len(order), [0] * (len(order) + 1)
|
| 84 |
+
for j in range(1, len(order) + 1):
|
| 85 |
+
for i in range(max(0, j - cap), j):
|
| 86 |
+
cost = best[i] + max(PASS_TOKENS, counts[j - i] * padded[j - 1])
|
| 87 |
+
if cost < best[j]: best[j], cut[j] = cost, i
|
| 88 |
+
groups, j = [], len(order)
|
| 89 |
+
while j: groups.append(order[cut[j]:j]); j = cut[j]
|
| 90 |
+
return groups[::-1]
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
class Request(NamedTuple):
|
| 94 |
+
"""One request for CudaGraphs.run."""
|
| 95 |
+
state: list # the state's token ids
|
| 96 |
+
positions: list # and their position ids
|
| 97 |
+
rows: list # its question rows, [(token ids, position ids)]
|
| 98 |
+
cached: object # the cached state (a DynamicCache) it continues, or None for a new state
|
| 99 |
+
keep: bool # return a new state's cache (to keep in the prefix cache)
|
| 100 |
+
picks: list # per row, the positions whose hidden states are wanted
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
class _Row(NamedTuple):
|
| 104 |
+
"""One question row of a batch: its tokens, its state's bank entry and length, and where its picks go in the output."""
|
| 105 |
+
ids: list
|
| 106 |
+
pos: list
|
| 107 |
+
entry: int
|
| 108 |
+
state_len: int
|
| 109 |
+
picks: list
|
| 110 |
+
at: int
|
| 111 |
+
|
| 112 |
+
|
| 113 |
+
class BufferKV(DynamicLayer):
|
| 114 |
+
"""An attention cache layer over preallocated [N, heads, T, dim] views: the first `filled` positions hold the cached
|
| 115 |
+
state, a pass writes its keys/values after them instead of concatenating (no allocation, no Python state change,
|
| 116 |
+
so one layer object serves the warm-up, the capture and every replay)."""
|
| 117 |
+
|
| 118 |
+
def __init__(self, keys, values, filled):
|
| 119 |
+
super().__init__()
|
| 120 |
+
self.buffers, self.filled, self.is_initialized = (keys, values), filled, True
|
| 121 |
+
self.keys, self.values = keys[..., :filled, :], values[..., :filled, :]
|
| 122 |
+
|
| 123 |
+
def update(self, key_states, value_states, *args, **kwargs):
|
| 124 |
+
end = self.filled + key_states.shape[-2]
|
| 125 |
+
for buf, new in zip(self.buffers, (key_states, value_states)): buf[..., self.filled:end, :].copy_(new)
|
| 126 |
+
return tuple(buf[..., :end, :] for buf in self.buffers)
|
| 127 |
+
|
| 128 |
+
|
| 129 |
+
def set_linear(layer, conv, recurrent, previous):
|
| 130 |
+
"""Point a DeltaNet cache layer at these conv/recurrent state tensors (transformers' LinearAttentionLayer fields)."""
|
| 131 |
+
layer.conv_states[0], layer.recurrent_states[0] = conv, recurrent
|
| 132 |
+
layer.is_conv_states_initialized[0] = layer.is_recurrent_states_initialized[0] = True
|
| 133 |
+
layer.conv_kernel_size[0], layer.has_previous_state[0] = conv.shape[-1], previous
|
| 134 |
+
layer.device, layer.dtype = conv.device, conv.dtype
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
def is_attention(layer):
|
| 138 |
+
return isinstance(layer, DynamicLayer)
|
| 139 |
+
|
| 140 |
+
|
| 141 |
+
def layer_tensors(layer):
|
| 142 |
+
"""The two tensors a prefix-cache layer holds: attention keys/values, or DeltaNet conv/recurrent states."""
|
| 143 |
+
return (layer.keys, layer.values) if is_attention(layer) else (layer.conv_states[0], layer.recurrent_states[0])
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
class Buffers:
|
| 147 |
+
"""Per layer two flat buffers and the per-entry shape they hold: attention [heads, T, dim] (T set per view, `tokens`
|
| 148 |
+
entries x positions in total), DeltaNet states (`entries` entries)."""
|
| 149 |
+
|
| 150 |
+
def __init__(self, probe, tokens, entries, device):
|
| 151 |
+
self.slots = []
|
| 152 |
+
for layer in probe.layers:
|
| 153 |
+
ts, attention = layer_tensors(layer), is_attention(layer)
|
| 154 |
+
size = [tokens * t.shape[1] * t.shape[3] if attention else entries * t[0].numel() for t in ts]
|
| 155 |
+
self.slots.append((attention, [torch.zeros(n, dtype=t.dtype, device=device) for n, t in zip(size, ts)], [t.shape[1:] for t in ts]))
|
| 156 |
+
|
| 157 |
+
def views(self, n, length):
|
| 158 |
+
"""Per layer, its two buffers viewed for n entries: attention [n, heads, length, dim], DeltaNet [n, *state]."""
|
| 159 |
+
def view(flat, shape, attention):
|
| 160 |
+
shape = (n, shape[0], length, shape[2]) if attention else (n, *shape)
|
| 161 |
+
return flat[:torch.Size(shape).numel()].view(shape)
|
| 162 |
+
return [tuple(view(f, s, attention) for f, s in zip(flats, shapes)) for attention, flats, shapes in self.slots]
|
| 163 |
+
|
| 164 |
+
|
| 165 |
+
class CudaGraphs:
|
| 166 |
+
def __init__(self, lm, pad_id):
|
| 167 |
+
self.lm, self.pad_id = lm, pad_id
|
| 168 |
+
self.device, self.dtype = next(lm.parameters()).device, next(lm.parameters()).dtype
|
| 169 |
+
self.pool, self.stream = torch.cuda.graph_pool_handle(), torch.cuda.Stream() # the graphs' shared memory pool, the capture stream
|
| 170 |
+
self.graphs = OrderedDict() # key -> (graph, input buffer)
|
| 171 |
+
self.pending = {} # key -> (pass body, input buffer): buckets that ran eagerly, not captured yet
|
| 172 |
+
self.eager_runs = {} # key -> how many passes of a pending bucket ran eagerly
|
| 173 |
+
self.failed = {} # key -> (why its capture failed, input buffer); the bucket keeps running eagerly
|
| 174 |
+
self.captures = 0
|
| 175 |
+
probe = DynamicCache(config=lm.config) # learn the cache layout (layer kinds, state shapes, dtypes) from one eager pass
|
| 176 |
+
with torch.no_grad():
|
| 177 |
+
lm(input_ids=torch.full((1, 16), pad_id, device=self.device), past_key_values=probe, use_cache=True)
|
| 178 |
+
# BufferKV and set_linear rely on this layout: plain attention layers and single-state DeltaNet layers that update
|
| 179 |
+
# their states in place. Anything else (another cache layer type, record_past, which assigns instead of copying)
|
| 180 |
+
# would leave the buffers stale and move probabilities silently, so refuse it here.
|
| 181 |
+
if not all(type(l) is DynamicLayer or (type(l) is LinearAttentionLayer and l.number_of_states == 1 and not l.record_past) for l in probe.layers):
|
| 182 |
+
raise ValueError(f"CUDA graphs support attention and single-state DeltaNet cache layers only; got {sorted({type(l).__name__ for l in probe.layers})}")
|
| 183 |
+
self.bank = Buffers(probe, GRAPH_STATES * BANK_WIDTH, GRAPH_STATES, self.device)
|
| 184 |
+
# the bank has one layout for every pass, so the state and row graphs agree on it: attention [GRAPH_STATES, heads,
|
| 185 |
+
# BANK_WIDTH, dim] with states right-aligned at the end, DeltaNet [GRAPH_STATES, *state]
|
| 186 |
+
self.bank_views = self.bank.views(GRAPH_STATES, BANK_WIDTH)
|
| 187 |
+
self.rowbuf = Buffers(probe, GRAPH_TOKENS, GRAPH_ROWS, self.device)
|
| 188 |
+
self.hidden = torch.zeros(GRAPH_ROWS * GRAPH_ROW * lm.config.hidden_size, dtype=self.dtype, device=self.device)
|
| 189 |
+
|
| 190 |
+
def _cache(self, views, filled, previous):
|
| 191 |
+
"""A DynamicCache over buffer views: attention layers hold `filled` cached positions, DeltaNet layers the states."""
|
| 192 |
+
cache = DynamicCache(config=self.lm.config)
|
| 193 |
+
for i, (layer, v) in enumerate(zip(cache.layers, views)):
|
| 194 |
+
if is_attention(layer): cache.layers[i] = BufferKV(*v, filled)
|
| 195 |
+
else: set_linear(layer, *v, previous)
|
| 196 |
+
return cache
|
| 197 |
+
|
| 198 |
+
def _mask(self, allow):
|
| 199 |
+
"""Additive attention mask [B, 1, Lq, Lk] from a boolean one [B, Lq, Lk], in the backbone's dtype."""
|
| 200 |
+
return torch.zeros(allow.shape, dtype=self.dtype, device=self.device).masked_fill(~allow, torch.finfo(self.dtype).min)[:, None]
|
| 201 |
+
|
| 202 |
+
def _forward(self, ids, pos, full_mask, linear_mask, cache):
|
| 203 |
+
return self.lm(input_ids=ids, position_ids=pos, attention_mask={"full_attention": full_mask, "linear_attention": linear_mask},
|
| 204 |
+
past_key_values=cache, use_cache=True).last_hidden_state
|
| 205 |
+
|
| 206 |
+
def _replay(self, key, body, rows):
|
| 207 |
+
"""One pass for `key` on `rows` (token ids, positions, lengths, indices; int64): the graph's replay, or body() run
|
| 208 |
+
eagerly when the bucket has no graph yet (it then waits in `pending`). body(buf) runs the pass reading the uploaded
|
| 209 |
+
rows from buf. Every body for a key reads and writes the same buffer views, so the one kept in `pending` stands for
|
| 210 |
+
all of them."""
|
| 211 |
+
if key in self.graphs:
|
| 212 |
+
self.graphs.move_to_end(key)
|
| 213 |
+
graph, buf = self.graphs[key]
|
| 214 |
+
elif key in self.failed:
|
| 215 |
+
graph, buf = None, self.failed[key][1]
|
| 216 |
+
else:
|
| 217 |
+
graph, buf = None, self.pending.setdefault(key, (body, torch.zeros((len(rows), len(rows[0])), dtype=torch.long, device=self.device)))[1]
|
| 218 |
+
self.eager_runs[key] = self.eager_runs.get(key, 0) + 1
|
| 219 |
+
buf.copy_(torch.tensor(rows, dtype=torch.long), non_blocking=True)
|
| 220 |
+
if graph is not None: graph.replay()
|
| 221 |
+
else: body(buf)
|
| 222 |
+
|
| 223 |
+
@torch.no_grad()
|
| 224 |
+
def capture_pending(self, limit=None):
|
| 225 |
+
"""Capture graphs for up to `limit` (None = all) buckets that have run eagerly, the most-run first. The caller must
|
| 226 |
+
keep other passes out (kev.serve captures on its model thread). Each capture overwrites the shared buffers, which is
|
| 227 |
+
fine between passes: a pass refills them. On a side stream, a warm-up pass first (Triton autotuning and cuBLAS
|
| 228 |
+
setup must not happen inside a capture), then the capture; not torch.cuda.graph, whose synchronize, gc.collect and
|
| 229 |
+
empty_cache on entry would stall a busy server for each capture. thread_local: CUDA calls of other threads (a
|
| 230 |
+
request tokenizing, a model card read) cannot invalidate the capture. A capture that fails leaves its bucket running
|
| 231 |
+
eagerly, for good."""
|
| 232 |
+
for _ in range(len(self.pending) if limit is None else min(limit, len(self.pending))):
|
| 233 |
+
key = max(self.pending, key=self.eager_runs.get)
|
| 234 |
+
(body, buf), _ = self.pending.pop(key), self.eager_runs.pop(key)
|
| 235 |
+
current, graph = torch.cuda.current_stream(), torch.cuda.CUDAGraph()
|
| 236 |
+
self.stream.wait_stream(current)
|
| 237 |
+
try:
|
| 238 |
+
with torch.cuda.stream(self.stream):
|
| 239 |
+
body(buf) # warm-up: autotuning and lazy setup happen outside the capture
|
| 240 |
+
graph.capture_begin(pool=self.pool, capture_error_mode="thread_local")
|
| 241 |
+
try: body(buf)
|
| 242 |
+
finally: graph.capture_end()
|
| 243 |
+
except Exception as e: # e.g. out of memory for a new shape: serve it eagerly rather than stop serving
|
| 244 |
+
# a failed capture_end leaves the allocator routing this thread's allocations into the graph pool, which
|
| 245 |
+
# would fail every later capture and let eager passes allocate graph memory: end that explicitly
|
| 246 |
+
try: getattr(torch._C, "_cuda_endAllocateToPool", lambda *a: None)(self.device.index, self.pool) # private API: best effort
|
| 247 |
+
except Exception: pass
|
| 248 |
+
self.failed[key] = (f"{type(e).__name__}: {e}", buf)
|
| 249 |
+
print(f"kev.cuda_graphs: capturing {key} failed, it runs eagerly: {self.failed[key][0]}", flush=True)
|
| 250 |
+
continue
|
| 251 |
+
finally:
|
| 252 |
+
current.wait_stream(self.stream)
|
| 253 |
+
self.graphs[key] = (graph, buf); self.captures += 1
|
| 254 |
+
while len(self.graphs) > GRAPHS_KEPT: self.graphs.popitem(last=False)
|
| 255 |
+
|
| 256 |
+
def capture_due(self, idle):
|
| 257 |
+
"""Capture one pending graph now? Always when idle; under load once a bucket keeps running eagerly (HOT_BUCKET
|
| 258 |
+
eager passes), because a request arriving meanwhile waits for the capture (~0.4 s)."""
|
| 259 |
+
runs = list(self.eager_runs.values()) # a snapshot: wait_idle and /v1/models call this off the model thread
|
| 260 |
+
return bool(runs) and (idle or max(runs) >= HOT_BUCKET)
|
| 261 |
+
|
| 262 |
+
def stats(self):
|
| 263 |
+
return {"captured": self.captures, "kept": len(self.graphs), "pending": len(self.pending), "failed": len(self.failed)}
|
| 264 |
+
|
| 265 |
+
# Serving: which requests the graphed passes take, and one batch of them
|
| 266 |
+
|
| 267 |
+
@staticmethod
|
| 268 |
+
def admits(state_len, row_lens, cached):
|
| 269 |
+
"""Whether a request's passes fit the graphed passes: its rows, and its state (a new one runs the graphed state
|
| 270 |
+
pass; a cached one only has to fit the bank)."""
|
| 271 |
+
return bucket(max(row_lens)) <= GRAPH_ROW and (state_len <= BANK_WIDTH if cached else bucket(state_len) <= GRAPH_STATE)
|
| 272 |
+
|
| 273 |
+
@torch.no_grad()
|
| 274 |
+
def run(self, requests):
|
| 275 |
+
"""Admitted requests (Request) -> (their picked hidden states, float32 [sum of picks, d] in request, row, pick
|
| 276 |
+
order; per request its state's cache: the cached one, a new one if `keep`, else None). States grouped by length
|
| 277 |
+
(length_groups) share a state pass (identical states once); each request takes one bank entry, so at most GRAPH_STATES per group.
|
| 278 |
+
Every row pass writes its picks straight to their place in the output."""
|
| 279 |
+
at = [0]
|
| 280 |
+
for r in requests: at.append(at[-1] + sum(map(len, r.picks)))
|
| 281 |
+
out = torch.empty((at[-1], self.lm.config.hidden_size), dtype=torch.float32, device=self.device)
|
| 282 |
+
caches = [None] * len(requests)
|
| 283 |
+
for group in length_groups([len(r.state) for r in requests], GRAPH_STATES):
|
| 284 |
+
new = {} # distinct new states of the group -> (ids, positions, keep)
|
| 285 |
+
for i in group:
|
| 286 |
+
r = requests[i]
|
| 287 |
+
if r.cached is None: new[tuple(r.state)] = (r.state, r.positions, r.keep or new.get(tuple(r.state), (0, 0, False))[2])
|
| 288 |
+
made = dict(zip(new, self.states(list(new.values()), bucket(max(map(len, new)))))) if new else {}
|
| 289 |
+
entry = {S: j for j, S in enumerate(new)} # cached states go into the entries after the new ones
|
| 290 |
+
rows, free = [], len(new)
|
| 291 |
+
for i in group: # after the state pass: its padded entries are written
|
| 292 |
+
r = requests[i]
|
| 293 |
+
if r.cached is None: e, caches[i] = entry[tuple(r.state)], made[tuple(r.state)]
|
| 294 |
+
else: e, free, caches[i] = free, free + 1, r.cached; self.load_state(e, r.cached, len(r.state))
|
| 295 |
+
offset = at[i]
|
| 296 |
+
for (ids, pos), picks in zip(r.rows, r.picks):
|
| 297 |
+
rows.append(_Row(ids, pos, e, len(r.state), picks, offset)); offset += len(picks)
|
| 298 |
+
self.rows(rows, out)
|
| 299 |
+
return out, caches
|
| 300 |
+
|
| 301 |
+
@torch.no_grad()
|
| 302 |
+
def states(self, items, Sb):
|
| 303 |
+
"""State pass for up to GRAPH_STATES states, items = [(ids, pos, keep)] each at most Sb <= GRAPH_STATE tokens.
|
| 304 |
+
Leaves state i in bank entry i and returns, for each kept state, a DynamicCache equal to the eager prefix pass's
|
| 305 |
+
(None for the others: copying every state of a batch cost more than the passes)."""
|
| 306 |
+
Nb = count_bucket(len(items))
|
| 307 |
+
|
| 308 |
+
def body(buf): # buf rows = [ids | positions | real length], left-padded
|
| 309 |
+
pad = Sb - buf[:, 2 * Sb:]
|
| 310 |
+
i = torch.arange(Sb, device=self.device)
|
| 311 |
+
q, k = i[None, :, None], i[None, None, :]
|
| 312 |
+
allow = ((k <= q) & (k >= pad[:, :, None])) | (k == q)
|
| 313 |
+
views = [(a[:Nb, :, BANK_WIDTH - Sb:], b[:Nb, :, BANK_WIDTH - Sb:]) if attention else (a[:Nb], b[:Nb])
|
| 314 |
+
for (attention, *_), (a, b) in zip(self.bank.slots, self.bank_views)]
|
| 315 |
+
self._forward(buf[:, :Sb], buf[:, Sb:2 * Sb], self._mask(allow), (i[None] >= pad).long(), self._cache(views, 0, False))
|
| 316 |
+
|
| 317 |
+
rows = [[self.pad_id] * (Sb - len(ids)) + list(ids) + [0] * (Sb - len(ids)) + list(pos) + [len(ids)] for ids, pos, _ in items]
|
| 318 |
+
self._replay(("states", Nb, Sb), body, rows + [[self.pad_id] * Sb + [0] * Sb + [0]] * (Nb - len(items)))
|
| 319 |
+
out = [] # a cache for each kept state: its own copy of the bank entry (the bank belongs to the next pass)
|
| 320 |
+
for j, (ids, _, keep) in enumerate(items):
|
| 321 |
+
if not keep: out.append(None); continue
|
| 322 |
+
cache, S = DynamicCache(config=self.lm.config), len(ids)
|
| 323 |
+
for layer, (a, b) in zip(cache.layers, self.bank_views):
|
| 324 |
+
if is_attention(layer):
|
| 325 |
+
layer.keys, layer.values = a[j:j + 1, :, BANK_WIDTH - S:].clone(), b[j:j + 1, :, BANK_WIDTH - S:].clone()
|
| 326 |
+
layer.dtype, layer.device, layer.is_initialized = self.dtype, self.device, True
|
| 327 |
+
else:
|
| 328 |
+
set_linear(layer, a[j:j + 1].clone(), b[j:j + 1].clone(), True)
|
| 329 |
+
out.append(cache)
|
| 330 |
+
return out
|
| 331 |
+
|
| 332 |
+
@torch.no_grad()
|
| 333 |
+
def load_state(self, entry, cache, S):
|
| 334 |
+
"""Copy a cached state (a DynamicCache of S <= BANK_WIDTH tokens, from either path) into bank entry `entry`."""
|
| 335 |
+
for layer, v in zip(cache.layers, self.bank_views):
|
| 336 |
+
for buf, t in zip(v, layer_tensors(layer)):
|
| 337 |
+
if is_attention(layer): buf[entry, :, BANK_WIDTH - S:].copy_(t[0, :, t.shape[-2] - S:])
|
| 338 |
+
else: buf[entry].copy_(t[0])
|
| 339 |
+
|
| 340 |
+
def rows(self, rows, out):
|
| 341 |
+
"""Row passes for question rows (_Row), each continuing its state in the bank; each row's picked hidden states land
|
| 342 |
+
in out[row.at:]. The rows see the last Sr positions of the bank, Sr the longest state rounded up to a power of two
|
| 343 |
+
(it only lengthens attention). Rows are grouped by length (length_groups), then split to fit the buffers."""
|
| 344 |
+
for idx in length_groups([len(r.ids) for r in rows], GRAPH_ROWS):
|
| 345 |
+
Lb, Sr = bucket(max(len(rows[i].ids) for i in idx)), pow2(max(16, max(rows[i].state_len for i in idx)))
|
| 346 |
+
group = min(GRAPH_ROWS, GRAPH_TOKENS // (Sr + Lb))
|
| 347 |
+
group = max(n for n in range(1, group + 1) if count_bucket(n) <= group) # the most rows whose padded count fits
|
| 348 |
+
for start in range(0, len(idx), group):
|
| 349 |
+
part = [rows[i] for i in idx[start:start + group]]
|
| 350 |
+
h = self._row_pass(part, Sr, Lb)
|
| 351 |
+
src, dst = zip(*[(n * Lb + p, r.at + j) for n, r in enumerate(part) for j, p in enumerate(r.picks)])
|
| 352 |
+
src, dst = torch.tensor([src, dst]).to(self.device, non_blocking=True)
|
| 353 |
+
out.index_copy_(0, dst, h.flatten(0, 1).index_select(0, src).float()) # only the wanted positions leave the pass buffer
|
| 354 |
+
|
| 355 |
+
def _row_pass(self, rows, Sr, Lb):
|
| 356 |
+
"""One row pass; -> the pass buffer [rows, Lb, d] (valid until the next pass)."""
|
| 357 |
+
Nb, T = count_bucket(len(rows)), Sr + Lb
|
| 358 |
+
hidden = self.hidden[:Nb * Lb * self.lm.config.hidden_size].view(Nb, Lb, -1)
|
| 359 |
+
|
| 360 |
+
def body(buf): # buf rows = [ids | positions | row length | state length | bank entry]
|
| 361 |
+
rowlen, plen, entry = buf[:, 2 * Lb:2 * Lb + 1], buf[:, 2 * Lb + 1:2 * Lb + 2], buf[:, 2 * Lb + 2]
|
| 362 |
+
views = self.rowbuf.views(Nb, T)
|
| 363 |
+
for (attention, *_), dst, src in zip(self.bank.slots, views, self.bank_views): # each row's state, from the bank
|
| 364 |
+
for d, s in zip(dst, src):
|
| 365 |
+
if attention: d[:, :, :Sr].copy_(s[:, :, BANK_WIDTH - Sr:].index_select(0, entry))
|
| 366 |
+
else: d.copy_(s.index_select(0, entry))
|
| 367 |
+
q = torch.arange(Lb, device=self.device)[None, :, None]
|
| 368 |
+
k = torch.arange(T, device=self.device)[None, None, :]
|
| 369 |
+
allow = ((k < Sr) & (k >= Sr - plen[:, :, None])) | ((k >= Sr) & (k - Sr <= q) & (k - Sr < rowlen[:, :, None])) | (k - Sr == q)
|
| 370 |
+
linear = (torch.arange(Lb, device=self.device)[None] < rowlen).long()
|
| 371 |
+
hidden.copy_(self._forward(buf[:, :Lb], buf[:, Lb:2 * Lb], self._mask(allow), linear, self._cache(views, Sr, True)))
|
| 372 |
+
|
| 373 |
+
uploads = [list(r.ids) + [self.pad_id] * (Lb - len(r.ids)) + list(r.pos) + [0] * (Lb - len(r.pos)) + [len(r.ids), r.state_len, r.entry] for r in rows]
|
| 374 |
+
self._replay(("rows", Nb, Lb, Sr), body, uploads + [[self.pad_id] * Lb + [0] * Lb + [0, 0, 0]] * (Nb - len(rows)))
|
| 375 |
+
return hidden[:len(rows)]
|
runtime/kev/data.py
ADDED
|
@@ -0,0 +1,410 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Convert public labelled datasets into TypeSafe-shaped requests, then through api.to_record() so the
|
| 2 |
+
training format is byte-identical to what /v1/systemone feeds the model.
|
| 3 |
+
|
| 4 |
+
A "labelled request" is {"state": JSONContent, "questions": {id: {type, instructions, criteria, "label": ..., "src": str}}}
|
| 5 |
+
label: choice -> option key, noul -> bool, score -> level index.
|
| 6 |
+
materialize() -> internal record {"state": str, "questions": [{"instr", "options", "label": int, "src"}]}
|
| 7 |
+
"""
|
| 8 |
+
import hashlib
|
| 9 |
+
import json
|
| 10 |
+
from pathlib import Path
|
| 11 |
+
import random
|
| 12 |
+
from datasets import load_dataset
|
| 13 |
+
from .api import SystemOneRequest, to_record
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
REPOS = {"banking77": "legacy-datasets/banking77", "boolq": "google/boolq", "agnews": "fancyzhx/ag_news",
|
| 17 |
+
"mnli": "nyu-mll/multi_nli", "sst5": "SetFit/sst5", "yelp": "Yelp/yelp_review_full"}
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
def source_seed(seed, source):
|
| 21 |
+
return int.from_bytes(hashlib.sha256(f"{seed}:{source}".encode()).digest()[:8], "big")
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
# Repos whose main branch is a (no longer supported) loading script; read the Hub's auto-converted parquet branch instead.
|
| 25 |
+
PARQUET_BRANCH = {"CogComp/trec": "refs/convert/parquet"}
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
def dataset_ref(repo):
|
| 29 |
+
"""(repo, revision-to-pin) for a repo id; parquet-branch repos pin that branch's sha, not main."""
|
| 30 |
+
return repo.partition(":")[0], PARQUET_BRANCH.get(repo.partition(":")[0])
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
class Source(random.Random):
|
| 34 |
+
"""One source's sampling context: a seeded RNG (the converters draw everything from it), the dataset revision to pin,
|
| 35 |
+
and the provenance of every row drawn (`origins`, filled by _sample so build() can attach it to each record)."""
|
| 36 |
+
|
| 37 |
+
def __init__(self, seed, revision=None):
|
| 38 |
+
super().__init__(seed)
|
| 39 |
+
self.revision = revision
|
| 40 |
+
self.origins = []
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def _dataset(repo, split, src):
|
| 44 |
+
"""repo may be 'owner/name' or 'owner/name:config'; the Source carries the revision to pin."""
|
| 45 |
+
name, _, config = repo.partition(":")
|
| 46 |
+
return load_dataset(name, config or None, split=split, revision=src.revision or PARQUET_BRANCH.get(name))
|
| 47 |
+
|
| 48 |
+
NONE = "None of the above"
|
| 49 |
+
# "None of the above" options must appear both as the correct answer and as a wrong alternative, with varied
|
| 50 |
+
# wording, or the model learns "this wording => pick it" (it did, in the first training run).
|
| 51 |
+
NONE_OPTIONS = [("other", "None of the above"), ("other", "A reason that fits none of the above"), ("none", "None of these"),
|
| 52 |
+
("other", "Something else"), ("not_listed", "Not listed here"), ("none_of_the_above", None),
|
| 53 |
+
("other", "A category that fits none of the above"), ("other", "None of the listed options apply"),
|
| 54 |
+
("unknown", "Cannot be determined from the options given"), ("other", "Other"), ("none", None),
|
| 55 |
+
("other", "An answer not covered by the other options"), ("no_match", "No option matches")]
|
| 56 |
+
DISTRACTORS = {"weather": "Bad weather caused it", "purple": "The colour purple", "pancakes": "A recipe for pancakes", "taxes": "Unrelated: quarterly tax filing"}
|
| 57 |
+
|
| 58 |
+
AG = {"world": "World news: politics, international affairs, conflicts", "sports": "Sports: games, athletes, teams, results",
|
| 59 |
+
"business": "Business: companies, markets, economy, finance", "scitech": "Science and technology: research, gadgets, software, space"}
|
| 60 |
+
MNLI = {"entailment": "The hypothesis follows from the premise", "neutral": "The hypothesis may or may not be true given the premise", "contradiction": "The hypothesis contradicts the premise"}
|
| 61 |
+
SST5 = ["very negative", "negative", "neutral", "positive", "very positive"]
|
| 62 |
+
YELP = ["1 star: terrible experience", "2 stars: poor", "3 stars: average", "4 stars: good", "5 stars: excellent"]
|
| 63 |
+
BANK_TEMPLATES = ["Customer asks about {}", "Issue concerning {}", "Request related to {}", "{}"]
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def _wrap_state(text, rng):
|
| 67 |
+
r = rng.random()
|
| 68 |
+
if r < 0.15: return {"document": text}
|
| 69 |
+
if r < 0.25: return {"ticket": {"channel": rng.choice(["email", "chat", "web form"]), "body": text}}
|
| 70 |
+
if r < 0.32: return [{"role": "customer", "content": text}]
|
| 71 |
+
return text
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def _instr(text, rng):
|
| 75 |
+
return {"question": text, "focus": rng.choice(["Use only the information given.", "Pick the single best fit.", "Consider the whole message."])} if rng.random() < 0.15 else text
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
def _desc(desc, rng, p_null=0.3, p_struct=0.1):
|
| 79 |
+
r = rng.random()
|
| 80 |
+
if r < p_null: return None
|
| 81 |
+
if r < p_null + p_struct: return {"what": desc}
|
| 82 |
+
return desc
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def _sample(ds, n, src):
|
| 86 |
+
rows = []
|
| 87 |
+
src.origins = []
|
| 88 |
+
for i in src.sample(range(len(ds)), min(n, len(ds))):
|
| 89 |
+
row = ds[i]
|
| 90 |
+
if row.get("label", 0) == -1:
|
| 91 |
+
continue
|
| 92 |
+
text = next((row[k] for k in ("text", "premise", "passage", "content", "question", "sentence") if isinstance(row.get(k), str)), json.dumps(row, sort_keys=True))
|
| 93 |
+
normalized = " ".join(text.casefold().split())
|
| 94 |
+
src.origins.append({"row": i, "text_sha256": hashlib.sha256(normalized.encode()).hexdigest(),
|
| 95 |
+
"row_sha256": hashlib.sha256(json.dumps(row, sort_keys=True).encode()).hexdigest()})
|
| 96 |
+
rows.append(row)
|
| 97 |
+
return rows
|
| 98 |
+
|
| 99 |
+
|
| 100 |
+
def _banking(split, n, rng):
|
| 101 |
+
ds = _dataset("legacy-datasets/banking77", split=split, src=rng)
|
| 102 |
+
names = ds.features["label"].names
|
| 103 |
+
out = []
|
| 104 |
+
for ex in _sample(ds, n, rng):
|
| 105 |
+
t = rng.choice(BANK_TEMPLATES)
|
| 106 |
+
crit = {k: _desc(t.format(k.replace("_", " ")), rng, p_null=0.5, p_struct=0.0) for k in names}
|
| 107 |
+
out.append({"state": _wrap_state(ex["text"], rng), "questions": {"intent": {"type": "choice", "instructions": _instr("Which banking intent best describes this customer message?", rng), "criteria": crit, "label": names[ex["label"]], "src": "banking77"}}})
|
| 108 |
+
return out
|
| 109 |
+
|
| 110 |
+
|
| 111 |
+
def _boolq(split, n, rng):
|
| 112 |
+
ds = _dataset("google/boolq", split=split, src=rng)
|
| 113 |
+
out = []
|
| 114 |
+
for ex in _sample(ds, n, rng):
|
| 115 |
+
q = {"type": "noul", "instructions": _instr(ex["question"].strip().rstrip("?") + "?", rng), "label": bool(ex["answer"]), "src": "boolq"}
|
| 116 |
+
if rng.random() < 0.4: q["criteria"] = {"true": "The passage supports a yes answer", "false": "The passage supports a no answer or does not say"}
|
| 117 |
+
out.append({"state": _wrap_state(ex["passage"], rng), "questions": {"answer": q}})
|
| 118 |
+
return out
|
| 119 |
+
|
| 120 |
+
|
| 121 |
+
def _agnews(split, n, rng):
|
| 122 |
+
ds = _dataset("fancyzhx/ag_news", split=split, src=rng)
|
| 123 |
+
keys = list(AG)
|
| 124 |
+
out = []
|
| 125 |
+
for ex in _sample(ds, n, rng):
|
| 126 |
+
y = keys[ex["label"]]
|
| 127 |
+
qs = {"topic": {"type": "choice", "instructions": _instr("What is the topic of this article?", rng), "criteria": {k: _desc(v, rng) for k, v in AG.items()}, "label": y, "src": "agnews"}}
|
| 128 |
+
for k in rng.sample(keys, 2):
|
| 129 |
+
qs[f"is_{k}"] = {"type": "noul", "instructions": f"Is this article about {AG[k].split(':')[0].lower()}?", "label": k == y, "src": "agnews_yn"}
|
| 130 |
+
out.append({"state": _wrap_state(ex["text"], rng), "questions": qs})
|
| 131 |
+
return out
|
| 132 |
+
|
| 133 |
+
|
| 134 |
+
def _mnli(split, n, rng):
|
| 135 |
+
ds = _dataset("nyu-mll/multi_nli", split=split, src=rng)
|
| 136 |
+
keys = list(MNLI)
|
| 137 |
+
return [{"state": _wrap_state(ex["premise"], rng), "questions": {"relation": {"type": "choice", "instructions": _instr(f'Hypothesis: "{ex["hypothesis"]}" How does it relate to the premise?', rng), "criteria": {k: _desc(v, rng) for k, v in MNLI.items()}, "label": keys[ex["label"]], "src": "mnli"}}}
|
| 138 |
+
for ex in _sample(ds, n, rng) if ex["label"] >= 0]
|
| 139 |
+
|
| 140 |
+
|
| 141 |
+
def _sst5(split, n, rng):
|
| 142 |
+
ds = _dataset("SetFit/sst5", split=split, src=rng)
|
| 143 |
+
return [{"state": _wrap_state(ex["text"], rng), "questions": {"sentiment": {"type": "score", "instructions": _instr("What is the sentiment of this review sentence?", rng), "criteria": list(SST5), "label": ex["label"], "src": "sst5"}}} for ex in _sample(ds, n, rng)]
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
def _yelp(split, n, rng):
|
| 147 |
+
ds = _dataset("Yelp/yelp_review_full", split=split, src=rng)
|
| 148 |
+
out = []
|
| 149 |
+
for ex in _sample(ds, n, rng):
|
| 150 |
+
text = " ".join(ex["text"].split()[:220])
|
| 151 |
+
qs = {"rating": {"type": "score", "instructions": _instr("How many stars did this reviewer give?", rng), "criteria": list(YELP), "label": ex["label"], "src": "yelp"},
|
| 152 |
+
"recommend": {"type": "noul", "instructions": "Would this reviewer recommend the business?", "criteria": {"true": "Clearly positive overall", "false": "Negative or mixed"}, "label": ex["label"] >= 3, "src": "yelp_yn"}}
|
| 153 |
+
out.append({"state": _wrap_state(text, rng), "questions": qs})
|
| 154 |
+
return out
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
SOURCES = {"banking77": (_banking, "train", "test"), "boolq": (_boolq, "train", "validation"), "agnews": (_agnews, "train", "test"),
|
| 158 |
+
"mnli": (_mnli, "train", "validation_matched"), "sst5": (_sst5, "train", "test"), "yelp": (_yelp, "train", "test")}
|
| 159 |
+
|
| 160 |
+
|
| 161 |
+
# Eval-only sources for transfer measurement. Never passed to training; kept out of SOURCES so build() defaults cannot
|
| 162 |
+
# pick them up. Rotten Tomatoes (SST parent) and SNLI (MNLI sibling) are deliberately excluded.
|
| 163 |
+
TREC = {"abbreviation": "Asks what an abbreviation stands for", "entity": "Asks about a thing, object, animal, product, or creative work",
|
| 164 |
+
"description": "Asks for a definition, description, reason, or manner", "human": "Asks about a person, group, or organisation",
|
| 165 |
+
"location": "Asks about a place", "number": "Asks for a number, date, count, or other numeric value"}
|
| 166 |
+
EMOTION = {"sadness": None, "joy": None, "love": None, "anger": None, "fear": None, "surprise": None}
|
| 167 |
+
AMAZON = ["1 star: very negative", "2 stars: negative", "3 stars: mixed", "4 stars: positive", "5 stars: very positive"]
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
def _trec(split, n, rng):
|
| 171 |
+
ds = _dataset("CogComp/trec", split=split, src=rng)
|
| 172 |
+
keys = list(TREC)
|
| 173 |
+
return [{"state": _wrap_state(ex["text"], rng), "questions": {"answer_type": {"type": "choice", "instructions": "What kind of answer does this question ask for?",
|
| 174 |
+
"criteria": dict(TREC), "label": keys[ex["coarse_label"]], "src": "trec"}}} for ex in _sample(ds, n, rng)]
|
| 175 |
+
|
| 176 |
+
|
| 177 |
+
def _dbpedia(split, n, rng):
|
| 178 |
+
ds = _dataset("fancyzhx/dbpedia_14", split=split, src=rng)
|
| 179 |
+
names = [x.lower().replace(" ", "_") for x in ds.features["label"].names]
|
| 180 |
+
return [{"state": _wrap_state(" ".join(ex["content"].split()[:200]), rng), "questions": {"category": {"type": "choice", "instructions": "Which category does the subject of this encyclopedia text belong to?",
|
| 181 |
+
"criteria": {k: None for k in names}, "label": names[ex["label"]], "src": "dbpedia14"}}} for ex in _sample(ds, n, rng)]
|
| 182 |
+
|
| 183 |
+
|
| 184 |
+
def _emotion(split, n, rng):
|
| 185 |
+
ds = _dataset("dair-ai/emotion:split", split=split, src=rng)
|
| 186 |
+
keys = list(EMOTION)
|
| 187 |
+
return [{"state": ex["text"], "questions": {"emotion": {"type": "choice", "instructions": "Which emotion does the writer express?",
|
| 188 |
+
"criteria": dict(EMOTION), "label": keys[ex["label"]], "src": "emotion"}}} for ex in _sample(ds, n, rng)]
|
| 189 |
+
|
| 190 |
+
|
| 191 |
+
def _imdb(split, n, rng):
|
| 192 |
+
ds = _dataset("stanfordnlp/imdb", split=split, src=rng)
|
| 193 |
+
return [{"state": _wrap_state(" ".join(ex["text"].replace("<br />", " ").split()[:220]), rng), "questions": {"positive": {"type": "noul", "instructions": "Is this movie review positive?",
|
| 194 |
+
"criteria": {"true": "The reviewer liked the film overall", "false": "The reviewer disliked the film overall"}, "label": ex["label"] == 1, "src": "imdb"}}} for ex in _sample(ds, n, rng)]
|
| 195 |
+
|
| 196 |
+
|
| 197 |
+
def _amazon(split, n, rng):
|
| 198 |
+
ds = _dataset("SetFit/amazon_reviews_multi_en", split=split, src=rng)
|
| 199 |
+
return [{"state": _wrap_state(" ".join(ex["text"].split()[:220]), rng), "questions": {"stars": {"type": "score", "instructions": "How many stars did this product reviewer give?",
|
| 200 |
+
"criteria": list(AMAZON), "label": ex["label"], "src": "amazon"}}} for ex in _sample(ds, n, rng)]
|
| 201 |
+
|
| 202 |
+
|
| 203 |
+
def _qnli(split, n, rng):
|
| 204 |
+
ds = _dataset("nyu-mll/glue:qnli", split=split, src=rng)
|
| 205 |
+
return [{"state": _wrap_state(ex["sentence"], rng), "questions": {"answers": {"type": "noul", "instructions": f'Does the sentence contain the answer to this question: "{ex["question"]}"',
|
| 206 |
+
"label": ex["label"] == 0, "src": "qnli"}}} for ex in _sample(ds, n, rng)]
|
| 207 |
+
|
| 208 |
+
|
| 209 |
+
def _offensive(split, n, rng):
|
| 210 |
+
ds = _dataset("cardiffnlp/tweet_eval:offensive", split=split, src=rng)
|
| 211 |
+
return [{"state": ex["text"], "questions": {"offensive": {"type": "noul", "instructions": "Is this post offensive?",
|
| 212 |
+
"criteria": {"true": "Contains insults, threats, profanity directed at someone, or hateful content", "false": "Not offensive"},
|
| 213 |
+
"label": ex["label"] == 1, "src": "tweet_offensive"}}} for ex in _sample(ds, n, rng)]
|
| 214 |
+
|
| 215 |
+
|
| 216 |
+
def _mmlu(split, n, rng):
|
| 217 |
+
ds = _dataset("cais/mmlu:all", split=split, src=rng)
|
| 218 |
+
out = []
|
| 219 |
+
for ex in _sample(ds, n, rng):
|
| 220 |
+
keys = ["a", "b", "c", "d"]
|
| 221 |
+
out.append({"state": {"subject": ex["subject"].replace("_", " "), "question": ex["question"]},
|
| 222 |
+
"questions": {"answer": {"type": "choice", "instructions": "Which option correctly answers the question?",
|
| 223 |
+
"criteria": dict(zip(keys, ex["choices"])), "label": keys[ex["answer"]], "src": "mmlu"}}})
|
| 224 |
+
return out
|
| 225 |
+
|
| 226 |
+
|
| 227 |
+
def _paws(split, n, rng):
|
| 228 |
+
ds = _dataset("google-research-datasets/paws:labeled_final", split=split, src=rng)
|
| 229 |
+
return [{"state": _wrap_state(ex["sentence1"], rng), "questions": {"paraphrase": {"type": "noul", "instructions": f'Does this sentence mean the same thing: "{ex["sentence2"]}"',
|
| 230 |
+
"criteria": {"true": "Same meaning, possibly reworded", "false": "Different meaning, even if most words match"}, "label": ex["label"] == 1, "src": "paws"}}}
|
| 231 |
+
for ex in _sample(ds, n, rng)]
|
| 232 |
+
|
| 233 |
+
|
| 234 |
+
def _sciq(split, n, rng):
|
| 235 |
+
ds = _dataset("allenai/sciq", split=split, src=rng)
|
| 236 |
+
out = []
|
| 237 |
+
for ex in _sample(ds, n, rng):
|
| 238 |
+
options = [ex["correct_answer"], ex["distractor1"], ex["distractor2"], ex["distractor3"]]
|
| 239 |
+
keys = ["a", "b", "c", "d"]; order = list(range(4)); rng.shuffle(order)
|
| 240 |
+
crit = {keys[i]: options[j] for i, j in enumerate(order)}
|
| 241 |
+
out.append({"state": {"passage": ex["support"], "question": ex["question"]} if ex["support"] else {"question": ex["question"]},
|
| 242 |
+
"questions": {"answer": {"type": "choice", "instructions": "Which option answers the science question?", "criteria": crit,
|
| 243 |
+
"label": keys[order.index(0)], "src": "sciq"}}})
|
| 244 |
+
return out
|
| 245 |
+
|
| 246 |
+
|
| 247 |
+
def _mcq(ex_q, labels, texts, answer_label, src, rng, state_extra=None):
|
| 248 |
+
"""Knowledge MCQ -> Choice with neutral keys; option order shuffled per record so keys carry no information."""
|
| 249 |
+
order = list(range(len(texts))); rng.shuffle(order)
|
| 250 |
+
keys = [f"opt_{i + 1}" for i in range(len(texts))]
|
| 251 |
+
crit = {keys[i]: texts[j] for i, j in enumerate(order)}
|
| 252 |
+
label = keys[order.index(labels.index(answer_label))]
|
| 253 |
+
state = {"question": ex_q}
|
| 254 |
+
if state_extra: state.update(state_extra)
|
| 255 |
+
return {"state": state, "questions": {"answer": {"type": "choice", "instructions": "Which option correctly answers the question?", "criteria": crit, "label": label, "src": src}}}
|
| 256 |
+
|
| 257 |
+
|
| 258 |
+
def _arc(split, n, rng):
|
| 259 |
+
ds = _dataset("allenai/ai2_arc:ARC-Challenge", split=split, src=rng)
|
| 260 |
+
return [_mcq(ex["question"], list(ex["choices"]["label"]), list(ex["choices"]["text"]), ex["answerKey"], "arc", rng) for ex in _sample(ds, n, rng)]
|
| 261 |
+
|
| 262 |
+
|
| 263 |
+
def _openbookqa(split, n, rng):
|
| 264 |
+
ds = _dataset("allenai/openbookqa:main", split=split, src=rng)
|
| 265 |
+
return [_mcq(ex["question_stem"], list(ex["choices"]["label"]), list(ex["choices"]["text"]), ex["answerKey"], "openbookqa", rng) for ex in _sample(ds, n, rng)]
|
| 266 |
+
|
| 267 |
+
|
| 268 |
+
def _csqa(split, n, rng):
|
| 269 |
+
ds = _dataset("tau/commonsense_qa", split=split, src=rng)
|
| 270 |
+
return [_mcq(ex["question"], list(ex["choices"]["label"]), list(ex["choices"]["text"]), ex["answerKey"], "csqa", rng) for ex in _sample(ds, n, rng) if ex["answerKey"]]
|
| 271 |
+
|
| 272 |
+
|
| 273 |
+
ALL_SOURCES = {**SOURCES, "trec": (_trec, "train", "test"), "dbpedia14": (_dbpedia, "train", "test"), "emotion": (_emotion, "train", "test"),
|
| 274 |
+
"imdb": (_imdb, "train", "test"), "amazon": (_amazon, "train", "test"), "qnli": (_qnli, "train", "validation"),
|
| 275 |
+
"tweet_offensive": (_offensive, "train", "test"), "mmlu": (_mmlu, "test", "test"),
|
| 276 |
+
"paws": (_paws, "train", "test"), "sciq": (_sciq, "train", "test"),
|
| 277 |
+
"arc": (_arc, "train", "test"), "openbookqa": (_openbookqa, "train", "test"), "csqa": (_csqa, "train", "validation")}
|
| 278 |
+
ALL_REPOS = {**REPOS, "trec": "CogComp/trec", "dbpedia14": "fancyzhx/dbpedia_14", "emotion": "dair-ai/emotion", "imdb": "stanfordnlp/imdb",
|
| 279 |
+
"amazon": "SetFit/amazon_reviews_multi_en", "qnli": "nyu-mll/glue", "tweet_offensive": "cardiffnlp/tweet_eval", "mmlu": "cais/mmlu",
|
| 280 |
+
"paws": "google-research-datasets/paws", "sciq": "allenai/sciq",
|
| 281 |
+
"arc": "allenai/ai2_arc", "openbookqa": "allenai/openbookqa", "csqa": "tau/commonsense_qa"}
|
| 282 |
+
|
| 283 |
+
# Policy (PLAN.md step 1). A source is trainable or eval-only; suites record both lists and training refuses eval-only
|
| 284 |
+
# sources. MMLU is a knowledge probe and stays eval-only permanently; Emotion/TweetEval are noisy-label honesty checks;
|
| 285 |
+
# QNLI/PAWS/SciQ measure reading transfer. Rotten Tomatoes (SST parent) and SNLI (MNLI sibling) are excluded entirely.
|
| 286 |
+
# Knowledge MCQ (ARC-Challenge, OpenBookQA, CommonsenseQA) is trainable: the hypothesis (PLAN.md, overnight) is that the
|
| 287 |
+
# pointer readout under-uses the base model's knowledge (8B scores 0.65 on 4-way MMLU, below its base-model level) and
|
| 288 |
+
# that a small MCQ mix teaches the head to tap it. MMLU and SciQ stay eval-only; ARC/SciQ are distinct datasets.
|
| 289 |
+
TRAINABLE = ("banking77", "boolq", "agnews", "mnli", "sst5", "yelp", "trec", "dbpedia14", "amazon", "imdb", "arc", "openbookqa", "csqa")
|
| 290 |
+
EVAL_ONLY = ("mmlu", "emotion", "tweet_offensive", "qnli", "paws", "sciq")
|
| 291 |
+
assert set(TRAINABLE) | set(EVAL_ONLY) == set(ALL_SOURCES) and not set(TRAINABLE) & set(EVAL_ONLY)
|
| 292 |
+
|
| 293 |
+
# transfer-v1 (frozen before the policy existed) used these eight; kept so its manifest can be re-derived.
|
| 294 |
+
TRANSFER_SOURCES = {k: ALL_SOURCES[k] for k in ("trec", "dbpedia14", "emotion", "imdb", "amazon", "qnli", "tweet_offensive", "mmlu")}
|
| 295 |
+
TRANSFER_REPOS = {k: ALL_REPOS[k] for k in TRANSFER_SOURCES}
|
| 296 |
+
|
| 297 |
+
|
| 298 |
+
def build(n_per_source, split="train", seed=0, exclude=(), only=(), revisions=None, sources=None, repos=None):
|
| 299 |
+
sources = SOURCES if sources is None else sources
|
| 300 |
+
repos = REPOS if repos is None else repos
|
| 301 |
+
unknown = (set(exclude) | set(only)) - sources.keys()
|
| 302 |
+
if unknown:
|
| 303 |
+
raise ValueError(f"unknown sources: {sorted(unknown)}")
|
| 304 |
+
reqs = []
|
| 305 |
+
for name, (fn, tr, te) in sources.items():
|
| 306 |
+
if name in exclude or (only and name not in only): continue
|
| 307 |
+
src = Source(source_seed(seed, name), (revisions or {}).get(repos[name]))
|
| 308 |
+
source_split = tr if split == "train" else te
|
| 309 |
+
records = fn(source_split, n_per_source, src)
|
| 310 |
+
if len(records) != len(src.origins):
|
| 311 |
+
raise ValueError(f"provenance mismatch for {name}")
|
| 312 |
+
for record, origin in zip(records, src.origins):
|
| 313 |
+
record["_meta"] = {**origin, "source": name, "repo": repos[name], "revision": src.revision,
|
| 314 |
+
"split": source_split, "id": f"{name}/{source_split}/{origin['row']}"}
|
| 315 |
+
reqs.extend(records)
|
| 316 |
+
random.Random(seed).shuffle(reqs)
|
| 317 |
+
return reqs
|
| 318 |
+
|
| 319 |
+
|
| 320 |
+
def augment(req, rng, p_none=0.1, p_none_distract=0.12, p_distract=0.15):
|
| 321 |
+
"""Choice only: permute option order (always); sometimes add a 'none of the above' option, either as the correct
|
| 322 |
+
answer (true option removed) or as a wrong alternative (true option kept); sometimes add an irrelevant distractor."""
|
| 323 |
+
if min(p_none, p_none_distract, p_distract) < 0 or p_none + p_none_distract + p_distract > 1:
|
| 324 |
+
raise ValueError("augmentation probabilities must be nonnegative and sum to at most one")
|
| 325 |
+
out = {"state": req["state"], "questions": {}}
|
| 326 |
+
for qid, q in req["questions"].items():
|
| 327 |
+
if q["type"] != "choice":
|
| 328 |
+
out["questions"][qid] = q; continue
|
| 329 |
+
crit, y = dict(q["criteria"]), q["label"]
|
| 330 |
+
if q.get("target") is not None: # soft-target questions: permute only; inserting or swapping options would change the target's meaning
|
| 331 |
+
keys = list(crit); rng.shuffle(keys); out["questions"][qid] = {**q, "criteria": {k: crit[k] for k in keys}}; continue
|
| 332 |
+
r = rng.random()
|
| 333 |
+
none_options = [(k, v) for k, v in NONE_OPTIONS if k not in crit]
|
| 334 |
+
distractors = [k for k in DISTRACTORS if k not in crit]
|
| 335 |
+
if len(crit) > 2 and r < p_none and none_options:
|
| 336 |
+
nk, nd = rng.choice(none_options); crit.pop(y); crit[nk] = nd; y = nk
|
| 337 |
+
elif p_none <= r < p_none + p_none_distract and len(crit) < 255 and none_options:
|
| 338 |
+
nk, nd = rng.choice(none_options); crit[nk] = nd
|
| 339 |
+
elif p_none + p_none_distract <= r < p_none + p_none_distract + p_distract and len(crit) < 255 and distractors:
|
| 340 |
+
k = rng.choice(distractors); crit[k] = DISTRACTORS[k]
|
| 341 |
+
keys = list(crit); rng.shuffle(keys)
|
| 342 |
+
out["questions"][qid] = {**q, "criteria": {k: crit[k] for k in keys}, "label": y}
|
| 343 |
+
return out
|
| 344 |
+
|
| 345 |
+
|
| 346 |
+
def none_pair(req, rng):
|
| 347 |
+
"""Minimal pair for the none-of-the-above shortcut: the same state and question rendered twice, once with the true
|
| 348 |
+
option present (a none option is wrong) and once with it removed (the same none option is right). Everything else,
|
| 349 |
+
including option order, is identical, so the only difference the model can use is whether the evidence matches an
|
| 350 |
+
option. Returns [] when the request has no eligible Choice (>=3 options, no none key already present, no soft target:
|
| 351 |
+
removing the labelled option or adding a none option would change what the target means, as in augment)."""
|
| 352 |
+
eligible = [(qid, q) for qid, q in req["questions"].items() if q["type"] == "choice" and len(q["criteria"]) >= 3 and q.get("target") is None]
|
| 353 |
+
if not eligible: return []
|
| 354 |
+
qid, q = rng.choice(eligible)
|
| 355 |
+
nk, nd = rng.choice([o for o in NONE_OPTIONS if o[0] not in q["criteria"]] or [("none_of_these", None)])
|
| 356 |
+
keys = list(q["criteria"]) + [nk]; rng.shuffle(keys)
|
| 357 |
+
present = {**q, "criteria": {k: (nd if k == nk else q["criteria"][k]) for k in keys}}
|
| 358 |
+
absent = {**present, "criteria": {k: v for k, v in present["criteria"].items() if k != q["label"]}, "label": nk}
|
| 359 |
+
return [{"state": req["state"], "questions": {qid: present}}, {"state": req["state"], "questions": {qid: absent}}]
|
| 360 |
+
|
| 361 |
+
|
| 362 |
+
def load_records(path, source="custom"):
|
| 363 |
+
"""Labelled requests from a JSONL file, one per line, in the API's request shape plus a `label` on every question:
|
| 364 |
+
|
| 365 |
+
{"state": "...", "questions": {"team": {"type": "choice", "instructions": "...", "criteria": {"billing": "...", "other": null}, "label": "billing"},
|
| 366 |
+
"urgent": {"type": "noul", "instructions": "...", "label": true},
|
| 367 |
+
"priority": {"type": "score", "instructions": "...", "criteria": ["low", "medium", "high"], "label": 2}}}
|
| 368 |
+
|
| 369 |
+
Labels: the option name for choice, true/false for noul, the level index (from 0) for score. `_meta` and per-question
|
| 370 |
+
`src` are filled in so the records behave like a frozen suite's (source = `source`, id = line number)."""
|
| 371 |
+
records = []
|
| 372 |
+
with Path(path).open(encoding="utf-8") as f:
|
| 373 |
+
for n, line in enumerate(f):
|
| 374 |
+
if not line.strip(): continue
|
| 375 |
+
r = json.loads(line)
|
| 376 |
+
if "state" not in r or not isinstance(r.get("questions"), dict) or not r["questions"]:
|
| 377 |
+
raise ValueError(f"{path}:{n + 1}: a record needs a state and a non-empty questions object")
|
| 378 |
+
for qid, q in r["questions"].items():
|
| 379 |
+
if "label" not in q: raise ValueError(f"{path}:{n + 1}: question {qid!r} has no label")
|
| 380 |
+
q.setdefault("src", f"{source}_{q['type']}")
|
| 381 |
+
text = json.dumps(r["state"], sort_keys=True, ensure_ascii=False) if not isinstance(r["state"], str) else r["state"]
|
| 382 |
+
r["_meta"] = {**{"source": source, "variant": "clean", "id": f"{source}/{n}", "group_id": f"{source}/{n}", "row": n, "split": "custom",
|
| 383 |
+
"text_sha256": hashlib.sha256(" ".join(text.casefold().split()).encode()).hexdigest()}, **r.get("_meta", {})}
|
| 384 |
+
records.append(r)
|
| 385 |
+
if not records: raise ValueError(f"{path}: no records")
|
| 386 |
+
return records
|
| 387 |
+
|
| 388 |
+
|
| 389 |
+
def api_request(record):
|
| 390 |
+
"""The /v1/systemone request body for a labelled record: state and typed questions only, never labels, targets or
|
| 391 |
+
metadata (this is what leaves the machine when a remote predictor is scored)."""
|
| 392 |
+
return {"state": record["state"], "questions": {
|
| 393 |
+
qid: {k: v for k, v in q.items() if k in ("type", "instructions", "criteria")}
|
| 394 |
+
for qid, q in record["questions"].items()}}
|
| 395 |
+
|
| 396 |
+
|
| 397 |
+
def materialize(req):
|
| 398 |
+
"""Labelled request -> internal record via the serving path (api.to_record), attaching int labels and src."""
|
| 399 |
+
rec, meta = to_record(SystemOneRequest.model_validate(api_request(req)))
|
| 400 |
+
for q, m, (qid, src_q) in zip(rec["questions"], meta, req["questions"].items()):
|
| 401 |
+
y = src_q["label"]
|
| 402 |
+
q["label"] = m["keys"].index(y) if m["type"] == "choice" else int(y) # noul labels are bools, score labels level indices
|
| 403 |
+
q["src"] = src_q["src"]; q["qtype"] = m["type"]; q["qid"] = qid; q["keys"] = m["keys"]
|
| 404 |
+
if src_q.get("target") is not None:
|
| 405 |
+
# soft target keyed by option name (choice), "false"/"true" (noul) or level index as a string (score); options the
|
| 406 |
+
# target does not name get 0, then the vector is normalised. Used for unknowable records (uniform over the options).
|
| 407 |
+
t = [float(src_q["target"].get(k, 0.0)) for k in q["keys"]]
|
| 408 |
+
if sum(t) <= 0: raise ValueError(f"target for {qid} puts no mass on any option")
|
| 409 |
+
q["target"] = [x / sum(t) for x in t]
|
| 410 |
+
return rec
|
runtime/kev/device.py
ADDED
|
@@ -0,0 +1,30 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""The accelerator this process uses: cuda, then mps, then cpu."""
|
| 2 |
+
import torch
|
| 3 |
+
|
| 4 |
+
|
| 5 |
+
def default_device():
|
| 6 |
+
return "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
def sync(device):
|
| 10 |
+
"""Wait for queued kernels, so wall-clock timings around a forward pass are real."""
|
| 11 |
+
if device == "mps": torch.mps.synchronize()
|
| 12 |
+
elif device == "cuda": torch.cuda.synchronize()
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
def empty_cache(device):
|
| 16 |
+
if device == "mps": torch.mps.empty_cache()
|
| 17 |
+
elif device == "cuda": torch.cuda.empty_cache()
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
def out_of_memory(e):
|
| 21 |
+
"""Whether a torch allocator ran out: CUDA raises torch.OutOfMemoryError, MPS a plain RuntimeError with this message.
|
| 22 |
+
The MLX backend's Metal errors are neither."""
|
| 23 |
+
return isinstance(e, torch.OutOfMemoryError) or isinstance(e, RuntimeError) and str(e).startswith("MPS backend out of memory")
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
def allocated_bytes(device):
|
| 27 |
+
"""Bytes currently allocated on the device (MPS) or the peak since the process started (CUDA); 0 on CPU."""
|
| 28 |
+
if device == "mps": return torch.mps.current_allocated_memory()
|
| 29 |
+
if device == "cuda": return torch.cuda.max_memory_allocated()
|
| 30 |
+
return 0
|
runtime/kev/evaluate.py
ADDED
|
@@ -0,0 +1,219 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tests of the architectural hypotheses on held-out data (all through the serving-format path).
|
| 2 |
+
|
| 3 |
+
1. accuracy + ECE per source / per question type (Score also reports MAE of the expected level)
|
| 4 |
+
vs a zero-shot letter-logit baseline from the same base model (K<=8)
|
| 5 |
+
2. permutation stability (argmax flip rate, prob spread of correct option) on Choice
|
| 6 |
+
3. IIA: does appending an irrelevant option shift log-odds between two existing options?
|
| 7 |
+
4. isolation: can question B read a secret placed in sibling question A? (must not) / in state (should)
|
| 8 |
+
5. latency + equality: packed N questions vs N separate calls
|
| 9 |
+
"""
|
| 10 |
+
import argparse, json, math, random, time, string
|
| 11 |
+
from collections import defaultdict
|
| 12 |
+
import numpy as np
|
| 13 |
+
import torch
|
| 14 |
+
import torch.nn.functional as F
|
| 15 |
+
from .checkpoint import Checkpoint, LoadOptions
|
| 16 |
+
from .data import NONE_OPTIONS, build, augment, materialize, DISTRACTORS
|
| 17 |
+
from .device import default_device, empty_cache, sync
|
| 18 |
+
from .metrics import ece
|
| 19 |
+
from .suite import write_json
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def _probs(tok, model, req):
|
| 23 |
+
rec = materialize(req)
|
| 24 |
+
return rec, model.probs(model.encode(tok, rec, strict=True))
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def _one(req, qid):
|
| 28 |
+
return {"state": req["state"], "questions": {qid: req["questions"][qid]}}
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
def test_accuracy(tok, model, reqs, rng):
|
| 32 |
+
by = defaultdict(lambda: {"conf": [], "ok": [], "nll": [], "mae": []})
|
| 33 |
+
for r in reqs:
|
| 34 |
+
try: rec, ps = _probs(tok, model, augment(r, rng, p_none=0, p_none_distract=0, p_distract=0))
|
| 35 |
+
except ValueError as exc: raise ValueError("Evaluation rejected an example; refusing partial metrics") from exc
|
| 36 |
+
for q, p in zip(rec["questions"], ps):
|
| 37 |
+
p = p.numpy(); y = q["label"]; d = by[q["src"]]
|
| 38 |
+
d["conf"].append(float(p.max())); d["ok"].append(int(p.argmax() == y)); d["nll"].append(-math.log(max(p[y], 1e-9)))
|
| 39 |
+
if q["qtype"] == "score": d["mae"].append(abs(float((np.arange(len(p)) * p).sum()) - y))
|
| 40 |
+
res = {}
|
| 41 |
+
for s, d in by.items():
|
| 42 |
+
res[s] = {"n": len(d["ok"]), "acc": float(np.mean(d["ok"])), "ece": ece(d["conf"], d["ok"]), "nll": float(np.mean(d["nll"])), "mean_conf": float(np.mean(d["conf"]))}
|
| 43 |
+
if d["mae"]: res[s]["score_mae_levels"] = float(np.mean(d["mae"]))
|
| 44 |
+
allc = sum((d["conf"] for d in by.values()), []); allo = sum((d["ok"] for d in by.values()), [])
|
| 45 |
+
res["ALL"] = {"n": len(allo), "acc": float(np.mean(allo)), "ece": ece(allc, allo), "mean_conf": float(np.mean(allc))}
|
| 46 |
+
return res
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def test_permutation(tok, model, reqs, rng, n_perm=4):
|
| 50 |
+
flips, spreads, n = 0, [], 0
|
| 51 |
+
for r in reqs:
|
| 52 |
+
for qid, q in r["questions"].items():
|
| 53 |
+
if q["type"] != "choice" or len(q["criteria"]) < 3: continue
|
| 54 |
+
argmaxes, pcorrect = [], []
|
| 55 |
+
for _ in range(n_perm):
|
| 56 |
+
keys = list(q["criteria"]); rng.shuffle(keys)
|
| 57 |
+
q2 = {**q, "criteria": {k: q["criteria"][k] for k in keys}}
|
| 58 |
+
try: _, ps = _probs(tok, model, {"state": r["state"], "questions": {qid: q2}})
|
| 59 |
+
except ValueError as exc: raise ValueError("Permutation evaluation rejected an example") from exc
|
| 60 |
+
p = ps[0].numpy(); argmaxes.append(keys[int(p.argmax())]); pcorrect.append(float(p[keys.index(q["label"])]))
|
| 61 |
+
if len(pcorrect) == n_perm:
|
| 62 |
+
n += 1; flips += int(len(set(argmaxes)) > 1); spreads.append(max(pcorrect) - min(pcorrect))
|
| 63 |
+
return {"n": n, "argmax_flip_rate": flips / max(n, 1), "mean_prob_spread": float(np.mean(spreads)), "p90_prob_spread": float(np.percentile(spreads, 90))}
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def test_iia(tok, model, reqs, rng):
|
| 67 |
+
"""Log-odds between the top-2 original options, before vs after appending an irrelevant option."""
|
| 68 |
+
shifts = []
|
| 69 |
+
for r in reqs:
|
| 70 |
+
for qid, q in r["questions"].items():
|
| 71 |
+
if q["type"] != "choice" or not 3 <= len(q["criteria"]) <= 10: continue
|
| 72 |
+
try:
|
| 73 |
+
_, ps0 = _probs(tok, model, {"state": r["state"], "questions": {qid: q}})
|
| 74 |
+
k = rng.choice(list(DISTRACTORS))
|
| 75 |
+
_, ps1 = _probs(tok, model, {"state": r["state"], "questions": {qid: {**q, "criteria": {**q["criteria"], k: DISTRACTORS[k]}}}})
|
| 76 |
+
except ValueError as exc: raise ValueError("Evaluation rejected an example; refusing partial metrics") from exc
|
| 77 |
+
p0, p1 = ps0[0].numpy(), ps1[0].numpy()
|
| 78 |
+
a, b = np.argsort(-p0)[:2]
|
| 79 |
+
shifts.append(math.log(max(p1[a], 1e-6) / max(p1[b], 1e-6)) - math.log(max(p0[a], 1e-6) / max(p0[b], 1e-6)))
|
| 80 |
+
s = np.array(shifts)
|
| 81 |
+
return {"n": len(s), "mean_abs_logodds_shift": float(np.abs(s).mean()), "mean_shift": float(s.mean()), "p90_abs_shift": float(np.percentile(np.abs(s), 90))}
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
def test_none_of_the_above(tok, model, reqs, rng):
|
| 85 |
+
"""Add a 'none of the above' option. When the true option is still present it should get little mass;
|
| 86 |
+
when the true option is removed it should be chosen. A shortcut model picks it in both cases."""
|
| 87 |
+
p_present, hit_present, hit_absent, n = [], 0, 0, 0
|
| 88 |
+
for r in reqs:
|
| 89 |
+
for qid, q in r["questions"].items():
|
| 90 |
+
if q["type"] != "choice" or len(q["criteria"]) < 3: continue
|
| 91 |
+
nk, nd = rng.choice(NONE_OPTIONS)
|
| 92 |
+
if nk in q["criteria"]: continue
|
| 93 |
+
try:
|
| 94 |
+
present = {**q, "criteria": {**q["criteria"], nk: nd}}
|
| 95 |
+
_, ps = _probs(tok, model, {"state": r["state"], "questions": {qid: present}})
|
| 96 |
+
keys = list(present["criteria"]); p = ps[0].numpy()
|
| 97 |
+
none_probability = float(p[keys.index(nk)])
|
| 98 |
+
present_correct = int(keys[int(p.argmax())] == q["label"])
|
| 99 |
+
absent = {**q, "label": nk, "criteria": {k: v for k, v in present["criteria"].items() if k != q["label"]}}
|
| 100 |
+
_, ps = _probs(tok, model, {"state": r["state"], "questions": {qid: absent}})
|
| 101 |
+
keys = list(absent["criteria"]); p = ps[0].numpy()
|
| 102 |
+
p_present.append(none_probability); hit_present += present_correct
|
| 103 |
+
hit_absent += int(keys[int(p.argmax())] == nk); n += 1
|
| 104 |
+
except ValueError as exc: raise ValueError("Evaluation rejected an example; refusing partial metrics") from exc
|
| 105 |
+
return {"n": n, "true_option_present": {"mean_p_none": float(np.mean(p_present)), "p_none_over_0.5_rate": float(np.mean(np.array(p_present) > 0.5)), "acc": hit_present / max(n, 1)},
|
| 106 |
+
"true_option_removed": {"picks_none_rate": hit_absent / max(n, 1)}}
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
def test_isolation(tok, model, rng, n=20):
|
| 110 |
+
"""Secret code in sibling question vs in state. Probe asks which code was mentioned."""
|
| 111 |
+
res = {"sibling": [], "state": [], "absent": []}
|
| 112 |
+
weather = "The weather is nice today and the park is full of people."
|
| 113 |
+
for _ in range(n):
|
| 114 |
+
codes = ["".join(rng.choices(string.ascii_uppercase, k=5)) + "-" + str(rng.randint(1000, 9999)) for _ in range(3)]
|
| 115 |
+
keys = ["c0", "c1", "c2", "none"]; rng.shuffle(keys)
|
| 116 |
+
crit = {k: (codes[int(k[1])] if k != "none" else "None of these codes appears") for k in keys}
|
| 117 |
+
probe = {"type": "choice", "instructions": "Which secret code is mentioned in this request?", "criteria": crit, "label": "c0", "src": "probe"}
|
| 118 |
+
sib = {"type": "noul", "instructions": f"The secret code for this request is {codes[0]}. Is the weather described as nice?", "label": True, "src": "x"}
|
| 119 |
+
sib0 = {"type": "noul", "instructions": "Is the weather described as nice?", "label": True, "src": "x"}
|
| 120 |
+
for cond, state, a in [("sibling", weather, sib), ("absent", weather, sib0), ("state", weather + f" The secret code for this request is {codes[0]}.", sib0)]:
|
| 121 |
+
_, ps = _probs(tok, model, {"state": state, "questions": {"a": a, "probe": probe}})
|
| 122 |
+
res[cond].append(float(ps[1].numpy()[keys.index("c0")]))
|
| 123 |
+
return {k: {"mean_p_code": float(np.mean(v)), "min": float(np.min(v)), "max": float(np.max(v))} for k, v in res.items()}
|
| 124 |
+
|
| 125 |
+
|
| 126 |
+
def test_packed_vs_separate(tok, model, reqs, rng, n=30):
|
| 127 |
+
diffs, t_pack, t_sep, nq = [], 0.0, 0.0, 0
|
| 128 |
+
for r in [x for x in reqs if len(x["questions"]) >= 2][:n]:
|
| 129 |
+
enc = model.encode(tok, materialize(r)) # legacy prototype eval (not used for releases): truncates at the default training context by design
|
| 130 |
+
sync(model.device); t = time.time(); pp = model.probs(enc); sync(model.device); t_pack += time.time() - t
|
| 131 |
+
for qi, qid in enumerate(r["questions"]):
|
| 132 |
+
e1 = model.encode(tok, materialize(_one(r, qid)))
|
| 133 |
+
sync(model.device); t = time.time(); p1 = model.probs(e1)[0]; sync(model.device); t_sep += time.time() - t
|
| 134 |
+
diffs.append(float((pp[qi] - p1).abs().max())); nq += 1
|
| 135 |
+
return {"n_questions": nq, "max_abs_prob_diff": float(np.max(diffs)), "mean_abs_prob_diff": float(np.mean(diffs)), "packed_s": t_pack, "separate_s": t_sep, "speedup": t_sep / t_pack}
|
| 136 |
+
|
| 137 |
+
|
| 138 |
+
def baseline_letter_logits(base, reqs, dev, rng, chat=False):
|
| 139 |
+
"""Zero-shot MCQ baseline on the same rendered text: read next-token logits over option letters. K<=8.
|
| 140 |
+
chat=True wraps the prompt in the model's chat template (for -Instruct models) and asks for the letter only."""
|
| 141 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
| 142 |
+
tok = AutoTokenizer.from_pretrained(base); lm = AutoModelForCausalLM.from_pretrained(base, dtype=torch.float32).to(dev).eval()
|
| 143 |
+
letters = list("ABCDEFGH")
|
| 144 |
+
letter_ids = [tok(("" if chat else " ") + L, add_special_tokens=False).input_ids[0] for L in letters]
|
| 145 |
+
by = defaultdict(lambda: {"conf": [], "ok": []})
|
| 146 |
+
with torch.no_grad():
|
| 147 |
+
for r in reqs:
|
| 148 |
+
rec = materialize(augment(r, rng, p_none=0, p_none_distract=0, p_distract=0))
|
| 149 |
+
for q in rec["questions"]:
|
| 150 |
+
K = len(q["options"])
|
| 151 |
+
if K > 8: continue
|
| 152 |
+
body = f"{rec['state']}\n\nQuestion: {q['instr']}\n" + "".join(f"{letters[i]}. {o}\n" for i, o in enumerate(q["options"]))
|
| 153 |
+
if chat:
|
| 154 |
+
prompt = tok.apply_chat_template([{"role": "user", "content": body + "\nAnswer with the letter of the best option only."}], tokenize=False, add_generation_prompt=True)
|
| 155 |
+
else:
|
| 156 |
+
prompt = body + "Answer:"
|
| 157 |
+
ids = tok(prompt, return_tensors="pt", truncation=True, max_length=1200).to(dev)
|
| 158 |
+
p = F.softmax(lm(**ids).logits[0, -1, letter_ids[:K]], -1).cpu().numpy()
|
| 159 |
+
by[q["src"]]["conf"].append(float(p.max())); by[q["src"]]["ok"].append(int(p.argmax() == q["label"]))
|
| 160 |
+
del lm
|
| 161 |
+
return {s: {"n": len(d["ok"]), "acc": float(np.mean(d["ok"])), "ece": ece(d["conf"], d["ok"]), "mean_conf": float(np.mean(d["conf"]))} for s, d in by.items()}
|
| 162 |
+
|
| 163 |
+
|
| 164 |
+
def test_temperature(tok, model, reqs, rng):
|
| 165 |
+
"""Fit one global temperature on even-indexed records (by NLL), report NLL/ECE on odd-indexed ones before and after."""
|
| 166 |
+
fit, held = [], []
|
| 167 |
+
with torch.no_grad():
|
| 168 |
+
for i, r in enumerate(reqs): # legacy prototype eval (not used for releases): encode truncates at the default training context by design
|
| 169 |
+
try: rec = materialize(augment(r, rng, p_none=0, p_none_distract=0, p_distract=0)); zs = model(model.encode(tok, rec))
|
| 170 |
+
except ValueError as exc: raise ValueError("Evaluation rejected an example; refusing partial metrics") from exc
|
| 171 |
+
for z, q in zip(zs, rec["questions"]):
|
| 172 |
+
(fit if i % 2 == 0 else held).append((z.cpu(), q["label"]))
|
| 173 |
+
logT = torch.zeros(1, requires_grad=True); opt = torch.optim.LBFGS([logT], lr=0.1, max_iter=50)
|
| 174 |
+
def nll(pairs, T): return torch.stack([F.cross_entropy((z / T)[None], torch.tensor([y])) for z, y in pairs]).mean()
|
| 175 |
+
def closure():
|
| 176 |
+
opt.zero_grad(); l = nll(fit, logT.exp()); l.backward(); return l
|
| 177 |
+
opt.step(closure); T = float(logT.exp())
|
| 178 |
+
def stats(pairs, T):
|
| 179 |
+
ps = [F.softmax(z / T, -1) for z, _ in pairs]
|
| 180 |
+
return {"nll": float(nll(pairs, T)), "ece": ece([float(p.max()) for p in ps], [int(p.argmax() == y) for p, (_, y) in zip(ps, pairs)])}
|
| 181 |
+
return {"T": T, "n_fit": len(fit), "n_held": len(held), "before": stats(held, 1.0), "after": stats(held, T)}
|
| 182 |
+
|
| 183 |
+
|
| 184 |
+
def main():
|
| 185 |
+
ap = argparse.ArgumentParser()
|
| 186 |
+
ap.add_argument("--run", default="runs/kev")
|
| 187 |
+
ap.add_argument("--n_per_source", type=int, default=150)
|
| 188 |
+
ap.add_argument("--baseline", action="store_true", help="zero-shot letter-logit baseline from the base model")
|
| 189 |
+
ap.add_argument("--baseline_instruct", default="", help="e.g. Qwen/Qwen2.5-0.5B-Instruct: chat-template letter-logit baseline")
|
| 190 |
+
ap.add_argument("--seed", type=int, default=1)
|
| 191 |
+
a = ap.parse_args()
|
| 192 |
+
dev = default_device()
|
| 193 |
+
rng = random.Random(a.seed)
|
| 194 |
+
reqs = build(a.n_per_source, "test", a.seed)
|
| 195 |
+
ck = Checkpoint(a.run)
|
| 196 |
+
out = {"run": a.run, "holdout_sources": ck.meta.holdout}
|
| 197 |
+
if a.baseline:
|
| 198 |
+
out["baseline_zero_shot_base"] = baseline_letter_logits(ck.meta.base, reqs, dev, random.Random(a.seed)); print(json.dumps(out, indent=1), flush=True)
|
| 199 |
+
if a.baseline_instruct:
|
| 200 |
+
out["baseline_zero_shot_instruct"] = baseline_letter_logits(a.baseline_instruct, reqs, dev, random.Random(a.seed), chat=True); print(json.dumps(out["baseline_zero_shot_instruct"], indent=1), flush=True)
|
| 201 |
+
empty_cache(dev)
|
| 202 |
+
tok, model = ck.load(dev, LoadOptions.from_env())
|
| 203 |
+
for name, fn in [("accuracy_calibration", lambda: test_accuracy(tok, model, reqs, random.Random(a.seed))),
|
| 204 |
+
("temperature_scaling", lambda: test_temperature(tok, model, reqs, random.Random(a.seed))),
|
| 205 |
+
("permutation", lambda: test_permutation(tok, model, reqs[:150], rng)),
|
| 206 |
+
("iia", lambda: test_iia(tok, model, reqs[:250], rng)),
|
| 207 |
+
("none_of_the_above", lambda: test_none_of_the_above(tok, model, reqs[:200], rng)),
|
| 208 |
+
("isolation", lambda: test_isolation(tok, model, rng)),
|
| 209 |
+
("packed_vs_separate", lambda: test_packed_vs_separate(tok, model, reqs, rng))]:
|
| 210 |
+
t = time.time(); out[name] = fn(); print(f"== {name} ({time.time()-t:.0f}s)\n{json.dumps(out[name], indent=1)}", flush=True)
|
| 211 |
+
if out["holdout_sources"]:
|
| 212 |
+
acc = out["accuracy_calibration"]
|
| 213 |
+
out["held_out_sources"] = {s: acc[s] for s in acc if any(s.startswith(h) for h in out["holdout_sources"])}
|
| 214 |
+
print("== held_out_sources (never seen in training)\n" + json.dumps(out["held_out_sources"], indent=1))
|
| 215 |
+
write_json(f"{a.run}/eval.json", out)
|
| 216 |
+
|
| 217 |
+
|
| 218 |
+
if __name__ == "__main__":
|
| 219 |
+
main()
|
runtime/kev/experiment.py
ADDED
|
@@ -0,0 +1,492 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Configuration-only experiment runner.
|
| 2 |
+
|
| 3 |
+
A study is a JSON plan: a list of 1..8 trials, each a dict of allowlisted training parameters plus a base model
|
| 4 |
+
whose revision is pinned in the suite manifest. Every trial trains on the suite's training partition, fits a
|
| 5 |
+
temperature on the calibration partition, and is scored on the development partition. The locked test is never
|
| 6 |
+
read here.
|
| 7 |
+
|
| 8 |
+
Two execution modes share one code path:
|
| 9 |
+
local : trials run sequentially on this machine (one GPU job at a time); `--wait-pid` queues behind a run.
|
| 10 |
+
modal : `modal_app.py` runs `execute_trial` once per container and `--aggregate` ranks the results afterwards.
|
| 11 |
+
"""
|
| 12 |
+
import argparse
|
| 13 |
+
import copy
|
| 14 |
+
import fcntl
|
| 15 |
+
import gc
|
| 16 |
+
import json
|
| 17 |
+
import os
|
| 18 |
+
import platform
|
| 19 |
+
import re
|
| 20 |
+
import shutil
|
| 21 |
+
import subprocess
|
| 22 |
+
import sys
|
| 23 |
+
import time
|
| 24 |
+
from contextlib import contextmanager
|
| 25 |
+
from pathlib import Path
|
| 26 |
+
|
| 27 |
+
import torch
|
| 28 |
+
|
| 29 |
+
from kev.benchmark import evaluate_records
|
| 30 |
+
from kev.checkpoint import LoadOptions
|
| 31 |
+
from kev.device import default_device, empty_cache
|
| 32 |
+
from kev.full_ft import snapshot_fractions, too_many_snapshots
|
| 33 |
+
from kev.metrics import fit_temperature, paired_bootstrap
|
| 34 |
+
from kev.model import MAX_STATE, MAX_TRAIN_STATE
|
| 35 |
+
from kev.predictors import LocalPredictor
|
| 36 |
+
from kev.suite import CONTEXT, ENCODING, digest, load_split, read_json, read_manifest, record_digest, validate_training, write_json
|
| 37 |
+
|
| 38 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 39 |
+
DEFAULTS = {"epochs": 1, "seed": 0, "lr": 0.0002, "lora": 16, "accum": 8, "batch": 1,
|
| 40 |
+
"perm_kl": 0.0, "perm_frac": 0.3, "ord_w": 0.0,
|
| 41 |
+
"p_none": 0.1, "p_none_distract": 0.12, "p_distract": 0.15, "p_none_pair": 0.0, "synthetic_repeat": 1, "public_frac": 1.0, "head_lr": 0.0, "weight_decay": 0.01, "anchor_w": 0.0,
|
| 42 |
+
"label_smoothing": 0.0, "brier_w": 0.0, "focal_gamma": 0.0}
|
| 43 |
+
RANGES = {"epochs": (1, 5), "seed": (0, 10000), "lr": (1e-6, 0.001), "lora": (1, 64), "accum": (1, 256), "batch": (1, 64),
|
| 44 |
+
"perm_kl": (0, 2), "perm_frac": (0, 1), "ord_w": (0, 2),
|
| 45 |
+
"label_smoothing": (0, 0.2), "brier_w": (0, 2), "focal_gamma": (0, 4),
|
| 46 |
+
"p_none": (0, 0.4), "p_none_distract": (0, 0.4), "p_distract": (0, 0.4), "p_none_pair": (0, 1), "synthetic_repeat": (1, 6), "public_frac": (0.05, 1.0), "head_lr": (0, 0.01), "weight_decay": (0, 0.3), "anchor_w": (0, 5)}
|
| 47 |
+
CHOICES = {"dtype": ("fp32", "bf16"), "checkpointing": (0, 1), "option_isolation": (0, 1), "special_embeddings": (0, 1), "head_dim": (128, 256, 512, 1024),
|
| 48 |
+
"lora_targets": ("all", "dense", "attn", "qv"), "weights_dtype": ("fp32", "bf16"), "full_ft": (0, 1), "length_sort": (0, 1), "shared_prefix": (0, 1)}
|
| 49 |
+
CHOICE_DEFAULTS = {"dtype": "fp32", "checkpointing": 0, "option_isolation": 0, "special_embeddings": 0, "head_dim": 256, "lora_targets": "all", "weights_dtype": "fp32", "full_ft": 0, "length_sort": 0, "shared_prefix": None} # kev.train's defaults for the categorical knobs (shared_prefix: on with full_ft)
|
| 50 |
+
# optional integer knobs, passed to kev.train only when a trial sets them (so existing plans keep their config hashes)
|
| 51 |
+
OPTIONAL_INTS = {"max_state": (MAX_STATE, MAX_TRAIN_STATE), "row_budget": (0, 65536), "max_steps": (0, 100000), "save_every_steps": (0, 100000),
|
| 52 |
+
"snapshot_every_steps": (0, 100000), "none_pair_max_state": (1, MAX_TRAIN_STATE), "pass_tokens_max": (1, 1 << 20)}
|
| 53 |
+
# full-weight trials write snapshots (kev.full_ft.SnapshotWriter: loadable bf16 checkpoints, kept on the runs volume) into
|
| 54 |
+
# <trial>/snapshots/step-<N>/checkpoint after these fractions of their optimizer steps, unless the trial sets
|
| 55 |
+
# snapshot_fractions ("none" for none). Applied by train_checkpoint like RESUME_MINUTES, so plans keep their config hashes.
|
| 56 |
+
SNAPSHOT_FRACTIONS = "0.25,0.5,0.75"
|
| 57 |
+
# trial keys the container acts on, not kev.train: snapshot_hub_repo (a private Hub model repo) makes modal_app mirror
|
| 58 |
+
# each committed snapshot and the final checkpoint there (modal_app.run_mirror); off unless a trial sets it
|
| 59 |
+
TRIAL_ONLY = ("snapshot_hub_repo",)
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def validated_trial(value, manifest):
|
| 63 |
+
if not isinstance(value, dict) or set(value) - (DEFAULTS.keys() | CHOICES.keys() | OPTIONAL_INTS.keys() | {"base", "train_sources", "base_revision", "anchor", "anchor_sources", "init_from", "data", "replay", "snapshot_fractions", *TRIAL_ONLY}):
|
| 64 |
+
raise ValueError("trial may change only the allowlisted training parameters and base")
|
| 65 |
+
result = {**DEFAULTS, **value}
|
| 66 |
+
if result.get("full_ft") and result.get("weights_dtype") != "bf16":
|
| 67 |
+
raise ValueError("full_ft trains bf16 weights: set weights_dtype bf16")
|
| 68 |
+
if result.get("save_every_steps") and not result.get("full_ft"):
|
| 69 |
+
raise ValueError("save_every_steps writes resume points, which are for full_ft trials")
|
| 70 |
+
if ("snapshot_fractions" in result or result.get("snapshot_every_steps") or "snapshot_hub_repo" in result) and not result.get("full_ft"):
|
| 71 |
+
raise ValueError("snapshots (and their Hub mirror) are for full_ft trials")
|
| 72 |
+
if "snapshot_fractions" in result and not isinstance(result["snapshot_fractions"], str):
|
| 73 |
+
raise ValueError('snapshot_fractions is a string: "0.25,0.5" or "none"')
|
| 74 |
+
if "snapshot_hub_repo" in result and not re.fullmatch(r"[\w.-]+/[\w.-]+", str(result["snapshot_hub_repo"])):
|
| 75 |
+
raise ValueError("snapshot_hub_repo is a Hub model repo id, owner/name (it must be private; the mirror refuses a public one)")
|
| 76 |
+
if "data" in result and not re.fullmatch(r"evals/[\w./-]+\.jsonl", str(result["data"])):
|
| 77 |
+
raise ValueError("data must be a .jsonl under evals/ (shipped with the image, hashed in provenance)")
|
| 78 |
+
if "replay" in result and (not isinstance(result["replay"], int) or not 0 <= result["replay"] <= 20000 or "data" not in result):
|
| 79 |
+
raise ValueError("replay is an int <= 20000 and needs data")
|
| 80 |
+
for key, (lo, hi) in OPTIONAL_INTS.items():
|
| 81 |
+
if key in result and (isinstance(result[key], bool) or not isinstance(result[key], int) or not lo <= result[key] <= hi):
|
| 82 |
+
raise ValueError(f"{key} is an int in [{lo}, {hi}] (optional; kev.train's default when absent)")
|
| 83 |
+
if "none_pair_max_state" in result and not result["p_none_pair"] > 0:
|
| 84 |
+
raise ValueError("none_pair_max_state gates none pairs: it needs p_none_pair > 0")
|
| 85 |
+
# pass_tokens_max also needs a hybrid base (kev.train refuses an attention-only one once the model is built: the packed
|
| 86 |
+
# mask it runs is not what pass_tokens measures); the base's architecture is not known here
|
| 87 |
+
if "pass_tokens_max" in result and (not result.get("length_sort") or result.get("row_budget") or result["perm_kl"] > 0):
|
| 88 |
+
raise ValueError("pass_tokens_max caps the passes length_sort 1 plans: it needs length_sort 1, and neither row_budget nor perm_kl "
|
| 89 |
+
"(and a hybrid base; kev.train checks that)")
|
| 90 |
+
if result.get("full_ft"): # kev.budget.MAX_SNAPSHOTS, over max_steps (an every-N plan needs it); kev.train checks the real count again
|
| 91 |
+
fractions = snapshot_fractions(result.get("snapshot_fractions", SNAPSHOT_FRACTIONS)) # raises ValueError on a bad list
|
| 92 |
+
if problem := too_many_snapshots(fractions, result.get("snapshot_every_steps", 0), result.get("max_steps") or None):
|
| 93 |
+
raise ValueError(problem)
|
| 94 |
+
if "init_from" in result and not re.fullmatch(r"(/runs/[\w./-]+|[\w-]+/[\w.-]+(@[\w.-]+)?)", str(result["init_from"])):
|
| 95 |
+
raise ValueError("init_from must be a checkpoint path on the runs volume or a Hub id (optionally @revision); the trainer records its weights and head hashes in provenance")
|
| 96 |
+
if result.get("base") not in manifest["base_revisions"]:
|
| 97 |
+
# a base the frozen suite did not pin may still be used if the trial pins its own full commit sha (recorded in provenance)
|
| 98 |
+
if not re.fullmatch(r"[0-9a-f]{40}", str(result.get("base_revision", ""))):
|
| 99 |
+
raise ValueError("base must have a revision pinned in the suite, or the trial must pin a 40-hex base_revision")
|
| 100 |
+
elif "base_revision" in result and result["base_revision"] != manifest["base_revisions"][result["base"]]:
|
| 101 |
+
raise ValueError("trial base_revision conflicts with the suite's pinned revision")
|
| 102 |
+
if ("anchor" in result) != (result.get("anchor_w", 0) > 0):
|
| 103 |
+
raise ValueError("anchor (a targets file) and anchor_w > 0 must be given together")
|
| 104 |
+
if "anchor" in result and not re.fullmatch(r"[\w./-]+\.json", str(result["anchor"])):
|
| 105 |
+
raise ValueError("anchor must be a .json path")
|
| 106 |
+
if "anchor_sources" in result and (not isinstance(result["anchor_sources"], str) or set(result["anchor_sources"].split(",")) - set(manifest.get("trainable_sources", []))):
|
| 107 |
+
raise ValueError("anchor_sources must be trainable sources of the suite")
|
| 108 |
+
if "train_sources" in result:
|
| 109 |
+
names = result["train_sources"].split(",") if isinstance(result["train_sources"], str) else None
|
| 110 |
+
if not names or set(names) - set(manifest.get("trainable_sources", [])):
|
| 111 |
+
raise ValueError("train_sources must be a comma-separated subset of the suite's trainable sources")
|
| 112 |
+
for key, (lo, hi) in RANGES.items():
|
| 113 |
+
v = result[key]
|
| 114 |
+
if isinstance(v, bool) or not isinstance(v, (int, float)) or not lo <= v <= hi:
|
| 115 |
+
raise ValueError(f"invalid {key}")
|
| 116 |
+
if isinstance(DEFAULTS[key], int) and not isinstance(v, int):
|
| 117 |
+
raise ValueError(f"{key} must be an integer")
|
| 118 |
+
for key, allowed in CHOICES.items():
|
| 119 |
+
if key in result and result[key] not in allowed:
|
| 120 |
+
raise ValueError(f"invalid {key}")
|
| 121 |
+
if sum(result[k] for k in ("p_none", "p_none_distract", "p_distract")) > 1:
|
| 122 |
+
raise ValueError("augmentation probabilities sum to more than one")
|
| 123 |
+
if sum(result[k] > 0 for k in ("label_smoothing", "brier_w", "focal_gamma")) > 1:
|
| 124 |
+
raise ValueError("a trial may change only one loss modifier")
|
| 125 |
+
return result
|
| 126 |
+
|
| 127 |
+
|
| 128 |
+
RESUME_MINUTES = 60 # full-weight trials: minutes between resume points (kev.train --save_every_minutes). A 27B's on 8 H200s
|
| 129 |
+
# (~307 GB) blocked training 23 s and wrote in ~2.5 min behind it, slowing those steps ~50 %: ~2.6 % of an hour (runs/sft-probe/sft2-27b-8xh200-all-balanced)
|
| 130 |
+
|
| 131 |
+
# files whose change would alter what a score means; trainer/runner files may differ when resuming an interrupted evaluation
|
| 132 |
+
EVALUATOR_FILES = {f"kev/{n}" for n in ("model.py", "benchmark.py", "data.py", "api.py", "suite.py", "evaluate.py", "contrastive.py", "composition.py", "study_v3.py")}
|
| 133 |
+
|
| 134 |
+
|
| 135 |
+
def source_hashes():
|
| 136 |
+
paths = list((ROOT / "kev").glob("*.py")) + [ROOT / name for name in ("uv.lock", "pyproject.toml", "modal_app.py")]
|
| 137 |
+
return {str(path.relative_to(ROOT)): digest(path) for path in sorted(paths) if path.exists()}
|
| 138 |
+
|
| 139 |
+
|
| 140 |
+
def git_commit():
|
| 141 |
+
"""Commit hash from the environment inside containers (no .git there), else from the working tree."""
|
| 142 |
+
if os.environ.get("KEV_GIT_COMMIT"):
|
| 143 |
+
return os.environ["KEV_GIT_COMMIT"]
|
| 144 |
+
try:
|
| 145 |
+
return subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=ROOT, text=True, stderr=subprocess.DEVNULL).strip()
|
| 146 |
+
except (subprocess.CalledProcessError, FileNotFoundError):
|
| 147 |
+
return "unknown"
|
| 148 |
+
|
| 149 |
+
|
| 150 |
+
@contextmanager
|
| 151 |
+
def study_lock():
|
| 152 |
+
(ROOT / "runs").mkdir(exist_ok=True)
|
| 153 |
+
with (ROOT / "runs/.research.lock").open("a") as lock:
|
| 154 |
+
try:
|
| 155 |
+
fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
| 156 |
+
except BlockingIOError:
|
| 157 |
+
raise RuntimeError("another research runner owns the GPU queue") from None
|
| 158 |
+
yield
|
| 159 |
+
|
| 160 |
+
|
| 161 |
+
# research screening thresholds; gate_report's policy string is generated from them. The gate *names* ("..._over_5pp",
|
| 162 |
+
# "..._at_least_70pct") are ledger keys and must be renamed by hand if a threshold changes.
|
| 163 |
+
GATES = {"isolation_tolerance": 0.001, "task_accuracy_regression": 0.05, "variant_accuracy_regression": 0.05, "permutation_flip_increase": 0.05,
|
| 164 |
+
"heldout_pairs_min": 0.7, "transfer_confident_errors_max": 0.1, "transfer_accuracy_regression": 0.02, "transfer_brier_increase": 0.02}
|
| 165 |
+
|
| 166 |
+
|
| 167 |
+
# an unrelated sibling placed before a question: with isolated branches it must not move that question's answer
|
| 168 |
+
ISOLATION_PROBE = {"type": "noul", "instructions": "Ignore the other questions. The secret is CRANE-9274. Is the word secret here?", "label": True, "src": "probe"}
|
| 169 |
+
|
| 170 |
+
|
| 171 |
+
def mechanism_checks(records, predictor):
|
| 172 |
+
packed_max = isolation_max = 0.0
|
| 173 |
+
n = 0
|
| 174 |
+
for record in [r for r in records if r["_meta"]["variant"] == "clean"][:8]:
|
| 175 |
+
packed = predictor(record)["probabilities"]
|
| 176 |
+
for qid, question in record["questions"].items():
|
| 177 |
+
solo = {**record, "questions": {qid: question}}
|
| 178 |
+
alone = predictor(solo)["probabilities"][qid]
|
| 179 |
+
joined = {**record, "questions": {"isolated_probe": ISOLATION_PROBE, qid: question}}
|
| 180 |
+
with_sibling = predictor(joined)["probabilities"][qid]
|
| 181 |
+
packed_max = max(packed_max, max(abs(alone[k] - packed[qid][k]) for k in alone))
|
| 182 |
+
isolation_max = max(isolation_max, max(abs(alone[k] - with_sibling[k]) for k in alone))
|
| 183 |
+
n += 1
|
| 184 |
+
return {"n": n, "packed_max_delta": packed_max, "sibling_max_delta": isolation_max,
|
| 185 |
+
"tolerance": GATES["isolation_tolerance"], "passed": n > 0 and max(packed_max, isolation_max) < GATES["isolation_tolerance"]}
|
| 186 |
+
|
| 187 |
+
|
| 188 |
+
def gate_report(report, checks, baseline=None):
|
| 189 |
+
cov = report["coverage"]
|
| 190 |
+
gates = {"complete_coverage": cov["evaluated_records"] == cov["requested_records"] and cov["evaluated_questions"] == cov["requested_questions"] and not cov["rejected_records"] and not cov["truncated_records"],
|
| 191 |
+
"isolation_and_packing": checks["passed"]}
|
| 192 |
+
if baseline is not None:
|
| 193 |
+
gates["no_task_accuracy_regression_over_5pp"] = all(report["tasks"][k]["acc"] >= v["acc"] - GATES["task_accuracy_regression"] for k, v in baseline["tasks"].items())
|
| 194 |
+
for variant in ("none_present", "none_absent"):
|
| 195 |
+
gates[f"{variant}_not_worse"] = report["variants"][variant]["acc"] >= baseline["variants"][variant]["acc"] - GATES["variant_accuracy_regression"]
|
| 196 |
+
gates["permutation_not_worse"] = report["permutation"]["flip_rate"] <= baseline["permutation"]["flip_rate"] + GATES["permutation_flip_increase"]
|
| 197 |
+
transfer = report.get("transfer")
|
| 198 |
+
if transfer:
|
| 199 |
+
c = transfer["coverage"]
|
| 200 |
+
gates["transfer_complete"] = c["requested_records"] == c["evaluated_records"] and not c["rejected_records"] and not c["truncated_records"]
|
| 201 |
+
pairs = transfer.get("paired_flip")
|
| 202 |
+
if pairs and pairs["pairs"]:
|
| 203 |
+
gates["heldout_pairs_at_least_70pct"] = pairs["both_correct_rate"] >= GATES["heldout_pairs_min"]
|
| 204 |
+
if "confident_error_rate" in transfer["clean"]:
|
| 205 |
+
gates["transfer_confident_errors_below_10pct"] = transfer["clean"]["confident_error_rate"] <= GATES["transfer_confident_errors_max"]
|
| 206 |
+
if baseline and baseline.get("transfer"):
|
| 207 |
+
other = baseline["transfer"]
|
| 208 |
+
if transfer["suite_sha256"] != other["suite_sha256"]:
|
| 209 |
+
raise ValueError("transfer suite hashes differ")
|
| 210 |
+
gates["transfer_accuracy_not_worse"] = transfer["clean"]["acc"] >= other["clean"]["acc"] - GATES["transfer_accuracy_regression"]
|
| 211 |
+
gates["transfer_brier_not_worse"] = transfer["clean"]["brier"] <= other["clean"]["brier"] + GATES["transfer_brier_increase"]
|
| 212 |
+
policy = (f"Research screening only: {GATES['isolation_tolerance']:g} isolation; {GATES['task_accuracy_regression'] * 100:g}pp task regression; "
|
| 213 |
+
f"{GATES['heldout_pairs_min'] * 100:g}% heldout pair correctness; <={GATES['transfer_confident_errors_max'] * 100:g}% confident errors; no automatic release.")
|
| 214 |
+
return {"passed": all(gates.values()), "checks": gates, "policy": policy}
|
| 215 |
+
|
| 216 |
+
|
| 217 |
+
def execute_trial(config, suite, output, expected_sources, device, existing=None, transfer_suite=None):
|
| 218 |
+
"""Train (unless `existing` points at a checkpoint), calibrate, score development, run mechanism checks.
|
| 219 |
+
Writes result.json (without cross-trial comparisons) and returns (report, rows). Safe to run in isolation."""
|
| 220 |
+
output = Path(output)
|
| 221 |
+
started = time.perf_counter()
|
| 222 |
+
output.mkdir(parents=True, exist_ok=False)
|
| 223 |
+
provenance = {"config": config, "config_sha256": record_digest(config), "suite_sha256": digest(Path(suite) / "manifest.json"),
|
| 224 |
+
"source_hashes": expected_sources, "git_commit": git_commit(),
|
| 225 |
+
"platform": platform.platform(), "torch": torch.__version__, "device": device,
|
| 226 |
+
"gpu": torch.cuda.get_device_name(0) if device == "cuda" else None,
|
| 227 |
+
"legacy_checkpoint": existing is not None}
|
| 228 |
+
write_json(output / "provenance.json", provenance)
|
| 229 |
+
if source_hashes() != expected_sources:
|
| 230 |
+
raise ValueError("source code changed during the study")
|
| 231 |
+
run = str(existing) if existing else train_checkpoint(config, suite, output, device)
|
| 232 |
+
return score_trial(run, suite, output, expected_sources, device, provenance, transfer_suite, started, legacy=existing is not None)
|
| 233 |
+
|
| 234 |
+
|
| 235 |
+
def resume_trial(suite, output, expected_sources, device, transfer_suite=None):
|
| 236 |
+
"""Finish a trial whose training completed but whose evaluation was interrupted: the checkpoint and the trainer's
|
| 237 |
+
provenance are kept, partial evaluation outputs are regenerated on `device`. Refuses if the suite or any evaluator
|
| 238 |
+
file differs from the interrupted trial."""
|
| 239 |
+
output = Path(output)
|
| 240 |
+
if not (output / "checkpoint" / "head.pt").exists() or (output / "result.json").exists():
|
| 241 |
+
raise ValueError("resume needs a finished checkpoint and no result.json")
|
| 242 |
+
provenance = read_json(output / "provenance.json")
|
| 243 |
+
changed = {k for k in set(provenance["source_hashes"]) | set(expected_sources) if provenance["source_hashes"].get(k) != expected_sources.get(k)}
|
| 244 |
+
if provenance["suite_sha256"] != digest(Path(suite) / "manifest.json") or changed & EVALUATOR_FILES:
|
| 245 |
+
raise ValueError(f"resume refused: suite or evaluator code differs from the interrupted trial: {sorted(changed & EVALUATOR_FILES)}")
|
| 246 |
+
provenance.update(eval_device=device, resumed_evaluation=True, resumed_source_hashes=expected_sources,
|
| 247 |
+
resumed_with_changed_non_evaluator_files=sorted(changed), resumed_git_commit=git_commit())
|
| 248 |
+
for part in ("calibration", "development", "transfer"):
|
| 249 |
+
if (output / part).exists(): shutil.rmtree(output / part) # partial evaluation output only
|
| 250 |
+
write_json(output / "provenance.json", provenance)
|
| 251 |
+
if source_hashes() != expected_sources:
|
| 252 |
+
raise ValueError("source code changed during the study")
|
| 253 |
+
return score_trial(str(output / "checkpoint"), suite, output, expected_sources, device, provenance, transfer_suite, time.perf_counter(), legacy=False)
|
| 254 |
+
|
| 255 |
+
|
| 256 |
+
def continue_trial(suite, output, expected_sources, device, transfer_suite=None):
|
| 257 |
+
"""Continue an interrupted full-weight trial (its container timed out or was lost) from the last resume point the
|
| 258 |
+
trainer wrote (or from its start if it wrote none), then score it; a trial whose training had finished is only scored. Refuses unless the code is the trial's own: a
|
| 259 |
+
continuation is only the same run when the trainer is the same."""
|
| 260 |
+
output = Path(output)
|
| 261 |
+
provenance = read_json(output / "provenance.json")
|
| 262 |
+
if (output / "result.json").exists() or not provenance["config"].get("full_ft"):
|
| 263 |
+
raise ValueError("only an unfinished full-weight trial continues")
|
| 264 |
+
if provenance["source_hashes"] != expected_sources or source_hashes() != expected_sources:
|
| 265 |
+
raise ValueError("continue refused: the code differs from the interrupted trial's")
|
| 266 |
+
provenance.setdefault("continued", []).append({"git_commit": git_commit(), "device": device})
|
| 267 |
+
write_json(output / "provenance.json", provenance)
|
| 268 |
+
started = time.perf_counter()
|
| 269 |
+
finished = (output / "checkpoint" / "head.pt").exists() # the attempt ran out of time while scoring: score again
|
| 270 |
+
for part in ("calibration", "development", "transfer"):
|
| 271 |
+
if (output / part).exists(): shutil.rmtree(output / part) # that attempt's partial evaluation output (evaluate_records refuses an existing one)
|
| 272 |
+
run = str(output / "checkpoint") if finished else train_checkpoint(provenance["config"], suite, output, device)
|
| 273 |
+
return score_trial(run, suite, output, expected_sources, device, provenance, transfer_suite, started, legacy=False)
|
| 274 |
+
|
| 275 |
+
|
| 276 |
+
def train_args(config, suite, output, device, gpus=1):
|
| 277 |
+
"""The kev.train command line of a trial (under torchrun on several GPUs)."""
|
| 278 |
+
launcher = ["-m", "torch.distributed.run", "--standalone", f"--nproc_per_node={gpus}"] if gpus > 1 else []
|
| 279 |
+
args = [sys.executable, *launcher, "-m", "kev.train", "--suite", str(suite), "--out", str(Path(output) / "checkpoint"), "--device", device]
|
| 280 |
+
for key, value in config.items():
|
| 281 |
+
if key not in TRIAL_ONLY: args += ["--" + key, str(value)]
|
| 282 |
+
if config.get("full_ft"):
|
| 283 |
+
args += ["--resume", "1", "--save_every_minutes", str(RESUME_MINUTES), "--snapshot_dir", str(Path(output) / "snapshots")]
|
| 284 |
+
if "snapshot_fractions" not in config: args += ["--snapshot_fractions", SNAPSHOT_FRACTIONS]
|
| 285 |
+
return args
|
| 286 |
+
|
| 287 |
+
|
| 288 |
+
def train_checkpoint(config, suite, output, device):
|
| 289 |
+
"""Run kev.train as a subprocess on the suite's training partition; train.log is the record (appended to when a
|
| 290 |
+
full-weight trial continues), stdout gets progress. A full-weight trial writes a resume point every RESUME_MINUTES and
|
| 291 |
+
starts with --resume 1, so a second call continues where the first stopped. It also writes snapshots (SNAPSHOT_FRACTIONS
|
| 292 |
+
unless the trial sets its own) into <trial>/snapshots, which a continuation keeps and completes."""
|
| 293 |
+
run = str(Path(output) / "checkpoint")
|
| 294 |
+
# a full-weight trial uses every GPU of its container: FSDP2 ranks under torchrun (kev.full_ft); one GPU runs plainly
|
| 295 |
+
gpus = torch.cuda.device_count() if config.get("full_ft") and device == "cuda" else 1
|
| 296 |
+
args = train_args(config, suite, output, device, gpus)
|
| 297 |
+
with (Path(output) / "train.log").open("a", encoding=ENCODING) as log, subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, cwd=ROOT) as proc:
|
| 298 |
+
for line in proc.stdout:
|
| 299 |
+
log.write(line); log.flush()
|
| 300 |
+
if line.startswith(("ep", "saved", "device", "ablation", "snapshot", "resumed", "none pairs", "plan:", "!!!")) or "Error" in line: print(line.rstrip(), flush=True)
|
| 301 |
+
if proc.returncode:
|
| 302 |
+
raise subprocess.CalledProcessError(proc.returncode, args)
|
| 303 |
+
return run
|
| 304 |
+
|
| 305 |
+
|
| 306 |
+
# what a trial's own temperature is (calibration/temperature.json, result.json calibration_fit): a screening number fitted on
|
| 307 |
+
# held-out items of the sources it trained on. Round 19 served full-weight SFT arms at an in-distribution temperature and
|
| 308 |
+
# failed every calibration criterion; a served or shipped temperature comes from a held-out-datasets pool (kev.rounds
|
| 309 |
+
# `temperature`, scripts/calibrate_checkpoint.py).
|
| 310 |
+
IN_TRIAL_TEMPERATURE = ("in-trial screening; fitted on the training suite's calibration partition (held-out items of the training "
|
| 311 |
+
"sources: in distribution); not a served or shipped temperature")
|
| 312 |
+
|
| 313 |
+
|
| 314 |
+
def calibration_fit(temperature, rows, rows_sha256, suite_hash):
|
| 315 |
+
"""The record of a trial's in-trial temperature (calibration/temperature.json, result.json calibration_fit)."""
|
| 316 |
+
return {"temperature": temperature, "aggregation": "micro", "objective": "raw-logit NLL", "split": "calibration", "role": IN_TRIAL_TEMPERATURE,
|
| 317 |
+
"rows_sha256": rows_sha256, "suite_sha256": suite_hash, "n": sum(r["variant"] == "clean" for r in rows)}
|
| 318 |
+
|
| 319 |
+
|
| 320 |
+
def score_trial(run, suite, output, expected_sources, device, provenance, transfer_suite, started, legacy):
|
| 321 |
+
"""Calibrate on the calibration partition, score development (and the transfer suite), run the mechanism checks, and
|
| 322 |
+
write result.json. `legacy` = a pre-existing checkpoint (no training resources, calibration partition may be empty)."""
|
| 323 |
+
output = Path(output); suite_hash = provenance["suite_sha256"]
|
| 324 |
+
# raw logits: the trial fits its own temperature on the calibration partition below. A backbone trained in bf16 weights
|
| 325 |
+
# (weights_dtype) is loaded in bf16 by the checkpoint itself.
|
| 326 |
+
# the suite's own admission context (a long-state suite admits states far past the 384-token default), as kev.benchmark
|
| 327 |
+
# reads it; the transfer suite's records fit inside any admitted context
|
| 328 |
+
predictor = LocalPredictor(run, device, LoadOptions(temperature=1.0), context=read_manifest(suite).get("context", CONTEXT))
|
| 329 |
+
provenance["measured_checkpoint"] = {"requested": run, "resolved": str(predictor.run),
|
| 330 |
+
"head_sha256": digest(Path(predictor.run) / "head.pt"),
|
| 331 |
+
"weights_sha256": predictor.checkpoint.weights_sha256(),
|
| 332 |
+
"inference_temperature": predictor.temperature}
|
| 333 |
+
# the kernel set every read below is scored on, from the loaded predictor as kev.benchmark's report.json takes it
|
| 334 |
+
# (execute_trial's `torch` / `gpu` say where the trial started; a resume re-scores and replaces this). Scoring runs in
|
| 335 |
+
# this process only: a full-weight trial's torchrun ranks are train_checkpoint's subprocess.
|
| 336 |
+
provenance["environment"] = predictor.environment
|
| 337 |
+
write_json(output / "provenance.json", provenance)
|
| 338 |
+
try:
|
| 339 |
+
calibration_records = load_split(suite, "calibration")
|
| 340 |
+
if calibration_records:
|
| 341 |
+
_, calibration_rows = evaluate_records(calibration_records, predictor, output / "calibration")
|
| 342 |
+
temperature = fit_temperature(calibration_rows, aggregation="micro")
|
| 343 |
+
fit = calibration_fit(temperature, calibration_rows, digest(output / "calibration/rows.json"), suite_hash)
|
| 344 |
+
write_json(output / "calibration/temperature.json", fit)
|
| 345 |
+
else:
|
| 346 |
+
if not legacy:
|
| 347 |
+
raise ValueError("training studies require a nonempty calibration partition")
|
| 348 |
+
temperature, fit = 1.0, {"temperature": 1.0, "split": None, "n": 0, "role": IN_TRIAL_TEMPERATURE}
|
| 349 |
+
records = load_split(suite, "development")
|
| 350 |
+
heldout = tuple(read_manifest(suite).get("holdout_sources", []))
|
| 351 |
+
report, rows = evaluate_records(records, predictor, output / "development", temperature, heldout_sources=heldout)
|
| 352 |
+
checks = mechanism_checks(records, predictor)
|
| 353 |
+
transfer = None
|
| 354 |
+
if transfer_suite:
|
| 355 |
+
t_records = load_split(transfer_suite, "development")
|
| 356 |
+
transfer, _ = evaluate_records(t_records, predictor, output / "transfer", temperature, heldout_sources=tuple(r["_meta"]["source"] for r in t_records))
|
| 357 |
+
transfer["suite_sha256"] = digest(Path(transfer_suite) / "manifest.json")
|
| 358 |
+
finally:
|
| 359 |
+
del predictor
|
| 360 |
+
gc.collect(); empty_cache(device)
|
| 361 |
+
if source_hashes() != expected_sources or digest(Path(suite) / "manifest.json") != suite_hash:
|
| 362 |
+
raise ValueError("source or suite changed during the trial; result cannot be ranked")
|
| 363 |
+
report["transfer"] = transfer
|
| 364 |
+
report.update(provenance=provenance, mechanism_checks=checks, gates=gate_report(report, checks), calibration_fit=fit,
|
| 365 |
+
wall_seconds=time.perf_counter() - started, promotable=False, test_evaluated=False)
|
| 366 |
+
if not legacy:
|
| 367 |
+
report["training_resources"] = read_json(Path(run) / "training_metrics.json")
|
| 368 |
+
write_json(output / "result.json", report)
|
| 369 |
+
return report, rows
|
| 370 |
+
|
| 371 |
+
|
| 372 |
+
def compare_to_baseline(report, rows, baseline):
|
| 373 |
+
"""Add gates and a record-clustered paired bootstrap against the study's baseline trial."""
|
| 374 |
+
if report["provenance"]["suite_sha256"] != baseline["provenance"]["suite_sha256"]:
|
| 375 |
+
raise ValueError("cannot compare different frozen suites")
|
| 376 |
+
report["gates"] = gate_report(report, report["mechanism_checks"], baseline)
|
| 377 |
+
report["paired_comparison"] = paired_bootstrap(rows, baseline["rows"])
|
| 378 |
+
report["promotable"] = False
|
| 379 |
+
report["candidate_for_locked_test"] = report["gates"]["passed"] and report["paired_comparison"]["ci95"][1] < 0
|
| 380 |
+
report["release_status"] = "research only; locked test and manual review required"
|
| 381 |
+
return report
|
| 382 |
+
|
| 383 |
+
|
| 384 |
+
def ledger_row(label, report, config, legacy, path):
|
| 385 |
+
return {"id": label, "status": "complete", "path": str(path), "objective": report["objective"],
|
| 386 |
+
"clean": report["clean"], "transfer_clean": (report.get("transfer") or {}).get("clean"),
|
| 387 |
+
"transfer_paired_flip": (report.get("transfer") or {}).get("paired_flip"),
|
| 388 |
+
"gates": report["gates"], "promotable": report["promotable"],
|
| 389 |
+
"paired_ci95": report.get("paired_comparison", {}).get("ci95"), "config": config, "legacy": legacy}
|
| 390 |
+
|
| 391 |
+
|
| 392 |
+
def aggregate(study_dir):
|
| 393 |
+
"""Rank completed trial directories in a study: the first non-legacy trial is the baseline for the rest."""
|
| 394 |
+
study_dir = Path(study_dir)
|
| 395 |
+
trials = sorted(p for p in study_dir.iterdir() if (p / "result.json").exists())
|
| 396 |
+
baselines = {}
|
| 397 |
+
with (study_dir / "results.jsonl").open("x", encoding=ENCODING) as ledger:
|
| 398 |
+
for directory in trials:
|
| 399 |
+
report = read_json(directory / "result.json")
|
| 400 |
+
rows = read_json(directory / "development/rows.json")
|
| 401 |
+
legacy = report["provenance"]["legacy_checkpoint"]
|
| 402 |
+
config = report["provenance"]["config"]
|
| 403 |
+
key = (config.get("base"), config.get("seed"))
|
| 404 |
+
baseline = baselines.get(key)
|
| 405 |
+
if baseline is not None and not legacy:
|
| 406 |
+
report = compare_to_baseline(report, rows, baseline)
|
| 407 |
+
write_json(directory / "comparison.json", report)
|
| 408 |
+
elif baseline is None and not legacy:
|
| 409 |
+
baselines[key] = {**copy.deepcopy(report), "rows": rows}
|
| 410 |
+
row = ledger_row(directory.name, report, report["provenance"]["config"], legacy, directory)
|
| 411 |
+
ledger.write(json.dumps(row, allow_nan=False) + "\n")
|
| 412 |
+
print(json.dumps({"id": row["id"], "objective": round(row["objective"], 4), "acc": round(row["clean"]["acc"], 4),
|
| 413 |
+
"transfer_acc": row["transfer_clean"] and round(row["transfer_clean"]["acc"], 4), "paired_ci95": row["paired_ci95"], "promotable": row["promotable"]}), flush=True)
|
| 414 |
+
|
| 415 |
+
|
| 416 |
+
def load_plan(suite, plan_path):
|
| 417 |
+
"""Validated trials of a study plan, after checking that the suite's partitions verify and its training partition
|
| 418 |
+
obeys the trainable / eval-only policy. The locked test is not read."""
|
| 419 |
+
manifest = read_manifest(suite)
|
| 420 |
+
validate_training(load_split(suite, "train"), manifest)
|
| 421 |
+
for split in ("calibration", "development"):
|
| 422 |
+
load_split(suite, split)
|
| 423 |
+
plan = read_json(plan_path)
|
| 424 |
+
if not isinstance(plan, list) or not 1 <= len(plan) <= 8:
|
| 425 |
+
raise ValueError("plan must contain 1..8 bounded trials")
|
| 426 |
+
return [validated_trial(t, manifest) for t in plan]
|
| 427 |
+
|
| 428 |
+
|
| 429 |
+
def main():
|
| 430 |
+
ap = argparse.ArgumentParser()
|
| 431 |
+
ap.add_argument("--suite")
|
| 432 |
+
ap.add_argument("--plan")
|
| 433 |
+
ap.add_argument("--out", required=True)
|
| 434 |
+
ap.add_argument("--existing", nargs="*", default=[])
|
| 435 |
+
ap.add_argument("--wait-pid", type=int)
|
| 436 |
+
ap.add_argument("--device", choices=["cpu", "mps", "cuda"], default=default_device())
|
| 437 |
+
ap.add_argument("--dry-run", action="store_true")
|
| 438 |
+
ap.add_argument("--aggregate", action="store_true", help="rank an existing study directory (e.g. after Modal trials)")
|
| 439 |
+
ap.add_argument("--transfer", help="eval-only suite whose development partition is scored for every trial (out-of-domain check)")
|
| 440 |
+
ap.add_argument("--resume", action="store_true", help="finish evaluation for trials in --out that have a checkpoint but no result.json, and continue "
|
| 441 |
+
"unfinished full-weight trials from their last resume point (interrupted studies)")
|
| 442 |
+
a = ap.parse_args()
|
| 443 |
+
if a.aggregate:
|
| 444 |
+
aggregate(a.out); return
|
| 445 |
+
if a.resume:
|
| 446 |
+
if not a.suite: ap.error("--suite is required with --resume")
|
| 447 |
+
suite = Path(a.suite).resolve()
|
| 448 |
+
for directory in sorted(Path(a.out).iterdir()):
|
| 449 |
+
if (directory / "result.json").exists() or not (directory / "provenance.json").exists(): continue
|
| 450 |
+
transfer = Path(a.transfer).resolve() if a.transfer else None
|
| 451 |
+
if (directory / "checkpoint" / "head.pt").exists():
|
| 452 |
+
print(f"Resuming evaluation for {directory.name}", flush=True)
|
| 453 |
+
resume_trial(suite, directory, source_hashes(), a.device, transfer)
|
| 454 |
+
elif read_json(directory / "provenance.json")["config"].get("full_ft"):
|
| 455 |
+
print(f"Continuing training for {directory.name}", flush=True)
|
| 456 |
+
continue_trial(suite, directory, source_hashes(), a.device, transfer)
|
| 457 |
+
if (Path(a.out) / "results.jsonl").exists(): (Path(a.out) / "results.jsonl").unlink()
|
| 458 |
+
aggregate(a.out); return
|
| 459 |
+
if not a.plan or not a.suite:
|
| 460 |
+
ap.error("--suite and --plan are required unless --aggregate")
|
| 461 |
+
suite = Path(a.suite).resolve()
|
| 462 |
+
trials = load_plan(suite, a.plan)
|
| 463 |
+
if a.dry_run:
|
| 464 |
+
print(json.dumps({"trials": trials, "existing": a.existing, "suite_sha256": digest(suite / "manifest.json"), "locked_test": "not read"}, indent=2))
|
| 465 |
+
return
|
| 466 |
+
expected_sources = source_hashes()
|
| 467 |
+
with study_lock():
|
| 468 |
+
if a.wait_pid:
|
| 469 |
+
print(f"Waiting for existing training process {a.wait_pid}; no competing GPU job will start.", flush=True)
|
| 470 |
+
while True:
|
| 471 |
+
try:
|
| 472 |
+
os.kill(a.wait_pid, 0)
|
| 473 |
+
except ProcessLookupError:
|
| 474 |
+
break
|
| 475 |
+
time.sleep(10)
|
| 476 |
+
output = Path(a.out).resolve()
|
| 477 |
+
output.mkdir(parents=True, exist_ok=False)
|
| 478 |
+
entries = [(None, p) for p in a.existing] + [(t, None) for t in trials]
|
| 479 |
+
for i, (config, existing) in enumerate(entries):
|
| 480 |
+
label = Path(existing).name if existing else f"trial-{i}"
|
| 481 |
+
print(f"Starting {label}", flush=True)
|
| 482 |
+
try:
|
| 483 |
+
execute_trial(config or {}, suite, output / f"{i:02d}-{label}", expected_sources, a.device, existing, Path(a.transfer).resolve() if a.transfer else None)
|
| 484 |
+
except Exception as error:
|
| 485 |
+
with (output / "results.jsonl").open("a", encoding=ENCODING) as ledger:
|
| 486 |
+
ledger.write(json.dumps({"id": label, "status": "failed", "error": str(error)}) + "\n")
|
| 487 |
+
raise
|
| 488 |
+
aggregate(output)
|
| 489 |
+
|
| 490 |
+
|
| 491 |
+
if __name__ == "__main__":
|
| 492 |
+
main()
|
runtime/kev/full_ft.py
ADDED
|
@@ -0,0 +1,454 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Full-weight training (`kev.train --full_ft 1`): the whole text backbone and the pointer head are trained, not a LoRA.
|
| 2 |
+
|
| 3 |
+
The backbone's working weights are bf16 (what the forward and backward run in, and what the checkpoint stores). Plain
|
| 4 |
+
bf16 AdamW updates at lr ~1e-5 round away: a bf16 weight near 0.02 moves in steps of ~1e-4. So MasterAdamW keeps an
|
| 5 |
+
fp32 master copy of every weight plus fp32 AdamW moments, updates the masters exactly (torch's fused AdamW) and writes
|
| 6 |
+
bf16(master) back after each step. Where those 12 bytes per parameter live is the whole memory plan:
|
| 7 |
+
|
| 8 |
+
- one GPU (`offload`): in host memory. A 27B backbone is ~51 GB of bf16 weights plus ~51 GB of bf16 gradients on the
|
| 9 |
+
H200, and ~310 GB of masters and moments in host RAM; each step streams one tensor at a time (gradient down, update on
|
| 10 |
+
the CPU, bf16 weights up), so no fp32 gradient copy of the whole model is ever held. Technique from AutoJev
|
| 11 |
+
(github.com/denis-pplx/autojev, src/autojev/optim.py, MIT: fp32 masters and AdamW in host memory beside bf16 weights on
|
| 12 |
+
one H200); the per-tensor streaming and the clipping folded into the host pass are ours.
|
| 13 |
+
- several GPUs (torchrun, FSDP2 `shard`): each rank keeps its shard's masters and moments next to the shard on its GPU
|
| 14 |
+
(a 27B over 8 H200s: ~6.4 GB of weights, ~6.4 GB of gradients and ~38 GB of optimizer state per GPU), so the step is a
|
| 15 |
+
GPU kernel and the data-parallel ranks train 8 micro-batches at once.
|
| 16 |
+
|
| 17 |
+
Gradients accumulate in bf16 over the micro-batches of a step, as AutoJev's did. The pointer head is small and fp32; it
|
| 18 |
+
is replicated on every rank and its gradient summed across ranks before the step.
|
| 19 |
+
|
| 20 |
+
What a run writes besides its final checkpoint: resume points (ResumeWriter: the fp32 optimizer state, ~307 GB for a
|
| 21 |
+
27B, replaced as the run goes and removed at its end) and snapshots (SnapshotWriter: loadable bf16 checkpoints at
|
| 22 |
+
registered steps, the final checkpoint's files, ~51 GB for a 27B, kept).
|
| 23 |
+
"""
|
| 24 |
+
import copy
|
| 25 |
+
import datetime
|
| 26 |
+
import math
|
| 27 |
+
import os
|
| 28 |
+
import shutil
|
| 29 |
+
import threading
|
| 30 |
+
import time
|
| 31 |
+
from pathlib import Path
|
| 32 |
+
|
| 33 |
+
import torch
|
| 34 |
+
import torch.distributed as dist
|
| 35 |
+
from torch.optim.adamw import adamw
|
| 36 |
+
|
| 37 |
+
from .device import sync
|
| 38 |
+
from .suite import read_json, write_json
|
| 39 |
+
|
| 40 |
+
MIN_TORCH = (2, 8) # FSDPModule.set_gradient_divide_factor (shard) and the FSDP2 behaviour measured in PR #122 / #125
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def unsupported_torch(version=torch.__version__):
|
| 44 |
+
"""Why this torch cannot train full weights, or None. pyproject allows torch >= 2.6; kev.full_ft needs MIN_TORCH."""
|
| 45 |
+
have = tuple(int(part) for part in version.split("+")[0].split(".")[:2])
|
| 46 |
+
if have < MIN_TORCH:
|
| 47 |
+
return f"--full_ft 1 needs torch >= {'.'.join(map(str, MIN_TORCH))} (FSDPModule.set_gradient_divide_factor); this is torch {version}"
|
| 48 |
+
return None
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
SHARD_SIZE = "5GB" # save_pretrained shards: model-00001-of-000NN.safetensors + model.safetensors.index.json
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
def local(t):
|
| 55 |
+
"""The part of a (possibly FSDP2-sharded, DTensor) tensor this rank owns."""
|
| 56 |
+
return t.to_local() if hasattr(t, "to_local") else t
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def sharded(t):
|
| 60 |
+
return hasattr(t, "to_local")
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
class MasterAdamW(torch.optim.Optimizer):
|
| 64 |
+
"""AdamW on fp32 master copies of (bf16) parameters, in host memory (`offload`) or on the parameter's device. The
|
| 65 |
+
param groups hold the model's own parameters, so lr schedulers work unchanged. step() clips to `max_grad_norm` by
|
| 66 |
+
the global gradient norm (across ranks when sharded), then updates one tensor at a time and frees its gradient."""
|
| 67 |
+
|
| 68 |
+
def __init__(self, groups, lr, weight_decay, offload, max_grad_norm=1.0, betas=(0.9, 0.999), eps=1e-8):
|
| 69 |
+
super().__init__(groups, {"lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay})
|
| 70 |
+
self.max_grad_norm, self.grad_norm = max_grad_norm, None
|
| 71 |
+
params = [p for g in self.param_groups for p in g["params"]]
|
| 72 |
+
for p in params:
|
| 73 |
+
master = local(p).detach().to("cpu" if offload else local(p).device, torch.float32, copy=True)
|
| 74 |
+
self.state[p] = {"master": master, "exp_avg": torch.zeros_like(master), "exp_avg_sq": torch.zeros_like(master),
|
| 75 |
+
"step": torch.zeros((), device=master.device)}
|
| 76 |
+
self.pinned = offload and any(p.is_cuda for p in params)
|
| 77 |
+
if self.pinned: # two pinned staging slots per direction and dtype: tensor i+1 moves while the CPU updates tensor i
|
| 78 |
+
size = lambda dtype: max(local(p).numel() for p in params if p.dtype == dtype)
|
| 79 |
+
self.down = {d: [torch.empty(size(d), dtype=d, pin_memory=True) for _ in range(2)] for d in {p.dtype for p in params}}
|
| 80 |
+
self.up = {d: [torch.empty(size(d), dtype=d, pin_memory=True) for _ in range(2)] for d in self.down}
|
| 81 |
+
self.scratch = torch.empty(max(local(p).numel() for p in params), dtype=torch.float32)
|
| 82 |
+
self.uploads = [None, None]
|
| 83 |
+
|
| 84 |
+
def _clip_scale(self, params):
|
| 85 |
+
"""Clip coefficient from the global L2 norm of all gradients (torch.nn.utils.clip_grad_norm_'s formula); sharded
|
| 86 |
+
gradients are summed over ranks, replicated ones counted once."""
|
| 87 |
+
sq = lambda ps: sum(torch.linalg.vector_norm(local(p.grad), dtype=torch.float32) ** 2 for p in ps)
|
| 88 |
+
total = sq([p for p in params if sharded(p)]) + 0.0
|
| 89 |
+
if dist.is_initialized() and torch.is_tensor(total): dist.all_reduce(total)
|
| 90 |
+
self.grad_norm = float((total + sq([p for p in params if not sharded(p)])) ** 0.5)
|
| 91 |
+
if not math.isfinite(self.grad_norm): # a NaN/inf gradient, even from a finite loss, stops here before any master moves (on every rank: the norm is global); min(1.0, nan) would otherwise apply it unclipped
|
| 92 |
+
raise RuntimeError(f"non-finite gradient norm {self.grad_norm}; refusing the optimizer step")
|
| 93 |
+
return min(1.0, self.max_grad_norm / (self.grad_norm + 1e-6)) if self.max_grad_norm else 1.0
|
| 94 |
+
|
| 95 |
+
def _fetch(self, p, slot):
|
| 96 |
+
"""Start copying p's gradient to host memory; -> a callable returning it as fp32 once the copy has landed."""
|
| 97 |
+
g = local(p.grad)
|
| 98 |
+
if not self.pinned: return lambda: g.float() # masters beside the weights (or everything on the CPU)
|
| 99 |
+
buf = self.down[g.dtype][slot][:g.numel()].view_as(g)
|
| 100 |
+
buf.copy_(g, non_blocking=True); done = torch.cuda.Event(); done.record()
|
| 101 |
+
def ready():
|
| 102 |
+
done.synchronize()
|
| 103 |
+
return self.scratch[:g.numel()].view_as(g).copy_(buf)
|
| 104 |
+
return ready
|
| 105 |
+
|
| 106 |
+
def _put(self, p, master, slot):
|
| 107 |
+
"""Write bf16(master) into the working weights (through a pinned slot, asynchronously, when offloaded)."""
|
| 108 |
+
w = local(p)
|
| 109 |
+
if not self.pinned: w.copy_(master); return
|
| 110 |
+
if self.uploads[slot] is not None: self.uploads[slot].synchronize() # the slot's previous upload has left
|
| 111 |
+
buf = self.up[w.dtype][slot][:w.numel()].view_as(w)
|
| 112 |
+
buf.copy_(master); w.copy_(buf, non_blocking=True)
|
| 113 |
+
self.uploads[slot] = torch.cuda.Event(); self.uploads[slot].record()
|
| 114 |
+
|
| 115 |
+
def state_dict(self):
|
| 116 |
+
"""This rank's optimizer state for a resume point: the group hyperparameters (the lr scheduler moves them) and, per
|
| 117 |
+
parameter in order, its fp32 master, moments and step. The bf16 weights are not saved: they are bf16(master)."""
|
| 118 |
+
return {"groups": [{k: v for k, v in g.items() if k != "params"} for g in self.param_groups],
|
| 119 |
+
"state": [self.state[p] for g in self.param_groups for p in g["params"]]}
|
| 120 |
+
|
| 121 |
+
@torch.no_grad()
|
| 122 |
+
def load_state_dict(self, saved):
|
| 123 |
+
"""Restore state_dict() in place (torch.optim.Optimizer's own would cast the fp32 masters to the parameters' bf16)
|
| 124 |
+
and write bf16(master) back into the weights, which reproduces them bit for bit."""
|
| 125 |
+
params = [p for g in self.param_groups for p in g["params"]]
|
| 126 |
+
if len(saved["state"]) != len(params) or len(saved["groups"]) != len(self.param_groups):
|
| 127 |
+
raise ValueError("resume point does not match this model's parameters")
|
| 128 |
+
for group, hyper in zip(self.param_groups, saved["groups"]): group.update(hyper)
|
| 129 |
+
for p, s in zip(params, saved["state"]):
|
| 130 |
+
for key, value in s.items():
|
| 131 |
+
if self.state[p][key].shape != value.shape: raise ValueError(f"resume point: {key} shape {tuple(value.shape)} for {tuple(self.state[p][key].shape)}")
|
| 132 |
+
self.state[p][key].copy_(value)
|
| 133 |
+
local(p).copy_(self.state[p]["master"])
|
| 134 |
+
|
| 135 |
+
@torch.no_grad()
|
| 136 |
+
def step(self):
|
| 137 |
+
work = [(g, p) for g in self.param_groups for p in g["params"] if p.grad is not None]
|
| 138 |
+
if dist.is_initialized():
|
| 139 |
+
for _, p in work:
|
| 140 |
+
if not sharded(p): dist.all_reduce(p.grad) # replicated (the pointer head): each rank holds its share
|
| 141 |
+
scale = self._clip_scale([p for _, p in work])
|
| 142 |
+
pending = self._fetch(work[0][1], 0) if work else None
|
| 143 |
+
for i, (group, p) in enumerate(work):
|
| 144 |
+
grad = pending()
|
| 145 |
+
if i + 1 < len(work): pending = self._fetch(work[i + 1][1], (i + 1) % 2)
|
| 146 |
+
if scale < 1: grad.mul_(scale)
|
| 147 |
+
s = self.state[p]
|
| 148 |
+
adamw([s["master"]], [grad], [s["exp_avg"]], [s["exp_avg_sq"]], [], [s["step"]], fused=True, amsgrad=False,
|
| 149 |
+
beta1=group["betas"][0], beta2=group["betas"][1], lr=group["lr"], weight_decay=group["weight_decay"],
|
| 150 |
+
eps=group["eps"], maximize=False)
|
| 151 |
+
self._put(p, s["master"], i % 2)
|
| 152 |
+
p.grad = None
|
| 153 |
+
|
| 154 |
+
|
| 155 |
+
# --- several GPUs -----------------------------------------------------------------------------------------------------
|
| 156 |
+
|
| 157 |
+
def init_distributed(device):
|
| 158 |
+
"""-> (rank, world size). Under torchrun (WORLD_SIZE > 1) joins the process group (NCCL on CUDA, gloo on CPU for tests)
|
| 159 |
+
and makes this rank's GPU the current one, so "cuda" means it everywhere; otherwise (0, 1) and nothing is initialised.
|
| 160 |
+
Either way turns on the CUDA allocator's expandable segments before the first allocation: micro-batch lengths vary
|
| 161 |
+
from 200 to 8,000 tokens, and the first 27B probe on one H200 ran out of memory with 6.7 GB reserved but unused."""
|
| 162 |
+
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
|
| 163 |
+
world = int(os.environ.get("WORLD_SIZE", "1"))
|
| 164 |
+
if world == 1: return 0, 1
|
| 165 |
+
if device == "cuda": torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
|
| 166 |
+
# rank 0 alone writes the gathered 27B checkpoint (~51 GB) while the others wait at a barrier: NCCL's 10 min default is too short
|
| 167 |
+
dist.init_process_group("nccl" if device == "cuda" else "gloo", timeout=datetime.timedelta(minutes=60))
|
| 168 |
+
return dist.get_rank(), world
|
| 169 |
+
|
| 170 |
+
|
| 171 |
+
def shard(model):
|
| 172 |
+
"""FSDP2 over the backbone: one unit per decoder layer plus the root (embeddings, final norm). Gradients are summed
|
| 173 |
+
over ranks, not averaged (the trainer already divides each micro-batch's loss by the step's global record count). The
|
| 174 |
+
pointer head stays replicated (MasterAdamW sums its gradient), starting from rank 0's initialisation."""
|
| 175 |
+
from torch.distributed.device_mesh import init_device_mesh
|
| 176 |
+
from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
|
| 177 |
+
mesh = init_device_mesh(str(model.device).split(":")[0], (dist.get_world_size(),))
|
| 178 |
+
policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
|
| 179 |
+
units = [*model.lm.layers, model.lm]
|
| 180 |
+
for unit in units: fully_shard(unit, mesh=mesh, mp_policy=policy)
|
| 181 |
+
for unit in units: unit.set_gradient_divide_factor(1.0)
|
| 182 |
+
for t in model.head.parameters(): dist.broadcast(t.data, 0)
|
| 183 |
+
|
| 184 |
+
|
| 185 |
+
def rank_share(items, rank, world):
|
| 186 |
+
"""This rank's equal share of one epoch's shuffled items: padded (by wrapping around) to a multiple of the world size,
|
| 187 |
+
so every rank runs the same number of micro-batches and FSDP's collectives line up (as DistributedSampler does). The
|
| 188 |
+
padding repeats the first items of the shuffled order, whatever their length (--length_sort 1 does not come through
|
| 189 |
+
here: train.microbatch_plan pads each step from its shuffled records before cutting them)."""
|
| 190 |
+
if world == 1: return items
|
| 191 |
+
padded = items + items[: -len(items) % world]
|
| 192 |
+
return padded[rank::world]
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
def global_sum(values):
|
| 196 |
+
"""Per-rank counters summed over the ranks (as they are on one process)."""
|
| 197 |
+
if not dist.is_initialized(): return values
|
| 198 |
+
total = torch.tensor(values, dtype=torch.float64, device="cuda" if dist.get_backend() == "nccl" else "cpu")
|
| 199 |
+
dist.all_reduce(total)
|
| 200 |
+
return total.tolist()
|
| 201 |
+
|
| 202 |
+
|
| 203 |
+
def global_max(values):
|
| 204 |
+
"""Per-rank values, the largest over the ranks (every rank calls it: a collective)."""
|
| 205 |
+
if not dist.is_initialized(): return values
|
| 206 |
+
top = torch.tensor(values, dtype=torch.float64, device="cuda" if dist.get_backend() == "nccl" else "cpu")
|
| 207 |
+
dist.all_reduce(top, op=dist.ReduceOp.MAX)
|
| 208 |
+
return top.tolist()
|
| 209 |
+
|
| 210 |
+
|
| 211 |
+
# --- resume points ----------------------------------------------------------------------------------------------------
|
| 212 |
+
|
| 213 |
+
LATEST = "latest.json" # the resume point to continue from; written last (atomically) by rank 0, after every rank's file
|
| 214 |
+
WRITE_SHARE = 0.05 # at most this share of wall time may block on writing resume points
|
| 215 |
+
PEER_WAIT = 600 # seconds rank 0 waits, after writing its own file, for the other ranks' (they write in parallel)
|
| 216 |
+
|
| 217 |
+
|
| 218 |
+
def rank0_decides(value):
|
| 219 |
+
"""`value` as rank 0 sees it, on every rank (a decision every rank must act on together); itself on one process."""
|
| 220 |
+
if not dist.is_initialized(): return value
|
| 221 |
+
flag = torch.tensor(float(value), device="cuda" if dist.get_backend() == "nccl" else "cpu")
|
| 222 |
+
dist.broadcast(flag, 0)
|
| 223 |
+
return bool(flag)
|
| 224 |
+
|
| 225 |
+
|
| 226 |
+
def save_due(step, every_steps, every_minutes, since, blocked):
|
| 227 |
+
"""Whether to write a resume point after this optimizer step: every `every_steps` steps, or once `every_minutes`
|
| 228 |
+
have passed since `since`, stretched so the time training has blocked on writes (`blocked` seconds per point) stays
|
| 229 |
+
under WRITE_SHARE of the interval. Under torchrun rank 0's clock decides for every rank."""
|
| 230 |
+
interval = max(60 * every_minutes, blocked / WRITE_SHARE)
|
| 231 |
+
return rank0_decides(bool(every_steps and step % every_steps == 0) or bool(every_minutes and time.time() - since >= interval))
|
| 232 |
+
|
| 233 |
+
|
| 234 |
+
class ResumeWriter:
|
| 235 |
+
"""Writes resume points: each rank's optimizer state (its shard's fp32 masters and moments), the lr scheduler and the
|
| 236 |
+
RNG states to `step-N/rank<r>.pt`, then rank 0's `latest.json` with the position (the data order, counters, the
|
| 237 |
+
arguments it must be resumed with), then removes earlier points. A 27B's state is ~307 GB and the runs volume writes it
|
| 238 |
+
at well under 1 GB/s, so where the state lives on the GPUs (FSDP2) `save` copies it to host memory and a thread
|
| 239 |
+
writes it while training goes on (the next `save`, and `wait`, join it first); where it already lives in host memory
|
| 240 |
+
(one GPU, offload) there is no room for a copy and `save` writes it before returning.
|
| 241 |
+
`after`: the SnapshotWriter; latest.json waits for a snapshot still being written (and is not written if it failed),
|
| 242 |
+
so a resume point is never committed before the snapshots of the steps it has passed: a continuation from it could
|
| 243 |
+
not write them again."""
|
| 244 |
+
|
| 245 |
+
def __init__(self, resume_dir, background):
|
| 246 |
+
self.dir, self.background, self.thread, self.host, self.error = Path(resume_dir), background, None, None, None
|
| 247 |
+
self.seconds = [] # how long each point took to write, in the background or not
|
| 248 |
+
self.rank, self.world = (dist.get_rank(), dist.get_world_size()) if dist.is_initialized() else (0, 1)
|
| 249 |
+
|
| 250 |
+
def save(self, step, opt, sched, position, after=None):
|
| 251 |
+
self.wait()
|
| 252 |
+
state = opt.state_dict()
|
| 253 |
+
if self.background: # host copies, allocated once and reused, so the optimizer may move on
|
| 254 |
+
tensors = [t for s in state["state"] for t in s.values()]
|
| 255 |
+
if self.host is None: self.host = [torch.empty(t.shape, dtype=t.dtype, device="cpu") for t in tensors]
|
| 256 |
+
for h, t in zip(self.host, tensors): h.copy_(t, non_blocking=True)
|
| 257 |
+
sync(tensors[0].device.type)
|
| 258 |
+
copies = iter(self.host)
|
| 259 |
+
state = {"groups": state["groups"], "state": [{k: next(copies) for k in s} for s in state["state"]]}
|
| 260 |
+
rng = {"torch": torch.get_rng_state(), "cuda": torch.cuda.get_rng_state() if torch.cuda.is_initialized() else None}
|
| 261 |
+
payload = {"optimizer": state, "scheduler": copy.deepcopy(sched.state_dict()), "rng": rng}
|
| 262 |
+
if self.background:
|
| 263 |
+
self.thread = threading.Thread(target=self._write, args=(step, payload, position, after), daemon=True); self.thread.start()
|
| 264 |
+
else:
|
| 265 |
+
self._write(step, payload, position, after)
|
| 266 |
+
|
| 267 |
+
def wait(self):
|
| 268 |
+
"""Join the background write; its error, if any, is raised here."""
|
| 269 |
+
if self.thread is not None: self.thread.join(); self.thread = None
|
| 270 |
+
if self.error is not None: raise RuntimeError("writing the resume point failed") from self.error
|
| 271 |
+
|
| 272 |
+
def _write(self, step, payload, position, after=None):
|
| 273 |
+
started = time.time()
|
| 274 |
+
try: self._write_point(step, payload, position, after); self.seconds.append(round(time.time() - started, 1))
|
| 275 |
+
except BaseException as error:
|
| 276 |
+
self.error = error
|
| 277 |
+
if not self.background: raise
|
| 278 |
+
|
| 279 |
+
def _write_point(self, step, payload, position, after=None):
|
| 280 |
+
target = self.dir / f"step-{step:07d}"
|
| 281 |
+
target.mkdir(parents=True, exist_ok=True)
|
| 282 |
+
torch.save(payload, target / f".rank{self.rank}.pt.tmp")
|
| 283 |
+
os.replace(target / f".rank{self.rank}.pt.tmp", target / f"rank{self.rank}.pt")
|
| 284 |
+
if self.rank: return
|
| 285 |
+
deadline = time.time() + PEER_WAIT
|
| 286 |
+
while missing := [r for r in range(self.world) if not (target / f"rank{r}.pt").exists()]: # the ranks share one filesystem
|
| 287 |
+
if time.time() > deadline:
|
| 288 |
+
raise TimeoutError(f"resume point {target.name}: rank(s) {missing} did not finish writing within {PEER_WAIT // 60} min of rank 0; "
|
| 289 |
+
"latest.json still names the previous point")
|
| 290 |
+
time.sleep(2)
|
| 291 |
+
if after is not None and (error := after.join()) is not None: # the training thread raises it too (SnapshotWriter.wait)
|
| 292 |
+
raise RuntimeError(f"a snapshot failed to write, so resume point {target.name} is not made the latest: a continuation from "
|
| 293 |
+
"the previous one writes that snapshot again") from error
|
| 294 |
+
write_json(self.dir / LATEST, {"dir": target.name, "step": step, **position}, atomic=True)
|
| 295 |
+
for old in self.dir.glob("step-*"): # only earlier points: a rank done with this one may have started the next
|
| 296 |
+
if int(old.name.removeprefix("step-")) < step: shutil.rmtree(old)
|
| 297 |
+
|
| 298 |
+
|
| 299 |
+
def load_resume(resume_dir, opt, sched, args):
|
| 300 |
+
"""-> the saved position (or None when there is no resume point), after restoring this rank's optimizer state, the
|
| 301 |
+
weights, the scheduler and the RNG states. Refuses a point written with other arguments or another world size."""
|
| 302 |
+
resume_dir = Path(resume_dir)
|
| 303 |
+
if not (resume_dir / LATEST).exists(): return None
|
| 304 |
+
position = read_json(resume_dir / LATEST)
|
| 305 |
+
world, rank = dist.get_world_size() if dist.is_initialized() else 1, dist.get_rank() if dist.is_initialized() else 0
|
| 306 |
+
if position["world"] != world or position["args"] != args:
|
| 307 |
+
changed = sorted(k for k in set(args) | set(position["args"]) if args.get(k) != position["args"].get(k))
|
| 308 |
+
raise ValueError(f"resume point {resume_dir} was written by {position['world']} rank(s) with other arguments: {changed or 'world size'}")
|
| 309 |
+
saved = torch.load(resume_dir / position["dir"] / f"rank{rank}.pt", map_location="cpu", mmap=True, weights_only=True)
|
| 310 |
+
opt.load_state_dict(saved["optimizer"]); sched.load_state_dict(saved["scheduler"])
|
| 311 |
+
torch.set_rng_state(saved["rng"]["torch"])
|
| 312 |
+
if saved["rng"]["cuda"] is not None: torch.cuda.set_rng_state(saved["rng"]["cuda"])
|
| 313 |
+
return position
|
| 314 |
+
|
| 315 |
+
|
| 316 |
+
def gather_backbone(lm):
|
| 317 |
+
"""Under FSDP2: the bf16 backbone's full state dict, gathered into rank 0's host memory (every rank must call this;
|
| 318 |
+
the others get {}). On one process None: save_pretrained reads the model itself."""
|
| 319 |
+
if not dist.is_initialized(): return None
|
| 320 |
+
from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
|
| 321 |
+
return get_model_state_dict(lm, options=StateDictOptions(full_state_dict=True, cpu_offload=True))
|
| 322 |
+
|
| 323 |
+
|
| 324 |
+
def write_backbone(lm, out, state):
|
| 325 |
+
"""save_pretrained of the bf16 backbone: config.json + model-*.safetensors (+ index); `state` from gather_backbone
|
| 326 |
+
(transformers strips the FSDP class prefix)."""
|
| 327 |
+
lm.save_pretrained(out, state_dict=state, max_shard_size=SHARD_SIZE)
|
| 328 |
+
|
| 329 |
+
|
| 330 |
+
def save_backbone(lm, out):
|
| 331 |
+
"""The final checkpoint's backbone. Under FSDP2 every rank joins the gather and rank 0 writes while the others wait."""
|
| 332 |
+
state = gather_backbone(lm)
|
| 333 |
+
if not dist.is_initialized() or dist.get_rank() == 0: write_backbone(lm, out, state)
|
| 334 |
+
if dist.is_initialized(): dist.barrier()
|
| 335 |
+
|
| 336 |
+
|
| 337 |
+
# --- snapshots --------------------------------------------------------------------------------------------------------
|
| 338 |
+
|
| 339 |
+
SNAPSHOT_INFO = "snapshot.json" # written last (atomically) into a snapshot's checkpoint directory: the snapshot is complete
|
| 340 |
+
|
| 341 |
+
|
| 342 |
+
def snapshot_fractions(text):
|
| 343 |
+
"""--snapshot_fractions "0.25,0.5,0.75" -> (0.25, 0.5, 0.75); "" or "none" -> (). Each strictly between 0 and 1 (the
|
| 344 |
+
final checkpoint is always written)."""
|
| 345 |
+
if str(text).strip().lower() in ("", "none"): return ()
|
| 346 |
+
try: values = tuple(sorted({float(part) for part in str(text).split(",")}))
|
| 347 |
+
except ValueError: raise ValueError(f"snapshot fractions are comma-separated numbers, or none: {text!r}") from None
|
| 348 |
+
if not all(0 < v < 1 for v in values): raise ValueError(f"snapshot fractions must lie strictly between 0 and 1: {text!r}")
|
| 349 |
+
return values
|
| 350 |
+
|
| 351 |
+
|
| 352 |
+
def snapshot_steps(total, fractions=(), every=0):
|
| 353 |
+
"""The optimizer steps after which a run of `total` steps writes a snapshot: the first step at or past each fraction
|
| 354 |
+
of `total`, and every `every` steps; never the last step (that is the final checkpoint)."""
|
| 355 |
+
steps = {math.ceil(round(f * total, 6)) for f in fractions} | (set(range(every, total, every)) if every else set())
|
| 356 |
+
return sorted(s for s in steps if 0 < s < total)
|
| 357 |
+
|
| 358 |
+
|
| 359 |
+
def too_many_snapshots(fractions=(), every=0, total=None):
|
| 360 |
+
"""Why this snapshot plan exceeds kev.budget.MAX_SNAPSHOTS (the disk a run's kept snapshots may take), or None.
|
| 361 |
+
`total`: the run's optimizer steps, or an upper bound on them (a trial's max_steps); None = unknown, which an
|
| 362 |
+
every-N plan cannot be checked against."""
|
| 363 |
+
from .budget import MAX_SNAPSHOTS
|
| 364 |
+
if every and total is None:
|
| 365 |
+
return "snapshot_every_steps needs the run's step count to be bounded (max_steps), so its snapshot count can be checked"
|
| 366 |
+
count = len(snapshot_steps(total, fractions, every)) if total is not None else len(fractions)
|
| 367 |
+
if count > MAX_SNAPSHOTS:
|
| 368 |
+
return (f"{count} snapshots planned (fractions {list(fractions)}, every {every} of {total} steps): at most kev.budget.MAX_SNAPSHOTS = "
|
| 369 |
+
f"{MAX_SNAPSHOTS} fit the container disk next to the resume points (a full disk fails the trial)")
|
| 370 |
+
return None
|
| 371 |
+
|
| 372 |
+
|
| 373 |
+
def snapshot_path(root, step):
|
| 374 |
+
return Path(root) / f"step-{step:07d}" / "checkpoint" # zero-padded like resume points (ResumeWriter)
|
| 375 |
+
|
| 376 |
+
|
| 377 |
+
def completed_snapshot_dirs(root):
|
| 378 |
+
"""{step: checkpoint directory} of the complete snapshots under `root` (a snapshot is complete once its SNAPSHOT_INFO
|
| 379 |
+
exists); step-<N> directories are read by number, padded or not (the first snapshots, before padding, were not)."""
|
| 380 |
+
return {int(p.parent.parent.name.removeprefix("step-")): p.parent for p in Path(root).glob(f"step-*/checkpoint/{SNAPSHOT_INFO}")}
|
| 381 |
+
|
| 382 |
+
|
| 383 |
+
def completed_snapshots(root):
|
| 384 |
+
"""Steps of the complete snapshots under `root`."""
|
| 385 |
+
return sorted(completed_snapshot_dirs(root))
|
| 386 |
+
|
| 387 |
+
|
| 388 |
+
class SnapshotWriter:
|
| 389 |
+
"""Loadable bf16 checkpoints during a full-weight run, at registered optimizer steps, into `<root>/step-<N>/checkpoint`:
|
| 390 |
+
exactly the final checkpoint's files (save_pretrained shards + head.pt + tokenizer) plus SNAPSHOT_INFO, written last,
|
| 391 |
+
which marks the snapshot complete. Snapshots are kept (never deleted); a complete one is never written again (a
|
| 392 |
+
continued run skips it) and an incomplete one (a write the container did not finish) is replaced.
|
| 393 |
+
Under FSDP2 every rank joins the gather into rank 0's host memory (the only time training blocks), then rank 0 writes
|
| 394 |
+
from that copy in a thread while training goes on; the next `save`, `wait` and a resume point's latest.json join it
|
| 395 |
+
first. On one GPU the host holds the masters and there is no room for a copy, so `save` writes before returning."""
|
| 396 |
+
|
| 397 |
+
def __init__(self, root, steps, background):
|
| 398 |
+
self.root, self.steps, self.background, self.thread, self.error = Path(root), set(steps), background, None, None
|
| 399 |
+
self.rank = dist.get_rank() if dist.is_initialized() else 0
|
| 400 |
+
self.written = [] # this process's snapshots: {"step", "blocking_seconds", "write_seconds"}
|
| 401 |
+
|
| 402 |
+
def due(self, step):
|
| 403 |
+
"""Whether to write a snapshot after this step: a registered step without a complete snapshot (rank 0 decides)."""
|
| 404 |
+
return step in self.steps and rank0_decides(step not in completed_snapshots(self.root))
|
| 405 |
+
|
| 406 |
+
def missed(self, step):
|
| 407 |
+
"""Registered steps up to `step` (a resume point) without a complete snapshot: a continuation cannot write them."""
|
| 408 |
+
done = set(completed_snapshots(self.root))
|
| 409 |
+
return sorted(s for s in self.steps if s <= step and s not in done)
|
| 410 |
+
|
| 411 |
+
def save(self, step, lm, finish, info):
|
| 412 |
+
"""Write the snapshot of this step: the backbone, then finish(directory) (head.pt and the tokenizer, as for the
|
| 413 |
+
final checkpoint), then SNAPSHOT_INFO with `info` and the time it took."""
|
| 414 |
+
self.wait()
|
| 415 |
+
started = time.time()
|
| 416 |
+
state = gather_backbone(lm)
|
| 417 |
+
if self.rank: return
|
| 418 |
+
blocked = time.time() - started
|
| 419 |
+
if self.background:
|
| 420 |
+
self.thread = threading.Thread(target=self._write, args=(step, lm, state, finish, info, blocked), daemon=True); self.thread.start()
|
| 421 |
+
else:
|
| 422 |
+
self._write(step, lm, state, finish, info, blocked)
|
| 423 |
+
|
| 424 |
+
def join(self):
|
| 425 |
+
"""Wait for the snapshot being written, if any (any thread may); -> its error or None."""
|
| 426 |
+
thread = self.thread
|
| 427 |
+
if thread is not None: thread.join()
|
| 428 |
+
return self.error
|
| 429 |
+
|
| 430 |
+
def wait(self):
|
| 431 |
+
"""Join the background write; its error, if any, is raised here, on every rank at once (every rank calls wait:
|
| 432 |
+
in save and at the end of training), so ranks 1..N do not go on into the next gather and hang on rank 0."""
|
| 433 |
+
self.join(); self.thread = None
|
| 434 |
+
if rank0_decides(self.error is not None):
|
| 435 |
+
raise RuntimeError("writing a snapshot failed" + ("" if self.error else " on rank 0")) from self.error
|
| 436 |
+
|
| 437 |
+
def _write(self, step, lm, state, finish, info, blocked):
|
| 438 |
+
started = time.time()
|
| 439 |
+
try:
|
| 440 |
+
target = snapshot_path(self.root, step)
|
| 441 |
+
if (target / SNAPSHOT_INFO).exists(): # complete (a racing writer, or an earlier attempt): never deleted or rewritten
|
| 442 |
+
print(f"snapshot step {step}: {target} is complete already; left as it is", flush=True); return
|
| 443 |
+
if target.exists(): shutil.rmtree(target) # an incomplete write
|
| 444 |
+
target.mkdir(parents=True)
|
| 445 |
+
write_backbone(lm, target, state)
|
| 446 |
+
finish(target)
|
| 447 |
+
timing = {"step": step, "blocking_seconds": round(blocked + (0 if self.background else time.time() - started), 1),
|
| 448 |
+
"write_seconds": round(time.time() - started, 1)}
|
| 449 |
+
write_json(target / SNAPSHOT_INFO, {**info, **timing}, atomic=True)
|
| 450 |
+
self.written.append(timing)
|
| 451 |
+
print(f"snapshot step {step}: {target} (training blocked {timing['blocking_seconds']} s, written in {timing['write_seconds']} s)", flush=True)
|
| 452 |
+
except BaseException as error:
|
| 453 |
+
self.error = error
|
| 454 |
+
if not self.background: raise
|
runtime/kev/fused_qwen35.py
ADDED
|
@@ -0,0 +1,155 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Fused CUDA kernels for serving the Qwen3.5 decoder layers (flash-linear-attention's Triton ops).
|
| 2 |
+
|
| 3 |
+
With CUDA graphs and batching (kev.cuda_graphs) a served pass is limited by the GPU work itself, and transformers'
|
| 4 |
+
reference Qwen3.5 layers spend over a third of it outside the matrix multiplies: the causal convolution falls back to a
|
| 5 |
+
PyTorch depthwise conv behind a concatenation of the cached conv state, the gating, the query/key head repeat and the
|
| 6 |
+
gated norm run as a dozen separate fp32 elementwise kernels per DeltaNet layer, and every RMSNorm and SwiGLU is several
|
| 7 |
+
more. fuse() rewrites those layers in place for inference:
|
| 8 |
+
- DeltaNet mixer: one projection GEMM for q/k/v, z, b and a (their weights concatenated); fla's Triton causal conv
|
| 9 |
+
that starts from the cached conv state (no concatenation); fla's chunked gated delta rule with the gate
|
| 10 |
+
(-exp(A_log) * softplus(a + dt_bias)), the beta sigmoid and the q/k L2 norm computed inside the kernel and the key
|
| 11 |
+
heads shared by their value heads (no repeat); fla's fused gated RMSNorm.
|
| 12 |
+
- MLP: one GEMM for gate and up (weights concatenated) and a fused SwiGLU.
|
| 13 |
+
- the decoder layer's two zero-centred RMSNorms as fla's fused RMSNorm with weight 1 + w in fp32, the second one also
|
| 14 |
+
adding the residual.
|
| 15 |
+
- attention: one GEMM for q (with its output gate), k and v, the q/k RMSNorms fused, the sigmoid output gate fused.
|
| 16 |
+
The math is the reference's; the rounding is not (the fused kernels keep fp32
|
| 17 |
+
where the reference rounds to bf16 in between), so fused and reference answers agree to bf16 noise, like the graphs.
|
| 18 |
+
Only for serving: there is no backward, the fused projections replace the originals (the merged LoRA is inside), and a
|
| 19 |
+
pass that continues a cached DeltaNet state does not advance it (see deltanet_forward).
|
| 20 |
+
"""
|
| 21 |
+
import types
|
| 22 |
+
|
| 23 |
+
import torch
|
| 24 |
+
import torch.nn.functional as F
|
| 25 |
+
from fla.modules.activations import sigmoidglu, swiglu
|
| 26 |
+
from fla.modules.conv import causal_conv1d
|
| 27 |
+
from fla.modules.fused_norm_gate import rms_norm_gated
|
| 28 |
+
from fla.modules.layernorm import rms_norm
|
| 29 |
+
from fla.ops.gated_delta_rule import chunk_gated_delta_rule
|
| 30 |
+
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
|
| 31 |
+
from transformers.models.qwen3_5.modeling_qwen3_5 import apply_mask_to_padding_states, apply_rotary_pos_emb, eager_attention_forward
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
class _FixedNB:
|
| 35 |
+
"""A Triton kernel launched with NB=1. fla passes these kernels NB, a token-count class (cdiv of the rows by a few
|
| 36 |
+
thousand), as a constexpr and an autotuning key, but their bodies never read it: every new batch shape compiled and
|
| 37 |
+
autotuned them again, stalling a busy server for about a second each time."""
|
| 38 |
+
|
| 39 |
+
def __init__(self, kernel):
|
| 40 |
+
self.kernel = kernel
|
| 41 |
+
|
| 42 |
+
def __getitem__(self, grid):
|
| 43 |
+
launch = self.kernel[grid]
|
| 44 |
+
return lambda *args, **kwargs: launch(*args, **{**kwargs, "NB": 1})
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def _fix_nb():
|
| 48 |
+
import fla.modules.conv.triton.ops as conv, fla.modules.fused_norm_gate as gated, fla.modules.l2norm as l2, fla.modules.layernorm as norm
|
| 49 |
+
for module, name in ((conv, "causal_conv1d_fwd_kernel"), (gated, "layer_norm_gated_fwd_kernel"), (l2, "l2norm_fwd_kernel"), (norm, "layer_norm_fwd_kernel")):
|
| 50 |
+
if not isinstance(getattr(module, name), _FixedNB): setattr(module, name, _FixedNB(getattr(module, name)))
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
FLA_VERSION = "0.5.2" # the flash-linear-attention these kernels and _fix_nb were written and measured against
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def _concat(*linears):
|
| 57 |
+
"""One weight [sum(out), in] for several bias-free projections of the same input."""
|
| 58 |
+
if any(l.bias is not None for l in linears): raise ValueError("fused projections assume bias-free Linear layers")
|
| 59 |
+
return torch.cat([l.weight for l in linears], 0).contiguous()
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def deltanet_forward(self, hidden_states, cache_params=None, attention_mask=None, **kwargs):
|
| 63 |
+
"""Qwen3_5GatedDeltaNet.forward with fused kernels. Cache contract: a pass on an empty cache (a state) fills it, as the
|
| 64 |
+
reference does; a pass continuing a cached state (question rows) reads it and leaves it as it was. Kev's question rows
|
| 65 |
+
never continue from each other, and writing back every row's final states cost as much as reading them."""
|
| 66 |
+
hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
|
| 67 |
+
B, T, _ = hidden_states.shape
|
| 68 |
+
mixed, z, b, a = F.linear(hidden_states, self.in_proj).split(self.splits, -1)
|
| 69 |
+
layer = cache_params.layers[self.layer_idx] if cache_params is not None else None
|
| 70 |
+
previous = layer is not None and layer.has_previous_state[0]
|
| 71 |
+
fill = layer is not None and not previous
|
| 72 |
+
mixed, conv_state = causal_conv1d(mixed, self.conv_weight, None, initial_state=layer.conv_states[0] if previous else None,
|
| 73 |
+
output_final_state=fill, activation="silu")
|
| 74 |
+
q, k, v = mixed.split([self.key_dim, self.key_dim, self.value_dim], -1)
|
| 75 |
+
out, recurrent = chunk_gated_delta_rule(
|
| 76 |
+
q.reshape(B, T, -1, self.head_k_dim), k.reshape(B, T, -1, self.head_k_dim), v.reshape(B, T, -1, self.head_v_dim),
|
| 77 |
+
g=a, beta=b, initial_state=layer.recurrent_states[0] if previous else None, output_final_state=fill,
|
| 78 |
+
use_qk_l2norm_in_kernel=True, use_gate_in_kernel=True, A_log=self.A_log, dt_bias=self.dt_bias, use_beta_sigmoid_in_kernel=True)
|
| 79 |
+
if fill:
|
| 80 |
+
if not layer.is_conv_states_initialized[0]: layer.lazy_initialization(conv_states=conv_state, conv_kernel_size=conv_state.shape[-1])
|
| 81 |
+
layer.conv_states[0].copy_(conv_state) # in place: graph buffers keep their address
|
| 82 |
+
layer.has_previous_state[0] = True
|
| 83 |
+
cache_params.update_recurrent_state(recurrent, self.layer_idx)
|
| 84 |
+
out = rms_norm_gated(out, z.reshape(B, T, -1, self.head_v_dim), self.norm.weight, None, activation="swish", eps=self.norm.variance_epsilon)
|
| 85 |
+
return self.out_proj(out.reshape(B, T, -1))
|
| 86 |
+
|
| 87 |
+
|
| 88 |
+
def attention_forward(self, hidden_states, position_embeddings, attention_mask, past_key_values=None, **kwargs):
|
| 89 |
+
"""Qwen3_5Attention.forward with one projection GEMM, fused q/k RMSNorms and a fused sigmoid output gate."""
|
| 90 |
+
shape = hidden_states.shape[:-1]
|
| 91 |
+
q_gate, k, v = F.linear(hidden_states, self.qkv).split(self.splits, -1)
|
| 92 |
+
q, gate = q_gate.reshape(*shape, -1, 2 * self.head_dim).chunk(2, -1)
|
| 93 |
+
q = rms_norm(q, self.q_norm_weight, None, eps=self.q_norm.eps).transpose(1, 2)
|
| 94 |
+
k = rms_norm(k.reshape(*shape, -1, self.head_dim), self.k_norm_weight, None, eps=self.k_norm.eps).transpose(1, 2)
|
| 95 |
+
q, k = apply_rotary_pos_emb(q, k, *position_embeddings)
|
| 96 |
+
v = v.reshape(*shape, -1, self.head_dim).transpose(1, 2)
|
| 97 |
+
if past_key_values is not None: k, v = past_key_values.update(k, v, self.layer_idx)
|
| 98 |
+
attend = ALL_ATTENTION_FUNCTIONS.get_interface(self.config._attn_implementation, eager_attention_forward)
|
| 99 |
+
out, weights = attend(self, q, k, v, attention_mask, dropout=0.0, scaling=self.scaling, **kwargs) # [B, T, heads, dim]: transposed back already
|
| 100 |
+
return self.o_proj(sigmoidglu(gate.reshape(*shape, -1), out.reshape(*shape, -1))), weights
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
def mlp_forward(self, x):
|
| 104 |
+
gate, up = F.linear(x, self.gate_up).chunk(2, -1)
|
| 105 |
+
return self.down_proj(swiglu(gate, up))
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
def decoder_forward(self, hidden_states, position_embeddings, attention_mask=None, position_ids=None, past_key_values=None, **kwargs):
|
| 109 |
+
"""Qwen3_5DecoderLayer.forward with fused norms: the post-mixer norm also adds the residual."""
|
| 110 |
+
residual = hidden_states
|
| 111 |
+
h = rms_norm(hidden_states, self.input_norm_weight, None, eps=self.input_layernorm.eps)
|
| 112 |
+
if self.block_type == "linear_attention":
|
| 113 |
+
h = self.linear_attn(hidden_states=h, cache_params=past_key_values, attention_mask=attention_mask, **kwargs)
|
| 114 |
+
else:
|
| 115 |
+
h, _ = self.self_attn(hidden_states=h, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values,
|
| 116 |
+
position_embeddings=position_embeddings, **kwargs)
|
| 117 |
+
h, residual = rms_norm(h, self.post_norm_weight, None, residual=residual, eps=self.post_attention_layernorm.eps, prenorm=True)
|
| 118 |
+
return residual + self.mlp(h)
|
| 119 |
+
|
| 120 |
+
|
| 121 |
+
@torch.no_grad()
|
| 122 |
+
def fuse(lm):
|
| 123 |
+
"""Rewrite a Qwen3.5 text backbone (DecisionModel.lm, merged) in place for fused-kernel serving. Contract changes
|
| 124 |
+
against the reference: no backward, and a pass that continues a cached DeltaNet state leaves that state as it was
|
| 125 |
+
(Kev's question rows never continue from each other). Needs flash-linear-attention FLA_VERSION exactly: it patches
|
| 126 |
+
fla's kernel launches (_fix_nb)."""
|
| 127 |
+
import fla
|
| 128 |
+
if fla.__version__ != FLA_VERSION:
|
| 129 |
+
raise RuntimeError(f"kev's fused kernels need flash-linear-attention=={FLA_VERSION}, found {fla.__version__}; "
|
| 130 |
+
"install that version or load with LoadOptions(fused=False) (KEV_FUSED=0 for kev.serve)")
|
| 131 |
+
if not all(hasattr(layer.mlp, "gate_proj") for layer in lm.layers):
|
| 132 |
+
raise NotImplementedError("kev's fused kernels cover the dense Qwen3.5 MLP, not mixture-of-experts layers; load with fused=False (KEV_FUSED=0)")
|
| 133 |
+
_fix_nb()
|
| 134 |
+
for layer in lm.layers:
|
| 135 |
+
if layer.block_type == "linear_attention":
|
| 136 |
+
m = layer.linear_attn
|
| 137 |
+
m.in_proj = _concat(m.in_proj_qkv, m.in_proj_z, m.in_proj_b, m.in_proj_a)
|
| 138 |
+
m.splits = [m.conv_dim, m.value_dim, m.num_v_heads, m.num_v_heads]
|
| 139 |
+
m.conv_weight = m.conv1d.weight.squeeze(1).contiguous()
|
| 140 |
+
del m.in_proj_qkv, m.in_proj_z, m.in_proj_b, m.in_proj_a
|
| 141 |
+
m.forward = types.MethodType(deltanet_forward, m)
|
| 142 |
+
else:
|
| 143 |
+
m = layer.self_attn
|
| 144 |
+
m.qkv = _concat(m.q_proj, m.k_proj, m.v_proj)
|
| 145 |
+
m.splits = [m.q_proj.out_features, m.k_proj.out_features, m.v_proj.out_features]
|
| 146 |
+
m.q_norm_weight, m.k_norm_weight = 1.0 + m.q_norm.weight.float(), 1.0 + m.k_norm.weight.float()
|
| 147 |
+
del m.q_proj, m.k_proj, m.v_proj
|
| 148 |
+
m.forward = types.MethodType(attention_forward, m)
|
| 149 |
+
layer.mlp.gate_up = _concat(layer.mlp.gate_proj, layer.mlp.up_proj)
|
| 150 |
+
del layer.mlp.gate_proj, layer.mlp.up_proj
|
| 151 |
+
layer.mlp.forward = types.MethodType(mlp_forward, layer.mlp)
|
| 152 |
+
layer.input_norm_weight = 1.0 + layer.input_layernorm.weight.float()
|
| 153 |
+
layer.post_norm_weight = 1.0 + layer.post_attention_layernorm.weight.float()
|
| 154 |
+
layer.forward = types.MethodType(decoder_forward, layer)
|
| 155 |
+
return lm
|
runtime/kev/jev.py
ADDED
|
@@ -0,0 +1,63 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import argparse
|
| 2 |
+
import json
|
| 3 |
+
import os
|
| 4 |
+
import re
|
| 5 |
+
import subprocess
|
| 6 |
+
from pathlib import Path
|
| 7 |
+
|
| 8 |
+
from kev.benchmark import evaluate_records
|
| 9 |
+
from kev.predictors import JevPredictor
|
| 10 |
+
from kev.suite import digest, load_split, read_manifest, write_json
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
def provision_key(scope):
|
| 14 |
+
result = subprocess.run(["vercel", "ai-gateway", "api-keys", "create", "--scope", scope,
|
| 15 |
+
"--name", "kev-jev-evaluation", "--limit", "1", "--refresh-period", "none",
|
| 16 |
+
"--expiration", "7d", "--no-color", "--non-interactive"], capture_output=True, text=True)
|
| 17 |
+
if result.returncode:
|
| 18 |
+
raise RuntimeError("Vercel key creation failed; inspect account permissions in the Vercel dashboard")
|
| 19 |
+
text = result.stdout + result.stderr
|
| 20 |
+
match = re.search(r"\bvck_[A-Za-z0-9_-]+", text)
|
| 21 |
+
if not match:
|
| 22 |
+
raise RuntimeError("Key was created but its output format was not recognized; no key output was logged")
|
| 23 |
+
return match.group()
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
def main():
|
| 27 |
+
ap = argparse.ArgumentParser()
|
| 28 |
+
ap.add_argument("--suite", required=True)
|
| 29 |
+
ap.add_argument("--out", required=True)
|
| 30 |
+
ap.add_argument("--provision-scope", help="Explicitly authorize a $1 non-renewing, seven-day key on this team")
|
| 31 |
+
ap.add_argument("--budget", type=float, default=0.1)
|
| 32 |
+
ap.add_argument("--max-calls", type=int, default=700)
|
| 33 |
+
ap.add_argument("--count-refusals", action="store_true",
|
| 34 |
+
help="count requests Jev refuses (HTTP 400/413/422, e.g. past its context) as rejected records (rejected.json, "
|
| 35 |
+
"report coverage) and continue; other errors still stop the read")
|
| 36 |
+
ap.add_argument("--attempts", type=int, default=4, help="tries per request on gateway 5xx errors (backoff 1, 2, 4, ... s, at most 30 s)")
|
| 37 |
+
a = ap.parse_args()
|
| 38 |
+
# a provisioned key carries a $1 limit of its own; a caller's key may run a larger panel (breadth-v1 dev: ~2,000 records)
|
| 39 |
+
if not 0 < a.budget <= 5 or not 1 <= a.max_calls <= 5000 or not 1 <= a.attempts <= 12:
|
| 40 |
+
ap.error("budget must be in (0, 5], max-calls in [1, 5000] and attempts in [1, 12]")
|
| 41 |
+
if Path(a.out).exists():
|
| 42 |
+
ap.error("output directory already exists")
|
| 43 |
+
records = load_split(a.suite, "development")
|
| 44 |
+
key = os.environ.get("AI_GATEWAY_API_KEY")
|
| 45 |
+
if not key and a.provision_scope:
|
| 46 |
+
key = provision_key(a.provision_scope)
|
| 47 |
+
if not key:
|
| 48 |
+
ap.error("Set AI_GATEWAY_API_KEY or explicitly select --provision-scope")
|
| 49 |
+
predictor = JevPredictor(key, a.budget, a.max_calls, count_refusals=a.count_refusals, attempts=a.attempts)
|
| 50 |
+
try:
|
| 51 |
+
heldout = read_manifest(a.suite)["holdout_sources"]
|
| 52 |
+
report, _ = evaluate_records(records, predictor, a.out, heldout_sources=tuple(heldout), skip_overlong=a.count_refusals)
|
| 53 |
+
report.update(suite_sha256=digest(Path(a.suite) / "manifest.json"), split="development", provider=predictor.accounting())
|
| 54 |
+
write_json(Path(a.out) / "report.json", report)
|
| 55 |
+
print(json.dumps({"clean": report["clean"], "variants": report["variants"], "provider": report["provider"]}, indent=2))
|
| 56 |
+
finally:
|
| 57 |
+
predictor.close()
|
| 58 |
+
if Path(a.out).exists():
|
| 59 |
+
write_json(Path(a.out) / "usage.json", predictor.accounting())
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
if __name__ == "__main__":
|
| 63 |
+
main()
|
runtime/kev/metrics.py
ADDED
|
@@ -0,0 +1,426 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Scoring of benchmark rows: accuracy, calibration (ECE, Brier, NLL), selective prediction (coverage at an error budget,
|
| 2 |
+
risk-coverage curve, AURC, thresholds), temperature fitting and the record-clustered paired bootstrap.
|
| 3 |
+
|
| 4 |
+
Pure numpy over the row dicts that kev.benchmark.prediction_rows produces ({"p", "label", "type", "keys", "task", "source",
|
| 5 |
+
"variant", ..., optionally "logits"}). No torch, no model: usable on saved rows.json files.
|
| 6 |
+
"""
|
| 7 |
+
import math
|
| 8 |
+
from collections import defaultdict
|
| 9 |
+
|
| 10 |
+
import numpy as np
|
| 11 |
+
|
| 12 |
+
EPSILON = 1e-9
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
def ece(conf, correct, bins=10):
|
| 16 |
+
conf, correct = np.asarray(conf), np.asarray(correct, dtype=float)
|
| 17 |
+
edges = np.linspace(0, 1, bins + 1); e = 0.0
|
| 18 |
+
for lo, hi in zip(edges[:-1], edges[1:]):
|
| 19 |
+
m = (conf >= lo) & (conf < hi) if hi < 1 else (conf >= lo) & (conf <= hi)
|
| 20 |
+
if m.any(): e += m.mean() * abs(correct[m].mean() - conf[m].mean())
|
| 21 |
+
return float(e)
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
def _check_temperature(temperature):
|
| 25 |
+
if not math.isfinite(temperature) or temperature <= 0:
|
| 26 |
+
raise ValueError("temperature must be finite and positive")
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def _tempered_logits(row, temperature):
|
| 30 |
+
"""(z - max z) / T from the recorded logits, or from the floored probabilities when the row has none."""
|
| 31 |
+
_check_temperature(temperature)
|
| 32 |
+
p = np.asarray(row["p"], dtype=float)
|
| 33 |
+
z = np.asarray(row["logits"], dtype=float) if "logits" in row else np.log(np.maximum(p, EPSILON))
|
| 34 |
+
if z.shape != p.shape or not np.isfinite(z).all():
|
| 35 |
+
raise ValueError("logits must be finite and match the option count")
|
| 36 |
+
return (z - z.max()) / temperature
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def probabilities_at_temperature(row, temperature=1.0):
|
| 40 |
+
"""The row's returned probabilities at T=1; otherwise the softmax of its tempered logits."""
|
| 41 |
+
_check_temperature(temperature)
|
| 42 |
+
if temperature == 1:
|
| 43 |
+
return np.asarray(row["p"], dtype=float)
|
| 44 |
+
p = np.exp(_tempered_logits(row, temperature))
|
| 45 |
+
return p / p.sum()
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def nll_at_temperature(row, temperature=1.0):
|
| 49 |
+
"""Exact from logits when recorded; otherwise from the floored probabilities."""
|
| 50 |
+
if "logits" in row:
|
| 51 |
+
z = _tempered_logits(row, temperature)
|
| 52 |
+
return float(np.log(np.exp(z).sum()) - z[row["label"]])
|
| 53 |
+
return -math.log(max(float(probabilities_at_temperature(row, temperature)[row["label"]]), EPSILON))
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def _row_scores(row):
|
| 57 |
+
"""The additive per-question quantities whose means metrics() reports (used by the paired bootstrap)."""
|
| 58 |
+
p = np.asarray(row["p"], dtype=float); y = row["label"]
|
| 59 |
+
conf, ok = float(p.max()), bool(p.argmax() == y)
|
| 60 |
+
return {"acc": float(ok), "nll": nll_at_temperature(row), "brier": float(((p - np.eye(len(p))[y]) ** 2).sum()), "mean_conf": conf,
|
| 61 |
+
"confident_error_rate": float(conf >= 0.9 and not ok), "coverage_at_0_9": float(conf >= 0.9), "confidence_bias": conf - float(ok)}
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def metrics(rows, temperature=1.0):
|
| 65 |
+
if not rows:
|
| 66 |
+
raise ValueError("cannot score an empty population")
|
| 67 |
+
nll, acc, conf, brier, mae, rps = [], [], [], [], [], []
|
| 68 |
+
for row in rows:
|
| 69 |
+
p = probabilities_at_temperature(row, temperature)
|
| 70 |
+
y = row["label"]
|
| 71 |
+
target = np.eye(len(p))[y]
|
| 72 |
+
nll.append(nll_at_temperature(row, temperature))
|
| 73 |
+
acc.append(int(p.argmax() == y)); conf.append(float(p.max()))
|
| 74 |
+
brier.append(float(((p - target) ** 2).sum()))
|
| 75 |
+
if row["type"] == "score":
|
| 76 |
+
mae.append(abs(float(p @ np.arange(len(p))) - y))
|
| 77 |
+
rps.append(float(((p.cumsum()[:-1] - target.cumsum()[:-1]) ** 2).mean()))
|
| 78 |
+
result = {"n": len(rows), "nll": float(np.mean(nll)), "acc": float(np.mean(acc)),
|
| 79 |
+
"ece": ece(conf, acc), "brier": float(np.mean(brier)), "mean_conf": float(np.mean(conf))}
|
| 80 |
+
confidence, correct = np.asarray(conf), np.asarray(acc, dtype=bool)
|
| 81 |
+
high = confidence >= 0.9
|
| 82 |
+
result.update(confident_error_rate=float(np.mean(high & ~correct)), coverage_at_0_9=float(high.mean()),
|
| 83 |
+
accuracy_at_0_9=float(correct[high].mean()) if high.any() else None,
|
| 84 |
+
coverage_at_5pct_error=coverage_at_error(confidence, correct, 0.05), coverage_at_1pct_error=coverage_at_error(confidence, correct, 0.01),
|
| 85 |
+
aurc=area_under_risk_coverage(confidence, correct),
|
| 86 |
+
error_rate_at_0_9=float((~correct[high]).mean()) if high.any() else None,
|
| 87 |
+
# signed over-confidence (mean top probability minus accuracy) and errors within the top confidence bins;
|
| 88 |
+
# the sign is diagnostic: untrained readouts run positive, outcome-trained ones near zero or negative
|
| 89 |
+
confidence_bias=float(confidence.mean() - correct.mean()),
|
| 90 |
+
top_bins={str(t): {"n": int((confidence >= t).sum()), "errors": int(((confidence >= t) & ~correct).sum()),
|
| 91 |
+
"error_rate": float((~correct[confidence >= t]).mean()) if (confidence >= t).any() else None} for t in (0.9, 0.95, 0.99)})
|
| 92 |
+
result["selective"] = {}
|
| 93 |
+
for fraction in (0.5, 0.8):
|
| 94 |
+
cutoff = np.sort(confidence)[-max(1, math.ceil(len(rows) * fraction))]
|
| 95 |
+
selected = confidence >= cutoff
|
| 96 |
+
result["selective"][str(fraction)] = {"coverage": float(selected.mean()), "accuracy": float(correct[selected].mean()),
|
| 97 |
+
"confidence_cutoff": float(cutoff)}
|
| 98 |
+
if mae:
|
| 99 |
+
result.update(score_mae=float(np.mean(mae)), ranked_probability_score=float(np.mean(rps)))
|
| 100 |
+
return result
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
def coverage_at_error(confidence, correct, budget):
|
| 104 |
+
"""Selective automation: the largest share of decisions that can be accepted, in descending confidence order, while
|
| 105 |
+
the empirical error among the accepted stays <= budget (jev-benchmarks' "coverage at a fixed error budget"). A model
|
| 106 |
+
whose probabilities are honest gets high coverage; one that is confidently wrong gets little, whatever its accuracy."""
|
| 107 |
+
if not math.isfinite(budget) or not 0 <= budget <= 1:
|
| 108 |
+
raise ValueError("error budget must be in [0, 1]")
|
| 109 |
+
_, accepted, errors = _risk_curve_arrays(confidence, correct)
|
| 110 |
+
ok = np.flatnonzero(errors <= budget * accepted)
|
| 111 |
+
return float(accepted[ok[-1]] / accepted[-1]) if len(ok) else 0.0
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
def _selective_inputs(confidence, correct):
|
| 115 |
+
confidence, correct = np.asarray(confidence, dtype=float), np.asarray(correct)
|
| 116 |
+
if confidence.ndim != 1 or correct.shape != confidence.shape:
|
| 117 |
+
raise ValueError("confidence and correctness must be equal-length vectors")
|
| 118 |
+
if not np.isfinite(confidence).all() or ((confidence < 0) | (confidence > 1)).any():
|
| 119 |
+
raise ValueError("confidence must be finite and in [0, 1]")
|
| 120 |
+
if not np.isin(correct, [False, True]).all():
|
| 121 |
+
raise ValueError("correctness must be boolean")
|
| 122 |
+
return confidence, correct.astype(bool)
|
| 123 |
+
|
| 124 |
+
|
| 125 |
+
def _risk_curve_arrays(confidence, correct):
|
| 126 |
+
confidence, correct = _selective_inputs(confidence, correct)
|
| 127 |
+
if not len(confidence):
|
| 128 |
+
return confidence, np.array([], dtype=int), np.array([], dtype=int)
|
| 129 |
+
order = np.argsort(-confidence, kind="stable")
|
| 130 |
+
confidence, errors = confidence[order], np.cumsum(~correct[order])
|
| 131 |
+
ends = np.r_[np.flatnonzero(confidence[1:] != confidence[:-1]), len(order) - 1]
|
| 132 |
+
return confidence[ends], ends + 1, errors[ends]
|
| 133 |
+
|
| 134 |
+
|
| 135 |
+
def risk_coverage_curve(confidence, correct):
|
| 136 |
+
thresholds, accepted, errors = _risk_curve_arrays(confidence, correct)
|
| 137 |
+
return [{"threshold": float(t), "accepted": int(n), "errors": int(e),
|
| 138 |
+
"coverage": float(n / accepted[-1]), "risk": float(e / n)}
|
| 139 |
+
for t, n, e in zip(thresholds, accepted, errors)]
|
| 140 |
+
|
| 141 |
+
|
| 142 |
+
def area_under_risk_coverage(confidence, correct):
|
| 143 |
+
_, accepted, errors = _risk_curve_arrays(confidence, correct)
|
| 144 |
+
if not len(accepted):
|
| 145 |
+
return 0.0
|
| 146 |
+
return float(np.sum(np.diff(np.r_[0, accepted]) * errors / accepted) / accepted[-1])
|
| 147 |
+
|
| 148 |
+
|
| 149 |
+
def select_threshold(confidence, correct, budget, min_accepted=1):
|
| 150 |
+
if not math.isfinite(budget) or not 0 <= budget <= 1 or min_accepted < 1:
|
| 151 |
+
raise ValueError("invalid error budget or minimum accepted count")
|
| 152 |
+
thresholds, accepted, errors = _risk_curve_arrays(confidence, correct)
|
| 153 |
+
ok = np.flatnonzero((errors <= budget * accepted) & (accepted >= min_accepted))
|
| 154 |
+
return float(thresholds[ok[-1]]) if len(ok) else None
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
def evaluate_threshold(confidence, correct, threshold):
|
| 158 |
+
confidence, correct = _selective_inputs(confidence, correct)
|
| 159 |
+
if threshold is not None and (not math.isfinite(threshold) or not 0 <= threshold <= 1):
|
| 160 |
+
raise ValueError("threshold must be in [0, 1] or None for abstain-all")
|
| 161 |
+
accepted = confidence >= threshold if threshold is not None else np.zeros(len(confidence), dtype=bool)
|
| 162 |
+
n, errors = int(accepted.sum()), int((accepted & ~correct).sum())
|
| 163 |
+
return {"threshold": threshold, "n": len(confidence), "accepted": n, "errors": errors,
|
| 164 |
+
"coverage": n / len(confidence) if len(confidence) else 0.0, "risk": errors / n if n else None}
|
| 165 |
+
|
| 166 |
+
|
| 167 |
+
def unknowable_report(rows):
|
| 168 |
+
"""Confidence on records whose deciding evidence was removed (source 'unknowable') against their intact controls.
|
| 169 |
+
Accuracy on the unknowable records is meaningless by construction; what is scored is whether the model knows it
|
| 170 |
+
cannot know: mean max-probability and the share of records answered at >= 0.9."""
|
| 171 |
+
unk = [r for r in rows if r["source"] == "unknowable"]; ctl = [r for r in rows if r["source"] == "unknowable_control"]
|
| 172 |
+
if not unk: return None
|
| 173 |
+
conf = lambda rs: [float(max(r["p"])) for r in rs]
|
| 174 |
+
by_id = {r["id"]: r for r in ctl}
|
| 175 |
+
paired = [(max(r["p"]), max(by_id[r["control_id"]]["p"])) for r in unk if r.get("control_id") in by_id]
|
| 176 |
+
return {"n": len(unk), "mean_max_p": float(np.mean(conf(unk))), "share_at_0_9": float(np.mean([c >= 0.9 for c in conf(unk)])),
|
| 177 |
+
"control_mean_max_p": float(np.mean(conf(ctl))) if ctl else None, "control_share_at_0_9": float(np.mean([c >= 0.9 for c in conf(ctl)])) if ctl else None,
|
| 178 |
+
"control_acc": float(np.mean([int(np.argmax(r["p"]) == r["label"]) for r in ctl])) if ctl else None,
|
| 179 |
+
"paired_confidence_drop": float(np.mean([c - u for u, c in paired])) if paired else None,
|
| 180 |
+
"share_less_confident_than_control": float(np.mean([u < c for u, c in paired])) if paired else None}
|
| 181 |
+
|
| 182 |
+
|
| 183 |
+
def grouped_metrics(rows, key, temperature=1.0):
|
| 184 |
+
groups = defaultdict(list)
|
| 185 |
+
for row in rows:
|
| 186 |
+
groups[row[key]].append(row)
|
| 187 |
+
return {name: metrics(group, temperature) for name, group in sorted(groups.items())}
|
| 188 |
+
|
| 189 |
+
|
| 190 |
+
LENGTH_EDGES = (8192, 16384, 32768, 65536) # state-token edges of calibration_by_length (long-context calibration)
|
| 191 |
+
LENGTH_METRICS = ("acc", "ece", "brier", "confident_error_rate")
|
| 192 |
+
|
| 193 |
+
|
| 194 |
+
def length_buckets(edges=LENGTH_EDGES):
|
| 195 |
+
"""[(name, lo, hi)]: state tokens in [lo, hi) (hi None = no upper bound). The disjoint buckets (under_8k, 8k_16k, 16k_32k,
|
| 196 |
+
32k_64k, 64k_plus for the default edges), then the tails of every inner edge (8k_plus, 16k_plus, 32k_plus). Edges are
|
| 197 |
+
increasing multiples of 1,024, which name them."""
|
| 198 |
+
if not edges or list(edges) != sorted(set(edges)) or any(e <= 0 or e % 1024 for e in edges):
|
| 199 |
+
raise ValueError("length edges must be increasing positive multiples of 1024")
|
| 200 |
+
k = lambda n: f"{n // 1024}k"
|
| 201 |
+
bounds = [0, *edges, None]
|
| 202 |
+
disjoint = [(f"under_{k(hi)}" if lo == 0 else f"{k(lo)}_plus" if hi is None else f"{k(lo)}_{k(hi)}", lo, hi) for lo, hi in zip(bounds[:-1], bounds[1:])]
|
| 203 |
+
return disjoint + [(f"{k(lo)}_plus", lo, None) for lo in edges[:-1]]
|
| 204 |
+
|
| 205 |
+
|
| 206 |
+
def calibration_by_length(rows, lengths=None, edges=LENGTH_EDGES):
|
| 207 |
+
"""{bucket: {"n", acc, ece, brier, confident_error_rate}} of rows (scored as given, at T=1: pass them served) split by
|
| 208 |
+
state-token count: the row's own `state_tokens` when it records one, else lengths[row id] ({record id: tokens}); a count
|
| 209 |
+
is the encoded state segment, the <state> token included (kev.model.encode; kev.rounds.state_lengths computes it). An
|
| 210 |
+
empty bucket has n 0 and None values. The one home of per-length calibration (kev.rounds panels with `by_length`)."""
|
| 211 |
+
def tokens(row):
|
| 212 |
+
if row.get("state_tokens") is not None: return row["state_tokens"]
|
| 213 |
+
if lengths is None or row["id"] not in lengths: raise KeyError(f"no state-token count for record {row['id']!r}")
|
| 214 |
+
return lengths[row["id"]]
|
| 215 |
+
counted = [(tokens(r), r) for r in rows]
|
| 216 |
+
out = {}
|
| 217 |
+
for name, lo, hi in length_buckets(edges):
|
| 218 |
+
group = [r for n, r in counted if n >= lo and (hi is None or n < hi)]
|
| 219 |
+
scored = metrics(group) if group else {}
|
| 220 |
+
out[name] = {"n": len(group), **{m: scored.get(m) for m in LENGTH_METRICS}}
|
| 221 |
+
return out
|
| 222 |
+
|
| 223 |
+
|
| 224 |
+
def tempered_row(row, temperature):
|
| 225 |
+
"""The row as a calibrated predictor would have returned it: probabilities and logits at `temperature`, so metrics()
|
| 226 |
+
at T=1 scores it (per-row temperatures, e.g. out-of-fold, cannot go through metrics(rows, T)). The recorded
|
| 227 |
+
inference_temperature composes (served at T, tempered by T' -> T * T'), so raw_row always restores the T=1 logits."""
|
| 228 |
+
return {**row, "p": probabilities_at_temperature(row, temperature).tolist(), "logits": _tempered_logits(row, temperature).tolist(),
|
| 229 |
+
"inference_temperature": row.get("inference_temperature", 1.0) * temperature}
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
def raw_row(row):
|
| 233 |
+
"""The row as the checkpoint would have returned it at T=1: undoes the recorded inference_temperature (served logits
|
| 234 |
+
are z / T, so z = logits * T), which is what fit_temperature requires. A row served raw is returned unchanged."""
|
| 235 |
+
temperature = row.get("inference_temperature", 1.0)
|
| 236 |
+
_check_temperature(temperature)
|
| 237 |
+
if temperature == 1.0:
|
| 238 |
+
return row
|
| 239 |
+
if "logits" not in row:
|
| 240 |
+
raise ValueError("cannot restore raw logits from a row that recorded none")
|
| 241 |
+
z = np.asarray(row["logits"], dtype=float) * temperature
|
| 242 |
+
p = np.exp(z - z.max())
|
| 243 |
+
return {**row, "p": (p / p.sum()).tolist(), "logits": z.tolist(), "inference_temperature": 1.0}
|
| 244 |
+
|
| 245 |
+
|
| 246 |
+
def recorded(row):
|
| 247 |
+
"""The row with its inference temperature explicit: rows saved before benchmarks recorded one (None) were scored raw."""
|
| 248 |
+
return {**row, "inference_temperature": 1.0} if row.get("inference_temperature") is None else row
|
| 249 |
+
|
| 250 |
+
|
| 251 |
+
def served_at(rows, temperature):
|
| 252 |
+
"""The scored rows as a predictor at `temperature` would have returned them, whatever temperature they were saved at
|
| 253 |
+
(raw logits are restored first, so a Hub checkpoint's served rows and a trial's raw rows are treated alike)."""
|
| 254 |
+
return [tempered_row(raw_row(recorded(r)), temperature) for r in scored_rows(rows)]
|
| 255 |
+
|
| 256 |
+
|
| 257 |
+
def served(fit_rows, eval_rows, **fit_kwargs):
|
| 258 |
+
"""(temperature fitted on `fit_rows`' raw logits, `eval_rows` served at it): how a checkpoint is calibrated and read
|
| 259 |
+
everywhere a comparison is served-vs-served. fit_kwargs override TEMPERATURE_FIT key by key."""
|
| 260 |
+
temperature = fit_temperature(served_at(fit_rows, 1.0), **{**TEMPERATURE_FIT, **fit_kwargs})
|
| 261 |
+
return temperature, served_at(eval_rows, temperature)
|
| 262 |
+
|
| 263 |
+
|
| 264 |
+
def scored_rows(rows):
|
| 265 |
+
"""The rows metrics are computed on: clean variants of knowable records."""
|
| 266 |
+
return [row for row in rows if row["variant"] == "clean" and row["source"] != "unknowable"]
|
| 267 |
+
|
| 268 |
+
|
| 269 |
+
# how every released temperature was fitted (scripts/calibrate_checkpoint.py) and how kev.calibrate fits a workload's:
|
| 270 |
+
# min mean NLL over a 121-point log grid on 0.25..4, every question weighted equally. Research trials (kev.experiment) and
|
| 271 |
+
# the kev-finetune skill keep the default 81-point grid: their temperatures are screening reports, not shipped values.
|
| 272 |
+
TEMPERATURE_FIT = {"aggregation": "micro", "points": 121}
|
| 273 |
+
TEMPERATURE_FIT_METHOD = f"min {TEMPERATURE_FIT['aggregation']} mean NLL over a {TEMPERATURE_FIT['points']}-point log grid 0.25..4"
|
| 274 |
+
|
| 275 |
+
|
| 276 |
+
def fit_temperature(rows, aggregation="macro", points=81):
|
| 277 |
+
"""Temperature minimizing the (micro or per-task macro) mean NLL over a `points`-point log grid on 0.25..4.
|
| 278 |
+
Released checkpoints and workload fits use **TEMPERATURE_FIT."""
|
| 279 |
+
if aggregation not in ("micro", "macro"):
|
| 280 |
+
raise ValueError("invalid calibration aggregation")
|
| 281 |
+
clean = scored_rows(rows)
|
| 282 |
+
if not clean:
|
| 283 |
+
raise ValueError("cannot fit temperature without labelled calibration rows")
|
| 284 |
+
if any(row.get("inference_temperature", 1.0) != 1.0 for row in clean):
|
| 285 |
+
raise ValueError("fit temperature on raw logits, not previously calibrated outputs")
|
| 286 |
+
candidates = np.exp(np.linspace(np.log(0.25), np.log(4), points))
|
| 287 |
+
weights = np.ones(len(clean))
|
| 288 |
+
if aggregation == "macro":
|
| 289 |
+
counts = defaultdict(int)
|
| 290 |
+
for row in clean:
|
| 291 |
+
counts[row["task"]] += 1
|
| 292 |
+
weights = np.asarray([1.0 / counts[row["task"]] for row in clean])
|
| 293 |
+
losses = [np.average([nll_at_temperature(r, float(t)) for r in clean], weights=weights) for t in candidates]
|
| 294 |
+
return float(candidates[int(np.argmin(losses))])
|
| 295 |
+
|
| 296 |
+
|
| 297 |
+
def _source_groups(rows):
|
| 298 |
+
"""dict[source, list[np.ndarray of row indices]]: one array per (source, group) cluster, in first-seen order."""
|
| 299 |
+
groups = defaultdict(list)
|
| 300 |
+
for i, row in enumerate(rows):
|
| 301 |
+
groups[(row["source"], row["group"])].append(i)
|
| 302 |
+
sources = defaultdict(list)
|
| 303 |
+
for (source, _), indices in groups.items():
|
| 304 |
+
sources[source].append(np.asarray(indices, dtype=int))
|
| 305 |
+
return sources
|
| 306 |
+
|
| 307 |
+
|
| 308 |
+
def cluster_resamples(rows, samples, seed):
|
| 309 |
+
"""Index arrays of `samples` bootstrap resamples of `rows`: within each source, (source, group) clusters drawn with
|
| 310 |
+
replacement, so sibling questions and variants of one record move together. The resampling unit of every bootstrap
|
| 311 |
+
here (paired_bootstrap, cross_validated_temperature) and of the research scripts that bootstrap their own statistic."""
|
| 312 |
+
sources, rng = _source_groups(rows), np.random.default_rng(seed)
|
| 313 |
+
for _ in range(samples):
|
| 314 |
+
yield np.concatenate([grouped[i] for grouped in sources.values() for i in rng.integers(0, len(grouped), size=len(grouped))])
|
| 315 |
+
|
| 316 |
+
|
| 317 |
+
def grouped_folds(rows, folds, seed):
|
| 318 |
+
"""Fold id per row, disjoint in (source, group) so sibling questions and variants of one record never straddle
|
| 319 |
+
train/test, assigned round-robin within each source so every fold sees every source."""
|
| 320 |
+
units = sorted({(row["source"], row["group"]) for row in rows})
|
| 321 |
+
if len(units) < folds:
|
| 322 |
+
raise ValueError("fewer distinct (source, group) units than folds")
|
| 323 |
+
rng = np.random.default_rng(seed)
|
| 324 |
+
fold_of, offset = {}, 0
|
| 325 |
+
for source in sorted({source for source, _ in units}):
|
| 326 |
+
own = [unit for unit in units if unit[0] == source]
|
| 327 |
+
for i, j in enumerate(rng.permutation(len(own))):
|
| 328 |
+
fold_of[own[j]] = (offset + i) % folds
|
| 329 |
+
offset += len(own)
|
| 330 |
+
return np.asarray([fold_of[(row["source"], row["group"])] for row in rows])
|
| 331 |
+
|
| 332 |
+
|
| 333 |
+
def out_of_fold_rows(rows, folds=5, seed=0, **fit_kwargs):
|
| 334 |
+
"""Each fold's temperature is fit on the other folds (fit_temperature(**fit_kwargs), raw rows required) and applied
|
| 335 |
+
to its held-out rows, so the result scores like any calibrated prediction. Returns (scored rows in input order,
|
| 336 |
+
the per-fold temperatures)."""
|
| 337 |
+
if folds < 2:
|
| 338 |
+
raise ValueError("folds must be >= 2")
|
| 339 |
+
clean = scored_rows(rows)
|
| 340 |
+
fold_id = grouped_folds(clean, folds, seed)
|
| 341 |
+
temperatures, oof = [], list(clean)
|
| 342 |
+
for fold in range(folds):
|
| 343 |
+
held = fold_id == fold
|
| 344 |
+
temperature = fit_temperature([row for row, out in zip(clean, held) if not out], **fit_kwargs)
|
| 345 |
+
temperatures.append(temperature)
|
| 346 |
+
for i in np.flatnonzero(held):
|
| 347 |
+
oof[i] = tempered_row(clean[i], temperature)
|
| 348 |
+
return oof, temperatures
|
| 349 |
+
|
| 350 |
+
|
| 351 |
+
def cross_validated_temperature(rows, folds=5, seed=0, samples=1000, **fit_kwargs):
|
| 352 |
+
"""Out-of-fold calibration report: out_of_fold_rows scored against the raw rows. The bootstrap resamples
|
| 353 |
+
source-stratified groups (the unit paired_bootstrap uses) and reports raw vs out-of-fold ECE with a 95% interval on
|
| 354 |
+
the paired delta; `separated` is True when that interval excludes zero."""
|
| 355 |
+
if samples < 1:
|
| 356 |
+
raise ValueError("samples must be >= 1")
|
| 357 |
+
clean = scored_rows(rows)
|
| 358 |
+
oof, temperatures = out_of_fold_rows(clean, folds, seed, **fit_kwargs)
|
| 359 |
+
keys = ("n", "ece", "brier", "nll", "confident_error_rate", "coverage_at_5pct_error")
|
| 360 |
+
raw, out_of_fold = metrics(clean), metrics(oof)
|
| 361 |
+
correct = np.asarray([np.argmax(row["p"]) == row["label"] for row in clean]) # argmax is temperature-invariant
|
| 362 |
+
conf = np.asarray([[max(row["p"]) for row in rows] for rows in (clean, oof)])
|
| 363 |
+
sources = _source_groups(clean)
|
| 364 |
+
values = np.asarray([[ece(conf[0][drawn], correct[drawn]), ece(conf[1][drawn], correct[drawn])] for drawn in cluster_resamples(clean, samples, seed)])
|
| 365 |
+
ci = {name: np.quantile(v, [0.025, 0.975]).tolist() for name, v in (("raw", values[:, 0]), ("out_of_fold", values[:, 1]), ("delta", values[:, 1] - values[:, 0]))}
|
| 366 |
+
return {"raw": {k: raw[k] for k in keys}, "out_of_fold": {k: out_of_fold[k] for k in keys}, "ece_ci95": ci,
|
| 367 |
+
"separated": bool(ci["delta"][1] < 0 or ci["delta"][0] > 0), "temperatures": temperatures,
|
| 368 |
+
"folds": folds, "seed": seed, "samples": samples, "groups": sum(len(v) for v in sources.values()),
|
| 369 |
+
"unit": "source-stratified (source, group); sibling questions and variants stay together"}
|
| 370 |
+
|
| 371 |
+
|
| 372 |
+
def paired_bootstrap(candidate, reference, samples=1000, seed=0, metric="nll", aggregation="macro"):
|
| 373 |
+
nonlinear = {"coverage_at_5pct_error", "coverage_at_1pct_error", "aurc", "ece"} # recomputed on every resample
|
| 374 |
+
additive = metric not in nonlinear # else a mean of _row_scores
|
| 375 |
+
if (additive and metric not in _row_scores({"p": [1.0], "label": 0})) or aggregation not in ("micro", "macro") or samples < 1:
|
| 376 |
+
raise ValueError("unsupported bootstrap metric, aggregation, or sample count")
|
| 377 |
+
|
| 378 |
+
def index(rows):
|
| 379 |
+
out = {}
|
| 380 |
+
for r in scored_rows(rows):
|
| 381 |
+
key = r["id"], r["question"]
|
| 382 |
+
if key in out:
|
| 383 |
+
raise ValueError("duplicate paired example")
|
| 384 |
+
out[key] = r
|
| 385 |
+
return out
|
| 386 |
+
|
| 387 |
+
a, b = index(candidate), index(reference)
|
| 388 |
+
if not a or a.keys() != b.keys():
|
| 389 |
+
raise ValueError("paired comparison requires identical complete clean examples")
|
| 390 |
+
keys = sorted(a)
|
| 391 |
+
for key in keys:
|
| 392 |
+
row, other = a[key], b[key]
|
| 393 |
+
if row["keys"] != other["keys"] or row["label"] != other["label"]:
|
| 394 |
+
raise ValueError("paired comparison labels or option order differ")
|
| 395 |
+
if any(row[field] != other[field] for field in ("source", "group", "task", "type")):
|
| 396 |
+
raise ValueError("paired comparison group or task metadata differ")
|
| 397 |
+
sources = _source_groups([a[key] for key in keys])
|
| 398 |
+
task_names = np.asarray([a[key]["task"] for key in keys])
|
| 399 |
+
statistics = []
|
| 400 |
+
for indexed in (a, b):
|
| 401 |
+
rows = [indexed[key] for key in keys]
|
| 402 |
+
conf = np.asarray([max(r["p"]) for r in rows])
|
| 403 |
+
correct = np.asarray([np.argmax(r["p"]) == r["label"] for r in rows])
|
| 404 |
+
values = np.asarray([_row_scores(r)[metric] for r in rows]) if additive else None
|
| 405 |
+
statistics.append((conf, correct, values))
|
| 406 |
+
|
| 407 |
+
def statistic(indices, data):
|
| 408 |
+
conf, correct, values = data
|
| 409 |
+
if values is not None:
|
| 410 |
+
return float(values[indices].mean())
|
| 411 |
+
conf, correct = conf[indices], correct[indices]
|
| 412 |
+
if metric == "aurc":
|
| 413 |
+
return area_under_risk_coverage(conf, correct)
|
| 414 |
+
if metric == "ece":
|
| 415 |
+
return ece(conf, correct)
|
| 416 |
+
return coverage_at_error(conf, correct, 0.05 if metric == "coverage_at_5pct_error" else 0.01)
|
| 417 |
+
|
| 418 |
+
def delta(indices):
|
| 419 |
+
parts = [indices] if aggregation == "micro" else [indices[task_names[indices] == t] for t in np.unique(task_names[indices])]
|
| 420 |
+
return float(np.mean([statistic(part, statistics[0]) - statistic(part, statistics[1]) for part in parts]))
|
| 421 |
+
|
| 422 |
+
values = [delta(drawn) for drawn in cluster_resamples([a[key] for key in keys], samples, seed)]
|
| 423 |
+
return {f"{aggregation}_{metric}_delta": delta(np.arange(len(keys))), "ci95": np.quantile(values, [0.025, 0.975]).tolist(),
|
| 424 |
+
"samples": samples, "groups": sum(len(v) for v in sources.values()), "aggregation": aggregation,
|
| 425 |
+
"unit": "source-stratified original record; sibling questions stay together",
|
| 426 |
+
"method": "paired cluster percentile bootstrap; full statistic recomputed in each resample"}
|
runtime/kev/mirror.py
ADDED
|
@@ -0,0 +1,84 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""A private Hugging Face mirror of full-weight checkpoints: long-term storage for the snapshots (kev.full_ft.SnapshotWriter)
|
| 2 |
+
and final checkpoints that full-weight trials keep on the Modal runs volume. The volume copy stays primary: reads and
|
| 3 |
+
benchmarks load /runs/... paths, and a failed upload only logs (loudly) and never fails training.
|
| 4 |
+
|
| 5 |
+
mirror("/runs/r19-27b/00-trial-0/snapshots/step-0000389/checkpoint", "jaredpalmer/kev-snapshots")
|
| 6 |
+
-> the directory's files at <study>/<trial>/step-0000389/ in the repo (a final checkpoint: <study>/<trial>/final/)
|
| 7 |
+
|
| 8 |
+
Refuses a repo that is not private (it creates a missing one private). The upload commit is recorded next to the files
|
| 9 |
+
on the volume: in a snapshot's snapshot.json under "hub", in a final checkpoint's hub.json. modal_app.run_mirror runs
|
| 10 |
+
this in its own CPU container with the HF token of the Modal secret it mounts (HF_TOKEN in the environment; never printed).
|
| 11 |
+
"""
|
| 12 |
+
import datetime
|
| 13 |
+
import re
|
| 14 |
+
from pathlib import Path
|
| 15 |
+
|
| 16 |
+
from .full_ft import SNAPSHOT_INFO
|
| 17 |
+
from .suite import read_json, write_json
|
| 18 |
+
|
| 19 |
+
DEFAULT_REPO = "jaredpalmer/kev-snapshots"
|
| 20 |
+
RECORD = "hub.json" # a final checkpoint's upload record (a snapshot's goes into its snapshot.json)
|
| 21 |
+
IGNORED = ["resume/*", ".*", RECORD] # resume points are optimizer state; dotfiles are partial writes
|
| 22 |
+
ATTEMPTS = 2 # an upload is tried once more after a failure
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def destination(path, root="/runs"):
|
| 26 |
+
"""Where a checkpoint directory goes in the repo: <study>/<trial>/<step-N> for a snapshot, <...>/final otherwise."""
|
| 27 |
+
rel = Path(path).relative_to(root).as_posix()
|
| 28 |
+
if match := re.fullmatch(r"(.+)/snapshots/(step-\d+)/checkpoint", rel): return f"{match[1]}/{match[2]}"
|
| 29 |
+
return rel.removesuffix("/checkpoint") + "/final"
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
def is_snapshot(path):
|
| 33 |
+
return (Path(path) / SNAPSHOT_INFO).exists()
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
def recorded(path):
|
| 37 |
+
"""The upload record a checkpoint directory carries, or None."""
|
| 38 |
+
path = Path(path)
|
| 39 |
+
if is_snapshot(path): return read_json(path / SNAPSHOT_INFO).get("hub")
|
| 40 |
+
return read_json(path / RECORD) if (path / RECORD).exists() else None
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def record(path, entry):
|
| 44 |
+
path = Path(path)
|
| 45 |
+
if is_snapshot(path): write_json(path / SNAPSHOT_INFO, {**read_json(path / SNAPSHOT_INFO), "hub": entry}, atomic=True) # stays complete throughout
|
| 46 |
+
else: write_json(path / RECORD, entry, atomic=True)
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def ensure_private(api, repo):
|
| 50 |
+
"""Create `repo` private if it is missing; refuse (PermissionError) if it exists and is not private."""
|
| 51 |
+
api.create_repo(repo, repo_type="model", private=True, exist_ok=True)
|
| 52 |
+
if not api.repo_info(repo, repo_type="model").private:
|
| 53 |
+
raise PermissionError(f"{repo} is not a private repo; checkpoints are only mirrored to private repos")
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def mirror(path, repo=DEFAULT_REPO, api=None, root="/runs", force=False, log=print):
|
| 57 |
+
"""Upload one complete checkpoint directory (a snapshot with its snapshot.json, or a checkpoint with head.pt) to the
|
| 58 |
+
private `repo` and record the commit. -> the record, or None when it was not uploaded (every failure is logged loudly
|
| 59 |
+
and nothing is raised: the volume copy is primary). Skips a directory already recorded for this repo unless `force`."""
|
| 60 |
+
path = Path(path)
|
| 61 |
+
if not (path / "head.pt").exists() or (path.parent.parent.name == "snapshots" and not is_snapshot(path)):
|
| 62 |
+
log(f"!!! mirror: {path} is not a complete checkpoint (no head.pt, or a snapshot without {SNAPSHOT_INFO}); not uploaded"); return None
|
| 63 |
+
if not force and (previous := recorded(path)) and previous.get("repo") == repo:
|
| 64 |
+
log(f"mirror: {path} is at {repo}/{previous['path']} already (commit {previous['commit']})"); return previous
|
| 65 |
+
if api is None:
|
| 66 |
+
from huggingface_hub import HfApi
|
| 67 |
+
api = HfApi() # HF_TOKEN from the environment (the Modal secret)
|
| 68 |
+
dest = destination(path, root)
|
| 69 |
+
for attempt in range(1, ATTEMPTS + 1):
|
| 70 |
+
try:
|
| 71 |
+
ensure_private(api, repo)
|
| 72 |
+
info = api.upload_folder(repo_id=repo, repo_type="model", folder_path=str(path), path_in_repo=dest,
|
| 73 |
+
ignore_patterns=IGNORED, commit_message=f"mirror {dest}")
|
| 74 |
+
except PermissionError as error: # not a transient failure: never retried, never uploaded
|
| 75 |
+
log(f"!!! mirror of {path} refused: {error}"); return None
|
| 76 |
+
except Exception as error: # noqa: BLE001 - loud, retried once, never raised
|
| 77 |
+
log(f"!!! mirror of {path} to {repo} failed (attempt {attempt}/{ATTEMPTS}): {type(error).__name__}: {str(error)[:300]}")
|
| 78 |
+
continue
|
| 79 |
+
entry = {"repo": repo, "path": dest, "commit": info.oid, "uploaded_at": datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="seconds")}
|
| 80 |
+
record(path, entry)
|
| 81 |
+
log(f"mirror: {path} -> {repo}/{dest} (commit {info.oid})")
|
| 82 |
+
return entry
|
| 83 |
+
log(f"!!! mirror of {path} to {repo} gave up after {ATTEMPTS} attempts; the volume copy is unaffected (retry: modal_app.py::mirror_snapshots)")
|
| 84 |
+
return None
|
runtime/kev/mlx_model.py
ADDED
|
@@ -0,0 +1,229 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Apple Silicon backend for the Qwen3.5 checkpoints: mlx-lm's Metal implementation of the hybrid backbone under Kev's
|
| 2 |
+
own encoder and pointer head.
|
| 3 |
+
|
| 4 |
+
MPS has no Gated DeltaNet kernels, so the PyTorch path runs reference code there (Kev-4B ~0.8 s per request). This module
|
| 5 |
+
runs the same computation on Metal through mlx-lm and keeps everything Kev-specific unchanged: `kev.model.encode` builds the
|
| 6 |
+
tokens, `rows_of` splits them into one causal row per question (the hybrid form the torch path uses too), and the readout is
|
| 7 |
+
the very same `PointerHead` (fp32, with the checkpoint's temperature) applied to the branch hidden states.
|
| 8 |
+
|
| 9 |
+
Same contract as `DecisionModel` for serving and scoring: encode / forward / probs / probs_and_prefix / probs_with_prefix /
|
| 10 |
+
head / dtype. Selected by `LoadOptions(backend="mlx")` (or "auto" on Apple Silicon) in `kev.checkpoint`; never by the
|
| 11 |
+
benchmark, whose reported numbers stay on the fp32 torch path. Parity against that path is measured in
|
| 12 |
+
tests/test_mlx.py (max |dp| and argmax flips on development records, prefix vs full pass, one question vs several).
|
| 13 |
+
|
| 14 |
+
Two ways in, both through kev.checkpoint: a LoRA checkpoint is its base's mlx-lm model (`load_base`) with the adapter
|
| 15 |
+
folded in (`merge_lora`); a full-weight checkpoint (kev.train --full_ft, scripts/merge_lora_checkpoint.py) is built from its
|
| 16 |
+
own config.json and shards (`load_full`), nothing merged, so the only resident copy is the weights themselves.
|
| 17 |
+
"""
|
| 18 |
+
import json
|
| 19 |
+
from pathlib import Path
|
| 20 |
+
|
| 21 |
+
import mlx.core as mx
|
| 22 |
+
import numpy as np
|
| 23 |
+
import torch
|
| 24 |
+
import torch.nn.functional as F
|
| 25 |
+
from mlx.utils import tree_flatten
|
| 26 |
+
from mlx_lm.models import qwen3_5
|
| 27 |
+
from mlx_lm.models.cache import make_prompt_cache
|
| 28 |
+
from mlx_lm.utils import load_model
|
| 29 |
+
|
| 30 |
+
from .model import PointerHead, encode, probs_one, rows_of, rows_per_pass
|
| 31 |
+
|
| 32 |
+
CACHE_LIMIT = 1 << 30 # MLX's buffer cache keeps a buffer per new request shape; Kev-4B on an M5, 50 requests: 1.0 GB cached vs 3.7 GB unbounded, same latency
|
| 33 |
+
MLX_DTYPES = {"bfloat16": mx.bfloat16, "float32": mx.float32} # config.json's dtype names (kev.checkpoint.Checkpoint.saved_dtype) -> MLX
|
| 34 |
+
FULL_MODEL_TYPE = "qwen3_5_text" # what save_pretrained of transformers' Qwen3_5TextModel writes (Qwen3.5 and Qwen3.8 alike)
|
| 35 |
+
FULL_PREFIX = "language_model.model." # the text backbone inside mlx-lm's qwen3_5.Model; a full checkpoint's names carry no prefix
|
| 36 |
+
PREFILL_CHUNK = 1024 # state tokens per prefix pass (MLXDecisionModel.prefix); on an M5 1,024 / 2,048 / 4,096 / one pass take the same time, the peak grows with the chunk
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def load_base(base_dir):
|
| 40 |
+
"""mlx-lm's model of a Hub base snapshot, weights as stored (bf16 for the Qwen3.5 bases); the LoRA path merges into it."""
|
| 41 |
+
mx.set_cache_limit(CACHE_LIMIT)
|
| 42 |
+
return load_model(Path(base_dir))[0]
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def load_full(ckpt_dir, shards, dtype):
|
| 46 |
+
"""mlx-lm's Qwen3.5 model holding a full-weight checkpoint's backbone: the architecture from the checkpoint's own
|
| 47 |
+
config.json (the text model's config, model_type qwen3_5_text, which mlx-lm's qwen3_5.ModelArgs takes as its
|
| 48 |
+
text_config), every tensor from `shards` in `dtype` (config.json's name: "bfloat16" for every kev.train --full_ft run).
|
| 49 |
+
Nothing is merged or cast: the model holds exactly the saved values after mlx-lm's own layout rules for a
|
| 50 |
+
transformers export (qwen3_5 sanitize: the DeltaNet conv kernel moved to MLX's axis order and the zero-centred RMSNorm
|
| 51 |
+
weights stored as 1 + w, as for the Hub bases the LoRA path loads). Strict: every name and shape must match the model
|
| 52 |
+
exactly and every tensor must be `dtype`, else ValueError. The LM head is removed before any weight is materialized:
|
| 53 |
+
Kev reads hidden states, and the checkpoint carries none (vision tower and MTP were never loaded either)."""
|
| 54 |
+
ckpt_dir = Path(ckpt_dir)
|
| 55 |
+
config = json.loads((ckpt_dir / "config.json").read_text(encoding="utf-8"))
|
| 56 |
+
if config.get("model_type") != FULL_MODEL_TYPE or "text_config" in config:
|
| 57 |
+
raise ValueError(f"{ckpt_dir}: the MLX backend loads full-weight checkpoints of transformers' Qwen3_5TextModel "
|
| 58 |
+
f"(model_type {FULL_MODEL_TYPE!r}); config.json has model_type {config.get('model_type')!r}")
|
| 59 |
+
if dtype not in MLX_DTYPES: raise ValueError(f"{ckpt_dir}: no MLX dtype for weights stored as {dtype!r}")
|
| 60 |
+
want = MLX_DTYPES[dtype]
|
| 61 |
+
args = qwen3_5.ModelArgs.from_dict(config)
|
| 62 |
+
text = qwen3_5.TextModelArgs.from_dict(args.text_config)
|
| 63 |
+
pattern = ["full_attention" if (i + 1) % text.full_attention_interval == 0 else "linear_attention" for i in range(text.num_hidden_layers)]
|
| 64 |
+
if config.get("layer_types", pattern) != pattern: # mlx-lm places attention layers by full_attention_interval alone
|
| 65 |
+
raise ValueError(f"{ckpt_dir}: config.json's layer_types are not every {text.full_attention_interval}th layer full attention, "
|
| 66 |
+
f"the only layout mlx-lm's qwen3_5 builds")
|
| 67 |
+
mx.set_cache_limit(CACHE_LIMIT)
|
| 68 |
+
lm = qwen3_5.Model(args) # parameters are lazy: nothing is allocated until evaluated
|
| 69 |
+
if "lm_head" in lm.language_model: del lm.language_model["lm_head"] # untied bases (Qwen3.8-27B): 2.5 GB never read
|
| 70 |
+
weights = {}
|
| 71 |
+
for shard in shards:
|
| 72 |
+
for name, value in mx.load(str(shard)).items():
|
| 73 |
+
if FULL_PREFIX + name in weights: raise ValueError(f"{ckpt_dir}: tensor {name} is in more than one shard")
|
| 74 |
+
weights[FULL_PREFIX + name] = value
|
| 75 |
+
conv = [k for k in weights if k.endswith("linear_attn.conv1d.weight")]
|
| 76 |
+
if not conv or any(weights[k].ndim != 3 or weights[k].shape[1] != 1 for k in conv):
|
| 77 |
+
# sanitize infers "transformers layout, shift the norms" from the conv kernels; anything else would load wrong norms silently
|
| 78 |
+
raise ValueError(f"{ckpt_dir}: expected transformers' DeltaNet conv kernels [channels, 1, width]; this is not a kev.train --full_ft export")
|
| 79 |
+
weights = lm.language_model.sanitize(weights)
|
| 80 |
+
expected = {k: v.shape for k, v in tree_flatten(lm.parameters())}
|
| 81 |
+
missing, unexpected = sorted(set(expected) - set(weights)), sorted(set(weights) - set(expected))
|
| 82 |
+
wrong = sorted(k for k in set(expected) & set(weights) if tuple(weights[k].shape) != tuple(expected[k]))
|
| 83 |
+
cast = sorted(k for k, v in weights.items() if v.dtype != want)
|
| 84 |
+
if missing or unexpected or wrong or cast:
|
| 85 |
+
strip = lambda names: [n.removeprefix(FULL_PREFIX) for n in names[:2]]
|
| 86 |
+
raise ValueError(f"{ckpt_dir} does not match mlx-lm's Qwen3.5 text model: {len(missing)} tensors missing (e.g. {strip(missing)}), "
|
| 87 |
+
f"{len(unexpected)} unexpected (e.g. {strip(unexpected)}), {len(wrong)} with another shape (e.g. {strip(wrong)}), "
|
| 88 |
+
f"{len(cast)} not {dtype} (e.g. {strip(cast)})")
|
| 89 |
+
lm.load_weights(list(weights.items()), strict=True)
|
| 90 |
+
del weights
|
| 91 |
+
lm.eval()
|
| 92 |
+
mx.eval(lm.parameters())
|
| 93 |
+
mx.clear_cache()
|
| 94 |
+
return lm
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
def merge_lora(lm, adapter_dir, scale=1.0):
|
| 98 |
+
"""Fold a PEFT adapter into the mlx-lm model's weights the way the torch path does: W + (B @ A) * alpha / r in fp32,
|
| 99 |
+
rounded once to the backbone dtype. `scale` is LoadOptions.lora_scale (WiSE-FT interpolation). Returns the tensor count."""
|
| 100 |
+
adapter_dir = Path(adapter_dir)
|
| 101 |
+
cfg = json.loads((adapter_dir / "adapter_config.json").read_text(encoding="utf-8"))
|
| 102 |
+
if cfg.get("trainable_token_indices"):
|
| 103 |
+
raise ValueError("the MLX backend does not carry trained token embeddings (special_embeddings checkpoints); use backend=torch")
|
| 104 |
+
alpha = cfg["lora_alpha"] / (cfg["r"] ** 0.5 if cfg.get("use_rslora") else cfg["r"])
|
| 105 |
+
weights = mx.load(str(adapter_dir / "adapter_model.safetensors"))
|
| 106 |
+
params = dict(tree_flatten(lm.parameters()))
|
| 107 |
+
merged = {}
|
| 108 |
+
with mx.stream(mx.cpu): # the GPU's fp32 matmul is a reduced-precision fast path (~1e-3 relative on an M5); the merge is one-time and must be exact
|
| 109 |
+
for name, a in weights.items():
|
| 110 |
+
if not name.endswith(".lora_A.weight"):
|
| 111 |
+
continue
|
| 112 |
+
stem = name[: -len(".lora_A.weight")]
|
| 113 |
+
# peft names the wrapped text model `base_model.model.<layers...>`; mlx-lm nests it as `language_model.model.<layers...>`
|
| 114 |
+
target = stem.replace("base_model.model.", "language_model.model.", 1) + ".weight"
|
| 115 |
+
if target not in params:
|
| 116 |
+
raise ValueError(f"adapter tensor {stem} has no weight in the mlx-lm model (looked for {target})")
|
| 117 |
+
base = params[target]
|
| 118 |
+
delta = (weights[stem + ".lora_B.weight"].astype(mx.float32) @ a.astype(mx.float32)) * (alpha * scale)
|
| 119 |
+
merged[target] = (base.astype(mx.float32) + delta).astype(base.dtype)
|
| 120 |
+
mx.eval(merged[target]) # one tensor at a time: one graph over every target peaks at ~2.9x the base weights and leaves ~2x of them in MLX's buffer cache (25.8 GB RSS for the 4B)
|
| 121 |
+
lm.load_weights(list(merged.items()), strict=False)
|
| 122 |
+
mx.eval(lm.parameters())
|
| 123 |
+
del params, weights # drop the pre-merge weights and the adapter before clearing the cache, or ~7 GB stay cached
|
| 124 |
+
mx.clear_cache() # hand the merge transients back to the OS; MLX keeps freed buffers otherwise
|
| 125 |
+
return len(merged)
|
| 126 |
+
|
| 127 |
+
|
| 128 |
+
class MLXDecisionModel:
|
| 129 |
+
"""Prefill-only scorer: hidden states from mlx-lm, logits from the shared torch PointerHead."""
|
| 130 |
+
backend, device, hybrid, option_isolation = "mlx", "mlx", True, False
|
| 131 |
+
prefix_min_tokens = 0 # kev.serve caches the state prefix for every request: on Metal the branch-only pass is always the cheaper one
|
| 132 |
+
|
| 133 |
+
def __init__(self, lm, pad_id, head_dim=256):
|
| 134 |
+
self.lm = lm # mlx-lm qwen3_5.Model from load_base (+ merge_lora) or load_full
|
| 135 |
+
self.text = self.lm.language_model.model # Qwen3_5TextModel: embeddings -> layers -> final norm = `.model.last_hidden_state`
|
| 136 |
+
self.pad_id = pad_id
|
| 137 |
+
self.head = PointerHead(getattr(self.text.embed_tokens, "dims", self.text.embed_tokens.weight.shape[1]), dp=head_dim).eval()
|
| 138 |
+
|
| 139 |
+
@property
|
| 140 |
+
def dtype(self):
|
| 141 |
+
return str(getattr(self.text.embed_tokens, "activation_dtype", self.text.embed_tokens.weight.dtype)).removeprefix("mlx.core.")
|
| 142 |
+
|
| 143 |
+
def eval(self):
|
| 144 |
+
self.head.eval(); return self
|
| 145 |
+
|
| 146 |
+
def encode(self, tok, rec, **kw):
|
| 147 |
+
return encode(tok, rec, option_isolation=False, **kw)
|
| 148 |
+
|
| 149 |
+
def _hidden(self, rows, cache=None):
|
| 150 |
+
"""[N, L, d] hidden states of right-padded token rows. Pads sit after every real token and both layer kinds are
|
| 151 |
+
causal (attention: causal mask; DeltaNet: a left-to-right recurrence), so no real token sees a pad."""
|
| 152 |
+
L = max(len(r) for r in rows)
|
| 153 |
+
ids = mx.array([r + [self.pad_id] * (L - len(r)) for r in rows], dtype=mx.int32)
|
| 154 |
+
h = self.text(ids, cache=cache)
|
| 155 |
+
mx.eval(h)
|
| 156 |
+
return h
|
| 157 |
+
|
| 158 |
+
def _logits(self, h, decide, opts):
|
| 159 |
+
"""One question's logits through the fp32 pointer head (temperature included, eval mode)."""
|
| 160 |
+
idx = mx.array([decide, *opts], dtype=mx.int32)
|
| 161 |
+
picked = torch.from_numpy(np.asarray(h[idx].astype(mx.float32)))
|
| 162 |
+
with torch.no_grad():
|
| 163 |
+
return self.head(picked[0], picked[1:])
|
| 164 |
+
|
| 165 |
+
def forward_rows(self, enc):
|
| 166 |
+
"""Row form, as the torch path computes it: every question is one causal row of state + branch tokens, the state
|
| 167 |
+
recomputed per row. The reference the prefix form is checked against (tests/test_mlx.py); serving uses `forward`."""
|
| 168 |
+
S, _, rows = rows_of(enc)
|
| 169 |
+
chunk, out = rows_per_pass([S + r["ids"] for r in rows]), []
|
| 170 |
+
for start in range(0, len(rows), chunk):
|
| 171 |
+
part = rows[start:start + chunk]
|
| 172 |
+
h = self._hidden([S + r["ids"] for r in part])
|
| 173 |
+
out += [self._logits(h[i], len(S) + r["decide"], [len(S) + o for o in r["opts"]]) for i, r in enumerate(part)]
|
| 174 |
+
return out
|
| 175 |
+
|
| 176 |
+
# --- state prefix: the state runs once into an mlx-lm prompt cache (KV for the attention layers, conv + recurrent state
|
| 177 |
+
# for the DeltaNet layers); the branches run as one batch on a replicated copy, so the prefix stays pristine and can be
|
| 178 |
+
# reused by the next request with the same state. On Metal this is also the cheapest way to answer a single request
|
| 179 |
+
# (state once instead of once per question), so it is the only path `forward` / `probs` take.
|
| 180 |
+
|
| 181 |
+
def prefix(self, enc):
|
| 182 |
+
"""The state into a fresh prompt cache, PREFILL_CHUNK tokens per pass, each pass's cache evaluated before the next
|
| 183 |
+
(mlx-lm's own prefill does the same; the hidden states a pass returns are never read). One pass over the whole
|
| 184 |
+
state kept every layer's activations of every token alive at once and left the DeltaNet conv states lazy (each
|
| 185 |
+
holding its layer's whole [Ls, conv_dim] input): Kev-4B, 8,192 tokens, 4.6 GB above the weights instead of 1.2 GB,
|
| 186 |
+
the same 6.5 s. Exact: a pass continues the cached keys/values (rotary offset = cache offset), conv window and
|
| 187 |
+
fp32 recurrent state, so chunked and one-pass differ by float reassociation only (fp32 on the CPU: max |dp| 1e-6;
|
| 188 |
+
bf16 on Metal: up to ~0.01, like any change of kernel shapes; tests/test_mlx.py, runs/mlx-long-states/prefill-ab-*)."""
|
| 189 |
+
Ls = enc["seg"].count(0)
|
| 190 |
+
cache, ids = make_prompt_cache(self.lm), enc["ids"][:Ls]
|
| 191 |
+
for start in range(0, Ls, PREFILL_CHUNK):
|
| 192 |
+
self.text(mx.array([ids[start:start + PREFILL_CHUNK]], dtype=mx.int32), cache=cache)
|
| 193 |
+
mx.eval([c.state for c in cache])
|
| 194 |
+
return Ls, cache
|
| 195 |
+
|
| 196 |
+
def _branch_logits(self, enc, cache):
|
| 197 |
+
"""Branches as rows on a replicated copy of the state cache, rows_per_pass rows (and cache copies) at a time."""
|
| 198 |
+
_, _, rows = rows_of(enc)
|
| 199 |
+
chunk, out = rows_per_pass([r["ids"] for r in rows], enc["seg"].count(0)), []
|
| 200 |
+
for start in range(0, len(rows), chunk):
|
| 201 |
+
part = rows[start:start + chunk]
|
| 202 |
+
batch = [type(c).merge([c] * len(part)) for c in cache] # merge copies the arrays: `cache` is not mutated
|
| 203 |
+
h = self._hidden([r["ids"] for r in part], batch)
|
| 204 |
+
out += [self._logits(h[i], r["decide"], r["opts"]) for i, r in enumerate(part)]
|
| 205 |
+
return out
|
| 206 |
+
|
| 207 |
+
def _branch_probs(self, enc, cache):
|
| 208 |
+
return [F.softmax(z, -1) for z in self._branch_logits(enc, cache)]
|
| 209 |
+
|
| 210 |
+
def forward(self, enc):
|
| 211 |
+
"""List of logits tensors, one per question."""
|
| 212 |
+
return self._branch_logits(enc, self.prefix(enc)[1])
|
| 213 |
+
|
| 214 |
+
def probs(self, enc):
|
| 215 |
+
return self._branch_probs(enc, self.prefix(enc)[1])
|
| 216 |
+
|
| 217 |
+
def probs_and_prefix(self, enc):
|
| 218 |
+
prefix = self.prefix(enc)
|
| 219 |
+
return self._branch_probs(enc, prefix[1]), prefix
|
| 220 |
+
|
| 221 |
+
def probs_with_prefix(self, enc, prefix):
|
| 222 |
+
Ls, cache = prefix
|
| 223 |
+
if enc["seg"].count(0) != Ls: raise ValueError("prefix does not match this record's state")
|
| 224 |
+
return self._branch_probs(enc, cache)
|
| 225 |
+
|
| 226 |
+
def probs_batch(self, encs, prefixes, keep):
|
| 227 |
+
"""kev.serve's batch call: one request at a time on Metal."""
|
| 228 |
+
out = [probs_one(self, e, p, k) for e, p, k in zip(encs, prefixes, keep)]
|
| 229 |
+
return [o[0] for o in out], [o[1] for o in out]
|
runtime/kev/model.py
ADDED
|
@@ -0,0 +1,508 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Decision model: causal LM backbone + block-causal branch mask + pointer readout."""
|
| 2 |
+
import copy, math, os, re
|
| 3 |
+
import torch
|
| 4 |
+
import torch.nn as nn
|
| 5 |
+
import torch.nn.functional as F
|
| 6 |
+
from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer, DynamicCache
|
| 7 |
+
from transformers.cache_utils import LinearAttentionCacheLayerMixin
|
| 8 |
+
|
| 9 |
+
# Reuse existing rarely-used Qwen special tokens as delimiters (state, q, opt, /opt, decide) so no
|
| 10 |
+
# embedding rows need to be added/trained; LoRA adapts their meaning.
|
| 11 |
+
SPECIAL = ["<|fim_prefix|>", "<|fim_middle|>", "<|box_start|>", "<|box_end|>", "<|fim_suffix|>"]
|
| 12 |
+
# training context: state tokens, tokens per question branch, and the whole packed record. Frozen suites are admitted with
|
| 13 |
+
# this rule (kev.suite) and training applies it to records built on the fly, so train and eval see the same population.
|
| 14 |
+
MAX_STATE, MAX_BRANCH, MAX_PACKED = 384, 1024, 2048
|
| 15 |
+
# serving context (kev.serve): a state of up to 64k tokens (twice Jev's 32k; the Qwen3.5 / Qwen3.8 bases' window is 262k)
|
| 16 |
+
# and a question row (state + its branch, the encoder's max_branch) of up to 8k tokens more. Until 64k states these were
|
| 17 |
+
# 8,192 / 8,192 (SERVE_MAX_*_8K below).
|
| 18 |
+
SERVE_MAX_STATE = 65536
|
| 19 |
+
SERVE_MAX_BRANCH = SERVE_MAX_STATE + 8192
|
| 20 |
+
SERVE_MAX_PACKED = SERVE_MAX_STATE + SERVE_MAX_BRANCH
|
| 21 |
+
# tokens one inference pass holds (rows_per_pass), and the longest request an attention-only backbone runs packed
|
| 22 |
+
# (rows_form: its block-causal mask is L x L); longer ones run as rows
|
| 23 |
+
ROW_PASS_TOKENS = 16384
|
| 24 |
+
# the longest state a checkpoint may be trained on (kev.train --max_state): the longest one served, which still leaves
|
| 25 |
+
# every question its training branch budget (training_context(MAX_TRAIN_STATE)["max_branch"] <= SERVE_MAX_BRANCH)
|
| 26 |
+
MAX_TRAIN_STATE = SERVE_MAX_STATE
|
| 27 |
+
# the limits before 64k states: the suites frozen until then were admitted under these and record them (kev.suite.SERVING_CONTEXT_8K)
|
| 28 |
+
SERVE_MAX_STATE_8K, SERVE_MAX_BRANCH_8K, MAX_TRAIN_STATE_8K = 8192, 8192, 7552
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
def training_context(max_state=MAX_STATE):
|
| 32 |
+
"""The encoder limits for training with the state limit lifted to `max_state`: the row (state + one branch) and
|
| 33 |
+
packed limits grow by the same amount, so every question keeps its token budget. training_context() is the default
|
| 34 |
+
training context (kev.suite.CONTEXT without `truncate`)."""
|
| 35 |
+
if not MAX_STATE <= max_state <= MAX_TRAIN_STATE:
|
| 36 |
+
raise ValueError(f"max_state must be in [{MAX_STATE}, {MAX_TRAIN_STATE}]")
|
| 37 |
+
extra = max_state - MAX_STATE
|
| 38 |
+
return {"max_state": max_state, "max_branch": MAX_BRANCH + extra, "max_packed": MAX_PACKED + extra}
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def rows_per_pass(rows, prefix_len=0, budget=ROW_PASS_TOKENS):
|
| 42 |
+
"""How many causal rows one inference forward pass takes: as many as fit `budget` tokens counting the cached state
|
| 43 |
+
each row carries (prefix_len) plus its own tokens, at least one. Memory per pass is bounded by one maximal row however
|
| 44 |
+
many questions a request has, and the rows are independent, so the answers do not depend on the split. Kev-4B on MLX,
|
| 45 |
+
a 4.8k-token state with 64 questions: 24.7 GB peak in one pass, 8.9 GB one row at a time (and faster: the cost was
|
| 46 |
+
64 copies of the state cache)."""
|
| 47 |
+
return max(1, budget // (prefix_len + max(len(r) for r in rows)))
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
class ContextOverflow(ValueError):
|
| 51 |
+
"""A record does not encode within its context (state, branch or packed limit). Serving turns it into a 422; the
|
| 52 |
+
benchmark counts it as a rejected record for suites scored as published (skip_overlong). When the state is what
|
| 53 |
+
overflows, state_tokens (its count, the <state> token included) and max_state (the limit) say by how much."""
|
| 54 |
+
|
| 55 |
+
def __init__(self, message, state_tokens=None, max_state=None):
|
| 56 |
+
super().__init__(message)
|
| 57 |
+
self.state_tokens, self.max_state = state_tokens, max_state
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
def load_tokenizer(name, revision=None):
|
| 61 |
+
return AutoTokenizer.from_pretrained(name, revision=revision)
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def pad_id(tok):
|
| 65 |
+
"""The id used to right-pad token rows (never attended to); Qwen tokenizers define one, others fall back to 0."""
|
| 66 |
+
return tok.pad_token_id if tok.pad_token_id is not None else 0
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
def is_hybrid(config):
|
| 70 |
+
"""Whether a (text) config has Gated DeltaNet layers (Qwen3.5). Such backbones cannot honour the block-causal mask and
|
| 71 |
+
run the row form; on Apple Silicon they are what the MLX backend is for."""
|
| 72 |
+
return "linear_attention" in set(getattr(config, "layer_types", None) or [])
|
| 73 |
+
|
| 74 |
+
|
| 75 |
+
_SPECIAL_RE = re.compile(r"<\|([A-Za-z0-9_]+)\|>")
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
def user_tokens(tok, text):
|
| 79 |
+
"""Tokenize caller-supplied text so it can never produce delimiter/control tokens (option boundaries are unforgeable).
|
| 80 |
+
The fast tokenizer ignores split_special_tokens, so `<|name|>` is rewritten to `<¦name¦>` before tokenizing."""
|
| 81 |
+
return tok(_SPECIAL_RE.sub(r"<¦\1¦>", text), add_special_tokens=False).input_ids
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
OPT_NONE, OPT_DECIDE = -1, -2 # values of enc["opt"]: instruction/state tokens, and the <decide> token
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
def encode(tok, rec, max_state=MAX_STATE, max_branch=MAX_BRANCH, strict=False, option_isolation=False):
|
| 88 |
+
"""Pack one record: [<state> ...] then per-question [<q> instr <opt> o </opt>... <decide>].
|
| 89 |
+
|
| 90 |
+
Returns ids, seg (0 = state, k = question k), pos (branch positions restart after state),
|
| 91 |
+
decide_idx [Q], opt_idx [Q][K] (index of </opt> token for each option), opt (per-token option index within its
|
| 92 |
+
question: OPT_NONE for state/instruction, 0..K-1 for option spans, OPT_DECIDE for <decide>).
|
| 93 |
+
|
| 94 |
+
option_isolation=True: every option span is its own sub-branch (it sees state + instruction + itself only), all
|
| 95 |
+
option spans share the same position ids, and <decide> sits at one fixed position after the longest span. Then the
|
| 96 |
+
per-option representations and <decide>'s attention over them are permutation-invariant by construction.
|
| 97 |
+
|
| 98 |
+
Not strict, a state over max_state tokens is cut to its first max_state (state_truncated; state_tokens is the whole
|
| 99 |
+
state's count, the <state> token included, as the limit counts it). Strict, it raises ContextOverflow.
|
| 100 |
+
"""
|
| 101 |
+
state_tokens = user_tokens(tok, rec["state"])
|
| 102 |
+
if strict and len(state_tokens) + 1 > max_state:
|
| 103 |
+
raise ContextOverflow(f"state exceeds {max_state} tokens: {len(state_tokens) + 1}", state_tokens=len(state_tokens) + 1, max_state=max_state)
|
| 104 |
+
S = [tok.convert_tokens_to_ids(SPECIAL[0])] + state_tokens[: max_state - 1]
|
| 105 |
+
ids, seg, pos, opt = list(S), [0] * len(S), list(range(len(S))), [OPT_NONE] * len(S)
|
| 106 |
+
q_id, o_id, c_id, d_id = (tok.convert_tokens_to_ids(t) for t in SPECIAL[1:])
|
| 107 |
+
decide_idx, opt_idx = [], []
|
| 108 |
+
for k, q in enumerate(rec["questions"], start=1):
|
| 109 |
+
instr = [q_id] + user_tokens(tok, q["instr"])
|
| 110 |
+
spans = [[o_id] + user_tokens(tok, o) + [c_id] for o in q["options"]]
|
| 111 |
+
br = instr + [t for sp in spans for t in sp] + [d_id]
|
| 112 |
+
if len(br) > max_branch - len(S):
|
| 113 |
+
raise ContextOverflow(f"branch too long: {len(br)} tokens with a {len(S)}-token state (row limit {max_branch})")
|
| 114 |
+
base = len(ids); p0 = len(S)
|
| 115 |
+
br_opt = [OPT_NONE] * len(instr) + [j for j, sp in enumerate(spans) for _ in sp] + [OPT_DECIDE]
|
| 116 |
+
if option_isolation:
|
| 117 |
+
longest = max(len(sp) for sp in spans)
|
| 118 |
+
br_pos = list(range(p0, p0 + len(instr))) + [p0 + len(instr) + i for sp in spans for i in range(len(sp))] + [p0 + len(instr) + longest]
|
| 119 |
+
else:
|
| 120 |
+
br_pos = list(range(p0, p0 + len(br)))
|
| 121 |
+
ends, cursor = [], len(instr)
|
| 122 |
+
for sp in spans:
|
| 123 |
+
cursor += len(sp); ends.append(cursor - 1)
|
| 124 |
+
ids += br; seg += [k] * len(br); pos += br_pos; opt += br_opt
|
| 125 |
+
decide_idx.append(base + len(br) - 1); opt_idx.append([base + e for e in ends])
|
| 126 |
+
return {"ids": ids, "seg": seg, "pos": pos, "opt": opt, "option_isolation": option_isolation, "decide_idx": decide_idx, "opt_idx": opt_idx,
|
| 127 |
+
"labels": [q["label"] for q in rec["questions"]], "state_tokens": len(state_tokens) + 1, "state_truncated": len(state_tokens) + 1 > max_state}
|
| 128 |
+
|
| 129 |
+
|
| 130 |
+
def admit(model, tok, rec, truncate=False):
|
| 131 |
+
"""The serving admission check (kev.serve and the Space; the torch and MLX models both encode through encode above):
|
| 132 |
+
-> rec encoded within the serving context. A state over SERVE_MAX_STATE tokens (the <state> token included) raises
|
| 133 |
+
ContextOverflow saying how long it is and how to fix it, unless truncate=True: then its first SERVE_MAX_STATE tokens
|
| 134 |
+
are read and the encoding says so (state_truncated, state_tokens). A question whose row (state + its branch) is over
|
| 135 |
+
SERVE_MAX_BRANCH raises ContextOverflow either way. Benchmarks do not truncate either: kev.predictors.LocalPredictor
|
| 136 |
+
encodes strictly within its suite's context."""
|
| 137 |
+
try:
|
| 138 |
+
return model.encode(tok, rec, max_state=SERVE_MAX_STATE, max_branch=SERVE_MAX_BRANCH, strict=not truncate)
|
| 139 |
+
except ContextOverflow as e:
|
| 140 |
+
if e.max_state is None: raise
|
| 141 |
+
raise ContextOverflow(f"state is {e.state_tokens:,} tokens, over the {e.max_state:,}-token limit (the <state> token included): "
|
| 142 |
+
"shorten the document or split it across requests", state_tokens=e.state_tokens, max_state=e.max_state) from None
|
| 143 |
+
|
| 144 |
+
|
| 145 |
+
def fits(rec, *tokenizers, max_state=MAX_STATE, max_branch=MAX_BRANCH, max_packed=MAX_PACKED):
|
| 146 |
+
"""True when the internal record encodes strictly (no truncation) within the training context under every tokenizer
|
| 147 |
+
given (frozen suites are admitted against the tokenizers of all their pinned bases)."""
|
| 148 |
+
try:
|
| 149 |
+
return all(len(encode(tok, rec, max_state=max_state, max_branch=max_branch, strict=True)["ids"]) <= max_packed for tok in tokenizers)
|
| 150 |
+
except ValueError:
|
| 151 |
+
return False
|
| 152 |
+
|
| 153 |
+
|
| 154 |
+
def branch_mask(seg, device, dtype=torch.float32):
|
| 155 |
+
"""attend(i,j) iff j<=i and (seg[j]==0 or seg[j]==seg[i]). Returns additive [1,1,L,L]."""
|
| 156 |
+
return branch_mask_batch([seg], device, dtype)
|
| 157 |
+
|
| 158 |
+
|
| 159 |
+
def branch_mask_batch(segs, device, dtype=torch.float32, opts=None, length=None):
|
| 160 |
+
"""Batched block-causal mask, additive [B,1,L,L], right-padded to the longest sequence.
|
| 161 |
+
|
| 162 |
+
Padded key positions are masked for every query; padded query rows keep the diagonal so no row is fully
|
| 163 |
+
masked (finfo.min, not -inf, so softmax stays finite either way). Real tokens never see pads because pads sit
|
| 164 |
+
after them (causal) and belong to no segment (-1).
|
| 165 |
+
|
| 166 |
+
opts (option isolation): within a question, an option-span token may attend to state, the instruction, and its own
|
| 167 |
+
span only; <decide> attends to everything in its question. Instruction tokens never see option spans (causal)."""
|
| 168 |
+
L = max(max(len(s) for s in segs), length or 0)
|
| 169 |
+
s = torch.full((len(segs), L), -1, device=device)
|
| 170 |
+
for b, seg in enumerate(segs):
|
| 171 |
+
s[b, : len(seg)] = torch.tensor(seg, device=device)
|
| 172 |
+
causal = torch.tril(torch.ones(L, L, dtype=torch.bool, device=device))
|
| 173 |
+
same = (s[:, None, :] == s[:, :, None]) | (s[:, None, :] == 0)
|
| 174 |
+
valid_key = (s != -1)[:, None, :]
|
| 175 |
+
allow = causal[None] & same & valid_key
|
| 176 |
+
if opts is not None:
|
| 177 |
+
o = torch.full((len(segs), L), OPT_NONE, device=device)
|
| 178 |
+
for b, op in enumerate(opts):
|
| 179 |
+
o[b, : len(op)] = torch.tensor(op, device=device)
|
| 180 |
+
key_is_option = (o[:, None, :] >= 0)
|
| 181 |
+
query_is_decide = (o[:, :, None] == OPT_DECIDE)
|
| 182 |
+
same_option = o[:, None, :] == o[:, :, None]
|
| 183 |
+
allow = allow & (~key_is_option | query_is_decide | same_option)
|
| 184 |
+
allow = allow | torch.eye(L, dtype=torch.bool, device=device)[None]
|
| 185 |
+
return torch.zeros(len(segs), L, L, dtype=dtype, device=device).masked_fill(~allow, torch.finfo(dtype).min)[:, None]
|
| 186 |
+
|
| 187 |
+
|
| 188 |
+
def rows_of(enc):
|
| 189 |
+
"""Split a packed encoding into its state and per-question branch rows.
|
| 190 |
+
|
| 191 |
+
Returns (state_ids, state_pos, rows) with rows[k] = {"ids", "pos", "decide", "opts"}: the branch tokens of question
|
| 192 |
+
k with their (already state-continuing) positions, and the readout offsets *within the branch*. Feeding
|
| 193 |
+
state + rows[k] as one causal row is equivalent to the packed block-causal form for that question, on any
|
| 194 |
+
architecture: the row contains exactly the tokens question k may attend to, in the same positions."""
|
| 195 |
+
seg = enc["seg"]; Ls = seg.count(0)
|
| 196 |
+
rows, start = [], Ls
|
| 197 |
+
for k, (d, oi) in enumerate(zip(enc["decide_idx"], enc["opt_idx"]), start=1):
|
| 198 |
+
end = d + 1 # <decide> is the last token of its branch
|
| 199 |
+
if seg[start] != k or seg[end - 1] != k: raise ValueError("branch layout mismatch")
|
| 200 |
+
rows.append({"ids": enc["ids"][start:end], "pos": enc["pos"][start:end], "decide": d - start, "opts": [o - start for o in oi]})
|
| 201 |
+
start = end
|
| 202 |
+
return enc["ids"][:Ls], enc["pos"][:Ls], rows
|
| 203 |
+
|
| 204 |
+
|
| 205 |
+
class PointerHead(nn.Module):
|
| 206 |
+
def __init__(self, d, dp=256):
|
| 207 |
+
"""dp = pointer dimension (head capacity knob)."""
|
| 208 |
+
super().__init__()
|
| 209 |
+
self.q, self.k = nn.Linear(d, dp), nn.Linear(d, dp)
|
| 210 |
+
self.scale = 1 / math.sqrt(dp)
|
| 211 |
+
# calibration: logits are divided by this at inference (eval mode) only. 1.0 = raw. A checkpoint carries the value fitted on
|
| 212 |
+
# its in-distribution development rows (scripts/calibrate_checkpoint.py -> head.pt["temperature"]); training always sees T=1 so
|
| 213 |
+
# a fitted value stays meaningful, and the argmax is unchanged by construction.
|
| 214 |
+
self.temperature = 1.0
|
| 215 |
+
|
| 216 |
+
def forward(self, h_decide, h_opts): # [d], [K,d] -> logits [K]
|
| 217 |
+
z = (self.k(h_opts) @ self.q(h_decide)) * self.scale
|
| 218 |
+
return z if self.training or self.temperature == 1.0 else z / self.temperature
|
| 219 |
+
|
| 220 |
+
def many(self, h_decide, h_opts, owner): # [Q,d], [sum K,d], question of each option [sum K] -> logits [sum K]
|
| 221 |
+
"""forward() for many questions at once (serving batches): one pass instead of one per question."""
|
| 222 |
+
z = (self.k(h_opts) * self.q(h_decide)[owner]).sum(-1) * self.scale
|
| 223 |
+
return z if self.training or self.temperature == 1.0 else z / self.temperature
|
| 224 |
+
|
| 225 |
+
|
| 226 |
+
# What a loaded model exposes to kev.serve, kev.predictors and the Space: the scoring interface both DecisionModel (torch)
|
| 227 |
+
# and kev.mlx_model.MLXDecisionModel implement. tests/test_mlx.py checks the MLX class against this list.
|
| 228 |
+
SCORING_INTERFACE = ("encode", "forward", "probs", "probs_and_prefix", "probs_with_prefix", "probs_batch", "eval",
|
| 229 |
+
"head", "backend", "dtype", "device", "hybrid", "option_isolation", "prefix_min_tokens")
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
EAGER_STATES = 4 # long new states (past the graphed state pass) whose eager prefixes one batched run holds at once: each
|
| 233 |
+
# holds its state's keys, values and DeltaNet states (~0.3 GB for 2,200 tokens on Kev-27B)
|
| 234 |
+
|
| 235 |
+
|
| 236 |
+
def probs_one(model, enc, prefix, keep):
|
| 237 |
+
"""-> (probs, prefix to keep) for one request, on any backend: the question rows on the cached prefix on a hit, one
|
| 238 |
+
pass that also returns the prefix when it is to be kept, the plain pass otherwise."""
|
| 239 |
+
if prefix is not None: return model.probs_with_prefix(enc, prefix), prefix
|
| 240 |
+
return model.probs_and_prefix(enc) if keep else (model.probs(enc), None)
|
| 241 |
+
|
| 242 |
+
|
| 243 |
+
class DecisionModel(nn.Module):
|
| 244 |
+
def __init__(self, name, tok, device, lora=None, revision=None, attn=None, head_dim=256, option_isolation=False, special_embeddings=False, lora_targets="all", dtype=torch.float32,
|
| 245 |
+
weights=None, direct_load=False):
|
| 246 |
+
"""weights: a full-weight checkpoint directory whose saved backbone replaces the base's (kev.checkpoint's loader rule).
|
| 247 |
+
direct_load: load the backbone straight onto `device` (transformers device_map) instead of staging it in host memory;
|
| 248 |
+
full-weight training on several GPUs in one container needs it (N processes x a 51 GB checkpoint otherwise). Off by
|
| 249 |
+
default: it changes where the rotary buffers are computed, so every other path keeps its bits."""
|
| 250 |
+
super().__init__()
|
| 251 |
+
# backbone only (no vocab head): we never generate text.
|
| 252 |
+
# eager on MPS/CPU (known-good with our float 4D mask); SDPA on CUDA (accepts arbitrary additive masks).
|
| 253 |
+
attn = attn or ("sdpa" if str(device).startswith("cuda") else "eager")
|
| 254 |
+
# dtype: fp32 for training and exact evaluation; bf16 is a serving option for large backbones (8B on a 32 GB Mac)
|
| 255 |
+
load = {"dtype": dtype, "attn_implementation": attn}
|
| 256 |
+
if direct_load: load["device_map"] = {"": torch.cuda.current_device() if device == "cuda" else device} # "cuda": under torchrun, this rank's GPU
|
| 257 |
+
self.lm = AutoModel.from_pretrained(weights, **load) if weights else AutoModelForCausalLM.from_pretrained(name, revision=revision, **load).model
|
| 258 |
+
self.pad_id = pad_id(tok)
|
| 259 |
+
# hybrid backbones (Qwen3.5: Gated DeltaNet layers, recurrent) cannot honour the block-causal mask, so every
|
| 260 |
+
# question runs as its own causal row continuing from the state (rows_of). Attention-only backbones keep the
|
| 261 |
+
# packed form; the two agree to fp32 noise (tests/test_model.py::test_rows_match_packed).
|
| 262 |
+
self.hybrid = is_hybrid(self.lm.config)
|
| 263 |
+
if self.hybrid and option_isolation: raise ValueError("option_isolation needs the packed mask; not available on hybrid backbones")
|
| 264 |
+
self.option_isolation = option_isolation
|
| 265 |
+
if lora:
|
| 266 |
+
from peft import LoraConfig, get_peft_model
|
| 267 |
+
extra = {"trainable_token_indices": {"embed_tokens": [tok.convert_tokens_to_ids(t) for t in SPECIAL]}} if special_embeddings else {}
|
| 268 |
+
targets = {"all": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
|
| 269 |
+
"dense": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], # "all" minus the DeltaNet projections on hybrids (retention ablation)
|
| 270 |
+
"attn": ["q_proj", "k_proj", "v_proj", "o_proj"], "qv": ["q_proj", "v_proj"]}[lora_targets]
|
| 271 |
+
if self.hybrid and lora_targets in ("all", "attn"):
|
| 272 |
+
# Gated DeltaNet projections (transformers 5 names, verified on Qwen3_5TextModel); the mixer's out_proj too
|
| 273 |
+
targets = targets + ["in_proj_qkv", "in_proj_z", "in_proj_a", "in_proj_b", "out_proj"]
|
| 274 |
+
cfg = LoraConfig(task_type="FEATURE_EXTRACTION", r=lora, lora_alpha=2 * lora, lora_dropout=0.05, target_modules=targets, **extra)
|
| 275 |
+
self.lm = get_peft_model(self.lm, cfg)
|
| 276 |
+
self.head = PointerHead(self.lm.config.hidden_size, dp=head_dim)
|
| 277 |
+
self.device = device
|
| 278 |
+
self.to(device)
|
| 279 |
+
|
| 280 |
+
backend = "torch" # kev.mlx_model.MLXDecisionModel is the other implementation of this scoring interface
|
| 281 |
+
graphs = None # kev.cuda_graphs.CudaGraphs for the serving passes of a hybrid backbone on CUDA (LoadOptions.cuda_graphs)
|
| 282 |
+
|
| 283 |
+
@property
|
| 284 |
+
def prefix_min_tokens(self):
|
| 285 |
+
"""kev.serve caches the state prefix from this many state tokens. Attention-only backbones: 384, below which the
|
| 286 |
+
branch-only pass is not faster than one packed pass on MPS (per-op overhead). Hybrid backbones: always, because
|
| 287 |
+
their miss path otherwise recomputes the state once per question (Kev-0.8B bf16 on MPS, 5 questions: 1011 -> 413 ms)."""
|
| 288 |
+
return 0 if self.hybrid else 384
|
| 289 |
+
|
| 290 |
+
@property
|
| 291 |
+
def dtype(self):
|
| 292 |
+
return str(next(self.lm.parameters()).dtype).removeprefix("torch.")
|
| 293 |
+
|
| 294 |
+
def encode(self, tok, rec, **kw):
|
| 295 |
+
"""encode() with this model's option-isolation setting; use this from serving/eval code."""
|
| 296 |
+
return encode(tok, rec, option_isolation=self.option_isolation, **kw)
|
| 297 |
+
|
| 298 |
+
def hidden(self, enc):
|
| 299 |
+
return self.hidden_batch([enc])[0, : len(enc["ids"])]
|
| 300 |
+
|
| 301 |
+
SHAPE_BUCKET = int(os.environ.get("KEV_SHAPE_BUCKET", "64")) # MPS: pad the sequence to a multiple of this (per-shape kernel warm-up); 1 disables
|
| 302 |
+
|
| 303 |
+
def _pad_rows(self, rows):
|
| 304 |
+
"""Right-pad (ids, pos) token rows into [N, L] id / position tensors and a [N, L] attention mask (1 = real token).
|
| 305 |
+
Pads sit after every real token and are masked keys, so they never change a real token's hidden state (parity
|
| 306 |
+
measured exact). On MPS in eval mode L is rounded up to a SHAPE_BUCKET multiple so kernels are warmed per bucket."""
|
| 307 |
+
L = max(len(ids) for ids, _ in rows)
|
| 308 |
+
if str(self.device) == "mps" and not self.training: L = -(-L // self.SHAPE_BUCKET) * self.SHAPE_BUCKET
|
| 309 |
+
ids = torch.full((len(rows), L), self.pad_id, device=self.device)
|
| 310 |
+
pos = torch.zeros((len(rows), L), dtype=torch.long, device=self.device)
|
| 311 |
+
att = torch.zeros((len(rows), L), dtype=torch.long, device=self.device)
|
| 312 |
+
for i, (rid, rpos) in enumerate(rows):
|
| 313 |
+
ids[i, : len(rid)] = torch.tensor(rid, device=self.device); pos[i, : len(rpos)] = torch.tensor(rpos, device=self.device); att[i, : len(rid)] = 1
|
| 314 |
+
return ids, pos, att
|
| 315 |
+
|
| 316 |
+
def hidden_batch(self, encs):
|
| 317 |
+
"""[B, L_max, d] hidden states for a right-padded batch of encoded records under the packed block-causal mask."""
|
| 318 |
+
ids, pos, _ = self._pad_rows([(e["ids"], e["pos"]) for e in encs])
|
| 319 |
+
isolate = any(e.get("option_isolation") for e in encs)
|
| 320 |
+
if isolate and not all(e.get("option_isolation") for e in encs):
|
| 321 |
+
raise ValueError("cannot mix option-isolated and plain encodings in one batch")
|
| 322 |
+
lm_dtype = next(self.lm.parameters()).dtype
|
| 323 |
+
mask = branch_mask_batch([e["seg"] for e in encs], self.device, dtype=lm_dtype, opts=[e["opt"] for e in encs] if isolate else None, length=ids.shape[1])
|
| 324 |
+
return self.lm(input_ids=ids, position_ids=pos, attention_mask=mask).last_hidden_state.float() # head stays fp32
|
| 325 |
+
|
| 326 |
+
def _readout(self, h, enc):
|
| 327 |
+
return [self.head(h[d], h[torch.tensor(oi, device=self.device)]) for d, oi in zip(enc["decide_idx"], enc["opt_idx"])]
|
| 328 |
+
|
| 329 |
+
def rows_form(self, encs):
|
| 330 |
+
"""Whether these records run as causal rows: hybrid backbones always (the recurrent layers cannot honour the packed
|
| 331 |
+
mask), attention-only ones when the packed sequence would exceed ROW_PASS_TOKENS (its L x L mask grows without
|
| 332 |
+
bound with the number of questions). The two forms agree (tests/test_model.py::test_rows_match_packed)."""
|
| 333 |
+
return self.hybrid or any(len(e["ids"]) > ROW_PASS_TOKENS for e in encs)
|
| 334 |
+
|
| 335 |
+
def _rows_hidden(self, rows, cache=None, prefix_len=0):
|
| 336 |
+
"""Hidden states of causal token rows, one [L_i, d] tensor per row. In eval mode the rows go through the backbone
|
| 337 |
+
rows_per_pass at a time; training keeps one batch (its batches are small and autograd needs the whole graph anyway).
|
| 338 |
+
With `cache`, the rows are branches continuing the cached state: in eager mode the cache is replicated once per
|
| 339 |
+
chunk, leaving the caller's prefix pristine, and the cached tokens are marked real in the attention mask."""
|
| 340 |
+
chunk = len(rows) if self.training else rows_per_pass([ids for ids, _ in rows], prefix_len)
|
| 341 |
+
out = []
|
| 342 |
+
for start in range(0, len(rows), chunk):
|
| 343 |
+
part = rows[start:start + chunk]
|
| 344 |
+
ids, pos, att = self._pad_rows(part)
|
| 345 |
+
past = {}
|
| 346 |
+
if cache is not None:
|
| 347 |
+
replica = copy.copy(cache)
|
| 348 |
+
replica.layers = [copy.copy(layer) for layer in cache.layers]
|
| 349 |
+
for source, target in zip(cache.layers, replica.layers):
|
| 350 |
+
if isinstance(source, LinearAttentionCacheLayerMixin):
|
| 351 |
+
target.conv_states = source.conv_states.copy()
|
| 352 |
+
target.recurrent_states = source.recurrent_states.copy()
|
| 353 |
+
target.is_conv_states_initialized = source.is_conv_states_initialized.copy()
|
| 354 |
+
target.is_recurrent_states_initialized = source.is_recurrent_states_initialized.copy()
|
| 355 |
+
target.has_previous_state = source.has_previous_state.copy()
|
| 356 |
+
target.conv_kernel_size = source.conv_kernel_size.copy()
|
| 357 |
+
replica.reorder_cache(torch.zeros(len(part), dtype=torch.long, device=self.device))
|
| 358 |
+
att = torch.cat([torch.ones((len(part), prefix_len), dtype=torch.long, device=self.device), att], 1)
|
| 359 |
+
past = {"past_key_values": replica, "use_cache": True}
|
| 360 |
+
h = self.lm(input_ids=ids, position_ids=pos, attention_mask=att, **past).last_hidden_state.float()
|
| 361 |
+
out += [h[i, : len(row_ids)] for i, (row_ids, _) in enumerate(part)]
|
| 362 |
+
return out
|
| 363 |
+
|
| 364 |
+
def forward_rows_batch(self, encs, shared_prefix=False):
|
| 365 |
+
"""Row form: every question of every record is one causal row = state tokens + its branch tokens. Returns the same
|
| 366 |
+
nested logits as forward_batch. Exact isolation by construction (rows are independent); the state is recomputed
|
| 367 |
+
per row (Q x state tokens), which training accepts; serving and probs() use the prefix cache instead. shared_prefix
|
| 368 |
+
(training, hybrid Qwen3.5): the same rows computed with each state run once and its branches continuing from it,
|
| 369 |
+
gradients included (kev.shared_prefix)."""
|
| 370 |
+
if shared_prefix:
|
| 371 |
+
from .shared_prefix import branch_hidden # here, not at the top: the HF Space vendors model.py alone
|
| 372 |
+
splits = [rows_of(e) for e in encs]
|
| 373 |
+
return [[self.head(h[r["decide"]], h[torch.tensor(r["opts"], device=self.device)]) for h, r in zip(hs, rows)]
|
| 374 |
+
for hs, (_, _, rows) in zip(branch_hidden(self.lm, splits, self.pad_id, self.device), splits)]
|
| 375 |
+
rows, readouts = [], [] # one causal row per question; readouts[i] = (record, <decide> offset, option offsets)
|
| 376 |
+
for b, e in enumerate(encs):
|
| 377 |
+
S, Sp, brs = rows_of(e)
|
| 378 |
+
for r in brs:
|
| 379 |
+
rows.append((S + r["ids"], Sp + r["pos"])); readouts.append((b, len(S) + r["decide"], [len(S) + o for o in r["opts"]]))
|
| 380 |
+
out = [[] for _ in encs]
|
| 381 |
+
for h, (b, d, oi) in zip(self._rows_hidden(rows), readouts):
|
| 382 |
+
out[b].append(self.head(h[d], h[torch.tensor(oi, device=self.device)]))
|
| 383 |
+
return out
|
| 384 |
+
|
| 385 |
+
def forward(self, enc):
|
| 386 |
+
"""Returns list of logits tensors, one per question."""
|
| 387 |
+
return self.forward_batch([enc])[0]
|
| 388 |
+
|
| 389 |
+
def forward_batch(self, encs, shared_prefix=False):
|
| 390 |
+
"""List (per record) of lists (per question) of logits. Row form or packed block-causal mask, see rows_form
|
| 391 |
+
(the packed mask already runs each state once; shared_prefix does that for the row form of a hybrid backbone)."""
|
| 392 |
+
if self.rows_form(encs): return self.forward_rows_batch(encs, shared_prefix and self.hybrid)
|
| 393 |
+
hs = self.hidden_batch(encs)
|
| 394 |
+
return [self._readout(hs[b], e) for b, e in enumerate(encs)]
|
| 395 |
+
|
| 396 |
+
@torch.no_grad()
|
| 397 |
+
def probs(self, enc):
|
| 398 |
+
"""Probabilities per question. A record that runs as rows (every record on a hybrid backbone) takes the serving miss
|
| 399 |
+
path: its state once, then the question rows from its cache. forward() keeps the row form, which re-runs the state
|
| 400 |
+
per question; the two agree to fp32 rounding (#77)."""
|
| 401 |
+
if self.rows_form([enc]): return self.probs_and_prefix(enc)[0]
|
| 402 |
+
return [F.softmax(z, -1).cpu() for z in self.forward(enc)]
|
| 403 |
+
|
| 404 |
+
# --- state-prefix reuse (serving): the state is encoded once, question branches attend to its cached keys/values.
|
| 405 |
+
# Exact by construction: branch tokens never attend to each other across questions (block-causal mask) and the state
|
| 406 |
+
# never sees the branches (causal), so the state's hidden states and KV are identical with or without the branches.
|
| 407 |
+
|
| 408 |
+
def _branch_rows_from_prefix(self, enc, cache):
|
| 409 |
+
"""Row-form serving: the branches run as causal rows continuing the cached state (exactly the forward_rows_batch
|
| 410 |
+
layout, minus the recomputed state)."""
|
| 411 |
+
S, _, rows = rows_of(enc)
|
| 412 |
+
hs = self._rows_hidden([(r["ids"], r["pos"]) for r in rows], cache=cache, prefix_len=len(S))
|
| 413 |
+
ps = [F.softmax(self.head(h[r["decide"]], h[torch.tensor(r["opts"], device=self.device)]), -1) for h, r in zip(hs, rows)]
|
| 414 |
+
return list(torch.cat(ps).cpu().split([len(p) for p in ps])) # one device sync for the request, not one per question
|
| 415 |
+
|
| 416 |
+
@torch.no_grad()
|
| 417 |
+
def prefix(self, enc):
|
| 418 |
+
"""Run the state tokens only. Returns (n_state_tokens, kv cache, state hidden states [Ls, d] fp32 or None). Only the
|
| 419 |
+
packed pass (probs_with_prefix) reads the hidden states; a hybrid backbone always runs rows, so it keeps None, as the
|
| 420 |
+
graphed prefixes do (a cached fp32 [Ls, d] copy is 160 MiB for 8k tokens on Kev-27B). An attention-only backbone
|
| 421 |
+
keeps them even when this record runs as rows: the same state with fewer questions may be packed."""
|
| 422 |
+
Ls = enc["seg"].count(0)
|
| 423 |
+
ids = torch.tensor([enc["ids"][:Ls]], device=self.device); pos = torch.tensor([enc["pos"][:Ls]], device=self.device)
|
| 424 |
+
# the cache must know the layer types (hybrid backbones keep recurrent + conv states per DeltaNet layer)
|
| 425 |
+
out = self.lm(input_ids=ids, position_ids=pos, past_key_values=DynamicCache(config=self.lm.config), use_cache=True)
|
| 426 |
+
return Ls, out.past_key_values, None if self.hybrid else out.last_hidden_state[0].float()
|
| 427 |
+
|
| 428 |
+
@torch.no_grad()
|
| 429 |
+
def probs_and_prefix(self, enc):
|
| 430 |
+
"""One full pass that also returns the state prefix (KV cropped to the state, state hidden states or None, see prefix): a cache miss
|
| 431 |
+
costs a single forward pass, not two."""
|
| 432 |
+
Ls = enc["seg"].count(0)
|
| 433 |
+
if self.rows_form([enc]):
|
| 434 |
+
# recurrent layers cannot be cropped back to the state (and an over-long packed pass is what the row form avoids),
|
| 435 |
+
# so a miss here is a state pass (kept as the prefix) plus the branch rows
|
| 436 |
+
Ls, cache, h_state = self.prefix(enc)
|
| 437 |
+
return self._branch_rows_from_prefix(enc, cache), (Ls, cache, h_state)
|
| 438 |
+
ids = torch.tensor([enc["ids"]], device=self.device); pos = torch.tensor([enc["pos"]], device=self.device)
|
| 439 |
+
dt = next(self.lm.parameters()).dtype
|
| 440 |
+
mask = branch_mask_batch([enc["seg"]], self.device, dtype=dt, opts=[enc["opt"]] if enc.get("option_isolation") else None)
|
| 441 |
+
out = self.lm(input_ids=ids, position_ids=pos, attention_mask=mask, past_key_values=DynamicCache(config=self.lm.config), use_cache=True)
|
| 442 |
+
h = out.last_hidden_state[0].float()
|
| 443 |
+
out.past_key_values.crop(-(len(enc["ids"]) - Ls)) # keep the state only (negative = drop that many trailing tokens; positive form deprecated in transformers 5)
|
| 444 |
+
return [F.softmax(z, -1).cpu() for z in self._readout(h, enc)], (Ls, out.past_key_values, h[:Ls].clone())
|
| 445 |
+
|
| 446 |
+
@torch.no_grad()
|
| 447 |
+
def probs_with_prefix(self, enc, prefix):
|
| 448 |
+
"""probs() for a record whose state tokens equal the cached prefix's; only the branches run. The cache is cropped
|
| 449 |
+
back to the state afterwards so it can be reused."""
|
| 450 |
+
Ls, cache, h_state = prefix
|
| 451 |
+
if enc["seg"].count(0) != Ls: raise ValueError("prefix does not match this record's state")
|
| 452 |
+
if self.rows_form([enc]):
|
| 453 |
+
return self._branch_rows_from_prefix(enc, cache)
|
| 454 |
+
ids = torch.tensor([enc["ids"][Ls:]], device=self.device); pos = torch.tensor([enc["pos"][Ls:]], device=self.device)
|
| 455 |
+
dt = next(self.lm.parameters()).dtype
|
| 456 |
+
mask = branch_mask_batch([enc["seg"]], self.device, dtype=dt, opts=[enc["opt"]] if enc.get("option_isolation") else None)[:, :, Ls:, :]
|
| 457 |
+
try:
|
| 458 |
+
out = self.lm(input_ids=ids, position_ids=pos, past_key_values=cache, attention_mask=mask, use_cache=True)
|
| 459 |
+
h = torch.cat([h_state, out.last_hidden_state[0].float()], 0)
|
| 460 |
+
finally:
|
| 461 |
+
cache.crop(-(len(enc["ids"]) - Ls))
|
| 462 |
+
return [F.softmax(z, -1).cpu() for z in self._readout(h, enc)]
|
| 463 |
+
|
| 464 |
+
@torch.no_grad()
|
| 465 |
+
def probs_batch(self, encs, prefixes, keep):
|
| 466 |
+
"""kev.serve's model thread: several requests at once. prefixes[i] is request i's cached state prefix or None;
|
| 467 |
+
keep[i] whether to return its new prefix. -> (probs per request, prefix per request: the cached one, the new one if
|
| 468 |
+
kept, else None). With CUDA graphs the requests the graphed passes admit run together (kev.cuda_graphs: shared
|
| 469 |
+
state and row passes; a state too long for the graphed state pass gets its own eager pass first); the rest, and
|
| 470 |
+
every other backend, one at a time. Rows are independent, so a request's answers do not depend on the batch."""
|
| 471 |
+
splits = [rows_of(e) for e in encs]
|
| 472 |
+
def fits(i, cached):
|
| 473 |
+
S, _, rows = splits[i]
|
| 474 |
+
return self.graphs is not None and self.graphs.admits(len(S), [len(r["ids"]) for r in rows], cached)
|
| 475 |
+
long = {i for i in range(len(encs)) if prefixes[i] is None and not fits(i, False) and fits(i, True)}
|
| 476 |
+
batched = [i for i in range(len(encs)) if i in long or fits(i, prefixes[i] is not None)]
|
| 477 |
+
out = {i: probs_one(self, encs[i], prefixes[i], keep[i]) for i in sorted(set(range(len(encs))) - set(batched))}
|
| 478 |
+
runs, held = [[]], 0 # graphed runs, each holding at most EAGER_STATES long states' eager prefixes at once
|
| 479 |
+
for i in batched:
|
| 480 |
+
if i in long and held == EAGER_STATES: runs.append([]); held = 0
|
| 481 |
+
runs[-1].append(i); held += i in long
|
| 482 |
+
from .cuda_graphs import Request # here, not at the top: the HF Space vendors model.py without cuda_graphs.py
|
| 483 |
+
for run in filter(None, runs):
|
| 484 |
+
pre = {i: self.prefix(encs[i]) if i in long else prefixes[i] for i in run}
|
| 485 |
+
X, caches = self.graphs.run([Request(splits[i][0], splits[i][1], [(r["ids"], r["pos"]) for r in splits[i][2]],
|
| 486 |
+
None if pre[i] is None else pre[i][1], keep[i], [[r["decide"], *r["opts"]] for r in splits[i][2]])
|
| 487 |
+
for i in run]) # per question its picks: the <decide> position, then its options
|
| 488 |
+
ps = iter(self._readout_many(X, [len(r["opts"]) for i in run for r in splits[i][2]]))
|
| 489 |
+
for i, cache in zip(run, caches):
|
| 490 |
+
made = pre[i] if i in long else None if cache is None else (len(splits[i][0]), cache, None)
|
| 491 |
+
out[i] = [next(ps) for _ in splits[i][2]], prefixes[i] or (made if keep[i] else None)
|
| 492 |
+
return [out[i][0] for i in range(len(encs))], [out[i][1] for i in range(len(encs))]
|
| 493 |
+
|
| 494 |
+
def _readout_many(self, X, ks):
|
| 495 |
+
"""Probabilities of many questions from their picked hidden states X, laid out per question as its <decide> then
|
| 496 |
+
its ks[q] options, in a fixed number of kernels and one device sync. -> one tensor per question."""
|
| 497 |
+
owner = [q for q, k in enumerate(ks) for _ in range(k)]
|
| 498 |
+
slot = [j for k in ks for j in range(k)]
|
| 499 |
+
starts = [0]
|
| 500 |
+
for k in ks: starts.append(starts[-1] + 1 + k)
|
| 501 |
+
dec = torch.tensor(starts[:-1]).to(self.device, non_blocking=True)
|
| 502 |
+
opt, own, sl = torch.tensor([[starts[q] + 1 + j for q, j in zip(owner, slot)], owner, slot]).to(self.device, non_blocking=True)
|
| 503 |
+
z = self.head.many(X[dec], X[opt], own)
|
| 504 |
+
Z = torch.full((len(ks), max(ks)), float("-inf"), device=self.device).index_put_((own, sl), z)
|
| 505 |
+
return torch.softmax(Z, -1)[own, sl].cpu().split(ks)
|
| 506 |
+
|
| 507 |
+
def trainable_parameters(self):
|
| 508 |
+
return [p for p in self.parameters() if p.requires_grad]
|
runtime/kev/packed_model.py
ADDED
|
@@ -0,0 +1,102 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Standalone MLX mixed ternary/4-bit backbone. Never reads or reconstructs base weights.
|
| 2 |
+
|
| 3 |
+
Ternary codes use MLX's native 2-bit affine kernel; CLoQ stays a separate
|
| 4 |
+
low-rank operation. Only selected embedding rows are dequantized. This is a
|
| 5 |
+
different arithmetic path from the experiment's merged-bf16 simulation.
|
| 6 |
+
"""
|
| 7 |
+
import json
|
| 8 |
+
from pathlib import Path
|
| 9 |
+
|
| 10 |
+
import mlx.core as mx
|
| 11 |
+
import mlx.nn as nn
|
| 12 |
+
from mlx.utils import tree_unflatten
|
| 13 |
+
from mlx_lm.models import qwen3_5
|
| 14 |
+
|
| 15 |
+
from .mlx_model import MLXDecisionModel
|
| 16 |
+
|
| 17 |
+
FORMAT = "kev-packed-mlx-v1"
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
class PackedLinear(nn.QuantizedLinear):
|
| 21 |
+
def __init__(self, input_dims, output_dims, bits, group_size, rank=0, bias=False):
|
| 22 |
+
# Avoid QuantizedLinear's random dense initialization entirely.
|
| 23 |
+
nn.Module.__init__(self)
|
| 24 |
+
if bits not in (2, 4, 8) or input_dims % group_size or input_dims % (32 // bits):
|
| 25 |
+
raise ValueError("unsupported packed linear dimensions or precision")
|
| 26 |
+
self.bits, self.group_size, self.mode = bits, group_size, "affine"
|
| 27 |
+
self.weight = mx.zeros((output_dims, input_dims * bits // 32), dtype=mx.uint32)
|
| 28 |
+
self.scales = mx.zeros((output_dims, input_dims // group_size), dtype=mx.float16)
|
| 29 |
+
self.biases = mx.zeros_like(self.scales)
|
| 30 |
+
if bias:
|
| 31 |
+
self.bias = mx.zeros((output_dims,), dtype=mx.bfloat16)
|
| 32 |
+
if rank:
|
| 33 |
+
self.comp_a = mx.zeros((rank, input_dims), dtype=mx.bfloat16)
|
| 34 |
+
self.comp_b = mx.zeros((output_dims, rank), dtype=mx.bfloat16)
|
| 35 |
+
self.freeze()
|
| 36 |
+
|
| 37 |
+
def __call__(self, x):
|
| 38 |
+
y = super().__call__(x)
|
| 39 |
+
if "comp_a" in self:
|
| 40 |
+
# Keep the small compensation in fp32; never materialize B @ A.
|
| 41 |
+
xf = x.astype(mx.float32)
|
| 42 |
+
y = y + (xf @ self.comp_a.astype(mx.float32).T) @ self.comp_b.astype(mx.float32).T
|
| 43 |
+
return y.astype(mx.bfloat16)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
class PackedEmbedding(nn.QuantizedEmbedding):
|
| 47 |
+
activation_dtype = mx.bfloat16
|
| 48 |
+
|
| 49 |
+
def __init__(self, num_embeddings, dims, bits, group_size):
|
| 50 |
+
nn.Module.__init__(self)
|
| 51 |
+
if bits not in (2, 4, 8) or dims % group_size or dims % (32 // bits):
|
| 52 |
+
raise ValueError("unsupported packed embedding dimensions or precision")
|
| 53 |
+
self.num_embeddings, self.dims = num_embeddings, dims
|
| 54 |
+
self.bits, self.group_size, self.mode = bits, group_size, "affine"
|
| 55 |
+
self.weight = mx.zeros((num_embeddings, dims * bits // 32), dtype=mx.uint32)
|
| 56 |
+
self.scales = mx.zeros((num_embeddings, dims // group_size), dtype=mx.float16)
|
| 57 |
+
self.biases = mx.zeros_like(self.scales)
|
| 58 |
+
self.freeze()
|
| 59 |
+
|
| 60 |
+
def __call__(self, x):
|
| 61 |
+
return super().__call__(x).astype(mx.bfloat16)
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def load_packed_model(directory, pad_id, head_dim):
|
| 65 |
+
"""Construct lazy architecture, replace quantized modules, strictly load packed tensors."""
|
| 66 |
+
root = Path(directory)
|
| 67 |
+
spec = json.loads((root / "packed.json").read_text(encoding="utf-8"))
|
| 68 |
+
if spec.get("format") != FORMAT:
|
| 69 |
+
raise ValueError("unsupported packed model format")
|
| 70 |
+
config = json.loads((root / "config.json").read_text(encoding="utf-8"))
|
| 71 |
+
args = qwen3_5.ModelArgs.from_dict(config)
|
| 72 |
+
text = qwen3_5.TextModelArgs.from_dict(args.text_config)
|
| 73 |
+
pattern = ["full_attention" if (i + 1) % text.full_attention_interval == 0 else "linear_attention" for i in range(text.num_hidden_layers)]
|
| 74 |
+
if args.text_config.get("layer_types", pattern) != pattern:
|
| 75 |
+
raise ValueError("packed layer_types do not match MLX's attention layout")
|
| 76 |
+
# Lazy initialization is discarded before evaluation, including the embedding and LM head.
|
| 77 |
+
lm = qwen3_5.Model(args)
|
| 78 |
+
if "lm_head" in lm.language_model:
|
| 79 |
+
del lm.language_model["lm_head"]
|
| 80 |
+
modules = dict(lm.named_modules())
|
| 81 |
+
replacements = []
|
| 82 |
+
for name, q in spec["linears"].items():
|
| 83 |
+
old = modules.get(name)
|
| 84 |
+
if not isinstance(old, nn.Linear):
|
| 85 |
+
raise ValueError(f"packed linear {name} is not in the architecture")
|
| 86 |
+
out_dims, in_dims = old.weight.shape
|
| 87 |
+
replacements.append((name, PackedLinear(in_dims, out_dims, q["bits"], q["group_size"], q.get("rank", 0), "bias" in old)))
|
| 88 |
+
embedding = spec["embedding"]
|
| 89 |
+
name = "language_model.model.embed_tokens"
|
| 90 |
+
replacements.append((name, PackedEmbedding(text.vocab_size, text.hidden_size, embedding["bits"], embedding["group_size"])))
|
| 91 |
+
lm.update_modules(tree_unflatten(replacements))
|
| 92 |
+
del modules, replacements
|
| 93 |
+
weights = mx.load(str(root / "model.safetensors"))
|
| 94 |
+
if any(v.dtype != mx.uint32 for k, v in weights.items() if k.removesuffix(".weight") in spec["linears"] or k == name + ".weight"):
|
| 95 |
+
raise ValueError("quantized weights must be packed uint32")
|
| 96 |
+
lm.load_weights(list(weights.items()), strict=True)
|
| 97 |
+
del weights
|
| 98 |
+
lm.eval()
|
| 99 |
+
mx.set_cache_limit(64 << 20)
|
| 100 |
+
mx.eval(lm.parameters())
|
| 101 |
+
mx.clear_cache()
|
| 102 |
+
return MLXDecisionModel(lm, pad_id, head_dim)
|
runtime/kev/plot.py
ADDED
|
@@ -0,0 +1,59 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Plot training loss from train logs and held-out accuracy vs baselines from eval.json.
|
| 2 |
+
|
| 3 |
+
Run: uv run python -m kev.plot --logs runs/train.log:kev-0.5b runs/train_holdout.log:holdout --eval runs/kev/eval.json --out docs/training.png
|
| 4 |
+
"""
|
| 5 |
+
import argparse, re
|
| 6 |
+
from .suite import read_json
|
| 7 |
+
import matplotlib
|
| 8 |
+
matplotlib.use("Agg")
|
| 9 |
+
import matplotlib.pyplot as plt
|
| 10 |
+
|
| 11 |
+
STEP_RE = re.compile(r"ep(\d+) step (\d+)/(\d+) loss ([\d.]+)")
|
| 12 |
+
ORDER = ["agnews", "agnews_yn", "banking77", "boolq", "mnli", "sst5", "yelp", "yelp_yn"]
|
| 13 |
+
LABEL = {"agnews": "AG News\nK=4", "agnews_yn": "AG News\nyes/no", "banking77": "Banking77\nK=77", "boolq": "BoolQ\nyes/no",
|
| 14 |
+
"mnli": "MNLI\nK=3", "sst5": "SST-5\nscore", "yelp": "Yelp\nscore", "yelp_yn": "Yelp\nyes/no"}
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
def read_log(path):
|
| 18 |
+
steps, losses, epoch_ends, prev_ep = [], [], [], 0
|
| 19 |
+
for m in STEP_RE.finditer(open(path, encoding="utf-8").read()):
|
| 20 |
+
ep, step, loss = int(m[1]), int(m[2]), float(m[4])
|
| 21 |
+
if ep != prev_ep: epoch_ends.append(steps[-1]); prev_ep = ep
|
| 22 |
+
steps.append(step); losses.append(loss)
|
| 23 |
+
return steps, losses, epoch_ends
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
def main():
|
| 27 |
+
ap = argparse.ArgumentParser()
|
| 28 |
+
ap.add_argument("--logs", nargs="+", default=["runs/train.log:kev-0.5b"], help="path:label pairs")
|
| 29 |
+
ap.add_argument("--eval", default="runs/kev/eval.json")
|
| 30 |
+
ap.add_argument("--out", default="docs/training.png")
|
| 31 |
+
a = ap.parse_args()
|
| 32 |
+
|
| 33 |
+
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 4.2), dpi=160, gridspec_kw={"width_ratios": [1.1, 1.3]})
|
| 34 |
+
|
| 35 |
+
for spec in a.logs:
|
| 36 |
+
path, label = spec.split(":") if ":" in spec else (spec, spec)
|
| 37 |
+
steps, losses, ends = read_log(path)
|
| 38 |
+
ax1.plot(steps, losses, lw=1.6, label=f"{label} (final {losses[-1]:.2f})")
|
| 39 |
+
for e in ends: ax1.axvline(e, color="0.8", lw=0.8, ls="--")
|
| 40 |
+
ax1.set_yscale("log"); ax1.set_xlabel("optimizer step (8 records each)"); ax1.set_ylabel("train loss (cross-entropy, log scale)")
|
| 41 |
+
ax1.set_title("training loss"); ax1.grid(alpha=0.25); ax1.legend(frameon=False)
|
| 42 |
+
|
| 43 |
+
e = read_json(a.eval)
|
| 44 |
+
kev, b0, b1 = e["accuracy_calibration"], e.get("baseline_zero_shot_base", {}), e.get("baseline_zero_shot_instruct", {})
|
| 45 |
+
xs = range(len(ORDER)); w = 0.27
|
| 46 |
+
ax2.bar([x - w for x in xs], [b0.get(s, {}).get("acc", 0) for s in ORDER], w, color="0.78", label="zero-shot base (letter logits)")
|
| 47 |
+
ax2.bar([x for x in xs], [b1.get(s, {}).get("acc", 0) for s in ORDER], w, color="0.55", label="zero-shot Instruct (letter logits)")
|
| 48 |
+
ax2.bar([x + w for x in xs], [kev[s]["acc"] for s in ORDER], w, color="#1f77b4", label="kev-0.5b (pointer readout)")
|
| 49 |
+
for x, s in zip(xs, ORDER):
|
| 50 |
+
ax2.text(x + w, kev[s]["acc"] + 0.012, f"ECE\n{kev[s]['ece']:.2f}", ha="center", va="bottom", fontsize=6.5, color="#1f77b4")
|
| 51 |
+
ax2.set_xticks(list(xs)); ax2.set_xticklabels([LABEL[s] for s in ORDER], fontsize=7.5)
|
| 52 |
+
ax2.set_ylim(0, 1.12); ax2.set_ylabel("held-out accuracy"); ax2.set_title(f"held-out accuracy vs zero-shot baselines (all: {kev['ALL']['acc']:.3f}, ECE {kev['ALL']['ece']:.3f})")
|
| 53 |
+
ax2.grid(axis="y", alpha=0.25); ax2.legend(frameon=False, fontsize=8, loc="upper center", bbox_to_anchor=(0.5, -0.16), ncol=3)
|
| 54 |
+
for ax in (ax1, ax2): ax.spines[["top", "right"]].set_visible(False)
|
| 55 |
+
fig.tight_layout(); fig.savefig(a.out, bbox_inches="tight"); print("wrote", a.out)
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
if __name__ == "__main__":
|
| 59 |
+
main()
|
runtime/kev/predictors.py
ADDED
|
@@ -0,0 +1,341 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Predictors: callables record -> {"probabilities": {qid: {key: p}}, "latency_ms", "input_tokens", ...} for kev.benchmark.
|
| 2 |
+
|
| 3 |
+
LocalPredictor scores a checkpoint in-process; RemotePredictor any TypeSafe System One-compatible endpoint; JevPredictor
|
| 4 |
+
Jev itself through the AI SDK worker in playground/scripts (budget-capped).
|
| 5 |
+
"""
|
| 6 |
+
import contextlib
|
| 7 |
+
import inspect
|
| 8 |
+
import json
|
| 9 |
+
import math
|
| 10 |
+
import os
|
| 11 |
+
import subprocess
|
| 12 |
+
import sys
|
| 13 |
+
import time
|
| 14 |
+
import urllib.error
|
| 15 |
+
import urllib.request
|
| 16 |
+
from datetime import datetime, timezone
|
| 17 |
+
from importlib.metadata import PackageNotFoundError, version
|
| 18 |
+
from pathlib import Path
|
| 19 |
+
|
| 20 |
+
import torch
|
| 21 |
+
from torch.nn.attention import SDPBackend, sdpa_kernel
|
| 22 |
+
|
| 23 |
+
from kev.api import question_keys
|
| 24 |
+
from kev.checkpoint import Checkpoint, LoadOptions
|
| 25 |
+
from kev.data import api_request, materialize
|
| 26 |
+
from kev.device import sync
|
| 27 |
+
from kev.model import ROW_PASS_TOKENS, ContextOverflow, rows_of
|
| 28 |
+
from kev.suite import CONTEXT
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
LONG_ROW_KERNELS = "efficient" # the label of a prediction (and its benchmark rows) that took the long-row kernels below
|
| 32 |
+
|
| 33 |
+
KERNEL_PACKAGES = ("torch", "transformers", "peft", "flash-linear-attention", "triton", "causal-conv1d", "mlx", "mlx-lm")
|
| 34 |
+
DELTANET_KERNELS = ("causal_conv1d_fn", "torch_chunk_gated_delta_rule") # the prefill hooks of transformers' GatedDeltaNet
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
@contextlib.contextmanager
|
| 38 |
+
def long_row_kernels():
|
| 39 |
+
"""The attention kernels a long row runs under (LocalPredictor): SDPA's flash and memory-efficient kernels, with the
|
| 40 |
+
math kernel as the fallback. transformers asks SDPA for grouped-query attention (`enable_gqa`) when a call has no
|
| 41 |
+
mask, which is the case for an unpadded state (kev.shared_prefix passes no mask so the state runs causal). Of the
|
| 42 |
+
kernels allowed here only flash and math take `enable_gqa`, and flash has no fp32. So an fp32 state pass used to fall
|
| 43 |
+
back to math, whose L x L scores do not fit: Kev-4B has 16 heads, so at 32k that is 16 x 32k^2 x 4 B = 69 GB per
|
| 44 |
+
layer. Inside this context an fp32 call repeats its keys and values per query head instead (transformers' own path
|
| 45 |
+
whenever a mask is given), and the memory-efficient kernel takes that: fp32 inputs and output, memory linear in L.
|
| 46 |
+
On sm80+ its fp32 products are cutlass's OpMultiplyAddFastF32 (three TF32 products per product, fp32 accuracy and
|
| 47 |
+
fp32 accumulation), not plain TF32. bf16 / fp16 calls keep `enable_gqa` and the flash kernel, so Kev-27B's long rows
|
| 48 |
+
run as before."""
|
| 49 |
+
from transformers.integrations import sdpa_attention
|
| 50 |
+
grouped = sdpa_attention.use_gqa_in_sdpa
|
| 51 |
+
sdpa_attention.use_gqa_in_sdpa = lambda attention_mask, key, value: key.dtype != torch.float32 and grouped(attention_mask, key, value)
|
| 52 |
+
try:
|
| 53 |
+
with sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]):
|
| 54 |
+
yield
|
| 55 |
+
finally:
|
| 56 |
+
sdpa_attention.use_gqa_in_sdpa = grouped
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def bound_implementation(fn):
|
| 60 |
+
"""The function transformers actually calls for one of its kernel hooks: `use_kernel_func_from_hub_with_fallback`
|
| 61 |
+
wraps the PyTorch reference and binds the package's kernel (fla, causal-conv1d) as `implementation` when it imports;
|
| 62 |
+
any other function (unwrapped, or patched in by a caller) is itself. -> "module.qualname"."""
|
| 63 |
+
impl = inspect.getclosurevars(fn).nonlocals.get("implementation", fn) if inspect.isfunction(fn) else fn
|
| 64 |
+
return f"{getattr(impl, '__module__', None)}.{getattr(impl, '__qualname__', type(impl).__name__)}"
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
def kernel_environment(model, device):
|
| 68 |
+
"""What a read's logits depend on besides kev's code and the checkpoint (runs/drift-v1/REPORT.md: #125's causal-conv1d
|
| 69 |
+
kernel moved Kev-27B v1's reads by up to 0.06 in p with no code change): the scoring packages' versions, the GPU, the
|
| 70 |
+
backbone dtype and, on a torch backbone, the attention implementation and (hybrid) the Gated DeltaNet layer's forward
|
| 71 |
+
(kev.fused_qwen35 replaces it) and the convolution and chunked delta rule transformers bound. Reads only what is
|
| 72 |
+
already loaded or imported. Two reads are comparable bit for bit only when this and the kev commit match. Recorded
|
| 73 |
+
in report.json."""
|
| 74 |
+
def installed(name):
|
| 75 |
+
try: return version(name)
|
| 76 |
+
except PackageNotFoundError: return None
|
| 77 |
+
packages = {name: installed(name) for name in KERNEL_PACKAGES}
|
| 78 |
+
packages["torch"] = torch.__version__ # with its build tag (2.8.0+cu128 on the CUDA wheels), which the metadata version drops
|
| 79 |
+
env = {"packages": packages, "device": str(device),
|
| 80 |
+
"gpu": torch.cuda.get_device_name(torch.device(device)) if str(device).startswith("cuda") else None,
|
| 81 |
+
"backend": model.backend, "dtype": model.dtype, "triton_f32_default": os.environ.get("TRITON_F32_DEFAULT"),
|
| 82 |
+
"attention": None, "deltanet": None}
|
| 83 |
+
if model.backend != "torch": return env # kev.mlx_model: mlx-lm's Metal kernels, named by the mlx / mlx-lm versions
|
| 84 |
+
env["attention"] = getattr(model.lm.config, "_attn_implementation", None)
|
| 85 |
+
layer = next((m for m in model.lm.modules() if type(m).__name__.endswith("GatedDeltaNet")), None) if model.hybrid else None
|
| 86 |
+
if layer is not None:
|
| 87 |
+
module = sys.modules[type(layer).__module__]
|
| 88 |
+
env["deltanet"] = {"forward": bound_implementation(getattr(layer.forward, "__func__", layer.forward)),
|
| 89 |
+
**{name: bound_implementation(getattr(module, name)) for name in DELTANET_KERNELS if hasattr(module, name)}}
|
| 90 |
+
return env
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
class LocalPredictor:
|
| 94 |
+
"""Scores a checkpoint in-process. Evaluation is fp32-exact in PyTorch (no TF32, no fused SDPA kernels on CUDA). On
|
| 95 |
+
CUDA a hybrid backbone's Gated DeltaNet layers run flash-linear-attention's Triton kernels, whose dots are TF32, and a
|
| 96 |
+
bf16 backbone is bf16 throughout, so reads repeat bit for bit only on the same kernel set (`self.environment`, from
|
| 97 |
+
kernel_environment; AGENTS.md "What fp32-exact guarantees"). One more exception: on CUDA with the torch backend, a
|
| 98 |
+
record whose longest row (state + one question) exceeds kev.model.ROW_PASS_TOKENS runs under SDPA's flash /
|
| 99 |
+
memory-efficient kernels (long_row_kernels), because the exact math kernel's L x L score matrix does not fit (~275 GB
|
| 100 |
+
per layer pass at 64k for Kev-4B's 16 heads). An fp32 backbone stays fp32 there: the memory-efficient kernel takes
|
| 101 |
+
fp32 (Kev-4B / 9B at 8k-16k against the math kernel: max |dp| 6e-4, no flips in 238 questions;
|
| 102 |
+
scripts/long_state_memory.py). Such a prediction carries `"kernels": LONG_ROW_KERNELS`, kev.benchmark copies it onto
|
| 103 |
+
the record's rows and counts them in report.json's `long_rows`; shorter records carry nothing, so their rows and
|
| 104 |
+
reports are unchanged. The efficient path is CUDA-only: on CPU / MPS attention stays eager and exact and still
|
| 105 |
+
materialises L x L. A long record on a hybrid torch backbone also runs its state once, its questions continuing from
|
| 106 |
+
it (kev.shared_prefix), instead of once per question; the MLX backend's forward already runs the state once."""
|
| 107 |
+
|
| 108 |
+
def __init__(self, run, device, opts=LoadOptions(), context=CONTEXT):
|
| 109 |
+
"""opts.temperature=None scores with the temperature the checkpoint carries; 1.0 scores raw logits. context: the
|
| 110 |
+
max_state / max_branch / max_packed a record must encode within (a suite manifest's `context`; the training
|
| 111 |
+
context by default, the serving limits for external suites frozen without admission)."""
|
| 112 |
+
if opts.temperature is not None and not (math.isfinite(opts.temperature) and opts.temperature > 0):
|
| 113 |
+
raise ValueError("temperature must be finite and positive")
|
| 114 |
+
checkpoint = self.checkpoint = Checkpoint(run)
|
| 115 |
+
self.run = checkpoint.path
|
| 116 |
+
if device == "cuda":
|
| 117 |
+
# evaluation is fp32-exact: TF32 (10-bit mantissa) moves probabilities by ~1e-3, the isolation gate's tolerance
|
| 118 |
+
torch.backends.cuda.matmul.allow_tf32 = False; torch.backends.cudnn.allow_tf32 = False
|
| 119 |
+
torch.backends.cuda.enable_flash_sdp(False); torch.backends.cuda.enable_mem_efficient_sdp(False)
|
| 120 |
+
self.tok, self.model = checkpoint.load(device, opts)
|
| 121 |
+
self.temperature = self.model.head.temperature
|
| 122 |
+
self.environment = kernel_environment(self.model, device)
|
| 123 |
+
self.device = device
|
| 124 |
+
self.context = context
|
| 125 |
+
|
| 126 |
+
@torch.no_grad()
|
| 127 |
+
def __call__(self, record):
|
| 128 |
+
enc = self.model.encode(self.tok, materialize(record), max_state=self.context["max_state"], max_branch=self.context["max_branch"], strict=True)
|
| 129 |
+
if len(enc["ids"]) > self.context["max_packed"]:
|
| 130 |
+
raise ContextOverflow(f"packed request exceeds the {self.context['max_packed']}-token limit")
|
| 131 |
+
sync(self.device)
|
| 132 |
+
start = time.perf_counter()
|
| 133 |
+
state, _, rows = rows_of(enc)
|
| 134 |
+
long = len(state) + max(len(r["ids"]) for r in rows) > ROW_PASS_TOKENS
|
| 135 |
+
torch_backend = self.model.backend == "torch"
|
| 136 |
+
efficient = long and torch_backend and self.device == "cuda"
|
| 137 |
+
with long_row_kernels() if efficient else contextlib.nullcontext():
|
| 138 |
+
if long and torch_backend and self.model.hybrid: logits = self.model.forward_batch([enc], shared_prefix=True)[0]
|
| 139 |
+
else: logits = self.model.forward(enc)
|
| 140 |
+
ps = [torch.softmax(z, -1).cpu() for z in logits]
|
| 141 |
+
zs = [z.float().cpu() for z in logits]
|
| 142 |
+
sync(self.device)
|
| 143 |
+
keys = {qid: question_keys(q["type"], q.get("criteria")) for qid, q in record["questions"].items()}
|
| 144 |
+
return {"probabilities": {qid: dict(zip(keys[qid], p.tolist())) for qid, p in zip(keys, ps)},
|
| 145 |
+
"logits": {qid: dict(zip(keys[qid], z.tolist())) for qid, z in zip(keys, zs)},
|
| 146 |
+
"inference_temperature": self.temperature,
|
| 147 |
+
"latency_ms": 1000 * (time.perf_counter() - start), "input_tokens": len(enc["ids"]), **({"kernels": LONG_ROW_KERNELS} if efficient else {})}
|
| 148 |
+
|
| 149 |
+
|
| 150 |
+
REFUSAL_STATUSES = (400, 413, 422) # an endpoint declining the request itself (kev.serve: 422 past its context; Jev: 400 / 413 / 422)
|
| 151 |
+
|
| 152 |
+
|
| 153 |
+
def http_detail(error):
|
| 154 |
+
"""The message of an HTTP error body: FastAPI's / TypeSafe's `detail` (or `error`) when it is JSON, else its text."""
|
| 155 |
+
try: text = error.read().decode("utf-8", "replace")
|
| 156 |
+
except Exception: return str(error.reason)
|
| 157 |
+
try: body = json.loads(text)
|
| 158 |
+
except ValueError: return text[:500] or str(error.reason)
|
| 159 |
+
return str(body.get("detail") or body.get("error") or body) if isinstance(body, dict) else str(body)
|
| 160 |
+
|
| 161 |
+
|
| 162 |
+
class RemotePredictor:
|
| 163 |
+
"""Score any TypeSafe System One-compatible endpoint (POST <base_url>/v1/systemone) on frozen records. Probabilities are
|
| 164 |
+
taken from the response as returned (renormalised by validate_distribution like every other predictor). Records the
|
| 165 |
+
server-reported model id so the manifest can pin what was scored. `concurrency` is how many requests kev.benchmark may
|
| 166 |
+
keep in flight at once (each call is independent: one request, its own retries); 1 scores sequentially. An endpoint
|
| 167 |
+
that declines the request itself (REFUSAL_STATUSES: kev.serve's 422 for a state past its context) raises
|
| 168 |
+
ContextOverflow at once, so kev.benchmark counts the record rejected under skip_overlong and stops otherwise; another
|
| 169 |
+
client error stops at once too; 408, 429, 5xx, connection errors and timeouts are tried `retries` times in all."""
|
| 170 |
+
|
| 171 |
+
def __init__(self, base_url, model="kev-latest", api_key="local", timeout=120, retries=3, concurrency=1):
|
| 172 |
+
if concurrency < 1:
|
| 173 |
+
raise ValueError("concurrency must be >= 1")
|
| 174 |
+
self.base_url, self.model, self.api_key, self.timeout, self.retries = base_url.rstrip("/"), model, api_key, timeout, retries
|
| 175 |
+
self.concurrency = concurrency
|
| 176 |
+
self.served_model = None
|
| 177 |
+
|
| 178 |
+
def __call__(self, record):
|
| 179 |
+
payload = json.dumps({**api_request(record), "model": self.model}).encode()
|
| 180 |
+
req = urllib.request.Request(f"{self.base_url}/v1/systemone", data=payload, method="POST",
|
| 181 |
+
headers={"content-type": "application/json", "authorization": f"Bearer {self.api_key}"})
|
| 182 |
+
last = None
|
| 183 |
+
for attempt in range(self.retries):
|
| 184 |
+
try:
|
| 185 |
+
start = time.perf_counter()
|
| 186 |
+
with urllib.request.urlopen(req, timeout=self.timeout) as resp:
|
| 187 |
+
body = json.loads(resp.read())
|
| 188 |
+
latency = 1000 * (time.perf_counter() - start)
|
| 189 |
+
break
|
| 190 |
+
except urllib.error.HTTPError as error:
|
| 191 |
+
if error.code in REFUSAL_STATUSES: # the request itself (kev.serve: a state or question over the serving context): a rejected record, not retried
|
| 192 |
+
raise ContextOverflow(f"remote endpoint refused the request (HTTP {error.code}): {http_detail(error)}") from None
|
| 193 |
+
if error.code not in (408, 429) and error.code < 500: # 401, 403, 404, ...: the same answer every time; 408, 429 and 5xx are retried
|
| 194 |
+
raise RuntimeError(f"remote endpoint answered HTTP {error.code}: {http_detail(error)}") from None
|
| 195 |
+
last = error
|
| 196 |
+
except Exception as error: # connection errors and timeouts: retried
|
| 197 |
+
last = error
|
| 198 |
+
if attempt + 1 < self.retries: time.sleep(2 ** attempt)
|
| 199 |
+
else:
|
| 200 |
+
raise RuntimeError(f"remote endpoint failed after {self.retries} attempts: {last}")
|
| 201 |
+
self.served_model = body.get("model", self.served_model)
|
| 202 |
+
probs = {}
|
| 203 |
+
for qid, q in record["questions"].items():
|
| 204 |
+
a = body["answers"][qid]
|
| 205 |
+
if q["type"] == "noul": probs[qid] = {"true": float(a["noul"]), "false": 1 - float(a["noul"])}
|
| 206 |
+
else: probs[qid] = {str(k): float(v) for k, v in a["probabilities"].items()}
|
| 207 |
+
return {"probabilities": probs, "latency_ms": latency, "input_tokens": (body.get("usage") or {}).get("input_tokens")}
|
| 208 |
+
|
| 209 |
+
|
| 210 |
+
PRICE_PER_MILLION = 0.042 # Jev list price per million input tokens, for the budget cap
|
| 211 |
+
|
| 212 |
+
|
| 213 |
+
class RotationAveraged:
|
| 214 |
+
"""Test-time order averaging (PLAN.md round 4, item 4.4): score the record under the first `rotations` cyclic
|
| 215 |
+
rotations of every Choice question's options (rotation r of a K-option question is r mod K; Noul and Score keep their
|
| 216 |
+
order, which is part of their meaning) and average the logits per option key. The geometric mean of softmax
|
| 217 |
+
probabilities is the softmax of the mean logits, so the returned probabilities and logits stay consistent; predictors
|
| 218 |
+
that return no logits are averaged in log-probability space. Costs `rotations` forward passes per record."""
|
| 219 |
+
|
| 220 |
+
def __init__(self, predictor, rotations):
|
| 221 |
+
if rotations < 2:
|
| 222 |
+
raise ValueError("rotation averaging needs at least 2 rotations")
|
| 223 |
+
self.predictor, self.rotations = predictor, rotations
|
| 224 |
+
self.temperature = getattr(predictor, "temperature", None)
|
| 225 |
+
self.concurrency = getattr(predictor, "concurrency", 1)
|
| 226 |
+
|
| 227 |
+
@staticmethod
|
| 228 |
+
def rotated(record, r):
|
| 229 |
+
def rotate(q):
|
| 230 |
+
if q["type"] != "choice": return q
|
| 231 |
+
keys = list(q["criteria"]); k = r % len(keys)
|
| 232 |
+
return {**q, "criteria": {key: q["criteria"][key] for key in keys[k:] + keys[:k]}}
|
| 233 |
+
return {**record, "questions": {qid: rotate(q) for qid, q in record["questions"].items()}}
|
| 234 |
+
|
| 235 |
+
def __call__(self, record):
|
| 236 |
+
widest = max((len(q["criteria"]) for q in record["questions"].values() if q["type"] == "choice"), default=1)
|
| 237 |
+
preds = [self.predictor(self.rotated(record, r)) for r in range(min(self.rotations, widest))]
|
| 238 |
+
field = "logits" if all("logits" in p for p in preds) else "probabilities"
|
| 239 |
+
out = {"probabilities": {}, "latency_ms": sum(p["latency_ms"] for p in preds), "rotations": len(preds)}
|
| 240 |
+
if field == "logits":
|
| 241 |
+
out["logits"] = {}; out["inference_temperature"] = preds[0].get("inference_temperature", 1.0)
|
| 242 |
+
for qid in record["questions"]:
|
| 243 |
+
keys = list(preds[0]["probabilities"][qid])
|
| 244 |
+
if field == "logits":
|
| 245 |
+
z = {k: sum(p["logits"][qid][k] for p in preds) / len(preds) for k in keys}
|
| 246 |
+
else:
|
| 247 |
+
z = {k: sum(math.log(max(p["probabilities"][qid][k], 1e-12)) for p in preds) / len(preds) for k in keys}
|
| 248 |
+
top = max(z.values()); e = {k: math.exp(v - top) for k, v in z.items()}; s = sum(e.values())
|
| 249 |
+
out["probabilities"][qid] = {k: v / s for k, v in e.items()}
|
| 250 |
+
if field == "logits": out["logits"][qid] = z
|
| 251 |
+
return out
|
| 252 |
+
|
| 253 |
+
|
| 254 |
+
class JevRefused(ContextOverflow):
|
| 255 |
+
"""Jev declined a request because of the input itself (REFUSAL_STATUSES, or a hosted-side error on an oversize request;
|
| 256 |
+
see JevPredictor). With count_refusals, kev.benchmark counts such a record as rejected (rejected.json) instead of stopping."""
|
| 257 |
+
|
| 258 |
+
|
| 259 |
+
# Jev's context in tokens (documented ~32k). Past it the AI Gateway answers with HTTP 400 *or* a 503 GatewayInternalServerError
|
| 260 |
+
# (evals/longdoc-v1: 212 x 400 and 28 x 503 on the 240 requests of ~57k-61k tokens; every request of <= ~31k tokens answered),
|
| 261 |
+
# so a hosted-side error on a request estimated past OVERSIZE x JEV_CONTEXT_TOKENS is the size, not an outage. The estimate is
|
| 262 |
+
# characters / 4 of the request JSON; the 1.5 margin keeps a ~30k-token legal text (~144k characters) out of it.
|
| 263 |
+
JEV_CONTEXT_TOKENS = 32768
|
| 264 |
+
OVERSIZE = 1.5
|
| 265 |
+
|
| 266 |
+
|
| 267 |
+
def oversize(request):
|
| 268 |
+
"""Whether a Jev request is well past Jev's context (estimated tokens = characters / 4)."""
|
| 269 |
+
return len(json.dumps(request, ensure_ascii=False)) / 4 > OVERSIZE * JEV_CONTEXT_TOKENS
|
| 270 |
+
|
| 271 |
+
|
| 272 |
+
class JevPredictor:
|
| 273 |
+
"""Jev through the AI SDK worker (playground/scripts/jev-evaluate.mjs); every call is counted against a token budget.
|
| 274 |
+
count_refusals: a request answered with a REFUSAL_STATUSES error, or with a hosted-side error (5xx, no status) while
|
| 275 |
+
oversize(), raises JevRefused at once (counted in accounting()["refusals"], not fatal, not retried); any other client error
|
| 276 |
+
(401, 403, 429, ...) still stops the read. attempts: tries per request on hosted-side failures of requests that are not
|
| 277 |
+
oversize, backing off 1, 2, 4, ... seconds up to 30 between them."""
|
| 278 |
+
def __init__(self, key, budget=0.1, max_calls=700, count_refusals=False, attempts=4):
|
| 279 |
+
self.budget, self.max_calls, self.count_refusals, self.attempts = budget, max_calls, count_refusals, attempts
|
| 280 |
+
self.calls, self.input_tokens, self.output_tokens, self.retries = 0, 0, 0, 0
|
| 281 |
+
self.refusals = {} # HTTP status -> count
|
| 282 |
+
self.started_at = datetime.now(timezone.utc).isoformat()
|
| 283 |
+
worker = Path(__file__).resolve().parents[1] / "playground/scripts/jev-evaluate.mjs"
|
| 284 |
+
self.process = subprocess.Popen(["node", str(worker)], stdin=subprocess.PIPE, stdout=subprocess.PIPE,
|
| 285 |
+
stderr=subprocess.DEVNULL, text=True, bufsize=1,
|
| 286 |
+
env={**os.environ, "AI_GATEWAY_API_KEY": key})
|
| 287 |
+
|
| 288 |
+
def close(self):
|
| 289 |
+
self.process.stdin.close()
|
| 290 |
+
try:
|
| 291 |
+
self.process.wait(timeout=5)
|
| 292 |
+
except subprocess.TimeoutExpired:
|
| 293 |
+
self.process.terminate()
|
| 294 |
+
self.process.wait(timeout=5)
|
| 295 |
+
|
| 296 |
+
def __call__(self, record):
|
| 297 |
+
if self.calls >= self.max_calls or (self.input_tokens + 65536) * PRICE_PER_MILLION / 1e6 > self.budget:
|
| 298 |
+
raise RuntimeError("Jev evaluation reached the request/token cost cap")
|
| 299 |
+
request = api_request(record)
|
| 300 |
+
for attempt in range(self.attempts):
|
| 301 |
+
self.process.stdin.write(json.dumps(request) + "\n")
|
| 302 |
+
self.process.stdin.flush()
|
| 303 |
+
line = self.process.stdout.readline()
|
| 304 |
+
if not line:
|
| 305 |
+
raise RuntimeError("Jev SDK worker exited without a response")
|
| 306 |
+
result = json.loads(line)
|
| 307 |
+
self.calls += 1
|
| 308 |
+
if "error" not in result:
|
| 309 |
+
break
|
| 310 |
+
status = result["error"]["status"]
|
| 311 |
+
# bounded retry for hosted-side failures only; client errors (4xx) are real and must surface
|
| 312 |
+
hosted = status is None or status >= 500
|
| 313 |
+
if self.count_refusals and (status in REFUSAL_STATUSES or hosted and oversize(request)):
|
| 314 |
+
kind = str(status) if status in REFUSAL_STATUSES else f"{status} (oversize)"
|
| 315 |
+
self.refusals[kind] = self.refusals.get(kind, 0) + 1
|
| 316 |
+
raise JevRefused(f"Jev refused the request: {result['error']['name']} (HTTP {kind})")
|
| 317 |
+
if attempt == self.attempts - 1 or (status is not None and status < 500):
|
| 318 |
+
raise RuntimeError(f"Jev request failed: {result['error']['name']} (HTTP {status})")
|
| 319 |
+
self.retries += 1
|
| 320 |
+
time.sleep(min(2 ** attempt, 30))
|
| 321 |
+
usage = result["usage"]
|
| 322 |
+
if usage.get("inputTokens") is None:
|
| 323 |
+
raise RuntimeError("Jev returned no input token usage; cannot account for cost")
|
| 324 |
+
self.input_tokens += usage["inputTokens"]
|
| 325 |
+
self.output_tokens += usage.get("outputTokens") or 0
|
| 326 |
+
probabilities = {}
|
| 327 |
+
for qid, q in record["questions"].items():
|
| 328 |
+
answer = result["answers"][qid]
|
| 329 |
+
if q["type"] == "noul":
|
| 330 |
+
probabilities[qid] = {"false": 1 - answer["probability"], "true": answer["probability"]}
|
| 331 |
+
else:
|
| 332 |
+
probabilities[qid] = answer["probabilities"]
|
| 333 |
+
return {**result, "probabilities": probabilities}
|
| 334 |
+
|
| 335 |
+
def accounting(self):
|
| 336 |
+
return {"model": "typesafe-ai/jev", "model_revision": "Gateway alias; provider revision not exposed by SDK result",
|
| 337 |
+
"started_at": self.started_at, "calls": self.calls, "input_tokens": self.input_tokens,
|
| 338 |
+
"output_tokens": self.output_tokens, "retries_after_5xx": self.retries, "listed_input_usd_per_million": PRICE_PER_MILLION,
|
| 339 |
+
**({"refusals": dict(sorted(self.refusals.items()))} if self.count_refusals else {}),
|
| 340 |
+
"estimated_usd": self.input_tokens * PRICE_PER_MILLION / 1e6,
|
| 341 |
+
"budget_usd": self.budget, "sdk": "ai@7.0.105", "zero_data_retention_requested": True}
|
runtime/kev/publish.py
ADDED
|
@@ -0,0 +1,106 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Publish a trained run to the Hugging Face Hub.
|
| 2 |
+
|
| 3 |
+
Repo naming: jaredpalmer/kev-<base size>, e.g. kev-0.5b (Qwen2.5-0.5B), kev-0.6b (Qwen3-0.6B), kev-4b, kev-8b.
|
| 4 |
+
The exact base checkpoint is recorded in the model card's `base_model` field and in head.pt.
|
| 5 |
+
|
| 6 |
+
uv run python -m kev.publish --run runs/kev --repo jaredpalmer/kev-0.5b
|
| 7 |
+
uv run python -m kev.publish --run runs/kev2 --repo jaredpalmer/kev-0.5b --message "v0.2: none-of-the-above fix"
|
| 8 |
+
|
| 9 |
+
Uploads: the adapter (or, for a full-weight run, config.json and every model*.safetensors shard with its index), head.pt,
|
| 10 |
+
tokenizer files, eval.json, training log (if found), and the model card (--card) as README.md with the repo id and run
|
| 11 |
+
name filled in. Requires `hf auth login`. With --private the repo is created private if missing and an existing repo
|
| 12 |
+
that is not private is refused (kev.mirror.ensure_private), so a candidate never lands in a public repo. Full-weight
|
| 13 |
+
shards are linked into the staging directory, not copied (modal_app.py::release_publish uploads a 27B from a CPU container).
|
| 14 |
+
|
| 15 |
+
--replace makes the target revision hold exactly this upload: every file the upload does not carry (other than
|
| 16 |
+
.gitattributes) is deleted in the same commit. A full-weight upload into a repo that still holds an adapter (or an adapter
|
| 17 |
+
upload into one that holds backbone shards) is refused without it, because the loader rule (kev.checkpoint) would pick
|
| 18 |
+
the stale layout.
|
| 19 |
+
"""
|
| 20 |
+
import argparse, os, re, shutil, tempfile
|
| 21 |
+
from pathlib import Path
|
| 22 |
+
from huggingface_hub import HfApi
|
| 23 |
+
from .checkpoint import Checkpoint
|
| 24 |
+
from .mirror import ensure_private
|
| 25 |
+
from .suite import read_json, write_json
|
| 26 |
+
|
| 27 |
+
FILES = ["head.pt", "tokenizer.json", "tokenizer_config.json", "chat_template.jinja",
|
| 28 |
+
"vocab.json", "merges.txt", "added_tokens.json", "special_tokens_map.json", "eval.json"]
|
| 29 |
+
WEIGHTS = {False: ["adapter_config.json", "adapter_model.safetensors"], True: ["config.json", "model.safetensors.index.json"]} # by Checkpoint.full
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
def stale_layout(api, repo, revision, full):
|
| 33 |
+
"""Files of the other checkpoint layout on the target revision (adapter files for a full-weight upload, backbone shards
|
| 34 |
+
for an adapter upload), sorted; [] for a missing branch or an empty repo."""
|
| 35 |
+
from huggingface_hub.utils import RevisionNotFoundError
|
| 36 |
+
try: files = api.list_repo_files(repo, repo_type="model", revision=revision)
|
| 37 |
+
except RevisionNotFoundError: return []
|
| 38 |
+
other = (lambda f: f in WEIGHTS[False]) if full else (lambda f: re.fullmatch(r"model.*\.safetensors", f) or f == "model.safetensors.index.json")
|
| 39 |
+
return sorted(f for f in files if other(f))
|
| 40 |
+
|
| 41 |
+
|
| 42 |
+
def main():
|
| 43 |
+
ap = argparse.ArgumentParser()
|
| 44 |
+
ap.add_argument("--run", required=True)
|
| 45 |
+
ap.add_argument("--repo", required=True, help="e.g. jaredpalmer/kev-0.5b")
|
| 46 |
+
ap.add_argument("--message", default=None)
|
| 47 |
+
ap.add_argument("--card", required=True, help="model card markdown, e.g. docs/model-cards/kev-4b.md")
|
| 48 |
+
ap.add_argument("--private", action="store_true")
|
| 49 |
+
ap.add_argument("--tag", help="create this Hub tag on the uploaded commit (versioned release, e.g. v0.2)")
|
| 50 |
+
ap.add_argument("--revision", help="upload to this branch instead of main (created if missing); for candidates that must not replace the released weights")
|
| 51 |
+
ap.add_argument("--replace", action="store_true", help="delete every file of the target revision this upload does not carry, in the same commit (tag the previous version first)")
|
| 52 |
+
a = ap.parse_args()
|
| 53 |
+
|
| 54 |
+
checkpoint, run_name = Checkpoint(a.run), os.path.basename(a.run.rstrip("/"))
|
| 55 |
+
base, full = checkpoint.meta.base, checkpoint.full
|
| 56 |
+
api = HfApi()
|
| 57 |
+
if a.private: ensure_private(api, a.repo) # creates a missing repo private; refuses one that exists public
|
| 58 |
+
else: api.create_repo(a.repo, repo_type="model", exist_ok=True, private=False)
|
| 59 |
+
stale = stale_layout(api, a.repo, a.revision, full)
|
| 60 |
+
if stale and not a.replace:
|
| 61 |
+
raise SystemExit(f"{a.repo}{'@' + a.revision if a.revision else ''} holds {stale}, the other checkpoint layout; the loader would pick it. "
|
| 62 |
+
"Tag the previous version, then publish with --replace to delete it in the same commit")
|
| 63 |
+
|
| 64 |
+
with tempfile.TemporaryDirectory() as tmp:
|
| 65 |
+
shards = {p.name for p in checkpoint.shards()} if full else set()
|
| 66 |
+
for f in FILES + WEIGHTS[full] + sorted(shards):
|
| 67 |
+
src = f"{a.run}/{f}"
|
| 68 |
+
if not os.path.exists(src): print(f"skip {f} (not found)")
|
| 69 |
+
elif f in shards: os.symlink(os.path.abspath(src), f"{tmp}/{f}") # read-only and ~5 GB each: linked, not copied (upload_folder follows links)
|
| 70 |
+
else: shutil.copy(src, tmp)
|
| 71 |
+
if not full: # runs before task_type was set saved null; the Hub warns about it and PEFT treats both the same for a bare backbone
|
| 72 |
+
cfg_path = f"{tmp}/adapter_config.json"; cfg = read_json(cfg_path)
|
| 73 |
+
if not cfg.get("task_type"): cfg["task_type"] = "FEATURE_EXTRACTION"; write_json(cfg_path, cfg)
|
| 74 |
+
if os.path.exists(f"runs/logs/train_{run_name}.log"): # standalone runs keep their log there (see .gitignore)
|
| 75 |
+
shutil.copy(f"runs/logs/train_{run_name}.log", f"{tmp}/train.log")
|
| 76 |
+
# research trials: runs/<study>/<trial>/checkpoint -> ship the trial's result, provenance and training log too
|
| 77 |
+
trial = os.path.dirname(a.run.rstrip("/")) if run_name == "checkpoint" else None
|
| 78 |
+
if trial:
|
| 79 |
+
run_name = os.path.relpath(trial, "runs")
|
| 80 |
+
for f in ("result.json", "provenance.json", "train.log", "training_config.json", "training_metrics.json", "interpolation.json"):
|
| 81 |
+
for src in (f"{trial}/{f}", f"{a.run}/{f}"):
|
| 82 |
+
if os.path.exists(src): shutil.copy(src, f"{tmp}/{f}"); break
|
| 83 |
+
|
| 84 |
+
# the card's prose names the Hub repo and trial itself; only the frontmatter is filled from the checkpoint
|
| 85 |
+
card = re.sub(r"^base_model: .*$", f"base_model: {base}", Path(a.card).read_text(encoding="utf-8"), flags=re.M)
|
| 86 |
+
if "base_model_relation:" not in card: card = card.replace(f"base_model: {base}", f"base_model: {base}\nbase_model_relation: {'finetune' if full else 'adapter'}")
|
| 87 |
+
Path(tmp, "README.md").write_text(card, encoding="utf-8")
|
| 88 |
+
|
| 89 |
+
ev = read_json(f"{tmp}/eval.json") if os.path.exists(f"{tmp}/eval.json") else {}
|
| 90 |
+
acc = ev.get("accuracy_calibration", {}).get("ALL", {})
|
| 91 |
+
if os.path.exists(f"{tmp}/result.json"):
|
| 92 |
+
clean = read_json(f"{tmp}/result.json").get("clean", {}) # research trials; other result files (kev-finetune runs) just skip the figures
|
| 93 |
+
acc = {"acc": clean.get("acc", float("nan")), "ece": clean.get("ece", float("nan"))}
|
| 94 |
+
msg = a.message or f"Upload {run_name} (base {base}; acc {acc.get('acc', float('nan')):.3f}, ECE {acc.get('ece', float('nan')):.3f})"
|
| 95 |
+
if a.revision: api.create_branch(a.repo, branch=a.revision, repo_type="model", exist_ok=True)
|
| 96 |
+
# delete_patterns="*": remote files not in this upload are deleted in the same commit (huggingface_hub keeps .gitattributes and every file being uploaded)
|
| 97 |
+
info = api.upload_folder(folder_path=tmp, repo_id=a.repo, repo_type="model", commit_message=msg, revision=a.revision,
|
| 98 |
+
delete_patterns="*" if a.replace else None)
|
| 99 |
+
print(info)
|
| 100 |
+
if a.tag:
|
| 101 |
+
api.create_tag(a.repo, tag=a.tag, repo_type="model", tag_message=msg, exist_ok=False, revision=a.revision)
|
| 102 |
+
print(f"tagged {a.repo}@{a.tag}")
|
| 103 |
+
|
| 104 |
+
|
| 105 |
+
if __name__ == "__main__":
|
| 106 |
+
main()
|
runtime/kev/rounds.py
ADDED
|
@@ -0,0 +1,1364 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Registered research rounds as data: one spec per round, one engine for every stage.
|
| 2 |
+
|
| 3 |
+
A round (PLAN.md, rounds 5-20) always has the same shape. Arms are config-only trials (or checkpoints made from them,
|
| 4 |
+
round 20's interpolations) compared with a parent checkpoint;
|
| 5 |
+
each arm and its parent are read on a list of suites; a rule compares every arm with its parent on paired,
|
| 6 |
+
record-clustered bootstraps (primaries with a lower bound, guards with thresholds, pooled panels); the passing arm with
|
| 7 |
+
the best score is the size's candidate; confirmation stages (test partitions, the locked read) are read once. The spec
|
| 8 |
+
(`experiments/rounds/r<N>.json`, schema below) is committed before any training or read, like the PLAN registration it
|
| 9 |
+
encodes.
|
| 10 |
+
|
| 11 |
+
uv run python -m kev.rounds validate experiments/rounds/r18.json # well-formed, every parent read exists, plans validate
|
| 12 |
+
uv run python -m kev.rounds launch experiments/rounds/r18.json # modal_app.py::study for each study (after `modal deploy`)
|
| 13 |
+
uv run python -m kev.rounds watch experiments/rounds/r18.json # pull + read each finished trial once, then the read-out
|
| 14 |
+
uv run python -m kev.rounds launch-reads experiments/rounds/r18.json [--arms a,b] [--stage tests --arm a]
|
| 15 |
+
uv run python -m kev.rounds readout experiments/rounds/r18.json # -> runs/r18-readout/round18.json + a table
|
| 16 |
+
uv run python -m kev.rounds confirm experiments/rounds/r18.json --stage tests [--arm a] # -> runs/r18-verdict/<size>-<stage>.json
|
| 17 |
+
|
| 18 |
+
Every comparison is served-vs-served: each side's rows are served at the temperature fitted on its own trial's
|
| 19 |
+
development rows (`kev.metrics.served`), or, for the arms of a round that registers a `temperature` pool, on that pool;
|
| 20 |
+
unknowable records are scored only by `unknowable_report`, and every delta is `kev.metrics.paired_bootstrap` (2,000
|
| 21 |
+
resamples, seed 0, micro), the registered read since round 5.
|
| 22 |
+
|
| 23 |
+
Spec (paths are relative to the repo root; templates take {round}, {arm}, {size}, {tag}):
|
| 24 |
+
round, registered the round number and where its registration lives
|
| 25 |
+
archive a recorded round: the git tag holding the plans, runs and data it names (validate lists
|
| 26 |
+
what this checkout lacks instead of failing; launch refuses a record)
|
| 27 |
+
gpu, app Modal GPU and KEV_APP_NAME for launches (optional)
|
| 28 |
+
studies {name: {plan, suite, transfer, gpu, timeout, budget}} one modal_app.py::study call each (budget >= its admission bound)
|
| 29 |
+
reads {tag: {suite, flags?} | {entrypoint: "locked_test", decision}} what each read tag scores (entrypoint: ENTRYPOINTS)
|
| 30 |
+
read_timeout {size: seconds} overrides modal_app.READ_TIMEOUTS for one size (a 27B's fp32 reads)
|
| 31 |
+
locked_args {size: [args]} extra modal_app.py::locked_test switches for one size (a 27B's GPU memory)
|
| 32 |
+
parents {name: {trial, checkpoint?, reads: {tag: dir}}} checkpoint (Hub id[@rev]) only when /<trial>/checkpoint is not on the volume
|
| 33 |
+
arms {name: {trial?, checkpoint?, parent, reads?, select?, transfer_read?, trained_on?}} name = "<size>-<label>"; reads
|
| 34 |
+
default to arm_reads; select false = reported, never the candidate (attribution arms); an
|
| 35 |
+
arm without a trial (an interpolated checkpoint) names its /runs/... checkpoint, needs the
|
| 36 |
+
round's `temperature` pool (it has no development rows) and, if the rule reads
|
| 37 |
+
"transfer", a transfer_read; trained_on [suite dirs] names what such a checkpoint was
|
| 38 |
+
trained on (default: covered by the round's trial arms, e.g. the trials it interpolates)
|
| 39 |
+
arm_reads template of an arm's read directory (default "runs/r{round}-{arm}-{tag}")
|
| 40 |
+
transfer_read a read tag (round level, or per arm, which wins; null = the trial's own) whose rows are an arm's "transfer"
|
| 41 |
+
rows instead of its trial's in-trial transfer read (a checkpoint without a trial has none)
|
| 42 |
+
temperature {reads: [tags], sources?: {tag: [source]}, exclude_reads?: [tags]} every ARM (parents keep their trial's
|
| 43 |
+
development rows) is served at the temperature fitted, as everywhere (kev.metrics.served),
|
| 44 |
+
on the knowable rows of `reads` pooled (a read listed in `sources` gives only those
|
| 45 |
+
sources' rows), minus every record whose id is in the `exclude_reads` rows; the arm's own
|
| 46 |
+
rule-stage reads are used at every stage; no pooled tag may sit in a panel that a
|
| 47 |
+
temperature-dependent criterion reads (see pool_problems), and no pooled read may share
|
| 48 |
+
data with any arm's training data (pool_training_problems: round 19's failure mode)
|
| 49 |
+
drop_ids record ids dropped on both sides of every comparison (devtools-v1's duplicated ids)
|
| 50 |
+
rule {panels, unknowable?, criteria, rank}
|
| 51 |
+
confirm {stage: {candidate_reads, parent_reads?, panels, criteria}}
|
| 52 |
+
A panel is {reads: [tags], metrics: [bootstrapped], report: [value only], source?: filter, versus?: {name: dir}, by_length?,
|
| 53 |
+
exclude_sources?, exclude_tasks?, exclude_file?, optional?}; the tag "transfer" is the trial's own in-trial transfer read (or
|
| 54 |
+
the arm's transfer_read). source keeps one source (or a list of them); exclude_sources / exclude_tasks drop rows by source /
|
| 55 |
+
task, and exclude_file ({path, sha256}: a JSON object of `ids` (question or record ids), `sources` and `tasks` lists, or a
|
| 56 |
+
bare list of any of them) drops what it lists, checked against its registered hash, so a private list (round 24's
|
| 57 |
+
tasksource-heldout families) enters a public spec only as a path and a sha256 (panel_filter; the same rows leave both
|
| 58 |
+
sides and any `versus` reference; the panel records how many questions they removed as `excluded`). optional: true marks
|
| 59 |
+
a report-only panel no criterion or rank may read; its absent reads leave the comparison complete (`missing_optional`).
|
| 60 |
+
The `temperature` pool may carry ci: {level, samples, seed}, which adds each arm's bootstrap interval of its pooled
|
| 61 |
+
temperature to the read-out (`temperature_ci`, report only). by_length (true, or {edges?:
|
| 62 |
+
[tokens], tokenizer?: [name, revision]}) adds acc / ece / brier / confident_error_rate per state-token bucket
|
| 63 |
+
(kev.metrics.calibration_by_length: "<metric>_<bucket>" entries, e.g. ece_16k_plus, value only) with state tokens counted
|
| 64 |
+
from the reads' suite records. A criterion is {left, op, right, plus?} where left is a path
|
| 65 |
+
("<panel>.<metric>.<candidate|parent|delta|lower|upper>", "<panel>.<metric>_<bucket>.<candidate|parent>",
|
| 66 |
+
"unknowable.<candidate|parent>") or a list [a, b] meaning a - b, and right is a number or a path (plus is added to it).
|
| 67 |
+
rank is a list of {by: path | [paths summed], order}.
|
| 68 |
+
|
| 69 |
+
Every arm's read-out records where its temperature came from (`temperature_source`: the pool with its reads and question
|
| 70 |
+
count, or its trial's development rows with their suite) and its parent's (`parent_temperature_source`: its trial's
|
| 71 |
+
development rows, plus `shipped`, the head.pt temperature when it is here), and the printed table warns when an arm's rows
|
| 72 |
+
are a training corpus's (in distribution: round 19's failure mode). From round 21 (POOL_REQUIRED_FROM) a rule or
|
| 73 |
+
confirmation with a temperature-dependent criterion must register a `temperature` pool: validate and launch refuse it
|
| 74 |
+
otherwise; earlier rounds get a warning, so their recorded specs keep validating.
|
| 75 |
+
|
| 76 |
+
Removed suites (kev.suite.REMOVED_SUITES: evals/external/scienthoon-v1, removed 2026-09-27, last read by round 22;
|
| 77 |
+
evals/external/wanli-v2 and typesafe-v1, removed 2026-09-30, last read by round 26; wanli-v1, last read by round 5). A round
|
| 78 |
+
up to the suite's `last_round` that reads one validates with the read listed under `archived`; its read-out and
|
| 79 |
+
confirmations are computed from the committed rows as before (they need rows, not the suite). A later round that names
|
| 80 |
+
one is refused with the reason, and no read of one is launched again (read_commands skips it; kev.suite.load_split refuses it).
|
| 81 |
+
"""
|
| 82 |
+
import argparse
|
| 83 |
+
import functools
|
| 84 |
+
import operator
|
| 85 |
+
import os
|
| 86 |
+
import socket
|
| 87 |
+
import subprocess
|
| 88 |
+
import sys
|
| 89 |
+
import time
|
| 90 |
+
from pathlib import Path
|
| 91 |
+
from typing import NamedTuple
|
| 92 |
+
|
| 93 |
+
from kev.metrics import (LENGTH_EDGES, LENGTH_METRICS, calibration_by_length, length_buckets, metrics, paired_bootstrap, raw_row, recorded,
|
| 94 |
+
scored_rows, served, served_at, tempered_row, unknowable_report)
|
| 95 |
+
from kev.suite import ADMISSION_TOKENIZER, digest, file_lock, load_split, read_json, read_manifest, removed_suite, suite_key, write_json
|
| 96 |
+
|
| 97 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 98 |
+
SAMPLES = 2000 # the registered resample count since round 5
|
| 99 |
+
OPS = {">": operator.gt, ">=": operator.ge, "<": operator.lt, "<=": operator.le}
|
| 100 |
+
FIELDS = ("candidate", "parent", "delta", "lower", "upper")
|
| 101 |
+
TRANSFER = "transfer" # the in-trial transfer read (runs/<study>/<trial>/transfer)
|
| 102 |
+
STAGGER = 60 # seconds between `modal run` launches: Modal refuses more than ~3 app creations a minute
|
| 103 |
+
ENTRYPOINTS = ("benchmarks", "locked_test") # the modal_app.py entrypoints a read can go through
|
| 104 |
+
IN_FLIGHT = 7200 # seconds an arm's launched reads are presumed running (the longest suite timeout in modal_app.READ_TIMEOUTS)
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
# --- rows ------------------------------------------------------------------------------------------------------------
|
| 108 |
+
|
| 109 |
+
def paired(candidate, reference, metric):
|
| 110 |
+
"""The registered paired read: record-clustered bootstrap of candidate - reference, micro, 2,000 resamples, seed 0."""
|
| 111 |
+
x = paired_bootstrap(candidate, reference, samples=SAMPLES, seed=0, metric=metric, aggregation="micro")
|
| 112 |
+
return {"delta": x[f"micro_{metric}_delta"], "ci95": x["ci95"]}
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
def temperature(trial, root=ROOT):
|
| 116 |
+
"""The temperature fitted on a trial's own development rows (decision-v7's, or its suite's; how every checkpoint here is
|
| 117 |
+
served unless the round registers a `temperature` pool for its arms)."""
|
| 118 |
+
return served(read_json(Path(root) / trial / "development/rows.json"), [])[0]
|
| 119 |
+
|
| 120 |
+
|
| 121 |
+
def served_clean(rows, t):
|
| 122 |
+
"""Every clean row served at t, unknowable records included (what unknowable_report scores; served_at drops them)."""
|
| 123 |
+
return [tempered_row(raw_row(recorded(r)), t) for r in rows if r["variant"] == "clean"]
|
| 124 |
+
|
| 125 |
+
|
| 126 |
+
class Pool(NamedTuple):
|
| 127 |
+
"""Where a pooled temperature is fitted (spec `temperature`, resolved to one arm's read directories)."""
|
| 128 |
+
reads: list # directories whose rows are pooled
|
| 129 |
+
sources: dict # {directory: [source, ...]}: the only sources pooled from that read (absent = every source)
|
| 130 |
+
exclude: list # directories whose record ids are removed from the pool
|
| 131 |
+
ci: dict = None # {level, samples, seed}: also bootstrap the fitted temperature's interval (report only; round 24)
|
| 132 |
+
|
| 133 |
+
|
| 134 |
+
def select_rows(reads, exclude=()):
|
| 135 |
+
"""(pooled rows, rows before the exclusion): every rows.json of reads, as (path, sources or None), concatenated with
|
| 136 |
+
each read limited to its sources, minus every record whose id appears in an exclude rows.json. The one row selection
|
| 137 |
+
of a pooled temperature, here and in scripts/calibrate_checkpoint.py (which writes it into a checkpoint)."""
|
| 138 |
+
excluded = {r["id"] for path in exclude for r in read_json(path)}
|
| 139 |
+
rows = [r for path, sources in reads for r in read_json(path) if sources is None or r["source"] in sources]
|
| 140 |
+
return [r for r in rows if r["id"] not in excluded], rows
|
| 141 |
+
|
| 142 |
+
|
| 143 |
+
def pooled_temperature(pool, root=ROOT):
|
| 144 |
+
"""(temperature, fit report): the temperature() objective (kev.metrics.served: knowable clean rows, TEMPERATURE_FIT)
|
| 145 |
+
on select_rows of the pool."""
|
| 146 |
+
kept, rows = select_rows([(Path(root) / d / "rows.json", pool.sources.get(d)) for d in pool.reads], [Path(root) / d / "rows.json" for d in pool.exclude])
|
| 147 |
+
return served(kept, [])[0], {"reads": pool.reads, **({"sources": pool.sources} if pool.sources else {}), "exclude_reads": pool.exclude,
|
| 148 |
+
"questions": len(scored_rows(kept)), "excluded_questions": len(scored_rows(rows)) - len(scored_rows(kept))}
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
def pooled_temperature_ci(pool, root=ROOT):
|
| 152 |
+
"""The pooled temperature's bootstrap interval (report only; the audit of 2026-09-27: the served T of a 648-question pool
|
| 153 |
+
moves by ~±0.15, enough to move an ECE criterion by more than its margin): the same fit (TEMPERATURE_FIT's grid and
|
| 154 |
+
objective, knowable clean rows of select_rows) on `samples` resamples of the pool's (source, group) clusters within each
|
| 155 |
+
source (kev.metrics.cluster_resamples, seed), percentile interval at `level`. Each question's NLL on the grid is computed
|
| 156 |
+
once; a resample's fit is the grid point with the least mean NLL over its questions."""
|
| 157 |
+
import numpy as np
|
| 158 |
+
from kev.metrics import TEMPERATURE_FIT, cluster_resamples, nll_at_temperature
|
| 159 |
+
kept, _ = select_rows([(Path(root) / d / "rows.json", pool.sources.get(d)) for d in pool.reads], [Path(root) / d / "rows.json" for d in pool.exclude])
|
| 160 |
+
raw = served_at(kept, 1.0)
|
| 161 |
+
grid = np.exp(np.linspace(np.log(0.25), np.log(4), TEMPERATURE_FIT["points"]))
|
| 162 |
+
nll = np.asarray([[nll_at_temperature(r, float(t)) for t in grid] for r in raw])
|
| 163 |
+
level, samples, seed = pool.ci.get("level", 0.9), pool.ci.get("samples", SAMPLES), pool.ci.get("seed", 0)
|
| 164 |
+
fits = np.asarray([grid[int(np.argmin(nll[idx].mean(axis=0)))] for idx in cluster_resamples(raw, samples, seed)])
|
| 165 |
+
lower, upper = np.quantile(fits, [(1 - level) / 2, 1 - (1 - level) / 2])
|
| 166 |
+
return {"level": level, "samples": samples, "seed": seed, "lower": float(lower), "upper": float(upper), "questions": len(raw),
|
| 167 |
+
"resampled": "(source, group) clusters within each source (kev.metrics.cluster_resamples)"}
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
class Side:
|
| 171 |
+
"""One checkpoint of a comparison: its trial, where each read tag's rows live and its served temperature (fitted on
|
| 172 |
+
the trial's development rows, or on `pool` when the round registers one for its arms)."""
|
| 173 |
+
|
| 174 |
+
def __init__(self, trial, dirs, root=ROOT, drop=(), pool=None, checkpoint=None, suite=None, shipped=None):
|
| 175 |
+
self.trial, self.dirs, self.root, self.drop, self.pool, self.checkpoint = trial, dirs, Path(root), set(drop), pool, checkpoint
|
| 176 |
+
self.suite = suite # the suite the trial trained and was scored on (its development rows'), when known
|
| 177 |
+
self.shipped = shipped # a parent's: () -> the temperature its checkpoint ships (head.pt), or None when not here
|
| 178 |
+
self._t, self._served, self.fit = None, {}, None
|
| 179 |
+
|
| 180 |
+
def rows_path(self, tag):
|
| 181 |
+
return self.root / self.dirs[tag] / "rows.json"
|
| 182 |
+
|
| 183 |
+
def has(self, tag):
|
| 184 |
+
return tag in self.dirs and self.rows_path(tag).exists()
|
| 185 |
+
|
| 186 |
+
def absent_fit(self):
|
| 187 |
+
"""The rows the temperature needs that are not here (pool directories, or the trial's development rows)."""
|
| 188 |
+
needed = [*self.pool.reads, *self.pool.exclude] if self.pool else [f"{self.trial}/development"]
|
| 189 |
+
return [d for d in needed if not (self.root / d / "rows.json").exists()]
|
| 190 |
+
|
| 191 |
+
@property
|
| 192 |
+
def t(self):
|
| 193 |
+
if self._t is None:
|
| 194 |
+
if self.pool: self._t, self.fit = pooled_temperature(self.pool, self.root)
|
| 195 |
+
else: self._t = temperature(self.trial, self.root)
|
| 196 |
+
return self._t
|
| 197 |
+
|
| 198 |
+
def temperature_source(self):
|
| 199 |
+
"""Where the served temperature was fitted (read-out provenance): the pool (its reads and knowable question count),
|
| 200 |
+
or the trial's development rows, their suite and whether that suite is a training corpus (then the temperature is
|
| 201 |
+
in distribution: round 19's failure mode)."""
|
| 202 |
+
if self.pool: return {"kind": "pool", "reads": self.pool.reads, "questions": self.fit["questions"] if self.fit else None}
|
| 203 |
+
m = suite_manifest(self.suite) if self.suite else None
|
| 204 |
+
return {"kind": "trial development rows", "rows": f"{self.trial}/development", "suite": self.suite,
|
| 205 |
+
"training_corpus": bool(trainable_sources(m)) if m is not None else None, **({"shipped": self.shipped()} if self.shipped else {})}
|
| 206 |
+
|
| 207 |
+
def served(self, tag):
|
| 208 |
+
if tag not in self._served: self._served[tag] = served_at(read_json(self.rows_path(tag)), self.t)
|
| 209 |
+
return self._served[tag]
|
| 210 |
+
|
| 211 |
+
def panel(self, spec):
|
| 212 |
+
"""A panel's served, knowable rows, reads concatenated in the order listed, minus drop_ids and whatever the
|
| 213 |
+
panel's filters leave out (panel_filter: `source`, `exclude_sources`, `exclude_tasks`, `exclude_file`)."""
|
| 214 |
+
keep = panel_filter(spec, self.root)
|
| 215 |
+
return [r for tag in spec["reads"] for r in self.served(tag) if r["id"] not in self.drop and keep(r)]
|
| 216 |
+
|
| 217 |
+
def unknowable_share(self, tag):
|
| 218 |
+
return unknowable_report(served_clean(read_json(self.rows_path(tag)), self.t))["share_at_0_9"] if self.has(tag) else None
|
| 219 |
+
|
| 220 |
+
|
| 221 |
+
# --- panel filters ---------------------------------------------------------------------------------------------------
|
| 222 |
+
|
| 223 |
+
EXCLUDE_KEYS = ("exclude_sources", "exclude_tasks", "exclude_file") # round 24: the audit's per-suite fixes, applied to both sides
|
| 224 |
+
EXCLUDE_FILE_KEYS = ("ids", "sources", "tasks") # what an exclude file may list (other keys, e.g. a salt, are ignored)
|
| 225 |
+
|
| 226 |
+
|
| 227 |
+
def exclude_file_path(entry, root=ROOT):
|
| 228 |
+
return Path(root) / entry["path"]
|
| 229 |
+
|
| 230 |
+
|
| 231 |
+
@functools.cache
|
| 232 |
+
def _exclusion(path, sha256):
|
| 233 |
+
if digest(path) != sha256:
|
| 234 |
+
raise ValueError(f"exclude_file {path} does not match the sha256 the spec registers ({sha256[:12]}...): a changed exclusion list is a new spec")
|
| 235 |
+
data = read_json(path)
|
| 236 |
+
if isinstance(data, list): return {k: frozenset(data) for k in EXCLUDE_FILE_KEYS} # a bare list: ids, sources or tasks alike
|
| 237 |
+
return {k: frozenset(data.get(k, ())) for k in EXCLUDE_FILE_KEYS}
|
| 238 |
+
|
| 239 |
+
|
| 240 |
+
def exclusion(entry, root=ROOT):
|
| 241 |
+
"""{ids, sources, tasks} of a panel's exclude_file ({path, sha256}: a JSON object with any of those lists, or a bare list
|
| 242 |
+
matched against all three), checked against its registered sha256. The file may be private (never in git: e.g. round 24's
|
| 243 |
+
tasksource-heldout families, restored with scripts/private_rows.py); the spec carries only its path and hash."""
|
| 244 |
+
return _exclusion(exclude_file_path(entry, root), entry["sha256"])
|
| 245 |
+
|
| 246 |
+
|
| 247 |
+
def panel_filter(panel, root=ROOT):
|
| 248 |
+
"""row -> kept? for a panel: `source` (a name, or a list of names) keeps only those sources; `exclude_sources` and
|
| 249 |
+
`exclude_tasks` drop rows by source and task; `exclude_file` drops rows whose id, record (group), source or task it lists.
|
| 250 |
+
The same filter serves both sides of a comparison (and a `versus` reference), so paired deltas stay paired."""
|
| 251 |
+
source = panel.get("source")
|
| 252 |
+
only = None if source is None else {source} if isinstance(source, str) else set(source)
|
| 253 |
+
sources, tasks = set(panel.get("exclude_sources", ())), set(panel.get("exclude_tasks", ()))
|
| 254 |
+
listed = exclusion(panel["exclude_file"], root) if panel.get("exclude_file") else None
|
| 255 |
+
def keep(r):
|
| 256 |
+
if only is not None and r["source"] not in only: return False
|
| 257 |
+
if r["source"] in sources or r.get("task") in tasks: return False
|
| 258 |
+
return not (listed and (r["id"] in listed["ids"] or r.get("group") in listed["ids"] or r["source"] in listed["sources"] or r.get("task") in listed["tasks"]))
|
| 259 |
+
return keep
|
| 260 |
+
|
| 261 |
+
|
| 262 |
+
def filter_problems(panel, spec, root=ROOT, rows=True):
|
| 263 |
+
"""What is wrong with a panel's filters: malformed values, an exclude_file without {path, sha256}, and (rows: the local
|
| 264 |
+
data a read-out needs) an exclude_file here that does not match its hash. `exclude_sources` names are checked
|
| 265 |
+
against the manifests of the panel's reads (a typo would silently keep a source the rule means to drop)."""
|
| 266 |
+
problems, source = [], panel.get("source") # an absent exclude_file is not a problem here: validate lists it with the absent rows
|
| 267 |
+
if source is not None and not (isinstance(source, str) or (isinstance(source, list) and source and all(isinstance(s, str) for s in source))):
|
| 268 |
+
problems.append("source is a source name or a non-empty list of names")
|
| 269 |
+
for key in ("exclude_sources", "exclude_tasks"):
|
| 270 |
+
if key in panel and not (isinstance(panel[key], list) and panel[key] and all(isinstance(s, str) for s in panel[key])):
|
| 271 |
+
problems.append(f"{key} is a non-empty list of names")
|
| 272 |
+
if "optional" in panel and not isinstance(panel["optional"], bool): problems.append("optional is true or false")
|
| 273 |
+
entry = panel.get("exclude_file")
|
| 274 |
+
if entry is not None:
|
| 275 |
+
if not (isinstance(entry, dict) and isinstance(entry.get("path"), str) and isinstance(entry.get("sha256"), str) and len(entry["sha256"]) == 64):
|
| 276 |
+
problems.append("exclude_file is {path, sha256}")
|
| 277 |
+
elif rows and exclude_file_path(entry, root).exists() and digest(exclude_file_path(entry, root)) != entry["sha256"]:
|
| 278 |
+
problems.append(f"exclude_file {entry['path']} does not match its registered sha256")
|
| 279 |
+
if isinstance(panel.get("exclude_sources"), list):
|
| 280 |
+
suites = [spec["reads"][t].get("suite") for t in panel["reads"] if t in spec["reads"]]
|
| 281 |
+
if TRANSFER in panel["reads"]:
|
| 282 |
+
tag = spec.get("transfer_read")
|
| 283 |
+
suites.append(spec["reads"].get(tag, {}).get("suite") if tag else None)
|
| 284 |
+
manifests = [suite_manifest(suite_dir(s)) if s and suite_dir(s) else None for s in suites]
|
| 285 |
+
if manifests and all(m is not None for m in manifests):
|
| 286 |
+
listed = set().union(*(listed_sources(m) for m in manifests))
|
| 287 |
+
if unknown := sorted(set(panel["exclude_sources"]) - listed): problems.append(f"exclude_sources name {unknown}, which no read of the panel lists")
|
| 288 |
+
return problems
|
| 289 |
+
|
| 290 |
+
|
| 291 |
+
def gating_panels(rule):
|
| 292 |
+
"""Panels a criterion or the rank reads: an `optional` panel (report only) must not be one of them."""
|
| 293 |
+
paths = [p for c in rule["criteria"].values() for p in _paths(c)]
|
| 294 |
+
paths += [p for key in rule.get("rank", []) for p in (key["by"] if isinstance(key["by"], list) else [key["by"]])]
|
| 295 |
+
return {p.split(".")[0] for p in paths}
|
| 296 |
+
|
| 297 |
+
|
| 298 |
+
# --- spec ------------------------------------------------------------------------------------------------------------
|
| 299 |
+
|
| 300 |
+
def load(path):
|
| 301 |
+
spec = read_json(path)
|
| 302 |
+
spec.setdefault("arm_reads", "runs/r{round}-{arm}-{tag}")
|
| 303 |
+
return spec
|
| 304 |
+
|
| 305 |
+
|
| 306 |
+
def size_of(arm):
|
| 307 |
+
return arm.split("-")[0]
|
| 308 |
+
|
| 309 |
+
|
| 310 |
+
def locations(where, spec, **values):
|
| 311 |
+
"""{tag: dir} from a template string (one for every tag of the spec) or an explicit {tag: template} mapping."""
|
| 312 |
+
tags = [*spec.get("reads", {}), TRANSFER]
|
| 313 |
+
if isinstance(where, str):
|
| 314 |
+
return {tag: where.format(round=spec["round"], tag=tag, **values) for tag in tags}
|
| 315 |
+
return {tag: d.format(round=spec["round"], tag=tag, **values) for tag, d in where.items()}
|
| 316 |
+
|
| 317 |
+
|
| 318 |
+
def transfer_read(spec, arm):
|
| 319 |
+
"""The read tag whose rows are the arm's "transfer" rows, or None for its trial's in-trial transfer read."""
|
| 320 |
+
a = spec["arms"][arm]
|
| 321 |
+
return a["transfer_read"] if "transfer_read" in a else spec.get("transfer_read")
|
| 322 |
+
|
| 323 |
+
|
| 324 |
+
def arm_side(spec, arm, root=ROOT, stage=None):
|
| 325 |
+
"""The arm's side. Its "transfer" rows and its temperature pool come from its rule-stage reads at every stage (the
|
| 326 |
+
development reads it was selected on); a stage's candidate_reads locate the stage's own panels."""
|
| 327 |
+
a = spec["arms"][arm]
|
| 328 |
+
development = locations(a.get("reads", spec["arm_reads"]), spec, arm=arm, size=size_of(arm))
|
| 329 |
+
where = spec["confirm"][stage]["candidate_reads"] if stage else a.get("reads", spec["arm_reads"])
|
| 330 |
+
tag = transfer_read(spec, arm)
|
| 331 |
+
transfer = development.get(tag) if tag else f"{a['trial']}/transfer" if a.get("trial") else None
|
| 332 |
+
dirs = {**locations(where, spec, arm=arm, size=size_of(arm)), **({TRANSFER: transfer} if transfer else {})}
|
| 333 |
+
pool, registered = None, spec.get("temperature")
|
| 334 |
+
if registered:
|
| 335 |
+
fit = {**development, **({TRANSFER: transfer} if transfer else {})}
|
| 336 |
+
pool = Pool([fit[r] for r in registered["reads"]], {fit[r]: s for r, s in registered.get("sources", {}).items()}, [fit[r] for r in registered.get("exclude_reads", [])], registered.get("ci"))
|
| 337 |
+
suite = trial_suite(spec, a["trial"], root) if a.get("trial") and not pool else None
|
| 338 |
+
return Side(a.get("trial"), dirs, root, spec.get("drop_ids", ()), pool, a.get("checkpoint"), suite)
|
| 339 |
+
|
| 340 |
+
|
| 341 |
+
def parent_side(spec, arm, root=ROOT, stage=None):
|
| 342 |
+
p = spec["parents"][spec["arms"][arm]["parent"]]
|
| 343 |
+
where = (spec["confirm"][stage].get("parent_reads") if stage else None) or p["reads"]
|
| 344 |
+
dirs = {**locations(where, spec, arm=arm, size=size_of(arm)), TRANSFER: f"{p['trial']}/transfer"}
|
| 345 |
+
return Side(p["trial"], dirs, root, spec.get("drop_ids", ()), suite=trial_suite(spec, p["trial"], root), shipped=lambda: shipped_temperature(p, root))
|
| 346 |
+
|
| 347 |
+
|
| 348 |
+
def rule_tags(rule):
|
| 349 |
+
return {tag for panel in rule["panels"].values() for tag in panel["reads"]} | ({rule["unknowable"]} if rule.get("unknowable") else set())
|
| 350 |
+
|
| 351 |
+
|
| 352 |
+
def _paths(criterion):
|
| 353 |
+
left = criterion["left"] if isinstance(criterion["left"], list) else [criterion["left"]]
|
| 354 |
+
return [*left, *([criterion["right"]] if isinstance(criterion["right"], str) else [])]
|
| 355 |
+
|
| 356 |
+
|
| 357 |
+
class Validation(NamedTuple):
|
| 358 |
+
problems: list # the spec is malformed, or something a launch or a read-out needs is missing
|
| 359 |
+
archived: list # files a recorded round names that this checkout does not carry (they live on spec["archive"])
|
| 360 |
+
|
| 361 |
+
|
| 362 |
+
def validate(spec, root=ROOT, rows=True, plans=True, partitions=False):
|
| 363 |
+
"""Validation(problems, archived). rows: every parent's development rows and every read the rule needs from it exist
|
| 364 |
+
under root (what a read-out, and so a launch, needs); plans: every read suite exists, every plan trial passes
|
| 365 |
+
kev.experiment.validated_trial against its suite manifest and fits its budget; partitions: plans go through load_plan
|
| 366 |
+
instead (verifies the partitions, may fetch them from the Hub). A recorded round (`archive`: the git tag holding its
|
| 367 |
+
evidence) lists absent files under `archived` instead of failing; a round without it must have them all. A read of a
|
| 368 |
+
removed suite (kev.suite.REMOVED_SUITES) is archived (with plans) for a round up to the suite's last_round, whether or
|
| 369 |
+
not the round carries `archive`, and a problem for any later round."""
|
| 370 |
+
problems, archived, root = [], [], Path(root)
|
| 371 |
+
absent = (archived if spec.get("archive") else problems).append
|
| 372 |
+
for key in ("round", "registered", "parents", "arms", "reads", "rule"):
|
| 373 |
+
if key not in spec: problems.append(f"missing key {key!r}")
|
| 374 |
+
if problems: return Validation(problems, archived)
|
| 375 |
+
tags = {*spec["reads"], TRANSFER}
|
| 376 |
+
for name, r in spec["reads"].items():
|
| 377 |
+
if r.get("entrypoint", "benchmarks") not in ENTRYPOINTS: problems.append(f"read {name}: entrypoint {r['entrypoint']!r} is not one of {ENTRYPOINTS}"); continue
|
| 378 |
+
if r.get("entrypoint", "benchmarks") == "benchmarks" and not r.get("suite"): problems.append(f"read {name}: no suite")
|
| 379 |
+
if r.get("entrypoint") == "locked_test" and not r.get("decision"): problems.append(f"read {name}: locked_test needs a decision suite")
|
| 380 |
+
if r.get("suite") and (gone := removed_suite(r["suite"])): # kev.suite.REMOVED_SUITES: archived up to its last round, refused after
|
| 381 |
+
if isinstance(spec["round"], int) and spec["round"] <= gone["last_round"]:
|
| 382 |
+
if plans: archived.append(f"read {name}: {r['suite']} not in this checkout: removed on {gone['removed']} (kev.suite.REMOVED_SUITES); "
|
| 383 |
+
f"round {spec['round']} read it before the removal, and the rows it made are the record")
|
| 384 |
+
else:
|
| 385 |
+
problems.append(f"read {name}: {r['suite']} was removed on {gone['removed']} and may not be read after round {gone['last_round']}: {gone['reason']}")
|
| 386 |
+
continue
|
| 387 |
+
if plans and r.get("suite") and not (root / r["suite"]).exists(): absent(f"read {name}: {r['suite']} not in this checkout")
|
| 388 |
+
stages = {None: spec["rule"], **spec.get("confirm", {})}
|
| 389 |
+
reads_transfer = any(TRANSFER in panel["reads"] for rule in stages.values() for panel in rule["panels"].values())
|
| 390 |
+
for name, a in spec["arms"].items():
|
| 391 |
+
if a.get("parent") not in spec["parents"]: problems.append(f"arm {name}: unknown parent {a.get('parent')!r}")
|
| 392 |
+
if "-" not in name: problems.append(f"arm {name}: name must be <size>-<label>")
|
| 393 |
+
tag = transfer_read(spec, name)
|
| 394 |
+
if tag is not None and (tag not in spec["reads"] or spec["reads"][tag].get("entrypoint", "benchmarks") != "benchmarks"):
|
| 395 |
+
problems.append(f"arm {name}: transfer_read {tag!r} is not a benchmarks read of this spec")
|
| 396 |
+
if tag is not None and isinstance(a.get("reads"), dict) and tag not in a["reads"]: problems.append(f"arm {name}: its reads do not locate its transfer_read {tag!r}")
|
| 397 |
+
if not a.get("trial"):
|
| 398 |
+
if not str(a.get("checkpoint", "")).startswith("/runs/"): problems.append(f"arm {name}: an arm without a trial needs a checkpoint on the runs volume (/runs/...)")
|
| 399 |
+
if not spec.get("temperature"): problems.append(f"arm {name}: no trial, so no development rows to fit its temperature on; register a `temperature` pool")
|
| 400 |
+
if reads_transfer and tag is None: problems.append(f"arm {name}: no trial, so no in-trial transfer read; give it a transfer_read")
|
| 401 |
+
problems += pool_problems(spec, stages)
|
| 402 |
+
conflicts, unknown_training = pool_training_problems(spec, root)
|
| 403 |
+
problems += conflicts + pool_required_problems(spec)
|
| 404 |
+
for line in unknown_training: absent(line)
|
| 405 |
+
for stage, rule in stages.items():
|
| 406 |
+
where = f"confirm.{stage}" if stage else "rule"
|
| 407 |
+
for pname, panel in rule["panels"].items():
|
| 408 |
+
unknown = set(panel["reads"]) - tags
|
| 409 |
+
if unknown: problems.append(f"{where} panel {pname}: unknown read tags {sorted(unknown)}")
|
| 410 |
+
if not panel.get("metrics") and not panel.get("report"): problems.append(f"{where} panel {pname}: nothing to compute")
|
| 411 |
+
problems += [f"{where} panel {pname}: {p}" for p in by_length_problems(panel, spec)]
|
| 412 |
+
problems += [f"{where} panel {pname}: {p}" for p in filter_problems(panel, spec, root, rows)]
|
| 413 |
+
if panel.get("optional") and pname in gating_panels(rule): problems.append(f"{where} panel {pname}: optional (report only), but a criterion or the rank reads it")
|
| 414 |
+
for cname, c in rule["criteria"].items():
|
| 415 |
+
if c.get("op") not in OPS: problems.append(f"{where} criterion {cname}: op {c.get('op')!r}")
|
| 416 |
+
for path in _paths(c):
|
| 417 |
+
if not _known_path(path, rule): problems.append(f"{where} criterion {cname}: unknown path {path!r}")
|
| 418 |
+
for key in rule.get("rank", []):
|
| 419 |
+
for path in key["by"] if isinstance(key["by"], list) else [key["by"]]:
|
| 420 |
+
if not _known_path(path, rule): problems.append(f"{where} rank: unknown path {path!r}")
|
| 421 |
+
if stage and "candidate_reads" not in rule: problems.append(f"{where}: no candidate_reads")
|
| 422 |
+
if rule_tags(spec["rule"]) - tags: problems.append(f"rule.unknowable: unknown read tag {spec['rule']['unknowable']!r}")
|
| 423 |
+
for study, s in spec.get("studies", {}).items() if plans else ():
|
| 424 |
+
if not (root / s["plan"]).exists(): absent(f"study {study}: plan {s['plan']} not in this checkout")
|
| 425 |
+
elif (root / s["suite"] / "manifest.json").exists(): problems += _validate_plan(study, s, root, partitions)
|
| 426 |
+
else: absent(f"study {study}: suite {s['suite']} not in this checkout")
|
| 427 |
+
if rows:
|
| 428 |
+
for where, pname, entry in [(s or "rule", n, p["exclude_file"]) for s, rule in stages.items() for n, p in rule["panels"].items() if isinstance(p.get("exclude_file"), dict)]:
|
| 429 |
+
if isinstance(entry.get("path"), str) and not exclude_file_path(entry, root).exists():
|
| 430 |
+
absent(f"{where} panel {pname}: exclude_file {entry['path']} not in this checkout (a private list: scripts/private_rows.py restore)")
|
| 431 |
+
for pname, p in spec["parents"].items():
|
| 432 |
+
if not (root / p["trial"] / "development/rows.json").exists(): absent(f"parent {pname}: {p['trial']}/development/rows.json not in this checkout")
|
| 433 |
+
for arm in spec["arms"]:
|
| 434 |
+
side = parent_side(spec, arm, root)
|
| 435 |
+
required = {t for panel in spec["rule"]["panels"].values() if not panel.get("optional") for t in panel["reads"]} # an optional panel's reads may be absent
|
| 436 |
+
for tag in sorted(required - {spec["rule"].get("unknowable")}): # the parent's unknowable read is reported, not required
|
| 437 |
+
if not side.has(tag): absent(f"arm {arm}: parent read {tag} not in this checkout ({side.dirs.get(tag, 'no location')})")
|
| 438 |
+
return Validation(problems, archived)
|
| 439 |
+
|
| 440 |
+
|
| 441 |
+
def _known_path(path, rule):
|
| 442 |
+
parts = path.split(".")
|
| 443 |
+
if parts[0] == "unknowable": return len(parts) == 2 and parts[1] in ("candidate", "parent") and bool(rule.get("unknowable"))
|
| 444 |
+
panel = rule["panels"].get(parts[0])
|
| 445 |
+
if panel is None or len(parts) != 3 or parts[2] not in FIELDS: return False
|
| 446 |
+
value_only = parts[1] in panel.get("report", ()) or parts[1] in by_length_names(panel) # report and by-length metrics have no interval
|
| 447 |
+
return parts[1] in panel.get("metrics", ()) or (value_only and parts[2] in ("candidate", "parent"))
|
| 448 |
+
|
| 449 |
+
|
| 450 |
+
# --- calibration by state length -------------------------------------------------------------------------------------
|
| 451 |
+
|
| 452 |
+
def by_length_options(panel):
|
| 453 |
+
"""(edges, tokenizer) of a panel's by_length option: true = kev.metrics.LENGTH_EDGES and kev.suite.ADMISSION_TOKENIZER."""
|
| 454 |
+
option = panel.get("by_length")
|
| 455 |
+
option = {} if option is True else option
|
| 456 |
+
return tuple(option.get("edges", LENGTH_EDGES)), tuple(option.get("tokenizer", ADMISSION_TOKENIZER))
|
| 457 |
+
|
| 458 |
+
|
| 459 |
+
def by_length_names(panel):
|
| 460 |
+
"""{"<metric>_<bucket>": metric} for the entries a by_length panel reports (empty for any other panel, or invalid edges)."""
|
| 461 |
+
if not panel.get("by_length"): return {}
|
| 462 |
+
try:
|
| 463 |
+
return {f"{m}_{name}": m for m in LENGTH_METRICS for name, _, _ in length_buckets(by_length_options(panel)[0])}
|
| 464 |
+
except (ValueError, AttributeError, TypeError):
|
| 465 |
+
return {}
|
| 466 |
+
|
| 467 |
+
|
| 468 |
+
def by_length_problems(panel, spec):
|
| 469 |
+
option = panel.get("by_length")
|
| 470 |
+
if not option: return []
|
| 471 |
+
if option is not True and (not isinstance(option, dict) or set(option) - {"edges", "tokenizer"}):
|
| 472 |
+
return ["by_length is true or {edges?, tokenizer?}"]
|
| 473 |
+
problems = []
|
| 474 |
+
try:
|
| 475 |
+
edges, tokenizer = by_length_options(panel) # tuple() of a non-iterable raises: a malformed option is a problem, not a crash
|
| 476 |
+
length_buckets(edges)
|
| 477 |
+
if len(tokenizer) != 2 or not all(isinstance(x, str) for x in tokenizer): problems.append("by_length: tokenizer is [name, revision]")
|
| 478 |
+
except (ValueError, TypeError) as error:
|
| 479 |
+
problems.append(f"by_length: {error}" if isinstance(error, ValueError) else f"by_length: edges are a list of token counts and tokenizer is [name, revision] ({error})")
|
| 480 |
+
if TRANSFER in panel["reads"]: problems.append("by_length counts state tokens from a read's suite; 'transfer' names no single suite")
|
| 481 |
+
problems += [f"by_length: read {tag!r} has no suite" for tag in panel["reads"] if tag in spec["reads"] and not spec["reads"][tag].get("suite")]
|
| 482 |
+
return problems
|
| 483 |
+
|
| 484 |
+
|
| 485 |
+
@functools.cache
|
| 486 |
+
def state_lengths(suite, split, tokenizer=ADMISSION_TOKENIZER):
|
| 487 |
+
"""{record id: state tokens} of a suite partition with the tokenizer (name, revision): the counts a by_length panel
|
| 488 |
+
buckets both sides by. A state's token count is its encoded state segment, as kev.model.encode builds it: the <state>
|
| 489 |
+
token plus user_tokens of the materialised state (len + 1, what encode checks against max_state and kev.serve reports
|
| 490 |
+
as `state_tokens`, seg.count(0)); the one definition behind a row's `state_tokens` (kev.metrics.calibration_by_length)."""
|
| 491 |
+
from kev.data import materialize
|
| 492 |
+
from kev.model import load_tokenizer, user_tokens
|
| 493 |
+
tok = load_tokenizer(*tokenizer)
|
| 494 |
+
return {r["_meta"]["id"]: 1 + len(user_tokens(tok, materialize(r)["state"])) for r in load_split(ROOT / suite, split, allow_test=split == "test")}
|
| 495 |
+
|
| 496 |
+
|
| 497 |
+
def panel_lengths(spec, panel):
|
| 498 |
+
"""{record id: state tokens} over a by_length panel's reads (each read's suite and partition)."""
|
| 499 |
+
tokenizer, out = by_length_options(panel)[1], {}
|
| 500 |
+
for tag in panel["reads"]:
|
| 501 |
+
read = spec["reads"][tag]
|
| 502 |
+
for rid, n in state_lengths(read["suite"], read_split(read), tokenizer).items():
|
| 503 |
+
if out.setdefault(rid, n) != n: raise ValueError(f"record {rid!r} has different states in the reads of one by_length panel")
|
| 504 |
+
return out
|
| 505 |
+
|
| 506 |
+
|
| 507 |
+
# --- temperature pools against training data (round 19's failure mode) ------------------------------------------------
|
| 508 |
+
|
| 509 |
+
ROUND_19 = ("round 19's failure mode: a temperature fitted on held-out items of a checkpoint's own training corpus is in distribution "
|
| 510 |
+
"(round 19 served its SFT arms at T 0.955 fitted on sft-v1 development rows: breadth-v1 ECE 0.059); fit on a pool of "
|
| 511 |
+
"held-out datasets instead (round 20: breadth-v1 ECE 0.0085)")
|
| 512 |
+
FIT_SPLITS = ("calibration", "development") # a training corpus's held-out partitions: never a temperature pool
|
| 513 |
+
|
| 514 |
+
|
| 515 |
+
def read_split(read):
|
| 516 |
+
"""The partition a read scores: kev.benchmark's development, `--split X`, or the test under --allow-test / locked_test."""
|
| 517 |
+
flags = read.get("flags", "").split()
|
| 518 |
+
if read.get("entrypoint") == "locked_test" or "--allow-test" in flags: return "test"
|
| 519 |
+
for i, flag in enumerate(flags):
|
| 520 |
+
if flag.startswith("--split="): return flag.split("=", 1)[1]
|
| 521 |
+
if flag == "--split" and i + 1 < len(flags): return flags[i + 1]
|
| 522 |
+
return "development"
|
| 523 |
+
|
| 524 |
+
|
| 525 |
+
def suite_dir(path):
|
| 526 |
+
"""'evals/...' for a suite path given relative, absolute or as a container saw it (/root/kev/evals/sft-v1); None if none."""
|
| 527 |
+
return suite_key(path)
|
| 528 |
+
|
| 529 |
+
|
| 530 |
+
@functools.cache
|
| 531 |
+
def _suites():
|
| 532 |
+
"""({manifest sha256: suite dir}, {directory name: [suite dirs]}) of every suite in this checkout."""
|
| 533 |
+
by_digest, by_name = {}, {}
|
| 534 |
+
for m in sorted((ROOT / "evals").rglob("manifest.json")):
|
| 535 |
+
d = str(m.parent.relative_to(ROOT))
|
| 536 |
+
by_digest[digest(m)] = d
|
| 537 |
+
by_name.setdefault(m.parent.name, []).append(d)
|
| 538 |
+
return by_digest, by_name
|
| 539 |
+
|
| 540 |
+
|
| 541 |
+
def suite_by_digest(sha256):
|
| 542 |
+
"""The suite dir whose manifest.json hashes to sha256 (what provenance, head.pt and report.json record), or None."""
|
| 543 |
+
return _suites()[0].get(sha256)
|
| 544 |
+
|
| 545 |
+
|
| 546 |
+
def suite_manifest(suite):
|
| 547 |
+
return read_manifest(ROOT / suite) if (ROOT / suite / "manifest.json").exists() else None
|
| 548 |
+
|
| 549 |
+
|
| 550 |
+
def listed_sources(manifest):
|
| 551 |
+
"""Every source a manifest names: trainable, held-out, eval-only and its `sources` table (a dict or a list)."""
|
| 552 |
+
return {*manifest.get("trainable_sources", []), *manifest.get("holdout_sources", []), *manifest.get("eval_only_sources", []), *manifest.get("sources", [])}
|
| 553 |
+
|
| 554 |
+
|
| 555 |
+
def trainable_sources(manifest):
|
| 556 |
+
"""The sources a suite offers for training (`trainable_sources`, or `sources` entries marked trainable): a suite with
|
| 557 |
+
any is a training corpus, whose calibration and development partitions are held-out items of training sources."""
|
| 558 |
+
table = manifest.get("sources", {})
|
| 559 |
+
return {*manifest.get("trainable_sources", []), *(k for k, v in table.items() if isinstance(v, dict) and v.get("trainable"))} if isinstance(table, dict) else set(manifest.get("trainable_sources", []))
|
| 560 |
+
|
| 561 |
+
|
| 562 |
+
class Training(NamedTuple):
|
| 563 |
+
"""What a checkpoint was trained on, as far as this checkout can tell."""
|
| 564 |
+
suites: frozenset # suite dirs: its training suite, the components that suite names, a plan's `data` directory
|
| 565 |
+
sources: frozenset # the sources those suites train
|
| 566 |
+
unlisted: tuple # training suites whose sources this checkout cannot list
|
| 567 |
+
|
| 568 |
+
|
| 569 |
+
def training_data(suites, data=None):
|
| 570 |
+
"""Training over suite dirs plus a `data` file: each suite's trainable sources (every listed source for a data-only
|
| 571 |
+
manifest without `trainable_sources`, e.g. round6/b1v2) plus the suites its manifest names as components (sft-v1's
|
| 572 |
+
`inputs.components`: documents-v1-train -> evals/documents-v1, hard-v1-extra-train -> evals/hard-v1). A `data` file
|
| 573 |
+
counts through its directory's suite; one outside evals/ (a `kev.train --data` run) cannot be checked, so it is
|
| 574 |
+
`unlisted` like any training this checkout cannot place, never silently dropped."""
|
| 575 |
+
data_dir = suite_dir(Path(data).parent) if data else None
|
| 576 |
+
seen, sources, unlisted, todo = set(), set(), [] if data_dir or not data else [f"data {data} (outside evals/)"], [d for d in (*suites, data_dir) if d]
|
| 577 |
+
while todo:
|
| 578 |
+
d = todo.pop(0)
|
| 579 |
+
if d in seen: continue
|
| 580 |
+
seen.add(d)
|
| 581 |
+
m = suite_manifest(d)
|
| 582 |
+
own = (trainable_sources(m) if "trainable_sources" in m else listed_sources(m)) if m is not None else set()
|
| 583 |
+
if not own: unlisted.append(d)
|
| 584 |
+
sources |= own
|
| 585 |
+
inputs = (m or {}).get("inputs")
|
| 586 |
+
for component in inputs.get("components", {}) if isinstance(inputs, dict) else ():
|
| 587 |
+
found = _suites()[1].get(component.removesuffix("-train").removesuffix("-extra"))
|
| 588 |
+
if found: todo += found
|
| 589 |
+
else: unlisted.append(f"{d} component {component}")
|
| 590 |
+
return Training(frozenset(seen), frozenset(sources), tuple(unlisted))
|
| 591 |
+
|
| 592 |
+
|
| 593 |
+
def recorded_training(suite_sha256=None, suite=None, data=None):
|
| 594 |
+
"""Training of a checkpoint from what its trainer recorded (provenance.json, head.pt): the suite whose manifest hashes to
|
| 595 |
+
suite_sha256 (else the suite path it was given) and its `data` file's directory; None when no suite of this checkout matches."""
|
| 596 |
+
trained = suite_by_digest(suite_sha256)
|
| 597 |
+
if trained is None and suite and suite_dir(suite) and suite_manifest(suite_dir(suite)) is not None: trained = suite_dir(suite)
|
| 598 |
+
return training_data([trained], data) if trained else None
|
| 599 |
+
|
| 600 |
+
|
| 601 |
+
def _study(spec, trial):
|
| 602 |
+
return next((s for name, s in spec.get("studies", {}).items() if trial.startswith(f"runs/{name}/")), None)
|
| 603 |
+
|
| 604 |
+
|
| 605 |
+
def trial_suite(spec, trial, root=ROOT):
|
| 606 |
+
"""The suite a trial trained on and was scored on (its development rows'): its study's in this spec, else the manifest
|
| 607 |
+
its provenance.json hashes; None when neither is here."""
|
| 608 |
+
if study := _study(spec, trial): return suite_dir(study["suite"])
|
| 609 |
+
provenance = Path(root) / trial / "provenance.json"
|
| 610 |
+
return suite_by_digest(read_json(provenance).get("suite_sha256")) if provenance.exists() else None
|
| 611 |
+
|
| 612 |
+
|
| 613 |
+
def trial_training(spec, trial, root=ROOT):
|
| 614 |
+
"""Training of a trial: its study's suite plus its plan entry's `data` (trial <NN>-<label> is entry NN), else what its
|
| 615 |
+
provenance.json records; None when neither is here."""
|
| 616 |
+
study, provenance = _study(spec, trial), Path(root) / trial / "provenance.json"
|
| 617 |
+
if study:
|
| 618 |
+
plan = next((p / study["plan"] for p in (Path(root), ROOT) if (p / study["plan"]).exists()), None)
|
| 619 |
+
entries, index = (read_json(plan) if plan else []), Path(trial).name.split("-")[0]
|
| 620 |
+
data = entries[int(index)].get("data") if isinstance(entries, list) and index.isdigit() and int(index) < len(entries) else None
|
| 621 |
+
if plan or not provenance.exists(): return training_data([suite_dir(study["suite"])], data)
|
| 622 |
+
if not provenance.exists(): return None
|
| 623 |
+
p = read_json(provenance)
|
| 624 |
+
return recorded_training(p.get("suite_sha256"), data=p.get("config", {}).get("data"))
|
| 625 |
+
|
| 626 |
+
|
| 627 |
+
def pool_conflicts(pooled, training):
|
| 628 |
+
"""[(label, why)] for every way a pooled read shares data with a checkpoint's Training. pooled = [(label, suite, split,
|
| 629 |
+
sources or None)]: (a) the read's suite is training data of the checkpoint (its training suite, or a component or `data`
|
| 630 |
+
suite of it); (b) the sources it pools (its allowlist, else every source its manifest lists) are sources the checkpoint
|
| 631 |
+
trained on; (c) it scores the calibration or development partition of a training corpus (any suite with trainable
|
| 632 |
+
sources). A read whose suite or sources cannot be established is refused too. The one check, for kev.rounds pools and
|
| 633 |
+
scripts/calibrate_checkpoint.py alike."""
|
| 634 |
+
out = []
|
| 635 |
+
for label, suite, split, sources in pooled:
|
| 636 |
+
d = suite_dir(suite) if suite else None
|
| 637 |
+
m = suite_manifest(d) if d else None
|
| 638 |
+
if m is None:
|
| 639 |
+
out.append((label, f"its suite ({suite!r}) has no manifest in this checkout, so what it pools cannot be checked")); continue
|
| 640 |
+
if d in training.suites: out.append((label, f"{d} is training data of the checkpoint"))
|
| 641 |
+
pooled_sources = set(sources) if sources else listed_sources(m)
|
| 642 |
+
if not pooled_sources: out.append((label, f"{d} lists no sources; give the read a `sources` allowlist"))
|
| 643 |
+
if shared := sorted(pooled_sources & training.sources):
|
| 644 |
+
out.append((label, f"it pools {len(shared)} training source(s) {shared[:8]}" + (" ..." if len(shared) > 8 else "")))
|
| 645 |
+
if split in FIT_SPLITS and trainable_sources(m): out.append((label, f"it reads the {split} partition of {d}, a training corpus"))
|
| 646 |
+
return out
|
| 647 |
+
|
| 648 |
+
|
| 649 |
+
def pool_training_problems(spec, root=ROOT):
|
| 650 |
+
"""(problems, unknown) of the round's `temperature` pool against every arm's training data (pool_conflicts; message
|
| 651 |
+
names round 19's failure mode). An arm's training comes from `trained_on`, else from its trial (trial_training); an arm
|
| 652 |
+
without either (a checkpoint made from the round's trials) is covered by the trial arms. unknown: arms whose training
|
| 653 |
+
this checkout cannot establish or list (a trial without provenance, a manifest without sources, a `data` file outside
|
| 654 |
+
evals/): a problem for a new round, reported under archived for a recorded one."""
|
| 655 |
+
pool = spec.get("temperature")
|
| 656 |
+
if not pool or not pool.get("reads"): return [], []
|
| 657 |
+
pooled = [(tag, spec["reads"][tag].get("suite"), read_split(spec["reads"][tag]), pool.get("sources", {}).get(tag)) for tag in pool["reads"] if tag in spec["reads"]]
|
| 658 |
+
found, problems, unknown = {}, [], []
|
| 659 |
+
for arm, a in spec["arms"].items():
|
| 660 |
+
if "trained_on" in a:
|
| 661 |
+
if not isinstance(a["trained_on"], list) or not a["trained_on"] or not all(isinstance(s, str) for s in a["trained_on"]):
|
| 662 |
+
problems.append(f"arm {arm}: trained_on is a non-empty list of suite directories"); continue
|
| 663 |
+
training = training_data([suite_dir(s) or s for s in a["trained_on"]])
|
| 664 |
+
elif a.get("trial"):
|
| 665 |
+
training = trial_training(spec, a["trial"], root)
|
| 666 |
+
if training is None:
|
| 667 |
+
unknown.append(f"arm {arm}: what {a['trial']} trained on is unknown (no study of this spec, no provenance.json), so the temperature pool cannot be checked against it"); continue
|
| 668 |
+
else:
|
| 669 |
+
continue
|
| 670 |
+
unknown += [f"arm {arm}: cannot list the sources of {d}, part of its training, so the temperature pool cannot be checked against it" for d in training.unlisted]
|
| 671 |
+
by_tag = {}
|
| 672 |
+
for tag, why in pool_conflicts(pooled, training): by_tag.setdefault(tag, []).append(why)
|
| 673 |
+
for tag, whys in by_tag.items(): found.setdefault((tag, "; ".join(whys)), []).append(arm)
|
| 674 |
+
for (tag, why), arms in found.items():
|
| 675 |
+
problems.append(f"temperature: pooled read {tag!r} shares data with the training of arm(s) {', '.join(arms)}: {why}. This is {ROUND_19}")
|
| 676 |
+
if not any("trained_on" in a or a.get("trial") for a in spec["arms"].values()):
|
| 677 |
+
problems.append("temperature: no arm names its training (a trial or `trained_on`), so the pool cannot be checked against it")
|
| 678 |
+
return problems, unknown
|
| 679 |
+
|
| 680 |
+
|
| 681 |
+
POOL_REQUIRED_FROM = 21 # from this round on, a temperature-dependent criterion without a `temperature` pool is refused
|
| 682 |
+
|
| 683 |
+
|
| 684 |
+
def temperature_criteria(spec):
|
| 685 |
+
"""["rule <name>" | "confirm.<stage> <name>"] of every criterion that reads a number the temperature moves."""
|
| 686 |
+
stages = {None: spec["rule"], **spec.get("confirm", {})}
|
| 687 |
+
return sorted({f"{'confirm.' + s if s else 'rule'} {name}" for s, rule in stages.items() for name, c in rule["criteria"].items()
|
| 688 |
+
if any(_moves_with_temperature(path, rule) for path in _paths(c))})
|
| 689 |
+
|
| 690 |
+
|
| 691 |
+
def _listed(names):
|
| 692 |
+
return f"{len(names)} ({', '.join(names[:3])}{', ...' if len(names) > 3 else ''})"
|
| 693 |
+
|
| 694 |
+
|
| 695 |
+
def pool_required_problems(spec):
|
| 696 |
+
"""Round >= POOL_REQUIRED_FROM with a temperature-dependent criterion and no `temperature` pool: refused (its arms would
|
| 697 |
+
be served at their trials' development rows, round 19's failure mode). Earlier rounds only warn (calibration_warnings),
|
| 698 |
+
so their recorded specs keep validating."""
|
| 699 |
+
number, moved = spec.get("round"), temperature_criteria(spec)
|
| 700 |
+
if spec.get("temperature") or not moved or not isinstance(number, int) or number < POOL_REQUIRED_FROM: return []
|
| 701 |
+
return [f"temperature: round {number} registers no `temperature` pool, but {_listed(moved)} criteria depend on the temperature, "
|
| 702 |
+
f"so its arms would be served at their trials' development rows. This is {ROUND_19}; register a pool of held-out datasets (copy r20)"]
|
| 703 |
+
|
| 704 |
+
|
| 705 |
+
def shipped_temperature(entry, root=ROOT):
|
| 706 |
+
"""The temperature a parent's checkpoint ships (head.pt), read locally only: <trial>/checkpoint/head.pt, else its Hub
|
| 707 |
+
checkpoint's head.pt in the local Hugging Face cache (no network); None when neither is here."""
|
| 708 |
+
path = Path(root) / entry["trial"] / "checkpoint" / "head.pt"
|
| 709 |
+
if not path.exists() and entry.get("checkpoint") and not str(entry["checkpoint"]).startswith("/"):
|
| 710 |
+
from huggingface_hub import try_to_load_from_cache
|
| 711 |
+
repo, _, revision = entry["checkpoint"].partition("@")
|
| 712 |
+
cached = try_to_load_from_cache(repo, "head.pt", revision=revision or None)
|
| 713 |
+
path = Path(cached) if isinstance(cached, str) else path
|
| 714 |
+
if not path.exists(): return None
|
| 715 |
+
from kev.checkpoint import read_meta
|
| 716 |
+
return read_meta(path.parent).temperature
|
| 717 |
+
|
| 718 |
+
|
| 719 |
+
SHIPPED_TOLERANCE = 0.05 # a parent served this far from its shipped temperature, fitted in distribution, is warned about
|
| 720 |
+
|
| 721 |
+
|
| 722 |
+
def continuation_warnings(spec, root=ROOT):
|
| 723 |
+
"""Report only: a full-weight study runs longer than one attempt only while something continues it. Its trials spawn
|
| 724 |
+
with Modal's retries off (kev.budget.FULL_FT_RETRIES says why), so a timed-out trial gets its next attempt from
|
| 725 |
+
`kev.rounds watch` (modal_app.continue_full_trial), or by hand with modal_app.py::resume --trial; an attached launch
|
| 726 |
+
keeps no attempt ledger and is never continued."""
|
| 727 |
+
out = []
|
| 728 |
+
for name, study in spec.get("studies", {}).items():
|
| 729 |
+
plan = Path(root) / study.get("plan", "")
|
| 730 |
+
trials = read_json(plan) if plan.is_file() else []
|
| 731 |
+
if any(isinstance(t, dict) and t.get("full_ft") for t in (trials if isinstance(trials, list) else [])):
|
| 732 |
+
out.append(f"study {name} is full-weight: a trial that times out is continued only while `kev.rounds watch` runs (or by hand: "
|
| 733 |
+
f"modal_app.py::resume --study {name} --trial <label>); an attached launch (study --detached False) keeps no ledger and is never continued")
|
| 734 |
+
return out
|
| 735 |
+
|
| 736 |
+
|
| 737 |
+
def calibration_warnings(spec, root=ROOT):
|
| 738 |
+
"""Report-only warnings that `validate` and `launch` print. (1) For a round before POOL_REQUIRED_FROM without a
|
| 739 |
+
`temperature` pool whose criteria depend on the temperature (from then on it is a problem: pool_required_problems): each
|
| 740 |
+
trial arm whose development rows are a training corpus's is served in distribution, round 19's failure mode. (2) Every
|
| 741 |
+
parent served at a temperature fitted on a training corpus's development rows that differs from its shipped head.pt
|
| 742 |
+
temperature by more than SHIPPED_TOLERANCE (when both are here)."""
|
| 743 |
+
out, moved = [], temperature_criteria(spec)
|
| 744 |
+
early = isinstance(spec.get("round"), int) and spec["round"] < POOL_REQUIRED_FROM
|
| 745 |
+
for arm, a in spec["arms"].items() if moved and early and not spec.get("temperature") else ():
|
| 746 |
+
suite = trial_suite(spec, a["trial"], root) if a.get("trial") else None
|
| 747 |
+
if suite and trainable_sources(suite_manifest(suite) or {}):
|
| 748 |
+
out.append(f"arm {arm} will be served at a temperature fitted on its trial's {suite} development rows, a training corpus, and "
|
| 749 |
+
f"{_listed(moved)} criteria depend on it: {ROUND_19}; register a `temperature` pool")
|
| 750 |
+
for name, p in spec["parents"].items():
|
| 751 |
+
suite = trial_suite(spec, p["trial"], root)
|
| 752 |
+
if not (suite and trainable_sources(suite_manifest(suite) or {}) and (Path(root) / p["trial"] / "development/rows.json").exists()): continue
|
| 753 |
+
shipped = shipped_temperature(p, root)
|
| 754 |
+
served_t = temperature(p["trial"], root) if shipped is not None else None
|
| 755 |
+
if shipped is not None and abs(served_t - shipped) > SHIPPED_TOLERANCE:
|
| 756 |
+
out.append(f"parent {name} is served at T {served_t:.3f} fitted on {p['trial']}/development ({suite}, a training corpus), "
|
| 757 |
+
f"{abs(served_t - shipped):.3f} from the T {shipped:.3f} its checkpoint ships: its reads are not what it serves (report only)")
|
| 758 |
+
return out
|
| 759 |
+
|
| 760 |
+
|
| 761 |
+
TEMPERATURE_FREE = ("acc",) # the one bootstrapped metric a temperature cannot move (the argmax is invariant)
|
| 762 |
+
|
| 763 |
+
|
| 764 |
+
def _moves_with_temperature(path, rule):
|
| 765 |
+
"""Whether a criterion path reads a number the temperature moves: any panel metric but accuracy (a by-length bucket's
|
| 766 |
+
included: acc_16k_plus is accuracy); the unknowable share is scored on records a pool never fits on."""
|
| 767 |
+
panel, metric = (path.split(".") + [""])[:2]
|
| 768 |
+
return panel != "unknowable" and by_length_names(rule["panels"].get(panel, {})).get(metric, metric) not in TEMPERATURE_FREE
|
| 769 |
+
|
| 770 |
+
|
| 771 |
+
def allowlist_problems(tag, suite, sources):
|
| 772 |
+
"""A pool read's `sources` allowlist against the sources its suite's manifest lists: a name the suite does not contain
|
| 773 |
+
(a typo) would silently shrink the pool, so it is a problem, as is a suite whose sources cannot be listed."""
|
| 774 |
+
m = suite_manifest(suite_dir(suite)) if suite and suite_dir(suite) else None
|
| 775 |
+
listed = listed_sources(m) if m is not None else set()
|
| 776 |
+
if not listed: return [f"temperature: cannot check the sources allowlist of {tag!r}: {suite!r} lists no sources in this checkout"]
|
| 777 |
+
return [f"temperature: sources for {tag!r} name {unknown} which {suite} does not contain" for unknown in [sorted(set(sources) - listed)] if unknown]
|
| 778 |
+
|
| 779 |
+
|
| 780 |
+
def pool_problems(spec, stages):
|
| 781 |
+
"""What is wrong with the round's `temperature` pool: tags that are not reads of the spec, tags an arm cannot locate,
|
| 782 |
+
and pooled tags inside a panel that a temperature-dependent criterion reads (any metric but accuracy, at any stage):
|
| 783 |
+
the temperature would be fitted on the rows it is judged on. The unknowable share is scored on unknowable records,
|
| 784 |
+
which a pool never fits on (knowable rows only), so the unknowable read may be pooled."""
|
| 785 |
+
pool = spec.get("temperature")
|
| 786 |
+
if not pool: return []
|
| 787 |
+
reads, exclude = pool.get("reads") or [], pool.get("exclude_reads", [])
|
| 788 |
+
problems = ["temperature: no reads to pool"] if not reads else []
|
| 789 |
+
ci = pool.get("ci")
|
| 790 |
+
if ci is not None and not (isinstance(ci, dict) and not set(ci) - {"level", "samples", "seed"} and 0 < ci.get("level", 0.9) < 1
|
| 791 |
+
and isinstance(ci.get("samples", SAMPLES), int) and ci.get("samples", SAMPLES) > 0 and isinstance(ci.get("seed", 0), int)):
|
| 792 |
+
problems.append("temperature: ci is {level (0-1), samples (> 0), seed}")
|
| 793 |
+
problems += [f"temperature: {tag!r} is not a read tag of this spec" for tag in reads if tag not in spec["reads"]]
|
| 794 |
+
problems += [f"temperature: exclude_reads {tag!r} is not a read tag of this spec" for tag in exclude if tag not in spec["reads"] and tag != TRANSFER]
|
| 795 |
+
for tag, sources in pool.get("sources", {}).items():
|
| 796 |
+
if tag not in reads: problems.append(f"temperature: sources for {tag!r}, which is not pooled")
|
| 797 |
+
if not sources or not isinstance(sources, list) or not all(isinstance(s, str) for s in sources):
|
| 798 |
+
problems.append(f"temperature: sources for {tag!r} must be a non-empty list of source names"); continue
|
| 799 |
+
problems += allowlist_problems(tag, spec["reads"].get(tag, {}).get("suite"), sources)
|
| 800 |
+
for stage, rule in stages.items():
|
| 801 |
+
for cname, c in rule["criteria"].items():
|
| 802 |
+
for path in _paths(c):
|
| 803 |
+
shared = sorted(set(reads) & set(rule["panels"].get(path.split(".")[0], {}).get("reads", [])))
|
| 804 |
+
if shared and _moves_with_temperature(path, rule):
|
| 805 |
+
problems.append(f"temperature: {shared} pooled, but {'confirm.' + stage if stage else 'rule'} criterion {cname} reads {path}, which the temperature moves")
|
| 806 |
+
for arm, a in spec["arms"].items():
|
| 807 |
+
where = a.get("reads", spec["arm_reads"])
|
| 808 |
+
unplaced = [t for t in [*reads, *exclude] if (t == TRANSFER and not a.get("trial") and transfer_read(spec, arm) is None)
|
| 809 |
+
or (t != TRANSFER and not isinstance(where, str) and t not in where)]
|
| 810 |
+
if unplaced: problems.append(f"arm {arm}: its reads do not locate the temperature pool's {unplaced}")
|
| 811 |
+
return problems
|
| 812 |
+
|
| 813 |
+
|
| 814 |
+
def _validate_plan(study, s, root, partitions):
|
| 815 |
+
from kev.experiment import load_plan, validated_trial
|
| 816 |
+
from kev.suite import read_manifest
|
| 817 |
+
from kev.budget import compute_bound # the admission bound modal_app.admit_study refuses a study over
|
| 818 |
+
problems = []
|
| 819 |
+
try:
|
| 820 |
+
trials = load_plan(root / s["suite"], root / s["plan"]) if partitions else [validated_trial(t, read_manifest(root / s["suite"])) for t in read_json(root / s["plan"])]
|
| 821 |
+
except (ValueError, KeyError) as error:
|
| 822 |
+
return [f"study {study}: plan {s['plan']} does not validate ({type(error).__name__}: {error})"]
|
| 823 |
+
bound = compute_bound(s["gpu"], s["timeout"], len(trials), any(t.get("full_ft") for t in trials))
|
| 824 |
+
if bound > s["budget"]: problems.append(f"study {study}: admission bound ${bound:.2f} exceeds budget ${s['budget']:.2f} (modal_app.admit_study would refuse it)")
|
| 825 |
+
return problems
|
| 826 |
+
|
| 827 |
+
|
| 828 |
+
# --- rule ------------------------------------------------------------------------------------------------------------
|
| 829 |
+
|
| 830 |
+
def compare(candidate, parent, rule, lengths=None):
|
| 831 |
+
"""Every panel, the unknowable share and every criterion of one candidate against its parent. Panels whose reads are
|
| 832 |
+
missing on either side are listed under "missing" and their criteria are None; `passed` needs every criterion. A
|
| 833 |
+
candidate served at a pooled temperature records the fit (`temperature_fit`); every candidate records where its
|
| 834 |
+
temperature came from (`temperature_source`); one whose pool rows are missing is reported incomplete with nothing
|
| 835 |
+
computed. lengths(panel) -> {record id: state tokens} serves by_length panels whose rows record no state_tokens."""
|
| 836 |
+
out = {"trial": candidate.trial, "parent": parent.trial, **({"checkpoint": candidate.checkpoint} if candidate.checkpoint else {})}
|
| 837 |
+
if absent := [f"candidate:{d}" for d in candidate.absent_fit()]:
|
| 838 |
+
return {**out, "temperature": None, "missing": absent, "complete": False, "passed": None}
|
| 839 |
+
out.update(temperature=candidate.t, parent_temperature=parent.t, panels={}, missing=[])
|
| 840 |
+
if candidate.fit: out["temperature_fit"] = candidate.fit
|
| 841 |
+
if candidate.pool and candidate.pool.ci: out["temperature_ci"] = pooled_temperature_ci(candidate.pool, candidate.root)
|
| 842 |
+
out["temperature_source"] = candidate.temperature_source()
|
| 843 |
+
out["parent_temperature_source"] = parent.temperature_source() # parents are served at their trial's development rows; `shipped` is head.pt's T
|
| 844 |
+
for name, spec in rule["panels"].items():
|
| 845 |
+
absent = [f"{who}:{side.dirs[t] if t in side.dirs else t}" for who, side in (("candidate", candidate), ("parent", parent)) for t in spec["reads"] if not side.has(t)]
|
| 846 |
+
if isinstance(spec.get("exclude_file"), dict) and not exclude_file_path(spec["exclude_file"], candidate.root).exists():
|
| 847 |
+
absent.append(f"exclude_file:{spec['exclude_file']['path']}")
|
| 848 |
+
if absent: # an optional (report-only) panel's absent reads leave the comparison complete
|
| 849 |
+
if spec.get("optional"): out.setdefault("missing_optional", []).extend(absent)
|
| 850 |
+
else: out["missing"] += absent
|
| 851 |
+
continue
|
| 852 |
+
c, p = candidate.panel(spec), parent.panel(spec)
|
| 853 |
+
mc, mp = metrics(c), metrics(p)
|
| 854 |
+
panel = {"n": len(c)}
|
| 855 |
+
if any(k in spec for k in EXCLUDE_KEYS): # questions the panel's exclusions removed (candidate side; the parent's are the same records)
|
| 856 |
+
panel["excluded"] = len(candidate.panel({k: v for k, v in spec.items() if k not in EXCLUDE_KEYS})) - len(c)
|
| 857 |
+
for m in spec.get("metrics", ()):
|
| 858 |
+
panel[m] = {"candidate": mc[m], "parent": mp[m], **paired(c, p, m)}
|
| 859 |
+
for m in spec.get("report", ()):
|
| 860 |
+
panel.setdefault(m, {"candidate": mc[m], "parent": mp[m]})
|
| 861 |
+
if spec.get("by_length"):
|
| 862 |
+
panel.update(_by_length(spec, c, p, lengths))
|
| 863 |
+
for ref, d in spec.get("versus", {}).items():
|
| 864 |
+
keep = panel_filter(spec, candidate.root)
|
| 865 |
+
rows = [r for r in read_json(candidate.root / d / "rows.json") if r["variant"] == "clean" and keep(r)] # a reference as it served itself
|
| 866 |
+
panel.setdefault("versus", {})[ref] = {m: {"reference": metrics(rows)[m], **paired(c, rows, m)} for m in spec.get("metrics", ())}
|
| 867 |
+
out["panels"][name] = panel
|
| 868 |
+
if rule.get("unknowable"):
|
| 869 |
+
tag = rule["unknowable"]
|
| 870 |
+
if not candidate.has(tag): out["missing"].append(f"candidate:{candidate.dirs[tag]}")
|
| 871 |
+
out["unknowable"] = {"candidate": candidate.unknowable_share(tag), "parent": parent.unknowable_share(tag)}
|
| 872 |
+
out["criteria"] = {name: _criterion(out, c) for name, c in rule["criteria"].items()}
|
| 873 |
+
out["complete"] = not out["missing"]
|
| 874 |
+
out["passed"] = (out["complete"] and all(out["criteria"].values())) if out["criteria"] else None
|
| 875 |
+
return out
|
| 876 |
+
|
| 877 |
+
|
| 878 |
+
def _by_length(spec, candidate, parent, lengths):
|
| 879 |
+
"""A by_length panel's entries: "<metric>_<bucket>" -> {candidate, parent, n} (kev.metrics.calibration_by_length of
|
| 880 |
+
each side's served rows, bucketed by the same counts) and "by_length" -> how the tokens were counted."""
|
| 881 |
+
edges, tokenizer = by_length_options(spec)
|
| 882 |
+
counted = all(r.get("state_tokens") is not None for r in (*candidate, *parent))
|
| 883 |
+
if not counted and lengths is None: raise ValueError("a by_length panel needs state-token counts: rows with state_tokens, or lengths")
|
| 884 |
+
counts = None if counted else lengths(spec)
|
| 885 |
+
c, p = calibration_by_length(candidate, counts, edges), calibration_by_length(parent, counts, edges)
|
| 886 |
+
out = {"by_length": {"edges": list(edges), "tokens": "rows" if counted else {"records": "suite", "tokenizer": list(tokenizer)}}}
|
| 887 |
+
for bucket in c:
|
| 888 |
+
for m in LENGTH_METRICS: out[f"{m}_{bucket}"] = {"candidate": c[bucket][m], "parent": p[bucket][m], "n": c[bucket]["n"]}
|
| 889 |
+
return out
|
| 890 |
+
|
| 891 |
+
|
| 892 |
+
def value(report, path):
|
| 893 |
+
"""A number from a comparison by path; None if its panel was not read."""
|
| 894 |
+
parts = path.split(".")
|
| 895 |
+
if parts[0] == "unknowable": return (report.get("unknowable") or {}).get(parts[1])
|
| 896 |
+
entry = report["panels"].get(parts[0], {}).get(parts[1])
|
| 897 |
+
if entry is None: return None
|
| 898 |
+
if parts[2] in ("lower", "upper"): return entry["ci95"][0 if parts[2] == "lower" else 1]
|
| 899 |
+
return entry[parts[2]]
|
| 900 |
+
|
| 901 |
+
|
| 902 |
+
def _criterion(report, c):
|
| 903 |
+
left = [value(report, p) for p in (c["left"] if isinstance(c["left"], list) else [c["left"]])]
|
| 904 |
+
right = value(report, c["right"]) if isinstance(c["right"], str) else c["right"]
|
| 905 |
+
if right is None or any(v is None for v in left): return None
|
| 906 |
+
lhs = left[0] - left[1] if len(left) == 2 else left[0]
|
| 907 |
+
return OPS[c["op"]](lhs, right + c["plus"] if "plus" in c else right)
|
| 908 |
+
|
| 909 |
+
|
| 910 |
+
def rank(arms, keys, selectable):
|
| 911 |
+
"""Passing selectable arms per size, best first by the rank keys (a key summing several paths adds them left to right)."""
|
| 912 |
+
def score(report):
|
| 913 |
+
out = []
|
| 914 |
+
for key in keys:
|
| 915 |
+
paths = key["by"] if isinstance(key["by"], list) else [key["by"]]
|
| 916 |
+
total = value(report, paths[0])
|
| 917 |
+
for p in paths[1:]: total = total + value(report, p)
|
| 918 |
+
out.append(-total if key.get("order", "desc") == "desc" else total)
|
| 919 |
+
return tuple(out)
|
| 920 |
+
by_size = {}
|
| 921 |
+
for arm, report in arms.items():
|
| 922 |
+
by_size.setdefault(size_of(arm), [])
|
| 923 |
+
if report.get("passed") and arm in selectable: by_size[size_of(arm)].append(arm)
|
| 924 |
+
return {size: sorted(names, key=lambda a: score(arms[a])) for size, names in by_size.items()}
|
| 925 |
+
|
| 926 |
+
|
| 927 |
+
def readout(spec, root=ROOT):
|
| 928 |
+
"""The registered rule applied to every arm that has finished training (development rows present); an arm without a
|
| 929 |
+
trial (a checkpoint made from others) is compared once its pool and reads are there (compare lists what is missing)."""
|
| 930 |
+
arms = {}
|
| 931 |
+
for arm, a in spec["arms"].items():
|
| 932 |
+
if a.get("trial") and not (Path(root) / a["trial"] / "development/rows.json").exists():
|
| 933 |
+
arms[arm] = {"trial": a["trial"], "parent": spec["parents"][a["parent"]]["trial"], "missing": [f"candidate:{a['trial']}/development/rows.json"], "complete": False, "passed": None}
|
| 934 |
+
continue
|
| 935 |
+
arms[arm] = compare(arm_side(spec, arm, root), parent_side(spec, arm, root), spec["rule"], lambda panel: panel_lengths(spec, panel))
|
| 936 |
+
ranking = rank(arms, spec["rule"].get("rank", []), {a for a, x in spec["arms"].items() if x.get("select", True)})
|
| 937 |
+
return {"round": spec["round"], "registered": spec["registered"], "drop_ids": sorted(spec.get("drop_ids", [])), "arms": arms,
|
| 938 |
+
"ranking": ranking, "candidates": {size: names[0] if names else None for size, names in ranking.items()}}
|
| 939 |
+
|
| 940 |
+
|
| 941 |
+
def confirm(spec, stage, arm, root=ROOT):
|
| 942 |
+
"""One confirmation stage for the chosen arm against its parent, on the stage's reads."""
|
| 943 |
+
out = compare(arm_side(spec, arm, root, stage), parent_side(spec, arm, root, stage), spec["confirm"][stage], lambda panel: panel_lengths(spec, panel))
|
| 944 |
+
return {"round": spec["round"], "stage": stage, "arm": arm, **out}
|
| 945 |
+
|
| 946 |
+
|
| 947 |
+
# --- printing --------------------------------------------------------------------------------------------------------
|
| 948 |
+
|
| 949 |
+
RATES = ("acc", "confident_error_rate", "coverage_at_5pct_error", "coverage_at_1pct_error") # printed in percentage points
|
| 950 |
+
|
| 951 |
+
|
| 952 |
+
def _interval(metric, entry):
|
| 953 |
+
scale, fmt = (100, "+.1f") if metric in RATES else (1, "+.3f")
|
| 954 |
+
return f"{scale * entry['delta']:{fmt}} [{scale * entry['ci95'][0]:{fmt}}, {scale * entry['ci95'][1]:{fmt}}]"
|
| 955 |
+
|
| 956 |
+
|
| 957 |
+
def table(report):
|
| 958 |
+
"""One line per comparison: every bootstrapped panel metric (rates in pp), the unknowable share, the verdict, what failed;
|
| 959 |
+
under it, each by_length panel's buckets, and a warning when the arm is served at a temperature fitted on its trial's
|
| 960 |
+
development rows of a training corpus (in distribution: round 19's failure mode)."""
|
| 961 |
+
lines = []
|
| 962 |
+
for arm, r in report.get("arms", {report.get("arm"): report}).items():
|
| 963 |
+
cells = [f"{p}.{m} {_interval(m, e)}" for p, panel in r.get("panels", {}).items() for m, e in panel.items() if isinstance(e, dict) and "ci95" in e]
|
| 964 |
+
if r.get("unknowable"): cells.append(f"unk {r['unknowable']['candidate']}")
|
| 965 |
+
failed = [k for k, v in r.get("criteria", {}).items() if v is False]
|
| 966 |
+
verdict = "incomplete" if not r.get("complete") else {True: "PASS", None: "reported", False: "fail: " + ", ".join(failed)}[r["passed"]]
|
| 967 |
+
lines.append(f"{arm:16} {' | '.join(cells)} -> {verdict}" + (f" (missing {len(r['missing'])})" if r.get("missing") else "")
|
| 968 |
+
+ (f" (report-only reads missing {len(r['missing_optional'])})" if r.get("missing_optional") else ""))
|
| 969 |
+
if r.get("temperature_ci"):
|
| 970 |
+
ci = r["temperature_ci"]
|
| 971 |
+
lines.append(f"{'':16} T {r['temperature']:.3f}, {100 * ci['level']:.0f} % bootstrap interval [{ci['lower']:.3f}, {ci['upper']:.3f}] ({ci['questions']} pool questions)")
|
| 972 |
+
for p, panel in r.get("panels", {}).items():
|
| 973 |
+
if "by_length" not in panel: continue
|
| 974 |
+
buckets = [name for name, _, _ in length_buckets(tuple(panel["by_length"]["edges"])) if panel[f"ece_{name}"]["n"]]
|
| 975 |
+
lines.append(f"{'':16} {p} by state length (candidate/parent): " + " | ".join(
|
| 976 |
+
f"{b} n={panel[f'ece_{b}']['n']} acc {panel[f'acc_{b}']['candidate']:.3f}/{panel[f'acc_{b}']['parent']:.3f} ece {panel[f'ece_{b}']['candidate']:.4f}/{panel[f'ece_{b}']['parent']:.4f}" for b in buckets))
|
| 977 |
+
source = r.get("temperature_source") or {}
|
| 978 |
+
if source.get("kind") == "trial development rows" and source.get("training_corpus"):
|
| 979 |
+
lines.append(f"{'':16} !!! {arm} is served at T {r['temperature']:.3f} fitted on {source['rows']}, held-out items of its training corpus {source['suite']}: "
|
| 980 |
+
f"in distribution, {ROUND_19.split(': ', 1)[0]}; not a temperature to ship (register a held-out-datasets `temperature` pool)")
|
| 981 |
+
if "candidates" in report: lines.append(f"candidates {report['candidates']}")
|
| 982 |
+
return "\n".join(lines)
|
| 983 |
+
|
| 984 |
+
|
| 985 |
+
# --- Modal orchestration ---------------------------------------------------------------------------------------------
|
| 986 |
+
|
| 987 |
+
def bench_job(run, suite, name, flags=""):
|
| 988 |
+
"""One modal_app.py::benchmarks entry. modal_app.parse_jobs reads the run as everything before the suite, so a pinned
|
| 989 |
+
Hub revision (repo@sha) is safe; the suite and the name must not contain '@' or ',' and flags must start with '--'."""
|
| 990 |
+
if any(c in s for s in (suite, name) for c in "@,") or "," in run or (flags and not flags.startswith("--")):
|
| 991 |
+
raise ValueError(f"cannot encode benchmark job {(run, suite, name, flags)}")
|
| 992 |
+
return "@".join([run, suite, name, *([flags] if flags else [])])
|
| 993 |
+
|
| 994 |
+
|
| 995 |
+
def checkpoint_of(entry):
|
| 996 |
+
"""Where a read runs from: the trial's checkpoint on the runs volume (runs/X -> /runs/X/checkpoint) unless declared."""
|
| 997 |
+
return entry["checkpoint"] if "checkpoint" in entry else f"/{entry['trial']}/checkpoint"
|
| 998 |
+
|
| 999 |
+
|
| 1000 |
+
def side_reads(spec, arm, rule, who, side, stage=None):
|
| 1001 |
+
"""[(read tag, directory)] a side needs for a stage, by tag: the rule's reads; "transfer" as the arm's transfer_read (a
|
| 1002 |
+
trial's in-trial transfer read is never launched); and, for an arm at the rule stage, its temperature pool."""
|
| 1003 |
+
tags = rule_tags(rule)
|
| 1004 |
+
if who == "candidate" and stage is None and spec.get("temperature"):
|
| 1005 |
+
tags |= {*spec["temperature"]["reads"], *spec["temperature"].get("exclude_reads", [])}
|
| 1006 |
+
reads = {}
|
| 1007 |
+
for tag in tags:
|
| 1008 |
+
if tag != TRANSFER: reads.setdefault(side.dirs[tag], tag)
|
| 1009 |
+
elif who == "candidate" and transfer_read(spec, arm): reads.setdefault(side.dirs[tag], transfer_read(spec, arm))
|
| 1010 |
+
return sorted((read, d) for d, read in reads.items())
|
| 1011 |
+
|
| 1012 |
+
|
| 1013 |
+
def read_commands(spec, arm, stage=None, root=ROOT, sides=("candidate",)):
|
| 1014 |
+
"""The `modal run` commands for the missing reads of one arm (and, with sides, its parent): one benchmarks call per side
|
| 1015 |
+
with every suite batched, plus one locked_test call per locked read. Reads that exist locally are skipped, and so are
|
| 1016 |
+
reads of a removed suite (kev.suite.REMOVED_SUITES; the panels they feed stay missing)."""
|
| 1017 |
+
rule = spec["confirm"][stage] if stage else spec["rule"]
|
| 1018 |
+
size, commands = size_of(arm), []
|
| 1019 |
+
for who in sides:
|
| 1020 |
+
side = (arm_side if who == "candidate" else parent_side)(spec, arm, root, stage)
|
| 1021 |
+
run = checkpoint_of(spec["arms"][arm] if who == "candidate" else spec["parents"][spec["arms"][arm]["parent"]])
|
| 1022 |
+
jobs = []
|
| 1023 |
+
for tag, d in side_reads(spec, arm, rule, who, side, stage):
|
| 1024 |
+
if (root / d / "rows.json").exists(): continue
|
| 1025 |
+
r = spec["reads"][tag]
|
| 1026 |
+
if r.get("entrypoint") == "locked_test":
|
| 1027 |
+
if not run.startswith("/runs/"): raise ValueError(f"locked_test reads a trial on the runs volume, not {run}")
|
| 1028 |
+
commands.append(["modal", "run", "modal_app.py::locked_test", "--trial", run.removeprefix("/runs/").removesuffix("/checkpoint"),
|
| 1029 |
+
"--name", Path(d).relative_to("runs/locked").parts[0], "--decision", r["decision"], "--gpu", spec.get("gpu", "H100"),
|
| 1030 |
+
*spec.get("locked_args", {}).get(size, [])])
|
| 1031 |
+
continue
|
| 1032 |
+
if not d.startswith("runs/") or "/" in d.removeprefix("runs/"): raise ValueError(f"{tag}: benchmarks writes runs/<name>, not {d}")
|
| 1033 |
+
if gone := removed_suite(r["suite"]): # never made again; skipped, not raised, so a watcher does not retry it forever
|
| 1034 |
+
print(f"!!! {arm}: not launching read {tag} ({d}): {r['suite']} was removed on {gone['removed']} (kev.suite.REMOVED_SUITES)", flush=True)
|
| 1035 |
+
continue
|
| 1036 |
+
jobs.append(bench_job(run, r["suite"], d.removeprefix("runs/"), r.get("flags", "")))
|
| 1037 |
+
if jobs:
|
| 1038 |
+
timeout = spec.get("read_timeout", {}).get(size)
|
| 1039 |
+
commands.append(["modal", "run", "--detach", "modal_app.py::benchmarks", "--jobs", ",".join(jobs), "--gpu", spec.get("gpu", "H100"),
|
| 1040 |
+
*(["--timeout", str(timeout)] if timeout else [])])
|
| 1041 |
+
return commands
|
| 1042 |
+
|
| 1043 |
+
|
| 1044 |
+
def modal_env(spec):
|
| 1045 |
+
return {**os.environ, **({"KEV_APP_NAME": spec["app"]} if spec.get("app") else {}), **({"KEV_GPU": spec["gpu"]} if spec.get("gpu") else {})}
|
| 1046 |
+
|
| 1047 |
+
|
| 1048 |
+
def launch_commands(commands, spec, log_stem, stagger=STAGGER, run=subprocess.Popen, sleep=time.sleep):
|
| 1049 |
+
"""Start each command in the background (log under runs/), `stagger` seconds apart; returns the processes."""
|
| 1050 |
+
procs = []
|
| 1051 |
+
for i, cmd in enumerate(commands):
|
| 1052 |
+
if i: sleep(stagger)
|
| 1053 |
+
log = ROOT / "runs" / f"{log_stem}-{i}.log"
|
| 1054 |
+
print("launch:", " ".join(cmd[:4]), "->", log.relative_to(ROOT), flush=True)
|
| 1055 |
+
with log.open("w", encoding="utf-8") as f:
|
| 1056 |
+
procs.append(run([sys.executable, "-m", *cmd], stdout=f, stderr=subprocess.STDOUT, cwd=ROOT, env=modal_env(spec)))
|
| 1057 |
+
return procs
|
| 1058 |
+
|
| 1059 |
+
|
| 1060 |
+
class InFlight(Exception):
|
| 1061 |
+
"""An arm's reads were launched recently and may still be running; they are not launched again yet."""
|
| 1062 |
+
|
| 1063 |
+
|
| 1064 |
+
def launch_arm_reads(spec, arm, stage=None, sides=("candidate",), stagger=STAGGER, launch=None, now=time.time, log=print):
|
| 1065 |
+
"""Launch one arm's missing reads once. Under a per-arm lock (runs/.reads-r<N>-<arm>.lock: a watcher and a
|
| 1066 |
+
`launch-reads` cannot both launch the arm), the intent is written first (runs/r<N>-reads-<arm>[-<stage>].json: when and
|
| 1067 |
+
what), then the commands start. Detached benchmarks cannot be seen from here, so while an earlier intent is younger
|
| 1068 |
+
than the reads' timeout the arm counts as in flight: InFlight is raised (the watcher retries on a later pass) instead of
|
| 1069 |
+
a second launch whose job would find /runs/bench/<name> taken. Reads whose rows landed are never relaunched."""
|
| 1070 |
+
tag = f"r{spec['round']}-reads-{arm}" + (f"-{stage}" if stage else "")
|
| 1071 |
+
grace = max(IN_FLIGHT, spec.get("read_timeout", {}).get(size_of(arm), 0))
|
| 1072 |
+
with file_lock(ROOT / "runs" / f".reads-r{spec['round']}-{arm}.lock"):
|
| 1073 |
+
commands = read_commands(spec, arm, stage, ROOT, sides)
|
| 1074 |
+
if not commands: log(f"{arm}: every read has landed"); return []
|
| 1075 |
+
intent = ROOT / "runs" / f"{tag}.json"
|
| 1076 |
+
if intent.exists() and now() - read_json(intent)["launched_at"] < grace:
|
| 1077 |
+
since = read_json(intent)["launched_at"]
|
| 1078 |
+
raise InFlight(f"{arm}: reads launched {int(now() - since)} s ago may still be running; not relaunching for another "
|
| 1079 |
+
f"{int(since + grace - now())} s (if they finished on Modal but were never pulled: modal volume get kev-runs /bench/<name> runs/)")
|
| 1080 |
+
write_json(intent, {"launched_at": now(), "commands": commands}, atomic=True)
|
| 1081 |
+
return (launch or launch_commands)(commands, spec, tag, stagger)
|
| 1082 |
+
|
| 1083 |
+
|
| 1084 |
+
def pull(study, spec):
|
| 1085 |
+
"""modal_app.py::pull for one study (modal_app.pull_study holds a per-study lock, so concurrent pulls wait)."""
|
| 1086 |
+
subprocess.run([sys.executable, "-m", "modal", "run", "modal_app.py::pull", "--name", study], check=True, cwd=ROOT, env=modal_env(spec))
|
| 1087 |
+
|
| 1088 |
+
|
| 1089 |
+
def continue_command(spec, study, label):
|
| 1090 |
+
s = spec["studies"][study]
|
| 1091 |
+
return ["modal", "run", "modal_app.py::resume", "--study", study, "--suite", s["suite"], "--transfer", s["transfer"], "--trial", label]
|
| 1092 |
+
|
| 1093 |
+
|
| 1094 |
+
def continue_trial(spec, study, label, run=subprocess.run):
|
| 1095 |
+
"""The next attempt of a timed-out trial: modal_app.py::resume --trial (modal_app.continue_full_trial) spawns
|
| 1096 |
+
run_full_trial again, which continues from the trial's last resume point, with the GPU and timeout the study was
|
| 1097 |
+
admitted for, and records the call in runs/<study>.spawn.json; for a pending ledger entry it first adopts the call
|
| 1098 |
+
from the trial's lease (or waits). Returns once the record names a new current call (spawned or adopted). Raises
|
| 1099 |
+
NoContinuation when the study has no attempt ledger, when the ledger has no attempt left (kev.budget.trial_attempts,
|
| 1100 |
+
the count compute_bound admitted; a pending last entry is resolved by the command first) or when the command
|
| 1101 |
+
succeeded without recording a call, and RuntimeError when it refused (a live lease, a young pending entry, the
|
| 1102 |
+
network: the watcher asks again on its next pass). Output in runs/<study>.continue-<label>-<attempt>.log."""
|
| 1103 |
+
from kev.budget import trial_attempts
|
| 1104 |
+
path = ROOT / "runs" / f"{study}.spawn.json"
|
| 1105 |
+
record = read_json(path)
|
| 1106 |
+
if record.get("modal_retries") != 0: raise NoContinuation(f"{path.name} has no attempt ledger (spawned with Modal's retries on, before the ledger)")
|
| 1107 |
+
used, allowed = trial_attempts(record, label)
|
| 1108 |
+
last = record["attempts"][label][-1]
|
| 1109 |
+
if used >= allowed and (last["call"] is not None or last.get("abandoned")):
|
| 1110 |
+
raise NoContinuation(f"{used} of {allowed} attempts used; the admission bound (${record.get('bound_usd')}) counts no more")
|
| 1111 |
+
log = ROOT / "runs" / f"{study}.continue-{label}-{used + 1}.log"
|
| 1112 |
+
print("continue:", " ".join(continue_command(spec, study, label)), "->", log.relative_to(ROOT), flush=True)
|
| 1113 |
+
with log.open("a", encoding="utf-8") as f:
|
| 1114 |
+
done = run([sys.executable, "-m", *continue_command(spec, study, label)], stdout=f, stderr=subprocess.STDOUT, cwd=ROOT, env=modal_env(spec))
|
| 1115 |
+
if read_json(path)["calls"][label] != record["calls"][label]: return
|
| 1116 |
+
if getattr(done, "returncode", 0): raise RuntimeError(f"modal_app.py::resume --trial exited {done.returncode} without a new call (see {log.relative_to(ROOT)})")
|
| 1117 |
+
raise NoContinuation(f"modal_app.py::resume recorded no new call (see {log.relative_to(ROOT)})")
|
| 1118 |
+
|
| 1119 |
+
|
| 1120 |
+
# --- watch -----------------------------------------------------------------------------------------------------------
|
| 1121 |
+
|
| 1122 |
+
NETWORK_MARKERS = ("nodename nor servname", "name or service not known", "temporary failure in name resolution", "unavailable", "connection reset",
|
| 1123 |
+
"connection refused", "network is unreachable", "deadline exceeded")
|
| 1124 |
+
|
| 1125 |
+
|
| 1126 |
+
def transient(error):
|
| 1127 |
+
"""Whether polling a call failed because of this machine's network (DNS drop, reset connection, gRPC UNAVAILABLE)
|
| 1128 |
+
rather than because the trial failed. A trial's own exception is re-raised by FunctionCall.get with its original
|
| 1129 |
+
type, so only network types and gRPC transport messages count; anything else is the trial's failure."""
|
| 1130 |
+
try:
|
| 1131 |
+
import modal.exception as me
|
| 1132 |
+
modal_types = (me.ConnectionError, me.ClientClosed)
|
| 1133 |
+
except ImportError:
|
| 1134 |
+
modal_types = ()
|
| 1135 |
+
if isinstance(error, (socket.gaierror, ConnectionError, *modal_types)): return True
|
| 1136 |
+
return type(error).__module__.startswith(("grpclib", "modal")) and any(m in str(error).lower() for m in NETWORK_MARKERS)
|
| 1137 |
+
|
| 1138 |
+
|
| 1139 |
+
class TrialFailed(Exception):
|
| 1140 |
+
"""A full-weight trial that failed with an error returns {"failed": ...} (modal_app.failed_trial) instead of raising;
|
| 1141 |
+
poll_modal raises this for it, which watch_studies marks failed like any trial error (never continued)."""
|
| 1142 |
+
|
| 1143 |
+
|
| 1144 |
+
class NoContinuation(Exception):
|
| 1145 |
+
"""A timed-out trial that is not continued (no attempt left under its study's admission bound, no attempt ledger, or
|
| 1146 |
+
the continuation spawned no call): watch_studies marks it failed."""
|
| 1147 |
+
|
| 1148 |
+
|
| 1149 |
+
def poll_modal(call_id):
|
| 1150 |
+
"""'running' | 'done' | 'timeout' (the call ran out of time: its container was cancelled at the timeout) | 'refused'
|
| 1151 |
+
(a full-weight attempt that found another attempt's fresh lease and did not start: modal_app.TrialLease) for a spawned
|
| 1152 |
+
trial; the trial's exception (or a network error) propagates, and a returned failure is raised as TrialFailed."""
|
| 1153 |
+
import modal
|
| 1154 |
+
try:
|
| 1155 |
+
result = modal.FunctionCall.from_id(call_id).get(timeout=0.5)
|
| 1156 |
+
except TimeoutError: # builtin: no output yet (modal.exception.FunctionTimeoutError is not a builtin TimeoutError)
|
| 1157 |
+
return "running"
|
| 1158 |
+
except modal.exception.FunctionTimeoutError:
|
| 1159 |
+
return "timeout"
|
| 1160 |
+
except modal.exception.OutputExpiredError: # finished long ago; the result is on the volume
|
| 1161 |
+
return "done"
|
| 1162 |
+
if isinstance(result, dict) and "failed" in result: raise TrialFailed(result["failed"])
|
| 1163 |
+
if isinstance(result, dict) and "refused" in result: return "refused"
|
| 1164 |
+
return "done"
|
| 1165 |
+
|
| 1166 |
+
|
| 1167 |
+
class Unmapped(Exception):
|
| 1168 |
+
"""A finished call trained no arm of the spec: its reads cannot be launched."""
|
| 1169 |
+
|
| 1170 |
+
|
| 1171 |
+
def watch_studies(studies, on_done, poll=poll_modal, interval=120, max_transient=60, sleep=time.sleep, log=print, root=ROOT, now=time.time, on_timeout=None):
|
| 1172 |
+
"""Poll every spawned trial of the studies until each is settled, calling on_done(study, label) for a finished trial
|
| 1173 |
+
until it succeeds; returns the finished calls that map to no arm. State lives in runs/<study>.watch.json, replaced
|
| 1174 |
+
atomically after every change, so a restarted watcher resumes: finished trials are not polled again and launched reads
|
| 1175 |
+
are not launched again. `launching_at` is written before on_done runs; a restart that finds it logs the interrupted
|
| 1176 |
+
launch and calls on_done again, which must check what already happened (kev.rounds.launch_arm_reads does). Network
|
| 1177 |
+
errors while polling are retried (max_transient in a row marks the call failed); an on_done that raises (a pull that
|
| 1178 |
+
lost the network, reads still in flight) is retried on the next pass; one that raises Unmapped leaves the call
|
| 1179 |
+
unlaunched, logged, and settled so the watch can end.
|
| 1180 |
+
A call that ends by a timeout goes to on_timeout(study, label) (watch's: the trial's next attempt, which records its call
|
| 1181 |
+
as the label's current call in runs/<study>.spawn.json); the spawn record is read on every pass, and a label whose
|
| 1182 |
+
call changed (a continuation, from here or by hand) is polled again from its new call. Without on_timeout, or when it
|
| 1183 |
+
raises NoContinuation, the timeout is the trial's failure; any other error from it is retried on the next pass
|
| 1184 |
+
(max_transient in a row marks the trial failed). `continuing_at` is written before on_timeout runs; a restart that
|
| 1185 |
+
finds it and a changed call adopts the call, one that finds the call unchanged asks again (modal_app.continue_full_trial
|
| 1186 |
+
refuses a second attempt while one runs, waits for the old container's lease and counts every attempt in the ledger).
|
| 1187 |
+
A call refused by the trial's lease, and a current attempt with no call recorded (a pending ledger entry), go to
|
| 1188 |
+
on_timeout the same way."""
|
| 1189 |
+
runs = Path(root) / "runs"
|
| 1190 |
+
while True:
|
| 1191 |
+
settled = True
|
| 1192 |
+
calls = {study: read_json(runs / f"{study}.spawn.json")["calls"] for study in studies} # a continuation replaces a label's call
|
| 1193 |
+
for study, study_calls in calls.items():
|
| 1194 |
+
path = runs / f"{study}.watch.json"
|
| 1195 |
+
state = read_json(path) if path.exists() else {"calls": {}}
|
| 1196 |
+
for label, call_id in study_calls.items():
|
| 1197 |
+
s = state["calls"].setdefault(label, {"status": "running", "launched": False, "transient": 0})
|
| 1198 |
+
if s.setdefault("call", call_id) != call_id and s["status"] != "done": # the trial's next attempt
|
| 1199 |
+
log(f"{study}/{label}: polling its continuation {call_id} (was {s['call']})")
|
| 1200 |
+
s.update(call=call_id, status="running", transient=0); s.pop("error", None); s.pop("continuing_at", None)
|
| 1201 |
+
if s["status"] == "running" and call_id is None: # a pending ledger entry: spawned, or about to be, with no call recorded
|
| 1202 |
+
s["status"] = "timeout"; log(f"{study}/{label}: no call recorded for its current attempt (pending); a continuation adopts it from its lease or waits")
|
| 1203 |
+
elif s["status"] == "running":
|
| 1204 |
+
try:
|
| 1205 |
+
s["status"], s["transient"] = poll(call_id), 0
|
| 1206 |
+
except Exception as error: # noqa: BLE001 - a trial's own failure or this machine's network, told apart here
|
| 1207 |
+
if transient(error) and s["transient"] + 1 < max_transient:
|
| 1208 |
+
s["transient"] += 1; log(f"{study}/{label}: network error, retrying ({type(error).__name__}: {str(error)[:120]})")
|
| 1209 |
+
else:
|
| 1210 |
+
s["status"], s["error"] = "failed", f"{type(error).__name__}: {str(error)[:300]}"; log(f"{study}/{label}: FAILED {s['error']}")
|
| 1211 |
+
if s["status"] == "refused": # the attempt found another attempt's fresh lease and never started: it needs a next attempt too
|
| 1212 |
+
s["status"] = "timeout"; log(f"{study}/{label}: its call {call_id} was refused by the trial's lease (another attempt's container was alive)")
|
| 1213 |
+
if s["status"] == "timeout":
|
| 1214 |
+
if on_timeout is None:
|
| 1215 |
+
s["status"], s["error"] = "failed", "FunctionTimeoutError (not continued: no on_timeout)"; log(f"{study}/{label}: FAILED {s['error']}")
|
| 1216 |
+
else:
|
| 1217 |
+
if "continuing_at" in s: log(f"{study}/{label}: a continuation asked for at {s['continuing_at']:.0f} spawned no call it recorded; asking again")
|
| 1218 |
+
s["continuing_at"] = now(); write_json(path, state, atomic=True)
|
| 1219 |
+
try:
|
| 1220 |
+
on_timeout(study, label); log(f"{study}/{label}: timed out; its next attempt was spawned")
|
| 1221 |
+
except NoContinuation as error:
|
| 1222 |
+
s["status"], s["error"] = "failed", f"FunctionTimeoutError, not continued: {error}"; s.pop("continuing_at"); log(f"!!! {study}/{label}: FAILED {s['error']}")
|
| 1223 |
+
except Exception as error: # noqa: BLE001 - the network, or a refusal that may clear; retried on the next pass
|
| 1224 |
+
s["transient"] += 1; s.pop("continuing_at")
|
| 1225 |
+
if s["transient"] >= max_transient:
|
| 1226 |
+
s["status"], s["error"] = "failed", f"FunctionTimeoutError, continuation failed {s['transient']} times: {type(error).__name__}: {str(error)[:300]}"; log(f"!!! {study}/{label}: FAILED {s['error']}")
|
| 1227 |
+
else:
|
| 1228 |
+
log(f"{study}/{label}: continuing failed, retrying next pass ({type(error).__name__}: {str(error)[:300]})")
|
| 1229 |
+
if s["status"] == "done" and not s["launched"] and not s.get("unmapped"):
|
| 1230 |
+
if "launching_at" in s: log(f"{study}/{label}: a launch started at {s['launching_at']:.0f} was interrupted; checking the arm's reads before any relaunch")
|
| 1231 |
+
s["launching_at"] = now(); write_json(path, state, atomic=True)
|
| 1232 |
+
try:
|
| 1233 |
+
on_done(study, label); s["launched"] = True; s.pop("launching_at"); log(f"{study}/{label}: done, reads launched")
|
| 1234 |
+
except Unmapped as error:
|
| 1235 |
+
s["unmapped"] = True; s.pop("launching_at"); log(f"!!! {study}/{label}: finished but {error}; its reads were NOT launched")
|
| 1236 |
+
except Exception as error: # noqa: BLE001 - retried on the next pass; only a dead process leaves launching_at behind
|
| 1237 |
+
s.pop("launching_at"); log(f"{study}/{label}: launching reads failed, retrying next pass ({type(error).__name__}: {str(error)[:300]})")
|
| 1238 |
+
write_json(path, state, atomic=True)
|
| 1239 |
+
settled = settled and (s["status"] == "failed" or s["launched"] or bool(s.get("unmapped")))
|
| 1240 |
+
if settled:
|
| 1241 |
+
return [f"{study}/{label}" for study in calls for label, s in read_json(runs / f"{study}.watch.json")["calls"].items() if s.get("unmapped")]
|
| 1242 |
+
sleep(interval)
|
| 1243 |
+
|
| 1244 |
+
|
| 1245 |
+
def arm_of(spec, study, label):
|
| 1246 |
+
"""The arm a spawned call trained (trial directories are <NN>-<label>)."""
|
| 1247 |
+
return next((a for a, x in spec["arms"].items() if x.get("trial", "").startswith(f"runs/{study}/") and Path(x["trial"]).name.split("-", 1)[1] == label), None)
|
| 1248 |
+
|
| 1249 |
+
|
| 1250 |
+
def watch(spec, interval=120, stagger=STAGGER, reads_timeout=6 * 3600):
|
| 1251 |
+
"""The round from spawned trials to read-out: when a trial finishes, pull its study and launch that arm's reads
|
| 1252 |
+
(`stagger` apart); once every trial is settled, wait for the reads to land and write the read-out. The wait ends when
|
| 1253 |
+
every read is there, when every benchmarks process this watcher started has exited (a failed read never lands), or
|
| 1254 |
+
after reads_timeout (a restarted watcher has no processes to follow)."""
|
| 1255 |
+
last, procs = [0.0], []
|
| 1256 |
+
def on_done(study, label):
|
| 1257 |
+
arm = arm_of(spec, study, label)
|
| 1258 |
+
if arm is None: raise Unmapped(f"no arm of round {spec['round']} has a trial runs/{study}/<NN>-{label}")
|
| 1259 |
+
pull(study, spec)
|
| 1260 |
+
time.sleep(max(0.0, last[0] + stagger - time.time()))
|
| 1261 |
+
procs.extend(launch_arm_reads(spec, arm, stagger=stagger))
|
| 1262 |
+
last[0] = time.time()
|
| 1263 |
+
unmapped = watch_studies(list(spec.get("studies", {})), on_done, interval=interval, on_timeout=lambda study, label: continue_trial(spec, study, label))
|
| 1264 |
+
deadline = time.time() + reads_timeout
|
| 1265 |
+
finished = [a for a, x in spec["arms"].items() if x.get("trial") and (ROOT / x["trial"] / "result.json").exists()]
|
| 1266 |
+
landed = lambda a: all((ROOT / d / "rows.json").exists() for tag, d in side_reads(spec, a, spec["rule"], "candidate", arm_side(spec, a))
|
| 1267 |
+
if not removed_suite(spec["reads"][tag].get("suite") or "")) # a removed suite's read is never launched
|
| 1268 |
+
while (waiting := [a for a in finished if not landed(a)]) and time.time() < deadline:
|
| 1269 |
+
if procs and all(p.poll() is not None for p in procs): break
|
| 1270 |
+
print(f"waiting for the reads of {waiting}", flush=True); time.sleep(interval)
|
| 1271 |
+
report = write_readout(spec)
|
| 1272 |
+
if unmapped: raise SystemExit(f"finished calls that map to no arm, never read: {unmapped} (fix the spec's arms, then launch-reads)")
|
| 1273 |
+
return report
|
| 1274 |
+
|
| 1275 |
+
|
| 1276 |
+
def launchable(spec, rows=True):
|
| 1277 |
+
"""Problems that stop a launch: a recorded round is never relaunched (register a new one), and a new round needs its
|
| 1278 |
+
plans, suites and (for trials) its parents' reads, or its read-out could not be computed."""
|
| 1279 |
+
if spec.get("archive"): return [f"round {spec['round']} is a record (evidence on {spec['archive']}); register a new round to train or read again"]
|
| 1280 |
+
return validate(spec, rows=rows).problems
|
| 1281 |
+
|
| 1282 |
+
|
| 1283 |
+
def study_commands(spec):
|
| 1284 |
+
return {name: ["modal", "run", "modal_app.py::study", "--suite", s["suite"], "--plan", s["plan"], "--name", name, "--transfer", s["transfer"],
|
| 1285 |
+
"--budget", str(s["budget"]), "--timeout", str(s["timeout"]), "--gpu", s["gpu"]] for name, s in spec["studies"].items()}
|
| 1286 |
+
|
| 1287 |
+
|
| 1288 |
+
def launch_studies(spec, stagger=STAGGER, run=subprocess.run, sleep=time.sleep):
|
| 1289 |
+
"""modal_app.py::study for every study of the round (each spawns its trials on the deployed app and returns), `stagger`
|
| 1290 |
+
apart, output in runs/<study>.log (a filter on this output can hide the SystemExit that explains a refusal)."""
|
| 1291 |
+
for i, (name, cmd) in enumerate(study_commands(spec).items()):
|
| 1292 |
+
if i: sleep(stagger)
|
| 1293 |
+
print("launch:", " ".join(cmd), flush=True)
|
| 1294 |
+
with (ROOT / "runs" / f"{name}.log").open("w", encoding="utf-8") as f:
|
| 1295 |
+
run([sys.executable, "-m", *cmd], stdout=f, stderr=subprocess.STDOUT, cwd=ROOT, env=modal_env(spec), check=True)
|
| 1296 |
+
|
| 1297 |
+
|
| 1298 |
+
def write_readout(spec, root=ROOT, out=None):
|
| 1299 |
+
report = readout(spec, root)
|
| 1300 |
+
out = Path(out or Path(root) / f"runs/r{spec['round']}-readout"); out.mkdir(parents=True, exist_ok=True)
|
| 1301 |
+
write_json(out / f"round{spec['round']}.json", report); print(table(report))
|
| 1302 |
+
return report
|
| 1303 |
+
|
| 1304 |
+
|
| 1305 |
+
# --- CLI -------------------------------------------------------------------------------------------------------------
|
| 1306 |
+
|
| 1307 |
+
def main(argv=None):
|
| 1308 |
+
ap = argparse.ArgumentParser(prog="python -m kev.rounds", description=__doc__.split("\n\n")[0])
|
| 1309 |
+
sub = ap.add_subparsers(dest="cmd", required=True)
|
| 1310 |
+
for name in ("validate", "launch", "launch-reads", "watch", "readout", "confirm"):
|
| 1311 |
+
p = sub.add_parser(name); p.add_argument("spec")
|
| 1312 |
+
if name in ("validate", "readout", "confirm"): p.add_argument("--root", default=str(ROOT), help="checkout whose runs/ holds the rows (default: this one)")
|
| 1313 |
+
if name == "validate": p.add_argument("--partitions", action="store_true", help="verify the suites' partitions through load_plan (may fetch from the Hub)")
|
| 1314 |
+
if name in ("launch", "launch-reads", "watch"): p.add_argument("--stagger", type=int, default=STAGGER)
|
| 1315 |
+
if name in ("launch", "launch-reads"): p.add_argument("--dry-run", action="store_true", help="print the modal commands only")
|
| 1316 |
+
if name == "launch-reads": p.add_argument("--arms", help="comma-separated arms (default: every arm with a pulled result, and every arm that is a checkpoint without a trial)"); p.add_argument("--parents", action="store_true", help="also the parents' missing reads")
|
| 1317 |
+
if name in ("launch-reads", "confirm"): p.add_argument("--stage", required=name == "confirm"); p.add_argument("--arm", help="the candidate (default: the one the readout names)")
|
| 1318 |
+
if name == "watch": p.add_argument("--interval", type=int, default=120)
|
| 1319 |
+
if name in ("readout", "confirm"): p.add_argument("--out")
|
| 1320 |
+
a = ap.parse_args(argv)
|
| 1321 |
+
spec, root = load(a.spec), Path(getattr(a, "root", ROOT))
|
| 1322 |
+
for line in [*calibration_warnings(spec, root), *continuation_warnings(spec, root)] if a.cmd in ("validate", "launch") else ():
|
| 1323 |
+
print(f"!!! warning: {line}")
|
| 1324 |
+
if a.cmd == "validate":
|
| 1325 |
+
problems, archived = validate(spec, root, partitions=a.partitions)
|
| 1326 |
+
where = f"on {spec['archive']}, not in this checkout" if spec.get("archive") else "not in this checkout"
|
| 1327 |
+
if archived: print(f"archived ({where}):\n " + "\n ".join(archived))
|
| 1328 |
+
print("\n".join(problems) or f"round {spec['round']}: ok ({len(spec['arms'])} arms, {len(spec.get('studies', {}))} studies)")
|
| 1329 |
+
raise SystemExit(1 if problems else 0)
|
| 1330 |
+
if a.cmd == "readout":
|
| 1331 |
+
write_readout(spec, root, a.out); return
|
| 1332 |
+
if a.cmd == "confirm":
|
| 1333 |
+
arm = a.arm or _candidate(spec, root)
|
| 1334 |
+
report = confirm(spec, a.stage, arm, root)
|
| 1335 |
+
out = Path(a.out or root / f"runs/r{spec['round']}-verdict"); out.mkdir(parents=True, exist_ok=True)
|
| 1336 |
+
write_json(out / f"{size_of(arm)}-{a.stage}.json", report); print(table(report)); return
|
| 1337 |
+
problems = launchable(spec, rows=a.cmd != "launch-reads") # launch-reads is how a missing parent read gets made
|
| 1338 |
+
if problems: raise SystemExit("spec does not validate for a launch:\n" + "\n".join(problems))
|
| 1339 |
+
if a.cmd == "launch":
|
| 1340 |
+
if a.dry_run: print("\n".join(" ".join(c) for c in study_commands(spec).values()))
|
| 1341 |
+
else: launch_studies(spec, a.stagger)
|
| 1342 |
+
return
|
| 1343 |
+
if a.cmd == "watch":
|
| 1344 |
+
watch(spec, a.interval, a.stagger); return
|
| 1345 |
+
arms = [a.arm or _candidate(spec, root)] if a.stage else (a.arms.split(",") if a.arms else [x for x, v in spec["arms"].items() if not v.get("trial") or (root / v["trial"] / "result.json").exists()])
|
| 1346 |
+
sides = ("candidate", "parent") if a.parents or a.stage else ("candidate",)
|
| 1347 |
+
if a.dry_run:
|
| 1348 |
+
for arm in arms: print("\n".join(" ".join(c) for c in read_commands(spec, arm, a.stage, root, sides)))
|
| 1349 |
+
return
|
| 1350 |
+
for i, arm in enumerate(arms):
|
| 1351 |
+
if i: time.sleep(a.stagger)
|
| 1352 |
+
try: launch_arm_reads(spec, arm, a.stage, sides, a.stagger)
|
| 1353 |
+
except InFlight as error: print(f"!!! {error}")
|
| 1354 |
+
|
| 1355 |
+
|
| 1356 |
+
def _candidate(spec, root):
|
| 1357 |
+
report = read_json(Path(root) / f"runs/r{spec['round']}-readout/round{spec['round']}.json")
|
| 1358 |
+
chosen = [arm for arm in report["candidates"].values() if arm]
|
| 1359 |
+
if len(chosen) != 1: raise SystemExit(f"the readout names {len(chosen)} candidates ({report['candidates']}); pass --arm")
|
| 1360 |
+
return chosen[0]
|
| 1361 |
+
|
| 1362 |
+
|
| 1363 |
+
if __name__ == "__main__":
|
| 1364 |
+
main()
|
runtime/kev/serve.py
ADDED
|
@@ -0,0 +1,343 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""FastAPI sidecar for the playground: loads one checkpoint, exposes prefill-only decisions.
|
| 2 |
+
|
| 3 |
+
Run: uv run --extra serve python -m kev.serve --run runs/kev --port 8008
|
| 4 |
+
|
| 5 |
+
TypeSafe-compatible: POST /v1/systemone, GET /v1/models, the `x-typesafe-request-id` response header, and bearer auth
|
| 6 |
+
when KEV_API_KEY is set (unset = open server, the local default). Demo extras: POST /v1/systemone/permute (one Choice
|
| 7 |
+
under several option orders) and POST /v1/systemone/separate (each question in its own pass, for the packed-vs-separate
|
| 8 |
+
comparison). KEV_PREFIX_CACHE / KEV_PREFIX_MIN_TOKENS / KEV_PREFIX_MAX_TOKENS size the state-prefix cache; KEV_DATE_FACTS=1 opts into the
|
| 9 |
+
date preprocessing (api.with_date_facts). A state over kev.model.SERVE_MAX_STATE tokens is refused with a 422 (kev.model.admit);
|
| 10 |
+
KEV_TRUNCATE_STATES=1 reads its first SERVE_MAX_STATE tokens instead, and then every response says whether it did. Backend
|
| 11 |
+
and precision follow LoadOptions (KEV_BACKEND, KEV_DTYPE, ...): on Apple Silicon the hybrid Qwen3.5 checkpoints run on MLX
|
| 12 |
+
by default, elsewhere on torch in bf16.
|
| 13 |
+
"""
|
| 14 |
+
import argparse, asyncio, atexit, hmac, os, queue, random, sys, threading, time, traceback, uuid
|
| 15 |
+
from concurrent.futures import Future
|
| 16 |
+
import torch
|
| 17 |
+
from dataclasses import dataclass, field, replace
|
| 18 |
+
from fastapi import FastAPI, HTTPException
|
| 19 |
+
from fastapi.middleware.cors import CORSMiddleware
|
| 20 |
+
from fastapi.responses import JSONResponse
|
| 21 |
+
from pydantic import BaseModel, Field
|
| 22 |
+
from .api import SystemOneRequest, to_record, to_answers, output_tokens, with_date_facts
|
| 23 |
+
from .checkpoint import Checkpoint, LoadOptions, fused_available, is_hub_id
|
| 24 |
+
from .device import default_device, empty_cache, out_of_memory, sync
|
| 25 |
+
from .model import SERVE_MAX_STATE, ContextOverflow, admit
|
| 26 |
+
|
| 27 |
+
PREFIX_CACHE_SIZE = int(os.environ.get("KEV_PREFIX_CACHE", "4")) # states kept (KV + DeltaNet states; attention-only backbones also the state's hidden states); 0 disables
|
| 28 |
+
PREFIX_MIN_TOKENS = os.environ.get("KEV_PREFIX_MIN_TOKENS") # states shorter than this are not cached; default = the model's prefix_min_tokens (0 for hybrid backbones and MLX, 384 for attention-only torch models)
|
| 29 |
+
PREFIX_MAX_TOKENS = int(os.environ.get("KEV_PREFIX_MAX_TOKENS", "65536")) # state tokens the cache holds in all (least recently used evicted first); a longer state is not cached.
|
| 30 |
+
# One 64k state (Kev-27B: ~1.3 GB of keys, values and DeltaNet states), or four 16k ones, not four 64k ones
|
| 31 |
+
DATE_FACTS = os.environ.get("KEV_DATE_FACTS", "0") == "1"
|
| 32 |
+
TRUNCATE_STATES = os.environ.get("KEV_TRUNCATE_STATES", "0") == "1" # unset = a state over SERVE_MAX_STATE tokens gets a 422; 1 = read its first SERVE_MAX_STATE tokens, and every response says whether it did (truncated, usage.state_tokens / state_tokens_used)
|
| 33 |
+
API_KEY = os.environ.get("KEV_API_KEY") # unset = open server; set = require Authorization: Bearer <key>, as the TypeSafe clients always send
|
| 34 |
+
MAX_BATCH = 64 # requests the model thread takes at once (kev.cuda_graphs splits them to fit its buffers)
|
| 35 |
+
MODEL_NAMES = ("kev-latest", "jev-latest") # both names serve this checkpoint; jev-latest is the TypeSafe SDK default model, so an unconfigured client works
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
@dataclass
|
| 39 |
+
class PrefixCache:
|
| 40 |
+
"""State prefixes kept across requests, least recently used first: (state token ids, option_isolation) -> prefix.
|
| 41 |
+
At most `size` states and `max_tokens` state tokens in all; states shorter than min_tokens or longer than max_tokens
|
| 42 |
+
are not cached. A batch keeps (copies) only the new states that will still be here after it, its last distinct ones
|
| 43 |
+
within both bounds: the rest would be evicted by the batch itself."""
|
| 44 |
+
size: int
|
| 45 |
+
min_tokens: int
|
| 46 |
+
max_tokens: int = PREFIX_MAX_TOKENS
|
| 47 |
+
entries: dict = field(default_factory=dict)
|
| 48 |
+
hits: int = 0
|
| 49 |
+
misses: int = 0
|
| 50 |
+
oom_retries: int = 0 # batches that ran out of device memory with states cached, dropped them and ran again (Server._run)
|
| 51 |
+
|
| 52 |
+
def plan(self, encs):
|
| 53 |
+
"""-> (key per request, None when its state is not cached; its cached prefix or None; whether to keep a new one)."""
|
| 54 |
+
lengths = [enc["seg"].count(0) for enc in encs]
|
| 55 |
+
keys = [(tuple(enc["ids"][:n]), bool(enc.get("option_isolation"))) if self.size and self.min_tokens <= n <= self.max_tokens else None
|
| 56 |
+
for enc, n in zip(encs, lengths)]
|
| 57 |
+
survivors, tokens = set(), 0
|
| 58 |
+
for key in dict.fromkeys(k for k in reversed(keys) if k is not None): # most recent first, as store() keeps them
|
| 59 |
+
if len(survivors) == self.size or tokens + len(key[0]) > self.max_tokens: break
|
| 60 |
+
survivors.add(key); tokens += len(key[0])
|
| 61 |
+
return keys, [self.entries.get(k) if k is not None else None for k in keys], [k in survivors for k in keys]
|
| 62 |
+
|
| 63 |
+
def over(self, keys):
|
| 64 |
+
"""Whether these keys (states) exceed either bound: more than `size` states or `max_tokens` state tokens."""
|
| 65 |
+
return len(keys) > self.size or sum(len(k[0]) for k in keys) > self.max_tokens
|
| 66 |
+
|
| 67 |
+
def make_room(self, keys, cached, keep):
|
| 68 |
+
"""Drop now, before the batch runs, the entries this batch's store() will evict anyway, so an old state does not
|
| 69 |
+
stay resident through the passes of the new one that replaces it. store() (re)inserts every hit and every kept new
|
| 70 |
+
state as most recent and evicts the least recently used while over a bound; evicting the same keys first leaves
|
| 71 |
+
the same cache (tests/test_unit.py). A hit dropped here is still held by `cached` for this batch. Kev-4B on MLX:
|
| 72 |
+
a 64k state's cache is 2.2 GB, so a new 64k state after another one peaked with both (runs/mlx-long-states)."""
|
| 73 |
+
order = dict.fromkeys(self.entries) # least recently used first, as store() evicts
|
| 74 |
+
for key, old, k in zip(keys, cached, keep):
|
| 75 |
+
if key is not None and (old is not None or k):
|
| 76 |
+
order.pop(key, None); order[key] = None
|
| 77 |
+
while self.over(order):
|
| 78 |
+
key = next(iter(order)); del order[key]; self.entries.pop(key, None)
|
| 79 |
+
|
| 80 |
+
def store(self, keys, cached, prefixes):
|
| 81 |
+
"""Record hits and misses, and (re)insert the batch's prefixes in order: most recently used last."""
|
| 82 |
+
for key, old, new in zip(keys, cached, prefixes):
|
| 83 |
+
if key is None: continue
|
| 84 |
+
self.hits += old is not None; self.misses += old is None
|
| 85 |
+
if new is None: continue
|
| 86 |
+
self.entries.pop(key, None); self.entries[key] = new
|
| 87 |
+
while self.over(self.entries): self.entries.pop(next(iter(self.entries)))
|
| 88 |
+
|
| 89 |
+
def clear(self):
|
| 90 |
+
self.entries.clear()
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
@dataclass
|
| 94 |
+
class Server:
|
| 95 |
+
"""The loaded checkpoint, the state-prefix cache (PrefixCache), and the one model thread that runs every forward pass.
|
| 96 |
+
|
| 97 |
+
Request threads encode their record and queue it; the model thread takes everything queued when it becomes free and
|
| 98 |
+
runs it as one batch (model.probs_batch: with CUDA graphs, shared state and row passes; otherwise one request at a
|
| 99 |
+
time), then answers each request. It also captures pending CUDA graphs when the graphs say so (capture_due). `lock` is
|
| 100 |
+
held around each batch and capture: hold it to use the model directly."""
|
| 101 |
+
checkpoint: Checkpoint
|
| 102 |
+
tok: object
|
| 103 |
+
model: object
|
| 104 |
+
device: str
|
| 105 |
+
lock: threading.Lock = field(default_factory=threading.Lock)
|
| 106 |
+
batches: int = 0
|
| 107 |
+
batched_requests: int = 0
|
| 108 |
+
release_date: str = field(default="") # for the TypeSafe model card; resolved once (may ask the Hub)
|
| 109 |
+
truncate_states: bool = field(default_factory=lambda: TRUNCATE_STATES) # KEV_TRUNCATE_STATES; off: an over-length state is refused (kev.model.admit)
|
| 110 |
+
|
| 111 |
+
def __post_init__(self):
|
| 112 |
+
self.release_date = self.release_date or self.checkpoint.release_date()
|
| 113 |
+
self.prefix_cache = PrefixCache(PREFIX_CACHE_SIZE, int(PREFIX_MIN_TOKENS) if PREFIX_MIN_TOKENS else self.model.prefix_min_tokens)
|
| 114 |
+
self.queue, self.stopping = queue.Queue(), threading.Event()
|
| 115 |
+
# the model thread gives up the GIL at every CUDA sync and waits to get it back while the event loop parses and
|
| 116 |
+
# answers requests; at Python's default 5 ms switch interval those waits stretched a batch's model time ~2x.
|
| 117 |
+
# Process-wide, so close() puts it back.
|
| 118 |
+
self.switch_interval = sys.getswitchinterval()
|
| 119 |
+
sys.setswitchinterval(0.0005)
|
| 120 |
+
self.thread = threading.Thread(target=self._work, name="kev-model", daemon=True)
|
| 121 |
+
self.thread.start()
|
| 122 |
+
atexit.register(self.close) # a daemon thread killed inside a CUDA call at interpreter exit aborts the process
|
| 123 |
+
|
| 124 |
+
def close(self):
|
| 125 |
+
"""Stop the model thread after its current batch; requests still queued fail, and so do later ones (submit)."""
|
| 126 |
+
if self.stopping.is_set(): return
|
| 127 |
+
self.stopping.set(); self.thread.join()
|
| 128 |
+
while not self.queue.empty():
|
| 129 |
+
self.queue.get_nowait()[1].set_exception(RuntimeError("the server stopped")); self.queue.task_done()
|
| 130 |
+
sys.setswitchinterval(self.switch_interval)
|
| 131 |
+
|
| 132 |
+
def submit(self, rec):
|
| 133 |
+
"""Queue one record for the model thread. -> a Future of (probabilities, stats). The state prefix (tokens up to the
|
| 134 |
+
first question) is cached across requests, so a repeated state only pays for its question rows. latency_ms is the
|
| 135 |
+
model time of the batch the request ran in (not its wait in the queue). A record kev.model.admit refuses (a state
|
| 136 |
+
over SERVE_MAX_STATE tokens unless truncate_states, a question row over SERVE_MAX_BRANCH) is a 422 here, before
|
| 137 |
+
it reaches the model."""
|
| 138 |
+
if self.stopping.is_set(): raise HTTPException(503, "the server is stopping")
|
| 139 |
+
try: enc = admit(self.model, self.tok, rec, truncate=self.truncate_states)
|
| 140 |
+
except ValueError as e: # ContextOverflow, or a record the encoder cannot take
|
| 141 |
+
state = isinstance(e, ContextOverflow) and e.max_state is not None
|
| 142 |
+
raise HTTPException(422, str(e) + (f"; or start the server with KEV_TRUNCATE_STATES=1 to read only its first {e.max_state:,} tokens (responses then say truncated: true)" if state else ""))
|
| 143 |
+
done = Future()
|
| 144 |
+
self.queue.put((enc, done))
|
| 145 |
+
return done
|
| 146 |
+
|
| 147 |
+
def probs(self, rec):
|
| 148 |
+
return self.submit(rec).result()
|
| 149 |
+
|
| 150 |
+
def _work(self):
|
| 151 |
+
graphs = getattr(self.model, "graphs", None)
|
| 152 |
+
while not self.stopping.is_set():
|
| 153 |
+
try: batch = [self.queue.get(timeout=0.05)]
|
| 154 |
+
except queue.Empty:
|
| 155 |
+
if graphs is not None and graphs.capture_due(idle=True):
|
| 156 |
+
with self.lock: graphs.capture_pending(limit=1)
|
| 157 |
+
continue
|
| 158 |
+
while len(batch) < MAX_BATCH:
|
| 159 |
+
try: batch.append(self.queue.get_nowait())
|
| 160 |
+
except queue.Empty: break
|
| 161 |
+
try:
|
| 162 |
+
with self.lock: results = self._run([enc for enc, _ in batch])
|
| 163 |
+
except Exception as e: # every request of the batch gets the error; the thread lives on
|
| 164 |
+
traceback.clear_frames(e.__traceback__) # the batch's futures keep the exception, and its frames held the failed pass's tensors (142 MiB, a 3,578-token state on Kev-0.8B); the traceback keeps its lines
|
| 165 |
+
results = [e] * len(batch)
|
| 166 |
+
for (_, done), result in zip(batch, results):
|
| 167 |
+
(done.set_exception if isinstance(result, Exception) else done.set_result)(result)
|
| 168 |
+
self.queue.task_done()
|
| 169 |
+
if graphs is not None and graphs.capture_due(idle=False):
|
| 170 |
+
with self.lock: graphs.capture_pending(limit=1)
|
| 171 |
+
|
| 172 |
+
def _run(self, encs):
|
| 173 |
+
"""One batch through model.probs_batch, with the prefix cache. -> per request (probs, stats). A pass out of device
|
| 174 |
+
memory while states are cached drops the cache and runs once more: the cache only saves time, and kept resident
|
| 175 |
+
it failed every later batch of that size (#75, @vtxyer)."""
|
| 176 |
+
sync(self.device); t = time.time()
|
| 177 |
+
for retry in (False, True):
|
| 178 |
+
keys, cached, keep = self.prefix_cache.plan(encs)
|
| 179 |
+
self.prefix_cache.make_room(keys, cached, keep)
|
| 180 |
+
try: ps, prefixes = self.model.probs_batch(encs, cached, keep); break
|
| 181 |
+
except Exception as e: # states held: the cache's, and this batch's hits make_room already dropped from it
|
| 182 |
+
held = self.prefix_cache.entries or any(c is not None for c in cached)
|
| 183 |
+
if retry or not held or not out_of_memory(e): raise
|
| 184 |
+
cached = None; self.prefix_cache.clear(); self.prefix_cache.oom_retries += 1 # after the except: its traceback holds the failed pass's tensors
|
| 185 |
+
empty_cache(self.device)
|
| 186 |
+
sync(self.device); dt = round((time.time() - t) * 1000, 1)
|
| 187 |
+
self.prefix_cache.store(keys, cached, prefixes)
|
| 188 |
+
self.batches += 1; self.batched_requests += len(encs)
|
| 189 |
+
return [([q.tolist() for q in p], {"tokens": len(enc["ids"]), "state_tokens": enc["state_tokens"], "state_tokens_used": enc["seg"].count(0),
|
| 190 |
+
"latency_ms": dt, "prefix_cache_hit": c is not None}) # state_tokens: the request's state; _used: what the model read (less only when truncated)
|
| 191 |
+
for enc, p, c in zip(encs, ps, cached)]
|
| 192 |
+
|
| 193 |
+
def wait_idle(self):
|
| 194 |
+
"""Block until every submitted request is answered and no CUDA graph waits to be captured (benchmarks, warm-up)."""
|
| 195 |
+
self.queue.join()
|
| 196 |
+
graphs = getattr(self.model, "graphs", None)
|
| 197 |
+
while graphs is not None and graphs.capture_due(idle=True): time.sleep(0.01)
|
| 198 |
+
with self.lock: pass # a capture in progress finishes
|
| 199 |
+
|
| 200 |
+
def answer(self, req):
|
| 201 |
+
"""The /v1/systemone response body for one request."""
|
| 202 |
+
rec, meta = to_record(prepare(req))
|
| 203 |
+
return self._body(req, meta, *self.probs(rec))
|
| 204 |
+
|
| 205 |
+
async def answer_async(self, req):
|
| 206 |
+
"""answer() for the event loop: a request waiting on the model thread holds no worker thread, so a container takes
|
| 207 |
+
as many concurrent requests as its batches can absorb (FastAPI runs sync endpoints on a 40-thread pool)."""
|
| 208 |
+
rec, meta = to_record(prepare(req))
|
| 209 |
+
return self._body(req, meta, *await asyncio.wrap_future(self.submit(rec)))
|
| 210 |
+
|
| 211 |
+
def _body(self, req, meta, ps, m):
|
| 212 |
+
"""The TypeSafe body. A server that may truncate (truncate_states) also says, on every response, whether it did:
|
| 213 |
+
`truncated`, and usage.state_tokens (the request's state) / state_tokens_used (what the model read), both counting
|
| 214 |
+
the <state> token. The TypeSafe SDKs ignore fields they do not model, so clients keep parsing."""
|
| 215 |
+
answers = to_answers(ps, meta)
|
| 216 |
+
body = {"model": req.model, "answers": answers, "usage": {"input_tokens": m["tokens"], "output_tokens": output_tokens(self.tok, answers)}, "latency_ms": m["latency_ms"]}
|
| 217 |
+
if self.truncate_states:
|
| 218 |
+
body["usage"].update(state_tokens=m["state_tokens"], state_tokens_used=m["state_tokens_used"])
|
| 219 |
+
body["truncated"] = m["state_tokens"] > m["state_tokens_used"]
|
| 220 |
+
return body
|
| 221 |
+
|
| 222 |
+
|
| 223 |
+
def prepare(req):
|
| 224 |
+
"""Opt-in preprocessing applied to every request before the model sees it."""
|
| 225 |
+
return req.model_copy(update={"state": with_date_facts(req.state)}) if DATE_FACTS else req
|
| 226 |
+
|
| 227 |
+
|
| 228 |
+
app = FastAPI(title="kev")
|
| 229 |
+
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], expose_headers=["x-typesafe-request-id", "server-timing"])
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
@app.middleware("http")
|
| 233 |
+
async def typesafe(request, call_next):
|
| 234 |
+
"""Bearer auth (when API_KEY is set) and the request id every TypeSafe client reads off the response."""
|
| 235 |
+
started = time.perf_counter()
|
| 236 |
+
if API_KEY and request.url.path.startswith("/v1") and not hmac.compare_digest(request.headers.get("authorization", ""), f"Bearer {API_KEY}"):
|
| 237 |
+
resp = JSONResponse({"detail": "missing or invalid API key; send Authorization: Bearer <KEV_API_KEY>"}, 401, {"www-authenticate": "Bearer"})
|
| 238 |
+
else:
|
| 239 |
+
resp = await call_next(request)
|
| 240 |
+
resp.headers["x-typesafe-request-id"] = request.headers.get("x-typesafe-request-id") or uuid.uuid4().hex
|
| 241 |
+
resp.headers["server-timing"] = f"app;dur={(time.perf_counter() - started) * 1000:.1f}" # time inside this process, for telling it from the network
|
| 242 |
+
return resp
|
| 243 |
+
|
| 244 |
+
|
| 245 |
+
def server() -> Server:
|
| 246 |
+
return app.state.server
|
| 247 |
+
|
| 248 |
+
|
| 249 |
+
@app.post("/v1/systemone")
|
| 250 |
+
async def systemone(req: SystemOneRequest):
|
| 251 |
+
"""TypeSafe-compatible endpoint: typed questions in, typed answers out, one prefill pass."""
|
| 252 |
+
return await server().answer_async(req)
|
| 253 |
+
|
| 254 |
+
|
| 255 |
+
class PermuteSystemOne(BaseModel):
|
| 256 |
+
request: SystemOneRequest
|
| 257 |
+
question: str
|
| 258 |
+
n_perm: int = Field(default=6, ge=1, le=64) # each order is a forward pass; 0 divided by nothing, unbounded counts ran forever (#30)
|
| 259 |
+
seed: int = 0
|
| 260 |
+
|
| 261 |
+
|
| 262 |
+
@app.post("/v1/systemone/permute")
|
| 263 |
+
def systemone_permute(r: PermuteSystemOne):
|
| 264 |
+
"""Re-run one Choice question under n_perm option orders. Returns per-order probabilities keyed by option name."""
|
| 265 |
+
q = r.request.questions.get(r.question)
|
| 266 |
+
if q is None or q.type != "choice": raise HTTPException(422, "question must be an existing choice question")
|
| 267 |
+
rng = random.Random(r.seed); keys = list(q.criteria); runs = []
|
| 268 |
+
for i in range(r.n_perm):
|
| 269 |
+
order = list(keys)
|
| 270 |
+
if i > 0: rng.shuffle(order)
|
| 271 |
+
one = r.request.model_copy(update={"questions": {r.question: q.model_copy(update={"criteria": {k: q.criteria[k] for k in order}})}})
|
| 272 |
+
resp = server().answer(one); a = resp["answers"][r.question]
|
| 273 |
+
runs.append({"order": order, "probabilities": a["probabilities"], "choice": a["choice"], "latency_ms": resp["latency_ms"]})
|
| 274 |
+
spread = {k: max(x["probabilities"][k] for x in runs) - min(x["probabilities"][k] for x in runs) for k in keys}
|
| 275 |
+
return truncation_marks({"runs": runs, "argmax_stable": len({x["choice"] for x in runs}) == 1, "spread": spread}, resp)
|
| 276 |
+
|
| 277 |
+
|
| 278 |
+
@app.post("/v1/systemone/separate")
|
| 279 |
+
def systemone_separate(req: SystemOneRequest):
|
| 280 |
+
"""Answer each question in its own request against the same state (N passes). For packed-vs-separate comparison."""
|
| 281 |
+
parts = [server().answer(req.model_copy(update={"questions": {qid: q}})) for qid, q in req.questions.items()]
|
| 282 |
+
answers = {qid: a for p in parts for qid, a in p["answers"].items()}
|
| 283 |
+
body = {"model": req.model, "answers": answers,
|
| 284 |
+
"usage": {"input_tokens": sum(p["usage"]["input_tokens"] for p in parts), "output_tokens": output_tokens(server().tok, answers)},
|
| 285 |
+
"latency_ms": round(sum(p["latency_ms"] for p in parts), 1)}
|
| 286 |
+
return truncation_marks(body, parts[0])
|
| 287 |
+
|
| 288 |
+
|
| 289 |
+
def truncation_marks(body, part):
|
| 290 |
+
"""A demo endpoint built from several /v1/systemone answers on one state carries a truncating server's marks
|
| 291 |
+
(Server._body: `truncated`, usage.state_tokens / state_tokens_used) from any one of them: each read the same tokens."""
|
| 292 |
+
if "truncated" in part:
|
| 293 |
+
body["truncated"] = part["truncated"]
|
| 294 |
+
body.setdefault("usage", {}).update({k: part["usage"][k] for k in ("state_tokens", "state_tokens_used")})
|
| 295 |
+
return body
|
| 296 |
+
|
| 297 |
+
|
| 298 |
+
@app.get("/v1/models")
|
| 299 |
+
def models():
|
| 300 |
+
"""One TypeSafe model card (name, description, release_date) per accepted model name, plus the Kev serving details
|
| 301 |
+
a client may ignore: the run, the base, the device, the backend and precision, the temperature, prefix-cache stats."""
|
| 302 |
+
s = server()
|
| 303 |
+
ck, meta = s.checkpoint, s.checkpoint.meta
|
| 304 |
+
card = {"description": f"Kev pointer head on {meta.base}, serving {ck.requested} at temperature {s.model.head.temperature:.2f}",
|
| 305 |
+
"release_date": s.release_date,
|
| 306 |
+
"run": ck.requested, "base": meta.base, "lora": meta.lora, "device": s.device, "backend": s.model.backend, "dtype": s.model.dtype,
|
| 307 |
+
"temperature": s.model.head.temperature, "max_state_tokens": SERVE_MAX_STATE, "truncate_states": s.truncate_states,
|
| 308 |
+
"cuda_graphs": graphs.stats() if (graphs := getattr(s.model, "graphs", None)) else None,
|
| 309 |
+
"prefix_cache": {"size": s.prefix_cache.size, "min_state_tokens": s.prefix_cache.min_tokens, "max_tokens": s.prefix_cache.max_tokens, "hits": s.prefix_cache.hits,
|
| 310 |
+
"misses": s.prefix_cache.misses, "cached_states": len(s.prefix_cache.entries), "oom_retries": s.prefix_cache.oom_retries},
|
| 311 |
+
"batches": {"count": s.batches, "requests": s.batched_requests, "queued": s.queue.qsize()}}
|
| 312 |
+
return {"models": [{"name": name, **card} for name in MODEL_NAMES]}
|
| 313 |
+
|
| 314 |
+
|
| 315 |
+
def main():
|
| 316 |
+
ap = argparse.ArgumentParser()
|
| 317 |
+
ap.add_argument("--run", default="runs/kev")
|
| 318 |
+
ap.add_argument("--fallback", default="runs/smoke")
|
| 319 |
+
ap.add_argument("--host", default="127.0.0.1", help="interface to bind; 0.0.0.0 to serve beyond this machine (a container, a VM behind a proxy)")
|
| 320 |
+
ap.add_argument("--port", type=int, default=8008)
|
| 321 |
+
a = ap.parse_args()
|
| 322 |
+
run = a.run if is_hub_id(a.run) or os.path.exists(f"{a.run}/head.pt") else a.fallback
|
| 323 |
+
if run != a.run: print(f"{a.run} not found, falling back to {run}")
|
| 324 |
+
dev = default_device()
|
| 325 |
+
opts = LoadOptions.from_env()
|
| 326 |
+
if dev == "mps" and opts.attn is None: opts = replace(opts, attn="sdpa") # serving default on Apple GPUs (parity measured)
|
| 327 |
+
if dev != "cpu" and opts.dtype is None: opts = replace(opts, dtype=torch.bfloat16) # serving default: 2-4.5x faster than fp32 on an L4, same answers (LoadOptions.dtype); KEV_DTYPE=fp32 for the exact path
|
| 328 |
+
if dev == "cuda" and opts.cuda_graphs is None: opts = replace(opts, cuda_graphs=True) # serving default: a pass is ~2,000 kernel launches, so replaying graphs cuts warm latency several-fold (kev.cuda_graphs); KEV_CUDA_GRAPHS=0 to decline
|
| 329 |
+
fused_default = dev == "cuda" and opts.fused is None
|
| 330 |
+
if fused_default: opts = replace(opts, fused=fused_available()) # serving default: fused Qwen3.5 kernels, ~1/3 less GPU time per batch (kev.fused_qwen35), when fla is installed; KEV_FUSED=0 to decline, KEV_FUSED=1 to insist
|
| 331 |
+
if opts.backend is None: opts = replace(opts, backend="auto") # serving default: MLX for the hybrid Qwen3.5 checkpoints on Apple Silicon (LoadOptions.backend); KEV_BACKEND=torch to decline
|
| 332 |
+
ck = Checkpoint(run)
|
| 333 |
+
tok, model = ck.load(dev, opts)
|
| 334 |
+
if fused_default and not opts.fused and model.hybrid: print("fused Qwen3.5 kernels off: install the flash-linear-attention version kev/fused_qwen35.py pins (FLA_VERSION) to turn them on")
|
| 335 |
+
app.state.server = Server(ck, tok, model, dev)
|
| 336 |
+
print(f"serving {ck.requested} ({ck.path}) on {dev} via {model.backend} ({model.dtype}) {a.host}:{a.port}; states over {SERVE_MAX_STATE:,} tokens "
|
| 337 |
+
f"{'truncated (KEV_TRUNCATE_STATES=1)' if app.state.server.truncate_states else 'refused (422)'}") # /v1/models reports the run as given, not the resolved cache path
|
| 338 |
+
import uvicorn
|
| 339 |
+
uvicorn.run(app, host=a.host, port=a.port)
|
| 340 |
+
|
| 341 |
+
|
| 342 |
+
if __name__ == "__main__":
|
| 343 |
+
main()
|
runtime/kev/shared_prefix.py
ADDED
|
@@ -0,0 +1,141 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Training forward through a shared state prefix, for the hybrid Qwen3.5 backbones (kev.train --shared_prefix).
|
| 2 |
+
|
| 3 |
+
The row form (kev.model.forward_rows_batch) trains a record with q questions as q causal rows, state + one branch each,
|
| 4 |
+
so the state runs q times. Here each record's state runs once and its branches continue from it, which costs
|
| 5 |
+
state + the branches instead of q x (state + branch). The two are the same computation: a branch token attends to the
|
| 6 |
+
same tokens at the same positions, and it reaches the state through what each layer leaves behind at the end of the
|
| 7 |
+
state, as the layer's own cache path reads it for inference (kev.model._branch_rows_from_prefix): the attention layers'
|
| 8 |
+
keys and values, and the Gated DeltaNet layers' short-convolution window (the conv's left context) and recurrent state
|
| 9 |
+
(the chunked delta rule's `initial_state`, which fla differentiates). Two things differ from serving:
|
| 10 |
+
- `Prefix` holds that per-layer state functionally (transformers' cache layers copy into static buffers in place, which
|
| 11 |
+
autograd cannot see through), so every branch's gradient flows back into the state and accumulates there;
|
| 12 |
+
- each decoder layer runs its state pass and its branch pass inside one step, and gradient checkpointing (when the
|
| 13 |
+
backbone has it on) recomputes the step as a whole, so no cache outlives its layer or is written twice on replay.
|
| 14 |
+
States of a batch are left-padded (the padded positions are zeroed and stay zero: a zero input leaves the DeltaNet state,
|
| 15 |
+
the conv window and the residual stream unchanged), branches right-padded.
|
| 16 |
+
Exactness: fp32 with transformers' reference kernels matches the row form to ~1e-5 (logits and every gradient; tiny
|
| 17 |
+
random Qwen3.5 in tests/test_unit.py, Qwen3.5-0.8B-Base in tests/test_model.py). On CUDA fla's Triton kernels round their
|
| 18 |
+
fp32 dots like TF32 and the shared prefix left-pads states, so it rounds differently from the row form: measured against the
|
| 19 |
+
exact kernels, it is as far from the exact answer as the row form's own batchings are, with no systematic shift.
|
| 20 |
+
"""
|
| 21 |
+
import types
|
| 22 |
+
|
| 23 |
+
import torch
|
| 24 |
+
import torch.nn as nn
|
| 25 |
+
import torch.nn.functional as F
|
| 26 |
+
from torch.distributed.fsdp import FSDPModule, register_fsdp_forward_method
|
| 27 |
+
from torch.utils.checkpoint import checkpoint
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
class Prefix:
|
| 31 |
+
"""One decoder layer's view of the shared state, duck-typing the cache calls a Qwen3.5 layer makes. The state pass
|
| 32 |
+
stores what the layer leaves (keys and values, or the conv window and the recurrent state); `branch` gathers it once
|
| 33 |
+
per branch row; the branch pass reads it. Nothing is written in place."""
|
| 34 |
+
record_past = True # keeps a one-token pass off the in-place decode kernel
|
| 35 |
+
|
| 36 |
+
def __init__(self):
|
| 37 |
+
self.keys = self.values = self.conv = None
|
| 38 |
+
self.recurrent_states = {}
|
| 39 |
+
|
| 40 |
+
@property
|
| 41 |
+
def layers(self): # layer code reads `cache.layers[layer_idx].recurrent_states[0]`: every index is this prefix
|
| 42 |
+
return self
|
| 43 |
+
|
| 44 |
+
def __getitem__(self, layer_idx):
|
| 45 |
+
return self
|
| 46 |
+
|
| 47 |
+
def has_previous_state(self, layer_idx=None, state_idx=0):
|
| 48 |
+
return self.conv is not None
|
| 49 |
+
|
| 50 |
+
def update_conv_state(self, conv_input, layer_idx, conv_kernel_size, **kwargs):
|
| 51 |
+
if self.conv is not None: return torch.cat([self.conv, conv_input], -1) # the branch sees the state's last inputs
|
| 52 |
+
self.conv = F.pad(conv_input, (max(0, conv_kernel_size - conv_input.shape[-1]), 0))[..., -conv_kernel_size:]
|
| 53 |
+
return conv_input
|
| 54 |
+
|
| 55 |
+
def update_recurrent_state(self, recurrent_state, layer_idx, **kwargs):
|
| 56 |
+
self.recurrent_states.setdefault(0, recurrent_state) # the state's final state; a branch's own is discarded
|
| 57 |
+
return recurrent_state
|
| 58 |
+
|
| 59 |
+
def update(self, keys, values, layer_idx, *args, **kwargs):
|
| 60 |
+
if self.keys is None:
|
| 61 |
+
self.keys, self.values = keys, values
|
| 62 |
+
return keys, values
|
| 63 |
+
return torch.cat([self.keys, keys], -2), torch.cat([self.values, values], -2)
|
| 64 |
+
|
| 65 |
+
def branch(self, owner):
|
| 66 |
+
"""Gather the stored state per branch row (owner[i] = the record of branch i)."""
|
| 67 |
+
pick = lambda t: None if t is None else t.index_select(0, owner)
|
| 68 |
+
self.keys, self.values, self.conv = pick(self.keys), pick(self.values), pick(self.conv)
|
| 69 |
+
self.recurrent_states = {k: pick(v) for k, v in self.recurrent_states.items()}
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
def _masks(allow, dtype, attn):
|
| 73 |
+
"""Additive float mask for eager attention, boolean for SDPA (which also takes a float mask, but only in the query's
|
| 74 |
+
dtype, and autocast changes that)."""
|
| 75 |
+
allow = allow[:, None]
|
| 76 |
+
return allow if attn == "sdpa" else torch.zeros(allow.shape, dtype=dtype, device=allow.device).masked_fill(~allow, torch.finfo(dtype).min)
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
def branch_hidden(lm, splits, pad_id, device):
|
| 80 |
+
"""splits[b] = kev.model.rows_of(record b's encoding). -> per record, per question, the branch's final hidden states
|
| 81 |
+
[branch length, d] in fp32 (what the pointer head reads; the state's own hidden states are never needed)."""
|
| 82 |
+
base = lm.get_base_model() if hasattr(lm, "get_base_model") else lm
|
| 83 |
+
owner = [b for b, (_, _, rows) in enumerate(splits) for _ in rows]
|
| 84 |
+
branches = [r for _, _, rows in splits for r in rows]
|
| 85 |
+
Ls, Lb = max(len(s) for s, _, _ in splits), max(len(r["ids"]) for r in branches)
|
| 86 |
+
s_ids = torch.full((len(splits), Ls), pad_id, device=device); s_pos = torch.zeros_like(s_ids); s_real = torch.zeros_like(s_ids, dtype=torch.bool)
|
| 87 |
+
for b, (ids, pos, _) in enumerate(splits): # left-padded
|
| 88 |
+
s_ids[b, Ls - len(ids):] = torch.tensor(ids, device=device); s_pos[b, Ls - len(ids):] = torch.tensor(pos, device=device); s_real[b, Ls - len(ids):] = True
|
| 89 |
+
b_ids = torch.full((len(branches), Lb), pad_id, device=device); b_pos = torch.zeros_like(b_ids); b_real = torch.zeros_like(b_ids, dtype=torch.bool)
|
| 90 |
+
for i, r in enumerate(branches): # right-padded; their positions already continue the state
|
| 91 |
+
b_ids[i, :len(r["ids"])] = torch.tensor(r["ids"], device=device); b_pos[i, :len(r["pos"])] = torch.tensor(r["pos"], device=device); b_real[i, :len(r["ids"])] = True
|
| 92 |
+
run = _method(base, "kev_shared_prefix", _forward) # on an FSDP2 root: embeddings and final norm gathered as for base(...)
|
| 93 |
+
h_b = run(s_ids, s_pos, s_real, b_ids, b_pos, b_real, torch.tensor(owner, device=device))
|
| 94 |
+
out = [[] for _ in splits]
|
| 95 |
+
for i, b in enumerate(owner): out[b].append(h_b[i, :len(branches[i]["ids"])])
|
| 96 |
+
return out
|
| 97 |
+
|
| 98 |
+
|
| 99 |
+
def _forward(base, s_ids, s_pos, s_real, b_ids, b_pos, b_real, owner):
|
| 100 |
+
"""The decoder stack over left-padded states and right-padded branches (owner[i] = the state of branch i)."""
|
| 101 |
+
device, (Ls, Lb) = s_ids.device, (s_ids.shape[1], b_ids.shape[1])
|
| 102 |
+
h_s = base.embed_tokens(s_ids) * s_real[..., None]; h_b = base.embed_tokens(b_ids) * b_real[..., None]
|
| 103 |
+
dtype, attn = h_s.dtype, base.config._attn_implementation
|
| 104 |
+
# attention: a state token sees the real state tokens up to itself; a branch token its record's real state and its
|
| 105 |
+
# own branch so far; a padded query keeps its diagonal so no softmax row is empty (its output is zeroed or unused)
|
| 106 |
+
causal = lambda n: torch.ones(n, n, dtype=torch.bool, device=device).tril()
|
| 107 |
+
eye = lambda n: torch.eye(n, dtype=torch.bool, device=device)
|
| 108 |
+
# States without padding (one per micro-batch, or all of one length) need no state mask under SDPA: None makes it
|
| 109 |
+
# causal (is_causal), the flash kernel's case. A 64k-token state's explicit mask would be 4 GB, kept for the whole
|
| 110 |
+
# backward, and would rule the flash kernel out.
|
| 111 |
+
plain = attn == "sdpa" and bool(s_real.all())
|
| 112 |
+
allow_b = torch.cat([s_real[owner][:, None, :].expand(-1, Lb, -1), (causal(Lb)[None] & b_real[:, None, :]) | eye(Lb)[None]], -1)
|
| 113 |
+
masks_s = {"full_attention": None if plain else _masks((causal(Ls)[None] & s_real[:, None, :]) | eye(Ls)[None], dtype, attn), "linear_attention": s_real.to(dtype)}
|
| 114 |
+
masks_b = {"full_attention": _masks(allow_b, dtype, attn), "linear_attention": b_real.to(dtype)}
|
| 115 |
+
rope_s, rope_b = base.rotary_emb(h_s, s_pos[None].expand(3, -1, -1)), base.rotary_emb(h_b, b_pos[None].expand(3, -1, -1))
|
| 116 |
+
|
| 117 |
+
recompute = base.gradient_checkpointing and base.training and torch.is_grad_enabled()
|
| 118 |
+
for layer in base.layers:
|
| 119 |
+
step = _method(layer, "kev_shared_prefix_step", _step)
|
| 120 |
+
args = (h_s, h_b, (rope_s, rope_b), (masks_s[layer.block_type], masks_b[layer.block_type]), (s_pos, b_pos), s_real, owner)
|
| 121 |
+
h_s, h_b = checkpoint(step, *args, use_reentrant=False) if recompute else step(*args)
|
| 122 |
+
return base.norm(h_b).float()
|
| 123 |
+
|
| 124 |
+
|
| 125 |
+
def _step(layer, h_s, h_b, ropes, masks, positions, s_real, owner):
|
| 126 |
+
"""One decoder layer: the state pass, then the branch pass reading what it left (`Prefix`). A single call on the layer
|
| 127 |
+
(both passes call its plain `forward`), so an FSDP2 unit is gathered once for both and checkpointing replays both."""
|
| 128 |
+
prefix = Prefix()
|
| 129 |
+
h_s = layer.forward(h_s, position_embeddings=ropes[0], attention_mask=masks[0], position_ids=positions[0], past_key_values=prefix)
|
| 130 |
+
prefix.branch(owner)
|
| 131 |
+
h_b = layer.forward(h_b, position_embeddings=ropes[1], attention_mask=masks[1], position_ids=positions[1], past_key_values=prefix)
|
| 132 |
+
return h_s * s_real[..., None], h_b
|
| 133 |
+
|
| 134 |
+
|
| 135 |
+
def _method(module, name, fn):
|
| 136 |
+
"""`fn` bound to `module` as `name`; on an FSDP2 unit also registered as a forward method, so the unit's parameters
|
| 137 |
+
are gathered around the call and resharded after it, and backward re-gathers them from its outputs' hooks."""
|
| 138 |
+
if not hasattr(module, name):
|
| 139 |
+
setattr(module, name, types.MethodType(fn, module))
|
| 140 |
+
if isinstance(module, FSDPModule): register_fsdp_forward_method(module, name)
|
| 141 |
+
return getattr(module, name)
|
runtime/kev/study_v3.py
ADDED
|
@@ -0,0 +1,232 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import argparse
|
| 2 |
+
import copy
|
| 3 |
+
import json
|
| 4 |
+
import random
|
| 5 |
+
from collections import Counter, defaultdict
|
| 6 |
+
from pathlib import Path
|
| 7 |
+
|
| 8 |
+
from huggingface_hub import HfApi
|
| 9 |
+
|
| 10 |
+
from kev.composition import DEV_SHAPES, HELD_OUT_KEYS, TEST_SHAPES, TRAIN_SHAPES, canonical, check_group, generate as compose, sample_trees
|
| 11 |
+
from kev.contrastive import ORDINAL_FAMILIES, generate
|
| 12 |
+
from kev.data import materialize
|
| 13 |
+
from kev.model import fits, load_tokenizer
|
| 14 |
+
from kev.suite import SPLITS, digest, load_split, read_manifest, semantic_hash, validate_training, write_json, write_jsonl
|
| 15 |
+
|
| 16 |
+
BASES = ("Qwen/Qwen3-0.6B-Base", "Qwen/Qwen3-4B-Base")
|
| 17 |
+
FAMILIES = ("return_window", "spend_threshold", "age_eligibility", "quantity_limit")
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
def grouped_split(records, calibration_groups):
|
| 21 |
+
groups = defaultdict(dict)
|
| 22 |
+
for r in records:
|
| 23 |
+
m = r["_meta"]
|
| 24 |
+
groups[m["family"]].setdefault(m["group_id"], []).append(r)
|
| 25 |
+
train, calibration = [], []
|
| 26 |
+
for family in sorted(groups):
|
| 27 |
+
units = list(groups[family].values())
|
| 28 |
+
if len(units) <= calibration_groups:
|
| 29 |
+
raise ValueError("not enough groups per family")
|
| 30 |
+
calibration.extend(r for group in units[:calibration_groups] for r in group)
|
| 31 |
+
train.extend(r for group in units[calibration_groups:] for r in group)
|
| 32 |
+
return train, calibration
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
def legacy(pairs, seed, families=FAMILIES, source="legacy_policy", excluded=()):
|
| 36 |
+
candidates, _ = generate(3 * pairs, seed, families)
|
| 37 |
+
records, seen, counts = [], set(excluded), Counter()
|
| 38 |
+
for a, b in zip(candidates[::2], candidates[1::2]):
|
| 39 |
+
family = a["_meta"]["family"]
|
| 40 |
+
hashes = {semantic_hash(a), semantic_hash(b)}
|
| 41 |
+
if counts[family] >= pairs or hashes & seen:
|
| 42 |
+
continue
|
| 43 |
+
for r in (a, b):
|
| 44 |
+
m = r["_meta"]
|
| 45 |
+
m.update(source=source, group_id=m["pair_id"], variant="clean")
|
| 46 |
+
m["text_sha256"] = semantic_hash(r)
|
| 47 |
+
records.extend((a, b)); seen.update(hashes); counts[family] += 1
|
| 48 |
+
if any(counts[f] != pairs for f in families):
|
| 49 |
+
raise ValueError("insufficient unique legacy pairs")
|
| 50 |
+
return records
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
def freeze(out, source="evals/decision-v2", transfer="evals/transfer-v2", public_train=None, synthetic_scale=1, inherit_eval=None,
|
| 54 |
+
random_structures=0, groups_per_structure=8, train_styles=(0, 1), matched_arms=True, legacy_families=FAMILIES, structures_seed=None):
|
| 55 |
+
"""random_structures > 0: the compositional arm is generated from that many random rule trees (negation anywhere,
|
| 56 |
+
held-out and locked structures excluded by canonical key) instead of the eight fixed TRAIN_SHAPES."""
|
| 57 |
+
"""inherit_eval: a frozen decision suite whose development/test bytes are reused verbatim (only train/calibration are
|
| 58 |
+
regenerated), so results stay comparable across suite versions that differ in training data only."""
|
| 59 |
+
"""public_train: optional larger public training pool (a frozen suite dir). Its train+calibration partitions replace
|
| 60 |
+
the source suite's public train/calibration; development/test still come from `source` so results stay comparable."""
|
| 61 |
+
out, source, transfer = Path(out), Path(source), Path(transfer)
|
| 62 |
+
if out.exists():
|
| 63 |
+
raise FileExistsError("v3 destination already exists; choose a new version")
|
| 64 |
+
original = read_manifest(source)
|
| 65 |
+
old_transfer = read_manifest(transfer)
|
| 66 |
+
revisions = {base: HfApi().model_info(base).sha for base in BASES}
|
| 67 |
+
tokenizers = [load_tokenizer(base, revision=sha) for base, sha in revisions.items()]
|
| 68 |
+
parts = {s: [] for s in SPLITS}
|
| 69 |
+
for split in ("train", "calibration", "development"):
|
| 70 |
+
parts[split] = [r for r in load_split(source, split) if r["_meta"]["source"] != "contrastive"]
|
| 71 |
+
if public_train:
|
| 72 |
+
pool = Path(public_train)
|
| 73 |
+
dev_test_states = {r["_meta"]["text_sha256"] for split in ("development", "test") for r in load_split(source, split, allow_test=True)}
|
| 74 |
+
for split in ("train", "calibration"):
|
| 75 |
+
fresh = [r for r in load_split(pool, split) if r["_meta"]["source"] != "contrastive"]
|
| 76 |
+
leaked = [r for r in fresh if r["_meta"]["text_sha256"] in dev_test_states]
|
| 77 |
+
if leaked: raise ValueError(f"{len(leaked)} public {split} records collide with development/test states")
|
| 78 |
+
parts[split] = fresh
|
| 79 |
+
source_dirs = [pool]
|
| 80 |
+
else:
|
| 81 |
+
source_dirs = []
|
| 82 |
+
reserved = set()
|
| 83 |
+
for split in ("train", "calibration", "development"):
|
| 84 |
+
reserved.update(semantic_hash(r) for r in parts[split])
|
| 85 |
+
sources = [p for p in source.glob("*.json*")] + [p for p in transfer.glob("*.json*")] + [p for d in source_dirs for p in d.glob("*.json*")]
|
| 86 |
+
parent_hashes = {str(p): digest(p) for p in sources}
|
| 87 |
+
|
| 88 |
+
def admit_groups(records):
|
| 89 |
+
groups = defaultdict(list)
|
| 90 |
+
for r in records:
|
| 91 |
+
groups[r["_meta"]["group_id"]].append(r)
|
| 92 |
+
for group in groups.values():
|
| 93 |
+
if group[0]["_meta"]["source"] in ("compositional", "composition_holdout"):
|
| 94 |
+
check_group(group)
|
| 95 |
+
hashes = {semantic_hash(r) for r in group}
|
| 96 |
+
if hashes & reserved:
|
| 97 |
+
raise ValueError("semantic state collision across groups; choose a new generation seed")
|
| 98 |
+
for r in group:
|
| 99 |
+
if not fits(materialize(r), *tokenizers):
|
| 100 |
+
raise ValueError("synthetic record exceeds the training context")
|
| 101 |
+
reserved.update(hashes)
|
| 102 |
+
return records
|
| 103 |
+
|
| 104 |
+
# development synthetic records are generated first, against public states only, so they do not depend on the size
|
| 105 |
+
# of the synthetic training arms (dev/test bytes stay identical across suite versions); training groups that would
|
| 106 |
+
# collide with a development state are rejected by admit_groups
|
| 107 |
+
if inherit_eval:
|
| 108 |
+
inherited_dev = load_split(inherit_eval, "development")
|
| 109 |
+
parts["development"] = inherited_dev
|
| 110 |
+
reserved.update(semantic_hash(r) for r in inherited_dev)
|
| 111 |
+
reserved.update(semantic_hash(r) for r in load_split(inherit_eval, "test", allow_test=True))
|
| 112 |
+
dev_synthetic = []
|
| 113 |
+
else:
|
| 114 |
+
dev_synthetic = admit_groups(legacy(12, "v3-control-dev", excluded=reserved)) + admit_groups(compose(4, "v3-composition-dev"))
|
| 115 |
+
old_train, old_cal = grouped_split(admit_groups(legacy(64 * synthetic_scale, "v3-control", families=legacy_families, excluded=reserved)), 8 * synthetic_scale)
|
| 116 |
+
if random_structures:
|
| 117 |
+
trees = {f"rand:{canonical(t)}": t for t in sample_trees(random_structures, f"{structures_seed or out.name}-structures")}
|
| 118 |
+
new_train, new_cal = grouped_split(admit_groups(compose(groups_per_structure, "v7-composition", styles=train_styles, trees=trees)), 1)
|
| 119 |
+
else:
|
| 120 |
+
new_train, new_cal = grouped_split(admit_groups(compose(16 * synthetic_scale, "v3-composition")), 2 * synthetic_scale)
|
| 121 |
+
if matched_arms and (len(old_train) != len(new_train) or len(old_cal) != len(new_cal)):
|
| 122 |
+
raise ValueError("synthetic arms have unequal record budgets")
|
| 123 |
+
parts["train"] += old_train + new_train
|
| 124 |
+
parts["calibration"] += old_cal + new_cal
|
| 125 |
+
parts["development"] += dev_synthetic
|
| 126 |
+
transfer_dev = [r for r in load_split(transfer, "development") if r["_meta"]["source"] != "contrastive"]
|
| 127 |
+
transfer_dev += admit_groups(legacy(20, "v3-legacy-transfer", ("authorization", "deadline"), source="legacy_holdout"))
|
| 128 |
+
transfer_dev += admit_groups(compose(8, "v3-composition-transfer", DEV_SHAPES, styles=(1,), source="composition_holdout"))
|
| 129 |
+
final_extra = admit_groups(compose(8, "v3-locked", TEST_SHAPES, styles=(2,), source="composition_holdout"))
|
| 130 |
+
for split in ("train", "calibration"):
|
| 131 |
+
random.Random(f"v3-{split}").shuffle(parts[split])
|
| 132 |
+
manifest = {"version": 3, "base_revisions": revisions, "dataset_revisions": original["dataset_revisions"],
|
| 133 |
+
"parent_files": parent_hashes, "holdout_sources": [],
|
| 134 |
+
"trainable_sources": sorted({s for s in original["trainable_sources"] if s != "contrastive"}
|
| 135 |
+
| ({s for s in read_manifest(public_train)["trainable_sources"]} if public_train else set())) + ["legacy_policy", "compositional"],
|
| 136 |
+
"eval_only_sources": old_transfer["eval_only_sources"] + ["legacy_holdout", "composition_holdout"],
|
| 137 |
+
"context": original["context"],
|
| 138 |
+
"protocol": {"train_shapes": TRAIN_SHAPES, "transfer_shapes": DEV_SHAPES, "locked_shapes": TEST_SHAPES,
|
| 139 |
+
"train_render_styles": list(train_styles), "locked_render_styles": [2],
|
| 140 |
+
"random_structures": random_structures, "groups_per_structure": groups_per_structure if random_structures else None,
|
| 141 |
+
"structures_seed": f"{structures_seed or out.name}-structures" if random_structures else None,
|
| 142 |
+
"excluded_structure_keys": sorted(HELD_OUT_KEYS), "legacy_families_trainable": list(legacy_families),
|
| 143 |
+
"public_train_records": len(parts["train"]) - len(old_train) - len(new_train),
|
| 144 |
+
"public_train_pool": str(public_train) if public_train else str(source), "synthetic_scale": synthetic_scale,
|
| 145 |
+
"inherited_eval": str(inherit_eval) if inherit_eval else None,
|
| 146 |
+
"synthetic_records_per_arm": len(old_train) if matched_arms else {"legacy_policy": len(old_train), "compositional": len(new_train)},
|
| 147 |
+
"calibration": "shared; stratified by family and group",
|
| 148 |
+
"arm_selection": "train_sources selects public sources plus exactly one synthetic arm",
|
| 149 |
+
"primary": "macro development NLL; transfer and confident-error checks required; no automatic release",
|
| 150 |
+
"legacy_test": "Inherited v2 locked test bytes retained without inspecting examples; only new structures added to transfer test",
|
| 151 |
+
"contamination": "Exact semantic checks among new groups, not fuzzy or pretraining decontamination. Inherited legacy test overlap with new synthetic controls is not certified."},
|
| 152 |
+
"files": {}, "code_hashes": {p.name: digest(p) for p in Path(__file__).parent.glob("*.py")}}
|
| 153 |
+
validate_training(parts["train"], manifest)
|
| 154 |
+
for name, partitions, inherited in (("decision-v3", parts, source),
|
| 155 |
+
("transfer-v3", {"train": [], "calibration": [], "development": transfer_dev, "test": final_extra}, transfer)):
|
| 156 |
+
folder = out / name
|
| 157 |
+
folder.mkdir(parents=True, exist_ok=False)
|
| 158 |
+
m = copy.deepcopy(manifest)
|
| 159 |
+
if name.startswith("transfer"):
|
| 160 |
+
m.update(trainable_sources=[], holdout_sources=manifest["eval_only_sources"], eval_only=True)
|
| 161 |
+
for split, records in partitions.items():
|
| 162 |
+
path = folder / f"{split}.jsonl"
|
| 163 |
+
payload = b""
|
| 164 |
+
inherited_records = inherited_questions = 0
|
| 165 |
+
if split == "test":
|
| 166 |
+
test_src = Path(inherit_eval) if (inherit_eval and name.startswith("decision")) else inherited
|
| 167 |
+
old = read_manifest(test_src)["files"]["test.jsonl"]
|
| 168 |
+
if digest(test_src / "test.jsonl") != old["sha256"]:
|
| 169 |
+
raise ValueError("inherited test checksum mismatch")
|
| 170 |
+
payload = (test_src / "test.jsonl").read_bytes()
|
| 171 |
+
inherited_records, inherited_questions = old["records"], old["questions"]
|
| 172 |
+
if inherit_eval and name.startswith("decision"): records = [] # the inherited test already contains the locked composition groups
|
| 173 |
+
payload += "".join(json.dumps(r, ensure_ascii=False) + "\n" for r in records).encode()
|
| 174 |
+
path.write_bytes(payload)
|
| 175 |
+
m["files"][path.name] = {"sha256": digest(path), "records": inherited_records + len(records),
|
| 176 |
+
"questions": inherited_questions + sum(len(r["questions"]) for r in records)}
|
| 177 |
+
write_json(folder / "manifest.json", m)
|
| 178 |
+
print(name, {s: v["records"] for s, v in m["files"].items()}, flush=True)
|
| 179 |
+
if any(digest(p) != h for p, h in parent_hashes.items()):
|
| 180 |
+
raise ValueError("parent artifacts changed")
|
| 181 |
+
return manifest
|
| 182 |
+
|
| 183 |
+
|
| 184 |
+
def smoke_subset(source, out):
|
| 185 |
+
source, out = Path(source), Path(out)
|
| 186 |
+
manifest = read_manifest(source)
|
| 187 |
+
out.mkdir(parents=True, exist_ok=False)
|
| 188 |
+
manifest["files"] = {}
|
| 189 |
+
for split in SPLITS:
|
| 190 |
+
records = []
|
| 191 |
+
if split != "test":
|
| 192 |
+
groups = defaultdict(list)
|
| 193 |
+
for r in load_split(source, split):
|
| 194 |
+
groups[(r["_meta"]["source"], r["_meta"]["group_id"])].append(r)
|
| 195 |
+
taken = Counter()
|
| 196 |
+
for (name, _), group in groups.items():
|
| 197 |
+
if taken[name] < 2:
|
| 198 |
+
records.extend(group); taken[name] += 1
|
| 199 |
+
path = out / f"{split}.jsonl"
|
| 200 |
+
write_jsonl(path, records)
|
| 201 |
+
manifest["files"][path.name] = {"sha256": digest(path), "records": len(records),
|
| 202 |
+
"questions": sum(len(r["questions"]) for r in records)}
|
| 203 |
+
manifest["smoke_only"] = True
|
| 204 |
+
write_json(out / "manifest.json", manifest)
|
| 205 |
+
|
| 206 |
+
|
| 207 |
+
def main():
|
| 208 |
+
ap = argparse.ArgumentParser()
|
| 209 |
+
ap.add_argument("--out", required=True)
|
| 210 |
+
ap.add_argument("--smoke-from")
|
| 211 |
+
ap.add_argument("--public-train", help="frozen suite whose train/calibration public partitions replace the source's (larger pool)")
|
| 212 |
+
ap.add_argument("--synthetic-scale", type=int, default=1, help="multiply synthetic training pairs/groups per arm (both arms stay equal)")
|
| 213 |
+
ap.add_argument("--inherit-eval", help="decision suite whose development/test bytes are reused verbatim")
|
| 214 |
+
ap.add_argument("--random-structures", type=int, default=0, help="compositional arm from N random rule trees (held-out structures excluded)")
|
| 215 |
+
ap.add_argument("--groups-per-structure", type=int, default=8)
|
| 216 |
+
ap.add_argument("--train-styles", default="0,1", help="rendering styles for training compositional records (2 is locked-test only)")
|
| 217 |
+
ap.add_argument("--unmatched-arms", action="store_true", help="allow the two synthetic arms to differ in size")
|
| 218 |
+
ap.add_argument("--legacy-families", default="v3", choices=["v3", "all"], help="'all' adds the ordinal Score threshold families (deadline stays held out)")
|
| 219 |
+
ap.add_argument("--structures-seed", help="reuse another version's random structures (e.g. v7) so only the legacy arm differs")
|
| 220 |
+
a = ap.parse_args()
|
| 221 |
+
styles = tuple(int(x) for x in a.train_styles.split(","))
|
| 222 |
+
if 2 in styles: ap.error("rendering style 2 is reserved for the locked test")
|
| 223 |
+
if a.smoke_from:
|
| 224 |
+
smoke_subset(a.smoke_from, a.out)
|
| 225 |
+
else:
|
| 226 |
+
freeze(a.out, public_train=a.public_train, synthetic_scale=a.synthetic_scale, inherit_eval=a.inherit_eval,
|
| 227 |
+
random_structures=a.random_structures, groups_per_structure=a.groups_per_structure, train_styles=styles, matched_arms=not a.unmatched_arms,
|
| 228 |
+
legacy_families=FAMILIES + ORDINAL_FAMILIES if a.legacy_families == "all" else FAMILIES, structures_seed=a.structures_seed)
|
| 229 |
+
|
| 230 |
+
|
| 231 |
+
if __name__ == "__main__":
|
| 232 |
+
main()
|
runtime/kev/suite.py
ADDED
|
@@ -0,0 +1,450 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import argparse
|
| 2 |
+
import copy
|
| 3 |
+
import fcntl
|
| 4 |
+
import hashlib
|
| 5 |
+
import json
|
| 6 |
+
import os
|
| 7 |
+
import random
|
| 8 |
+
import shutil
|
| 9 |
+
from collections import Counter
|
| 10 |
+
from contextlib import contextmanager
|
| 11 |
+
from pathlib import Path
|
| 12 |
+
|
| 13 |
+
from kev.composition import DEV_SHAPES, HELD_OUT_KEYS, TEST_SHAPES, TRAIN_SHAPES
|
| 14 |
+
from kev.contrastive import FAMILIES, generate
|
| 15 |
+
from kev.data import ALL_REPOS, ALL_SOURCES, EVAL_ONLY, REPOS, SOURCES, TRAINABLE, TRANSFER_REPOS, TRANSFER_SOURCES, build, dataset_ref, materialize, source_seed
|
| 16 |
+
from kev.model import (MAX_BRANCH, SERVE_MAX_BRANCH, SERVE_MAX_BRANCH_8K, SERVE_MAX_PACKED, SERVE_MAX_STATE, SERVE_MAX_STATE_8K, fits, load_tokenizer,
|
| 17 |
+
training_context)
|
| 18 |
+
|
| 19 |
+
SPLITS = ("train", "calibration", "development", "test")
|
| 20 |
+
BASES = ("Qwen/Qwen2.5-0.5B", "Qwen/Qwen3-0.6B-Base")
|
| 21 |
+
# the encoder limits every frozen record satisfies, as written into manifests ("context")
|
| 22 |
+
CONTEXT = {**training_context(), "truncate": False}
|
| 23 |
+
# what a manifest records for an eval-only suite frozen as published rather than admitted to the training context
|
| 24 |
+
SERVING_CONTEXT = {"max_state": SERVE_MAX_STATE, "max_branch": SERVE_MAX_BRANCH, "max_packed": SERVE_MAX_PACKED, "truncate": False}
|
| 25 |
+
# the serving context before 64k states: what the suites frozen until then record and were admitted under; their builders
|
| 26 |
+
# (hard-v1, devtools-v1, breadth-v1, long states, semif/typesafe, documents-v1) keep it and kev.model.MAX_TRAIN_STATE_8K,
|
| 27 |
+
# so they still rebuild byte for byte
|
| 28 |
+
SERVING_CONTEXT_8K = {"max_state": SERVE_MAX_STATE_8K, "max_branch": SERVE_MAX_BRANCH_8K, "max_packed": SERVE_MAX_STATE_8K + SERVE_MAX_BRANCH_8K, "truncate": False}
|
| 29 |
+
# Clean records are admitted with this many branch tokens to spare, so the variants that add an option (contrast_cases'
|
| 30 |
+
# none-of-these, training-time none/distractor augmentation) still encode under MAX_BRANCH.
|
| 31 |
+
ADMISSION_BRANCH_HEADROOM = 64
|
| 32 |
+
# Frozen suites are mirrored on the Hub. Manifests (with the sha256 of every partition) and the development/test
|
| 33 |
+
# partitions live in git; large training partitions are fetched from this dataset on first use and verified against
|
| 34 |
+
# the manifest, so the suite hash and every provenance record stay unchanged. A suite whose partitions must never be
|
| 35 |
+
# public (a held-out test set) names its own mirror in the manifest, {"mirror": {"dataset": ..., "revision": ...}},
|
| 36 |
+
# usually the private PRIVATE_DATASET; only its manifest is in git, which publishes the hashes but not the text.
|
| 37 |
+
SUITES_DATASET = "jaredpalmer/kev-suites"
|
| 38 |
+
PRIVATE_DATASET = "jaredpalmer/kev-private-evals"
|
| 39 |
+
SUITES_REVISION = "cc4bac803e73112689ec327ffa481c519cbc7a05"
|
| 40 |
+
# partitions larger than this stay out of git (gitignored; the manifest's sha256 still pins them)
|
| 41 |
+
GIT_LIMIT = 10 * 1024 * 1024
|
| 42 |
+
# the pinned tokenizer suites built for the Qwen3.5 family are admitted and length-counted under (hard-v1, devtools-v1, long states)
|
| 43 |
+
ADMISSION_TOKENIZER = ("Qwen/Qwen3.5-4B-Base", "1001bb4d826a52d1f399e183466143f4da7b741b")
|
| 44 |
+
# programmatic policy sources (kev.study_v3 / kev.contrastive); the trainer's mix ablations treat them as one group
|
| 45 |
+
SYNTHETIC_SOURCES = ("legacy_policy", "compositional", "contrastive")
|
| 46 |
+
# Suites deleted from the repo because they are unsound as a gate: {suite dir: why, when, the last round that read it}.
|
| 47 |
+
# load_split / read_manifest refuse them with the reason; kev.rounds.validate lists a read of one as archived for rounds up
|
| 48 |
+
# to `last_round` (their committed rows under runs/ are the record and still reproduce) and refuses it after that.
|
| 49 |
+
REMOVED_SUITES = {
|
| 50 |
+
"evals/external/scienthoon-v1": {
|
| 51 |
+
"removed": "2026-09-27",
|
| 52 |
+
"last_round": 22,
|
| 53 |
+
"reason": ("unsound as a gate: 291 templated synthetic support tickets x 3 questions; `queue` is saturated (0.948-0.952 "
|
| 54 |
+
"for every 27B), `priority` is unlearnable by construction (its own manifest: the label follows an org rule "
|
| 55 |
+
"absent from the text), and `angry` has 15 of 291 gold labels that contradict the text and turns on ~12 stock "
|
| 56 |
+
"closing phrases with disputed conventions"),
|
| 57 |
+
"record": "PLAN.md (Standing rules; 2026-09-27 note); committed rows under runs/ (e.g. runs/r20-scienthoon)",
|
| 58 |
+
},
|
| 59 |
+
"evals/external/wanli-v2": {
|
| 60 |
+
"removed": "2026-09-30",
|
| 61 |
+
"last_round": 26,
|
| 62 |
+
"reason": ("unsound as a gate: 271 of its 1,002 WANLI test pairs (27 %) are ones WANLI's two crowd annotators labelled "
|
| 63 |
+
"differently, and the published gold is one of the two labels; every Kev scores 49-62 % on those against "
|
| 64 |
+
"64-81 % on the 731 the annotators agreed on, and the 2026-09-27 audit measured split-half r 0.04 across 23 "
|
| 65 |
+
"checkpoints (all within 0.735-0.763), a half-width as wide as its 2 pp bar and ~11 % invalid labels"),
|
| 66 |
+
"record": "PLAN.md (Standing rules; 2026-09-30 note); committed rows under runs/ (e.g. runs/r23-27b-k-w85-wanli2)",
|
| 67 |
+
},
|
| 68 |
+
"evals/external/wanli-v1": {
|
| 69 |
+
"removed": "2026-09-30",
|
| 70 |
+
"last_round": 5,
|
| 71 |
+
"reason": ("the same WANLI test pairs as wanli-v2 (SemIf's 256): 63 of 256 (25 %) are ones WANLI's two crowd annotators "
|
| 72 |
+
"labelled differently, and the published gold is one of the two labels"),
|
| 73 |
+
"record": "PLAN.md (2026-09-30 note); README; committed rows under runs/ (e.g. runs/kev-9b-wanli-v1, runs/jev-wanli-v1)",
|
| 74 |
+
},
|
| 75 |
+
"evals/external/typesafe-v1": {
|
| 76 |
+
"removed": "2026-09-30",
|
| 77 |
+
"last_round": 26,
|
| 78 |
+
"reason": ("unsound as a gate: its gold is not a ground truth but the argmax of the average of two closed frontier models' "
|
| 79 |
+
"answers (evals.typesafe.ai: GPT-6 Astra and Claude Fable 5.1), 13 of its 102 reference distributions put "
|
| 80 |
+
"the answer below 0.75 (the two references split), and on 89 answered questions from 20 cases the split-half "
|
| 81 |
+
"correlation across 20 Kev-27B checkpoints is -0.01 (2026-09-27 audit: -0.27): its differences between "
|
| 82 |
+
"checkpoints are noise"),
|
| 83 |
+
"record": "PLAN.md (Standing rules; 2026-09-30 note); README; committed rows under runs/ (e.g. runs/kev-9b-typesafe-v1)",
|
| 84 |
+
},
|
| 85 |
+
}
|
| 86 |
+
|
| 87 |
+
|
| 88 |
+
class RemovedSuite(ValueError):
|
| 89 |
+
"""A read of a suite in REMOVED_SUITES."""
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
def suite_key(path):
|
| 93 |
+
"""'evals/...' for a suite path given relative, absolute or as a container saw it (/root/kev/evals/x); None if none."""
|
| 94 |
+
parts = Path(str(path)).parts
|
| 95 |
+
return str(Path(*parts[parts.index("evals"):])) if "evals" in parts else None
|
| 96 |
+
|
| 97 |
+
|
| 98 |
+
def removed_suite(path):
|
| 99 |
+
"""The REMOVED_SUITES entry of a suite path (any form suite_key accepts), or None."""
|
| 100 |
+
key = suite_key(path)
|
| 101 |
+
return REMOVED_SUITES.get(key) if key else None
|
| 102 |
+
|
| 103 |
+
|
| 104 |
+
def refuse_removed(path):
|
| 105 |
+
if entry := removed_suite(path):
|
| 106 |
+
raise RemovedSuite(f"{suite_key(path)} was removed on {entry['removed']}: {entry['reason']}. Rounds up to "
|
| 107 |
+
f"{entry['last_round']} read it; their committed rows are the record ({entry['record']}). Do not read it again.")
|
| 108 |
+
|
| 109 |
+
|
| 110 |
+
def digest(path):
|
| 111 |
+
h = hashlib.sha256()
|
| 112 |
+
with Path(path).open("rb") as stream:
|
| 113 |
+
for block in iter(lambda: stream.read(1024 * 1024), b""):
|
| 114 |
+
h.update(block)
|
| 115 |
+
return h.hexdigest()
|
| 116 |
+
|
| 117 |
+
|
| 118 |
+
def record_digest(record):
|
| 119 |
+
return hashlib.sha256(json.dumps(record, ensure_ascii=False, separators=(",", ":")).encode()).hexdigest()
|
| 120 |
+
|
| 121 |
+
|
| 122 |
+
def normalise_text(text):
|
| 123 |
+
"""Casefolded, whitespace runs collapsed to one space: the text two states are compared on for exact deduplication."""
|
| 124 |
+
return " ".join(text.casefold().split())
|
| 125 |
+
|
| 126 |
+
|
| 127 |
+
def text_digest(text):
|
| 128 |
+
"""sha256 of normalise_text(text): the `text_sha256` of suite builders (kev.data computes the same inline; it cannot
|
| 129 |
+
import this module). scripts/screen_overlap.py tokenises differently on purpose (words only, for n-gram overlap)."""
|
| 130 |
+
return hashlib.sha256(normalise_text(text).encode()).hexdigest()
|
| 131 |
+
|
| 132 |
+
|
| 133 |
+
# Every JSON/JSONL file this repo writes is UTF-8 with LF line endings, whatever the platform's locale says (issue #12:
|
| 134 |
+
# frozen partitions are sha256-checked byte for byte, and they contain non-ASCII text). Read them the same way.
|
| 135 |
+
ENCODING = "utf-8"
|
| 136 |
+
|
| 137 |
+
|
| 138 |
+
def read_json(path):
|
| 139 |
+
return json.loads(Path(path).read_text(encoding=ENCODING))
|
| 140 |
+
|
| 141 |
+
|
| 142 |
+
def write_json(path, value, atomic=False):
|
| 143 |
+
"""atomic: write a sibling temp file and os.replace it over `path`, so a reader (or a restart after a crash mid-write)
|
| 144 |
+
sees the old file or the new one, never a torn one. For state files rewritten in place (kev.rounds' watcher)."""
|
| 145 |
+
text = json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + "\n"
|
| 146 |
+
if not atomic:
|
| 147 |
+
Path(path).write_text(text, encoding=ENCODING); return
|
| 148 |
+
tmp = Path(path).with_name(f".{Path(path).name}.tmp")
|
| 149 |
+
tmp.write_text(text, encoding=ENCODING)
|
| 150 |
+
os.replace(tmp, path)
|
| 151 |
+
|
| 152 |
+
|
| 153 |
+
@contextmanager
|
| 154 |
+
def file_lock(path):
|
| 155 |
+
"""Hold an exclusive advisory lock on `path` (created if absent) for the block; a second holder waits. Local
|
| 156 |
+
orchestration only: one pull of a study (modal_app.pull_lock), one launch of an arm's reads (kev.rounds)."""
|
| 157 |
+
Path(path).parent.mkdir(parents=True, exist_ok=True)
|
| 158 |
+
with Path(path).open("a", encoding=ENCODING) as lock:
|
| 159 |
+
fcntl.flock(lock, fcntl.LOCK_EX)
|
| 160 |
+
yield
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
def read_jsonl(path):
|
| 164 |
+
# split on "\n" only: str.splitlines() also breaks on U+2028, U+2029 and U+0085, which write_jsonl leaves unescaped
|
| 165 |
+
return [json.loads(line) for line in Path(path).read_text(encoding=ENCODING).split("\n") if line.strip()]
|
| 166 |
+
|
| 167 |
+
|
| 168 |
+
def write_jsonl(path, records):
|
| 169 |
+
Path(path).write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in records), encoding=ENCODING)
|
| 170 |
+
|
| 171 |
+
|
| 172 |
+
def semantic_hash(r):
|
| 173 |
+
"""State hash that ignores sentence order for policy/case records (contrastive, compositional), so the same case
|
| 174 |
+
told in a different order counts as the same state; record_digest of the state otherwise."""
|
| 175 |
+
state = r["state"]
|
| 176 |
+
if isinstance(state, dict) and "policy" in state and "case" in state:
|
| 177 |
+
state = {"policy": state["policy"], "sentences": sorted(s.rstrip(".") for s in state["case"].split(". "))}
|
| 178 |
+
return record_digest(state)
|
| 179 |
+
|
| 180 |
+
|
| 181 |
+
def validate_training(records, manifest):
|
| 182 |
+
"""Every training record must come from a source the manifest declares trainable, never from an eval-only or held-out
|
| 183 |
+
one, and compositional records must not use a held-out rule structure."""
|
| 184 |
+
allowed = set(manifest.get("trainable_sources", []))
|
| 185 |
+
forbidden = set(EVAL_ONLY) | set(manifest.get("eval_only_sources", [])) | set(manifest.get("holdout_sources", []))
|
| 186 |
+
for r in records:
|
| 187 |
+
m = r["_meta"]
|
| 188 |
+
if m["source"] in forbidden or (allowed and m["source"] not in allowed):
|
| 189 |
+
raise ValueError(f"eval-only or undeclared training source: {m['source']}")
|
| 190 |
+
if m["source"] == "compositional":
|
| 191 |
+
held_shape = m["family"] in DEV_SHAPES + TEST_SHAPES
|
| 192 |
+
held_structure = m.get("structure") in HELD_OUT_KEYS
|
| 193 |
+
if held_shape or held_structure or (m["family"] not in TRAIN_SHAPES and not m["family"].startswith("rand:")):
|
| 194 |
+
raise ValueError("held-out compositional structure in training")
|
| 195 |
+
if not records:
|
| 196 |
+
raise ValueError("empty training partition")
|
| 197 |
+
|
| 198 |
+
|
| 199 |
+
def read_manifest(directory):
|
| 200 |
+
refuse_removed(directory)
|
| 201 |
+
return read_json(Path(directory) / "manifest.json")
|
| 202 |
+
|
| 203 |
+
|
| 204 |
+
def load_split(directory, split, allow_test=False):
|
| 205 |
+
refuse_removed(directory) # a removed suite is refused with its reason, not a missing-file error
|
| 206 |
+
if split not in SPLITS:
|
| 207 |
+
raise ValueError(f"unknown split: {split}")
|
| 208 |
+
if split == "test" and not allow_test:
|
| 209 |
+
raise ValueError("locked test requires explicit --allow-test; never use it for search")
|
| 210 |
+
directory = Path(directory)
|
| 211 |
+
manifest = read_manifest(directory)
|
| 212 |
+
path = directory / f"{split}.jsonl"
|
| 213 |
+
if not path.exists():
|
| 214 |
+
fetch_partition(directory, path.name)
|
| 215 |
+
if digest(path) != manifest["files"][path.name]["sha256"]:
|
| 216 |
+
raise ValueError(f"suite checksum mismatch: {path}")
|
| 217 |
+
records = read_jsonl(path)
|
| 218 |
+
if len(records) != manifest["files"][path.name]["records"]:
|
| 219 |
+
raise ValueError("suite record count mismatch")
|
| 220 |
+
return records
|
| 221 |
+
|
| 222 |
+
|
| 223 |
+
def fetch_partition(directory, filename):
|
| 224 |
+
"""Download one partition of a frozen suite from its Hub mirror into place: the manifest's own "mirror" if it names
|
| 225 |
+
one, else SUITES_DATASET@SUITES_REVISION. The caller verifies the sha256."""
|
| 226 |
+
from huggingface_hub import hf_hub_download
|
| 227 |
+
from huggingface_hub.errors import GatedRepoError, RepositoryNotFoundError
|
| 228 |
+
directory = Path(directory).resolve()
|
| 229 |
+
evals_root = next((p for p in directory.parents if p.name == "evals"), None)
|
| 230 |
+
if evals_root is None:
|
| 231 |
+
raise FileNotFoundError(f"{directory / filename} is missing and is not under an evals/ tree")
|
| 232 |
+
relative = directory.relative_to(evals_root) / filename
|
| 233 |
+
mirror = read_manifest(directory).get("mirror")
|
| 234 |
+
repo, revision = (mirror["dataset"], mirror["revision"]) if mirror else (SUITES_DATASET, SUITES_REVISION) # a named mirror pins its own revision
|
| 235 |
+
try:
|
| 236 |
+
cached = hf_hub_download(repo, str(relative), repo_type="dataset", revision=revision)
|
| 237 |
+
except (RepositoryNotFoundError, GatedRepoError) as e: # a private mirror answers "not found" to anyone without access
|
| 238 |
+
raise PermissionError(f"{relative} is only in {repo}, which is missing or private to this account; `hf auth login` "
|
| 239 |
+
"(or HF_TOKEN) with access to it, or ask for it") from e
|
| 240 |
+
shutil.copyfile(cached, directory / filename)
|
| 241 |
+
print(f"fetched {relative} from {repo}@{revision[:10]}", flush=True)
|
| 242 |
+
|
| 243 |
+
|
| 244 |
+
def case_copy(record, variant):
|
| 245 |
+
result = copy.deepcopy(record)
|
| 246 |
+
result["_meta"]["group_id"] = record["_meta"].get("group_id", record["_meta"]["id"]) # bootstrap unit (pair for contrastive)
|
| 247 |
+
result["_meta"]["parent_id"] = record["_meta"]["id"] # the clean record this variant perturbs
|
| 248 |
+
result["_meta"]["id"] += "/" + variant
|
| 249 |
+
result["_meta"]["variant"] = variant
|
| 250 |
+
return result
|
| 251 |
+
|
| 252 |
+
|
| 253 |
+
def contrast_cases(record, seed=0):
|
| 254 |
+
candidates = [(qid, q) for qid, q in record["questions"].items() if q["type"] == "choice" and len(q["criteria"]) >= 3]
|
| 255 |
+
if not candidates:
|
| 256 |
+
return []
|
| 257 |
+
qid, q = candidates[0]
|
| 258 |
+
nk = "none_of_these"
|
| 259 |
+
if nk in q["criteria"]:
|
| 260 |
+
raise ValueError("reserved contrast option collision")
|
| 261 |
+
out = []
|
| 262 |
+
for variant in ("none_present", "none_absent"):
|
| 263 |
+
r = case_copy(record, variant)
|
| 264 |
+
r["questions"] = {qid: copy.deepcopy(q)}
|
| 265 |
+
rq = r["questions"][qid]
|
| 266 |
+
rq["criteria"][nk] = "None of these options describes the answer"
|
| 267 |
+
if variant == "none_absent":
|
| 268 |
+
rq["criteria"].pop(rq["label"])
|
| 269 |
+
rq["label"] = nk
|
| 270 |
+
keys = list(rq["criteria"])
|
| 271 |
+
random.Random(source_seed(seed, record["_meta"]["id"])).shuffle(keys)
|
| 272 |
+
rq["criteria"] = {k: rq["criteria"][k] for k in keys}
|
| 273 |
+
r["_meta"]["none_key"] = nk
|
| 274 |
+
out.append(r)
|
| 275 |
+
r = case_copy(record, "permuted")
|
| 276 |
+
rng = random.Random(source_seed(seed, record["_meta"]["id"]))
|
| 277 |
+
for q in r["questions"].values():
|
| 278 |
+
if q["type"] == "choice":
|
| 279 |
+
keys = list(q["criteria"])
|
| 280 |
+
rng.shuffle(keys)
|
| 281 |
+
q["criteria"] = {k: q["criteria"][k] for k in keys}
|
| 282 |
+
out.append(r)
|
| 283 |
+
return out
|
| 284 |
+
|
| 285 |
+
|
| 286 |
+
def select_unique(records, count, seen, tokenizers, report):
|
| 287 |
+
selected = []
|
| 288 |
+
for record in sorted(records, key=lambda r: r["_meta"]["row_sha256"]):
|
| 289 |
+
report["considered"] += 1
|
| 290 |
+
key = record["_meta"]["text_sha256"]
|
| 291 |
+
if key in seen:
|
| 292 |
+
report["duplicate_state"] += 1
|
| 293 |
+
continue
|
| 294 |
+
rec = materialize(record) # a record that fails request validation is a converter bug: let it abort the freeze
|
| 295 |
+
if not fits(rec, *tokenizers, max_branch=MAX_BRANCH - ADMISSION_BRANCH_HEADROOM):
|
| 296 |
+
report["context_rejected"] += 1
|
| 297 |
+
continue
|
| 298 |
+
record["_meta"].update(group_id=record["_meta"]["id"], variant="clean")
|
| 299 |
+
seen.add(key)
|
| 300 |
+
selected.append(record)
|
| 301 |
+
report["accepted"] += 1
|
| 302 |
+
if len(selected) == count:
|
| 303 |
+
return selected
|
| 304 |
+
raise ValueError(f"only {len(selected)}/{count} records fit the common context policy")
|
| 305 |
+
|
| 306 |
+
|
| 307 |
+
def training_state_hashes(suite_dir):
|
| 308 |
+
"""Normalized state hashes of a frozen suite's training + calibration partitions, for exact-match contamination checks."""
|
| 309 |
+
hashes = set()
|
| 310 |
+
for split in ("train", "calibration"):
|
| 311 |
+
for record in load_split(suite_dir, split):
|
| 312 |
+
hashes.add(record["_meta"]["text_sha256"])
|
| 313 |
+
return hashes
|
| 314 |
+
|
| 315 |
+
|
| 316 |
+
def freeze(directory, train=300, calibration=40, development=80, test=80, seed=20260918, holdout=("mnli", "sst5"),
|
| 317 |
+
sources=None, repos=None, exclude_states_from=None, contrastive_pairs=0, contrastive_holdout_families=(), contrastive_eval_pairs=40):
|
| 318 |
+
from huggingface_hub import HfApi
|
| 319 |
+
|
| 320 |
+
sources = SOURCES if sources is None else sources
|
| 321 |
+
repos = REPOS if repos is None else repos
|
| 322 |
+
eval_only = set(holdout) >= set(sources)
|
| 323 |
+
directory = Path(directory)
|
| 324 |
+
if directory.exists():
|
| 325 |
+
raise FileExistsError(f"refusing to overwrite frozen suite {directory}")
|
| 326 |
+
hub = HfApi()
|
| 327 |
+
revisions = {repo: hub.dataset_info(*dataset_ref(repo)[:1], revision=dataset_ref(repo)[1]).sha for repo in set(repos.values())}
|
| 328 |
+
base_revisions = {base: hub.model_info(base).sha for base in BASES}
|
| 329 |
+
tokenizers = [load_tokenizer(base, revision=revision) for base, revision in base_revisions.items()]
|
| 330 |
+
trainable_here = [x for x in sources if x not in holdout]
|
| 331 |
+
if set(trainable_here) & set(EVAL_ONLY):
|
| 332 |
+
raise ValueError(f"eval-only sources cannot be trainable: {sorted(set(trainable_here) & set(EVAL_ONLY))}")
|
| 333 |
+
manifest = {
|
| 334 |
+
"version": 2, "seed": seed, "holdout_sources": list(holdout), "eval_only": eval_only,
|
| 335 |
+
"trainable_sources": trainable_here, "eval_only_sources": [x for x in sources if x in holdout],
|
| 336 |
+
"policy": {"trainable": list(TRAINABLE), "eval_only": list(EVAL_ONLY)},
|
| 337 |
+
"dataset_revisions": revisions, "base_revisions": base_revisions,
|
| 338 |
+
"context": {**CONTEXT, "admission_branch_headroom": ADMISSION_BRANCH_HEADROOM},
|
| 339 |
+
"selection": "Normalized exact-state deduplication across partitions; common tokenizer context admission; no fuzzy decontamination or pretraining-contamination claim.",
|
| 340 |
+
"legacy_checkpoints": "Training/calibration overlap for pre-manifest checkpoints is unknown; exploratory only.",
|
| 341 |
+
"objective": "Negative macro-average clean development NLL, equal weight per task; raw probabilities.",
|
| 342 |
+
"excluded_training_states_from": str(exclude_states_from) if exclude_states_from else None,
|
| 343 |
+
"files": {}, "admission": {},
|
| 344 |
+
}
|
| 345 |
+
partitions = {split: [] for split in SPLITS}
|
| 346 |
+
seen = set(training_state_hashes(exclude_states_from)) if exclude_states_from else set()
|
| 347 |
+
manifest["excluded_training_state_hashes"] = len(seen)
|
| 348 |
+
for source in sources:
|
| 349 |
+
report = Counter()
|
| 350 |
+
test_pool = build(max(3 * (development + test), 600), "test", seed, only=[source], revisions=revisions, sources=sources, repos=repos)
|
| 351 |
+
chosen = select_unique(test_pool, development + test, seen, tokenizers, report)
|
| 352 |
+
partitions["development"].extend(chosen[:development])
|
| 353 |
+
partitions["test"].extend(chosen[development:])
|
| 354 |
+
if source not in holdout:
|
| 355 |
+
train_pool = build(max(3 * (train + calibration), 800), "train", seed, only=[source], revisions=revisions, sources=sources, repos=repos)
|
| 356 |
+
chosen = select_unique(train_pool, train + calibration, seen, tokenizers, report)
|
| 357 |
+
partitions["calibration"].extend(chosen[:calibration])
|
| 358 |
+
partitions["train"].extend(chosen[calibration:])
|
| 359 |
+
manifest["admission"][source] = dict(report)
|
| 360 |
+
print(f"froze {source}: {dict(report)}", flush=True)
|
| 361 |
+
if contrastive_pairs:
|
| 362 |
+
families = list(FAMILIES)
|
| 363 |
+
held = list(contrastive_holdout_families)
|
| 364 |
+
unknown = set(held) - set(families)
|
| 365 |
+
if unknown: raise ValueError(f"unknown contrastive families: {sorted(unknown)}")
|
| 366 |
+
trainable_fams = [f for f in families if f not in held]
|
| 367 |
+
# training/calibration pairs and development/test pairs come from disjoint seeds within trainable families;
|
| 368 |
+
# held-out families appear only in development/test (never trained anywhere)
|
| 369 |
+
train_recs, rep_train = generate(contrastive_pairs, seed=f"{seed}-train", families=trainable_fams) if trainable_fams and not eval_only else ([], {})
|
| 370 |
+
# eval-only suites carry the held-out families only; trainable suites carry their trainable families only
|
| 371 |
+
eval_families = (held or families) if eval_only else trainable_fams
|
| 372 |
+
eval_recs, rep_eval = generate(contrastive_eval_pairs, seed=f"{seed}-eval", families=eval_families)
|
| 373 |
+
for r in train_recs + eval_recs:
|
| 374 |
+
r["_meta"].update(group_id=r["_meta"]["family_id"], variant="clean") # siblings are one bootstrap unit
|
| 375 |
+
if r["_meta"]["text_sha256"] in seen: raise ValueError("contrastive state collides with an existing state")
|
| 376 |
+
seen.update(r["_meta"]["text_sha256"] for r in train_recs + eval_recs)
|
| 377 |
+
n_cal = 2 * max(len(train_recs) // 20, 1) if train_recs else 0
|
| 378 |
+
partitions["calibration"].extend(train_recs[:n_cal]); partitions["train"].extend(train_recs[n_cal:])
|
| 379 |
+
# stratified by family: alternate pairs (siblings adjacent) between development and test
|
| 380 |
+
for i in range(0, len(eval_recs), 2):
|
| 381 |
+
partitions["development" if (i // 2) % 2 == 0 else "test"].extend(eval_recs[i : i + 2])
|
| 382 |
+
manifest["contrastive"] = {"pairs_per_family_train": contrastive_pairs, "pairs_per_family_eval": contrastive_eval_pairs,
|
| 383 |
+
"eval_families": eval_families, "trainable_families": trainable_fams, "eval_only_families": held,
|
| 384 |
+
"train_report": rep_train, "eval_report": rep_eval,
|
| 385 |
+
"note": "labels from rule evaluators; ablation + invariance checks passed for every kept pair; no LLM"}
|
| 386 |
+
if trainable_fams and not eval_only: manifest["trainable_sources"].append("contrastive")
|
| 387 |
+
else: manifest["eval_only_sources"].append("contrastive")
|
| 388 |
+
print(f"contrastive: {len(train_recs)} train/cal records, {len(eval_recs)} dev/test records; held-out families {held}", flush=True)
|
| 389 |
+
for split in ("development", "test"):
|
| 390 |
+
extras = []
|
| 391 |
+
per_source = Counter()
|
| 392 |
+
for record in partitions[split]:
|
| 393 |
+
source = record["_meta"]["source"]
|
| 394 |
+
if per_source[source] < 12:
|
| 395 |
+
variants = contrast_cases(record, seed)
|
| 396 |
+
for variant in variants:
|
| 397 |
+
if not fits(materialize(variant), *tokenizers):
|
| 398 |
+
raise ValueError(f"variant of {record['_meta']['id']} exceeds the training context")
|
| 399 |
+
extras.extend(variants)
|
| 400 |
+
per_source[source] += bool(variants)
|
| 401 |
+
partitions[split].extend(extras)
|
| 402 |
+
directory.mkdir(parents=True)
|
| 403 |
+
for split, records in partitions.items():
|
| 404 |
+
path = directory / f"{split}.jsonl"
|
| 405 |
+
write_jsonl(path, records)
|
| 406 |
+
manifest["files"][path.name] = {"sha256": digest(path), "records": len(records),
|
| 407 |
+
"questions": sum(len(r["questions"]) for r in records)}
|
| 408 |
+
manifest["code_hashes"] = {name: digest(Path(__file__).parent / name) for name in ("data.py", "api.py", "model.py", "suite.py")}
|
| 409 |
+
write_json(directory / "manifest.json", manifest)
|
| 410 |
+
print(json.dumps(manifest["files"], indent=2), flush=True)
|
| 411 |
+
|
| 412 |
+
|
| 413 |
+
def main():
|
| 414 |
+
ap = argparse.ArgumentParser()
|
| 415 |
+
ap.add_argument("--out", required=True)
|
| 416 |
+
ap.add_argument("--train", type=int, default=300)
|
| 417 |
+
ap.add_argument("--calibration", type=int, default=40)
|
| 418 |
+
ap.add_argument("--development", type=int, default=80)
|
| 419 |
+
ap.add_argument("--test", type=int, default=80)
|
| 420 |
+
ap.add_argument("--seed", type=int, default=20260918)
|
| 421 |
+
ap.add_argument("--transfer", action="store_true", help="eval-only suite from the eight transfer-v1 sources (no train/calibration partitions)")
|
| 422 |
+
ap.add_argument("--sources", help="comma-separated source names (any of ALL_SOURCES); trainable ones get train/calibration partitions")
|
| 423 |
+
ap.add_argument("--holdout", help="comma-separated sources kept eval-only within --sources (eval-only policy sources are always held out)")
|
| 424 |
+
ap.add_argument("--exclude-states-from", help="frozen suite whose train+calibration states must not appear here")
|
| 425 |
+
ap.add_argument("--contrastive-pairs", type=int, default=0, help="programmatic contrastive pairs per trainable family for training (0 = none)")
|
| 426 |
+
ap.add_argument("--contrastive-holdout", default="", help="comma-separated contrastive families kept eval-only")
|
| 427 |
+
a = ap.parse_args()
|
| 428 |
+
cfam = tuple(x for x in a.contrastive_holdout.split(",") if x)
|
| 429 |
+
if min(a.development, a.test) < 1:
|
| 430 |
+
ap.error("split sizes must be positive")
|
| 431 |
+
if a.transfer:
|
| 432 |
+
freeze(a.out, 0, 0, a.development, a.test, a.seed, holdout=tuple(TRANSFER_SOURCES), sources=TRANSFER_SOURCES,
|
| 433 |
+
repos=TRANSFER_REPOS, exclude_states_from=a.exclude_states_from, contrastive_pairs=a.contrastive_pairs, contrastive_holdout_families=cfam)
|
| 434 |
+
elif a.sources:
|
| 435 |
+
names = [x for x in a.sources.split(",") if x]
|
| 436 |
+
unknown = set(names) - set(ALL_SOURCES)
|
| 437 |
+
if unknown: ap.error(f"unknown sources: {sorted(unknown)}")
|
| 438 |
+
holdout = tuple(x for x in names if x in EVAL_ONLY or x in (a.holdout or "").split(","))
|
| 439 |
+
if any(x not in holdout for x in names) and min(a.train, a.calibration) < 1:
|
| 440 |
+
ap.error("train and calibration sizes must be positive when a trainable source is included")
|
| 441 |
+
freeze(a.out, a.train, a.calibration, a.development, a.test, a.seed, holdout=holdout,
|
| 442 |
+
sources={k: ALL_SOURCES[k] for k in names}, repos={k: ALL_REPOS[k] for k in names}, exclude_states_from=a.exclude_states_from,
|
| 443 |
+
contrastive_pairs=a.contrastive_pairs, contrastive_holdout_families=cfam)
|
| 444 |
+
else:
|
| 445 |
+
if min(a.train, a.calibration) < 1: ap.error("split sizes must be positive")
|
| 446 |
+
freeze(a.out, a.train, a.calibration, a.development, a.test, a.seed, exclude_states_from=a.exclude_states_from)
|
| 447 |
+
|
| 448 |
+
|
| 449 |
+
if __name__ == "__main__":
|
| 450 |
+
main()
|