VAELoaderDisTorch2MultiGPU Crashed on VAE Encode/Decode Audio

[van-thieu]

With virtual_vram_gb config at 96, or more than 0 the workflow crashed upon reaching the default ComfyUI VAE Audio encoding/decoding nodes:

[MultiGPU DisTorch V2] DisTorch loading completed.
[MultiGPU DisTorch V2] Total memory: 321.70MB
!!! Exception during processing !!! Input type (torch.cuda.HalfTensor) and bias type (torch.HalfTensor) should be the same
Traceback (most recent call last):
File "/home/ai/comfyui/execution.py", line 534, in execute
output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ai/comfyui/execution.py", line 334, in get_output_data
return_values = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ai/comfyui/execution.py", line 308, in _async_map_node_over_list
await process_inputs(input_dict, i)
File "/home/ai/comfyui/execution.py", line 296, in process_inputs
result = f(**inputs)
File "/home/ai/comfyui/comfy_api/internal/init.py", line 149, in wrapped_func
return method(locked_class, **inputs)
File "/home/ai/comfyui/comfy_api/latest/_io.py", line 1826, in EXECUTE_NORMALIZED
to_return = cls.execute(*args, **kwargs)
File "/home/ai/comfyui/comfy_extras/nodes_audio.py", line 106, in execute
t = vae.encode(waveform.movedim(1, -1))
File "/home/ai/comfyui/comfy/sd.py", line 1075, in encode
model_management.raise_non_oom(e)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^
File "/home/ai/comfyui/comfy/model_management.py", line 289, in raise_non_oom
raise e
File "/home/ai/comfyui/comfy/sd.py", line 1068, in encode
out = self.first_stage_model.encode(pixels_in)
File "/home/ai/comfyui/comfy/ldm/audio/autoencoder.py", line 272, in encode
return self.bottleneck.encode(self.encoder(x))
~~~~~~~~~~~~^^^
File "/home/ai/comfyui/lib/python3.14/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "/home/ai/comfyui/lib/python3.14/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
return forward_call(*args, **kwargs)
File "/home/ai/comfyui/comfy/ldm/audio/autoencoder.py", line 207, in forward
return self.layers(x)
~~~~~~~~~~~^^^
File "/home/ai/comfyui/lib/python3.14/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "/home/ai/comfyui/lib/python3.14/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
return forward_call(*args, **kwargs)
File "/home/ai/comfyui/lib/python3.14/site-packages/torch/nn/modules/container.py", line 253, in forward
input = module(input)
File "/home/ai/comfyui/lib/python3.14/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "/home/ai/comfyui/lib/python3.14/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
return forward_call(*args, **kwargs)
File "/home/ai/comfyui/comfy/ops.py", line 425, in forward
return super().forward(*args, **kwargs)
~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "/home/ai/comfyui/lib/python3.14/site-packages/torch/nn/modules/conv.py", line 375, in forward
return self._conv_forward(input, self.weight, self.bias)
~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ai/comfyui/lib/python3.14/site-packages/torch/nn/modules/conv.py", line 370, in _conv_forward
return F.conv1d(
~~~~~~~~^
input, weight, bias, self.stride, self.padding, self.dilation, self.groups
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
)
^
RuntimeError: Input type (torch.cuda.HalfTensor) and bias type (torch.HalfTensor) should be the same

[Booyaka101]

Did some triage on this — leaving notes in case it helps you (or whoever picks this up). The signature Input type (torch.cuda.HalfTensor) and bias type (torch.HalfTensor) should be the same is PyTorch's way of saying the input tensor is on CUDA but the layer's bias is still on CPU. With virtual_vram_gb > 0 for VAE this is consistent with a block-swap path that's moving weights but not biases (or moving them to the wrong device for that specific module class).

Two specific places in distorch_2.py look load-bearing for the audio-VAE case:

