��
# >��� Memorizing Transformer with Grouped Query Attention
An extended GPT-2-style large language model (LLM) that implements core components from the research paper Memorizing Transformers (Wu et al., 2022).
This project incorporates Grouped Query Attention (GQA), KNN-based memory retrieval, XL-style attention, and Rotary Positional Encoding (RoPE).
The training pipeline supports distributed training, data parallelism, and sharded dataset streaming.
---
## =�,� Key Features
- ' Grouped Query Attention: Efficient query representation by grouping multiple attention heads for shared K/V access
- ' KNN-based Memory: Long-term memory retrieval from past activations using a learned KNN mechanism
- ' XL-style Attention: Recurrence-based memory layers adapted for KNN and grouped attention logic
- ' Rotary Positional Encoding: More efficient and generalizable positional representation than vanilla sin-cos encoding
- ' Sharded Dataset Loader: Handles large datasets with sharding and supports data parallelism via PyTorch DDP
- ' Custom Memory Clearing Logic: Memory reset and lifespan mechanisms tuned for stability and performance during training
- ' Mixed Precision & DDP Training: Efficient large-scale training using torch.autocast and torchrun
---
## =��� Project Structure
bash MEM_TRANSFORMER/ %%% configs/ % %%% config.json # Model + training hyperparameters % %%% data/ % %%% edu_fineweb/ # Token-sharded training data % % %%% train_000001.npy % % %%% train_000002.npy % % %%% test_000001.npy % %%% hellaswag/ % % %%% hellaswag_val.jsonl % %%% fineweb.py # Sharding logic with memory-aligned sequence control % %%% model_core/ % %%% __init__.py % %%% attention.py # Grouped Query Attention, KNN & XL attention logic.Rotary Positional Encoding implementation % %%% model.py # Transformer model with memory and RoPE support % %%% dataloader.py # Memory-aware DataLoader % %%% training.py # train_memgpt function % %%% scripts/ % %%% train.py # Training script (DDP-compatible) % %%% evaluate.py # Evaluation on benchmarks % %%% generate.py # Text generation from trained model % %%% evaluation/ % %%% __init__.py % %%% hellaswag.py # HellaSwag data loader % %%% val_hellaswag.py # Evaluation logic with loss-based scoring % %%% logs/ % %%% log.txt # Training logs % %%% model_*.pt # Checkpoints % %%% .gitignore %%% README.md %%% requirements.txt
## �&� Configuration
Edit the config file at configs/config.json to adjust model and training hyperparameters:
json { "model": { "block_size": 1024, "vocab_size": 50304, "n_layer": 12, "n_head": 12, "n_embd": 768, "n_kv_head": 4, "max_knn_memories": 81920 }, "training": { "max_steps": 19073, "log_dir": "log", "total_batch_size": 2048, "B": 64, "T": 1024, "max_lr": 0.0006, "min_lr": 0.00006, "warmup_steps": 715, "weight_decay": 0.1, "learning_rate": 0.0006 } }
## =� Training
�%� Single-GPU Training
bash python scripts/train.py
�%� Distributed Training (Multi-GPU with DDP)
bash torchrun --nproc_per_node=NUM_GPUS scripts/train.py
Replace NUM_GPUS with the number of GPUs available.
##=��� Evaluation Evaluate on the HellaSwag benchmark
=��� Evaluation
Evaluate on the HellaSwag benchmark:
python scripts/evaluate.py
Make sure the file data/hellaswag/hellaswag_val.jsonl is present.
The evaluation uses completion scoring based on masked loss comparisons across candidate endings.
>��� Attention Mechanism Notes
>��� Grouped Query Attention (GQA)
n_head query heads
n_kv_head shared key/value heads
Query heads are grouped and averaged before memory lookup
More efficient than per-head K/V for large models
>��� KNN Memory Integration
A maximum memory buffer of 81920 tokens (max_knn_memories)
Query vectors are projected and grouped for efficient KNN search
Careful shape transformations ensure fast grouped matching
>��� XL-style Attention + Memory Clearing
Recurrence with cached memory states
Implements custom memory clearing to avoid stale token influence
Helps stability in long training runs
=ء� Positional Encoding
Rotary Positional Encoding (RoPE) replaces standard sin/cos
RoPE improves generalization over longer contexts
Implemented in model_core/rotary.py
>��� Dataloader & Dataset Handling
Sharded training data using .npy files
Matching stride and memory alignment logic
Optimized for DDP compatibility and large-scale throughput
Code in model_core/dataloader.py and data/fineweb.py
=��� Requirements
Install dependencies:
bash pip install -r requirements.txt
Ensure PyTorch and CUDA versions match your GPU setup.