Vexion-LM 🧠

Vexion-LM is a family of open-source language models built from scratch. The project includes a complete pipeline: from writing a custom architecture in PyTorch to pretraining base models and then retraining them for the conversational format using LoRA adapters.

🚀 Model Features

• Custom Architecture: A completely independent and optimized transformer with gradient checkpointing support for drastic VRAM savings.
• Efficient Training: Uses an 8-bit optimizer (AdamW8bit) from the bitsandbytes library, allowing models to be trained locally on consumer GPUs.
• Flexible LoRA System: Built-in mechanism for freezing base weights and training compact adapters for the conversational format without the risk of catastrophic forgetting.
• Smart Generation: Custom inference script with support for repetition penalties, temperature management, and early stopping for special tokens.
• Efficient RoPE: Hand-written RoPE - allows the model to better understand the location of each token.
• Smart Sliding Window: Custom attention mask with an "anchor" system. The router permanently keeps the system prompt (e.g., the first 64 tokens) and the local dialog window in memory, allowing it to digest gigantic contexts (2048+ tokens) with minimal VRAM consumption.
• Dynamic Experts (VexionMoE): Innovative Mixture of Experts system with a "genetic mutation" mechanism. The architecture automatically tracks "dead" (unused) experts, kills them, and clones the knowledge of the best neurons with added noise for continuous evolution during training. * GQA (Grouped-Query Attention): Radically reduces video memory consumption when generating long contexts (up to 2048+ tokens), making VRAM growth linear rather than quadratic.
• SwiGLU & VexNorm (L1 Absolute Norm): Abandoning the cumbersome industry standard (RMSNorm) in favor of our own ultra-fast L1 normalization. Without tensor squaring or root extraction, it perfectly protects gradients from outliers and is mathematically optimized for quantization.
• Custom BPE Tokenizer: A dictionary of exactly 40,960 tokens—an ideal multiple of 64, ensuring 100% efficiency of the GPU's Tensor Cores.
• Differential Attention: Abandoning classic Causal Attention in favor of an architecture with active noise reduction. The model computes two parallel attention maps (Signal and Noise) and subtracts junk information, leaving only crystal-clear reasoning.

Create_Vexion-LM - allows you to train this model from scratch or further train it. All files, from generate.py to train.py, are available here. Create your own models using Custom_architecture!

⚠️ 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 cannot load it through standard classes like 'AutoModelForCausalLM'. For inference and retraining, use exclusively the scripts provided in this repository ('model.py', 'generate.py', 'generation.py').

⚙️ Weight Format and Precision

• File format: '.safetensors' (safe and fast loading format).
• Base parameter precision: FP32.
• When running generation via generate.py, the code automatically uses automatic mixed-precision (AMP via torch.amp.autocast), switching to FP16 or BF16 on supported GPUs to save video memory and speed up execution. Manual weight conversion is not required.

🛠 How to use the inference model - 2 options

1. LoRA: Since the model is built on a custom architecture, the generate.py file is included in the repository, which is used to run it. It is launched via the terminal/command line (CMD, PowerShell, VSCode terminal). The LoRA model (pre-trained using dialogs) can only be run using generate.py.

2. PreTrained: If the model is pre-trained, generate.py is NOT SUITABLE. Use generate.py to communicate with the model. Keep in mind that a pre-trained model cannot communicate; it functions as a "text extension." ##

1. Model Training

This guide will help you prepare data and train a language model on your home graphics card.

Step 1: Preparing the Workspace:

1. Create a main folder for your project, for example, Models/. Place all scripts (train.py, Prepare_data.py, etc.) in it.
2. Inside, create a checkpoints/ folder. This is critical: the model will save its weights (checkpoints) here during training, preventing you from losing progress.
3. Place two text files with your dataset in the main folder:

• train.txt - the main, large text file on which the model will train.
• val.txt - a short validation text needed to monitor quality and prevent overfitting.

Step 2: Creating a Custom Tokenizer

The model doesn't understand letters; it understands tokens (word fragments). We need to train the BPE tokenizer on your text so that it understands the language perfectly.

Open the train_tokenizer.py file and enter your tokenizer's value in the vocab_size= field. This ensures that the tokenizer is created 100% of the time.

Then open the console (CMD), navigate to the project folder, and enter the following command:

python train.py --data_path train.txt --total_steps 40000 --embed_dim 768 --n_layers 12 --n_heads 12 --vocab_size 40960

(Replace train.txt with your file name if different). What will happen? The script will parse your text and create a tokenizer.json file with a vocabulary size of 40960 tokens (a multiple of 64 is ideal for GPU performance). After creating the tokenizer, the script will return an error – this is absolutely normal! The error occurs because the script itself requires binary formats, not .txt, for training. The main thing is that the tokenizer is ready!

Step 3: Converting the dataset to binary format (.bin)

To prevent the graphics card and RAM from choking on gigabytes of text, we convert it to a special format, np.memmap.

1. Open the Prepare_data.py script in any code editor.
2. In the TXT_FILE = line, specify your text file (e.g., 'train.txt').
3. In the BIN_FILE = line, specify the desired output name (e.g., 'train.bin').
4. Tip: The script is configured to accumulate 9 million tokens in RAM and write them to the hard drive in chunks. This saves your PC's resources.
5. Run this script. Once completed, do the same for the val.txt file to obtain val.bin.

Step 4: Running Training**

Now you have everything you need. Return to the console and run the final command:

python train.py --data_path train.bin --val_path val.bin --total_steps 40000 --save_every 1000 --batch_size 4 --accumulate_steps 16 --embed_dim 768 --n_layers 12 --n_heads 12 --max_seq_len 1024 --lr 1e-4 --vocab_size 40960

⚙️Analysis of launch parameters Carefully adjust these parameters for your graphics card, otherwise you risk getting an OUT OF MEMORY error (insufficient VRAM):

• --data_path and --val_path are the paths to your generated .bin files (training and validation). Validation is strongly recommended!
• --total_steps is the total number of gradient steps the model will take before completion.
• --save_every - how often to save weights. For example, 1000 means that every thousandth step a safe save will be created in the checkpoints/ folder.
• --batch_size - how many texts the video card processes at a time. For 8 GB of VRAM, a value of 4 with a context of 1024 is the sweet spot.
• --accumulate_steps - gradient accumulation. The effective batch size is equal to batch_size * accumulate_steps. The industry standard is to make the final number equal to 64. That is, with a batch of 4, the accumulation should be 16 (16 * 4 = 64).
• --max_seq_len - context window size (how many tokens the model "sees" simultaneously). Increasing this value significantly eats up VRAM.
• --embed_dim - hidden state dimension (the "brain" width). The larger the value, the more complex the logic the model can understand. 768 is an excellent starting point for a compact model.
• --n_layers and --n_heads - the number of layers and attention heads. Important rule: embed_dim divided by n_heads must always equal 64 (e.g., 768 / 12 = 64). Otherwise, the model will be slow and dumber.
• --lr - learning rate. For pretraining, the default is 1e-4.
• --vocab_size - the size of your tokenizer. It must strictly match the size you specified in Step 2 (e.g., 40960).

2. Inference

Create a checkpoints folder in the project directory and place the downloaded model file there (e.g., model.safetensors). Open a terminal and navigate to the project folder:

cd C:\Users\YourName\Desktop\FileName

### 2. Run
For plain text: python generate.py --checkpoint checkpoints/model.safetensors --prompt "Is artificial intelligence dangerous?" --temperature 0.7 --rep_penalty 1.2 --max_new_tokens 400 --device cuda

For dialog text: python generate.py --checkpoint checkpoints/Vexion-LM_mini_lora.safetensors --prompt "[CLS] What is a human [SEP]" --temperature 0.7 --device cuda --use_lora

📝 Prompt writing rules:
[CLS] is the special token at the beginning of your request. Write your question after it.

[SEP] is the special token at the beginning of the AI's response. No text should be written after this token, otherwise the model will break the response logic!

The --use_lora flag is required when running dialog versions of the model so that the script includes additional adapter weights.