feat: Add Swanlab logging backend#605
feat: Add Swanlab logging backend#605wantbook-book wants to merge 5 commits intodecoderesearch:mainfrom
Conversation
Signed-off-by: akai <2780896955@qq.com>
|
chanind
left a comment
There was a problem hiding this comment.
Thank you for contributing this! However, I'm inclined to hold off on merging this until the next major release of SAELens. Ideally we should have a generic "Logger" interface, and then the user should specify the loggers they want in the config. That should make it possible for users to register their own loggers if they want as well. I would envision something like the following:
cfg = LanguageModelSAERunnerConfig(
...
logger=[
WandBLoggingConfig(...),
SwanLabLoggingConfig(...),
],
)This would be a breaking change though so we'd need to de a major release, and requires a bit more refactoring than is ideal in just this PR.
| if backend not in {"auto", "wandb", "swanlab"}: | ||
| backend = "auto" | ||
|
|
||
| if backend in {"auto", "swanlab"}: |
There was a problem hiding this comment.
I'd rather the default be wandb, not swanlab. If a user just happens to have SAELens and swanlab deps in the same project their logging will seem to break and this will be very hard to debug. IMO if users want to use swanlab they should need to opt-in specifically with a ENV var.
| ) | ||
| for key, value in self.sae.log_histograms().items(): | ||
| eval_metrics[key] = wandb.Histogram(value) # type: ignore | ||
| if BACKEND != "swanlab": |
There was a problem hiding this comment.
I'd rather not need to have if-statements like this in code. Can you provide an interface so swanlab works identical to wandb, and we don't need any of these checks? Or maybe make a generic Logger interface and have both Swanlab and WandB versions of that same interface? It shouldn't be the responsibility of the SAETrainer to keep track of what logger is being used.
| run_minimal_standard_sae_training(output_dir) | ||
|
|
||
|
|
||
| if __name__ == "__main__": |
There was a problem hiding this comment.
nit: don't use main functions in tests. Just run the test using pytest.
Description
Summary:
Motivation:
Dependencies:
swanlab(whenSAE_LENS_LOGGING_BACKEND=swanlab)Fixes #604
Type of change
Please delete options that are not relevant.
Checklist:
You have tested formatting, typing and tests
make check-cito check format and linting. (you can runmake formatto format code if needed.)Performance Check.
If you have implemented a training change, please indicate precisely how performance changes with respect to the following metrics:
Please links to wandb dashboards with a control and test group.