feat(train): default to Eagle 3.1 settings (llama arch + norms)

docs/cli/train.md

@@ -136,9 +136,9 @@ torchrun --standalone --nproc_per_node=4 scripts/train.py \
 
 - **`--embed-requires-grad` / `--no-embed-requires-grad`** (flag, default: `False`) Whether to train embedding layer weights.
 
-- **`--norm-before-fc`** (flag, default: `False`) Use RMSNorm before FC layer in draft path (e.g., for Eagle 3.1 / gpt-oss models).
+- **`--norm-before-fc` / `--no-norm-before-fc`** (flag, default: `True` for eagle3, `False` otherwise) Apply RMSNorm before the FC layer in the draft path.
 
-- **`--norm-output`** (flag, default: `False`) Feed post-norm hidden states back across TTT steps to stabilize magnitude drift across speculation depths (Eagle 3.1).
+- **`--norm-output` / `--no-norm-output`** (flag, default: `True` for eagle3, `False` otherwise) Feed post-norm hidden states back across TTT steps to stabilize magnitude drift across speculation depths.
 
 - **`--ttt-steps`** (int, default: `3`) Number of test-time training steps
 

scripts/train.py

@@ -800,9 +800,10 @@ def parse_args():
     parser.add_argument(
         "--draft-arch",
         type=str,
-        default="qwen3",
+        default=None,
         choices=list(DRAFT_ARCH_CONFIGS.keys()),
-        help="Architecture for draft decoder layers. Defaults to 'qwen3'.",
+        help="Architecture for draft decoder layers "
+        "(default: 'llama' for eagle3, 'qwen3' otherwise).",
     )
     parser.add_argument(
         "--draft-hidden-act",
@@ -913,17 +914,20 @@ def parse_args():
     )
     parser.add_argument(
         "--norm-before-fc",
-        action="store_true",
-        default=False,
-        help="Use RMSNorm before FC layer in draft path "
-        "(e.g., for Eagle 3.1 / gpt-oss models).",
+        action=argparse.BooleanOptionalAction,
+        default=None,
+        help="Apply RMSNorm before the FC layer in the draft path "
+        "(default: True for eagle3, False otherwise). "
+        "Disable with --no-norm-before-fc.",
     )
     parser.add_argument(
         "--norm-output",
-        action="store_true",
-        default=False,
+        action=argparse.BooleanOptionalAction,
+        default=None,
         help="Feed post-norm hidden states back across TTT steps to stabilize "
-        "magnitude drift across speculation depths (Eagle 3.1).",
+        "magnitude drift across speculation depths "
+        "(default: True for eagle3, False otherwise). "
+        "Disable with --no-norm-output.",
     )
     # D-Flash specific parameters
     parser.add_argument(
@@ -1100,6 +1104,15 @@ def parse_args():
     )
 
     args = parser.parse_args()
+
+    is_eagle3 = args.speculator_type == "eagle3"
+    if args.draft_arch is None:
+        args.draft_arch = "llama" if is_eagle3 else "qwen3"
+    if args.norm_before_fc is None:
+        args.norm_before_fc = is_eagle3
+    if args.norm_output is None:
+        args.norm_output = is_eagle3
+
     provided = explicitly_provided_dests(parser, DECODER_SHAPING_FLAGS)
     validate_draft_init_args(parser, args, provided)
     resolve_loss_config(args.loss_fn)

tests/unit/train/test_cli_args.py

@@ -124,3 +124,48 @@ def test_dspark_confidence_head_alpha(monkeypatch):
     train_kw, val_kw = DSparkDraftModel.get_trainer_kwargs(**vars(args))
     assert train_kw["confidence_head_alpha"] == 0.5
     assert val_kw["confidence_head_alpha"] == 0.5
+
+
+# ---------------------------------------------------------------------------
+# Per-speculator-type defaults for draft_arch, norm_before_fc, norm_output
+# ---------------------------------------------------------------------------
+
+
+def test_eagle3_defaults_to_llama_arch(monkeypatch):
+    args = _parse(monkeypatch, [])
+    assert args.draft_arch == "llama"
+
+
+def test_eagle3_defaults_norm_before_fc_true(monkeypatch):
+    args = _parse(monkeypatch, [])
+    assert args.norm_before_fc is True
+
+
+def test_eagle3_defaults_norm_output_true(monkeypatch):
+    args = _parse(monkeypatch, [])
+    assert args.norm_output is True
+
+
+def test_dflash_defaults_to_qwen3_arch(monkeypatch):
+    args = _parse(monkeypatch, ["--speculator-type", "dflash"])
+    assert args.draft_arch == "qwen3"
+
+
+def test_dflash_defaults_norm_before_fc_false(monkeypatch):
+    args = _parse(monkeypatch, ["--speculator-type", "dflash"])
+    assert args.norm_before_fc is False
+
+
+def test_dflash_defaults_norm_output_false(monkeypatch):
+    args = _parse(monkeypatch, ["--speculator-type", "dflash"])
+    assert args.norm_output is False
+
+
+def test_no_norm_before_fc_flag(monkeypatch):
+    args = _parse(monkeypatch, ["--no-norm-before-fc"])
+    assert args.norm_before_fc is False
+
+
+def test_no_norm_output_flag(monkeypatch):
+    args = _parse(monkeypatch, ["--no-norm-output"])
+    assert args.norm_output is False

tests/unit/train/test_draft_config_init.py

@@ -413,7 +413,7 @@ def _create_layer_config_for(verifier: SimpleNamespace):
         return create_transformer_layer_config(
             "target",
             num_layers=2,
-            draft_arch="qwen3",
+            draft_arch="llama",
             hidden_act=None,
             sliding_window=2048,
             sliding_window_indices=[],