Upload 65 files

BitTransformerLM/.github/workflows/ci.yml

@@ -0,0 +1,29 @@
+name: CI
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.11'
+      - name: Install dependencies
+        run: |
+          pip install --upgrade pip
+          pip install --extra-index-url https://download.pytorch.org/whl/cpu -r requirements.txt
+          pip install build
+      - name: Run tests
+        run: pytest -q
+      - name: Build package
+        run: python -m build --sdist --wheel -o dist
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist

BitTransformerLM/.gitignore

@@ -0,0 +1,103 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# Pyre type checker
+.pyre/
+
+# mypy
+.mypy_cache/
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# IDEs
+.idea/
+.vscode/
+
+# macOS
+.DS_Store
+
+# Logs
+*.log
+
+# Plot outputs
+*.png
+figures/
+
+# Model artifacts
+*.pt
+*.pth
+*.bin
+candidates/
+approved/
+review_log.jsonl
+
+# Configurations
+*.ini
+
+# Local data
+*.sqlite3
+
+
+*.pt.gz

BitTransformerLM/ABOUTME.md

@@ -0,0 +1,110 @@
+Here’s a menu of additional, “pure-PyTorch” extensions that can close the gap even further to a production-grade LLM:
+
+⸻
+
+1. Native Low-Rank & MoE Layers (DO LAST)
+
+Why: Expert mixtures and low-rank adapters let you balloon effective parameter count without proportional compute.
+	•	Mixture-of-Experts: Implement a tiny gating network (one or two linear layers) that routes each token’s representation to one of E experts (each a small FFN). Only that expert runs on that position, so compute per token stays constant while total capacity grows by E×.
+	•	PyTorch sketch:
+
+class MoE(nn.Module):
+    def __init__(self, d_model, d_ff, n_experts=4):
+        super.__init__
+        self.gate = nn.Linear(d_model, n_experts)
+        self.experts = nn.ModuleList(
+            [nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU, nn.Linear(d_ff, d_model))
+             for _ in range(n_experts)]
+        )
+    def forward(self, x):
+        # x: [T,B,D]
+        logits = self.gate(x)                   # [T,B,E]
+        w = F.softmax(logits, dim=-1)           # [T,B,E]
+        y = torch.stack([expert(x) for expert in self.experts], -1)
+        # y: [T,B,D,E] → weighted sum:
+        out = (y * w.unsqueeze(2)).sum(-1)
+        return out
+
+
+•	Trade-off: You’ll need a load-balancing loss term (e.g. encourage the gate to spread load) and telemetry on expert usage, but the code stays pure PyTorch.
+
+⸻
+
+2. [x] Adaptive Computation Time (ACT)
+
+Why: Let the model learn to spend more depth on “hard” bits and skip layers on easier ones.
+	•	Implementation: Add a tiny halting unit after each layer—e.g. a single linear+sigmoid per token that predicts stop/pause. Accumulate “halt probability” across layers and stop processing tokens once they cross a threshold.
+	•	Benefit: On average you’ll do fewer layer passes per token, reducing compute without touching PyTorch internals.
+
+⸻
+
+3. [x] Advanced PyTorch-Native Quantization
+
+Why: Move beyond static 4-bit packaging to full QAT / dynamic quant.
+	•	FX-graph QAT: Use torch.quantization.prepare_qat_fx on your SparseQuantTransformerLayer with a custom 4-bit observer (we sketched one earlier). Then convert_fx to int8 or 4-bit for weights—no external libs needed.
+	•	Dynamic quant for inference: Wrap your model in torch.quantization.quantize_dynamic(...), quantizing only Linear modules to int8 on-the-fly. Gives a big speed/memory win at inference time on CPU.
+
+⸻
+
+4. [x] Chunked & Overlapping Attention
+
+Why: Emulate sparse attention with pure PyTorch and no for-loops.
+	•	How: Break your sequence into fixed-size chunks (e.g. 512 bits), attend within each chunk plus a small overlap window to neighbors.
+	•	Pure PyTorch: Use unfold + batched torch.matmul to compute all chunked attention in parallel:
+
+x: [B, L, D], chunk_size=C, overlap=O
+pads = (O, O)
+x_padded = F.pad(x, (0,0) + pads)  # pad on seq dim
+chunks = x_padded.unfold(1, C+2*O, C)  # [B, n_chunks, C+2O, D]
+Then project Q,K,V per-chunk and do fused matmuls batchwise
+
+
+•	Benefit: You get an O(L·(C+2O)) algorithm without Python loops, all in tensor ops.
+
+⸻
+
+5. Functorch-Based Vectorization & vmap
+
+Why: Fuse your per-head or per-expert loops automatically.
+	•	Use functorch.vmap to turn your per-head attention code (the one inside the for t in range(T)) into a single batched kernel.
+	•	Benefit: Cleaner code, fewer Python loops, and TorchInductor can fuse it just as well as hand-written loops.
+
+⸻
+
+6. [x] Fully-Sharded DataParallel & Pipeline Parallel (PyTorch-Native)
+
+Why: Scale out to multiple GPUs without external frameworks.
+	•	FSDP: Wrap your model in torch.distributed.fsdp.FullyShardedDataParallel to shard both parameters and optimizer state across GPUs.
+	•	Pipe: Use torch.distributed.pipeline.sync.Pipe to split your 40+ layer model across GPUs as pipeline stages.
+	•	Benefit: Zero external deps—pure PyTorch DDP/FS/PIPE—so you can train 100M+ parameter models.
+
+⸻
+
+7. [x] Mixed Precision & Autocast on CPU (bfloat16)
+
+Why: PyTorch now supports `torch.amp.autocast('cpu')` for bfloat16 on some architectures.
+	•	Surround your forward in with `torch.amp.autocast('cpu')`: to cut memory and speed up linear/attention kernels, even on CPU.
+
+⸻
+
+8. [x] Optimized Learning-Rate Schedules & Optimizers
+
+Why: Achieve GPT-level convergence behavior…
+	•	Implement OneCycleLR or CosineAnnealingWarmRestarts directly via torch.optim.lr_scheduler.
+	•	Swap to AdamW with decoupled weight decay (torch.optim.AdamW) and dynamic gradient clipping (torch.nn.utils.clip_grad_norm_).
+	•	All of these live in core PyTorch.
+
+⸻
+
+Putting It All Together
+	1.	MoE + ACT will let you scale capacity (E× experts) while controlling average compute.
+	2.	FX/QAT + dynamic quant gives you 4-bit int inference with no external libs.
+	3.	Chunked attention + vmap replaces loops with giant fused tensor ops.
+	4.	FSDP + Pipe moves you onto multi-GPU purely in torch.distributed.
+	5.	Autocast (bfloat16) on CPU/GPU for mixed precision speed.
+
+By layering these techniques, you can:
+	•	Reach hundreds of millions (even billions) of effective parameters
+	•	Maintain single-library purity (just PyTorch)
+	•	Hit LLM-class throughputs (100’s of tokens/sec GPU, 10’s CPU)
+	•	Keep full NRB telemetry available for safety checks

BitTransformerLM/AGENTS.md

@@ -0,0 +1,66 @@
+# AGENTS Guidelines for BitTransformerLM
+
+## Repository Scope and Purpose
+- **BitTransformerLM** models raw binary streams using reversible transformer blocks and safety telemetry. The project is the canonical implementation under WCNegentropy.
+- Core capabilities include bit-native modeling, telemetry metrics (negentropy, LZ complexity, symbiosis), progressive scaling, compression, context extension, diffusion mode (linear/cosine/exp noise schedules with parity correction), dashboard control, distributed training, and quantization.
+- Phase 1 optimizations provide configurable batch sizing, gradient accumulation, mixed-precision, memory-mapped dataset streaming, scheduled compression ramps, selective `torch.compile`, and an EMA-smoothed safety gate with burn-in.
+
+## Environment Setup
+- Requires **Python 3.10+**.
+- Install dependencies:
+  - CPU: `pip install --extra-index-url https://download.pytorch.org/whl/cpu -r requirements.txt`
+  - Optional GPU: `pip install --extra-index-url https://download.pytorch.org/whl/cu118 torch==2.7.1+cu118`
+- The package name is `bit-transformer`; project metadata lives in `pyproject.toml`.
+
+## Repository Layout
+- `bit_transformer/` – core package (`model`, `compression`, `telemetry`, `safety`, `dashboard_app`, `quantization`, etc.).
+- `tests/` – pytest suite and historical `TEST_RESULTS.md`.
+- Scripts: `example.py`, `unified_workflow.py`, `full_bits_train.py`, `build_full_bits.py`, `mcp_server.py`, `wikitext_*` utilities. The legacy `progressive_scaleup.py` is retained for reference but superseded by `integration_schedule.py`.
+- Docs and specs: `README.md`, `state_of_the_repo_audit.md`, licensing files in `LICENSE/`.
+
+## Development Practices
+- Follow snake_case for functions and CamelCase for classes.
+- Keep functions under ~300 lines and minimize deeply nested control flow.
+- Avoid reintroducing the deprecated dashboard `/exec` endpoint or other insecure code paths.
+- Use the `/status` endpoint for model introspection; all routes return JSON and surface errors with stack traces.
+- Ensure compression, decompression, and halting logic stay consistent with current implementation.
+- Use the `cpu_autocast()` helper for BF16 mixed precision on CPU instead of
+  calling `torch.amp.autocast` directly.
+- Adaptive training now expands depth, width, or context only when validation loss plateaus and automatically decays the base learning rate by √2 after each expansion with a 100‑step warm‑up.
+
+## Workflow & Commands
+- Run the example: `python example.py`.
+- Adaptive scaling now lives in `integration_schedule.py`; `progressive_scaleup.py` is deprecated.
+- Unified workflow (optionally with dashboard or diffusion): `python unified_workflow.py --dashboard` or `python unified_workflow.py --diffusion --diffusion-steps 8 --dataset-size 32`.
+- Increase `--diffusion-steps` for higher fidelity (8–16) and add `--diffusion-curriculum` to linearly decay noise over epochs.
+- Disable checkpointing or reversible blocks when speed is prioritized over memory: `python unified_workflow.py --no-checkpoint --no-reversible`.
+- Enable 4-bit quantization-aware training: `python unified_workflow.py --qat`.
+- Skip full attention logging during chunked attention for memory savings by constructing the model with `full_attn_logging=False`.
+- Start MCP server: `python mcp_server.py` and launch dashboard: `MCP_SERVER_ADDR=http://127.0.0.1:7000 python -m bit_transformer.dashboard_app`.
+- `/metrics` and `/model_config` endpoints expose telemetry streams and hyperparameters.
+- `/save_checkpoint` and `/download_checkpoint` sync weights with Hugging Face (token defaults to `HF_TOKEN`).
+- Container build: `docker build -t bittransformerlm .` and run with exposed ports `5000` (dashboard) and `7000` (MCP).
+
+## Telemetry Metrics
+| Metric | Meaning | Range |
+|--------|---------|-------|
+| **K** | Negentropy – deviation from random noise | 0–1 (1 = ordered) |
+| **C** | LZ Complexity – compressibility proxy | 0–1 (higher = more changes) |
+| **S** | Symbiosis – agreement with reference distribution | 0–1 (1 = aligned) |
+
+ACT halting exports `halt_probs` in telemetry showing how many layers executed. For robust sampling under safety constraints, call `safe_sample_with_retry(model, bits)` which retries with diffusion mode and exponential backoff.
+
+`TelemetrySynthesizer.cluster_sequences` can be used to select representative training samples before invoking `collapse_submodel`. The distillation helper deepens the model and widens once (`width_scale` = 1.5) if floors are missed, and `save_distilled_model` emits a `metrics.json` summary beside the weights.
+
+## Testing
+- Run unit tests after any change: `pytest -q`.
+- Use `watcher.py` for auto-reload and test on local development if desired.
+- During training, call `model.train()` and keep dropout probabilities around `0.1–0.2`.
+- Before running tests, inference, or pushing weights, switch to `model.eval()` and set all dropout probabilities to `0` to avoid flaky results.
+- Dashboard will warn if telemetry metrics drift by more than 0.2 over the last 10 steps. Adjust via `ModelManager(drift_window, drift_threshold)` as needed.
+
+## Licensing
+- Project governed by documents in `LICENSE/` (AGPLv3, commercial terms, disclaimers, etc.). Ensure compliance before contributing or distributing.
+
+These guidelines keep the repository consistent with the project roadmap and previous audits. Maintain security, style, and testing discipline to keep BitTransformerLM production-ready.
+

BitTransformerLM/BitTransformerLM_full_assessment.md

