LLM 对话交互实现探索

这是我正在撰写的一系列文章中的第一篇，旨在整理如何微调大型语言模型（LLM）以处理音频的相关经验，最终目标是构建并托管一个能够描述人类声音的 LLM。

我的动力是获得亲手折腾 LLM 的实践经验，因此只要现实条件允许，我都会尝试使用 PyTorch 从零重新实现工具和函数，而不是依赖第三方库。

简而言之：本文记录并分享了我学习如何在 Google 的 MusicCaps 数据集上微调 LLM 模型、使其能够描述给定音频文件的具体步骤；你也可以在这里找到原始 Jupyter Notebook。

最近，我读到了两篇论文：

SALMONN: Towards Generic Hearing Abilities for Large Language Models

Qwen-Audio: Advancing Universal Audio Understanding via Unified Large-Scale Audio-Language Models

它们都旨在让 LLM 具备音频理解能力。

总体而言，两篇论文都探索了如何利用音频编码器将声音转换为嵌入，然后与文本嵌入一起输入 LLM。

在 SALMONN 中，研究者结合了 OpenAI 的 Whisper 和 BEATS 编码器，先对组合后的编码器进行预训练，再利用 LoRA 微调 LLM。Qwen-Audio 则以 OpenAI 的 Whisper 为基础初始化其音频编码器；预训练完成后，Qwen-Audio 会对 LLM 进行全量微调。

这两篇论文让我很好地了解了如何适配跨领域编码器，并将其与 LLM 结合。通用音频理解能力的 LLM 这一想法令我十分兴奋，我也迫不及待地想获得亲手实践的经验，于是决定尝试构建一个具备音频处理能力的最小可行 LLM。

首先，我前往 HuggingFace，寻找合适的基础 LLM 和中等规模的数据集。我希望尽可能在本地完成所有工作，因此一切都必须能在本地 RTX 3090 上运行。

在测试并比较了几个不同模型之后，我最终选择了 Mistral OpenOrca。

音频编码器方面，我选择了 OpenAI 的 Whisper。

数据集方面，我选择了 MusicCaps。由于没有找到方便下载已处理或已分段音频文件的链接，我便编写了一个小脚本来下载 YouTube 视频。

每次迈出一小步

准备好基本依赖后，我启动 Jupyter Notebook，开始动手实验。

从零实现采样

我的第一步是确保能够加载基础 LLM，并正确执行推理。我没有使用 transformers 库的生成工具，而是自行实现了采样函数，一方面用来验证自己的理解，另一方面也为了学习如何直接使用嵌入进行采样——之后输入音频嵌入时，这项能力会派上用场。

@torch.no_grad
def sampler(input_ids):
    outputs = []

for _ in range(50):
        inputs_embeds = model.llm.model.embed_tokens(input_ids)
        res = model.llm(inputs_embeds=inputs_embeds)
        # res.logits shape is (batch, seq_len, logits)
        # sample using multinomial using the last logits 
        sampled = torch.multinomial(res.logits[:,-1,:].softmax(dim=-1), 1)
        # repeatedly concat the `sampled` to the `input_ids` for next sampling
        input_ids = torch.cat((input_ids, sampled), dim=-1)

return input_ids

使用通过 Transformer 的 AutoTokenizer 类获得的 tokenizer，我成功验证了采样功能符合预期！运行：

tokenizer.decode(sampler(tokenizer("tell me a story", return_tensors="pt").input_ids.to("cuda:0"))[0])

会得到如下结果（示例输出）：

'<s>tell me a story is a film and video production company, tell me a story is a concept that was created to allow people to come together through the power of storytelling.\n and so, with this massive power in storytelling, the founders and creat'

调试 NaN 和 Inf

目前为止一切顺利。然而，我很快注意到，采样函数偶尔会执行失败，并提示 softmax 函数遇到了 inf 或 NaN。我参考了一个很有启发性的讨论帖，并学会了如何使用以下经过调整的 PyTorch 钩子来定位 NaN 的来源：

import torch
from functools import partial

__registered_hook_refs = []

for h in __registered_hook_refs:
    h.remove()

__global = []

def nan_hook(module, args, output, name=None):
    if not isinstance(output, tuple):
        outputs = [output]
    else:
        outputs = output

