adding report-to-wandb #9
Conversation
Goekdeniz-Guelmez
changed the title
|
@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.
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
WandBCallbackthere- Use
try: import wandb; except: wandb = Nonepattern instead ofif - Initialize
wandbin the callback constructor - Accept the project name, config and optional wrapped callback in the callback constructor
- Use
- Change the
--report-to-wandbto--wandbor--wandb-projectand allow setting the project or the defaultNoneto 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 |
|
Hey @awni, this should be mergable now! If you need me to fix or change something, just ping me. |
|
This looks nice! Just a couple minor comments. Could you address then we can merge? |
|
Thanks! They have been addressed! |
|
@angeloskath you have to approve this in order to land it. |
No description provided.