MusicGen Fine-Tuning Pipeline

Overview

What Is This Pipeline?

This pipeline fine-tunes Meta's MusicGen — a state-of-the-art text-to-music generation model with 3.3 billion parameters — using LoRA (Low-Rank Adaptation) to specialize the model for specific music genres while keeping the base model frozen.

Instead of training all 3.3B parameters (which would require massive compute and risk catastrophic forgetting), LoRA injects small trainable adapters into the transformer's attention layers. This reduces trainable parameters by ~99.8% while still achieving strong genre specialization.

Key Capabilities

• Multi-genre training — Train on multiple genres simultaneously or focus on a single genre
• Stereo output — Full stereo generation at 32 kHz sample rate
• Data augmentation — Built-in pitch shift, time stretch, gaussian noise, and gain augmentation
• Automatic metadata — AI-generated text descriptions via Gemini or OpenAI for conditioning
• Vocal removal — Automatic instrumental extraction using HT-Demucs
• W&B integration — Real-time experiment tracking with Weights & Biases
• Docker deployment — Containerized for serverless GPU training (RunPod, Lambda)
• Checkpoint management — Automatic best-model saving, early stopping, and run isolation

Architecture

Pipeline Flow

The complete training pipeline follows this flow:

Model Architecture

Under the hood, MusicGen consists of two main components:

LoRA adapters are injected into the Transformer LM's attention layers: specifically the q_proj, k_proj, v_proj, and out_proj linear layers in each transformer block.

Prerequisites

GPU Requirements

Software Requirements

• Python 3.10+
• CUDA 12.1+ with cuDNN 8
• PyTorch 2.1.0+ (cu121)
• Docker (optional, for containerized training)
• NVIDIA Container Toolkit (for Docker GPU passthrough)

API Keys (Optional)

Installation

Option A: Bare Metal (Direct GPU Access)

# Clone lumina-musicgen
git clone https://github.com/FoldArtists/lumina-musicgen.git
cd lumina-musicgen

python3 -m venv .venv
source .venv/bin/activate

# Install PyTorch with CUDA 12.1
pip install torch==2.1.0 torchaudio==2.1.0 \
  --index-url https://download.pytorch.org/whl/cu121

# Install project dependencies
pip install -e .

python scripts/verify_deps.py
# Should output:
# ✓ PyTorch 2.1.0+cu121
# ✓ CUDA available
# ✓ AudioCraft 1.3.0
# ✓ PEFT 0.18.1

Option B: Docker (Recommended for Serverless)

cd lumina-musicgen
docker build -f docker/Dockerfile.trainer \
  -t lumina-musicgen-trainer:1.0 .

# Dry-run test (validates GPU, imports, config)
docker run --gpus all \
  -v /path/to/your/audio:/data \
  -v /tmp/test-output:/output \
  -e DRY_RUN=true \
  lumina-musicgen-trainer:1.0

Data Preparation

Directory Structure

Organize your audio files by genre in subdirectories:

/your/audio/data/
├── blues/
│   ├── track001.wav
│   ├── track002.wav
│   └── ...
├── jazz/
│   ├── track001.wav
│   └── ...
├── rock/
│   └── ...
└── classical/
    └── ...

Audio Requirements

python scripts/prepare_gtzan.py --output-dir /data/gtzan

Automatic Data Processing

The pipeline automatically handles these preprocessing steps during training:

1. Manifest creation — Scans audio files and creates train.jsonl / val.jsonl splits
2. Vocal removal — Uses HT-Demucs to extract instrumental stems (configurable)
3. Metadata generation — Creates text descriptions using Gemini AI for text-conditioning
4. Segmentation — Splits audio into 30-second segments at 32 kHz for training

Configuration Reference

Configuration uses OmegaConf YAML with a base + experiment override pattern. The base config (configs/base.yaml) defines all defaults. Experiment configs in configs/experiments/ override specific values.

Creating an Experiment Config

Create a YAML file in configs/experiments/ that overrides only the settings you want to change:

# configs/experiments/my_experiment.yaml
data:
  source_dir: "/path/to/my/audio"
  dataset_dir: "/path/to/processed"

lora:
  rank: 64
  alpha: 192        # 3× rank

training:
  epochs: 50
  batch_size: 2      # For 24 GB GPUs
  optimizer:
    lr: 1.0e-4

logging:
  wandb:
    name: "my-experiment"

Training Guide

Running a Training Job

export WANDB_API_KEY="your_wandb_key"
export GEMINI_API_KEY="your_gemini_key"  # optional

# Dry run first (prints config, no training)
python scripts/train_adapter.py \
  --config configs/experiments/my_experiment.yaml \
  --dry-run

# Actual training
python scripts/train_adapter.py \
  --config configs/experiments/my_experiment.yaml

Training Tips

Docker Deployment

The Docker image packages the entire training environment with frozen, proven dependencies. This is the recommended approach for serverless GPU platforms.

Volume Mounts

Environment Variables

Full Docker Run Command

docker run --gpus all \
  -v /home/user/audio:/data \
  -v /home/user/configs:/config \
  -v /home/user/output:/output \
  -e WANDB_API_KEY=your_key \
  -e EXPERIMENT_CONFIG=docker_preset_c.yaml \
  lumina-musicgen-trainer:1.0

Serverless Platforms

Tested on these serverless GPU providers:

• RunPod — Use GPU Pod with Docker image. Mount network volumes for data persistence.
• Lambda Labs — Push image to ECR/DockerHub, launch via API.
• Vast.ai — Upload image, select GPU tier, configure volume mounts.

Outputs & Evaluation

Output Structure

/output/runs/experiment_20260305-151233/
├── config.yaml          # Frozen config snapshot
├── training.log         # Full training log
├── metadata.json        # Data metadata copy
├── checkpoints/
│   ├── epoch_1/         # Per-epoch adapter checkpoints
│   ├── epoch_2/
│   └── best/            # Best model (lowest val_loss)
│       ├── lora_A.pt
│       └── lora_B.pt
├── final/
│   └── adapter_final.pt # Final merged adapter
└── samples/
    ├── epoch_5_sample_0.wav
    └── epoch_10_sample_0.wav

Generated Audio Samples

The pipeline generates audio samples at configurable intervals during training. These allow you to aurally monitor how the model adapts to your target genre. Each sample is a 30-second stereo WAV at 32 kHz.

Evaluation Metrics

Troubleshooting

# Effective batch size = batch_size × gradient_accumulation_steps
training:
  batch_size: 1
  gradient_accumulation_steps: 4  # Effective batch size = 4

sudo usermod -aG docker $USER
# Log out and back in for changes to take effect