飞桨大模型统一存储文档#
1. 背景#
在大模型背景下,通常我们需要进行多卡分布式的训练,在保存Checkpoint时所得到的模型权重通常是分片放置的,例如根据张量并行、流水线并行进行切分保存。这种根据分布式策略直接存储Checkpoint的方式非常直接明了,但也存在如下的问题:
对下游推理不够友好,当用户希望获取中间阶段保存的Checkpoint做下游推理时,需要手动对模型权重进行合并。
不利于应对做恢复训练时,可能会面临的分布式策略改变、训练节点数发生变化的情况。用户往往需要手动对Checkpoint进行处理,增加了操作复杂度。
为了最大程度地解决上述的问题,降低用户操作难度,我们提出了大模型统一存储方案——Unified Checkpoint。Unified Checkpoint的核心思想是将模型权重、优化器权重等进行统一safetensors格式存储,在Checkpoint存储时不再对分布式策略进行区分,提高大模型存储的通用性。以下将首先介绍Unified Checkpoint具体存储格式以及如何使用,随后再简要介绍统一存储的实现原理。
2. 统一存储 Unified Checkpoint 使用介绍#
2.1 使用命令与配置项说明#
使用示例
python -u -m paddle.distributed.launch \
--gpus "0,1,2,3,4,5,6,7" \
run_pretrain.py \
--unified_checkpoint 1 \
--unified_checkpoint_config "enable_all_options"
总开关
unified_checkpoint
用于控制是否使用Unified Checkpoint存储格式。
unified_checkpoint: Optional[bool] = field(
default=False,
metadata={"help": "Whether to unify hybrid parallel checkpoint."},
)
配置项说明
unified_checkpoint_config: Optional[str] = field(
default="",
metadata={
"help": (
"Configs to unify hybrid parallel checkpoint.\n"
"Following options are supports:\n"
"- skip_save_model_weight: do not save model weights when the masters weight exist\n"
"- master_weight_compatible: 1. if the master weights exist, only load when needed\n"
" 2. if master weights does not exist, convert model weights to master weights when needed\n"
"- enable_all_options: enable all optimization configurations\n"
)
},
)
介绍如下:
skip_save_model_weight:当optimizer具有master weight时,跳过model weight保存,重启时将master weight作为model weight加载。在PaddleNLP中,仅fp16_opt_level=O1时,optimizer不存在master weight。
master_weight_compatible:仅当optimizer需要master weight时,才加载master weight; 如果ckpt中不存在master weight,将model weight作为master weight加载。
enable_all_options:上述参数均开启。
2.2 Unified Checkpoint存储格式介绍#
这里以facebook/llama-7b的pretrain checkpoint保存为例进行说明。以TP=4,PP=2的分布式训练为例,原始的存储格式举例如下代码片段。无论是模型参数,异或是优化器参数,均按照TP、PP训练方式进行了分片存储。
-rw-r--r-- 1 root root 1015 Dec 21 11:27 config.json
-rw-r--r-- 1 root root 1.6G Dec 21 11:27 model_state.tp00_pp00.pdparams
-rw-r--r-- 1 root root 1.6G Dec 21 11:27 model_state.tp00_pp01.pdparams
-rw-r--r-- 1 root root 1.6G Dec 21 11:27 model_state.tp01_pp00.pdparams
-rw-r--r-- 1 root root 1.6G Dec 21 11:27 model_state.tp01_pp01.pdparams
-rw-r--r-- 1 root root 1.6G Dec 21 11:27 model_state.tp02_pp00.pdparams
-rw-r--r-- 1 root root 1.6G Dec 21 11:27 model_state.tp02_pp01.pdparams
-rw-r--r-- 1 root root 1.6G Dec 21 11:27 model_state.tp03_pp00.pdparams
-rw-r--r-- 1 root root 1.6G Dec 21 11:27 model_state.tp03_pp01.pdparams
-rw-r--r-- 1 root root 9.5G Dec 21 11:27 optimizer.tp00_pp00.pdopt
-rw-r--r-- 1 root root 9.5G Dec 21 11:27 optimizer.tp00_pp01.pdopt
-rw-r--r-- 1 root root 9.5G Dec 21 11:27 optimizer.tp01_pp00.pdopt
-rw-r--r-- 1 root root 9.5G Dec 21 11:27 optimizer.tp01_pp01.pdopt
-rw-r--r-- 1 root root 9.5G Dec 21 11:27 optimizer.tp02_pp00.pdopt
-rw-r--r-- 1 root root 9.5G Dec 21 11:27 optimizer.tp02_pp01.pdopt
-rw-r--r-- 1 root root 9.5G Dec 21 11:27 optimizer.tp03_pp00.pdopt
-rw-r--r-- 1 root root 9.5G Dec 21 11:27 optimizer.tp03_pp01.pdopt
-rw-r--r-- 1 root root 54K Dec 21 11:27 rng_state_8.pth
-rw-r--r-- 1 root root 317 Dec 21 11:27 scaler.pdparams
-rw-r--r-- 1 root root 50 Dec 21 11:27 scheduler.pdparams
-rw-r--r-- 1 root root 489K Dec 21 11:27 sentencepiece.bpe.model
-rw-r--r-- 1 root root 63 Dec 21 11:27 special_tokens_map.json
-rw-r--r-- 1 root root 207 Dec 21 11:27 tokenizer_config.json
-rw-r--r-- 1 root root 3.1K Dec 21 11:27 trainer_state.json
-rw-r--r-- 1 root root 2.3K Dec 21 11:27 training_args.bin
采用Unified Checkpoint进行统一存储后,新格式如下面代码片段。可以看到,无论是模型参数、优化器参数,我们均采用了safetensors格式进行存储,不再区分TP、PP策略;进一步地,我们将优化器参数区分为了optimizer与master_weights(如果有的话),而master_weights本身就是模型参数的FP32版本。其中,model.safetensors.index.json
等json文件用于记录参数对应所在的文件部分。
-rw-r--r-- 1 root root 1015 Dec 21 11:24 config.json
-rw-r--r-- 1 root root 3.1G Dec 21 11:25 master_weights-00001-of-00008.safetensors
-rw-r--r-- 1 root root 3.2G Dec 21 11:25 master_weights-00002-of-00008.safetensors
-rw-r--r-- 1 root root 3.2G Dec 21 11:25 master_weights-00003-of-00008.safetensors
-rw-r--r-- 1 root root 3.2G Dec 21 11:25 master_weights-00004-of-00008.safetensors
-rw-r--r-- 1 root root 3.1G Dec 21 11:25 master_weights-00005-of-00008.safetensors
-rw-r--r-- 1 root root 3.2G Dec 21 11:25 master_weights-00006-of-00008.safetensors
-rw-r--r-- 1 root root 3.1G Dec 21 11:25 master_weights-00007-of-00008.safetensors
-rw-r--r-- 1 root root 3.3G Dec 21 11:25 master_weights-00008-of-00008.safetensors
-rw-r--r-- 1 root root 28K Dec 21 11:25 master_weights.safetensors.index.json
-rw-r--r-- 1 root root 1.6G Dec 21 11:24 model-00001-of-00008.safetensors
-rw-r--r-- 1 root root 1.6G Dec 21 11:24 model-00002-of-00008.safetensors
-rw-r--r-- 1 root root 1.6G Dec 21 11:24 model-00003-of-00008.safetensors
-rw-r--r-- 1 root root 1.6G Dec 21 11:24 model-00004-of-00008.safetensors
-rw-r--r-- 1 root root 1.6G Dec 21 11:24 model-00005-of-00008.safetensors
-rw-r--r-- 1 root root 1.6G Dec 21 11:24 model-00006-of-00008.safetensors
-rw-r--r-- 1 root root 1.6G Dec 21 11:24 model-00007-of-00008.safetensors
-rw-r--r-- 1 root root 1.7G Dec 21 11:24 model-00008-of-00008.safetensors
-rw-r--r-- 1 root root 25K Dec 21 11:24 model.safetensors.index.json
-rw-r--r-- 1 root root 6.2G Dec 21 11:25 optimizer-00001-of-00008.safetensors
-rw-r--r-- 1 root root 6.4G Dec 21 11:25 optimizer-00002-of-00008.safetensors
-rw-r--r-- 1 root root 6.2G Dec 21 11:25 optimizer-00003-of-00008.safetensors
-rw-r--r-- 1 root root 6.4G Dec 21 11:25 optimizer-00004-of-00008.safetensors
-rw-r--r-- 1 root root 6.3G Dec 21 11:25 optimizer-00005-of-00008.safetensors
-rw-r--r-- 1 root root 6.4G Dec 21 11:25 optimizer-00006-of-00008.safetensors
-rw-r--r-- 1 root root 6.3G Dec 21 11:25 optimizer-00007-of-00008.safetensors
-rw-r--r-- 1 root root 6.4G Dec 21 11:25 optimizer-00008-of-00008.safetensors
-rw-r--r-- 1 root root 118K Dec 21 11:25 optimizer.safetensors.index.json
-rw-r--r-- 1 root root 54K Dec 21 11:25 rng_state_8.pth
-rw-r--r-- 1 root root 317 Dec 21 11:25 scaler.pdparams
-rw-r--r-- 1 root root 50 Dec 21 11:25 scheduler.pdparams
-rw-r--r-- 1 root root 489K Dec 21 11:24 sentencepiece.bpe.model
-rw-r--r-- 1 root root 63 Dec 21 11:24 special_tokens_map.json
-rw-r--r-- 1 root root 207 Dec 21 11:24 tokenizer_config.json
-rw-r--r-- 1 root root 3.1K Dec 21 11:25 trainer_state.json
-rw-r--r-- 1 root root 2.3K Dec 21 11:24 training_args.bin
其中,safetensors是由huggingface开发的一种新序列化格式,旨在简化和精简大型复杂张量的存储和加载。使用Safetensors有很多好处,这里简要列举部分如下:
速度快:Safetensors采用Zero-copy技术进行了速度优化,可以高效处理大型张量的序列化和反序列化。
大小优化:混合使用了有效的序列化和压缩算法,以减少大型张量的大小,相比于其他序列化格式(如pickle),性能更快、更高效。
懒惰加载:Safetensors在加载参数时只需要加载文件中需要的部分张量即可,效率更高。
安全性:为了防止序列化张量在存储或传输过程中出现损坏,Safetensors使用了校验和机制。这保证了额外的安全性,确保存储在Safetensors中的所有数据都准确可靠。
2.3 训练分布式策略发生变化时怎么办?#
在Unified checkpoint统一存储格式下,当训练分布式策略不变时,我们直接原地加载Checkpoint进行训练即可。那么,当分布式策略发生变化时,应当怎么做?以下区分两种情况进行讨论。
2.3.1 机器不变#
在训练机器不变的情况下,进行分布式策略的改变有多种情况,简单举例如下:
例如单机训练,希望TP=8转为TP=4、Sharding=2进行训练;
例如两机训练时,希望TP=8、Sharding=2转为PP=8、Sharding=2训练;
又或者是希望在相同机器的情况下减少参与训练的进程数(GPU卡数)。 在这些情况下,我们都不需要对checkpoint进行处理,只需要进行重启操作即可,Unified checkpoint会自动加载文件并执行相应的Tensor切分、发送接收等操作。
2.3.2 机器数量发生变化,例如1->多,多->1,多->多#
尽管机器数量发生变化的情况很多,用户在处理Checkpoint时原则上只需要保证:新的训练机器上至少需要有一份完整的Checkpoint参数,这份完整参数可以放置在同一台机器上,也可以分散放置在多台机器上(如果为多机训练的话)。
1->多:例如,原先我们在机器A上训练,Checkpoint存储在A上,接下来想用A、B两台机器同时训练,此时需要确保机器A、B上有一份完整Checkpoint参数即可。
多->1:例如,原先我们在两台机器A、B训练,Checkpoint可能分了两台机器进行存储,接下来想只在机器A上进行训练,那么需要将Checkpoint文件完整放置在机器A上。
多->多:例如,原先我们在机器A、B上训练,Checkpoint存储在A、B上,接下来想用四台机器(A、B、C、D)训练,此时需要确保参与训练的四台机器上具备一份完整Checkpoint参数即可。 用户只需要确保参与训练的机器上具备一份完整的Checkpoint参数,即可进行重启训练。
2.4 旧格式的Checkpoint如何兼容?#
在打开unified_checkpoint开关后,我们会对Checkpoint文件夹中的内容进行检查。
如果文件夹中具备旧格式的参数文件等,我们会按照旧格式的方式进行参数加载,在后续保存新的Checkpoint时会保存成Unified Checkpoint的格式。
如果文件夹中不含旧格式参数文件,则默认采用Unified Checkpoint格式进行加载。我们会检查参与训练的机器中的参数文件是否完整,如完整则直接加载训练,否则会进行报错。