mmap 行为优化 by sfiisf · Pull Request #17 · siliconflow/ComfyUI

sfiisf · 2026-04-23T08:54:22Z

~~使用方法与 #13 保持一致，设置环境变量 MMAP_MEM_THRESHOLD_GB=x 表示若 cpu mem 小于 xG 时，遇到 offload 会 offload 到 mmap~~
考虑到 gds 对性能存在影响，额外增加了环境变量 USE_GDS_OFFLOAD 用于控制行为:

MMAP_MEM_THRESHOLD_GB=x, USE_GDS_OFFLOAD=False 时行为同 offload to mmap #13 一致，表示若 cpu mem 小于 xG 时，遇到 offload 使用 torch.save 卸载到磁盘再 torch.load(mmap=True) 读回 cpu
MMAP_MEM_THRESHOLD_GB=x, USE_GDS_OFFLOAD=True 时表示若 cpu mem 小于 xG 时，遇到 offload 使用 gds 卸载到磁盘再 mmap 读回 cpu

USE_GDS_OFFLOAD 默认为 False

mmap 当前存在一些问题

OOM 发生时 offload 行为异常:
OOM 时将 memory_to_free 置为 1e30, 若模型 model_loaded_size > available_memory - mmap_mem_threshold, 则会始终走 partially_unload; 应该是希望 memory_to_free 够大时 partially_unload 也可以将模型完整 offload, 但实际仅能卸载较小一部分内容, 导致显存一直在被占用, 再次请求仍旧发生 OOM 然后重复上述逻辑.

        if min(memory_to_free, model_loaded_size) > available_memory - mmap_mem_threshold or memory_to_free < model_loaded_size:
            partially_unload = True
        else:
            partially_unload = False

https://github.com/siliconflow/cce/issues/176#issuecomment-4249756325

仍旧会发生暴 CPU 内存的情况
目前是注意到 to_mmap 过程中有个 .cpu 会存在较明显的内存峰值

def to_mmap(t: torch.Tensor, filename: Optional[str] = None) -> torch.Tensor:
...
    cpu_tensor = t.cpu()
    torch.save(cpu_tensor, temp_file)
...
    mmap_tensor = torch.load(temp_file, map_location='cpu', mmap=True, weights_only=False)

https://github.com/siliconflow/ComfyGridRuntime/issues/181#issuecomment-4220391298

API Node PR Checklist

Scope

Is API Node Change

Pricing & Billing

Need pricing update
No pricing update

If Need pricing update:

Metronome rate cards updated
Auto‑billing tests updated and passing

QA

QA done
QA not required

Comms

Informed Kosinkadink

sfiisf · 2026-04-23T09:00:41Z

torch2.7 后新增了 torch.cuda.gds 对 NVIDIA cuFile / GPUDirect Storage (类似于 DMA, GPU<->DISK 不经过 CPU)进行了薄封装 https://docs.pytorch.org/docs/stable/generated/torch.cuda.gds.GdsFile.html
但需注意 save_storage 存储内容为裸的二进制数据, 无法通过 torch.load 等方式直接读取. 读回 GPU 配套使用 load_storage, 但失去了 offload 的语义, 因此尝试使用 mmap + frombuffer 读回 CPU

用形如下的样例进行测试,

SHAPE = (1024, 1024, 1024)
DEVICE = "cuda:2"

# gds
src = torch.randn(*SHAPE, device=DEVICE)
file = torch.cuda.gds.GdsFile("temp1.pt", os.O_CREAT | os.O_RDWR)
file.save_storage(src.untyped_storage(), offset=0)
num_bytes = src.numel() * src.element_size()
src_dtype = src.dtype
src_shape = src.shape

fo = open("temp1.pt", "rb")
mm = mmap.mmap(fo.fileno(), length=num_bytes, access=mmap.ACCESS_READ)
dest = torch.frombuffer(mm, dtype=src_dtype).reshape(src_shape)

# old
src = torch.randn(*SHAPE, device=DEVICE)
src_cpu = src.cpu()
torch.save(src_cpu, "temp2.pt")
dest = torch.load("temp2.pt", map_location="cpu", mmap=True, weights_only=False)

可得到结果, 发现提升还是较大的

gds mmap time cost: 3614.706 ms
gds mmap cpu rss start: 0.498 GB
gds mmap cpu rss end:   0.641 GB
gds mmap cpu rss peak:  0.641 GB
gds mmap cpu rss current delta: 0.143 GB
gds mmap cpu rss peak delta:    0.142 GB
cpu save mmap time cost: 9936.152 ms
cpu save mmap cpu rss start: 0.641 GB
cpu save mmap cpu rss end:   4.642 GB
cpu save mmap cpu rss peak:  4.642 GB

