Skip to content

Add wandb logging callback for Trainer#173

Open
zywilliamli wants to merge 4 commits intomagic-wmdp-attributionfrom
add-wandb-log-fn
Open

Add wandb logging callback for Trainer#173
zywilliamli wants to merge 4 commits intomagic-wmdp-attributionfrom
add-wandb-log-fn

Conversation

@zywilliamli
Copy link
Collaborator

Summary

  • Adds bergson/utils/logging.py with a wandb_log_fn() factory that creates a log_fn callback compatible with Trainer.train()
  • Lazily imports wandb and initializes a run if one isn't already active
  • Used by scripts/reconstruct_loss.py for optional loss logging during checkpoint reconstruction

Test plan

  • Verify from bergson.utils.logging import wandb_log_fn imports cleanly
  • Test with a short training run using wandb_log_fn("test-project") as the log_fn

🤖 Generated with Claude Code

@zywilliamli zywilliamli requested a review from luciaquirke March 3, 2026 23:05
- Add bergson/utils/logging.py with wandb_log_fn() factory
- Wire into magic_wmdp.py with --wandb_project flag (default: "magic-wmdp")
- Logs per-step training loss on rank 0, set --wandb_project "" to disable

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
pre-commit-ci bot and others added 2 commits March 3, 2026 23:17
Stores per-step loss on the Trainer instance during step() and calls
the optional log_fn(step, loss) callback after each training step.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@luciaquirke
Copy link
Collaborator

@claude could you please fix the pre-commit issues

@claude
Copy link

claude bot commented Mar 6, 2026

Claude encountered an error —— View job


Fixing pre-commit issues

  • Gather context and run pre-commit
  • Fix pre-commit issues
  • Commit and push changes

@claude
Copy link

claude bot commented Mar 6, 2026

Claude encountered an error —— View job


Fixing pre-commit issues

Note: I'm unable to increase my own max turns — that's a system-level configuration set when invoking Claude Code. You can find more information in the FAQ.

  • Gather context and run pre-commit
  • Fix pre-commit issues
  • Commit and push changes

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@luciaquirke
Copy link
Collaborator

you probably want to merge this into the main magic branch

@luciaquirke
Copy link
Collaborator

@zywilliamli this can merge into main now

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants