Terjemahan disediakan oleh mesin penerjemah. Jika konten terjemahan yang diberikan bertentangan dengan versi bahasa Inggris aslinya, utamakan versi bahasa Inggris.

Pengecekan Aktivasi

Checkpoint aktivasi (atau pemeriksaan gradien) adalah teknik untuk mengurangi penggunaan memori dengan membersihkan aktivasi lapisan tertentu dan mengkomputasi ulang selama lintasan mundur. Secara efektif, ini memperdagangkan waktu komputasi ekstra untuk mengurangi penggunaan memori. Jika modul diperiksa, pada akhir pass maju, input dan output dari modul tetap dalam memori. Setiap tensor perantara yang akan menjadi bagian dari komputasi di dalam modul itu dibebaskan selama pass maju. Selama kelulusan mundur modul checkpoint, tensor ini dihitung ulang. Pada titik ini, lapisan di luar modul checkpoint ini telah menyelesaikan lompatan mundurnya, sehingga penggunaan memori puncak dengan checkpointing bisa lebih rendah.

Cara Menggunakan Checkpointing Aktivasi

Dengansmdistributed.modelparallel, Anda dapat menggunakan checkpointing aktivasi pada granularitas modul. Untuk semua torch.nn modul kecualitorch.nn.Sequential, Anda hanya dapat memeriksa pohon modul jika terletak di dalam satu partisi dari perspektif paralelisme pipa. Dalam kasus torch.nn.Sequential modul, setiap pohon modul di dalam modul berurutan harus terletak sepenuhnya dalam satu partisi agar checkpointing aktivasi berfungsi. Saat Anda menggunakan partisi manual, perhatikan batasan ini.

Saat Anda menggunakan partisi model otomatis, Anda dapat menemukan log penugasan partisi dimulai dengan Partition assignments: di log pekerjaan pelatihan. Jika modul dipartisi di beberapa peringkat (misalnya, dengan satu keturunan pada satu peringkat dan keturunan lain pada peringkat yang berbeda), pustaka mengabaikan upaya untuk memeriksa modul dan memunculkan pesan peringatan bahwa modul tidak akan diperiksa.

Contoh 1: Kode contoh berikut menunjukkan cara menggunakan checkpointing aktivasi saat Anda memiliki definisi model dalam skrip Anda.

import torch.nn as nn
import torch.nn.functional as F

from smdistributed.modelparallel.torch.patches.checkpoint import checkpoint

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = F.max_pool2d(x, 2)
        x = torch.flatten(x, 1)
        # This call of fc1 will be checkpointed
        x = checkpoint(self.fc1, x)
        x = self.fc2(x)
        return F.log_softmax(x, 1)

Contoh 2: Kode contoh berikut menunjukkan cara menggunakan checkpointing aktivasi saat Anda memiliki model sekuensial dalam skrip Anda.

import torch.nn as nn
from smdistributed.modelparallel.torch.patches.checkpoint import checkpoint_sequential

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.seq = nn.Sequential(
            nn.Conv2d(1,20,5),
            nn.ReLU(),
            nn.Conv2d(20,64,5),
            nn.ReLU()
        )

    def forward(self, x):
        # This call of self.seq will be checkpointed
        x = checkpoint_sequential(self.seq, x)
        return F.log_softmax(x, 1)

Contoh 3: Kode contoh berikut menunjukkan cara menggunakan checkpointing aktivasi saat Anda mengimpor model bawaan dari pustaka, seperti PyTorch dan Hugging Face Transformers. Apakah Anda memeriksa modul berurutan atau tidak, lakukan hal berikut:

import smdistributed.modelparallel.torch as smp
from transformers import AutoModelForCausalLM

smp.init()
model = AutoModelForCausalLM(*args, **kwargs)
model = smp.DistributedModel(model)

# Call set_activation_checkpointing API
transformer_layers = model.module.module.module.transformer.seq_layers
smp.set_activation_checkpointing(
    transformer_layers, pack_args_as_tuple=True, strategy='each')