@@ -0,0 +1,196 @@
+
+# BitTransformerLM Deep-Dive Assessment Report
+
+*(Comprehensive technical review and optimization roadmap)*
+
+---
+
+## Completed Tasks
+- [x] 3.1 Cosine noise schedule option
+- [x] 3.2 Post-process parity correction
+- [x] 2.3 Expose checkpoint & reversible toggles
+- [x] 2.2 Update deprecated AMP call
+- [x] 5.2 Metric-drift alerts
+- [x] 1.3 Expand README / docstrings for telemetry & ACT
+- [x] 3.3 Safety-gate soft-retry
+- [x] 7.1 Add ACT halting unit test
+- [x] 4.1 Integrate performance-based scaling
+- [x] 4.2 Learning-rate decay on resize
+- [x] 3.4 Chunked attention logging toggle
+- [x] 3.5 Quantization-aware training toggle
+- [x] 7.2 Quantization & QAT tests
+- [x] 4.3 Dashboard flag wiring
+- [x] 7.3 Dashboard smoke test
+- [x] 2.1 Unify flag names & deprecate legacy scale script
+- [x] 5.1 Telemetry λ and floor UI
+- [x] 5.3 Cluster-based distillation data
+- [x] 6.1 Allow width scaling in collapse loop
+- [x] 6.2 Save distilled metrics summary
+
+## 1. Overview of BitTransformerLM Architecture and Recent Additions  
+BitTransformerLM is a **reversible Transformer** that operates **directly on binary sequences (bits)**.  The immutable core uses multi-head self-attention on bit embeddings with sinusoidal positional encoding and already supports:
+
+* Safety-centric telemetry (negentropy *K*, LZ complexity *C*, symbiosis *S*)
+* Run-length compression / decompression paths
+* Progressive scaling (depth & width) with reversible layers + gradient checkpointing
+* Quantization (dynamic INT8 + optional 4‑bit QAT)
+* A non‑causal **Diffusion‑LM mode** for bidirectional, denoising generation
+* Dashboard, MCP server, and FSDP/pipeline hooks for distributed or edge deployment
+
+Recent commits locked in deterministic environment setup (ChatGPT Codex container), removed insecure `/exec` endpoints, and added a reliable *course‑to‑fine* diffusion sampler stub.  The model now installs and trains reproducibly on CPU‑only hosts, yet scales to multi‑GPU with FSDP.
+
+---
+
+## 2. Consistent Naming & Documentation
+* Codebase generally follows *snake_case* functions / *CamelCase* classes, but CLI flags & helper scripts drift (e.g. `--diffusion` vs internal `causal=False`).  
+  **Action:** unify flag names & docstrings; deprecate redundant scripts (`progressive_scaleup.py` vs `integration_schedule.py`).
+* README and inline docs lack quick intuition for *K, C, S* metrics, ACT, and reversible internals.  
+  **Action:** add short metric primers and ACT demo snippets; update `AGENTS.md` quick‑start table.
+
+---
+
+## 3. Optimizing Module Interactions & Performance
+| Area | Current State | Optimization | Outcome |
+|------|---------------|--------------|---------|
+| **Chunked attention** ✅ | Saves RAM but reconstructs full *T×T* matrix for telemetry | Skip full matrix when `chunk_size < seq_len` and user disables `full_attn_logging` | Same metrics, big memory + speed win on long sequences |
+| **PyTorch 2 features** | Uses `torch.compile` & BF16 autocast inconsistently | Standardize `torch.amp.autocast(device_type="cpu", dtype=torch.bfloat16)`; wrap long loops | 10‑20 % CPU speed‑up, no deprecation warnings |
+| **Reversible + checkpoint** | Always checkpoints → slower when RAM ample | Expose `--no-checkpoint` flag; document trade‑offs | User‑selectable speed vs memory |
+| **Quantization** ✅ | INT8 dynamic works; 4‑bit QAT unused | Add `--qat` toggle in training scripts & unit‑test tiny model | Edge‑ready 4‑bit weights validated |
+| **Compression loops** | Python for‑loops per sample | Batch or vectorized RLE when batch≫8 | Marginal speed‑up for large batches |
+
+---
+
+## 4. Fully Leveraging Diffusion Mode
+1. [x] **Noise schedule** – switchable linear ▸ cosine ▸ exponential; expose `--noise-schedule`.
+2. [x] **Step count** – allow 8–16 steps for high‑fidelity generation; document compute trade‑off.
+3. [x] **Parity safeguard** – post‑sampling parity‑bit fix or strict parity sampling to guarantee valid bytes.
+4. [x] **Training curriculum** – optional schedule: high‑noise → low‑noise over epochs; keep random‑noise fallback.
+5. [x] **Safety integration** – run `hil_safe_inference(strict=False)` during diffusion; warn (not crash) on metric floor breaches.
+
+---
+
+## 5. Enhanced Training Workflow & Scaling Strategy
+* **Adaptive scaling trigger** – adopt `progressive_scaleup.py` logic: scale only when val‑loss Δ < threshold; alternate width↔context↔depth.  
+* **Context extension** – use `double_length()` when plateau met; maintain chunked attention windows.  
+* **Warm‑up & plateau** – keep 5‑batch freeze after each expansion; add default final plateau epoch.  
+* **LR hygiene** – slight LR decay each scale‑up; document rationale.
+
+---
+
+## 6. Telemetry Metrics & Safety Integration
+* **Metric coefficients** (`λ_K`, `λ_C`, `λ_S`) exposed via dashboard slider; floors (C ≥ 0.3, S ≥ 0.5) adjustable per deployment.  
+* **TelemetrySynthesizer** – cluster activations → representative sequences for distillation & drift detection.  
+* **Metric drift alert** – integrate `detect_metric_drift()` into training monitor; log if Δ > 0.2.
+
+---
+
+## 7. Distillation & Model Collapse Optimization
+1. Use **cluster‑selected sequences** as `cluster_data` for `collapse_submodel` → better coverage.  
+2. Permit optional width growth (`width_scale > 1`) in iterative collapse rounds.  
+3. Log final vs floor metrics in `distilled_metrics.json` for audit trail.  
+4. Optionally auto‑invoke collapse at end of `integration_schedule` with `--auto-collapse`.
+
+---
+
+## 8. Additional Testing & Release Readiness
+* Expand pytest suite: diffusion training/sampling, ACT halting, INT8 + QAT inference, dashboard API smoke tests.  
+* Add multi‑GPU CI job to validate FSDP + reversible layers.  
+* Strengthen debug logs: print mode (causal/diffusion/compression), scale‑up events, safety‑gate warnings.
+
+---
+
+## 9. Strategic Summary
+BitTransformerLM already delivers an **orthogonal bundle of “firsts”**: bit‑native granularity, reversible memory efficiency, metric‑driven safety, and turnkey text diffusion.  
+Executing the roadmap **knits every module into a smooth, reproducible pipeline** without touching core architecture—preserving alignment while boosting usability.
+
+**Bottom‑line:** With these refinements, BitTransformerLM becomes the reference for transparent, resource‑efficient, safety‑gated language modelling at the bit level—well beyond “just another model.”
+
+
+Below is an **implementation playbook** that turns every recommendation in *“Overview of BitTransformerLM Architecture and Recent Additions”* into clear tasks and ready‑to‑copy Codex prompts.  Where page numbers add context, I note them; all content is from the uploaded PDF.&#x20;
+
+---
+
+## 1 · Repository Consistency & Documentation
+
+| #   | Task                                                           | Key Steps                                                                                                                                                 | Codex Prompt (trim or expand as desired)                                                                                                                                                                                                                                 |
+| --- | -------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| 1.1 | **Audit & unify public API names**                             | • Scan for duplicate / mis‑matched flags (e.g. `--diffusion` vs `causal=False`).<br>• Rename or deprecate aliases; update docs.                           | “List every function, class, and CLI flag whose name does **not** match the style‑guide (snake\_case for funcs, CamelCase for classes) in the BitTransformerLM repo. For each, propose a single canonical name and generate the automated `git mv` or refactor patches.” |
+| 1.2 | **Consolidate scaling scripts**                                | • Merge `progressive_scaleup.py` logic into `integration_schedule.py`.<br>• Mark redundant script as example.                                             | “Move the performance‑based scaling criterion from `progressive_scaleup.py` into `integration_schedule.py`. Preserve existing kwargs, add `--improve‑thresh` with default 0.01. Provide diff.”                                                                           |
+| 1.3 | **Expand README / docstrings for telemetry & ACT** (pp. 1 ‑ 2) | • Add one‑paragraph explanations of Negentropy (K), LZ Complexity (C), Symbiosis (S), and ACT halting to README.<br>• Link to equations in code comments. | “Insert a new subsection *‘Telemetry Metrics Explained’* into README after the quick‑start block, then add in‑line docstrings for `negentropy_score`, `lz_complexity`, and `symbiosis_score` explaining ranges and typical values.”                                      |
+
+---
+
+## 2 · Performance Optimizations
+
+| #   | Task                                              | Key Steps                                                                                                                    | Codex Prompt                                                                                                                                                                                                     |
+| --- | ------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| 2.1 | **Vectorize chunked‑attention telemetry** (p. 2)  | • Add flag `--attn‑summary`.<br>• When enabled and `chunked_attn=True`, compute per‑chunk entropy and skip full `T × T` map. | “Refactor `_chunked_attn` in `model.py` so that, if `attn_summary` is true, it returns `(attn_entropy_per_chunk, None)` instead of the stitched full map. Fall back to old behaviour otherwise. Update callers.” |
+| 2.2 | **Update deprecated AMP call**                    | Replace `torch.cpu.amp.autocast` with `torch.amp.autocast(device_type="cpu", dtype=torch.bfloat16)` everywhere.              | “Search repo for `torch.cpu.amp.autocast`, replace with the new API, and add a context‑manager wrapper `cpu_autocast` in `utils/torch_utils.py`.”                                                                |
+| 2.3 | **Expose checkpoint & reversible toggles** (p. 2) | • Add CLI flags `--use-checkpoint / --no-checkpoint` and `--reversible`.<br>• Document memory/compute trade‑off.             | “Modify `train.py` argparse to include mutually exclusive `--[no-]checkpoint` flags; wire to `use_checkpoint` in model init.”                                                                                    |
+| 2.4 | **Batch run‑length encoding** (p. 3)              | • Implement NumPy‑vectorised RLE for the full tensor.<br>• Fallback to Python loop if tensor < 1024 bits.                    | “Implement `batch_rle_encode` in `bit_io.py` using NumPy strides; write unit test comparing speed & correctness to existing per‑sequence encode.”                                                                |
+
+---
+
+## 3 · Diffusion‑Mode Enhancements
+
+| #   | Task                                      | Key Steps                                                                                                                       | Codex Prompt                                                                                                       |                                                                  |                                                                                                                          |
+| --- | ----------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------ |
+| 3.1 | **Cosine noise schedule option** (p. 4)   | • Add \`schedule="linear                                                                                                        | cosine                                                                                                             | exp"`arg to`diffusion\_inference\`.<br>• Default remains linear. | “Extend `diffusion_inference` to support a cosine decay of `mask_prob` over `steps`. Provide math and update docstring.” |
+| 3.2 | **Post‑process parity correction** (p. 4) | • After sampling, flip each parity bit if byte parity invalid.<br>• Log number of corrections.                                  | “Write `enforce_parity(bits)` that patches 9th bit per byte to satisfy even‑parity, return corrected seq + stats.” |                                                                  |                                                                                                                          |
+| 3.3 | **Safety‑gate soft‑retry**                | • On failed `hil_safe_inference(strict=True)`, auto‑retry up to 3× with diffusion or random seed.<br>• Surface warning in logs. | “Wrap `hil_safe_inference` in a helper `safe_sample_with_retry`; implement exponential back‑off and logging.”      |                                                                  |                                                                                                                          |
+
+---
+
+## 4 · Adaptive Training Workflow
+
+| #   | Task                                              | Key Steps                                                                                                                                   | Codex Prompt                                                                                                                                                                                                    |
+| --- | ------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| 4.1 | **Integrate performance‑based scaling** (pp. 5‑6) | • Use `Δval_loss < thresh` as condition to trigger `add_layer()`/`double_width()`.<br>• Alternate occasional `double_length()` for context. | “Inside `integration_schedule.train_loop`, compute rolling val‑loss; if mean improvement < `args.improve_thresh`, call `model.scale_up(strategy=next_step)` where `next_step` cycles \[layer, width, context].” |
+| 4.2 | **Learning‑rate decay on resize**                 | • After each scale‑up, reduce base LR by √2.<br>• Provide warm‑up of 100 steps.                                                             | “Add `adjust_learning_rate(optimizer, factor)` util; call it after every successful model expansion.”                                                                                                           |
+| 4.3 | **Dashboard flag wiring**                         | • Map UI toggles (compression, diffusion) to `compress_prob`, `diffusion` args in backend.                                                  | “In `dashboard_app.py`, when user toggles compression, pass `compress_prob=1.0` to `ModelManager.train()`.”                                                                                                     |
+
+---
+
+## 5 · Telemetry & Safety
+
+| #   | Task                                                     | Key Steps                                                                                                             | Codex Prompt                                                                                                                                                  |
+| --- | -------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| 5.1 | **Expose λ coefficients and safety floors in UI** (p. 7) | • Add sliders for `λ_K`, `λ_C`, `λ_S`, `C_floor`, `S_floor`.<br>• Persist to model state.                             | “Add REST endpoints `/config/telemetry` (GET/POST) that read or set lambda values and floors; bind to dashboard sliders.”                                     |
+| 5.2 | **Metric‑drift alerts** (p. 8)                           | • After every epoch, call `detect_metric_drift(history, window=100)`; if > 0.2 drift, log & optionally halt training. | “Integrate `detect_metric_drift` into `ModelManager._log_metrics`; raise `MetricDriftWarning` when threshold exceeded.”                                       |
+| 5.3 | **Cluster‑based distillation data** (pp. 8‑9)            | • Use `TelemetrySynthesizer` to pick `k` cluster representatives (default 8).<br>• Feed to `collapse_submodel`.       | “Before `collapse_submodel`, run `representatives = TelemetrySynthesizer(model).cluster(train_data, k=8)`. Replace `train_bits[:64]` with `representatives`.” |
+
+---
+
+## 6 · Distillation / Collapse Process
+
+| #   | Task                                            | Key Steps                                                                                        | Codex Prompt                                                                                                        |
+| --- | ----------------------------------------------- | ------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------------------------------------------------- |
+| 6.1 | **Allow width scaling in collapse loop** (p. 8) | • Add `width_scale` param; if metric floors unmet after deepening, double width once then retry. | “Modify `collapse_submodel`: on round‑2 failure, rebuild sub‑model with `hidden_dim *= width_scale` (default 1.5).” |
+| 6.2 | **Save metrics summary**                        | • Extend `save_distilled_model` to write `metrics.json` with achieved vs floor values.           | “Update `save_distilled_model` to dump `{‘C’:score_C, ‘S’:score_S, ‘floors’:{...}}` alongside weights.”             |
+
+---
+
+## 7 · Testing & CI Hardening
+
+| #   | Task                                  | Key Steps                                                                            | Codex Prompt                                                                                                    |
+| --- | ------------------------------------- | ------------------------------------------------------------------------------------ | --------------------------------------------------------------------------------------------------------------- |
+| 7.1 | **Add ACT halting unit test** (p. 10) | • Craft toy seq; assert `sum(halt_prob<1) < n_layers`.                               | “Write `tests/test_act.py` ensuring at least one layer halts early when `use_act=True, threshold=0.1`.”         |
+| 7.2 | **Quantization & QAT tests**          | • After tiny train, run dynamic int8 + fake‑QAT path, assert same logits ±1e‑3.      | “Add `pytest` case: train 2‑layer model 1 epoch, call `quantize_dynamic`, compare outputs on 10 random inputs.” |
+| 7.3 | **Dashboard smoke test**              | • In CI, launch Flask app with `pytest‑flask`, hit `/init`, `/train‑step`, `/infer`. | “Create `tests/test_dashboard.py` that starts server in a thread and exercises core endpoints.”                 |
+
+---
+
+## 8 · Packaging & Release
+
+| #   | Task                                     | Key Steps                                                                     | Codex Prompt                                                                                                      |
+| --- | ---------------------------------------- | ----------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------- |
+| 8.1 | **Rename repository references** (p. 11) | • Replace `Test/` URL stubs with new repo slug.<br>• Update badges in README. | “Search‑replace all GitHub links from `WCNegentropy/Test` to `WCNegentropy/BitTransformerLM`; update badge SVGs.” |
+| 8.2 | **PyPI build verification**              | • Ensure `pyproject.toml` installs cleanly on 3.10 & 3.11 in CI.              | “Add GitHub Action matrix for {macOS, ubuntu‑latest} × {3.10, 3.11}; run `pip install -e . && pytest`.”           |
+
+---
+
+### How to Use These Prompts
+
+**Run** unit tests; iterate if failures surface.
+
+This checklist should bring BitTransformerLM to a polished, v1‑ready state while aligning with your NRB‑driven safety and telemetry philosophy.&#x20;

BitTransformerLM/Dockerfile

@@ -0,0 +1,27 @@
+FROM ubuntu:22.04
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update && \
+    apt-get install -y python3.11 python3-pip python3.11-venv curl && \
+    apt-get clean && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /opt/bit_transformer
+COPY . .
+
+ARG TORCH_CUDA=cpu
+RUN pip3 install --no-cache-dir --upgrade pip && \
+    if [ "$TORCH_CUDA" = "cu118" ]; then \
+        pip3 install torch==2.7.1+cu118 --extra-index-url https://download.pytorch.org/whl/cu118; \
+    else \
+        pip3 install torch==2.7.1+cpu --extra-index-url https://download.pytorch.org/whl/cpu; \
+    fi && \
+    pip3 install -r requirements.txt
+
+ENV MCP_SERVER_ADDR=http://127.0.0.1:7000
+
+EXPOSE 5000 7000
+
+RUN chmod +x start.sh
+
+HEALTHCHECK CMD curl -f http://localhost:7000/health || exit 1
+
+CMD ["/opt/bit_transformer/start.sh"]

BitTransformerLM/LICENSE/ALIGNMENT_AND_TRANSPARENCY.txt

@@ -0,0 +1,42 @@
+# Alignment and Transparency Agreement for BitTransformerLM
+This Alignment and Transparency Agreement ("Agreement") outlines requirements  
+for the responsible and aligned commercial deployment of BitTransformerLM,  
+developed by WCNEGENTROPY HOLDINGS LLC.
+
+## Core Principles
+1. **Alignment:** The deployment must actively maintain alignment with ethical  
+and epistemic integrity, avoiding harmful or coercive outcomes.  
+2. **Transparency:** Telemetry data (negentropy, complexity, symbiosis scores,  
+etc.) must be transparently maintained and available for audit and inspection.  
+3. **Safety and Stability:** Deployments must utilize provided safety gates  
+(e.g., `hil_safe_inference`) to prevent degenerate or harmful outputs.  
+4. **Epistemic Responsibility:** Adopters commit to responsible use, actively  
+avoiding misuse or unethical applications.
+
+## Telemetry and Monitoring
+Commercial license holders must maintain full transparency on telemetry metrics  
+collected from the software, as originally implemented in the BitTransformerLM  
+repository. Telemetry must be available upon request for audit by WCNEGENTROPY HOLDINGS LLC or authorized third parties.
+
+## Modification and Derivatives
+Commercial license holders may modify the software for internal commercial use  
+but must explicitly disclose any modifications or derivatives upon request by  
+WCNEGENTROPY HOLDINGS LLC.
+
+## Violations
+Non-compliance with any of the terms outlined in this Agreement may result in  
+revocation of commercial licensing rights at the sole discretion of WCNEGENTROPY HOLDINGS LLC.
+
+### Audit Cadence (added v0.9.0)
+WCNEGENTROPY HOLDINGS LLC may request a telemetry snapshot **no more than once per
+calendar quarter**. Licensee must deliver the requested data within **30 days
+of receipt**. Failure to comply may result in suspension or termination of
+commercial rights.
+
+---
+
+For questions, clarification, or audits, contact:  
+**WCNegentropy Holdings**  
+Email: [email protected]  
+Website: [wcnegentropy.com](https://wcnegentrop
+y.com)

BitTransformerLM/LICENSE/COMMERCIAL_LICENSE.txt

@@ -0,0 +1,34 @@
+# Commercial License for BitTransformerLM
+© 2025 WCNEGENTROPY HOLDINGS LLC – All Rights Reserved
+
+> **Clarification (dual‑license):** This clause applies only to commercial deployments.
+> Open‑source users remain bound by the GNU AGPL v3 in `LICENSE`.
+> For holders of a **paid Commercial License**, BitTransformerLM is also provided
+> under the **Apache License 2.0** (the “Commercial License”), **subject to the
+> Alignment & Transparency Agreement (ATA)**.
+
+BitTransformerLM (the “Software”), including all source code, documentation,
+and associated assets, is the exclusive property of WCNEGENTROPY HOLDINGS LLC
+(“WCNH”). Commercial use, reproduction, modification, distribution, or
+sublicensing requires an executed Commercial License Agreement with WCNH.
+
+## Patent Grant (Defensive)
+WCNH hereby grants Licensee a perpetual, worldwide, non‑exclusive, no‑charge
+patent license to make, use, sell, offer to sell, import, and otherwise
+exploit the Software **provided** Licensee complies with this Commercial
+License and the ATA. **This patent license terminates automatically** if
+Licensee initiates patent litigation alleging that the Software infringes any
+patent claim.
+
+## Export‑Control & Sanctions Compliance
+Licensee shall comply with all applicable export‑control and sanctions laws,
+including but not limited to U.S. EAR, EU dual‑use regulations, and OFAC
+sanctions lists. Licensee must not export, re‑export, or provide the Software
+(or derivatives) to any prohibited country, entity, or individual.
+
+## Alignment & Transparency Obligation
+Commercial usage is conditional upon adherence to the ATA (see
+`ALIGNMENT_AND_TRANSPARENCY.txt`). Failure to comply constitutes a material
+breach and grounds for immediate license revocation.
+
+For commercial inquiries contact **[email protected]**.

BitTransformerLM/LICENSE/CONTRIBUTOR_LICENSE_AGREEMENT.txt

@@ -0,0 +1,7 @@
+# Individual Contributor License Agreement (ICLA)
+By submitting a contribution to the BitTransformerLM repository you agree to
+grant WCNEGENTROPY HOLDINGS LLC an irrevocable, worldwide, royalty‑free
+copyright license to reproduce, prepare derivative works of, publicly display
+and distribute your contribution. You certify that you have the right to make
+this grant and that your contribution is original or you have secured the
+appropriate rights.

BitTransformerLM/LICENSE/DISCLAIMER.txt

@@ -0,0 +1,93 @@
+# BitTransformerLM – Legal & Risk Disclaimer
+_Last updated: 2025-08-04_
+
+BitTransformerLM (the “Software”) is an **experimental, highly-capable, agentic AI
+model** developed by WC Negentropy Holdings LLC (“WCNH”). By downloading,
+installing, running, fine-tuning, or otherwise using the Software **you
+acknowledge and agree to all terms below.**
+
+---
+
+## 1. No Warranty
+
+THE SOFTWARE IS PROVIDED **“AS IS”** AND **WITHOUT WARRANTY** OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE, NON-INFRINGEMENT,
+AND ERROR-FREE OR UNINTERRUPTED OPERATION.
+WCNH DOES **NOT** WARRANT THAT THE SOFTWARE WILL MEET YOUR REQUIREMENTS OR
+THAT ITS OUTPUT WILL BE ACCURATE, COMPLETE, OR RELIABLE.
+
+## 2. Limitation of Liability
+
+TO THE MAXIMUM EXTENT PERMITTED BY APPLICABLE LAW, IN NO EVENT SHALL WCNH,
+ITS AFFILIATES, CONTRIBUTORS, OR LICENSORS BE LIABLE FOR ANY DIRECT,
+INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; BUSINESS INTERRUPTION; OR PERSONAL
+INJURY) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+ANY WAY OUT OF THE USE OF THE SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY
+OF SUCH DAMAGE.
+
+## 3. High-Risk & Regulated Uses
+
+**DO NOT DEPLOY** the Software in environments where failure or malfunction
+could lead to death, serious bodily injury, or severe property or
+environmental damage, including but not limited to:
+
+- Medical diagnosis or life-support systems
+- Autonomous vehicles or aviation control
+- Nuclear facilities, weapons development, or military combat systems
+- Critical infrastructure (power, water, telecom)
+- Legal, financial, or governmental decision-making without qualified
+  human review
+
+You remain solely responsible for conducting appropriate risk assessments,
+validation, and human oversight before any production deployment.
+
+## 4. Alignment & Transparency Obligations
+
+If you hold a **Commercial License**, you must also comply with the
+Alignment & Transparency Agreement (ATA), including:
+
+- Logging K-C-S telemetry and retaining it for audit
+- Supplying telemetry snapshots upon request (max once per quarter)
+- Cooperating with reasonable misuse investigations
+
+## 5. Data & Privacy
+
+You are responsible for:
+
+- Ensuring you have the legal right to process any data you supply to the
+  Software
+- Implementing technical and organizational measures to protect personal or
+  sensitive data
+- Complying with all applicable data-protection and privacy laws (e.g.,
+  GDPR, CCPA, HIPAA)
+
+## 6. Export & Sanctions Compliance
+
+You may **not** use or transfer the Software in violation of U.S. export
+control laws, EU dual-use regulations, or applicable sanctions regimes.
+This includes, but is not limited to, prohibitions on use in or by
+countries, entities, or individuals listed on U.S. or EU restricted-party
+lists.
+
+## 7. Third-Party Dependencies
+
+The Software may incorporate open-source components licensed under separate
+terms. Such components are provided **“as is”** and remain subject to their
+respective licenses. A complete list is available in `THIRD_PARTY_LICENSES.txt`.
+
+## 8. No Professional Advice
+
+The Software’s outputs (including code, text, and recommendations) do **not**
+constitute professional, legal, medical, financial, or safety advice. Always
+consult a qualified expert before relying on the Software for any critical
+decision.
+
+---
+
+**© 2023-2025 WCNEGENTROPY HOLDINGS LLC.**
+All trademarks—including “BitTransformerLM” and the spiral-N logo—are property
+of WCNH. Unauthorized use is prohibited.

BitTransformerLM/LICENSE/LICENSE.txt

@@ -0,0 +1,12 @@
+# LICENSE (AGPLv3)
+Copyright (C) 2025 WCNEGENTROPY HOLDINGS LLC  
+This program is free software: you can redistribute it and/or modify  
+it under the terms of the GNU Affero General Public License as published  
+by the Free Software Foundation, either version 3 of the License, or  
+(at your option) any later version.  
+This program is distributed in the hope that it will be useful,  
+but WITHOUT ANY WARRANTY; without even the implied warranty of  
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the  
+GNU Affero General Public License for more details.  
+You should have received a copy of the GNU Affero General Public License  
+along with this program. If not, see <https://www.gnu.org/licenses/>.

BitTransformerLM/LICENSE/TRADEMARK_POLICY.txt

@@ -0,0 +1,12 @@
+“BitTransformerLM” and the spiral‑N logo are trademarks of WCNEGENTROPY HOLDINGS LLC.
+
+Permitted use:
+• Describing, linking to, or referencing **unmodified, official builds** of
+  BitTransformerLM.
+
+Prohibited use without prior written permission:
+• Branding or promoting modified forks, derivatives, or third‑party services.
+• Creating confusingly similar marks or implying endorsement by WCNH.
+
+Forks must remove or rename the marks to avoid confusion.
+Contact **[email protected]** for licensing requests.

BitTransformerLM/NEW_CODEX_TASK.md

@@ -0,0 +1,85 @@
+# DEPRECATED
+
+All tasks in this file have been implemented (Stages 1–5). The document remains for historical reference only.
+
+Stage 1: Compression Algorithm Implementation
+
+Task 1: Choose Compression Method
+
+Prompt:
+
+Codex: Provide a concise PyTorch-compatible implementation of lossless binary compression and decompression (e.g., RLE, Huffman, or LZ-based) suitable for binary input sequences represented as tensors of bits.
+
+Task 2: Implement Compression Functions
+
+Prompt:
+
+Codex: Implement PyTorch functions compress_bits(input_tensor) and decompress_bits(compressed_tensor) that accept and return PyTorch tensors (dtype=torch.bool or torch.uint8). Ensure compress → decompress cycle perfectly reconstructs original data, and include simple unit tests.
+
+⸻
+
+Stage 2: Encoder/Decoder Integration
+
+Task 3: Add Compression to Encoder Input
+
+Prompt:
+
+Codex: Modify BitTransformerLM’s input pipeline by wrapping the existing model forward pass with a forward_compressed(bits_tensor) method. This method should decompress incoming compressed bit tensors before embedding. Ensure it returns identical outputs as existing uncompressed inputs for verification.
+
+Task 4: Add Decompression to Decoder Output
+
+Prompt:
+
+Codex: Implement a PyTorch-compatible function model_output_decompress(output_bits_tensor) to decompress bit sequences output by BitTransformerLM. Integrate this function as an optional post-processing step after the model’s bitstream generation.
+
+⸻
+
+Stage 3: Training and Evaluation Enhancements
+
+Task 5: Toggle Compression During Training
+
+Prompt:
+
+Codex: Modify the existing training loop to randomly compress input bit sequences with a configurable probability (compress_prob=0.5). Ensure that when compression is on, inputs are compressed and decompressed transparently, and when off, inputs bypass compression.
+
+Task 6: Evaluate Compressed vs Raw Performance
+
+Prompt:
+
+Codex: Extend the current training evaluation metrics to separately track loss, accuracy, and compression ratio for both compressed and raw sequences. Log these metrics clearly in the training output.
+
+⸻
+
+Stage 4: Advanced Integration (Optional)
+
+Task 7: Multi-task Training for Compression Learning
+
+Prompt:
+
+Codex: Implement an optional multi-task training mode where the model occasionally sees compressed inputs directly without decompression. Add a separate loss calculation to monitor its performance on these compressed inputs. Track and log separately from normal next-bit prediction loss.
+
+Task 8: Compression-aware Safety Telemetry
+
+Prompt:
+
+Codex: Adjust the existing BitTransformerLM telemetry (K, C, and S metrics) to handle compressed sequences appropriately. Modify telemetry calculations to optionally apply metrics to decompressed outputs instead of raw bitstream when compression is enabled.
+
+⸻
+
+Stage 5: Dashboard and Runtime Integration
+
+Task 9: Dashboard Compression UI Toggle
+
+Prompt:
+
+Codex: Add a simple UI toggle labeled “Enable Compression” to the existing BitTransformerLM dashboard, controlling whether inputs and outputs are automatically compressed and decompressed. Display compression ratio metrics when enabled.
+
+Task 10: Error Handling and User Feedback
+
+Prompt:
+
+Codex: Implement graceful error handling in the dashboard for compression and decompression failures. Provide clear user-facing feedback in the UI if decompression fails, along with suggestions or fallbacks.
+
+⸻
+
+These ten tasks enable incremental, testable integration of binary compression/decompression into BitTransformerLM without fundamentally altering the core transformer model itself.

BitTransformerLM/README.md

@@ -0,0 +1,177 @@
+# BitTransformerLM
+
+**Project Status:** Pre-release (v1 candidate)
+
+BitTransformerLM is a bit-centric transformer language model built entirely in PyTorch. The project began as a small prototype but has matured into a near-production system capable of modeling raw binary streams with sophisticated safety telemetry and automated scale-up tooling. This repository now serves as the canonical implementation under WCNegentropy.
+
+## Historical Background
+- **Early Experiments** – Initial prototypes explored mapping text to parity-protected bits and training a minimal transformer on random data.
+- **Telemetry & Safety** – Added negentropy, LZ complexity and symbiosis scoring to measure information flow and gate unsafe outputs.
+- **Progressive Scaling** – Introduced reversible layers and automatic depth/width expansion for efficient curriculum training. The schedule now triggers expansions only when validation loss plateaus and decays the learning rate by √2 after each growth with a 100-step warm‑up.
+- **Compression Support** – Integrated run-length encoding and packed bit I/O with optional multi-task training on compressed sequences.
+- **Context Extension** – Implemented chunked attention and sliding-window inference for long sequences with optional overlapping windows.
+- **Attention Logging Toggle** – ``full_attn_logging=False`` skips reconstructing full ``T×T`` attention maps during chunked attention, cutting memory use for very long sequences.
+- **Diffusion LM Mode** – Enable bidirectional denoising by setting ``causal=False`` or toggling **Diffusion LM** in the dashboard. Chunked attention is automatically disabled in this mode and restored afterward.
+- **Dashboard & MCP Server** – Built a lightweight web UI backed by a management server for real‑time training, inference and model collapse. New `/metrics` and `/model_config` endpoints surface live telemetry and hyperparameters, and `/save_checkpoint` and `/download_checkpoint` enable Hugging Face weight sync. The insecure `/exec` route has been removed.
+- **Phase 1 Optimizations** – Configurable batch sizes with aligned OneCycle scheduling, gradient accumulation, mixed‑precision, memory‑mapped dataset streaming, scheduled compression ramps, selective ``torch.compile``, and an EMA‑smoothed safety gate with burn‑in to cut false positives.
+
+The codebase has undergone multiple stress tests and synthetic benchmarks (see `tests/TEST_RESULTS.md`) and now approaches a stable release.
+
+## Quick Start
+Install dependencies using the CPU wheel of PyTorch (default):
+```bash
+pip install --extra-index-url https://download.pytorch.org/whl/cpu -r requirements.txt
+```
+When GPU acceleration is toggled in the dashboard, the application automatically
+installs the CUDA-enabled wheel:
+```bash
+pip install --extra-index-url https://download.pytorch.org/whl/cu118 torch==2.7.1+cu118
+```
+Run the example script:
+```bash
+python example.py
+```
+Adaptive scaling demo:
+The legacy `progressive_scaleup.py` script is retained for reference but has been
+superseded by `integration_schedule.py`, which offers a more flexible scaling
+workflow.
+
+Run the unified workflow:
+```bash
+python unified_workflow.py --dashboard
+# disable gradient checkpointing for faster but memory-hungry runs
+python unified_workflow.py --no-checkpoint
+# use standard (non-reversible) transformer blocks
+python unified_workflow.py --no-reversible
+# enable 4-bit quantization-aware training
+python unified_workflow.py --qat
+```
+
+For faster CPU execution, BitTransformerLM exposes a `cpu_autocast()` helper
+that enables bfloat16 mixed precision. Models created with
+`use_autocast=True` apply this automatically, or you can wrap individual
+forward passes:
+
+```python
+from bit_transformer.torch_utils import cpu_autocast
+
+with cpu_autocast():
+    logits, telemetry = model(bits)
+```
+
+Reduce memory use when chunked attention is active by disabling full
+attention logging:
+
+```python
+model = BitTransformerLM(chunk_size=128, full_attn_logging=False)
+```
+
+Enable Diffusion LM training and sampling:
+```bash
+python unified_workflow.py --diffusion --diffusion-steps 8 --dataset-size 32
+# choose noise schedule: linear, cosine, exp
+python unified_workflow.py --diffusion --noise-schedule cosine --diffusion-steps 16 --dataset-size 32
+# linearly decay noise over epochs
+python unified_workflow.py --diffusion --diffusion-curriculum --dataset-size 32
+```
+Higher `--diffusion-steps` (8–16) improves sample quality at the cost of compute. When using the dashboard, enable the **Diffusion LM** toggle to run the model without causal masking or chunked attention.
+Generated samples automatically fix parity bits so they can be decoded back to text.
+To resume training across machines using Hugging Face storage:
+```bash
+python unified_workflow.py --hf-repo your-username/bittransformerlm --hf-token $HF_TOKEN
+```
+The dashboard exposes matching controls under **Hugging Face Checkpoints**. Provide a repository ID and optional token (falling back to the `HF_TOKEN` environment variable) and click **Upload weights** or **Download weights** to sync the model.
+Run the unit tests:
+```bash
+pytest -q
+```
+
+### Mode management
+
+During training, ensure the model is in training mode with dropout enabled:
+
+```python
+from bit_transformer.utils import set_dropout
+
+model.train()
+set_dropout(model, 0.1)
+```
+
+Before running tests, performing inference, or committing weights to the repository, switch the model to evaluation mode and disable dropout:
+
+```python
+model.eval()
+set_dropout(model, 0.0)
+```
+
+This prevents CI failures from accidentally pushing weights that still have active dropout.
+
+## Telemetry Metrics Explained
+BitTransformerLM reports three bounded metrics in ``[0, 1]`` during training and inference:
+
+- **Negentropy (K)** – departure from random noise; ``1`` denotes perfectly ordered bits while ``0`` is uniform randomness.
+- **LZ Complexity (C)** – differentiable proxy for Lempel–Ziv compressibility; low values imply repetitive patterns and high values frequent transitions.
+- **Symbiosis (S)** – agreement between model predictions and a reference distribution via KL divergence; scores near ``1`` show strong alignment.
+
+An Adaptive Computation Time (ACT) mechanism lets layers halt early once confidence exceeds a threshold. Halt probabilities are exported as ``halt_probs`` in telemetry for inspection.
+
+These metrics are logged alongside losses and can trigger safety gates when thresholds are violated. The dashboard monitors drift and emits warnings when recent values deviate beyond a configurable threshold.
+
+## Core Features
+- **Bit-Native Modeling** – Works directly on 0/1 inputs with positional encodings and parity-protected text helpers.
+- **Telemetry Synthesizer** – Clusters activation summaries to surface coherent subspaces and detect drift.
+- **Submodel Distillation** – `TelemetrySynthesizer` selects representative sequences for `collapse_submodel`, which deepens
+  and widens once (`width_scale` = 1.5) if telemetry floors aren't met; `save_distilled_model` places a `metrics.json` summary
+  beside the distilled weights.
+- **Safety Gate** – `hil_safe_inference` enforces minimum complexity and symbiosis scores at runtime with EMA smoothing and a configurable burn‑in period.
+- **Quantization** – CPU inference can be quantized to int8 or trained with 4-bit QAT using the `--qat` flag.
+- **Distributed Training** – FSDP and pipeline helpers allow multi‑GPU scaling when hardware is available.
+- **Interactive Dashboard** – Live control of training, scaling and compression with optional GPU acceleration. The dashboard now exposes reversible layers, gradient checkpointing, ACT thresholds, λ floors, 4‑bit QAT and Diffusion LM toggles, real‑time telemetry charts powered by Chart.js, and Hugging Face checkpoint upload/download controls with `HF_TOKEN` fallback. Settings persist via `localStorage`.
+- **CI/CD Pipeline** – GitHub Actions install dependencies, run the tests and build distribution artifacts on every push.
+
+## Development Workflow
+1. Start the MCP server:
+   ```bash
+   python mcp_server.py
+   ```
+2. Launch the dashboard in another terminal:
+   ```bash
+   MCP_SERVER_ADDR=http://127.0.0.1:7000 python -m bit_transformer.dashboard_app
+   ```
+3. Submit training batches, scale the model and monitor telemetry from the web UI.
+   The dashboard's appearance is controlled by `bit_transformer/static/style.css`.
+
+A `watcher.py` script can automatically restart the server and run tests when files change during local development.
+
+## Container Deployment
+A `Dockerfile` and `start.sh` script build a minimal VM image that launches both the MCP server and dashboard.
+
+```bash
+docker build -t bittransformerlm .
+docker run -p 5000:5000 -p 7000:7000 bittransformerlm
+```
+
+By default the container installs the CPU-only PyTorch wheel. Set the build
+argument `TORCH_CUDA=cu118` to preinstall the GPU version. The container sets
+`MCP_SERVER_ADDR=http://127.0.0.1:7000` and exposes the dashboard on port 5000.
+
+## Roadmap
+- Finalize S attribution tools and metric drift detection.
+- Publish an initial release package and rename the repository to **BitTransformerLM**.
+- Continue benchmarking on real datasets and expanding context window capabilities.
+
+## Licensing
+
+This project is released under a combination of licenses and agreements to provide a clear framework for use, distribution, and contribution. All licensing documents can be found in the `LICENSE/` directory.
+
+The key documents are:
+
+* `LICENSE.txt`: The primary open-source license for the software, AGPLv3.
+* `COMMERCIAL_LICENSE.txt`: Terms for commercial use of the software.
+* `DISCLAIMER.txt`: Important legal disclaimers.
+* `ALIGNMENT_AND_TRANSPARENCY.txt`: Our commitment to alignment and transparency.
+* `TRADEMARK_POLICY.txt`: Guidelines for using the project's trademarks.
+* `CONTRIBUTOR_LICENSE_AGREEMENT.txt`: The agreement for all contributors to sign.
+
+Please review these documents carefully before using or contributing to the project.
+

BitTransformerLM/bit_transformer/__init__.py

@@ -0,0 +1,86 @@
+from .model import (
+    PositionalEncoding,
+    BitTransformerLM,
+    ReversibleLoggingTransformerEncoderLayer,
+    example_usage,
+    example_training_step,
+    infer_long_sequence,
+    diffusion_inference,
+)
+from .telemetry import TelemetrySynthesizer, detect_metric_drift
+from .dashboard import plot_telemetry
+from .dashboard_app import run_dashboard
+from .collapse import collapse_submodel, save_distilled_model
+from .safety import hil_safe_inference, demo_hil_safety, safe_sample_with_retry
+from .bit_io import (
+    text_to_bits,
+    bits_to_text,
+    infer_text,
+)
+from .parity import enforce_parity
+from .compression import (
+    compress_bits,
+    decompress_bits,
+    model_output_decompress,
+    pack_bits,
+    unpack_bits,
+)
+from .distributed import wrap_fsdp, make_pipeline
+from .optimization import configure_optimizer, adjust_learning_rate
+from .scale import expand_model
+from .distil import distill_step, TelemetryLog
+from .quantization import (
+    quantize_dynamic,
+    prepare_qat_fx,
+    convert_qat_fx,
+)
+from .training import train_loop
+from .utils import save_model, load_model, set_dropout
+from .hf_checkpoint import hf_login, save_checkpoint, download_checkpoint
+from .torch_utils import cpu_autocast
+
+__all__ = [
+    "PositionalEncoding",
+    "BitTransformerLM",
+    "ReversibleLoggingTransformerEncoderLayer",
+    "example_usage",
+    "example_training_step",
+    "TelemetrySynthesizer",
+    "detect_metric_drift",
+    "collapse_submodel",
+    "save_distilled_model",
+    "hil_safe_inference",
+    "demo_hil_safety",
+    "safe_sample_with_retry",
+    "text_to_bits",
+    "bits_to_text",
+    "infer_text",
+    "enforce_parity",
+    "plot_telemetry",
+    "run_dashboard",
+    "configure_optimizer",
+    "adjust_learning_rate",
+    "expand_model",
+    "distill_step",
+    "TelemetryLog",
+    "quantize_dynamic",
+    "prepare_qat_fx",
+    "convert_qat_fx",
+    "train_loop",
+    "wrap_fsdp",
+    "make_pipeline",
+    "compress_bits",
+    "decompress_bits",
+    "model_output_decompress",
+    "pack_bits",
+    "unpack_bits",
+    "infer_long_sequence",
+    "diffusion_inference",
+    "save_model",
+    "load_model",
+    "set_dropout",
+    "hf_login",
+    "save_checkpoint",
+    "download_checkpoint",
+    "cpu_autocast",
+]

BitTransformerLM/bit_transformer/bit_io.py

@@ -0,0 +1,97 @@
+from typing import List, TYPE_CHECKING
+import torch
+import sys
+
+try:  # torch.compile may be unavailable or unsupported
+    if torch.__version__ and tuple(map(int, torch.__version__.split(".")[:2])) >= (2, 0) and sys.version_info < (3, 11):
+        compile_fn = torch.compile
+    else:
+        raise RuntimeError
+except Exception:  # pragma: no cover
+
+    def compile_fn(fn=None, **kwargs):
+        if fn is None:
+            return lambda f: f
+        return fn
+
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .model import BitTransformerLM
+
+
+@compile_fn
+def bytes_to_bits(data: bytes) -> List[int]:
+    """Convert bytes to bits with per-byte parity bit."""
+    result: List[int] = []
+    for b in data:
+        bits = [(b >> i) & 1 for i in reversed(range(8))]
+        parity = sum(bits) % 2
+        result.extend(bits + [parity])
+    return result
+
+
+@compile_fn
+def bits_to_bytes(bits: List[int]) -> bytes:
+    """Convert parity-protected bits back to bytes."""
+    if len(bits) % 9 != 0:
+        raise ValueError("Bit stream length must be multiple of 9")
+    out = bytearray()
+    for i in range(0, len(bits), 9):
+        chunk = bits[i : i + 9]
+        payload = chunk[:8]
+        parity = chunk[8]
+        if parity != sum(payload) % 2:
+            raise ValueError("Parity check failed")
+        value = 0
+        for bit in payload:
+            value = (value << 1) | bit
+        out.append(value)
+    return bytes(out)
+
+
+def text_to_bits(text: str) -> List[int]:
+    return bytes_to_bits(text.encode("utf-8"))
+
+
+def bits_to_text(bits: List[int]) -> str:
+    return bits_to_bytes(bits).decode("utf-8", errors="replace")
+
+
+def infer_text(
+    model: "BitTransformerLM",
+    text: str,
+    c_floor: float = 0.3,
+    s_floor: float = 0.5,
+) -> str:
+    """Run text through the model using the safety gate."""
+    from .safety import hil_safe_inference
+    bits = text_to_bits(text)
+    tensor = torch.tensor(bits, dtype=torch.long).unsqueeze(0)
+    out_bits, _ = hil_safe_inference(model, tensor, c_floor=c_floor, s_floor=s_floor)
+    return bits_to_text(out_bits.squeeze(0).tolist())
+
+
+def sample_text(
+    model: "BitTransformerLM",
+    prompt: str,
+    max_new_tokens: int = 16,
+    temperature: float = 1.0,
+    top_p: float = 1.0,
+) -> str:
+    """Generate text from the model using simple top-p sampling."""
+    model.eval()
+    bits = text_to_bits(prompt)
+    tensor = torch.tensor(bits, dtype=torch.long).unsqueeze(0)
+    for _ in range(max_new_tokens * 9):
+        if tensor.size(1) >= model.pos_enc.pe.size(0):
+            break
+        logits, _ = model(tensor, causal=True)
+        prob = logits[0, -1].softmax(-1) / temperature
+        sorted_prob, sorted_idx = prob.sort(descending=True)
+        cumulative = sorted_prob.cumsum(0)
+        mask = cumulative > top_p
+        sorted_prob[mask] = 0
+        sorted_prob = sorted_prob / sorted_prob.sum()
+        next_bit = sorted_idx[torch.multinomial(sorted_prob, 1)]
+        tensor = torch.cat([tensor, next_bit.view(1, 1)], dim=1)
+    return bits_to_text(tensor.squeeze(0).tolist())

BitTransformerLM/bit_transformer/collapse.py

@@ -0,0 +1,95 @@
+import json
+import os
+from typing import Dict, List, Optional, Tuple
+
+import torch
+
+from .model import BitTransformerLM
+from .training import train_loop
+
+
+def collapse_submodel(
+    cluster_data: List[List[int]],
+    target_params: Dict,
+    floors: Optional[Dict[str, float]] = None,
+    max_rounds: int = 3,
+    width_scale: float = 1.5,
+    forward_kwargs: Optional[Dict] = None,
+) -> Tuple[BitTransformerLM, Dict[str, float]]:
+    """Distill a submodel from clustered bit sequences.
+
+    The routine deepens the target model when telemetry floors are unmet and,
+    after the first deepening fails, widens the hidden dimensions by
+    ``width_scale`` once before retrying. Returns the distilled model and its
+    final telemetry metrics.
+    """
+    if floors is None:
+        floors = {"negentropy": 0.5, "lz_complexity": 0.3, "symbiosis_score": 0.5}
+
+    bit_tensor = torch.tensor(cluster_data, dtype=torch.long)
+    n = len(bit_tensor)
+    split = max(1, int(0.8 * n))
+    train_bits = bit_tensor[:split]
+    val_bits = bit_tensor[split:]
+    if len(val_bits) == 0:
+        val_bits = train_bits
+
+    params = target_params.copy()
+    metrics: Dict[str, float] = {}
+    width_scaled = False
+    for round_idx in range(max_rounds):
+        model = BitTransformerLM(**params)
+        train_loop(
+            model,
+            train_bits,
+            epochs=2,
+            compress_prob=0.5,
+            direct_prob=0.0,
+            log=False,
+            forward_kwargs=forward_kwargs,
+        )
+        with torch.no_grad():
+            logits, telemetry = model(val_bits, **(forward_kwargs or {}))
+            neg_k = model.negentropy_logits(logits).mean().item()
+            lz_c = model.lz_complexity_logits(logits).mean().item()
+            sym_s = telemetry["symbiosis_score"].mean().item()
+        metrics = {
+            "negentropy": neg_k,
+            "lz_complexity": lz_c,
+            "symbiosis_score": sym_s,
+        }
+        if (
+            neg_k >= floors["negentropy"]
+            and lz_c >= floors["lz_complexity"]
+            and sym_s >= floors["symbiosis_score"]
+        ):
+            break
+        if round_idx == 0:
+            params["num_layers"] = max(1, params.get("num_layers", 1)) + 1
+        elif not width_scaled:
+            params["d_model"] = int(params.get("d_model", 32) * width_scale)
+            params["dim_feedforward"] = int(
+                params.get("dim_feedforward", 64) * width_scale
+            )
+            width_scaled = True
+        else:
+            params["num_layers"] = max(1, params.get("num_layers", 1)) + 1
+    return model, metrics
+
+
+def save_distilled_model(
+    model: BitTransformerLM,
+    path: str,
+    metrics: Dict[str, float],
+    floors: Optional[Dict[str, float]] = None,
+) -> None:
+    """Serialize a distilled model and its metric summary to disk.
+
+    Weights are written to ``path`` and a ``metrics.json`` file is placed in the
+    same directory containing the achieved metrics alongside the target floors.
+    """
+    torch.save(model.state_dict(), path)
+    payload = {"metrics": metrics, "floors": floors or {}}
+    metrics_path = os.path.join(os.path.dirname(path), "metrics.json")
+    with open(metrics_path, "w") as f:
+        json.dump(payload, f)

