Activation Checkpointing

Gradient accumulation memangkas memori dengan memperkecil batch, tetapi ada sumber konsumsi lain yang tetap tumbuh mengikuti kedalaman model: aktivasi yang disimpan autograd untuk menghitung gradien. Makin banyak layer, makin banyak tensor antara yang harus ditahan sejak forward() hingga backward() selesai. Activation checkpointing menyerang komponen ini dengan kesepakatan yang terang: buang aktivasi mayoritas saat forward, simpan hanya tensor di batas segmen, lalu hitung ulang bagian dalam saat gradien dibutuhkan. Memori turun, hitungan naik — dan pada model dalam yang sering kalah memori sebelum kalah waktu, tukar ini sangat menguntungkan. Chapter ini mengukur kedua sisi pertukaran itu secara nyata di CPU dengan torch.utils.checkpoint.

Mekanisme Simpan di Batas, Hitung Ulang di Dalam #

Cara kerjanya bersandar pada struktur graph autograd. Tanpa checkpoint, setiap operasi menyimpan inputnya agar rantai aturan gradien bisa dijalankan mundur. Dengan checkpoint, PyTorch membangun segmen dari beberapa operasi, menyimpan hanya input segmen, membuang semua tensor antara di dalamnya, lalu pada saat backward menjalankan ulang forward segmen tersebut untuk membangkitkan aktivasi yang dibutuhkan. Perhatikan alurnya:

Konsekuensinya presisi: setiap segmen dijalankan dua kali (sekali saat forward asli, sekali saat recompute), sehingga total hitungan maju menjadi dua kali lipat bila seluruh model di-checkpoint per segmen. Berkas activation_checkpoint.py mengukur kedua sisi pertukaran ini dengan penghitung forward per blok dan pemindaian memori tensor sungguhan:

"""Activation checkpointing: tukar hitungan ulang dengan memori aktivasi."""

import gc
import warnings

import torch
import torch.utils.checkpoint as cp

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


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


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)
        self.hitung = 0  # penghitung forward (bukan parameter)

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


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


def jalankan(blok, x, target, pakai_checkpoint):
    awal = memori_tensor_hidup()
    hitung = 0
    h = x
    for b in blok:
        if pakai_checkpoint:
            h = cp.checkpoint(b, h, use_reentrant=False)
        else:
            h = b(h)
    sesudah_fwd = memori_tensor_hidup()
    loss = ((h - target) ** 2).mean()
    loss.backward()
    sesudah_bwd = memori_tensor_hidup()
    hitung = sum(b.hitung for b in blok)
    return awal, sesudah_fwd, sesudah_bwd, hitung


D, B = 256, 32
x = torch.randn(B, D)
target = torch.randn(B, D)

print("=== Tanpa checkpoint (12 blok residual, batch 32) ===")
a, f, b, n = jalankan(bangun(), x, target, pakai_checkpoint=False)
print(f"memori tensor sebelum forward : {a:>10,} byte")
print(f"memori tensor setelah forward : {f:>10,} byte")
print(f"aktivasi tertahan forward     : {f - a:>10,} byte")
print(f"memori tensor setelah backward: {b:>10,} byte")
print(f"total forward per blok        : {n} kali (12 blok x 1)")

print()
print("=== Dengan checkpoint per blok (use_reentrant=False) ===")
a2, f2, b2, n2 = jalankan(bangun(), x, target, pakai_checkpoint=True)
print(f"memori tensor sebelum forward : {a2:>10,} byte")
print(f"memori tensor setelah forward : {f2:>10,} byte")
print(f"aktivasi tertahan forward     : {f2 - a2:>10,} byte")
print(f"memori tensor setelah backward: {b2:>10,} byte")
print(f"total forward per blok        : {n2} kali (12 blok x 2: recompute)")

print()
print(f"hemat aktivasi tertahan       : {(f - a) - (f2 - a2):,} byte "
      f"({(f - a) / (f2 - a2):.1f}x lebih kecil)")
