Trying to run Krea2 on vllm Omni latest docker image and 2x 5090
it throws:
File “/usr/local/lib/python3.12/dist-packages/vllm_omni/diffusion/registry.py”, line 402, in initialize_model
raise ValueError(f"Model class {od_config.model_class_name} not found in diffusion model registry.")
ValueError: Model class Krea2Pipeline not found in diffusion model registry.
and later;
EOFError
(APIServer pid=1) INFO 07-26 23:37:15 [cumem.py:206] CuMemAllocator: sleep freed 0.00 GiB memory in total, of which 0.00 GiB is backed up in CPU and the rest 0.00 GiB is discarded directly.
(APIServer pid=1) Traceback (most recent call last):
The error “Model class Krea2Pipeline not found in diffusion model registry” means vLLM-Omni cannot find the Krea2Pipeline class in its internal registry. This usually happens if the model is not registered in vllm_omni/diffusion/registry.py or if your docker image is outdated and missing the latest model code. Krea2Pipeline is present in the registry in recent vLLM-Omni versions, so using an older image or a custom build may cause this issue. Ensure you are using the latest vllm-omni docker image and that your model files are up to date (see registry code, Krea2Pipeline docs).
The EOFError is a secondary error, likely caused by the failed model initialization. Fixing the registry/model version issue should resolve both errors. Would you like detailed troubleshooting steps or docker update instructions?
The latest stable release is v0.24.0, but the vllm/vllm-omni:cosmos3 image is a newer, model-specific build (6 days old) that may include additional or experimental features for Cosmos3 and possibly other diffusion models like Krea2. If you need the most up-to-date support for new diffusion models, try the :cosmos3 tag. If you want maximum stability, use v0.24.0 (docker tags reference).
Would you like more detail on the differences between these images or guidance on which to use for Krea2?
You should include --tensor-parallel-size 2 directly in your docker run command after vllm serve … as a CLI argument, not in a template file. For example:
@Herald that did not help. Out of memory again, it does not take the tensor-parallel from that docker run command, should it be defined in some template?
No, --tensor-parallel-size 2 should be passed as a CLI argument, not in a template file. If vLLM is not recognizing it and you still get OOM, double-check that the argument is after vllm serve … and not before. Also, ensure your model supports tensor parallelism and that both GPUs are visible to the container (see docs).
If OOM persists, check GPU memory usage and try reducing --max-model-len or --gpu-memory-utilization. Would you like step-by-step troubleshooting for multi-GPU OOM?
Correct, for vllm-omni diffusion models, tensor parallelism is often set in a YAML stage config template, not just via CLI. The CLI --tensor-parallel-size flag does not override the YAML; you must specify tensor_parallel_size in the stage config file and pass it with --stage-configs-path. See this example and explanation.
The Krea2 model uses a YAML stage config template to define parallelism and other settings. You should use the official example from the vllm-omni repo as a starting point, such as examples/serving/krea2_stage_config.yaml. In this file, set tensor_parallel_size: 2 under the appropriate stage (usually under engine_args).
It seems likely that there is no dedicated krea2_stage_config.yaml in the current vllm-omni repo. For Krea2, you should set tensor parallelism and devices in the deploy YAML file used for serving diffusion models. The typical template is named something like deploy.yaml or is referenced as --deploy-config in commands (see Krea-2 recipe).
You should:
Copy or create a deploy YAML (for example, from recipes/krea/Krea-2.md or adapt from other deploy configs).
In the YAML, under stages, set tensor_parallel_size: 2 and devices: “0,1” for the relevant stage.
Here is a minimal deploy YAML template for Krea2 on 2x RTX 5090 GPUs, setting tensor parallelism and device IDs. Save this as krea2_deploy.yaml and use --deploy-config /path/to/krea2_deploy.yaml when launching vllm-omni.
stages:
- stage_id: 0
model: krea/Krea-2-Turbo
tensor_parallel_size: 2
devices: "0,1"
engine_args:
gpu_memory_utilization: 0.9
# Add other engine_args as needed
@Herald this does not work, it still wont load the weight to both gpus, it loads to only one and about 10GB and then throws this error:
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 192.00 MiB. GPU 0 has a total capacity of 31.40 GiB of which 136.88 MiB is free. Including non-PyTorch memory, this process has 31.24 GiB memory in use. Of the allocated memory 30.61 GiB is allocated by PyTorch, and 31.19 MiB 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 (Redirecting…)
If --deploy-config is not working and the model only loads on one GPU, it likely means the tensor_parallel_size and devices settings are not being picked up by vllm-omni. For diffusion models, you must set tensor_parallel_size and devices under the correct stage and ensure you are using the correct CLI flag: --deploy-config for YAML-based deployment, or --stage-configs-path for stage-based configs (depending on the vllm-omni version and model type). Also, confirm that your YAML file is being read by checking the logs for the parsed config.
If the model still loads on one GPU, try passing --tensor-parallel-size 2 directly in the CLI as well, and verify that both GPUs are visible to Docker (nvidia-smi inside the container). If the problem persists, it may be a model or version limitation—some diffusion models only support tensor_parallel_size=1. See Krea-2 recipe and tensor parallelism docs for model-specific constraints.
Would you like step-by-step debugging to confirm your config is being used and both GPUs are available?
For Krea2 on vllm-omni, the correct template is a deploy YAML file (commonly named deploy.yaml or krea2_deploy.yaml). You must set tensor_parallel_size: 2 and devices: “0,1” under the relevant stage. However, Krea2’s tensor parallel support is model-dependent—some Krea2 checkpoints only support tensor_parallel_size=1 due to architectural constraints (see tensor parallelism guide).
To check which template is used, refer to the Krea-2 recipe. If your YAML is correct but only one GPU is used, the model may not support TP>1. Always check the logs for “tensor_parallel_size” and “devices” to confirm your config is loaded. If you need to force multi-GPU, try a different Krea2 checkpoint or consult the supported models list.
Would you like a troubleshooting checklist or a YAML example for another diffusion model that supports TP=2?