Metadata-Version: 2.4
Name: dabsn
Version: 0.1.2
Summary: DABSN recurrent modeling framework with native CPU and CUDA runtimes
Author: Nicholas Rosdahl
License-Expression: Apache-2.0
Project-URL: Homepage, https://github.com/BleedingXiko/dabsn
Project-URL: Documentation, https://github.com/BleedingXiko/dabsn#readme
Project-URL: Paper, https://github.com/BleedingXiko/dabsn/blob/main/paper1/main.pdf
Project-URL: Source, https://github.com/BleedingXiko/dabsn
Project-URL: Issues, https://github.com/BleedingXiko/dabsn/issues
Keywords: deep-learning,neural-network,recurrent-neural-network,sequence-modeling,pytorch,triton
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.6
Requires-Dist: numpy>=1.24
Requires-Dist: ninja>=1.11
Requires-Dist: safetensors>=0.4.5
Provides-Extra: test
Requires-Dist: pytest>=8; extra == "test"
Provides-Extra: cuda-turing
Requires-Dist: torch==2.6.0; platform_system == "Linux" and extra == "cuda-turing"
Requires-Dist: triton==3.2.0; platform_system == "Linux" and extra == "cuda-turing"
Provides-Extra: paper
Requires-Dist: matplotlib>=3.8; extra == "paper"
Dynamic: license-file

# DABSN

**A recurrent modeling framework for persistent state and admitted memory.**

DABSN is a PyTorch architecture for causal sequences, whole fields, and
structured data that mixes both. Its core maintains a nonlinear recurrent
state while its read system combines admitted short memory, successor
induction, permanent associative memory, and a recurrent long-memory channel.
The same block design is used across all three geometries.

The framework includes native C++/OpenMP CPU kernels and Triton/CUDA kernels
for forward and backward execution, task-owned input and output adapters,
structured checkpoints, training and inference helpers, and the complete
source and result tables for the accompanying paper. Multi-GPU training uses
PyTorch DDP or block-wrapped FSDP with full parameter, gradient, and optimizer
sharding.

## Paper

### One Layer, Both Gaps

**A Persistent-Modulation Recurrence that Generalizes Copy and Tracks
Non-Solvable Group State**

The paper tests one architecture, trained separately per task, against two
regimes commonly treated as opposing requirements:

| Task | Train length | Evaluation length | DABSN result |
| --- | ---: | ---: | ---: |
| Copy, vocabulary 64 | 64 | 3,200 (50x) | 0.961 +/- 0.035, three seeds |
| A5/60 word problem | 256 | 16,384 (64x) | 1.000, two seeds |

These are separately trained models, not one checkpoint reused across tasks.
The paper includes causal ablations of the nonlinear state and read pathways;
the machine-readable tables used for every reported result are included with
the source.

- [Read the paper](https://github.com/BleedingXiko/dabsn/blob/main/paper1/main.pdf)
- [Inspect the paper source and result tables](https://github.com/BleedingXiko/dabsn/tree/main/paper1)
- [Read the architecture contract](https://github.com/BleedingXiko/dabsn/blob/main/ARCHITECTURE.md)

## Installation

Install DABSN into an environment containing the PyTorch build appropriate for
your machine:

```bash
pip install dabsn
```

Turing GPUs such as the GTX 1660 Ti use the final compatible Torch/Triton
combination:

```bash
pip install 'dabsn[cuda-turing]'
```

Python 3.10 or newer and PyTorch 2.6 or newer are required. Linux CUDA builds
of PyTorch provide the matching Triton runtime. Native backend
selection is explicit: `required=True` raises instead of silently switching to
another runtime family.

## Quick start

```python
import torch

from dabsn import DABSNLayerSpec, DABSNModel, dabsn_adamw_param_groups
from dabsn.kernels import enable, status
from dabsn.runtime import train_step

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
enable(device.type, required=True)

model = DABSNModel(
    input_dim=24,
    out_dim=10,
    layers=[
        DABSNLayerSpec(128, 96, "seq"),
        DABSNLayerSpec(192, 128, "seq"),
    ],
    output_adapter="token",
).to(device)

inputs = torch.randn(8, 256, 24, device=device)
targets = torch.randint(0, 10, (8, 256), device=device)
optimizer = torch.optim.AdamW(
    dabsn_adamw_param_groups(model, weight_decay=0.1),
    lr=1e-3,
)

loss = train_step(
    model,
    inputs,
    targets,
    optimizer,
    clip_grad_norm=1.0,
)

print({"loss": loss, "backend": status()["active_backend"]})
```

Every model is a normal `torch.nn.Module`. Use the supplied runtime helpers or
an ordinary PyTorch training loop:
[`examples/minimal_train.py`](https://github.com/BleedingXiko/dabsn/blob/main/examples/minimal_train.py)
is the same step written with no DABSN helpers and no native backend.

## Model geometry

Each layer owns an output width, recurrent-state width, and read geometry.
Widths may change across a stack.

| Geometry | Memory eligibility | Typical structure |
| --- | --- | --- |
| `seq` | causal prefix | language, events, control streams |
| `field` | whole object | images, boards, sets, spatial state |
| `hybrid` | learned sequence/field mixture | structured streams with both relations |

Layer stacks can be written directly or parsed from compact specifications:

```python
from dabsn import parse_dabsn_layer_specs

layers = parse_dabsn_layer_specs(
    "seq:128:96,hybrid:192:128,field:128:96"
)
```

The outer model API is the same for every geometry. Geometry changes memory
eligibility, not the recurrent block or checkpoint format.

## Task adapters

DABSN owns the recurrent body; applications own the meaning of their data.
Input adapters transform raw task records into model-width features, and
output heads transform hidden states into task predictions. Registered
adapters become construction and checkpoint metadata rather than notebook-only
glue.

This example handles industrial telemetry with continuous measurements,
elapsed time, sensor identity, and missingness. Its output jointly predicts an
event class and a log-normal time-to-event distribution.

```python
import torch
import torch.nn as nn
import torch.nn.functional as F

from dabsn import DABSNLayerSpec, DABSNTaskModel
from dabsn.adapters import register_input_adapter, register_output_head


class TelemetryInput(nn.Module):
    def __init__(self, raw_dim: int, model_dim: int, sensors: int = 32):
        super().__init__()
        if raw_dim != 8:
            raise ValueError("expected five values, elapsed time, sensor ID, and mask")
        self.output_dim = model_dim
        self.sensors = sensors
        self.value_norm = nn.LayerNorm(5)
        self.sensor = nn.Embedding(sensors, 12)
        self.missing = nn.Embedding(2, 4)
        self.register_buffer("frequencies", torch.tensor([1., 2., 4., 8.]))
        self.fuse = nn.Sequential(
            nn.Linear(5 + 8 + 12 + 4, model_dim * 2),
            nn.SiLU(),
            nn.Linear(model_dim * 2, model_dim),
            nn.LayerNorm(model_dim),
        )

    def forward(self, x):
        values = self.value_norm(torch.nan_to_num(x[..., :5].float()))
        elapsed = x[..., 5].float().clamp_min(0)
        sensor = x[..., 6].long().clamp(0, self.sensors - 1)
        missing = x[..., 7].long().clamp(0, 1)
        phase = torch.log1p(elapsed).unsqueeze(-1) * self.frequencies
        time = torch.cat([phase.sin(), phase.cos()], dim=-1)
        return self.fuse(torch.cat([
            values, time, self.sensor(sensor), self.missing(missing)
        ], dim=-1))


class EventForecast(nn.Module):
    def __init__(self, hidden_dim: int, out_dim: int):
        super().__init__()
        self.classes = out_dim - 2
        self.norm = nn.LayerNorm(hidden_dim)
        self.event_logits = nn.Linear(hidden_dim, self.classes)
        self.time_parameters = nn.Linear(hidden_dim, 2)

    def forward(self, hidden):
        hidden = self.norm(hidden)
        return torch.cat([
            self.event_logits(hidden), self.time_parameters(hidden)
        ], dim=-1)

    def unpack(self, output):
        logits = output[..., :self.classes]
        log_time_mean = output[..., -2]
        log_time_scale = F.softplus(output[..., -1]) + 1e-4
        return logits, log_time_mean, log_time_scale


register_input_adapter(
    "telemetry",
    lambda raw_dim, model_dim: TelemetryInput(raw_dim, model_dim or raw_dim),
)
register_output_head("event_forecast", EventForecast)

model = DABSNTaskModel(
    raw_input_dim=8,
    model_input_dim=96,
    out_dim=6,  # four event classes plus two distribution parameters
    layers=[
        DABSNLayerSpec(96, 64, "seq"),
        DABSNLayerSpec(128, 96, "seq"),
    ],
    input_adapter="telemetry",
    output_adapter="event_forecast",
)
```

The [complete telemetry example](https://github.com/BleedingXiko/dabsn/blob/main/examples/custom_adapter.py)
includes synthetic data, the joint classification/distribution loss, and an
optimizer step. A separate [local 2D field adapter](https://github.com/BleedingXiko/dabsn/blob/main/examples/field2d_local_adapter.py)
demonstrates native neighborhood gather/scatter for spatial models.

## Native runtimes

```python
from dabsn.kernels import enable, status

enable("cuda", required=True)   # CUDA/Triton + batched-GEMM forward/backward
# enable("cpu", required=True)  # C++/OpenMP forward and backward
# enable("reference")           # explicit PyTorch reference runtime

print(status())
```

Backend activation is process-wide because it installs model dispatch hooks.
Requested native execution never silently falls back. The status report names
the active implementation for the core scan, admitted read, permanent memory,
long-memory recurrence, and local-field gather.

CUDA training dispatch is execution-shape aware. Small batches use the
persistent Triton scan. Batches of 64 or more use the batched recurrent runtime,
which shares each recurrent-matrix read across the device batch through GEMM;
this changes neither the DABSN equations nor model depth. For modest training
score tensors (up to 8,388,608 `[B,T,N]` entries by default), admitted-read
forward and backward use native BMM instead of pairing a tiled forward with the
older serial-query backward. The controls are explicit when a benchmark needs
to pin them:

```bash
export DABSN_CORE_BACKEND=batched       # auto | batched | persistent | batched_fused
export DABSN_BATCHED_STEP_COMPILE=1     # compile only the pure pointwise step
export DABSN_TRAIN_DENSE_MAX_SCORES=8388608
```

The complete DABSN model or backbone is never compiled by this dispatch. The
batched custom-autograd recurrence has a separately tested explicit backward,
including every parameter and carried-state gradient.

The release gates cover:

- `seq`, `field`, and `hybrid` model forward/backward parity;
- single-block and stacked execution;
- recurrent execution with and without an explicit initial core state;
- gradients through inputs, parameters, and carried state;
- admitted, permanent, long-memory, and local-field primitives;
- configuration-aware checkpoint reload.

The repository does not claim that its fused kernels outperform every existing
sequence runtime. Their contract is native DABSN execution with explicit
forward/backward parity and no hidden backend switch.

## Performance and scaling

Most of the throughput-relevant behavior is automatic once a native runtime is
enabled; nothing below changes the DABSN equations. All of it is geometry
agnostic (`seq`, `field`, `hybrid`) because it lives at the core-scan and
admitted-read level, not in any task head.

**Automatic (no configuration):**

- **Sub-quadratic admitted read.** The read scores each query position against
  the *admitted* bank, whose width is the data-dependent admitted count, so the
  cost is `O(T * admitted)` — not `O(T^2)`. The width is sized dynamically from
  the learned admission; it only approaches `seq_len` (quadratic) if the model
  genuinely learns to admit almost every position, which is the correct cost for
  a task that needs it. Inference and ordinary GPU training both use this
  dynamic width. A static width is used only while a CUDA graph is actively being
  captured (where a host sync is illegal), and even then the capture path pins a
  measured, padded, still-sub-quadratic cap.
- **Work-aware core dispatch.** `select_core_backend` picks the persistent Triton
  scan for small work and the batched tensor-core GEMM scan once the batch is
  large enough to fill the device (`B >= 64` or `B*H >=
  DABSN_BATCHED_CORE_MIN_WORK`, default 4096), so wide/large-batch training uses
  tensor cores automatically.
- **Tensor-core compute dtype.** With a BF16/FP16 model the recurrent GEMMs run
  on tensor cores; pointwise state stays FP32.

**Opt-in:**

- **Fused single-launch core scan** (`DABSN_CORE_BACKEND=batched_fused`, hidden
  width `<= 256`). Runs the whole `T`-step scan in one Triton launch with state
  carried in registers, removing the per-step launch overhead. Wider cores use
  the batched per-step GEMM, which has no such width bound. `auto` never selects
  the fused backend on its own — request it explicitly.
- **CUDA-graph training** (`make_graphed_train_callable`, or `cuda_graph=True` in
  `DABSNPretrainConfig`). Captures the forward+backward once and replays it,
  removing kernel-launch overhead — the dominant cost of the sequential scan at
  small microbatches. Single-process CUDA only; pair each replay with
  `ManualGradientAccumulator` for exact microbatch accumulation. Capture failure
  raises rather than silently degrading.

```python
from dabsn.runtime import make_graphed_train_callable, ManualGradientAccumulator
```

**Batch vs. context.** The core is a *sequential* recurrence: it advances one
position at a time and cannot parallelize across context the way attention does.
Its device parallelism therefore comes from the **batch**, not the sequence
length — a tiny microbatch leaves the GPU idle on every step regardless of
context. Raise the microbatch as high as memory allows and use gradient
accumulation for the effective batch. Because the read is sub-quadratic in `T`,
context length scales close to linearly, so long-context training is bounded by
the (linear) number of scan steps rather than a quadratic read.

## Gradient preflight

Before a long training run, verify the complete model stack:

```python
from dabsn.runtime import verify_gradients

rows = verify_gradients(model, sample_input, compile_forward=True)
print(rows)
```

This compiles the outer forward boundary, runs one backward pass, and raises if
any block has missing, zero, or non-finite representative gradients.

## Distributed training

Launch two or more CUDA workers with `torchrun` and select FSDP explicitly:

```bash
torchrun --standalone --nproc-per-node=2 -m dabsn.cli train \
  --config model.json \
  --data batch.pt \
  --output run/model.safetensors \
  --device cuda \
  --backend cuda \
  --distributed fsdp \
  --precision bf16 \
  --grad-checkpoint \
  --grad-accum-steps 4 \
  --verify-gradients
```

Use `--precision fp16` on Turing GPUs such as the T4 or GTX 1660 Ti. The input
file contains one global batch; its first dimension must be divisible by the
number of workers. Each rank receives a distinct batch shard. FSDP uses
`FULL_SHARD`, wraps each `DABSNBlock`, retains original parameters for the
optimizer, and uses the FSDP-aware gradient scaler and global gradient clip.

Portable mode writes a self-describing SafeTensors model to
`run/model.safetensors`. Optimizer, AMP scaler, and completed-step state are
stored in the trusted local sidecar
`run/model.safetensors.optimizer.pt`. Add `--resume` to continue the same run.
Resume rejects a missing model or sidecar instead of silently starting over.

For a checkpoint too large to gather on rank zero, use distributed checkpoint
mode:

```bash
torchrun --nnodes=2 --nproc-per-node=8 \
  --rdzv-id=dabsn-pretrain-01 \
  --rdzv-backend=c10d \
  --rdzv-endpoint=trainer-0.example:29400 \
  -m dabsn.cli train \
  --config model.json \
  --data batch.pt \
  --output run/checkpoint \
  --device cuda --backend cuda --distributed fsdp --precision bf16 \
  --checkpoint-mode sharded --steps 10000 --resume
```

The sharded directory contains reshardable model and optimizer files plus
`dabsn-training.json`. It avoids a full rank-zero state gather. If a complete
model can fit in rank-zero host memory, add
`--final-export run/model.safetensors` to consolidate a shareable inference
file.

FSDP is parameter, gradient, optimizer-state, and batch parallelism. It does
not split one sequence, one oversized matrix, or the recurrent context across
GPUs. DABSN does not currently ship tensor, pipeline, or context parallelism;
therefore this repository does not claim that FSDP alone can train an
arbitrary one-trillion-parameter configuration.

Programmatic users can access the same implementation through
`setup_distributed`, `prepare_distributed_model`, `save_distributed_dabsn`,
`save_sharded_training_checkpoint`, and their matching load functions from
`dabsn.runtime`.

## Checkpoints and export

```python
from dabsn import load_dabsn, save_dabsn
from dabsn.runtime import export_dabsn

save_dabsn(model, "model.safetensors")
restored = load_dabsn("model.safetensors", map_location="cpu")

export_dabsn(model, "weights.safetensors", format="safetensors")
export_dabsn(
    model,
    "program.pt2",
    sample_input=sample_input,
    format="torch-export",
)
```

Model checkpoints are atomic, non-pickle SafeTensors files. They embed the full
clean DABSN construction config and preserve tied weights. Custom adapter
implementations remain application-owned and must be registered before loading
a checkpoint that names them. Optimizer sidecars and distributed training
directories are trusted run state, not files to accept from an untrusted
source.

## Train, pretrain, fine-tune, and resume

These commands are deliberately separate:

- `train` creates a new model from `model.json` and prepared input/target
  tensors. `train --resume` continues that exact run with its optimizer and
  completed step.
- `finetune` loads model weights but intentionally creates a new optimizer and
  starts at step zero. Its output must differ from its input checkpoint.
- `pretrain` builds a `DABSNSequenceLM` and learns next-token prediction from a
  token corpus. A binary corpus is memory-mapped: batches are sliced from disk
  without loading the entire corpus into RAM.

`AMP` means automatic mixed precision (`fp16` or `bf16`). It reduces tensor
memory and compute cost while the supplied scaler protects fp16 gradients.
Gradient accumulation divides each update across several smaller batches.

A minimal pretraining config is:

```json
{
  "corpus_bin": "/data/tokens.uint16",
  "corpus_dtype": "uint16",
  "vocab": 50257,
  "hidden_dim": 768,
  "depth": 12,
  "layer_geometries": ["seq"],
  "train_context": 2048,
  "steps": 16000,
  "batch_size": 4,
  "precision": "bf16",
  "distributed": "fsdp",
  "grad_checkpoint": true,
  "grad_accum_steps": 8,
  "checkpoint_every": 1000
}
```

Launch it with:

```bash
torchrun --standalone --nproc-per-node=8 -m dabsn.cli pretrain \
  --config pretrain.json \
  --output run/checkpoint \
  --device cuda --backend cuda \
  --checkpoint-mode sharded \
  --final-export run/model.safetensors \
  --verify-gradients
```

`steps` counts corpus microsteps, matching the canonical training loop. One
optimizer update occurs every `grad_accum_steps`; `checkpoint_every` must land
on an update boundary. The checkpoint records every rank's corpus RNG stream.
Bitwise data-stream continuation therefore requires the same worker count.
Changing the worker count may reshard model/optimizer state, but it is a new
global data trajectory and is not called exact continuation.

Fine-tuning uses a prepared tensor payload and a fresh output path:

```bash
dabsn finetune \
  --checkpoint base.safetensors \
  --data task-batch.pt \
  --output task-model.safetensors \
  --device cuda --backend cuda --precision bf16 --steps 2000
```

## Language modeling

```python
from dabsn import DABSNSequenceLM

model = DABSNSequenceLM(
    vocab=50_257,
    hidden_dim=512,
    depth=4,
    layers="seq:256:256,seq:768:512,seq:768:512,seq:256:256",
    tie_embeddings=False,
)

logits = model.forward_sequence(token_ids)
```

## CLI and reproductions

```bash
dabsn --help
dabsn kernels --enable cuda --required
dabsn doctor

dabsn-reproduce-copy --help
dabsn-reproduce-mqar --help
dabsn-reproduce-keyvalue --help
dabsn-reproduce-a5 --help
```

The full reproduction defaults correspond to the checked-in result tables.
Reduced settings are available for local execution checks and are not presented
as replacements for the reported experiments.

## Development

```bash
git clone https://github.com/BleedingXiko/dabsn.git
cd dabsn
pip install -e '.[test]'
pytest
```

Native release gates are available for a fresh wheel-installed checkout:

```bash
bash tools/cpu_check.sh
bash tools/gpu_check.sh
bash tools/fsdp_check.sh  # requires two NVIDIA GPUs
```

## Citation

If DABSN or its native runtimes contribute to your work, cite the paper:

```bibtex
@misc{rosdahl2026onelayer,
  title     = {One Layer, Both Gaps: A Persistent-Modulation Recurrence that
               Generalizes Copy and Tracks Non-Solvable Group State},
  author    = {Rosdahl, Nicholas},
  year      = {2026},
  publisher = {Zenodo},
  doi       = {10.5281/zenodo.21391204},
  url       = {https://github.com/BleedingXiko/dabsn}
}
```

[`CITATION.cff`](https://github.com/BleedingXiko/dabsn/blob/main/CITATION.cff)
carries the same metadata in machine-readable form, and GitHub's "Cite this
repository" control reads it directly.

## License

DABSN is released under the [Apache License 2.0](https://github.com/BleedingXiko/dabsn/blob/main/LICENSE).
