feat: Upgrade Weights & Biases callback

[parambharat]

What does this PR do?

• Continuation of old PR feat: Upgrade Weights & Biases callback #29125

This PR adds a few new functionalities to the Weights & Biases Callback

• Logs Peft and Lora Config to wandb if present

• Adds model parameter counts to wandb config and artifact metadata

• Adds on_predict methods to log prediction metrics

• Prints the model architecture to a file alongside the wandb artifact

• Logs initial and final models to the wandb artifact to full reproducibility

• Adds steps and epoch aliases to checkpoint artifacts

• Here's a link to the what the logged artifacts look like

• Here's a run overview page with added config and metadata for the run with peft configs logged

Before submitting

• Did you read the contributor guideline,
  Pull Request section?

Who can review?

• @muellerzr
• @pacman100
• @amyeroberts

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[parambharat]

@amyeroberts: I think I was finally able to get all ci checks to pass. The issues were related to type annotations and lazy imports. I have fixed these. Can you please take a look ?

[amyeroberts]

Thanks for all of the work adding this!

[morganmcg1]

Hey @amyeroberts , do you know if this PR is going into the next release?

[amyeroberts]

@morganmcg1 Yes, all commits merged into main will be part of the next minor release - v4.41. If there are patch releases before this e.g. 4.40.2, then it won't be part of that release.

If you want to use this immediately, you can install from source, getting all the most recent commits from main: pip install git+https://github.com/huggingface/transformers

[jonghwanhyeon]

This PR causes an error when a custom model with no save_pretrained() method uses Trainer class.

[amyeroberts]

@jonghwanhyeon Could you open an issue, detailing the errors and linking to this PR for reference? This way we'll be able to properly track if/when it's resolved

[mgerstgrasser]

@amyeroberts @parambharat It seems that this always logs the initial model checkpoint to wandb, with no way of turning it off, and uploads the entire model when doing a full fine-tune. This is not intended, is it? See #30896 and #30897

[tomaarsen]

@ArthurZucker I appreciate you reverting this; my W&B storage was immediately full and now I can't access my W&B logs anymore without being immediately redirected to the billing page :/
Freeing up space doesn't immediately give me access again either, it's super frustrating

• Tom Aarasen