Skip to content

adding report-to-wandb - #9

Merged
awni merged 25 commits into
ml-explore:mainfrom
Goekdeniz-Guelmez:adding-report-to-wandb
May 10, 2025
Merged

adding report-to-wandb#9
awni merged 25 commits into
ml-explore:mainfrom
Goekdeniz-Guelmez:adding-report-to-wandb

Conversation

@Goekdeniz-Guelmez

Copy link
Copy Markdown
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
Copy Markdown
Contributor Author

@awni should be mergable now.

@ivanfioravanti

Copy link
Copy Markdown
Contributor

It would be great having this integrated OOB 🙏

@Goekdeniz-Guelmez

Copy link
Copy Markdown
Contributor Author

@awni should be mergeable now!

@angeloskath angeloskath left a comment

Copy link
Copy Markdown
Member

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

Goekdeniz-Guelmez commented Apr 5, 2025

Copy link
Copy Markdown
Contributor Author

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
Copy Markdown
Contributor Author

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

@Goekdeniz-Guelmez

Copy link
Copy Markdown
Contributor Author

Hey @awni, this should be mergable now! If you need me to fix or change something, just ping me.

Comment thread mlx_lm/LORA.md Outdated
Comment thread mlx_lm/examples/lora_config.yaml Outdated
Comment thread mlx_lm/examples/lora_config.yaml
@awni

awni commented May 9, 2025

Copy link
Copy Markdown
Member

This looks nice! Just a couple minor comments. Could you address then we can merge?

@Goekdeniz-Guelmez

Copy link
Copy Markdown
Contributor Author

Thanks! They have been addressed!

@awni awni left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Awesome!

@awni

awni commented May 10, 2025

Copy link
Copy Markdown
Member

@angeloskath you have to approve this in order to land it.

@angeloskath angeloskath left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍

@awni
awni merged commit 4b48477 into ml-explore:main May 10, 2025
@Goekdeniz-Guelmez
Goekdeniz-Guelmez deleted the adding-report-to-wandb branch May 10, 2025 09:53
nitky pushed a commit to nitky/mlx-lm that referenced this pull request May 16, 2026
Round 2 fixes from cross-model adversarial review (B/C strata):

- B1: drop private workspace path from module docstring
- B2: invert network-test polarity (MLX_LM_RUN_NETWORK_TESTS=1 opt-in,
  default-skip in CI); use public huggingface_hub.snapshot_download
  instead of private _download; rename test to describe behavior
- I1: remove unused List import
- I2: rename quant_predicate's _module to _ (matches sibling style)
- I3: pass position_ids[-1] as mx.array offset to RoPE (no .item() sync;
  MTP hot path was paying a GPU->host roundtrip per forward)
- I4: build MaskedEmbedder fill on-device (no .item() sync inside the
  per-draft-step path)
- ml-explore#6: rename masked_embedding.token_ordering -> _token_ordering so the
  int32 gather buffer is excluded from Module.parameters(); install via
  Model.sanitize side-channel so load_weights (which rejects underscored
  keys) doesn't trip; eliminates the fragile post-tree_map int32 restore
  in the synthetic Tier 1 test
- ml-explore#9: make_cache() now raises NotImplementedError instead of returning []
  (zip-based cache iteration would silently no-op the previous return)

77 model tests pass (10 gemma4 tests + 1 opt-in real-weight test pass).
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.

4 participants