feat: add SwanLab as alternative logging backend#640
feat: add SwanLab as alternative logging backend#640puranikyashaswin wants to merge 6 commits intodecoderesearch:mainfrom
Conversation
- Create logging_compat.py compatibility layer for wandb/swanlab - Update config.py, llm_sae_training_runner.py, sae_trainer.py to use compat layer - Add swanlab as optional dependency in pyproject.toml - Add unit tests for logging compatibility layer Users can switch to SwanLab by setting: export SAE_LENS_LOGGING_BACKEND=swanlab SwanLab is useful for users with network restrictions who cannot access wandb.
95ddddd to
60df55f
Compare
wandb.Artifact.add_file() validates file existence. Use pytest's tmp_path fixture to create actual files during tests.
60df55f to
7a988fc
Compare
|
This is the second PR adding swanlab after #605, so I'll post the same response as to that PR here: 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. I'll aim to get something like this in the next few weeks. |
|
@chanind Thanks for the thoughtful feedback! The generic Logger interface is a much cleaner approach I can see how it would make the codebase more maintainable and give users flexibility. I'd be glad to help out with the refactoring whenever you start working on it. Happy to take on whatever part would be most useful whether that's the interface design, migrating existing WandB logic, or adding SwanLab support after. Should I close this PR for now, or leave it open as a reference? Either way, feel free to tag me when the work kicks off. |
|
let's just leave it open for now, I'll let you know once the refactor of the logging is further along! Thank you again for this! |
Description
Adds SwanLab as an optional alternative logging backend to wandb for users with network restrictions. This implements a compatibility layer that abstracts the logging backend, allowing users to switch between wandb and SwanLab via an environment variable.
Fixes #604
Type of change
Checklist:
You have tested formatting, typing and tests
make check-cito check format and linting.Performance Check.
Not applicable - this is a logging infrastructure change only, no training behavior is modified.