From 28ae947161d4670d4f865cbaad84397d47215a53 Mon Sep 17 00:00:00 2001 From: zhouwei <363232733@qq.com> Date: Mon, 6 May 2024 13:29:59 +0800 Subject: [PATCH 1/4] The training efficiency of the Ascend 910A has been significantly enhanced, leveraging the full computational power of the NPU (Neural Processing Unit) and the capabilities of torch_npu, a PyTorch library optimized for NPUs. This improvement has resulted in a remarkable tenfold increase in efficiency. --- src/train.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/train.py b/src/train.py index 6a3212cb..e2609b66 100644 --- a/src/train.py +++ b/src/train.py @@ -1,3 +1,7 @@ +import os +import torch +import torch_npu +from torch_npu.contrib import transfer_to_npu from llmtuner.train.tuner import run_exp @@ -11,4 +15,6 @@ def _mp_fn(index): if __name__ == "__main__": + use_jit_compile = os.getenv('JIT_COMPILE', 'False').lower() in ['true', '1'] + torch.npu.set_compile_mode(jit_compile=use_jit_compile) main() From 0ac6e73f9971a9310026ddc609b5266cb1639b64 Mon Sep 17 00:00:00 2001 From: hoshi-hiyouga Date: Tue, 14 May 2024 20:44:04 +0800 Subject: [PATCH 2/4] Apply suggestions from code review Co-authored-by: Huazhong Ji --- src/train.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/train.py b/src/train.py index e2609b66..098ec1b5 100644 --- a/src/train.py +++ b/src/train.py @@ -1,6 +1,4 @@ import os -import torch -import torch_npu from torch_npu.contrib import transfer_to_npu from llmtuner.train.tuner import run_exp @@ -15,6 +13,7 @@ def _mp_fn(index): if __name__ == "__main__": - use_jit_compile = os.getenv('JIT_COMPILE', 'False').lower() in ['true', '1'] - torch.npu.set_compile_mode(jit_compile=use_jit_compile) + if is_torch_npu_available(): + use_jit_compile = os.getenv('JIT_COMPILE', 'False').lower() in ['true', '1'] + torch.npu.set_compile_mode(jit_compile=use_jit_compile) main() From 9089bc70c8838cb80473e557a750855f7b7a7695 Mon Sep 17 00:00:00 2001 From: hoshi-hiyouga Date: Tue, 14 May 2024 20:44:21 +0800 Subject: [PATCH 3/4] Apply suggestions from code review Co-authored-by: Huazhong Ji --- src/train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/train.py b/src/train.py index 098ec1b5..00a7fa26 100644 --- a/src/train.py +++ b/src/train.py @@ -1,5 +1,5 @@ import os -from torch_npu.contrib import transfer_to_npu +from transformers import is_torch_npu_available from llmtuner.train.tuner import run_exp From 1c3c4989022025db756965350ae0381fc9db32e5 Mon Sep 17 00:00:00 2001 From: hoshi-hiyouga Date: Tue, 14 May 2024 20:47:52 +0800 Subject: [PATCH 4/4] Update train.py --- src/train.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/train.py b/src/train.py index 00a7fa26..4cc21194 100644 --- a/src/train.py +++ b/src/train.py @@ -1,5 +1,8 @@ import os + +import torch from transformers import is_torch_npu_available + from llmtuner.train.tuner import run_exp