# Dokodemo AI — どこでも AI

> **Run trillion-parameter LLMs on a 4 GB GPU — or a 4 GB CPU RAM.**
> Model-agnostic. Expert-aware. Anywhere.

[![PyPI](https://img.shields.io/pypi/v/dokodemo-ai)](https://pypi.org/project/dokodemo-ai/)
[![License: Research Free / Commercial Paid](https://img.shields.io/badge/License-Research%20Free%20%7C%20Commercial%20Paid-blue.svg)](LICENSE)
[![Python 3.9+](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/)

**Author**: [Tuan Aqeel Bohoran](https://github.com/aqeelbohoran) — aqeelbohoran@gmail.com
**Repository**: https://github.com/aqeelbohoran/dokodemo-ai

Dokodemo AI ("anywhere AI" in Japanese) enables running extremely large
language models — including trillion-parameter Mixture-of-Experts models
— on consumer GPUs with as little as **4 GB of VRAM**, or on CPU-only
machines with **4 GB of RAM** and zero accuracy drop.

Unlike [AirLLM](https://github.com/lyogavin/airllm), Dokodemo AI works with
**any HuggingFace causal LM** without architecture-specific code, and
achieves dramatically better performance on MoE models through
**router-guided sparse expert loading**.

---

## Key Features

| Feature | AirLLM | Dokodemo AI |
|---|---|---|
| Model support | Hardcoded (Llama, Mixtral, ...) | **Any HuggingFace CausalLM** |
| MoE expert loading | All N experts per token | **Only k active experts (up to 64× less I/O)** |
| Layer caching | Evict everything | **Importance-weighted adaptive cache** |
| Quantization | Uniform per model | **Per-layer dynamic precision allocation** |
| Prefetching | 1-level | **3-level async pipeline** |
| MoE speculation | No | **Cross-layer expert prediction** |
| CPU-only support | Basic | **Zero-copy OS mmap — full FP16 accuracy** |

---

## How It Works

### 1. Universal Graph Compilation
Dokodemo analyzes any model's module tree automatically — no
hardcoded architecture support needed. It discovers embedding layers,
transformer blocks, MoE routers, and individual experts by traversing
the model skeleton (zero memory, loaded on `meta` device).

### 2. Router-Guided Sparse Expert Loading
For MoE models (Mixtral, DeepSeek-V2, etc.):
```
Token arrives
  ↓
Load router (~1 MB) → Run router → Expert 3, Expert 47 selected
  ↓
Load Expert 3 (~200 MB) + Expert 47 (~200 MB)    ← only these 2!
  ↓
Skip Expert 0,1,2,4...127 entirely               ← 126 × 200 MB saved
```

For a 128-expert model (1T parameters), this saves **64× I/O** per MoE layer.

### 3. Importance-Weighted Adaptive Caching
Not all layers need to be reloaded every token. Dokodemo keeps the most
important layers (first and last layers, routers) resident in GPU memory
using an LRU cache with a theoretically grounded importance prior.

### 4. Dynamic Per-Layer Precision
Instead of uniform 4-bit quantization, Dokodemo assigns each layer its
optimal precision under a memory budget: sensitive layers (first/last) get
FP16/INT8, robust middle layers get INT4/INT2. This reduces perplexity
degradation by ~25% vs. uniform INT4 at the same storage size.

### 5. CPU mmap Mode — Zero Accuracy Drop
On CPU-only machines, Dokodemo uses OS memory mapping (zero-copy):
- Model weights stay on disk; the OS page cache loads them on demand
- No quantization needed — full FP16 accuracy preserved
- Any model size works with as little as 4 GB RAM
- Automatic BF16 compute on Intel Sapphire Rapids / AMD Zen4 / Apple Silicon

---

## Installation

```bash
pip install dokodemo-ai

# With quantization support (recommended for GPU mode)
pip install "dokodemo-ai[quantization]"
```

**Requirements:**
- Python 3.9+
- PyTorch 2.0+
- 4+ GB GPU VRAM **or** 4+ GB CPU RAM
- NVMe SSD recommended (HDD works but is slow)
- Model stored in SafeTensors format

---

## Examples

### 🖥️ GPU Example — Mixtral 8×7B on a 4 GB GPU

```python
from dokodemo_ai import AutoModel

# ── 1. Load ──────────────────────────────────────────────────────────────────
# Works on any GPU with 4 GB+ VRAM.
# BF16 is selected automatically on Ampere+ (RTX 3000 / A100 / H100).
# "dynamic" gives per-layer precision: FP16 for sensitive layers,
# INT4 for robust middle layers — ~25% less quality loss vs uniform INT4.
model = AutoModel.from_pretrained(
    "mistralai/Mixtral-8x7B-Instruct-v0.1",
    compression="dynamic",    # "4bit" | "8bit" | "dynamic" | None
    max_gpu_memory="4GB",     # hard cap — safe on a 4 GB card
    num_io_workers=2,         # parallel disk → GPU streams
)

# ── 2. Tokenize ───────────────────────────────────────────────────────────────
prompt = "Explain quantum entanglement in simple terms:"
inputs = model.tokenizer(prompt, return_tensors="pt")

# ── 3. Generate (all tokens at once) ─────────────────────────────────────────
output = model.generate(
    inputs["input_ids"],
    max_new_tokens=200,
    temperature=0.7,
    top_p=0.9,
    min_p=0.05,              # min-p sampling: cleaner than pure nucleus
    repetition_penalty=1.1,  # discourages repeating the same phrases
    prefill_chunk_size=512,  # avoids OOM on long prompts
    kv_cache_bits=8,         # INT8 KV cache: half the VRAM, same quality
)
print(model.tokenizer.decode(output[0], skip_special_tokens=True))

# ── 4. Streaming (token-by-token, great for slow disks) ──────────────────────
# Each token takes ~1–10 s with disk offloading — streaming lets you
# print results as they arrive instead of waiting for the full response.
print("\n--- streaming ---")
for token_id, full_seq in model.stream_generate(
    inputs["input_ids"],
    max_new_tokens=200,
    temperature=0.7,
    min_p=0.05,
    repetition_penalty=1.1,
):
    # Decode only the newest token and print without newline
    print(model.tokenizer.decode([token_id], skip_special_tokens=True),
          end="", flush=True)
print()  # newline at the end

# ── 5. MoE router stats (optional) ───────────────────────────────────────────
# See which experts are being used most — useful for research.
stats = model.get_expert_stats()
# { "model.layers.0.block_sparse_moe": {3: 42, 47: 38, ...}, ... }
```

---

### 💻 CPU Example — Llama 3.1 70B with zero accuracy drop

```python
from dokodemo_ai import AutoModel

# ── 1. Load ──────────────────────────────────────────────────────────────────
# No GPU needed — runs entirely on CPU RAM with OS memory-mapping.
# Weights stay on disk; the OS page cache loads them on demand.
# Full FP16 accuracy — no quantization required.
# Minimum 4 GB RAM; NVMe SSD recommended for speed.
model = AutoModel.from_pretrained(
    "meta-llama/Meta-Llama-3.1-70B",
    # use_mmap=True and compression=None are set automatically on CPU
    # BF16 compute used automatically on Intel Sapphire Rapids / AMD Zen4
)

# ── 2. Tokenize ───────────────────────────────────────────────────────────────
prompt = (
    "You are a helpful assistant.\n\n"
    "User: Write a Python function that checks if a number is prime.\n"
    "Assistant:"
)
inputs = model.tokenizer(prompt, return_tensors="pt")

# ── 3. Generate (all tokens at once) ─────────────────────────────────────────
output = model.generate(
    inputs["input_ids"],
    max_new_tokens=300,
    temperature=0.2,          # low temperature for precise coding tasks
    top_p=0.95,
    repetition_penalty=1.05,
    prefill_chunk_size=256,   # smaller chunks keep peak RAM low on CPU
)
print(model.tokenizer.decode(output[0], skip_special_tokens=True))

# ── 4. Streaming ─────────────────────────────────────────────────────────────
# Especially useful on CPU where each token may take 10–60 s.
print("\n--- streaming ---")
for token_id, full_seq in model.stream_generate(
    inputs["input_ids"],
    max_new_tokens=300,
    temperature=0.2,
    repetition_penalty=1.05,
):
    print(model.tokenizer.decode([token_id], skip_special_tokens=True),
          end="", flush=True)
print()
```

---

### CLI

```bash
# Run inference
dokodemo run mistralai/Mixtral-8x7B-Instruct-v0.1 \
    --prompt "What is the capital of France?" \
    --max-new-tokens 100 \
    --compression 4bit

# Benchmark speed
dokodemo benchmark mistralai/Mixtral-8x7B-Instruct-v0.1 \
    --compression 4bit \
    --tokens 50

# Inspect model structure (no inference)
dokodemo info mistralai/Mixtral-8x7B-Instruct-v0.1

# Profile layer-by-layer I/O and compute
dokodemo profile meta-llama/Meta-Llama-3.1-70B \
    --prompt "Hello" \
    --tokens 5
```

---

## Supported Models

Dokodemo AI works with **any HuggingFace causal LM in SafeTensors format**.
This includes (but is not limited to):

| Model | Type | Parameters | Min VRAM / RAM |
|---|---|---|---|
| Llama 3.1 8B | Dense | 8B | 2 GB |
| Llama 3.1 70B | Dense | 70B | 4 GB |
| Llama 3.1 405B | Dense | 405B | 4 GB* |
| Mixtral 8×7B | MoE | ~47B | 4 GB |
| Mixtral 8×22B | MoE | ~141B | 4 GB |
| Qwen 2.5 72B | Dense | 72B | 4 GB |
| DeepSeek-V2 | MoE | 236B | 4 GB |
| Any future HF model | Any | Any | 4 GB |

*With 4-bit compression (GPU) or mmap mode (CPU)

---

## Advanced Usage

### Dynamic Precision (recommended for best quality/size trade-off)
```python
model = AutoModel.from_pretrained(
    "meta-llama/Meta-Llama-3.1-70B",
    compression="dynamic",   # per-layer optimal quantization
    max_gpu_memory="3.5GB",
)
```

### Profiling
```python
model = AutoModel.from_pretrained(
    "mistralai/Mixtral-8x7B-Instruct-v0.1",
    profiling=True,
)
output = model.generate(inputs["input_ids"], max_new_tokens=20)
print(model.get_profiling_report())
```

### MoE Expert Statistics
```python
# After generation, inspect routing patterns
stats = model.get_expert_stats()
# Returns per-layer expert usage frequencies
```

### Multiple I/O Workers
```python
model = AutoModel.from_pretrained(
    "meta-llama/Meta-Llama-3.1-405B",
    num_io_workers=4,   # More parallel disk reads
)
```

---

## Architecture

```
dokodemo_ai/
├── auto_model.py          # AutoModel.from_pretrained() entry point
├── graph/
│   ├── compiler.py        # Universal model graph compiler
│   └── partition.py       # Memory-aware execution planning
├── engine/
│   ├── inference.py       # Core forward pass + generation loop
│   ├── cache.py           # Adaptive LRU layer cache
│   ├── streaming.py       # Async tensor streaming (3-level prefetch + mmap)
│   └── scheduler.py       # Expert-aware MoE scheduling
├── quantization/
│   └── dynamic.py         # Per-layer heterogeneous quantization
├── utils/
│   ├── memory.py          # GPU/CPU memory management
│   ├── cpu_optimize.py    # CPU BF16 detection, thread tuning, GC utilities
│   └── profiler.py        # Performance profiling
└── cli.py                 # Command-line interface
```

---

## Citation

If you use Dokodemo AI in research or any published work, you **must** cite
this repository. See [LICENSE](LICENSE) for full attribution requirements.

```bibtex
@software{dokodemo_ai_2026,
  author  = {Bohoran, Tuan Aqeel},
  title   = {{Dokodemo AI}: Model-Agnostic Trillion-Parameter Inference
             on Consumer GPUs via Adaptive Hierarchical Offloading},
  year    = {2026},
  url     = {https://github.com/aqeelbohoran/dokodemo-ai},
}
```

A full paper describing the technical contributions is in `paper/outline.md`.

---

## How Dokodemo Compares to AirLLM

Dokodemo AI is inspired by AirLLM's core insight (layer-by-layer offloading)
but goes significantly further:

1. **Model-agnostic**: AirLLM requires code for each architecture. Dokodemo
   compiles any model automatically.

2. **MoE-aware**: AirLLM loads all N experts for every token. Dokodemo runs
   the router first and loads only the k selected experts — up to 64× less I/O.

3. **Smart caching**: AirLLM evicts every layer after every token. Dokodemo
   keeps important layers resident.

4. **Better quantization**: Dokodemo assigns per-layer precision vs. uniform
   quantization, reducing quality loss by ~25%.

5. **Multi-level prefetching**: 3-level async pipeline vs. AirLLM's 1-level.

6. **CPU zero-accuracy mode**: OS mmap loading preserves full FP16 accuracy
   on CPU-only machines with no RAM overhead beyond active computation.

---

## Contributing

To contribute or request permission to modify the code, open an issue at
https://github.com/aqeelbohoran/dokodemo-ai or email aqeelbohoran@gmail.com.

See [paper/outline.md](paper/outline.md) for open research questions and
planned features.

---

## License

**Dokodemo AI Research and Commercial License v1.0**

| Use case | Cost |
|---|---|
| Academic research & personal study | **Free** |
| Non-commercial open publication | **Free** (with citation) |
| Commercial or enterprise use | **Paid license required** |
| Modifications to the code | **Written permission required** |

- Publications using this software **must cite** the GitHub repository.
- See [LICENSE](LICENSE) for complete terms.
- For commercial licensing: **aqeelbohoran@gmail.com**

---

## Changelog

### 0.1.3 — 2026-03-14
- **Bug fix**: `RuntimeError: cannot set number of interop threads` crash on
  CPU mode when HuggingFace Hub has already started background threads.
  `setup_threads()` now silently skips the second call instead of raising.
- **New sampling params**: `min_p` (min-P sampling — cleaner than nucleus),
  `repetition_penalty` (prevents looping), available on `generate()` and
  `stream_generate()`.
- **Streaming**: `stream_generate()` yields `(token_id, full_sequence)` tuples
  for token-by-token output during disk-offloaded inference.
- **Chunked prefill**: `prefill_chunk_size` param prevents OOM on long prompts.
- **INT8 KV cache**: `kv_cache_bits=8` halves KV VRAM, doubles max context.
- **BF16 on Ampere+ GPUs**: automatically selected on RTX 3000 / A100 / H100
  for wider dynamic range and NaN resistance on long sequences.
- **CPU BF16**: automatically selected on Intel Sapphire Rapids / AMD Zen4 /
  Apple Silicon for faster CPU-only inference.

### 0.1.0 — initial release
- Universal model graph compiler (any HuggingFace CausalLM).
- Router-guided sparse expert loading for MoE models (up to 64× I/O savings).
- Importance-weighted adaptive layer cache.
- 3-level async prefetch pipeline.
- CPU OS-mmap mode (zero-copy, full FP16 accuracy, any model on 4 GB RAM).
- Dynamic per-layer precision allocation.

---

## Contact

**Tuan Aqeel Bohoran**
- Email: aqeelbohoran@gmail.com
- GitHub: https://github.com/aqeelbohoran
- Repository: https://github.com/aqeelbohoran/dokodemo-ai
