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

Code before:
            rank_zero_warn(
                f"FSDPPrecision `{precision}` runs computations in reduced precision "
                "(e.g., float16/bfloat16) while keeping model weights stored in full precision. "
                "These modes are still experimental and may produce slightly different accuracy or stability "
                "compared to full precision (`precision='32-true'`)."
            )
        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", "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

Review comment (suggestion):
let's instead add a warning here for 16-true and bf16-true that fsdp always keeps a full precision reference of the weights

Code after:
        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
    def forward_context(self) -> AbstractContextManager:
        if "mixed" in self.precision: