Torch对接X1框架

cd X1安装目录/Megatron-LM/megatron

1. 打开checkpointing.py文件。

   vim checkpointing.py

2. 按“i”进入编辑模式，修改如下内容。
   • 在文件首行加入以下内容。

     import torch_mindio

   • 将torch.load函数替换为torch_mindio.load函数。
     替换前：

     optim_checkpoint = torch.load(optim_load_path,
                                   map_location=torch.device('cpu'))

     替换后：

     optim_checkpoint = torch_mindio.load(optim_load_path, map_location='cpu')

   • 将普通的torch.save函数替换为torch_mindio.save函数。
     替换前：

     torch.save(state, save_path)

     替换后：

     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))

3. 按“Esc”键，输入:wq!，按“Enter”保存并退出编辑。