A from-scratch multimodal AI stack (text + image + speech in/out) trainable on a single GPU. Uses a Thinker-Talker architecture inspired by recent multimodal LLMs.

Image ──→ ViT Encoder ──→ Projector ──┐
Audio ──→ Audio Encoder ─→ Projector ──┤
Text  ──→ Token Embeddings ───────────┤
                                       ├──→ Thinker (LLM) ──→ Text Output
                                       └──→ Talker ──→ RVQ ──→ Vocoder ──→ Speech

~13.9M parameters (synthetic config) | 16GB VRAM (RTX 5070 Ti) | Reference learning repo — compact and readable

Training Dashboard — live metrics, pipeline management, and GPU monitoring

Hyperparameter Tuning — search space configuration (left) and parameter slice analysis (right)

Architecture (synthetic config):

Training time (RTX 5070 Ti Laptop GPU, synthetic 2000 samples):

Note: B+C can run in parallel. First-time setup with config tuning may take 2+ hours.

# 1. Install
pip install -r requirements.txt

# 2. Generate synthetic data (for testing)
python -m scripts.generate_synthetic_data

# 3. Train (any order for A/B/C, then D needs A, E needs all)
python -m train.train_thinker --config configs/synthetic_thinker.json    # Stage A: Thinker LLM
python -m train.train_audio_enc --config configs/synthetic_audio_enc.json # Stage B: Audio Encoder
python -m train.train_vision --config configs/synthetic_vision.json       # Stage C: Vision Encoder
python -m train.train_talker --config configs/synthetic_talker.json       # Stage D: Talker + RVQ
python -m train.sft_omni --config configs/synthetic_omni_sft.json         # Stage E: Multimodal SFT

# 4. Inference
python -m test.infer_chat --ckpt_dir checkpoints/thinker_tiny                                    # Text chat
python -m test.infer_chat --ckpt_dir checkpoints/omni_sft_tiny --image photo.jpg "describe this"  # Image QA
python -m test.infer_chat --ckpt_dir checkpoints/omni_sft_tiny --audio_in speech.wav              # Audio transcription
python -m test.infer_chat --ckpt_dir checkpoints/omni_sft_tiny --image doc.jpg --ocr              # OCR

Train/test scripts (except test/infer_chat.py) now emit structured records to logs/metrics/*.jsonl.

• Required fields: timestamp, script, phase, run_id, epoch, batch, step, split, metric_name, metric_value
• Optional fields: lr, loss, checkpoint, device, extra
• Resume safety: duplicate keys are upserted by (run_id, phase, epoch, batch, step, split, metric_name)

Use the viewer:

start scripts/metrics_viewer.html

Then choose one or more files from logs/metrics/ and filter by script/run/phase/metric.

Configs in configs/ — one per training stage (synthetic_*.json). These use smaller vocab/steps for quick iteration.

Key settings across all configs:

python -m scripts.generate_synthetic_data

python -m scripts.download_production_text --combine
python -m scripts.download_production_audio --combine
python -m scripts.download_production_image --combine
python -m scripts.download_production_ocr --combine

Data formats:

• Text: Plain .txt, one sample per line
• ASR Audio: CSV with wav,text columns
• TTS Audio: CSV with text,wav columns (reversed!)
• Images: JSON manifest with image + caption fields
• OCR: CSV with image,text columns

• Multi-Token Prediction (MTP) — predict t+2, t+3 during training for richer gradients
• Sliding Window Attention — O(n*w) complexity for alternating layers, enabling longer sequences
• YaRN RoPE — context extension beyond training length via rope_scaling_factor
• Cached RoPE cos/sin tables (not recomputed every forward pass)
• Pre-allocated causal masks (sliced, not recreated)
• Zero-copy GQA expansion (expand+reshape instead of repeat_interleave)
• Sorted MoE dispatch (batched, not nested Python loops)
• Efficient RVQ encoding (torch.cdist)
• No forward-pass NaN checks (removed GPU-syncing isnan/isinf)

• Mixed precision (AMP float16/bfloat16)
• Label smoothing (0.1) for better calibration and reduced overconfidence
• zero_grad(set_to_none=True) (frees gradient memory)
• pin_memory=True on DataLoaders
• cudnn.benchmark=True (auto-tuned convolutions)
• Flash Attention (PyTorch 2.0+ scaled_dot_product_attention)
• TF32 matmul enabled globally
• Single gradient norm check (clip + threshold in one pass)
• Streaming IterableDataset (90%+ RAM reduction)

• ThinkerLM.generate() with temperature, top-k, top-p (nucleus), and repetition penalty
• Repetition penalty for preventing degenerate loops and improving output quality
• Top-k / top-p (nucleus) sampling for controlled diversity during text generation

# Merge all trained components into HuggingFace-compatible format
python -m scripts.export \
  --thinker_ckpt checkpoints/thinker_tiny \
  --audio_ckpt checkpoints/audio_enc_tiny \
  --vision_ckpt checkpoints/vision_tiny \
  --talker_ckpt checkpoints/talker_tiny \
  --ocr_ckpt checkpoints/ocr_tiny \
  --output_dir export

This creates:

• export/model.safetensors — HF-compatible (flat thinker keys, works with from_pretrained)
• export/model_full.safetensors — Full multimodal model (all components with prefixed keys)
• export/config.json — HuggingFace config with auto_map for custom model discovery
• export/modeling_muomni.py — Self-contained model class (no omni/ dependency)

import sys; sys.path.insert(0, "export")
from modeling_muomni import MuOmniForCausalLM
from omni.tokenizer import BPETokenizer
import torch

# Load model (works exactly like any HF model)
model = MuOmniForCausalLM.from_pretrained("./export", trust_remote_code=True)
model = model.cuda().eval()

# Tokenize and generate
tok = BPETokenizer("export/tokenizer.model")
ids = [1] + tok.encode("The red cat")
x = torch.tensor([ids], device="cuda")
out = model.generate(x, max_new_tokens=20, temperature=0.7, top_k=40)
print(tok.decode(out[0].tolist()))
# Output: "The red cat sits in the park."

from modeling_muomni import MuOmniMultimodalModel
model = MuOmniMultimodalModel.from_pretrained_safetensors("./export")
out = model(input_ids=text_ids, pixel_values=image, mel_spectrogram=audio)

huggingface-cli login
huggingface-cli upload prskid1000/micro-Omni export/

# Single component
python -m test.test_thinker --checkpoint checkpoints/thinker_tiny

# Multimodal SFT test
python -m test.test_sft --checkpoint checkpoints/omni_sft_tiny --config configs/synthetic_omni_sft.json

# All tests (PowerShell)
Get-ChildItem -Filter 'test_*.py' -Recurse | ForEach-Object { python $_.FullName }

# Export validation
python -m export.test_safetensor

# HF model tests
python -m export.test_hf_text           # HF text model scored test
python -m export.test_hf_multimodal     # HF multimodal scored test

Flags: --device cpu (no GPU), --num_samples N (limit test size), --config path/to/config.json

See study/ for a complete zero-to-master tutorial (25 chapters + 5 appendices). Covers everything from "What is AI?" to deployment, with real-life analogies and ASCII diagrams.

MIT. Replace datasets with those compatible with your needs.