Repository: Lightning-AI/pytorch-lightning
File: src/lightning/fabric/plugins/precision/fsdp.py (Python)

Code before:
        self.precision = precision

        precision_to_type = {
            "bf16-mixed": torch.float32,
            "16-mixed": torch.float32,
            "bf16-true": torch.bfloat16,
            "16-true": torch.float16,
            "32-true": torch.float32,
        }
        self._desired_input_dtype = precision_to_type[self.precision]

    @override
    def convert_module(self, module: Module) -> Module:
        if "true" in self.precision:
            return module.to(dtype=self._desired_input_dtype)
        return module

    @property
    def mixed_precision_config(self) -> "TorchMixedPrecision":
        from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision as TorchMixedPrecision

        if self.precision in ("16-true", "bf16-true"):
            rank_zero_warn(
                f"FSDPPrecision `{self.precision}` enables mixed-precision execution. "
                "Model parameters remain in full precision `torch.float32`, while forward and backward passes "
                f"run with reduced precision `{self._desired_input_dtype}` for speed and memory efficiency."
            )

        if self.precision in ("16-true", "16-mixed"):
            param_dtype = reduce_dtype = buffer_dtype = torch.float16
        elif self.precision in ("bf16-true", "bf16-mixed"):
            param_dtype = reduce_dtype = buffer_dtype = torch.bfloat16
        elif self.precision == "32-true":
            param_dtype = reduce_dtype = buffer_dtype = torch.float32
        else:
            raise ValueError(f"Was unable to infer precision type, received {self.precision!r}.")

        return TorchMixedPrecision(
            param_dtype=param_dtype,
            reduce_dtype=reduce_dtype,
            buffer_dtype=buffer_dtype,
        )

    @override
    def tensor_init_context(self) -> AbstractContextManager:
        return _DtypeContextManager(self._desired_input_dtype)

    @override
    def module_init_context(self) -> AbstractContextManager:
        return _DtypeContextManager(self.mixed_precision_config.param_dtype or torch.float32)


Review comment (suggestion):
```suggestion
                f"FSDP with `{self.precision}` enables computation in lower precision. "
                "FSDP will always retain a full-precision copy of the model parameters for sharding."
```

Code after:
        self.precision = precision

        precision_to_type = {
            "bf16-mixed": torch.float32,
            "16-mixed": torch.float32,
            "bf16-true": torch.bfloat16,
            "16-true": torch.float16,
            "32-true": torch.float32,
        }
        self._desired_input_dtype = precision_to_type[self.precision]

    @override
    def convert_module(self, module: Module) -> Module:
        if "true" in self.precision:
            return module.to(dtype=self._desired_input_dtype)
        return module

    @property
    def mixed_precision_config(self) -> "TorchMixedPrecision":
        from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision as TorchMixedPrecision

        if self.precision in ("16-true", "bf16-true"):
            rank_zero_warn(
                f"FSDP with `{self.precision}` enables computation in lower precision. "
                "FSDP will always retain a full-precision copy of the model parameters for sharding."
            )

        if self.precision in ("16-true", "16-mixed"):
            param_dtype = reduce_dtype = buffer_dtype = torch.float16
        elif self.precision in ("bf16-true", "bf16-mixed"):
            param_dtype = reduce_dtype = buffer_dtype = torch.bfloat16
        elif self.precision == "32-true":
            param_dtype = reduce_dtype = buffer_dtype = torch.float32
        else:
            raise ValueError(f"Was unable to infer precision type, received {self.precision!r}.")

        return TorchMixedPrecision(
            param_dtype=param_dtype,
            reduce_dtype=reduce_dtype,
            buffer_dtype=buffer_dtype,
        )

    @override
    def tensor_init_context(self) -> AbstractContextManager:
        return _DtypeContextManager(self._desired_input_dtype)

    @override
    def module_init_context(self) -> AbstractContextManager:
        return _DtypeContextManager(self.mixed_precision_config.param_dtype or torch.float32)

    @override