22 lines
714 B
Bash
22 lines
714 B
Bash
# Copyright Huawei Technologies Co., Ltd. 2024. All rights reserved.
|
|
export ASCEND_RT_VISIBLE_DEVICES=2,3,4,5
|
|
export MASTER_PORT=20030
|
|
|
|
export ATB_LAYER_INTERNAL_TENSOR_REUSE=1
|
|
export INF_NAN_MODE_ENABLE=0
|
|
export ATB_OPERATION_EXECUTE_ASYNC=1
|
|
export TASK_QUEUE_ENABLE=1
|
|
export ATB_CONVERT_NCHW_TO_ND=1
|
|
export LCCL_ENABLE_FALLBACK=1
|
|
export ATB_WORKSPACE_MEM_ALLOC_GLOBAL=1
|
|
export ATB_CONTEXT_WORKSPACE_SIZE=0
|
|
export INT8_FORMAT_NZ_ENABLE=1
|
|
|
|
world_size=$(($(echo "${ASCEND_RT_VISIBLE_DEVICES}" | grep -o , | wc -l) +1))
|
|
|
|
if [ "$TP_WORLD_SIZE" == "1" ]; then
|
|
python -m examples.run_pa --model_path $1
|
|
else
|
|
torchrun --nproc_per_node $world_size --master_port $MASTER_PORT -m examples.run_pa --model_path $1
|
|
fi
|