@Borda FYA

@jyothisambolu Isn't PyTorch's FSDP only compatible with CUDA? Do you have a custom FSDPStrategy? I can't find one in lightning-habana.

@jyothisambolu Isn't PyTorch's FSDP only compatible with CUDA? Do you have a custom FSDPStrategy? I can't find one in lightning-habana.

@awaelchli - to push HPUFSDPStrategy, we need this change to be in place else we get the error mentioned in the issue above. Once this is pushed in, we'll push the FSDP PR to lightning-habana. I can share a draft if it helps.

@jerome-habana Yes please, that would be very insightful. I'd like to know whether your FSDP strategy subclasses ours and if so, to what extend.

@jerome-habana Yes please, that would be very insightful. I'd like to know whether your FSDP strategy subclasses ours and if so, to what extend.

@awaelchli checkout Lightning-AI/lightning-Habana#174 . It mostly includes changes wrt initialization and device usage instead of index

@jerome-habana Thanks for sharing. I opened an alternative approach in #19781 that would unblock you. #19781 is closer to the original intent of the error checking and doesn't require a check against the integration. After the FSDP strategy in habana-lightning is released, we can still add more error checking if we want to.

Regarding the HPU FSDP implementation: I see there are several overrides of internals. Please note that we don't provide backward compatibility for internals of strategies, especially the protected methods and members. We will make changes to the FSDP strategy that may break your overrides. If this is undesired, I recommend not subclassing FSDPStrategy and instead subclassing ParallelStrategy.

@jerome-habana Thanks for sharing. I opened an alternative approach in #19781 that would unblock you. #19781 is closer to the original intent of the error checking and doesn't require a check against the integration. After the FSDP strategy in habana-lightning is released, we can still add more error checking if we want to.

Regarding the HPU FSDP implementation: I see there are several overrides of internals. Please note that we don't provide backward compatibility for internals of strategies, especially the protected methods and members. We will make changes to the FSDP strategy that may break your overrides. If this is undesired, I recommend not subclassing FSDPStrategy and instead subclassing ParallelStrategy.

@awaelchli in the long run, we want to move away from inheriting native strategies. But for now, we are dependent on base class fsdp strategy with the gpu checks removed to enable fsdp with hpu. This is to avoid some code duplication and have accelerator specific change alone inside hpufsdpstrategy.