Optimasi Penggunaan Memori

"""Toolkit optimasi memori: pemantauan nyata dan dua konfigurasi pelatihan."""

import gc
import tracemalloc
import warnings
from contextlib import nullcontext

import torch
import torch.utils.checkpoint as cp

warnings.filterwarnings("ignore")
torch.manual_seed(3)


def memori_tensor_hidup():
    """Total byte seluruh tensor torch yang masih direferensikan Python."""
    total = 0
    for obj in gc.get_objects():
        if isinstance(obj, torch.Tensor):
            total += obj.element_size() * obj.numel()
    return total


# ---------- 1. Pemantauan: tracemalloc vs tensor torch ----------
print("=== Pemantauan memori: tracemalloc (alat Python) ===")
tracemalloc.start()
daftar = [0.0] * 1_000_000  # list Python ~8 MB
_, puncak_list = tracemalloc.get_traced_memory()
t = torch.zeros(2_000_000)  # tensor 8 MB di luar alokator Python
_, puncak_tensor = tracemalloc.get_traced_memory()
print(f"list 1 juta float : tracemalloc melihat {puncak_list:,} byte")
print(f"tensor 8 MB       : tracemalloc hanya melihat {puncak_tensor - puncak_list:,} byte")
print(f"tensor sesungguhnya: {t.numel() * t.element_size():,} byte (tak terlacak)")
tracemalloc.stop()
del daftar, t
gc.collect()  # lepas referensi sebelum mengukur pelatihan
print("padanan CPU untuk tensor: element_size() * numel() lewat pemindaian gc")


# ---------- 2. Dua konfigurasi pelatihan diukur nyata ----------
class Blok(torch.nn.Module):
    def __init__(self, d):
        super().__init__()
        self.fc1 = torch.nn.Linear(d, d)
        self.fc2 = torch.nn.Linear(d, d)

    def forward(self, x):
        return x + self.fc2(torch.nn.functional.gelu(self.fc1(x)))


def bangun(n_blok=10, d=256):
    torch.manual_seed(3)
    return torch.nn.ModuleList(Blok(d) for _ in range(n_blok))


N = 128
X = torch.randn(N, 256)
T = torch.randn(N, 256)


def latih_dan_ukur(autocast, checkpoint, micro, n_micro, n_epoch=2):
    blok = bangun()
    opt = torch.optim.SGD(blok.parameters(), lr=0.01)
    dasar = memori_tensor_hidup()
    puncak = dasar
    for _ in range(n_epoch):
        for awal in range(0, N, micro * n_micro):  # batch efektif 32
            opt.zero_grad()
            for langkah in range(n_micro):
                idx = slice(awal + langkah * micro, awal + (langkah + 1) * micro)
                konteks = (
                    torch.autocast("cpu", dtype=torch.bfloat16)
                    if autocast
                    else nullcontext()
                )
                with konteks:
                    h = X[idx]
                    for b in blok:
                        h = (
                            cp.checkpoint(b, h, use_reentrant=False)
                            if checkpoint
                            else b(h)
                        )
                puncak = max(puncak, memori_tensor_hidup())
                loss = ((h.float() - T[idx]) ** 2).mean() / n_micro
                loss.backward()
            opt.step()
    with torch.no_grad():
        h = X
        for b in blok:
            h = b(h)
        mse = float(((h - T) ** 2).mean())
    return dasar, puncak, mse


print()
print("=== Dua konfigurasi (batch efektif 32, 10 blok, 2 epoch) ===")
d1, p1, mse1 = latih_dan_ukur(autocast=False, checkpoint=False, micro=32, n_micro=1)
print(f"baseline fp32 batch 32      : dasar {d1:,} | puncak {p1:,} "
      f"| aktivasi {p1 - d1:,} byte | MSE {mse1:.4f}")

d2, p2, mse2 = latih_dan_ukur(autocast=True, checkpoint=True, micro=8, n_micro=4)
print(f"hemat bf16+ckpt+akum 4x8    : dasar {d2:,} | puncak {p2:,} "
      f"| aktivasi {p2 - d2:,} byte | MSE {mse2:.4f}")

print()
print(f"aktivasi tertahan turun     : {(p1 - d1) / (p2 - d2):.1f}x lebih kecil")
print(f"MSE kedua konfigurasi       : {mse1:.4f} vs {mse2:.4f} (kualitas terjaga)")
print(f"gradien akumulasi disimpan di sini; state SGD hanya 8P tanpa momentum")

=== Pemantauan memori: tracemalloc (alat Python) ===
list 1 juta float : tracemalloc melihat 8,000,064 byte
tensor 8 MB       : tracemalloc hanya melihat 128 byte
tensor sesungguhnya: 8,000,000 byte (tak terlacak)
padanan CPU untuk tensor: element_size() * numel() lewat pemindaian gc

=== Dua konfigurasi (batch efektif 32, 10 blok, 2 epoch) ===
baseline fp32 batch 32      : dasar 5,525,504 | puncak 6,541,316 | aktivasi 1,015,812 byte | MSE 2.4730
hemat bf16+ckpt+akum 4x8    : dasar 5,525,504 | puncak 5,666,180 | aktivasi 140,676 byte | MSE 2.4730

aktivasi tertahan turun     : 7.2x lebih kecil
MSE kedua konfigurasi       : 2.4730 vs 2.4730 (kualitas terjaga)
gradien akumulasi disimpan di sini; state SGD hanya 8P tanpa momentum