模型系统工程:模型分布式训练并行策略
引言
以前在训练模型的时候大部分使用的是单机单卡进行训练测试,真正使用单机多卡和多机多卡的时候,很少去实操到分布式的训练推理,本身这块对应硬件的成本高,对于个人是很少可以去把玩上的,更何况现在训练大模型大部分使用的PT好的模型,进行微调和蒸馏,或者量化部署,随着deepseek系列模型的开源,对模型的训练微调和部署的需求增多,里面涉及到的分布式训练策略,怎样对模型和数据进行拆分进行多卡sm并行处理, 这里简单介绍下相关的分布式训练推理并行策略,以及对应的代码,代码主要使用pytorch进行实现,采用单节点的2个gpu。(如果有足够多的资源,可以复现下ZeRO或者Megatron中的实现)
💡使用 PyTorch 进行模型训练的过程:
PyTorch 将值组织成Tensor , Tensor是具有丰富数据操作操作的通用 n 维数组。Module 定义从输入值到输出值的转换,其在正向传递期间的行为由其forward成员函数指定。Module 可以包含Tensor作为参数。例如,线性模块包含权重参数和偏差参数,其正向函数通过将输入与权重相乘并添加偏差来生成输出。应用程序通过在自定义正向函数中将本机Module (例如线性、卷积等)和Function (例如relu、pool 等)拼接在一起来组成自己的Module 。典型的训练迭代包含使用输入和标签生成损失的前向传递、用于计算参数梯度的后向传递以及使用梯度更新参数的优化器步骤。更具体地说,在正向传递期间,PyTorch 会构建一个自动求导图来记录执行的操作。然后,在反向传播中,它使用自动梯度图进行反向传播以生成梯度。最后,优化器应用梯度来更新参数。训练过程重复这三个步骤,直到模型收敛。
并行策略主要分为以下几种:
Data Parallelism (DP) 包括:
- Distributed Data Parallelism(DDP)
- Fully-Shared Data Parallelism (FSDP)
Model Parallelism(MP)包括:
-
Tensor Parallelism(TP, 有些论文中将TP描述成MP,比如Megatron-LM, ZeRO)
-
Pipeline Parallelism (PP)
-
Expert Parallelism (EP)
-
Activation Partitioning 包括:(SP和CP两者通常和TP一起使用)
- Sequence Parallelism (SP): 针对序列切分,在模块的输入和输出侧对序列进行切分,常使用在LayerNorm,Dropout(通过all_gather,reduce-scatter 通信计算,减少训练时激活值 gpu内存占用,以前是每个卡上单独的副本LayerNorm,Dropout,冗余gpu内存)
- Context Parallelism (CP): 主要针对Transformer模型长序列的训练。
DP和MP的比较:
-
DP 的扩展效率比 MP 更好,因为 MP 降低了计算的粒度,同时也增加了通信开销。超过某个点后,较低的计算粒度会降低每个 GPU 的效率,而增加的通信开销会阻碍跨 GPU 的可扩展性,尤其是跨越节点边界时。相反,DP 既具有更高的计算粒度,又具有更低的通信量,从而可以实现更高的效率。
-
DP 内存(GPU HBM)效率低,因为模型状态在所有数据并行进程中冗余存储。相反,MP 对模型状态进行分区,以提高内存效率。
-
DP和MP都保存了整个训练过程中需要的所有模型状态,但并不是所有状态都是一直需要的,比如每一层对应的参数只在该层的前向传播和后向传播时才需要。
模型训练的并行策略同样适用于推理测的并行策略,但是调度和执行的方式稍有不同。这里主要是结合pytorch 实现简单的多卡模型并行策略训练。
相关的分布式训练推理并行策略,以及对应的代码运行操作demo,可以查看这个PR:
https://github.com/ai-bot-pro/achatbot/pull/127 (如有不对,欢迎指出~)
论文解读:Fish-Speech: Leveraging Large Language Models for Advanced Multilingual Text-to-Speech Synthesis
相关论文
base: 基础普适研究
- hourglass transformers: 2021. Hierarchical Transformers Are More Efficient Language Models | lucidrains/simple-hierarchical-transformer vanilla layers and shortened layers use GPT AR GLM 🤞
- 2023.5 MEGABYTE: Predicting Million-byte Sequences with Multiscale Transformers (分别在文本,图片,语音建模)| lucidrains/MEGABYTE-pytorch
PS: 想法和自己整理的simple LM相似~
扩展阅读:2024.12 Byte Latent Transformer: Patches Scale Better Than Tokens | paper code
audio speech: 场景研究
-
⭐️2023.10 UniAudio: An Audio Foundation Model Toward Universal Audio Generation (灵感来自 MEGABYTE,将其应用于语音模型)| paper code (代码可扩展任务进行训练, 已扩展了音乐数据)
-
⭐️2024. Fish-Speech: Leveraging Large Language Models for Advanced Multilingual Text-to-Speech Synthesis (论文中介绍的Daul-AR, GFSQ, Firefly-GAN(FF-GAN) 对EVA-GAN改版,细节需要结合代码理解) | paper code
附:
- achatbot 接入fishspeech tts: https://github.com/ai-bot-pro/achatbot/pull/108
- achatbot 接入fishspeech tts colab 笔记: https://github.com/weedge/doraemon-nb/blob/main/achatbot_fishspeech_tts.ipynb
- 操作笔记: https://github.com/weedge/doraemon-nb/blob/main/fish_speech_tts.ipynb
论文解读:CosyVoice2: Scalable Streaming Speech Synthesis with Large Language Models
CosyVoice2 论文
- 2024.12 CosyVoice 2: Scalable Streaming Speech Synthesis with Large Language Models(流式合成)
- paper code: 公开推理和权重,训练过程需要在CosyVoice的基础上修改下。
- achatbot TTS 集成 CosyVoice2: https://github.com/ai-bot-pro/achatbot/pull/107
- Colab achatbot_CosyVoice 操作笔记: https://github.com/weedge/doraemon-nb/blob/main/achatbot_CosyVoice.ipynb
- Colab CosyVoice 操作笔记: https://github.com/weedge/doraemon-nb/blob/main/CosyVoice.ipynb
扩展阅读
zero-shot TTS models 零样本 TTS 模型
codec language models 编解码器语言模型
-
speech codec model to extract discrete speech representation:
- 2021.7 SoundStream: An End-to-End Neural Audio Codec | code
- 2022.10 High Fidelity Neural Audio Compression | facebookresearch/encodec
- 2023.5 HiFi-Codec: Group-residual Vector quantization for High Fidelity Audio Codec | paper code
- 2023.10 FunCodec: A Fundamental, Reproducible and Integrable Open-source Toolkit for Neural Speech Codec | modelscope/FunCodec
-
speech codec model + autoregressive model to predict the speech tokens (acoustic tokens):
-
speech codec model (speech semantics Codec) + non-autoregressive masked model to predict the speech tokens (acoustic tokens):
-
codec model (speech acoustic Codec) or vocoder to synthesize waveforms from mel-spectrograms:
-
2023.6 Vocos: Closing the gap between time-domain and Fourier-based neural vocoders for high-quality audio synthesis | paper code | 推理速度快:运行速度比 HiFi-GAN 快约 13 倍,比 BigVGAN 快近 70 倍。在没有 GPU 加速的情况下运行时,这种速度优势尤其明显。这主要是由于使用了短时傅里叶逆变换(ISTFT)算法而不是转置卷积。还评估了 Vocos 的一个变体,它利用 ResBlock 的扩张卷积而不是 ConvNeXt 块。在 GPU 上执行时,深度可分离卷积可提供额外的加速。
-
2023.12 WaveNeXt: ConvNeXt-Based Fast Neural Vocoder Without ISTFT layer | demo samples | paper code基于 ESPNet2-TTS | 一种新型的基于ConvNeXt的快速神经声码器WaveNeXt,它通过替换Vocos中的逆短时傅里叶变换(iSTFT)层为可训练的线性层,直接预测语音波形样本,而不依赖于STFT频谱。这一改进不仅保持了Vocos的快速推理速度,还提高了语音合成的质量。文章还探讨了如何将WaveNeXt与基于JETS的端到端文本到语音(E2E TTS)框架集成,并研究了采样频率为48kHz的全带模型(Full-band Model:能够处理和生成覆盖整个音频频谱范围的模型,通常是指能够处理从最低频到最高频的完整音频信号的模型)。实验结果表明,WaveNeXt在分析-合成和E2E TTS条件下均优于Vocos,同时保持了快速推理的能力。
-
feature diffusion models 特征扩散模型
DDPM + CFM + NAR(non-autoregressive) model, 没有 codec
-
Base module:
- Denoising Diffusion Probabilistic Model(DDPM): 2020.6 Denoising Diffusion Probabilistic Models | paper code
- Conditional Flow Matching (CFM): 2022.10 Flow Matching for Generative Modeling | CFM lib
-
the alignment modeling between input text and synthesized speech:
-
phoneme-level duration model:
- 2024.5 NaturalSpeech 3 and 2023.6 Voicebox use frame-wise phoneme alignment;
- 2023.9 Matcha-TTS adopts monotonic alignment search(MAS) and relies on phoneme-level duration model;
- 2024.6 E2 TTS 和2024.6 Seed-TTS 研究表明在文本和语音之间引入这种僵化和不灵活的对齐方式会阻碍模型生成更自然的结果。
-
E3 TTS 放弃音素级持续时间并对输入序列应用交叉注意力,但产生的音频质量有限;
-
DiTTo-TTS 使用扩散变换器 (DiT) ,并以来自预训练语言模型的编码文本为条件进行交叉注意。为了进一步增强对齐,它使用预训练的语言模型来微调神经音频编解码器,将语义信息注入生成的表示中;
-
相比之下,基于 Voicebox的 E2 TTS采用了更简单的方法,删除了音素和持续时间预测器,直接使用填充token填充到梅尔频谱图长度的字符作为输入。这个简单的方案也实现了非常自然和真实的合成结果。然而,F5-TTS 发现 E2 TTS 中文本和语音对齐存在鲁棒性问题。
-
2024.6 Seed-TTS 采用了类似的策略并取得了优异的结果,尽管没有详细说明模型细节。在这些未明确建模音素级持续时间的方法中,模型学习根据给定的总序列长度分配每个单词或音素的长度,从而改进韵律和节奏。
-
-
2024.10 F5-TTS: A fairytaler that fakes fluent and faithful speech with flow matching 保持了管道的简单性,无需音素对齐、持续时间预测器、文本编码器和语义注入编解码器模型,利用带有 ConvNeXt V2|paper code 的Diffusion Transformer(DiT)来更好地解决上下文学习期间的文本语音对齐问题。
codec language and feature diffusion hybrid systems 混合系统
text-to-codec language model 和 codec-to-feature diffusion model
语言模型解决文本和语音之间的对齐以及话语持续时间预测,而编解码器到特征扩散模型则根据生成的编解码器和其他条件合成语音特征(梅尔谱)。通过利用两种生成模型的优势,混合系统实现了高度多样性、韵律一致性和语音质量。
language model-based zero-shot TTS models streaming synthesis
- 2024.2 BASE TTS: Lessons from building a billion-parameter Text-to-Speech model on 100K hours of data | 小红书的FireRedTTS 来源于此 FireRedTTS paper code
- 2024.6 LiveSpeech: Low-Latency Zero-shot Text-to-Speech via Autoregressive Modeling of Audio Discrete Codes
- 2024.9 Speak While You Think: Streaming Speech Synthesis During Text Generation
- 2024.10 Zero-Shot Text-to-Speech from Continuous Text Streams
论文解读:CosyVoice: A Scalable Multilingual Zero-shot Text-to-speech Synthesizer based on Supervised Semantic Tokens
CosyVoice
- 2024.7 FunAudioLLM: Voice Understanding and Generation Foundation Models for Natural Interaction Between Humans and LLMs (主要介绍ASR SenseVoice 和 TTS CosyVoice,其中 SenseVoice 没有单独论文,相关CosyVoice 和单独论文是重复的, SenseVoice Large的工作可以用于 CosyVoice 在多语言上, Supervised speech tokenizer 模块的训练和推理)
- 2024.7 CosyVoice: A Scalable Multilingual Zero-shot Text-to-speech Synthesizer based on Supervised Semantic Tokens
- 2024.12 CosyVoice 2: Scalable Streaming Speech Synthesis with Large Language Models (流式合成)
- paper code: 公开推理和权重,在openslr公开数据集英文数据集LibriSpeech 和中文数据集 MAGICDATA 对模型进行训练代码; 无supervised Speech Tokenizer (对SenseVoice ASR的改造微调) 和Speaker Embedding model(context-aware masking CAM++) 的训练过程代码。
创新点:
- 将监督语音token集成到TTS 模型,增强了零样本语音克隆中的内容一致性和说话者相似性。
- 一种可扩展的零样本 TTS 合成系统,它将用于文本到token生成的 LLM 与用于token到语音合成的条件流匹配模型(conditional flow matching model(CFM))相结合,不依赖于音素持续时间预测(Duration predictor),不需要使用补充音素器(phonemizers)和强制对齐器aligners (比如:Glow-TTS中 Monotonic Alignment Search(MAS))。
- 为了进一步细化生成语音的质量,将 x-vector 合并到 LLM 中,将语音建模分为语义、说话者和韵律(semantic, speaker, and prosody)组件。 LLM 对语义(semantic)内容和韵律(prosody)进行建模,而条件流匹配模型(CFM)则捕获音色(timbre)和环境信息。我们使用Classifier-Free Guidance(2022. Classifier-free diffusion guidance)、余弦调度器(cosine scheduler)和屏蔽条件(masked conditions)等技术来优化流匹配过程。
CosyVoice由四个组件组成,即文本编码器(text encoder)、语音分词器(speech tokenizer)、大语言模型(large language model)和条件流匹配模型(conditional flow matching model)
- 文本编码器(text encoder)用于对齐文本和语音token的语义空间;
- 语音标记器(speech tokenizer)用于提取语义token;
- LLM(GLM)学习文本编码和语音标记的整个序列,将 TTS 重新表述为以文本作为提示的自回归序列生成问题;
- 利用条件流匹配模型(conditional flow matching model), 通过最优路径上的去噪处理,将语音标记转换为梅尔谱图(Mel spectrogram); 通过Classifier-Free Guidance(Classifier-free diffusion guidance CFG)提高扩散概率模型的生成质量, 将CFG适应到条件流匹配模型中;
- 获得人类耳朵可感知的声音信号,声码器(vocoder)使用 Hifi-GAN Generator 用于将生成的梅尔频谱图作为输入来合成波形(waveform)。
其中:
- conditional flow matching model (OT-CFM) 来自 2023.9 Matcha-TTS: A fast TTS architecture with conditional flow matching(CFM的改进版本OT-CFM)
- Classifier-free diffusion guidance (CFG) 来自 2022. Classifier-free diffusion guidance
- vocoder 来自 2020. HiFi-GAN: Generative Adversarial Networks for Efficient and High Fidelity Speech Synthesis Generator。
附:
论文解读 Matcha-TTS: A fast TTS architecture with conditional flow matching
🍵 matcha-tts
摘要:
一种用于快速 TTS 声学建模的新编码器-解码器架构,使用最佳传输条件流匹配 (OT-CFM) 进行训练。与使用分数匹配训练的模型相比,这产生了基于 ODE 的解码器,能够以更少的合成步骤实现高输出质量。仔细的设计选择还确保每个合成步骤都能快速运行。该方法是概率性的、非自回归的,并且无需外部对齐即可从头开始学习说话。与强大的预训练基线模型相比,Matcha-TTS 系统具有最小的内存占用,可以与长语音上最快的模型相媲美,并在听力测试中获得最高的平均意见得分。
一种非自回归神经 TTS 的新方法,它使用条件流匹配(CFM from 2022.10 Flow Matching for Generative Modeling)(类似于校正流(Rectified Flow 2022.9 Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow)来加速基于 ODE 的语音合成。
- Is probabilistic 是概率性的
- Has compact memory footprint 具有紧凑的内存占用
- Sounds highly natural 听起来非常自然
- Is very fast to synthesise from 合成速度非常快
简洁的结构,训练推理快,使用更少额内存空间,
一种基于连续归一化流的概率性、非自回归、快速采样的 TTS 声学模型。主要有两个创新点:
- 提出了一种改进的编码器-解码器 TTS 结构,该结构在解码器中结合使用 1D CNN 和 Transformer。这减少了内存消耗并且可以快速评估,从而提高综合速度。
- 使用最优传输条件流匹配 optimal-transport conditional flow matching(OT-CFM) 来训练这些模型,这是一种学习从数据分布中采样的 ODE 的新方法。与传统的 连续时间归一化流 CNF(continuous-time normalising flows ) 和分数匹配概率流 ODE (probability flow ODE) 相比,OT-CFM 定义了从源到目标的更简单的路径,从而能够以比 DPM(Diffusion probabilistic model) 更少的步骤进行准确合成。
实验结果表明,这两种创新都加速了合成,减少了速度和合成质量之间的权衡。尽管速度快且轻量级,Matcha-TTS能够在不需要外部对齐器的情况下学习说话和对齐。与强大的预训练基线模型相比,Matcha-TTS实现了快速合成,并获得了更好的自然度评分。
附:
-
使用LJ Speech数据集训练202 epochs, 在1xGPU 上训练了 79,372 step 的ckpt(论文中是2x3080GPU 500K step):
-
笔记地址: https://github.com/weedge/doraemon-nb/blob/main/matcha_tts.ipynb
论文解读: VALL-E 系列
前文讲到VITS,采用的端到端的NAR模型,这篇文章记录下微软提出的VALL-E系列,从 AR+NAR 到 AR 模型的转变,以及后面MELLE引入的VAE和Mel-Spectorgram,将neural codec text speech LM (AR+NAR Transformer Decoder) 转变为 autoregressive mel-spectrogram text speech LM (AR Transformer Decoder) ;由于LM生成的是mel-spectrogram 需要通过vocoder 转换成 waveform; 生成的内容采样模块:从top-p random sampling 变成 Latent Sampling潜在采样模块(思想源自VAE, 从预测的高斯分布中采样潜在嵌入,然后将其投影回频谱图空间)
VALL-E 系列
https://www.microsoft.com/en-us/research/project/vall-e-x/
Vall-E: https://ar5iv.labs.arxiv.org/html/2301.02111
具体来说,我们使用从现成的神经音频编解码器模型派生的离散代码来训练神经编解码器语言模型(neural codec language model)(称为VALL-E ),并将 TTS 视为条件语言建模任务,而不是像之前的工作那样将 TTS 视为连续信号回归。 在预训练阶段,我们将 TTS 训练数据扩展到 60K 小时的英语语音,比现有系统大数百倍。 VALL-E具有情境学习功能,可用于合成高质量的个性化语音,只需对看不见的说话者进行 3 秒的注册录音作为声音提示。实验结果表明, VALL-E在语音自然度和说话人相似度方面显着优于最先进的零样本 TTS 系统。此外,我们发现VALL-E可以在合成时保留说话者的情感和声音提示的声学环境。
与之前的管道不同(例如,音素 → 梅尔谱图 → 波形), VALL-E的管线是音素 → 离散码 → 波形。
VALL-E根据音素和声学代码提示生成与目标内容和说话者的声音相对应的离散音频编解码器代码。 VALL-E直接支持各种语音合成应用,例如零样本 TTS、语音编辑以及与 GPT-3 等其他生成式 AI 模型相结合的内容创建。
VALL-E系列:2023年的1月份开始 - 2024年的7月份
- VALL-E 使用从现成的神经音频编解码器模型派生的离散代码来训练神经编解码器语言模型,并将 TTS 视为条件语言建模任务,而不是像之前的工作那样将 TTS 视为连续信号回归。 VALL-E 具有情境学习功能,可用于合成高质量的个性化语音,只需录制未见过的讲话者的 3 秒注册录音作为提示。在语音自然度和说话人相似度方面,VALL-E 显着优于最先进的零样本 TTS 系统。此外,VALL-E可以在合成时保留说话者的情绪和声音提示的声学环境。
- VALL-E X 扩展其能力,适应多语言场景,促进跨语言零样本 TTS。
- VALL-E R 引入了音素单调对齐策略,增强了语音生成的鲁棒性。
- VALL-E 2 通过集成重复感知采样和分组代码建模技术, 实现了一个突破性的里程碑:在 LibriSpeech 和 VCTK 数据集上的零样本 TTS 性能与人类相当。这标志着此类成就的首次实例,为该领域树立了新标准。
- MELLE 是一种新颖的基于连续值标记的语言建模方法,用于文本到语音合成 (TTS)。 MELLE 直接从文本条件自回归生成连续的梅尔频谱图帧,绕过了矢量量化的需要,矢量量化最初是为音频压缩而设计的,与梅尔频谱图相比,牺牲了保真度。