Bugfix: WandbCallback uploads initial model checkpoint
Author: mgerstgrasser
Pull request: yes

# What does this PR do?
It seems that #30135 inadvertently makes WandbCallback _always_ upload the initial model checkpoint to wandb, without any way to disabling this, and irrespective of model size. I think this is obviously not intended behaviour.

This also breaks training with FSDP, as apparently calling `model.save_pretrained()` messes things up for later `model.forward()` calls.

<!-- Remove if not applicable -->

Fixes #30896 #30895


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@pacman100 @muellerzr @amyeroberts

Comments:
@pacman100 @muellerzr @amyeroberts
I think this might be high priority, actually. It seems that right now, just starting a training script with wandb logging will *immediately* upload a full model checkpoint to wandb. After just a few quick tests I'm at over 100GB in my wandb account, and getting a warning message.


@amyeroberts As per discussion in the linked issue, I've added a check to keep this but make it optional. I still am unsure what the intended (or best) behavior here is. For now I've made it so that the initial model is logged if `WANDB_LOG_MODEL` is set to either `checkpoint` or `end`. I'm unsure about the `end` case though - uploading the initial model will double the bandwidth and storage impact compared to not doing so in this case, and I can see how that may be undesirable. That said, enabling the upload in the `end` case is closer to the original PR, so that seems like the "safer" option for now. (In the `checkpoint` case, uploading seems pretty safe - you'd already be doing so many times, so one more is a small impact, relatively speaking.)

> @amyeroberts As per discussion in the linked issue, I've added a check to keep this but make it optional. I still am unsure what the intended (or best) behavior here is. For now I've made it so that the initial model is logged if `WANDB_LOG_MODEL` is set to either `checkpoint` or `end`. I'm unsure about the `end` case though - uploading the initial model will double the bandwidth and storage impact compared to not doing so in this case, and I can see how that may be undesirable. That said, enabling the upload in the `end` case is closer to the original PR, so that seems like the "safer" option for now. (In the `checkpoint` case, uploading seems pretty safe - you'd already be doing so many times, so one more is a small impact, relatively speaking.)

Caling in @muellerzr for his opinion, as he has far more experience with trainer and what might be the canonical way to control this 

> Caling in @muellerzr for his opinion, as he has far more experience with trainer and what might be the canonical way to control this

Okay, if we are looking for a proper long-term solution, then two more thoughts:
1. The current fix will still break FSDP if model upload is enabled. Maybe this won't affect too many people, since perhaps models large enough to require FSDP are also too large to be uploaded anyway, but still. Maybe someone more familiar with the internals like @muellerz can see what the issue is. I suspect it's either that (a) there's a different way to save an FSDP model than just calling `model.save_pretrained()`, or (b) it's something like "you need to call `model.forward()` before calling `model.save_pretrained()` so that FSDP can do all its setup correctly". If it's (b), perhaps that fix should happen elsewhere in the trainer, rather than here, or even in `accelerate`?
2. I'm not sure how all this works with PEFT/LORA models. Does `save_pretrained()` save the entire model, or just the adapter? If the latter, I could see how uploading the _base_ model at the start of training once would make sense potentially, if only the adapters are uploaded later on. 

Also, I see that in later `on_save()` and `on_train_end()` calls, the callback creates a fake `Trainer` object and then uses `save_model()` from there, instead of calling `model.save_pretrained()`. Does that make a difference for either of the above?


Thank You @mgerstgrasser
I had your same issue and your fix worked for me!!  

> The current fix will still break FSDP if model upload is enabled. Maybe this won't affect too many people, since perhaps models large enough to require FSDP are also too large to be uploaded anyway, but still. Maybe someone more familiar with the internals like @muellerzr can see what the issue is. I suspect it's either that (a) there's a different way to save an FSDP model than just calling model.save_pretrained(), or (b) it's something like "you need to call model.forward() before calling model.save_pretrained() so that FSDP can do all its setup correctly". If it's (b), perhaps that fix should happen elsewhere in the trainer, rather than here, or even in accelerate?

Yes. IMO we should not be uploading models by default, and esp with FSDP as generally if you're doing that we're doing our own checkpointing locally because otherwise we can be talking about transferring and uploading hundreds of gigs of weights, which will slow down things a ton.

With FSDP you end up in two scenarios depending on if the user has a sharded state dict or not, generally PyTorch is leading to the former nowadays so each shard has its own saved file. 

IMO: we don't upload FSDP models, and we can either add that later, or only do so if a user specifically requests it. 

> I'm not sure how all this works with PEFT/LORA models. Does save_pretrained() save the entire model, or just the adapter? If the latter, I could see how uploading the base model at the start of training once would make sense potentially, if only the adapters are uploaded later on.

Just the adapter. And again, if we have a large model we're doing PEFT on, we're at a similar situation uploading the base weights. 

Personally I'm not sure this should be a thing in general is my take unless we should have to explicitly enable it. 

@muellerzr @mgerstgrasser OK, shall we just completely remove uploading the model then? And then we can add this back in if it gets requested as a feature? 


> With FSDP you end up in two scenarios depending on if the user has a sharded state dict or not, generally PyTorch is leading to the former nowadays so each shard has its own saved file.

Ah, no, I meant that apparently calling `model.save()` here breaks FSDP training altogether, i.e. leads to a crash. See [#38095](https://github.com/huggingface/transformers/issues/30895) 

> IMO: we don't upload FSDP models, and we can either add that later, or only do so if a user specifically requests it.

But if we don't do this for FSDP anyway, then the above doesn't matter. :) 

> OK, shall we just completely remove uploading the model then? And then we can add this back in if it gets requested as a feature?

Fine with me, of course. Alternatively I could add an additional configuration option specifically for uploading the initial model, but I agree that it's unlikely to be something that many people would use, given the tradeoffs.

Yes, I agree let's remove it then bring it back if folks ask :) 

OK, I've reverted back to the version that just removes initial model upload altogether. @amyeroberts 

@mgerstgrasser For the quality checks - could you rebase to include the recent pushes from main? This should resolve these 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_30897). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.