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 fine-tuning them for a conversational format using LoRA adapters.

⚠️ IMPORTANT: Base models are intended solely for further fine-tuning to suit your needs. Models marked LoRA have already been fine-tuned on a conversational dataset and are ready for conversation (further training is not recommended!).

🚀 Model Features

• Custom Architecture: Optimized transformer with gradient checkpointing support to save VRAM. * Efficient Training: Using the 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 dialog 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.
• MoE (Mixture of Experts): Built-in system of 4 independent experts with a routing mechanism. Allows the model to be "smart" by activating only some of them during generation to save resources. * GQA (Grouped-Query Attention): Drastically reduces video memory consumption when generating long contexts (up to 2048+ tokens), making VRAM growth linear rather than quadratic.
• SwiGLU & RMSNorm: Using the latest activation and normalization functions (industry standard, a la Llama 3) for maximum logic and convergence speed.
• Custom BPE Tokenizer: A dictionary of exactly 40960 tokens—an ideal multiple of 64, ensuring 100% efficiency of the GPU's Tensor Cores.
• FlashAttention-2 & Memory Management: Native support for sdpa_kernel and work with a dynamic memory allocator (expandable_segments) to prevent OOMs and VRAM fragmentation.

📊 Vexion Model Family

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

Vexion-LM Family

Model	Parameters	Context	Status	Train_Loss	Val_Loss	Vocab_Size	Total_Steps	Sliding window
Vexion-LM Mini	317M	1024	base model	2.89	2.69	40960	40000	no
Vexion-LM V1	688M	2048	-	-	-	40960	-	1024
Vexion-LM V2	~850M	1024	-	-	-	40960	-	512

Technical details:

• Hardware: 1x RTX 3060 Ti 8GB
• Settings: batch_size - 8, accumulate_steps - 8
• Number of tokens in the Base Mini Model pre-train - 2.6 BILLION
• Total training time: 130 hours
• Pre-train dataset: OSCAR

⚠️ 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 only 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 graphics cards to save video memory and speed up execution. Manual weight conversion is not required.

🛠 How to use the model - 2 options

1. LoRA: Since the model is built on a custom architecture, the repository includes a file called generate.py for launching. Launching is done via a terminal/command line (CMD, PowerShell, VSCode terminal). A LoRA model (pre-trained using dialogs) can only be launched using generate.py.

2. PreTrained: For a pre-trained model, 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 add-on."

1. Preparation

Create a checkpoints folder in the project directory and place the downloaded model file (e.g., model.safetensors) there. 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

📝 Prompt writing rules:
[CLS] is the special token that begins your request. Write your question after it.

[SEP] is the special token that begins 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 can include additional adapter weights.