文档首页/ AI开发平台ModelArts/ 常见问题/ Standard Notebook/ 更多功能咨询/ 在Notebook中,如何使用昇腾多卡进行调试?
更新时间:2024-10-28 GMT+08:00

在Notebook中,如何使用昇腾多卡进行调试?

昇腾多卡训练任务是多进程多卡模式,跑几卡需要起几个python进程。昇腾底层会读取环境变量:RANK_TABLE_FILE,开发环境已经设置,用户无需关注。比如跑八卡,可以如下片段代码:

 export RANK_SIZE=8
 current_exec_path=$(pwd)
 echo 'start training'
 for((i=0;i<=$RANK_SIZE-1;i++));
 do
 echo 'start rank '$i
 mkdir ${current_exec_path}/device$i
 cd ${current_exec_path}/device$i
 echo $i
 export RANK_ID=$i
 dev=`expr $i + 0`
 echo $dev
 export DEVICE_ID=$dev
 python train.py > train.log 2>&1 &
 done

其中,train.py中设置环境变量DEVICE_ID:

devid = int(os.getenv('DEVICE_ID'))
 context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=devid)