diff --git a/docs/cli/train.md b/docs/cli/train.md index e127e6af7..b4d4f718c 100644 --- a/docs/cli/train.md +++ b/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 diff --git a/scripts/train.py b/scripts/train.py index bb1bf51d8..ba83a3332 100644 --- a/scripts/train.py +++ b/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) diff --git a/tests/unit/train/test_cli_args.py b/tests/unit/train/test_cli_args.py index 328926fa7..38e30ad56 100644 --- a/tests/unit/train/test_cli_args.py +++ b/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 diff --git a/tests/unit/train/test_draft_config_init.py b/tests/unit/train/test_draft_config_init.py index 7bb749dad..35aeb33e4 100644 --- a/tests/unit/train/test_draft_config_init.py +++ b/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=[],