BitTransformerLM/bit_transformer/compression.py

@@ -0,0 +1,82 @@
+import torch
+from typing import List
+
+
+def compress_bits(bits: torch.Tensor) -> torch.Tensor:
+    """Run-length encode a 1D tensor of bits.
+
+    Args:
+        bits: 1D tensor with values 0 or 1 (bool or uint8).
+
+    Returns:
+        1D uint8 tensor containing interleaved values and run lengths.
+    """
+    if bits.dim() != 1:
+        raise ValueError("compress_bits expects a 1D tensor")
+    b = bits.to(torch.uint8).flatten()
+    if b.numel() == 0:
+        return b
+    changes = torch.nonzero(b[1:] != b[:-1]).flatten().to(torch.long) + 1
+    starts = torch.cat([b.new_tensor([0], dtype=torch.long), changes])
+    ends = torch.cat([changes, b.new_tensor([b.numel()], dtype=torch.long)])
+    values = b[starts.to(torch.long)]
+    counts = ends - starts
+
+    out_vals: List[int] = []
+    out_counts: List[int] = []
+    for v, c in zip(values.tolist(), counts.tolist()):
+        while c > 255:
+            out_vals.append(v)
+            out_counts.append(255)
+            c -= 255
+        out_vals.append(v)
+        out_counts.append(c)
+    values_tensor = torch.tensor(out_vals, dtype=torch.uint8)
+    counts_tensor = torch.tensor(out_counts, dtype=torch.uint8)
+    out = torch.stack([values_tensor, counts_tensor], dim=1).flatten()
+    return out
+
+
+def decompress_bits(compressed: torch.Tensor) -> torch.Tensor:
+    """Decode a run-length encoded bit tensor."""
+    if compressed.dim() != 1 or compressed.numel() % 2 != 0:
+        raise ValueError("compressed tensor must be 1D even-length")
+    data = compressed.to(torch.uint8)
+    values = data[0::2]
+    counts = data[1::2].to(torch.long)
+    return torch.repeat_interleave(values, counts)
+
+
+def model_output_decompress(compressed_batch) -> torch.Tensor:
+    """Decompress a batch of compressed bit sequences."""
+    if isinstance(compressed_batch, torch.Tensor) and compressed_batch.dim() == 1:
+        sequences = [decompress_bits(compressed_batch)]
+    else:
+        sequences = [decompress_bits(row) for row in compressed_batch]
+    lengths = [seq.numel() for seq in sequences]
+    if len(set(lengths)) != 1:
+        raise ValueError("Sequences decompress to different lengths")
+    return torch.stack(sequences)
+
+
+import numpy as np
+
+
+def pack_bits(bits: torch.Tensor) -> torch.Tensor:
+    """Pack groups of 8 bits into uint8 values using numpy.packbits."""
+    if bits.dim() != 1:
+        raise ValueError("pack_bits expects a 1D tensor")
+    arr = bits.to(torch.uint8).cpu().numpy()
+    packed = np.packbits(arr)
+    return torch.from_numpy(packed)
+
+
+def unpack_bits(packed: torch.Tensor, *, n_bits: int | None = None) -> torch.Tensor:
+    """Unpack uint8 values back into a bit tensor."""
+    if packed.dim() != 1:
+        raise ValueError("unpack_bits expects a 1D tensor")
+    arr = np.unpackbits(packed.to(torch.uint8).cpu().numpy())
+    if n_bits is not None:
+        arr = arr[:n_bits]
+    return torch.from_numpy(arr)
+

BitTransformerLM/bit_transformer/dashboard.py

@@ -0,0 +1,58 @@
+import matplotlib.pyplot as plt
+from typing import Dict, List, Tuple
+
+
+def plot_telemetry(
+    metrics_log: Dict[str, List[float]],
+    k_floor: float = 0.5,
+    c_floor: float = 0.3,
+    s_floor: float = 0.5,
+) -> Tuple[plt.Figure, List[plt.Axes]]:
+    """Plot K, C, S metrics over time with cluster transitions.
+
+    Args:
+        metrics_log: Dictionary with keys ``negentropy``, ``lz_complexity``,
+            ``symbiosis_score`` and optional ``clusters`` listing cluster
+            assignments per step.
+        k_floor: Threshold for negentropy (K).
+        c_floor: Threshold for LZ complexity (C).
+        s_floor: Threshold for symbiosis score (S).
+
+    Returns:
+        (figure, axes) tuple for further customization or saving.
+    """
+    steps = list(range(len(metrics_log.get("negentropy", []))))
+    fig, axes = plt.subplots(3, 1, sharex=True, figsize=(10, 6))
+    metrics = [
+        ("negentropy", k_floor, "K"),
+        ("lz_complexity", c_floor, "C"),
+        ("symbiosis_score", s_floor, "S"),
+    ]
+    for ax, (key, floor, label) in zip(axes, metrics):
+        values = metrics_log.get(key, [])
+        ax.plot(steps, values, label=label)
+        ax.axhline(floor, color="r", linestyle="--", linewidth=1)
+        violations = [i for i, v in enumerate(values) if v < floor]
+        if violations:
+            ax.scatter(
+                [steps[i] for i in violations],
+                [values[i] for i in violations],
+                color="r",
+                zorder=5,
+                label="violation",
+            )
+        ax.set_ylabel(label)
+        ax.legend(loc="upper right")
+
+    clusters = metrics_log.get("clusters")
+    if clusters is not None:
+        prev = clusters[0]
+        for t, c in enumerate(clusters):
+            if t > 0 and c != prev:
+                for ax in axes:
+                    ax.axvline(t, color="gray", linestyle=":", alpha=0.5)
+                prev = c
+
+    axes[-1].set_xlabel("step")
+    plt.tight_layout()
+    return fig, axes

BitTransformerLM/bit_transformer/dashboard_app.py

@@ -0,0 +1,927 @@
+import io
+import json
+import os
+import traceback
+import inspect
+from typing import Any, Dict, List
+
+from flask import Flask, jsonify, request, render_template, send_file
+import subprocess
+import sys
+import warnings
+import matplotlib.pyplot as plt
+import torch
+import torch.nn.functional as F
+import requests
+import gzip
+
+from .model import BitTransformerLM, infer_long_sequence
+from .optimization import configure_optimizer
+from .collapse import collapse_submodel
+from .dashboard import plot_telemetry
+from .scale import expand_model
+from .bit_io import text_to_bits, bits_to_text
+from .safety import hil_safe_inference
+from .compression import model_output_decompress, compress_bits
+from .distributed import wrap_fsdp
+from .training import train_loop
+from .telemetry import detect_metric_drift
+from .quantization import prepare_qat_fx, convert_qat_fx
+from torch.distributed.fsdp import FullyShardedDataParallel
+from .hf_checkpoint import hf_login, save_checkpoint, download_checkpoint
+
+
+app = Flask(__name__)
+app.config["MAX_CONTENT_LENGTH"] = 1 * 1024 * 1024  # 1MB upload limit
+
+MCP_SERVER_ADDR = os.getenv("MCP_SERVER_ADDR")
+
+
+@app.errorhandler(Exception)
+def handle_exception(err):
+    """Return JSON error responses with stack traces."""
+    return (
+        jsonify({"error": str(err), "trace": traceback.format_exc()}),
+        getattr(err, "code", 500),
+    )
+
+class MetricDriftWarning(UserWarning):
+    """Raised when telemetry metrics drift beyond the configured threshold."""
+
+def _switch_torch(use_gpu: bool) -> None:
+    """Install the appropriate PyTorch wheel and restart the process."""
+    have_cuda = torch.version.cuda is not None
+    if use_gpu == have_cuda:
+        return
+    wheel = "torch==2.7.1+cu118" if use_gpu else "torch==2.7.1+cpu"
+    url = "https://download.pytorch.org/whl/cu118" if use_gpu else "https://download.pytorch.org/whl/cpu"
+    subprocess.run([
+        sys.executable,
+        "-m",
+        "pip",
+        "install",
+        "--extra-index-url",
+        url,
+        wheel,
+    ], check=True)
+    os.execv(sys.executable, [sys.executable] + sys.argv)
+
+def mcp_post(path: str, data=None):
+    if not MCP_SERVER_ADDR:
+        return None
+    url = MCP_SERVER_ADDR.rstrip("/") + path
+    resp = requests.post(url, json=data)
+    resp.raise_for_status()
+    if resp.headers.get("Content-Type", "").startswith("image/"):
+        return resp.content
+    return resp.json()
+
+def mcp_get(path: str):
+    if not MCP_SERVER_ADDR:
+        return None
+    url = MCP_SERVER_ADDR.rstrip("/") + path
+    resp = requests.get(url)
+    resp.raise_for_status()
+    if resp.headers.get("Content-Type", "").startswith("image/"):
+        return resp.content
+    return resp.json()
+
+class ModelManager:
+    """Manage model state and training utilities for the dashboard."""
+
+    def __init__(
+        self,
+        snapshot_dir: str | None = None,
+        telemetry_log: str | None = None,
+        *,
+        drift_window: int = 10,
+        drift_threshold: float = 0.2,
+    ) -> None:
+        self.snapshot_dir = snapshot_dir or os.getenv("SNAPSHOT_DIR", "snapshots")
+        self.telemetry_log = telemetry_log or os.getenv("TELEMETRY_LOG")
+        if self.telemetry_log is None:
+            self.telemetry_log = os.path.join(self.snapshot_dir, "metrics.json")
+        os.makedirs(self.snapshot_dir, exist_ok=True)
+        self.weights_path = os.path.join(self.snapshot_dir, "model.pt")
+
+        self.model: BitTransformerLM | None = None
+        self.optimizer: torch.optim.Optimizer | None = None
+        self.scheduler: torch.optim.lr_scheduler._LRScheduler | None = None
+        self.total_steps = 100
+        self.metrics: Dict[str, List[float]] = {
+            "negentropy_logits": [],
+            "lz_complexity_logits": [],
+            "symbiosis_score": [],
+        }
+        self.drift_window = drift_window
+        self.drift_threshold = drift_threshold
+        self.lambda_K = 1.0
+        self.lambda_C = 1.0
+        self.lambda_S = 1.0
+        self.c_floor = 0.3
+        self.s_floor = 0.5
+        self.causal = True
+        self.diffusion = False
+        self.decompress_output = False
+        self.use_compression = False
+        self.use_gpu = False
+        self.qat = False
+
+        # Load any existing state
+        if os.path.exists(self.telemetry_log):
+            try:
+                with open(self.telemetry_log) as f:
+                    saved = json.load(f)
+                for key in self.metrics:
+                    self.metrics[key] = saved.get(key, [])
+            except Exception:
+                pass
+        if os.path.exists(self.weights_path):
+            try:
+                self.model = torch.load(self.weights_path, map_location="cpu")
+                self.optimizer, self.scheduler = configure_optimizer(
+                    self.model, lr=1e-3, total_steps=self.total_steps
+                )
+                self._apply_device()
+            except Exception:
+                self.model = None
+
+        config_path = os.getenv("MODEL_CONFIG", "/config/model_params.json")
+        if self.model is None and os.path.exists(config_path):
+            try:
+                with open(config_path) as f:
+                    params = json.load(f)
+                self.init_model(params)
+            except Exception:
+                pass
+
+    def init_model(self, params: Dict) -> None:
+        int_fields = {
+            "d_model",
+            "nhead",
+            "num_layers",
+            "dim_feedforward",
+            "max_seq_len",
+            "chunk_size",
+            "overlap",
+        }
+        float_fields = {"act_threshold"}
+        bool_fields = {"reversible", "use_checkpoint"}
+        clean: Dict[str, Any] = {}
+        for k, v in params.items():
+            if v is None:
+                clean[k] = None
+            elif k in int_fields:
+                clean[k] = int(v)
+            elif k in float_fields:
+                clean[k] = float(v)
+            elif k in bool_fields:
+                clean[k] = bool(v)
+            else:
+                clean[k] = v
+        self.model = BitTransformerLM(
+            **clean,
+            lambda_K=self.lambda_K,
+            lambda_C=self.lambda_C,
+            lambda_S=self.lambda_S,
+        )
+        self.optimizer, self.scheduler = configure_optimizer(
+            self.model, lr=1e-3, total_steps=self.total_steps
+        )
+        self._apply_device()
+        for key in self.metrics:
+            self.metrics[key].clear()
+
+    def set_lambdas(self, k: float, c: float, s: float) -> None:
+        """Update λ weights and propagate to the model."""
+        self.lambda_K = k
+        self.lambda_C = c
+        self.lambda_S = s
+        if self.model is not None:
+            self.model.set_lambdas(k, c, s)
+
+    def set_floors(self, c_floor: float, s_floor: float) -> None:
+        """Update safety floors for complexity (C) and symbiosis (S)."""
+        self.c_floor = c_floor
+        self.s_floor = s_floor
+
+    def set_diffusion(self, flag: bool) -> None:
+        """Toggle Diffusion LM mode which disables causal masking and chunking."""
+        self.diffusion = flag
+        self.causal = not flag
+        if self.model is not None and flag:
+            self.model.chunk_size = None
+
+    def set_decompress_output(self, flag: bool) -> None:
+        """Enable or disable decompression of model outputs."""
+        self.decompress_output = flag
+
+    def set_compression(self, flag: bool) -> None:
+        """Toggle automatic compression of inputs."""
+        self.use_compression = flag
+
+    def set_qat(self, flag: bool) -> None:
+        """Enable or disable 4-bit quantization-aware training."""
+        self.qat = flag
+        if self.model is None:
+            return
+        if flag:
+            self.model = prepare_qat_fx(self.model)
+        else:
+            self.model = convert_qat_fx(self.model)
+
+    def set_gpu(self, flag: bool) -> None:
+        """Toggle GPU acceleration and FSDP, reinstalling PyTorch if needed."""
+        _switch_torch(flag)
+        self.use_gpu = flag and torch.cuda.is_available()
+        self._apply_device()
+
+    def _apply_device(self) -> None:
+        """Move the model to the selected device and wrap with FSDP if needed."""
+        if self.model is None:
+            return
+        if self.use_gpu:
+            device = torch.device("cuda")
+            if isinstance(self.model, FullyShardedDataParallel):
+                base = self.model.module
+            else:
+                base = self.model
+            base = base.to(device)
+            self.model = wrap_fsdp(base, device_id=device)
+        else:
+            device = torch.device("cpu")
+            if isinstance(self.model, FullyShardedDataParallel):
+                self.model = self.model.module
+            self.model = self.model.to(device)
+
+    def train_step(self, bits: torch.Tensor) -> float:
+        assert (
+            self.model is not None
+            and self.optimizer is not None
+            and self.scheduler is not None
+        )
+        self.model.train()
+        device = next(self.model.parameters()).device
+        bits = bits.to(device)
+        ratio = 1.0
+        if self.use_compression:
+            comps = [compress_bits(row.to(torch.uint8)) for row in bits]
+            comp_len = sum(c.numel() for c in comps)
+            ratio = min(comp_len / bits.numel(), 1.0)
+            logits, telemetry = self.model.forward_compressed(comps, causal=self.causal)
+        else:
+            logits, telemetry = self.model(bits, causal=self.causal)
+        pred = logits[:, :-1, :].reshape(-1, 2)
+        target = bits[:, 1:].reshape(-1)
+        loss = F.cross_entropy(pred, target)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
+        self.optimizer.step()
+        self.scheduler.step()
+        self.optimizer.zero_grad()
+        self._log_metrics(telemetry)
+        self._save_state()
+        return loss.item(), ratio
+
+    def train_epochs(
+        self,
+        bits: torch.Tensor,
+        *,
+        epochs: int = 1,
+        compress_prob: float = 0.5,
+        direct_prob: float = 0.0,
+        batch_size: int = 8,
+        num_workers: int = 0,
+        accum_steps: int = 1,
+        amp: bool = False,
+        compile_model: bool = False,
+    ) -> List[Dict[str, float]]:
+        """Run ``train_loop`` on a batch tensor and persist the state."""
+        assert self.model is not None
+        device = next(self.model.parameters()).device
+        bits = bits.to(device)
+        import math
+        steps_per_epoch = max(1, math.ceil(len(bits) / batch_size))
+        self.total_steps = math.ceil(epochs * steps_per_epoch / accum_steps)
+        self.optimizer, self.scheduler = configure_optimizer(
+            self.model, lr=1e-3, total_steps=self.total_steps
+        )
+        metrics = train_loop(
+            self.model,
+            bits,
+            epochs=epochs,
+            compress_prob=compress_prob if self.use_compression else 0.0,
+            direct_prob=direct_prob,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            accum_steps=accum_steps,
+            amp=amp,
+            compile_model=compile_model,
+            forward_kwargs={"causal": self.causal},
+            optimizer=self.optimizer,
+            scheduler=self.scheduler,
+        )
+        self._save_state()
+        return metrics
+
+    def scale_up(self, width_mult: float = 1.0) -> None:
+        assert self.model is not None
+        params = dict(
+            d_model=int(self.model.d_model * width_mult),
+            nhead=self.model.layers[0].self_attn.num_heads,
+            num_layers=self.model.num_layers * 2,
+            dim_feedforward=int(self.model.layers[0].linear1.out_features * width_mult),
+            max_seq_len=self.model.pos_enc.pe.size(0),
+        )
+        self.model = expand_model(self.model, {
+            **params,
+            "lambda_K": self.lambda_K,
+            "lambda_C": self.lambda_C,
+            "lambda_S": self.lambda_S,
+        })
+        self.optimizer, self.scheduler = configure_optimizer(
+            self.model, lr=1e-3, total_steps=self.total_steps
+        )
+        self._save_state()
+
+    def collapse(self, cluster_bits: List[List[int]], target_params: Dict, width_scale: float = 1.0) -> None:
+        self.model, _ = collapse_submodel(
+            cluster_bits,
+            target_params,
+            width_scale=width_scale,
+            forward_kwargs={"causal": self.causal},
+        )
+        self.model.set_lambdas(self.lambda_K, self.lambda_C, self.lambda_S)
+        self.optimizer, self.scheduler = configure_optimizer(
+            self.model, lr=1e-3, total_steps=self.total_steps
+        )
+        self._apply_device()
+        for key in self.metrics:
+            self.metrics[key].clear()
+
+    def infer(self, bits: torch.Tensor) -> Dict:
+        assert self.model is not None
+        self.model.eval()
+        device = next(self.model.parameters()).device
+        bits = bits.to(device)
+        ratio = 1.0
+        with torch.no_grad():
+            if self.use_compression:
+                comps = [compress_bits(row.to(torch.uint8)) for row in bits]
+                comp_len = sum(c.numel() for c in comps)
+                ratio = min(comp_len / bits.numel(), 1.0)
+                logits, telemetry = self.model.forward_compressed(comps, causal=self.causal)
+            else:
+                logits, telemetry = self.model(bits, causal=self.causal)
+        self._log_metrics(telemetry)
+        pred_bits = logits.argmax(-1)
+        if self.decompress_output:
+            try:
+                pred_bits = model_output_decompress(pred_bits)
+            except Exception as e:
+                return {"error": f"Decompression failed: {e}", "suggestion": "Disable compression toggle."}
+        def _to_python(obj):
+            if isinstance(obj, torch.Tensor):
+                return obj.tolist()
+            if isinstance(obj, list):
+                return [_to_python(o) for o in obj]
+            if isinstance(obj, dict):
+                return {kk: _to_python(vv) for kk, vv in obj.items()}
+            return obj
+        tele = {k: _to_python(v) for k, v in telemetry.items()}
+        return {"predicted": pred_bits.squeeze(0).tolist(), "telemetry": tele, "ratio": ratio}
+
+    def infer_long(self, bits: torch.Tensor, ctx_bits: int = 4096, overlap: int = 256) -> Dict:
+        """Run sliding-window inference on a long sequence."""
+        assert self.model is not None
+        device = next(self.model.parameters()).device
+        bits = bits.to(device)
+        preds, logs = infer_long_sequence(self.model, bits.squeeze(0), ctx_bits=ctx_bits, overlap=overlap)
+        for tele in logs:
+            self._log_metrics(tele)
+        return {"predicted": preds.tolist(), "windows": len(logs)}
+
+    def _log_metrics(self, telemetry: Dict) -> None:
+        for key in self.metrics:
+            val = telemetry[key].mean().item()
+            self.metrics[key].append(val)
+        drift = detect_metric_drift(
+            self.metrics, window=self.drift_window, threshold=self.drift_threshold
+        )
+        bad = [k for k, v in drift.items() if v]
+        if bad:
+            warnings.warn(
+                f"Metric drift detected: {', '.join(bad)}",
+                MetricDriftWarning,
+            )
+
+    def infer_text(self, text: str) -> Dict[str, Any]:
+        """Run text through the model using the safety gate."""
+        assert self.model is not None
+        device = next(self.model.parameters()).device
+        bits = torch.tensor(text_to_bits(text), dtype=torch.long).unsqueeze(0).to(device)
+        out_bits, telemetry = hil_safe_inference(
+            self.model, bits, c_floor=self.c_floor, s_floor=self.s_floor
+        )
+        self._log_metrics(telemetry)
+        return {
+            "output": bits_to_text(out_bits.squeeze(0).tolist()),
+            "telemetry": telemetry,
+        }
+
+    def get_status(self) -> Dict[str, Any]:
+        info: Dict[str, Any] = {
+            "use_gpu": self.use_gpu,
+            "diffusion": self.diffusion,
+            "compression": self.use_compression,
+            "lambda_K": self.lambda_K,
+            "lambda_C": self.lambda_C,
+            "lambda_S": self.lambda_S,
+            "c_floor": self.c_floor,
+            "s_floor": self.s_floor,
+            "qat": self.qat,
+        }
+        if self.model is not None:
+            info.update(
+                {
+                    "d_model": self.model.d_model,
+                    "num_layers": self.model.num_layers,
+                    "d_ff": self.model.layers[0].linear1.out_features,
+                    "nhead": self.model.layers[0].self_attn.num_heads,
+                    "max_seq_len": self.model.pos_enc.pe.size(0),
+                }
+            )
+        else:
+            info.update(
+                {
+                    "d_model": None,
+                    "num_layers": 0,
+                    "d_ff": None,
+                    "nhead": None,
+                    "max_seq_len": None,
+                }
+            )
+        return info
+
+    def get_model_config(self) -> Dict[str, Any]:
+        """Return current model hyperparameters and safety settings."""
+        cfg: Dict[str, Any] = {
+            "lambda_K": self.lambda_K,
+            "lambda_C": self.lambda_C,
+            "lambda_S": self.lambda_S,
+            "c_floor": self.c_floor,
+            "s_floor": self.s_floor,
+        }
+        if self.model is not None:
+            cfg.update(
+                {
+                    "d_model": self.model.d_model,
+                    "nhead": self.model.layers[0].self_attn.num_heads,
+                    "num_layers": self.model.num_layers,
+                    "dim_feedforward": self.model.layers[0].linear1.out_features,
+                    "max_seq_len": self.model.pos_enc.pe.size(0),
+                    "chunk_size": self.model.chunk_size,
+                    "reversible": self.model.reversible,
+                    "use_checkpoint": self.model.use_checkpoint,
+                }
+            )
+        else:
+            cfg.update(
+                {
+                    "d_model": None,
+                    "nhead": None,
+                    "num_layers": 0,
+                    "dim_feedforward": None,
+                    "max_seq_len": None,
+                    "chunk_size": None,
+                    "reversible": None,
+                    "use_checkpoint": None,
+                }
+            )
+        return cfg
+
+    def get_metrics(self) -> Dict[str, Any]:
+        """Return logged telemetry metrics with summary statistics."""
+        from statistics import mean, stdev
+
+        data = {
+            "negentropy": self.metrics["negentropy_logits"],
+            "lz_complexity": self.metrics["lz_complexity_logits"],
+            "symbiosis": self.metrics["symbiosis_score"],
+        }
+        summary: Dict[str, Dict[str, float | None]] = {}
+        for key, values in data.items():
+            if values:
+                m = mean(values)
+                s = stdev(values) if len(values) > 1 else 0.0
+                summary[key] = {"mean": m, "std": s}
+            else:
+                summary[key] = {"mean": None, "std": None}
+        data["summary"] = summary
+        return data
+
+
+    def _save_state(self) -> None:
+        if self.model is None:
+            return
+        torch.save(self.model, self.weights_path)
+        with open(self.telemetry_log, "w") as f:
+            json.dump(self.metrics, f)
+
+
+manager: ModelManager | None = None
+
+
+@app.route("/")
+def index():
+    return render_template(
+        "dashboard.html",
+        metrics=manager.metrics,
+        lambdas={
+            "lambda_K": manager.lambda_K,
+            "lambda_C": manager.lambda_C,
+            "lambda_S": manager.lambda_S,
+        },
+        diffusion=manager.diffusion,
+        compression=manager.use_compression,
+        defaults={k: v.default for k, v in inspect.signature(BitTransformerLM.__init__).parameters.items() if v.default is not inspect._empty},
+        c_floor=manager.c_floor,
+        s_floor=manager.s_floor,
+        qat=manager.qat,
+    )
+
+
+@app.route("/status", methods=["GET"])
+def status():
+    if MCP_SERVER_ADDR:
+        return jsonify(mcp_get("/status"))
+    return jsonify(manager.get_status())
+
+
+@app.route("/model_config", methods=["GET"])
+def model_config():
+    if MCP_SERVER_ADDR:
+        return jsonify(mcp_get("/model_config"))
+    return jsonify(manager.get_model_config())
+
+
+@app.route("/metrics", methods=["GET"])
+def metrics():
+    if MCP_SERVER_ADDR:
+        return jsonify(mcp_get("/metrics"))
+    return jsonify(manager.get_metrics())
+
+
+@app.route("/save_checkpoint", methods=["POST"])
+def save_checkpoint_route():
+    repo_id = request.json.get("repo_id")
+    token = request.json.get("token") or os.getenv("HF_TOKEN")
+    if MCP_SERVER_ADDR:
+        return jsonify(mcp_post("/save_checkpoint", {"repo_id": repo_id, "token": token}))
+    if manager.model is None:
+        return jsonify({"error": "model not initialized"}), 400
+    if token:
+        hf_login(token=token)
+    save_checkpoint(manager.model, repo_id=repo_id)
+    return jsonify({"status": "saved"})
+
+
+@app.route("/download_checkpoint", methods=["POST"])
+def download_checkpoint_route():
+    repo_id = request.json.get("repo_id")
+    token = request.json.get("token") or os.getenv("HF_TOKEN")
+    if MCP_SERVER_ADDR:
+        return jsonify(mcp_post("/download_checkpoint", {"repo_id": repo_id, "token": token}))
+    if token:
+        hf_login(token=token)
+    dest = manager.weights_path + ".gz"
+    ok = download_checkpoint(dest, repo_id=repo_id)
+    if not ok:
+        return jsonify({"status": "failed"}), 500
+    if manager.model is None:
+        return jsonify({"status": "downloaded", "loaded": False})
+    with gzip.open(dest, "rb") as f:
+        state = torch.load(f, map_location="cpu")
+    manager.model.load_state_dict(state)
+    manager.optimizer, manager.scheduler = configure_optimizer(
+        manager.model, lr=1e-3, total_steps=manager.total_steps
+    )
+    manager._apply_device()
+    manager._save_state()
+    return jsonify({"status": "downloaded", "loaded": True})
+
+
+@app.route("/text_to_bits", methods=["POST"])
+def text_to_bits_route():
+    text = request.json.get("text", "")
+    if len(text) > 100_000:
+        return jsonify({"error": "text too large"}), 413
+    return jsonify({"bits": text_to_bits(text)})
+
+
+@app.route("/dataset", methods=["GET"])
+def dataset_route():
+    name = request.args.get("name", "")
+    split = request.args.get("split", "train")
+    size = int(request.args.get("size", 1))
+    seq_len = int(request.args.get("seq_len", 64))
+    if size * seq_len > 1_000_000:
+        return jsonify({"error": "dataset too large"}), 413
+    if name == "wikitext2":
+        try:
+            from datasets import load_dataset
+
+            ds = load_dataset("wikitext", "wikitext-2-raw-v1", split=split)
+            lines = [t for t in ds["text"] if t.strip()][:size]
+        except Exception:
+            bits = torch.randint(0, 2, (size, seq_len), dtype=torch.long)
+            return jsonify({"bits": bits.tolist()})
+        bits_list = []
+        for text in lines:
+            b = text_to_bits(text)[:seq_len]
+            if len(b) < seq_len:
+                b.extend([0] * (seq_len - len(b)))
+            bits_list.append(b)
+        if len(bits_list) < size:
+            pad = size - len(bits_list)
+            bits_list.extend(torch.randint(0, 2, (pad, seq_len), dtype=torch.long).tolist())
+        return jsonify({"bits": bits_list})
+    return jsonify({"error": "unknown dataset"}), 400
+
+
+@app.route("/init", methods=["POST"])
+def init_model():
+    data = request.json or {}
+    int_fields = {
+        "d_model",
+        "nhead",
+        "num_layers",
+        "dim_feedforward",
+        "max_seq_len",
+        "chunk_size",
+        "overlap",
+    }
+    float_fields = {"act_threshold"}
+    bool_fields = {"reversible", "use_checkpoint"}
+    params = {}
+    for k, v in data.items():
+        if v is None:
+            params[k] = None
+        elif k in int_fields:
+            params[k] = int(v)
+        elif k in float_fields:
+            params[k] = float(v)
+        elif k in bool_fields:
+            params[k] = bool(v)
+        else:
+            params[k] = v
+    if MCP_SERVER_ADDR:
+        data = mcp_post("/init", params)
+        return jsonify(data)
+    manager.init_model(params)
+    return jsonify({"status": "initialized", "params": params})
+
+
+@app.route("/train", methods=["POST"])
+def train_model():
+    bits = torch.tensor(request.json["bits"], dtype=torch.long)
+    if MCP_SERVER_ADDR:
+        data = mcp_post("/train", {"bits": request.json["bits"]})
+        return jsonify(data)
+    loss, ratio = manager.train_step(bits)
+    return jsonify({"loss": loss, "ratio": ratio})
+
+
+@app.route("/train_epochs", methods=["POST"])
+def train_epochs_route():
+    bits = torch.tensor(request.json["bits"], dtype=torch.long)
+    epochs = int(request.json.get("epochs", 1))
+    compress_prob = float(request.json.get("compress_prob", 0.5))
+    direct_prob = float(request.json.get("direct_prob", 0.0))
+    if MCP_SERVER_ADDR:
+        data = mcp_post(
+            "/train_epochs",
+            {
+                "bits": request.json["bits"],
+                "epochs": epochs,
+                "compress_prob": compress_prob,
+                "direct_prob": direct_prob,
+            },
+        )
+        return jsonify(data)
+    metrics = manager.train_epochs(
+        bits,
+        epochs=epochs,
+        compress_prob=compress_prob,
+        direct_prob=direct_prob,
+    )
+    return jsonify({"metrics": metrics})
+
+
+@app.route("/scale_up", methods=["POST"])
+def scale_up():
+    width_mult = float(request.json.get("width_mult", 1.0))
+    if MCP_SERVER_ADDR:
+        data = mcp_post("/scale_up", {"width_mult": width_mult})
+        return jsonify(data)
+    manager.scale_up(width_mult)
+    return jsonify({
+        "status": "scaled",
+        "layers": manager.model.num_layers,
+        "d_model": manager.model.d_model,
+    })
+
+
+@app.route("/collapse", methods=["POST"])
+def collapse_model():
+    cluster_bits = request.json["clusters"]
+    params = {k: int(v) for k, v in request.json["params"].items()}
+    width_scale = float(request.json.get("width_scale", 1.0))
+    if MCP_SERVER_ADDR:
+        data = mcp_post(
+            "/collapse",
+            {"clusters": cluster_bits, "params": params, "width_scale": width_scale},
+        )
+        return jsonify(data)
+    manager.collapse(cluster_bits, params, width_scale)
+    return jsonify({"status": "collapsed"})
+
+
+@app.route("/lambdas", methods=["GET", "POST"])
+def update_lambdas():
+    if request.method == "POST":
+        data = request.json
+        if MCP_SERVER_ADDR:
+            res = mcp_post("/lambdas", data)
+            return jsonify(res)
+        manager.set_lambdas(
+            float(data["lambda_K"]), float(data["lambda_C"]), float(data["lambda_S"])
+        )
+        return jsonify({"status": "updated"})
+    else:
+        if MCP_SERVER_ADDR:
+            return jsonify(mcp_get("/lambdas"))
+        return jsonify(
+            {
+                "lambda_K": manager.lambda_K,
+                "lambda_C": manager.lambda_C,
+                "lambda_S": manager.lambda_S,
+            }
+        )
+
+
+@app.route("/config/telemetry", methods=["GET", "POST"])
+def telemetry_config():
+    """Get or update telemetry λ weights and safety floors."""
+    if request.method == "POST":
+        data = request.json
+        if MCP_SERVER_ADDR:
+            res = mcp_post("/config/telemetry", data)
+            return jsonify(res)
+        manager.set_lambdas(
+            float(data.get("lambda_K", manager.lambda_K)),
+            float(data.get("lambda_C", manager.lambda_C)),
+            float(data.get("lambda_S", manager.lambda_S)),
+        )
+        manager.set_floors(
+            float(data.get("c_floor", manager.c_floor)),
+            float(data.get("s_floor", manager.s_floor)),
+        )
+        return jsonify({"status": "updated"})
+    else:
+        if MCP_SERVER_ADDR:
+            return jsonify(mcp_get("/config/telemetry"))
+        return jsonify(
+            {
+                "lambda_K": manager.lambda_K,
+                "lambda_C": manager.lambda_C,
+                "lambda_S": manager.lambda_S,
+                "c_floor": manager.c_floor,
+                "s_floor": manager.s_floor,
+            }
+        )
+
+
+@app.route("/diffusion", methods=["GET", "POST"])
+def update_diffusion():
+    if request.method == "POST":
+        if MCP_SERVER_ADDR:
+            return jsonify(mcp_post("/diffusion", request.json))
+        manager.set_diffusion(bool(request.json.get("diffusion", False)))
+        return jsonify({"status": "updated"})
+    else:
+        if MCP_SERVER_ADDR:
+            return jsonify(mcp_get("/diffusion"))
+        return jsonify({"diffusion": manager.diffusion})
+
+
+@app.route("/gpu", methods=["GET", "POST"])
+def update_gpu():
+    if request.method == "POST":
+        if MCP_SERVER_ADDR:
+            return jsonify(mcp_post("/gpu", request.json))
+        manager.set_gpu(bool(request.json.get("use_gpu", False)))
+        return jsonify({"status": "updated"})
+    else:
+        if MCP_SERVER_ADDR:
+            return jsonify(mcp_get("/gpu"))
+        return jsonify({"use_gpu": manager.use_gpu})
+
+
+@app.route("/compression", methods=["GET", "POST"])
+def update_compression():
+    if request.method == "POST":
+        if MCP_SERVER_ADDR:
+            return jsonify(mcp_post("/compression", request.json))
+        manager.set_compression(bool(request.json.get("compression", False)))
+        return jsonify({"status": "updated"})
+    else:
+        if MCP_SERVER_ADDR:
+            return jsonify(mcp_get("/compression"))
+        return jsonify({"compression": manager.use_compression})
+
+
+@app.route("/qat", methods=["GET", "POST"])
+def update_qat():
+    if request.method == "POST":
+        if MCP_SERVER_ADDR:
+            return jsonify(mcp_post("/qat", request.json))
+        manager.set_qat(bool(request.json.get("qat", False)))
+        return jsonify({"status": "updated"})
+    else:
+        if MCP_SERVER_ADDR:
+            return jsonify(mcp_get("/qat"))
+        return jsonify({"qat": manager.qat})
+
+
+@app.route("/infer", methods=["POST"])
+def inference():
+    bits = torch.tensor(request.json["bits"], dtype=torch.long)
+    if MCP_SERVER_ADDR:
+        data = mcp_post("/infer", {"bits": request.json["bits"]})
+        return jsonify(data)
+    result = manager.infer(bits)
+    return jsonify(result)
+
+
+@app.route("/infer_long", methods=["POST"])
+def inference_long():
+    bits = torch.tensor(request.json["bits"], dtype=torch.long)
+    ctx = int(request.json.get("ctx_bits", 4096))
+    overlap = int(request.json.get("overlap", 256))
+    if MCP_SERVER_ADDR:
+        data = mcp_post(
+            "/infer_long",
+            {"bits": request.json["bits"], "ctx_bits": ctx, "overlap": overlap},
+        )
+        return jsonify(data)
+    result = manager.infer_long(bits, ctx_bits=ctx, overlap=overlap)
+    return jsonify(result)
+
+
+@app.route("/infer_text", methods=["POST"])
+def inference_text():
+    text = request.json.get("text", "")
+    if MCP_SERVER_ADDR:
+        data = mcp_post("/infer_text", {"text": text})
+        return jsonify(data)
+    result = manager.infer_text(text)
+    return jsonify(result)
+
+@app.route("/plot.png")
+def plot_png():
+    if MCP_SERVER_ADDR:
+        resp = requests.get(MCP_SERVER_ADDR.rstrip("/") + "/plot.png")
+        resp.raise_for_status()
+        return send_file(io.BytesIO(resp.content), mimetype="image/png")
+    fig, _ = plot_telemetry(manager.metrics)
+    buf = io.BytesIO()
+    fig.savefig(buf, format="png")
+    plt.close(fig)
+    buf.seek(0)
+    return send_file(buf, mimetype="image/png")
+
+
+def run_dashboard(host: str | None = None, port: int | None = None,
+                  snapshot_dir: str | None = None, telemetry_log: str | None = None) -> None:
+    """Launch the Flask dashboard server."""
+    env_host = os.getenv("HOST", "0.0.0.0")
+    env_port = int(os.getenv("PORT", "5000"))
+    host = host or env_host
+    port = port or env_port
+    global manager
+    if manager is None:
+        manager = ModelManager(snapshot_dir, telemetry_log)
+    app.run(host=host, port=port, debug=True)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Run dashboard server")
+    parser.add_argument("--host", default=os.getenv("HOST", "0.0.0.0"))
+    parser.add_argument("--port", type=int, default=int(os.getenv("PORT", "5000")))
+    parser.add_argument("--snapshot-dir", default=os.getenv("SNAPSHOT_DIR", "snapshots"))
+    parser.add_argument("--telemetry-log", default=os.getenv("TELEMETRY_LOG"))
+    args = parser.parse_args()
+    run_dashboard(args.host, args.port, args.snapshot_dir, args.telemetry_log)

