use fused Adam on CUDA

LoicGrobol · Nov 17, 2024 · faad953 · faad953
1 parent 7596537
commit faad953
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 20 deletions.
diff --git a/.markdownlint → .markdownlint.yaml b/.markdownlint → .markdownlint.yaml
diff --git a/tox.ini b/tox.ini
diff --git a/tox.toml b/tox.toml
@@ -0,0 +1,26 @@
+minversion = "4"
+env_list = ["py312", "py311", "py310", "py39"]
+isolated_build = true
+skip_missing_interpreters = true
+
+[env_run_base]
+extras = [tests]
+commands = [
+    [
+        'pytest',
+        "tests",
+        "--basetemp",
+        "{envtmpdir}",
+        { replace = "posargs", extend = true },
+    ],
+]
+
+[pytest]
+script_launch_mode = "subprocess"
+
+[gh-actions]
+python = """
+3.9: py39
+3.10: py310
+3.11: py311
+3.12: py312"""
diff --git a/zeldarose/common.py b/zeldarose/common.py
@@ -87,6 +87,7 @@ def configure_optimizers(self):
         optimizer = torch.optim.AdamW(
             optimizer_grouped_parameters,
             betas=self.training_config.betas,
+            fused=True,
             lr=self.training_config.learning_rate,
             eps=self.training_config.epsilon,
             weight_decay=decay_rate,

diff --git a/zeldarose/tasks/rtd.py b/zeldarose/tasks/rtd.py
@@ -347,6 +347,7 @@ def configure_optimizers(self):
         optimizer = torch.optim.AdamW(
             optimizer_grouped_parameters,
             betas=self.training_config.betas,
+            fused=True,
             lr=self.training_config.learning_rate,
             eps=self.training_config.epsilon,
             weight_decay=decay_rate,