New Memory Managment is a Disaster for me. OOM on my Lora Trainer WF where comfy 0.9.x still works

### Custom Node Testing

- [ ] I have tried disabling custom nodes and the issue persists (see [how to disable custom nodes](https://docs.comfy.org/troubleshooting/custom-node-issues#step-1%3A-test-with-all-custom-nodes-disabled) if you need help)

### Your question

ComfyUi is runing in a WSL on my Windows System. I am running my own custom node, a Lora Training workflow with my own block swapping logic. This was for Flux.2 and it was a tight task getting this to work with a 5090 but it ran succsessfully on the older version of comfy. I updated a few days ago to 0.18.1 and since then the whole thing collapses at the beginning og the the blockswapping phase where comfy allocates pinned memory and before the actual training logic starts.

Actually there is enough RAM left on the WSL side [DEBUG MEM] step=0:before_backward | rss=33.66GB vms=213.66GB ram_used=34.54GB/75.88GB | cuda_alloc=19.35GB cuda_reserved=26.11GB cuda_max_alloc=22.26GB cuda_max_reserved=26.11GB. On the Windows side it looks tighter but still ~5GB of RAM left. God knows why the pressure on Windows is more aggressive than on WSL.

I tried disableing the new memory stuf but that didn't help either.
--disable-async-offload
--disable-pinned-memory

I tried to debug the difference between old comfy and new comfy but I coud not pin it down. I made a script that simulated the workflow and bootstrapped comfy.memory_management and that went trough. So it must be somewhere in the depths of comfyui. I ran the complete training run on v0.9.x over 21 hours with no issue. On 0.18 it crashes everytime at the same position. The exception is a plain PyTorch device OOM raised inside loss.backward(). But enough VRAM and RAM is actually availbale imho.

### Logs

```powershell
[2026-04-17 22:44:34.540] INFO [DEBUG MEM] step=0:before_backward | rss=33.66GB vms=213.66GB flux2_train_comfy.py:98
ram_used=34.54GB/75.88GB | cuda_alloc=19.35GB cuda_reserved=26.11GB
cuda_max_alloc=22.26GB cuda_max_reserved=26.11GB
[2026-04-17 22:44:34.562] INFO DEBUG LAST PHASE: step=0:before_backward flux2_train_comfy.py:1492
[2026-04-17 22:44:34.564] INFO [DEBUG MEM] training_loop_exception phase=step=0:before_backward | flux2_train_comfy.py:98
rss=33.66GB vms=213.73GB ram_used=34.54GB/75.88GB | cuda_alloc=19.35GB
cuda_reserved=26.11GB cuda_max_alloc=22.26GB cuda_max_reserved=26.11GB
[2026-04-17 22:44:34.566] INFO CRITICAL ERROR in training_loop: Allocation on device flux2_train_comfy.py:1497
[2026-04-17 22:44:34.567] INFO Initiating emergency cleanup... flux2_train_comfy.py:1500
[2026-04-17 22:44:35.910] [04/17/26 22:44:35] INFO Offloader cleaned up. flux2_train_comfy.py:1506
[2026-04-17 22:44:35.910] INFO Moving models to CPU... flux2_train_comfy.py:1512
[2026-04-17 22:44:42.502] [04/17/26 22:44:42] INFO GPU memory cache cleared. flux2_train_comfy.py:1528
[2026-04-17 22:44:42.505] Traceback (most recent call last):
[2026-04-17 22:44:42.505] File "/opt/ComfyUI/custom_nodes/Flux2-Trainer/flux2_train_comfy.py", line 1369, in training_loop
self.accelerator.backward(loss)
~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^
[2026-04-17 22:44:42.506] File "/home/comfy/.conda/envs/cenv13/lib/python3.13/site-packages/accelerate/accelerator.py", line 2740, in backward
loss.backward(**kwargs)
~~~~~~~~~~~~~^^^^^^^^^^
[2026-04-17 22:44:42.506] File "/home/comfy/.conda/envs/cenv13/lib/python3.13/site-packages/torch/_tensor.py", line 625, in backward
torch.autograd.backward(
~~~~~~~~~~~~~~~~~~~~~~~^
self, gradient, retain_graph, create_graph, inputs=inputs
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
)
^
[2026-04-17 22:44:42.506] File "/home/comfy/.conda/envs/cenv13/lib/python3.13/site-packages/torch/autograd/__init__.py", line 354, in backward
_engine_run_backward(
~~~~~~~~~~~~~~~~~~~~^
tensors,
^^^^^^^^
...<5 lines>...
accumulate_grad=True,
^^^^^^^^^^^^^^^^^^^^^
)
^
[2026-04-17 22:44:42.506] File "/home/comfy/.conda/envs/cenv13/lib/python3.13/site-packages/torch/autograd/graph.py", line 841, in _engine_run_backward
return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
t_outputs, *args, **kwargs
^^^^^^^^^^^^^^^^^^^^^^^^^^
) # Calls into the C++ engine to run the backward pass
^
[2026-04-17 22:44:42.506] torch.OutOfMemoryError: Allocation on device
[2026-04-17 22:44:42.506] [2026-04-17 22:44:42.508] INFO [WORKFLOW DEBUG] train_loop:training_loop_exception break_at=1000 | nodes_flux2.py:60
rss=50.12GB | vms=230.31GB | ram_used=50.70GB/75.88GB | ram_avail=25.19GB |
cuda_alloc=2.60GB | cuda_reserved=4.26GB | cuda_max_alloc=22.26GB |
cuda_max_reserved=26.11GB | global_step=0 | epoch=1 | has_comfy_pbar=True
[2026-04-17 22:44:42.509] ERROR !!! Exception during processing !!! Allocation on device execution.py:621
[2026-04-17 22:44:42.515] ERROR Traceback (most recent call last): execution.py:622
File "/opt/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 "/opt/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 "/opt/ComfyUI/execution.py", line 308, in _async_map_node_over_list
await process_inputs(input_dict, i)
File "/opt/ComfyUI/execution.py", line 296, in process_inputs
result = f(**inputs)
File "/opt/ComfyUI/custom_nodes/Flux2-Trainer/nodes_flux2.py", line 719, in
train
steps_done = training_loop(
break_at_steps=break_at,
epoch=network_trainer_obj.current_epoch.value,
)
File "/opt/ComfyUI/custom_nodes/Flux2-Trainer/flux2_train_comfy.py", line
1531, in training_loop
raise e
File "/opt/ComfyUI/custom_nodes/Flux2-Trainer/flux2_train_comfy.py", line
1369, in training_loop
self.accelerator.backward(loss)
~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^
File
"/home/comfy/.conda/envs/cenv13/lib/python3.13/site-packages/accelerate/accel
erator.py", line 2740, in backward
loss.backward(**kwargs)
~~~~~~~~~~~~~^^^^^^^^^^
File
"/home/comfy/.conda/envs/cenv13/lib/python3.13/site-packages/torch/_tensor.py
", line 625, in backward
torch.autograd.backward(
~~~~~~~~~~~~~~~~~~~~~~~^
self, gradient, retain_graph, create_graph, inputs=inputs
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
)
^
File
"/home/comfy/.conda/envs/cenv13/lib/python3.13/site-packages/torch/autograd/_
_init__.py", line 354, in backward
_engine_run_backward(
~~~~~~~~~~~~~~~~~~~~^
tensors,
^^^^^^^^
...<5 lines>...
accumulate_grad=True,
^^^^^^^^^^^^^^^^^^^^^
)
^
File
"/home/comfy/.conda/envs/cenv13/lib/python3.13/site-packages/torch/autograd/g
raph.py", line 841, in _engine_run_backward
return Variable._execution_engine.run_backward( # Calls into the C++
engine to run the backward pass
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
t_outputs, *args, **kwargs
^^^^^^^^^^^^^^^^^^^^^^^^^^
) # Calls into the C++ engine to run the backward pass
^
torch.OutOfMemoryError: Allocation on device

[2026-04-17 22:44:42.518] INFO Memory summary: execution.py:627
|===========================================================================|
| PyTorch CUDA memory summary, device ID 0 |
|---------------------------------------------------------------------------|
| CUDA OOMs: 0 | cudaMalloc retries: 0 |
|===========================================================================|
| Metric | Cur Usage | Peak Usage | Tot Alloc | Tot Freed |
|---------------------------------------------------------------------------|
| Allocated memory | 2484 MiB | 21227 MiB | 0 B | 0 B |
| from large pool | 0 MiB | 0 MiB | 0 B | 0 B |
| from small pool | 0 MiB | 0 MiB | 0 B | 0 B |
|---------------------------------------------------------------------------|
| Active memory | 2484 MiB | 21227 MiB | 0 B | 0 B |
| from large pool | 0 MiB | 0 MiB | 0 B | 0 B |
| from small pool | 0 MiB | 0 MiB | 0 B | 0 B |
|---------------------------------------------------------------------------|
| Requested memory | 0 B | 0 B | 0 B | 0 B |
| from large pool | 0 B | 0 B | 0 B | 0 B |
| from small pool | 0 B | 0 B | 0 B | 0 B |
|---------------------------------------------------------------------------|
| GPU reserved memory | 4064 MiB | 24896 MiB | 0 B | 0 B |
| from large pool | 0 MiB | 0 MiB | 0 B | 0 B |
| from small pool | 0 MiB | 0 MiB | 0 B | 0 B |
|---------------------------------------------------------------------------|
| Non-releasable memory | 0 B | 0 B | 0 B | 0 B |
| from large pool | 0 B | 0 B | 0 B | 0 B |
| from small pool | 0 B | 0 B | 0 B | 0 B |
|---------------------------------------------------------------------------|
| Allocations | 0 | 0 | 0 | 0 |
| from large pool | 0 | 0 | 0 | 0 |
| from small pool | 0 | 0 | 0 | 0 |
|---------------------------------------------------------------------------|
| Active allocs | 0 | 0 | 0 | 0 |
| from large pool | 0 | 0 | 0 | 0 |
| from small pool | 0 | 0 | 0 | 0 |
|---------------------------------------------------------------------------|
| GPU reserved segments | 0 | 0 | 0 | 0 |
| from large pool | 0 | 0 | 0 | 0 |
| from small pool | 0 | 0 | 0 | 0 |
|---------------------------------------------------------------------------|
| Non-releasable allocs | 0 | 0 | 0 | 0 |
| from large pool | 0 | 0 | 0 | 0 |
| from small pool | 0 | 0 | 0 | 0 |
|---------------------------------------------------------------------------|
| Oversize allocations | 0 | 0 | 0 | 0 |
|---------------------------------------------------------------------------|
| Oversize GPU segments | 0 | 0 | 0 | 0 |
|===========================================================================|

[2026-04-17 22:44:42.519] ERROR Got an OOM, unloading all loaded models. execution.py:628
```

### Other

_No response_