Fix issues for saving checkpointing steps by leisuzz · Pull Request #9891 · huggingface/diffusers

leisuzz · 2024-11-08T14:43:09Z

What does this PR do?

These modification can help to save the checkpoint steps while training. Otherwise it will just stuck for too long and timeout.
Fixes get stuck when save_state using DeepSpeed backend under training train_text_to_image_lora #2606
Bug fix for weight pop from empty list

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 · 2024-11-13T04:43:56Z

@sayakpaul Please take a look at this PR, thanks for your help!

sayakpaul

Thanks for your PR. Can you please modify a single file first and discuss the changes first?

leisuzz · 2024-11-13T13:34:19Z

When I did the deambooth flux without Lora, and save the check pointing. It stuck for a while and break. So I think they all need these modifications. I can only do the flux ones if you want

sayakpaul · 2024-11-13T13:36:21Z

Yeah let's change a single file first and then we can discuss the changes first.

leisuzz · 2024-11-13T13:37:21Z

Sure

leisuzz · 2024-11-14T00:51:59Z

@sayakpaul I already changed the modifications only on FLUX models

sayakpaul · 2024-11-14T22:36:13Z

examples/dreambooth/train_dreambooth_flux.py

+                if global_step % args.checkpointing_steps == 0:
+                    # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`


Well, there is a better way to handle it:

if accelerator.is_main_process or accelerator.distributed_type == DistributedType.DEEPSPEED:

sayakpaul

Thaks for your PR!

You can refer to the following scripts:

To see how we handle saving and loading from checkpoints when using DeepSpeed.

Search for DistributedType.DEEPSPEED.

leisuzz · 2024-11-15T02:38:25Z

@sayakpaul I've changed it based on the reference

sayakpaul

Thanks! Left more comments. LMK if they're clear.

sayakpaul · 2024-11-15T08:40:12Z

examples/dreambooth/train_dreambooth_flux.py

+
                    model.load_state_dict(load_model.state_dict())
-                except Exception:
+                elif isinstance(unwrap_model(model), (CLIPTextModelWithProjection, T5EncoderModel)):


We don't support fine-tuning the T5 model. So, this seems wrong. It should just be CLIPTextModelWithProjection, no?

sayakpaul · 2024-11-15T08:40:27Z

examples/dreambooth/train_dreambooth_flux.py

+                        try:
+                            load_model = T5EncoderModel.from_pretrained(input_dir, subfolder="text_encoder_2")
+                            model(**load_model.config)
+                            model.load_state_dict(load_model.state_dict())
+                        except Exception:
+                            raise ValueError(f"Couldn't load the model of type: ({type(model)}).")
+                else:
+                    raise ValueError(f"Unsupported model found: {type(model)=}")


Same for this.

sayakpaul · 2024-11-15T08:41:36Z

examples/dreambooth/train_dreambooth_flux.py

-                try:
-                    load_model = CLIPTextModelWithProjection.from_pretrained(input_dir, subfolder="text_encoder")
-                    model(**load_model.config)
+        if not accelerator.distributed_type == DistributedType.DEEPSPEED:


We also need to handle the case when we're actually doing DeepSpeed training. Similar to:
https://github.com/a-r-r-o-w/cogvideox-factory/blob/d63a826f37758eccf226710f94f6c3a4d4ee7a25/training/cogvideox_text_to_video_sft.py#L385

sayakpaul · 2024-11-15T08:42:37Z

examples/dreambooth/train_dreambooth_lora_flux.py


-        while len(models) > 0:
-            model = models.pop()
+        if not accelerator.distributed_type == DistributedType.DEEPSPEED:


Same. We also need to handle the case when we are not doing DeepSpeed training. Reference:
https://github.com/a-r-r-o-w/cogvideox-factory/blob/d63a826f37758eccf226710f94f6c3a4d4ee7a25/training/cogvideox_text_to_video_lora.py#L396

sayakpaul · 2024-11-15T10:52:42Z

examples/dreambooth/train_dreambooth_flux.py

+                if weights:
+                    weights.pop()

    def load_model_hook(models, input_dir):


Seems like we're not handling the loading case appropriately here. I repeated this multiple times now but please refer to the changes here to get an idea of what is required.

In summary, we're not dealing with the changes required to load the state dict in the models being trained when DeepSpeed is enabled.

sayakpaul · 2024-11-25T13:03:56Z

Gentle ping @leisuzz

leisuzz · 2024-11-26T01:41:04Z

@sayakpaul I tested yesterday, with the DistributedType function, the checkpointing can be saved. And the if weights: condition will remove the issue cause I encountered once, that's why I added this

sayakpaul · 2024-11-26T02:13:26Z

@sayakpaul I tested yesterday, with the DistributedType function, the checkpointing can be saved. And the if weights: condition will remove the issue cause I encountered once, that's why I added this

Yeah but we have not addressed the comments fully yet. Specifically, we haven't addressed what's needed to enable loading of state dicts properly when using DeepSpeed. We need to address that.

leisuzz · 2024-11-26T02:16:44Z

I will check that

leisuzz mentioned this pull request Nov 8, 2024

get stuck when save_state using DeepSpeed backend under training train_text_to_image_lora #2606

Closed

sayakpaul reviewed Nov 13, 2024

View reviewed changes

sayakpaul reviewed Nov 14, 2024

View reviewed changes

sayakpaul reviewed Nov 15, 2024

View reviewed changes

leisuzz closed this Dec 17, 2024

leisuzz force-pushed the main branch from 7eae07b to 7667cfc Compare December 17, 2024 01:42

		if global_step % args.checkpointing_steps == 0:
		# _before_ saving state, check if this save would set us over the `checkpoints_total_limit`

Conversation

leisuzz commented Nov 8, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

leisuzz commented Nov 13, 2024

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

leisuzz commented Nov 13, 2024

Uh oh!

sayakpaul commented Nov 13, 2024

Uh oh!

leisuzz commented Nov 13, 2024

Uh oh!

leisuzz commented Nov 14, 2024

Uh oh!

sayakpaul Nov 14, 2024

Choose a reason for hiding this comment

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

leisuzz commented Nov 15, 2024

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

sayakpaul Nov 15, 2024

Choose a reason for hiding this comment

Uh oh!

sayakpaul Nov 15, 2024

Choose a reason for hiding this comment

Uh oh!

sayakpaul Nov 15, 2024

Choose a reason for hiding this comment

Uh oh!

sayakpaul Nov 15, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul Nov 15, 2024

Choose a reason for hiding this comment

Uh oh!

sayakpaul commented Nov 25, 2024

Uh oh!

leisuzz commented Nov 26, 2024

Uh oh!

sayakpaul commented Nov 26, 2024

Uh oh!

leisuzz commented Nov 26, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

leisuzz commented Nov 8, 2024 •

edited

Loading

sayakpaul Nov 15, 2024 •

edited

Loading