迁移修改

导入torch_npu框架

导入torch_npu框架,涉及已适配修改的训练脚本位置如下所示:

源码包根目录下的“train.py”“evaluate.py”脚本。

新增如下内容:

import torch_npu

导入自动迁移脚本

导入自动迁移脚本,涉及已适配修改的训练脚本位置如下所示:

源码包根目录下的“train.py”“evaluate.py”脚本。

新增如下内容:

from torch_npu.contrib import transfer_to_npu

适配Megatron优化器

使用Megatron-LM中的AdamW优化器替换torch.optim.AdamW优化器。

修改“lavis/runners/runner_base.py”脚本,改动如下:

...
from megatron_npu.adaptor_optimizer_optimizer import AdamW   # 新增导入

@registry.register_runner("runner_base")
class RunnerBase:
...
    @property
    def optimizer(self):
        # TODO make optimizer class and configurations
        if self._optimizer is None:
            lr_scale = self.config.run_cfg.get("lr_layer_decay", 1)
            weight_decay = self.config.run_cfg.get("weight_decay", 0.05)
            optim_params = self._model.get_optimizer_params(weight_decay,lr_scale)

            num_parameters = 0
            for p_group in optim_params:
                for p in p_group["params"]:
                    num_parameters += p.data.nelement()    
            logging.info("number of trainable parameters: {}".format(num_parameters))      

            beta2 = self.config.run_cfg.get("beta2", 0.999)
            # 使用Megatron-LM中的AdamW优化器替换torch.optim.AdamW优化器
            self._optimizer = AdamW(
                optim_params,
                lr=float(self.config.run_cfg.init_lr),
                betas=(0.9, beta2),
            )    
        return self._optimizer
...