From 52765ac58829ceb3e740122a87cbfd8599fbd9b4 Mon Sep 17 00:00:00 2001
From: NanoCode012 <kevinvong@rocketmail.com>
Date: Thu, 8 Jun 2023 23:41:12 +0900
Subject: [PATCH] Set matmul tf32

---
 scripts/finetune.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/scripts/finetune.py b/scripts/finetune.py
index 7c4d865fa..898f88c2c 100644
--- a/scripts/finetune.py
+++ b/scripts/finetune.py
@@ -183,6 +183,9 @@ def train(
             cfg.fp16 = True
         cfg.bf16 = False
 
+    if cfg.tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+
     # load the tokenizer first
     tokenizer_config = cfg.tokenizer_config or cfg.base_model_config
     logging.info(f"loading tokenizer... {tokenizer_config}")