1. distorch_2.py:474 — block_list = [b for b in distributable_all_blocks if (b[3] >= MIN_BLOCK_THRESHOLD and hasattr(b[1], "bias"))]. The hasattr(b[1], "bias") filter means modules without an explicit bias attribute land in tiny_block_list and don't get the full DisTorch2 treatment. Audio VAEs (Wan2.2 / Hunyuan / MMAudio) use Conv1d + custom WaveNet-style residual blocks with norms that may not have bias in the same shape image VAEs do, so the audio-side modules might be silently routed differently than image-side modules.

2. distorch_2.py:285-295 — the device-cast block:

   bias_key = f"{module_name}.bias"
   ...
   if bias_key in self.patches:
       self.patch_weight_to_device(bias_key, device_to=device_to)

   This relies on bias_key being present in self.patches. If the audio VAE has biases as plain nn.Parameters registered under non-standard names (or as part of a larger weight tensor like in WaveNet dilated_conv), they'll skip the device move, stay on CPU, and then explode the moment Conv1d(input) runs against them with input on CUDA.

To narrow it down precisely, two pieces of info from the original repro would help:

• Which audio VAE node specifically — Wan2.2 audio, Hunyuan audio, MMAudio, something else? (The class name in the workflow / a snippet of the workflow JSON would do.)
• The full traceback past the line shown — specifically which module class (Conv1d? WeightNorm-wrapped Conv? a custom AudioVAE block?) is the one whose bias is on CPU. The module type at the failing frame tells us whether the fix lives at line 474 (filter), line 285 (key matching), or in module_object.to(device_to) at line 281 not recursing into a wrapped submodule.

Happy to put up a PR once we know which of those it is. I don't have an audio-VAE workflow to repro locally — if @van-thieu or anyone else hitting this can share the workflow JSON and the rest of the traceback, I can target the fix accurately.

[van-thieu]

@Booyaka101 it's the new AceStep 1.5 XL Turbo template workflow, instead of the default clip/vae/diffusion loader, I've swapped with your multigpu dist loaders.

Additional info:

uv run ~/comfyui/main.py --enable-manager --listen 0.0.0.0 --port 8888 --cuda-malloc --cache-ram 32 --async-offload 32 --enable-dynamic-vram

Total VRAM 32607 MB, total RAM 128812 MB
pytorch version: 2.11.0+cu130
Set vram state to: NORMAL_VRAM
Device: cuda:0 NVIDIA GeForce RTX 5090 : cudaMallocAsync
Using async weight offloading with 32 streams
Enabled pinned memory 115930.0
Using pytorch attention
aimdo: /project/src-posix/cuda-funchooks.c:52:DEBUG:aimdo_setup_hooks: hooks successfully installed
aimdo: /project/src/control.c:151:INFO:comfy-aimdo inited for GPU: NVIDIA GeForce RTX 5090 (VRAM: 32606 MB)
DynamicVRAM support detected and enabled
Python version: 3.14.4 (main, Apr 8 2026, 04:02:31) [GCC 15.2.0]
ComfyUI version: 0.20.1
comfy-aimdo version: 0.3.0
comfy-kitchen version: 0.2.8
ComfyUI frontend version: 1.42.15
[Prompt Server] web root: /home/ai/comfyui/lib/python3.14/site-packages/comfyui_frontend_package/static
Asset seeder disabled
[START] ComfyUI-Manager
[ComfyUI-Manager] Using GitPython backend
[ComfyUI-Manager] network_mode: personal_cloud
[AudioSeparation] Registering 6 node(s) for version 1.1.3.

[Booyaka101]

Thanks for the workflow ID and env details, @van-thieu — that confirms AceStep's audio VAE (music_dcae/music_vocoder.py), which has three candidate module families (custom LayerNorm, ConvNeXtBlock, and weight_norm-wrapped Conv1d inside ResBlock1). Each would route through a different layer of distorch_2.py and need a different fix.

