Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / runtime / random.py: 26%

1# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15"""Trainer-side seeding and numerical-determinism controls. 

16 

17Split out of the former ``auto_models/components/utils/helper.py`` in stage 7 

18(05 §10.4); function names and signatures are unchanged. 

19""" 

20 

21import os 

22import random 

23 

24import numpy as np 

25import torch 

26from transformers import set_seed as set_seed_func 

27 

28from hyper_parallel.models.build_options import IS_CUDA_AVAILABLE, IS_NPU_AVAILABLE 

29 

30if IS_NPU_AVAILABLE: 

31 # Importing torch_npu preserves its torch.npu and HCCL registrations. 

32 import torch_npu # noqa: F401 

33 

34 

35def enable_high_precision_for_bf16() -> None: 

36 """ 

37 Set high accumulation dtype for matmul and reduction. 

38 """ 

39 if IS_CUDA_AVAILABLE: 

40 torch.backends.cuda.matmul.allow_tf32 = False 

41 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False 

42 

43 if IS_NPU_AVAILABLE: 

44 torch.npu.matmul.allow_tf32 = False 

45 torch.npu.matmul.allow_bf16_reduced_precision_reduction = False 

46 

47 

48def enable_full_determinism(seed: int) -> None: 

49 """ 

50 Helper function for reproducibility in distributed training. 

51 See https://pytorch.org/docs/stable/notes/randomness.html for details. 

52 """ 

53 

54 os.environ["PYTHONHASHSEED"] = str(seed) 

55 os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" 

56 os.environ["NCCL_DETERMINISTIC"] = "1" 

57 os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1" 

58 if IS_NPU_AVAILABLE: 

59 # The environment variable required to enable deterministic mode on Ascend NPUs. 

60 os.environ["NCCL_DETERMINISTIC"] = "true" 

61 os.environ["CLOSE_MATMUL_K_SHIFT"] = "1" 

62 

63 random.seed(seed) 

64 np.random.seed(seed) 

65 torch.manual_seed(seed) 

66 torch.cuda.manual_seed(seed) 

67 torch.cuda.manual_seed_all(seed) 

68 torch.use_deterministic_algorithms(True, warn_only=True) 

69 # Enable CUDNN deterministic mode 

70 torch.backends.cudnn.deterministic = True 

71 torch.backends.cudnn.benchmark = False 

72 torch.backends.cudnn.enabled = False 

73 

74 if IS_NPU_AVAILABLE: 

75 torch.npu.manual_seed(seed) 

76 torch.npu.manual_seed_all(seed) 

77 

78 

79def set_seed(seed: int | None, full_determinism: bool = False) -> None: 

80 """ 

81 Sets a manual seed on all devices. 

82 """ 

83 if seed is None: 

84 return 

85 if full_determinism: 

86 enable_full_determinism(seed) 

87 else: 

88 set_seed_func(seed) 

89 

90 

91__all__ = [ 

92 "enable_full_determinism", 

93 "enable_high_precision_for_bf16", 

94 "set_seed", 

95]