Cache function implementation for TracInAttributor

[brightXian]

Implement .cache() for TracIn family attributors. Currently, .cache() is an empty method and training gradients are recomputed on every .attribute() call.

Proposed changes:

1. .cache():

def cache(
    self,
    full_train_dataloader: torch.utils.data.DataLoader,
) -> None:

Pre-compute and store training gradients for all checkpoints.

2. Modify .attribute() signature to support cached mode:

def attribute(
    self,
    test_dataloader: torch.utils.data.DataLoader,
    train_dataloader: Optional[torch.utils.data.DataLoader] = None,
) -> Tensor:

train_dataloader becomes optional. When .cache() has been called, .attribute() can be called with only test_dataloader.

[TheaperDeng]

Maybe we can also support offloading in TracInAttributor (since it is the easiest one). Offloading means we want to offload some large data chunk (e.g., per-example training data gradients in this case) to CPU or disk.

Currently we have a offloader here: https://github.com/TRAIS-Lab/dattri/blob/main/dattri/algorithm/block_projected_if/offload/offload.py (though it is not designed to be used by other attributors). We may take this as a "trial" to see if there is any difficulties / tech block when we want to use the offloader in other attributor / as a modularized low-level API.

[brightXian]

I think I can add an offload parameter to TracInAttributor - something like offload="none" / "cpu" / "disk". none means current behavior (all in GPU), while cpu / disk let users offload gradients to save GPU memory. The idea is to reuse the existing offloader from block_projected_if.

With offload, we can keep the same .cache() + .attribute() workflow but save GPU memory:

# default: offload="none", all in GPU
attributor = TracInAttributor(...)
attributor.cache(train_loader)
score = attributor.attribute(test_loader)

# offload="cpu": gradients on CPU
attributor = TracInAttributor(..., offload="cpu")

# offload="disk": gradients on disk (requires cache_dir)
attributor = TracInAttributor(..., offload="disk", cache_dir="./cache")

attributor.cache(train_loader)
score = attributor.attribute(test_loader)