mmap 行为优化#17
Draft
sfiisf wants to merge 4 commits into
Draft
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
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
3 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 到 mmapmmap 当前存在一些问题
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