Readme.md

��

# >��� 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.