AXIOM-1 (EGen-Core/AXIOM-1) is a radically different approach to Large Language Modeling. Built entirely from scratch by EGen Labs, it abandons the traditional feed-forward paradigm of stacking 32+ specialized layers. Instead, AXIOM-1 utilizes Recurrent Depth (LoopLLM): learning a single, highly-generalized block and passing thoughts through it iteratively to refine understanding.

This repository serves as the definitive reference implementation and training grounds for Phase 0 of the EGen Core architecture.

⚡ Core Architecture (Upgraded V1.1)

AXIOM-1 has recently undergone a major architectural upgrade for extreme memory efficiency, stability, and multilingual scale:

• Universal LoopBlock: 1 shared Transformer block applied autoregressively 12 times per forward pass.
• Scale: d_model=512, heads=8q/2kv (GQA), d_ff=1365.
• Multilingual Expansion: Tokenizer vocabulary expanded from 32K to 132,000, enabling seamless multi-language processing (English, French, Arabic, coding languages).
• Loop Stability Interventions: Features mid-loop re-normalization (forces strict magnitude bounds every 4 iterations) and dynamic consistency loss warmup to prevent loop collapse.
• AirLLM Integrations: Engineered to run comfortably on sub-4GB consumer hardware via component-level layer streaming, dynamic INT8 FFN quantization, and memory-tracked per-step KV caches.

📦 Installation (PyPI Ready)

You can install AXIOM-1 as a python package directly from the source.

# Clone the Core repository
git clone https://github.com/EGen-V/EGen-Core.git
cd EGen-Core

# Install the axiom1 package and dependencies
pip install .

Once officially published, AXIOM-1 will be installable via pip install axiom1.

🔬 Running the Google Colab Prototype

We actively support barrier-free testing of the recurrent architecture. Use our ready-to-use Google Colab notebook: axiom1_colab_training.ipynb

Features natively handled by the notebook:

1. Zero-Setup Environment: Installs axiom1 and imports the 132K dataset builder.
2. Validation: Automatically runs the strict 6-Rule Sanity Check (gradient flow tracking).
3. End-to-End Pipeline: Trains the prototype using PyTorch's native loops.
4. Export: Automatically saves and pushes the weight-tied checkpoint to your Hugging Face Hub.

Quick Start in Colab:

1. Upload axiom1_colab_training.ipynb to Google Colab.
2. Set Runtime > Change runtime type > T4 GPU (or better).
3. Run all blocks (ensure you have a Hugging Face write token ready for the final cell).

🛠️ Manual Usage / API Guide

1. The 6-Rule Sanity Check (Mandatory)

Because recurrent networks are uniquely vulnerable to gradient explosion, you must run the 6-rule sanity check before any training session.

from axiom1.model.config import AXIOM1Config
from axiom1.model.model import AXIOM1
from axiom1.evals.sanity_check import run_sanity_check

cfg = AXIOM1Config()
model = AXIOM1(cfg).to('cuda')

# Validates params, initial loss (11.79), LSE norms, and verifies
# that recurrent gradients successfully flow through all 12 iterations.
run_sanity_check(model, cfg) 

2. High-Efficiency Generation (AirLLM Style)

AXIOM-1 is built for edge hardware. Use AXIOM1_Offloaded to stream the single transformer block continuously between CPU and GPU.

import torch
from axiom1.model.config import AXIOM1Config
from axiom1.model.model import AXIOM1
from axiom1.inference.offload import AXIOM1_Offloaded
from axiom1.tokenizer.tokenizer import AXIOM1Tokenizer

# 1. Load your trained model to CPU
cfg = AXIOM1Config()
model = AXIOM1(cfg)
# model.load_state_dict(...)
model.eval()

# 2. Wrap via AirLLM offload wrapper
offloaded_model = AXIOM1_Offloaded(model, device='cuda')

# 3. Generate with top-k, top-p, temperature
prompt_ids = torch.tensor([[1, 234, 555]])
generated = offloaded_model.generate(
    prompt_ids, 
    cfg, 
    max_new_tokens=256,
    top_k=50,
    top_p=0.9
)

3. VRAM Budget Estimation

AXIOM-1 scales infinitely depending on batch size and context window. Check your exact memory footprint before testing:

from axiom1.inference.memory import print_memory_report

# Will output a highly detailed console report verifying if
# your intended configuration fits within 2GB / 4GB / 8GB VRAM pools.
print_memory_report(cfg, batch_size=4, seq_len=1024)

4. Full Documentation

For a complete breakdown of the 3-Phase training curriculum, sanity check diagnostics, and INT8 quantization features, please see the AXIOM-1 Comprehensive Usage Guide.