neuropress 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (184) hide show
  1. neuropress-0.1.0/.editorconfig +30 -0
  2. neuropress-0.1.0/.github/workflows/ci.yml +67 -0
  3. neuropress-0.1.0/.github/workflows/publish.yml +183 -0
  4. neuropress-0.1.0/.gitignore +100 -0
  5. neuropress-0.1.0/.pre-commit-config.yaml +16 -0
  6. neuropress-0.1.0/.prospector.yaml +43 -0
  7. neuropress-0.1.0/.python-version +1 -0
  8. neuropress-0.1.0/ARCHITECTURE.md +110 -0
  9. neuropress-0.1.0/Dockerfile.vast +82 -0
  10. neuropress-0.1.0/LICENSE +201 -0
  11. neuropress-0.1.0/PKG-INFO +473 -0
  12. neuropress-0.1.0/README.md +420 -0
  13. neuropress-0.1.0/benchmarks/bench_tiered_vq.py +501 -0
  14. neuropress-0.1.0/benchmarks/bench_vptq_distribution.py +207 -0
  15. neuropress-0.1.0/benchmarks/test_bench_packing.py +187 -0
  16. neuropress-0.1.0/configs/lowering/auto.yaml +9 -0
  17. neuropress-0.1.0/configs/lowering/compile.yaml +10 -0
  18. neuropress-0.1.0/configs/lowering/cuda_jit.yaml +11 -0
  19. neuropress-0.1.0/configs/lowering/cuda_lut.yaml +10 -0
  20. neuropress-0.1.0/configs/lowering/cuda_wmma.yaml +10 -0
  21. neuropress-0.1.0/neuropress/__init__.py +127 -0
  22. neuropress-0.1.0/neuropress/adapter.py +125 -0
  23. neuropress-0.1.0/neuropress/calibration/__init__.py +29 -0
  24. neuropress-0.1.0/neuropress/calibration/data/__init__.py +26 -0
  25. neuropress-0.1.0/neuropress/calibration/data/audio.py +147 -0
  26. neuropress-0.1.0/neuropress/calibration/data/llm.py +415 -0
  27. neuropress-0.1.0/neuropress/calibration/data/registry.py +22 -0
  28. neuropress-0.1.0/neuropress/calibration/probe/__init__.py +36 -0
  29. neuropress-0.1.0/neuropress/calibration/probe/channel_max.py +37 -0
  30. neuropress-0.1.0/neuropress/calibration/probe/collectors.py +71 -0
  31. neuropress-0.1.0/neuropress/calibration/probe/hessian.py +216 -0
  32. neuropress-0.1.0/neuropress/calibration/probe/hessian_diag.py +143 -0
  33. neuropress-0.1.0/neuropress/calibration/probe/min_max.py +27 -0
  34. neuropress-0.1.0/neuropress/calibration/probe/registry.py +225 -0
  35. neuropress-0.1.0/neuropress/cloud/__init__.py +1 -0
  36. neuropress-0.1.0/neuropress/cloud/cli.py +38 -0
  37. neuropress-0.1.0/neuropress/cloud/direct_runner.py +123 -0
  38. neuropress-0.1.0/neuropress/cloud/runner.py +1573 -0
  39. neuropress-0.1.0/neuropress/cloud/tui.py +511 -0
  40. neuropress-0.1.0/neuropress/cloud/vast_runner.py +2936 -0
  41. neuropress-0.1.0/neuropress/estimator.py +103 -0
  42. neuropress-0.1.0/neuropress/estimator_configs.py +629 -0
  43. neuropress-0.1.0/neuropress/estimators.py +3308 -0
  44. neuropress-0.1.0/neuropress/eval/__init__.py +26 -0
  45. neuropress-0.1.0/neuropress/eval/harness.py +346 -0
  46. neuropress-0.1.0/neuropress/eval/llm_cli.py +1102 -0
  47. neuropress-0.1.0/neuropress/eval/size_estimate.py +600 -0
  48. neuropress-0.1.0/neuropress/eval/timing.py +121 -0
  49. neuropress-0.1.0/neuropress/graph/__init__.py +7 -0
  50. neuropress-0.1.0/neuropress/graph/builder.py +93 -0
  51. neuropress-0.1.0/neuropress/graph/graph.py +43 -0
  52. neuropress-0.1.0/neuropress/graph/node.py +28 -0
  53. neuropress-0.1.0/neuropress/kernel_lowering.py +357 -0
  54. neuropress-0.1.0/neuropress/kernels/__init__.py +1006 -0
  55. neuropress-0.1.0/neuropress/kernels/kmeans_assign.cu +468 -0
  56. neuropress-0.1.0/neuropress/kernels/q4_linear.cu +274 -0
  57. neuropress-0.1.0/neuropress/kernels/rms_norm.cu +123 -0
  58. neuropress-0.1.0/neuropress/kernels/vq_linear.cu +489 -0
  59. neuropress-0.1.0/neuropress/kernels/vq_linear_lut.cu +211 -0
  60. neuropress-0.1.0/neuropress/kernels/vq_linear_ntier.cu +699 -0
  61. neuropress-0.1.0/neuropress/kernels/vq_linear_tiered.cu +632 -0
  62. neuropress-0.1.0/neuropress/kernels/vq_linear_wmma.cu +557 -0
  63. neuropress-0.1.0/neuropress/kernels/vq_linear_zero_escape.cu +559 -0
  64. neuropress-0.1.0/neuropress/meta.py +93 -0
  65. neuropress-0.1.0/neuropress/methods/__init__.py +67 -0
  66. neuropress-0.1.0/neuropress/methods/base.py +147 -0
  67. neuropress-0.1.0/neuropress/methods/correction_mode.py +15 -0
  68. neuropress-0.1.0/neuropress/methods/gptq.py +558 -0
  69. neuropress-0.1.0/neuropress/methods/gptq_panel.py +497 -0
  70. neuropress-0.1.0/neuropress/methods/hqq.py +235 -0
  71. neuropress-0.1.0/neuropress/methods/smooth_quant.py +250 -0
  72. neuropress-0.1.0/neuropress/methods/vptq.py +348 -0
  73. neuropress-0.1.0/neuropress/methods/vptq_analysis.py +926 -0
  74. neuropress-0.1.0/neuropress/methods/vq_gptq.py +258 -0
  75. neuropress-0.1.0/neuropress/methods/zp_scale_calib.py +240 -0
  76. neuropress-0.1.0/neuropress/nn/__init__.py +54 -0
  77. neuropress-0.1.0/neuropress/nn/backends.py +159 -0
  78. neuropress-0.1.0/neuropress/nn/np_embedding.py +599 -0
  79. neuropress-0.1.0/neuropress/nn/np_linear.py +851 -0
  80. neuropress-0.1.0/neuropress/nn/np_rms_norm.py +180 -0
  81. neuropress-0.1.0/neuropress/nn/np_vq_embedding.py +289 -0
  82. neuropress-0.1.0/neuropress/nn/np_vq_linear.py +1423 -0
  83. neuropress-0.1.0/neuropress/pipeline.py +142 -0
  84. neuropress-0.1.0/neuropress/plan.py +32 -0
  85. neuropress-0.1.0/neuropress/preprocessing.py +279 -0
  86. neuropress-0.1.0/neuropress/qscheme/__init__.py +15 -0
  87. neuropress-0.1.0/neuropress/qscheme/descriptor.py +77 -0
  88. neuropress-0.1.0/neuropress/qscheme/group.py +44 -0
  89. neuropress-0.1.0/neuropress/qscheme/zp_scale.py +144 -0
  90. neuropress-0.1.0/neuropress/quantization/__init__.py +27 -0
  91. neuropress-0.1.0/neuropress/quantization/cache.py +211 -0
  92. neuropress-0.1.0/neuropress/quantization/spans.py +86 -0
  93. neuropress-0.1.0/neuropress/recipe.py +58 -0
  94. neuropress-0.1.0/neuropress/recipe_config.py +488 -0
  95. neuropress-0.1.0/neuropress/t2n/__init__.py +59 -0
  96. neuropress-0.1.0/neuropress/t2n/_qtensor_patch.py +103 -0
  97. neuropress-0.1.0/neuropress/t2n/model_io.py +563 -0
  98. neuropress-0.1.0/neuropress/t2n/packing.py +224 -0
  99. neuropress-0.1.0/neuropress/t2n/qtensors.py +526 -0
  100. neuropress-0.1.0/neuropress/t2n/registry.py +952 -0
  101. neuropress-0.1.0/neuropress/targets.py +28 -0
  102. neuropress-0.1.0/neuropress/tools/__init__.py +1 -0
  103. neuropress-0.1.0/neuropress/tools/bench.py +1649 -0
  104. neuropress-0.1.0/neuropress/tools/vq_probe.py +115 -0
  105. neuropress-0.1.0/neuropress/transforms/__init__.py +34 -0
  106. neuropress-0.1.0/neuropress/transforms/cluster_backends.py +1142 -0
  107. neuropress-0.1.0/neuropress/transforms/rotation.py +452 -0
  108. neuropress-0.1.0/neuropress/transforms/rotation_surgery.py +154 -0
  109. neuropress-0.1.0/neuropress/transforms/split_mixed_linear.py +306 -0
  110. neuropress-0.1.0/neuropress/transforms/vq.py +529 -0
  111. neuropress-0.1.0/neuropress/transforms/vq_packing.py +1384 -0
  112. neuropress-0.1.0/neuropress/transforms/weight_only.py +593 -0
  113. neuropress-0.1.0/pyproject.toml +150 -0
  114. neuropress-0.1.0/run +72 -0
  115. neuropress-0.1.0/scripts/bench_assign.py +222 -0
  116. neuropress-0.1.0/scripts/bench_kmeans.py +317 -0
  117. neuropress-0.1.0/scripts/bench_mstep.py +169 -0
  118. neuropress-0.1.0/scripts/benchmark_inference.py +467 -0
  119. neuropress-0.1.0/scripts/plot_bench.py +324 -0
  120. neuropress-0.1.0/scripts/profile_layers.py +202 -0
  121. neuropress-0.1.0/scripts/profile_mstep.py +262 -0
  122. neuropress-0.1.0/scripts/test_save.py +233 -0
  123. neuropress-0.1.0/scripts/vast_runner.py +9 -0
  124. neuropress-0.1.0/tests/__init__.py +0 -0
  125. neuropress-0.1.0/tests/calibration/__init__.py +0 -0
  126. neuropress-0.1.0/tests/calibration/test_audio_data.py +334 -0
  127. neuropress-0.1.0/tests/calibration/test_hessian.py +140 -0
  128. neuropress-0.1.0/tests/calibration/test_llm_data.py +397 -0
  129. neuropress-0.1.0/tests/calibration/test_stat_field.py +143 -0
  130. neuropress-0.1.0/tests/integration/test_gptq_nnef_export.py +147 -0
  131. neuropress-0.1.0/tests/kernels/__init__.py +0 -0
  132. neuropress-0.1.0/tests/kernels/test_kmeans_assign_grouped.py +148 -0
  133. neuropress-0.1.0/tests/methods/__init__.py +0 -0
  134. neuropress-0.1.0/tests/methods/test_down_proj_smoothing.py +157 -0
  135. neuropress-0.1.0/tests/methods/test_ecvq_annealing.py +348 -0
  136. neuropress-0.1.0/tests/methods/test_global_correction_modes.py +337 -0
  137. neuropress-0.1.0/tests/methods/test_gptq.py +152 -0
  138. neuropress-0.1.0/tests/methods/test_gptq_panel.py +265 -0
  139. neuropress-0.1.0/tests/methods/test_gptq_pipeline.py +158 -0
  140. neuropress-0.1.0/tests/methods/test_hqq.py +498 -0
  141. neuropress-0.1.0/tests/methods/test_smooth_quant.py +300 -0
  142. neuropress-0.1.0/tests/methods/test_vptq_analysis.py +192 -0
  143. neuropress-0.1.0/tests/methods/test_zp_scale.py +212 -0
  144. neuropress-0.1.0/tests/nn/test_np_embedding.py +149 -0
  145. neuropress-0.1.0/tests/nn/test_np_rms_norm.py +240 -0
  146. neuropress-0.1.0/tests/nn/test_np_vq_linear.py +414 -0
  147. neuropress-0.1.0/tests/t2n/__init__.py +0 -0
  148. neuropress-0.1.0/tests/t2n/test_model_io.py +352 -0
  149. neuropress-0.1.0/tests/t2n/test_model_size.py +191 -0
  150. neuropress-0.1.0/tests/t2n/test_qtensors.py +211 -0
  151. neuropress-0.1.0/tests/t2n/test_vq_serde.py +171 -0
  152. neuropress-0.1.0/tests/test_adapter.py +50 -0
  153. neuropress-0.1.0/tests/test_cache_geometry_validation.py +95 -0
  154. neuropress-0.1.0/tests/test_calibration.py +99 -0
  155. neuropress-0.1.0/tests/test_cloud_cli_entrypoint.py +77 -0
  156. neuropress-0.1.0/tests/test_cloud_runner.py +160 -0
  157. neuropress-0.1.0/tests/test_cuda_kmeans.py +112 -0
  158. neuropress-0.1.0/tests/test_eval.py +243 -0
  159. neuropress-0.1.0/tests/test_graph.py +131 -0
  160. neuropress-0.1.0/tests/test_per_group_centroids.py +37 -0
  161. neuropress-0.1.0/tests/test_pipeline.py +95 -0
  162. neuropress-0.1.0/tests/test_plan.py +80 -0
  163. neuropress-0.1.0/tests/test_preprocessing.py +289 -0
  164. neuropress-0.1.0/tests/test_qscheme.py +140 -0
  165. neuropress-0.1.0/tests/test_readme_references.py +48 -0
  166. neuropress-0.1.0/tests/test_recipe_config.py +382 -0
  167. neuropress-0.1.0/tests/test_size_estimate_cli.py +73 -0
  168. neuropress-0.1.0/tests/test_t2n_register.py +335 -0
  169. neuropress-0.1.0/tests/test_targets.py +40 -0
  170. neuropress-0.1.0/tests/test_transforms.py +330 -0
  171. neuropress-0.1.0/tests/test_vptq_estimator.py +424 -0
  172. neuropress-0.1.0/tests/test_vptq_freq_sort.py +144 -0
  173. neuropress-0.1.0/tests/test_weight_only_transforms.py +120 -0
  174. neuropress-0.1.0/tests/transforms/__init__.py +0 -0
  175. neuropress-0.1.0/tests/transforms/test_cluster_backend_selection.py +42 -0
  176. neuropress-0.1.0/tests/transforms/test_cluster_backends.py +333 -0
  177. neuropress-0.1.0/tests/transforms/test_ntier_packing.py +303 -0
  178. neuropress-0.1.0/tests/transforms/test_rotation.py +332 -0
  179. neuropress-0.1.0/tests/transforms/test_rotation_cuda.py +104 -0
  180. neuropress-0.1.0/tests/transforms/test_rotation_surgery.py +230 -0
  181. neuropress-0.1.0/tests/transforms/test_tiered_packing.py +296 -0
  182. neuropress-0.1.0/tests/transforms/test_vq_packing.py +120 -0
  183. neuropress-0.1.0/tests/transforms/test_vq_packing_extended.py +207 -0
  184. neuropress-0.1.0/tests/transforms/test_zero_escape_packing.py +370 -0
