Gradient accumulation in PyTorch - Practice Problems & Coding Challenges

Challenge - 5 Problems

❓ Predict Output

intermediate

Output of gradient accumulation with batch size 2

Consider the following PyTorch code snippet that uses gradient accumulation with accumulation_steps=2. What will be the printed value of model.linear.weight.grad after the loop?

PyTorch

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1, bias=False)

model = SimpleModel()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

inputs = [torch.tensor([[1.0]]), torch.tensor([[2.0]])]
targets = [torch.tensor([[2.0]]), torch.tensor([[4.0]])]

accumulation_steps = 2
optimizer.zero_grad()

for i in range(2):
    output = model.linear(inputs[i])
    loss = (output - targets[i]).pow(2).mean()
    loss.backward()
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

print(model.linear.weight.grad)

A tensor([[1.]])

B None

C tensor([[0.]])

D tensor([[0.5]])

💡 Hint

Remember that after optimizer.step() and optimizer.zero_grad(), gradients are reset.

After the loop, optimizer.step() and optimizer.zero_grad() are called, so gradients are cleared. Therefore, model.linear.weight.grad is a tensor of zeros.

🧠 Conceptual

intermediate

Why use gradient accumulation?

Why do we use gradient accumulation in training deep learning models?

A To avoid using an optimizer during training.

B To reduce the number of model parameters by accumulating gradients.

C To speed up training by skipping backward passes.

D To simulate a larger batch size than fits in memory by accumulating gradients over smaller batches before updating weights.

💡 Hint

Think about memory limits and batch sizes.

Gradient accumulation allows training with an effective larger batch size by summing gradients over multiple smaller batches before updating model weights, which helps when GPU memory is limited.

❓ Hyperparameter

advanced

Effect of accumulation steps on learning rate

If you increase the number of gradient accumulation steps from 1 to 4 without changing the learning rate, what is the effective change in the learning rate per sample?

A The effective learning rate per sample is 4 times smaller.

B The effective learning rate per sample is 4 times larger.

C The effective learning rate per sample stays the same.

D The effective learning rate per sample becomes zero.

💡 Hint

Think about how many updates happen per sample when accumulating gradients.

With 4 accumulation steps, the model updates weights once every 4 batches, so each sample's gradient contributes less frequently, effectively reducing the learning rate per sample by 4 times.

🔧 Debug

advanced

Bug in gradient accumulation code

What error will the following PyTorch code raise?

import torch
import torch.nn as nn

model = nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

accumulation_steps = 2
optimizer.zero_grad()

for i in range(3):
    x = torch.tensor([[float(i + 1)]])
    y = torch.tensor([[2.0 * (i + 1)]])
    output = model(x)
    loss = (output - y).pow(2).mean()
    loss.backward()
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

optimizer.step()

A RuntimeError: optimizer.step() called twice without backward()

B RuntimeError: optimizer.step() called without zero_grad()

C RuntimeError: Trying to backward through the graph a second time

D No error, code runs fine

💡 Hint

Check how many times optimizer.step() is called relative to loss.backward().

The last optimizer.step() is called after the loop without a preceding loss.backward(), causing a RuntimeError because optimizer.step() expects gradients from a backward pass.

❓ Model Choice

expert

Choosing batch size and accumulation steps for limited GPU memory

You want to train a large transformer model but your GPU memory only fits batch size 8. You want an effective batch size of 32. Which approach is best?

A Use batch size 4 with accumulation steps = 8 to simulate batch size 32.

B Use batch size 8 with gradient accumulation steps = 4 to simulate batch size 32.

C Use batch size 8 and update weights after every batch without accumulation.

D Use batch size 32 and reduce model size to fit GPU memory.

💡 Hint

Think about memory limits and how accumulation simulates larger batch sizes.

Using batch size 8 with accumulation steps 4 allows simulating batch size 32 without exceeding GPU memory. Batch size 32 may not fit, and batch size 4 with accumulation 8 is less efficient due to smaller batches.

Practice

1. What is the main purpose of gradient accumulation in PyTorch training?

easy

A. To reduce the number of model parameters

B. To increase the learning rate automatically

C. To change the model architecture dynamically

D. To split a large batch into smaller steps to save memory

Solution

1. Step 1: Understand gradient accumulation concept

   Gradient accumulation splits a large batch into smaller mini-batches to avoid memory overflow.

2. Step 2: Identify the main benefit

   This technique helps save GPU memory by processing smaller batches sequentially before updating weights.

3. Final Answer:

   To split a large batch into smaller steps to save memory -> Option D

4. Quick Check:

   Gradient accumulation = batch splitting for memory saving [OK]

Hint: Gradient accumulation saves memory by splitting batches [OK]

Common Mistakes:

• Thinking it changes learning rate
• Confusing with model size reduction
• Assuming it alters model layers

2. Which PyTorch code snippet correctly implements gradient accumulation over 4 steps?

easy

A. for i, data in enumerate(dataloader):
outputs = model(data)
loss = criterion(outputs, labels)
optimizer.step()
loss.backward()
optimizer.zero_grad()

B. for i, data in enumerate(dataloader):
outputs = model(data)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
if (i + 1) % 4 == 0:
optimizer.zero_grad()

C. for i, data in enumerate(dataloader):
outputs = model(data)
loss = criterion(outputs, labels)
loss.backward()
if (i + 1) % 4 == 0:
optimizer.step()
optimizer.zero_grad()

D. for i, data in enumerate(dataloader):
outputs = model(data)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
if (i + 1) % 4 == 0:
optimizer.step()

Solution

1. Step 1: Identify correct order of operations

   Gradient accumulation requires calling loss.backward() each step, then optimizer.step() and zero_grad() after accumulating gradients.

2. Step 2: Check modulo condition for accumulation

   Only after 4 mini-batches (i+1 % 4 == 0) should optimizer.step() and zero_grad() be called to update weights and reset gradients.

3. Final Answer:

   Code that calls loss.backward() every step and optimizer.step() plus zero_grad() every 4 steps -> Option C

4. Quick Check:

   Backward every step, step and zero_grad every 4 steps [OK]

Hint: Call backward every step, step and zero_grad every N steps [OK]

Common Mistakes:

• Calling optimizer.step() before loss.backward()
• Resetting gradients too early
• Not using modulo condition correctly

3. Given this PyTorch code with gradient accumulation steps=2, what is the printed loss after 4 batches?

optimizer.zero_grad()
for i, (inputs, labels) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, labels) / 2
    loss.backward()
    if (i + 1) % 2 == 0:
        optimizer.step()
        optimizer.zero_grad()
    print(f"Batch {i+1} loss: {loss.item()}")

medium

A. Loss values cause runtime error due to division

B. Loss values divided by 2 printed for each batch

C. Loss values accumulated and printed only every 2 batches

D. Loss values printed without division

Solution

1. Step 1: Understand loss scaling in code

   Loss is divided by 2 before backward to average gradients over 2 accumulation steps.

2. Step 2: Check print statement inside loop

   Loss.item() prints the scaled loss for each batch, so printed values are divided by 2 every batch.

3. Final Answer:

   Loss values divided by 2 printed for each batch -> Option B

4. Quick Check:

   Loss scaled by 2 and printed each batch [OK]

Hint: Loss is scaled before backward, printed as scaled every batch [OK]

Common Mistakes:

• Thinking loss prints only every 2 batches
• Ignoring loss division by 2
• Assuming runtime error from division

4. Identify the error in this gradient accumulation code snippet:

optimizer.zero_grad()
for i, (inputs, labels) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    if i % 4 == 0:
        optimizer.step()
        optimizer.zero_grad()

medium

A. optimizer.step() and zero_grad() are called too early at i=0

B. loss.backward() is missing gradient scaling

C. optimizer.zero_grad() should be called before loss.backward()

D. The loop should use (i + 1) % 4 == 0 instead of i % 4 == 0

Solution

1. Step 1: Analyze the modulo condition

   Condition i % 4 == 0 triggers optimizer.step() and zero_grad() at i=0, which is before any gradients accumulate.

2. Step 2: Understand correct timing for optimizer step

   Optimizer step should happen after 4 batches, so condition should be (i + 1) % 4 == 0 to avoid early update.

3. Final Answer:

   optimizer.step() and zero_grad() are called too early at i=0 -> Option A

4. Quick Check:

   Modulo condition triggers step too early [OK]

Hint: Use (i + 1) % N == 0 to delay step until after N batches [OK]

Common Mistakes:

• Calling optimizer.step() at first batch
• Confusing zero_grad() timing
• Ignoring off-by-one in modulo condition

5. You want to train a large model on a GPU with limited memory using gradient accumulation with batch size 64 split into 8 steps. Which approach correctly adjusts the learning rate and loss scaling to keep training stable?

hard

A. Divide the loss by 8 before backward and keep the original learning rate

B. Multiply the learning rate by 8 and do not scale the loss

C. Divide the learning rate by 8 and multiply the loss by 8

D. Keep loss and learning rate unchanged, just accumulate gradients

Solution

1. Step 1: Understand gradient accumulation effect on loss

   Splitting batch into 8 steps means gradients accumulate over 8 mini-batches, so loss must be divided by 8 to average gradients correctly.

2. Step 2: Adjust learning rate accordingly

   Learning rate stays the same because gradient accumulation simulates a larger batch size without changing optimization dynamics.

3. Final Answer:

   Divide the loss by 8 before backward and keep the original learning rate -> Option A

4. Quick Check:

   Loss scaled down, learning rate unchanged for stable accumulation [OK]

Hint: Scale loss down by steps, keep learning rate same [OK]

Common Mistakes:

• Scaling learning rate instead of loss
• Not scaling loss causing gradient explosion
• Assuming no changes needed for accumulation