Repository navigation
Qwen-Image-2.1 2K输出，多参考图 <=5, 推理时OOM #1701
Description
（1）7B的Qwen-Image-2.1推理 180G显存的B200还显存不够，错误堆栈如下：
（2）9B的FLux2-klein-base-9B 推理相同的数据，显存都够用，没有仔细统计，但是一般不超过80G. 为什么Qwen-Image-2.1 OOM？
433 [rank0]: File "/usr/local/miniconda3/envs/qwen21/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
434 [rank0]: return func(*args, **kwargs)
435 [rank0]: File "/xxx/diffsynth/pipelines/qwen_image_21.py", line 106, in __call__
436 [rank0]: noise_pred = self.cfg_guided_model_fn(
437 [rank0]: File "/xxx/diffsynth/diffusion/base_pipeline.py", line 343, in cfg_guided_model_fn
438 [rank0]: noise_pred_nega = model_fn(**inputs_nega, **inputs_shared, **inputs_others)
439 [rank0]: File "/xxx/diffsynth/pipelines/qwen_image_21.py", line 336, in model_fn_qwen_image_21
440 [rank0]: model_output = dit(
441 [rank0]: File "/usr/local/miniconda3/envs/qwen21/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
442 [rank0]: return self._call_impl(*args, **kwargs)
443 [rank0]: File "/usr/local/miniconda3/envs/qwen21/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
444 [rank0]: return forward_call(*args, **kwargs)
445 [rank0]: File "/xxx/diffsynth/models/qwen_image_21_dit.py", line 534, in forward
446 [rank0]: attention_mask = build_qwenimage21_block_causal_mask(
447 [rank0]: File "/xxx/diffsynth/models/qwen_image_21_dit.py", line 189, in build_qwenimage21_block_causal_mask
448 [rank0]: return create_block_mask(
449 [rank0]: File "/usr/local/miniconda3/envs/qwen21/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 897, in create_block_mask
450 [rank0]: partial_block_mask, full_block_mask = _convert_mask_to_block_mask(
451 [rank0]: File "/usr/local/miniconda3/envs/qwen21/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 697, in _convert_mask_to_block_mask
452 [rank0]: mask_block_sum = mask.sum(
453 [rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 50.16 GiB. GPU 0 has a total capacity of 178.35 GiB of which 32.23 GiB is free. Including non-PyTorch memory, this
process has 146.10 GiB memory in use. Of the allocated memory 76.70 GiB is allocated by PyTorch, and 68.62 GiB is reserved by PyTorch but unallocated.
983 [rank2]: image = pipe(
984 [rank2]: File "/usr/local/miniconda3/envs/qwen21/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
985 [rank2]: return func(*args, **kwargs)
986 [rank2]: File "/xxx/diffsynth/pipelines/qwen_image_21.py", line 106, in __call__
987 [rank2]: noise_pred = self.cfg_guided_model_fn(
988 [rank2]: File "/xxx/diffsynth/diffusion/base_pipeline.py", line 335, in cfg_guided_model_fn
989 [rank2]: noise_pred_posi = model_fn(**inputs_posi, **inputs_shared, **inputs_others)
990 [rank2]: File "/xxx/diffsynth/pipelines/qwen_image_21.py", line 336, in model_fn_qwen_image_21
991 [rank2]: model_output = dit(
992 [rank2]: File "/usr/local/miniconda3/envs/qwen21/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
993 [rank2]: return self._call_impl(*args, **kwargs)
994 [rank2]: File "/usr/local/miniconda3/envs/qwen21/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
995 [rank2]: return forward_call(*args, **kwargs)
996 [rank2]: File "/xxx/diffsynth/models/qwen_image_21_dit.py", line 534, in forward
997 [rank2]: attention_mask = build_qwenimage21_block_causal_mask(
998 [rank2]: File "/xxx/diffsynth/models/qwen_image_21_dit.py", line 189, in build_qwenimage21_block_causal_mask
999 [rank2]: return create_block_mask(
1000 [rank2]: File "/usr/local/miniconda3/envs/qwen21/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 897, in create_block_mask
1001 [rank2]: partial_block_mask, full_block_mask = _convert_mask_to_block_mask(
1002 [rank2]: File "/usr/local/miniconda3/envs/qwen21/lib/python3.10/site-packages/torch/nn/attention/flex_attention.py", line 697, in _convert_mask_to_block_mask
1003 [rank2]: mask_block_sum = mask.sum(
1004 [rank2]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 72.38 GiB. GPU 2 has a total capacity of 178.35 GiB of which 21.79 GiB is free. Including non-PyTorch memory,
this process has 156.54 GiB memory in use. Of the allocated memory 50.40 GiB is allocated by PyTorch, and 105.36 GiB is reserved by PyTorch but unallocated. If reserved but unallocated
memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/
cuda.html#environment-variables)
Activity
gobigrassland commentedon Sep 22, 2026
@mi804 多谢。 但是问题还是没有完全解决
Qwen-Image-2.1
代码中，将create_block_mask中_compile=True, 还是没有完全解决。
我的输入输出设置：
发现cfg, use_kv_cache参数对显存和耗时产生了很大的影响
而且 在完成denosie step循环后，应该是latent解码过程中，峰值显存突然增大约10G
虽然单卡测试上述输入和输出条件下没有OOM, 但是我是8卡分布式批量推理时，遇到这个case仍然在出现了OOM. （设置cfg=4.0情况下出现的， cfg=1.0没有）
Qwen-Image-2.1 可以支持10张参考图，在cfg=4.0时，使用上述输入*2, 单卡推理也OOM了。
FLUX.2-klein-Base-9B
针对这两个模型结构，生图部分DiT参数量差异没有那么大的情况下
fix(qwen-image-2.1): avoid OOM when building block-causal mask at hig… (