[kandinsky][train_text_to_image_prior.py] Fix LR scheduler when num_train_epochs is passed in a distributed training env

[adi-IL]

Fixes #8384.

What was wrong

examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py built the LR scheduler with num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps and num_training_steps=args.max_train_steps * args.gradient_accumulation_steps, and computed max_train_steps as num_train_epochs * ceil(len(dataloader)/gradient_accumulation) without sharding. In distributed training (accelerator.num_processes > 1) the dataloader length per process is smaller, so the scheduler saw the wrong total steps and warmup.

Same issue fixed for other trainers in #8312 and follow-ups (#8450, #9316, etc.). This file was one of the remaining unchecked items in #8384.

What changed

Mirrored the fixed pattern from examples/text_to_image/train_text_to_image.py:

• num_warmup_steps_for_scheduler = args.lr_warmup_steps * accelerator.num_processes
• when max_train_steps is None, compute len_train_dataloader_after_sharding = ceil(len(dataloader)/num_processes), then num_training_steps_for_scheduler = num_epochs * ceil(len_sharded/grad_accum) * num_processes; else max_train_steps * num_processes
• pass those to get_scheduler
• after accelerator.prepare, recalc max_train_steps when it was derived from epochs and warn if the expected sharded length does not match the actual length

Single file, single issue.

Coordination

• Issue: [Community] Help us fix the LR schedulers when num_train_epochs is passed in a distributed training env #8384 (standing community checklist, Good second issue)
• This is the next unchecked item after prior PRs [textual_inversion.py] Fix the LR scheduler when num_train_epochs is passed in a distributed training env #14546, [train_t2i_adapter_sdxl.py] Fix the LR scheduler when num_train_epochs is passed in a distributed training env #14540, [train_text_to_image_lora_sdxl.py] Fix the LR scheduler when num_train_epochs is passed in a distributed training env #14528, [train_text_to_image_sdxl.py] Fix the LR scheduler when num_train_epochs is passed in a distributed training env #14527, [train_text_to_image_decoder.py] Fix the LR scheduler when num_train_epochs is passed in a distributed training env #14542. No open PR targets this file.

Tests

• python -m py_compile examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py -> OK
• Pattern matches already-merged sibling fixes, so make style/make quality relevant checks are limited to this single script.

Self-review

• Chose single-file scope per the guide. Chameleon copy of the sibling fix. No drive-by refactors. No public API change. Warning branch preserves existing behavior for non-sharded case.