TrAIner Beta

Key Features

Installation and Setup

# Clone the repository
git clone https://github.com/CPScript/TrAIner.git
cd TrAIner

# Install dependencies
pip install torch safetensors tqdm numpy

# For CUDA support (adjust based on your CUDA version)
pip install torch --extra-index-url https://download.pytorch.org/whl/cu118

# Verify installation
python -c "import torch; print(f'PyTorch version: {torch.__version__}, CUDA available: {torch.cuda.is_available()}')"

# config.json
{
  "model": {
    "hidden_size": 768,
    "num_layers": 12,
    "num_heads": 12,
    "num_kv_heads": 4,
    "feed_forward_dim": 3072,
    "max_seq_length": 2048,
    "dropout": 0.1,
    "use_moe": true,
    "num_experts": 8,
    "use_rotary": true,
    "use_rmsnorm": true
  },
  "training": {
    "data_path": "./data/dialogues.txt",
    "data_format": "txt",
    "output_dir": "./model",
    "batch_size": 4,
    "epochs": 30,
    "learning_rate": 5e-5,
    "weight_decay": 0.01,
    "warmup_steps": 1000,
    "fp16": true
  }
}

Quickstart Guide

Training Your First Model

# Create a sample training data file
mkdir -p data
cat > data/dialogues.txt << EOF
<|user|>What is artificial intelligence?<|assistant|>Artificial intelligence refers to systems or machines that can perform tasks that typically require human intelligence. This includes learning from examples and experience, recognizing objects, understanding and responding to language, making decisions, and solving problems.<|end|>
<|user|>Explain neural networks<|assistant|>Neural networks are computing systems inspired by the biological neural networks in animal brains. They consist of artificial neurons that can learn from and make decisions based on input data. Deep learning, a subset of machine learning, uses multiple layers of these neural networks to progressively extract higher-level features from raw input.<|end|>
EOF

# Train a model with default parameters
python train.py --data_path ./data/dialogues.txt --epochs 3 --batch_size 2

# The model will be saved in the ./model directory

Running Inference

# Start the interactive chat interface
python chat.py --model_path ./model/final

# You can now chat with your trained model
> Tell me about machine learning
> exit  # Type 'exit' to quit

Non-Interactive Generation

# Generate a response for a specific input
python chat.py --model_path ./model/final --input "What are transformer models?" --no_interactive

Customizing Model Architecture

# Train with a customized model architecture
python train.py \
  --data_path ./data/dialogues.txt \
  --hidden_size 512 \
  --num_layers 8 \
  --num_heads 8 \
  --feed_forward_dim 2048 \
  --use_moe \
  --num_experts 4 \
  --batch_size 2 \
  --epochs 5

Model Architecture

TrAIner implements a modern transformer-based architecture with several advanced features:

+--------------------------------+
|         Input Embedding        |
+--------------------------------+
              |
              v
+--------------------------------+
|     Positional Encoding        |
+--------------------------------+
              |
              v
+--------------------------------+
|  Transformer Layer 1           |
|  +--------------------------+  |
|  |  Self-Attention          |  |
|  |  (GQA + RoPE)            |  |
|  +--------------------------+  |
|              |                 |
|              v                 |
|  +--------------------------+  |
|  |  Mixture of Experts      |  |
|  |  Feed Forward Network    |  |
|  +--------------------------+  |
+--------------------------------+
              |
              v
+--------------------------------+
|  Transformer Layer 2...N       |
+--------------------------------+
              |
              v
+--------------------------------+
|     Layer Normalization        |
+--------------------------------+
              |
              v
+--------------------------------+
|     Language Model Head        |
+--------------------------------+

Core Components

class MultiHeadAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_heads
        self.num_kv_heads = config.num_kv_heads if config.use_gqa else config.num_heads
        self.head_dim = config.hidden_size // config.num_heads
        
        # For GQA, we have different numbers of q and kv heads
        self.kv_repeats = self.num_heads // self.num_kv_heads if config.use_gqa else 1
        
        # Q, K, V projections
        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias)
        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias)
        self.out_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)

class MoEFeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.experts = nn.ModuleList([
            GEGLU(
                config.hidden_size, 
                config.feed_forward_dim, 
                bias=True
            ) for _ in range(config.num_experts)
        ])
        self.gate = nn.Linear(config.hidden_size, config.num_experts)
        self.dropout = nn.Dropout(config.dropout)
        self.num_experts = config.num_experts
        self.expert_capacity = 2  # Number of tokens per expert

Training Procedures

TrAIner provides comprehensive training capabilities with support for various optimizations:

# Enable mixed precision training with FP16
python train.py --data_path ./data/dialogues.txt --fp16

# Use BF16 for better numerical stability on supported hardware
python train.py --data_path ./data/dialogues.txt --optimized_bf16

