更新时间:2025-03-04 gmt 08:00
模型训练存储加速-九游平台
针对ai训练场景中大模型checkpoint保存和加载带来的i/o挑战,华为云提供了基于对象存储服务obs 高性能弹性文件服务sfs turbo的ai云存储九游平台的解决方案,如下图所示。
sfs turbo hpc型支持和obs数据联动,您可以通过sfs turbo hpc型文件系统来加速对obs对象存储中的数据访问,并将生成的结果数据异步持久化到obs对象存储中长期低成本保存。
图1 基于obs sfs turbo的存储九游平台的解决方案

obs sfs turbo存储加速的具体方案请参见面向ai场景使用obs sfs turbo的存储加速实践。
当训练程序从已有checkpoint恢复时,每张卡都需要从持久化存储中加载,在训练集群规模较大,存储带宽较低的场景下,加载耗时可能会达到小时级,严重影响训练恢复。因此,通过aiturbo sdk的快速保存和加载checkpoint的功能,可以有效提升训练恢复速度。具体方案请参见。
- 在保存checkpoint的时候,利用两阶段写、内存副本、异步持久化等技术保证checkpoint的快速、高可靠存储。
- 在加载checkpoint的时候,利用内存快恢、checkpoint广播等技术,大大减少后端存储的带宽压力,提升加载效率。
设置训练存储加速
当完成步骤后,在modelarts standard中创建训练作业时,设置训练“sfs turbo”,在“文件系统”中选择sfs turbo实例名称,并指定“存储位置”和“云上挂载路径”。系统会在训练作业启动前,自动将存储位置中的文件目录挂载到训练容器中指定路径。
图2 设置训练“sfs turbo”


当前训练作业支持挂载多个弹性文件服务sfs turbo,文件系统支持重复挂载,但挂载路径不可重复。文件系统目录需指定已存在的目录,否则会导致训练作业异常。
然后在超参或者环境变量中设置checkpoint和数据的挂载路径。
图3 在超参或者环境变量中设置checkpoint和数据的挂载路径

