Training LoRA for Flux Multi GPU not working with more than 2 GPUs

Hi.
Training LoRA for Flux.
This configuration works fine on 2 GPUs (RTX4090), but not with more GPUs (there are 8 Gpus in the machine).
It work on any pair of GPUs, but not with more than 2.
(running on WSL2/Ubuntu, it uses 22GB per GPU)

**The start command :**
```
CUDA_LAUNCH_BLOCKING=1 \
accelerate launch \
--config_file accelerate_config_L03.yaml \
--num_cpu_threads_per_process 2 \
--gpu_ids 3,4 \
/workspace/apps/sd-scripts/flux_train_network.py \
--fp8_base \
--highvram \
--network_train_unet_only \
--skip_cache_check \
--vae_batch_size 1 \
--pretrained_model_name_or_path sd_models/flux1-dev-fp8-e4m3fn.safetensors \
--clip_l sd_models/clip_l.safetensors \
--t5xxl sd_models/t5xxl_fp8_e4m3fn.safetensors \
--ae sd_models/ae.safetensors \
--cache_latents_to_disk \
--save_model_as safetensors \
--sdpa \
--persistent_data_loader_workers \
--max_data_loader_n_workers 8 \
--seed 42 \
--gradient_checkpointing \
--mixed_precision bf16 \
--save_precision bf16 \
--network_module networks.lora_flux \
--network_dim 32 \
--network_alpha 32 \
--optimizer_type adamw8bit \
--learning_rate 1e-4 \
--lr_scheduler constant_with_warmup \
--lr_warmup_steps 20 \
--cache_text_encoder_outputs \
--cache_text_encoder_outputs_to_disk \
--max_train_steps 5000 \
--save_every_n_epochs 2 \
--dataset_config lsvttnsnkrtrnr_L03_dataset01.toml \
--output_dir sd_output/lsvttnsnkrtrnr_sds_L03 \
--output_name lsvttnsnkrtrnr_sds_L03 \
--timestep_sampling sigmoid \
--discrete_flow_shift 3.1582 \
--model_prediction_type raw \
--guidance_scale 1.0 \
--max_grad_norm 1.0 \
--gradient_accumulation_steps 4 \
--clip_skip 2 \
--min_snr_gamma 5 \
--noise_offset 0.1 \
--sample_sampler="euler" \
--sample_prompts SneakerLVTrainer_prompts01.txt \
--sample_every_n_steps 25 \
--log_config
```

**The accelerate config file : accelerate_config_L03.yaml**
```
compute_environment: LOCAL_MACHINE
debug: true
distributed_type: MULTI_GPU
downcast_bf16: 'no'
enable_cpu_affinity: false
gpu_ids: 3,4
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 3
rdzv_backend: c10d
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

```

**The error with 3 or more GPUs :**

