Changed the way we get the parameters' device for fsdp

facebookresearch / esm

Evolutionary Scale Modeling (esm): Pretrained language models for proteins

MIT License

3.16k stars 627 forks source link

Changed the way we get the parameters' device for fsdp #248

Closed zzk1st closed 2 years ago

zzk1st commented 2 years ago

Currently if we want to get the device of the parameters of a module, we do: next(self.parameters()).device

However this is not possible if the model is FSDP-wrapped, since FSDP removes all parameters from self._parameters and next(self.parameters()).device will only get a StopIteration error.

The goal is to get the device in a way that whether the model is FSDP-wrapped, we can get the correct device.

zzk1st commented 2 years ago

Changes made in other places, closing this one.