-
Notifications
You must be signed in to change notification settings - Fork 60
adding report-to-wandb #9
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?
adding report-to-wandb #9
Conversation
@awni should be mergable now. |
It would be great having this integrated OOB 🙏 |
@awni should be mergeable now! |
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.
Hey @Goekdeniz-Guelmez thanks for adding that!
I think it would be better to move this out of lora.py
. I would do the following changes
- Add a file
mlx_lm/tuner/callbacks.py
- Add
WandBCallback
there- Use
try: import wandb; except: wandb = None
pattern instead ofif
- Initialize
wandb
in the callback constructor - Accept the project name, config and optional wrapped callback in the callback constructor
- Use
- Change the
--report-to-wandb
to--wandb
or--wandb-project
and allow setting the project or the defaultNone
to disable the callback
- Updated WandB reporting mechanism to use a project name argument instead of a boolean flag. - Removed the old TrainingCallback class definition from trainer.py and imported it from callbacks. - Adjusted argument parsing to accommodate the new WandB configuration.
- Added log_dir parameter to WandBCallback constructor for specifying the logging directory. - Updated lora.py to pass adapter_path as log_dir when initializing WandBCallback.
Hey @angeloskath your suggestions are really helpful! I did what you said and it works as expected. Also the run config and data files are save in the |
You can now call it like |
No description provided.