sfiisf · 2026-04-23T10:55:26Z

当前 gds 实现: 00c0028
测试工作流: https://github.com/anveshane/Comfyui_turbodiffusion/blob/main/turbowan_workflow.json
对比 MMAP_MEM_THRESHOLD_GB=0 即不开启 mmap; MMAP_MEM_THRESHOLD_GB=512 即全部通过 mmap 且使用 .cpu() .save; 全部通过 mmap 且使用 gds 的测试结果如下

可以看到耗时有比较明显的增加，但在降低内存峰值上同样也有较明显的提高

进一步降低峰值可能考虑使用 save_storage 与 load_storage 配套使用完全代替 offload 2 cpu, 即 disk_offload 完全绕开 CPU, 但需要改动的内容可能较大, 进一步测试看需求再做考虑.

MMAP_MEM_THRESHOLD_GB=0:
Prompt af7337e5-7a82-4cc8-b48d-d3a13bfbd108 memory stats (109.56s):
  RSS start/end/peak: 1397.63 MiB / 37158.70 MiB / 37158.70 MiB
  USS start/end/peak: 1322.99 MiB / 37093.86 MiB / 37093.86 MiB
Prompt 0e10e1e5-ee24-4360-bd1f-30af953bf119 memory stats (32.95s):
  RSS start/end/peak: 37158.70 MiB / 37527.13 MiB / 37527.13 MiB
  USS start/end/peak: 37093.91 MiB / 37454.35 MiB / 37454.44 MiB
Prompt 0445ec03-6a6a-42d8-8eea-8c7ea07d4f16 memory stats (29.47s):
  RSS start/end/peak: 37527.13 MiB / 37719.08 MiB / 37719.08 MiB
  USS start/end/peak: 37454.45 MiB / 37654.56 MiB / 37654.56 MiB



MMAP_MEM_THRESHOLD_GB=512 all_2cpu_save:
Prompt 71859d6b-ebae-4586-b75b-71bd0b01600f memory stats (699.34s):
  RSS start/end/peak: 1396.25 MiB / 24430.46 MiB / 30129.74 MiB
  USS start/end/peak: 955.58 MiB / 23880.52 MiB / 29583.76 MiB
Prompt a4550d38-d314-4ff1-a241-1adde9bec94f memory stats (693.21s):
  RSS start/end/peak: 24430.52 MiB / 24423.28 MiB / 24639.08 MiB
  USS start/end/peak: 23880.59 MiB / 23880.72 MiB / 24093.73 MiB
Prompt fa46e5f0-bf95-43ee-8137-c478d13da96f memory stats (725.49s):
  RSS start/end/peak: 24423.35 MiB / 22552.54 MiB / 24611.89 MiB
  USS start/end/peak: 23880.73 MiB / 22018.51 MiB / 24068.34 MiB
Prompt 51cf9b12-1820-4042-998d-10c20f6729e3 memory stats (694.32s):
  RSS start/end/peak: 21853.77 MiB / 21121.22 MiB / 21853.77 MiB
  USS start/end/peak: 21319.45 MiB / 20590.23 MiB / 21319.45 MiB
Prompt 263d6b45-8009-4f73-9ac4-d04362bce335 memory stats (711.84s):
  RSS start/end/peak: 21121.36 MiB / 21091.86 MiB / 21335.42 MiB
  USS start/end/peak: 20590.22 MiB / 20559.59 MiB / 20803.21 MiB


MMAP_MEM_THRESHOLD_GB=512 all_gds:
Prompt 7f0618c0-a778-49f3-bb4b-9844724a637c memory stats (739.28s):
  RSS start/end/peak: 1400.79 MiB / 17361.43 MiB / 23111.11 MiB
  USS start/end/peak: 957.23 MiB / 16823.97 MiB / 22572.88 MiB
Prompt 24dfc9db-9bb9-4ad6-9cbd-d89ab6a0909b memory stats (762.83s):
  RSS start/end/peak: 17361.59 MiB / 10944.14 MiB / 17478.92 MiB
  USS start/end/peak: 16823.98 MiB / 10398.46 MiB / 16988.87 MiB
