from dataclasses import dataclass from functools import partial from collections import defaultdict import math from typing import Optional from abc import ABC, abstractmethod import torch import torch.nn as nn from torch.nn import functional as F from torch.nn.attention.flex_attention import create_block_mask from utils.bct import PackedBlockSequence from .base import ( Block, LayerNorm, configure_optimizers, estimate_mfu, init_weights_simple, print_with_time_master, ) Z_LOSS = True @dataclass(eq=False) class IOModule(nn.Module, ABC): config: "GPTConfig" layer: int = 0 # 0 = first layer, -1 = last layer def __post_init__(self): super().__init__() @property @abstractmethod def pad_token(self): raise NotImplementedError @abstractmethod def forward(self, x): raise NotImplementedError @abstractmethod def init_weights(self): raise NotImplementedError @dataclass class GPTConfig: # Architecture parameters n_layer: int = 24 n_head: int = 16 n_kv_head: Optional[int] = 4 d_head: int = 64 block_size: int = 14_080 bias: bool = False # Text parameters text_vocab_size: int = 60_032 text_codebook_size: int = 60_001 text_pad_token: Optional[int] = None text_infer_token: Optional[int] = None hoot_vocab_size: int = 20544 hoot_pad_token: Optional[int] = 20480 block_type_vocab_size: int = ( 128 # number of different block types (now includes 5 semantic skip levels) ) # Semantic parameters semantic_vocab_size: int = 4032 semantic_codebook_size: int = 4000 semantic_n_codebooks: int = 1 semantic_rate_hz: int = 25 semantic_shift_factor: int = 5 semantic_pad_token: Optional[int] = None semantic_infer_token: Optional[int] = None semantic_mask_token: Optional[int] = None semantic_cover_token: Optional[int] = None semantic_artist_token: Optional[int] = None semantic_future_token: Optional[int] = None semantic_history_token: Optional[int] = None semantic_overpaint_token: Optional[int] = None semantic_underpaint_token: Optional[int] = None semantic_vox_token: Optional[int] = None semantic_stem_token: Optional[int] = None semantic_playlist_token: Optional[int] = None semantic_sample_token: Optional[int] = None semantic_cond_audio_token: Optional[int] = None # Shared token for text conditioning pairs semantic_type: Optional[str] = "mert" # semantic encoder type: "mert" or "musicfm" # Coarse parameters coarse_vocab_size: int = 2112 coarse_codebook_size: int = 2048 coarse_n_codebooks: int = 0 coarse_rate_hz: int = 25 coarse_shift_factor: int = 5 coarse_pad_token: Optional[int] = None coarse_infer_token: Optional[int] = None # Diffusion parameters vae_embed_dim: int = 128 ditto_embed_dim: int = 128 # Other parameters t_text: int = 2048 t_audio: int = 12_000 use_text_loss: bool = False use_mmbert: bool = False use_vae_input: bool = False output_paradigm: str = "gpt" # "gpt" or "diffusion" output_distribution: str = "semantic" # "semantic" or "vae" use_hoot: bool = False use_ditto: bool = False use_mt5: bool = False # no longer used, for backwards compatibility use_vae_output: bool = False # no longer used, for backwards compatibility use_rotary_pos_emb: bool = True rope_theta: int = 500_000 use_qk_norm: bool = True activation_f: str = "silu" embed_scale_factor: float = 10.0 attention_sliding_window_size: int = 1024 global_every_n_layers: int = 1 use_continuous_semantic_input: bool = False use_block_type_embeddings: bool = False use_reward_head: bool = False # enable reward model mode repa_semantic: bool = False # Add continuous semantic as auxiliary target repa_mixed_semantic: bool = False # Add mixed semantic as auxiliary target for stem conditioning repa_hoot: bool = False # Add hoot encoder embeddings as auxiliary target repa_midi: bool = False # Add midi encoder embeddings as auxiliary target repa_layer: int = -1 # Layer for all repa heads: -1 = final layer, 0-23 = specific layer def __post_init__(self): # Default to multi head attention if self.n_kv_head is None: self.n_kv_head = self.n_head # Initialize text tokens if not provided if self.text_pad_token is None: self.text_pad_token = self.text_codebook_size if self.text_infer_token is None: self.text_infer_token = self.text_codebook_size + 1 # Initialize semantic tokens if not provided if self.semantic_pad_token is None: self.semantic_pad_token = self.semantic_codebook_size if self.semantic_infer_token is None: self.semantic_infer_token = self.semantic_codebook_size + 1 if self.semantic_mask_token is None: self.semantic_mask_token = self.semantic_codebook_size + 2 if self.semantic_cover_token is None: self.semantic_cover_token = self.semantic_codebook_size + 3 if self.semantic_artist_token is None: self.semantic_artist_token = self.semantic_codebook_size + 4 if self.semantic_future_token is None: self.semantic_future_token = self.semantic_codebook_size + 5 if self.semantic_history_token is None: self.semantic_history_token = self.semantic_codebook_size + 6 if self.semantic_overpaint_token is None: self.semantic_overpaint_token = self.semantic_codebook_size + 7 if self.semantic_underpaint_token is None: self.semantic_underpaint_token = self.semantic_codebook_size + 8 if self.semantic_playlist_token is None: self.semantic_playlist_token = self.semantic_codebook_size + 9 if self.semantic_stem_token is None: self.semantic_stem_token = self.semantic_codebook_size + 10 if self.semantic_vox_token is None: self.semantic_vox_token = self.semantic_codebook_size + 11 if self.semantic_sample_token is None: self.semantic_sample_token = self.semantic_codebook_size + 12 if self.semantic_cond_audio_token is None: self.semantic_cond_audio_token = self.semantic_codebook_size + 13 # Initialize coarse tokens if not provided if self.coarse_pad_token is None: self.coarse_pad_token = self.coarse_codebook_size if self.coarse_infer_token is None: self.coarse_infer_token = self.coarse_codebook_size + 1 # Validate token indices assert self.text_pad_token < self.text_vocab_size assert self.text_infer_token < self.text_vocab_size assert self.semantic_pad_token < self.semantic_vocab_size assert self.semantic_infer_token < self.semantic_vocab_size assert self.semantic_mask_token < self.semantic_vocab_size assert self.semantic_cover_token < self.semantic_vocab_size assert self.semantic_artist_token < self.semantic_vocab_size assert self.semantic_future_token < self.semantic_vocab_size assert self.semantic_history_token < self.semantic_vocab_size assert self.semantic_overpaint_token < self.semantic_vocab_size assert self.semantic_underpaint_token < self.semantic_vocab_size assert self.semantic_playlist_token < self.semantic_vocab_size assert self.semantic_stem_token < self.semantic_vocab_size assert self.semantic_vox_token < self.semantic_vocab_size assert self.semantic_cond_audio_token < self.semantic_vocab_size assert self.semantic_sample_token < self.semantic_vocab_size assert self.coarse_pad_token < self.coarse_vocab_size assert self.coarse_infer_token < self.coarse_vocab_size @property def n_embd(self): """The width of the residual stream""" return self.n_head * self.d_head @property def use_learned_pos_emb(self): return not self.use_rotary_pos_emb @property def n_streams(self): return self.semantic_n_codebooks + self.coarse_n_codebooks + 1 @dataclass class GPTTrainConfig: dropout: float = 0.0 attention_type: str = "tao" # "torch", "tao" layer_init: bool = True def causal( b, h, q_idx, kv_idx, document_id: torch.Tensor, block_id: torch.Tensor, block_is_causal: torch.Tensor, window_size: int = -1, ): within_document = document_id[q_idx] == document_id[kv_idx] causal = q_idx >= kv_idx noncausal_within_block = (block_id[q_idx] == block_id[kv_idx]) & ~block_is_causal[kv_idx] # dont use abs cause it OOMs within_window = ((q_idx - kv_idx) <= window_size) & ((kv_idx - q_idx) <= window_size) within_window = within_window | (window_size == -1) # global attention return (causal | noncausal_within_block) & within_document & within_window class MMBertTextInput(IOModule): """Text input module using mmBERT embeddings (hidden_size=768)""" def __init__(self, config: GPTConfig, layer: int = 0): super().__init__(config, layer=layer) self.proj = nn.Linear(768, config.n_embd) # mmBERT hidden size self.ln_mmbert = LayerNorm(config.n_embd) @property def pad_token(self): return torch.zeros(768, dtype=torch.bfloat16).unsqueeze(0) def forward(self, x): x = self.proj(x) x = self.ln_mmbert(x) return x def init_weights(self): nn.init.normal_(self.proj.weight) nn.init.normal_(self.ln_mmbert.weight) class HootTextInput(IOModule): def __init__(self, config: GPTConfig, layer: int = 0): super().__init__(config, layer=layer) self.wte_hoot = nn.Embedding(config.hoot_vocab_size, config.n_embd) self.ln_hoot = LayerNorm(config.n_embd) @property def pad_token(self): return torch.tensor([self.config.hoot_pad_token]) def forward(self, x): x = self.wte_hoot(x) x = self.ln_hoot(x) return x def init_weights(self): nn.init.normal_(self.wte_hoot.weight) nn.init.normal_(self.ln_hoot.weight) class TextInput(IOModule): def __init__(self, config: GPTConfig, layer: int = 0): super().__init__(config, layer=layer) self.wte_text = nn.Embedding(config.text_vocab_size, config.n_embd) self.ln_text = LayerNorm(config.n_embd) @property def pad_token(self): return torch.tensor([self.config.text_pad_token]) def forward(self, x): x = self.wte_text(x) x = self.ln_text(x) return x def init_weights(self): nn.init.normal_(self.wte_text.weight) class TextOutput(IOModule): def __init__(self, config: GPTConfig, layer: int = -1): super().__init__(config, layer=layer) self.lm_heads = nn.Linear(config.n_embd, config.text_vocab_size, bias=False) @property def pad_token(self): return torch.tensor([self.config.text_pad_token]) def forward(self, x): return self.lm_heads(x) def loss(self, x, y, mask=None): loss = F.cross_entropy( x.reshape(-1, x.size(-1)), y.reshape(-1), reduction="none", ) if mask is not None: loss = loss * mask.reshape(-1) return loss def init_weights(self): final_out_std = self.config.n_embd**-0.5 cutoff_factor = 3 nn.init.trunc_normal_( self.lm_heads.weight, mean=0.0, std=final_out_std, a=-cutoff_factor * final_out_std, b=cutoff_factor * final_out_std, ) class SemanticInput(IOModule): def __init__(self, config: GPTConfig, layer: int = 0): super().__init__(config, layer=layer) self.wte_semantic = nn.Embedding(config.semantic_vocab_size, config.n_embd) self.ln_semantic = LayerNorm(config.n_embd) @property def pad_token(self): return torch.tensor([self.config.semantic_pad_token]) def forward(self, x): x = self.wte_semantic(x) x = self.ln_semantic(x) return x def init_weights(self): nn.init.normal_(self.wte_semantic.weight) class ContinuousSemanticInput(IOModule): def __init__(self, config: GPTConfig, layer: int = 0): super().__init__(config, layer=layer) self.wte_semantic = nn.Linear(768, config.n_embd) self.ln_semantic = LayerNorm(config.n_embd) @property def pad_token(self): return torch.zeros(768, dtype=torch.bfloat16).unsqueeze(0) def forward(self, x): x = self.wte_semantic(x) x = self.ln_semantic(x) return x def init_weights(self): nn.init.normal_(self.wte_semantic.weight) nn.init.normal_(self.ln_semantic.weight) class SemanticOutput(IOModule): def __init__(self, config: GPTConfig, layer: int = -1): super().__init__(config, layer=layer) self.lm_head = nn.Linear(config.n_embd, config.semantic_vocab_size, bias=False) @property def pad_token(self): return torch.tensor([self.config.semantic_pad_token]) def forward(self, x): return self.lm_head(x) def init_weights(self): final_out_std = self.config.n_embd**-0.5 cutoff_factor = 3 nn.init.trunc_normal_( self.lm_head.weight, mean=0.0, std=final_out_std, a=-cutoff_factor * final_out_std, b=cutoff_factor * final_out_std, ) def loss(self, x, y, mask=None): loss = F.cross_entropy( x.reshape(-1, x.size(-1)), y.reshape(-1), reduction="none", ) if mask is not None: loss = loss * mask.reshape(-1) return loss class ContinuousSemanticOutput(IOModule): def __init__(self, config: GPTConfig, layer: int = -1): super().__init__(config, layer=layer) self.proj = nn.Linear(config.n_embd, 768) # Output to 768-dim continuous semantic space @property def pad_token(self): return torch.zeros(768, dtype=torch.bfloat16).unsqueeze(0) def forward(self, x): return self.proj(x) def init_weights(self): nn.init.normal_(self.proj.weight, std=0.01) def loss(self, x, y, mask=None): x = x.squeeze(0) # remove extra batch dimension assert x.shape == y.shape, f"Shape mismatch: x={x.shape}, y={y.shape}" loss = F.mse_loss(x, y, reduction="none") loss = loss.mean(dim=-1) # Average over feature dimension if mask is not None: loss = loss * mask.reshape(-1) return loss class RewardHead(nn.Module): """Reward head for preference modeling and RLHF. Takes hidden states and outputs a scalar reward value. Uses layer normalization for stability. """ def __init__(self, config: GPTConfig): super().__init__() self.config = config self.ln = LayerNorm(config.n_embd) self.reward_proj = nn.Linear(config.n_embd, 1, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: """Project hidden states to scalar reward. Args: x: Hidden states of shape (batch, seq_len, n_embd) or (batch, n_embd) Returns: Scalar reward of shape (batch, seq_len, 1) or (batch, 1) """ x = self.ln(x) reward = self.reward_proj(x) return reward def init_weights(self): """Initialize with small weights to start near zero rewards.""" nn.init.normal_(self.ln.weight, mean=1.0, std=0.02) nn.init.normal_(self.reward_proj.weight, mean=0.0, std=0.01) class RepaContinuousSemanticOutput(IOModule): """2-layer MLP for repa semantic head with layer-specific target""" def __init__(self, config: GPTConfig): super().__init__(config, layer=config.repa_layer) # 2-layer MLP: n_embd -> hidden -> 768 hidden_dim = config.n_embd # Use same dimension as input self.mlp = nn.Sequential( nn.Linear(config.n_embd, hidden_dim), nn.SiLU(), # Use same activation as in the config nn.Linear(hidden_dim, 768), ) @property def pad_token(self): return torch.zeros(768, dtype=torch.bfloat16).unsqueeze(0) def forward(self, x): return self.mlp(x) def init_weights(self): for module in self.mlp.modules(): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, std=0.01) if module.bias is not None: nn.init.zeros_(module.bias) def loss(self, x, y, mask=None): x = x.squeeze(0) # remove extra batch dimension assert x.shape == y.shape, f"Shape mismatch: x={x.shape}, y={y.shape}" loss = F.mse_loss(x, y, reduction="none") loss = loss.mean(dim=-1) # Average over feature dimension if mask is not None: loss = loss * mask.reshape(-1) return loss class RepaMixedSemanticOutput(IOModule): """2-layer MLP for repa mixed semantic head with layer-specific target (for stem conditioning)""" def __init__(self, config: GPTConfig): super().__init__(config, layer=config.repa_layer) # 2-layer MLP: n_embd -> hidden -> 768 hidden_dim = config.n_embd # Use same dimension as input self.mlp = nn.Sequential( nn.Linear(config.n_embd, hidden_dim), nn.SiLU(), # Use same activation as in the config nn.Linear(hidden_dim, 768), ) @property def pad_token(self): return torch.zeros(768, dtype=torch.bfloat16).unsqueeze(0) def forward(self, x): return self.mlp(x) def init_weights(self): for module in self.mlp.modules(): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, std=0.01) if module.bias is not None: nn.init.zeros_(module.bias) def loss(self, x, y, mask=None): """MSE loss for continuous mix latent prediction.""" x = x.squeeze(0) # remove extra batch dimension assert x.shape == y.shape, f"Shape mismatch: x={x.shape}, y={y.shape}" loss = F.mse_loss(x, y, reduction="none") loss = loss.mean(dim=-1) # Average over feature dimension if mask is not None: loss = loss * mask.reshape(-1) return loss class RepaHootOutput(IOModule): """2-layer MLP for repa hoot head with layer-specific target""" def __init__(self, config: GPTConfig): super().__init__(config, layer=config.repa_layer) # 2-layer MLP: n_embd -> hidden -> 512 (hoot d_model) hidden_dim = config.n_embd # Use same dimension as input self.mlp = nn.Sequential( nn.Linear(config.n_embd, hidden_dim), nn.SiLU(), # Use same activation as in the config nn.Linear(hidden_dim, 512), # Hoot encoder d_model = 512 ) @property def pad_token(self): return torch.zeros(512, dtype=torch.bfloat16).unsqueeze(0) def forward(self, x): return self.mlp(x) def init_weights(self): for module in self.mlp.modules(): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, std=0.01) if module.bias is not None: nn.init.zeros_(module.bias) def loss(self, x, y, mask=None): """MSE loss for continuous hoot embedding prediction.""" x = x.squeeze(0) # remove extra batch dimension assert x.shape == y.shape, f"Shape mismatch: x={x.shape}, y={y.shape}" loss = F.mse_loss(x, y, reduction="none") loss = loss.mean(dim=-1) # Average over feature dimension if mask is not None: loss = loss * mask.reshape(-1) return loss class RepaMidiOutput(IOModule): """2-layer MLP for repa midi head with layer-specific target""" def __init__(self, config: GPTConfig): super().__init__(config, layer=config.repa_layer) # 2-layer MLP: n_embd -> hidden -> 1024 (midi model d_model) hidden_dim = config.n_embd # Use same dimension as input self.mlp = nn.Sequential( nn.Linear(config.n_embd, hidden_dim), nn.SiLU(), # Use same activation as in the config nn.Linear(hidden_dim, 1024), # MIDI model d_model = 1024 (16 heads * 64 d_head) ) @property def pad_token(self): return torch.zeros(1024, dtype=torch.bfloat16).unsqueeze(0) def forward(self, x): return self.mlp(x) def init_weights(self): for module in self.mlp.modules(): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, std=0.01) if module.bias is not None: nn.init.zeros_(module.bias) def loss(self, x, y, mask=None): """MSE loss for continuous midi embedding prediction.""" x = x.squeeze(0) # remove extra batch dimension assert x.shape == y.shape, f"Shape mismatch: x={x.shape}, y={y.shape}" loss = F.mse_loss(x, y, reduction="none") loss = loss.mean(dim=-1) # Average over feature dimension if mask is not None: loss = loss * mask.reshape(-1) return loss class VaeInput(IOModule): def __init__(self, config: GPTConfig, layer: int = 0): super().__init__(config, layer=layer) self.proj = nn.Linear(config.vae_embed_dim, config.n_embd) self.ln = LayerNorm(config.n_embd) @property def pad_token(self): return torch.zeros(self.config.vae_embed_dim, dtype=torch.bfloat16).unsqueeze(0) def forward(self, x): x = self.proj(x) x = self.ln(x) return x def init_weights(self): nn.init.normal_(self.proj.weight) nn.init.normal_(self.ln.weight) class VaeOutput(IOModule): def __init__(self, config: GPTConfig, layer: int = -1): super().__init__(config, layer=layer) self.proj = nn.Linear(config.n_embd, config.vae_embed_dim) @property def pad_token(self): return torch.zeros(self.config.vae_embed_dim, dtype=torch.bfloat16).unsqueeze(0) def forward(self, x): return self.proj(x) def init_weights(self): nn.init.normal_(self.proj.weight, std=0.01) def loss(self, x, y, mask=None): x = x.squeeze(0) # remove extra batch dimension assert x.shape == y.shape loss = F.mse_loss(x, y, reduction="none") loss = loss.mean(dim=-1) # Average over feature dimension if mask is not None: loss = loss * mask.reshape(-1) return loss class TimestepInput(IOModule): def __init__(self, config: GPTConfig, layer: int = 0): super().__init__(config, layer=layer) self.timestep_feature_dim = 256 self.weight = nn.Parameter(torch.zeros(self.timestep_feature_dim // 2, 1)) self.in_proj = nn.Linear(self.timestep_feature_dim, config.n_embd, bias=True) self.out_proj = nn.Linear(config.n_embd, config.n_embd, bias=True) @property def pad_token(self): return torch.zeros(1, dtype=torch.bfloat16).unsqueeze(-1) def forward(self, input): f = 2 * math.pi * input @ self.weight.T x = torch.cat([f.cos(), f.sin()], dim=-1) x = self.in_proj(x) x = F.silu(x) x = self.out_proj(x) return x def init_weights(self): pass class DittoInput(IOModule): def __init__(self, config: GPTConfig, layer: int = 0): super().__init__(config, layer=layer) self.proj = nn.Linear(config.ditto_embed_dim, config.n_embd) @property def pad_token(self): return torch.zeros(self.config.ditto_embed_dim, dtype=torch.bfloat16).unsqueeze(0) def forward(self, x): x = self.proj(x) return x def init_weights(self): nn.init.normal_(self.proj.weight) class BlockTypeInput(IOModule): """Block type input module that provides embeddings for block types""" def __init__(self, config: GPTConfig, layer: int = 0): super().__init__(config, layer=layer) self.wte_block_type = nn.Embedding(config.block_type_vocab_size, config.n_embd) @property def pad_token(self): return torch.tensor([0]) # Default to block type 0 (text) def forward(self, block_type_ids): """ Args: block_type_ids: tensor of shape (seq_len,) with block type indices Returns: block_type_embeddings: tensor of shape (seq_len, n_embd) """ return self.wte_block_type(block_type_ids) def init_weights(self): nn.init.normal_(self.wte_block_type.weight, std=0.02) class GPT(nn.Module): def __init__(self, config: GPTConfig, train_config: GPTTrainConfig): super().__init__() self.config = config self.train_config = train_config model_dict = dict( drop=nn.Dropout(0.0), h=nn.ModuleList( [Block(config, train_config, layer_id) for layer_id in range(config.n_layer)] ), ln_f=LayerNorm(config.n_embd), ) # Input and output modules input_modules = { "text_input": TextInput(config), } if config.use_continuous_semantic_input: input_modules["continuous_semantic_input"] = ContinuousSemanticInput(config) if config.use_mmbert: input_modules["mmbert_text_input"] = MMBertTextInput(config) if config.use_vae_input: input_modules["vae_input"] = VaeInput(config) if config.output_paradigm == "diffusion": input_modules["timestep_input"] = TimestepInput(config) if config.use_hoot: input_modules["hoot_input"] = HootTextInput(config) if config.use_ditto: input_modules["ditto_input"] = DittoInput(config) if config.use_block_type_embeddings: input_modules["block_type_input"] = BlockTypeInput(config) # Output modules based on output_paradigm and output_distribution output_modules = {} # Add text output if text loss is enabled if config.use_text_loss: output_modules["text_output"] = TextOutput(config) # Add output heads based on paradigm and distribution if config.output_paradigm == "gpt": # GPT paradigm: autoregressive generation if config.output_distribution == "semantic": # Discrete semantic tokens - create per-codebook modules for n in range(config.semantic_n_codebooks): input_modules[f"semantic_input_{n}"] = SemanticInput(config) output_modules[f"semantic_output_{n}"] = SemanticOutput(config) # Add auxiliary continuous semantic output for layer-specific target (repa) if config.repa_semantic: output_modules["repa_semantic_output"] = RepaContinuousSemanticOutput(config) elif config.output_distribution == "continuous_semantic": # Continuous semantic embeddings input_modules["continuous_semantic_input"] = ContinuousSemanticInput(config) output_modules["continuous_semantic_output"] = ContinuousSemanticOutput(config) elif config.output_distribution == "vae": # VAE latents (continuous) input_modules["vae_input"] = VaeInput(config) output_modules["vae_output"] = VaeOutput(config) else: raise ValueError( f"Unknown output_distribution for GPT paradigm: {config.output_distribution}" ) elif config.output_paradigm == "diffusion": # Diffusion paradigm: denoising if config.output_distribution == "vae": # VAE latents (standard for diffusion) input_modules["vae_input"] = VaeInput(config) output_modules["vae_output"] = VaeOutput(config) elif config.output_distribution == "continuous_semantic": # Continuous semantic (alternative diffusion target) input_modules["continuous_semantic_input"] = ContinuousSemanticInput(config) output_modules["continuous_semantic_output"] = ContinuousSemanticOutput(config) else: raise ValueError( f"output_distribution '{config.output_distribution}' not supported for diffusion paradigm. Use 'vae' or 'continuous_semantic'." ) else: raise ValueError(f"Unknown output_paradigm: {config.output_paradigm}") # Add reward head if enabled (for reward modeling / RLHF) if config.use_reward_head: output_modules["reward_head"] = RewardHead(config) # Add auxiliary mixed semantic output for stem conditioning if enabled (repa) if config.repa_mixed_semantic: output_modules["repa_mixed_semantic_output"] = RepaMixedSemanticOutput(config) # Add auxiliary hoot output if enabled (repa) if config.repa_hoot: output_modules["repa_hoot_output"] = RepaHootOutput(config) # Add auxiliary midi output if enabled (repa) if config.repa_midi: output_modules["repa_midi_output"] = RepaMidiOutput(config) self.input_modules = nn.ModuleDict(input_modules) self.output_modules = nn.ModuleDict(output_modules) if config.coarse_n_codebooks > 0: model_dict["wte_coarse"] = nn.ModuleList( [ nn.Embedding(config.coarse_vocab_size, config.n_embd) for _ in range(config.coarse_n_codebooks) ] ) model_dict["ln_coarse"] = LayerNorm(config.n_embd) if self.config.use_learned_pos_emb: model_dict["wpe"] = nn.Embedding(config.block_size, config.n_embd) self.transformer = nn.ModuleDict(model_dict) # init all weights if train_config.layer_init: self.init_weights() else: self.apply(self._init_weights_simple) for pn, p in self.named_parameters(): if pn.endswith("c_proj.weight"): torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer)) print_with_time_master(f"number of parameters: {self.get_num_params() / 1e6:.0f}M") if config.use_rotary_pos_emb: # pip install "git+https://github.com/Dao-AILab/flash-attention.git#subdirectory=csrc/rotary from .rotary import RotaryEmbedding self.rope_module = RotaryEmbedding(config.d_head, base=config.rope_theta) def _resolve_layer(self, value: int) -> int: """Convert layer index to absolute position. None or -1 means final layer.""" n_layers = self.config.n_layer if value < 0: return n_layers + value return value def init_weights(self): """ [Note: On ``init_weights`` vs. ``reset_parameters``] Modules may define ``reset_parameters`` to initialize parameter values. ``reset_parameters`` is meant to only initialize directly owned parameters/buffers, not those of their child modules, and it can be used to give the initial values for these tensors. Separately, users may want custom initialization for their modules, different from that in ``reset_parameters``. For this, we define ``init_weights``. We only call it in the constructor of this ``Transformer`` root module to avoid reinitializing tensors. """ for module in self.input_modules.values(): module.init_weights() for module in self.output_modules.values(): module.init_weights() for layer in self.transformer.h: layer.init_weights() def _build_rope_cache(self, *, seq_lens, seqlen, device, dtype): if not self.config.use_rotary_pos_emb or not hasattr(self, "rope_module"): return None if seq_lens is not None: self.rope_module._update_varlen_cos_sin_cache(seq_lens, device=device, dtype=dtype) else: self.rope_module._update_cos_sin_cache(seqlen, device=device, dtype=dtype) return ( self.rope_module._cos_cached, self.rope_module._sin_cached, getattr(self.rope_module, "_cos_k_cached", None), getattr(self.rope_module, "_sin_k_cached", None), self.rope_module.interleaved, ) def forward_old( self, x, y=None, seq_lens: list[int] = None, return_logits=False, last_only=True, ): """This is deprecated, use forward instead. Only for DPO for now.""" device = x.device b, ns, t = x.size() assert ns == 1 + self.config.semantic_n_codebooks + self.config.coarse_n_codebooks # Compute sequence info once if seq_lens is None: seq_lens = [t] cu_seqlen = torch.tensor([0] + seq_lens, device=device, dtype=torch.int32) cu_seqlen = torch.cumsum(cu_seqlen, dim=0, dtype=torch.int32) # embed text x_emb = self.input_modules["text_input"](x[:, 0, :]) # embed semantic for n in range(self.config.semantic_n_codebooks): x_emb = x_emb + self.input_modules[f"semantic_input_{n}"](x[:, 1 + n, :]) # embed coarse if self.config.coarse_n_codebooks > 0: for n in range(self.config.coarse_n_codebooks): n2 = 1 + n + self.config.semantic_n_codebooks x_emb = x_emb + self.transformer.ln_coarse(self.transformer.wte_coarse[n](x[:, n2, :])) # embed ditto if self.config.use_ditto: ditto_tokens = self.input_modules["ditto_input"].pad_token.to(device) x_emb = x_emb + self.input_modules["ditto_input"](ditto_tokens) # embed mmbert if self.config.use_mmbert: mmbert_tokens = self.input_modules["mmbert_text_input"].pad_token.to(device) x_emb = x_emb + self.input_modules["mmbert_text_input"](mmbert_tokens) # embed hoot if self.config.use_hoot: hoot_tokens = self.input_modules["hoot_input"].pad_token.to(device) x_emb = x_emb + self.input_modules["hoot_input"](hoot_tokens) if self.config.use_vae_input: vae_tokens = self.input_modules["vae_input"].pad_token.to(device) x_emb = x_emb + self.input_modules["vae_input"](vae_tokens) # embed timestep if self.config.output_paradigm == "diffusion": timestep_tokens = self.input_modules["timestep_input"].pad_token.to(device) x_emb = x_emb + self.input_modules["timestep_input"](timestep_tokens) # x_emb (b, t, n_embd) if self.config.use_learned_pos_emb: # Create position indices for each sequence # torch dataloader can turn ints into tensors so we handle both pos = torch.cat( [ torch.arange(seq_len if isinstance(seq_len, int) else seq_len.item(), device=device) for seq_len in seq_lens ] ) assert len(pos) == t pos_emb = self.transformer.wpe(pos) # (1, t, n_embd) x_emb = x_emb + pos_emb x = self.transformer.drop(x_emb) * self.config.embed_scale_factor rope_cache = None if self.config.use_rotary_pos_emb: rope_cache = self._build_rope_cache( seq_lens=seq_lens, seqlen=t, device=device, dtype=x.dtype ) max_seqlen = max(seq_lens) cu_seqlen = torch.tensor([0] + seq_lens, device=x.device).cumsum(dim=0).to(torch.int32) for block in self.transformer.h: x = block(x, rope_cache=rope_cache, cu_seqlen=cu_seqlen, causal=True, max_seqlen=max_seqlen) x = self.transformer.ln_f(x) # Only remove last token for standard training, not reward model if not self.config.use_reward_head: x = x[:, :-1] # remove last token for next-token prediction noffs = 1 if self.config.use_text_loss else 0 if y is not None: assert ( y.shape[1] == noffs + self.config.semantic_n_codebooks + self.config.coarse_n_codebooks ) if return_logits: # For reward model, return dict with BOTH rewards and generation logits # This allows model to be used for both generation and reward evaluation if self.config.use_reward_head: # Compute rewards reward_logits = self.output_modules["reward_head"](x) # (batch, seq_len, 1) reward_logits = reward_logits.squeeze(-1) # (batch, seq_len) # Also compute generation logits (for music generation) if last_only: x_for_gen = x[:, -1, :] else: x_for_gen = x text_logits = None if self.config.use_text_loss: text_logits = self.output_modules["text_output"](x_for_gen) semantic_logits_list = [] for n in range(self.config.semantic_n_codebooks): semantic_logits_list.append(self.output_modules[f"semantic_output_{n}"](x_for_gen)) semantic_logits = torch.stack(semantic_logits_list).swapaxes(0, 1) coarse_logits = None if self.config.coarse_n_codebooks > 0: coarse_logits_list = [] for n in range(self.config.coarse_n_codebooks): n2 = n + self.config.semantic_n_codebooks coarse_logits_list.append(self.lm_heads[n2 + noffs](x_for_gen)) coarse_logits = torch.stack(coarse_logits_list).swapaxes(0, 1) # Return dict with both rewards and logits return { "reward_logits": reward_logits, "text_logits": text_logits, "semantic_logits": semantic_logits, "coarse_logits": coarse_logits, } # Standard path (no reward head) if last_only: x = x[:, -1, :] text_logits = None if self.config.use_text_loss: text_logits = self.output_modules["text_output"](x) # Semantic logits per codebook semantic_logits_list = [ self.output_modules[f"semantic_output_{n}"](x) for n in range(self.config.semantic_n_codebooks) ] semantic_logits = torch.stack(semantic_logits_list).swapaxes(0, 1) coarse_logits = None if self.config.coarse_n_codebooks > 0: coarse_logits_list = [] for n in range(self.config.coarse_n_codebooks): n2 = n + self.config.semantic_n_codebooks coarse_logits_list.append(self.lm_heads[n2 + noffs](x)) coarse_logits = torch.stack(coarse_logits_list).swapaxes(0, 1) return text_logits, semantic_logits, coarse_logits loss_dict = {} if Z_LOSS: loss_dict["z_loss"] = 0 if self.config.use_text_loss: logits = self.output_modules["text_output"](x) loss_dict["text"] = F.cross_entropy( logits.reshape(-1, logits.size(-1)), y[:, 0, :].reshape(-1), ignore_index=-1, ) if Z_LOSS: loss_dict["z_loss"] += (torch.logsumexp(logits, dim=-1) ** 2).mean() # Semantic loss per codebook for n in range(self.config.semantic_n_codebooks): logits = self.output_modules[f"semantic_output_{n}"](x) loss_dict[f"semantic_output_{n}"] = F.cross_entropy( logits.reshape(-1, logits.size(-1)), y[:, n + noffs, :].reshape(-1), ignore_index=-1, ) if Z_LOSS: loss_dict["z_loss"] += (torch.logsumexp(logits, dim=-1) ** 2).mean() if self.config.coarse_n_codebooks > 0: for n in range(self.config.coarse_n_codebooks): n2 = n + self.config.semantic_n_codebooks logits = self.lm_heads[n2 + noffs](x) loss_dict[f"coarse_{n}"] = F.cross_entropy( logits.reshape(-1, logits.size(-1)), y[:, n2 + noffs, :].reshape(-1), ignore_index=-1, ) if Z_LOSS: loss_dict["z_loss"] += (torch.logsumexp(logits, dim=-1) ** 2).mean() return loss_dict def forward( self, packed_sequences: PackedBlockSequence | torch.Tensor, return_logits=False, last_only=True, device="cuda", y=None, seq_lens: list[int] = None, return_block_losses=False, ): if isinstance(packed_sequences, torch.Tensor): # old return self.forward_old( packed_sequences, y=y, seq_lens=seq_lens, return_logits=return_logits, last_only=last_only, ) t = packed_sequences.n_tokens # Initialize tensors to track document, block IDs and causality document_id = torch.full((t,), -1, device=device, dtype=torch.int32) block_id = torch.full((t,), -1, device=device, dtype=torch.int32) block_is_causal = torch.full((t,), False, device=device, dtype=torch.bool) # Tracking variables b_id = 0 pos = 0 seq_lens = [] x = [] # Input tensors y = [] # Target tensors # Process each block sequence (document) for doc_id, sample in enumerate(packed_sequences): # Calculate total length of all blocks in this sequence doc_length = sum(len(block) for block in sample) # Assign document ID to all positions in this sequence document_id[pos : pos + doc_length] = doc_id seq_lens.append(doc_length) # Process each block within the sequence for block in sample: block_length = len(block) # Assign block ID and causality flag block_id[pos : pos + block_length] = b_id block_is_causal[pos : pos + block_length] = block.spec.is_causal if block.spec.chunk_size is not None: # split into block causal chunks for i in range(0, block_length, block.spec.chunk_size): block_id[pos + i : pos + i + block.spec.chunk_size] = b_id # block_is_causal[pos + i : pos + i + block.spec.chunk_size] = False b_id += 1 # Update position and block ID pos += block_length b_id += 1 # Check that all block inputs have corresponding input modules for sample in packed_sequences: for block in sample: for input_name in block.inputs: if input_name not in self.input_modules: raise ValueError(f"Block input '{input_name}' has no corresponding input module") # process each input type x_emb = 0 layer_injections = defaultdict(float) for input_name, input_module in self.input_modules.items(): inputs = [] for sample in packed_sequences: for block in sample: if input_name in block.inputs: inputs.append(block.inputs[input_name]) else: if input_module.pad_token.dim() <= 1: inputs.append(input_module.pad_token.repeat(len(block))) elif input_module.pad_token.dim() == 2: inputs.append(input_module.pad_token.repeat(len(block), 1)) else: raise ValueError( f"Invalid pad token dimension: {input_module.pad_token.dim()}" ) inputs[-1] = inputs[-1].to(device) if input_name == "semantic_input" and inputs[-1].dtype != torch.int64: print(inputs[-1].dtype) print(block) # check all inputs have the same dtype assert inputs[0].dtype == inputs[-1].dtype, ( input_name, inputs[0].dtype, inputs[-1].dtype, block, ) inputs = torch.cat(inputs, dim=0) module_value = input_module(inputs) target_layer = self._resolve_layer(input_module.layer) if target_layer == 0: x_emb = x_emb + module_value else: layer_injections[target_layer] += module_value.unsqueeze(0) x_emb = x_emb.unsqueeze(0) assert x_emb.shape == (1, t, self.config.n_embd), (x_emb.shape, t, self.config.n_embd) is_causal = None if block_is_causal.all(): # this will use flash attention for speed block_mask_local = None block_mask_global = None is_causal = True elif all(len(seq) == 1 and not seq[0].spec.is_causal for seq in packed_sequences): is_causal = False block_mask_local = None block_mask_global = None else: causal_f = partial( causal, document_id=document_id, block_id=block_id, block_is_causal=block_is_causal, window_size=self.config.attention_sliding_window_size, ) causal_f_global = partial( causal, document_id=document_id, block_id=block_id, block_is_causal=block_is_causal, window_size=-1, ) block_mask_local = create_block_mask( causal_f, B=None, H=None, Q_LEN=t, KV_LEN=t, _compile=True, device=device ) block_mask_global = create_block_mask( causal_f_global, B=None, H=None, Q_LEN=t, KV_LEN=t, _compile=True, device=device ) # embed coarse (keeping original code since there's no coarse_input in io_dict) if self.config.coarse_n_codebooks > 0: for n in range(self.config.coarse_n_codebooks): n2 = 1 + n + self.config.semantic_n_codebooks x_emb = x_emb + self.transformer.ln_coarse(self.transformer.wte_coarse[n](x[:, n2, :])) # x_emb (b, t, n_embd) if self.config.use_learned_pos_emb: # Create position indices for each sequence # torch dataloader can turn ints into tensors so we handle both pos = torch.cat( [ torch.arange(seq_len if isinstance(seq_len, int) else seq_len.item(), device=device) for seq_len in seq_lens ] ) assert len(pos) == t pos_emb = self.transformer.wpe(pos) # (1, t, n_embd) x_emb = x_emb + pos_emb x = self.transformer.drop(x_emb) * self.config.embed_scale_factor rope_cache = None if self.config.use_rotary_pos_emb: rope_cache = self._build_rope_cache( seq_lens=seq_lens, seqlen=t, device=device, dtype=x.dtype ) # Determine which layers we need to store (for memory efficiency) needed_layers = set() n_layers = len(self.transformer.h) for output_module in self.output_modules.values(): target_layer = self._resolve_layer(output_module.layer) needed_layers.add(target_layer) # Store only needed intermediate layer outputs for layer-specific targets layer_outputs = [None] * n_layers max_seqlen = max(seq_lens) cu_seqlen = torch.tensor([0] + seq_lens, device=x.device).cumsum(dim=0).to(torch.int32) for i, block in enumerate(self.transformer.h): if i in layer_injections: x = x + layer_injections[i] * self.config.embed_scale_factor block_mask = ( block_mask_local if i % self.config.global_every_n_layers != 0 else block_mask_global ) x = block( x, rope_cache=rope_cache, cu_seqlen=cu_seqlen, block_mask=block_mask, causal=is_causal, max_seqlen=max_seqlen, ) # Store this layer if needed if i in needed_layers: layer_outputs[i] = self.transformer.ln_f(x) if return_logits: # For reward model, return dict with BOTH rewards and generation logits if self.config.use_reward_head: # Compute rewards reward_logits = self.output_modules["reward_head"](x) # (batch, seq_len, 1) reward_logits = reward_logits.squeeze(-1) # (batch, seq_len) # Also compute generation logits if last_only: x_for_gen = x[:, -1, :] else: x_for_gen = x text_logits = None if self.config.use_text_loss: text_logits = self.output_modules["text_output"](x_for_gen) semantic_logits_list = [ self.output_modules[f"semantic_output_{n}"](x_for_gen) for n in range(self.config.semantic_n_codebooks) ] semantic_logits = torch.stack(semantic_logits_list, dim=-1) coarse_logits = None # Return dict with both rewards and logits return { "reward_logits": reward_logits, "text_logits": text_logits, "semantic_logits": semantic_logits, "coarse_logits": coarse_logits, } # Standard path (no reward head) if last_only: x = layer_outputs[-1] x = x[:, -1, :] text_logits = None if self.config.use_text_loss: text_logits = self.output_modules["text_output"](x) # Semantic logits per codebook semantic_logits_list = [ self.output_modules[f"semantic_output_{n}"](x) for n in range(self.config.semantic_n_codebooks) ] semantic_logits = torch.stack(semantic_logits_list, dim=-1) coarse_logits = None return text_logits, semantic_logits, coarse_logits loss_dict = {} per_block_sequence_losses = [] # List of tuples (block_sequence_idx, block_type, loss_value) if Z_LOSS: loss_dict["z_loss"] = 0 for output_name, output_module in self.output_modules.items(): tgts = [] block_is_pad = [] for sample in packed_sequences: for block in sample: if output_name in block.targets: tgts.append(block.targets[output_name]) block_is_pad += [False] * len(block) else: block_is_pad += [True] * len(block) if output_module.pad_token.dim() <= 1: tgts.append(output_module.pad_token.repeat(len(block))) elif output_module.pad_token.dim() == 2: tgts.append(output_module.pad_token.repeat(len(block), 1)) else: raise ValueError( f"Invalid pad token dimension: {output_module.pad_token.dim()}" ) tgts = torch.cat(tgts, dim=0) tgts = tgts.to(device) block_is_pad = torch.tensor(block_is_pad, device=device, dtype=torch.bool) if len(tgts) != len(block_is_pad): print( f"DEBUG: output_name={output_name}, len(tgts)={len(tgts)}, len(block_is_pad)={len(block_is_pad)}" ) print(f"DEBUG: tgts.shape={tgts.shape}") for i, sample in enumerate(packed_sequences): for j, block in enumerate(sample): if output_name in block.targets: target_shape = block.targets[output_name].shape print( f"DEBUG: sample {i}, block {j}, target shape={target_shape}, block len={len(block)}" ) assert len(tgts) == len(block_is_pad) # Use layer-specific output if specified, otherwise use final layer target_layer = self._resolve_layer(output_module.layer) x_for_output = layer_outputs[target_layer] y = output_module(x_for_output) # Get unreduced losses and slice them unreduced_loss = output_module.loss(y, tgts, mask=~block_is_pad) # Compute regular aggregated loss if (~block_is_pad).sum() > 0: loss_dict[output_name] = unreduced_loss[~block_is_pad].mean() else: loss_dict[output_name] = torch.zeros_like(unreduced_loss[0]) # If we need block losses, slice the unreduced loss if return_block_losses: current_pos = 0 for seq_idx, sample in enumerate(packed_sequences): for block in sample: block_len = len(block) if output_name in block.targets: # Get mask for this block block_mask = ~block_is_pad[current_pos : current_pos + block_len] if block_mask.sum() > 0: # Get loss for this block block_loss = unreduced_loss[current_pos : current_pos + block_len][ block_mask ].mean() per_block_sequence_losses.append( (sample, block.spec.name, block_loss.item()) ) current_pos += block_len # Z_LOSS only applies to discrete outputs (with logits), not continuous outputs if Z_LOSS and output_name not in [ "continuous_semantic_output", "repa_semantic_output", "repa_mixed_semantic_output", "vae_output", ]: loss_dict["z_loss"] += (torch.logsumexp(y, dim=-1) ** 2).mean() if self.config.coarse_n_codebooks > 0: raise NotImplementedError("coarse loss not implemented") # Add tiny regularization for repa modules that didn't receive gradients # This ensures all parameters get gradients on every batch, preventing # optimizer state inconsistency across distributed ranks repa_module_names = [ "repa_semantic_output", "repa_mixed_semantic_output", "repa_hoot_output", "repa_midi_output", ] repa_reg = 0.0 for module_name in repa_module_names: if module_name in self.output_modules: # Add tiny L2 penalty (1e-10 is negligible for training but ensures gradients) for param in self.output_modules[module_name].parameters(): if param.requires_grad: repa_reg = repa_reg + (param**2).sum() loss_dict["repa_reg"] = 1e-10 * repa_reg if return_block_losses: return loss_dict, per_block_sequence_losses return loss_dict def get_num_params(self, non_embedding=True): n_params = sum(p.numel() for p in self.parameters()) if non_embedding: for module in self.input_modules.values(): n_params -= sum(p.numel() for p in module.parameters()) if self.config.use_learned_pos_emb: n_params -= self.transformer.wpe.weight.numel() return n_params def _init_weights_simple(self, module): init_weights_simple(self, module) def configure_optimizers(self, weight_decay, learning_rate, betas, device_type): return configure_optimizers(self, weight_decay, learning_rate, betas, device_type) def estimate_mfu(self, fwdbwd_per_iter, dt): return estimate_mfu(self, fwdbwd_per_iter, dt)