flux-2-klein-9B lora training error

[wuyuetiger]

when i train flux-2-klein-9B lora, this message is shown and program is blocked. error shows Qwen3_4B, i think flux-2-klein needs Qwen3_8B. i have no idea. help!

[rank0]: Traceback (most recent call last):
[rank0]: File "/mnt/i/AI/diffusion-pipe/train.py", line 495, in
[rank0]: dataset_manager.cache()
[rank0]: File "/mnt/i/AI/diffusion-pipe/utils/dataset.py", line 1148, in cache
[rank0]: self._handle_task(task)
[rank0]: File "/mnt/i/AI/diffusion-pipe/venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]: return func(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/utils/dataset.py", line 1187, in _handle_task
[rank0]: self.submodels[id].load_model_if_needed()
[rank0]: File "/mnt/i/AI/diffusion-pipe/models/base.py", line 362, in load_model_if_needed
[rank0]: self._model = self._load_fn()
[rank0]: ^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/models/base.py", line 414, in load_fn
[rank0]: return comfy.sd.load_clip(ckpt_paths=paths, clip_type=clip_type)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/models/../submodules/ComfyUI/comfy/sd.py", line 970, in load_clip
[rank0]: return load_text_encoder_state_dicts(clip_data, embedding_directory=embedding_directory, clip_type=clip_type, model_options=model_options)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/models/../submodules/ComfyUI/comfy/sd.py", line 1221, in load_text_encoder_state_dicts
[rank0]: m, u = clip.load_sd(c)
[rank0]: ^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/models/../submodules/ComfyUI/comfy/sd.py", line 264, in load_sd
[rank0]: return self.cond_stage_model.load_sd(sd)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/models/../submodules/ComfyUI/comfy/sd1_clip.py", line 711, in load_sd
[rank0]: return getattr(self, self.clip).load_sd(sd)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/models/../submodules/ComfyUI/comfy/sd1_clip.py", line 305, in load_sd
[rank0]: return self.transformer.load_state_dict(sd, strict=False)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 2593, in load_state_dict
[rank0]: raise RuntimeError(
[rank0]: RuntimeError: Error(s) in loading state_dict for Qwen3_4B:
[rank0]: size mismatch for model.embed_tokens.weight: copying a param with shape torch.Size([151936, 4096]) from checkpoint, the shape in current model is torch.Size([151936, 2560]).

[nlienard]

you are using Qwen3_4B instead of Qwen3_8B.

[wuyuetiger]

use Qwen3_4B instead of Qwen3_8B. still error.
i have updated latest version.

[rank0]: Traceback (most recent call last):
[rank0]: File "/mnt/i/ai/diffusion-pipe/train.py", line 502, in
[rank0]: model.load_diffusion_model()
[rank0]: File "/mnt/i/ai/diffusion-pipe/models/flux2.py", line 77, in load_diffusion_model
[rank0]: super().load_diffusion_model()
[rank0]: File "/mnt/i/ai/diffusion-pipe/models/base.py", line 422, in load_diffusion_model
[rank0]: self.model_patcher = comfy.sd.load_diffusion_model(self.model_config['diffusion_model'], model_options=model_options)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/ai/diffusion-pipe/models/../submodules/ComfyUI/comfy/sd.py", line 1442, in load_diffusion_model
[rank0]: model = load_diffusion_model_state_dict(sd, model_options=model_options, metadata=metadata)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/ai/diffusion-pipe/models/../submodules/ComfyUI/comfy/sd.py", line 1431, in load_diffusion_model_state_dict
[rank0]: model = model_config.get_model(new_sd, "")
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/ai/diffusion-pipe/models/../submodules/ComfyUI/comfy/supported_models.py", line 767, in get_model
[rank0]: out = model_base.Flux2(self, device=device)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/ai/diffusion-pipe/models/../submodules/ComfyUI/comfy/model_base.py", line 849, in init
[rank0]: super().init(model_config, model_type, device=device, unet_model=unet_model)
[rank0]: File "/mnt/i/ai/diffusion-pipe/models/../submodules/ComfyUI/comfy/model_base.py", line 140, in init
[rank0]: self.diffusion_model = unet_model(**unet_config, device=device, operations=operations)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/ai/diffusion-pipe/models/../submodules/ComfyUI/comfy/ldm/flux/model.py", line 58, in init
[rank0]: raise ValueError(
[rank0]: ValueError: Hidden size 4096 must be divisible by num_heads 48
[rank0]:[W120 07:32:45.866113010 ProcessGroupNCCL.cpp:1479] 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())

[tdrussell]

In the first error, the ComfyUI code thinks it has qwen3_4b but then when it tries to load the weights they are for 8b so the size doesn't match. The second error is different but it is also a problem with the comfy backend failing to load the diffusion model.

What model files are you using for the diffusion model and text encoder? They should be the same ones that work in ComfyUI.

[wuyuetiger]

i try the models from hugging-face black-forest-labs/FLUX.2-klein-base-9B including flux-2-klein-base-9b.safetensors, qwen_3_8b, flux2-vae.safetensors
but error: ValueError: Hidden size 4096 must be divisible by num_heads 48
then i try the models from Comfy-Org/vae-text-encorder-for-flux-klein-9b
but still error: ValueError: Hidden size 4096 must be divisible by num_heads 48
how can i do?

[wuyuetiger]

i find flux-2-klein-base-9b.safetensors and flux-2-klein-9b.safetensors size is both 18,157,185,168 bytes. is that right?

[tdrussell]

Is your ComfyUI submodule up to date? You should do git submodule update after pull (if a submodule was updated), or just always do git pull --recurse-submodules

[wuyuetiger]

Is your ComfyUI submodule up to date? You should do git submodule update after pull (if a submodule was updated), or just always do git pull --recurse-submodules

you are right. i updated submodule comfyui, then i have solved above two problems. thanks.

but i meet new problem, the following content is new error:

[rank0]: Traceback (most recent call last):
[rank0]: File "/mnt/i/ai/diffusion-pipe/train.py", line 852, in
[rank0]: loss = model_engine.train_batch(iterator).item()
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/venv/lib/python3.12/site-packages/deepspeed/runtime/pipe/engine.py", line 388, in train_batch
[rank0]: self._exec_schedule(sched)
[rank0]: File "/mnt/i/AI/diffusion-pipe/venv/lib/python3.12/site-packages/deepspeed/runtime/pipe/engine.py", line 1368, in _exec_schedule
[rank0]: self._exec_instr(**cmd.kwargs)
[rank0]: File "/mnt/i/AI/diffusion-pipe/venv/lib/python3.12/site-packages/deepspeed/runtime/pipe/engine.py", line 737, in _exec_forward_pass
[rank0]: outputs = super().forward(inputs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/venv/lib/python3.12/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
[rank0]: ret_val = func(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/venv/lib/python3.12/site-packages/deepspeed/runtime/engine.py", line 2236, in forward
[rank0]: loss = self.module(*inputs, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1857, in _call_impl
[rank0]: return inner()
[rank0]: ^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1805, in inner
[rank0]: result = forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/venv/lib/python3.12/site-packages/deepspeed/runtime/pipe/module.py", line 390, in forward
[rank0]: x = exec_range_func(start_idx, end_idx)(*x)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/venv/lib/python3.12/site-packages/deepspeed/runtime/pipe/module.py", line 365, in exec_func
[rank0]: inputs = layer(inputs)
[rank0]: ^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/venv/lib/python3.12/site-packages/torch/amp/autocast_mode.py", line 44, in decorate_autocast
[rank0]: return func(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/ai/diffusion-pipe/models/flux2.py", line 305, in forward
[rank0]: txt = self.txt_in(txt)
[rank0]: ^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/AI/diffusion-pipe/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/ai/diffusion-pipe/utils/../submodules/ComfyUI/comfy/ops.py", line 164, in forward
[rank0]: return self.forward_comfy_cast_weights(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/i/ai/diffusion-pipe/utils/../submodules/ComfyUI/comfy/ops.py", line 157, in forward_comfy_cast_weights
[rank0]: x = torch.nn.functional.linear(input, weight, bias)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: RuntimeError: mat1 and mat2 shapes cannot be multiplied (85x2560 and 12288x4096)

[wuyuetiger]

this is my config file

change this

output_dir = '/mnt/i/AI/diffusion-pipe/output'

and this

dataset = '/mnt/i/AI/diffusion-pipe/examples/songxiaorui_dataset.toml'

training settings

epochs = 100
micro_batch_size_per_gpu = 1
pipeline_stages = 1
gradient_accumulation_steps = 1
gradient_clipping = 1
warmup_steps = 10

eval settings

eval_every_n_epochs = 1
eval_before_first_step = true
eval_micro_batch_size_per_gpu = 1
eval_gradient_accumulation_steps = 1

misc settings

save_every_n_epochs = 5
checkpoint_every_n_minutes = 120
activation_checkpointing = 'unsloth'
partition_method = 'parameters'
save_dtype = 'bfloat16'
caching_batch_size = 1
steps_per_print = 1
video_clip_mode = 'single_beginning'
blocks_to_swap = 8

[model]
type = 'flux2'
diffusion_model = '/mnt/i/AI/diffusion-pipe/model/FLUX.2-klein-base-9B/flux-2-klein-base-9b.safetensors'
vae = '/mnt/i/AI/diffusion-pipe/model/FLUX.2-klein-base-9B/flux2-vae.safetensors'
text_encoders = [
{path = '/mnt/i/AI/diffusion-pipe/model/FLUX.2-klein-base-9B/qwen_3_8b.safetensors', type = 'flux2'}
]
dtype = 'bfloat16'
diffusion_model_dtype = 'float8'
timestep_sample_method = 'logit_normal'
shift = 3

[adapter]
type = 'lora'
rank = 128
dtype = 'bfloat16'

[optimizer]
type = 'AdamW8bitKahan'
lr = 2e-5
betas = [0.9, 0.99]
weight_decay = 0.01
stabilize = false

[tdrussell]

[rank0]: RuntimeError: mat1 and mat2 shapes cannot be multiplied (85x2560 and 12288x4096)

It looks like the text embeddings have dimension 2560 which means qwen3 4b. If you ever ran the caching with the wrong model it will just load those embeddings each time. It has no way to know you changed the text encoder and rerun the caching.

Try running with --regenerate_cache flag and see if that fixes it.

[wuyuetiger]

using --regenerate_cache, it works! thanks.