对接DeepSpeed框架
- 使用业务用户登录到计算节点。
业务用户非<MindIO-install-user>、HwHiAiUser、hwMindX用户,由用户根据实际情况决定。
- 进入DeepSpeed安装目录。
cd deepspeed安装目录/runtime
- 修改engine.py文件。
- 执行以下命令,打开engine.py文件。
vim engine.py
- 按“i”进入编辑模式,修改如下内容。
- 在文件首行加入以下内容。
import torch_mindio
- 将torch.load函数替换为torch_mindio.load函数。
替换后:
optim_checkpoint = torch_mindio.load(optim_load_path)
- 将普通的torch.save函数替换torch_mindio.save函数。
替换后:
torch_mindio.save(state, save_path)
- 将包含torch.save函数的with open语句替换为torch_mindio.save函数。
替换前:
with open(self._get_optimizer_ckpt_name(save_dir, tag, expp_rank), 'wb') as fd: torch.save(optimizer_state, fd) fd.flush()
替换后:
torch_mindio.save(optimizer_state, self._get_optimizer_ckpt_name(save_dir, tag, expp_rank))
- 替换DeepSpeedEngine._get_expert_ckpt_name函数。
expert_state_dict = torch.load(DeepSpeedEngine._get_expert_ckpt_name( checkpoint_path, -1, # -1 means ignore layer_id global_expert_id, tag, mpu), map_location=torch.device('cpu'))
替换后:
expert_state_dict = torch_mindio.load(DeepSpeedEngine._get_expert_ckpt_name( checkpoint_path, -1, # -1 means ignore layer_id global_expert_id, tag, mpu))
- 在文件首行加入以下内容。
- 按“Esc”键,输入:wq!,按“Enter”保存并退出编辑。
- 执行以下命令,打开engine.py文件。
- 修改module.py文件。
- 修改state_dict_factory.py文件。
- 完成3~5修改.py文件,DeepSpeed即可使用MindIO服务。
父主题: 使用指导