# Use gradient accumulation for larger effective batch size
python train.py --data_path ./data/dialogues.txt --batch_size 2 --gradient_accumulation_steps 8
# Effective batch size: 2 * 8 = 16

Inference

TrAIner provides efficient inference capabilities with several optimization techniques:

# Enable speculative decoding
python chat.py --model_path ./model/final --use_speculative --spec_len 5

# Use a separate draft model
python chat.py --model_path ./model/final --use_speculative --draft_model ./model/smaller

# Enable 8-bit quantization for inference
python chat.py --model_path ./model/final --quantize

Data Processing

TrAIner supports flexible data formats and processing pipelines:

<|user|>What is machine learning?<|assistant|>Machine learning is a branch of artificial intelligence that focuses on developing systems that can learn from and make decisions based on data. Instead of being explicitly programmed to perform a task, these systems learn from experience.<|end|>
<|user|>How does deep learning differ from machine learning?<|assistant|>Deep learning is a subset of machine learning that uses neural networks with many layers (hence "deep"). While traditional machine learning requires feature extraction by humans, deep learning automatically discovers the representations needed for feature detection or classification from raw data. It's particularly powerful for complex tasks like image and speech recognition.<|end|>

{"messages": [{"role": "user", "content": "Explain the concept of transformer models."}, {"role": "assistant", "content": "Transformer models are a type of neural network architecture introduced in the paper 'Attention Is All You Need' by Vaswani et al. They rely primarily on self-attention mechanisms rather than recurrence or convolution, allowing them to process input sequences in parallel rather than sequentially. This parallelization enables more efficient training on large datasets. The architecture typically consists of an encoder and decoder, each made up of attention layers and feed-forward neural networks. Transformers have become the foundation for many state-of-the-art language models like GPT, BERT, and T5."}]}
{"messages": [{"role": "system", "content": "You are a helpful AI assistant."}, {"role": "user", "content": "What are the limitations of current AI systems?"}, {"role": "assistant", "thinking": "I should cover limitations like data quality dependence, lack of common sense, explainability issues, and resource requirements.", "content": "Current AI systems face several important limitations: 1) They're highly dependent on training data quality and can perpetuate biases present in that data. 2) They lack true understanding or common sense reasoning. 3) Many modern AI systems, especially deep learning ones, function as 'black boxes' with limited explainability. 4) They typically require significant computational resources and energy to train. 5) They struggle with transferring knowledge between different tasks (though transfer learning helps). 6) They're vulnerable to adversarial attacks and can be fooled by carefully crafted inputs."}]}

Mixture of Experts

The Mixture of Experts (MoE) module in TrAIner enables training models with significantly increased parameter count while maintaining reasonable computational requirements. MoE works by activating only a subset of the model's parameters for each input token, which enables models to scale to billions of parameters even on consumer hardware.

This section will include detailed explanations of:

• How MoE architecture works in transformer models
• TrAIner's implementation of routing algorithms
• Expert capacity and load balancing techniques
• Performance benchmarks comparing standard vs MoE models
• Step-by-step guide to configuring and training MoE models
• Advanced tuning parameters and best practices

Distributed Training

TrAIner supports distributed training across multiple GPUs and even multiple machines, allowing you to scale up your training for faster completion or to accommodate larger models. This is implemented using PyTorch's DistributedDataParallel (DDP) framework.

This section will cover:

• Setting up multi-GPU training on a single machine
• Configuring distributed training across multiple machines
• Scaling strategies and performance optimizations
• Handling checkpoints and model saving in distributed environments
• Common issues and troubleshooting tips
• Advanced configurations for different network topologies

Quantization

Quantization is a technique that reduces the precision of model weights from 32-bit floating point (FP32) to lower precision formats like 8-bit integers (INT8). This can significantly reduce memory usage and improve inference speed with minimal impact on generation quality.

This section will include:

• Post-training quantization techniques supported in TrAIner
• Quantization-aware training for better results
• Performance benchmarks comparing FP32, FP16, and INT8 models
• Effects of quantization on different model components
• Advanced calibration techniques to minimize accuracy loss
• Hardware-specific optimizations for quantized models

Speculative Decoding

Speculative decoding is an advanced technique that significantly speeds up text generation by using a smaller "draft" model to predict multiple tokens at once, which are then verified by the main model. This can provide a 2-3x speedup in generation while maintaining output quality.

This section will cover:

• Theoretical foundation of speculative decoding
• Implementation details in TrAIner
• Creating and training effective draft models
• Performance benchmarks and quality comparisons
• Tuning parameters to optimize the speed/quality tradeoff
• Advanced techniques like tree-based speculation

Text Generation

TrAIner provides flexible text generation capabilities with various options to control the output style, creativity, and generation speed. This section will detail how to use the generation API for different use cases.

