[FLUX] Add single prompt mode for FLUX inference and inference CI test #1969
Conversation
meta-cla
bot
added
the
CLA Signed
wwwjn
changed the title
| # Create mapping from local indices to global prompt indices | ||
| global_ids = list(range(global_rank, total_prompts, world_size)) | ||
| if single_prompt_mode: | ||
| # In single prompt mode, all ranks process the same prompt (index 0) |
There was a problem hiding this comment.
Sorry I somehow feel this PR is doing some overkill stuff. Would it be simpler if we just provide a truncated version of prompts.txt in test assets, say 1/2 prompts on each rank? Our goal is to make CI job lighter. Users should be fine if they always have to specify prompts in a .txt file? I just feel dealing with two paths doesn't seem to be elegant.
There was a problem hiding this comment.
Thanks! That would be easier. I think specify prompts in a .txt file is easier, I will remove the single prompt path. However, there are a minor bug to fix: If the number of prompts < number of ranks, some rank will get 0 prompts. During forward path of T5/clips encoder, the program will hang because FSDP is applied on encoder, and some ranks didn't run forward so all_gather will hang.
I will modify the PR to fix the bug and always using prompts.txt file.
There was a problem hiding this comment.
Sounds good! I think we can just error out early in that case, for simplicity.
as titled