krishnah27/AriaLM

Music Generation LLM

A LLaMA-style Transformer model for symbolic music generation, trained on MIDI data.

Architecture

• Model: LLaMA-style Transformer with RoPE, GQA, SwiGLU, RMSNorm
• Tokenizer: REMI (REvamped MIDI-derived) — SOTA for symbolic music
• Dataset: drengskapur/midi-classical-music — 4,796 classical MIDI files (~50MB)
• Training: AdamW + Cosine LR warmup + AMP + Gradient Checkpointing

Key Features

• Memory Efficient: Grouped Query Attention, gradient checkpointing, mixed precision
• OOM Safe: Conservative batch sizes, AMP, lazy data loading
• SOTA Techniques: RoPE, SwiGLU, RMSNorm, KV-cache, top-p/top-k sampling

Quick Start

# Install dependencies
pip install -r requirements.txt

# Train + Generate (default)
python3 -m src.s00_main train+generate

# Train only
python3 -m src.s00_main train --epochs 20 --batch-size 4

# Generate from checkpoint
python3 -m src.s00_main generate --temperature 0.85

Project Structure

music_gen_llm/
├── src/
│   ├── s00_main.py          # Entry point — orchestrates pipeline
│   ├── s01_config.py         # All configuration dataclasses
│   ├── s02_tokenizer.py      # REMI MIDI tokenizer
│   ├── s03_dataset.py        # Data download + tokenization + DataLoader
│   ├── s04_model.py          # MusicTransformer (LLaMA-style)
│   ├── s05_trainer.py        # Training loop with AMP + checkpointing
│   ├── s06_generator.py      # Autoregressive generation with KV-cache
│   └── s07_utils.py          # Logging, memory monitoring, seeding
├── tests/
│   └── test_pipeline.py      # 7 unit tests covering all components
├── scripts/
│   ├── download_data.sh      # Dataset setup
│   ├── train.sh              # Training launcher
│   └── generate.sh           # Generation launcher
├── docs/
│   ├── README.md             # This file
│   ├── HLD.md                # High-Level Design
│   ├── LLD.md                # Low-Level Design
│   └── flow_diagram.drawio   # Execution flow diagram
├── data/                     # Downloaded MIDI + tokenized cache
├── checkpoints/              # Saved model weights
├── output/                   # Generated MIDI files
├── requirements.txt
├── Dockerfile
└── .gitignore

Execution Flow

s00_main.py → s01_config.py → s02_tokenizer.py → s03_dataset.py → s04_model.py → s05_trainer.py → s06_generator.py
     │              │                │                  │                │               │                │
  Entry point   Load configs   Init tokenizer   Download & tokenize  Build model  Train loop     Generate MIDI

Model Specifications

Parameter	Value
Dim	256
Layers	6
Heads	8 (Q) / 4 (KV) — GQA
Hidden (FFN)	448 (SwiGLU)
Max Seq Len	1024
Vocab Size	485 (REMI tokens)
Parameters	~5M
Precision	BF16/FP16 (AMP)

Algorithms Used

1. Rotary Position Embeddings (RoPE) — Su et al. 2021
2. Grouped Query Attention (GQA) — Ainslie et al. 2023, from LLaMA-2
3. SwiGLU Activation — Shazeer 2020, from LLaMA
4. RMS Layer Normalization — Zhang & Sennrich 2019
5. REMI Tokenization — Huang & Yang 2020
6. Cosine Annealing with Warmup — Loshchilov & Hutter 2017
7. Gradient Checkpointing — Chen et al. 2016
8. KV-Cache — Standard for efficient autoregressive decoding
9. Nucleus (Top-p) + Top-k Sampling — Holtzman et al. 2020
10. Repetition Penalty — Keskar et al. 2019