@@ -0,0 +1,30 @@
1
+ root = true
2
+
3
+ [*]
4
+ charset = utf-8
5
+ end_of_line = lf
6
+ insert_final_newline = true
7
+ trim_trailing_whitespace = true
8
+
9
+ [*.py]
10
+ indent_style = space
11
+ indent_size = 4
12
+ max_line_length = 88
13
+
14
+ [*.{toml,cfg,ini}]
15
+ indent_style = space
16
+ indent_size = 4
17
+
18
+ [*.{yml,yaml}]
19
+ indent_style = space
20
+ indent_size = 2
21
+
22
+ [*.{json,jsonc}]
23
+ indent_style = space
24
+ indent_size = 2
25
+
26
+ [*.md]
27
+ trim_trailing_whitespace = false
28
+
29
+ [Makefile]
30
+ indent_style = tab
@@ -0,0 +1,67 @@
1
+ name: ci
2
+
3
+ # Lint and test on every push to main and every pull request.
4
+ #
5
+ # The release workflow (publish.yml) verifies that the *artifact* installs and
6
+ # imports; this one verifies that the *code* is correct. Both gate a release:
7
+ # publish.yml runs on a tag, by which point these checks have already run on the
8
+ # commit being tagged.
9
+
10
+ on:
11
+ push:
12
+ branches: [main]
13
+ pull_request:
14
+
15
+ permissions:
16
+ contents: read
17
+
18
+ concurrency:
19
+ group: ci-${{ github.ref }}
20
+ cancel-in-progress: true
21
+
22
+ jobs:
23
+ lint:
24
+ name: ruff
25
+ runs-on: ubuntu-latest
26
+ steps:
27
+ - uses: actions/checkout@v5
28
+
29
+ - uses: astral-sh/setup-uv@v6
30
+ with:
31
+ python-version: "3.13"
32
+
33
+ # Mirrors the blocking half of `./run check`. Pinned so a new ruff release
34
+ # cannot fail CI on a commit that did not change; bump it deliberately.
35
+ - name: Lint
36
+ run: uvx ruff@0.15.2 check .
37
+
38
+ - name: Format
39
+ run: uvx ruff@0.15.2 format --check .
40
+
41
+ test:
42
+ name: pytest
43
+ runs-on: ubuntu-latest
44
+ steps:
45
+ - uses: actions/checkout@v5
46
+
47
+ - uses: astral-sh/setup-uv@v6
48
+ with:
49
+ python-version: "3.13"
50
+
51
+ - name: Install
52
+ run: |
53
+ set -euo pipefail
54
+ uv venv --python 3.13
55
+ # CPU torch first: the default PyPI wheel bundles CUDA (~2.5 GB) and
56
+ # fills the runner disk. Installing it up front leaves the project
57
+ # install below with the requirement already satisfied.
58
+ uv pip install --python .venv torch --index-url https://download.pytorch.org/whl/cpu
59
+ # `cloud` is needed for collection, not just for coverage:
60
+ # tests/test_cloud_runner.py imports neuropress.cloud.vast_runner,
61
+ # which imports rich at module scope.
62
+ uv pip install --python .venv -e ".[dev,cloud]"
63
+
64
+ - name: Test
65
+ # Integration tests need HuggingFace downloads and several minutes; they
66
+ # are additionally gated behind RUN_INTEGRATION inside the module.
67
+ run: .venv/bin/python -m pytest --ignore=tests/integration -q
@@ -0,0 +1,183 @@
1
+ name: publish
2
+
3
+ # Build neuropress and upload it to PyPI.
4
+ #
5
+ # release (published) -> PyPI (the tag must match the package version)
6
+ # workflow_dispatch -> TestPyPI (rehearsal; no tag check)
7
+ #
8
+ # Authentication is PyPI Trusted Publishing (OIDC). No API token is stored in
9
+ # this repository, so there is nothing to rotate or leak. One-time setup, on
10
+ # both pypi.org and test.pypi.org, before the first upload:
11
+ #
12
+ # Your projects -> Publishing -> Add a pending publisher
13
+ # PyPI project name : neuropress
14
+ # Owner : DreamerMind
15
+ # Repository : neuropress
16
+ # Workflow name : publish.yml
17
+ # Environment : pypi (use "testpypi" on test.pypi.org)
18
+ #
19
+ # Cutting a release:
20
+ # 1. bump `version` in pyproject.toml and commit it
21
+ # 2. gh release create v<version> --generate-notes
22
+ #
23
+ # PyPI is append-only: a version number can never be reused, not even after the
24
+ # file is deleted. Rehearse with "Run workflow" against TestPyPI first.
25
+
26
+ on:
27
+ release:
28
+ types: [published]
29
+ workflow_dispatch:
30
+
31
+ permissions:
32
+ contents: read
33
+
34
+ concurrency:
35
+ group: publish-${{ github.ref }}
36
+ cancel-in-progress: false
37
+
38
+ jobs:
39
+ build:
40
+ name: build and verify artifacts
41
+ runs-on: ubuntu-latest
42
+ steps:
43
+ - uses: actions/checkout@v5
44
+
45
+ - uses: astral-sh/setup-uv@v6
46
+ with:
47
+ python-version: "3.13"
48
+
49
+ - name: Build sdist and wheel
50
+ run: uv build
51
+
52
+ - name: Validate package metadata
53
+ run: uvx twine check --strict dist/*
54
+
55
+ - name: Tag must match the built version
56
+ if: github.event_name == 'release'
57
+ run: |
58
+ # Compare as versions, not as strings: a wheel filename is always
59
+ # PEP 440-normalized, so `version = "1.0.0-rc1"` builds as 1.0.0rc1
60
+ # and would never string-match the natural tag v1.0.0-rc1.
61
+ set -euo pipefail
62
+ built=$(basename "$(ls dist/*.whl)" | cut -d- -f2)
63
+ uv run --no-project --with packaging python - "$built" "${GITHUB_REF_NAME#v}" <<'PY'
64
+ import sys
65
+
66
+ from packaging.version import InvalidVersion, Version
67
+
68
+ built, tag = sys.argv[1], sys.argv[2]
69
+ try:
70
+ match = Version(built) == Version(tag)
71
+ except InvalidVersion:
72
+ match = False
73
+ if not match:
74
+ print(
75
+ f"::error::release tag {tag!r} and the built version {built!r} disagree — "
76
+ "correct whichever is wrong (the tag, or `version` in pyproject.toml)"
77
+ )
78
+ raise SystemExit(1)
79
+ print(f"publishing neuropress {built}")
80
+ PY
81
+
82
+ - name: Artifacts must carry no research-only material
83
+ run: |
84
+ # Defence in depth. pyproject.toml already excludes these from both
85
+ # build targets, so a hit means an exclude pattern regressed. Uploads
86
+ # are permanent, so fail closed rather than warn.
87
+ #
88
+ # Build the listing before grepping it: `unzip -l ... | grep` inside an
89
+ # `if` swallows a failing unzip, and a check that cannot read the
90
+ # artifact must not report it clean.
91
+ #
92
+ # `unzip -Z1` lists bare paths. `unzip -l` pads each name into a
93
+ # right-hand column, which leaves a top-level `reports/` preceded by
94
+ # spaces rather than by start-of-line, so the anchor below could never
95
+ # match it and the wheel half of this check would be dead code.
96
+ set -euo pipefail
97
+ wheel=$(ls dist/*.whl)
98
+ sdist=$(ls dist/*.tar.gz)
99
+ listing=$(unzip -Z1 "$wheel"; tar -tzf "$sdist")
100
+ # Scoped to scripts/: a bare `_diag\.py` also matches shipped modules
101
+ # such as neuropress/calibration/probe/hessian_diag.py.
102
+ pattern='research_registry|research_jobs|sync_public|public-manifest'
103
+ pattern="$pattern"'|(^|/)scripts/(qwen35|gemma4)_'
104
+ pattern="$pattern"'|(^|/)(reports|plans|docs|papers|sweeps)/'
105
+ if grep -Ei "$pattern" <<<"$listing"; then
106
+ echo "::error::research-only paths found in the distribution artifacts"
107
+ exit 1
108
+ fi
109
+ echo "artifacts clean: $wheel $sdist"
110
+
111
+ - uses: actions/upload-artifact@v4
112
+ with:
113
+ name: dist
114
+ path: dist/
115
+
116
+ smoke:
117
+ name: built wheel installs and imports
118
+ needs: build
119
+ runs-on: ubuntu-latest
120
+ steps:
121
+ - uses: actions/download-artifact@v4
122
+ with:
123
+ name: dist
124
+ path: dist
125
+
126
+ - uses: astral-sh/setup-uv@v6
127
+ with:
128
+ python-version: "3.13"
129
+
130
+ - name: Install the wheel with base dependencies only
131
+ run: |
132
+ uv venv --python 3.13 .venv
133
+ # torch's default PyPI wheel bundles CUDA (~2.5 GB) and fills the
134
+ # runner disk. The CPU build is all an import check needs.
135
+ uv pip install --python .venv torch --index-url https://download.pytorch.org/whl/cpu
136
+ uv pip install --python .venv dist/*.whl
137
+
138
+ - name: Import from outside the workspace
139
+ run: |
140
+ # Never run this from a directory containing neuropress/: python puts
141
+ # the working directory on sys.path, so the import would resolve to
142
+ # the source tree and prove nothing about the wheel.
143
+ mkdir -p "$RUNNER_TEMP/smoke"
144
+ cd "$RUNNER_TEMP/smoke"
145
+ "$GITHUB_WORKSPACE/.venv/bin/python" - <<'PY'
146
+ from importlib.metadata import version
147
+
148
+ import importlib.util
149
+
150
+ import neuropress
151
+ from neuropress.t2n import NEUROPRESS_COMPRESSION as registry
152
+
153
+ print(f"neuropress {version('neuropress')} from {neuropress.__file__}")
154
+ assert registry, "compression registry is empty"
155
+
156
+ # The invariant is that the research module is absent, not that its
157
+ # entries happen to share a key prefix — a prefix nothing enforces.
158
+ spec = importlib.util.find_spec("neuropress.t2n.research_registry")
159
+ assert spec is None, f"research module reached the wheel: {spec.origin}"
160
+ print(f"{len(registry)} compression entries, research module absent")
161
+ PY
162
+
163
+ publish:
164
+ name: upload to ${{ github.event_name == 'release' && 'PyPI' || 'TestPyPI' }}
165
+ needs: [build, smoke]
166
+ runs-on: ubuntu-latest
167
+ environment:
168
+ name: ${{ github.event_name == 'release' && 'pypi' || 'testpypi' }}
169
+ url: ${{ github.event_name == 'release' && 'https://pypi.org/p/neuropress' || 'https://test.pypi.org/p/neuropress' }}
170
+ permissions:
171
+ id-token: write # mints the OIDC token trusted publishing exchanges for an upload token
172
+ steps:
173
+ - uses: actions/download-artifact@v4
174
+ with:
175
+ name: dist
176
+ path: dist
177
+
178
+ - uses: pypa/gh-action-pypi-publish@release/v1
179
+ with:
180
+ repository-url: ${{ github.event_name == 'release' && 'https://upload.pypi.org/legacy/' || 'https://test.pypi.org/legacy/' }}
181
+ # Rehearsals re-upload the same version repeatedly; a real release
182
+ # must still fail loudly if that version is already on PyPI.
183
+ skip-existing: ${{ github.event_name != 'release' }}
@@ -0,0 +1,100 @@
1
+ # === Python ===
2
+ __pycache__/
3
+ *.py[cod]
4
+ *.pyo
5
+ *.pyd
6
+ *.so
7
+ *.egg
8
+ *.egg-info/
9
+ dist/
10
+ build/
11
+ wheels/
12
+ MANIFEST
13
+ pip-wheel-metadata/
14
+ share/python-wheels/
15
+
16
+ # === Virtual environments ===
17
+ .venv/
18
+ venv/
19
+ env/
20
+ .env/
21
+
22
+ # === uv / pip ===
23
+ uv.lock
24
+ pip-selfcheck.json
25
+
26
+ # === Testing & coverage ===
27
+ .pytest_cache/
28
+ .coverage
29
+ .coverage.*
30
+ coverage.xml
31
+ htmlcov/
32
+ .tox/
33
+ .nox/
34
+
35
+ # === Type checkers ===
36
+ .mypy_cache/
37
+ .pyright/
38
+ .dmypy.json
39
+ pyrightconfig.json
40
+
41
+ # === PyTorch artefacts ===
42
+ *.pt
43
+ *.pth
44
+ *.ckpt
45
+ *.bin
46
+ *.safetensors
47
+
48
+ # === ONNX / NNEF artefacts ===
49
+ *.onnx
50
+ *.nnef
51
+ *.nnef.tar
52
+ *.rgg
53
+
54
+ # === Jupyter ===
55
+ .ipynb_checkpoints/
56
+ *.ipynb
57
+
58
+ # === IDEs ===
59
+ .idea/
60
+ .vscode/
61
+ *.sublime-project
62
+ *.sublime-workspace
63
+ *.code-workspace
64
+ *.iml
65
+ .metals/
66
+ .bloop/
67
+
68
+ # === macOS ===
69
+ .DS_Store
70
+ .AppleDouble
71
+ .LSOverride
72
+ ._*
73
+ .Spotlight-V100
74
+ .Trashes
75
+
76
+ # === Linux ===
77
+ *~
78
+ .fuse_hidden*
79
+ .Trash-*
80
+
81
+ # === Windows ===
82
+ Thumbs.db
83
+ ehthumbs.db
84
+ Desktop.ini
85
+
86
+ # === Logs & temp ===
87
+ *.log
88
+ *.tmp
89
+ *.bak
90
+ *.swp
91
+ *.swo
92
+
93
+ # === neuropress output ===
94
+ /data/
95
+ vast_output/
96
+ .ruff_cache/
97
+ .ty_cache/
98
+ papers/
99
+ .cachesize
100
+ .tmp_inst.json
@@ -0,0 +1,16 @@
1
+ repos:
2
+ - repo: https://github.com/astral-sh/ruff-pre-commit
3
+ rev: v0.11.5
4
+ hooks:
5
+ - id: ruff
6
+ args: [--fix, --exit-non-zero-on-fix]
7
+ - id: ruff-format
8
+
9
+ - repo: local
10
+ hooks:
11
+ - id: ty
12
+ name: ty
13
+ entry: uv run ty check
14
+ language: system
15
+ types: [python]
16
+ pass_filenames: false
@@ -0,0 +1,43 @@
1
+ strictness: medium
2
+ doc-warnings: false
3
+ test-warnings: false
4
+
5
+ # ruff already handles pycodestyle and pyflakes
6
+ without-tools:
7
+ - pycodestyle
8
+ - pyflakes
9
+
10
+ ignore-patterns:
11
+ - .venv
12
+
13
+ mccabe:
14
+ options:
15
+ # GPTQ and hessian methods are inherently complex multi-branch algorithms.
16
+ max-complexity: 20
17
+
18
+ pylint:
19
+ options:
20
+ # Scientific / ML code naturally has many parameters and local variables.
21
+ max-args: 10
22
+ max-locals: 28
23
+ max-branches: 20
24
+ # Allow extension-pkg access (torch._C, torch.cuda, etc.)
25
+ extension-pkg-allow-list:
26
+ - torch
27
+
28
+ disable:
29
+ # Lazy / conditional imports are intentional in this codebase
30
+ # (OffloadedTensor, datasets, GranularityType, os, …).
31
+ - import-outside-toplevel
32
+ # Hook callbacks must match PyTorch signatures; unused args are structural.
33
+ - unused-argument
34
+ # Accessing _clone_weights / _LinAlgError on peer classes is intentional.
35
+ - protected-access
36
+ # torch.linalg.cholesky is callable; pylint static analysis is wrong.
37
+ - not-callable
38
+ # torch._C members are C-extension; unavailable at static analysis time.
39
+ - c-extension-no-member
40
+ # Project does not use Django.
41
+ - django-not-available
42
+ # Scientific APIs naturally have more than 5 positional parameters.
43
+ - too-many-positional-arguments
@@ -0,0 +1 @@
1
+ 3.13
@@ -0,0 +1,110 @@
1
+ # neuropress architecture
2
+
3
+ ## Commands
4
+
5
+ ```bash
6
+ # Install dependencies (Python 3.13, managed with uv)
7
+ uv sync
8
+
9
+ # Install with optional LLM data loader (HuggingFace datasets)
10
+ uv sync --extra llm
11
+
12
+ # Run all unit tests
13
+ uv run pytest
14
+
15
+ # Run a single test file
16
+ uv run pytest tests/test_pipeline.py -v
17
+
18
+ # Run a single test by name
19
+ uv run pytest tests/test_pipeline.py::TestQuantizationPipeline::test_pipeline_runs -v
20
+
21
+ # Run integration tests (downloads ~500 MB model, requires RUN_INTEGRATION=1)
22
+ RUN_INTEGRATION=1 uv run pytest tests/integration/ -v
23
+ ```
24
+
25
+ ## Architecture
26
+
27
+ `neuropress` is a graph-aware neural network quantization framework targeting NNEF export via `torch_to_nnef`. The full pipeline proceeds in six steps (see `pipeline.py`):
28
+
29
+ 1. **QGraph** – `RuntimeModuleGraphProvider` traces a single forward pass using PyTorch module hooks to build a `QGraph` (`nodes: Dict[str, QNode]`, `values: Dict[str, QValue]`). No symbolic tracing (FX/Dynamo) is used.
30
+
31
+ 2. **Targets** – `WeightTarget` (`module_id/param_name`) and `ActivationTarget` (`value_id`) identify what to quantize.
32
+
33
+ 3. **StatisticsRegistry** – `QuantizationParamEstimator.declare_statistics()` returns `StatisticRequest` objects. The registry deduplicates them (key = `target_key::stat_type`), installs hooks, runs a forward pass over calibration data, and stores results. Supported stat types: `"min_max"` and `"hessian"`.
34
+
35
+ 4. **QuantizationRecipe** – ties together `allowed_weight_qschemes`, `allowed_activation_qschemes`, a `QuantizationParamEstimator`, and a `ModuleQuantizationAdapter`.
36
+
37
+ 5. **QuantizationPlan** – the output of estimation: dicts mapping target keys to `ZpScaleScheme` or `QSchemeGroup` instances.
38
+
39
+ 6. **ModuleQuantizationAdapter** – applies the plan to the model (abstract; user-supplied).
40
+
41
+ ### Key abstractions
42
+
43
+ | Class | Role |
44
+ |---|---|
45
+ | `QSchemeDescriptor` | Static backend constraint (nbits, symmetric, granularity) |
46
+ | `ZpScaleScheme` | Concrete affine scheme (zero_point, scale, nbits); extends `torch_to_nnef.QScheme` |
47
+ | `QSchemeGroup` / `PanelSpec` / `TensorSlice` | Mixed-precision grouping of weight regions |
48
+ | `CalibrationMethod` | ABC for algorithms; subclasses implement `update_statistics` + `compute_qscheme` |
49
+ | `CalibrationMethodOnActivation` | Subclass of above for Hessian-based methods (GPTQ) |
50
+ | `GPTQCalibration` | Full GPTQ implementation (Frantar et al. ICLR 2023) |
51
+ | `GPTQMixedPrecisionCalibration` | GPTQ with per-column-group bit-width allocation via `MixedAllocationStrategy` |
52
+ | `MixedPrecisionLinear` | Drop-in replacement for `nn.Linear` produced by `split_mixed_precision_linear` |
53
+
54
+ ### Calibration data shapes (4-D convention)
55
+
56
+ Weights are reshaped to `[rows, 1, n_groups, group_size]` before quantization. Odd dims (1, group_size) are "reduction" dims; even dims (rows, n_groups) are "kept". `qscheme_data_shape` is the same rank with reduction dims set to 1.
57
+
58
+ ### `StatField` annotation pattern
59
+
60
+ `CalibrationMethod` subclasses declare statistic dependencies via `Annotated` type hints. The registry auto-discovers and injects them — no manual wiring needed.
61
+
62
+ ```python
63
+ hessian: Annotated[Optional[torch.Tensor], StatField("hessian", "hessian")] = None
64
+ ```
65
+
66
+ `StatField(stat_type, field_key)`:
67
+ - `stat_type` — matches the string key in `StatisticRequest` (`"hessian"`, `"hessian_diag"`, `"min_max"`).
68
+ - `field_key` — the attribute name to populate on the method instance.
69
+
70
+ Flow: `StatisticsRegistry.register_method(method, target_key)` scans `Annotated` hints to build `StatisticRequest` objects. After collection, `registry.inject(method, target_key)` sets the annotated fields from stored results. The pattern lives across three files: `calibration/probe/collectors.py` (StatField), `calibration/probe/registry.py` (discover + inject), and each `CalibrationMethod` subclass.
71
+
72
+ ### Estimator config / implementation split
73
+
74
+ `VptqGroupEstimatorConfig` (`estimator_configs.py`) is the user-facing Pydantic model; `VPTQGroupEstimator` (`estimators.py`) is the runtime dataclass. The config's `build()` passes every field explicitly, so **`VPTQGroupEstimator` is the single source of truth for defaults** — `VptqGroupEstimatorConfig` derives its defaults from the dataclass fields via `_vd()` at class definition time. If you add a field to `VPTQGroupEstimator`, add a corresponding `_vd("field_name")` default in the config.
75
+
76
+ ### Pipeline wiring: `_cache` and `_spans`
77
+
78
+ `QuantizationParamEstimator` declares `_cache: CacheManager` and `_spans: SpanCollector` as class attributes with no-op defaults. The pipeline calls `set_cache()` and `set_span_collector()` before `prepare()`. Subclasses use `self._cache` / `self._spans` directly — no `getattr` fallback needed.
79
+
80
+ ### Compression registry (`neuropress/t2n/registry.py`)
81
+
82
+ Follows the `torch_to_nnef` registry pattern: a dict (`NEUROPRESS_COMPRESSION`) maps string keys to callables with signature `fn(model, tokenizer=..., **kwargs) -> model`. Pass it to `LLMExporter.prepare`:
83
+
84
+ ```python
85
+ exporter.prepare(
86
+ compression_method="open_text_c4_gptq_q4_0",
87
+ compression_registry="neuropress.t2n.registry.NEUROPRESS_COMPRESSION",
88
+ )
89
+ ```
90
+
91
+ Every entry is built by `_make_entry(config)`, which binds a `QuantizationRecipeConfig` into a `functools.partial` of `run_quantization_pipeline` and stashes the config JSON on the partial as `__config__` so the manifest can record it. Calling an entry:
92
+
93
+ 1. Reads `QuantizationRecipe.calibration` and loads batches from `calibration.DATASET_REGISTRY`
94
+ 2. Transfers batches to the model device one at a time (no OOM on large calibration sets)
95
+ 3. Delegates to `QuantizationPipeline`, which runs the six steps above
96
+ 4. Applies kernel lowering when the recipe configures one
97
+
98
+ Recipes can also be built from JSON via `QuantizationRecipeConfig.from_json(...).build()`, which resolves component types from the registries at deserialization time.
99
+
100
+ Layers whose `in_features` is not a multiple of the group size are skipped (Q4_0 group size = 32).
101
+
102
+ Registry entries: `"open_text_c4_gptq_q4_0"` (128 samples), `"open_text_c4_gptq_q4_0_32samples"`, `"open_text_c4_gptq_q4_0_256samples"`.
103
+
104
+ ### LLM calibration data
105
+
106
+ `neuropress.calibration.data.llm.load_c4_calibration(tokenizer, ...)` loads C4 samples as `{"input_ids": Tensor[batch_size, seq_len]}` dicts. Requires the `llm` optional dependency (`datasets`).
107
+
108
+ ### Integration test pattern
109
+
110
+ The integration test (`tests/integration/test_gptq_nnef_export.py`) demonstrates the end-to-end workflow: HuggingFace model load → C4 calibration → layer-by-layer GPTQ → `torch_to_nnef` NNEF export. It is gated behind `RUN_INTEGRATION=1` and also requires `transformers` and `datasets`.
@@ -0,0 +1,82 @@
1
+ # Dockerfile.vast — pre-baked base image for neuropress vast.ai GPU jobs.
2
+ #
3
+ # Cuts cold-start time from ~15 min to ~3 min by pre-warming the uv package
4
+ # cache inside the image layer. On a fresh instance, uv sync only needs to
5
+ # install the neuropress source itself (a few seconds), not re-download torch
6
+ # or any other dependency.
7
+ #
8
+ # Build and push once:
9
+ #
10
+ # DOCKER_BUILDKIT=1 docker build -f Dockerfile.vast -t <user>/neuropress-base:latest .
11
+ # docker push <user>/neuropress-base:latest
12
+ #
13
+ # Subsequent rebuilds after a pyproject.toml change are fast (~30 s) because
14
+ # BuildKit caches the downloaded packages on the build host. The first build
15
+ # still downloads everything (~9 min).
16
+ #
17
+ # Then tell neuropress-vast to use it:
18
+ #
19
+ # export NEUROPRESS_DOCKER_IMAGE=<user>/neuropress-base:latest
20
+ # neuropress-vast --job neuropress-eval-compare ...
21
+ #
22
+ # The image must be public (vast.ai pulls without auth by default).
23
+ # Re-build any time pyproject.toml dependencies change.
24
+
25
+ # Slim CUDA 12.4 runtime — ~800 MB compressed vs ~5 GB for pytorch/pytorch.
26
+ # PyTorch is installed via uv from PyPI below, so we don't need it in the base.
27
+ FROM nvidia/cuda:13.0.2-cudnn-runtime-ubuntu24.04
28
+
29
+
30
+ ENV DEBIAN_FRONTEND=noninteractive
31
+
32
+ # System tools needed by setup commands:
33
+ # tmux — disconnect-resilient job executor
34
+ # build-essential — g++/make for torch-to-nnef C++ extension
35
+ # ninja-build — required by torch.utils.cpp_extension for CUDA JIT kernels
36
+ # cuda-nvcc — nvcc compiler for CUDA JIT (runtime images lack it)
37
+ # curl — uv installer
38
+ # git — some HuggingFace libs call git
39
+ #
40
+ # Cache mounts keep downloaded .deb files on the build host so subsequent
41
+ # builds skip re-downloading them. The mounts are not part of the image
42
+ # layer, so no rm -rf /var/lib/apt/lists/* is needed.
43
+ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
44
+ --mount=type=cache,target=/var/lib/apt,sharing=locked \
45
+ apt-get update -qq && \
46
+ apt-get install -y --no-install-recommends \
47
+ tmux build-essential ninja-build cuda-nvcc curl ca-certificates git
48
+
49
+ # Install uv
50
+ RUN curl -LsSf https://astral.sh/uv/install.sh | sh
51
+ ENV PATH="/root/.local/bin:$PATH"
52
+ ENV UV_CACHE_DIR=/root/.cache/uv
53
+
54
+ # Pre-install Python 3.13.
55
+ # Cache mount keeps the CPython tarball on the build host; the installed
56
+ # binary goes to ~/.local/share/uv/python (baked into the image layer).
57
+ RUN --mount=type=cache,target=/root/.cache/uv \
58
+ uv python install 3.13
59
+
60
+ # ── Pre-warm the uv package cache ────────────────────────────────────────────
61
+ # Copy only the dependency manifest (not the source) so this layer is rebuilt
62
+ # only when dependencies change, not on every code change.
63
+ WORKDIR /tmp/neuropress-seed
64
+ COPY pyproject.toml ./
65
+ # hatchling needs a minimal package layout to resolve the project
66
+ RUN mkdir -p neuropress && \
67
+ printf '"""neuropress."""\n' > neuropress/__init__.py && \
68
+ touch README.md
69
+
70
+ # Download packages to the BuildKit cache (/tmp/uv-cache, preserved on the
71
+ # build host between runs). On first build this takes ~9 min; on subsequent
72
+ # rebuilds after a dependency change uv resolves from cache in ~30 s.
73
+ # After sync, cp bakes the populated cache into UV_CACHE_DIR (image layer) so
74
+ # that `uv sync` on a fresh vast.ai instance needs no network downloads.
75
+ RUN --mount=type=cache,target=/tmp/uv-cache,sharing=locked \
76
+ UV_CACHE_DIR=/tmp/uv-cache \
77
+ uv sync --extra llm --extra eval --extra dev && \
78
+ rm -rf .venv && \
79
+ mkdir -p "$UV_CACHE_DIR" && \
80
+ cp -a /tmp/uv-cache/. "$UV_CACHE_DIR/"
81
+
82
+ WORKDIR /workspace