```
100%|███████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:00<00:00, 1290555.08it/s]
INFO bucket 2: resolution (640, 384), count: 4 train_util.py:997
INFO make buckets train_util.py:946
INFO bucket 3: resolution (704, 320), count: 8 train_util.py:997
INFO number of images (including repeats) / 各bucketの画像枚数（繰り返し回数を含む） train_util.py:992
INFO mean ar error (without repeats): 0.04492798538437573 train_util.py:1002
INFO bucket 0: resolution (512, 512), count: 19 train_util.py:997
INFO bucket 1: resolution (576, 448), count: 1 train_util.py:997
INFO bucket 2: resolution (640, 384), count: 4 train_util.py:997
INFO preparing accelerator train_network.py:369
INFO bucket 3: resolution (704, 320), count: 8 train_util.py:997
INFO mean ar error (without repeats): 0.04492798538437573 train_util.py:1002
INFO preparing accelerator train_network.py:369
[rank1]: Traceback (most recent call last):
[rank1]: File "/workspace/apps/sd-scripts/flux_train_network.py", line 564, in
[rank1]: trainer.train(args)
[rank1]: File "/mnt/d/TODAI/apps/sd-scripts/train_network.py", line 370, in train
[rank1]: accelerator = train_util.prepare_accelerator(args)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/mnt/d/TODAI/apps/sd-scripts/library/train_util.py", line 5182, in prepare_accelerator
[rank1]: accelerator = Accelerator(
[rank1]: ^^^^^^^^^^^^
[rank1]: File "/mnt/d/TODAI/apps/ai-toolkit/venv/lib/python3.12/site-packages/accelerate/accelerator.py", line 383, in __init__
[rank1]: self.state = AcceleratorState(
[rank1]: ^^^^^^^^^^^^^^^^^
[rank1]: File "/mnt/d/TODAI/apps/ai-toolkit/venv/lib/python3.12/site-packages/accelerate/state.py", line 846, in __init__
[rank1]: PartialState(cpu, **kwargs)
[rank1]: File "/mnt/d/TODAI/apps/ai-toolkit/venv/lib/python3.12/site-packages/accelerate/state.py", line 275, in __init__
[rank1]: self.set_device()
[rank1]: File "/mnt/d/TODAI/apps/ai-toolkit/venv/lib/python3.12/site-packages/accelerate/state.py", line 786, in set_device
[rank1]: device_module.set_device(self.device)
[rank1]: File "/mnt/d/TODAI/apps/ai-toolkit/venv/lib/python3.12/site-packages/torch/cuda/__init__.py", line 420, in set_device
[rank1]: torch._C._cuda_setDevice(device)
[rank1]: RuntimeError: CUDA error: CUDA-capable device(s) is/are busy or unavailable
[rank1]: CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
[rank1]: For debugging consider passing CUDA_LAUNCH_BLOCKING=1
[rank1]: Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

accelerator device: cuda:0
2024-10-24 19:32:29 INFO Checking the state dict: Diffusers or BFL, dev or schnell flux_utils.py:62
INFO Building Flux model dev from BFL checkpoint flux_utils.py:116
INFO Loading state dict from sd_models/flux1-dev.safetensors flux_utils.py:133
W1024 19:32:31.076000 140229863211136 torch/distributed/elastic/multiprocessing/api.py:858] Sending process 12125 closing signal SIGTERM
E1024 19:32:31.190000 140229863211136 torch/distributed/elastic/multiprocessing/api.py:833] failed (exitcode: 1) local_rank: 1 (pid: 12126) of binary: /mnt/d/TODAI/apps/ai-toolkit/venv/bin/python
Traceback (most recent call last):
File "/mnt/d/TODAI/apps/ai-toolkit/venv/bin/accelerate", line 8, in
sys.exit(main())
^^^^^^
File "/mnt/d/TODAI/apps/ai-toolkit/venv/lib/python3.12/site-packages/accelerate/commands/accelerate_cli.py", line 48, in main
args.func(args)
File "/mnt/d/TODAI/apps/ai-toolkit/venv/lib/python3.12/site-packages/accelerate/commands/launch.py", line 1097, in launch_command
multi_gpu_launcher(args)
File "/mnt/d/TODAI/apps/ai-toolkit/venv/lib/python3.12/site-packages/accelerate/commands/launch.py", line 734, in multi_gpu_launcher
distrib_run.run(args)
File "/mnt/d/TODAI/apps/ai-toolkit/venv/lib/python3.12/site-packages/torch/distributed/run.py", line 892, in run
elastic_launch(
File "/mnt/d/TODAI/apps/ai-toolkit/venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 133, in __call__
return launch_agent(self._config, self._entrypoint, list(args))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/d/TODAI/apps/ai-toolkit/venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 264, in launch_agent
raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:
============================================================
/workspace/apps/sd-scripts/flux_train_network.py FAILED
------------------------------------------------------------
Failures:

------------------------------------------------------------
Root Cause (first observed failure):
[0]:
time : 2024-10-24_19:32:31
host : BRMLCLXL0001.
rank : 1 (local_rank: 1)
exitcode : 1 (pid: 12126)
error_file:
traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
```

Any help appreciated.
Thanks.

Start with accelerate_config_L03.yaml and the multi-GPU setup in train_network.py, then inspect train_util.py:5182 where the Accelerator is prepared. Reproduce the failure using flux_train_network.py and the supplied launch command, comparing the configured process count with the selected GPUs. Done means identifying and resolving the cause of startup failure for more than two GPUs, with a verified multi-GPU run.

Written by the indexing model from the issue text.