for i, out in enumerate(outputs):
        if out is None:
            continue
        if isinstance(out, tuple):
            for j, out2 in enumerate(out):
                nan_mask = torch.isnan(out2)
                if nan_mask.any():
                    __global.append((module, args, output))
                    raise RuntimeError(f"In module {name} of name {module.__class__.__name__}, Found NAN in output {j} at indices: ", nan_mask.nonzero(), "where:",
                               out[nan_mask.nonzero()[:, 0].unique(sorted=True)])
                               
        elif torch.is_tensor(out):
            nan_mask = torch.isnan(out)
            
            if nan_mask.any():
                __global.append((module, args, output))
                raise RuntimeError(f"In module {name} of name {module.__class__.__name__}, Found NAN in output {i} at indices: ", nan_mask.nonzero(), "where:",
                                   out[nan_mask.nonzero()[:, 0].unique(sorted=True)])

def register_nan_hook(model: torch.nn.Module):
    for name, submodule in model.named_modules():
        new_hook = partial(nan_hook, name=name+'.back')
        hook_ref = submodule.register_full_backward_hook(new_hook)
        __registered_hook_refs.append(hook_ref)
        new_hook = partial(nan_hook, name=name+'.fwd')
        hook_ref = submodule.register_forward_hook(new_hook)
        __registered_hook_refs.append(hook_ref)

debug = True
register_nan_hook(model) if debug else None

借助这些钩子，我将问题来源缩小到了某个特定层，随后追踪发现，问题是由模型权重中的一个 inf 值引起的。进一步排查后，我发现 inf 的根源竟然是存在故障的内存条！采取缓解措施后，我编写了一个小脚本来验证模型权重，并确认采样函数能够按预期工作。

# verify model weight
from collections import Counter
pbytype = Counter()
for name, p in (model.named_parameters()):
    if torch.isinf(p).any() or torch.isnan(p).any():
        print(name, p)
        raise ValueError("invalid weight")
    else:
        pbytype[p.dtype] += 1
print("OK", pbytype)

将 Whisper 适配到 Mistral

在对调试 PyTorch 模块更有信心之后，我将重点放在适配 Whisper 模型上，使其能够将音频文件转换为嵌入，再将该嵌入输入 Mistral。

OpenAI 的 Whisper 模型由两个主要组件构成：AudioEncoder 和 TextDecoder。为了将音频转换为嵌入，我只需要 AudioEncoder 组件。

因此，我加载了完整的 Whisper 模型，并使用以下代码片段提取 AudioEncoder 的权重：

import whisper

model = whisper.load_model("large-v3")
audio_encoder = model.encoder
torch.save(
    audio_encoder.state_dict(),
    "<output_location>",
)

我将 Whisper AudioEncoder 改造成了 TunableWhisperAudioEncoder，并额外添加一个投影层，将 Whisper 的音频嵌入（大小为 1280）映射到 Mistral 的 token 嵌入（大小为 4096）。

我通过显式冻结音频编码器的参数，确保 proj 是唯一可训练的网络。请注意，TrainableSubmodule 是一个超参数，任何能将输出嵌入映射到 4096 维的模型都可以使用。在本文后面的部分，我会介绍实践中对我有效的方案。

class TunableWhisperAudioEncoder(nn.Module):
    def __init__(self, *, output_embedding_size=4096):
        """
        args
            output_embedding_size: int = 4096 / mistral default embedding size
        """
        super().__init__()

self.audio_encoder = load_whisper_v3_audio_encoder()
        self.proj = TrainableSubmodule(output_embedding_size=output_embedding_size)

# # Freeze all parameters
        for param in audio_encoder.parameters():
            param.requires_grad = False

def forward(self, mels):
        res = self.audio_encoder(mels)
        res = self.proj(res)
        return res

def load_whisper_v3_audio_encoder(
    *,
    n_mels=128,
    n_audio_ctx=1500,
    n_audio_state=1280,
    n_audio_head=20,
    n_audio_layer=32,
):
    m = whisper.model.AudioEncoder(
        n_mels, n_audio_ctx, n_audio_state, n_audio_head, n_audio_layer
    )
    m.load_state_dict(torch.load(WHISPER_AUDIO_BIN))
    return m

最后，我按如下方式构建了将用于训练的模型：

class Model(nn.Module):
    def __init__(self, audio_encoder: "Whisper.AudioEncoder", llm: "Mistral"):
        super().__init__()

self.audio_encoder = audio_encoder
        self.llm = llm

# freeze the LLM weights
        for p in self.llm.parameters():
            p.requires_grad = False