Topics to be covered include:

• Interactive chat interface usage
• Programmatic text generation with the API
• Controlling parameters like temperature, top-p, and top-k
• Streaming generation for real-time applications
• Advanced sampling techniques and their effects
• Best practices for prompt engineering

Fine-tuning

Fine-tuning allows you to take a pre-trained model and adapt it to specific domains or tasks by training on your own datasets. This often achieves better results than training from scratch, especially for specialized applications.

This section will provide detailed information on:

• Preparing your dataset for fine-tuning
• Selecting appropriate hyperparameters
• Techniques to prevent catastrophic forgetting
• Parameter-efficient fine-tuning methods
• Evaluating fine-tuned model performance
• Advanced fine-tuning for specific applications

JSONL Format

TrAIner supports structured JSONL (JSON Lines) format for training data, which provides more flexibility than plain text formats. This allows for more complex conversation structures, metadata, and special formatting.

This section will include:

• JSONL format specification for TrAIner
• Examples of various conversation patterns
• Supporting system prompts and multi-turn conversations
• Including metadata and auxiliary information
• Tools for converting between formats
• Performance considerations for large datasets

CLI Usage

TrAIner provides a comprehensive command-line interface for training, inference, and model management. This section will serve as a complete reference for all available options.

Topics to be covered include:

• Complete listing of all CLI commands and parameters
• Train command options and syntax
• Chat command options for inference
• Model conversion and export utilities
• Data processing command-line tools
• Automation and scripting examples

Model API

The Model API provides programmatic access to TrAIner's model architecture, allowing developers to integrate, extend, and customize models for their specific needs.

This section will include:

• ChatModel class API reference
• ChatConfig options and parameter details
• Methods for loading and saving models
• Generation API and parameter explanations
• Forward pass customization options
• Extending the model with custom layers

Tokenizer API

The Tokenizer API handles the conversion between text and token IDs that the model can process. TrAIner includes a custom tokenizer designed to efficiently handle multiple languages and special tokens.

This documentation will cover:

• ChatTokenizer class API reference
• Building and managing custom vocabularies
• Adding special tokens and handling
• Tokenization patterns and regular expressions
• Encoding and decoding methods
• Performance optimization techniques

Training API

The Training API provides programmatic access to TrAIner's training pipeline, allowing developers to customize the training process, integrate custom datasets, and implement advanced training techniques.

This section will include:

• Training loop implementation details
• Optimizer and scheduler configuration
• Gradient accumulation and checkpointing
• Mixed precision training API
• Distributed training setup
• Custom training callbacks and hooks

Dataset API

The Dataset API handles loading, processing, and batching training data for efficient model training. This includes support for various data formats, preprocessing techniques, and data augmentation.

This documentation will cover:

• ChatDataset and JsonlChatDataset class references
• Creating custom dataset implementations
• Efficient data loading and preprocessing
• Data augmentation techniques
• Handling large datasets efficiently
• Implementing custom collation functions

Performance Benchmarks

Optimization

This section will provide comprehensive information about performance optimization techniques for both training and inference. It will help you maximize the efficiency of TrAIner on your hardware.

Topics to be covered include:

• Memory optimization techniques
• Computational efficiency improvements
• Hardware-specific optimizations
• Speed/quality tradeoffs and how to balance them
• Advanced profiling and bottleneck identification
• Case studies and benchmark analysis

Hardware Requirements

This section will provide detailed hardware recommendations for different model sizes and use cases, helping you plan your hardware requirements based on your specific goals.

Information will include:

• Detailed GPU recommendations by model size
• RAM and CPU requirements
• Storage considerations for datasets and checkpoints
• Cost-effective hardware configurations
• Cloud computing options and cost analysis
• Future hardware roadmap considerations

Frequently Asked Questions

Troubleshooting

This section will provide step-by-step guidance for diagnosing and resolving common issues you might encounter when using TrAIner.

Topics to be covered include:

• Common installation issues and solutions
• CUDA compatibility problems
• Memory errors during training
• Slow training or inference performance
• Model quality issues and fine-tuning problems
• Data formatting and preprocessing errors

Community

TrAIner has a growing community of researchers, developers, and enthusiasts. This section will provide information about how to connect with others, contribute to the project, and share your work.

This section will include:

• Discord server and discussion forums
• Contributing guidelines for developers
• Model sharing platform and community models
• Research papers and publications
• Community showcase of TrAIner applications
• Tutorials and community resources

Roadmap

This section will outline the planned development roadmap for TrAIner, including upcoming features, improvements, and research directions.

Future plans include:

• Support for multi-modal models (text + images)
• Advanced LoRA and QLoRA fine-tuning support
• Improved quantization techniques for even better efficiency
• Web-based training and inference interface
• Mobile deployment options
• Pre-trained model collection and repository