add ascend npu accelerator support (#24879)

* Add Ascend NPU accelerator support

* fix style warining
This commit is contained in:
statelesshz
2023-07-18 20:20:32 +08:00
committed by GitHub
parent f14c7f999d
commit 9c875839c0
7 changed files with 68 additions and 1 deletions

View File

@@ -21,6 +21,7 @@ from transformers.testing_utils import (
get_torch_dist_unique_port,
require_torch_multi_gpu,
require_torch_neuroncore,
require_torch_npu,
)
from transformers.training_args import ParallelMode
from transformers.utils import logging
@@ -77,6 +78,20 @@ class TestTrainerDistributedNeuronCore(TestCasePlus):
# successful return here == success - any errors would have caused an error in the sub-call
class TestTrainerDistributedNPU(TestCasePlus):
@require_torch_npu
def test_trainer(self):
distributed_args = f"""--nproc_per_node=2
--master_port={get_torch_dist_unique_port()}
{self.test_file_dir}/test_trainer_distributed.py
""".split()
output_dir = self.get_auto_remove_tmp_dir()
args = f"--output_dir {output_dir}".split()
cmd = ["torchrun"] + distributed_args + args
execute_subprocess_async(cmd, env=self.get_env())
# successful return here == success - any errors would have caused an error in the sub-call
class TestTrainerDistributed(TestCasePlus):
@require_torch_multi_gpu
def test_trainer(self):