Gradient checkpointing s Acceleratorem

Pokračuješ v optimalizaci využití paměti, aby bylo možné trénovat model pro překlad jazyka přímo na tvém zařízení. Gradient accumulation ti pomohl efektivně trénovat s většími batch sizes. Navazuj na tuto práci a přidej gradient checkpointing, který sníží paměťové nároky modelu.

Objekty model, train_dataloader a accelerator jsou již předem definované.

Toto cvičení je součástí kurzu

Efficient AI Model Training with PyTorch

Zobrazit kurz

Pokyny k cvičení

Aktivuj gradient checkpointing na objektu model.
Nastav context manager Acceleratoru pro povolení gradient accumulation na objektu model.

Interaktivní cvičení na vyzkoušení si v praxi

Vyzkoušejte si toto cvičení dokončením tohoto ukázkového kódu.

# Enable gradient checkpointing on the model
____.____()

for batch in train_dataloader:
    with accelerator.accumulate(model):
        inputs, targets = batch["input_ids"], batch["labels"]
        # Get the outputs from a forward pass of the model
        ____ = ____(____, labels=targets)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        print(f"Loss = {loss}")

Upravit a spustit kód