Upload folder using huggingface_hub

30e9297 verified 5 months ago

3.39 kB

# 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()
```

	# 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()
	```