diff --git a/src/transformers/trainer_utils.py b/src/transformers/trainer_utils.py index d17113091777be..a52560a77cf339 100644 --- a/src/transformers/trainer_utils.py +++ b/src/transformers/trainer_utils.py @@ -71,6 +71,9 @@ def enable_full_determinism(seed: int, warn_only: bool = False): # depending on the CUDA version, so we set them both here os.environ["CUDA_LAUNCH_BLOCKING"] = "1" os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" + # The environment variable required to enable deterministic mode on Ascend NPUs. + os.environ["ASCEND_LAUNCH_BLOCKING"] = 1 + os.environ["HCCL_DETERMINISTIC"] = 1 torch.use_deterministic_algorithms(True, warn_only=warn_only) # Enable CUDNN deterministic mode