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

Currently, the repository contains the architecture source code and base training graphs. Model weights will be published after completing extensive pretraining runs on large datasets.

βš™οΈ Key Features of the Engine (Under the Hood)

Unlike many "training" models, Vexion-GPT is designed for real-world big data 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 be easily trained 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 code. Model configuration is implemented via config.json according to Hugging Face standards (full support for hidden_size, num_hidden_layers, etc.).

πŸ“‰ System Requirements and VRAM Consumption (Pre-training)

Because Vexion-gpt is a classic, dense transformer without overloaded modern add-ons, it is incredibly resource-efficient. Flash Attention integration and proper memory management allow model training from scratch even on budget home graphics cards.

The main secret to its low power consumption lies in its support for gradient accumulation. Training with batch_size = 1 and accumulate_steps = 256 results in a huge effective batch size (256), but the physical memory consumption remains at the level of a single pass. Gradient accumulation consumes computational time but does not inflate VRAM.

Memory Measurements (with a 1024-token context) Below are actual video memory consumption measurements for training from scratch (including weights, gradients, AdamW optimizer states, and PyTorch cache):

Nano / 117M parameters (Hidden size: 768, Layers: 12, Heads: 12)

Consumption: 2 GB of VRAM

Suitable for: RTX 3050 / RTX 4050 Laptop / Any NVIDIA RTX graphics card with 4 GB or more.

Medium / 345M parameters (Hidden size: 1024, Layers: 24, Heads: 16)

Consumption: 4.5 GB of VRAM

Suitable for: RTX 3060 Ti / 4060 (8 GB) / Any NVIDIA RTX graphics card with 6 GB or more.

Large / 646M parameters (Hidden size: 1280, Layers: 30, Heads: 20)

Consumption: 7.5 GB of VRAM

Suitable for: Cards with at least 8 GB (RTX 3060 Ti 8GB / 4070). Better with 10 GB or more.

XL / 1B parameters (Hidden size: 1536, Layers: 36, Heads: 12)

Consumption: 11.5 GB VRAM

Suitable for: Cards with at least 12 GB (RTX 3060 12GB / 3080). Better with 14-16 GB.

XXL / 2.5B parameters (Hidden size: 2048, Layers: 48, Heads: 16)

Consumption: 28 GB VRAM

Suitable for: Cards with 32 GB (TESLA V100 32GB / 5090).

πŸ’‘ Note: Power consumption is given for mixed precision (bfloat16 / float16). Monitor output typically consumes around 200-400 MB of your GPU, so always leave some headroom. Better yet, switch to CPU output if possible to free up all your VRAM.

Now, let's run tests with the Nano model, but only increasing the context. We'll start at 1024 and work our way up. How much VRAM do we need to, say, train it on a larger context than the initial one?

Model: Vexion-gpt-NANO

Context Minimum VRAM better on VRAM
1024 2GB 4GB
2048 3GB 6GB
3072 4GB 6GB
4096 5GB 8GB
5120 6.5GB 8GB
6144 7.5GB 10GB
7168 8.5GB 10GB
8192 9.5GB 11GB
9216 9.5GB 11GB
10240 10.5GB 12GB
11264 11.5GB 14GB
12288 12.5GB 14GB
13312 13.5GB 16GB
14336 14.5GB 16GB

⚠️ In this case, we're ONLY looking at the context consumption itself. We've intentionally used a small, but at least somewhat functional, model with batch 1 and a ton of accumulation. Training this way is obviously not possible, as it would take 2000 hours or more. I'm showing the MINIMUM threshold for video cards so you can understand how much "clean" VRAM the model needs. I also admit that context consumption may vary on larger models. But let's check, for example, on the Large model:

Context Minimum VRAM better on VRAM
1024 7.5GB 10GB
2048 9GB 11GB
3072 11.5GB 14GB
4096 13.5GB 16GB

XL Model:

ctx 1024: 11.5 GB VRAM. Entry threshold: 14 GB VRAM.

ctx 2048: 14 GB VRAM. Entry threshold: 16 GB VRAM.

You can use this table to understand what context you can run on your hardware to avoid suffering from a small context. BUT!!!! For at least proper training, I RECOMMEND a VRAM reserve of 4-6 GB. This is necessary to allow for a larger batch size. This prevents the machine from running at an excessively low speed and wasting energy on a six-month model training period. For example, if you have 8 GB VRAM, you can run a context of 3072 tokens on the NANO model. This initial size is more than enough, but it only uses half the memory, leaving the other half for a larger batch.

πŸ“Š Current training status (Phase 1: Wikipedia)

The attached graphs show the first stage of pre-training the base model.

  • Current build parameters: 117M parameters (Hidden size: 768, Layers: 12, Heads: 12, ctx: 256).
  • Dataset: Clean corpus of the Russian-language Wikipedia (~331 MILLION tokens).
  • Goal of this stage: To develop basic syntax, ideal grammar, and punctuation for the Russian language before moving to "dirty" datasets (CulturaX).

Π‘Π½ΠΈΠΌΠΎΠΊ экрана (427)

  • The graph shows a steady drop in loss without spikes, confirming the absolute mathematical stability of the custom transformer and gradient mechanism. At control checkpoints, the model successfully generates grammatically correct encyclopedic text.

πŸš€ Code Usage

The architecture is ready for experimentation. You can easily scale the model from Nano to 1B+ parameters by simply changing config.json.

Example of initializing a pure model graph:

import torch
from model import GPT, GPTConfig

# Loading a standardized HF config
config = GPTConfig.from_json("config.json")

# Initializing a pure model with bfloat16 support
device = "cuda" if torch.cuda.is_available() else "cpu"
model = GPT(config).to(device)

print(f"Vexion-GPT model initialized. Parameters: {sum(p.numel() for p in model.parameters()) / 1e6:.2f} M")

Running model training/retraining:

  • 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

Instructions: In the generate_base.py file, start changing parameters starting at line 73. In the CHECKPOINT_PATH = line, specify the path to the checkpoint. From lines 78-82, change the model settings. Don't touch anything ABOVE.

To disable the scheduler for custom LR speed tuning, add the command --scheduler_type constant to the other settings in the CMD. To re-enable the scheduler, remove the command.

Downloads last month
976
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support