← लैब पर वापस

Mac पर PyTorch: .to("cpu") के बाद मेमोरी खाली क्यों दिखती है

pytorchapple-siliconmemoryguide

Mac पर, mps से CPU पर ले जाया गया टेंसर मेमोरी से गायब होता दिखता है। प्रोसेस का footprint, यानी Activity Monitor के Memory कॉलम वाला नंबर, लगभग शून्य हो जाता है, और MPS ड्राइवर की गिनती भी। फिर भी कॉपी वहीं है और उतनी ही मेमोरी लेती है। यह वह छोटा उदाहरण है जो मैंने ऑफ़लोडिंग पर diffusers की चर्चा के लिए बनाया था, जिसमें 4 GiB का टेंसर मॉडल के वेट्स की जगह है (torch 2.14.0, macOS 26.5.2, M5 Pro जिसमें 48 GB है):

Mac पर CPU और GPU एक ही मेमोरी इस्तेमाल करते हैं: .to("cpu") के बाद वेट्स उसी पूल में रहते हैं, जबकि Activity Monitor 0.1 GB दिखाता है

import ctypes, os, time
import torch
 
# rusage_info_v2 from <libproc.h>; phys_footprint is the Memory column of Activity Monitor
FIELDS = ("user_time", "system_time", "pkg_idle_wkups", "interrupt_wkups", "pageins",
          "wired_size", "resident_size", "phys_footprint")
class RUsageInfoV2(ctypes.Structure):
    _fields_ = [("uuid", ctypes.c_uint8 * 16)] + [(f, ctypes.c_uint64) for f in FIELDS] \
        + [(f"rest{i}", ctypes.c_uint64) for i in range(10)]
 
libproc = ctypes.CDLL("/usr/lib/libproc.dylib")
 
def show(step):
    torch.mps.synchronize()
    info = RUsageInfoV2()
    libproc.proc_pid_rusage(os.getpid(), 2, ctypes.byref(info))
    print(f"{step:<14} footprint {info.phys_footprint / 2**30:4.1f} GiB, "
          f"mps driver {torch.mps.driver_allocated_memory() / 2**30:4.1f} GiB")
 
show("start")
w = torch.randn(2**30, device="mps")      # 4 GiB of float32, standing in for a model's weights
show("on mps")
w = w.to("cpu")                           # what model CPU offload does when a model's turn ends
torch.mps.empty_cache(); time.sleep(2)
show("moved to cpu")
float(w.sum())                            # read every page of the cpu copy
show("cpu copy read")
start          footprint  0.1 GiB, mps driver  0.0 GiB
on mps         footprint  4.2 GiB, mps driver  4.0 GiB
moved to cpu   footprint  0.1 GiB, mps driver  0.0 GiB
cpu copy read  footprint  4.2 GiB, mps driver  0.0 GiB

मैंने इसे तीन बार चलाया और हर बार वही आउटपुट मिला।

नंबर क्या बताते हैं

ले जाने के ठीक बाद दोनों काउंटर लगभग शून्य पर आ जाते हैं, लेकिन w में पूरे 4 GiB मान हैं। जैसे ही CPU कॉपी का हर पेज पढ़ता है, footprint फिर 4.2 GiB हो जाता है। यानी .to("cpu") और CPU के अगली बार उन वेट्स को छूने के बीच लिया गया footprint उन्हें नहीं दिखाता।

ले जाने से कुछ भी खाली न होने की वजह हार्डवेयर है। Apple Silicon पर CPU और GPU भौतिक मेमोरी का एक ही पूल इस्तेमाल करते हैं। PC पर वेट्स को GPU से CPU पर ले जाने से अगले मॉडल के लिए VRAM खाली होती है। Mac पर "CPU पर" रखे वेट्स उसी RAM में रहते हैं जो GPU को चाहिए।

model CPU offload के लिए यह क्यों मायने रखता है

diffusers का enable_model_cpu_offload() हर मॉडल को CPU पर रखता है और उसे सिर्फ़ चलते वक़्त GPU पर ले जाता है। मैंने उसी Mac पर इसी तरह Qwen-Image-2.1 को 1024 × 1024 पर चलाया (ट्रेस और विवरण)। प्रोसेस का footprint कभी 18.5 GB से ऊपर नहीं गया, जो 48 GB वाले Mac में आराम से आ जाना चाहिए। डीनॉइज़िंग के दौरान स्वैप 3.5 GB बढ़ा और VAE डिकोड तक 9.6 GB बढ़ चुका था, और मेरे मेमोरी गार्ड ने इमेज लिखने से पहले रन रोक दिया।

सिर्फ़ footprint का चार्ट इस रन को फ़िट मान लेता। ऊपर का उदाहरण एक तरीका दिखाता है जिससे वेट्स मेमोरी में रहते हुए भी footprint कम बना रहता है। offload वाले रन में कौन-से पेज स्वैप में गए, यह मैंने नहीं मापा।

अपना रन कैसे जाँचें

सिर्फ़ footprint नहीं, स्वैप देखें:

sysctl vm.swapusage             # compare "used" before and after the run
sysctl kern.memorystatus_level  # system-wide free memory, in percent

अगर जॉब चलते वक़्त स्वैप बढ़ता है, तो वह फ़िट नहीं होता, footprint चाहे जो कहे। stageload (pip install stageload), इसके लिए मेरी लिखी छोटी लाइब्रेरी, MemoryMeter से दोनों को हर 0.5 s पर ट्रेस में दर्ज करती है, और stageload guard स्वैप के बजट से ज़्यादा बढ़ते ही कमांड रोक देता है:

stageload guard --wait-free 40 --swap-budget 8G -- python generate.py

मेमोरी असल में क्या खाली करता है

जो पाइपलाइन मॉडल एक के बाद एक चलाती है, उसके लिए हल यह है कि मॉडल को CPU पर रखने के बजाय उसकी स्टेज शुरू होने पर लोड करें और स्टेज ख़त्म होने पर छोड़ दें। stageload का release(module) हर पैरामीटर और बफ़र को PyTorch के meta डिवाइस पर ले जाता है, गार्बेज इकट्ठा करता है और MPS कैश खाली करता है; meta पर टेंसर का कोई स्टोरेज नहीं होता, इसलिए मेमोरी तब भी लौट आती है जब कोई और कोड मॉड्यूल का रेफ़रेंस पकड़े हो। टेक्स्ट एनकोडर (16.3 GB) और ट्रांसफ़ॉर्मर (13.3 GB) को इसी तरह एक-एक करके लोड करने पर उसी Mac पर Qwen-Image-2.1 का पीक 19.0 GB रहा, और दोनों में से किसी भी रन में स्वैप नहीं बढ़ा। पूरा सेटअप 48 GB वाले Mac पर बिना स्वैप के Qwen-Image-2.1 कैसे चलाएँ में है।

मैंने क्या नहीं मापा