AriaLM / src /01_config.py

Upload folder using huggingface_hub

30e9297 verified 5 months ago

3 kB

"""
Configuration for the Music Generation LLM.
Tuned for constrained hardware (<=8GB VRAM, <=16GB RAM).
"""
from dataclasses import dataclass, field
from pathlib import Path
import torch


@dataclass
class ModelConfig:
    vocab_size: int = 0  # Set dynamically from tokenizer
    dim: int = 256
    n_layers: int = 6
    n_heads: int = 8
    n_kv_heads: int = 4  # Grouped Query Attention: fewer KV heads saves memory
    max_seq_len: int = 1024
    hidden_dim: int = 0  # Auto-calculated as 4 * dim * 2/3 rounded to multiple of 64
    dropout: float = 0.1
    rope_theta: float = 10000.0

    def __post_init__(self):
        if self.hidden_dim == 0:
            # SwiGLU hidden dim: 4 * dim * 2/3 (LLaMA convention)
            self.hidden_dim = int(2 * (4 * self.dim) / 3)
            # Round to nearest multiple of 64 for hardware efficiency
            self.hidden_dim = 64 * ((self.hidden_dim + 63) // 64)


@dataclass
class TrainConfig:
    batch_size: int = 8
    grad_accum_steps: int = 4  # Effective batch = 32
    learning_rate: float = 3e-4
    weight_decay: float = 0.1
    max_epochs: int = 50
    warmup_steps: int = 200
    max_grad_norm: float = 1.0
    use_amp: bool = True  # Mixed precision to save memory
    grad_checkpoint: bool = True  # Gradient checkpointing for OOM prevention
    eval_interval: int = 500
    save_interval: int = 1000
    log_interval: int = 50
    patience: int = 10  # Early stopping patience (epochs)
    min_delta: float = 0.001  # Minimum improvement for early stopping
    num_workers: int = 2  # DataLoader workers (low for constrained RAM)
    pin_memory: bool = True
    prefetch_factor: int = 2


@dataclass
class DataConfig:
    dataset_name: str = "drengskapur/midi-classical-music"
    max_seq_len: int = 1024
    train_split: float = 0.9
    val_split: float = 0.1
    tokenizer_params: str = "REMI"  # REMI tokenization — SOTA for symbolic music


@dataclass
class GenConfig:
    temperature: float = 0.85
    top_k: int = 40
    top_p: float = 0.92
    max_tokens: int = 1024
    repetition_penalty: float = 1.15
    seed: int = 42


@dataclass
class PathConfig:
    base_dir: Path = field(default_factory=lambda: Path(__file__).resolve().parent.parent)
    data_dir: Path = field(init=False)
    checkpoint_dir: Path = field(init=False)
    output_dir: Path = field(init=False)
    log_dir: Path = field(init=False)
    tokenizer_path: Path = field(init=False)

    def __post_init__(self):
        self.data_dir = self.base_dir / "data"
        self.checkpoint_dir = self.base_dir / "checkpoints"
        self.output_dir = self.base_dir / "output"
        self.log_dir = self.base_dir / "runs"
        self.tokenizer_path = self.data_dir / "tokenizer.json"
        for d in [self.data_dir, self.checkpoint_dir, self.output_dir, self.log_dir]:
            d.mkdir(parents=True, exist_ok=True)


def get_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")

	"""
	Configuration for the Music Generation LLM.
	Tuned for constrained hardware (<=8GB VRAM, <=16GB RAM).
	"""
	from dataclasses import dataclass, field
	from pathlib import Path
	import torch


	@dataclass
	class ModelConfig:
	vocab_size: int = 0 # Set dynamically from tokenizer
	dim: int = 256
	n_layers: int = 6
	n_heads: int = 8
	n_kv_heads: int = 4 # Grouped Query Attention: fewer KV heads saves memory
	max_seq_len: int = 1024
	hidden_dim: int = 0 # Auto-calculated as 4 * dim * 2/3 rounded to multiple of 64
	dropout: float = 0.1
	rope_theta: float = 10000.0

	def __post_init__(self):
	if self.hidden_dim == 0:
	# SwiGLU hidden dim: 4 * dim * 2/3 (LLaMA convention)
	self.hidden_dim = int(2 * (4 * self.dim) / 3)
	# Round to nearest multiple of 64 for hardware efficiency
	self.hidden_dim = 64 * ((self.hidden_dim + 63) // 64)


	@dataclass
	class TrainConfig:
	batch_size: int = 8
	grad_accum_steps: int = 4 # Effective batch = 32
	learning_rate: float = 3e-4
	weight_decay: float = 0.1
	max_epochs: int = 50
	warmup_steps: int = 200
	max_grad_norm: float = 1.0
	use_amp: bool = True # Mixed precision to save memory
	grad_checkpoint: bool = True # Gradient checkpointing for OOM prevention
	eval_interval: int = 500
	save_interval: int = 1000
	log_interval: int = 50
	patience: int = 10 # Early stopping patience (epochs)
	min_delta: float = 0.001 # Minimum improvement for early stopping
	num_workers: int = 2 # DataLoader workers (low for constrained RAM)
	pin_memory: bool = True
	prefetch_factor: int = 2


	@dataclass
	class DataConfig:
	dataset_name: str = "drengskapur/midi-classical-music"
	max_seq_len: int = 1024
	train_split: float = 0.9
	val_split: float = 0.1
	tokenizer_params: str = "REMI" # REMI tokenization — SOTA for symbolic music


	@dataclass
	class GenConfig:
	temperature: float = 0.85
	top_k: int = 40
	top_p: float = 0.92
	max_tokens: int = 1024
	repetition_penalty: float = 1.15
	seed: int = 42


	@dataclass
	class PathConfig:
	base_dir: Path = field(default_factory=lambda: Path(__file__).resolve().parent.parent)
	data_dir: Path = field(init=False)
	checkpoint_dir: Path = field(init=False)
	output_dir: Path = field(init=False)
	log_dir: Path = field(init=False)
	tokenizer_path: Path = field(init=False)

	def __post_init__(self):
	self.data_dir = self.base_dir / "data"
	self.checkpoint_dir = self.base_dir / "checkpoints"
	self.output_dir = self.base_dir / "output"
	self.log_dir = self.base_dir / "runs"
	self.tokenizer_path = self.data_dir / "tokenizer.json"
	for d in [self.data_dir, self.checkpoint_dir, self.output_dir, self.log_dir]:
	d.mkdir(parents=True, exist_ok=True)


	def get_device() -> torch.device:
	if torch.cuda.is_available():
	return torch.device("cuda")
	return torch.device("cpu")