提案：为 Twinkle 新增 TTS 训练支持（多码本 Preprocessor + Loss）

[Wallbreazzz]

背景

Twinkle 现在支持文本 LLM 和多模态，但没有 TTS 训练支持（cookbook 里没 TTS 示例，模型支持表也没列 TTS 变体）。我想给它加上 TTS 的 LoRA SFT 能力，先发个方案对一下，如果没问题我会按工作项实现并提交PR。

目标

让 Twinkle 能跑这两类 TTS 模型的 LoRA SFT：

1. 多码本 AR-LM（如 Qwen3-TTS-12Hz-1.7B-Base）：输入"文本+参考音频"，输出"目标音频的多码本 token"
2. 非 AR 局部 Transformer（如 MOSS-TTS-Local-Transformer-v1.5）

两类架构不同，但 LoRA 都挂在 q/k/v 这些线性层，组件能复用。

设计

仿 cookbook/mm/fsdp2.py 里 LatexOCRProcessor 的范式（自定义 Preprocessor 子类 + Trajectory/Message），加两个组件：

TTSMultiCodebookProcessor（Preprocessor 子类）

把 manifest 里的多码本 audio_codes 组织成训练样本：

• 文本 + 参考音频部分：label 置 -100（不学）
• 目标音频部分：保留多码本 label（学）
• 输出形如 Trajectory(messages=[Message(role='user', content=text, audio=ref_codes), Message(role='assistant', audio_codes=target_codes)])

TTSMultiCodebookLoss（Loss 子类）

多码本 next-token 交叉熵：每个位置对每个码本算 CE 再求和，只在 label != -100 的位置算。多码本 loss 的写法参考 OpenMOSS 官方 MOSS-TTS 仓库 community/norwegian-lora/train_lora.py。

工作项拆解

• W1 调研多码本 token 在序列里怎么组织（多头输出 / flatten / 残差量化），看 Qwen3-TTS 官方 tokenizer 代码确认。这决定 Preprocessor 的 encode 逻辑。
• W2 设计 TTSMultiCodebookProcessor 接口：仿 LatexOCRProcessor，定 preprocess(row) -> Trajectory 的字段（Message 要不要加 audio_codes 字段、参考音频怎么放）。
• W3 实现 TTSMultiCodebookProcessor + 单测：用极小 manifest 跑通 encode，断言 input_ids/labels 的 shape 和 -100 位置正确。
• W4 设计+实现 TTSMultiCodebookLoss：多码本 CE，仿现成 causal_lm 的接口风格，加 task='multi_codebook_lm'。
• W5 加 cookbook 示例 cookbook/tts/qwen3_tts_lora.py：单卡 + LoRA + bf16，仿 cookbook/transformers/fsdp2.py 结构。
• W6 MOSS-TTS-Local-Transformer 适配：非 AR 局部 Transformer 训练格式不同，可能要再写个 Preprocessor 变体（W1-W5 跑通后再碰）。
• W7 文档：README 教程表加 TTS 行、docs 加一节"TTS 训练"。

不在范围内

• vocoder / codec 本身训练（冻结用官方训好的）
• 全量 SFT（先做 LoRA，全量以后再说）

想确认的几点

1. Message 类要不要扩展支持 audio_codes 字段？还是复用现成的 images 之类？
2. 多码本 loss 是写个新 Loss 类，还是给现成 CE 加个 multi-codebook 模式？
3. task='multi_codebook_lm' 这个 task 名合不合适？要不要走现成 task 机制？
4. cookbook 放 cookbook/tts/ 还是 cookbook/transformers/ 下？

硬件

我在昇腾 NPU 单卡上做（bf16 LoRA），但组件本身不绑死硬件，CUDA 也能跑。

[Wallbreazzz]

cc @tastelikefeet @tpx818 could u please review this plan, thx

[tastelikefeet]

感谢你的代码计划！
先回复你的 plan
首先这个实现大致是合理的，有几个点需要注意：

1. TTSMultiCodebookProcessor：预处理一般与数据集相关，定义的这个数据集可以类似于一个基类，可以让所有类似数据集继承复用
2. Trajectory 结构中有 user_data 字段，可以将其他字段放在其中
3. 可以是一个新 loss
4. 目前的 task 的定义是针对模型结构有修改（例如 embedding、reranker），因此需要额外的 task 提示 forward 的时候进行处理，因此如果只是 loss 不同，可以复用现有 task
5. cookbook位置这个命题相对随意
6. 注意多 lora 部分是否可以兼容，毕竟 twinkle 是支持 c/s 模式的
   其次说一下其他相关的问题：
   swift 目前在和 twinkle 进行底层架构融合，可能对你的 plan 有所影响。但我们仍然欢迎你将自己的计划贡献进来，因此就有两个模式：
7. 等待我们的重构完成（约 9 月末稳定），然后基于我们的代码继续提升
8. 先行合入 twinkle，我们会基于你的代码和我们的已有代码组合，形成更完善和鲁棒的代码版本

再次感谢你的 plan！如果代码完成，直接提交 PR 即可。

[Wallbreazzz]

感谢 review！确认理解以下几点：

1. audio_codes 放 Trajectory.user_data（不改 Message），Processor 做成基类让类似 TTS 数据集继承复用
2. 写新的 TTSMultiCodebookLoss 类
3. 复用现有 task（causal_lm），不新增 task，loss 差异在 Loss 层处理
4. cookbook 放 cookbook/tts/

另外会注意多 LoRA 兼容（c/s 模式下 Processor/Loss 不写死 adapter_name）。

时机选择：模式2 先行合入，9 月末 swift 融合如有适配再调整。预计近期提 PR。