训练存储加速的代码样例(pytorch版reload ckpt)
- pytorch模型保存有两种方式。
- 仅保存模型参数
state_dict = model.state_dict() torch.save(state_dict, path)
- 保存整个model(不推荐)
torch.save(model, path)
- 仅保存模型参数
- 可根据step步数、时间等周期性保存模型的训练过程的产物。
将模型训练过程中的网络权重、优化器权重、以及epoch进行保存,便于中断后继续训练恢复。
checkpoint = { "net": model.state_dict(), "optimizer": optimizer.state_dict(), "epoch": epoch } if not os.path.isdir('model_save_dir'): os.makedirs('model_save_dir') torch.save(checkpoint,'model_save_dir/ckpt_{}.pth'.format(str(epoch)))
- 完整代码示例。
import os import argparse parser = argparse.argumentparser() parser.add_argument("--train_url", type=str) args, unparsed = parser.parse_known_args() args = parser.parse_known_args() # train_url 将被赋值为"/home/ma-user/modelarts/outputs/train_url_0" train_url = args.train_url # 判断输出路径中是否有模型文件。如果无文件则默认从头训练,如果有模型文件,则加载epoch值最大的ckpt文件当做预训练模型。 if os.listdir(train_url): print('> load last ckpt and continue training!!') last_ckpt = sorted([file for file in os.listdir(train_url) if file.endswith(".pth")])[-1] local_ckpt_file = os.path.join(train_url, last_ckpt) print('last_ckpt:', last_ckpt) # 加载断点 checkpoint = torch.load(local_ckpt_file) # 加载模型可学习参数 model.load_state_dict(checkpoint['net']) # 加载优化器参数 optimizer.load_state_dict(checkpoint['optimizer']) # 获取保存的epoch,模型会在此epoch的基础上继续训练 start_epoch = checkpoint['epoch'] start = datetime.now() total_step = len(train_loader) for epoch in range(start_epoch 1, args.epochs): for i, (images, labels) in enumerate(train_loader): images = images.cuda(non_blocking=true) labels = labels.cuda(non_blocking=true) # forward pass outputs = model(images) loss = criterion(outputs, labels) # backward and optimize optimizer.zero_grad() loss.backward() optimizer.step() ... # 保存模型训练过程中的网络权重、优化器权重、以及epoch checkpoint = { "net": model.state_dict(), "optimizer": optimizer.state_dict(), "epoch": epoch } if not os.path.isdir(train_url): os.makedirs(train_url) torch.save(checkpoint, os.path.join(train_url, 'ckpt_best_{}.pth'.format(epoch)))
训练存储加速的代码样例(mindspore版reload ckpt)
import os import argparse from resnet import resnet50 from mindspore.nn.optim.momentum import momentum from mindspore.nn.loss import softmaxcrossentropywithlogits from mindspore import load_checkpoint, load_param_into_net from mindspore.train import model, checkpointconfig, modelcheckpoint from mindspore.train.callback import lossmonitor parser = argparse.argumentparser() parser.add_argument("--train_url", type=str) parser.add_argument("--batch_size", type=int, default=32, help="batch size.") parser.add_argument("--num_classes", type=int, default=10, help="num classes.") parser.add_argument("--do_train", type=bool, default=true, help="do train or not.") args_opt, unparsed = parser.parse_known_args() # train_url 将被赋值为"/home/ma-user/modelarts/outputs/train_url_0" 。 train_url = args_opt.train_url # 初始定义的网络、损失函数及优化器,详细请参见。 # 1.初始定义的网络,以“resnet50”为例。详细请参见。 net = resnet50(args_opt.batch_size, args_opt.num_classes) # 2.定义损失函数,详细请参见。 ls = softmaxcrossentropywithlogits(sparse=true, reduction="mean") # 3.定义优化器,详细请参见。 opt = momentum(filter(lambda x: x.requires_grad, net.get_parameters()), 0.01, 0.9) # 首次训练的epoch初始值,mindspore1.3及以后版本会支持定义epoch_size初始值。 # cur_epoch_num = 0 # 判断输出obs路径中是否有模型文件。如果无文件则默认从头训练,如果有模型文件,则加载epoch值最大的ckpt文件当做预训练模型。 if os.listdir(train_url): last_ckpt = sorted([file for file in os.listdir(train_url) if file.endswith(".ckpt")])[-1] print('last_ckpt:', last_ckpt) last_ckpt_file = os.path.join(train_url, last_ckpt) # 加载断点,详细请参见。 param_dict = load_checkpoint(last_ckpt_file) print('> load last ckpt and continue training!!') # 加载模型参数到net。 load_param_into_net(net, param_dict) # 加载模型参数到opt。 load_param_into_net(opt, param_dict) # 获取保存的epoch值,模型会在此epoch的基础上继续训练,此参数在mindspore1.3及以后版本会支持。 # if param_dict.get("epoch_num"): # cur_epoch_num = int(param_dict["epoch_num"].data.asnumpy()) model = model(net, loss_fn=ls, optimizer=opt, metrics={'acc'}) # as for train, users could use model.train if args_opt.do_train: dataset = create_dataset() batch_num = dataset.get_dataset_size() config_ck = checkpointconfig(save_checkpoint_steps=batch_num, keep_checkpoint_max=35) # append_info=[{"epoch_num": cur_epoch_num}],mindspore1.3及以后版本会支持append_info参数,保存当前时刻的epoch值。 # 保存网络参数,详细请参见。 ckpoint_cb = modelcheckpoint(prefix="train_resnet_cifar10", directory=args_opt.train_url, config=config_ck) loss_cb = lossmonitor() model.train(epoch_size, dataset, callbacks=[ckpoint_cb, loss_cb]) # model.train(epoch_size-cur_epoch_num, dataset, callbacks=[ckpoint_cb, loss_cb]),mindspore1.3及以后版本支持从断点恢复训练。
相关文档
意见反馈
文档内容是否对您有帮助?
提交成功!非常感谢您的反馈,我们会继续努力做到更好!
您可在查看反馈及问题处理状态。
系统繁忙,请稍后重试
如您有其它疑问,您也可以通过华为云社区问答频道来与我们联系探讨