fix(tellama): infer

ChatTTS/model/cuda/patch.py

@@ -0,0 +1,17 @@
+import torch
+
+class LlamaRMSNorm(torch.nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        LlamaRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight.to(hidden_states.device) * hidden_states.to(input_dtype)

ChatTTS/model/cuda/te_llama.py

@@ -6,15 +6,12 @@
 #
 # Edited by fumiama.
 
-import os
 import re
-import gc
 from contextlib import contextmanager
 from typing import Dict
 
 import transformer_engine as te
 from transformer_engine.pytorch.attention import RotaryPositionEmbedding
-from transformer_engine.pytorch.fp8 import fp8_model_init
 
 import torch
 
@@ -23,17 +20,13 @@ from transformers.models.llama.modeling_llama import (
     LlamaModel,
     LlamaConfig,
 )
-from transformers.modeling_utils import (
-    _add_variant,
-    load_state_dict,
-    _load_state_dict_into_model,
-)
-from transformers.utils import WEIGHTS_INDEX_NAME
-from transformers.utils.hub import get_checkpoint_shard_files
+from transformers.modeling_utils import _load_state_dict_into_model
+
+from .patch import LlamaRMSNorm
 
 
 @contextmanager
-def replace_decoder(te_decoder_cls):
+def replace_decoder(te_decoder_cls, llama_rms_norm_cls):
     """
     Replace `LlamaDecoderLayer` with custom `TELlamaDecoderLayer`.
     """
@@ -41,12 +34,17 @@ def replace_decoder(te_decoder_cls):
         transformers.models.llama.modeling_llama.LlamaDecoderLayer
     )
     transformers.models.llama.modeling_llama.LlamaDecoderLayer = te_decoder_cls
+    original_llama_rms_norm_cls = (
+        transformers.models.llama.modeling_llama.LlamaRMSNorm
+    )
+    transformers.models.llama.modeling_llama.LlamaRMSNorm = llama_rms_norm_cls
     try:
         yield
     finally:
         transformers.models.llama.modeling_llama.LlamaDecoderLayer = (
             original_llama_decoder_cls
         )
+        transformers.models.llama.modeling_llama.LlamaRMSNorm = original_llama_rms_norm_cls
 
 
 class TELlamaDecoderLayer(te.pytorch.TransformerLayer):
@@ -106,7 +104,7 @@ class TELlamaModel:
     """
 
     def __new__(cls, config: LlamaConfig):
-        with replace_decoder(te_decoder_cls=TELlamaDecoderLayer):
+        with replace_decoder(te_decoder_cls=TELlamaDecoderLayer, llama_rms_norm_cls=LlamaRMSNorm):
             model = LlamaModel(config)
         return model
 

ChatTTS/model/gpt.py

@@ -43,6 +43,7 @@ class GPT(nn.Module):
         self.use_flash_attn = use_flash_attn
 
         self.gpt, self.llama_config = self._build_llama(gpt_config, self.device_gpt)
+        self.is_te_llama = False
         self.model_dim = int(self.gpt.config.hidden_size)
         self.emb_code = nn.ModuleList(
             [
@@ -99,6 +100,7 @@ class GPT(nn.Module):
                 del state_dict, self.gpt
                 gc.collect()
                 self.gpt = vanilla
+                self.is_te_llama = True
             except Exception as e:
                 self.logger.warning(
                     f"use default LlamaModel for importing TELlamaModel error: {e}"
@@ -139,7 +141,7 @@ class GPT(nn.Module):
     def prepare(self, compile=False):
         if self.use_flash_attn and is_flash_attn_2_available():
             self.gpt = self.gpt.to(dtype=torch.float16)
-        if compile:
+        if compile and not self.is_te_llama:
             try:
                 self.compile(backend="inductor", dynamic=True)
                 self.gpt.compile(backend="inductor", dynamic=True)
@@ -217,9 +219,10 @@ class GPT(nn.Module):
         # TODO joao: standardize interface for the different Cache classes and remove of this if
         has_static_cache = False
         if past_key_values is None:
-            past_key_values = getattr(
-                self.gpt.layers[0].self_attn, "past_key_value", None
-            )
+            if hasattr(self.gpt.layers[0], "self_attn"):
+                past_key_values = getattr(
+                    self.gpt.layers[0].self_attn, "past_key_value", None
+                )
             has_static_cache = past_key_values is not None
 
         past_length = 0
@@ -418,7 +421,7 @@ class GPT(nn.Module):
                 inputs_ids,
                 past_key_values,
                 attention_mask_cache.narrow(1, 0, inputs_ids.shape[1]),
-                use_cache=True,
+                use_cache=not self.is_te_llama,
             )
 
             if i > 0: