mmap 行为优化

[sfiisf]

使用方法与 #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 当前存在一些问题

1. 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

2. 仍旧会发生暴 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