[Activation-checkpointing] add tensor dedup and param offloading #4247
Conversation
|
@sywangyi would you be kind enough to test this on your hardware and give me some feedback? thank you! |
|
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. |
trl/models/activation_offloading.py
Outdated
|
|
||
| # Try to import DTensor for FSDP v2 support | ||
| try: | ||
| from torch.distributed._tensor import DTensor |
There was a problem hiding this comment.
do you know in which version DTensor was introduced? I'm wondering is this try/expect is needed
There was a problem hiding this comment.
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 DTensorbut we also need to check for torch.distributed.is_available(), otherwise you might get import issue.
|
|
||
| # Check if tensor is a parameter or buffer | ||
| if isinstance(activation, torch.nn.Parameter) or ( | ||
| hasattr(torch.nn, "Buffer") and isinstance(activation, torch.nn.Buffer) |
There was a problem hiding this comment.
same question here, is Buffer a recent addition?
There was a problem hiding this comment.
no buffer has always been there from the start, I can clean this up
trl/models/activation_offloading.py
Outdated
|
|
||
| # Try to import DTensor for FSDP v2 support | ||
| try: | ||
| from torch.distributed._tensor import DTensor |
There was a problem hiding this comment.
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 DTensorbut we also need to check for torch.distributed.is_available(), otherwise you might get import issue.
trl/models/activation_offloading.py
Outdated
| Returns: | ||
| A tuple of (storage_pointer, dtype) that uniquely identifies the tensor's storage | ||
| """ | ||
| storage_ptr = tensor.untyped_storage().data_ptr() + tensor.storage_offset() |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Edit: here is an example (i.e. float8linear, which can sometimes happen). https://github.com/huggingface/accelerate/blob/f0313a64a2f3de359924c85a98ee010c47b846ec/src/accelerate/accelerator.py#L3842
| # For FSDP v2: extract local tensor from DTensor | ||
| actual_tensor = p | ||
| if DTensor is not None and isinstance(p, DTensor) and hasattr(p, "_local_tensor"): | ||
| actual_tensor = p._local_tensor |
There was a problem hiding this comment.
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
pytest tests/test_activation_offloading.py::TestActivationOffloading::test_parameter_filtering these two cases pass in intel xpu |
|
@S1ro1 ok i'll just skip FP8 activations |
|
@SunMarc I have added support for broadcast and non-contiguous tensors |
7 participants
What does this PR do?