Gradient Checkpointing in Pure PyTorch

[pharringtonp19]

I was wondering if anyone had some suggestions about how to incorporate gradient checkpointing into the following code. I'm working with qlora LLMs

@dataclass
class DecoderTrainer:
    model: torch.nn.Module
    tokenizer: PreTrainedTokenizerBase
    optimizer: torch.optim.Optimizer
    scheduler: torch.optim.lr_scheduler
    criterion: torch.nn.Module
    device: torch.device
    verbose: bool = True   
    gradient_accumulation: int = 4
    threshold: int = 10 

    def process_batch(self, batch, token_type = 'target_token'):
        """Compute logits and labels for a batch of data."""
        input_ids = batch['input_ids'].to(self.device)
        attention_mask = batch['attention_mask'].to(self.device)
        logits = self.model(input_ids, attention_mask).logits
        return input_ids, logits, batch['type_indicator'].to(self.device), batch[token_type].to(self.device)

    def compute_loss(self, input_ids, logits, type_indicator, token_type):

        # Focus on the final 'threshold' tokens
        labels = input_ids[:, -self.threshold:].contiguous()
        logits = logits[:, -self.threshold-1:-1, :].contiguous()
        
        """Compute the loss for the given logits and labels."""
        criterion_mode = getattr(self.criterion, 'mode', None)
        if criterion_mode == 'input':
            type_indicator = type_indicator.repeat_interleave(self.threshold)
            loss = self.criterion(logits.view(-1, logits.size(-1)), labels.view(-1), type_indicator)
        else:
            loss = self.criterion(logits.view(-1, logits.size(-1)), labels.view(-1))

        # Extract log probabilities for the second-to-last token
        second_to_last_logits = logits[:, -2, :]
        log_probs = torch.log_softmax(second_to_last_logits, dim=-1)

        # Get log probability for the target token
        log_prob_target = log_probs.gather(1, token_type.unsqueeze(1)).squeeze(1)

        return loss, log_prob_target


    def train(self, data_loader):
        self.model.train()
        total_loss = 0
        total_samples = 0
        total_neg_log_prob = 0
        self.optimizer.zero_grad()

        for i, batch in enumerate(data_loader):
            input_ids, logits, type_indicator, target_token = self.process_batch(batch, 'target_token')
            loss, log_prob_target = self.compute_loss(input_ids, logits, type_indicator, target_token)
            neg_log_prob = -log_prob_target.mean().item()

            total_samples += len(batch)
            total_loss += loss.item() * len(batch)
            total_neg_log_prob += neg_log_prob * len(batch)

            loss = loss / self.gradient_accumulation  # Normalizing loss for accumulation steps
            loss.backward()

            if (i + 1) % self.gradient_accumulation == 0:  # Optimizer step after 'accumulation_steps' batches
                self.optimizer.step()
                self.scheduler.step()
                self.optimizer.zero_grad()
                if self.verbose:
                    print(f"Batch {i+1} Loss: {loss.item() * self.gradient_accumulation}")
                    print(f"Batch {i+1} Negative Log Probability: {neg_log_prob}")

                torch.cuda.empty_cache()

        average_loss = total_loss / total_samples
        average_neg_log_prob = total_neg_log_prob / total_samples

        return average_loss, average_neg_log_prob

[Advait251206]

Hi @pharringtonp19,

Since you're working with QLoRA LLMs and a custom PyTorch training loop, you can incorporate gradient checkpointing without needing to switch to Trainer or SFTTrainer.

The important distinction is that gradient checkpointing and gradient accumulation solve different memory problems:

• Gradient accumulation reduces the effective batch size held in memory during each forward/backward pass.
• Gradient checkpointing reduces activation memory by recomputing selected intermediate activations during backpropagation.

For your implementation, I would recommend enabling gradient checkpointing directly through the model's supported interface.

1. Enable gradient checkpointing

If your model is based on Hugging Face Transformers, you can generally enable checkpointing using:

model.gradient_checkpointing_enable()

Since you're using QLoRA, make sure the model is prepared appropriately for k-bit training.

For example, with PEFT:

from peft import prepare_model_for_kbit_training

model = prepare_model_for_kbit_training(
    model,
    use_gradient_checkpointing=True,
)

This prepares the quantized model for adapter training and enables gradient checkpointing.

If you have already configured checkpointing through prepare_model_for_kbit_training, you generally do not need to enable it a second time.

2. Disable the KV cache

One important detail when using gradient checkpointing with decoder-only language models is disabling the KV cache during training.

Add:

model.config.use_cache = False

The KV cache is primarily useful for autoregressive generation, where previously computed keys and values are reused.

During training with gradient checkpointing, it is generally unnecessary and can cause compatibility issues.

You can re-enable it for inference later if needed.

3. Your training loop can remain mostly unchanged

You do not need to manually wrap every forward pass with torch.utils.checkpoint.checkpoint() if the model already supports gradient checkpointing.

Your existing code:

logits = self.model(input_ids, attention_mask).logits

can remain as it is.

The model handles checkpointing internally during the forward and backward passes.

Your gradient accumulation logic can also remain unchanged.

4. One important consideration in your implementation

I noticed that your code calculates logits for the entire input sequence:

logits = self.model(input_ids, attention_mask).logits

But your loss only uses the final threshold tokens:

labels = input_ids[:, -self.threshold:].contiguous()

logits = logits[:, -self.threshold-1:-1, :].contiguous()

For long sequences, the full logits tensor can consume substantial GPU memory.

Gradient checkpointing reduces the memory required for intermediate activations, but it does not automatically eliminate the memory required to produce and retain the full logits tensor.

If your model architecture supports computing logits only for selected positions, you could investigate that optimization.

Otherwise, consider reducing sequence length or physical batch size if memory remains a bottleneck.

5. A small improvement to your training loop

Your current accumulation logic performs an optimizer step only when the number of batches is divisible by gradient_accumulation.

If the final batch group is incomplete, those accumulated gradients are never applied.

You can handle this by checking whether the current batch is either a complete accumulation group or the final batch:

for i, batch in enumerate(data_loader):

    input_ids, logits, type_indicator, target_token = self.process_batch(
        batch,
        "target_token"
    )

    loss, log_prob_target = self.compute_loss(
        input_ids,
        logits,
        type_indicator,
        target_token
    )

    (loss / self.gradient_accumulation).backward()

    is_accumulation_step = (
        (i + 1) % self.gradient_accumulation == 0
    )

    is_last_batch = (i + 1) == len(data_loader)

    if is_accumulation_step or is_last_batch:
        self.optimizer.step()
        self.scheduler.step()
        self.optimizer.zero_grad()

This is a simplified illustration. If the final accumulation group contains fewer batches than the configured accumulation count, you should adjust the loss scaling for that group to obtain the intended gradient average.

6. Recommended setup

For your QLoRA model, I would start with:

model.config.use_cache = False

model.gradient_checkpointing_enable()

model.train()

If you are using PEFT's k-bit preparation utility, use that to enable checkpointing instead of duplicating the configuration.

Then retain your existing custom training loop and benchmark GPU memory usage.

One final consideration: gradient checkpointing introduces additional computation during backward propagation. It should reduce memory consumption, but training may become somewhat slower.

Conclusion

You can use gradient checkpointing directly with your existing custom PyTorch training loop. The simplest approach is to enable it on the Hugging Face model, disable the KV cache, and leave your forward/backward logic intact.

For your particular implementation, I would also investigate the memory consumed by the full-sequence logits tensor, since your loss only uses a small number of final tokens.

Hope this helps!

---

If this answer helped or pointed you in the right direction, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206