print(f"biaya hitungan                : {n2 / n:.1f}x lebih banyak forward")

Output yang Dihasilkan:

=== Tanpa checkpoint (12 blok residual, batch 32) ===
memori tensor sebelum forward :  6,381,568 byte
memori tensor setelah forward :  7,561,216 byte
aktivasi tertahan forward     :  1,179,648 byte
memori tensor setelah backward:  6,414,340 byte
total forward per blok        : 12 kali (12 blok x 1)

=== Dengan checkpoint per blok (use_reentrant=False) ===
memori tensor sebelum forward :  6,381,568 byte
memori tensor setelah forward :  6,835,456 byte
aktivasi tertahan forward     :    453,888 byte
memori tensor setelah backward:  6,414,340 byte
total forward per blok        : 24 kali (12 blok x 2: recompute)

hemat aktivasi tertahan       : 725,760 byte (2.6x lebih kecil)
biaya hitungan                : 2.0x lebih banyak forward

Mari kita telaah kode di atas baris per baris:

1. Baris 13-19 (memori_tensor_hidup): fungsi memindai objek Python lewat gc.get_objects() dan menjumlahkan element_size() × numel() setiap tensor torch — padanan CPU untuk torch.cuda.memory_allocated() karena tracemalloc tidak melihat memori tensor (alat ini kita bedah di chapter penutup section ini).
2. Baris 22-31 (Blok): blok residual sederhana x + fc2(gelu(fc1(x))) meniru blok transformer; atribut self.hitung menghitung berapa kali forward dipanggil — bukan parameter, jadi tidak terdampak parameters().
3. Baris 39-53 (jalankan): fungsi menjalankan 12 blok berurutan; varian checkpoint membungkus tiap blok dengan cp.checkpoint(b, h, use_reentrant=False) — API modern PyTorch 2.x yang merekam konteks autocast dan aman untuk kontrol alur. Memori dipindai sebelum forward, sesudah forward, dan sesudah backward.
4. Baris 60-66 (varian tanpa checkpoint): aktivasi yang tertahan setelah forward mencapai 1.179.648 byte — seluruh tensor antara 12 blok ditahan untuk backward.
5. Baris 69-75 (varian checkpoint): aktivasi tertahan hanya 453.888 byte — 2,6x lebih kecil — karena hanya input tiap blok yang disimpan; total forward per blok menjadi 24 (12 blok × 2) sebagai bukti recompute benar-benar terjadi.
6. Baris 78-80 (ringkasan pertukaran): hemat 725.760 byte aktivasi dengan biaya 2,0x forward. Pada model 12 blok ini penghematan masih moderat; pada model ratusan blok, rasio aktivasi terhadap parameter jauh lebih besar sehingga penghematannya jauh lebih dramatis.

Kapan Pertukaran Ini Menguntungkan #

Biaya waktunya nyata dan kami mengukurnya langsung: di mesin uji CPU dua inti ini, iterasi forward-backward model 12 blok memakan sekitar 5,8 ms tanpa checkpoint dan sekitar 92,9 ms dengan checkpoint per blok pada konfigurasi yang sama — overhead checkpoint terasa besar pada model kecil CPU karena biaya mekanismenya sebanding dengan hitungan bloknya. Pada model GPU besar yang benar-benar kehabisan memori, overhead lazim jauh lebih kecil (kisaran 20-30%) dan sering menjadi satu-satunya cara pelatihan tetap berjalan; angka pastinya selalu harus diukur ulang di perangkat Anda.

Checkpointing dan accumulation sama-sama menekan aktivasi dari arah berbeda. Pertanyaan berikutnya memindah fokus dari satu perangkat ke banyak: bagaimana bila kita memiliki beberapa GPU dan ingin membagi data — bukan model — ke semuanya? Itulah pokok bahasan chapter Data Parallelism.

Rujukan resmi API torch.utils.checkpoint tersedia di docs.pytorch.org/docs/stable/checkpoint.html .