mmap 行为优化 - #17
Open
sfiisf wants to merge 92 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
Collaborator
sfiisf
force-pushed
the
feat/mmap-optimization
branch
from
July 24, 2026 04:27
112233a to
0b9536b
Compare
sfiisf
force-pushed
the
feat/mmap-optimization
branch
from
July 24, 2026 07:29
8a3cd7d to
7878e48
Compare
…5063) These were alll non-dynamic (some non-ModelPatcher) code path calling FreeMemory for management requiring up-front memory freeing. Convert it to dynamic to avoid legacy free behaviour mixing into otherwise dynamic workflows.
Some long running chaos testing on a 512GB RAM RTX6000 pro showed that this is a little bit too low for common template workflows switching around. The original number was just a guess from me, so go with the scientific result instead.
…Comfy-Org#15075) Signed-off-by: bigcat88 <bigcat88@icloud.com>
Changes: Remove sequential scan hint Prefer NVML pressure on windows Add async malloc clamp option (unused by comfy so far) Workaround AMD windows GPU virtual address space leak The largest change is the NVML pressure, which works around a cuMemGetInfo drift from actual VRAM in some circumstances.
…y-Org#15227) Signed-off-by: Alexander Piskun <bigcat88@icloud.com>
This priority scheme was broken in the case where you have pin registration exhaustion while loading a VBAR that gets a big evicition. The weight would stay in the loaded set but inherit the MRU priority against other workflow models WRT pin registration which leads to async offload without pinning. Fix by universally promiting active pin registration above workflow pins without concern for the weights/weights-loaded split. This diverges from the actual budgeting where the split still makes sense.
This avoids name collision (circular imports) for external custom nodes, for which the comfy path is pushed into sys.path so Python's own logging module is shadowed otherwise. fixes: Comfy-Org#15229
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
…l Try-On API (Comfy-Org#15249) Signed-off-by: bigcat88 <bigcat88@icloud.com>
This is performance critical for ACE step.
…ort (Comfy-Org#15294) Signed-off-by: Alexander Piskun <bigcat88@icloud.com>
* [Partner Nodes] feat(BFL): add Flux 3 video model Signed-off-by: Alexander Piskun <bigcat88@icloud.com> * [Partner Nodes] fix(BFL): harden keyframe value check Signed-off-by: Alexander Piskun <bigcat88@icloud.com> --------- Signed-off-by: Alexander Piskun <bigcat88@icloud.com> Co-authored-by: Jedrzej Kosinski <kosinkadink1@gmail.com>
Co-authored-by: Jedrzej Kosinski <kosinkadink1@gmail.com>
Comfy-aimdo 0.4.12 increases error logging reliablity to help root cause os errors in some of the C APIs that are causing issues for some users. The log is also unified with python logging, so non-terminal users see the logs properly. Aimdo 0.4.13 fixes a bug in async-offload + MRU primary weights allocation. Comfy-Org#15284
Collaborator
|
5090 回归测试结果： 云函数：fnki5mlg84 测试版本v31 CI：https://github.com/siliconflow/cce/actions/runs/30423523771/job/90487820035
测试正常 |
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
20 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