[Activation-checkpointing] add tensor dedup and param offloading

[kashif]

What does this PR do?

• Prevents redundant offloading when multiple tensor views share the same storage
• Tracks and filters out model parameters during offloading

[kashif]

@sywangyi would you be kind enough to test this on your hardware and give me some feedback? thank you!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[qgallouedec]

do you know in which version DTensor was introduced? I'm wondering is this try/expect is needed

[SunMarc]

For reference we have the following in accelerate:

if is_torch_version(">=", "2.5.0"):
        from torch.distributed.tensor import DTensor
    else:
        # from torch 2.0.0 (oldest supported accelerate torch version), DTensor is in torch.distributed._tensor
        from torch.distributed._tensor import DTensor

but we also need to check for torch.distributed.is_available(), otherwise you might get import issue.

[qgallouedec]

same question here, is Buffer a recent addition?

[kashif]

no buffer has always been there from the start, I can clean this up

[SunMarc]

LGTM for fsdpv2 part !

[SunMarc]

For reference we have the following in accelerate:

if is_torch_version(">=", "2.5.0"):
        from torch.distributed.tensor import DTensor
    else:
        # from torch 2.0.0 (oldest supported accelerate torch version), DTensor is in torch.distributed._tensor
        from torch.distributed._tensor import DTensor

but we also need to check for torch.distributed.is_available(), otherwise you might get import issue.

[S1ro1]

Using data_ptr() can be a bit tricky with for example TorchAO quantized tensors etc, as those can return 0 for data_ptr(). I don't have a concrete example, just something to be aware of.

[S1ro1]

Edit: here is an example (i.e. float8linear, which can sometimes happen). https://github.com/huggingface/accelerate/blob/f0313a64a2f3de359924c85a98ee010c47b846ec/src/accelerate/accelerator.py#L3842

[S1ro1]

Again, something to care for. If fp8 is used, it can return 0, viz here: https://github.com/huggingface/accelerate/blob/f0313a64a2f3de359924c85a98ee010c47b846ec/src/accelerate/accelerator.py#L3842

[sywangyi]

@sywangyi would you be kind enough to test this on your hardware and give me some feedback? thank you!

pytest tests/test_activation_offloading.py::TestActivationOffloading::test_parameter_filtering
pytest tests/test_activation_offloading.py::TestActivationOffloading::test_tensor_deduplication

these two cases pass in intel xpu

[kashif]

@S1ro1 ok i'll just skip FP8 activations

[kashif]

@SunMarc I have added support for broadcast and non-contiguous tensors

[sergiopaniego]

lgtm!