Add support for e5e2 and default to hybrid when launcher is used (#3640)

* add support for e5e2 and defaumt to hybrid when launcher is used

* style

Author: IlyasMoutawwakil

src/accelerate/commands/config/cluster.py

@@ -774,8 +774,8 @@ def get_cluster_input():
                     )
                     fp8_config["fp8_format"] = _ask_options(
                         "Which weight format should be used?",
-                        ["HYBRID", "E4M3"],
-                        lambda x: "HYBRID" if x == 0 else "E4M3",
+                        ["HYBRID", "E4M3", "E5M2"],
+                        lambda i: ["HYBRID", "E4M3", "E5M2"][i],
                         default=0,
                     )
                     fp8_config["amax_history_length"] = _ask_field(

src/accelerate/commands/launch.py

@@ -692,8 +692,8 @@ def launch_command_parser(subparsers=None):
     fp8_args.add_argument(
         "--fp8_format",
         type=str,
-        default="E4M3",
-        choices=["E4M3", "HYBRID"],
+        default="HYBRID",
+        choices=["HYBRID", "E4M3", "E5M2"],
         help="The format to use for the FP8 recipe (useful only when `--fp8_backend=te` is passed).",
     )
     fp8_args.add_argument(

src/accelerate/utils/dataclasses.py

@@ -289,7 +289,7 @@ def __post_init__(self):
 # Literals
 Backend = Literal["MSAMP", "TE"]
 OptLevel = Literal["O1", "O2"]
-FP8Format = Literal["E4M3", "HYBRID"]
+FP8Format = Literal["HYBRID", "E4M3", "E5M2"]
 AmaxComputeAlgorithm = Literal["max", "most_recent"]
 
 
@@ -342,8 +342,8 @@ class TERecipeKwargs(KwargsHandler):
         interval (`int`, *optional*, default to 1):
             The interval to use for how often the scaling factor is recomputed.
         fp8_format (`str`, *optional*, default to "HYBRID"):
-            The format to use for the FP8 recipe. Must be one of `HYBRID` or `E4M3`. (Generally `HYBRID` for training,
-            `E4M3` for evaluation)
+            The format to use for the FP8 recipe. Must be one of `HYBRID`, `E4M3` or `E5M2`. (Generally `HYBRID` for
+            training, `E4M3` or `E5M2` for evaluation)
         amax_history_len (`int`, *optional*, default to 1024):
             The length of the history to use for the scaling factor computation
         amax_compute_algo (`str`, *optional*, default to "most_recent"):