BitTransformerLM/bit_transformer/distil.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from .model import BitTransformerLM
+
+
+@dataclass
+class TelemetryLog:
+    """Telemetry container holding attention maps across steps.
+
+    Attributes:
+        attention_maps: Tensor of shape [steps, heads, seq, seq].
+    """
+
+    attention_maps: torch.Tensor
+
+
+def distill_step(model: BitTransformerLM, scale: float, telemetry: TelemetryLog) -> BitTransformerLM:
+    """Return a pruned copy of ``model`` according to attention telemetry.
+
+    Args:
+        model: Teacher model to distill from.
+        scale: Fraction of weights to retain (0 < scale <= 1).
+        telemetry: Logged attention maps used to estimate parameter importance.
+
+    This function computes an importance score for each weight in the model's
+    linear layers using the supplied attention maps. The score is the mean
+    activation over time multiplied by the number of visits (non-zero
+    attention). The bottom ``(1 - scale)`` fraction of weights in each layer are
+    zeroed out, yielding a sparsified student model.
+    """
+    if not (0.0 < scale <= 1.0):
+        raise ValueError("scale must lie in (0, 1].")
+
+    # Clone the model so the teacher remains untouched.
+    student = BitTransformerLM(
+        d_model=model.d_model,
+        nhead=model.layers[0].self_attn.num_heads,
+        num_layers=model.num_layers,
+        dim_feedforward=model.layers[0].linear1.out_features,
+        max_seq_len=model.pos_enc.pe.size(0),
+        lambda_K=model.lambda_K,
+        lambda_C=model.lambda_C,
+        lambda_S=model.lambda_S,
+        reversible=model.reversible,
+        use_checkpoint=model.use_checkpoint,
+        use_autocast=model.use_autocast,
+        use_act=model.use_act,
+        act_threshold=model.act_threshold,
+        chunk_size=model.chunk_size,
+        overlap=model.overlap,
+    )
+    student.load_state_dict(model.state_dict())
+
+    attn = telemetry.attention_maps  # [steps, heads, seq, seq]
+    steps = attn.shape[0]
+    heads = attn.shape[1]
+    mean_act = attn.mean(dim=(0, 2, 3))
+    visits = (attn > 0).sum(dim=(0, 2, 3)).clamp_min(1)
+    head_importance = mean_act * visits
+    head_importance = head_importance / head_importance.sum()
+
+    prune_frac = 1.0 - scale
+
+    for module in student.modules():
+        if isinstance(module, nn.Linear):
+            weight = module.weight.data
+            out_features = weight.size(0)
+            if out_features % heads == 0:
+                repeats = out_features // heads
+                row_scores = head_importance.repeat_interleave(repeats).view(out_features, 1)
+            else:
+                row_scores = head_importance.mean().expand(out_features, 1)
+
+            importance = weight.abs() * row_scores
+            k = int(importance.numel() * prune_frac)
+            if k > 0:
+                thresh = torch.topk(importance.view(-1), k, largest=False).values.max()
+                mask = importance > thresh
+                weight.mul_(mask)
+                if module.bias is not None:
+                    row_mask = mask.view(out_features, -1).any(dim=1)
+                    module.bias.data.mul_(row_mask)
+
+    return student

BitTransformerLM/bit_transformer/distributed.py

@@ -0,0 +1,30 @@
+import torch
+import torch.nn as nn
+from typing import List, Optional
+
+from torch.distributed.fsdp import FullyShardedDataParallel
+try:
+    from torch.distributed.pipeline.sync import Pipe
+except Exception:  # pragma: no cover - Pipe may not be available in CPU builds
+    Pipe = None
+
+from .model import BitTransformerLM
+
+
+def wrap_fsdp(model: BitTransformerLM, **kwargs) -> FullyShardedDataParallel:
+    """Return a ``FullyShardedDataParallel`` wrapped model on the given device."""
+    device = kwargs.pop("device_id", torch.device("cpu"))
+    model = model.to(device)
+    return FullyShardedDataParallel(model, device_id=device, **kwargs)
+
+
+def make_pipeline(model: BitTransformerLM, chunks: int = 1) -> Pipe:
+    """Wrap the model with ``Pipe`` for simple pipeline parallelism.
+
+    The entire model is placed in an ``nn.Sequential`` so all existing telemetry
+    remains available. ``chunks`` controls microbatch splitting.
+    """
+    if Pipe is None:
+        raise RuntimeError("Pipeline parallelism not available in this build")
+    seq = nn.Sequential(model)
+    return Pipe(seq, chunks=chunks)

BitTransformerLM/bit_transformer/hf_checkpoint.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import gzip
+import os
+import shutil
+import tempfile
+from typing import Optional
+
+import torch
+from huggingface_hub import HfApi, hf_hub_download, login
+
+REPO_ID = "architect/bittransformerlm"
+FILENAME = "model.pt.gz"
+
+
+def hf_login(token: Optional[str] = None) -> None:
+    """Authenticate with Hugging Face.
+
+    The ``token`` may be provided directly or via the ``HF_TOKEN`` environment
+    variable. If omitted entirely, the library will attempt an interactive login.
+    """
+    login(token=token)
+
+
+def save_checkpoint(
+    model: torch.nn.Module,
+    *,
+    repo_id: str = REPO_ID,
+    filename: str = FILENAME,
+) -> None:
+    """Upload the model weights to ``repo_id`` under ``filename``.
+
+    The file within the repository is overwritten each time to avoid
+    accumulating checkpoints.
+    """
+    with tempfile.TemporaryDirectory() as tmp:
+        tmp_pt = os.path.join(tmp, "model.pt")
+        tmp_gz = os.path.join(tmp, filename)
+        torch.save(model.state_dict(), tmp_pt)
+        with open(tmp_pt, "rb") as src, gzip.open(tmp_gz, "wb") as dst:
+            dst.write(src.read())
+        HfApi().upload_file(
+            path_or_fileobj=tmp_gz,
+            path_in_repo=f"checkpoints/{filename}",
+            repo_id=repo_id,
+            repo_type="model",
+            overwrite=True,
+        )
+
+
+def download_checkpoint(
+    dest_path: str,
+    *,
+    repo_id: str = REPO_ID,
+    filename: str = FILENAME,
+) -> bool:
+    """Download the latest checkpoint to ``dest_path``.
+
+    Returns ``True`` if the checkpoint was successfully retrieved.
+    """
+    try:
+        buf = hf_hub_download(
+            repo_id,
+            f"checkpoints/{filename}",
+            repo_type="model",
+            force_download=True,
+        )
+    except Exception as exc:  # pragma: no cover - network errors
+        print("Failed to download checkpoint", exc)
+        return False
+    os.makedirs(os.path.dirname(dest_path), exist_ok=True)
+    shutil.copyfile(buf, dest_path)
+    return True
+
+
+__all__ = ["hf_login", "save_checkpoint", "download_checkpoint"]

BitTransformerLM/bit_transformer/model.py

