Gradient checkpointing dan local SGD

Pelatihan Model AI Efisien dengan PyTorch

Dennis Lee

Data Engineer, Amazon

Meningkatkan efisiensi pelatihan

Ikon yang merepresentasikan efisiensi memori, komunikasi, dan komputasi.

Gradient checkpointing meningkatkan efisiensi memori

Ikon yang merepresentasikan efisiensi memori, komunikasi, dan komputasi.

Local SGD meningkatkan efisiensi komunikasi

Ikon yang merepresentasikan efisiensi memori, komunikasi, dan komputasi.

Apa itu gradient checkpointing?

Gradient checkpointing: kurangi memori dengan memilih aktivasi yang disimpan
Contoh: hitung A + B = C

Grafik yang menggambarkan gradient checkpointing dengan node dan edge

Apa itu gradient checkpointing?

Gradient checkpointing: kurangi memori dengan memilih aktivasi yang disimpan
Contoh: hitung A + B = C
- Hitung A, B dulu, lalu hitung C

Grafik yang menggambarkan gradient checkpointing dengan node dan edge

Apa itu gradient checkpointing?

Gradient checkpointing: kurangi memori dengan memilih aktivasi yang disimpan
Contoh: hitung A + B = C
- Hitung A, B dulu, lalu hitung C
- A, B tidak dibutuhkan untuk sisa forward pass
Haruskah kita simpan atau hapus A dan B?

Grafik yang menggambarkan gradient checkpointing dengan node dan edge

Apa itu gradient checkpointing?

Gradient checkpointing: kurangi memori dengan memilih aktivasi yang disimpan
Contoh: hitung A + B = C
- Hitung A, B dulu, lalu hitung C
- A, B tidak dibutuhkan untuk sisa forward pass
Haruskah kita simpan atau hapus A dan B?
- Tanpa gradient checkpointing: simpan A, B

Grafik yang menggambarkan gradient checkpointing dengan node dan edge

Apa itu gradient checkpointing?

Gradient checkpointing: kurangi memori dengan memilih aktivasi yang disimpan
Contoh: hitung A + B = C
- Hitung A, B dulu, lalu hitung C
- A, B tidak dibutuhkan untuk sisa forward pass
Haruskah kita simpan atau hapus A dan B?
- Tanpa gradient checkpointing: simpan A, B
- Dengan gradient checkpointing: hapus A, B

Grafik yang menggambarkan gradient checkpointing dengan node dan edge

Apa itu gradient checkpointing?

Gradient checkpointing: kurangi memori dengan memilih aktivasi yang disimpan
Contoh: hitung A + B = C
- Hitung A, B dulu, lalu hitung C
- A, B tidak dibutuhkan untuk sisa forward pass
Haruskah kita simpan atau hapus A dan B?
- Tanpa gradient checkpointing: simpan A, B
- Dengan gradient checkpointing: hapus A, B
- Hitung ulang A, B saat backward pass

Grafik yang menggambarkan gradient checkpointing dengan node dan edge

Apa itu gradient checkpointing?

Gradient checkpointing: kurangi memori dengan memilih aktivasi yang disimpan
Contoh: hitung A + B = C
- Hitung A, B dulu, lalu hitung C
- A, B tidak dibutuhkan untuk sisa forward pass
Haruskah kita simpan atau hapus A dan B?
- Tanpa gradient checkpointing: simpan A, B
- Dengan gradient checkpointing: hapus A, B
- Hitung ulang A, B saat backward pass
- Jika B mahal dihitung ulang, simpan B

Grafik yang menggambarkan gradient checkpointing dengan node dan edge

Trainer dan Accelerator

Bagan yang membandingkan kemudahan penggunaan vs. kemampuan kustomisasi untuk Accelerator dan Trainer.

Trainer dan Accelerator

Bagan yang membandingkan kemudahan penggunaan vs. kemampuan kustomisasi untuk Accelerator dan Trainer.

Gradient checkpointing dengan Trainer

training_args = TrainingArguments(output_dir="./results",
                                  evaluation_strategy="epoch",
                                  gradient_accumulation_steps=4)

Gradient checkpointing dengan Trainer

training_args = TrainingArguments(output_dir="./results",
                                  evaluation_strategy="epoch",
                                  gradient_accumulation_steps=4,
                                  gradient_checkpointing=True)

trainer = Trainer(model=model,
                  args=training_args,
                  train_dataset=dataset["train"],
                  eval_dataset=dataset["validation"],
                  compute_metrics=compute_metrics)

trainer.train()

{'epoch': 1.0, 'eval_loss': 0.73, 'eval_accuracy': 0.03, 'eval_f1': 0.05}

Dari Trainer ke Accelerator

Bagan yang membandingkan kemudahan penggunaan vs. kemampuan kustomisasi untuk Accelerator dan Trainer.

Gradient checkpointing dengan Accelerator

accelerator = Accelerator(gradient_accumulation_steps=2)


for index, batch in enumerate(dataloader):
    with accelerator.accumulate(model):
        inputs, targets = batch["input_ids"], batch["labels"]
        outputs = model(inputs, labels=targets)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

Gradient checkpointing dengan Accelerator

accelerator = Accelerator(gradient_accumulation_steps=2)
model.gradient_checkpointing_enable()

for index, batch in enumerate(dataloader):
    with accelerator.accumulate(model):
        inputs, targets = batch["input_ids"], batch["labels"]
        outputs = model(inputs, labels=targets)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

Local SGD meningkatkan efisiensi komunikasi

Ikon yang merepresentasikan efisiensi memori, komunikasi, dan komputasi.

Apa itu local SGD?

Diagram yang menunjukkan cara kerja local SGD dengan menyinkronkan gradien setelah sejumlah langkah.

Tiap perangkat menghitung gradien secara paralel

Apa itu local SGD?

Diagram yang menunjukkan cara kerja local SGD dengan menyinkronkan gradien setelah sejumlah langkah.

Tiap perangkat menghitung gradien secara paralel
Sinkronisasi gradien: node penggerak memperbarui parameter model di tiap perangkat
Local SGD: kurangi frekuensi sinkronisasi gradien

Local SGD dengan Accelerator





for index, batch in enumerate(dataloader):
    with accelerator.accumulate(model):
        inputs, targets = batch["input_ids"], batch["labels"]
        outputs = model(inputs, labels=targets)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

Local SGD dengan Accelerator

from accelerate.local_sgd import LocalSGD

with LocalSGD(accelerator=accelerator, model=model, local_sgd_steps=8, 
              enabled=True) as local_sgd:
    for index, batch in enumerate(dataloader):
        with accelerator.accumulate(model):
            inputs, targets = batch["input_ids"], batch["labels"]
            outputs = model(inputs, labels=targets)
            loss = outputs.loss
            accelerator.backward(loss)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()

Local SGD dengan Accelerator

from accelerate.local_sgd import LocalSGD

with LocalSGD(accelerator=accelerator, model=model, local_sgd_steps=8, 
              enabled=True) as local_sgd:
    for index, batch in enumerate(dataloader):
        with accelerator.accumulate(model):
            inputs, targets = batch["input_ids"], batch["labels"]
            outputs = model(inputs, labels=targets)
            loss = outputs.loss
            accelerator.backward(loss)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()
            local_sgd.step()

Ayo berlatih!

Pelatihan Model AI Efisien dengan PyTorch

Preparing Video For Download...