OMGFormers

Parallel Diffusion Language Model — 66 features

GitHub: https://github.com/fastloraoffical/OMGformers

Install

pip install omgformers

What is OMGFormers?

OMGFormers is a research-grade PyTorch library for building Parallel Diffusion Language Models. It provides a modular, composable set of building blocks covering:

• Attention mechanisms — GQA, MLA, Sliding Window, Linear, Block-Sparse, Flash Attention 2, RoPE variants (YaRN, NTK, LongRoPE), ALiBi, T5 relative bias
• Feed-forward layers — SwiGLU, GeGLU, ReGLU, Standard FFN
• Mixture of Experts — Dense MoE, Soft MoE, load-balancing loss
• Diffusion — Mask scheduler, Parallel decoder for masked diffusion LM training
• LoRA / DoRA — Parameter-efficient fine-tuning adapters with merge/save/load
• Training utilities — EMA, Lion optimizer, warm-up cosine schedules, FSDP, gradient checkpointing, checkpoint manager
• Tokenizer — HuggingFace-compatible tokenizer with char-level fallback, special token management, encode/decode batch, mask_tokens for diffusion
• Advanced — KV cache, multi-token prediction, model merging (SLERP, DARE, TIES), reward model, PPO, int8/int4 quantization, GGUF export, RAG context injection, dynamic batching, chunked long-doc attention

Quick Start

import torch
from omgformers import OMGConfig, OMGModel, create_base_model, OMGTokenizer

# Build a small model
cfg = OMGConfig(
    vocab_size=32000,
    hidden_size=512,
    num_layers=6,
    num_heads=8,
)
model = OMGModel(cfg)

# Or use the fast initializer
model, cfg = create_base_model(hidden_size=512, num_layers=6)

# Tokenizer
tok = OMGTokenizer.from_pretrained("gpt2")  # or char-level fallback
ids = tok.encode("Hello, world!")
print(tok.decode(ids))

Fine-tuning

from omgformers import FineTuneConfig, FineTuner

ft_cfg = FineTuneConfig(method="lora", lora_rank=16, steps=1000)
tuner  = FineTuner(model, tokenizer=tok, config=ft_cfg)
tuner.train(train_dataloader)

LoRA

from omgformers import add_lora, merge_lora, save_lora, LoRAConfig

lora_cfg = LoRAConfig(rank=16, alpha=32, target_modules=["q_proj", "v_proj"])
model    = add_lora(model, lora_cfg)
# ... train ...
model    = merge_lora(model)
save_lora(model, "my_lora_weights/")

Mixture of Experts

from omgformers import MoEConfig, OMGConfig

cfg = OMGConfig(
    hidden_size=1024,
    num_layers=12,
    moe=MoEConfig(num_experts=8, top_k=2, aux_loss_coeff=0.01),
)

What's New in v2.0.6-preview

Bug fixes: #T5, #T6, #T7, #Mo1–Mo4, #A1–A5, #M1–M3, #C1–C3

Requirements

• Python ≥ 3.9
• PyTorch ≥ 2.0
• Optional: transformers, safetensors, flash-attn, bitsandbytes

License

Apache License 2.0 — see LICENSE.