Closed alexghergh closed 1 month ago
cc @muellerzr @SunMarc
Are you supposed to change the state dict to a FULL_STATE_DICT before saving?
No, we support SHARDED but we may need to upstream this towards the saving, I only recently brought that into the accelerate side. Let me get to this today
Any updates on this?
Note that checkpoints do indeed seem to save normally
Do you see the checkpoints saved by the trainer? at least during the training period, i do not see any checkpoint saved by the trainer.
@katzurik Hmm, now that you mention, I'm double-checking in my mind whether that was really the case. I'm pretty sure the answer is yes (I'm quite sure I would've said otherwise in my initial post if that wasn't the case), but it was quite some time ago and I can't remember exactly.
In the meantime, my team and I moved to manually setting the state dict to FULL at the end (trainer.accelerator.state.fsdp_plugin.set_state_dict_type('FULL_STATE_DICT'
), which works exactly as expected.
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.
Please note that issues that do not follow the contributing guidelines are likely to be ignored.
I'd like to bump this issue up since it has not been updated.
The code explicitly ignores the case of SHARDED_STATE_DICT
and LOCAL_STATE_DICT
as you can see in the code here: https://github.com/huggingface/transformers/blob/main/src/transformers/trainer.py#L3533-3539
TBH I'm not sure what the correct thing here is since I'm new to the FSDP world, but this seems to be what's needed: https://huggingface.co/docs/accelerate/en/usage_guides/fsdp#saving-and-loading
System Info
transformers
version: 4.39.2Who can help?
@pacman100
Information
Tasks
examples
folder (such as GLUE/SQuAD, ...)Reproduction
Hey folks,
The issue seems to be simple enough. I tried to train an FSDP model using a multi-node setup, with transformers + accelerate. I'm launching my training script (which doesn't seem relevant to the issue for now, so I'll skip it for brevity) using
accelerate launch --config-file config.yaml ... python train.py
. Theconfig.yaml
looks like this:The issue, though, is when using transformers's built-in Trainer class and trying to save the model at the end of training. Trying to run:
This, however, doesn't seem to save any file in
save_dir
(as a matter of fact, it doesn't even create the directory).The problem seems to be in Trainer.save_model(), which tests for FSDP, and then tests for
FULL_STATE_DICT
. So this leads to the obvious question, what happens to theSHARDED_STATE_DICT
models? I didn't manage to find anything about this in the docs.Am I missing something? Are you supposed to change the state dict to a
FULL_STATE_DICT
before saving?Note that checkpoints do indeed seem to save normally (they are sharded, so it seems that every node saves it's part, which is expected). It is just the final
save_model
call which seems to faulter.Thanks a lot!
Expected behavior
The model should be saved in a safetensors format, with all the weights, inside
save_dir
, as described above, regardless of whether accelerate usesFULL_STATE_DICT
orSHARDED_STATE_DICT
with FSDP.