ramnet 0.1.0rc0__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.
- ramnet-0.1.0rc0/LICENSE +21 -0
- ramnet-0.1.0rc0/PKG-INFO +261 -0
- ramnet-0.1.0rc0/README.md +242 -0
- ramnet-0.1.0rc0/pyproject.toml +30 -0
- ramnet-0.1.0rc0/setup.cfg +4 -0
- ramnet-0.1.0rc0/src/ramnet/__init__.py +11 -0
- ramnet-0.1.0rc0/src/ramnet/layers/__init__.py +4 -0
- ramnet-0.1.0rc0/src/ramnet/layers/cache.py +140 -0
- ramnet-0.1.0rc0/src/ramnet/layers/ram_attn_block.py +351 -0
- ramnet-0.1.0rc0/src/ramnet/models/__init__.py +10 -0
- ramnet-0.1.0rc0/src/ramnet/models/config_ramnet.py +188 -0
- ramnet-0.1.0rc0/src/ramnet/models/modeling_ramnet.py +368 -0
- ramnet-0.1.0rc0/src/ramnet.egg-info/PKG-INFO +261 -0
- ramnet-0.1.0rc0/src/ramnet.egg-info/SOURCES.txt +15 -0
- ramnet-0.1.0rc0/src/ramnet.egg-info/dependency_links.txt +1 -0
- ramnet-0.1.0rc0/src/ramnet.egg-info/requires.txt +5 -0
- ramnet-0.1.0rc0/src/ramnet.egg-info/top_level.txt +1 -0
ramnet-0.1.0rc0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025-2026 Kaicheng Xiao, Haotian Li, Liran Dong
|
|
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.
|
ramnet-0.1.0rc0/PKG-INFO
ADDED
|
@@ -0,0 +1,261 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: ramnet
|
|
3
|
+
Version: 0.1.0rc0
|
|
4
|
+
Summary: RAM-Net layers and Hugging Face models
|
|
5
|
+
Maintainer: Kaicheng Xiao, Haotian Li, Liran Dong
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Source, https://github.com/MuonCat/RAM-Net
|
|
8
|
+
Project-URL: Issues, https://github.com/MuonCat/RAM-Net/issues
|
|
9
|
+
Project-URL: Documentation, https://github.com/MuonCat/RAM-Net/tree/main/packages/ramnet
|
|
10
|
+
Requires-Python: >=3.10
|
|
11
|
+
Description-Content-Type: text/markdown
|
|
12
|
+
License-File: LICENSE
|
|
13
|
+
Requires-Dist: ramnet-ops>=0.1.0b0
|
|
14
|
+
Requires-Dist: torch>=2.5
|
|
15
|
+
Requires-Dist: einops
|
|
16
|
+
Requires-Dist: transformers>=4.57.3
|
|
17
|
+
Requires-Dist: flash-linear-attention>=0.3.2
|
|
18
|
+
Dynamic: license-file
|
|
19
|
+
|
|
20
|
+
<div align="center">
|
|
21
|
+
|
|
22
|
+
<img src="https://raw.githubusercontent.com/MuonCat/RAM-Net/main/assets/logo.png" alt="RAM-Net" width="420">
|
|
23
|
+
|
|
24
|
+
# ramnet
|
|
25
|
+
|
|
26
|
+
[](https://arxiv.org/abs/2602.11958)
|
|
27
|
+
[](https://muoncat.github.io/ramnet_web/)
|
|
28
|
+
[](https://huggingface.co/kaichengx/RAM-Net-Models/tree/main)
|
|
29
|
+
[](https://github.com/MuonCat/RAM-Net)
|
|
30
|
+
[](https://github.com/MuonCat/RAM-Net/blob/main/packages/ramnet/README.md)
|
|
31
|
+
[](https://github.com/MuonCat/RAM-Net/blob/main/packages/ramnet-ops/README.md)
|
|
32
|
+
[](https://github.com/MuonCat/RAM-Net/blob/main/LICENSE)
|
|
33
|
+
|
|
34
|
+
</div>
|
|
35
|
+
|
|
36
|
+
RAM-Net attention layers and Hugging Face compatible language models.
|
|
37
|
+
This pure Python package uses [ramnet-ops](https://github.com/MuonCat/RAM-Net/blob/main/packages/ramnet-ops/README.md) for CUDA execution.
|
|
38
|
+
See the [repository overview](https://github.com/MuonCat/RAM-Net#overview) for the method and experimental results.
|
|
39
|
+
|
|
40
|
+
## Structure
|
|
41
|
+
|
|
42
|
+
```text
|
|
43
|
+
.
|
|
44
|
+
`-- src/
|
|
45
|
+
`-- ramnet/
|
|
46
|
+
|-- __init__.py # Public exports
|
|
47
|
+
|-- layers/
|
|
48
|
+
| |-- ram_attn_block.py # Attention layer and explicit state API
|
|
49
|
+
| `-- cache.py # Recurrent state and Hugging Face cache integration
|
|
50
|
+
`-- models/
|
|
51
|
+
|-- __init__.py # Hugging Face Auto class registration
|
|
52
|
+
|-- config_ramnet.py # RAMNetConfig
|
|
53
|
+
`-- modeling_ramnet.py # Blocks, backbone, causal LM, and generation
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
Public exports: `RAMNetConfig`, `RAMNetModel`, `RAMNetForCausalLM`, `RAMNetAttention`, `RAMState`, and `RAMCache`.
|
|
57
|
+
Importing `ramnet` registers its Hugging Face Auto classes.
|
|
58
|
+
|
|
59
|
+
## Installation
|
|
60
|
+
|
|
61
|
+
Requires:
|
|
62
|
+
|
|
63
|
+
```text
|
|
64
|
+
Python >= 3.10
|
|
65
|
+
|
|
66
|
+
ramnet
|
|
67
|
+
|-- torch >= 2.5
|
|
68
|
+
|-- einops
|
|
69
|
+
|-- ramnet-ops >= 0.1.0b0
|
|
70
|
+
| `-- numpy
|
|
71
|
+
|-- transformers >= 4.57.3
|
|
72
|
+
`-- flash-linear-attention >= 0.3.2
|
|
73
|
+
```
|
|
74
|
+
|
|
75
|
+
The attention layer uses the CUDA backend, so a compatible NVIDIA GPU and CUDA-enabled `ramnet-ops` build are **required**.
|
|
76
|
+
|
|
77
|
+
### 1. Install dependencies
|
|
78
|
+
|
|
79
|
+
First install PyTorch following the [official installation guide](https://pytorch.org/get-started/locally/).
|
|
80
|
+
Choose PyTorch ≥ 2.10 with CUDA 13.0 for the default ops wheels, then install the remaining dependencies from TestPyPI / PyPI:
|
|
81
|
+
|
|
82
|
+
```bash
|
|
83
|
+
python -m pip install einops numpy "transformers>=4.57.3" "flash-linear-attention>=0.3.2"
|
|
84
|
+
```
|
|
85
|
+
|
|
86
|
+
### 2. Install ramnet-ops
|
|
87
|
+
|
|
88
|
+
Select a CUDA build matching your PyTorch version, GPU, and platform using the [ops installation guide](../ramnet-ops/README.md#installation).
|
|
89
|
+
For environments matching the [default stable ABI wheel requirements](../ramnet-ops/README.md#2-install-ramnet-ops), install directly from TestPyPI / PyPI:
|
|
90
|
+
|
|
91
|
+
```bash
|
|
92
|
+
python -m pip install -i https://test.pypi.org/simple/ ramnet-ops
|
|
93
|
+
```
|
|
94
|
+
|
|
95
|
+
### 3. Install ramnet
|
|
96
|
+
|
|
97
|
+
The model package is pure Python and needs no compiler.
|
|
98
|
+
Install it from TestPyPI / PyPI:
|
|
99
|
+
|
|
100
|
+
```bash
|
|
101
|
+
python -m pip install -i https://test.pypi.org/simple/ ramnet
|
|
102
|
+
```
|
|
103
|
+
|
|
104
|
+
#### From source
|
|
105
|
+
|
|
106
|
+
With the dependencies above installed, choose one source.
|
|
107
|
+
|
|
108
|
+
**Download with pip** — install the published source archive from TestPyPI / PyPI:
|
|
109
|
+
|
|
110
|
+
```bash
|
|
111
|
+
python -m pip install -i https://test.pypi.org/simple/ ramnet --no-binary=ramnet
|
|
112
|
+
```
|
|
113
|
+
|
|
114
|
+
**Clone from GitHub and install locally**:
|
|
115
|
+
|
|
116
|
+
```bash
|
|
117
|
+
git clone https://github.com/MuonCat/RAM-Net.git
|
|
118
|
+
cd RAM-Net
|
|
119
|
+
python -m pip install ./packages/ramnet
|
|
120
|
+
```
|
|
121
|
+
|
|
122
|
+
If you also need to build the CUDA operators, follow [ops source installation](../ramnet-ops/README.md#from-source).
|
|
123
|
+
|
|
124
|
+
## Quick start
|
|
125
|
+
|
|
126
|
+
`ramnet` provides the Hugging Face model interface, including Auto classes, language-modeling loss, and cached generation with `generate`.
|
|
127
|
+
The table maps paper notation to model configuration parameters.
|
|
128
|
+
|
|
129
|
+
| Paper quantity / symbol | Model parameter |
|
|
130
|
+
| :---: | :---: |
|
|
131
|
+
| Address order $U$ (digits per address) | `addr_order` |
|
|
132
|
+
| Address radix $d_p$ (values per digit) | `addr_radix` |
|
|
133
|
+
| Value dimension $d_v$ | `value_chs` |
|
|
134
|
+
| Slots per head $d_p^U$ | `addr_radix ** addr_order` (derived) |
|
|
135
|
+
| Selected slots $K$ | `addr_top_k` |
|
|
136
|
+
|
|
137
|
+
### Training and generation
|
|
138
|
+
|
|
139
|
+
This example uses a small model with random weights, matching [model_smoke.py](../../examples/model_smoke.py):
|
|
140
|
+
|
|
141
|
+
```python
|
|
142
|
+
import torch
|
|
143
|
+
from transformers import AutoModelForCausalLM
|
|
144
|
+
from ramnet import RAMNetConfig
|
|
145
|
+
|
|
146
|
+
config = RAMNetConfig(
|
|
147
|
+
hidden_size=128, num_attention_heads=4, num_hidden_layers=2,
|
|
148
|
+
addr_order=3, addr_radix=4,
|
|
149
|
+
embed_num=2, addr_top_k=8,
|
|
150
|
+
output_norm=True,
|
|
151
|
+
hidden_act="SwiGLU",
|
|
152
|
+
vocab_size=1024, pad_token_id=0,
|
|
153
|
+
dtype="float32",
|
|
154
|
+
)
|
|
155
|
+
model = AutoModelForCausalLM.from_config(config).cuda()
|
|
156
|
+
input_ids = torch.randint(3, config.vocab_size, (2, 32), device="cuda")
|
|
157
|
+
attention_mask = torch.ones_like(input_ids)
|
|
158
|
+
|
|
159
|
+
model.train()
|
|
160
|
+
output = model(
|
|
161
|
+
input_ids=input_ids, attention_mask=attention_mask,
|
|
162
|
+
labels=input_ids, use_cache=False,
|
|
163
|
+
)
|
|
164
|
+
output.loss.backward()
|
|
165
|
+
print(output.logits.shape) # (2, 32, 1024)
|
|
166
|
+
|
|
167
|
+
model.eval()
|
|
168
|
+
with torch.inference_mode():
|
|
169
|
+
tokens = model.generate(
|
|
170
|
+
input_ids=input_ids, attention_mask=attention_mask,
|
|
171
|
+
max_new_tokens=8, do_sample=False, use_cache=True,
|
|
172
|
+
)
|
|
173
|
+
print(tokens.shape)
|
|
174
|
+
```
|
|
175
|
+
|
|
176
|
+
Run the same example from the repository root with `python examples/model_smoke.py`.
|
|
177
|
+
|
|
178
|
+
`embed_num=2` enables **Cyclic Address Positional Embedding (CAPE)** on the first two heads of each layer, shifting read and write slot indices by token position.
|
|
179
|
+
The other heads address by content alone; `0` disables CAPE and `num_attention_heads` enables it for all heads.
|
|
180
|
+
An integer applies to every layer; a list or tuple gives per-layer counts, with unspecified layers using `0`.
|
|
181
|
+
Set `embed_num` explicitly when changing the model depth or head count, since the default schedule targets 25 layers and 32 heads.
|
|
182
|
+
|
|
183
|
+
### Save and load
|
|
184
|
+
|
|
185
|
+
Use the standard Hugging Face checkpoint methods after training:
|
|
186
|
+
|
|
187
|
+
```python
|
|
188
|
+
model.save_pretrained("ramnet-checkpoint")
|
|
189
|
+
model = AutoModelForCausalLM.from_pretrained("ramnet-checkpoint").cuda().eval()
|
|
190
|
+
```
|
|
191
|
+
|
|
192
|
+
Import `ramnet` before using Auto classes to load a checkpoint in a new process.
|
|
193
|
+
|
|
194
|
+
## Standalone attention layer
|
|
195
|
+
|
|
196
|
+
```python
|
|
197
|
+
import torch
|
|
198
|
+
from ramnet import RAMNetAttention
|
|
199
|
+
|
|
200
|
+
attn = RAMNetAttention(
|
|
201
|
+
hidden_size=128, n_heads=4, addr_order=3, addr_radix=4,
|
|
202
|
+
embed_num=2, top_k=8, output_norm=True,
|
|
203
|
+
)
|
|
204
|
+
attn._init_weights()
|
|
205
|
+
attn = attn.cuda()
|
|
206
|
+
x = torch.randn(2, 32, 128, device="cuda")
|
|
207
|
+
slot = torch.zeros(2, 4, 4 ** 3, 32, device="cuda")
|
|
208
|
+
mass = torch.full((2, 4, 4 ** 3), 1 / (4 ** 3), device="cuda")
|
|
209
|
+
y, state = attn.forward_with_state(x, init_ram=(slot, mass))
|
|
210
|
+
slot, mass = state.slot, state.mass
|
|
211
|
+
print(y.shape) # (2, 32, 128)
|
|
212
|
+
print(slot.shape) # (2, 4, 64, 32)
|
|
213
|
+
print(mass.shape) # (2, 4, 64)
|
|
214
|
+
```
|
|
215
|
+
|
|
216
|
+
Standalone layers need the explicit initialization above; full models apply it automatically.
|
|
217
|
+
The standard attention `forward` accepts `(B, T, D)` inputs and returns `(output, None, past_key_values)`, following the FLA layer convention.
|
|
218
|
+
Set `layer_idx` when constructing attention that uses `RAMCache`.
|
|
219
|
+
|
|
220
|
+
Both attention and complete model blocks provide `forward_with_state`:
|
|
221
|
+
|
|
222
|
+
- `attn.forward_with_state(..., init_ram=state)` runs the attention layer.
|
|
223
|
+
- `model.model.layers[i].forward_with_state(..., init_state=state)` runs the full block, including normalization, MLP, and residual connections.
|
|
224
|
+
|
|
225
|
+
Both accept a `(slot, mass)` pair or `RAMState(slot, mass)` and return `(output, RAMState)`; these APIs also expose `t_dim` and `inplace`.
|
|
226
|
+
The default `inplace=False` preserves input states and supports differentiation through consecutive calls.
|
|
227
|
+
Explicit-state callers supply continuation `position_ids` for CAPE.
|
|
228
|
+
|
|
229
|
+
## Configuration
|
|
230
|
+
|
|
231
|
+
Blocks use pre-norm RMSNorm; `fuse_norm` selects its fused implementation.
|
|
232
|
+
`hidden_act="fuse_SwiGLU"` (default) or `"SwiGLU"` selects FLA's `GatedMLP`, with or without fusion.
|
|
233
|
+
`"GELU"`, `"SiLU"`, and `"Swish"` select a standard two-layer MLP; `"SiLU"` and `"Swish"` use the same activation.
|
|
234
|
+
Models support `input_ids` or `inputs_embeds`, `return_dict`, hidden-state outputs, and `logits_to_keep`.
|
|
235
|
+
|
|
236
|
+
- `addr_top_k` sets the number of selected slots for both writing and reading; a `(write_k, read_k)` pair specifies separate counts.
|
|
237
|
+
- `addr_top_k_re_norm` explicitly enables renormalization after Top-K: a bool for both sides, or a two-item tuple/list of `(write, read)` bools; defaults to `False`.
|
|
238
|
+
- `output_norm` enables head-wise RMSNorm; `mass_aware` enables mass-weighted reads.
|
|
239
|
+
|
|
240
|
+
See [RAMNetConfig](src/ramnet/models/config_ramnet.py) for the full parameter list and defaults.
|
|
241
|
+
|
|
242
|
+
## Cache and sequence handling
|
|
243
|
+
|
|
244
|
+
`config.use_cache` controls the evaluation default; training disables caching unless explicitly requested.
|
|
245
|
+
|
|
246
|
+
- `use_cache=True` carries recurrent state during generation.
|
|
247
|
+
The attention layer automatically uses inplace updates for initialized, gradient-free inference caches.
|
|
248
|
+
- `RAMCache()` allocates states on first use and stores per-layer `recurrent_state` entries.
|
|
249
|
+
It supports beam reordering, batch selection/repetition, and reset.
|
|
250
|
+
Logical CAPE positions are tracked per sample, so a full accumulated mask or a query-only mask can be used after padded prefill.
|
|
251
|
+
Explicit `position_ids` override automatic positions.
|
|
252
|
+
- `use_cache=False` reads a supplied cache without modifying or returning it.
|
|
253
|
+
Gradient checkpointing requires an uncached sequence and disables cache updates during training.
|
|
254
|
+
- Cache rollback, speculative decoding, static/offloaded KV-cache implementations, and full-graph compiled generation are not supported.
|
|
255
|
+
- Packed `cu_seqlens` inputs and attention-matrix outputs (`output_attentions=True`) are not supported.
|
|
256
|
+
|
|
257
|
+
## Development
|
|
258
|
+
|
|
259
|
+
See [RELEASING](RELEASING.md) for tests and publishing.
|
|
260
|
+
|
|
261
|
+
MIT License. See [LICENSE](LICENSE).
|
|
@@ -0,0 +1,242 @@
|
|
|
1
|
+
<div align="center">
|
|
2
|
+
|
|
3
|
+
<img src="https://raw.githubusercontent.com/MuonCat/RAM-Net/main/assets/logo.png" alt="RAM-Net" width="420">
|
|
4
|
+
|
|
5
|
+
# ramnet
|
|
6
|
+
|
|
7
|
+
[](https://arxiv.org/abs/2602.11958)
|
|
8
|
+
[](https://muoncat.github.io/ramnet_web/)
|
|
9
|
+
[](https://huggingface.co/kaichengx/RAM-Net-Models/tree/main)
|
|
10
|
+
[](https://github.com/MuonCat/RAM-Net)
|
|
11
|
+
[](https://github.com/MuonCat/RAM-Net/blob/main/packages/ramnet/README.md)
|
|
12
|
+
[](https://github.com/MuonCat/RAM-Net/blob/main/packages/ramnet-ops/README.md)
|
|
13
|
+
[](https://github.com/MuonCat/RAM-Net/blob/main/LICENSE)
|
|
14
|
+
|
|
15
|
+
</div>
|
|
16
|
+
|
|
17
|
+
RAM-Net attention layers and Hugging Face compatible language models.
|
|
18
|
+
This pure Python package uses [ramnet-ops](https://github.com/MuonCat/RAM-Net/blob/main/packages/ramnet-ops/README.md) for CUDA execution.
|
|
19
|
+
See the [repository overview](https://github.com/MuonCat/RAM-Net#overview) for the method and experimental results.
|
|
20
|
+
|
|
21
|
+
## Structure
|
|
22
|
+
|
|
23
|
+
```text
|
|
24
|
+
.
|
|
25
|
+
`-- src/
|
|
26
|
+
`-- ramnet/
|
|
27
|
+
|-- __init__.py # Public exports
|
|
28
|
+
|-- layers/
|
|
29
|
+
| |-- ram_attn_block.py # Attention layer and explicit state API
|
|
30
|
+
| `-- cache.py # Recurrent state and Hugging Face cache integration
|
|
31
|
+
`-- models/
|
|
32
|
+
|-- __init__.py # Hugging Face Auto class registration
|
|
33
|
+
|-- config_ramnet.py # RAMNetConfig
|
|
34
|
+
`-- modeling_ramnet.py # Blocks, backbone, causal LM, and generation
|
|
35
|
+
```
|
|
36
|
+
|
|
37
|
+
Public exports: `RAMNetConfig`, `RAMNetModel`, `RAMNetForCausalLM`, `RAMNetAttention`, `RAMState`, and `RAMCache`.
|
|
38
|
+
Importing `ramnet` registers its Hugging Face Auto classes.
|
|
39
|
+
|
|
40
|
+
## Installation
|
|
41
|
+
|
|
42
|
+
Requires:
|
|
43
|
+
|
|
44
|
+
```text
|
|
45
|
+
Python >= 3.10
|
|
46
|
+
|
|
47
|
+
ramnet
|
|
48
|
+
|-- torch >= 2.5
|
|
49
|
+
|-- einops
|
|
50
|
+
|-- ramnet-ops >= 0.1.0b0
|
|
51
|
+
| `-- numpy
|
|
52
|
+
|-- transformers >= 4.57.3
|
|
53
|
+
`-- flash-linear-attention >= 0.3.2
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
The attention layer uses the CUDA backend, so a compatible NVIDIA GPU and CUDA-enabled `ramnet-ops` build are **required**.
|
|
57
|
+
|
|
58
|
+
### 1. Install dependencies
|
|
59
|
+
|
|
60
|
+
First install PyTorch following the [official installation guide](https://pytorch.org/get-started/locally/).
|
|
61
|
+
Choose PyTorch ≥ 2.10 with CUDA 13.0 for the default ops wheels, then install the remaining dependencies from TestPyPI / PyPI:
|
|
62
|
+
|
|
63
|
+
```bash
|
|
64
|
+
python -m pip install einops numpy "transformers>=4.57.3" "flash-linear-attention>=0.3.2"
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
### 2. Install ramnet-ops
|
|
68
|
+
|
|
69
|
+
Select a CUDA build matching your PyTorch version, GPU, and platform using the [ops installation guide](../ramnet-ops/README.md#installation).
|
|
70
|
+
For environments matching the [default stable ABI wheel requirements](../ramnet-ops/README.md#2-install-ramnet-ops), install directly from TestPyPI / PyPI:
|
|
71
|
+
|
|
72
|
+
```bash
|
|
73
|
+
python -m pip install -i https://test.pypi.org/simple/ ramnet-ops
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
### 3. Install ramnet
|
|
77
|
+
|
|
78
|
+
The model package is pure Python and needs no compiler.
|
|
79
|
+
Install it from TestPyPI / PyPI:
|
|
80
|
+
|
|
81
|
+
```bash
|
|
82
|
+
python -m pip install -i https://test.pypi.org/simple/ ramnet
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
#### From source
|
|
86
|
+
|
|
87
|
+
With the dependencies above installed, choose one source.
|
|
88
|
+
|
|
89
|
+
**Download with pip** — install the published source archive from TestPyPI / PyPI:
|
|
90
|
+
|
|
91
|
+
```bash
|
|
92
|
+
python -m pip install -i https://test.pypi.org/simple/ ramnet --no-binary=ramnet
|
|
93
|
+
```
|
|
94
|
+
|
|
95
|
+
**Clone from GitHub and install locally**:
|
|
96
|
+
|
|
97
|
+
```bash
|
|
98
|
+
git clone https://github.com/MuonCat/RAM-Net.git
|
|
99
|
+
cd RAM-Net
|
|
100
|
+
python -m pip install ./packages/ramnet
|
|
101
|
+
```
|
|
102
|
+
|
|
103
|
+
If you also need to build the CUDA operators, follow [ops source installation](../ramnet-ops/README.md#from-source).
|
|
104
|
+
|
|
105
|
+
## Quick start
|
|
106
|
+
|
|
107
|
+
`ramnet` provides the Hugging Face model interface, including Auto classes, language-modeling loss, and cached generation with `generate`.
|
|
108
|
+
The table maps paper notation to model configuration parameters.
|
|
109
|
+
|
|
110
|
+
| Paper quantity / symbol | Model parameter |
|
|
111
|
+
| :---: | :---: |
|
|
112
|
+
| Address order $U$ (digits per address) | `addr_order` |
|
|
113
|
+
| Address radix $d_p$ (values per digit) | `addr_radix` |
|
|
114
|
+
| Value dimension $d_v$ | `value_chs` |
|
|
115
|
+
| Slots per head $d_p^U$ | `addr_radix ** addr_order` (derived) |
|
|
116
|
+
| Selected slots $K$ | `addr_top_k` |
|
|
117
|
+
|
|
118
|
+
### Training and generation
|
|
119
|
+
|
|
120
|
+
This example uses a small model with random weights, matching [model_smoke.py](../../examples/model_smoke.py):
|
|
121
|
+
|
|
122
|
+
```python
|
|
123
|
+
import torch
|
|
124
|
+
from transformers import AutoModelForCausalLM
|
|
125
|
+
from ramnet import RAMNetConfig
|
|
126
|
+
|
|
127
|
+
config = RAMNetConfig(
|
|
128
|
+
hidden_size=128, num_attention_heads=4, num_hidden_layers=2,
|
|
129
|
+
addr_order=3, addr_radix=4,
|
|
130
|
+
embed_num=2, addr_top_k=8,
|
|
131
|
+
output_norm=True,
|
|
132
|
+
hidden_act="SwiGLU",
|
|
133
|
+
vocab_size=1024, pad_token_id=0,
|
|
134
|
+
dtype="float32",
|
|
135
|
+
)
|
|
136
|
+
model = AutoModelForCausalLM.from_config(config).cuda()
|
|
137
|
+
input_ids = torch.randint(3, config.vocab_size, (2, 32), device="cuda")
|
|
138
|
+
attention_mask = torch.ones_like(input_ids)
|
|
139
|
+
|
|
140
|
+
model.train()
|
|
141
|
+
output = model(
|
|
142
|
+
input_ids=input_ids, attention_mask=attention_mask,
|
|
143
|
+
labels=input_ids, use_cache=False,
|
|
144
|
+
)
|
|
145
|
+
output.loss.backward()
|
|
146
|
+
print(output.logits.shape) # (2, 32, 1024)
|
|
147
|
+
|
|
148
|
+
model.eval()
|
|
149
|
+
with torch.inference_mode():
|
|
150
|
+
tokens = model.generate(
|
|
151
|
+
input_ids=input_ids, attention_mask=attention_mask,
|
|
152
|
+
max_new_tokens=8, do_sample=False, use_cache=True,
|
|
153
|
+
)
|
|
154
|
+
print(tokens.shape)
|
|
155
|
+
```
|
|
156
|
+
|
|
157
|
+
Run the same example from the repository root with `python examples/model_smoke.py`.
|
|
158
|
+
|
|
159
|
+
`embed_num=2` enables **Cyclic Address Positional Embedding (CAPE)** on the first two heads of each layer, shifting read and write slot indices by token position.
|
|
160
|
+
The other heads address by content alone; `0` disables CAPE and `num_attention_heads` enables it for all heads.
|
|
161
|
+
An integer applies to every layer; a list or tuple gives per-layer counts, with unspecified layers using `0`.
|
|
162
|
+
Set `embed_num` explicitly when changing the model depth or head count, since the default schedule targets 25 layers and 32 heads.
|
|
163
|
+
|
|
164
|
+
### Save and load
|
|
165
|
+
|
|
166
|
+
Use the standard Hugging Face checkpoint methods after training:
|
|
167
|
+
|
|
168
|
+
```python
|
|
169
|
+
model.save_pretrained("ramnet-checkpoint")
|
|
170
|
+
model = AutoModelForCausalLM.from_pretrained("ramnet-checkpoint").cuda().eval()
|
|
171
|
+
```
|
|
172
|
+
|
|
173
|
+
Import `ramnet` before using Auto classes to load a checkpoint in a new process.
|
|
174
|
+
|
|
175
|
+
## Standalone attention layer
|
|
176
|
+
|
|
177
|
+
```python
|
|
178
|
+
import torch
|
|
179
|
+
from ramnet import RAMNetAttention
|
|
180
|
+
|
|
181
|
+
attn = RAMNetAttention(
|
|
182
|
+
hidden_size=128, n_heads=4, addr_order=3, addr_radix=4,
|
|
183
|
+
embed_num=2, top_k=8, output_norm=True,
|
|
184
|
+
)
|
|
185
|
+
attn._init_weights()
|
|
186
|
+
attn = attn.cuda()
|
|
187
|
+
x = torch.randn(2, 32, 128, device="cuda")
|
|
188
|
+
slot = torch.zeros(2, 4, 4 ** 3, 32, device="cuda")
|
|
189
|
+
mass = torch.full((2, 4, 4 ** 3), 1 / (4 ** 3), device="cuda")
|
|
190
|
+
y, state = attn.forward_with_state(x, init_ram=(slot, mass))
|
|
191
|
+
slot, mass = state.slot, state.mass
|
|
192
|
+
print(y.shape) # (2, 32, 128)
|
|
193
|
+
print(slot.shape) # (2, 4, 64, 32)
|
|
194
|
+
print(mass.shape) # (2, 4, 64)
|
|
195
|
+
```
|
|
196
|
+
|
|
197
|
+
Standalone layers need the explicit initialization above; full models apply it automatically.
|
|
198
|
+
The standard attention `forward` accepts `(B, T, D)` inputs and returns `(output, None, past_key_values)`, following the FLA layer convention.
|
|
199
|
+
Set `layer_idx` when constructing attention that uses `RAMCache`.
|
|
200
|
+
|
|
201
|
+
Both attention and complete model blocks provide `forward_with_state`:
|
|
202
|
+
|
|
203
|
+
- `attn.forward_with_state(..., init_ram=state)` runs the attention layer.
|
|
204
|
+
- `model.model.layers[i].forward_with_state(..., init_state=state)` runs the full block, including normalization, MLP, and residual connections.
|
|
205
|
+
|
|
206
|
+
Both accept a `(slot, mass)` pair or `RAMState(slot, mass)` and return `(output, RAMState)`; these APIs also expose `t_dim` and `inplace`.
|
|
207
|
+
The default `inplace=False` preserves input states and supports differentiation through consecutive calls.
|
|
208
|
+
Explicit-state callers supply continuation `position_ids` for CAPE.
|
|
209
|
+
|
|
210
|
+
## Configuration
|
|
211
|
+
|
|
212
|
+
Blocks use pre-norm RMSNorm; `fuse_norm` selects its fused implementation.
|
|
213
|
+
`hidden_act="fuse_SwiGLU"` (default) or `"SwiGLU"` selects FLA's `GatedMLP`, with or without fusion.
|
|
214
|
+
`"GELU"`, `"SiLU"`, and `"Swish"` select a standard two-layer MLP; `"SiLU"` and `"Swish"` use the same activation.
|
|
215
|
+
Models support `input_ids` or `inputs_embeds`, `return_dict`, hidden-state outputs, and `logits_to_keep`.
|
|
216
|
+
|
|
217
|
+
- `addr_top_k` sets the number of selected slots for both writing and reading; a `(write_k, read_k)` pair specifies separate counts.
|
|
218
|
+
- `addr_top_k_re_norm` explicitly enables renormalization after Top-K: a bool for both sides, or a two-item tuple/list of `(write, read)` bools; defaults to `False`.
|
|
219
|
+
- `output_norm` enables head-wise RMSNorm; `mass_aware` enables mass-weighted reads.
|
|
220
|
+
|
|
221
|
+
See [RAMNetConfig](src/ramnet/models/config_ramnet.py) for the full parameter list and defaults.
|
|
222
|
+
|
|
223
|
+
## Cache and sequence handling
|
|
224
|
+
|
|
225
|
+
`config.use_cache` controls the evaluation default; training disables caching unless explicitly requested.
|
|
226
|
+
|
|
227
|
+
- `use_cache=True` carries recurrent state during generation.
|
|
228
|
+
The attention layer automatically uses inplace updates for initialized, gradient-free inference caches.
|
|
229
|
+
- `RAMCache()` allocates states on first use and stores per-layer `recurrent_state` entries.
|
|
230
|
+
It supports beam reordering, batch selection/repetition, and reset.
|
|
231
|
+
Logical CAPE positions are tracked per sample, so a full accumulated mask or a query-only mask can be used after padded prefill.
|
|
232
|
+
Explicit `position_ids` override automatic positions.
|
|
233
|
+
- `use_cache=False` reads a supplied cache without modifying or returning it.
|
|
234
|
+
Gradient checkpointing requires an uncached sequence and disables cache updates during training.
|
|
235
|
+
- Cache rollback, speculative decoding, static/offloaded KV-cache implementations, and full-graph compiled generation are not supported.
|
|
236
|
+
- Packed `cu_seqlens` inputs and attention-matrix outputs (`output_attentions=True`) are not supported.
|
|
237
|
+
|
|
238
|
+
## Development
|
|
239
|
+
|
|
240
|
+
See [RELEASING](RELEASING.md) for tests and publishing.
|
|
241
|
+
|
|
242
|
+
MIT License. See [LICENSE](LICENSE).
|
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=77.0.3", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "ramnet"
|
|
7
|
+
version = "0.1.0rc0"
|
|
8
|
+
description = "RAM-Net layers and Hugging Face models"
|
|
9
|
+
maintainers = [{name = "Kaicheng Xiao"}, {name = "Haotian Li"}, {name = "Liran Dong"}]
|
|
10
|
+
readme = "README.md"
|
|
11
|
+
license = "MIT"
|
|
12
|
+
license-files = ["LICENSE"]
|
|
13
|
+
requires-python = ">=3.10"
|
|
14
|
+
dependencies = [
|
|
15
|
+
"ramnet-ops>=0.1.0b0",
|
|
16
|
+
"torch>=2.5",
|
|
17
|
+
"einops",
|
|
18
|
+
"transformers>=4.57.3",
|
|
19
|
+
"flash-linear-attention>=0.3.2",
|
|
20
|
+
]
|
|
21
|
+
|
|
22
|
+
[project.urls]
|
|
23
|
+
Source = "https://github.com/MuonCat/RAM-Net"
|
|
24
|
+
Issues = "https://github.com/MuonCat/RAM-Net/issues"
|
|
25
|
+
Documentation = "https://github.com/MuonCat/RAM-Net/tree/main/packages/ramnet"
|
|
26
|
+
|
|
27
|
+
[tool.setuptools.packages.find]
|
|
28
|
+
where = ["src"]
|
|
29
|
+
include = ["ramnet", "ramnet.*"]
|
|
30
|
+
namespaces = false
|