adding report-to-wandb - #9
Conversation
|
@awni should be mergable now. |
|
It would be great having this integrated OOB 🙏 |
|
@awni should be mergeable now! |
angeloskath
left a comment
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. |
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).
No description provided.