Skip to content

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

Open
wants to merge 16 commits into
base: main
Choose a base branch
from

Conversation

Goekdeniz-Guelmez
Copy link
Contributor

No description provided.

@Goekdeniz-Guelmez Goekdeniz-Guelmez changed the title update lora_config.yaml + LORA.md + lora.py adding report-to-wandb Mar 14, 2025
@Goekdeniz-Guelmez
Copy link
Contributor Author

@awni should be mergable now.

@ivanfioravanti
Copy link
Contributor

It would be great having this integrated OOB 🙏

@Goekdeniz-Guelmez
Copy link
Contributor Author

@awni should be mergeable now!

Copy link
Member

@angeloskath angeloskath left a 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 of if
    • Initialize wandb in the callback constructor
    • Accept the project name, config and optional wrapped callback in the callback constructor
  • Change the --report-to-wandb to --wandb or --wandb-project and allow setting the project or the default None 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.
@Goekdeniz-Guelmez
Copy link
Contributor Author

Goekdeniz-Guelmez commented Apr 5, 2025

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 --adapter-path path.

@Goekdeniz-Guelmez
Copy link
Contributor Author

You can now call it like --wandb my-project-run

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.

3 participants