omgformers 2.0.6.dev0

Parallel Diffusion Language Model — 66 features

OMGFormers

Parallel Diffusion Language Model — 66 features

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

---

Installation

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

Feature	Description
#53	Fast base model initialization (create_base_model)
#54	Fine-tuning engine (FineTuner)
#55	Resume from checkpoint (Trainer.resume_from_checkpoint)
#56	Checkpoint manager (CheckpointManager)
#57	Flash Attention 2 real implementation (flash_attention_forward)
#58	MoE → OMGConfig full integration (MoEConfig)
#61	OMGTokenizer (HF + char-level fallback)
#62	Special token management
#63	Tokenizer save/load
#64	Tokenizer from_pretrained
#65	encode_batch / decode_batch
#66	mask_tokens for diffusion training

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.