docs/LLD.md

# Low-Level Design (LLD)

## Module Details

### s02_tokenizer.py — REMI Tokenizer

**Vocabulary Layout** (485 tokens total):
```
[0]        PAD
[1]        BOS (Beginning of Sequence)
[2]        EOS (End of Sequence)
[3]        SEP (Separator)
[4-131]    NoteOn (MIDI pitch 0-127)
[132-259]  NoteOff (MIDI pitch 0-127)
[260-291]  Velocity (32 bins, each = 4 MIDI velocity units)
[292-391]  TimeShift (10ms steps, 10ms-1000ms)
[392-451]  Tempo (40-200 BPM, 60 bins)
[452-483]  Position (32 positions per bar)
[484]      Bar (bar delimiter)
```

**Encoding Algorithm**:
1. Collect all non-drum notes across instruments
2. Sort by onset time, then pitch
3. For each note: emit [TimeShift, Position, Velocity, NoteOn, TimeShift(duration), NoteOff]
4. Insert Bar tokens at measure boundaries
5. Wrap with BOS/EOS

### s04_model.py — MusicTransformer

**Layer Stack** (per block):
```
Input → RMSNorm → GQA (with RoPE) → Residual Add
      → RMSNorm → SwiGLU FFN       → Residual Add → Output
```

**Grouped Query Attention**:
- 8 query heads, 4 key-value heads
- Each KV head serves 2 query heads (n_rep = 2)
- Head dim = 256/8 = 32
- Uses PyTorch 2.0 `scaled_dot_product_attention` (Flash Attention backend when available)

**RoPE Implementation**:
- Precompute sin/cos frequencies: `freq[i] = 1 / (θ^(2i/d))`
- Apply rotation: `q' = q * cos + rotate_half(q) * sin`
- Device-compatible real-valued implementation (no complex tensors)

**SwiGLU FFN**:
- `output = W2(SiLU(W1(x)) * W3(x))`
- Hidden dim = 448 (nearest multiple of 64 to `2/3 * 4 * 256`)

### s05_trainer.py — Training Loop

**Optimization Details**:
- AdamW: β₁=0.9, β₂=0.95, wd=0.1
- LR schedule: linear warmup (200 steps) → cosine decay to 1e-6
- Gradient accumulation: 4 steps (effective batch = 4 × 4 = 16)
- Max gradient norm: 1.0

**Memory Budget (estimated for 512 seq len)**:
| Component              | Memory     |
|------------------------|-----------|
| Model weights (FP32)   | ~20MB     |
| Gradients              | ~20MB     |
| Optimizer states        | ~40MB     |
| Activations (w/ ckpt)  | ~50MB     |
| Data batch             | ~4MB      |
| **Total**              | **~134MB**|

### s06_generator.py — Inference

**Decoding Pipeline**:
```
logits → temperature_scale → repetition_penalty → top_k_filter → top_p_filter → softmax → multinomial_sample
```

**KV-Cache**:
- Each layer stores (K, V) tensors after each forward pass
- New tokens only compute attention against cached KV + new KV
- Reset between generations to prevent cross-contamination

## Class Diagram

```
MusicTokenizer           MusicTransformer              Trainer
├── midi_to_tokens()     ├── TransformerBlock (×6)     ├── _train_epoch()
├── tokens_to_midi()     │   ├── GroupedQueryAttention  ├── _validate()
├── decode_token()       │   │   ├── wq, wk, wv, wo    ├── _save_checkpoint()
├── save()/load()        │   │   ├── RoPE application   └── load_checkpoint()
│                        │   │   └── KV-cache
└── Token constants      │   ├── SwiGLU FFN
                         │   └── RMSNorm (×2)
                         ├── token_emb (weight-tied)
                         ├── output projection
                         └── forward() / from_config()
```