Helper function for reproducible behavior during distributed training. See https://pytorch.org/docs/stable/notes/randomness.html for pytorch
(seed: int, warn_only: bool = False)
| 152 | |
| 153 | |
| 154 | def enable_full_determinism(seed: int, warn_only: bool = False): |
| 155 | """ |
| 156 | Helper function for reproducible behavior during distributed training. See |
| 157 | https://pytorch.org/docs/stable/notes/randomness.html for pytorch |
| 158 | """ |
| 159 | # set seed first |
| 160 | set_seed(seed) |
| 161 | |
| 162 | if is_torch_available(): |
| 163 | # Enable PyTorch deterministic mode. This potentially requires either the environment |
| 164 | # variable 'CUDA_LAUNCH_BLOCKING' or 'CUBLAS_WORKSPACE_CONFIG' to be set, |
| 165 | # depending on the CUDA version, so we set them both here |
| 166 | os.environ["CUDA_LAUNCH_BLOCKING"] = "1" |
| 167 | os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" |
| 168 | # The environment variable required to enable deterministic mode on Ascend NPUs. |
| 169 | os.environ["ASCEND_LAUNCH_BLOCKING"] = "1" |
| 170 | os.environ["HCCL_DETERMINISTIC"] = "1" |
| 171 | |
| 172 | os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1" |
| 173 | torch.use_deterministic_algorithms(True, warn_only=warn_only) |
| 174 | |
| 175 | # Enable CUDNN deterministic mode |
| 176 | torch.backends.cudnn.deterministic = True |
| 177 | torch.backends.cudnn.benchmark = False |
| 178 | |
| 179 | |
| 180 | def set_seed(seed: int, deterministic: bool = False): |
no test coverage detected