mmap 行为优化 - #17
Open
sfiisf wants to merge 39 commits into
Open
Conversation
Collaborator
Author
|
torch2.7 后新增了 用形如下的样例进行测试, 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)可得到结果, 发现提升还是较大的 |
Collaborator
Author
|
当前 gds 实现: 00c0028 可以看到耗时有比较明显的增加，但在降低内存峰值上同样也有较明显的提高 进一步降低峰值可能考虑使用 save_storage 与 load_storage 配套使用完全代替 offload 2 cpu, 即 disk_offload 完全绕开 CPU, 但需要改动的内容可能较大, 进一步测试看需求再做考虑. |
doombeaker
reviewed
Apr 24, 2026
|
|
||
| # Load with mmap - this doesn't load all data into RAM | ||
| mmap_tensor = torch.load(temp_file, map_location='cpu', mmap=True, weights_only=False) | ||
| file = torch.cuda.gds.GdsFile(temp_file, os.O_CREAT | os.O_RDWR) |
Collaborator
There was a problem hiding this comment.
这里可能需要做一个前置依赖检查，看是否 cuda 可用。
因为我们未来这个要合并到主分支的话，就会影响 musa，npu 等其它芯片。
另外，我觉得目前先不用做硬件抽象层（如果你看到这条有考虑的话），那种抽象，等真的适配了几家芯片后再做比较好。
| del t | ||
| gc.collect() | ||
|
|
||
| fo = open(temp_file, "rb") |
Collaborator
There was a problem hiding this comment.
这里应该会有资源（文件描述符）泄漏吧，需要被动等 gc 去回收。
函数返回后：
├── fd1 (来自 fo=open) → 泄漏！无人关闭
├── fd2 (mmap 内部 dup) → mmap 对象存活期间有效
│ └── mmap 对象被 Tensor 间接持有
│ └── Tensor GC → mmap GC → munmap + close(fd2)
│
└── 文件 temp_file → _cleanup 在 Tensor GC 时删除
└── 但 fd1 还开着！Windows 上可能无法删除文件
问 AI 给了个建议等版本：
def to_mmap(t: torch.Tensor, filename: Optional[str] = None) -> torch.Tensor:
# ... 创建文件、save_storage 等 ...
with open(temp_file, "rb") as fo:
mm = mmap.mmap(fo.fileno(), length=num, access=mmap.ACCESS_READ)
mmap_tensor = torch.frombuffer(mm, dtype=t_type).reshape(t_shape).cpu()
# ✅ fo 关闭，mm 存活（内部 dup 的 fd）
# 需要确保 mm 在 Tensor 销毁时被清理
# 包装 cleanup，同时处理 mm 和文件
def _cleanup():
try:
# 先确保 Tensor 不再引用 mm 内存
# 实际上需要 mm 被 GC，这里只能尽力
if os.path.exists(temp_file):
os.remove(temp_file)
except Exception:
pass
# 更好的做法：用弱引用跟踪 mm
weakref.finalize(mmap_tensor, _cleanup)
# 额外：保存 mm 引用，防止过早 GC（如果需要）
# 但通常 Tensor 内部持有足够信息
return mmap_tensor
Collaborator
sfiisf
marked this pull request as ready for review
June 16, 2026 09:31
strint
approved these changes
Jun 16, 2026
sfiisf
force-pushed
the
feat/mmap-optimization
branch
2 times, most recently
from
July 24, 2026 07:29
8a3cd7d to
7878e48
Compare
Collaborator
|
5090 回归测试结果： 云函数：fnki5mlg84 测试版本v31 CI：https://github.com/siliconflow/cce/actions/runs/30423523771/job/90487820035
测试正常 |
… extension (Comfy-Org#15579) Signed-off-by: Alexander Piskun <bigcat88@icloud.com>
Comfy-Org#15369) * Stop adding an opaque alpha channel to API node images bytesio_to_image_tensor converted every downloaded image to RGBA, so nodes whose API returns no transparency still emitted a 4 channel IMAGE. Keep the alpha when the decoded image has one, stay RGB when it does not. --------- Signed-off-by: bigcat88 <bigcat88@icloud.com> Co-authored-by: bigcat88 <bigcat88@icloud.com>
Signed-off-by: Alexander Piskun <bigcat88@icloud.com>
…logic (Comfy-Org#15655) Signed-off-by: Alexander Piskun <bigcat88@icloud.com>
Signed-off-by: Alexander Piskun <bigcat88@icloud.com>
…Comfy-Org#15684) Signed-off-by: Alexander Piskun <bigcat88@icloud.com>
…#15510) * Keep temp-directory assets visible while their files exist Assets written to the temp directory were flagged as missing and dropped from GET /api/assets, even with the file sitting on disk. One list of directories was answering two different questions -- where the scanner looks for new files, and which files ComfyUI considers its own -- and temp belongs only in the second, so every temp reference was disowned by the prune that runs at startup and on POST /api/assets/prune. Ownership now covers temp. Discovery still does not: the temp directory is wiped before the scan runs, and assets written there are already registered with a hash, mime type and dimensions, so walking it would find nothing. Temp references are instead reconciled against the filesystem directly, so a temp file that really is gone is still retired rather than lingering as a broken entry. get_prefixes_for_root becomes get_scan_prefixes_for_root so the two questions are told apart by name rather than by comment. * Cover the unhashed temp asset in the reconciliation tests The existing temp tests all registered hashed assets, so they never exercised the path an unhashed asset takes when its file is gone: the orphaned rows are removed rather than kept as missing, exactly as under any other root.
sfiisf
force-pushed
the
feat/mmap-optimization
branch
from
August 18, 2026 06:47
596cf8e to
977bf6e
Compare
- 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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
14 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
使用方法与 #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) 读回 cpuMMAP_MEM_THRESHOLD_GB=x, USE_GDS_OFFLOAD=True时表示若 cpu mem 小于 xG 时，遇到 offload 使用 gds 卸载到磁盘再 mmap 读回 cpuUSE_GDS_OFFLOAD默认为 Falsemmap 当前存在一些问题
OOM 时将
memory_to_free置为 1e30, 若模型model_loaded_size > available_memory - mmap_mem_threshold, 则会始终走 partially_unload; 应该是希望 memory_to_free 够大时 partially_unload 也可以将模型完整 offload, 但实际仅能卸载较小一部分内容, 导致显存一直在被占用, 再次请求仍旧发生 OOM 然后重复上述逻辑.https://github.com/siliconflow/cce/issues/176#issuecomment-4249756325
目前是注意到 to_mmap 过程中有个 .cpu 会存在较明显的内存峰值
https://github.com/siliconflow/ComfyGridRuntime/issues/181#issuecomment-4220391298
API Node PR Checklist
Scope
Pricing & Billing
If Need pricing update:
QA
Comms