🧠 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)

Have you noticed how the model parameters are changing? You’re right: I am slightly adjusting the dimensions—but which ones? The MLP parameters. The issue with MLPs is that a fourfold expansion increases model capacity but consumes a huge amount of memory and slows down training. But what if we took a different approach? What if we used a threefold MLP expansion instead? We lose about 25% of the model's capacity, but the model runs faster, allowing us to increase both the embedding dimension and the number of layers. Consequently, the 2-billion-parameter XXL model becomes an experiment: it incorporates RoPE to extend the future context window and reduces the MLP size from 8192 to 6144.

⚠️ 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