VRAM Usage Fluctuation with TeaCache Node in Retention Mode

[shuaixn]

After the Retention mode update in ComfyUI-TeaCache, I observed periodic fluctuations in VRAM usage. My system specs are: Windows 11, RTX 4090 (24GB VRAM), 64GB RAM, ComfyUI v0.3.27. Using the official ComfyUI i2v workflow with the wan2.1_720p_14b_fp16 model, I generated videos at 720×1280 or 1280×720 resolution (4 or 5 seconds). With the TeaCache node enabled and set to recommended parameters, VRAM usage overflows by 0.2-0.6GB into shared memory, causing inference slowdown. The overflow disappears after a while, and the VRAM usage fluctuates periodically. When disabling the TeaCache node, VRAM stays stable at around 23GB; when enabled, the periodic overflow occurs. I suspect the issue lies with the TeaCache node. Could the author please take a look when possible? Thank you!

[YunjieYu]

Hi @shuaixn, the slight fluctuation in VRAM usage is expected. I can tell you how to achieve the same VRAM usage by simply modifying the code in nodes.py as below:

in 578 line:
before: ori_x = x.clone()
after: ori_x = x.to(mm.unet_offload_device())

in 590 line:
before: self.teacache_state[k]['previous_residual'] = (x - ori_x)[i*b:(i+1)*b].to(mm.unet_offload_device())
after: self.teacache_state[k]['previous_residual'] = (x.to(mm.unet_offload_device()) - ori_x)[i*b:(i+1)*b]

However, I really don't want to make these changes in my release version, because it will bring about a 5%～10% slowdown, which is unacceptable for many people.

I hope you can understand my concerns.

[YunjieYu]

Another simple and feasible solution is to reduce the generation of 4 frames.

[shuaixn]

Another simple and feasible solution is to reduce the generation of 4 frames.另一个简单可行的解决方案是减少4帧的生成。
Okay, thank you.

[shuaixn]

Another simple and feasible solution is to reduce the generation of 4 frames.

I suddenly remembered that there are a few more details I need to add. When I open ComfyUI and run the entire workflow for the first time, the GPU memory does not overflow. However, during the second operation and subsequent ones, intermittent GPU memory overflow occurs. To avoid this issue, I need to restart ComfyUI, but the fluctuations still occur during the second run. Could this situation provide any useful insights for optimization? Or is it no different from the previous optimization plan?

[YunjieYu]

Hi @shuaixn, thank you for your careful observation and constructive insight. I will study this issue once I have time.

[shuaixn]

Hi @shuaixn, the slight fluctuation in VRAM usage is expected. I can tell you how to achieve the same VRAM usage by simply modifying the code in nodes.py as below:您好，VRAM 使用率略有波动是正常现象。我可以告诉您如何通过简单修改 nodes.py 中的代码来实现相同的 VRAM 使用率，如下所示：

in 578 line:  在第 578 行： before: ori_x = x.clone()之前： ori_x = x.clone() after: ori_x = x.to(mm.unet_offload_device())  之后： ori_x = x.to(mm.unet_offload_device())

in 590 line:  在第 590 行： before: self.teacache_state[k]['previous_residual'] = (x - ori_x)[i*b:(i+1)*b].to(mm.unet_offload_device())  之前： self.teacache_state[k]['previous_residual'] = (x - ori_x)[i*b:(i+1)*b].to(mm.unet_offload_device()) after: self.teacache_state[k]['previous_residual'] = (x.to(mm.unet_offload_device()) - ori_x)[i*b:(i+1)*b]  之后： self.teacache_state[k]['previous_residual'] = (x.to(mm.unet_offload_device()) - ori_x)[i*b:(i+1)*b]

However, I really don't want to make these changes in my release version, because it will bring about a 5%～10% slowdown, which is unacceptable for many people.不过我真的不想在我的发布版本中做这些改变，因为这会带来5%～10%的运行速度减慢，这对于很多人来说是难以接受的。

I hope you can understand my concerns.我希望你能理解我的担忧。

I've modified the code according to your suggestion. I feel like there is almost no change in speed. I roughly calculated that previously, each frame took an average of 0.19 minutes to generate, and after modifying the code, it became 0.2 minutes. Thanks again!

[YunjieYu]

Thanks for your useful feedback. The slowdown ​​may vary depending on GPU, CUDA or experimental settings. I also appreciate your attention for the ComfyUI-TeaCache project.

[shuaixn]

After such a long time, I see that TeaCache has been updated. However, even after updating to the latest version, I still encounter VRAM overflow issues. So, I would like to ask whether the changes to the code are still the same as previously mentioned?

[YunjieYu]

Yes, you can modify the code in nodes.py as below:

in 734 line:
before: ori_x = x.clone()
after: ori_x = x.to(mm.unet_offload_device())

in 746 line:
before: self.teacache_state[k]['previous_residual'] = (x - ori_x)[i*b:(i+1)*b].to(mm.unet_offload_device())
after: self.teacache_state[k]['previous_residual'] = (x.to(mm.unet_offload_device()) - ori_x)[i*b:(i+1)*b]

[YunjieYu]

After such a long time, I see that TeaCache has been updated. However, even after updating to the latest version, I still encounter VRAM overflow issues. So, I would like to ask whether the changes to the code are still the same as previously mentioned?

Hi @shuaixn, I have added a cache_device option in the latest 1.7.0 version. Now one can choose to offload cache on RAM instead of VRAM: