app/torchtune_models.py

"""Torchtune models for CSM-1B."""
import logging
from dataclasses import dataclass
import torch
import torch.nn as nn

# Set up logging
logger = logging.getLogger(__name__)

# First, try to import llama3_2 from torchtune directly
try:
    import torchtune
    logger.info(f"Torchtune version: {getattr(torchtune, '__version__', 'unknown')}")
    
    # Print available modules in torchtune.models
    try:
        import torchtune.models
        logger.info(f"Available modules in torchtune.models: {dir(torchtune.models)}")
    except Exception as e:
        logger.error(f"Error inspecting torchtune.models: {e}")
    
    # Try to import llama3_2 model
    try:
        from torchtune.models.llama3_2 import llama3_2
        logger.info("Successfully imported llama3_2 from torchtune")
    except ImportError as e:
        logger.warning(f"Could not import llama3_2: {e}")
        # Try to import regular llama as fallback
        try:
            from torchtune.models.llama import llama
            logger.info("Using llama from torchtune.models.llama as fallback")
            llama3_2 = llama  # Alias llama as llama3_2
        except ImportError:
            logger.error("Could not import llama model either. Will use custom implementation.")
            llama3_2 = None
except ImportError as e:
    logger.error(f"Torchtune not available: {e}")
    torchtune = None
    llama3_2 = None


# Define our own model implementations as fallbacks
def llama3_2_1B_custom():
    """Create a Llama 3.2 1B model."""
    from app.custom_transformer import CustomTransformerDecoder
    return CustomTransformerDecoder(
        vocab_size=128_256,
        num_layers=16,
        num_heads=32,
        num_kv_heads=8,
        embed_dim=2048,
        max_seq_len=2048,
        intermediate_dim=8192,
        attn_dropout=0.0,
        norm_eps=1e-5,
    )


def llama3_2_100M_custom():
    """Create a Llama 3.2 100M model."""
    from app.custom_transformer import CustomTransformerDecoder
    return CustomTransformerDecoder(
        vocab_size=128_256,
        num_layers=4,
        num_heads=8,
        num_kv_heads=2,
        embed_dim=1024,
        max_seq_len=2048,
        intermediate_dim=8192,
        attn_dropout=0.0,
        norm_eps=1e-5,
    )


# Setup fallback to our own implementations if needed
if llama3_2 is None:
    logger.warning("Using custom implementations for Llama models")
    FLAVORS = {
        "llama-1B": llama3_2_1B_custom,
        "llama-100M": llama3_2_100M_custom,
    }
else:
    logger.info("Using torchtune implementations for Llama models")
    FLAVORS = {
        "llama-1B": lambda: llama3_2(
            vocab_size=128_256,
            num_layers=16,
            num_heads=32,
            num_kv_heads=8,
            embed_dim=2048,
            max_seq_len=2048,
            intermediate_dim=8192,
            attn_dropout=0.0,
            norm_eps=1e-5,
            rope_base=500_000,
            scale_factor=32,
        ),
        "llama-100M": lambda: llama3_2(
            vocab_size=128_256,
            num_layers=4,
            num_heads=8,
            num_kv_heads=2,
            embed_dim=1024,
            max_seq_len=2048,
            intermediate_dim=8192,
            attn_dropout=0.0,
            norm_eps=1e-5,
            rope_base=500_000,
            scale_factor=32,
        ),
    }


def _prepare_transformer(model):
    """Prepare transformer for use."""
    embed_dim = model.tok_embeddings.embedding_dim
    model.tok_embeddings = nn.Identity()
    model.output = nn.Identity()
    return model, embed_dim


def _create_causal_mask(seq_len: int, device: torch.device):
    """Create causal mask."""
    return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device))


def _index_causal_mask(mask: torch.Tensor, input_pos: torch.Tensor):
    """Index causal mask.
    
    Args:
        mask: (max_seq_len, max_seq_len)
        input_pos: (batch_size, seq_len)
    
    Returns:
        (batch_size, seq_len, max_seq_len)
    """
    r = mask[input_pos, :]
    return r


def _multinomial_sample_one_no_sync(probs):
    """Do multinomial sampling without a cuda synchronization."""
    q = torch.empty_like(probs).exponential_(1)
    return torch.argmax(probs / q, dim=-1, keepdim=True).to(dtype=torch.int)


def sample_topk(logits: torch.Tensor, topk: int, temperature: float):
    """Sample from top-k logits."""
    logits = logits / temperature
    filter_value: float = -float("Inf")
    indices_to_remove = logits < torch.topk(logits, topk)[0][..., -1, None]
    scores_processed = logits.masked_fill(indices_to_remove, filter_value)
    scores_processed = torch.nn.functional.log_softmax(scores_processed, dim=-1)
    probs = torch.nn.functional.softmax(scores_processed, dim=-1)
    sample_token = _multinomial_sample_one_no_sync(probs)
    return sample_token


@dataclass
class ModelArgs:
    """Model arguments."""
    backbone_flavor: str
    decoder_flavor: str
    text_vocab_size: int
    audio_vocab_size: int
    audio_num_codebooks: int


