diff --git a/mindspeed_llm/mindspore/mindspore_adaptor.py b/mindspeed_llm/mindspore/mindspore_adaptor.py index cd3992a2b5dc57d05bf72382a0ae204708ce49bc..1b4196b15d20e733e0005871e6fd53bc855335fb 100644 --- a/mindspeed_llm/mindspore/mindspore_adaptor.py +++ b/mindspeed_llm/mindspore/mindspore_adaptor.py @@ -277,6 +277,9 @@ class MindSporeAdaptation(MegatronAdaptationABC): MindSporeAdaptation.register('mindspeed.core.pipeline_parallel.fb_overlap.modules.token_dispatcher.overlap_stream', overlap_stream) + from mindspeed_llm.core.pipeline_parallel.dualpipe.adaptor import dualpipe_forward_step_wrapper + MindSporeAdaptation.register('mindspeed.mindspore.core.pipeline_parallel.dualpipev.dualpipev_schedules.forward_step_with_model_graph', + dualpipe_forward_step_wrapper) if args.swap_optimizer: