⛔[SD3 branch] Huber+SNR RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

[deGENERATIVE-SQUAD]

Pytorch 2.7.0, CUDA 12.8

What: RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! when using Huber + SNR

Where: train_util.py, line 6025

Fix:

From
alphas_cumprod = torch.index_select(noise_scheduler.alphas_cumprod, 0, timesteps.cpu())
To
alphas_cumprod = torch.index_select(noise_scheduler.alphas_cumprod, 0, timesteps.to(noise_scheduler.alphas_cumprod.device))

[kohya-ss]

It seems to work without fix in my environment. Could you please share the command line arguments for training?

[deGENERATIVE-SQUAD]

@kohya-ss

Traceback (most recent call last):
  File "K:\sd-scripts\sd-scripts-sd3\sdxl_train_network.py", line 229, in <module>
    trainer.train(args)
  File "K:\sd-scripts\sd-scripts-sd3\train_network.py", line 1403, in train
    loss = self.process_batch(
           ^^^^^^^^^^^^^^^^^^^
  File "K:\sd-scripts\sd-scripts-sd3\train_network.py", line 463, in process_batch
    huber_c = train_util.get_huber_threshold_if_needed(args, timesteps, noise_scheduler)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "K:\sd-scripts\sd-scripts-sd3\library\train_util.py", line 6025, in get_huber_threshold_if_needed
    alphas_cumprod = torch.index_select(noise_scheduler.alphas_cumprod, 0, timesteps.cpu())
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper_CUDA__index_select)
steps:   0%|                                                                                  | 0/1000 [00:19<?, ?it/s]
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "C:\Users\user\AppData\Local\Programs\Python\Python311\Scripts\accelerate.exe\__main__.py", line 7, in <module>
  File "C:\Users\user\AppData\Local\Programs\Python\Python311\Lib\site-packages\accelerate\commands\accelerate_cli.py", line 50, in main
    args.func(args)
  File "C:\Users\user\AppData\Local\Programs\Python\Python311\Lib\site-packages\accelerate\commands\launch.py", line 1198, in launch_command
    simple_launcher(args)
  File "C:\Users\user\AppData\Local\Programs\Python\Python311\Lib\site-packages\accelerate\commands\launch.py", line 785, in simple_launcher
    raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)`


`accelerate launch --num_cpu_threads_per_process 8 sdxl_train_network.py ^
	--pretrained_model_name_or_path="K:/model.safetensors" ^
	--train_data_dir="data" ^
	--output_dir="output_dir" ^
	--output_name="000" ^
	--network_args "algo=boft" "rescaled=True" "constraint=1" "conv_dim=32" "conv_alpha=1" "preset=full" "train_norm=False" ^
	--resolution="864" ^
	--save_model_as="safetensors" ^
	--network_module="lycoris.kohya" ^
	--shuffle_caption ^
	--max_train_epochs=100 ^
	--save_every_n_epochs=1 ^
	--save_state_on_train_end ^
	--save_precision=bf16 ^
	--network_dim=32 ^
	--network_alpha=1 ^
	--train_batch_size=2 ^
	--gradient_accumulation_steps=2 ^
	--max_data_loader_n_workers=1 ^
	--enable_bucket ^
	--bucket_reso_steps=64 ^
	--min_bucket_reso=768 ^
	--max_bucket_reso=1280 ^
	--mixed_precision="bf16" ^
	--caption_extension=".txt" ^
	--gradient_checkpointing ^
	--optimizer_type="prodigyplus.prodigy_plus_schedulefree.ProdigyPlusScheduleFree" ^
	--optimizer_args "d0=1e-5" "d_limiter=False" "use_stableadamw=False" "weight_decay_by_lr=False" ^
	--learning_rate=1 ^
	--network_train_unet_only ^
	--loss_type="huber" ^
	--huber_schedule="snr" ^
	--huber_c=0.1 ^
	--xformers ^
	--vae="K:/sdxl_vae_fix.safetensors" ^
	--full_bf16 ^
	--mem_eff_attn ^
	--seed=1 ^
	--logging_dir="logs" ^
	--log_with="tensorboard" ^
	--persistent_data_loader_workers ^`

[kohya-ss]

Thank you for shariing the command. I don't have LyCORIS installed so I can't test that. The command gives the error: TypeError: ProdigyPlusScheduleFree.__init__() got an unexpected keyword argument 'd_limiter'.

I think prodigy_plus_schedulefree or LyCORIS might be the reason. Could you try removing them to see what is causing the error?

[deGENERATIVE-SQUAD]

@kohya-ss

This definitely doesn’t depend on the optimizer being used - the same issue occurs with Adafactor or any other optimizer. d_limiter is not the root of the problem, it is a Prodigy Plus argument (see: https://github.com/LoganBooker/prodigy-plus-schedule-free/blob/main/prodigyplus/prodigy_plus_schedulefree.py), and it works fine.

I haven’t tested without LyCORIS yet, since I’m using non-standard algorithms and haven’t used the standard LoRA module - but I can try it later.

[rockerBOO]

Yeah I have had this issue myself as well but haven't identified why it happens sometimes. Probably need to identify why the alpha_cumprod is on non-CPU device and then this would work as expected.

[kjerk]

Pytorch 2.7.0, CUDA 12.8

What: RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! when using Huber + SNR

Where: train_util.py, line 6025

Fix:

From alphas_cumprod = torch.index_select(noise_scheduler.alphas_cumprod, 0, timesteps.cpu()) To alphas_cumprod = torch.index_select(noise_scheduler.alphas_cumprod, 0, timesteps.to(noise_scheduler.alphas_cumprod.device))

Same problem and fix worked on my Linux box as well. I'm running standard LoRA training so I don't think LyCORIS is involved.

Torch 2.7 CU 12.6, ProdigyPlusScheduleFree, loss_type: Huber, Huber schedule: SNR, SDPA attention. Triton uninstalled entirely.

[deGENERATIVE-SQUAD]

@kohya-ss
I tested the regular mode of lora training, the problem repeats. Thus, the problem does not depend on custom network.module.

[deGENERATIVE-SQUAD]

@kjerk

Triton uninstalled entirely.

By the way, I have a triton, so it's not the triton's fault.

[kohya-ss]

@kohya-ss I tested the regular mode of lora training, the problem repeats. Thus, the problem does not depend on custom network.module.

Thank you for your investigation. It's very strange that works in my environment.

I'd like to fix the code, but I'd also like to find out the cause if possible.

[Nineball459]

I just updated and was suffering this bug, I spent hours trying to understand what was wrong until I finally stumbled across this thread. I'm doing standard Lora training with Adafactor. This bug was popping up despite using the same parameters I had before I updated. I even did a fresh install with no luck.

Swapping out that line of code suddenly made it start training, but it was training at a much slower rate than the previous version I'd been using. It went from 1.3 seconds per step to 2.3 seconds. I'm really confused why this bug is still here over a month later for me.

[deGENERATIVE-SQUAD]

@Nineball459

It went from 1.3 seconds per step to 2.3 seconds.

Strange, no difference on my side. Show your full config please.

[Nineball459]

@Nineball459

It went from 1.3 seconds per step to 2.3 seconds.

Strange, no difference on my side. Show your full config please.

I'm using Kohya_SS GUI, I dl'd the latest version and my PC is an RTX 4070 Super with a 7800X3D
I cant give more info than that unfortunately as ive already rolled back to a Kohya_SS GUI commit from december and no longer have any issues and my training speed has returned to normal.

[rockerBOO]

The issue is

sd-scripts/library/train_util.py

Lines 5977 to 6011 in 7c075a9

	def get_noise_noisy_latents_and_timesteps(
	args, noise_scheduler, latents: torch.FloatTensor
	) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.IntTensor]:
	# Sample noise that we'll add to the latents
	noise = torch.randn_like(latents, device=latents.device)
	if args.noise_offset:
	if args.noise_offset_random_strength:
	noise_offset = torch.rand(1, device=latents.device) * args.noise_offset
	else:
	noise_offset = args.noise_offset
	noise = custom_train_functions.apply_noise_offset(latents, noise, noise_offset, args.adaptive_noise_scale)
	if args.multires_noise_iterations:
	noise = custom_train_functions.pyramid_noise_like(
	noise, latents.device, args.multires_noise_iterations, args.multires_noise_discount
	)
	# Sample a random timestep for each image
	b_size = latents.shape[0]
	min_timestep = 0 if args.min_timestep is None else args.min_timestep
	max_timestep = noise_scheduler.config.num_train_timesteps if args.max_timestep is None else args.max_timestep
	timesteps = get_timesteps(min_timestep, max_timestep, b_size, latents.device)
	# Add noise to the latents according to the noise magnitude at each timestep
	# (this is the forward diffusion process)
	if args.ip_noise_gamma:
	if args.ip_noise_gamma_random_strength:
	strength = torch.rand(1, device=latents.device) * args.ip_noise_gamma
	else:
	strength = args.ip_noise_gamma
	noisy_latents = noise_scheduler.add_noise(latents, noise + strength * torch.randn_like(latents), timesteps)
	else:
	noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
	return noise, noisy_latents, timesteps

noise_scheduler.add_noise() moves the alphas_cumprod to the GPU

And you can see it how it is working in the DDIMScheduler

https://github.com/huggingface/diffusers/blob/a4df8dbc40e170ff828f8d8f79c2c861c9f1748d/src/diffusers/schedulers/scheduling_ddim.py#L474-L498

So this issue is for sd and sdxl models specifically. We can move the alphas_cumprod back to the CPU after.