Enable access to training_step output for Callback.on_before_backward

[zihaozou]

Description & Motivation

It will be very useful if I can access the training_output in on_before_backward. For training gaussian splatting, I can implement the density control as a Callback and to retain_grad() for means2d

Pitch

No response

Alternatives

No response

Additional context

No response

cc @lantiga @Borda

[SkafteNicki]

Thanks for raising this issue. I am glad to report that this is already possible, by saving the the training outputs as an attribute in your lightning module which can then be access in any callback using the pl_module argument that is available in the on_before_backward hook. An example to simulate gaussian splatting:

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset

import lightning.pytorch as pl
from lightning.pytorch import Trainer
from lightning.pytorch.callbacks import Callback


class GaussianSplattingModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(10, 1)
        self.current_training_outputs = None
        
    def forward(self, x):
        return self.layer(x)
        
    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        base_loss = F.mse_loss(y_hat, y)
        
        # Simulate gaussian splatting scenario
        means2d = torch.randn(100, 2, requires_grad=True, device=self.device)
        regularization = 0.001 * torch.mean(means2d ** 2)
        loss = base_loss + regularization
        
        outputs = {
            "loss": loss,
            "predictions": y_hat,
            "means2d": means2d,  # Needed for density control in callback
            "batch_idx": batch_idx,
        }
        
        # SAVE TRAINING OUTPUT
        self.current_training_outputs = outputs
        
        return outputs
    
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.01)


class DensityControlCallback(Callback):
    def on_before_backward(self, trainer, pl_module, loss):
        if hasattr(pl_module, 'current_training_outputs') and pl_module.current_training_outputs:
            outputs = pl_module.current_training_outputs
            
            print(f"Batch {outputs['batch_idx']}: Loss = {loss.item():.4f}")
            
            # Key use case: Apply retain_grad to intermediate tensors
            means2d = outputs['means2d']
            if means2d.requires_grad:
                means2d.retain_grad()
                print(f"  ✓ Applied retain_grad() to means2d (shape: {means2d.shape})")
            
    def on_after_backward(self, trainer, pl_module):
        """Verify gradients are available after backward."""
        if hasattr(pl_module, 'current_training_outputs') and pl_module.current_training_outputs:
            outputs = pl_module.current_training_outputs
            means2d = outputs['means2d']
            
            if means2d.grad is not None:
                print(f"  ✓ Gradients available for density control operations")

def main():
    x = torch.randn(20, 10)
    y = torch.randn(20, 1)
    train_loader = DataLoader(TensorDataset(x, y), batch_size=10)
    
    model = GaussianSplattingModel()
    callback = DensityControlCallback()
    
    trainer = Trainer(
        max_steps=2,
        callbacks=[callback],
        logger=False,
        enable_checkpointing=False,
        enable_progress_bar=False,
        enable_model_summary=False,
    )
    
    trainer.fit(model, train_loader)
    

if __name__ == "__main__":
    main()

Closing issue but feel free to ping and reopen if necessary.