DeepSpeed adaption for flux-kontext

[leisuzz]

What does this PR do?

1. When I use the accelerate with deepspeed zero2, I found the error:
   "ValueError: At least one of the dataloaders passed to accelerate.prepare() has None as batch size. Please set an integer value in train_micro_batch_size_per_gpu in the deepspeed config file or assign integer value to AcceleratorState().deepspeed_plugin.deepspeed_config['train_micro_batch_size_per_gpu']."

I think the current batch_sampler does not support deepspeed. So I added this function to support it.

2. I've added the function to support saving checkpoints when using deepspeed

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? 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, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[leisuzz]

Cc: @sayakpaul @a-r-r-o-w Please take a look at this PR. Thanks!

[leisuzz]

@sayakpaul I've also added ds support for load/save ckpt function. Thanks!

[sayakpaul]

Thanks a lot!

[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.

[leisuzz]

@sayakpaul Please approve the workflow, thx :)

[leisuzz]

Hi @sayakpaul, It looks weird but I think the original code does not need the empty line. I added the empty line to it. Thanks :)

[leisuzz]

Cc: @sayakpaul, I checked with ruff and make quality, they all passed. Sorry for the trouble and thanks for your help :)