From eb33b8bf1c979c4b414a3f126db8dc974d8f456c Mon Sep 17 00:00:00 2001 From: Alexander Pivovarov Date: Wed, 17 May 2023 03:26:48 +0000 Subject: [PATCH] Use bf16 only if supported --- bench.py | 2 +- sample.py | 2 +- train.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/bench.py b/bench.py index 25da9f2..33dd421 100644 --- a/bench.py +++ b/bench.py @@ -15,7 +15,7 @@ bias = False real_data = True seed = 1337 device = 'cuda' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1', etc. -dtype = 'bfloat16' # 'float32' or 'bfloat16' or 'float16' +dtype = 'bfloat16' if torch.cuda.is_bf16_supported() else 'float16' # 'float32' or 'bfloat16' or 'float16' compile = True # use PyTorch 2.0 to compile the model to be faster profile = False # use pytorch profiler, or just simple benchmarking? exec(open('configurator.py').read()) # overrides from command line or config file diff --git a/sample.py b/sample.py index 670759b..ecef5f9 100644 --- a/sample.py +++ b/sample.py @@ -18,7 +18,7 @@ temperature = 0.8 # 1.0 = no change, < 1.0 = less random, > 1.0 = more random, i top_k = 200 # retain only the top_k most likely tokens, clamp others to have 0 probability seed = 1337 device = 'cuda' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1', etc. -dtype = 'bfloat16' # 'float32' or 'bfloat16' or 'float16' +dtype = 'bfloat16' if torch.cuda.is_bf16_supported() else 'float16' # 'float32' or 'bfloat16' or 'float16' compile = False # use PyTorch 2.0 to compile the model to be faster exec(open('configurator.py').read()) # overrides from command line or config file # ----------------------------------------------------------------------------- diff --git a/train.py b/train.py index a0cd539..ca0912c 100644 --- a/train.py +++ b/train.py @@ -70,7 +70,7 @@ min_lr = 6e-5 # minimum learning rate, should be ~= learning_rate/10 per Chinchi backend = 'nccl' # 'nccl', 'gloo', etc. # system device = 'cuda' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1' etc., or try 'mps' on macbooks -dtype = 'bfloat16' # 'float32', 'bfloat16', or 'float16', the latter will auto implement a GradScaler +dtype = 'bfloat16' if torch.cuda.is_bf16_supported() else 'float16' # 'float32', 'bfloat16', or 'float16', the latter will auto implement a GradScaler compile = True # use PyTorch 2.0 to compile the model to be faster # ----------------------------------------------------------------------------- config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))]