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