Prompt ede03cc6-874b-4df1-9d22-975f568f7887 memory stats (749.50s):
  RSS start/end/peak: 10944.35 MiB / 10934.63 MiB / 11148.28 MiB
  USS start/end/peak: 10398.20 MiB / 10397.01 MiB / 10565.04 MiB
Prompt 90f6be16-2f29-476a-b214-db6b340e7f8d memory stats (760.65s):
  RSS start/end/peak: 10934.68 MiB / 10927.91 MiB / 11121.48 MiB
  USS start/end/peak: 10397.09 MiB / 10396.98 MiB / 10591.96 MiB
Prompt f256e8ba-a0d8-4c30-9457-708c2f728022 memory stats (764.62s):
  RSS start/end/peak: 10928.00 MiB / 10929.80 MiB / 11117.88 MiB
  USS start/end/peak: 10396.98 MiB / 10397.15 MiB / 10588.71 MiB

ccndcn · 2026-05-07T10:12:26Z

https://github.com/siliconflow/cce/actions/runs/25484792464/job/74780338746

结果正常

ccndcn · 2026-06-24T01:36:27Z

5090:
https://github.com/siliconflow/cce/actions/runs/28069062323
H200:
https://github.com/siliconflow/cce/actions/runs/28069063209

ccndcn · 2026-08-09T10:53:48Z

5090 回归测试结果：

云函数：fnki5mlg84 测试版本v31
comfyui镜像：hub.6scloud.com/d1s6t1r3jrms73f3ltpg/comfyui-gpu-torch2dot9:v202607291124-beta-default-929e3bf

CI：https://github.com/siliconflow/cce/actions/runs/30423523771/job/90487820035

cce-test-results-20260729_151016（CI日志）：https://gist.github.com/ccndcn/d4bf0d9f4aab95521780f7d8f6c63fcf
runs_20260729_151016（CI失败工作流二次测试日志）：https://gist.github.com/ccndcn/8fa7c2bd422549d20061609c81d836e9

测试正常

- keep mmap resources alive for shared tensor storage - use copy-on-write mappings and restrict GDS to CUDA tensors - base mmap pressure on loaded model weights - move custom-node directory dumps to debug logging

sfiisf · 2026-08-19T09:36:23Z

主要改动

新增 GDS offload to mmap，开启后会增加一定耗时，但可降低 CPU 内存峰值；简单测试见 mmap 行为优化 #17 (comment)
修复 OOM 场景下的 offload 策略：
- need_mmap 将本次待卸载数据量纳入判断，提前选择 mmap，避免普通 CPU offload 后突破内存阈值。
- 修正完整卸载与部分卸载的选择逻辑，避免因 CPU 可用内存不足错误进入 partial unload。
- mmap 转换失败时自动回退到普通 offload。
完善 mmap 资源与异常处理：
- 按临时文件所在分区检查磁盘可用空间。
- 完善临时文件描述符、GDS 文件和 mmap 映射的关闭与清理。
- 将清理生命周期绑定到底层 Tensor storage，确保临时文件在不再使用后删除，及防止 Parameter 路径的提前删除。
支持 QuantizedTensor offload to mmap：分别转换量化 Tensor 的内部数据并重建量化上下文，同时在 Parameter 重新封装时保留 _qdata 和 _params 元数据。

ccndcn · 2026-08-24T02:16:25Z

CI测试结果：
https://github.com/siliconflow/cce/actions/runs/32473383218/job/96745600619
二轮测试结果：
latest_36_failed.zip
结果正常

doombeaker · 2026-08-24T11:42:06Z

    for custom_node_path in node_paths:
-        possible_modules = os.listdir(custom_node_path)
+        possible_modules = sorted(os.listdir(custom_node_path))
+        logging.debug("!!!possible_modules1!!!:\n" + "\n".join(possible_modules))


这个是不是临时调试日志，可以删除或者改为正常文案了？

这部分为 beta 分支引入， @ccndcn 这部分还有留着的需求吗

需要留着，日志级别调整info或者print输出

doombeaker · 2026-08-24T12:13:05Z

+        t_type = t.dtype
+        t_shape = t.shape
+        num = t.numel() * t.element_size()
+        del file


这里改成 file.close() 更好吧（顺便注意下和上面那个 file.save_storage(t.untyped_storage(), offset=0) 做下合并。

Suggested change

del file

try:

file.save_storage(t.untyped_storage(), offset=0)

finally:

file.close()

