torchnative 0.0.1a0__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.
- torchnative-0.0.1a0/LICENSE +21 -0
- torchnative-0.0.1a0/PKG-INFO +293 -0
- torchnative-0.0.1a0/README.md +266 -0
- torchnative-0.0.1a0/pyproject.toml +87 -0
- torchnative-0.0.1a0/setup.cfg +4 -0
- torchnative-0.0.1a0/torchnative/src/main/torchnative/__init__.py +4 -0
- torchnative-0.0.1a0/torchnative/src/main/torchnative/adapt/__init__.py +8 -0
- torchnative-0.0.1a0/torchnative/src/main/torchnative/api/__init__.py +24 -0
- torchnative-0.0.1a0/torchnative/src/main/torchnative/delta/__init__.py +11 -0
- torchnative-0.0.1a0/torchnative/src/main/torchnative/kernels/__init__.py +7 -0
- torchnative-0.0.1a0/torchnative/src/main/torchnative/nn/__init__.py +0 -0
- torchnative-0.0.1a0/torchnative/src/main/torchnative/nn/federated/__init__.py +7 -0
- torchnative-0.0.1a0/torchnative/src/main/torchnative.egg-info/PKG-INFO +293 -0
- torchnative-0.0.1a0/torchnative/src/main/torchnative.egg-info/SOURCES.txt +15 -0
- torchnative-0.0.1a0/torchnative/src/main/torchnative.egg-info/dependency_links.txt +1 -0
- torchnative-0.0.1a0/torchnative/src/main/torchnative.egg-info/requires.txt +5 -0
- torchnative-0.0.1a0/torchnative/src/main/torchnative.egg-info/top_level.txt +1 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 thisisthepy
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,293 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: torchnative
|
|
3
|
+
Version: 0.0.1a0
|
|
4
|
+
Summary: Run the real PyTorch ecosystem on device -- not a reimplementation of it
|
|
5
|
+
Author: thisisthepy
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/thisisthepy/torchnative
|
|
8
|
+
Project-URL: Documentation, https://github.com/thisisthepy/torchnative/tree/develop/docs
|
|
9
|
+
Project-URL: Source, https://github.com/thisisthepy/torchnative
|
|
10
|
+
Keywords: pytorch,on-device,edge,federated-learning,test-time-adaptation
|
|
11
|
+
Classifier: Development Status :: 2 - Pre-Alpha
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
14
|
+
Classifier: Programming Language :: Rust
|
|
15
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
16
|
+
Classifier: Operating System :: Android
|
|
17
|
+
Classifier: Operating System :: iOS
|
|
18
|
+
Classifier: Operating System :: MacOS
|
|
19
|
+
Classifier: Operating System :: POSIX :: Linux
|
|
20
|
+
Requires-Python: >=3.13
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
License-File: LICENSE
|
|
23
|
+
Provides-Extra: federated
|
|
24
|
+
Provides-Extra: test
|
|
25
|
+
Requires-Dist: torch<2.14,>=2.13; extra == "test"
|
|
26
|
+
Dynamic: license-file
|
|
27
|
+
|
|
28
|
+
<div align="center">
|
|
29
|
+
|
|
30
|
+
# torchnative
|
|
31
|
+
|
|
32
|
+
**Run the real PyTorch ecosystem on device — not a reimplementation of it.**
|
|
33
|
+
|
|
34
|
+
[](https://pypi.org/project/torchnative/)
|
|
35
|
+
[](https://www.python.org/)
|
|
36
|
+
[](LICENSE)
|
|
37
|
+
[](#install)
|
|
38
|
+
[](#status)
|
|
39
|
+
|
|
40
|
+
</div>
|
|
41
|
+
|
|
42
|
+
---
|
|
43
|
+
|
|
44
|
+
`torchnative` replaces PyTorch's compiled core — `torch._C` — with a native extension, so the
|
|
45
|
+
genuine `torch` and `transformers` packages run on a phone the way they run on a workstation.
|
|
46
|
+
|
|
47
|
+
Models are not ported, converted, or re-expressed. They are imported.
|
|
48
|
+
|
|
49
|
+
```python
|
|
50
|
+
from transformers import AutoModelForCausalLM # the real one
|
|
51
|
+
model = AutoModelForCausalLM.from_pretrained("...")
|
|
52
|
+
model.generate(...) # on the device
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
> [!WARNING]
|
|
56
|
+
> **Pre-alpha.** The operator layer matches upstream PyTorch numerically and 15 of 20 tested
|
|
57
|
+
> architectures reach zero missing operators — but `import transformers` does not work yet, no
|
|
58
|
+
> checkpoint has ever been loaded, and no device has run the built artefact.
|
|
59
|
+
> See [Status](#status) before depending on this.
|
|
60
|
+
|
|
61
|
+
---
|
|
62
|
+
|
|
63
|
+
## Why not a reimplementation
|
|
64
|
+
|
|
65
|
+
Every other route to on-device inference re-expresses the model somewhere else.
|
|
66
|
+
|
|
67
|
+
| | approach | cost |
|
|
68
|
+
|---|---|---|
|
|
69
|
+
| llama.cpp | architectures rewritten in C++ | each new architecture is a porting task |
|
|
70
|
+
| ExecuTorch · CoreML | ahead-of-time compiled graph | export step, and what runs is not what you wrote |
|
|
71
|
+
| MLC | lowered to its own runtime | same |
|
|
72
|
+
| **torchnative** | **the real Python package** | **the substrate is hard; architectures are free** |
|
|
73
|
+
|
|
74
|
+
The reason nobody runs the real thing is that `torch._C` cannot be built for mobile. PyTorch's own
|
|
75
|
+
build sets `INTERN_BUILD_MOBILE` for any Android or iOS toolchain, and that path forces
|
|
76
|
+
`BUILD_PYTHON` off — so the mobile build is structurally incapable of producing the Python
|
|
77
|
+
extension module the Python package needs.
|
|
78
|
+
|
|
79
|
+
`torchnative` supplies that module instead. Everything above it is upstream source, unmodified.
|
|
80
|
+
|
|
81
|
+
---
|
|
82
|
+
|
|
83
|
+
## What it does
|
|
84
|
+
|
|
85
|
+
### 1 · LLM inference
|
|
86
|
+
|
|
87
|
+
Run `transformers` models directly. No conversion step, no per-architecture port — if
|
|
88
|
+
`transformers` supports it and the operators are covered, it runs.
|
|
89
|
+
|
|
90
|
+
```python
|
|
91
|
+
import torch
|
|
92
|
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
93
|
+
|
|
94
|
+
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
|
|
95
|
+
tok = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
|
|
96
|
+
|
|
97
|
+
out = model.generate(**tok("On-device inference is", return_tensors="pt"),
|
|
98
|
+
max_new_tokens=32, do_sample=True)
|
|
99
|
+
```
|
|
100
|
+
|
|
101
|
+
**Today:** the operator layer under this is complete for 15 of 20 architectures, and Llama and
|
|
102
|
+
GPT-2 match upstream **token for token and logit for logit** in both greedy and sampling mode.
|
|
103
|
+
The `from_pretrained` path itself is still blocked — see [Status](#status).
|
|
104
|
+
|
|
105
|
+
### 2 · Federated learning
|
|
106
|
+
|
|
107
|
+
Devices train locally and share updates, not data. Federated averaging *is* collective
|
|
108
|
+
communication, so this is built on `torch.distributed` rather than beside it — broadcast the
|
|
109
|
+
model, gather the updates, weighted all-reduce.
|
|
110
|
+
|
|
111
|
+
```python
|
|
112
|
+
from torchnative.nn import federated
|
|
113
|
+
|
|
114
|
+
engine = federated.Engine(model, rounds=..., aggregator=federated.FedAvg())
|
|
115
|
+
engine.participate() # local epochs, then contribute a delta
|
|
116
|
+
```
|
|
117
|
+
|
|
118
|
+
**Today:** planned. `torch.distributed` is being implemented from `world_size = 1` upward, which
|
|
119
|
+
is a truthful description of a single device rather than a stub.
|
|
120
|
+
|
|
121
|
+
### 3 · Test-time adaptation & training
|
|
122
|
+
|
|
123
|
+
A model that ships to a device meets data the training set never had. TTA, TTT and the wider
|
|
124
|
+
test-time learning family let it adapt in place — and every method reduces to the same thing: a
|
|
125
|
+
weight delta over base weights, differing only in lifetime and destination.
|
|
126
|
+
|
|
127
|
+
```python
|
|
128
|
+
from torchnative import adapt
|
|
129
|
+
|
|
130
|
+
model = adapt.wrap(model, method=adapt.Tent()) # or TTT, memory-based, entropy-based
|
|
131
|
+
model.online() # adapt as it serves
|
|
132
|
+
```
|
|
133
|
+
|
|
134
|
+
**Today:** planned; the delta abstraction is specified in [`docs/DESIGN.md`](docs/DESIGN.md) §3.
|
|
135
|
+
Lifetime is driven by system events — backgrounding, user switch, sync window — rather than by
|
|
136
|
+
the domain boundaries a benchmark hands you.
|
|
137
|
+
|
|
138
|
+
---
|
|
139
|
+
|
|
140
|
+
## How it works
|
|
141
|
+
|
|
142
|
+
```
|
|
143
|
+
your code · transformers · torch/*.py upstream Python, unmodified
|
|
144
|
+
──────────────────────────────────────────
|
|
145
|
+
torch._C ← replaced
|
|
146
|
+
├── _aten_dispatch the single door every operator passes
|
|
147
|
+
├── Python spellings torch.mm, x.softmax(), F.linear, ...
|
|
148
|
+
└── kernels Rust, backed by candle
|
|
149
|
+
──────────────────────────────────────────
|
|
150
|
+
CPU today · Metal, Vulkan, NPU planned
|
|
151
|
+
```
|
|
152
|
+
|
|
153
|
+
**One door.** Every operator reaches its kernel through `_aten_dispatch`, and nothing bypasses
|
|
154
|
+
it. That makes the surface measurable — an unimplemented operator names itself rather than
|
|
155
|
+
failing downstream — and it gives graph capture, which NPU backends will need, exactly one place
|
|
156
|
+
to attach.
|
|
157
|
+
|
|
158
|
+
**Demand-driven.** Nothing is implemented because it might be needed. The shim refuses by name,
|
|
159
|
+
the refusal names the next thing to build, and that list comes from running real models.
|
|
160
|
+
|
|
161
|
+
**Stable ABI.** Built against CPython's limited API (`abi3-py313`), so one binary per platform
|
|
162
|
+
loads on 3.13, 3.14 and later without a rebuild.
|
|
163
|
+
|
|
164
|
+
---
|
|
165
|
+
|
|
166
|
+
## Status
|
|
167
|
+
|
|
168
|
+
<table>
|
|
169
|
+
<tr><th align="left">Working</th><th align="left"></th></tr>
|
|
170
|
+
<tr><td>ATen operators</td><td><b>91</b>, each compared against upstream</td></tr>
|
|
171
|
+
<tr><td>Golden comparison cases</td><td><b>2095 / 2095</b> — values, shapes, dtypes</td></tr>
|
|
172
|
+
<tr><td>Python spellings</td><td><b>204</b> verified against upstream signatures</td></tr>
|
|
173
|
+
<tr><td>Architectures complete</td><td><b>15 of 20</b> measured</td></tr>
|
|
174
|
+
<tr><td>Build targets</td><td>macOS · Linux · Android arm64 · iOS arm64</td></tr>
|
|
175
|
+
</table>
|
|
176
|
+
|
|
177
|
+
Complete: Llama · GPT-2 · Qwen2 · Mistral · Gemma · GPT-NeoX · OPT · MPT · StarCoder2 ·
|
|
178
|
+
Persimmon · Cohere · StableLM · OLMo · Phi · BERT
|
|
179
|
+
|
|
180
|
+
`uniform_` and `normal_` are **bit-identical** to upstream, and `multinomial` consumes the same
|
|
181
|
+
generator stream — a seeded run reproduces exactly.
|
|
182
|
+
|
|
183
|
+
**Not working yet**
|
|
184
|
+
|
|
185
|
+
- `import transformers` fails — `torch.distributed` is unimplemented and an unguarded import
|
|
186
|
+
inside `torch._dynamo` reaches `dist.Store`. Every result above was measured against models
|
|
187
|
+
transcribed by hand.
|
|
188
|
+
- No checkpoint has ever been loaded. All weights so far are randomly initialised.
|
|
189
|
+
- Mobile is link-verified only — the artefacts build and link, but no device has loaded one.
|
|
190
|
+
- CPU only. No GPU or NPU backend.
|
|
191
|
+
|
|
192
|
+
Tracked with the measurements behind them in [`docs/DESIGN.md`](docs/DESIGN.md) §11.1.
|
|
193
|
+
|
|
194
|
+
---
|
|
195
|
+
|
|
196
|
+
## Verification
|
|
197
|
+
|
|
198
|
+
Correctness here means *agreeing with upstream PyTorch*, so the strategy is comparison rather
|
|
199
|
+
than assertion.
|
|
200
|
+
|
|
201
|
+
| | |
|
|
202
|
+
|---|---|
|
|
203
|
+
| **Golden comparison** | Every operator runs on both upstream torch and this shim, compared on value, shape and dtype. It has caught a `float16` GEMM accumulating in `float16` where torch accumulates in `float32`, `cumsum` routed through the wrong kernel, and integer overflow where torch refuses. |
|
|
204
|
+
| **The harness tests itself** | `--self-test` injects a fault shaped like a plausible misimplementation at each comparator and fails if the comparator accepts it — 11 comparators × 11 fault modes, with any comparator never exercised reported as failure. It found that the previous fault injection reached exactly one case out of 1781. |
|
|
205
|
+
| **Tokens are not enough** | A wrong `gelu` approximation produced *identical tokens* while logits differed by 5.9e-04. End-to-end tests compare logits too, with a tolerance measured to sit between normal float32 noise and that failure. |
|
|
206
|
+
|
|
207
|
+
```sh
|
|
208
|
+
sh rust/torch_c/pytests/run.sh # smoke tests + harness self-test
|
|
209
|
+
python tools/golden/compare.py # golden comparison against upstream
|
|
210
|
+
python rust/torch_c/pytests/verify_schemas.py # signature tables vs upstream
|
|
211
|
+
```
|
|
212
|
+
|
|
213
|
+
---
|
|
214
|
+
|
|
215
|
+
## Roadmap
|
|
216
|
+
|
|
217
|
+
The next milestone is the device abstraction, because everything waits on it — a distributed rank
|
|
218
|
+
needs a device to point at, and every accelerator attaches there.
|
|
219
|
+
|
|
220
|
+
```
|
|
221
|
+
torchnative.nn.federated rounds · client selection · aggregation · dropout
|
|
222
|
+
└ torch.distributed ProcessGroup · collectives (transport)
|
|
223
|
+
└ backends ours, via register_backend
|
|
224
|
+
└ devices CPU · Metal · Vulkan · NPU
|
|
225
|
+
```
|
|
226
|
+
|
|
227
|
+
| | |
|
|
228
|
+
|---|---|
|
|
229
|
+
| **Device abstraction** | `torch.device`, per-device dispatch. Everything else waits on it. |
|
|
230
|
+
| **Metal** | candle already has the backend; disabled here for build isolation, not absent. |
|
|
231
|
+
| **`torch.distributed`** | From `world_size = 1` upward. Unblocks `transformers` as a side effect. |
|
|
232
|
+
| **Android GPU** | No candle backend and no `vulkan` slot in the `kernels` contract — genuinely new work. |
|
|
233
|
+
| **NPU** | ANE, NNAPI and QNN take a whole graph ahead of time, so this needs a capture layer rather than another device. The single door is where it attaches. |
|
|
234
|
+
|
|
235
|
+
---
|
|
236
|
+
|
|
237
|
+
## Install
|
|
238
|
+
|
|
239
|
+
> Not published yet. The alpha will carry platform wheels for macOS, Linux, Android and iOS —
|
|
240
|
+
> the extension is native, so `py3-none-any` is not the shape this ships in.
|
|
241
|
+
|
|
242
|
+
```sh
|
|
243
|
+
pip install torchnative
|
|
244
|
+
```
|
|
245
|
+
|
|
246
|
+
### Building from source
|
|
247
|
+
|
|
248
|
+
Requires a Rust toolchain and CPython 3.13+.
|
|
249
|
+
|
|
250
|
+
```sh
|
|
251
|
+
bash vendor/vendor_torch.sh # assemble the vendored torch tree
|
|
252
|
+
bash vendor/install_shim.sh # build the extension and install it
|
|
253
|
+
```
|
|
254
|
+
|
|
255
|
+
Cross-compilation is documented in [`docs/RUST_CROSSBUILD.md`](docs/RUST_CROSSBUILD.md),
|
|
256
|
+
including the PyO3 configuration iOS needs in order not to link `libpython`.
|
|
257
|
+
|
|
258
|
+
---
|
|
259
|
+
|
|
260
|
+
## Repository layout
|
|
261
|
+
|
|
262
|
+
```
|
|
263
|
+
torchnative/ the Python library
|
|
264
|
+
rust/torch_c/ the torch._C replacement (Rust · PyO3 · candle)
|
|
265
|
+
tools/golden/ the upstream comparison harness
|
|
266
|
+
vendor/ scripts that assemble the vendored torch tree (not checked in)
|
|
267
|
+
docs/ design, measurements, and the reasoning behind open decisions
|
|
268
|
+
```
|
|
269
|
+
|
|
270
|
+
`docs/` is written to be read. It records what was measured, what was assumed, and where an
|
|
271
|
+
earlier conclusion turned out to be wrong — corrections are left visible rather than edited away.
|
|
272
|
+
Start with [`DESIGN.md`](docs/DESIGN.md); [`SURFACE_HONESTY.md`](docs/SURFACE_HONESTY.md) and
|
|
273
|
+
[`HARNESS.md`](docs/HARNESS.md) show the standard the rest aims for.
|
|
274
|
+
|
|
275
|
+
---
|
|
276
|
+
|
|
277
|
+
## Related
|
|
278
|
+
|
|
279
|
+
- [PythonMultiplatform](https://github.com/thisisthepy/PythonMultiplatform) — embeds CPython 3.13
|
|
280
|
+
into Kotlin Multiplatform; the deployment target for this library
|
|
281
|
+
- [pypackpack](https://github.com/thisisthepy/pypackpack) — the build and bundling tool
|
|
282
|
+
- [Hugging Face `kernels`](https://github.com/huggingface/kernels) — the fused-kernel contract
|
|
283
|
+
this adopts, with resolution moved from runtime download to build time, since downloading
|
|
284
|
+
executable code is not permitted on every target platform
|
|
285
|
+
|
|
286
|
+
---
|
|
287
|
+
|
|
288
|
+
## License
|
|
289
|
+
|
|
290
|
+
MIT — see [LICENSE](LICENSE).
|
|
291
|
+
|
|
292
|
+
PyTorch is vendored under its own BSD-3-Clause license. The vendored tree is assembled at build
|
|
293
|
+
time and is not redistributed in this repository.
|
|
@@ -0,0 +1,266 @@
|
|
|
1
|
+
<div align="center">
|
|
2
|
+
|
|
3
|
+
# torchnative
|
|
4
|
+
|
|
5
|
+
**Run the real PyTorch ecosystem on device — not a reimplementation of it.**
|
|
6
|
+
|
|
7
|
+
[](https://pypi.org/project/torchnative/)
|
|
8
|
+
[](https://www.python.org/)
|
|
9
|
+
[](LICENSE)
|
|
10
|
+
[](#install)
|
|
11
|
+
[](#status)
|
|
12
|
+
|
|
13
|
+
</div>
|
|
14
|
+
|
|
15
|
+
---
|
|
16
|
+
|
|
17
|
+
`torchnative` replaces PyTorch's compiled core — `torch._C` — with a native extension, so the
|
|
18
|
+
genuine `torch` and `transformers` packages run on a phone the way they run on a workstation.
|
|
19
|
+
|
|
20
|
+
Models are not ported, converted, or re-expressed. They are imported.
|
|
21
|
+
|
|
22
|
+
```python
|
|
23
|
+
from transformers import AutoModelForCausalLM # the real one
|
|
24
|
+
model = AutoModelForCausalLM.from_pretrained("...")
|
|
25
|
+
model.generate(...) # on the device
|
|
26
|
+
```
|
|
27
|
+
|
|
28
|
+
> [!WARNING]
|
|
29
|
+
> **Pre-alpha.** The operator layer matches upstream PyTorch numerically and 15 of 20 tested
|
|
30
|
+
> architectures reach zero missing operators — but `import transformers` does not work yet, no
|
|
31
|
+
> checkpoint has ever been loaded, and no device has run the built artefact.
|
|
32
|
+
> See [Status](#status) before depending on this.
|
|
33
|
+
|
|
34
|
+
---
|
|
35
|
+
|
|
36
|
+
## Why not a reimplementation
|
|
37
|
+
|
|
38
|
+
Every other route to on-device inference re-expresses the model somewhere else.
|
|
39
|
+
|
|
40
|
+
| | approach | cost |
|
|
41
|
+
|---|---|---|
|
|
42
|
+
| llama.cpp | architectures rewritten in C++ | each new architecture is a porting task |
|
|
43
|
+
| ExecuTorch · CoreML | ahead-of-time compiled graph | export step, and what runs is not what you wrote |
|
|
44
|
+
| MLC | lowered to its own runtime | same |
|
|
45
|
+
| **torchnative** | **the real Python package** | **the substrate is hard; architectures are free** |
|
|
46
|
+
|
|
47
|
+
The reason nobody runs the real thing is that `torch._C` cannot be built for mobile. PyTorch's own
|
|
48
|
+
build sets `INTERN_BUILD_MOBILE` for any Android or iOS toolchain, and that path forces
|
|
49
|
+
`BUILD_PYTHON` off — so the mobile build is structurally incapable of producing the Python
|
|
50
|
+
extension module the Python package needs.
|
|
51
|
+
|
|
52
|
+
`torchnative` supplies that module instead. Everything above it is upstream source, unmodified.
|
|
53
|
+
|
|
54
|
+
---
|
|
55
|
+
|
|
56
|
+
## What it does
|
|
57
|
+
|
|
58
|
+
### 1 · LLM inference
|
|
59
|
+
|
|
60
|
+
Run `transformers` models directly. No conversion step, no per-architecture port — if
|
|
61
|
+
`transformers` supports it and the operators are covered, it runs.
|
|
62
|
+
|
|
63
|
+
```python
|
|
64
|
+
import torch
|
|
65
|
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
66
|
+
|
|
67
|
+
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
|
|
68
|
+
tok = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
|
|
69
|
+
|
|
70
|
+
out = model.generate(**tok("On-device inference is", return_tensors="pt"),
|
|
71
|
+
max_new_tokens=32, do_sample=True)
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
**Today:** the operator layer under this is complete for 15 of 20 architectures, and Llama and
|
|
75
|
+
GPT-2 match upstream **token for token and logit for logit** in both greedy and sampling mode.
|
|
76
|
+
The `from_pretrained` path itself is still blocked — see [Status](#status).
|
|
77
|
+
|
|
78
|
+
### 2 · Federated learning
|
|
79
|
+
|
|
80
|
+
Devices train locally and share updates, not data. Federated averaging *is* collective
|
|
81
|
+
communication, so this is built on `torch.distributed` rather than beside it — broadcast the
|
|
82
|
+
model, gather the updates, weighted all-reduce.
|
|
83
|
+
|
|
84
|
+
```python
|
|
85
|
+
from torchnative.nn import federated
|
|
86
|
+
|
|
87
|
+
engine = federated.Engine(model, rounds=..., aggregator=federated.FedAvg())
|
|
88
|
+
engine.participate() # local epochs, then contribute a delta
|
|
89
|
+
```
|
|
90
|
+
|
|
91
|
+
**Today:** planned. `torch.distributed` is being implemented from `world_size = 1` upward, which
|
|
92
|
+
is a truthful description of a single device rather than a stub.
|
|
93
|
+
|
|
94
|
+
### 3 · Test-time adaptation & training
|
|
95
|
+
|
|
96
|
+
A model that ships to a device meets data the training set never had. TTA, TTT and the wider
|
|
97
|
+
test-time learning family let it adapt in place — and every method reduces to the same thing: a
|
|
98
|
+
weight delta over base weights, differing only in lifetime and destination.
|
|
99
|
+
|
|
100
|
+
```python
|
|
101
|
+
from torchnative import adapt
|
|
102
|
+
|
|
103
|
+
model = adapt.wrap(model, method=adapt.Tent()) # or TTT, memory-based, entropy-based
|
|
104
|
+
model.online() # adapt as it serves
|
|
105
|
+
```
|
|
106
|
+
|
|
107
|
+
**Today:** planned; the delta abstraction is specified in [`docs/DESIGN.md`](docs/DESIGN.md) §3.
|
|
108
|
+
Lifetime is driven by system events — backgrounding, user switch, sync window — rather than by
|
|
109
|
+
the domain boundaries a benchmark hands you.
|
|
110
|
+
|
|
111
|
+
---
|
|
112
|
+
|
|
113
|
+
## How it works
|
|
114
|
+
|
|
115
|
+
```
|
|
116
|
+
your code · transformers · torch/*.py upstream Python, unmodified
|
|
117
|
+
──────────────────────────────────────────
|
|
118
|
+
torch._C ← replaced
|
|
119
|
+
├── _aten_dispatch the single door every operator passes
|
|
120
|
+
├── Python spellings torch.mm, x.softmax(), F.linear, ...
|
|
121
|
+
└── kernels Rust, backed by candle
|
|
122
|
+
──────────────────────────────────────────
|
|
123
|
+
CPU today · Metal, Vulkan, NPU planned
|
|
124
|
+
```
|
|
125
|
+
|
|
126
|
+
**One door.** Every operator reaches its kernel through `_aten_dispatch`, and nothing bypasses
|
|
127
|
+
it. That makes the surface measurable — an unimplemented operator names itself rather than
|
|
128
|
+
failing downstream — and it gives graph capture, which NPU backends will need, exactly one place
|
|
129
|
+
to attach.
|
|
130
|
+
|
|
131
|
+
**Demand-driven.** Nothing is implemented because it might be needed. The shim refuses by name,
|
|
132
|
+
the refusal names the next thing to build, and that list comes from running real models.
|
|
133
|
+
|
|
134
|
+
**Stable ABI.** Built against CPython's limited API (`abi3-py313`), so one binary per platform
|
|
135
|
+
loads on 3.13, 3.14 and later without a rebuild.
|
|
136
|
+
|
|
137
|
+
---
|
|
138
|
+
|
|
139
|
+
## Status
|
|
140
|
+
|
|
141
|
+
<table>
|
|
142
|
+
<tr><th align="left">Working</th><th align="left"></th></tr>
|
|
143
|
+
<tr><td>ATen operators</td><td><b>91</b>, each compared against upstream</td></tr>
|
|
144
|
+
<tr><td>Golden comparison cases</td><td><b>2095 / 2095</b> — values, shapes, dtypes</td></tr>
|
|
145
|
+
<tr><td>Python spellings</td><td><b>204</b> verified against upstream signatures</td></tr>
|
|
146
|
+
<tr><td>Architectures complete</td><td><b>15 of 20</b> measured</td></tr>
|
|
147
|
+
<tr><td>Build targets</td><td>macOS · Linux · Android arm64 · iOS arm64</td></tr>
|
|
148
|
+
</table>
|
|
149
|
+
|
|
150
|
+
Complete: Llama · GPT-2 · Qwen2 · Mistral · Gemma · GPT-NeoX · OPT · MPT · StarCoder2 ·
|
|
151
|
+
Persimmon · Cohere · StableLM · OLMo · Phi · BERT
|
|
152
|
+
|
|
153
|
+
`uniform_` and `normal_` are **bit-identical** to upstream, and `multinomial` consumes the same
|
|
154
|
+
generator stream — a seeded run reproduces exactly.
|
|
155
|
+
|
|
156
|
+
**Not working yet**
|
|
157
|
+
|
|
158
|
+
- `import transformers` fails — `torch.distributed` is unimplemented and an unguarded import
|
|
159
|
+
inside `torch._dynamo` reaches `dist.Store`. Every result above was measured against models
|
|
160
|
+
transcribed by hand.
|
|
161
|
+
- No checkpoint has ever been loaded. All weights so far are randomly initialised.
|
|
162
|
+
- Mobile is link-verified only — the artefacts build and link, but no device has loaded one.
|
|
163
|
+
- CPU only. No GPU or NPU backend.
|
|
164
|
+
|
|
165
|
+
Tracked with the measurements behind them in [`docs/DESIGN.md`](docs/DESIGN.md) §11.1.
|
|
166
|
+
|
|
167
|
+
---
|
|
168
|
+
|
|
169
|
+
## Verification
|
|
170
|
+
|
|
171
|
+
Correctness here means *agreeing with upstream PyTorch*, so the strategy is comparison rather
|
|
172
|
+
than assertion.
|
|
173
|
+
|
|
174
|
+
| | |
|
|
175
|
+
|---|---|
|
|
176
|
+
| **Golden comparison** | Every operator runs on both upstream torch and this shim, compared on value, shape and dtype. It has caught a `float16` GEMM accumulating in `float16` where torch accumulates in `float32`, `cumsum` routed through the wrong kernel, and integer overflow where torch refuses. |
|
|
177
|
+
| **The harness tests itself** | `--self-test` injects a fault shaped like a plausible misimplementation at each comparator and fails if the comparator accepts it — 11 comparators × 11 fault modes, with any comparator never exercised reported as failure. It found that the previous fault injection reached exactly one case out of 1781. |
|
|
178
|
+
| **Tokens are not enough** | A wrong `gelu` approximation produced *identical tokens* while logits differed by 5.9e-04. End-to-end tests compare logits too, with a tolerance measured to sit between normal float32 noise and that failure. |
|
|
179
|
+
|
|
180
|
+
```sh
|
|
181
|
+
sh rust/torch_c/pytests/run.sh # smoke tests + harness self-test
|
|
182
|
+
python tools/golden/compare.py # golden comparison against upstream
|
|
183
|
+
python rust/torch_c/pytests/verify_schemas.py # signature tables vs upstream
|
|
184
|
+
```
|
|
185
|
+
|
|
186
|
+
---
|
|
187
|
+
|
|
188
|
+
## Roadmap
|
|
189
|
+
|
|
190
|
+
The next milestone is the device abstraction, because everything waits on it — a distributed rank
|
|
191
|
+
needs a device to point at, and every accelerator attaches there.
|
|
192
|
+
|
|
193
|
+
```
|
|
194
|
+
torchnative.nn.federated rounds · client selection · aggregation · dropout
|
|
195
|
+
└ torch.distributed ProcessGroup · collectives (transport)
|
|
196
|
+
└ backends ours, via register_backend
|
|
197
|
+
└ devices CPU · Metal · Vulkan · NPU
|
|
198
|
+
```
|
|
199
|
+
|
|
200
|
+
| | |
|
|
201
|
+
|---|---|
|
|
202
|
+
| **Device abstraction** | `torch.device`, per-device dispatch. Everything else waits on it. |
|
|
203
|
+
| **Metal** | candle already has the backend; disabled here for build isolation, not absent. |
|
|
204
|
+
| **`torch.distributed`** | From `world_size = 1` upward. Unblocks `transformers` as a side effect. |
|
|
205
|
+
| **Android GPU** | No candle backend and no `vulkan` slot in the `kernels` contract — genuinely new work. |
|
|
206
|
+
| **NPU** | ANE, NNAPI and QNN take a whole graph ahead of time, so this needs a capture layer rather than another device. The single door is where it attaches. |
|
|
207
|
+
|
|
208
|
+
---
|
|
209
|
+
|
|
210
|
+
## Install
|
|
211
|
+
|
|
212
|
+
> Not published yet. The alpha will carry platform wheels for macOS, Linux, Android and iOS —
|
|
213
|
+
> the extension is native, so `py3-none-any` is not the shape this ships in.
|
|
214
|
+
|
|
215
|
+
```sh
|
|
216
|
+
pip install torchnative
|
|
217
|
+
```
|
|
218
|
+
|
|
219
|
+
### Building from source
|
|
220
|
+
|
|
221
|
+
Requires a Rust toolchain and CPython 3.13+.
|
|
222
|
+
|
|
223
|
+
```sh
|
|
224
|
+
bash vendor/vendor_torch.sh # assemble the vendored torch tree
|
|
225
|
+
bash vendor/install_shim.sh # build the extension and install it
|
|
226
|
+
```
|
|
227
|
+
|
|
228
|
+
Cross-compilation is documented in [`docs/RUST_CROSSBUILD.md`](docs/RUST_CROSSBUILD.md),
|
|
229
|
+
including the PyO3 configuration iOS needs in order not to link `libpython`.
|
|
230
|
+
|
|
231
|
+
---
|
|
232
|
+
|
|
233
|
+
## Repository layout
|
|
234
|
+
|
|
235
|
+
```
|
|
236
|
+
torchnative/ the Python library
|
|
237
|
+
rust/torch_c/ the torch._C replacement (Rust · PyO3 · candle)
|
|
238
|
+
tools/golden/ the upstream comparison harness
|
|
239
|
+
vendor/ scripts that assemble the vendored torch tree (not checked in)
|
|
240
|
+
docs/ design, measurements, and the reasoning behind open decisions
|
|
241
|
+
```
|
|
242
|
+
|
|
243
|
+
`docs/` is written to be read. It records what was measured, what was assumed, and where an
|
|
244
|
+
earlier conclusion turned out to be wrong — corrections are left visible rather than edited away.
|
|
245
|
+
Start with [`DESIGN.md`](docs/DESIGN.md); [`SURFACE_HONESTY.md`](docs/SURFACE_HONESTY.md) and
|
|
246
|
+
[`HARNESS.md`](docs/HARNESS.md) show the standard the rest aims for.
|
|
247
|
+
|
|
248
|
+
---
|
|
249
|
+
|
|
250
|
+
## Related
|
|
251
|
+
|
|
252
|
+
- [PythonMultiplatform](https://github.com/thisisthepy/PythonMultiplatform) — embeds CPython 3.13
|
|
253
|
+
into Kotlin Multiplatform; the deployment target for this library
|
|
254
|
+
- [pypackpack](https://github.com/thisisthepy/pypackpack) — the build and bundling tool
|
|
255
|
+
- [Hugging Face `kernels`](https://github.com/huggingface/kernels) — the fused-kernel contract
|
|
256
|
+
this adopts, with resolution moved from runtime download to build time, since downloading
|
|
257
|
+
executable code is not permitted on every target platform
|
|
258
|
+
|
|
259
|
+
---
|
|
260
|
+
|
|
261
|
+
## License
|
|
262
|
+
|
|
263
|
+
MIT — see [LICENSE](LICENSE).
|
|
264
|
+
|
|
265
|
+
PyTorch is vendored under its own BSD-3-Clause license. The vendored tree is assembled at build
|
|
266
|
+
time and is not redistributed in this repository.
|
|
@@ -0,0 +1,87 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=77"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "torchnative"
|
|
7
|
+
version = "0.0.1a0"
|
|
8
|
+
description = "Run the real PyTorch ecosystem on device -- not a reimplementation of it"
|
|
9
|
+
readme = { file = "README.md", content-type = "text/markdown" }
|
|
10
|
+
license = "MIT"
|
|
11
|
+
license-files = ["LICENSE"]
|
|
12
|
+
requires-python = ">=3.13"
|
|
13
|
+
|
|
14
|
+
# Empty on purpose. This distribution **provides** `torch` -- the upstream
|
|
15
|
+
# Python tree with `_C` replaced (DESIGN.md §2) -- rather than consuming
|
|
16
|
+
# somebody else's, so requiring `torch` here would be a package declaring a
|
|
17
|
+
# dependency on itself.
|
|
18
|
+
#
|
|
19
|
+
# Two earlier attempts at this line were both wrong, in opposite directions.
|
|
20
|
+
# Empty with a module-scope `from torch import nn` installed cleanly and then
|
|
21
|
+
# failed on import. Requiring `torch` unconditionally was unresolvable on
|
|
22
|
+
# Android and iOS, where upstream publishes no wheel at all -- which is to say
|
|
23
|
+
# on exactly the platforms this project exists for. The mistake underneath both
|
|
24
|
+
# was reading upstream torch as a runtime dependency when it is a *comparison
|
|
25
|
+
# baseline*: the golden harness needs one installed to diff against, and nothing
|
|
26
|
+
# at runtime does.
|
|
27
|
+
#
|
|
28
|
+
# It follows that this cannot coexist with an installed PyTorch once the wheels
|
|
29
|
+
# carry the tree. That is not a defect to route around; there is one `torch` on
|
|
30
|
+
# a given interpreter, and this is a build of it.
|
|
31
|
+
dependencies = []
|
|
32
|
+
authors = [{ name = "thisisthepy" }]
|
|
33
|
+
keywords = ["pytorch", "on-device", "edge", "federated-learning", "test-time-adaptation"]
|
|
34
|
+
|
|
35
|
+
# Pre-alpha on purpose. The Python surface here is still a skeleton; the working
|
|
36
|
+
# part of this project is the `torch._C` replacement under `rust/torch_c`, which
|
|
37
|
+
# is not yet built into this wheel. See the README's Status section.
|
|
38
|
+
classifiers = [
|
|
39
|
+
"Development Status :: 2 - Pre-Alpha",
|
|
40
|
+
"Intended Audience :: Developers",
|
|
41
|
+
"Programming Language :: Python :: 3.13",
|
|
42
|
+
"Programming Language :: Rust",
|
|
43
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
44
|
+
"Operating System :: Android",
|
|
45
|
+
"Operating System :: iOS",
|
|
46
|
+
"Operating System :: MacOS",
|
|
47
|
+
"Operating System :: POSIX :: Linux",
|
|
48
|
+
]
|
|
49
|
+
|
|
50
|
+
[project.urls]
|
|
51
|
+
Homepage = "https://github.com/thisisthepy/torchnative"
|
|
52
|
+
Documentation = "https://github.com/thisisthepy/torchnative/tree/develop/docs"
|
|
53
|
+
Source = "https://github.com/thisisthepy/torchnative"
|
|
54
|
+
|
|
55
|
+
[project.optional-dependencies]
|
|
56
|
+
# Aggregation, transport and privacy live behind this extra so that using
|
|
57
|
+
# adaptation alone does not pull in the federated stack. See docs/DESIGN.md §10.
|
|
58
|
+
federated = []
|
|
59
|
+
|
|
60
|
+
# Upstream PyTorch, as the thing the golden harness diffs against -- not as a
|
|
61
|
+
# runtime dependency. Pinned to what is actually compared: this shim implements
|
|
62
|
+
# one release's `_C` surface, so a tree from another release expects different
|
|
63
|
+
# symbols from it. Installing this alongside the wheels that carry our own tree
|
|
64
|
+
# will conflict, which is why it is an extra and not a dependency.
|
|
65
|
+
test = ["torch>=2.13,<2.14"]
|
|
66
|
+
|
|
67
|
+
# Without this, setuptools auto-discovery treats `src` as the root and ships
|
|
68
|
+
# `main/` and `test/` as importable packages -- which is what it did, so a wheel
|
|
69
|
+
# built before this section answered `import main.torchnative` and not
|
|
70
|
+
# `import torchnative`. The source-set layout has to be spelled out.
|
|
71
|
+
[tool.setuptools]
|
|
72
|
+
package-dir = { "" = "torchnative/src/main" }
|
|
73
|
+
|
|
74
|
+
[tool.setuptools.packages.find]
|
|
75
|
+
where = ["torchnative/src/main"]
|
|
76
|
+
# `torch` is deliberately not shipped from this distribution. src/main/torch is
|
|
77
|
+
# the add-hook that grafts `torchnative.nn.federated` onto the torch namespace
|
|
78
|
+
# (DESIGN.md §2), and writing files into another distribution's package would
|
|
79
|
+
# collide with an installed PyTorch, break its uninstall, and conflict outright
|
|
80
|
+
# if upstream ever adds that path. It stays a build-time graft against the
|
|
81
|
+
# vendored tree we assemble ourselves.
|
|
82
|
+
include = ["torchnative*"]
|
|
83
|
+
|
|
84
|
+
[tool.ppp]
|
|
85
|
+
# Platform source sets follow the pypackpack layout: src/main is scanned for
|
|
86
|
+
# top-level Python packages, which is what lets one package provide both
|
|
87
|
+
# `torch` and `torchnative`. See docs/DESIGN.md §10.
|
|
@@ -0,0 +1,8 @@
|
|
|
1
|
+
"""Test-time learning methods.
|
|
2
|
+
|
|
3
|
+
TTL contains TTA contains TTT; the nesting is not flattened into sibling
|
|
4
|
+
modules. Each method declares its own differentiation requirement rather than
|
|
5
|
+
living in a directory named for one, because normalization calibration sits on
|
|
6
|
+
both sides of that line -- recomputing statistics needs no backward pass,
|
|
7
|
+
updating affine parameters by a loss does. See DESIGN.md §3.
|
|
8
|
+
"""
|
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
"""Deployment, lifetime policy and device orchestration.
|
|
2
|
+
|
|
3
|
+
The `torch` import is deferred rather than done at module scope. This
|
|
4
|
+
distribution *provides* `torch` -- the vendored tree plus our `_C` -- rather
|
|
5
|
+
than depending on someone else's, so at import time there may not be one yet:
|
|
6
|
+
the wheels that carry the tree are not wired up, and until they are, importing
|
|
7
|
+
`torchnative` should not require a PyTorch that this package is itself meant to
|
|
8
|
+
supply. See DESIGN.md §2 and the README's Status section.
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
from __future__ import annotations
|
|
12
|
+
|
|
13
|
+
from typing import TYPE_CHECKING
|
|
14
|
+
|
|
15
|
+
if TYPE_CHECKING:
|
|
16
|
+
from torch import nn
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class TorchNativeAPI(object):
|
|
20
|
+
def __init__(self, *args, **kwargs):
|
|
21
|
+
pass
|
|
22
|
+
|
|
23
|
+
def deploy(self, model: "nn.Module"):
|
|
24
|
+
pass
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
"""Weight deltas over base weights.
|
|
2
|
+
|
|
3
|
+
Every adaptation method produces one of these; they differ only in lifetime
|
|
4
|
+
and destination. The boundary is ``model(x)`` -- state a model manages inside
|
|
5
|
+
its own forward (fast weights, caches) is the model's, not ours.
|
|
6
|
+
|
|
7
|
+
Lifetime names are deliberately unnamed for now. They are driven by system
|
|
8
|
+
events -- backgrounding, user switch, sync window -- not by the domain
|
|
9
|
+
boundaries a benchmark hands you, and they get chosen once the first
|
|
10
|
+
integration shows the real usage. See DESIGN.md §3.
|
|
11
|
+
"""
|
|
@@ -0,0 +1,7 @@
|
|
|
1
|
+
"""Bundle-backed resolver for the Hugging Face ``kernels`` contract.
|
|
2
|
+
|
|
3
|
+
The contract is adopted as-is; only resolution is inverted. On mobile, kernel
|
|
4
|
+
variants are selected and compiled at build time and read from the app bundle,
|
|
5
|
+
because iOS will not execute downloaded native code. Desktop keeps resolving
|
|
6
|
+
against the Hub. See DESIGN.md §8.
|
|
7
|
+
"""
|
|
File without changes
|
|
@@ -0,0 +1,7 @@
|
|
|
1
|
+
"""Federated learning.
|
|
2
|
+
|
|
3
|
+
A layer above the single-device adaptation methods, not a sibling of them: the
|
|
4
|
+
local step is the same mechanism, with aggregation, transport and privacy on
|
|
5
|
+
top. Its dependencies stay behind the ``federated`` extra so that using
|
|
6
|
+
adaptation alone does not pull in the aggregation stack. See DESIGN.md §3.
|
|
7
|
+
"""
|
|
@@ -0,0 +1,293 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: torchnative
|
|
3
|
+
Version: 0.0.1a0
|
|
4
|
+
Summary: Run the real PyTorch ecosystem on device -- not a reimplementation of it
|
|
5
|
+
Author: thisisthepy
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/thisisthepy/torchnative
|
|
8
|
+
Project-URL: Documentation, https://github.com/thisisthepy/torchnative/tree/develop/docs
|
|
9
|
+
Project-URL: Source, https://github.com/thisisthepy/torchnative
|
|
10
|
+
Keywords: pytorch,on-device,edge,federated-learning,test-time-adaptation
|
|
11
|
+
Classifier: Development Status :: 2 - Pre-Alpha
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
14
|
+
Classifier: Programming Language :: Rust
|
|
15
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
16
|
+
Classifier: Operating System :: Android
|
|
17
|
+
Classifier: Operating System :: iOS
|
|
18
|
+
Classifier: Operating System :: MacOS
|
|
19
|
+
Classifier: Operating System :: POSIX :: Linux
|
|
20
|
+
Requires-Python: >=3.13
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
License-File: LICENSE
|
|
23
|
+
Provides-Extra: federated
|
|
24
|
+
Provides-Extra: test
|
|
25
|
+
Requires-Dist: torch<2.14,>=2.13; extra == "test"
|
|
26
|
+
Dynamic: license-file
|
|
27
|
+
|
|
28
|
+
<div align="center">
|
|
29
|
+
|
|
30
|
+
# torchnative
|
|
31
|
+
|
|
32
|
+
**Run the real PyTorch ecosystem on device — not a reimplementation of it.**
|
|
33
|
+
|
|
34
|
+
[](https://pypi.org/project/torchnative/)
|
|
35
|
+
[](https://www.python.org/)
|
|
36
|
+
[](LICENSE)
|
|
37
|
+
[](#install)
|
|
38
|
+
[](#status)
|
|
39
|
+
|
|
40
|
+
</div>
|
|
41
|
+
|
|
42
|
+
---
|
|
43
|
+
|
|
44
|
+
`torchnative` replaces PyTorch's compiled core — `torch._C` — with a native extension, so the
|
|
45
|
+
genuine `torch` and `transformers` packages run on a phone the way they run on a workstation.
|
|
46
|
+
|
|
47
|
+
Models are not ported, converted, or re-expressed. They are imported.
|
|
48
|
+
|
|
49
|
+
```python
|
|
50
|
+
from transformers import AutoModelForCausalLM # the real one
|
|
51
|
+
model = AutoModelForCausalLM.from_pretrained("...")
|
|
52
|
+
model.generate(...) # on the device
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
> [!WARNING]
|
|
56
|
+
> **Pre-alpha.** The operator layer matches upstream PyTorch numerically and 15 of 20 tested
|
|
57
|
+
> architectures reach zero missing operators — but `import transformers` does not work yet, no
|
|
58
|
+
> checkpoint has ever been loaded, and no device has run the built artefact.
|
|
59
|
+
> See [Status](#status) before depending on this.
|
|
60
|
+
|
|
61
|
+
---
|
|
62
|
+
|
|
63
|
+
## Why not a reimplementation
|
|
64
|
+
|
|
65
|
+
Every other route to on-device inference re-expresses the model somewhere else.
|
|
66
|
+
|
|
67
|
+
| | approach | cost |
|
|
68
|
+
|---|---|---|
|
|
69
|
+
| llama.cpp | architectures rewritten in C++ | each new architecture is a porting task |
|
|
70
|
+
| ExecuTorch · CoreML | ahead-of-time compiled graph | export step, and what runs is not what you wrote |
|
|
71
|
+
| MLC | lowered to its own runtime | same |
|
|
72
|
+
| **torchnative** | **the real Python package** | **the substrate is hard; architectures are free** |
|
|
73
|
+
|
|
74
|
+
The reason nobody runs the real thing is that `torch._C` cannot be built for mobile. PyTorch's own
|
|
75
|
+
build sets `INTERN_BUILD_MOBILE` for any Android or iOS toolchain, and that path forces
|
|
76
|
+
`BUILD_PYTHON` off — so the mobile build is structurally incapable of producing the Python
|
|
77
|
+
extension module the Python package needs.
|
|
78
|
+
|
|
79
|
+
`torchnative` supplies that module instead. Everything above it is upstream source, unmodified.
|
|
80
|
+
|
|
81
|
+
---
|
|
82
|
+
|
|
83
|
+
## What it does
|
|
84
|
+
|
|
85
|
+
### 1 · LLM inference
|
|
86
|
+
|
|
87
|
+
Run `transformers` models directly. No conversion step, no per-architecture port — if
|
|
88
|
+
`transformers` supports it and the operators are covered, it runs.
|
|
89
|
+
|
|
90
|
+
```python
|
|
91
|
+
import torch
|
|
92
|
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
93
|
+
|
|
94
|
+
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
|
|
95
|
+
tok = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
|
|
96
|
+
|
|
97
|
+
out = model.generate(**tok("On-device inference is", return_tensors="pt"),
|
|
98
|
+
max_new_tokens=32, do_sample=True)
|
|
99
|
+
```
|
|
100
|
+
|
|
101
|
+
**Today:** the operator layer under this is complete for 15 of 20 architectures, and Llama and
|
|
102
|
+
GPT-2 match upstream **token for token and logit for logit** in both greedy and sampling mode.
|
|
103
|
+
The `from_pretrained` path itself is still blocked — see [Status](#status).
|
|
104
|
+
|
|
105
|
+
### 2 · Federated learning
|
|
106
|
+
|
|
107
|
+
Devices train locally and share updates, not data. Federated averaging *is* collective
|
|
108
|
+
communication, so this is built on `torch.distributed` rather than beside it — broadcast the
|
|
109
|
+
model, gather the updates, weighted all-reduce.
|
|
110
|
+
|
|
111
|
+
```python
|
|
112
|
+
from torchnative.nn import federated
|
|
113
|
+
|
|
114
|
+
engine = federated.Engine(model, rounds=..., aggregator=federated.FedAvg())
|
|
115
|
+
engine.participate() # local epochs, then contribute a delta
|
|
116
|
+
```
|
|
117
|
+
|
|
118
|
+
**Today:** planned. `torch.distributed` is being implemented from `world_size = 1` upward, which
|
|
119
|
+
is a truthful description of a single device rather than a stub.
|
|
120
|
+
|
|
121
|
+
### 3 · Test-time adaptation & training
|
|
122
|
+
|
|
123
|
+
A model that ships to a device meets data the training set never had. TTA, TTT and the wider
|
|
124
|
+
test-time learning family let it adapt in place — and every method reduces to the same thing: a
|
|
125
|
+
weight delta over base weights, differing only in lifetime and destination.
|
|
126
|
+
|
|
127
|
+
```python
|
|
128
|
+
from torchnative import adapt
|
|
129
|
+
|
|
130
|
+
model = adapt.wrap(model, method=adapt.Tent()) # or TTT, memory-based, entropy-based
|
|
131
|
+
model.online() # adapt as it serves
|
|
132
|
+
```
|
|
133
|
+
|
|
134
|
+
**Today:** planned; the delta abstraction is specified in [`docs/DESIGN.md`](docs/DESIGN.md) §3.
|
|
135
|
+
Lifetime is driven by system events — backgrounding, user switch, sync window — rather than by
|
|
136
|
+
the domain boundaries a benchmark hands you.
|
|
137
|
+
|
|
138
|
+
---
|
|
139
|
+
|
|
140
|
+
## How it works
|
|
141
|
+
|
|
142
|
+
```
|
|
143
|
+
your code · transformers · torch/*.py upstream Python, unmodified
|
|
144
|
+
──────────────────────────────────────────
|
|
145
|
+
torch._C ← replaced
|
|
146
|
+
├── _aten_dispatch the single door every operator passes
|
|
147
|
+
├── Python spellings torch.mm, x.softmax(), F.linear, ...
|
|
148
|
+
└── kernels Rust, backed by candle
|
|
149
|
+
──────────────────────────────────────────
|
|
150
|
+
CPU today · Metal, Vulkan, NPU planned
|
|
151
|
+
```
|
|
152
|
+
|
|
153
|
+
**One door.** Every operator reaches its kernel through `_aten_dispatch`, and nothing bypasses
|
|
154
|
+
it. That makes the surface measurable — an unimplemented operator names itself rather than
|
|
155
|
+
failing downstream — and it gives graph capture, which NPU backends will need, exactly one place
|
|
156
|
+
to attach.
|
|
157
|
+
|
|
158
|
+
**Demand-driven.** Nothing is implemented because it might be needed. The shim refuses by name,
|
|
159
|
+
the refusal names the next thing to build, and that list comes from running real models.
|
|
160
|
+
|
|
161
|
+
**Stable ABI.** Built against CPython's limited API (`abi3-py313`), so one binary per platform
|
|
162
|
+
loads on 3.13, 3.14 and later without a rebuild.
|
|
163
|
+
|
|
164
|
+
---
|
|
165
|
+
|
|
166
|
+
## Status
|
|
167
|
+
|
|
168
|
+
<table>
|
|
169
|
+
<tr><th align="left">Working</th><th align="left"></th></tr>
|
|
170
|
+
<tr><td>ATen operators</td><td><b>91</b>, each compared against upstream</td></tr>
|
|
171
|
+
<tr><td>Golden comparison cases</td><td><b>2095 / 2095</b> — values, shapes, dtypes</td></tr>
|
|
172
|
+
<tr><td>Python spellings</td><td><b>204</b> verified against upstream signatures</td></tr>
|
|
173
|
+
<tr><td>Architectures complete</td><td><b>15 of 20</b> measured</td></tr>
|
|
174
|
+
<tr><td>Build targets</td><td>macOS · Linux · Android arm64 · iOS arm64</td></tr>
|
|
175
|
+
</table>
|
|
176
|
+
|
|
177
|
+
Complete: Llama · GPT-2 · Qwen2 · Mistral · Gemma · GPT-NeoX · OPT · MPT · StarCoder2 ·
|
|
178
|
+
Persimmon · Cohere · StableLM · OLMo · Phi · BERT
|
|
179
|
+
|
|
180
|
+
`uniform_` and `normal_` are **bit-identical** to upstream, and `multinomial` consumes the same
|
|
181
|
+
generator stream — a seeded run reproduces exactly.
|
|
182
|
+
|
|
183
|
+
**Not working yet**
|
|
184
|
+
|
|
185
|
+
- `import transformers` fails — `torch.distributed` is unimplemented and an unguarded import
|
|
186
|
+
inside `torch._dynamo` reaches `dist.Store`. Every result above was measured against models
|
|
187
|
+
transcribed by hand.
|
|
188
|
+
- No checkpoint has ever been loaded. All weights so far are randomly initialised.
|
|
189
|
+
- Mobile is link-verified only — the artefacts build and link, but no device has loaded one.
|
|
190
|
+
- CPU only. No GPU or NPU backend.
|
|
191
|
+
|
|
192
|
+
Tracked with the measurements behind them in [`docs/DESIGN.md`](docs/DESIGN.md) §11.1.
|
|
193
|
+
|
|
194
|
+
---
|
|
195
|
+
|
|
196
|
+
## Verification
|
|
197
|
+
|
|
198
|
+
Correctness here means *agreeing with upstream PyTorch*, so the strategy is comparison rather
|
|
199
|
+
than assertion.
|
|
200
|
+
|
|
201
|
+
| | |
|
|
202
|
+
|---|---|
|
|
203
|
+
| **Golden comparison** | Every operator runs on both upstream torch and this shim, compared on value, shape and dtype. It has caught a `float16` GEMM accumulating in `float16` where torch accumulates in `float32`, `cumsum` routed through the wrong kernel, and integer overflow where torch refuses. |
|
|
204
|
+
| **The harness tests itself** | `--self-test` injects a fault shaped like a plausible misimplementation at each comparator and fails if the comparator accepts it — 11 comparators × 11 fault modes, with any comparator never exercised reported as failure. It found that the previous fault injection reached exactly one case out of 1781. |
|
|
205
|
+
| **Tokens are not enough** | A wrong `gelu` approximation produced *identical tokens* while logits differed by 5.9e-04. End-to-end tests compare logits too, with a tolerance measured to sit between normal float32 noise and that failure. |
|
|
206
|
+
|
|
207
|
+
```sh
|
|
208
|
+
sh rust/torch_c/pytests/run.sh # smoke tests + harness self-test
|
|
209
|
+
python tools/golden/compare.py # golden comparison against upstream
|
|
210
|
+
python rust/torch_c/pytests/verify_schemas.py # signature tables vs upstream
|
|
211
|
+
```
|
|
212
|
+
|
|
213
|
+
---
|
|
214
|
+
|
|
215
|
+
## Roadmap
|
|
216
|
+
|
|
217
|
+
The next milestone is the device abstraction, because everything waits on it — a distributed rank
|
|
218
|
+
needs a device to point at, and every accelerator attaches there.
|
|
219
|
+
|
|
220
|
+
```
|
|
221
|
+
torchnative.nn.federated rounds · client selection · aggregation · dropout
|
|
222
|
+
└ torch.distributed ProcessGroup · collectives (transport)
|
|
223
|
+
└ backends ours, via register_backend
|
|
224
|
+
└ devices CPU · Metal · Vulkan · NPU
|
|
225
|
+
```
|
|
226
|
+
|
|
227
|
+
| | |
|
|
228
|
+
|---|---|
|
|
229
|
+
| **Device abstraction** | `torch.device`, per-device dispatch. Everything else waits on it. |
|
|
230
|
+
| **Metal** | candle already has the backend; disabled here for build isolation, not absent. |
|
|
231
|
+
| **`torch.distributed`** | From `world_size = 1` upward. Unblocks `transformers` as a side effect. |
|
|
232
|
+
| **Android GPU** | No candle backend and no `vulkan` slot in the `kernels` contract — genuinely new work. |
|
|
233
|
+
| **NPU** | ANE, NNAPI and QNN take a whole graph ahead of time, so this needs a capture layer rather than another device. The single door is where it attaches. |
|
|
234
|
+
|
|
235
|
+
---
|
|
236
|
+
|
|
237
|
+
## Install
|
|
238
|
+
|
|
239
|
+
> Not published yet. The alpha will carry platform wheels for macOS, Linux, Android and iOS —
|
|
240
|
+
> the extension is native, so `py3-none-any` is not the shape this ships in.
|
|
241
|
+
|
|
242
|
+
```sh
|
|
243
|
+
pip install torchnative
|
|
244
|
+
```
|
|
245
|
+
|
|
246
|
+
### Building from source
|
|
247
|
+
|
|
248
|
+
Requires a Rust toolchain and CPython 3.13+.
|
|
249
|
+
|
|
250
|
+
```sh
|
|
251
|
+
bash vendor/vendor_torch.sh # assemble the vendored torch tree
|
|
252
|
+
bash vendor/install_shim.sh # build the extension and install it
|
|
253
|
+
```
|
|
254
|
+
|
|
255
|
+
Cross-compilation is documented in [`docs/RUST_CROSSBUILD.md`](docs/RUST_CROSSBUILD.md),
|
|
256
|
+
including the PyO3 configuration iOS needs in order not to link `libpython`.
|
|
257
|
+
|
|
258
|
+
---
|
|
259
|
+
|
|
260
|
+
## Repository layout
|
|
261
|
+
|
|
262
|
+
```
|
|
263
|
+
torchnative/ the Python library
|
|
264
|
+
rust/torch_c/ the torch._C replacement (Rust · PyO3 · candle)
|
|
265
|
+
tools/golden/ the upstream comparison harness
|
|
266
|
+
vendor/ scripts that assemble the vendored torch tree (not checked in)
|
|
267
|
+
docs/ design, measurements, and the reasoning behind open decisions
|
|
268
|
+
```
|
|
269
|
+
|
|
270
|
+
`docs/` is written to be read. It records what was measured, what was assumed, and where an
|
|
271
|
+
earlier conclusion turned out to be wrong — corrections are left visible rather than edited away.
|
|
272
|
+
Start with [`DESIGN.md`](docs/DESIGN.md); [`SURFACE_HONESTY.md`](docs/SURFACE_HONESTY.md) and
|
|
273
|
+
[`HARNESS.md`](docs/HARNESS.md) show the standard the rest aims for.
|
|
274
|
+
|
|
275
|
+
---
|
|
276
|
+
|
|
277
|
+
## Related
|
|
278
|
+
|
|
279
|
+
- [PythonMultiplatform](https://github.com/thisisthepy/PythonMultiplatform) — embeds CPython 3.13
|
|
280
|
+
into Kotlin Multiplatform; the deployment target for this library
|
|
281
|
+
- [pypackpack](https://github.com/thisisthepy/pypackpack) — the build and bundling tool
|
|
282
|
+
- [Hugging Face `kernels`](https://github.com/huggingface/kernels) — the fused-kernel contract
|
|
283
|
+
this adopts, with resolution moved from runtime download to build time, since downloading
|
|
284
|
+
executable code is not permitted on every target platform
|
|
285
|
+
|
|
286
|
+
---
|
|
287
|
+
|
|
288
|
+
## License
|
|
289
|
+
|
|
290
|
+
MIT — see [LICENSE](LICENSE).
|
|
291
|
+
|
|
292
|
+
PyTorch is vendored under its own BSD-3-Clause license. The vendored tree is assembled at build
|
|
293
|
+
time and is not redistributed in this repository.
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
LICENSE
|
|
2
|
+
README.md
|
|
3
|
+
pyproject.toml
|
|
4
|
+
torchnative/src/main/torchnative/__init__.py
|
|
5
|
+
torchnative/src/main/torchnative.egg-info/PKG-INFO
|
|
6
|
+
torchnative/src/main/torchnative.egg-info/SOURCES.txt
|
|
7
|
+
torchnative/src/main/torchnative.egg-info/dependency_links.txt
|
|
8
|
+
torchnative/src/main/torchnative.egg-info/requires.txt
|
|
9
|
+
torchnative/src/main/torchnative.egg-info/top_level.txt
|
|
10
|
+
torchnative/src/main/torchnative/adapt/__init__.py
|
|
11
|
+
torchnative/src/main/torchnative/api/__init__.py
|
|
12
|
+
torchnative/src/main/torchnative/delta/__init__.py
|
|
13
|
+
torchnative/src/main/torchnative/kernels/__init__.py
|
|
14
|
+
torchnative/src/main/torchnative/nn/__init__.py
|
|
15
|
+
torchnative/src/main/torchnative/nn/federated/__init__.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
torchnative
|