arg

app.py

@@ -17,7 +17,10 @@ print("Done")
 
 def create_medusa_heads(model_id: str):
     parser = distributed_run.get_args_parser()
-    training_args = [
+    args = parser.parse_args([
+        "--nproc_per_node", "4",
+        "training_script", "medusa/medusa/train/train.py",
+        "training_script_args",
         "--model_name_or_path", model_id,
         "--data_path", "data/ShareGPT_V4.3_unfiltered_cleaned_split.json",
         "--bf16", "True",
@@ -38,10 +41,7 @@ def create_medusa_heads(model_id: str):
         "--lazy_preprocess", "True",
         "--medusa_num_heads", "3",
         "--medusa_num_layers", "1",
-    ]
-    args = parser.parse_args(
-        ["medusa/medusa/train/train.py", "training_script_args"] + training_args
-    )
+    ])
     distributed_run.run(args)
 
     # Upload the medusa heads to the Hub