Parse flex kernel options with argparse

Author: heiheiha798

scripts/train_dflash.py

@@ -70,7 +70,7 @@ def parse_args():
     )
     model_group.add_argument(
         "--flex-kernel-options-json",
-        type=str,
+        type=json.loads,
         default=None,
         help="JSON dict forwarded as kernel_options when attention-backend=flex_attention.",
     )
@@ -382,14 +382,13 @@ def main():
     )
 
     args = parse_args()
-    flex_kernel_options = None
-    if args.flex_kernel_options_json is not None:
+    flex_kernel_options = args.flex_kernel_options_json
+    if flex_kernel_options is not None:
         if args.attention_backend != "flex_attention":
             raise ValueError(
                 "--flex-kernel-options-json can only be used when "
                 "--attention-backend is 'flex_attention'."
             )
-        flex_kernel_options = json.loads(args.flex_kernel_options_json)
         if not isinstance(flex_kernel_options, dict):
             raise ValueError(
                 "--flex-kernel-options-json must decode to a JSON object."

scripts/train_domino.py

@@ -67,7 +67,7 @@ def parse_args():
     )
     model_group.add_argument(
         "--flex-kernel-options-json",
-        type=str,
+        type=json.loads,
         default=None,
         help="JSON dict forwarded as kernel_options when attention-backend=flex_attention.",
     )
@@ -451,14 +451,13 @@ def main():
     )
 
     args = parse_args()
-    flex_kernel_options = None
-    if args.flex_kernel_options_json is not None:
+    flex_kernel_options = args.flex_kernel_options_json
+    if flex_kernel_options is not None:
         if args.attention_backend != "flex_attention":
             raise ValueError(
                 "--flex-kernel-options-json can only be used when "
                 "--attention-backend is 'flex_attention'."
             )
-        flex_kernel_options = json.loads(args.flex_kernel_options_json)
         if not isinstance(flex_kernel_options, dict):
             raise ValueError(
                 "--flex-kernel-options-json must decode to a JSON object."