Improve the performance and suitable for NPU computing #9642
Conversation
leisuzz
changed the title
|
@sayakpaul Please refer to this one, thanks! |
| @@ -540,7 +541,7 @@ def compute_vae_encodings(batch, vae): | |||
| with torch.no_grad(): | |||
| model_input = vae.encode(pixel_values).latent_dist.sample() | |||
| model_input = model_input * vae.config.scaling_factor | |||
| return {"model_input": model_input.cpu()} | |||
| return {"model_input": accelerator.gather(model_input)} | |||
There was a problem hiding this comment.
By using the accelerator, the communication time can be reduced
There was a problem hiding this comment.
I think the reason may caused by the pixel_values, as it is in vae.device (accelerator). Therefore, by changing the code, the accelerator can distribute and reduce the time cost.
There was a problem hiding this comment.
But isn't an all-gather a more expensive op?
There was a problem hiding this comment.
In fact, I tested three different approaches, the first one is accelerator.gather(model_input), the average FPS is 29.13 with training duration 530; the second one is model_input.to(accelerator.device), the average FPS is 27.41 with training duration 544; the last one is the original model_input.cpu(), the average FPS is 28.56 with training duration 537. Overall, with same hardware, the FPS will increase a little with accelerator.gather. I tested multiple times with accelerator.gather and model_input.cpu(), the average FPS in accelerator.gather is larger than model_input.cpu().
There was a problem hiding this comment.
Hmm, thanks! Since the performance improvement seems to be minor, do you think it makes sense to not change this?
There was a problem hiding this comment.
Make sense, I will change it back to the .cpu
There was a problem hiding this comment.
But I welcome you to also add a note on your findings on accelerate.gather() so that users are aware. I think that'd still be quite valuable.
| @@ -1091,8 +1095,7 @@ def compute_time_ids(original_size, crops_coords_top_left): | |||
| # Adapted from pipeline.StableDiffusionXLPipeline._get_add_time_ids | |||
| target_size = (args.resolution, args.resolution) | |||
| add_time_ids = list(original_size + crops_coords_top_left + target_size) | |||
| add_time_ids = torch.tensor([add_time_ids]) | |||
| add_time_ids = add_time_ids.to(accelerator.device, dtype=weight_dtype) | |||
| add_time_ids = torch.tensor([add_time_ids], device=accelerator.device, dtype=weight_dtype) | |||
|
@sayakpaul I've changed the code based on your suggestions, thanks! |
|
@sayakpaul Is there anything I need to change to merge this PR? I saw there is something wrong with the 'Build PR Documentation', but I didn't change this section I think |
|
Code quality check needs to pass. |
|
@sayakpaul I couldn't see which line is needed to be changed to pass |
|
Can you follow https://github.com/huggingface/diffusers/actions/runs/11319895166/job/31477993314?pr=9642 and follow the instructions from the logs? |
|
@sayakpaul It didn't show the specific problem, but I figured it would be the formatting issue about my comment in line 544. I've changed it. Thanks |
|
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. |
|
Thanks for your contributions! |
* Improve the performance and suitable for NPU * Improve the performance and suitable for NPU computing * Improve the performance and suitable for NPU * Improve the performance and suitable for NPU * Improve the performance and suitable for NPU * Improve the performance and suitable for NPU --------- Co-authored-by: 蒋硕 <[email protected]> Co-authored-by: Sayak Paul <[email protected]>
What does this PR do?
Improve the performance (FPS) while training, and suitable for NPU computing.
Selection for free memory for CUDA or NPU
Before submitting
documentation guidelines, and
here are tips on formatting docstrings.
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.