feat(config): add decay ratio

OpenMOSS · Jul 18, 2024 · 7e8fc63 · 7e8fc63
1 parent bb86254
commit 7e8fc63
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/lm_saes/config.py b/src/lm_saes/config.py
@@ -286,7 +286,7 @@ class LanguageModelSAETrainingConfig(LanguageModelSAERunnerConfig):
     )
     lr_end: Optional[float] = 1 / 32
     lr_warm_up_steps: int | float = 0.1
-    lr_cool_down_steps: int = 10000
+    lr_cool_down_steps: int | float = 0.1
     train_batch_size: int = 4096
     clip_grad_norm: float = 0.0
     remove_gradient_parallel_to_decoder_directions: bool = False
@@ -338,6 +338,10 @@ def __post_init__(self):
             assert 0 <= self.sae.l1_coefficient_warmup_steps <= 1.0
             self.sae.l1_coefficient_warmup_steps = int(self.sae.l1_coefficient_warmup_steps * total_training_steps)
             print_once(f"L1 coefficient warm up steps: {self.sae.l1_coefficient_warmup_steps}")
+        if isinstance(self.lr_cool_down_steps, float):
+            assert 0 <= self.lr_cool_down_steps <= 1.0
+            self.lr_cool_down_steps = int(self.lr_cool_down_steps * total_training_steps)
+            print_once(f"Learning rate cool down steps: {self.lr_cool_down_steps}")
 
 @dataclass(kw_only=True)
 class LanguageModelSAEPruningConfig(LanguageModelSAERunnerConfig):