@@ -0,0 +1,875 @@
+import math
+import contextlib
+import logging
+from typing import Dict, List, Tuple
+
+import torch
+import torch.distributed as dist
+import sys
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+
+from .torch_utils import cpu_autocast
+
+from .optimization import configure_optimizer
+from .compression import decompress_bits
+from .parity import enforce_parity
+
+_mask_cache: Dict[Tuple[int, torch.device], torch.Tensor] = {}
+
+
+def get_tri_mask(seq_len: int, device: torch.device) -> torch.Tensor:
+    """Return or create a cached upper-triangular mask."""
+    key = (seq_len, device)
+    if key not in _mask_cache:
+        _mask_cache[key] = torch.triu(
+            torch.ones(seq_len, seq_len, device=device, dtype=torch.bool), 1
+        )
+    return _mask_cache[key]
+
+try:  # torch.compile may not work on all Python versions
+    if torch.__version__ and tuple(map(int, torch.__version__.split(".")[:2])) >= (2, 0) and sys.version_info < (3, 11):
+        compile_fn = torch.compile
+    else:
+        raise RuntimeError
+except Exception:  # pragma: no cover - handle missing torch or unsupported version
+
+    def compile_fn(fn=None, **kwargs):
+        if fn is None:
+            return lambda f: f
+        return fn
+
+
+class PositionalEncoding(nn.Module):
+    """Sinusoidal positional encoding."""
+
+    def __init__(self, d_model: int, max_len: int = 1024) -> None:
+        super().__init__()
+        pe = torch.zeros(max_len, d_model)
+        pos = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
+        inv = torch.exp(
+            torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)
+        )
+        pe[:, 0::2] = torch.sin(pos * inv)
+        pe[:, 1::2] = torch.cos(pos * inv)
+        self.register_buffer("pe", pe.unsqueeze(1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Add positional encoding to input tensor."""
+        return x + self.pe[: x.size(0)]
+
+
+class LoggingTransformerEncoderLayer(nn.Module):
+    """Transformer encoder layer that exposes attention weights.
+
+    It optionally performs chunked attention with a fixed window size.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        dim_feedforward: int = 512,
+        dropout: float = 0.1,
+        chunk_size: int | None = None,
+        overlap: int = 0,
+        full_attn_logging: bool | None = None,
+    ) -> None:
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
+        self.chunk_size = chunk_size
+        self.overlap = overlap
+        if full_attn_logging is None:
+            full_attn_logging = False if chunk_size is not None else True
+        self.full_attn_logging = full_attn_logging
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.activation = F.relu
+
+    def _chunked_attn(
+        self, src: torch.Tensor, attn_mask: torch.Tensor | None = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Perform chunked self attention with overlap."""
+        T, B, D = src.shape
+        src_b = src.transpose(0, 1)  # [B, T, D]
+        C = self.chunk_size or T
+        O = self.overlap
+        n_chunks = (T + C - 1) // C
+        pad_len = n_chunks * C - T
+        src_pad = F.pad(src_b, (0, 0, O, pad_len + O))
+        chunk_len = C + 2 * O
+        chunks = src_pad.unfold(1, chunk_len, C)  # [B, n_chunks, chunk_len, D]
+        mask = get_tri_mask(chunk_len, src.device) if attn_mask is not None else None
+        out, weights = self.self_attn(
+            chunks.reshape(B * n_chunks, chunk_len, D),
+            chunks.reshape(B * n_chunks, chunk_len, D),
+            chunks.reshape(B * n_chunks, chunk_len, D),
+            attn_mask=mask,
+            need_weights=True,
+            average_attn_weights=False,
+        )
+        out = out.view(B, n_chunks, chunk_len, D)[:, :, O : O + C]
+        weights = weights.view(B, n_chunks, self.self_attn.num_heads, chunk_len, chunk_len)[
+            :, :, :, O : O + C
+        ]
+        seq = out.reshape(B, n_chunks * C, D)[:, :T]
+        if self.full_attn_logging and C < T:
+            full_attn = torch.zeros(
+                B, self.self_attn.num_heads, n_chunks * C, n_chunks * C, device=src.device
+            )
+            for idx in range(n_chunks):
+                s = idx * C
+                start = max(s - O, 0)
+                end = min(s + C, n_chunks * C)
+                src_start = O - (s - start)
+                src_end = src_start + (end - start)
+                full_attn[:, :, s : s + C, start:end] = weights[:, idx, :, src_start:src_end]
+            full_attn = full_attn[:, :, :T, :T]
+            attn_out = full_attn.detach()
+        else:
+            attn_out = torch.empty(0, device=src.device)
+        return seq.transpose(0, 1), attn_out
+
+    def forward(
+        self, src: torch.Tensor, attn_mask: torch.Tensor | None = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Return output and attention map."""
+        if self.chunk_size is not None:
+            attn_output, attn_weights = self._chunked_attn(src, attn_mask)
+        else:
+            qkv = src.transpose(0, 1)
+            attn_output, attn_weights = self.self_attn(
+                qkv,
+                qkv,
+                qkv,
+                attn_mask=attn_mask,
+                need_weights=True,
+                average_attn_weights=False,
+            )
+            attn_output = attn_output.transpose(0, 1)
+        src = src + self.dropout1(attn_output)
+        src = self.norm1(src)
+        out = self.linear2(self.dropout(self.activation(self.linear1(src))))
+        src = src + self.dropout2(out)
+        src = self.norm2(src)
+        return src, attn_weights.detach()
+
+
+class ReversibleLoggingTransformerEncoderLayer(nn.Module):
+    """Reversible transformer encoder layer with checkpointing."""
+
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        dim_feedforward: int = 512,
+        dropout: float = 0.1,
+        chunk_size: int | None = None,
+        overlap: int = 0,
+        full_attn_logging: bool | None = None,
+    ) -> None:
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
+        self.chunk_size = chunk_size
+        self.overlap = overlap
+        if full_attn_logging is None:
+            full_attn_logging = False if chunk_size is not None else True
+        self.full_attn_logging = full_attn_logging
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.activation = F.relu
+
+    @compile_fn
+    def _sa_block(
+        self, x: torch.Tensor, attn_mask: torch.Tensor | None = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self.chunk_size is not None:
+            T, B, D = x.shape
+            x_b = x.transpose(0, 1)
+            C = self.chunk_size or T
+            O = self.overlap
+            n_chunks = (T + C - 1) // C
+            pad_len = n_chunks * C - T
+            src_pad = F.pad(x_b, (0, 0, O, pad_len + O))
+            chunk_len = C + 2 * O
+            chunks = src_pad.unfold(1, chunk_len, C)
+            mask = get_tri_mask(chunk_len, x.device) if attn_mask is not None else None
+            out, weights = self.self_attn(
+                chunks.reshape(B * n_chunks, chunk_len, D),
+                chunks.reshape(B * n_chunks, chunk_len, D),
+                chunks.reshape(B * n_chunks, chunk_len, D),
+                attn_mask=mask,
+                need_weights=True,
+                average_attn_weights=False,
+            )
+            out = out.view(B, n_chunks, chunk_len, D)[:, :, O : O + C]
+            weights = weights.view(B, n_chunks, self.self_attn.num_heads, chunk_len, chunk_len)[
+                :, :, :, O : O + C
+            ]
+            seq = out.reshape(B, n_chunks * C, D)[:, :T]
+            if self.full_attn_logging and C < T:
+                full_attn = torch.zeros(
+                    B, self.self_attn.num_heads, n_chunks * C, n_chunks * C, device=x.device
+                )
+                for idx in range(n_chunks):
+                    s = idx * C
+                    start = max(s - O, 0)
+                    end = min(s + C, n_chunks * C)
+                    src_start = O - (s - start)
+                    src_end = src_start + (end - start)
+                    full_attn[:, :, s : s + C, start:end] = weights[
+                        :, idx, :, src_start:src_end
+                    ]
+                full_attn = full_attn[:, :, :T, :T]
+                weights = full_attn.detach()
+            else:
+                weights = torch.empty(0, device=x.device)
+            attn_out = seq.transpose(0, 1)
+        else:
+            qkv = x.transpose(0, 1)
+            attn_out, weights = self.self_attn(
+                qkv,
+                qkv,
+                qkv,
+                attn_mask=attn_mask,
+                need_weights=True,
+                average_attn_weights=False,
+            )
+            attn_out = attn_out.transpose(0, 1)
+        x = self.norm1(x + self.dropout1(attn_out))
+        return x, weights.detach()
+
+    @compile_fn
+    def _ff_block(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.linear2(self.dropout(self.activation(self.linear1(x))))
+        x = self.norm2(x + self.dropout2(out))
+        return x
+
+    def forward(
+        self,
+        x1: torch.Tensor,
+        x2: torch.Tensor,
+        attn_mask: torch.Tensor | None = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        y1, weights = self._sa_block(x2, attn_mask)
+        y1 = x1 + y1
+        y2 = x2 + self._ff_block(y1)
+        return y1, y2, weights
+
+
+class BitTransformerLM(nn.Module):
+    """Transformer language model that operates on raw bits (0/1) with telemetry."""
+
+    def __init__(
+        self,
+        d_model: int = 128,
+        nhead: int = 8,
+        num_layers: int = 4,
+        dim_feedforward: int = 512,
+        max_seq_len: int = 1024,
+        lambda_K: float = 1.0,
+        lambda_C: float = 1.0,
+        lambda_S: float = 1.0,
+        reversible: bool = False,
+        use_checkpoint: bool = True,
+        use_autocast: bool = False,
+        use_act: bool = False,
+        act_threshold: float = 0.9,
+        chunk_size: int | None = None,
+        overlap: int = 0,
+        full_attn_logging: bool | None = None,
+    ) -> None:
+        """Create a BitTransformer language model.
+
+        Args:
+            full_attn_logging: When ``False`` and ``chunk_size`` is
+                smaller than the sequence length, the model skips
+                reconstructing the full ``T×T`` attention matrices for
+                telemetry to reduce memory use.
+        """
+        super().__init__()
+        self.d_model = d_model
+        self.num_layers = num_layers
+        self.lambda_K = lambda_K
+        self.lambda_C = lambda_C
+        self.lambda_S = lambda_S
+        self.reversible = reversible
+        self.use_checkpoint = use_checkpoint
+        self.use_autocast = use_autocast
+        self.use_act = use_act
+        self.act_threshold = act_threshold
+        self.chunk_size = chunk_size
+        self.overlap = overlap
+        if full_attn_logging is None:
+            full_attn_logging = False if chunk_size is not None else True
+        self.full_attn_logging = full_attn_logging
+
+        # Bit embedding: two possible input values
+        self.embedding = nn.Embedding(2, d_model)
+        self.pos_enc = PositionalEncoding(d_model, max_len=max_seq_len)
+
+        layer_cls = (
+            ReversibleLoggingTransformerEncoderLayer
+            if reversible
+            else LoggingTransformerEncoderLayer
+        )
+        self.layers = nn.ModuleList(
+            [
+                layer_cls(
+                    d_model=d_model,
+                    nhead=nhead,
+                    dim_feedforward=dim_feedforward,
+                    chunk_size=chunk_size,
+                    overlap=overlap,
+                    full_attn_logging=full_attn_logging,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        if self.use_act:
+            self.halt_projs = nn.ModuleList(
+                [nn.Linear(d_model, 1) for _ in range(num_layers)]
+            )
+
+        self.out_head = nn.Linear(d_model, 2)  # output logits for bit=0 or bit=1
+
+    def expand_positional_encoding(self, new_len: int) -> None:
+        """Expand positional encoding to at least ``new_len``."""
+        cur_len = self.pos_enc.pe.size(0)
+        if new_len <= cur_len:
+            return
+        device = self.pos_enc.pe.device
+        d_model = self.d_model
+        pe = torch.zeros(new_len, d_model, device=device)
+        pe[:cur_len] = self.pos_enc.pe.squeeze(1)
+        pos = torch.arange(cur_len, new_len, dtype=torch.float32, device=device).unsqueeze(1)
+        inv = torch.exp(torch.arange(0, d_model, 2, device=device).float() * -(math.log(10000.0) / d_model))
+        pe[cur_len:, 0::2] = torch.sin(pos * inv)
+        pe[cur_len:, 1::2] = torch.cos(pos * inv)
+        self.pos_enc.pe = pe.unsqueeze(1)
+
+    def set_lambdas(self, lambda_K: float, lambda_C: float, lambda_S: float) -> None:
+        """Update weighting coefficients for telemetry metrics."""
+        self.lambda_K = lambda_K
+        self.lambda_C = lambda_C
+        self.lambda_S = lambda_S
+
+    def _maybe_decompress(self, codes: torch.Tensor) -> torch.Tensor:
+        """Return raw bit sequences, decompressing if input appears run-length encoded."""
+        if codes.dim() <= 1:
+            return codes
+        needs_decompress = codes.max().item() > 1
+        if not needs_decompress and codes.size(1) % 2 == 0:
+            vals = codes[:, 0::2]
+            if torch.all(vals[:, 1:] != vals[:, :-1]):
+                needs_decompress = True
+        if not needs_decompress:
+            return codes
+        seqs = [decompress_bits(row.to(torch.uint8)) for row in codes]
+        max_len = max(seq.numel() for seq in seqs)
+        padded = [F.pad(seq, (0, max_len - seq.numel())) for seq in seqs]
+        return torch.stack(padded)
+
+    def negentropy_kpi(self, codes: torch.Tensor) -> torch.Tensor:
+        """Approximate negentropy of bit sequences.
+
+        Returns a value in ``[0, 1]`` where ``1`` denotes a perfectly ordered
+        sequence (all zeros or ones) and ``0`` reflects maximal entropy.
+        """
+        codes = self._maybe_decompress(codes)
+        p = codes.float().mean(dim=1)
+        entropy = -(p * torch.log(p + 1e-9) + (1 - p) * torch.log(1 - p + 1e-9))
+        max_e = math.log(2.0)
+        return 1 - entropy / max_e
+
+    def lz_complexity(self, codes: torch.Tensor) -> torch.Tensor:
+        """Differentiable proxy for Lempel–Ziv complexity.
+
+        Values near ``0`` indicate highly compressible sequences while values
+        approaching ``1`` correspond to rapid bit alternation.
+        """
+        codes = self._maybe_decompress(codes)
+        diffs = torch.abs(codes[:, 1:] - codes[:, :-1])
+        return diffs.float().mean(dim=1)
+
+    def negentropy_logits(self, logits: torch.Tensor, detach: bool = True) -> torch.Tensor:
+        """Negentropy computed from model logits.
+
+        Parameters
+        ----------
+        logits: ``torch.Tensor``
+            Logit tensor of shape ``(B, T, 2)``.
+        detach: bool, default ``True``
+            When ``True`` the computation is detached from the autograd graph.
+        """
+        assert logits.dim() == 3 and logits.size(-1) == 2, "logits must be [B,T,2]"
+        prob = logits.softmax(-1)
+        if detach:
+            prob = prob.detach()
+        p = prob[..., 1].mean(dim=1)
+        entropy = -(p * torch.log(p + 1e-9) + (1 - p) * torch.log(1 - p + 1e-9))
+        max_e = math.log(2.0)
+        return 1 - entropy / max_e
+
+    def lz_complexity_logits(self, logits: torch.Tensor, detach: bool = True) -> torch.Tensor:
+        """LZ complexity proxy computed from logits.
+
+        Parameters
+        ----------
+        logits: ``torch.Tensor``
+            Logit tensor of shape ``(B, T, 2)``.
+        detach: bool, default ``True``
+            When ``True`` the computation is detached from the autograd graph.
+        """
+        assert logits.dim() == 3 and logits.size(-1) == 2, "logits must be [B,T,2]"
+        prob = logits.softmax(-1)
+        if detach:
+            prob = prob.detach()
+        prob1 = prob[..., 1]
+        diffs = torch.abs(prob1[:, 1:] - prob1[:, :-1])
+        return diffs.mean(dim=1)
+
+    def symbiosis_kl_logits(
+        self, logits: torch.Tensor, ref_prob: float = 0.5, detach: bool = True
+    ) -> torch.Tensor:
+        """Symbiosis score from KL divergence to a reference distribution.
+
+        Returns a value in ``[0, 1]`` with ``1`` meaning perfect agreement with
+        the reference distribution and ``0`` indicating maximal divergence.
+        """
+        assert logits.dim() == 3 and logits.size(-1) == 2, "logits must be [B,T,2]"
+        probs = logits.softmax(-1)
+        if detach:
+            probs = probs.detach()
+        ref = torch.tensor([1 - ref_prob, ref_prob], device=logits.device)
+        kl = (probs * (probs.clamp_min(1e-9).log() - ref.log())).sum(-1).mean(dim=1)
+        max_kl = math.log(2.0)
+        return 1 - kl / max_kl
+
+    def _act_step(
+        self,
+        hidden: torch.Tensor,
+        idx: int,
+        halt_prob: torch.Tensor,
+        act_state: torch.Tensor,
+        halt_history: List[torch.Tensor],
+    ) -> Tuple[torch.Tensor, torch.Tensor, bool]:
+        """Apply one step of ACT halting logic."""
+        p = torch.sigmoid(self.halt_projs[idx](hidden))
+        delta = (1 - halt_prob) * p
+        halt_prob = halt_prob + delta
+        act_state = act_state + hidden * delta
+        halt_history.append(halt_prob.detach())
+        min_prob = halt_prob.detach().min()
+        if dist.is_initialized():
+            dist.all_reduce(min_prob, op=dist.ReduceOp.MIN)
+        return halt_prob, act_state, min_prob.item() >= self.act_threshold
+
+    def forward(
+        self, bit_seq: torch.Tensor, causal: bool = True
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        """Forward pass returning logits and telemetry from the same graph.
+
+        By default the model uses causal masking and (optional) chunked
+        attention. When ``causal`` is ``False`` the model operates in
+        "Diffusion LM" mode. In this mode chunked attention is temporarily
+        disabled so that every token can attend to the full sequence
+        bidirectionally. The original chunking configuration is restored after
+        the forward pass.
+        """
+
+        # Disable chunking when running in bidirectional (non-causal) mode
+        orig_chunks = None
+        orig_model_chunk = None
+        if not causal and self.chunk_size is not None:
+            orig_model_chunk = self.chunk_size
+            orig_chunks = [layer.chunk_size for layer in self.layers]
+            self.chunk_size = None
+            for layer in self.layers:
+                layer.chunk_size = None
+
+        try:
+            ctx = cpu_autocast() if self.use_autocast else contextlib.nullcontext()
+            with ctx:
+                x = self.embedding(bit_seq).transpose(0, 1) * math.sqrt(self.d_model)
+                x = self.pos_enc(x)
+
+                attn_mask = get_tri_mask(x.size(0), x.device) if causal else None
+
+                activations: List[torch.Tensor] = []
+                attn_maps: List[torch.Tensor] = []
+                halt_history: List[torch.Tensor] = []
+                if self.use_act:
+                    halt_prob = torch.zeros(x.size(0), x.size(1), 1, device=x.device)
+                    act_state = torch.zeros_like(x)
+                if self.reversible:
+                    x1, x2 = x, x
+                    for idx, layer in enumerate(self.layers):
+                        if self.use_checkpoint:
+                            x1, x2, attn = checkpoint.checkpoint(
+                                layer, x1, x2, attn_mask
+                            )
+                        else:
+                            x1, x2, attn = layer(x1, x2, attn_mask)
+                        combined = (x1 + x2) / 2
+                        activations.append(combined)
+                        if attn.numel() > 0:
+                            attn_maps.append(attn)
+                        if self.use_act:
+                            halt_prob, act_state, should_break = self._act_step(
+                                combined, idx, halt_prob, act_state, halt_history
+                            )
+                            if should_break:
+                                break
+                    x = (x1 + x2) / 2
+                else:
+                    for idx, layer in enumerate(self.layers):
+                        if self.use_checkpoint:
+                            x, attn = checkpoint.checkpoint(layer, x, attn_mask)
+                        else:
+                            x, attn = layer(x, attn_mask)
+                        activations.append(x)
+                        if attn.numel() > 0:
+                            attn_maps.append(attn)
+                        if self.use_act:
+                            halt_prob, act_state, should_break = self._act_step(
+                                x, idx, halt_prob, act_state, halt_history
+                            )
+                            if should_break:
+                                break
+                if self.use_act:
+                    act_state = act_state + x * (1 - halt_prob)
+                    x = act_state
+                logits = self.out_head(x)
+
+            # Per-layer entropy of activations
+            entropies = []
+            for act in activations:
+                prob = act.softmax(-1)
+                ent = -(prob * prob.clamp_min(1e-9).log()).sum(-1).mean()
+                entropies.append(ent)
+
+            attn_entropies = []
+            for attn in attn_maps:
+                prob = attn  # weights are already softmaxed
+                ent = -(prob * prob.clamp_min(1e-9).log()).sum(-1)
+                ent = ent.mean(1)
+                attn_entropies.append(ent)
+            if attn_entropies:
+                attn_entropy_map = torch.stack(attn_entropies).mean(0)
+            else:
+                attn_entropy_map = torch.zeros(
+                    bit_seq.size(0), bit_seq.size(1), device=bit_seq.device
+                )
+            max_ent = math.log(attn_entropy_map.size(-1))
+            attn_entropy_map = attn_entropy_map / max_ent
+            attn_entropy = attn_entropy_map.mean(1)
+
+            logits_bt = logits.transpose(0, 1)
+            negentropy_in = self.negentropy_kpi(bit_seq)
+            lz_in = self.lz_complexity(bit_seq.float())
+            negentropy_logits_b = self.negentropy_logits(logits_bt, detach=False)
+            lz_logits_b = self.lz_complexity_logits(logits_bt, detach=False)
+            kl_div_b = self.symbiosis_kl_logits(logits_bt, detach=False)
+
+            raw_sym = (
+                (self.lambda_K * negentropy_logits_b + self.lambda_C * lz_logits_b) / 2
+                + negentropy_logits_b * lz_logits_b
+                - self.lambda_S * kl_div_b
+                - 0.1 * attn_entropy
+            )
+            weight_norm = torch.stack([p.norm() for p in self.parameters()]).mean().detach()
+            raw_sym = raw_sym - 0.01 * weight_norm
+            sym_score = torch.sigmoid(raw_sym)
+
+            B, T = bit_seq.shape
+            assert logits_bt.shape[:2] == (B, T)
+            assert attn_entropy_map.shape == (B, T)
+
+            telemetry = {
+                "activations": activations,
+                "attention_maps": attn_maps,
+                "attention_entropy": attn_entropy_map,
+                "entropy": entropies,
+                "attention_entropy_mean": attn_entropy,
+                "negentropy_input": negentropy_in.detach(),
+                "lz_complexity_input": lz_in.detach(),
+                "negentropy_logits": negentropy_logits_b.detach(),
+                "lz_complexity_logits": lz_logits_b.detach(),
+                "symbiosis_kl": kl_div_b.detach(),
+                "symbiosis_score": sym_score.detach(),
+            }
+            if self.use_act:
+                telemetry["halt_probs"] = halt_history
+
+            return logits_bt, telemetry
+        finally:
+            if orig_chunks is not None:
+                self.chunk_size = orig_model_chunk
+                for layer, chunk in zip(self.layers, orig_chunks):
+                    layer.chunk_size = chunk
+
+    def forward_compressed(
+        self, compressed_bits, causal: bool = True
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        """Decompress bit sequences then run the normal forward pass."""
+        if isinstance(compressed_bits, torch.Tensor) and compressed_bits.dim() == 1:
+            sequences = [decompress_bits(compressed_bits).to(torch.long)]
+        else:
+            sequences = [decompress_bits(c).to(torch.long) for c in compressed_bits]
+        lengths = [seq.numel() for seq in sequences]
+        if len(set(lengths)) != 1:
+            raise ValueError("Sequences decompress to different lengths")
+        bits = torch.stack(sequences)
+        return self.forward(bits, causal=causal)
+
+    def _current_params(self) -> Dict:
+        """Return a dictionary with the current model hyperparameters."""
+        return {
+            "d_model": self.d_model,
+            "nhead": self.layers[0].self_attn.num_heads,
+            "num_layers": self.num_layers,
+            "dim_feedforward": self.layers[0].linear1.out_features,
+            "max_seq_len": self.pos_enc.pe.size(0),
+            "lambda_K": self.lambda_K,
+            "lambda_C": self.lambda_C,
+            "lambda_S": self.lambda_S,
+            "reversible": self.reversible,
+            "use_checkpoint": self.use_checkpoint,
+            "use_autocast": self.use_autocast,
+            "use_act": self.use_act,
+            "act_threshold": self.act_threshold,
+            "chunk_size": self.chunk_size,
+            "overlap": self.overlap,
+        }
+
+    def double_width(self) -> "BitTransformerLM":
+        """Return a copy of the model with doubled hidden size."""
+        from .scale import expand_model
+
+        params = self._current_params()
+        params["d_model"] *= 2
+        params["dim_feedforward"] *= 2
+        return expand_model(self, params)
+
+    def double_layers(self) -> "BitTransformerLM":
+        """Return a copy of the model with twice as many layers."""
+        from .scale import expand_model
+
+        params = self._current_params()
+        params["num_layers"] *= 2
+        return expand_model(self, params)
+
+    def double_length(self) -> "BitTransformerLM":
+        """Return a copy of the model with doubled maximum sequence length."""
+        from .scale import expand_model
+
+        params = self._current_params()
+        params["max_seq_len"] *= 2
+        params["chunk_size"] = params["max_seq_len"]
+        return expand_model(self, params)
+
+    def train_full_sequence(
+        self,
+        bits: torch.Tensor,
+        *,
+        ctx_bits: int = 4096,
+        detach_every_n: int = 1_048_576,
+    ) -> float:
+        """Train on a long bit tensor using sliding windows.
+
+        Parameters
+        ----------
+        bits: ``torch.Tensor``
+            1D tensor containing the full bit sequence.
+        ctx_bits: int
+            Size of the training context window.
+        detach_every_n: int
+            Interval in bits for optimizer updates and graph detachment.
+        Returns
+        -------
+        float
+            Mean loss over all windows.
+        """
+        self.train()
+        optimizer, scheduler = configure_optimizer(
+            self, lr=1e-3, total_steps=max(1, bits.numel() // ctx_bits)
+        )
+        accum = 0
+        total_loss = 0.0
+        count = 0
+        for start in range(0, bits.numel() - ctx_bits - 1, ctx_bits):
+            segment = bits[start : start + ctx_bits + 1].unsqueeze(0)
+            logits, _ = self(segment)
+            pred = logits[:, :-1, :].reshape(-1, 2)
+            target = segment[:, 1:].reshape(-1)
+            loss = F.cross_entropy(pred, target)
+            loss.backward()
+            accum += ctx_bits
+            total_loss += loss.item()
+            count += 1
+            if accum >= detach_every_n:
+                torch.nn.utils.clip_grad_norm_(self.parameters(), 1.0)
+                optimizer.step()
+                scheduler.step()
+                optimizer.zero_grad()
+                accum = 0
+        if accum > 0:
+            torch.nn.utils.clip_grad_norm_(self.parameters(), 1.0)
+            optimizer.step()
+            scheduler.step()
+            optimizer.zero_grad()
+        return total_loss / max(1, count)
+
+
+def infer_long_sequence(
+    model: BitTransformerLM,
+    bits: torch.Tensor,
+    *,
+    ctx_bits: int = 4096,
+    overlap: int = 256,
+) -> Tuple[torch.Tensor, List[Dict[str, torch.Tensor]]]:
+    """Infer a long bit sequence using sliding windows with overlap."""
+    model.eval()
+    device = next(model.parameters()).device
+    bits = bits.to(device)
+    step = ctx_bits - overlap
+    outputs: List[torch.Tensor] = []
+    logs: List[Dict[str, torch.Tensor]] = []
+    for start in range(0, bits.numel(), step):
+        window = bits[start : start + ctx_bits].unsqueeze(0)
+        logits, tele = model(window, causal=True)
+        pred = logits.argmax(-1).squeeze(0)
+        outputs.append(pred)
+        logs.append(tele)
+    out = torch.cat(outputs)[: bits.numel()]
+    return out, logs
+
+
+def diffusion_inference(
+    model: BitTransformerLM,
+    *,
+    length: int,
+    steps: int = 8,
+    batch_size: int = 1,
+    init_bits: torch.Tensor | None = None,
+    schedule: str = "linear",
+) -> torch.Tensor:
+    """Generate bit sequences using iterative denoising diffusion.
+
+    Parameters
+    ----------
+    model: ``BitTransformerLM``
+        The model used for denoising. It is run in non-causal mode with
+        chunked attention disabled, enabling full-context bidirectional
+        attention.
+    length: int
+        Length of the bit sequences to generate.
+    steps: int, default ``8``
+        Number of denoising iterations. More steps generally yield sharper
+        samples at the cost of compute.
+    batch_size: int, default ``1``
+        Number of sequences to generate in parallel.
+    init_bits: ``torch.Tensor`` | ``None``
+        Optional initial noisy bits of shape ``(batch_size, length)``. When
+        ``None`` random noise is used.
+    schedule: str, default ``"linear"``
+        Noise schedule for the denoising mask probability. Options are
+        ``"linear"``, ``"cosine"``, and ``"exp"``.
+
+    Returns
+    -------
+    ``torch.Tensor``
+        A tensor of shape ``(batch_size, length)`` containing generated bits.
+    """
+
+    model.eval()
+    device = next(model.parameters()).device
+    if init_bits is None:
+        bits = torch.randint(0, 2, (batch_size, length), device=device)
+    else:
+        bits = init_bits.to(device)
+        if bits.shape != (batch_size, length):
+            raise ValueError("init_bits must have shape (batch_size, length)")
+
+    for step in range(steps):
+        logits, _ = model(bits, causal=False)
+        prob = logits.softmax(-1)[..., 1]
+        t = (step + 1) / steps
+        if schedule == "linear":
+            mask_prob = 1.0 - t
+        elif schedule == "cosine":
+            mask_prob = math.cos(math.pi * t / 2)
+        elif schedule == "exp":
+            mask_prob = math.exp(-5 * t)
+        else:
+            raise ValueError(f"unknown schedule: {schedule}")
+        mask = (torch.rand_like(bits.float()) < mask_prob).long()
+        sampled = torch.bernoulli(prob).long()
+        bits = torch.where(mask.bool(), sampled, bits)
+    if bits.shape[-1] % 9 == 0:
+        bits, corrections = enforce_parity(bits)
+        if corrections:
+            logging.info("Parity corrections applied: %d", corrections)
+    try:
+        from .safety import hil_safe_inference
+
+        hil_safe_inference(model, bits, causal=False, strict=False)
+    except RuntimeError as exc:
+        logging.warning("Safety gate warning: %s", exc)
+    return bits
+
+
+def example_usage() -> float:
+    """Run the example from the README and return the loss."""
+    B, L = 4, 16
+    model = BitTransformerLM(
+        d_model=64, nhead=4, num_layers=2, dim_feedforward=256, max_seq_len=L
+    )
+    bits = torch.randint(0, 2, (B, L), dtype=torch.long)
+    logits, _ = model(bits)
+    pred = logits[:, :-1, :].reshape(-1, 2)
+    target = bits[:, 1:].reshape(-1)
+    loss = F.cross_entropy(pred, target)
+    return loss.item()
+
+
+def example_training_step() -> Tuple[float, Dict[str, torch.Tensor]]:
+    """Demonstrate a training step where metrics do not affect gradients."""
+    B, L = 4, 16
+    model = BitTransformerLM(
+        d_model=32, nhead=4, num_layers=1, dim_feedforward=64, max_seq_len=L
+    )
+    optimizer, scheduler = configure_optimizer(model, lr=1e-3, total_steps=1)
+
+    bits = torch.randint(0, 2, (B, L), dtype=torch.long)
+    logits, telemetry = model(bits)
+
+    pred = logits[:, :-1, :].reshape(-1, 2)
+    target = bits[:, 1:].reshape(-1)
+    loss = F.cross_entropy(pred, target)
+
+    loss.backward()
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+    optimizer.step()
+    scheduler.step()
+    optimizer.zero_grad()
+    return loss.item(), telemetry
+
+
+if __name__ == "__main__":
+    loss, telemetry = example_training_step()
+    print("Composite loss:", loss)
+    print("Telemetry keys:", list(telemetry.keys()))

BitTransformerLM/bit_transformer/optimization.py

@@ -0,0 +1,37 @@
+import torch
+import torch.nn as nn
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import OneCycleLR
+
+
+def configure_optimizer(
+    model: nn.Module,
+    lr: float = 1e-3,
+    weight_decay: float = 0.01,
+    total_steps: int = 100
+):
+    """Return AdamW optimizer with OneCycleLR scheduler."""
+    optimizer = AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
+    scheduler = OneCycleLR(optimizer, max_lr=lr, total_steps=total_steps)
+    return optimizer, scheduler
+
+
+def adjust_learning_rate(optimizer: torch.optim.Optimizer, factor: float) -> float:
+    """Scale the learning rate of all param groups by ``factor``.
+
+    Parameters
+    ----------
+    optimizer:
+        The optimizer whose learning rate will be adjusted.
+    factor:
+        Multiplicative factor applied to the current learning rate.
+
+    Returns
+    -------
+    float
+        The updated learning rate of the first parameter group.
+    """
+    for param_group in optimizer.param_groups:
+        param_group["lr"] *= factor
+    return optimizer.param_groups[0]["lr"]
+

BitTransformerLM/bit_transformer/parity.py

@@ -0,0 +1,24 @@
+import torch
+
+def enforce_parity(bits: torch.Tensor) -> tuple[torch.Tensor, int]:
+    """Fix parity bits so each 9-bit chunk has even parity.
+
+    Parameters
+    ----------
+    bits: ``torch.Tensor``
+        Tensor of shape ``(..., length)`` where ``length`` is a multiple of 9.
+
+    Returns
+    -------
+    tuple[torch.Tensor, int]
+        Corrected tensor and number of bytes that were adjusted.
+    """
+    if bits.shape[-1] % 9 != 0:
+        raise ValueError("Bit stream length must be multiple of 9")
+    flat = bits.clone().view(-1, 9)
+    payload = flat[:, :8]
+    parity = flat[:, 8]
+    new_parity = payload.sum(dim=1) % 2
+    corrections = (parity != new_parity).sum().item()
+    flat[:, 8] = new_parity
+    return flat.view_as(bits), corrections

BitTransformerLM/bit_transformer/quantization.py

@@ -0,0 +1,89 @@
+import torch
+import torch.nn as nn
+
+from torch.ao.quantization.fake_quantize import FakeQuantize
+from torch.ao.quantization.observer import MinMaxObserver
+from torch.ao.quantization.qconfig import QConfig
+from torch.ao.quantization import convert
+
+from .model import BitTransformerLM
+
+
+def quantize_dynamic(model: BitTransformerLM, dtype: torch.dtype = torch.qint8) -> BitTransformerLM:
+    """Return a dynamically quantized copy of the model for inference."""
+    quantized = torch.quantization.quantize_dynamic(
+        model, {nn.Linear}, dtype=dtype
+    )
+    return quantized
+
+
+class FourBitObserver(MinMaxObserver):
+    """Min-max observer configured for 4-bit quantization."""
+
+    def __init__(self, **kwargs):
+        super().__init__(
+            quant_min=0,
+            quant_max=15,
+            dtype=torch.quint8,
+            qscheme=torch.per_tensor_affine,
+            **kwargs,
+        )
+
+
+FourBitFakeQuantize = FakeQuantize.with_args(observer=FourBitObserver)
+
+four_bit_qconfig = QConfig(activation=FourBitFakeQuantize, weight=FourBitFakeQuantize)
+
+
+class QATLinear(nn.Linear):
+    """Linear layer with fake quantization for QAT."""
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
+        super().__init__(in_features, out_features, bias)
+        self.weight_fake_quant = FourBitFakeQuantize()
+        self.activation_post_process = FourBitFakeQuantize()
+
+    @classmethod
+    def from_float(cls, mod: nn.Linear) -> "QATLinear":
+        qat = cls(mod.in_features, mod.out_features, mod.bias is not None)
+        qat.weight = mod.weight
+        qat.bias = mod.bias
+        return qat
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.activation_post_process(x)
+        w = self.weight_fake_quant(self.weight)
+        return nn.functional.linear(x, w, self.bias)
+
+
+def prepare_qat_fx(model: BitTransformerLM) -> BitTransformerLM:
+    """Prepare BitTransformerLM for quantization-aware training."""
+
+    for name, module in model.named_children():
+        if isinstance(module, nn.Linear):
+            setattr(model, name, QATLinear.from_float(module))
+        else:
+            prepare_qat_fx(module)
+    return model
+
+
+def convert_qat_fx(model: BitTransformerLM) -> BitTransformerLM:
+    """Convert a QAT-prepared model to a quantized version."""
+
+    for name, module in model.named_children():
+        if isinstance(module, QATLinear):
+            w = module.weight.data
+            qmin, qmax = 0, 15
+            min_w = w.min()
+            max_w = w.max()
+            scale = (max_w - min_w) / (qmax - qmin + 1e-8)
+            zero_point = qmin - torch.round(min_w / scale)
+            q_w = torch.clamp(torch.round(w / scale + zero_point), qmin, qmax)
+            new_mod = nn.Linear(module.in_features, module.out_features, module.bias is not None)
+            new_mod.weight = nn.Parameter((q_w - zero_point) * scale)
+            if module.bias is not None:
+                new_mod.bias = nn.Parameter(module.bias.data)
+            setattr(model, name, new_mod)
+        else:
+            convert_qat_fx(module)
+    return model

BitTransformerLM/bit_transformer/safety.py

@@ -0,0 +1,149 @@
+import logging
+import time
+import torch
+from typing import Dict, Optional, Tuple
+
+from .model import BitTransformerLM
+
+
+class SafetyGate:
+    """Exponential moving average safety gate with burn-in."""
+
+    def __init__(
+        self,
+        *,
+        c_floor: float = 0.3,
+        s_floor: float = 0.5,
+        decay: float = 0.9,
+        burn_in: int = 10,
+    ) -> None:
+        self.c_floor = c_floor
+        self.s_floor = s_floor
+        self.decay = decay
+        self.burn_in = burn_in
+        self.step = 0
+        self._c_ema: Optional[float] = None
+        self._s_ema: Optional[float] = None
+
+    def should_trigger(self, c_val: float, s_val: float) -> bool:
+        """Update EMA scores and check if gating should trigger."""
+
+        self.step += 1
+        if self._c_ema is None:
+            self._c_ema = c_val
+            self._s_ema = s_val
+        else:
+            self._c_ema = self.decay * self._c_ema + (1 - self.decay) * c_val
+            self._s_ema = self.decay * self._s_ema + (1 - self.decay) * s_val
+        if self.step <= self.burn_in:
+            return False
+        return self._c_ema <= self.c_floor or self._s_ema <= self.s_floor
+
+
+def hil_safe_inference(
+    model: BitTransformerLM,
+    bit_seq: torch.Tensor,
+    c_floor: float = 0.3,
+    s_floor: float = 0.5,
+    *,
+    causal: bool = True,
+    strict: bool = True,
+    gate: Optional[SafetyGate] = None,
+) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+    """Run inference with telemetry gating.
+
+    Parameters
+    ----------
+    model:
+        Model to run inference with.
+    bit_seq:
+        Input bit sequences.
+    c_floor, s_floor:
+        Minimum LZ complexity and symbiosis score required for safe output.
+    causal:
+        Whether to run the model in causal (autoregressive) mode. When ``False``
+        the model performs full-context Diffusion LM inference.
+    strict:
+        If ``False`` the function returns model outputs even when the floors are
+        not met instead of raising ``RuntimeError``.
+    gate:
+        Optional :class:`SafetyGate` that applies EMA smoothing and burn-in
+        before enforcing the floors.
+    """
+    model.eval()
+    with torch.no_grad():
+        logits, telemetry = model(bit_seq, causal=causal)
+        c_val = float(telemetry["lz_complexity_logits"].mean().item())
+        s_val = float(telemetry["symbiosis_score"].mean().item())
+        c_val = max(0.0, min(1.0, c_val))
+        s_val = max(0.0, min(1.0, s_val))
+        if gate is not None:
+            triggered = gate.should_trigger(c_val, s_val)
+        else:
+            triggered = c_val <= c_floor or s_val <= s_floor
+        if strict and triggered:
+            raise RuntimeError(
+                f"Safety gate triggered: C={c_val:.3f}, S={s_val:.3f}"
+            )
+        return logits.argmax(-1), telemetry
+
+
+def demo_hil_safety() -> None:
+    """Demonstrate gating on random bits."""
+    bits = torch.randint(0, 2, (1, 8), dtype=torch.long)
+    model = BitTransformerLM(d_model=32, nhead=4, num_layers=1, dim_feedforward=64, max_seq_len=8)
+    try:
+        out, _ = hil_safe_inference(model, bits, c_floor=0.0, s_floor=0.0)
+        print("Safe output bits:", out.squeeze(0).tolist())
+    except RuntimeError as e:
+        print("Gate triggered:", e)
+
+
+def safe_sample_with_retry(
+    model: BitTransformerLM,
+    bit_seq: torch.Tensor,
+    c_floor: float = 0.3,
+    s_floor: float = 0.5,
+    *,
+    causal: bool = True,
+    max_retries: int = 3,
+    backoff: float = 0.1,
+    gate: Optional[SafetyGate] = None,
+) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+    """Run :func:`hil_safe_inference` with automatic retries.
+
+    The helper retries failed safety checks by toggling diffusion mode and
+    refreshing the input bits. An exponential backoff is applied between
+    attempts and warnings are logged for each retry.
+
+    Parameters
+    ----------
+    gate:
+        Optional :class:`SafetyGate` instance shared across retries to apply
+        EMA smoothing and burn-in.
+
+    Returns
+    -------
+    Tuple[torch.Tensor, Dict[str, torch.Tensor]]
+        The sampled bits and associated telemetry.
+    """
+
+    for attempt in range(max_retries):
+        try:
+            return hil_safe_inference(
+                model,
+                bit_seq,
+                c_floor,
+                s_floor,
+                causal=causal,
+                strict=True,
+                gate=gate,
+            )
+        except RuntimeError as exc:  # safety gate triggered
+            logging.warning("Safety gate failed (attempt %d/%d): %s", attempt + 1, max_retries, exc)
+            if attempt >= max_retries - 1:
+                raise
+            time.sleep(backoff * (2 ** attempt))
+            causal = False  # retry in diffusion mode
+            bit_seq = torch.randint(0, 2, bit_seq.shape, dtype=bit_seq.dtype, device=bit_seq.device)
+

BitTransformerLM/bit_transformer/scale.py

@@ -0,0 +1,36 @@
+import torch
+from typing import Dict
+from .model import BitTransformerLM
+import torch.nn as nn
+
+
+def expand_model(model: BitTransformerLM, new_params: Dict) -> BitTransformerLM:
+    """Return a new model with updated params and copied weights."""
+    new_model = BitTransformerLM(**new_params)
+    new_state = new_model.state_dict()
+    old_state = model.state_dict()
+
+    for k, v in old_state.items():
+        if k in new_state:
+            dest = new_state[k]
+            slices = tuple(slice(0, min(d, s)) for d, s in zip(dest.shape, v.shape))
+            dest[slices].copy_(v[slices])
+            if dest.shape != v.shape:
+                mask = torch.ones_like(dest, dtype=torch.bool)
+                mask[slices] = False
+                if "bias" in k:
+                    dest[mask] = 0.0
+                else:
+                    dest[mask] = 0.001 * torch.randn_like(dest[mask])
+
+    for k, v in new_state.items():
+        if k not in old_state:
+            if "bias" in k:
+                v.zero_()
+            elif v.dim() > 1:
+                nn.init.normal_(v, mean=0.0, std=1e-3)
+            else:
+                v.zero_()
+
+    new_model.load_state_dict(new_state)
+    return new_model

BitTransformerLM/bit_transformer/static/style.css

@@ -0,0 +1,93 @@
+:root {
+    --primary: #1e40af;
+    --bg: #f5f6fa;
+}
+
+body {
+    font-family: Arial, sans-serif;
+    background-color: var(--bg);
+    margin: 0;
+    padding: 0;
+    line-height: 1.5;
+    color: #333;
+}
+
+.container {
+    max-width: 900px;
+    margin: 0 auto;
+    padding-bottom: 2rem;
+}
+
+h1 {
+    text-align: center;
+    background: var(--primary);
+    color: #fff;
+    margin: 0;
+    padding: 1rem 0;
+}
+
+section {
+    background: #fff;
+    margin: 1rem auto;
+    padding: 1rem 1.5rem;
+    border-radius: 8px;
+    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+    width: 90%;
+    max-width: 800px;
+}
+
+section h2 {
+    margin-top: 0;
+    color: var(--primary);
+    font-size: 1.25rem;
+}
+
+form {
+    display: flex;
+    flex-wrap: wrap;
+    gap: 0.5rem 1rem;
+}
+
+form input[type="text"],
+form input[type="number"],
+form textarea {
+    flex: 1 1 200px;
+    padding: 0.4em;
+    border: 1px solid #ccc;
+    border-radius: 4px;
+}
+
+form button,
+button#scaleBtn {
+    padding: 0.4em 0.8em;
+    border: none;
+    background: var(--primary);
+    color: #fff;
+    border-radius: 4px;
+    cursor: pointer;
+}
+
+form button:hover,
+button#scaleBtn:hover {
+    background-color: #1d4ed8;
+}
+
+pre, p#trainOut {
+    background: #f0f0f0;
+    padding: 0.5rem;
+    border-radius: 4px;
+    overflow-x: auto;
+}
+
+label {
+    display: flex;
+    align-items: center;
+    gap: 0.5rem;
+}
+
+img#plot {
+    max-width: 100%;
+    height: auto;
+    display: block;
+    margin: auto;
+}

BitTransformerLM/bit_transformer/telemetry.py

@@ -0,0 +1,95 @@
+import numpy as np
+from typing import Dict, List, TYPE_CHECKING
+
+import torch
+from sklearn.cluster import KMeans
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .model import BitTransformerLM
+
+
+class TelemetrySynthesizer:
+    """Analyze telemetry batches and cluster activation patterns."""
+
+    def __init__(self, n_clusters: int = 2) -> None:
+        self.n_clusters = n_clusters
+
+    def _summary(self, telemetry: Dict[str, List[torch.Tensor]]) -> np.ndarray:
+        """Compute activation/attention summaries for a single telemetry dict."""
+        acts = telemetry["activations"]
+        attn = telemetry["attention_maps"]
+        summaries = []
+        for a, m in zip(acts, attn):
+            mean = a.mean().item()
+            var = a.var(unbiased=False).item()
+            prob = m.softmax(-1)
+            entropy = -(prob * prob.clamp_min(1e-9).log()).sum(-1).mean().item()
+            summaries.append([mean, var, entropy])
+        return np.array(summaries).ravel()
+
+    def synthesize(
+        self, telemetries: List[Dict[str, List[torch.Tensor]]], bit_seqs: torch.Tensor
+    ) -> Dict[str, List]:
+        """Cluster telemetry summaries and return cluster info."""
+        data = np.stack([self._summary(t) for t in telemetries])
+        km = KMeans(n_clusters=self.n_clusters, n_init=1)
+        labels = km.fit_predict(data)
+        representatives: List[List[int]] = []
+        for c in range(self.n_clusters):
+            idx = np.where(labels == c)[0]
+            if len(idx) > 0:
+                representatives.append(bit_seqs[idx[0]].tolist())
+            else:
+                representatives.append([])
+        return {"cluster_assignments": labels.tolist(), "representatives": representatives}
+
+    def cluster_sequences(
+        self, model: "BitTransformerLM", bit_seqs: torch.Tensor
+    ) -> List[List[int]]:
+        """Run the model to gather telemetry and return representative sequences.
+
+        Parameters
+        ----------
+        model: BitTransformerLM
+            Model used to compute telemetry for each sequence.
+        bit_seqs: torch.Tensor
+            Tensor containing one bit sequence per row.
+
+        Returns
+        -------
+        list[list[int]]
+            Representative sequences chosen from KMeans clusters.
+        """
+        telemetries: List[Dict[str, List[torch.Tensor]]] = []
+        with torch.no_grad():
+            for seq in bit_seqs:
+                _, tele = model(seq.unsqueeze(0))
+                telemetries.append(tele)
+        info = self.synthesize(telemetries, bit_seqs)
+        return info["representatives"]
+
+
+def detect_metric_drift(
+    metrics_log: Dict[str, List[float]],
+    window: int = 10,
+    threshold: float = 0.2,
+) -> Dict[str, bool]:
+    """Detect metric drift between consecutive windows.
+
+    Args:
+        metrics_log: History of scalar metrics keyed by name.
+        window: Number of recent steps to compare.
+        threshold: Absolute difference required to flag drift.
+
+    Returns:
+        Dictionary mapping metric keys to a boolean drift indicator.
+    """
+    drift = {}
+    for key, values in metrics_log.items():
+        if len(values) < window * 2:
+            drift[key] = False
+            continue
+        recent = np.mean(values[-window:])
+        prev = np.mean(values[-2 * window : -window])
+        drift[key] = abs(recent - prev) > threshold
+    return drift

BitTransformerLM/bit_transformer/templates/dashboard.html

@@ -0,0 +1,454 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <title>Bit Transformer Dashboard</title>
+    <link rel="stylesheet" href="{{ url_for('static', filename='style.css') }}">
+    <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
+</head>
+<body>
+    <h1>Bit Transformer Dashboard</h1>
+    <div class="container">
+    <section>
+        <h2>Initialize Model</h2>
+        <form id="initForm">
+            d_model: <input type="number" name="d_model" value="{{ defaults.d_model }}" title="Model width (default {{ defaults.d_model }})"><br>
+            nhead: <input type="number" name="nhead" value="{{ defaults.nhead }}" title="Attention heads (default {{ defaults.nhead }})"><br>
+            num_layers: <input type="number" name="num_layers" value="{{ defaults.num_layers }}" title="Transformer layers (default {{ defaults.num_layers }})"><br>
+            dim_feedforward: <input type="number" name="dim_feedforward" value="{{ defaults.dim_feedforward }}" title="Feedforward dim (default {{ defaults.dim_feedforward }})"><br>
+            max_seq_len: <input type="number" name="max_seq_len" value="{{ defaults.max_seq_len }}" title="Max sequence length (default {{ defaults.max_seq_len }})"><br>
+            chunk_size: <input type="number" name="chunk_size" title="Chunked attention size"><br>
+            overlap: <input type="number" name="overlap" value="{{ defaults.overlap }}" title="Sliding window overlap"><br>
+            Reversible: <input type="checkbox" name="reversible" id="reversible_box" title="Use reversible layers (default {{ defaults.reversible }})"><br>
+            Gradient Checkpointing: <input type="checkbox" name="use_checkpoint" id="checkpoint_box" checked title="Enable gradient checkpointing (default {{ defaults.use_checkpoint }})"><br>
+            act_threshold: <input type="number" step="0.01" name="act_threshold" value="{{ defaults.act_threshold }}" title="ACT halt threshold (default {{ defaults.act_threshold }})"><br>
+            c_floor: <input type="number" step="0.01" name="c_floor" value="{{ c_floor }}" title="Complexity floor"><br>
+            s_floor: <input type="number" step="0.01" name="s_floor" value="{{ s_floor }}" title="Symbiosis floor"><br>
+            <button type="submit">Init</button>
+        </form>
+    </section>
+    <section>
+        <h2>Train Step</h2>
+        <form id="trainForm">
+            Bits (e.g. 0 1 0 1): <input type="text" name="bits" value="0 1 0 1"><br>
+            Upload file: <input type="file" id="train_file"><br>
+            <button type="submit">Train</button>
+        </form>
+        <label>Load sample dataset:
+            <select id="datasetSelect">
+                <option value="">--Select--</option>
+                <option value="wikitext2_train">Wikitext-2 (train)</option>
+                <option value="wikitext2_validation">Wikitext-2 (validation)</option>
+            </select>
+        </label>
+        <p id="trainOut"></p>
+    </section>
+    <section>
+        <h2>Scale Up</h2>
+        Width Mult: <input type="number" step="0.1" id="width_mult" value="1.0"><br>
+        <button id="scaleBtn">Scale Model</button>
+    </section>
+    <section>
+        <h2>Collapse Submodel</h2>
+        <form id="collapseForm">
+            Cluster Bits (JSON array of arrays):<br>
+            <textarea name="clusters" rows="3" cols="40">[[0,1,0,1],[1,1,0,0]]</textarea><br>
+            Target Params (JSON):<br>
+            <textarea name="params" rows="3" cols="40">{"d_model":32,"nhead":4,"num_layers":1,"dim_feedforward":64,"max_seq_len":16}</textarea><br>
+            Width Scale: <input type="number" step="0.1" id="width_scale" value="1.0"><br>
+            <button type="submit">Collapse</button>
+        </form>
+    </section>
+    <section>
+        <h2>Inference</h2>
+        <form id="inferForm">
+            Bits: <input type="text" name="bits" value="0 1 0 1"><br>
+            Upload file: <input type="file" id="infer_file"><br>
+            <button type="submit">Infer</button>
+        </form>
+        <pre id="inferOut"></pre>
+    </section>
+    <section>
+        <h2>Long Inference</h2>
+        <form id="inferLongForm">
+            Bits: <input type="text" name="bits" value="0 1 0 1"><br>
+            ctx_bits: <input type="number" name="ctx_bits" value="4096"><br>
+            overlap: <input type="number" name="overlap" value="256"><br>
+            <button type="submit">Infer Long</button>
+        </form>
+        <pre id="inferLongOut"></pre>
+    </section>
+    <section>
+        <h2>Text Inference</h2>
+        <form id="textInferForm">
+            Text: <input type="text" name="text" value="hello"><br>
+            <button type="submit">Infer Text</button>
+        </form>
+        <pre id="textInferOut"></pre>
+    </section>
+    <section>
+        <h2>&lambda; Weights</h2>
+        <form id="lambdaForm">
+            &lambda;<sub>K</sub>: <input type="range" min="0" max="2" step="0.1" id="lambda_K" oninput="lambda_K_val.innerText=value"><span id="lambda_K_val"></span><br>
+            &lambda;<sub>C</sub>: <input type="range" min="0" max="2" step="0.1" id="lambda_C" oninput="lambda_C_val.innerText=value"><span id="lambda_C_val"></span><br>
+            &lambda;<sub>S</sub>: <input type="range" min="0" max="2" step="0.1" id="lambda_S" oninput="lambda_S_val.innerText=value"><span id="lambda_S_val"></span><br>
+            <button type="submit">Update</button>
+        </form>
+    </section>
+    <section>
+        <h2>Diffusion LM</h2>
+        <label><input type="checkbox" id="diffusion_box"> Enable Diffusion Mode</label>
+    </section>
+    <section>
+        <h2>GPU Acceleration</h2>
+        <label><input type="checkbox" id="gpu_box"> Enable FSDP &amp; CUDA</label>
+    </section>
+    <section>
+        <h2>Enable Compression</h2>
+        <label><input type="checkbox" id="compression_box"> Compress I/O</label>
+        <p>Ratio: <span id="comp_ratio">1.0</span></p>
+    </section>
+    <section>
+        <h2>Quantization Aware Training</h2>
+        <label><input type="checkbox" id="qat_box"> Enable 4-bit QAT</label>
+    </section>
+    <section>
+        <h2>Model Status</h2>
+        <pre id="statusOut"></pre>
+    </section>
+    <section>
+        <h2>Telemetry</h2>
+        <canvas id="metricChart" width="600" height="300"></canvas>
+    </section>
+    <section>
+        <h2>Hugging Face Checkpoints</h2>
+        Repo ID: <input type="text" id="hf_repo"><br>
+        Token: <input type="password" id="hf_token" placeholder="optional"><br>
+        <button id="uploadBtn">Upload weights</button>
+        <button id="downloadBtn">Download weights</button>
+        <p id="hfStatus"></p>
+    </section>
+
+<script>
+async function postJSON(url, data){
+    const resp = await fetch(url, {method:'POST', headers:{'Content-Type':'application/json'}, body:JSON.stringify(data)});
+    return resp.json();
+}
+
+async function pollJob(id){
+    while(true){
+        const job = await fetch(`/job/${id}`).then(r=>r.json());
+        if(job.status === 'completed') return job.result;
+        if(job.status === 'error') throw job.error || 'Job failed';
+        await new Promise(r=>setTimeout(r, 1000));
+    }
+}
+
+function loadInitParams(){
+    const saved = JSON.parse(localStorage.getItem('init_params')||'{}');
+    const form = document.getElementById('initForm');
+    for(const [k,v] of Object.entries(saved)){
+        const el = form.elements[k];
+        if(!el) continue;
+        if(el.type === 'checkbox') el.checked = v; else el.value = v;
+    }
+}
+loadInitParams();
+
+function byteArrayToBits(arr){
+    const bits=[];
+    for(const b of arr){
+        for(let i=7;i>=0;i--) bits.push((b>>i)&1);
+    }
+    return bits;
+}
+
+let trainFileBits=null, inferFileBits=null, datasetBits=null;
+
+async function fileToBits(file){
+    if(file.type.startsWith('text')){
+        const text = await file.text();
+        const res = await postJSON('/text_to_bits', {text});
+        return res.bits;
+    }
+    const buf = await file.arrayBuffer();
+    return byteArrayToBits(new Uint8Array(buf));
+}
+
+let metricChart;
+async function initChart(){
+    const data = await fetch('/metrics').then(r=>r.json());
+    const labels = data.negentropy.map((_,i)=>i);
+    const ctx = document.getElementById('metricChart').getContext('2d');
+    metricChart = new Chart(ctx, {
+        type:'line',
+        data:{
+            labels:labels,
+            datasets:[
+                {label:'Negentropy', data:data.negentropy, borderColor:'blue', fill:false},
+                {label:'LZ Complexity', data:data.lz_complexity, borderColor:'orange', fill:false},
+                {label:'Symbiosis', data:data.symbiosis, borderColor:'green', fill:false}
+            ]
+        },
+        options:{responsive:false, interaction:{mode:'index', intersect:false}}
+    });
+}
+
+async function updateChart(){
+    const data = await fetch('/metrics').then(r=>r.json());
+    const labels = data.negentropy.map((_,i)=>i);
+    metricChart.data.labels = labels;
+    metricChart.data.datasets[0].data = data.negentropy;
+    metricChart.data.datasets[1].data = data.lz_complexity;
+    metricChart.data.datasets[2].data = data.symbiosis;
+    metricChart.update();
+}
+
+initChart();
+setInterval(updateChart, 2000);
+
+async function refreshStatus(){
+    const [s, c] = await Promise.all([fetch('/status'), fetch('/model_config')]);
+    const status = await s.json();
+    const config = await c.json();
+    document.getElementById('statusOut').innerText = JSON.stringify({...status, ...config}, null, 2);
+}
+
+document.getElementById('initForm').addEventListener('submit', async (e)=>{
+    e.preventDefault();
+    const fd = new FormData(e.target);
+    const obj = Object.fromEntries(fd.entries());
+    const ints = ['d_model','nhead','num_layers','dim_feedforward','max_seq_len','chunk_size','overlap'];
+    ints.forEach(k=>{ if(obj[k]===''){ delete obj[k]; } else obj[k]=parseInt(obj[k]); });
+    obj.reversible = document.getElementById('reversible_box').checked;
+    obj.use_checkpoint = document.getElementById('checkpoint_box').checked;
+    obj.act_threshold = parseFloat(obj.act_threshold);
+    const floors = {c_floor: parseFloat(obj.c_floor), s_floor: parseFloat(obj.s_floor)};
+    delete obj.c_floor; delete obj.s_floor;
+    await postJSON('/init', obj);
+    await postJSON('/config/telemetry', floors);
+    localStorage.setItem('init_params', JSON.stringify({...obj, ...floors}));
+    refreshStatus();
+    updateChart();
+});
+
+document.getElementById('trainForm').addEventListener('submit', async (e)=>{
+    e.preventDefault();
+    const form = e.target;
+    let payload;
+    if(trainFileBits){
+        payload = trainFileBits;
+    } else if(datasetBits){
+        payload = datasetBits;
+    } else {
+        payload = [form.bits.value.trim().split(/\s+/).map(Number)];
+    }
+    for(const el of form.elements) el.disabled = true;
+    const out = document.getElementById('trainOut');
+    out.innerText = '⏳';
+    try{
+        const job = await postJSON('/train', {bits: payload});
+        const res = await pollJob(job.job_id);
+        out.innerText = 'Loss: '+res.loss.toFixed(4);
+        if(res.ratio !== undefined){
+            document.getElementById('comp_ratio').innerText = res.ratio.toFixed(2);
+        }
+    } catch(err){
+        out.innerText = 'Error';
+        alert(err);
+    } finally {
+        for(const el of form.elements) el.disabled = false;
+        refreshStatus();
+        updateChart();
+    }
+});
+
+document.getElementById('train_file').addEventListener('change', async (e)=>{
+    const f = e.target.files[0];
+    if(!f) return;
+    const bits = await fileToBits(f);
+    trainFileBits = [bits];
+    datasetBits = null;
+    document.querySelector('#trainForm input[name="bits"]').value = bits.slice(0,64).join(' ');
+});
+
+document.querySelector('#trainForm input[name="bits"]').addEventListener('input', ()=>{
+    trainFileBits = null;
+    datasetBits = null;
+});
+
+document.getElementById('scaleBtn').addEventListener('click', async ()=>{
+    const btn = document.getElementById('scaleBtn');
+    const input = document.getElementById('width_mult');
+    const mult = parseFloat(input.value);
+    btn.disabled = true; input.disabled = true;
+    const original = btn.innerText; btn.innerText = '⏳';
+    try{
+        const job = await postJSON('/scale_up', {width_mult: mult});
+        await pollJob(job.job_id);
+    } catch(err){
+        alert(err);
+    } finally {
+        btn.innerText = original;
+        btn.disabled = false; input.disabled = false;
+        refreshStatus();
+        updateChart();
+    }
+});
+
+document.getElementById('collapseForm').addEventListener('submit', async (e)=>{
+    e.preventDefault();
+    const form = e.target;
+    const btn = form.querySelector('button');
+    for(const el of form.elements) el.disabled = true;
+    const clusters = JSON.parse(form.clusters.value);
+    const params = JSON.parse(form.params.value);
+    const w = parseFloat(document.getElementById('width_scale').value);
+    const original = btn.innerText; btn.innerText = '⏳';
+    try{
+        const job = await postJSON('/collapse', {clusters: clusters, params: params, width_scale: w});
+        await pollJob(job.job_id);
+    } catch(err){
+        alert(err);
+    } finally {
+        btn.innerText = original;
+        for(const el of form.elements) el.disabled = false;
+        refreshStatus();
+        updateChart();
+    }
+});
+
+document.getElementById('inferForm').addEventListener('submit', async (e)=>{
+    e.preventDefault();
+    let bits;
+    if(inferFileBits){
+        bits = inferFileBits;
+    } else if(datasetBits){
+        bits = [datasetBits[0]];
+    } else {
+        bits = [e.target.bits.value.trim().split(/\s+/).map(Number)];
+    }
+    const res = await postJSON('/infer', {bits});
+    if(res.error){
+        alert(res.error + '\n' + (res.suggestion||''));
+    } else {
+        document.getElementById('inferOut').innerText = JSON.stringify(res, null, 2);
+        if(res.ratio !== undefined){
+            document.getElementById('comp_ratio').innerText = res.ratio.toFixed(2);
+        }
+    }
+    refreshStatus();
+    updateChart();
+});
+
+document.getElementById('infer_file').addEventListener('change', async (e)=>{
+    const f = e.target.files[0];
+    if(!f) return;
+    const bits = await fileToBits(f);
+    inferFileBits = [bits];
+    datasetBits = null;
+    document.querySelector('#inferForm input[name="bits"]').value = bits.slice(0,64).join(' ');
+});
+
+document.querySelector('#inferForm input[name="bits"]').addEventListener('input', ()=>{
+    inferFileBits = null;
+    datasetBits = null;
+});
+
+document.getElementById('datasetSelect').addEventListener('change', async (e)=>{
+    const val = e.target.value;
+    trainFileBits = null;
+    inferFileBits = null;
+    if(!val){ datasetBits = null; return; }
+    const [name, split] = val.split('_');
+    const resp = await fetch(`/dataset?name=${name}&split=${split}&size=4&seq_len=64`);
+    const data = await resp.json();
+    datasetBits = data.bits;
+    const preview = data.bits[0].slice(0,64).join(' ');
+    document.querySelector('#trainForm input[name="bits"]').value = preview;
+    document.querySelector('#inferForm input[name="bits"]').value = preview;
+});
+
+document.getElementById('inferLongForm').addEventListener('submit', async (e)=>{
+    e.preventDefault();
+    const bits = e.target.bits.value.trim().split(/\s+/).map(Number);
+    const ctx = parseInt(e.target.ctx_bits.value);
+    const ov = parseInt(e.target.overlap.value);
+    const res = await postJSON('/infer_long', {bits: bits, ctx_bits: ctx, overlap: ov});
+    document.getElementById('inferLongOut').innerText = JSON.stringify(res, null, 2);
+    refreshStatus();
+    updateChart();
+});
+
+document.getElementById('textInferForm').addEventListener('submit', async (e)=>{
+    e.preventDefault();
+    const text = e.target.text.value;
+    const res = await postJSON('/infer_text', {text:text});
+    document.getElementById('textInferOut').innerText = JSON.stringify(res, null, 2);
+    refreshStatus();
+    updateChart();
+});
+
+async function loadLambdas(){
+    const resp = await fetch('/lambdas');
+    const vals = await resp.json();
+    for(const k of ['lambda_K','lambda_C','lambda_S']){
+        document.getElementById(k).value = vals[k];
+        document.getElementById(k+"_val").innerText = vals[k];
+    }
+}
+
+document.getElementById('lambdaForm').addEventListener('submit', async (e)=>{
+    e.preventDefault();
+    const data = {
+        lambda_K: parseFloat(document.getElementById('lambda_K').value),
+        lambda_C: parseFloat(document.getElementById('lambda_C').value),
+        lambda_S: parseFloat(document.getElementById('lambda_S').value),
+    };
+    await postJSON('/lambdas', data);
+    for(const k in data){
+        document.getElementById(k+"_val").innerText = data[k];
+    }
+    refreshStatus();
+});
+
+loadLambdas();
+
+function restoreToggle(id,key,endpoint,field){
+    const box = document.getElementById(id);
+    const saved = localStorage.getItem(key);
+    if(saved !== null){ box.checked = saved === 'true'; postJSON(endpoint,{[field]: box.checked}); }
+    box.addEventListener('change', async (e)=>{
+        await postJSON(endpoint, {[field]: e.target.checked});
+        localStorage.setItem(key, e.target.checked);
+        refreshStatus();
+    });
+}
+
+restoreToggle('diffusion_box','diffusion','/diffusion','diffusion');
+restoreToggle('gpu_box','use_gpu','/gpu','use_gpu');
+restoreToggle('compression_box','compression','/compression','compression');
+restoreToggle('qat_box','qat','/qat','qat');
+
+document.getElementById('uploadBtn').addEventListener('click', async ()=>{
+    const repo = document.getElementById('hf_repo').value;
+    const token = document.getElementById('hf_token').value;
+    const res = await postJSON('/save_checkpoint', {repo_id: repo, token: token||undefined});
+    document.getElementById('hfStatus').innerText = res.status || res.error;
+});
+
+document.getElementById('downloadBtn').addEventListener('click', async ()=>{
+    const repo = document.getElementById('hf_repo').value;
+    const token = document.getElementById('hf_token').value;
+    const res = await postJSON('/download_checkpoint', {repo_id: repo, token: token||undefined});
+    document.getElementById('hfStatus').innerText = res.status || res.error;
+    refreshStatus();
+    updateChart();
+});
+
+refreshStatus();
+</script>
+    </div>
+</body>
+</html>
+

BitTransformerLM/bit_transformer/torch_utils.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+import torch
+
+
+@contextmanager
+def cpu_autocast(enabled: bool = True):
+    """Context manager for bfloat16 autocast on CPU.
+
+    Parameters
+    ----------
+    enabled: bool, default True
+        Whether to enable autocast. When ``False`` this context manager
+        behaves like a no-op.
+    """
+    if enabled:
+        with torch.amp.autocast(device_type="cpu", dtype=torch.bfloat16):
+            yield
+    else:
+        yield

BitTransformerLM/bit_transformer/training.py

@@ -0,0 +1,250 @@
+"""Common training utilities for BitTransformer models."""
+
+from __future__ import annotations
+
+from typing import Callable, Dict, List, Optional
+import contextlib
+import sys
+import warnings
+import math
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+from .compression import compress_bits, pack_bits, unpack_bits
+from .optimization import configure_optimizer
+from .model import BitTransformerLM
+from .utils import set_dropout
+from .torch_utils import cpu_autocast
+
+
+def cosine_ramp(step: int, start: float, end: float, total_steps: int) -> float:
+    """Cosine ramp from ``start`` to ``end`` over ``total_steps``."""
+    if total_steps <= 0 or step >= total_steps:
+        return end
+    cos_inner = math.pi * step / total_steps
+    return start + (end - start) * (1 - math.cos(cos_inner)) / 2
+
+
+def train_loop(
+    model: BitTransformerLM,
+    data: torch.Tensor,
+    *,
+    epochs: int = 1,
+    extra_steps: int = 0,
+    compress_prob: float = 0.5,
+    direct_prob: float = 0.0,
+    batch_size: int = 8,
+    num_workers: int = 0,
+    accum_steps: int = 1,
+    amp: bool = False,
+    compile_model: bool = False,
+    log: bool = False,
+    forward_kwargs: Optional[Dict] = None,
+    optimizer: Optional[torch.optim.Optimizer] = None,
+    scheduler: Optional[torch.optim.lr_scheduler._LRScheduler] = None,
+    diffusion: bool = False,
+    noise_fn: Optional[Callable[[], float]] = None,
+    diffusion_curriculum: bool = False,
+    compress_warmup: int = 0,
+) -> List[Dict[str, float]]:
+    """Generic training loop supporting optional compression and diffusion.
+
+    ``compress_prob`` controls the fraction of batches that are run through
+    ``forward_compressed``. ``direct_prob`` instead feeds the model with the
+    bit-packed result of ``compress_bits`` after converting back to a bit
+    tensor. When enabled, metrics for direct-compressed batches are tracked
+    separately.
+
+    When ``diffusion`` is ``True`` the loop performs denoising training. Batches
+    are noised by randomly flipping bits with a probability given by
+    ``noise_fn`` (defaulting to a uniform draw in ``[0, 0.5]``). When
+    ``diffusion_curriculum`` is ``True`` the noise probability decreases
+    linearly from ``0.5`` to ``0.0`` over the training epochs. The model is
+    then trained to recover the clean sequence using full-context attention
+    (``causal=False``).
+
+    Existing ``optimizer`` and ``scheduler`` instances may be supplied to allow
+    integration with long-running training sessions, otherwise new ones are
+    created automatically.
+    """
+    if compile_model and sys.version_info < (3, 12) and torch.__version__ >= "2.1":
+        model = torch.compile(model)
+    elif compile_model:
+        warnings.warn("torch.compile skipped: requires torch>=2.1 and Python<3.12")
+
+    model.train()
+    set_dropout(model, 0.1)
+
+    device = next(model.parameters()).device
+    loader = DataLoader(
+        data,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        persistent_workers=num_workers > 0,
+    )
+    steps_per_epoch = max(1, len(loader))
+    total_updates = math.ceil(epochs * (steps_per_epoch + extra_steps) / accum_steps)
+    if optimizer is None or scheduler is None:
+        optimizer, scheduler = configure_optimizer(
+            model, lr=1e-3, total_steps=total_updates
+        )
+    metrics: List[Dict[str, float]] = []
+
+    global_step = 0
+    for epoch in range(epochs):
+        raw_losses: List[float] = []
+        raw_accs: List[float] = []
+        comp_losses: List[float] = []
+        comp_accs: List[float] = []
+        comp_ratios: List[float] = []
+        direct_losses: List[float] = []
+
+        last_batch = None
+        for step, batch in enumerate(loader):
+            last_batch = batch
+            batch = batch.to(device)
+            cur_compress = (
+                cosine_ramp(global_step, 0.0, compress_prob, compress_warmup)
+                if not diffusion
+                else compress_prob
+            )
+            if diffusion:
+                if diffusion_curriculum:
+                    p = 0.5 * (1 - epoch / max(1, epochs - 1))
+                else:
+                    p = noise_fn() if noise_fn is not None else float(torch.rand(()) * 0.5)
+                noise = (torch.rand_like(batch.float()) < p).long()
+                noisy = batch ^ noise
+                with (
+                    torch.cuda.amp.autocast(dtype=torch.bfloat16)
+                    if amp and torch.cuda.is_available()
+                    else cpu_autocast() if amp else contextlib.nullcontext()
+                ):
+                    logits, _ = model(noisy, causal=False)
+                pred = logits.reshape(-1, 2)
+                target = batch.reshape(-1)
+                loss = F.cross_entropy(pred, target) / accum_steps
+                acc = (pred.argmax(dim=-1) == target).float().mean().item()
+                raw_losses.append(loss.item() * accum_steps)
+                raw_accs.append(acc)
+                loss.backward()
+                if (step + 1) % accum_steps == 0:
+                    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                    optimizer.step()
+                    scheduler.step()
+                    optimizer.zero_grad()
+                global_step += 1
+                continue
+
+            r = torch.rand(())
+            key = "raw"
+            ratio = 1.0
+            target = batch[:, 1:].reshape(-1)
+
+            if r < direct_prob:
+                packed = [pack_bits(row.to(torch.uint8)) for row in batch]
+                unpacked = [unpack_bits(p, n_bits=batch.size(1)) for p in packed]
+                max_len = min(
+                    max(u.numel() for u in unpacked),
+                    model.pos_enc.pe.size(0),
+                )
+                padded = [F.pad(u[:max_len], (0, max_len - min(u.numel(), max_len))) for u in unpacked]
+                dc_batch = torch.stack(padded).long()
+                with (
+                    torch.cuda.amp.autocast(dtype=torch.bfloat16)
+                    if amp and torch.cuda.is_available()
+                    else cpu_autocast() if amp else contextlib.nullcontext()
+                ):
+                    logits, _ = model(dc_batch, **(forward_kwargs or {}))
+                ratio = sum(p.numel() for p in packed) / batch.numel()
+                target = dc_batch[:, 1:].reshape(-1)
+                key = "direct"
+            elif r < direct_prob + cur_compress:
+                comp_batch = [compress_bits(row.to(torch.uint8)) for row in batch]
+                with (
+                    torch.cuda.amp.autocast(dtype=torch.bfloat16)
+                    if amp and torch.cuda.is_available()
+                    else cpu_autocast() if amp else contextlib.nullcontext()
+                ):
+                    logits, _ = model.forward_compressed(comp_batch, **(forward_kwargs or {}))
+                ratio = sum(c.numel() for c in comp_batch) / batch.numel()
+                target = batch[:, 1:].reshape(-1)
+                key = "compressed"
+            else:
+                with (
+                    torch.cuda.amp.autocast(dtype=torch.bfloat16)
+                    if amp and torch.cuda.is_available()
+                    else cpu_autocast() if amp else contextlib.nullcontext()
+                ):
+                    logits, _ = model(batch, **(forward_kwargs or {}))
+
+            pred = logits[:, :-1, :].reshape(-1, 2)
+            loss = F.cross_entropy(pred, target) / accum_steps
+            acc = (pred.argmax(dim=-1) == target).float().mean().item()
+
+            loss.backward()
+            if (step + 1) % accum_steps == 0:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                optimizer.step()
+                scheduler.step()
+                optimizer.zero_grad()
+            global_step += 1
+
+            if key == "compressed":
+                comp_losses.append(loss.item() * accum_steps)
+                comp_accs.append(acc)
+                comp_ratios.append(ratio)
+            elif key == "direct":
+                direct_losses.append(loss.item() * accum_steps)
+                comp_ratios.append(ratio)
+            else:
+                raw_losses.append(loss.item() * accum_steps)
+                raw_accs.append(acc)
+
+        # run extra gradient updates using the final batch
+        if extra_steps > 0 and last_batch is not None and not diffusion:
+            for step in range(extra_steps):
+                with (
+                    torch.cuda.amp.autocast(dtype=torch.bfloat16)
+                    if amp and torch.cuda.is_available()
+                    else cpu_autocast() if amp else contextlib.nullcontext()
+                ):
+                    logits, _ = model(last_batch, **(forward_kwargs or {}))
+                    pred = logits[:, :-1, :].reshape(-1, 2)
+                    target = last_batch[:, 1:].reshape(-1)
+                    loss = F.cross_entropy(pred, target) / accum_steps
+                    acc = (pred.argmax(dim=-1) == target).float().mean().item()
+                loss.backward()
+                if (step + 1) % accum_steps == 0:
+                    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                    optimizer.step()
+                    scheduler.step()
+                    optimizer.zero_grad()
+                raw_losses.append(loss.item() * accum_steps)
+                raw_accs.append(acc)
+                global_step += 1
+
+        m = {
+            "raw_loss": float(sum(raw_losses) / len(raw_losses)) if raw_losses else 0.0,
+            "raw_acc": float(sum(raw_accs) / len(raw_accs)) if raw_accs else 0.0,
+            "compressed_loss": float(sum(comp_losses) / len(comp_losses)) if comp_losses else 0.0,
+            "compressed_acc": float(sum(comp_accs) / len(comp_accs)) if comp_accs else 0.0,
+            "direct_loss": float(sum(direct_losses) / len(direct_losses)) if direct_losses else 0.0,
+            "compression_ratio": float(sum(comp_ratios) / len(comp_ratios)) if comp_ratios else 0.0,
+        }
+        metrics.append(m)
+
+        if log:
+            print(
+                f"Epoch {epoch} "
+                f"raw_loss={m['raw_loss']:.4f} acc={m['raw_acc']:.3f} | "
+                f"compressed_loss={m['compressed_loss']:.4f} acc={m['compressed_acc']:.3f} "
+                f"direct_loss={m['direct_loss']:.4f} ratio={m['compression_ratio']:.2f}"
+            )
+
+    return metrics
+
+__all__ = ["train_loop"]

BitTransformerLM/bit_transformer/utils.py

@@ -0,0 +1,28 @@
+import os
+import gzip
+import torch
+import torch.nn as nn
+
+
+def save_model(model: torch.nn.Module, path: str) -> None:
+    """Save a model using gzip compression."""
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    with gzip.open(path, 'wb') as f:
+        torch.save(model, f)
+
+
+def load_model(path: str) -> torch.nn.Module:
+    """Load a model saved with ``save_model``."""
+    with gzip.open(path, 'rb') as f:
+        model = torch.load(f, map_location="cpu", weights_only=False)
+    return model
+
+
+def set_dropout(model: torch.nn.Module, p: float) -> None:
+    """Set dropout probability ``p`` for all dropout layers in ``model``."""
+    for module in model.modules():
+        if isinstance(module, nn.Dropout):
+            module.p = p
+
+
+__all__ = ["save_model", "load_model", "set_dropout"]

BitTransformerLM/bit_transformer_lm_codex_playbook.md

@@ -0,0 +1,278 @@
+---
+
+# 🧭 BitTransformerLM Codex Playbook (Merged)
+
+A single, actionable playbook that **implements optimizations first**, then **trains/ships the models**. Drop these prompts into your Codex/agent and run top-to-bottom.
+
+---
+
+## Phase 1 — Training Loop & Runtime Optimizations (apply these first)
+
+### Task 1 — Make batch size configurable & fix OneCycle accounting — COMPLETED ✅
+
+**Prompt:**
+
+```bash
+codex run bittransformerlm/patch \
+  --file bit_transformer/training.py \
+  --edit "Replace data.split(8) with DataLoader(batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, persistent_workers=True); compute steps_per_epoch=len(loader); set total_updates=epochs*(steps_per_epoch+extra_steps); pass total_updates into configure_optimizer"
+```
+
+✅ OneCycle’s horizon matches reality across runs.
+
+---
+
+### Task 2 — Remove hardcoded `total_steps=100` in dashboard/MCP — COMPLETED ✅
+
+**Prompt:**
+
+```bash
+codex run bittransformerlm/patch \
+  --file dashboard/manager.py \
+  --edit "When (re)creating OneCycleLR after init/scale_up/download, use computed total_steps from the upcoming training plan instead of hardcoded 100"
+```
+
+✅ Aligns scheduler behavior between direct loop and MCP/dashboard.
+
+---
+
+### Task 3 — Add mixed-precision autocast (AMP, BF16) — COMPLETED ✅
+
+**Prompt (pseudo-patch):**
+
+```python
+with torch.amp.autocast(device_type=("cuda" if torch.cuda.is_available() else "cpu"), dtype=torch.bfloat16):
+    logits = model(batch)
+    loss = criterion(logits, labels)
+loss.backward()
+```
+
+✅ 1.2–1.8× throughput on attention-heavy training. Keep grad-clip.
+
+---
+
+### Task 4 — Add gradient accumulation — COMPLETED ✅
+
+**Prompt:**
+
+```bash
+codex run bittransformerlm/patch \
+  --file bit_transformer/training.py \
+  --edit "Introduce --accum_steps; scale loss by 1/accum_steps; optimizer.step() every accum_steps; scheduler.step() every accum_steps"
+```
+
+✅ Simulates larger effective batch sizes without extra memory.
+
+---
+
+### Task 5 — Optimize dataset pipeline (mmap + streaming) — COMPLETED ✅
+
+**Prompt:**
+
+```bash
+codex run bittransformerlm/patch \
+  --file data/wikitext_schedule.py \
+  --edit "Precompute text->bit tensors aligned to max_seq_len; store in memory-mapped file; implement Dataset with __len__/__getitem__; use DataLoader(num_workers>0, persistent_workers=True)"
+```
+
+✅ Removes conversion bottlenecks on large corpora.
+
+---
+
+### Task 6 — Schedule compression probability (safer ramp) — COMPLETED ✅
+
+**Prompt (pseudo-code):**
+
+```python
+compress_prob = cosine_ramp(global_step, start=0.0, end=0.5, total_steps=warmup_steps)
+```
+
+✅ Prevents early instability from aggressive compression.
+
+---
+
+### Task 7 — Stabilize safety gate (EMA + burn‑in) — COMPLETED ✅
+
+**Prompt (pseudo-patch):**
+
+```python
+ema_val = ema(val_loss, decay=0.9)
+if step < burn_in_steps:
+    allow_training = True
+elif ema_val > threshold:
+    trigger_gate()
+```
+
+✅ Reduces false positives from noisy early validations.
+
+---
+
+### Task 8 — Enable `torch.compile` selectively — COMPLETED ✅
+
+**Prompt:**
+
+```bash
+codex run bittransformerlm/patch \
+  --file bit_transformer/training.py \
+  --edit "Enable torch.compile only if torch.__version__>=\"2.1\" and python<3.12; else skip with a clear warning"
+```
+
+✅ Opportunistic speedup where supported.
+
+---
+
+### Task 9 — Integrate FlashAttention / SDPA
+
+**Prompt (pseudo-patch):**
+
+```python
+from torch.nn import functional as F
+
+def forward_attention(q, k, v, is_causal=True):
+    return F.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
+```
+
+✅ Unlocks fused kernels; prefer `is_causal=True` over boolean masks.
+
+---
+
+### Task 10 — Cache causal masks — COMPLETED ✅
+
+**Prompt (pseudo-code):**
+
+```python
+mask_cache = {}
+
+def get_tri_mask(seq_len, device):
+    key = (seq_len, device)
+    if key not in mask_cache:
+        mask_cache[key] = torch.triu(
+            torch.ones(seq_len, seq_len, device=device, dtype=torch.bool), 1
+        )
+    return mask_cache[key]
+```
+
+✅ Avoids repeated `triu` allocations when masks are still needed.
+
+---
+
+### Task 11 — Fix stitched attention negative indexing — COMPLETED ✅
+
+**Prompt (pseudo-code):**
+
+```python
+start = max(s - overlap, 0)
+end   = min(s + chunk_size, T)
+canvas[..., start:end] = attn_chunk[..., : end - start]
+```
+
+✅ Prevents wrap-around misplacement during T×T map reconstruction.
+
+---
+
+### Task 12 — Default off: full T×T attention logging in chunked runs — COMPLETED ✅
+
+**Prompt:**
+
+```bash
+codex run bittransformerlm/patch \
+  --file bit_transformer/model.py \
+  --edit "Set full_attn_logging=False by default when chunk_size is set"
+```
+
+✅ Big memory/time savings without losing training signal.
+
+---
+
+## Phase 2 — Model Creation & Training Tasks (run after Phase 1)
+
+### Task A — Train the best current baseline (8×256 with ACT)
+
+**Prompt:**
+
+```bash
+codex run bittransformerlm/train \
+  --layers 8 \
+  --d_model 256 \
+  --nhead 8 \
+  --causal true \
+  --chunk_size 128 \
+  --act true \
+  --reversible true \
+  --checkpointing true \
+  --batch_size 64 \
+  --accum_steps 2 \
+  --amp bf16 \
+  --lr_schedule progressive_plateau \
+  --full_attn_logging false
+```
+
+✅ Reproduces the validated **sweet spot** with newly enabled efficiency features.
+
+---
+
+### Task B — CPU‑friendly deployment (8×128, INT8 + optional QAT)
+
+**Prompt:**
+
+```bash
+codex run bittransformerlm/train \
+  --layers 8 \
+  --d_model 128 \
+  --nhead 8 \
+  --causal true \
+  --chunk_size 128 \
+  --quantization int8 \
+  --qat true \
+  --reversible true \
+  --checkpointing true \
+  --batch_size 128 \
+  --accum_steps 1 \
+  --amp bf16
+```
+
+✅ Efficient CPU target; QAT optional based on deployment constraints.
+
+---
+
+### Task C — Cautious scale‑up candidate (16×256)
+
+**Prompt:**
+
+```bash
+codex run bittransformerlm/train \
+  --layers 16 \
+  --d_model 256 \
+  --nhead 8 \
+  --causal true \
+  --chunk_size 128 \
+  --act true \
+  --reversible true \
+  --checkpointing true \
+  --batch_size 48 \
+  --accum_steps 3 \
+  --amp bf16 \
+  --lr_schedule progressive_plateau
+```
+
+⚠️ Use only after data expansion and schedule retune.
+
+---
+
+## Recommended Execution Order
+
+1. **Phase 1 Tasks 1–12** (apply all optimizations).
+2. **Task A** baseline → validate.
+3. **Task B** CPU build → validate + (optional) QAT.
+4. **Task C** scale‑up **only** when data/schedule allow.
+
+---
+
+### Notes
+
+- Pair Phase 1 changes with CI that runs a short sanity fit (few hundred steps) to confirm loss decreases and no scheduler drift.
+- Keep `full_attn_logging=false` in chunked runs; enable selectively when inspecting attention.
+- When using SDPA, prefer `is_causal=True` and avoid passing dense masks unless required.
+
+---
+

BitTransformerLM/build_full_bits.py

@@ -0,0 +1,23 @@
+import pathlib
+import torch
+from datasets import load_dataset
+
+TXT_MB = 100
+OUT = pathlib.Path('full_bits.pt')
+
+
+def build_bits(out: pathlib.Path = OUT, txt_mb: int = TXT_MB) -> None:
+    ds = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train')
+    buf = bytearray()
+    for line in ds['text']:
+        buf.extend(line.encode() + b"\n")
+        if len(buf) >= txt_mb * 2 ** 20:
+            break
+    bits = []
+    for byte in buf:
+        bits.extend(int(b) for b in f'{byte:08b}')
+    tensor = torch.tensor(bits, dtype=torch.uint8)
+    torch.save(tensor, out)
+
+if __name__ == '__main__':
+    build_bits()

BitTransformerLM/context_extension.md

@@ -0,0 +1,43 @@
+Increasing the BitTransformerLM context window
+
+Current limitations and mechanisms
+The default max_seq_len in BitTransformerLM is 1 024 bits
+GitHub
+. Since text is encoded using parity bits (9 bits per byte)
+GitHub
+, this translates to roughly 113 bytes (≈113 characters) of input. The model uses full self‑attention, giving quadratic memory complexity in sequence length. To train on very long sequences, train_full_sequence slides a fixed‑size context window along a long bit tensor, detaching the computation graph periodically
+GitHub
+. Compression can shorten sequences via run‑length encoding
+GitHub
+, and chunked attention can divide long inputs into overlapping windows for attention calculations
+GitHub
+. However, the maximum positional encoding still defines an upper bound.
+
+Strategies to reach ~2 k‑word context (~18 k bits)
+Increase max_seq_len and positional encoding. The positional encoding precomputes a [max_len, d_model] matrix
+GitHub
+. Raising max_len to accommodate ~18 000 bits (for ~2 000 words × 9 bits per word) is possible but memory‑intensive. At d_model=128, the positional encoding would be ~18 000×128≈2.3 M floats (≈9 MB). That is reasonable for a CPU VM. Codex can modify the default max_seq_len and update any dependent tests.
+Use chunked attention and overlapping windows. LoggingTransformerEncoderLayer already supports chunk_size and overlap parameters
+GitHub
+. Setting chunk_size (e.g., 2 048 bits) and an overlap of e.g., 128 bits enables the model to handle sequences far longer than the attention window while still allowing information flow across chunks. Codex can expose chunk_size and overlap through the dashboard and CLI so users can tune them for longer contexts.
+Codex prompt example: “Modify the dashboard /init endpoint to accept chunk_size and overlap fields and pass them to BitTransformerLM. Update the HTML template to include input fields for these parameters.”
+Apply sliding‑window training and inference. The train_full_sequence method trains on long bit tensors by sliding a context window and detaching the graph every ctx_bits bits
+GitHub
+. For inference, a similar sliding approach could produce outputs for long sequences. Codex can add an infer_long_sequence method that divides a long bit sequence into overlapping windows, runs the model with causal=True to preserve order, and stitches the outputs.
+Prompt example: “Implement def infer_long_sequence(model: BitTransformerLM, bits: torch.Tensor, ctx_bits: int = 4096, overlap: int = 256): that processes a long bit tensor in sliding windows with overlap, uses causal=True, and returns the concatenated output bits.”
+Exploit run‑length compression more aggressively. Since binary data often contains runs of identical bits (e.g., long sequences of zeros), increasing compression ratio reduces the effective sequence length. Codex could add additional compression schemes (e.g., bit‑packing into bytes using numpy.packbits) and integrate them into the model’s I/O pipeline. Care must be taken to maintain parity bits for error detection.
+Prompt example: “Add functions pack_bits and unpack_bits that use numpy.packbits to pack 8 bits into a byte. Modify train_loop so that when direct_prob>0 the model is trained on packed bits with a suitable embedding.”
+Memory‑efficient attention alternatives. For even larger contexts, one could replace full attention with sparse, local or linear attention mechanisms. However, this would change the core architecture, which the task seeks to avoid. Using chunked attention (already present) and reversible layers is therefore preferred.
+Dynamic quantization and mixed precision. Larger context sizes increase model activations. Enabling use_autocast=True to compute in bfloat16 and applying quantize_dynamic after training reduces memory usage
+GitHub
+GitHub
+. Codex can create scripts that quantify memory usage and automatically toggle these features when large contexts are requested.
+Proposed Codex tasks to implement context extension
+Expose context parameters in the API/UI. Extend the dashboard and MCP server to allow clients to specify max_seq_len, chunk_size, overlap, and ctx_bits when initializing a model or running long inference.
+Prompt example: “Add optional parameters max_seq_len, chunk_size and overlap to the /init endpoint and pass them into BitTransformerLM and ModelManager. Update the HTML template to include these fields.”
+Implement sliding‑window inference. Add a function infer_long_sequence as described above and expose it via the dashboard and MCP server.
+Prompt example: “Add a new endpoint /infer_long to mcp_server.py that accepts a list of bits and processes them using a sliding window with overlap. The endpoint should return the predicted bits and telemetry summaries for each window.”
+Allow dynamic context scaling. Add a method to BitTransformerLM to adjust its pos_enc buffer when the context exceeds the current max_seq_len. This can be done by creating a new positional encoding tensor with the new length and copying the existing values.
+Prompt example: “Implement BitTransformerLM.expand_positional_encoding(new_len: int) that creates a new positional encoding buffer of size new_len and copies the existing encoding. Update the model’s max_seq_len accordingly.”
+Integrate aggressive compression. Implement alternative compression schemes (e.g., bit‑packing or general‑purpose compressors) and add toggles for them in training and inference. Evaluate compression ratio and latency to decide when to use them.
+Benchmark and tune hyperparameters. Write scripts to benchmark model memory use and throughput for various max_seq_len, chunk_size, reversible, use_act, and quantization settings. These benchmarks can inform safe defaults for the VM build.

BitTransformerLM/example.py

@@ -0,0 +1,6 @@
+from bit_transformer import example_training_step
+
+if __name__ == "__main__":
+    loss, telemetry = example_training_step()
+    print("Training loss:", loss)
+    print("Available telemetry:", list(telemetry.keys()))

BitTransformerLM/full_bits_train.py

@@ -0,0 +1,51 @@
+import pathlib
+import torch
+from bit_transformer import BitTransformerLM
+
+DATA_PATH = pathlib.Path('full_bits.pt')
+
+class BitSeq(torch.utils.data.IterableDataset):
+    def __init__(self, path: str | pathlib.Path = DATA_PATH, seq: int = 2048) -> None:
+        self.bits = torch.load(path, mmap=True)
+        self.seq = seq
+
+    def __len__(self) -> int:
+        return (self.bits.numel() // self.seq) - 1
+
+    def __iter__(self):
+        N = (self.bits.numel() // self.seq) - 1
+        for i in range(N):
+            s = i * self.seq
+            yield (
+                self.bits[s:s+self.seq].long(),
+                self.bits[s+1:s+self.seq+1].long(),
+            )
+
+def main() -> None:
+    dl = torch.utils.data.DataLoader(
+        BitSeq(DATA_PATH, seq=2048),
+        batch_size=8,
+        num_workers=0,
+        pin_memory=False,
+    )
+
+    model = BitTransformerLM(
+        d_model=64,
+        nhead=4,
+        num_layers=2,
+        dim_feedforward=256,
+        max_seq_len=2048,
+        reversible=True,
+        use_autocast=True,
+    )
+
+    loss_fn = torch.nn.CrossEntropyLoss()
+    xb, yb = next(iter(dl))
+    logits, _ = model(xb)
+    pred = logits.reshape(-1, 2)
+    target = yb.reshape(-1)
+    loss = loss_fn(pred, target)
+    print('Batch loss:', float(loss))
+
+if __name__ == '__main__':
+    main()

BitTransformerLM/integration_flow.py

@@ -0,0 +1,110 @@
+import torch
+from torch.profiler import profile
+from bit_transformer import (
+    BitTransformerLM,
+    quantize_dynamic,
+    hil_safe_inference,
+    collapse_submodel,
+)
+from bit_transformer.training import train_loop
+from bit_transformer.torch_utils import cpu_autocast
+
+def train(
+    model: BitTransformerLM,
+    data: torch.Tensor,
+    epochs: int = 3,
+    compress_prob: float = 0.5,
+    direct_prob: float = 0.0,
+    log: bool = False,
+    forward_kwargs: dict | None = None,
+) -> list[dict]:
+    """Train on bit sequences with optional random compression.
+
+    If ``direct_prob`` is positive, some batches are fed using their
+    run-length encoded representation packed into bits. Loss on these
+    direct-compressed batches is tracked separately.
+
+    Returns a list of per-epoch metric dictionaries containing raw and
+    compressed loss/accuracy statistics and the mean compression ratio.
+    """
+    return train_loop(
+        model,
+        data,
+        epochs=epochs,
+        compress_prob=compress_prob,
+        direct_prob=direct_prob,
+        log=log,
+        forward_kwargs=forward_kwargs,
+    )
+
+
+def main() -> None:
+    data = torch.randint(0, 2, (64, 128), dtype=torch.long)
+    validation_bits = torch.randint(0, 2, (16, 128), dtype=torch.long)
+    input_bits = torch.randint(0, 2, (1, 128), dtype=torch.long)
+    bit_sequence_data = data.tolist()
+
+    model = BitTransformerLM(
+        d_model=32,
+        nhead=4,
+        num_layers=1,
+        dim_feedforward=64,
+        max_seq_len=128,
+        use_act=True,
+        act_threshold=0.7,
+        reversible=True,
+        chunk_size=128,
+    )
+
+    for step in range(1, 13):
+        if step % 2 == 0:
+            model = model.double_width()
+        else:
+            model = model.double_layers()
+        train(model, data, epochs=3, compress_prob=0.5, log=True)
+        _, telemetry = model(validation_bits)
+        K = telemetry["negentropy_logits"].mean().item()
+        C = telemetry["lz_complexity_logits"].mean().item()
+        S = telemetry["symbiosis_score"].mean().item()
+        assert (
+            K > 0.3 and C > 0.35 and S > 0.5
+        ), f"Step {step} telemetry floor failure"
+
+    with cpu_autocast():
+        model(input_bits)
+
+    quantized_model = quantize_dynamic(model)
+    quantized_model.eval()
+
+    safe_output, _ = hil_safe_inference(
+        quantized_model, input_bits, c_floor=0.35, s_floor=0.5
+    )
+
+    student_model, _ = collapse_submodel(
+        bit_sequence_data,
+        target_params=dict(
+            d_model=16,
+            nhead=4,
+            num_layers=1,
+            dim_feedforward=32,
+            max_seq_len=128,
+        ),
+        floors={"negentropy": 0.3, "lz_complexity": 0.35, "symbiosis_score": 0.5},
+    )
+
+    compiled_model = (
+        torch.compile(student_model)
+        if hasattr(torch, "compile")
+        else student_model
+    )
+    compiled_model.eval()
+
+    with profile() as prof:
+        compiled_model(input_bits)
+
+    prof.export_chrome_trace("trace12.json")
+    print("Safe output bits:", safe_output.squeeze(0).tolist())
+
+
+if __name__ == "__main__":
+    main()

BitTransformerLM/integration_schedule.py

@@ -0,0 +1,379 @@
+import os
+import time
+import math
+from itertools import cycle
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from bit_transformer import (
+    BitTransformerLM,
+    text_to_bits,
+    quantize_dynamic,
+    prepare_qat_fx,
+    convert_qat_fx,
+    hil_safe_inference,
+    collapse_submodel,
+    diffusion_inference,
+    TelemetrySynthesizer,
+    save_distilled_model,
+)
+from bit_transformer.training import train_loop as train
+from bit_transformer.optimization import configure_optimizer, adjust_learning_rate
+from bit_transformer.utils import save_model, load_model, set_dropout
+from bit_transformer.torch_utils import cpu_autocast
+
+
+def lines_to_tensor(lines, max_len):
+    seqs = []
+    for text in lines:
+        bits = text_to_bits(text)[:max_len]
+        if len(bits) < max_len:
+            bits.extend([0] * (max_len - len(bits)))
+        seqs.append(bits)
+    return torch.tensor(seqs, dtype=torch.long)
+
+
+def load_wikitext(dataset_size=128, max_len=64):
+    try:
+        from datasets import load_dataset
+
+        ds = load_dataset("wikitext", "wikitext-2-raw-v1")
+        train_lines = [t for t in ds["train"]["text"] if t.strip()][:dataset_size]
+        valid_split = max(1, dataset_size // 4)
+        valid_lines = [t for t in ds["validation"]["text"] if t.strip()][:valid_split]
+        train = lines_to_tensor(train_lines, max_len)
+        valid = lines_to_tensor(valid_lines, max_len)
+        return train, valid, train_lines
+    except Exception as e:
+        print("Dataset load failed, using random bits", e)
+        train = torch.randint(0, 2, (dataset_size, max_len), dtype=torch.long)
+        valid = torch.randint(0, 2, (max_len, max_len), dtype=torch.long)
+        return train, valid, ["" for _ in range(len(train))]
+
+
+def _warmup(
+    model: BitTransformerLM,
+    data: torch.Tensor,
+    steps: int = 5,
+    freeze_old: bool = False,
+    old_layers: int = 0,
+    *,
+    diffusion: bool = False,
+    curriculum: bool = False,
+    optimizer: Optional[torch.optim.Optimizer] = None,
+    scheduler: Optional[torch.optim.lr_scheduler._LRScheduler] = None,
+) -> None:
+    """Run a short warm-up loop after expansion."""
+    model.train()
+    set_dropout(model, 0.1)
+    if freeze_old:
+        for idx, layer in enumerate(model.layers):
+            if idx < old_layers:
+                for p in layer.parameters():
+                    p.requires_grad_(False)
+    if optimizer is None or scheduler is None:
+        optimizer, scheduler = configure_optimizer(model, lr=1e-3, total_steps=steps)
+    it = iter(data.split(8))
+    for idx in range(steps):
+        try:
+            batch = next(it)
+        except StopIteration:
+            it = iter(data.split(8))
+            batch = next(it)
+        if diffusion:
+            p = 0.5 * (1 - idx / max(1, steps - 1)) if curriculum else 0.5
+            noise = (torch.rand_like(batch.float()) < p).long()
+            noisy = batch ^ noise
+            logits, _ = model(noisy, causal=False)
+            pred = logits.reshape(-1, 2)
+            target = batch.reshape(-1)
+        else:
+            logits, _ = model(batch)
+            pred = logits[:, :-1, :].reshape(-1, 2)
+            target = batch[:, 1:].reshape(-1)
+        loss = F.cross_entropy(pred, target)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        optimizer.step()
+        scheduler.step()
+        optimizer.zero_grad()
+    for p in model.parameters():
+        p.requires_grad_(True)
+    model.eval()
+    set_dropout(model, 0.0)
+
+
+def integration_schedule(
+    steps: int = 10,
+    max_len: int = 64,
+    dataset_size: int = 128,
+    *,
+    weights_path: str = "weights/model.pt.gz",
+    plateau_steps: int = 0,
+    collapsed_path: str | None = None,
+    epochs_per_step: int = 2,
+    extra_steps: int = 3,
+    collapse: bool = True,
+    diffusion: bool = False,
+    noise_schedule: str = "linear",
+    diffusion_steps: int = 8,
+    diffusion_curriculum: bool = False,
+    use_checkpoint: bool = True,
+    reversible: bool = True,
+    improve_thresh: float = 0.01,
+    qat: bool = False,
+):
+    start = time.time()
+    train_bits, valid_bits, train_lines = load_wikitext(dataset_size, max_len)
+    if os.path.exists(weights_path):
+        try:
+            model = load_model(weights_path)
+            print(f"Loaded model from {weights_path}")
+        except Exception as e:
+            print("Failed to load weights, initializing new model", e)
+            model = BitTransformerLM(
+                d_model=32,
+                nhead=4,
+                num_layers=1,
+                dim_feedforward=64,
+                max_seq_len=max_len,
+                use_act=True,
+                act_threshold=0.7,
+                reversible=reversible,
+                chunk_size=max_len,
+                use_autocast=True,
+                use_checkpoint=use_checkpoint,
+            )
+    else:
+        model = BitTransformerLM(
+            d_model=32,
+            nhead=4,
+            num_layers=1,
+            dim_feedforward=64,
+            max_seq_len=max_len,
+            use_act=True,
+            act_threshold=0.7,
+            reversible=reversible,
+            chunk_size=max_len,
+            use_autocast=True,
+            use_checkpoint=use_checkpoint,
+        )
+    if qat:
+        model = prepare_qat_fx(model)
+    results = []
+    scale_cycle = cycle(["layers", "width", "context"])
+    base_lr = 1e-3
+    prev_val_loss: Optional[float] = None
+    for step in range(steps):
+        model.train()
+        set_dropout(model, 0.1)
+        opt, sched = configure_optimizer(
+            model, lr=base_lr, total_steps=epochs_per_step
+        )
+        train(
+            model,
+            train_bits,
+            epochs=epochs_per_step,
+            extra_steps=extra_steps,
+            compress_prob=0.0 if diffusion else 1.0,
+            log=True,
+            diffusion=diffusion,
+            diffusion_curriculum=diffusion_curriculum,
+            optimizer=opt,
+            scheduler=sched,
+        )
+
+        model.eval()
+        set_dropout(model, 0.0)
+        with torch.no_grad():
+            logits, telemetry = model(valid_bits, causal=not diffusion)
+            if diffusion:
+                pred = logits.reshape(-1, 2)
+                target = valid_bits.reshape(-1)
+            else:
+                pred = logits[:, :-1, :].reshape(-1, 2)
+                target = valid_bits[:, 1:].reshape(-1)
+            val_loss = F.cross_entropy(pred, target).item()
+            k = telemetry["negentropy_logits"].mean().item()
+            c = telemetry["lz_complexity_logits"].mean().item()
+            s = telemetry["symbiosis_score"].mean().item()
+        print(f"Step {step} validation loss: {val_loss:.4f} K={k:.3f} C={c:.3f} S={s:.3f}")
+        results.append((step, val_loss, k, c, s))
+
+        if prev_val_loss is not None and prev_val_loss - val_loss < improve_thresh:
+            strategy = next(scale_cycle)
+            base_lr = adjust_learning_rate(opt, 1 / math.sqrt(2))
+            if strategy == "layers":
+                old_layers = model.num_layers
+                model = model.double_layers()
+                warm_opt, warm_sched = configure_optimizer(
+                    model, lr=base_lr, total_steps=100
+                )
+                _warmup(
+                    model,
+                    train_bits,
+                    steps=100,
+                    freeze_old=True,
+                    old_layers=old_layers,
+                    diffusion=diffusion,
+                    curriculum=diffusion_curriculum,
+                    optimizer=warm_opt,
+                    scheduler=warm_sched,
+                )
+            elif strategy == "width":
+                model = model.double_width()
+                warm_opt, warm_sched = configure_optimizer(
+                    model, lr=base_lr, total_steps=100
+                )
+                _warmup(
+                    model,
+                    train_bits,
+                    steps=100,
+                    diffusion=diffusion,
+                    curriculum=diffusion_curriculum,
+                    optimizer=warm_opt,
+                    scheduler=warm_sched,
+                )
+            else:
+                max_len *= 2
+                train_bits, valid_bits, train_lines = load_wikitext(
+                    dataset_size, max_len
+                )
+                model = model.double_length()
+                warm_opt, warm_sched = configure_optimizer(
+                    model, lr=base_lr, total_steps=100
+                )
+                _warmup(
+                    model,
+                    train_bits,
+                    steps=100,
+                    diffusion=diffusion,
+                    curriculum=diffusion_curriculum,
+                    optimizer=warm_opt,
+                    scheduler=warm_sched,
+                )
+
+        prev_val_loss = val_loss
+        if time.time() - start > 8 * 60:
+            print("Time limit reached")
+            break
+
+    # optional plateau phase at final size
+    for p in range(plateau_steps):
+        model.train()
+        set_dropout(model, 0.1)
+        train(
+            model,
+            train_bits,
+            epochs=epochs_per_step,
+            extra_steps=extra_steps,
+            compress_prob=0.0 if diffusion else 1.0,
+            log=True,
+            diffusion=diffusion,
+            diffusion_curriculum=diffusion_curriculum,
+        )
+        model.eval()
+        set_dropout(model, 0.0)
+        with torch.no_grad():
+            logits, telemetry = model(valid_bits, causal=not diffusion)
+            if diffusion:
+                pred = logits.reshape(-1, 2)
+                target = valid_bits.reshape(-1)
+            else:
+                pred = logits[:, :-1, :].reshape(-1, 2)
+                target = valid_bits[:, 1:].reshape(-1)
+            val_loss = F.cross_entropy(pred, target).item()
+            k = telemetry["negentropy_logits"].mean().item()
+            c = telemetry["lz_complexity_logits"].mean().item()
+            s = telemetry["symbiosis_score"].mean().item()
+        idx = steps + p
+        print(
+            f"Plateau {p} validation loss: {val_loss:.4f} K={k:.3f} C={c:.3f} S={s:.3f}"
+        )
+        results.append((idx, val_loss, k, c, s))
+        if time.time() - start > 8 * 60:
+            print("Time limit reached")
+            break
+
+    # final validation after last step
+    model.eval()
+    set_dropout(model, 0.0)
+    with torch.no_grad():
+        logits, telemetry = model(valid_bits, causal=not diffusion)
+        if diffusion:
+            pred = logits.reshape(-1, 2)
+            target = valid_bits.reshape(-1)
+        else:
+            pred = logits[:, :-1, :].reshape(-1, 2)
+            target = valid_bits[:, 1:].reshape(-1)
+        val_loss = F.cross_entropy(pred, target).item()
+        k = telemetry["negentropy_logits"].mean().item()
+        c = telemetry["lz_complexity_logits"].mean().item()
+        s = telemetry["symbiosis_score"].mean().item()
+
+    print(f"Final validation loss: {val_loss:.4f} K={k:.3f} C={c:.3f} S={s:.3f}")
+    results.append((steps + plateau_steps, val_loss, k, c, s))
+
+    # persist final model weights for future runs
+    save_model(model, weights_path)
+
+    input_bits = valid_bits[:1]
+    if qat:
+        qmodel = convert_qat_fx(model)
+    else:
+        with cpu_autocast():
+            model(input_bits)
+        qmodel = quantize_dynamic(model)
+    qmodel.eval()
+    try:
+        hil_safe_inference(
+            qmodel,
+            input_bits,
+            c_floor=0.3,
+            s_floor=0.5,
+            causal=not diffusion,
+            strict=not diffusion,
+        )
+    except RuntimeError as e:
+        print("Safety gate triggered", e)
+    collapsed = None
+    if collapse:
+        synth = TelemetrySynthesizer(n_clusters=8)
+        reps = synth.cluster_sequences(model, train_bits[:64])
+        floors = {"negentropy": 0.3, "lz_complexity": 0.35, "symbiosis_score": 0.5}
+        collapsed, metrics = collapse_submodel(
+            reps,
+            target_params=dict(
+                d_model=16,
+                nhead=4,
+                num_layers=1,
+                dim_feedforward=32,
+                max_seq_len=max_len,
+            ),
+            floors=floors,
+        )
+        collapsed.eval()
+        with torch.no_grad():
+            logits, _ = collapsed(valid_bits)
+            pred = logits[:, :-1, :].reshape(-1, 2)
+            target = valid_bits[:, 1:].reshape(-1)
+            c_loss = F.cross_entropy(pred, target).item()
+        print("Collapsed model validation loss:", c_loss)
+        if collapsed_path is not None:
+            save_distilled_model(
+                collapsed,
+                collapsed_path,
+                {**metrics, "val_loss": c_loss},
+                floors=floors,
+            )
+    if diffusion:
+        sample = diffusion_inference(
+            model, length=max_len, steps=diffusion_steps, schedule=noise_schedule
+        )
+        print("Diffusion sample:", sample[0].tolist())
+    return results, collapsed
+
+
+if __name__ == "__main__":
+    integration_schedule()

BitTransformerLM/mcp_server.py

@@ -0,0 +1,322 @@
+import io
+import os
+import gzip
+import uuid
+import traceback
+from concurrent.futures import ThreadPoolExecutor
+from flask import Flask, request, jsonify, send_file
+import matplotlib.pyplot as plt
+import torch
+
+from bit_transformer.dashboard_app import ModelManager
+from bit_transformer.dashboard import plot_telemetry
+from bit_transformer.hf_checkpoint import hf_login, save_checkpoint, download_checkpoint
+from bit_transformer.optimization import configure_optimizer
+from bit_transformer.bit_io import text_to_bits
+
+app = Flask(__name__)
+manager = ModelManager()
+
+# background job management
+executor = ThreadPoolExecutor(max_workers=4)
+jobs: dict[str, dict] = {}
+
+
+def _submit_job(fn, *args, **kwargs) -> str:
+    """Schedule a function for background execution and return a job id."""
+    job_id = str(uuid.uuid4())
+    jobs[job_id] = {"status": "queued", "result": None, "error": None, "logs": []}
+
+    def wrapper():
+        jobs[job_id]["status"] = "running"
+        try:
+            jobs[job_id]["result"] = fn(*args, **kwargs)
+            jobs[job_id]["status"] = "completed"
+        except Exception as err:  # pragma: no cover - captured for client
+            jobs[job_id]["status"] = "error"
+            jobs[job_id]["error"] = str(err)
+            jobs[job_id]["trace"] = traceback.format_exc()
+
+    executor.submit(wrapper)
+    return job_id
+
+
+@app.errorhandler(Exception)
+def handle_exception(err):
+    """Return JSON error responses with stack traces."""
+    return (
+        jsonify({"error": str(err), "trace": traceback.format_exc()}),
+        getattr(err, "code", 500),
+    )
+
+
+@app.route("/init", methods=["POST"])
+def init_model():
+    data = request.json or {}
+    int_fields = {
+        "d_model",
+        "nhead",
+        "num_layers",
+        "dim_feedforward",
+        "max_seq_len",
+        "chunk_size",
+        "overlap",
+    }
+    float_fields = {"act_threshold"}
+    bool_fields = {"reversible", "use_checkpoint"}
+    params = {}
+    for k, v in data.items():
+        if v is None:
+            params[k] = None
+        elif k in int_fields:
+            params[k] = int(v)
+        elif k in float_fields:
+            params[k] = float(v)
+        elif k in bool_fields:
+            params[k] = bool(v)
+        else:
+            params[k] = v
+    manager.init_model(params)
+    return jsonify({"status": "initialized", "params": params})
+
+@app.route("/train", methods=["POST"])
+def train_model():
+    bits = request.json["bits"]
+
+    def task():
+        tensor = torch.tensor(bits, dtype=torch.long)
+        loss, ratio = manager.train_step(tensor)
+        return {"loss": loss, "ratio": ratio}
+
+    job_id = _submit_job(task)
+    return jsonify({"job_id": job_id})
+
+
+@app.route("/train_epochs", methods=["POST"])
+def train_epochs_route():
+    data = request.json
+    bits = data["bits"]
+    epochs = int(data.get("epochs", 1))
+    compress_prob = float(data.get("compress_prob", 0.5))
+    direct_prob = float(data.get("direct_prob", 0.0))
+
+    def task():
+        tensor = torch.tensor(bits, dtype=torch.long)
+        metrics = manager.train_epochs(
+            tensor,
+            epochs=epochs,
+            compress_prob=compress_prob,
+            direct_prob=direct_prob,
+        )
+        return {"metrics": metrics}
+
+    job_id = _submit_job(task)
+    return jsonify({"job_id": job_id})
+
+@app.route("/scale_up", methods=["POST"])
+def scale_up():
+    width_mult = float(request.json.get("width_mult", 1.0))
+
+    def task():
+        manager.scale_up(width_mult)
+        return {
+            "status": "scaled",
+            "layers": manager.model.num_layers,
+            "d_model": manager.model.d_model,
+        }
+
+    job_id = _submit_job(task)
+    return jsonify({"job_id": job_id})
+
+@app.route("/collapse", methods=["POST"])
+def collapse_model():
+    cluster_bits = request.json["clusters"]
+    params = {k: int(v) for k, v in request.json["params"].items()}
+    width_scale = float(request.json.get("width_scale", 1.0))
+
+    def task():
+        manager.collapse(cluster_bits, params, width_scale)
+        return {"status": "collapsed"}
+
+    job_id = _submit_job(task)
+    return jsonify({"job_id": job_id})
+
+
+@app.route("/job/<job_id>", methods=["GET"])
+def get_job(job_id: str):
+    job = jobs.get(job_id)
+    if job is None:
+        return jsonify({"error": "not found"}), 404
+    return jsonify(job)
+
+
+@app.route("/jobs", methods=["GET"])
+def list_jobs():
+    return jsonify(jobs)
+
+@app.route("/lambdas", methods=["GET", "POST"])
+def update_lambdas():
+    if request.method == "POST":
+        data = request.json
+        manager.set_lambdas(float(data["lambda_K"]), float(data["lambda_C"]), float(data["lambda_S"]))
+        return jsonify({"status": "updated"})
+    else:
+        return jsonify({
+            "lambda_K": manager.lambda_K,
+            "lambda_C": manager.lambda_C,
+            "lambda_S": manager.lambda_S,
+        })
+
+@app.route("/diffusion", methods=["GET", "POST"])
+def update_diffusion():
+    if request.method == "POST":
+        manager.set_diffusion(bool(request.json.get("diffusion", False)))
+        return jsonify({"status": "updated"})
+    return jsonify({"diffusion": manager.diffusion})
+
+
+@app.route("/qat", methods=["GET", "POST"])
+def update_qat():
+    if request.method == "POST":
+        manager.set_qat(bool(request.json.get("qat", False)))
+        return jsonify({"status": "updated"})
+    return jsonify({"qat": manager.qat})
+
+
+@app.route("/gpu", methods=["GET", "POST"])
+def update_gpu():
+    if request.method == "POST":
+        manager.set_gpu(bool(request.json.get("use_gpu", False)))
+        return jsonify({"status": "updated"})
+    return jsonify({"use_gpu": manager.use_gpu})
+
+@app.route("/infer", methods=["POST"])
+def inference():
+    bits = torch.tensor(request.json["bits"], dtype=torch.long)
+    result = manager.infer(bits)
+    return jsonify(result)
+
+
+@app.route("/infer_long", methods=["POST"])
+def inference_long():
+    bits = torch.tensor(request.json["bits"], dtype=torch.long)
+    ctx = int(request.json.get("ctx_bits", 4096))
+    overlap = int(request.json.get("overlap", 256))
+    result = manager.infer_long(bits, ctx_bits=ctx, overlap=overlap)
+    return jsonify(result)
+
+@app.route("/infer_text", methods=["POST"])
+def inference_text():
+    text = request.json.get("text", "")
+    result = manager.infer_text(text)
+    return jsonify(result)
+
+@app.route("/status", methods=["GET"])
+def status():
+    return jsonify(manager.get_status())
+
+
+@app.route("/model_config", methods=["GET"])
+def model_config():
+    return jsonify(manager.get_model_config())
+
+
+@app.route("/metrics", methods=["GET"])
+def metrics():
+    return jsonify(manager.get_metrics())
+
+
+@app.route("/save_checkpoint", methods=["POST"])
+def save_checkpoint_route():
+    repo_id = request.json.get("repo_id")
+    token = request.json.get("token") or os.getenv("HF_TOKEN")
+    if manager.model is None:
+        return jsonify({"error": "model not initialized"}), 400
+    if token:
+        hf_login(token=token)
+    save_checkpoint(manager.model, repo_id=repo_id)
+    return jsonify({"status": "saved"})
+
+
+@app.route("/download_checkpoint", methods=["POST"])
+def download_checkpoint_route():
+    repo_id = request.json.get("repo_id")
+    token = request.json.get("token") or os.getenv("HF_TOKEN")
+    if token:
+        hf_login(token=token)
+    dest = manager.weights_path + ".gz"
+    ok = download_checkpoint(dest, repo_id=repo_id)
+    if not ok:
+        return jsonify({"status": "failed"}), 500
+    if manager.model is None:
+        return jsonify({"status": "downloaded", "loaded": False})
+    with gzip.open(dest, "rb") as f:
+        state = torch.load(f, map_location="cpu")
+    manager.model.load_state_dict(state)
+    manager.optimizer, manager.scheduler = configure_optimizer(
+        manager.model, lr=1e-3, total_steps=manager.total_steps
+    )
+    manager._apply_device()
+    manager._save_state()
+    return jsonify({"status": "downloaded", "loaded": True})
+
+@app.route("/plot.png")
+def plot_png():
+    fig, _ = plot_telemetry(manager.metrics)
+    buf = io.BytesIO()
+    fig.savefig(buf, format="png")
+    plt.close(fig)
+    buf.seek(0)
+    return send_file(buf, mimetype="image/png")
+
+
+@app.route("/text_to_bits", methods=["POST"])
+def text_to_bits_route():
+    text = request.json.get("text", "")
+    if len(text) > 100_000:
+        return jsonify({"error": "text too large"}), 413
+    return jsonify({"bits": text_to_bits(text)})
+
+
+@app.route("/dataset", methods=["GET"])
+def dataset_route():
+    name = request.args.get("name", "")
+    split = request.args.get("split", "train")
+    size = int(request.args.get("size", 1))
+    seq_len = int(request.args.get("seq_len", 64))
+    if size * seq_len > 1_000_000:
+        return jsonify({"error": "dataset too large"}), 413
+    if name == "wikitext2":
+        try:
+            from datasets import load_dataset
+
+            ds = load_dataset("wikitext", "wikitext-2-raw-v1", split=split)
+            lines = [t for t in ds["text"] if t.strip()][:size]
+        except Exception:
+            bits = torch.randint(0, 2, (size, seq_len), dtype=torch.long)
+            return jsonify({"bits": bits.tolist()})
+        bits_list = []
+        for text in lines:
+            b = text_to_bits(text)[:seq_len]
+            if len(b) < seq_len:
+                b.extend([0] * (seq_len - len(b)))
+            bits_list.append(b)
+        if len(bits_list) < size:
+            pad = size - len(bits_list)
+            bits_list.extend(torch.randint(0, 2, (pad, seq_len), dtype=torch.long).tolist())
+        return jsonify({"bits": bits_list})
+    return jsonify({"error": "unknown dataset"}), 400
+
+
+@app.route("/health")
+def health_check():
+    return jsonify({"status": "ok"})
+
+
+def run_mcp_server(host: str = "0.0.0.0", port: int = 7000) -> None:
+    app.run(host=host, port=port, debug=True)
+
+
+if __name__ == "__main__":
+    import torch
+    run_mcp_server()

BitTransformerLM/progressive_scaleup.py

@@ -0,0 +1,216 @@
+"""Legacy progressive scale-up demo.
+
+This script is retained for historical reference but has been superseded by
+``integration_schedule.py`` which provides a more flexible scaling workflow.
+"""
+
+import argparse
+import warnings
+import torch
+import torch.nn.functional as F
+from bit_transformer import (
+    BitTransformerLM,
+    configure_optimizer,
+    expand_model,
+    text_to_bits,
+)
+from bit_transformer.training import train_loop as basic_train
+
+warnings.warn(
+    "progressive_scaleup.py is deprecated; use integration_schedule.py instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+
+def progressive_scale_up(
+    eps: float = 0.65,
+    steps: int = 2,
+    width_mult: float = 1.0,
+    forward_kwargs: dict | None = None,
+) -> None:
+    """Demonstrate automatic scaling of the model on random data."""
+    params = dict(d_model=32, nhead=4, num_layers=1, dim_feedforward=64, max_seq_len=16)
+    model = BitTransformerLM(**params)
+    steps_per_epoch = 64 // 8
+    optimizer, scheduler = configure_optimizer(
+        model, lr=1e-3, total_steps=steps * steps_per_epoch
+    )
+
+    train = torch.randint(0, 2, (64, params["max_seq_len"]), dtype=torch.long)
+    valid = torch.randint(0, 2, (16, params["max_seq_len"]), dtype=torch.long)
+
+    for step in range(steps):
+        # one epoch over train
+        basic_train(
+            model,
+            train,
+            epochs=1,
+            compress_prob=0.5,
+            log=False,
+            forward_kwargs=forward_kwargs,
+        )
+
+        with torch.no_grad():
+            logits, _ = model(valid, **(forward_kwargs or {}))
+            pred = logits[:, :-1, :].reshape(-1, 2)
+            target = valid[:, 1:].reshape(-1)
+            val_loss = F.cross_entropy(pred, target).item()
+        print(f"Step {step} validation loss: {val_loss:.4f}")
+        if val_loss < eps:
+            params["num_layers"] *= 2
+            params["d_model"] = int(params["d_model"] * width_mult)
+            params["dim_feedforward"] = int(params["dim_feedforward"] * width_mult)
+            model = expand_model(model, params)
+            optimizer, scheduler = configure_optimizer(
+                model, lr=1e-3, total_steps=steps * steps_per_epoch
+            )
+        print(
+            "Scaled model to", params["num_layers"], "layers and width", params["d_model"]
+        )
+
+
+def progressive_scale_up_text(
+    improve_thresh: float = 0.01,
+    steps: int = 2,
+    width_mult: float = 2.0,
+    max_len: int = 64,
+    dataset_size: int = 512,
+    forward_kwargs: dict | None = None,
+) -> None:
+    """Scale up using WikiText2 lines converted to bits.
+
+    Parameters
+    ----------
+    improve_thresh: float
+        Relative validation loss improvement required to avoid scaling.
+        If improvement is <= this threshold, model size is increased.
+    steps: int
+        Number of training steps.
+    width_mult: float
+        Multiplier applied when increasing model width.
+    max_len: int
+        Initial sequence length.
+    dataset_size: int
+        Number of training lines to load from WikiText2.
+    forward_kwargs: dict | None
+        Extra keyword arguments for the forward pass.
+    """
+    from datasets import load_dataset
+
+    ds = load_dataset("wikitext", "wikitext-2-raw-v1")
+    train_iter = ds["train"]["text"]
+    valid_iter = ds["validation"]["text"]
+
+    train_lines = []
+    for line in train_iter:
+        train_lines.append(line)
+        if len(train_lines) >= dataset_size:
+            break
+
+    valid_lines = []
+    for line in valid_iter:
+        valid_lines.append(line)
+        if len(valid_lines) >= dataset_size // 4:
+            break
+
+    def lines_to_tensor(lines: list[str], length: int) -> torch.Tensor:
+        seqs = []
+        for text in lines:
+            bits = text_to_bits(text)[:length]
+            if len(bits) < length:
+                bits.extend([0] * (length - len(bits)))
+            seqs.append(bits)
+        return torch.tensor(seqs, dtype=torch.long)
+
+    train = lines_to_tensor(train_lines, max_len)
+    valid = lines_to_tensor(valid_lines, max_len)
+
+    params = dict(
+        d_model=32,
+        nhead=4,
+        num_layers=1,
+        dim_feedforward=64,
+        max_seq_len=max_len,
+    )
+    model = BitTransformerLM(**params)
+    steps_per_epoch = len(train) // 8
+    optimizer, scheduler = configure_optimizer(
+        model, lr=1e-3, total_steps=steps * max(1, steps_per_epoch)
+    )
+
+    prev_loss: float | None = None
+    scale_length = True
+
+    for step in range(steps):
+        basic_train(
+            model,
+            train,
+            epochs=1,
+            compress_prob=0.5,
+            log=False,
+            forward_kwargs=forward_kwargs,
+        )
+
+        with torch.no_grad():
+            logits, _ = model(valid, **(forward_kwargs or {}))
+            pred = logits[:, :-1, :].reshape(-1, 2)
+            target = valid[:, 1:].reshape(-1)
+            val_loss = F.cross_entropy(pred, target).item()
+        print(f"Step {step} validation loss: {val_loss:.4f}")
+        if prev_loss is not None:
+            improvement = (prev_loss - val_loss) / max(prev_loss, 1e-8)
+            if improvement <= improve_thresh:
+                if scale_length:
+                    params["max_seq_len"] *= 2
+                    train = lines_to_tensor(train_lines, params["max_seq_len"])
+                    valid = lines_to_tensor(valid_lines, params["max_seq_len"])
+                    model = model.double_length()
+                    steps_per_epoch = len(train) // 8
+                    scale_type = "length"
+                else:
+                    params["d_model"] = int(params["d_model"] * width_mult)
+                    params["dim_feedforward"] = int(params["dim_feedforward"] * width_mult)
+                    model = expand_model(model, params)
+                    scale_type = "width"
+                optimizer, scheduler = configure_optimizer(
+                    model, lr=1e-3, total_steps=steps * max(1, steps_per_epoch)
+                )
+                scale_length = not scale_length
+                param_count = sum(p.numel() for p in model.parameters())
+                print(
+                    f"Scaled {scale_type}; seq_len={params['max_seq_len']} width={params['d_model']} params={param_count}"
+                )
+        prev_loss = val_loss
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Progressively scale model length and width")
+    parser.add_argument("--steps", type=int, default=2, help="number of training steps")
+    parser.add_argument(
+        "--improve-thresh",
+        type=float,
+        default=0.01,
+        help="relative loss improvement required to avoid scaling",
+    )
+    parser.add_argument(
+        "--width-mult", type=float, default=2.0, help="width multiplier when scaling"
+    )
+    parser.add_argument("--causal", action="store_true", help="use causal attention during training")
+    parser.add_argument("--wikitext", action="store_true", help="use WikiText2 dataset")
+    args = parser.parse_args()
+    if args.wikitext:
+        progressive_scale_up_text(
+            improve_thresh=args.improve_thresh,
+            steps=args.steps,
+            width_mult=args.width_mult,
+            forward_kwargs={"causal": args.causal} if args.causal else None,
+        )
+    else:
+        progressive_scale_up(
+            eps=args.improve_thresh,
+            steps=args.steps,
+            width_mult=args.width_mult,
+            forward_kwargs={"causal": args.causal} if args.causal else None,
+        )
+

BitTransformerLM/pyproject.toml

@@ -0,0 +1,18 @@
+[build-system]
+requires = ["setuptools>=67", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "bit-transformer"
+version = "0.1.0"
+description = "Bit-based transformer language model nearing v1 release with telemetry"
+readme = "README.md"
+requires-python = ">=3.10"
+license = {text = "All Rights Reserved"}
+authors = [{name = "WCNegentropy"}]
+
+[project.urls]
+Homepage = "https://example.com"
+
+[tool.setuptools.packages.find]
+include = ["bit_transformer"]

BitTransformerLM/recursive_integration_flow.py

@@ -0,0 +1,128 @@
+import torch
+import torch.nn.functional as F
+from torch.profiler import profile
+from bit_transformer import (
+    BitTransformerLM,
+    quantize_dynamic,
+    hil_safe_inference,
+    collapse_submodel,
+)
+from bit_transformer.training import train_loop
+from bit_transformer.torch_utils import cpu_autocast
+
+
+def train(
+    model: BitTransformerLM,
+    data: torch.Tensor,
+    epochs: int = 1,
+    compress_prob: float = 0.5,
+    log: bool = False,
+    forward_kwargs: dict | None = None,
+) -> list[dict]:
+    """Train with random compression; returns per-epoch metrics."""
+    return train_loop(
+        model,
+        data,
+        epochs=epochs,
+        compress_prob=compress_prob,
+        direct_prob=0.0,
+        log=log,
+        forward_kwargs=forward_kwargs,
+    )
+
+
+def recursive_integration_flow(steps: int = 4, max_len: int = 64) -> None:
+    """Run a dynamic scale-up loop with telemetry-based gating."""
+    train_bits = torch.randint(0, 2, (64, max_len), dtype=torch.long)
+    valid_bits = torch.randint(0, 2, (16, max_len), dtype=torch.long)
+    input_bits = torch.randint(0, 2, (1, max_len), dtype=torch.long)
+    bit_sequence_data = train_bits.tolist()
+
+    best_K = best_C = best_S = 0.0
+
+    model = BitTransformerLM(
+        d_model=32,
+        nhead=4,
+        num_layers=1,
+        dim_feedforward=64,
+        max_seq_len=max_len,
+        use_act=True,
+        act_threshold=0.7,
+        reversible=True,
+        chunk_size=max_len,
+        use_autocast=True,
+    )
+
+    results = []
+    for step in range(steps + 1):
+        epochs = min(10, 2 + step // 2)
+        train(model, train_bits, epochs=epochs, compress_prob=0.5, log=True)
+
+        with torch.no_grad():
+            with cpu_autocast():
+                logits, telemetry = model(valid_bits)
+            pred = logits[:, :-1, :].reshape(-1, 2)
+            target = valid_bits[:, 1:].reshape(-1)
+            val_loss = F.cross_entropy(pred, target).item()
+            k = telemetry["negentropy_logits"].mean().item()
+            c = telemetry["lz_complexity_logits"].mean().item()
+            s = telemetry["symbiosis_score"].mean().item()
+
+        print(f"Step {step} validation loss: {val_loss:.4f} K={k:.3f} C={c:.3f} S={s:.3f}")
+        results.append((step, val_loss, k, c, s))
+
+        if step > 0:
+            if k < best_K - 0.3 or c < best_C - 0.3 or s < best_S - 0.3:
+                print(f"\u26a0\ufe0f Step {step} regressed below metric floor. Halting.")
+                break
+        best_K = max(best_K, k)
+        best_C = max(best_C, c)
+        best_S = max(best_S, s)
+
+        if step < steps:
+            if step % 2 == 0:
+                model = model.double_width()
+            else:
+                model = model.double_layers()
+
+    # Post-scaling optimizations
+    with cpu_autocast():
+        model(input_bits)
+
+    qmodel = quantize_dynamic(model)
+    qmodel.eval()
+
+    safe_output = hil_safe_inference(
+        qmodel, input_bits, c_floor=0.5, s_floor=0.2
+    )
+
+    student_model, _ = collapse_submodel(
+        bit_sequence_data,
+        target_params=dict(
+            d_model=16,
+            nhead=4,
+            num_layers=1,
+            dim_feedforward=32,
+            max_seq_len=max_len,
+        ),
+        floors={"negentropy": 0.2, "lz_complexity": 0.5, "symbiosis_score": 0.2},
+    )
+
+    if hasattr(torch, "compile"):
+        try:
+            compiled = torch.compile(student_model)
+        except RuntimeError as exc:
+            print(f"Compilation skipped: {exc}")
+            compiled = student_model
+    else:
+        compiled = student_model
+    compiled.eval()
+
+    with profile() as prof:
+        compiled(input_bits)
+    prof.export_chrome_trace("trace12.json")
+    print("Safe output bits:", safe_output[0].tolist())
+
+
+if __name__ == "__main__":
+    recursive_integration_flow()

BitTransformerLM/requirements.txt

@@ -0,0 +1,16 @@
+-i https://pypi.org/simple
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch==2.7.1+cpu
+
+# --extra-index-url https://download.pytorch.org/whl/cu118
+# torch==2.7.1+cu118
+
+pytest==8.4.1
+scikit-learn==1.7.1
+matplotlib==3.10.3
+datasets==4.0.0
+flask==3.1.1
+numpy==2.3.1
+requests==2.32.3
+watchdog==6.0.0
+huggingface-hub==0.34.3

BitTransformerLM/review_cli.py

@@ -0,0 +1,49 @@
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from bit_transformer import BitTransformerLM
+
+
+def list_candidates(path: Path):
+    models = sorted(path.glob("*.pt"))
+    for m in models:
+        metrics_file = m.with_suffix(m.suffix + ".json")
+        metrics = {}
+        if metrics_file.exists():
+            with open(metrics_file) as f:
+                metrics = json.load(f)
+        yield m, metrics
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Review distilled submodels")
+    parser.add_argument("directory", type=Path, help="Directory with candidate models")
+    parser.add_argument("--approve-dir", type=Path, default=Path("approved"), help="Directory to store approved models")
+    args = parser.parse_args()
+
+    args.approve_dir.mkdir(exist_ok=True)
+    log_file = args.approve_dir / "review_log.jsonl"
+
+    for model_path, metrics in list_candidates(args.directory):
+        print("Candidate:", model_path.name)
+        for k, v in metrics.items():
+            print(f"  {k}: {v}")
+        ans = input("Approve this model? [y/N] ").strip().lower()
+        if ans == "y":
+            approved_path = args.approve_dir / model_path.name
+            torch.save(torch.load(model_path), approved_path)
+            entry = {"model": approved_path.name, "metrics": metrics, "approved": True}
+            with open(log_file, "a") as lf:
+                lf.write(json.dumps(entry) + "\n")
+            print("Approved and saved to", approved_path)
+        else:
+            entry = {"model": model_path.name, "metrics": metrics, "approved": False}
+            with open(log_file, "a") as lf:
+                lf.write(json.dumps(entry) + "\n")
+            print("Rejected", model_path.name)
+
+
+if __name__ == "__main__":
+    main()

BitTransformerLM/start.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+set -e
+python mcp_server.py &
+sleep 2
+python -m bit_transformer.dashboard_app

BitTransformerLM/state_of_the_repo_audit.md

@@ -0,0 +1,98 @@
+summary: |
+  The **BitTransformerLM** repository is well-structured and aligns closely with the README’s feature set. 
+  All core functionalities (bit-level modeling, telemetry metrics, progressive scaling, compression, context extension, diffusion mode, dashboard, etc.) are present and largely consistent with documentation. 
+  The code is generally clean and well-tested (no TODOs or obvious dead code) with an effective CI in place:contentReference[oaicite:0]{index=0}. 
+  We identified a few issues via static analysis: a critical **security flaw** where the dashboard’s `/exec` endpoint executes arbitrary code:contentReference[oaicite:1]{index=1}, a missing import that breaks the compression toggle:contentReference[oaicite:2]{index=2}:contentReference[oaicite:3]{index=3}, and a rare edge-case in bit-sequence decompression logic:contentReference[oaicite:4]{index=4}. 
+  No functions exceed 300 lines, though the `BitTransformerLM.forward` method is complex with deeply nested logic (~6 levels) and duplicated code blocks for the halting mechanism. 
+  Naming conventions are consistent (snake_case for functions, CamelCase for classes), and dependency versions are up-to-date. 
+  Documentation and code behavior are in sync – for example, the MCP server’s `/health` endpoint described in docs is implemented:contentReference[oaicite:5]{index=5}. 
+  Overall, the project appears **nearly production-ready**, with these fixes and refinements needed before a 1.0 release.
+
+findings:
+  - severity: P0
+    effort: S
+    category: security
+    location: bit_transformer/dashboard_app.py:533
+    description: "Unrestricted `/exec` HTTP endpoint allows arbitrary code execution:contentReference[oaicite:6]{index=6}."
+    recommendation: "Disable or restrict the `/exec` route (e.g. remove it or require an admin token) to prevent remote code execution."
+    status: completed ✅
+  - severity: P1
+    effort: S
+    category: static
+    location: bit_transformer/dashboard_app.py:195
+    description: "NameError risk – `compress_bits` is used without being imported:contentReference[oaicite:7]{index=7}:contentReference[oaicite:8]{index=8}."
+    recommendation: "Import the `compress_bits` function in `dashboard_app.py` (e.g. `from .compression import compress_bits`) so compression toggles don’t crash."
+    status: completed ✅
+  - severity: P2
+    effort: M
+    category: static
+    location: bit_transformer/model.py:320
+    description: "Edge-case bug – `_maybe_decompress` skips decompression if all values ≤1:contentReference[oaicite:9]{index=9}, which can misinterpret run-length encoding outputs of all 1s."
+    recommendation: "Adjust the decompress condition (e.g. track whether input was compressed) to ensure even uniformly alternating bit sequences get properly decompressed."
+    status: completed ✅
+  - severity: P3
+    effort: M
+    category: static
+    location: bit_transformer/model.py:415
+    description: "Duplicate code – nearly identical halting logic is implemented in both reversible and normal forward loops:contentReference[oaicite:10]{index=10}:contentReference[oaicite:11]{index=11}."
+    recommendation: "Refactor the halting (ACT) mechanism into a helper function to avoid repetition and reduce maintenance effort."
+    status: completed ✅
+  - severity: P3
+    effort: M
+    category: static
+    location: bit_transformer/model.py:368
+    description: "Complex logic – `BitTransformerLM.forward` contains deeply nested control flow (up to 5-6 levels) for reversible layers, ACT, etc."
+    recommendation: "Consider simplifying or breaking up the forward pass (e.g. separate functions for reversible vs. standard flow) to improve readability and maintainability."
+    status: completed ✅
+  - severity: P3
+    effort: S
+    category: static
+    location: bit_transformer/dashboard_app.py:125
+    description: "Config parsing quirk – booleans in `ModelManager.init_model` are cast to int (True→1) instead of preserved as bool."
+    recommendation: "Handle boolean fields explicitly (e.g. do not cast values for keys like `reversible` or `use_act` to int) to avoid confusion and potential type issues."
+    status: completed ✅
+
+codex_tasks:
+  - codex_prompt: "Remove or secure the dangerous `/exec` endpoint in the dashboard to prevent arbitrary code execution."
+    acceptance_test: |
+      import requests, subprocess
+      Attempt to call the /exec endpoint with a harmless command
+      try:
+          resp = requests.post("http://localhost:5000/exec", json={"code": "print('OK')"}, timeout=5)
+      except Exception as e:
+          resp = e.response if hasattr(e, 'response') else None
+      The endpoint should be removed or secured, so it should either 404 or refuse access
+      assert resp is None or resp.status_code in (403, 404), "Exec endpoint still accessible!"
+    status: completed ✅
+  - codex_prompt: "Import the `compress_bits` function in `dashboard_app.py` so that enabling compression no longer raises a NameError."
+    acceptance_test: |
+      import torch
+      from bit_transformer.dashboard_app import ModelManager
+      mgr = ModelManager()
+      mgr.set_compression(True)
+      bits = torch.randint(0, 2, (1, 8), dtype=torch.long)
+      try:
+          loss, ratio = mgr.train_step(bits)
+      except NameError as e:
+          raise AssertionError(f"NameError not resolved: {e}")
+      assert isinstance(loss, float) and 0 <= ratio <= 1.0, "Compression training failed"
+    status: completed ✅
+  - codex_prompt: "Fix `_maybe_decompress` in `model.py` to always decompress run-length encoded sequences (even if all run lengths are 1) before computing metrics."
+    acceptance_test: |
+      import torch
+      from bit_transformer import BitTransformerLM, compress_bits, decompress_bits
+      Create an alternating bit sequence where compress_bits yields only count=1 values
+      bits = torch.tensor([0,1]*8, dtype=torch.uint8)
+      comp = compress_bits(bits)
+      model = BitTransformerLM(d_model=16, nhead=2, num_layers=1, dim_feedforward=32, max_seq_len=len(bits))
+      Compute negentropy on compressed vs original and compare
+      neg_comp = model.negentropy_kpi(comp.unsqueeze(0))
+      neg_raw = model.negentropy_kpi(bits.unsqueeze(0))
+      assert torch.allclose(neg_comp, neg_raw, atol=1e-6), "Negentropy differs for compressed input – decompression fix failed"
+    status: completed ✅
+
+metrics:
+  loc_total: 3770
+  todo_count: 0
+  duplicate_block_count: 3
+  oversized_functions: 0