feat(trainer): support param groups

- Updated the `get_parameters` method in `AbstractSparseAutoEncoder` to separate parameters into groups for "others" and "jumprelu" with respective learning rates.
- Modified the `_initialize_optimizer` method in `Trainer` to apply different learning rates based on parameter group names and log detailed parameter information.
- Added a new configuration field `jumprelu_lr_factor` in `TrainerConfig` to control the learning rate for JumpReLU parameters.

Author: dest1n1s

src/lm_saes/abstract_sae.py

@@ -409,7 +409,11 @@ def init_encoder_with_decoder_transpose(self, factor: float = 1.0):
 
     def get_parameters(self) -> list[dict[str, Any]]:
         """Get the parameters of the model for optimization."""
-        return [{"params": self.parameters()}]
+        jumprelu_params = (
+            list(self.activation_function.parameters()) if isinstance(self.activation_function, JumpReLU) else []
+        )
+        other_params = [p for p in self.parameters() if not any(p is param for param in jumprelu_params)]
+        return [{"params": other_params, "name": "others"}, {"params": jumprelu_params, "name": "jumprelu"}]
 
     def load_full_state_dict(self, state_dict: dict[str, torch.Tensor], device_mesh: DeviceMesh | None = None) -> None:
         # Extract and set dataset_average_activation_norm if present

src/lm_saes/config.py

@@ -201,6 +201,7 @@ class TrainerConfig(BaseConfig):
     lr_end_ratio: float = 1 / 32
     lr_warm_up_steps: int | float = 5000
     lr_cool_down_steps: int | float = 0.2
+    jumprelu_lr_factor: float = 1.0
     clip_grad_norm: float = 0.0
     feature_sampling_window: int = 1000
     total_training_tokens: int = 300_000_000

src/lm_saes/sae.py

@@ -424,6 +424,3 @@ def prepare_label(self, batch: dict[str, torch.Tensor], **kwargs) -> torch.Tenso
         else:
             label = batch[self.cfg.hook_point_out]
         return label
-
-    def get_parameters(self) -> list[dict[str, Any]]:
-        return [{"params": self.parameters()}]

src/lm_saes/trainer.py

@@ -1,6 +1,6 @@
 import math
 import os
-from typing import Callable, Iterable
+from typing import Any, Callable, Iterable
 
 import torch
 import torch.optim.lr_scheduler as lr_scheduler
@@ -72,7 +72,31 @@ def calculate_warmup_steps(warmup_steps: float | int) -> int:
     @timer.time("initialize_optimizer")
     def _initialize_optimizer(self, sae: AbstractSparseAutoEncoder):
         assert isinstance(self.cfg.lr, float)
-        optimizer = Adam(sae.get_parameters(), lr=self.cfg.lr, betas=self.cfg.betas)
+
+        def _apply_lr(parameters: dict[str, Any]):
+            assert isinstance(self.cfg.lr, float)
+            if parameters["name"] == "jumprelu":
+                return {**parameters, "lr": self.cfg.jumprelu_lr_factor * self.cfg.lr}
+            return parameters
+
+        params = [_apply_lr(parameters) for parameters in sae.get_parameters()]
+
+        def _format_parameters(parameters: dict[str, Any]) -> str:
+            param_info = f"{parameters['name']}:"
+            for i, param in enumerate(parameters["params"]):
+                param_info += f"\n    [{i}] shape={list(param.shape)}, dtype={param.dtype}"
+                if param.requires_grad:
+                    param_info += ", trainable"
+                else:
+                    param_info += ", frozen"
+            if "lr" in parameters:
+                param_info += f"\n    lr={parameters['lr']}"
+            return param_info
+
+        param_str = "\n".join([_format_parameters(p) for p in params])
+        logger.info(f"\nParameter Groups: \n{param_str}\n")
+
+        optimizer = Adam(params, lr=self.cfg.lr, betas=self.cfg.betas)
         scheduler = get_scheduler(
             scheduler_name=self.cfg.lr_scheduler_name,
             optimizer=optimizer,