Skip to content

Commit

Permalink
Merge pull request #32 from OpenMOSS/AprilTricks
Browse files Browse the repository at this point in the history
April tricks
  • Loading branch information
StarConnor authored Jul 17, 2024
2 parents 4a14ae0 + 652bc19 commit 900b98c
Showing 1 changed file with 11 additions and 2 deletions.
13 changes: 11 additions & 2 deletions src/lm_saes/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -185,7 +185,7 @@ class SAEConfig(BaseModelConfig):
init_encoder_with_decoder_transpose: bool = True

l1_coefficient: float = 0.00008
l1_coefficient_warmup_steps: int = 0
l1_coefficient_warmup_steps: int | float = 0.1
lp: int = 1

use_ghost_grads: bool = False
Expand Down Expand Up @@ -286,7 +286,7 @@ class LanguageModelSAETrainingConfig(LanguageModelSAERunnerConfig):
"constantwithwarmup" # constant, constantwithwarmup, linearwarmupdecay, cosineannealing, cosineannealingwarmup, exponentialwarmup
)
lr_end: Optional[float] = 1 / 32
lr_warm_up_steps: int = 5000
lr_warm_up_steps: int | float = 0.1
lr_cool_down_steps: int = 10000
train_batch_size: int = 4096
clip_grad_norm: float = 0.0
Expand Down Expand Up @@ -331,6 +331,15 @@ def __post_init__(self):
total_training_steps = self.total_training_tokens // self.effective_batch_size
print_once(f"Total training steps: {total_training_steps}")

if self.lr_scheduler_name == "constantwithwarmup" and isinstance(self.lr_warm_up_steps, float):
assert 0 <= self.lr_warm_up_steps <= 1.0
self.lr_warm_up_steps = int(self.lr_warm_up_steps * total_training_steps)
print_once(f"Learning rate warm up steps: {self.lr_warm_up_steps}")
if isinstance(self.sae.l1_coefficient_warmup_steps, float):
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}")

@dataclass(kw_only=True)
class LanguageModelSAEPruningConfig(LanguageModelSAERunnerConfig):
"""
Expand Down

0 comments on commit 900b98c

Please sign in to comment.