Could you paste the next ~10-15 lines of the traceback after the RuntimeError: Input type (torch.cuda.HalfTensor) and bias type (torch.HalfTensor) should be the same line? Specifically I'm looking for the deepest frame's self class name — that tells us which module class is the one whose bias didn't get moved, so the fix targets the right layer instead of guessing.

[van-thieu]

Is this what you're looking for?

Requested to load AudioOobleckVAE

DisTorch2 Model Virtual VRAM Analysis

===============================================
Object Role Original(GB) Total(GB) Virt(GB)

cuda:0 recip 31.84GB 127.84GB +96.00GB
cpu donor 125.79GB 29.79GB -96.00GB

!!! Exception during processing !!! 'ParametrizationList' object has no attribute 'numel'
Traceback (most recent call last):
File "/home/ai/comfyui/execution.py", line 535, in execute
output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ai/comfyui/execution.py", line 335, in get_output_data
return_values = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ai/comfyui/execution.py", line 309, in _async_map_node_over_list
await process_inputs(input_dict, i)
File "/home/ai/comfyui/execution.py", line 297, in process_inputs
result = f(**inputs)
File "/home/ai/comfyui/comfy_api/internal/init.py", line 149, in wrapped_func
return method(locked_class, **inputs)
File "/home/ai/comfyui/comfy_api/latest/_io.py", line 1826, in EXECUTE_NORMALIZED
to_return = cls.execute(*args, **kwargs)
File "/home/ai/comfyui/comfy_extras/nodes_audio.py", line 192, in execute
return IO.NodeOutput(vae_decode_audio(vae, samples))
~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^
File "/home/ai/comfyui/comfy_extras/nodes_audio.py", line 145, in vae_decode_audio
model_management.load_models_gpu([vae.patcher])
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^
File "/home/ai/comfyui/custom_nodes/comfyui-multigpu/distorch_2.py", line 203, in patched_load_models_gpu
loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights)
~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ai/comfyui/comfy/model_management.py", line 579, in model_load
self.model_use_more_vram(use_more_vram, force_patch_weights=force_patch_weights)
~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ai/comfyui/comfy/model_management.py", line 607, in model_use_more_vram
return self.model.partially_load(self.device, extra_memory, force_patch_weights=force_patch_weights)
~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ai/comfyui/custom_nodes/comfyui-multigpu/distorch_2.py", line 253, in new_partially_load
device_assignments = analyze_safetensor_loading(self, allocations, is_clip=is_clip_model) ## This should be the only required line - that is how it worked previous release so if it doesn't it is Comfy changes
File "/home/ai/comfyui/custom_nodes/comfyui-multigpu/distorch_2.py", line 373, in analyze_safetensor_loading
distorch_alloc = calculate_safetensor_vvram_allocation(model_patcher, virtual_vram_str)
File "/home/ai/comfyui/custom_nodes/comfyui-multigpu/distorch_2.py", line 754, in calculate_safetensor_vvram_allocation
total_memory += module.weight.numel() * module.weight.element_size()
^^^^^^^^^^^^^^^^^^^
File "/home/ai/comfyui/.venv/lib/python3.14/site-packages/torch/nn/modules/module.py", line 1968, in getattr
raise AttributeError(
f"'{type(self).name}' object has no attribute '{name}'"
)
AttributeError: 'ParametrizationList' object has no attribute 'numel'

Prompt executed in 34.17 seconds

[Booyaka101]

That's exactly the trace I needed, thanks @van-thieu. The cause is the third candidate from the earlier triage — weight_norm-wrapped Conv1d inside AceStep's ResBlock1. PyTorch's modern register_parametrization machinery attaches a parametrizations ModuleDict whose .weight slot is a ParametrizationList, not a Tensor — so calling .numel() on it raises the AttributeError you saw.

Fix is up at #198 — guards the size calculation with isinstance(weight, torch.Tensor) so the synthetic ModuleDict is skipped while the real parametrized weight is still counted via the parent module's computed Tensor. Empirically reproduced and verified locally with a mini AceStep-style model. Ready for @pollockjj's review when convenient.