虽然名字叫做 file, 但实际 torch.cuda.gds.GdsFile 为一个无 close() 方法的 class, 对应句柄的关闭有在其析构函数中处理. 因此我将变量 file 改为 gds, 但仍保留显示的 del 操作

class GdsFile: ... def __del__(self) -> None: if self.handle is not None: self.deregister_handle() os.close(self.fd)

doombeaker reviewed Apr 24, 2026

View reviewed changes

Comment thread comfy/model_patcher.py Outdated

Comment thread comfy/model_patcher.py Outdated

sfiisf marked this pull request as ready for review June 16, 2026 09:31

sfiisf requested a review from strint June 16, 2026 09:32

strint approved these changes Jun 16, 2026

View reviewed changes

ccndcn and others added 4 commits July 6, 2026 16:16

fix

d8e1c93

fix

f7d05c5

Merge branch 'master' into beta

817ee0e

Merge branch 'master' of github.com:siliconflow/ComfyUI into beta

ec7b605

sfiisf force-pushed the feat/mmap-optimization branch 2 times, most recently from 8a3cd7d to 7878e48 Compare July 24, 2026 07:29

ccssu and others added 2 commits August 3, 2026 11:37

Merge branch 'master' into beta

9b6b85f

Merge branch 'master' into beta

2ba7128

ccndcn and others added 9 commits August 14, 2026 10:26

Merge branch 'master' into beta

861882e

修改 partially_unload 与 need_mmap 及 model_to_mmap 逻辑

ecb4026

use gds to mmap offload

ca38678

mmap 失败时回退到 normal offload

06539d5

增加前置依赖检查&关闭fo

0e61721

增加 _USE_GDS_OFFLOAD 控制是否启用 GDS

222045f

完善文件符处理

59ca53b

fix

65711d4

fix

977bf6e

sfiisf force-pushed the feat/mmap-optimization branch from 596cf8e to 977bf6e Compare August 18, 2026 06:47

sfiisf and others added 3 commits August 18, 2026 15:12

refine

ab02e26

Fix mmap offload safety and startup logging

97ecec3

- keep mmap resources alive for shared tensor storage - use copy-on-write mappings and restrict GDS to CUDA tensors - base mmap pressure on loaded model weights - move custom-node directory dumps to debug logging

Merge branch 'master' into beta

db5b21b

sfiisf force-pushed the feat/mmap-optimization branch from 66cf2e8 to 97ecec3 Compare August 19, 2026 04:13

sfiisf and others added 4 commits August 19, 2026 12:14

Merge branch 'refine_offload' into feat/mmap-optimization

e1a1610

support QuantizedTensor offload to mmap

0c3e66e

refine get_free_disk

f35becb

refine QuantizedTensor offload to mmap behavior

4a64eea

Merge branch 'beta' into feat/mmap-optimization

e3a6631

doombeaker approved these changes Aug 24, 2026

View reviewed changes

refine gds behavior&adjust node check logging level

a24bc15

sfiisf merged commit 0024445 into refine_offload Aug 25, 2026
7 of 12 checks passed

github-actions Bot locked and limited conversation to collaborators Aug 25, 2026

-        del file
+try:
+    file.save_storage(t.untyped_storage(), offset=0)
+finally:
+    file.close()

Uh oh!

Conversation

sfiisf commented Apr 23, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

API Node PR Checklist

Scope

Pricing & Billing

QA

Comms

Uh oh!

sfiisf commented Apr 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sfiisf commented Apr 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ccndcn commented May 7, 2026

Uh oh!

ccndcn commented Jun 24, 2026

Uh oh!

ccndcn commented Aug 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sfiisf commented Aug 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

主要改动

Uh oh!

ccndcn commented Aug 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

doombeaker Aug 24, 2026

Choose a reason for hiding this comment

Uh oh!

sfiisf Aug 25, 2026

Choose a reason for hiding this comment

Uh oh!

ccndcn Aug 25, 2026

Choose a reason for hiding this comment

Uh oh!

doombeaker Aug 24, 2026

Choose a reason for hiding this comment

Uh oh!

sfiisf Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

sfiisf commented Apr 23, 2026 •

edited by github-actions Bot

Loading

sfiisf commented Apr 23, 2026 •

edited

Loading

sfiisf commented Apr 23, 2026 •

edited

Loading

ccndcn commented Aug 9, 2026 •

edited

Loading

sfiisf commented Aug 19, 2026 •

edited

Loading

ccndcn commented Aug 24, 2026 •

edited

Loading

sfiisf Aug 25, 2026 •

edited

Loading