class Model(nn.Module):
    """CSM-1B model."""
    
    def __init__(self, args: ModelArgs):
        """Initialize model."""
        super().__init__()
        self.args = args
        logger.info(f"Creating model with backbone: {args.backbone_flavor}, decoder: {args.decoder_flavor}")
        
        # Load backbone and decoder
        self.backbone, backbone_dim = _prepare_transformer(FLAVORS[args.backbone_flavor]())
        self.decoder, decoder_dim = _prepare_transformer(FLAVORS[args.decoder_flavor]())
        
        # Embeddings
        self.text_embeddings = nn.Embedding(args.text_vocab_size, backbone_dim)
        self.audio_embeddings = nn.Embedding(args.audio_vocab_size * args.audio_num_codebooks, backbone_dim)
        
        # Projection and heads
        self.projection = nn.Linear(backbone_dim, decoder_dim, bias=False)
        self.codebook0_head = nn.Linear(backbone_dim, args.audio_vocab_size, bias=False)
        self.audio_head = nn.Parameter(torch.empty(args.audio_num_codebooks - 1, decoder_dim, args.audio_vocab_size))
        
        # Initialize audio head
        nn.init.normal_(self.audio_head, mean=0.0, std=0.02)
    
    def setup_caches(self, max_batch_size: int) -> torch.Tensor:
        """Setup KV caches and return a causal mask."""
        dtype = next(self.parameters()).dtype
        device = next(self.parameters()).device
        
        with device:
            self.backbone.setup_caches(max_batch_size, dtype)
            self.decoder.setup_caches(max_batch_size, dtype, decoder_max_seq_len=self.args.audio_num_codebooks)
        
        self.register_buffer("backbone_causal_mask", _create_causal_mask(self.backbone.max_seq_len, device))
        self.register_buffer("decoder_causal_mask", _create_causal_mask(self.args.audio_num_codebooks, device))
    
    def generate_frame(
        self,
        tokens: torch.Tensor,
        tokens_mask: torch.Tensor,
        input_pos: torch.Tensor,
        temperature: float,
        topk: int,
    ) -> torch.Tensor:
        """Generate a frame of audio tokens.
        
        Args:
            tokens: (batch_size, seq_len, audio_num_codebooks+1)
            tokens_mask: (batch_size, seq_len, audio_num_codebooks+1)
            input_pos: (batch_size, seq_len) positions for each token
        
        Returns:
            (batch_size, audio_num_codebooks) sampled tokens
        """
        dtype = next(self.parameters()).dtype
        b, s = tokens.size()[:2]
        
        assert self.backbone.caches_are_enabled(), "backbone caches are not enabled"
        
        curr_backbone_mask = _index_causal_mask(self.backbone_causal_mask, input_pos)
        embeds = self._embed_tokens(tokens)
        masked_embeds = embeds * tokens_mask.unsqueeze(-1)
        h = masked_embeds.sum(dim=2)
        
        h = self.backbone(h, input_pos=input_pos, mask=curr_backbone_mask).to(dtype=dtype)
        
        last_h = h[:, -1, :]
        c0_logits = self.codebook0_head(last_h)
        c0_sample = sample_topk(c0_logits, topk, temperature)
        c0_embed = self._embed_audio(0, c0_sample)
        
        curr_h = torch.cat([last_h.unsqueeze(1), c0_embed], dim=1)
        curr_sample = c0_sample.clone()
        curr_pos = torch.arange(0, curr_h.size(1), device=curr_h.device).unsqueeze(0).repeat(curr_h.size(0), 1)
        
        # Decoder caches must be reset every frame.
        self.decoder.reset_caches()
        
        for i in range(1, self.args.audio_num_codebooks):
            curr_decoder_mask = _index_causal_mask(self.decoder_causal_mask, curr_pos)
            decoder_h = self.decoder(self.projection(curr_h), input_pos=curr_pos, mask=curr_decoder_mask).to(
                dtype=dtype
            )
            ci_logits = torch.mm(decoder_h[:, -1, :], self.audio_head[i - 1])
            ci_sample = sample_topk(ci_logits, topk, temperature)
            ci_embed = self._embed_audio(i, ci_sample)
            
            curr_h = ci_embed
            curr_sample = torch.cat([curr_sample, ci_sample], dim=1)
            curr_pos = curr_pos[:, -1:] + 1
        
        return curr_sample
    
    def reset_caches(self):
        """Reset KV caches."""
        self.backbone.reset_caches()
        self.decoder.reset_caches()
    
    def _embed_audio(self, codebook: int, tokens: torch.Tensor) -> torch.Tensor:
        """Embed audio tokens."""
        return self.audio_embeddings(tokens + codebook * self.args.audio_vocab_size)
    
    def _embed_tokens(self, tokens: torch.Tensor) -> torch.Tensor:
        """Embed tokens."""
        text_embeds = self.text_embeddings(tokens[:, :, -1]).unsqueeze(-2)
        audio_tokens = tokens[:, :, :-1] + (
            self.args.audio_vocab_size * torch.arange(self.args.audio_num_codebooks, device=tokens.device)
        )
        audio_embeds = self.audio_embeddings(audio_tokens.view(-1)).reshape(
            tokens.size(0), tokens.size(1), self.args.audio_num_codebooks, -1
        )
        
        return torch.cat([audio_embeds, text_embeds], dim=-2)