From 53a99151f3e1540460016937851b9bc40dfae732 Mon Sep 17 00:00:00 2001 From: lit Date: Tue, 27 Jan 2026 19:45:45 -0800 Subject: [PATCH] mrope for mtp --- megatron/training/arguments.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 9dcc8f65e81..3108cd5dbe9 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1319,9 +1319,11 @@ def validate_args(args, defaults={}): if args.mtp_num_layers: assert not args.use_legacy_models, "The legacy Megatron models does not support Multi-Token Prediction (MTP)." - assert args.position_embedding_type == "rope" or args.position_embedding_type == "none", ( - f"Multi-Token Prediction (MTP) is not supported with {args.position_embedding_type} position embedding type." - + f"The supported position embedding types are rope and none." + # MTP is compatible with position embedding types that use position_ids. + supported_position_types = ["learned_absolute", "rope", "mrope", "none"] + assert args.position_embedding_type in supported_position_types, ( + f"Multi-Token Prediction (MTP) is not supported with '{args.position_embedding_type}' position embedding type. " + f"The supported position embedding types are: {', '.join(supported_position_types)}." ) if args.cpu_offloading_num_layers > 0: