🧠 Vexion-GPT: Custom Dense LLM Architecture

Vexion-GPT is a classic dense language model engine built from scratch on PyTorch. The project was created to deeply understand the architecture of Transformers, optimize memory, and pretraining processes without the use of heavy third-party frameworks.

⚙️ Key Features of the Engine (Under the Hood)

Unlike many "training" models, Vexion-GPT is designed for real-world big data work and maximum GPU utilization:

• Pure Dense Architecture: Classic, mathematically pure GPT architecture without additives (no MoE, no RoPE). Only the proven Causal Attention and GELU/SiLU activation functions.
• Flash Attention Integrated: Full support for Fused Kernels for on-the-fly attention computation. VRAM consumption has been dramatically reduced (the model can easily train on consumer GPUs with batch sizes that previously caused OOMs).
• Ultra-fast Custom DataLoader: The dataloader has been rewritten to stream binary data (.bin), bypassing Python's garbage collector and Windows system caching. The token feed rate is static and does not degrade over long distances.
• HF-Compatible Config: The architecture is completely decoupled from hardcoded data. Model configuration is implemented via config.json according to Hugging Face standards (full support for hidden_size, num_hidden_layers, etc.).
• Adaw: Standard Adaw allows the model to store all its values ​​in 32-bit format, which helps achieve lower LOSS and higher-quality models.

📊 Vexion Model Family

The project's development is divided into several stages, from compact test versions to fully-fledged billion-dollar vehicles.

📊 Model training dynamics

Technical details:

• Hardware: 1x RTX 3060 Ti 8GB or later on 1x RTX 4060 Ti 16GB.
• Settings: batch_size - 12, accumulate_steps - 8 at ctx 256. Batch_size - 8, accumulate_steps - 32 at ctx 1024.
• Number of tokens in pre-train - 39 BILLION
• Total training time: 570 hours
• Pre-train dataset: CulturaX

📊 News (temporarily)

Well, we can expect some fairly good news, and it really deserves attention:

1. The dialogue dataset with reflection is about 20% complete; currently, about 400 lines of dialogues have been written, and the task is to increase it to 1,000–8,000 dialogues, after which the models will be trained on it. It will also retain its power and reasoning capabilities.
2. The models are being switched to RoPE. No matter what I might say, the transformer standard will remain — it will, but the entire family is moving from WPE to RoPE. I realized that it makes more sense to work with RoPE, as the same loss metrics are much better. Plus, in the future, I can expand the context in the models, which will allow them to retain moments with you for a longer time.
3. Context — since the NANO and Medium models will need to be retrained for the new context, switching from WPE to RoPE, this will take time. Unfortunately, XL will be released later, but with a good update, because the models will receive the following context: NANO — 3072 tokens, Medium — within 4096–6144 tokens, XL — within 4096–8192 tokens, 2B — it’s still unknown, but I hope to reach 2048 tokens. The issue is that it’s hard to give exact figures, since I’m using additional code that forces the weights to be trained in FP32 — there’s an increase in memory usage, albeit within the range of 2–3 GB, but it makes you think about which context is better to use. But one thing is clear — all models with a context size below 2048 WILL NOT BE TRAINED (except for 2B; there will be a separate solution for it)!

⚠️ Important compatibility warning

This model is built on a completely custom architecture written in pure PyTorch. It DOES NOT support the Hugging Face transformers library (Transformer API). You will not be able to load it through standard classes like AutoModelForCausalLM. For inference and retraining, use only the scripts provided in this repository (model.py, generate.py, generation.py).

🚀 Code Usage

Running/Retraining the Model:

• Before training the model, launch the command prompt (CMD) as administrator and enter the following command: cd C:\Users\Username\Desktop\model folder

Training: python train.py --data_path train.bin --val_path val.bin --total_steps 30000 --save_every 1000 --batch_size 1 --accumulate_steps 256 --lr 1e-4 --dropout 0.0 --warmup_steps 500

Retraining: python train.py --data_path train.bin --val_path val.bin --total_steps 30000 --save_every 1000 --batch_size 1 --accumulate_steps 256 --lr 1e-4 --dropout 0.0 --warmup_steps 500 --resume checkpoints/gpt_step_1000.safetensors

Communication: python generate_base.py --checkpoint checkpoints/gpt_step_180000.safetensors --prompt "What's dangerous about artificial intelligence?" --temperature 0.7 --rep_penalty 1.2 --max_new_tokens 400 --device cuda