Skip to main content

A professional, modular Transformer framework with advanced RLHF, continual learning, and efficient fine-tuning capabilities.

Project description

🌌 Nexuss Transformer Framework (NTF)

From Blank Slate to Superintelligence.
A production-grade, end-to-end framework for training, fine-tuning, and aligning Decoder-Only Large Language Models from scratch.


🚀 Overview

Nexuss Transformer Framework (NTF) is a comprehensive architectural solution designed for researchers and engineers who need full control over the LLM lifecycle. Unlike wrapper libraries, NTF provides a blank-slate implementation of the transformer architecture while leveraging battle-tested backends like Hugging Face accelerate, peft, and trl for distributed training and optimization.

Whether you are pre-training a foundation model on terabytes of text, performing parameter-efficient fine-tuning (LoRA), or aligning models via RLHF (PPO/DPO), NTF offers a unified, modular, and scalable interface.

🔹 Native Ethiopian Tokenization

NTF seamlessly integrates with EthioBBPE, a specialized Byte-Pair Encoding tokenizer optimized for Amharic and Ge'ez scripts. This integration ensures superior compression ratios and cultural nuance preservation for Ethiopian languages right out of the box.

pip install ethiobbpe  # Optional: For enhanced Ethiopian language support

✨ Key Features

🏗️ Core Architecture

  • Pure Decoder-Only Stack: Custom implementation featuring RoPE embeddings, RMSNorm, SwiGLU activations, and Multi-Query Attention (MQA).
  • Blank Slate Training: Build knowledge from scratch with no pretrained weights bias.
  • Gradient Checkpointing: Optimized memory usage for training massive models on limited hardware.
  • Flash Attention Ready: Seamless integration with flash-attn for 2-3x speedups.

🎛️ Advanced Fine-Tuning

  • Full Fine-Tuning: Discriminative learning rates, layer-wise decay, and gradual unfreezing strategies.
  • PEFT & LoRA: Low-Rank Adaptation support via peft for efficient downstream task adaptation.
  • Flexible Freezing: Freeze top-k, bottom-k, or alternating layers to preserve base knowledge.

🧠 Alignment & RLHF

  • Reward Modeling: Train reward models to score model outputs.
  • PPO & DPO: Integrated Proximal Policy Optimization and Direct Preference Optimization using trl.
  • Human Feedback Loop: End-to-end pipeline for RLHF alignment.

🔄 Continual Learning

  • Catastrophic Forgetting Prevention: Built-in Elastic Weight Consolidation (EWC), Experience Replay, and Gradient Episodic Memory (GEM).
  • Lifelong Learning: Seamlessly continue training on new domains without losing previous capabilities.

🛠️ Production Engineering

  • Distributed Training: Native support for DeepSpeed ZeRO and FSDP via accelerate.
  • Smart Checkpointing: Versioned saves, auto-resume, and "best model" tracking based on validation loss.
  • Model Versioning: Semantic versioning system for releasing and rolling back model iterations.

⚡ Quick Start

Installation

# Install from PyPI (Coming Soon)
pip install ntf

# Or install from source
git clone https://github.com/nexuss0781/Nexuss-Transformer.git
cd Nexuss-Transformer
pip install -e .

CLI Usage (Recommended)

The framework provides a powerful command-line interface for all operations:

# Pre-training
ntf train --config pretrain_small --output_dir ./outputs/my-model

# Fine-tuning with LoRA
ntf finetune --model ./outputs/my-model --config finetune_lora --data instructions.jsonl

# RLHF Alignment (DPO)
ntf align --model ./outputs/my-model --config dpo_alignment --method dpo --data preferences.jsonl

# Evaluate model
ntf evaluate --model ./outputs/my-model --metrics perplexity,bleu,rouge

Use --override to customize any parameter:

ntf train --config pretrain_small \
    --override model.hidden_size=1024 training.learning_rate=1e-4 checkpoint.save_steps=1000

Python API

1. Pre-Training from Scratch

Initialize a blank slate model and train on your corpus:

from ntf.models import NTFConfig, NexussTransformer
from ntf.training import NTFTrainer, TrainingArguments

# Define Model Architecture (Small Preset)
config = NTFConfig.from_preset("small") 
model = NexussTransformer(config)

# Configure Training
train_config = TrainingConfig(
    output_dir="./outputs/my-first-llm",
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
    learning_rate=5e-4,
    num_train_epochs=10,
    fp16=True,
    logging_steps=10,
    save_strategy="steps",
    save_steps=500
)

# Initialize Trainer
trainer = Trainer(
    model=model,
    args=train_config,
    train_dataset=my_tokenized_dataset,
    eval_dataset=my_val_dataset
)

# Launch Training
trainer.train()

2. Parameter-Efficient Fine-Tuning (LoRA)

Adapt your pretrained model to a specific task (e.g., Medical QA) with minimal resources:

from ntf.finetuning import PeftFinetuner

finetuner = PeftFinetuner(
    model_path="./outputs/my-first-llm/checkpoint-5000",
    lora_r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"]
)

finetuner.train(dataset=medical_qa_dataset)
finetuner.save_adapter("./adapters/medical-lora")

3. RLHF Alignment (DPO)

Align your model with human preferences:

from ntf.reward import DPOTrainer

dpo = DPOTrainer(
    model_ref="./outputs/my-first-llm",
    beta=0.1,
    loss_type="sigmoid"
)

dpo.train(preference_dataset=human_prefs)

📚 Documentation

Module Description
TRAINING.md Comprehensive End-to-End Guide. Covers installation, configuration, pre-training, fine-tuning, RLHF, and continual learning strategies.
models/ Core Transformer architecture, configurations, and tokenization utilities.
training/ Distributed trainer, checkpoint manager, and data collators.
finetuning/ Full fine-tuning, LoRA/PEFT wrappers, and layer freezing logic.
reward/ Reward modeling and RLHF algorithms (PPO, DPO).
utils/ Continual learning helpers, metrics, and model versioning tools.

🏛️ Architecture Diagram

graph TD
    A[Raw Text Data] --> B[Tokenization]
    B --> C{Training Mode}
    
    C -->|Blank Slate| D[Pre-Training Loop]
    C -->|Adaptation| E[Fine-Tuning Full/LoRA]
    C -->|Alignment| F[RLHF PPO/DPO]
    
    D --> G[Base Model Weights]
    E --> G
    F --> H[Aligned Model]
    
    G --> I[Checkpoint Manager]
    H --> I
    
    I --> J[Model Registry / Versioning]
    J --> K[Inference / Deployment]
    
    subgraph Optimization
    L[Gradient Checkpointing]
    M[Mixed Precision FP16/BF16]
    N[DeepSpeed ZeRO]
    end
    
    D -.-> L
    D -.-> M
    D -.-> N

🛡️ Continual Learning Strategies

NTF includes specialized mechanisms to prevent Catastrophic Forgetting when updating models with new data:

  1. Elastic Weight Consolidation (EWC): Penalizes changes to weights important for previous tasks.
  2. Experience Replay: Interleaves new training batches with stored samples from previous distributions.
  3. Gradual Unfreezing: Slowly unfreezes layers from top to bottom to stabilize representations.
  4. Learning Rate Decay: Applies layer-wise decay to protect lower-level features.

📦 Model Zoo & Presets

NTF comes with predefined architectural presets ready for instantiation:

Preset Parameters Layers Heads Dim Use Case
nano ~10M 4 4 128 Testing / Debugging
small ~60M 6 8 256 Edge Devices / Mobile
medium ~350M 12 12 512 Research / Specialized Tasks
large ~1.2B 24 16 768 General Purpose Foundation
xl ~3.5B 32 20 1024 High-Performance Applications

🤝 Contributing

We welcome contributions from the community! Whether it's optimizing kernels, adding new RLHF algorithms, or improving documentation, please check our Contributing Guidelines.

  1. Fork the repo.
  2. Create your feature branch (git checkout -b feature/AmazingFeature).
  3. Commit your changes (git commit -m 'Add some AmazingFeature').
  4. Push to the branch (git push origin feature/AmazingFeature).
  5. Open a Pull Request.

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.


🙏 Acknowledgments

Built upon the shoulders of giants:


Nexuss Transformer Framework
Empowering the next generation of Language Models.

GitHubIssuesDiscussions

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

ntf-1.0.0.tar.gz (65.5 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

ntf-1.0.0-py3-none-any.whl (70.1 kB view details)

Uploaded Python 3

File details

Details for the file ntf-1.0.0.tar.gz.

File metadata

  • Download URL: ntf-1.0.0.tar.gz
  • Upload date:
  • Size: 65.5 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.11.0rc1

File hashes

Hashes for ntf-1.0.0.tar.gz
Algorithm Hash digest
SHA256 9b75b5107b85308943bb71e87003c4e1da624d728c8379855b7115e2edaa9b4a
MD5 08bc88374c715905e49c5f0522d4da3b
BLAKE2b-256 80472970705f0d32a82323659020299293a4dd2a91a110f9f129110111dc16a3

See more details on using hashes here.

File details

Details for the file ntf-1.0.0-py3-none-any.whl.

File metadata

  • Download URL: ntf-1.0.0-py3-none-any.whl
  • Upload date:
  • Size: 70.1 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.11.0rc1

File hashes

Hashes for ntf-1.0.0-py3-none-any.whl
Algorithm Hash digest
SHA256 f27221f3293d842f46b0d40e76d223367aa96ecf787d6b4199a2a7f9074e8235
MD5 f69849c7ed1e8e065ad48545b2f4c424
BLAKE2b-256 54ce77d461a7a291a10bae258aab71da687b2ab1e09d8f3ce13d3680fc84e6f8

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page