support fsdp2 string strategy (#996)

Summary:
Pull Request resolved: #996

Support "fsdp2" string as strategy

Reviewed By: galrotem

Differential Revision: D74410718

fbshipit-source-id: 4419cea3d1709ca89911b4743c53ae6536448015

Author: JKSenthil

torchtnt/utils/prepare_module.py

@@ -619,7 +619,7 @@ def prepare_module(
 
 def convert_str_to_strategy(
     strategy: str,
-) -> Union[DDPStrategy, FSDPStrategy, NOOPStrategy]:
+) -> Union[DDPStrategy, FSDPStrategy, FSDP2Strategy, NOOPStrategy]:
     """
     Converts strategy as a string to a default instance of the Strategy dataclass.
 
@@ -633,6 +633,7 @@ def convert_str_to_strategy(
     string_to_strategy_mapping = {
         "ddp": DDPStrategy(),
         "fsdp": FSDPStrategy(),
+        "fsdp2": FSDP2Strategy(),
         "noop": NOOPStrategy(),
     }