diff --git a/miles/backends/experimental/fsdp_utils/actor.py b/miles/backends/experimental/fsdp_utils/actor.py index 47c2540f98d..5e878129cac 100644 --- a/miles/backends/experimental/fsdp_utils/actor.py +++ b/miles/backends/experimental/fsdp_utils/actor.py @@ -21,7 +21,7 @@ from miles.utils.tracking_utils import init_tracking from ....utils.profile_utils import TrainProfiler -from ...training_utils.ci_utils import check_grad_norm +from ...training_utils.ci_utils import assert_rollout_engine_weight_versions, check_grad_norm from ...training_utils.data import DataIterator, get_batch, get_data_iterator, get_rollout_data from ...training_utils.log_utils import ( aggregate_forward_results, @@ -569,7 +569,12 @@ def update_weights(self) -> None: # type: ignore[override] self.weight_updater.update_weights() - if self.args.ci_test and len(rollout_engines) > 0: + if getattr(self.args, "check_all_engine_weight_versions", False): + assert_rollout_engine_weight_versions( + rollout_engines, + self.weight_updater.weight_version, + ) + elif self.args.ci_test and len(rollout_engines) > 0: engine = random.choice(rollout_engines) engine_version = ray.get(engine.get_weight_version.remote()) if str(engine_version) != str(self.weight_updater.weight_version): diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index c138ea87460..08f771614aa 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -29,6 +29,7 @@ from ...utils.profile_utils import TrainProfiler from ...utils.tensor_backper import TensorBackuper from ..training_utils.cp_utils import slice_with_cp +from ..training_utils.ci_utils import assert_rollout_engine_weight_versions from ..training_utils.data import DataIterator, get_data_iterator, get_rollout_data, sync_actor_critic_data from ..training_utils.log_utils import log_cpu_memory, log_perf_data, log_rollout_data from ..training_utils.loss import compute_advantages_and_returns, get_log_probs_and_entropy, get_values @@ -1013,13 +1014,19 @@ def update_weights(self) -> None: self.weight_updater.update_weights() print_memory("after update_weights") - if self.args.ci_test and len(rollout_engines) > 0 and not is_lora_enabled(self.args): - engine = random.choice(rollout_engines) - engine_version = ray.get(engine.get_weight_version.remote()) - if str(engine_version) != str(self.weight_updater.weight_version): - raise RuntimeError( - f"Weight version mismatch! Engine: {engine_version}, Updater: {self.weight_updater.weight_version}" + if not is_lora_enabled(self.args): + if getattr(self.args, "check_all_engine_weight_versions", False): + assert_rollout_engine_weight_versions( + rollout_engines, + self.weight_updater.weight_version, ) + elif self.args.ci_test and len(rollout_engines) > 0: + engine = random.choice(rollout_engines) + engine_version = ray.get(engine.get_weight_version.remote()) + if str(engine_version) != str(self.weight_updater.weight_version): + raise RuntimeError( + f"Weight version mismatch! Engine: {engine_version}, Updater: {self.weight_updater.weight_version}" + ) if getattr(self.args, "keep_old_actor", False): if self.args.update_weights_interval == 1: diff --git a/miles/backends/megatron_utils/megatron_to_hf/xllm.py b/miles/backends/megatron_utils/megatron_to_hf/xllm.py index 76c340977fa..6a05ae51443 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/xllm.py +++ b/miles/backends/megatron_utils/megatron_to_hf/xllm.py @@ -3,8 +3,133 @@ import torch +def _is_mova(args) -> bool: + return getattr(args, "mova_num_value_experts", 0) > 0 + + +def _attention_geometry(args) -> tuple[int, int, int, int]: + hidden_size = args.hidden_size + num_attention_heads = args.num_attention_heads + num_query_groups = args.num_query_groups + kv_channels = getattr(args, "kv_channels", None) + head_dim = kv_channels if kv_channels is not None else hidden_size // num_attention_heads + + if num_attention_heads % num_query_groups: + raise ValueError( + f"num_attention_heads={num_attention_heads} must be divisible by " + f"num_query_groups={num_query_groups}" + ) + if head_dim <= 0 or head_dim % 2: + raise ValueError(f"xLLM MoVA requires an even positive head dimension, got {head_dim}") + return hidden_size, num_attention_heads, num_query_groups, head_dim + + +def _permute_qk_to_hf( + weight: torch.Tensor, + *, + num_heads: int, + head_dim: int, + hidden_size: int, + name: str, +) -> torch.Tensor: + """Convert MCore's adjacent-complex-pair Q/K rows to xLLM HF layout.""" + + expected_shape = (num_heads * head_dim, hidden_size) + if tuple(weight.shape) != expected_shape: + raise ValueError(f"Invalid {name} shape: got {tuple(weight.shape)}, expected {expected_shape}") + return ( + weight.reshape(num_heads, head_dim // 2, 2, hidden_size) + .transpose(1, 2) + .reshape(expected_shape) + .contiguous() + ) + + +def _unpack_grouped_attention_projection( + args, + name: str, + param: torch.Tensor, + *, + include_value: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Unpack MCore's per-GQA-group Q/G/K[/V] projection.""" + + hidden_size, num_attention_heads, num_query_groups, head_dim = _attention_geometry(args) + query_heads_per_group = num_attention_heads // num_query_groups + segment_heads = [query_heads_per_group, query_heads_per_group, 1] + if include_value: + segment_heads.append(1) + + expected_rows = num_query_groups * sum(segment_heads) * head_dim + if tuple(param.shape) != (expected_rows, hidden_size): + raise ValueError( + f"Invalid {name} shape: got {tuple(param.shape)}, expected " + f"{(expected_rows, hidden_size)}" + ) + + packed = param.reshape(num_query_groups, sum(segment_heads), head_dim, hidden_size) + chunks = torch.split(packed, segment_heads, dim=1) + query = chunks[0].reshape(num_attention_heads * head_dim, hidden_size) + gate = chunks[1].reshape(num_attention_heads * head_dim, hidden_size) + key = chunks[2].reshape(num_query_groups * head_dim, hidden_size) + value = chunks[3].reshape(num_query_groups * head_dim, hidden_size) if include_value else None + + query = _permute_qk_to_hf( + query, + num_heads=num_attention_heads, + head_dim=head_dim, + hidden_size=hidden_size, + name=f"{name}.query", + ) + key = _permute_qk_to_hf( + key, + num_heads=num_query_groups, + head_dim=head_dim, + hidden_size=hidden_size, + name=f"{name}.key", + ) + return query, gate, key, value + + +def _convert_grouped_value_experts( + args, layer_idx: str, name: str, param: torch.Tensor +) -> list[tuple[str, torch.Tensor]]: + """Convert gathered MCore Wv [expert, hidden, value] to HF [value, hidden].""" + + if param.ndim != 3: + raise ValueError( + f"Invalid {name} shape: got {tuple(param.shape)}, expected " + "[num_value_experts, hidden_size, value_width]" + ) + expected_experts = getattr(args, "mova_num_value_experts", 0) + if expected_experts and param.shape[0] != expected_experts: + raise ValueError( + f"Invalid {name} expert count: got {param.shape[0]}, expected {expected_experts}" + ) + if param.shape[1] != args.hidden_size: + raise ValueError( + f"Invalid {name} hidden dimension: got {param.shape[1]}, expected {args.hidden_size}. " + "The grouped MoVA weight must be gathered across regular attention TP before conversion." + ) + + return [ + ( + f"model.layers.{layer_idx}.self_attn.v_experts.{expert_idx}.weight", + expert_weight.transpose(0, 1).contiguous(), + ) + for expert_idx, expert_weight in enumerate(param.unbind(dim=0)) + ] + + def convert_xllm_to_hf(args, name, param): - """Convert Megatron parameter names/tensors to HuggingFace xLLM format.""" + """Convert Megatron parameter names/tensors to HuggingFace xLLM format. + + MoVA uses MCore's interleaved Q/G/K[/V] projections and a grouped value + weight stored as ``[expert, hidden / TP, value_width]``. The caller first + gathers those tensors over regular attention TP; this function then emits + canonical xLLM HF names consumed by both SGLang broadcast and P2P loading. + """ + if name == "module.module.embedding.word_embeddings.weight": return [("model.embed_tokens.weight", param)] if name == "module.module.output_layer.weight": @@ -12,11 +137,8 @@ def convert_xllm_to_hf(args, name, param): if name == "module.module.decoder.final_layernorm.weight": return [("model.norm.weight", param)] - try: - head_dim = args.kv_channels if args.kv_channels is not None else args.hidden_size // args.num_attention_heads - except AttributeError: - head_dim = args.hidden_size // args.num_attention_heads - value_num_per_group = args.num_attention_heads // args.num_query_groups + hidden_size, num_attention_heads, num_query_groups, head_dim = _attention_geometry(args) + query_heads_per_group = num_attention_heads // num_query_groups decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)" match = re.match(decoder_layers_pattern, name) @@ -53,18 +175,60 @@ def convert_xllm_to_hf(args, name, param): if rest == "self_attention.linear_proj.weight": return [(f"model.layers.{layer_idx}.self_attn.o_proj.weight", param)] + if rest == "self_attention.linear_qkv.weight": - param = param.view(args.num_query_groups, -1, head_dim, args.hidden_size) - q_param, k_param, v_param = torch.split(param, [value_num_per_group, 1, 1], dim=1) - q_param = q_param.reshape(-1, args.hidden_size) - k_param = k_param.reshape(-1, args.hidden_size) - v_param = v_param.reshape(-1, args.hidden_size) + if _is_mova(args): + query, gate, key, value = _unpack_grouped_attention_projection( + args, name, param, include_value=True + ) + assert value is not None + return [ + (f"model.layers.{layer_idx}.self_attn.q_proj.weight", query), + (f"model.layers.{layer_idx}.self_attn.attn_gate_proj.weight", gate), + (f"model.layers.{layer_idx}.self_attn.k_proj.weight", key), + (f"model.layers.{layer_idx}.self_attn.v_proj.weight", value), + ] + + # Preserve the legacy xLLM Q/K/V contract when MoVA is disabled. + packed = param.view(num_query_groups, -1, head_dim, hidden_size) + query, key, value = torch.split(packed, [query_heads_per_group, 1, 1], dim=1) return [ - (f"model.layers.{layer_idx}.self_attn.q_proj.weight", q_param), - (f"model.layers.{layer_idx}.self_attn.k_proj.weight", k_param), - (f"model.layers.{layer_idx}.self_attn.v_proj.weight", v_param), + (f"model.layers.{layer_idx}.self_attn.q_proj.weight", query.reshape(-1, hidden_size)), + (f"model.layers.{layer_idx}.self_attn.k_proj.weight", key.reshape(-1, hidden_size)), + (f"model.layers.{layer_idx}.self_attn.v_proj.weight", value.reshape(-1, hidden_size)), ] + if rest == "self_attention.linear_qkg.weight": + if not _is_mova(args): + raise ValueError(f"Found MoVA Q/K/gate projection while MoVA is disabled: {name}") + query, gate, key, value = _unpack_grouped_attention_projection( + args, name, param, include_value=False + ) + assert value is None + return [ + (f"model.layers.{layer_idx}.self_attn.q_proj.weight", query), + (f"model.layers.{layer_idx}.self_attn.attn_gate_proj.weight", gate), + (f"model.layers.{layer_idx}.self_attn.k_proj.weight", key), + ] + + if rest == "self_attention.value_projection.experts.weight": + if not _is_mova(args): + raise ValueError(f"Found grouped MoVA value experts while MoVA is disabled: {name}") + return _convert_grouped_value_experts(args, layer_idx, name, param) + + sequential_value_expert_pattern = ( + r"self_attention\.value_projection\.experts\.experts\.(\d+)\.weight" + ) + match = re.match(sequential_value_expert_pattern, rest) + if match: + expert_idx = match.group(1) + return [(f"model.layers.{layer_idx}.self_attn.v_experts.{expert_idx}.weight", param)] + + if rest == "self_attention.value_projection.router.weight": + return [(f"model.layers.{layer_idx}.self_attn.v_router.weight", param)] + if rest == "self_attention.value_projection.router.expert_bias": + return [(f"model.layers.{layer_idx}.self_attn.v_router.bias", param)] + if rest == "mlp.linear_fc1.weight": gate_weight, up_weight = param.chunk(2, dim=0) return [ @@ -74,7 +238,11 @@ def convert_xllm_to_hf(args, name, param): if rest == "mlp.linear_fc2.weight": return [(f"model.layers.{layer_idx}.mlp.down_proj.weight", param)] - if rest in ("self_attention.linear_qkv.layer_norm_weight", "input_layernorm.weight"): + if rest in ( + "self_attention.linear_qkv.layer_norm_weight", + "self_attention.linear_qkg.layer_norm_weight", + "input_layernorm.weight", + ): return [(f"model.layers.{layer_idx}.input_layernorm.weight", param)] if rest in ("mlp.linear_fc1.layer_norm_weight", "pre_mlp_layernorm.weight"): return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)] diff --git a/miles/backends/megatron_utils/model_provider.py b/miles/backends/megatron_utils/model_provider.py index 617068333e4..2c3b120fe06 100644 --- a/miles/backends/megatron_utils/model_provider.py +++ b/miles/backends/megatron_utils/model_provider.py @@ -23,6 +23,24 @@ logger = logging.getLogger(__name__) +def _get_mova_model_components(): + """Import MoVA components only when the architecture is requested. + + Keeping this import lazy preserves compatibility with ordinary Megatron + installations that have not yet taken the native MoVA extension. + """ + + try: + from megatron.core.models.gpt.mova_layer_specs import get_mova_gpt_decoder_block_spec + from megatron.core.transformer.mova import MoVATransformerConfig + except ImportError as error: + raise RuntimeError( + "Native MoVA was requested, but this Megatron installation does not " + "provide MoVATransformerConfig/get_mova_gpt_decoder_block_spec" + ) from error + return MoVATransformerConfig, get_mova_gpt_decoder_block_spec + + # Adapt from https://github.com/volcengine/verl/blob/c3b20575d2bc815fcccd84bddb4c0401fc4b632b/verl/models/llama/megatron/layers/parallel_linear.py#L82 class LinearForLastLayer(torch.nn.Linear): def __init__( @@ -59,6 +77,14 @@ def get_model_provider_func( args: argparse.Namespace, role: Literal["actor", "critic"] = "actor", ): + is_mova = getattr(args, "mova_num_value_experts", 0) > 0 + if is_mova: + # Validate before any provider branch so MoVA can never silently build a + # custom/ordinary provider or select an unsupported converter. + from miles.utils.arguments import validate_mova_args + + validate_mova_args(args) + # Support custom model provider path (similar to --custom-rm-path for reward models) if getattr(args, "custom_model_provider_path", None): @@ -147,9 +173,21 @@ def model_provider( # Experimental loading arguments from yaml assert config is None, "miles builds the config from args, so it expects config to be None" - config = core_transformer_config_from_args(args) - - if args.spec is not None: + if is_mova: + mova_config_class, get_mova_block_spec = _get_mova_model_components() + config = core_transformer_config_from_args(args, mova_config_class) + else: + config = core_transformer_config_from_args(args) + + if is_mova: + transformer_layer_spec = get_mova_block_spec( + config, + use_transformer_engine=use_te, + moe_grouped_gemm=args.moe_grouped_gemm, + moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm, + vp_stage=vp_stage, + ) + elif args.spec is not None: transformer_layer_spec = import_module(args.spec) # Allow the spec to be a function so that user can use customized Megatron easier. if callable(transformer_layer_spec): diff --git a/miles/backends/megatron_utils/update_weight/common.py b/miles/backends/megatron_utils/update_weight/common.py index c61f2648766..2f2621f893b 100644 --- a/miles/backends/megatron_utils/update_weight/common.py +++ b/miles/backends/megatron_utils/update_weight/common.py @@ -17,6 +17,38 @@ logger = logging.getLogger(__name__) +def is_ffn_expert_parameter(name: str) -> bool: + """Return whether ``name`` belongs to the EP-sharded feed-forward experts. + + MoVA value projections also contain an ``experts`` path component, but + they are replicated over FFN expert parallelism and sharded over regular + attention tensor parallelism. Keep the predicate deliberately tied to the + MLP path so the two expert systems cannot be mixed accidentally. + """ + + return ".mlp.experts." in name + + +def is_mova_model(args: Namespace) -> bool: + return getattr(args, "mova_num_value_experts", 0) > 0 + + +def validate_weight_update_cache_mode(args: Namespace) -> None: + """Reject cache-preserving updates for MoVA. + + MoVA caches the routed value, which depends on both router and value-expert + weights. Reusing a KV cache after a policy update therefore mixes two model + versions even when keys are unchanged. + """ + + if is_mova_model(args) and getattr(args, "pause_generation_mode", None) == "in_place": + raise ValueError( + "MoVA weight updates do not support pause_generation_mode='in_place': " + "the routed-value KV cache must be flushed after every policy update. " + "Use 'retract' (recommended) or 'abort'." + ) + + def _gather_with_stride( param_partitions: list[torch.Tensor], partition_dim: int, partition_stride: int ) -> torch.Tensor: @@ -50,7 +82,8 @@ def _check_and_fix_partition(args: Namespace, name: str, partition_stride: int, def all_gather_param(args: Namespace, name: str, param: torch.nn.Parameter) -> torch.Tensor: """ All-gather TP-sharded param to full tensor. expert_bias→param, non-TP/duplicated→param.data. - Uses expert-TP for ".experts.", else regular-TP. Handles strided partitioning via partition_stride. + Uses expert-TP for FFN ``.mlp.experts.`` tensors, else regular-TP. Handles + strided partitioning via ``partition_stride``. """ if "expert_bias" in name: return param @@ -59,7 +92,7 @@ def all_gather_param(args: Namespace, name: str, param: torch.nn.Parameter) -> t if not param.tensor_model_parallel or getattr(param, "parallel_mode", None) == "duplicated": return param.data - if ".experts." in name: + if is_ffn_expert_parameter(name): tp_size = mpu.get_expert_tensor_parallel_world_size() tp_group = mpu.get_expert_tensor_parallel_group() else: @@ -99,7 +132,7 @@ def all_gather_params_async( handles.append(None) else: # Start async all_gather - if ".experts." in info.name: + if is_ffn_expert_parameter(info.name): tp_size = mpu.get_expert_tensor_parallel_world_size() tp_group = mpu.get_expert_tensor_parallel_group() else: @@ -270,7 +303,7 @@ def collect_named_tensors_for_weight_transfer( convert_to_global_name, translate_gpu_to_cpu, ): - if is_expert == (".experts." in name): + if is_expert == is_ffn_expert_parameter(name): yield name, tensor diff --git a/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py b/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py index ecdba3c8c77..152b212e270 100644 --- a/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py +++ b/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py @@ -12,7 +12,7 @@ from ..megatron_to_hf import convert_to_hf from ..sglang import monkey_patch_torch_reductions -from .common import all_gather_params_async, named_params_and_buffers +from .common import all_gather_params_async, is_ffn_expert_parameter, named_params_and_buffers from .hf_weight_iterator_base import HfWeightIteratorBase @@ -83,7 +83,7 @@ def _get_megatron_full_params( if ep_size > 1: handles = [] for info, param in zip(megatron_local_param_infos, params, strict=False): - if ".experts." in info.name: + if is_ffn_expert_parameter(info.name): src_rank = ( info.src_rank if info.src_rank in dist.get_process_group_ranks(mpu.get_expert_model_parallel_group()) @@ -118,7 +118,7 @@ def _get_megatron_local_param_info_buckets(args: Namespace, model: Sequence[torc for info in param_infos: # Expert params use expert-TP size, others use regular-TP size - if ".experts." in info.name: + if is_ffn_expert_parameter(info.name): tp_size = mpu.get_expert_tensor_parallel_world_size() else: tp_size = mpu.get_tensor_model_parallel_world_size() diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py index 6707993822f..85b931912d8 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py @@ -10,7 +10,12 @@ from miles.utils.distributed_utils import get_gloo_group from ...megatron_to_hf import convert_to_hf -from ..common import all_gather_param, collect_named_tensors_for_weight_transfer, post_process_weights +from ..common import ( + all_gather_param, + collect_named_tensors_for_weight_transfer, + post_process_weights, + validate_weight_update_cache_mode, +) class DistBucketedWeightUpdateMixin: @@ -138,6 +143,9 @@ def _update_expert_bucket_weights( def _pause_and_prepare_engines(self) -> None: """Pause rollout engines, flush cache, and run pre-process if needed.""" + # Validate on every training rank before rank 0 performs RPCs. Raising + # only on rank 0 would strand the other ranks at the following barrier. + validate_weight_update_cache_mode(self.args) if dist.get_rank() == 0: mode = self.args.pause_generation_mode ray.get([engine.pause_generation.remote(mode=mode) for engine in self.rollout_engines]) diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/p2p_transfer_utils.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/p2p_transfer_utils.py index 804d73ce043..5d75a35b1c9 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/p2p_transfer_utils.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/p2p_transfer_utils.py @@ -172,14 +172,32 @@ def submit_returning_future(self, fn: Callable, *args) -> torch.Future: return future def wait_transfers(self) -> None: - """Wait for all submitted tasks to complete.""" - for future in self.transfer_futures: - try: - future.result(timeout=self.transfer_timeout) - except Exception as e: - logger.error(f"[P2P] Transfer future failed: {e}") - - self.transfer_futures.clear() + """Drain submitted tasks and fail the update if any transfer failed.""" + failures: list[tuple[int, Exception]] = [] + futures = tuple(self.transfer_futures) + try: + for index, future in enumerate(futures): + try: + future.result(timeout=self.transfer_timeout) + except Exception as error: + failures.append((index, error)) + logger.error( + "[P2P] Transfer task %d/%d failed: %s", + index + 1, + len(futures), + error, + ) + finally: + self.transfer_futures.clear() + + if failures: + details = "; ".join( + f"task {index + 1}: {type(error).__name__}: {error}" for index, error in failures + ) + first_error = failures[0][1] + raise RuntimeError( + f"{len(failures)} of {len(futures)} P2P weight transfers failed ({details})" + ) from first_error def create_server_args_from_dict(data_dict: dict) -> ServerArgs: diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py index 86073d5ab3b..0cc74d94347 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py @@ -15,7 +15,7 @@ from miles.utils.distributed_utils import get_gloo_group from ..sglang import FlattenedTensorBucket, MultiprocessingSerializer -from .common import post_process_weights +from .common import post_process_weights, validate_weight_update_cache_mode from .hf_weight_iterator_base import HfWeightIteratorBase from .update_weight_from_distributed.broadcast import ( connect_rollout_engines_from_distributed, @@ -172,6 +172,9 @@ def update_weights(self) -> None: """ version++, flush caches, process buckets. Progress on rank 0. """ + # Validate on every rank before any RPC or collective. MoVA caches + # routed values, so retaining KV state across a policy update is stale. + validate_weight_update_cache_mode(self.args) self.weight_version += 1 rank = dist.get_rank() diff --git a/miles/backends/training_utils/ci_utils.py b/miles/backends/training_utils/ci_utils.py index 9124889beda..3bd4eaf0775 100644 --- a/miles/backends/training_utils/ci_utils.py +++ b/miles/backends/training_utils/ci_utils.py @@ -9,6 +9,31 @@ logger = logging.getLogger(__name__) +def assert_rollout_engine_weight_versions(rollout_engines, expected_version) -> None: + """Assert that every rollout engine exposes the synchronized version. + + This is intentionally metadata-only: it is cheap enough to enable for + acceptance runs and avoids the cost of full tensor equality checks. + """ + + if not rollout_engines: + return + + import ray + + versions = ray.get([engine.get_weight_version.remote() for engine in rollout_engines]) + mismatches = [ + (index, version) + for index, version in enumerate(versions) + if str(version) != str(expected_version) + ] + if mismatches: + raise RuntimeError( + "Rollout engine weight-version mismatch: " + f"expected {expected_version}, mismatches {mismatches}" + ) + + def check_kl(args: Namespace, log_dict: dict[str, float], step_id: int, accumulated_step_id: int) -> None: if step_id == 0 and "train/ppo_kl" in log_dict and "train/pg_clipfrac" in log_dict: if args.multi_latent_attention: diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index a9af7d4fb7b..ca683c17af2 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -35,6 +35,115 @@ def reset_arg(parser, name, **kwargs): parser.add_argument(name, **kwargs) +def add_mova_arguments(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: + """Register the native Megatron MoVA architecture arguments. + + These names and defaults intentionally mirror Megatron's ``pretrain_mova`` + entry point. Miles owns the registration because its train entry points do + not invoke that script's ``extra_args_provider``. + + Unlike ``pretrain_mova``, this function must not change shared Transformer + defaults globally: the same Miles parser also serves every non-MoVA model. + Exact MoVA runs therefore pass the required zero-dropout/base-attention + flags explicitly and are checked by :func:`validate_mova_args`. + """ + + group = parser.add_argument_group(title="mixture-of-value attention") + group.add_argument("--mova-num-value-experts", type=int, default=0) + group.add_argument("--mova-router-topk", type=int, default=1) + group.add_argument( + "--mova-router-score-function", + choices=("softmax", "sigmoid"), + default="sigmoid", + ) + group.add_argument("--mova-router-topk-scaling-factor", type=float, default=1.0) + group.add_argument("--mova-router-enable-expert-bias", action="store_true") + group.add_argument("--mova-router-bias-update-rate", type=float, default=1.0e-3) + group.add_argument("--mova-router-aux-loss-coeff", type=float, default=0.0) + group.add_argument( + "--mova-router-load-balancing-type", + choices=("none", "aux_loss"), + default="none", + ) + group.add_argument("--mova-num-dense-layers", type=int, default=0) + group.add_argument("--mova-norm-num-groups", type=int, default=1) + group.add_argument( + "--mova-attention-gate-function", + choices=("softplus", "silu"), + default="softplus", + ) + group.add_argument( + "--mova-value-backend", + choices=("sequential", "grouped_gemm"), + default="grouped_gemm", + ) + group.add_argument( + "--mova-use-torch-rms-norm", + action=argparse.BooleanOptionalAction, + default=False, + help="Use PyTorch's native RMSNorm primitive for grouped normalization.", + ) + group.add_argument( + "--xllm-router-compatibility", + action=argparse.BooleanOptionalAction, + default=True, + help="Match xLLM's BF16 router GEMM followed by FP32 top-k scoring.", + ) + group.add_argument( + "--xllm-router-gemm-partitions", + type=int, + default=1, + help="Original xLLM model-parallel partitions used by router GEMMs.", + ) + return parser + + +def validate_mova_args(args) -> None: + """Validate the Miles-specific boundary for native Megatron MoVA. + + Shape/routing invariants remain owned by ``MoVATransformerConfig``. The + checks here prevent selecting an incompatible Miles provider, conversion + path, or cache/model feature before distributed model construction starts. + """ + + if getattr(args, "mova_num_value_experts", 0) <= 0: + return + + if args.train_backend != "megatron": + raise ValueError("MoVA training requires --train-backend megatron") + if args.megatron_to_hf_mode != "raw": + raise ValueError("Native MoVA requires --megatron-to-hf-mode raw") + if getattr(args, "custom_model_provider_path", None) is not None: + raise ValueError("Native MoVA owns its model provider; remove --custom-model-provider-path") + if getattr(args, "use_legacy_models", False): + raise ValueError("MoVA is supported only by Megatron Core models") + if getattr(args, "yaml_cfg", None) is not None: + raise ValueError("MoVA's initial Miles integration supports CLI configuration only") + if getattr(args, "spec", None) is not None: + raise ValueError("MoVA owns its heterogeneous layer spec; --spec is not supported") + if getattr(args, "mtp_num_layers", None) is not None: + raise ValueError("MoVA does not currently support MTP layers") + if getattr(args, "multi_latent_attention", False): + raise ValueError("MoVA cannot be combined with multi-latent attention") + if getattr(args, "heterogeneous_layers_config_path", None) is not None: + raise ValueError("MoVA owns its heterogeneous layer layout") + if not getattr(args, "attention_output_gate", False): + raise ValueError("MoVA requires --attention-output-gate") + if not getattr(args, "rotary_interleaved", False): + raise ValueError("xLLM MoVA requires --rotary-interleaved") + if not getattr(args, "group_query_attention", False): + raise ValueError("MoVA requires --group-query-attention") + if getattr(args, "num_experts", 0) <= 0: + raise ValueError("MoVA sparse layers require --num-experts") + if getattr(args, "attention_dropout", None) != 0.0 or getattr(args, "hidden_dropout", None) != 0.0: + raise ValueError("xLLM MoVA requires --attention-dropout 0 and --hidden-dropout 0") + if ( + getattr(args, "tensor_model_parallel_size", 1) > 1 + and not getattr(args, "sequence_parallel", False) + ): + raise ValueError("MoVA with tensor parallelism requires --sequence-parallel") + + def get_miles_extra_args_provider(add_custom_arguments=None): def add_miles_arguments(parser): # Ray @@ -1373,6 +1482,14 @@ def add_debug_arguments(parser): default="torch", ) parser.add_argument("--check-weight-update-equal", action="store_true") + parser.add_argument( + "--check-all-engine-weight-versions", + action="store_true", + help=( + "After every live-weight update, query every rollout engine " + "and require its version to match the updater." + ), + ) parser.add_argument( "--env-report", type=str, @@ -1682,6 +1799,7 @@ def add_sglang_tp_size(): parser = add_cluster_arguments(parser) parser = add_train_arguments(parser) + parser = add_mova_arguments(parser) parser = add_rollout_arguments(parser) parser = add_fault_tolerance_arguments(parser) parser = add_data_arguments(parser) @@ -1854,6 +1972,7 @@ def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]: def miles_validate_args(args): + validate_mova_args(args) args.eval_datasets = _resolve_eval_datasets(args) # Normalize --tito-allowed-append-roles: lowercase + deduplicate. diff --git a/tests/fast/backends/megatron_utils/test_mova_model_provider.py b/tests/fast/backends/megatron_utils/test_mova_model_provider.py new file mode 100644 index 00000000000..cfecc01d6d3 --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_mova_model_provider.py @@ -0,0 +1,193 @@ +from argparse import Namespace +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from miles.backends.megatron_utils import model_provider as model_provider_module + + +def _provider_args(*, mova_num_value_experts: int) -> Namespace: + return Namespace( + custom_model_provider_path=None, + mova_num_value_experts=mova_num_value_experts, + train_backend="megatron", + megatron_to_hf_mode="raw", + use_legacy_models=False, + yaml_cfg=None, + spec=None, + mtp_num_layers=None, + multi_latent_attention=False, + heterogeneous_layers_config_path=None, + attention_output_gate=True, + rotary_interleaved=True, + group_query_attention=True, + num_experts=100 if mova_num_value_experts else 0, + attention_dropout=0.0, + hidden_dropout=0.0, + tensor_model_parallel_size=8, + sequence_parallel=True, + transformer_impl="transformer_engine", + moe_grouped_gemm=True, + moe_use_legacy_grouped_gemm=False, + qk_layernorm=False, + fp8_param_gather=False, + padded_vocab_size=256, + max_position_embeddings=32768, + fp16_lm_cross_entropy=False, + untie_embeddings_and_output_weights=True, + position_embedding_type="rope", + rotary_percent=1.0, + rotary_base=10_000_000, + use_rope_scaling=False, + use_rollout_routing_replay=False, + ) + + +def test_native_mova_provider_uses_mova_config_and_heterogeneous_spec(monkeypatch): + args = _provider_args(mova_num_value_experts=64) + fake_config_class = type("FakeMoVATransformerConfig", (), {}) + fake_config = SimpleNamespace(hidden_size=2560) + config_builder = MagicMock(return_value=fake_config) + spec_builder = MagicMock(return_value="mova-block-spec") + gpt_model = MagicMock(return_value=SimpleNamespace(config=fake_config)) + + monkeypatch.setattr(model_provider_module, "core_transformer_config_from_args", config_builder) + monkeypatch.setattr( + model_provider_module, + "_get_mova_model_components", + lambda: (fake_config_class, spec_builder), + ) + monkeypatch.setattr(model_provider_module, "GPTModel", gpt_model) + + provider = model_provider_module.get_model_provider_func(args, role="actor") + model = provider(pre_process=False, post_process=True, vp_stage=2) + + assert model.config is fake_config + config_builder.assert_called_once_with(args, fake_config_class) + spec_builder.assert_called_once_with( + fake_config, + use_transformer_engine=True, + moe_grouped_gemm=True, + moe_use_legacy_grouped_gemm=False, + vp_stage=2, + ) + gpt_model.assert_called_once_with( + config=fake_config, + transformer_layer_spec="mova-block-spec", + vocab_size=256, + max_sequence_length=32768, + pre_process=False, + post_process=True, + fp16_lm_cross_entropy=False, + parallel_output=True, + share_embeddings_and_output_weights=False, + position_embedding_type="rope", + rotary_percent=1.0, + rotary_base=10_000_000, + rope_scaling=False, + vp_stage=2, + ) + + +@pytest.mark.parametrize("role", ["actor", "critic"]) +def test_actor_and_critic_share_native_mova_provider(monkeypatch, role): + args = _provider_args(mova_num_value_experts=64) + fake_config_class = type("FakeMoVATransformerConfig", (), {}) + fake_config = SimpleNamespace(hidden_size=2560, sequence_parallel=True) + spec_builder = MagicMock(return_value="mova-block-spec") + model = SimpleNamespace(config=fake_config) + model.output_layer = "original-output-layer" + + monkeypatch.setattr( + model_provider_module, + "core_transformer_config_from_args", + MagicMock(return_value=fake_config), + ) + monkeypatch.setattr( + model_provider_module, + "_get_mova_model_components", + lambda: (fake_config_class, spec_builder), + ) + monkeypatch.setattr(model_provider_module, "GPTModel", MagicMock(return_value=model)) + critic_head = MagicMock(return_value="critic-output-layer") + monkeypatch.setattr(model_provider_module, "LinearForLastLayer", critic_head) + + provider = model_provider_module.get_model_provider_func(args, role=role) + result = provider() + + assert result is model + spec_builder.assert_called_once() + if role == "critic": + critic_head.assert_called_once_with( + input_size=2560, + output_size=1, + config=fake_config, + ) + assert model.output_layer == "critic-output-layer" + else: + critic_head.assert_not_called() + assert model.output_layer == "original-output-layer" + + +def test_non_mova_provider_preserves_standard_transformer_path(monkeypatch): + args = _provider_args(mova_num_value_experts=0) + fake_config = SimpleNamespace(hidden_size=2560) + config_builder = MagicMock(return_value=fake_config) + te_spec_builder = MagicMock(return_value="standard-te-layer-spec") + gpt_model = MagicMock(return_value=SimpleNamespace(config=fake_config)) + mova_components = MagicMock(side_effect=AssertionError("legacy path imported MoVA")) + + monkeypatch.setattr(model_provider_module, "core_transformer_config_from_args", config_builder) + monkeypatch.setattr(model_provider_module, "get_gpt_layer_with_transformer_engine_spec", te_spec_builder) + monkeypatch.setattr(model_provider_module, "_get_mova_model_components", mova_components) + monkeypatch.setattr(model_provider_module, "GPTModel", gpt_model) + + provider = model_provider_module.get_model_provider_func(args, role="actor") + provider(pre_process=True, post_process=False) + + config_builder.assert_called_once_with(args) + mova_components.assert_not_called() + te_spec_builder.assert_called_once_with( + num_experts=0, + moe_grouped_gemm=True, + qk_layernorm=False, + multi_latent_attention=False, + moe_use_legacy_grouped_gemm=False, + ) + assert gpt_model.call_args.kwargs["transformer_layer_spec"] == "standard-te-layer-spec" + assert "vp_stage" not in gpt_model.call_args.kwargs + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("megatron_to_hf_mode", "bridge", "raw"), + ("custom_model_provider_path", "some.module.provider", "owns its model provider"), + ("spec", "some.module.spec", "heterogeneous layer spec"), + ("mtp_num_layers", 1, "MTP"), + ], +) +def test_mova_provider_fails_before_selecting_incompatible_provider(field, value, message): + args = _provider_args(mova_num_value_experts=64) + setattr(args, field, value) + + with pytest.raises(ValueError, match=message): + model_provider_module.get_model_provider_func(args) + + +def test_mova_components_report_incompatible_megatron_cleanly(monkeypatch): + real_import = __import__ + + def reject_mova(name, *args, **kwargs): + if name in { + "megatron.core.models.gpt.mova_layer_specs", + "megatron.core.transformer.mova", + }: + raise ImportError("no native MoVA") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr("builtins.__import__", reject_mova) + + with pytest.raises(RuntimeError, match="does not provide MoVATransformerConfig"): + model_provider_module._get_mova_model_components() diff --git a/tests/fast/backends/megatron_utils/test_p2p_transfer_manager.py b/tests/fast/backends/megatron_utils/test_p2p_transfer_manager.py new file mode 100644 index 00000000000..a01dcc3627d --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_p2p_transfer_manager.py @@ -0,0 +1,66 @@ +from concurrent.futures import Future +from unittest.mock import MagicMock + +import pytest + +from miles.backends.megatron_utils.update_weight.update_weight_from_distributed.p2p_transfer_utils import ( + P2PTransferManager, +) + + +def _completed_future() -> Future: + future = Future() + future.set_result(None) + return future + + +def test_wait_transfers_drains_successful_tasks_and_clears_queue() -> None: + manager = P2PTransferManager(transfer_timeout=7.0) + first = MagicMock(wraps=_completed_future()) + second = MagicMock(wraps=_completed_future()) + manager.transfer_futures = [first, second] + + manager.wait_transfers() + + first.result.assert_called_once_with(timeout=7.0) + second.result.assert_called_once_with(timeout=7.0) + assert manager.transfer_futures == [] + + +def test_wait_transfers_drains_all_tasks_then_raises_aggregate_error() -> None: + manager = P2PTransferManager(transfer_timeout=3.0) + first = MagicMock() + second = MagicMock() + third = MagicMock() + first.result.side_effect = RuntimeError("session-a failed") + second.result.return_value = None + third.result.side_effect = ValueError("session-c failed") + manager.transfer_futures = [first, second, third] + + with pytest.raises(RuntimeError, match=r"2 of 3 P2P weight transfers failed") as exc_info: + manager.wait_transfers() + + first.result.assert_called_once_with(timeout=3.0) + second.result.assert_called_once_with(timeout=3.0) + third.result.assert_called_once_with(timeout=3.0) + assert "session-a failed" in str(exc_info.value) + assert "session-c failed" in str(exc_info.value) + assert isinstance(exc_info.value.__cause__, RuntimeError) + assert manager.transfer_futures == [] + + +def test_wait_transfers_propagates_timeout_and_clears_queue() -> None: + manager = P2PTransferManager(transfer_timeout=0.25) + timed_out = MagicMock() + completed = MagicMock() + timed_out.result.side_effect = TimeoutError("transfer timed out") + completed.result.return_value = None + manager.transfer_futures = [timed_out, completed] + + with pytest.raises(RuntimeError, match=r"1 of 2 P2P weight transfers failed") as exc_info: + manager.wait_transfers() + + timed_out.result.assert_called_once_with(timeout=0.25) + completed.result.assert_called_once_with(timeout=0.25) + assert isinstance(exc_info.value.__cause__, TimeoutError) + assert manager.transfer_futures == [] diff --git a/tests/fast/backends/megatron_utils/test_xllm_mova_weight_sync.py b/tests/fast/backends/megatron_utils/test_xllm_mova_weight_sync.py new file mode 100644 index 00000000000..438e4edd9a9 --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_xllm_mova_weight_sync.py @@ -0,0 +1,426 @@ +from argparse import Namespace +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch +from torch import nn + +from miles.backends.megatron_utils.megatron_to_hf.xllm import convert_xllm_to_hf +from miles.backends.megatron_utils.update_weight.common import ( + all_gather_param, + collect_named_tensors_for_weight_transfer, + is_ffn_expert_parameter, + named_params_and_buffers, + validate_weight_update_cache_mode, +) + + +def _args(*, num_value_experts: int = 3, pause_generation_mode: str = "retract") -> Namespace: + return Namespace( + hidden_size=8, + num_attention_heads=4, + num_query_groups=2, + kv_channels=2, + mova_num_value_experts=num_value_experts, + pause_generation_mode=pause_generation_mode, + ) + + +def _projection(rows: int, hidden_size: int, offset: int) -> torch.Tensor: + return torch.arange( + offset, + offset + rows * hidden_size, + dtype=torch.float32, + ).reshape(rows, hidden_size) + + +def _pack_by_query_group(*projections: torch.Tensor, args: Namespace) -> torch.Tensor: + query_heads_per_group = args.num_attention_heads // args.num_query_groups + head_dim = args.kv_channels + grouped = [] + for index, projection in enumerate(projections): + heads_per_group = query_heads_per_group if index < 2 else 1 + grouped.append( + projection.reshape( + args.num_query_groups, + heads_per_group, + head_dim, + args.hidden_size, + ) + ) + return torch.cat(grouped, dim=1).reshape(-1, args.hidden_size) + + +def _permute_qk_to_hf(weight: torch.Tensor, num_heads: int, args: Namespace) -> torch.Tensor: + return ( + weight.reshape(num_heads, args.kv_channels // 2, 2, args.hidden_size) + .transpose(1, 2) + .reshape_as(weight) + ) + + +def _restore_native_qk(weight: torch.Tensor, num_heads: int, args: Namespace) -> torch.Tensor: + """Independent inverse used by the Megatron HF checkpoint reader.""" + + return ( + weight.reshape(num_heads, 2, args.kv_channels // 2, args.hidden_size) + .transpose(1, 2) + .reshape_as(weight) + ) + + +def test_dense_mova_qkgv_conversion_preserves_gate_and_value() -> None: + args = _args() + query = _projection(8, args.hidden_size, 0) + gate = _projection(8, args.hidden_size, 1_000) + key = _projection(4, args.hidden_size, 2_000) + value = _projection(4, args.hidden_size, 3_000) + packed = _pack_by_query_group(query, gate, key, value, args=args) + + converted = dict( + convert_xllm_to_hf( + args, + "module.module.decoder.layers.0.self_attention.linear_qkv.weight", + packed, + ) + ) + + torch.testing.assert_close( + converted["model.layers.0.self_attn.q_proj.weight"], + _permute_qk_to_hf(query, args.num_attention_heads, args), + ) + torch.testing.assert_close( + _restore_native_qk( + converted["model.layers.0.self_attn.q_proj.weight"], + args.num_attention_heads, + args, + ), + query, + ) + torch.testing.assert_close(converted["model.layers.0.self_attn.attn_gate_proj.weight"], gate) + torch.testing.assert_close( + converted["model.layers.0.self_attn.k_proj.weight"], + _permute_qk_to_hf(key, args.num_query_groups, args), + ) + torch.testing.assert_close( + _restore_native_qk( + converted["model.layers.0.self_attn.k_proj.weight"], + args.num_query_groups, + args, + ), + key, + ) + torch.testing.assert_close(converted["model.layers.0.self_attn.v_proj.weight"], value) + + +def test_sparse_mova_qkg_conversion_emits_no_dense_value() -> None: + args = _args() + query = _projection(8, args.hidden_size, 0) + gate = _projection(8, args.hidden_size, 1_000) + key = _projection(4, args.hidden_size, 2_000) + packed = _pack_by_query_group(query, gate, key, args=args) + + converted = dict( + convert_xllm_to_hf( + args, + "module.module.decoder.layers.3.self_attention.linear_qkg.weight", + packed, + ) + ) + + assert set(converted) == { + "model.layers.3.self_attn.q_proj.weight", + "model.layers.3.self_attn.k_proj.weight", + "model.layers.3.self_attn.attn_gate_proj.weight", + } + torch.testing.assert_close( + converted["model.layers.3.self_attn.q_proj.weight"], + _permute_qk_to_hf(query, args.num_attention_heads, args), + ) + torch.testing.assert_close(converted["model.layers.3.self_attn.attn_gate_proj.weight"], gate) + torch.testing.assert_close( + converted["model.layers.3.self_attn.k_proj.weight"], + _permute_qk_to_hf(key, args.num_query_groups, args), + ) + + +def test_grouped_value_experts_convert_input_to_output_sharding_layout() -> None: + args = _args() + value_width = args.num_query_groups * args.kv_channels + # MCore all-gather result: [expert, full hidden input, full value output]. + grouped_weight = torch.arange( + args.mova_num_value_experts * args.hidden_size * value_width, + dtype=torch.float32, + ).reshape(args.mova_num_value_experts, args.hidden_size, value_width) + + converted = convert_xllm_to_hf( + args, + "module.module.decoder.layers.3.self_attention.value_projection.experts.weight", + grouped_weight, + ) + + assert [name for name, _ in converted] == [ + f"model.layers.3.self_attn.v_experts.{expert}.weight" + for expert in range(args.mova_num_value_experts) + ] + for expert, (_, weight) in enumerate(converted): + assert weight.is_contiguous() + torch.testing.assert_close(weight, grouped_weight[expert].transpose(0, 1)) + + +def test_grouped_value_experts_require_regular_tp_gather_first() -> None: + args = _args() + local_input_shard = torch.empty(args.mova_num_value_experts, args.hidden_size // 2, 4) + with pytest.raises(ValueError, match="gathered across regular attention TP"): + convert_xllm_to_hf( + args, + "module.module.decoder.layers.3.self_attention.value_projection.experts.weight", + local_input_shard, + ) + + +@pytest.mark.parametrize( + ("megatron_suffix", "hf_suffix"), + [ + ("router.weight", "v_router.weight"), + ("router.expert_bias", "v_router.bias"), + ], +) +def test_value_router_parameter_and_buffer_are_both_synchronized( + megatron_suffix: str, hf_suffix: str +) -> None: + args = _args() + tensor = torch.randn(args.mova_num_value_experts, args.hidden_size) + if megatron_suffix.endswith("expert_bias"): + tensor = torch.randn(args.mova_num_value_experts) + + converted = convert_xllm_to_hf( + args, + f"module.module.decoder.layers.3.self_attention.value_projection.{megatron_suffix}", + tensor, + ) + assert len(converted) == 1 + assert converted[0][0] == f"model.layers.3.self_attn.{hf_suffix}" + assert converted[0][1] is tensor + + +def test_legacy_xllm_qkv_conversion_is_unchanged_when_mova_is_disabled() -> None: + args = _args(num_value_experts=0) + query = _projection(8, args.hidden_size, 0) + key = _projection(4, args.hidden_size, 1_000) + value = _projection(4, args.hidden_size, 2_000) + # Legacy Q/K/V has Q first, so construct its own three-way group packing. + query_grouped = query.reshape(2, 2, 2, 8) + key_grouped = key.reshape(2, 1, 2, 8) + value_grouped = value.reshape(2, 1, 2, 8) + packed = torch.cat((query_grouped, key_grouped, value_grouped), dim=1).reshape(-1, 8) + + converted = dict( + convert_xllm_to_hf( + args, + "module.module.decoder.layers.0.self_attention.linear_qkv.weight", + packed, + ) + ) + + torch.testing.assert_close(converted["model.layers.0.self_attn.q_proj.weight"], query) + torch.testing.assert_close(converted["model.layers.0.self_attn.k_proj.weight"], key) + torch.testing.assert_close(converted["model.layers.0.self_attn.v_proj.weight"], value) + assert "model.layers.0.self_attn.attn_gate_proj.weight" not in converted + + +@pytest.mark.parametrize( + ("name", "expected"), + [ + ("module.module.decoder.layers.3.mlp.experts.linear_fc1.weight7", True), + ("module.module.mtp.layers.0.transformer_layer.mlp.experts.linear_fc2.weight7", True), + ("module.module.decoder.layers.3.self_attention.value_projection.experts.weight", False), + ( + "module.module.decoder.layers.3.self_attention.value_projection.experts.experts.7.weight", + False, + ), + ("model.layers.3.self_attn.v_experts.7.weight", False), + ], +) +def test_only_ffn_experts_use_expert_parallelism(name: str, expected: bool) -> None: + assert is_ffn_expert_parameter(name) is expected + + +@pytest.mark.parametrize( + ("name", "expected_group"), + [ + ( + "module.module.decoder.layers.3.self_attention.value_projection.experts.weight", + "regular-tp", + ), + ("module.module.decoder.layers.3.mlp.experts.weight", "expert-tp"), + ], +) +def test_tensor_gather_uses_the_correct_tp_group(name: str, expected_group: str) -> None: + param = torch.nn.Parameter(torch.arange(4, dtype=torch.float32).reshape(2, 2)) + param.tensor_model_parallel = True + param.parallel_mode = None + param.partition_dim = 0 + param.partition_stride = 1 + + def fake_all_gather(partitions, source, *, group): + assert group == expected_group + partitions[0].copy_(source) + + with ( + patch( + "miles.backends.megatron_utils.update_weight.common.mpu.get_tensor_model_parallel_world_size", + return_value=1, + ), + patch( + "miles.backends.megatron_utils.update_weight.common.mpu.get_tensor_model_parallel_group", + return_value="regular-tp", + ), + patch( + "miles.backends.megatron_utils.update_weight.common.mpu.get_expert_tensor_parallel_world_size", + return_value=1, + ), + patch( + "miles.backends.megatron_utils.update_weight.common.mpu.get_expert_tensor_parallel_group", + return_value="expert-tp", + ), + patch( + "miles.backends.megatron_utils.update_weight.common.dist.all_gather", + side_effect=fake_all_gather, + ), + ): + gathered = all_gather_param(Namespace(swiglu=False), name, param) + + torch.testing.assert_close(gathered, param) + + +def test_weight_transfer_partition_keeps_mova_values_with_non_experts() -> None: + tensors = [ + ("module.module.decoder.layers.3.mlp.experts.linear_fc1.weight7", torch.tensor(1)), + ( + "module.module.decoder.layers.3.self_attention.value_projection.experts.weight", + torch.tensor(2), + ), + ("module.module.decoder.layers.3.self_attention.linear_qkg.weight", torch.tensor(3)), + ] + with patch( + "miles.backends.megatron_utils.update_weight.common.named_params_and_buffers", + return_value=iter(tensors), + ): + regular = list(collect_named_tensors_for_weight_transfer(_args(), [], is_expert=False)) + with patch( + "miles.backends.megatron_utils.update_weight.common.named_params_and_buffers", + return_value=iter(tensors), + ): + experts = list(collect_named_tensors_for_weight_transfer(_args(), [], is_expert=True)) + + assert [name for name, _ in regular] == [tensors[1][0], tensors[2][0]] + assert [name for name, _ in experts] == [tensors[0][0]] + + +def test_value_and_ffn_expert_bias_buffers_are_enumerated_for_sync() -> None: + class ModuleWithRouterBuffers(nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("value_projection_router_expert_bias", torch.arange(3.0)) + self.register_buffer("mlp_router_expert_bias", torch.arange(4.0)) + self.register_buffer("unrelated_statistics", torch.arange(5.0)) + + names = [ + name + for name, _ in named_params_and_buffers( + _args(), + [ModuleWithRouterBuffers()], + convert_to_global_name=False, + ) + ] + assert names == [ + "vp_stages.0.value_projection_router_expert_bias", + "vp_stages.0.mlp_router_expert_bias", + ] + + +def test_p2p_staging_contract_requires_all_value_expert_shards() -> None: + from miles.backends.megatron_utils.update_weight.update_weight_from_distributed.p2p import ( + UpdateWeightP2P, + ) + + updater = object.__new__(UpdateWeightP2P) + num_value_experts = 3 + packed_name = "model.layers.3.self_attn.v_experts.weight" + updater._shared_params_dict = {packed_name: torch.empty(num_value_experts, 4, 8)} + updater._shared_param_mapper = MagicMock() + updater._shared_param_mapper.map.return_value = SimpleNamespace( + sglang_name=packed_name, + num_shards=num_value_experts, + num_local_experts=None, + ) + updater._staged_tensors = {} + updater._tensor_update_pending = {} + + ready_names = [] + ready_tensors = [] + for expert in range(num_value_experts): + names, tensors = updater._get_transfer_ready_params( + [(f"model.layers.3.self_attn.v_experts.{expert}.weight", torch.full((4, 8), expert))] + ) + ready_names.extend(names) + ready_tensors.extend(tensors) + + assert ready_names == [packed_name] + assert [name for name, _ in ready_tensors] == [ + f"model.layers.3.self_attn.v_experts.{expert}.weight" + for expert in range(num_value_experts) + ] + assert updater._staged_tensors == {} + assert updater._tensor_update_pending == {} + + +def test_broadcast_path_preserves_all_converted_value_expert_metadata() -> None: + from miles.backends.megatron_utils.update_weight.update_weight_from_distributed.broadcast import ( + update_weights_from_distributed, + ) + + tensors = [ + (f"model.layers.3.self_attn.v_experts.{expert}.weight", torch.full((4, 8), expert)) + for expert in range(3) + ] + engine = MagicMock() + engine.update_weights_from_distributed.remote.return_value = "engine-ref" + handle = MagicMock() + + with patch( + "miles.backends.megatron_utils.update_weight.update_weight_from_distributed.broadcast.dist.broadcast", + return_value=handle, + ) as broadcast: + refs = update_weights_from_distributed( + "mova-test", + MagicMock(), + 7, + [engine], + tensors, + ) + + assert refs == ["engine-ref"] + kwargs = engine.update_weights_from_distributed.remote.call_args.kwargs + assert kwargs["names"] == [name for name, _ in tensors] + assert kwargs["dtypes"] == [tensor.dtype for _, tensor in tensors] + assert kwargs["shapes"] == [tensor.shape for _, tensor in tensors] + assert kwargs["weight_version"] == "7" + assert broadcast.call_count == len(tensors) + assert handle.wait.call_count == len(tensors) + + +@pytest.mark.parametrize("mode", ["retract", "abort"]) +def test_mova_cache_safe_update_modes(mode: str) -> None: + validate_weight_update_cache_mode(_args(pause_generation_mode=mode)) + + +def test_mova_rejects_in_place_cache_preservation() -> None: + with pytest.raises(ValueError, match="routed-value KV cache must be flushed"): + validate_weight_update_cache_mode(_args(pause_generation_mode="in_place")) + + +def test_legacy_model_may_retain_existing_in_place_behavior() -> None: + validate_weight_update_cache_mode(_args(num_value_experts=0, pause_generation_mode="in_place")) diff --git a/tests/fast/backends/training_utils/test_weight_version_validation.py b/tests/fast/backends/training_utils/test_weight_version_validation.py new file mode 100644 index 00000000000..5a5ed6f2516 --- /dev/null +++ b/tests/fast/backends/training_utils/test_weight_version_validation.py @@ -0,0 +1,41 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from miles.backends.training_utils import ci_utils + + +def _engine(version): + remote = MagicMock(return_value=f"version-ref-{version}") + return SimpleNamespace(get_weight_version=SimpleNamespace(remote=remote)) + + +def test_all_rollout_engine_versions_are_queried(monkeypatch): + engines = [_engine("7"), _engine("7"), _engine("7")] + ray_get = MagicMock(return_value=["7", 7, "7"]) + monkeypatch.setattr("ray.get", ray_get) + + ci_utils.assert_rollout_engine_weight_versions(engines, expected_version=7) + + assert [engine.get_weight_version.remote.call_count for engine in engines] == [1, 1, 1] + ray_get.assert_called_once_with( + ["version-ref-7", "version-ref-7", "version-ref-7"] + ) + + +def test_all_rollout_engine_versions_report_every_mismatch(monkeypatch): + engines = [_engine("7"), _engine("6"), _engine("stale")] + monkeypatch.setattr("ray.get", MagicMock(return_value=["7", "6", "stale"])) + + with pytest.raises(RuntimeError, match=r"expected 7.*\(1, '6'\).*\(2, 'stale'\)"): + ci_utils.assert_rollout_engine_weight_versions(engines, expected_version=7) + + +def test_all_rollout_engine_versions_accepts_empty_engine_set(monkeypatch): + ray_get = MagicMock() + monkeypatch.setattr("ray.get", ray_get) + + ci_utils.assert_rollout_engine_weight_versions([], expected_version=7) + + ray_get.assert_not_called() diff --git a/tests/fast/utils/test_arguments.py b/tests/fast/utils/test_arguments.py index 37ed8046cfc..8635485a4d8 100644 --- a/tests/fast/utils/test_arguments.py +++ b/tests/fast/utils/test_arguments.py @@ -5,13 +5,159 @@ import pytest -from miles.utils.arguments import _maybe_apply_dumper_overrides, get_miles_extra_args_provider +from miles.utils.arguments import ( + _maybe_apply_dumper_overrides, + get_miles_extra_args_provider, + validate_mova_args, +) from miles.utils.misc import function_registry PATH_ARGS = ["--rollout-function-path", "--custom-generate-function-path"] REQUIRED_ARGS = ["--rollout-batch-size", "64"] +class TestMoVAArguments: + def _parse(self, flags): + with patch.object(sys, "argv", ["test", *flags]): + parser = argparse.ArgumentParser() + get_miles_extra_args_provider()(parser) + return parser.parse_args(flags) + + def test_defaults_match_native_megatron_entry_point(self): + args = self._parse(REQUIRED_ARGS) + + assert args.mova_num_value_experts == 0 + assert args.mova_router_topk == 1 + assert args.mova_router_score_function == "sigmoid" + assert args.mova_router_topk_scaling_factor == 1.0 + assert args.mova_router_enable_expert_bias is False + assert args.mova_router_bias_update_rate == 1.0e-3 + assert args.mova_router_aux_loss_coeff == 0.0 + assert args.mova_router_load_balancing_type == "none" + assert args.mova_num_dense_layers == 0 + assert args.mova_norm_num_groups == 1 + assert args.mova_attention_gate_function == "softplus" + assert args.mova_value_backend == "grouped_gemm" + assert args.mova_use_torch_rms_norm is False + assert args.xllm_router_compatibility is True + assert args.xllm_router_gemm_partitions == 1 + + def test_generated_k2mova_rl_flags_are_accepted(self): + args = self._parse( + REQUIRED_ARGS + + [ + "--mova-num-value-experts", + "64", + "--mova-router-topk", + "4", + "--mova-router-score-function", + "sigmoid", + "--mova-router-topk-scaling-factor", + "2.5", + "--mova-router-enable-expert-bias", + "--mova-router-bias-update-rate", + "0.001", + "--mova-router-aux-loss-coeff", + "0", + "--mova-router-load-balancing-type", + "none", + "--mova-num-dense-layers", + "3", + "--mova-norm-num-groups", + "2", + "--mova-attention-gate-function", + "softplus", + "--mova-value-backend", + "grouped_gemm", + "--no-mova-use-torch-rms-norm", + "--xllm-router-compatibility", + "--xllm-router-gemm-partitions", + "1", + ] + ) + + assert args.mova_num_value_experts == 64 + assert args.mova_router_topk == 4 + assert args.mova_router_topk_scaling_factor == 2.5 + assert args.mova_router_enable_expert_bias is True + assert args.mova_num_dense_layers == 3 + assert args.mova_norm_num_groups == 2 + assert args.mova_value_backend == "grouped_gemm" + assert args.mova_use_torch_rms_norm is False + assert args.xllm_router_compatibility is True + + def test_registration_does_not_change_shared_transformer_defaults(self): + with patch.object(sys, "argv", ["test", *REQUIRED_ARGS]): + parser = argparse.ArgumentParser() + parser.add_argument("--attention-dropout", type=float, default=0.1) + parser.add_argument("--hidden-dropout", type=float, default=0.1) + get_miles_extra_args_provider()(parser) + args, _ = parser.parse_known_args(REQUIRED_ARGS) + + assert args.attention_dropout == 0.1 + assert args.hidden_dropout == 0.1 + + def test_native_validation_accepts_exact_k2mova_contract(self): + args = SimpleNamespace( + mova_num_value_experts=64, + train_backend="megatron", + megatron_to_hf_mode="raw", + custom_model_provider_path=None, + use_legacy_models=False, + yaml_cfg=None, + spec=None, + mtp_num_layers=None, + multi_latent_attention=False, + heterogeneous_layers_config_path=None, + attention_output_gate=True, + rotary_interleaved=True, + group_query_attention=True, + num_experts=100, + attention_dropout=0.0, + hidden_dropout=0.0, + tensor_model_parallel_size=8, + sequence_parallel=True, + ) + + validate_mova_args(args) + + @pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("megatron_to_hf_mode", "bridge", "raw"), + ("attention_output_gate", False, "attention-output-gate"), + ("rotary_interleaved", False, "rotary-interleaved"), + ("attention_dropout", 0.1, "attention-dropout"), + ("sequence_parallel", False, "sequence-parallel"), + ], + ) + def test_native_validation_rejects_incompatible_rl_flags(self, field, value, message): + args = SimpleNamespace( + mova_num_value_experts=64, + train_backend="megatron", + megatron_to_hf_mode="raw", + custom_model_provider_path=None, + use_legacy_models=False, + yaml_cfg=None, + spec=None, + mtp_num_layers=None, + multi_latent_attention=False, + heterogeneous_layers_config_path=None, + attention_output_gate=True, + rotary_interleaved=True, + group_query_attention=True, + num_experts=100, + attention_dropout=0.0, + hidden_dropout=0.0, + tensor_model_parallel_size=8, + sequence_parallel=True, + ) + setattr(args, field, value) + + with pytest.raises(ValueError, match=message): + validate_mova_args(args) + + @pytest.mark.parametrize("reward_type", ["logr", "k3"]) def test_opd_reward_type_is_parsed(reward_type: str) -> None: with patch.object(sys, "argv", ["test", "--opd-reward-type", reward_type] + REQUIRED_ARGS):