-
Notifications
You must be signed in to change notification settings - Fork 217
feat: Add Swanlab logging backend #605
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
feat: Add Swanlab logging backend #605
Conversation
Signed-off-by: akai <2780896955@qq.com>
|
chanind
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.