RuntimeError: mat1 and mat2 must have the same dtype, but got Float and BFloat16

[Vincento-Wang]

[rank3]: ret_val = func(*args, **kwargs)
[rank3]: File "/home/anaconda3/envs/cogview/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2054, in forward
[rank3]: loss = self.module(*inputs, **kwargs)
[rank3]: File "/home/anaconda3/envs/cogview/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank3]: return self._call_impl(*args, **kwargs)
[rank3]: File "/home/anaconda3/envs/cogview/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1857, in _call_impl
[rank3]: return inner()
[rank3]: File "/home/anaconda3/envs/cogview/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1805, in inner
[rank3]: result = forward_call(*args, **kwargs)
[rank3]: File "/home/anaconda3/envs/cogview/lib/python3.10/site-packages/diffusers/models/transformers/transformer_cogview4.py", line 713, in forward
[rank3]: hidden_states, encoder_hidden_states = self.patch_embed(hidden_states, encoder_hidden_states)
[rank3]: File "/home/anaconda3/envs/cogview/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank3]: return self._call_impl(*args, **kwargs)
[rank3]: File "/home/anaconda3/envs/cogview/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
[rank3]: return forward_call(*args, **kwargs)
[rank3]: File "/home/anaconda3/envs/cogview/lib/python3.10/site-packages/diffusers/models/transformers/transformer_cogview4.py", line 59, in forward
[rank3]: hidden_states = self.proj(hidden_states)
[rank3]: File "/home/anaconda3/envs/cogview/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank3]: return self._call_impl(*args, **kwargs)
[rank3]: File "/home/anaconda3/envs/cogview/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
[rank3]: return forward_call(*args, **kwargs)
[rank3]: File "/home/anaconda3/envs/cogview/lib/python3.10/site-packages/torch/nn/modules/linear.py", line 125, in forward
[rank3]: return F.linear(input, self.weight, self.bias)
[rank3]: RuntimeError: mat1 and mat2 must have the same dtype, but got Float and BFloat16
Training steps: 0%| | 0/200 [00:04<?, ?it/s]
[rank0]:[W508 15:09:04.137948373 ProcessGroupNCCL.cpp:1476] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
W0508 15:09:05.675000 25220 site-packages/torch/distributed/elastic/multiprocessing/api.py:900] Sending process 25478 closing signal SIGTERM
W0508 15:09:05.675000 25220 site-packages/torch/distributed/elastic/multiprocessing/api.py:900] Sending process 25480 closing signal SIGTERM
W0508 15:09:05.677000 25220 site-packages/torch/distributed/elastic/multiprocessing/api.py:900] Sending process 25481 closing signal SIGTERM
E0508 15:09:06.356000 25220 site-packages/torch/distributed/elastic/multiprocessing/api.py:874] failed (exitcode: 1) local_rank: 1 (pid: 25479) of binary: /home/anaconda3/envs/cogview/bin/python3.10
Traceback (most recent call last):
File "/home/anaconda3/envs/cogview/bin/accelerate", line 8, in
sys.exit(main())
File "/home/anaconda3/envs/cogview/lib/python3.10/site-packages/accelerate/commands/accelerate_cli.py", line 50, in main
args.func(args)
File "/home/anaconda3/envs/cogview/lib/python3.10/site-packages/accelerate/commands/launch.py", line 1196, in launch_command
deepspeed_launcher(args)
File "/home/anaconda3/envs/cogview/lib/python3.10/site-packages/accelerate/commands/launch.py", line 878, in deepspeed_launcher
distrib_run.run(args)
File "/home/anaconda3/envs/cogview/lib/python3.10/site-packages/torch/distributed/run.py", line 883, in run
elastic_launch(
File "/home/anaconda3/envs/cogview/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 139, in call
return launch_agent(self._config, self._entrypoint, list(args))
File "/home/anaconda3/envs/cogview/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 270, in launch_agent
raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:

train.py FAILED

Failures:
<NO_OTHER_FAILURES>

Root Cause (first observed failure):
[0]:
time : 2025-05-08_15:09:05
host : 36aa4b45f54f
rank : 1 (local_rank: 1)
exitcode : 1 (pid: 25479)
error_file: <N/A>
traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html

[OleehyO]

Will fix in this pr

[Vincento-Wang]

found it. the computer_loss function have no tensor type specified.

[JasonHuangChina]

Will fix in this pr

@OleehyO When？ I face the same problem， when u fix it, pls kick me, thanks

[OleehyO]

Within three days

[Vincento-Wang]

Will fix in this pr

@OleehyO When？ I face the same problem， when u fix it, pls kick me, thanks

change the code like this will fix it.

def compute_loss(self, batch: dict[str, Any]) -> torch.Tensor:
dtype = self.get_training_dtype()
batch_size, text_seqlen, text_embedding_dim = batch["prompt_embedding"].shape
prompt_embeds = batch["prompt_embedding"]
latent = batch["encoded_image"]

    batch_size, num_channels, height, width = latent.shape
    image_height, image_width = self.state.train_resolution
    vae_scale_factor = 8
    image_seq_len = (
        (image_height // vae_scale_factor) * (image_width // vae_scale_factor)
    ) // (self.state.transformer_config.patch_size**2)
    image_seq_len = torch.tensor([image_seq_len], device=self.accelerator.device)

    text_attn_mask = batch["text_attn_mask"]

    num_train_timesteps = self.components.scheduler.config.num_train_timesteps
    sigmas = self.get_sigmas(batch_size, image_seq_len)
    timestep = self.get_timestep(batch_size, num_train_timesteps)

    noise = torch.randn_like(latent, dtype=dtype)
    model_input, model_label = self.add_noise(latent, noise, timestep, sigmas)

    original_size = torch.tensor(
        [[image_height, image_width] for _ in range(batch_size)],
        dtype=dtype,
        device=self.accelerator.device,
    )
    target_size = torch.tensor(
        [[image_height, image_width] for _ in range(batch_size)],
        dtype=dtype,
        device=self.accelerator.device,
    )
    crop_coords = torch.tensor(
        [[0, 0] for _ in range(batch_size)], dtype=dtype, device=self.accelerator.device
    )

    noise_pred_cond = self.components.transformer(
        hidden_states=model_input.to(dtype=dtype),
        encoder_hidden_states=prompt_embeds.to(dtype=dtype),
        timestep=timestep,
        original_size=original_size,
        target_size=target_size,
        crop_coords=crop_coords,
        return_dict=False,
        attention_kwargs={"text_attn_mask": text_attn_mask},
    )[0]

    loss = torch.mean((noise_pred_cond - model_label) ** 2, dim=(1, 2, 3))
    loss = loss.mean()

    return loss

[OleehyO]

Will fix in this pr

@OleehyO When？ I face the same problem， when u fix it, pls kick me, thanks

Have merged @JasonHuangChina