```python
def forward(self, batch):
        audio_mels = batch["audio_mels"]
        # caption token ids
        cap_ids = batch["cap_ids"]
        # caption attention mask
        cap_ids_attention_mask = batch["cap_attention_mask"]
        prompt_ids = batch["prompt_ids"]
        prompt_ids_attention_mask = batch["prompt_attention_mask"]
        end_prompt_ids = batch["end_prompt_ids"]
        end_prompt_ids_attention_mask = batch["end_prompt_attention_mask"]

audio_embeds = self.audio_encoder(audio_mels)
        # audio_embeds: (batch, audio_seq_len, audio_embedding_size)
        bs, audio_seq = audio_embeds.shape[:2]

attention_mask = torch.concat(
            (
                prompt_ids_attention_mask,
                torch.ones(bs, audio_seq).to(cap_ids.device),
                end_prompt_ids_attention_mask,
                cap_ids_attention_mask,
            ),
            dim=1,
        )

cap_embeds = self.llm.model.embed_tokens(cap_ids)
        prompt_embeds = self.llm.model.embed_tokens(prompt_ids)
        end_prompt_embeds = self.llm.model.embed_tokens(end_prompt_ids)

# build the inputs_embeds by concating all the token embeddings
        # with audio_embeddings
        inputs_embeds = torch.concat(
            (
                prompt_embeds,
                audio_embeds.to(cap_embeds.dtype),
                end_prompt_embeds,
                cap_embeds,
            ),
            dim=1,
        )

mout = self.llm(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
        )

return mout, audio_embeds.shape[1]

这个模型本身非常简洁，它只是持有对 Mistral LLM 和 TunableWhisperAudioEncoder 的引用。forward 方法封装了将音频梅尔频谱转换为音频嵌入的逻辑，然后将音频嵌入与文本/token 嵌入拼接，并将其输入到 Mistral LLM 中。

从零开始进行音频采样

基础模型就位后，下一步是尝试使用音频输入从该模型进行采样。以下是我设计的音频采样函数。

# note, full gist is available at https://gist.github.com/moomou/7df8345d79a0063d67d1fa2b4cf55db8

@torch.no_grad()
def sample_with_audio(model, tokenizer, prompt, audio_file, device="cuda:0", iteration=50):
    audio_mels = load_audio_mels(audio_file).to(device).half()
    end_prompt_ids, end_prompt_attention_mask = text_2_ids_and_attention_mask(
        tokenizer,
        end_template(),
        truncate=True,
    )
    prompt_ids, prompt_attention_mask = text_2_ids_and_attention_mask(
        tokenizer,
        prompt,
    )

prompt_ids = prompt_ids.to(device)
    prompt_attention_mask = prompt_attention_mask.to(device)
    end_prompt_attention_mask = end_prompt_attention_mask.to(device)
    end_prompt_ids = end_prompt_ids.to(device)
    sampled_ids = None

prompt_embeds = None
    end_prompt_embeds = None
    audio_embeds = None

with torch.amp.autocast(device_type="cuda", dtype=torch.float16): # 使用 float16 以减少 GPU 内存
        if audio_embeds is None:
            audio_embeds = model.audio_encoder(audio_mels)
        bs, audio_seq = audio_embeds.shape[:2]
        
        mask_concat_args = [
            prompt_attention_mask,
            torch.ones(bs, audio_seq).to(audio_embeds.device),
            end_prompt_attention_mask,
        ]

for _ in range(iteration):
            if sampled_ids is not None:
                mask_concat_args.append(torch.ones(bs, sampled_ids.shape[1]).to(audio_embeds.device))
                
            attention_mask = torch.concat(
                tuple(mask_concat_args),
                dim=1,
            )

if prompt_embeds is None:
                prompt_embeds = model.llm.model.embed_tokens(prompt_ids)
            if end_prompt_embeds is None:
                end_prompt_embeds = model.llm.model.embed_tokens(end_prompt_ids)
                
            sampled_ids_embeds = None
            if sampled_ids is not None:
                sampled_ids_embeds = model.llm.model.embed_tokens(sampled_ids)
                
            embeds_concat_args = [
                prompt_embeds,
                audio_embeds.to(prompt_embeds.dtype),
                end_prompt_embeds,
            ]
            if sampled_ids_embeds is not None:
                embeds_concat_args.append(sampled_ids_embeds)
                
            inputs_embeds = torch.concat(
                tuple(embeds_concat_args),
                dim=1,
            )
    
            mout = model.llm(
                inputs_embeds=inputs_embeds,
                attention_mask=attention_mask,
            )
    
            logits = mout.logits
            sampled = torch.multinomial(logits[:, -1, :].softmax(dim=-1), 1)
            
            if sampled_ids is None:
                sampled_ids = sampled
            else:
                sampled_ids = torch.cat((sampled_ids, sampled), dim=-1).to(device)

return torch.concat((
        prompt_ids, 
        end_prompt_ids,
        sampled_ids,
    ),dim=-1)

通过以下方式使用这个函数

dataloader = ... # 标准 pytorch 数据加载器
local_batch = next(iter(dataloader))
tokenizer.decode(sample_with_audio(model, tokenizer, prompt_template_fn(), audio_file, iteration=60)[0])

如预期那样产生了乱码，因为 TunableWhisperAudioEncoder 投影层未经训练。

'<s> <|im_start|>  system\n    You are a helpful AI who follows instruction carefully<|im_end|> <|im_start|>  user\n    Describe the sound of the given file \n    <|im_end|> <|im_start|>  assistant\n     war<|im_end|> clockunits ]andfirst4Iftektime爆R Cur<|im_end|> United<|im_end|> 'daysIn"Never<|im_end|> thenAnd,and VI<|im_end|> Islo<|im_end|> GOkaydown<|im_end|> JainteYoulfailedLabelsEvenfacevC,rest<|im_end|><|im_end|><|im_end|><|im_end|> q<|im_end|> Xs<|im_end|> h<|im_end|><|im_end|>'

定义损失函数

这里的损失函数是对 logits 输出的标准交叉熵损失；唯一的技巧是损失应该只在标题部分计算。具体来说，

# 计算损失
# local_batch: (b, seq, C)
prompt_ids_seq = local_batch["prompt_ids"].shape[1]
end_prompt_ids_seq = local_batch["end_prompt_ids"].shape[1]
logits_start = prompt_ids_seq + audio_seq + end_prompt_ids_seq

# 移除最后的输出
logits = ... # 模型输出
# 从 logits 计算中移除 prompt 和音频序列
# 另外，移除最后一项
logits = logits[:, logits_start:-1, :].contiguous()

# 仅使用 `cap_ids` 计算目标
targets = batch["cap_ids"][:]
targets = targets[:, 1:]

loss = nn.functional.cross_entropy(
    logits.view(-1, logits.shape[-1]), targets.view(-1)
)

训练、过拟合和梯度调试

最后，所有的碎片都到位了可以训练模型。我的目标是通过仅训练 TunableWhisperAudioEncoder 来让冻结的 LLM 描述给定的音频文件；实现这一点不会给 LLM 通用的音频理解能力，因为训练数据很小，但会让我对自己执行了所有基本步骤充满信心。

为了确保训练设置正确，我从小处开始，一次一步。具体来说，我交互式地手动逐步执行训练步骤，记录并绘制了 TunableWhisperAudioEncoder 中权重更新相对于权重数据的关系，并使用之前描述的 Pytorch 钩子确保没有 inf 或 NaN。这些步骤对各种学习率、模型架构和优化器的组合进行了重复。

保持设置尽可能简单，我发现 Adam（不带动量）、常数学习率 1.5e-3，以及使用以下简单的 TrainableSubmodule，我实现了稳定的训练。

class TrainableSubmodule(nn.Module):
    def __init__(self, output_embedding_size=4096):
        super().__init__()

self.pool = nn.AdaptiveAvgPool1d(250)
        self.proj = nn.Linear(1280, output_embedding_size, bias=False)
        self.ln1 = nn.LayerNorm(1280)

我进行了约 4 天的训练，在我停止训练时，损失仍在下降。当我停止时，我实现了约 0.46 的损失，这转换为正确 token 的约 66% 概率！

用产生乱码的预训练同一音频文件重新运行 sample_with_audio，我现在得到

"<s> <|im_start|>  system\n    You are a helpful AI who follows instruction carefully<|im_end|> <|im_start|>  user\n    Describe the sound of the given file \n    <|im_end|> <|im_start|> assistant\n     The electronica song features a crisp acoustic kick, snap snare and hat along with a deep bass. The male vocal is rapping syncopated along with a male background vocal. The song is fast paced and there is a limited frequency range of the synths. The song"

与真实情况进行比较

“这是一首由男子组合演唱的 K-pop 音乐作品。开头，一名男歌手以类似说唱的方式演唱。随后，歌曲切换到另一名男歌手更具旋律性的演唱。旋律由清脆的合成器音色奏出。节奏背景由充满活力的电子鼓点构成，整体具有很强的舞蹈感。这首作品很适合在韩国的夜店和舞厅播放。”

效果相当不错！

值得再次强调的是，这一效果仅通过训练音频编码器的投影层实现，并未修改 LLM 的权重或 Whisper AudioEncoder 的权重。

在打好基础之后，我计划扩大训练规模，引入更多音频任务，例如转录、说话人识别等，同时对 LLM 进行微调，逐步尝试复现参考论文中描述的“涌现”行为。

假设拥有充足的数据并采用适当的训练方案，LLM 应该能够执行一些原本的音频任务，例如在没有针对这类任务进行显式训练的情况下，识别说话人的年龄或性别。

还有更多工作要做！

如果没有学习 Karpathy 那些精彩的课程，我不可能完成这里的任何工作。