From 3ba93629c3bb434989b79191a69543435dd82009 Mon Sep 17 00:00:00 2001 From: Jett Date: Wed, 15 May 2024 13:46:09 +0200 Subject: [PATCH 1/2] handle CUBLAS_WORKSPACE_CONFIG env var --- src/delphi/train/training.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/src/delphi/train/training.py b/src/delphi/train/training.py index 85fda188..3d17a9fb 100644 --- a/src/delphi/train/training.py +++ b/src/delphi/train/training.py @@ -36,7 +36,25 @@ def setup_training(config: TrainingConfig): init_wandb(config=config) +def check_set_env_cublas_workspace_config(): + expected_val = ":4096:8" + actual_val = os.getenv("CUBLAS_WORKSPACE_CONFIG") + if actual_val is None: + # https://docs.nvidia.com/cuda/archive/12.4.0/cublas/index.html#results-reproducibility + logging.info( + f"Environment variable CUBLAS_WORKSPACE_CONFIG not set. Setting to '{expected_val}' to ensure reproducibility." + ) + os.environ["CUBLAS_WORKSPACE_CONFIG"] = expected_val + elif actual_val != expected_val: + logging.warning( + f"Environment variable CUBLAS_WORKSPACE_CONFIG is set to {actual_val}, which may affect reproducibility. " + f"We recommend setting it to '{expected_val}' to ensure reproducibility." + ) + + def run_training(config: TrainingConfig) -> tuple[ModelTrainingState, RunContext]: + if torch.cuda.is_available(): + check_set_env_cublas_workspace_config() setup_training(config) logging.info("Starting training...") logging.info("Config:") From 0cd5e29f7aa934aec7350183b53d2f3522eb71e0 Mon Sep 17 00:00:00 2001 From: Jett Date: Mon, 20 May 2024 12:36:23 +0200 Subject: [PATCH 2/2] warning -> assert --- src/delphi/train/training.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/delphi/train/training.py b/src/delphi/train/training.py index 3d17a9fb..2b29f77d 100644 --- a/src/delphi/train/training.py +++ b/src/delphi/train/training.py @@ -40,16 +40,16 @@ def check_set_env_cublas_workspace_config(): expected_val = ":4096:8" actual_val = os.getenv("CUBLAS_WORKSPACE_CONFIG") if actual_val is None: - # https://docs.nvidia.com/cuda/archive/12.4.0/cublas/index.html#results-reproducibility logging.info( f"Environment variable CUBLAS_WORKSPACE_CONFIG not set. Setting to '{expected_val}' to ensure reproducibility." ) os.environ["CUBLAS_WORKSPACE_CONFIG"] = expected_val - elif actual_val != expected_val: - logging.warning( - f"Environment variable CUBLAS_WORKSPACE_CONFIG is set to {actual_val}, which may affect reproducibility. " - f"We recommend setting it to '{expected_val}' to ensure reproducibility." - ) + correct_values = [expected_val, ":16:8"] + assert actual_val in correct_values, ( + f"Environment variable CUBLAS_WORKSPACE_CONFIG is set to {actual_val}, which is incompatibe with reproducible training. " + f"Please set it to one of the following values: {correct_values}. " + f"See https://docs.nvidia.com/cuda/archive/12.4.0/cublas/index.html#results-reproducibility for more information." + ) def run_training(config: TrainingConfig) -> tuple[ModelTrainingState, RunContext]: