如何在训练中集成并配置 DeepSpeed Transformer Kernel【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed用 DeepSpeed 训练 Transformer 类模型教程以 BERT 预训练为例时可以把模型中的 Transformer 层替换为 DeepSpeed 自研的 transformer kernel 来加速训练。整条操作路径分三步先把 DeepSpeed 集成进训练脚本并用deepspeed命令拉起训练再在脚本里加--deepspeed_transformer_kernel开关并在模型内部用DeepSpeedTransformerConfigDeepSpeedTransformerLayer实例化 kernel 版层最后按需打开内存优化 flags让每张卡能装下更大的 micro-batch。官方 DeepSpeed Transformer Kernel 教程给出参数含义与 V100 上的内存配置参考BERT Pre-training 教程给出接入 bing_bert 训练代码的完整改法两篇文档配合即可完成本文任务。准备把 DeepSpeed 集成进训练脚本使用 transformer kernel 的前提是训练脚本已经能跑通 DeepSpeed即完成 Getting Started 指南中的集成安装pip install deepspeed用deepspeed.initialize包装模型model_engine, optimizer, _, _ deepspeed.initialize(argscmd_args, modelmodel, model_parametersparams)返回的model_engine用于前向model_engine(batch)、反向model_engine.backward(loss)和权重更新model_engine.step()分布式与混合精度的底层设置由 DeepSpeed 自动完成。如果脚本里原本有torch.distributed.init_process_group(...)文档要求改为deepspeed.init_distributed()。一个硬性限制要先明确DeepSpeed Transformer Kernels 目前不支持 Sparse Attention两者不能同时开启要用 Sparse Attention 就必须禁用 Transformer Kernels。在训练脚本中加 --deepspeed_transformer_kernel 开关BERT Pre-training 教程以 DeepSpeedExamples 仓库中的 bing_bert 示例代码为例先在其utils.py中加一个开关参数方便随时开关 kernelparser.add_argument(--deepspeed_transformer_kernel, defaultFalse, actionstore_true, helpUse DeepSpeed transformer kernel to accelerate.)同时按 Getting Started 指南调用deepspeed.add_config_arguments(parser)让脚本能识别--deepspeed、--deepspeed_config等 DeepSpeed 参数。在模型中实例化 kernel 版 Transformer 层开关打开后在模型的编码器类bing_bert 示例中是BertEncoder里用 kernel 版层替代原来的BertLayer。文档示例代码if args.deepspeed_transformer_kernel: from deepspeed import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig, DeepSpeedConfig if hasattr(args, deepspeed_config) and args.deepspeed_config: ds_config DeepSpeedConfig(args.deepspeed_config) else: raise RuntimeError(deepspeed_config is not found in args.) cuda_config DeepSpeedTransformerConfig( batch_size ds_config.train_micro_batch_size_per_gpu, max_seq_length args.max_seq_length, hidden_size config.hidden_size, heads config.num_attention_heads, attn_dropout_ratio config.attention_probs_dropout_prob, hidden_dropout_ratio config.hidden_dropout_prob, num_hidden_layers config.num_hidden_layers, initializer_range config.initializer_range, local_rank args.local_rank if hasattr(args, local_rank) else -1, seed args.seed, fp16 ds_config.fp16_enabled, pre_layer_normTrue, attn_dropout_checkpointargs.attention_dropout_checkpoint, normalize_invertibleargs.normalize_invertible, gelu_checkpointargs.gelu_checkpoint, stochastic_modeTrue) layer DeepSpeedTransformerLayer(cuda_config) else: layer BertLayer(config) self.layer nn.ModuleList([copy.deepcopy(layer) for _ in range(config.num_hidden_layers)])参数值来自两处DeepSpeed JSON 配置文件通过DeepSpeedConfig读取和训练脚本的命令行参数args所以要把args传到模型构造入口。文档明确交代了四件事batch_size是 kernel 可接受的输入 batch 上限所有训练或预测数据不能超过它否则抛异常。例如配置文件里train_micro_batch_size_per_gpu是 8而预测用 12就要把 kernel 的 batch size 设成 12或用--predict_batch_size把预测 batch 压到 8 或更小。local_rank用于把 kernel 放到正确的 GPU如果模型在此之前已经调过set_device()就不需要再设置。stochastic_mode打开后性能更高但有一定非确定性不同 run 结果可能不同。文档在预训练场景启用它微调时建议关闭以保证可复现。checkpoint 兼容kernel 有自己的参数布局用 transformer kernel 生成的 checkpoint 必须由同样开启 transformer kernel 的模型加载比如微调时。Transformer kernel 教程 还给出了一个独立的构造示例Pre-LN BERT-Large24 层、hidden 1024、seq 128、batch 64、fp16True并把全部参数分为四类通用配置batch_size、max_seq_length、hidden_size、heads、两个 dropout ratio、num_hidden_layers、pre_layer_norm、环境参数local_rank、seed、fp16、initializer_range、高性能 flagstochastic_mode文档称打开后训练平均快 2%和内存优化 flags。准备 DeepSpeed JSON 配置文件用 kernel 训练需要一份 DeepSpeed 配置文件。文档示例deepspeed_bsz4096_adam_config.json{ train_batch_size: 4096, train_micro_batch_size_per_gpu: 64, steps_per_print: 1000, optimizer: { type: Adam, params: { lr: 2e-4, max_grad_norm: 1.0, weight_decay: 0.01, bias_correction: false } }, fp16: { enabled: true, loss_scale: 0, initial_scale_power: 16 } }文档对其中四项的解释train_batch_size为有效 batch 4096train_micro_batch_size_per_gpu为每卡 micro-batch 64——这个值就是上面传给DeepSpeedTransformerConfig的batch_sizeoptimizer用 Adamfp16开启 FP16 混合精度初始 loss scale 为 2^16。选择内存优化 flags 以放大 micro-batch当单卡显存放不下目标 micro-batch 时kernel 提供三个内存优化 flags用部分算力换激活内存从而支持更大 batch。文档的结论根据性能剖析后两项 checkpoint 的重新计算开销可忽略且更大 batch 带来的收益可以弥补。Flag文档说明的作用attn_dropout_checkpoint对 attention dropout 做 checkpoint省激活内存normalize_invertible可逆 LayerNorm 执行方式不保存 normalize 层的输入激活gelu_checkpoint对 Gelu 激活输出做 checkpoint省激活内存下面是文档给出的参考表在 NVIDIA V10032GB 显存上跑 BERT-Large 时不同 micro-batch 与序列长度需要打开的 flags数值为文档测量示例不是其他硬件上的固定预期Micro-batch size128 sequence-length512 sequence-length 12-attn_dropout_checkpoint 16-normalize_invertible,gelu_checkpoint 80normalize_invertibleOOM 112attn_dropout_checkpointOOM 128gelu_checkpointOOM也就是说micro-batch 超过第一行阈值后需要按对应列打开相应 flags 才能继续运行512 序列长度下超过 80 的 micro-batch 在文档中记录为 OOM。用 deepspeed 命令启动训练文档示例为 4 节点、每节点 4 卡的启动命令deepspeed --num_nodes 4 \ deepspeed_train.py \ --deepspeed \ --deepspeed_config deepspeed_bsz4096_adam_config.json \ --cf /path-to-deepspeed/examples/tests/bing_bert/bert_large_adam_seq128.json \ --train_batch_size 4096 \ --max_seq_length 128 \ --gradient_accumulation_steps 4 \ --max_grad_norm 1.0 \ --fp16 \ --loss_scale 0 \ --delay_allreduce \ --max_steps 32 \ --print_steps 1 \ --deepspeed_transformer_kernel \ --output_dir output_directory其中--cf指向模型自身的配置文件、--output_dir指向输出目录这两处文档示例路径需要替换为你环境中的实际路径。--deepspeed--deepspeed_config表示以 DeepSpeed 模式运行并传入上面的 JSON 配置--deepspeed_transformer_kernel是前文在脚本中定义的开关--max_steps 32与--print_steps 1把示例限制在 32 个训练步并每步打印一次便于先小步数验证启动。单节点启动不需要 hostfiledeepspeed会自动探测本机 GPU 数量见 Getting Started。Transformer kernel 教程 另外展示了一个带内存优化选项和断点恢复参数的启动方式用于 512 序列长度、每卡 micro-batch 16 的场景deepspeed deepspeed_train.py \ --cf bert_large_lamb.json \ --max_seq_length 512 \ --print_steps 100 \ --deepspeed \ --deepspeed_transformer_kernel \ --deepspeed_config deepspeed_bsz32K_lamb_config_seq512.json \ --rewarmup \ --lr_schedule EE \ --lr_offset 0.0 \ --attention_dropout_checkpoint \ --load_training_checkpoint ${CHECKPOINT_BASE_PATH} \ --load_checkpoint_id ${CHECKPOINT_EPOCH150_NAME}这条命令中${CHECKPOINT_BASE_PATH}和${CHECKPOINT_EPOCH150_NAME}是文档原例中的环境变量需要替换为你的 checkpoint 目录与 checkpoint 标识不需要恢复断点时删除最后两行即可。该命令使用 Lamb 优化器的配置文件deepspeed_bsz32K_lamb_config_seq512.json完整脚本与 JSON 在文档中指向 DeepSpeedExamples 仓库的ds_train_bert_bsz64k_seq128.sh、ds_train_bert_bsz32k_seq512.sh等脚本。结果判断与已知限制小步数试跑文档示例用--max_steps 32 --print_steps 1限制步数并逐步打印可作为启动验证的最小配置。显存预期对照上面的 V100 32GB 参考表判断 flags 是否够——超过阈值而未开对应 flags 时文档记录的现象是 OOM。batch 上限异常输入 batch 超过DeepSpeedTransformerConfig的batch_size时 kernel 会抛异常这是文档明确给出的报错场景。checkpoint 加载开启 kernel 生成的 checkpoint 必须由开启 kernel 的模型加载跨开关状态加载不成立。文档与源码的参数不一致两篇教程的示例代码都向DeepSpeedTransformerConfig传入max_seq_length而当前仓库源码 transformer.py 中的构造函数参数列表是intermediate_size、没有max_seq_length。这是文档与源码的滞后接入前请以你所装版本的构造函数签名为准核对参数名。性能数字文档示例为 BERT-Large 在 V100 上序列长度 128/512 时达到最高 64/53 teraflops较 NVIDIA BERT 最多高 28%、较 HuggingFace BERT 最多高 62%并支持最多 1.8 倍更大 batch 而不 OOM。这些是文档的测量数字不构成其他硬件上的预期值。下一步BERT Pre-training 教程指出预训练完成后可以按 BERT fine-tuning 教程微调由 transformer kernel 预训练的模型来复现 SQUAD F1 分数结合前文的 checkpoint 兼容性要求微调时需要保持 kernel 开启并按文档建议关闭stochastic_mode。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
