197 Star 1.3K Fork 1.2K

GVPAscend/MindSpeed-LLM

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
克隆/下载
training.py 1.29 KB
一键复制 编辑 原始数据 按行查看 历史
from mindspeed.features_manager.feature import MindSpeedFeature
class TrainingDefaultFeature(MindSpeedFeature):
def __init__(self):
super(TrainingDefaultFeature, self).__init__('train-default', optimization_level=0)
def register_patches(self, patch_manager, args):
from mindspeed_llm.training import train
from mindspeed_llm.training.checkpointing import load_checkpoint_wrapper
from mindspeed_llm.legacy.data import build_pretraining_data_loader
from mindspeed_llm.tasks.posttrain.lora.utils import is_enable_qlora
if is_enable_qlora(args):
from mindspeed_llm.tasks.posttrain.lora.qlora import get_model
patch_manager.register_patch('megatron.training.training.get_model', get_model)
else:
from mindspeed_llm.training import get_model_wrapper
patch_manager.register_patch('megatron.training.training.get_model', get_model_wrapper)
patch_manager.register_patch('megatron.training.training.build_pretraining_data_loader',
build_pretraining_data_loader)
patch_manager.register_patch('megatron.training.training.train', train)
patch_manager.register_patch('megatron.training.training.load_checkpoint', load_checkpoint_wrapper)
Loading...
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
Python
1
https://gitee.com/ascend/MindSpeed-LLM.git
git@gitee.com:ascend/MindSpeed-LLM.git
ascend
MindSpeed-LLM
MindSpeed-LLM
2.1.0

搜索帮助