diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index c2f3deb0fd6..a125be7e17e 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -21,7 +21,7 @@ from miles.utils.processing_utils import load_tokenizer from miles.utils.ray_utils import Box from miles.utils.reloadable_process_group import destroy_process_groups, monkey_patch_torch_dist, reload_process_groups -from miles.utils.replay_base import all_replay_managers +from miles.utils.replay_base import Replay, all_replay_managers from miles.utils.timer import Timer, inverse_timer, timer from miles.utils.tracking_utils import init_tracking from miles.utils.types import RolloutBatch @@ -39,7 +39,7 @@ from .lora_utils import is_lora_enabled from .model import forward_only, initialize_model_and_optimizer, save, train from .parallel import verify_megatron_parallel_state -from .replay_utils import get_register_replay_list_func +from .replay_utils import get_replay_columns from .update_weight.common import named_params_and_buffers from .update_weight.update_weight_from_distributed.broadcast import UpdateWeightFromDistributed from .update_weight.update_weight_from_distributed.p2p import UpdateWeightP2P @@ -749,8 +749,7 @@ def _fill_replay_data( num_microbatches, rollout_data, data_key: str, - replay_list: list, - register_replay_list_func, + replay_columns: list[tuple[Replay, int]], if_sp_region=True, ): if data_key not in rollout_data: @@ -807,7 +806,8 @@ def pad_func(data, pad): start, end = seqlen // tp_size * tp_rank, seqlen // tp_size * (tp_rank + 1) replay_data = replay_data[start:end] - register_replay_list_func(replay_list, replay_data, self.model) + for replay, column in replay_columns: + replay.record(replay_data[:, column]) del rollout_data[data_key] @@ -892,14 +892,17 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data) for m in all_replay_managers: - if self._use_rollout_replay(m): + if not self._use_rollout_replay(m): + continue + for kind, data_key in m.data_keys.items(): + if data_key not in rollout_data: + continue self._fill_replay_data( data_iterator, num_microbatches, rollout_data, - data_key=m.data_key, - replay_list=m.replays, - register_replay_list_func=get_register_replay_list_func(m), + data_key=data_key, + replay_columns=get_replay_columns(kind, self.model), if_sp_region=m.if_sp_region, ) diff --git a/miles/backends/megatron_utils/replay_utils.py b/miles/backends/megatron_utils/replay_utils.py index 2429964e236..30c5738019e 100644 --- a/miles/backends/megatron_utils/replay_utils.py +++ b/miles/backends/megatron_utils/replay_utils.py @@ -1,33 +1,17 @@ -from megatron.core.transformer.transformer_block import get_num_layers_to_build -from megatron.core.transformer.transformer_layer import get_transformer_layer_offset +from megatron.core.utils import unwrap_model -from miles.utils.replay_base import BaseReplayManager, RoutingReplayManager +from miles.utils.replay_base import Replay -def _register_replay_list_moe(replay_list, replay_data, models): - layer_indices = [] - replay_idx = 0 - for vp_stage, model in enumerate(models): - config = model.module.config - num_layers_to_build = get_num_layers_to_build(config, vp_stage=vp_stage) - offset = get_transformer_layer_offset(config, vp_stage=vp_stage) - for layer_id in range(offset, offset + num_layers_to_build): - if isinstance(config.moe_layer_freq, int): - if layer_id % config.moe_layer_freq != 0: - continue - elif isinstance(config.moe_layer_freq, list): - assert len(config.moe_layer_freq) == config.num_layers - if config.moe_layer_freq[layer_id] == 0: - continue - layer_indices.append(layer_id) +def get_replay_columns(kind: str, models) -> list[tuple[Replay, int]]: + is_value = kind == "value" + dense_prefix = unwrap_model(models[0]).config.mova_num_dense_layers if is_value else 0 - for replay_idx, layer_idx in enumerate(layer_indices): - layer_data = replay_data[:, layer_idx] - replay_list[replay_idx].record(layer_data) - - -def get_register_replay_list_func(manager: BaseReplayManager): - if isinstance(manager, RoutingReplayManager): - return _register_replay_list_moe - else: - raise ValueError(f"Unsupported manager type: {type(manager)}") + columns = [] + for model in unwrap_model(models): + for layer in model.decoder.layers: + owner = getattr(layer.self_attention, "value_projection", None) if is_value else layer.mlp + replay = getattr(getattr(owner, "router", None), "routing_replay", None) + if replay is not None: + columns.append((replay, layer.layer_number - 1 - dense_prefix)) + return columns diff --git a/miles/backends/training_utils/data.py b/miles/backends/training_utils/data.py index 5109676285c..3df820f2f35 100644 --- a/miles/backends/training_utils/data.py +++ b/miles/backends/training_utils/data.py @@ -79,8 +79,9 @@ def get_rollout_data(args: Namespace, rollout_data_ref: Box) -> RolloutBatch: ) ) ] - if "rollout_routed_experts" in rollout_data: - rollout_data["rollout_routed_experts"] = [torch.from_numpy(r) for r in rollout_data["rollout_routed_experts"]] + for key in ("rollout_routed_experts", "rollout_routed_value_experts"): + if key in rollout_data: + rollout_data[key] = [torch.from_numpy(r) for r in rollout_data[key]] return rollout_data diff --git a/miles/backends/training_utils/log_utils.py b/miles/backends/training_utils/log_utils.py index a20b2d6d8be..94c0acc9bdc 100644 --- a/miles/backends/training_utils/log_utils.py +++ b/miles/backends/training_utils/log_utils.py @@ -200,6 +200,7 @@ def log_rollout_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatc "loss_masks", "sample_indices", "rollout_routed_experts", + "rollout_routed_value_experts", "max_seq_lens", "dynamic_global_batch_size", "weight_versions", diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 0ca16e96686..93760b0e5fb 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -833,6 +833,8 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl if samples[0].rollout_routed_experts is not None: train_data["rollout_routed_experts"] = [sample.rollout_routed_experts for sample in samples] + if samples[0].rollout_routed_value_experts is not None: + train_data["rollout_routed_value_experts"] = [sample.rollout_routed_value_experts for sample in samples] if samples[0].train_metadata is not None: train_data["metadata"] = [sample.train_metadata for sample in samples] @@ -963,6 +965,7 @@ def _stat(xs): "rollout_log_probs", "rollout_sampling_masks", "rollout_routed_experts", + "rollout_routed_value_experts", "prompt", "teacher_log_probs", "weight_versions", diff --git a/miles/rollout/generate_utils/generate_endpoint_utils.py b/miles/rollout/generate_utils/generate_endpoint_utils.py index 0811753425d..979e7f3f017 100644 --- a/miles/rollout/generate_utils/generate_endpoint_utils.py +++ b/miles/rollout/generate_utils/generate_endpoint_utils.py @@ -97,6 +97,7 @@ async def update_sample_from_response( # TODO handle multi-turn cases (may need concat instead of assignment) sample.rollout_routed_experts = get_rollout_topk_from_response(args, output, sample, "routed_experts") + sample.rollout_routed_value_experts = get_rollout_topk_from_response(args, output, sample, "routed_value_experts") # TODO may unify (currently there are both methods inside Sample and separate functions) sample.update_from_meta_info(args, output["meta_info"]) @@ -107,4 +108,9 @@ def get_rollout_topk_from_response(args, output, sample, key): if info is None: return None x = np.frombuffer(pybase64.b64decode(info.encode("ascii")), dtype=np.int32) - return x.reshape(len(sample.tokens) - 1, args.num_layers, args.moe_router_topk) + num_layers = args.num_layers + router_topk = args.moe_router_topk + if key == "routed_value_experts": + num_layers -= args.mova_num_dense_layers + router_topk = args.mova_router_topk + return x.reshape(len(sample.tokens) - 1, num_layers, router_topk) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index 42449ae25c6..054ace11e76 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -644,6 +644,11 @@ def apply_merged_session_sample( if merged.rollout_routed_experts is not None: choice = {"meta_info": {"routed_experts": merged.rollout_routed_experts}} sample.rollout_routed_experts = get_rollout_topk_from_response(args, choice, sample, "routed_experts") + if merged.rollout_routed_value_experts is not None: + choice = {"meta_info": {"routed_value_experts": merged.rollout_routed_value_experts}} + sample.rollout_routed_value_experts = get_rollout_topk_from_response( + args, choice, sample, "routed_value_experts" + ) routed_experts_shape = None routed_experts_decoded_bytes = 0 @@ -770,6 +775,7 @@ def _compute_sample_from_openai_record( sample.response_length = len(output_token_ids) sample.loss_mask = [1] * len(output_token_ids) sample.rollout_routed_experts = get_rollout_topk_from_response(args, choice, sample, "routed_experts") + sample.rollout_routed_value_experts = get_rollout_topk_from_response(args, choice, sample, "routed_value_experts") if not hasattr(sample, "metadata") or sample.metadata is None: sample.metadata = {} @@ -817,6 +823,8 @@ def _strip_last_output_tokens_without_decode(sample: Sample, trim_count: int) -> sample.loss_mask = sample.loss_mask[:keep_tokens] if sample.rollout_routed_experts is not None: sample.rollout_routed_experts = sample.rollout_routed_experts[: len(sample.tokens) - 1] + if sample.rollout_routed_value_experts is not None: + sample.rollout_routed_value_experts = sample.rollout_routed_value_experts[: len(sample.tokens) - 1] def truncate_samples_by_total_tokens( @@ -868,4 +876,6 @@ def _truncate_sample_output(sample: Sample, keep_tokens: int, tokenizer) -> None sample.loss_mask = sample.loss_mask[:keep_tokens] if sample.rollout_routed_experts is not None: sample.rollout_routed_experts = sample.rollout_routed_experts[: len(sample.tokens) - 1] + if sample.rollout_routed_value_experts is not None: + sample.rollout_routed_value_experts = sample.rollout_routed_value_experts[: len(sample.tokens) - 1] sample.status = Sample.Status.TRUNCATED diff --git a/miles/rollout/generate_utils/sample_utils.py b/miles/rollout/generate_utils/sample_utils.py index 0c7b9913342..7e8ffc2fd02 100644 --- a/miles/rollout/generate_utils/sample_utils.py +++ b/miles/rollout/generate_utils/sample_utils.py @@ -72,6 +72,8 @@ def _fill_defaults(sample: Sample): assert obs_len > 0, f"obs_len must be > 0, got {obs_len}" if a.rollout_routed_experts is not None: assert a.rollout_routed_experts.shape[0] <= b.rollout_routed_experts.shape[0] + if a.rollout_routed_value_experts is not None: + assert a.rollout_routed_value_experts.shape[0] <= b.rollout_routed_value_experts.shape[0] assert a.status == Sample.Status.COMPLETED, f"a.status must be COMPLETED, got {a.status}" return _create_with_all_fields( @@ -91,6 +93,7 @@ def _fill_defaults(sample: Sample): rollout_log_probs=a.rollout_log_probs + [0.0] * obs_len + b.rollout_log_probs, rollout_sampling_mask=None, rollout_routed_experts=b.rollout_routed_experts, + rollout_routed_value_experts=b.rollout_routed_value_experts, remove_sample=_merge_equal_value("remove_sample"), status=b.status, metadata=metadata, diff --git a/miles/rollout/session/session_types.py b/miles/rollout/session/session_types.py index 53bf43eccea..b9ab94479f6 100644 --- a/miles/rollout/session/session_types.py +++ b/miles/rollout/session/session_types.py @@ -30,6 +30,7 @@ class MergedSessionSample(BaseModel): weight_versions: list[str] = Field(default_factory=list) prefix_cache_meta_infos: list[dict] = Field(default_factory=list) rollout_routed_experts: str | None = None + rollout_routed_value_experts: str | None = None class GetMergedSessionResponse(BaseModel): diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index 7c90c2ab4b2..3a31ffb53b8 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -164,9 +164,10 @@ def _compact_session_record(record: SessionRecord) -> SessionRecord: compact_meta_info[key] = meta_info[key] # Preserve routed experts if present. get_rollout_topk_from_response() - # may read this from choice or meta_info depending on implementation. - if "routed_experts" in meta_info: - compact_meta_info["routed_experts"] = meta_info["routed_experts"] + # may read these from choice or meta_info depending on implementation. + for key in ("routed_experts", "routed_value_experts"): + if key in meta_info: + compact_meta_info[key] = meta_info[key] compact_choice = { "prompt_token_ids": choice.get("prompt_token_ids", []), @@ -174,8 +175,9 @@ def _compact_session_record(record: SessionRecord) -> SessionRecord: "meta_info": compact_meta_info, } - if "routed_experts" in choice: - compact_choice["routed_experts"] = choice["routed_experts"] + for key in ("routed_experts", "routed_value_experts"): + if key in choice: + compact_choice[key] = choice[key] return SessionRecord( timestamp=record.timestamp, @@ -262,6 +264,7 @@ def _merge_session_records_to_sample( weight_versions: list[str] = [] prefix_cache_meta_infos: list[dict] = [] routed_experts = None + routed_value_experts = None prev_checkpoint: list[int] | None = None final_cursor = 0 final_status = "completed" @@ -345,6 +348,7 @@ def _merge_session_records_to_sample( ) routed_experts = meta_info.get("routed_experts", choice.get("routed_experts")) + routed_value_experts = meta_info.get("routed_value_experts", choice.get("routed_value_experts")) final_status = _status_from_finish_reason(choice.get("finish_reason")) if not is_last: assert ( @@ -401,6 +405,7 @@ def _merge_session_records_to_sample( weight_versions=weight_versions, prefix_cache_meta_infos=prefix_cache_meta_infos, rollout_routed_experts=routed_experts, + rollout_routed_value_experts=routed_value_experts, ) @app.get("/sessions/{session_id}") diff --git a/miles/rollout/sglang_rollout.py b/miles/rollout/sglang_rollout.py index 551021f53c5..e005f2560f3 100644 --- a/miles/rollout/sglang_rollout.py +++ b/miles/rollout/sglang_rollout.py @@ -8,7 +8,6 @@ from typing import Any import numpy as np -import pybase64 import sglang_router from packaging.version import parse from tqdm import tqdm @@ -16,6 +15,7 @@ from miles.backends.megatron_utils.lora_utils import LORA_ADAPTER_NAME, is_lora_enabled from miles.rollout.base_types import GenerateFnInput, RolloutFnEvalOutput, RolloutFnTrainOutput from miles.rollout.filter_hub.base_types import MetricGatherer, call_dynamic_filter +from miles.rollout.generate_utils.generate_endpoint_utils import get_rollout_topk_from_response from miles.rollout.inference_rollout.compatibility import load_generate_function from miles.utils import dumper_utils from miles.utils.async_utils import run @@ -216,15 +216,12 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A sample.rollout_log_probs = [] sample.rollout_log_probs += new_response_log_probs - if "routed_experts" in output["meta_info"]: - sample.rollout_routed_experts = np.frombuffer( - pybase64.b64decode(output["meta_info"]["routed_experts"].encode("ascii")), - dtype=np.int32, - ).reshape( - len(sample.tokens) - 1, - args.num_layers, - args.moe_router_topk, - ) + for key, attr in ( + ("routed_experts", "rollout_routed_experts"), + ("routed_value_experts", "rollout_routed_value_experts"), + ): + if key in output["meta_info"]: + setattr(sample, attr, get_rollout_topk_from_response(args, output, sample, key)) sample.update_from_meta_info(args, output["meta_info"]) diff --git a/miles/utils/replay_base.py b/miles/utils/replay_base.py index e4a003a7307..d7761dd7671 100644 --- a/miles/utils/replay_base.py +++ b/miles/utils/replay_base.py @@ -52,6 +52,7 @@ def clear_forward(self): class BaseReplayManager: name: str = "" filename: str = "" + data_keys: dict[str, str] = {} def __init__(self): self.replays: list[Replay] = [] @@ -220,7 +221,7 @@ def check_replay_result(self, old_topk_fn, scores, topk, top_indices, *args, **k class RoutingReplayManager(BaseReplayManager): name = "routing" filename = "routing_replay.pt" - data_key = "rollout_routed_experts" + data_keys = {"ffn": "rollout_routed_experts", "value": "rollout_routed_value_experts"} if_sp_region = True enable_check_replay_result = False replay_check_threshold = 1e-2 diff --git a/miles/utils/types.py b/miles/utils/types.py index 2c470704b5e..a5e7204fd48 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -92,6 +92,9 @@ class Sample: rollout_routed_experts: numpy.ndarray | None = ( None # Routed experts from rollout engine. shape: (num_tokens-1, num_layers, moe_router_topk), dtype=int32 ) + rollout_routed_value_experts: numpy.ndarray | None = ( + None # Value-routed experts. shape: (num_tokens-1, num_layers - mova_num_dense_layers, mova_router_topk), dtype=int32 + ) remove_sample: bool = False class Status(Enum): @@ -239,6 +242,10 @@ def validate(self): actual = len(self.rollout_routed_experts) expect = len(self.tokens) - 1 assert actual == expect, f"rollout_routed_experts length ({actual}) != len(tokens) - 1 ({expect})" + if self.rollout_routed_value_experts is not None: + actual = len(self.rollout_routed_value_experts) + expect = len(self.tokens) - 1 + assert actual == expect, f"rollout_routed_value_experts length ({actual}) != len(tokens) - 1 ({expect})" def strip_last_output_tokens(self, n: int, tokenizer) -> None: """Remove the last *n* output tokens and all associated per-token info.""" @@ -258,6 +265,8 @@ def strip_last_output_tokens(self, n: int, tokenizer) -> None: self.response = tokenizer.decode(self.tokens[-self.response_length :]) if self.response_length > 0 else "" if self.rollout_routed_experts is not None: self.rollout_routed_experts = self.rollout_routed_experts[:-n] + if self.rollout_routed_value_experts is not None: + self.rollout_routed_value_experts = self.rollout_routed_value_experts[:-n] def reset_for_retry(self) -> None: """Reset generated outputs so the original prompt can be re-sampled. @@ -276,6 +285,7 @@ def reset_for_retry(self) -> None: self.rollout_log_probs = None self.rollout_sampling_mask = None self.rollout_routed_experts = None + self.rollout_routed_value_experts = None self.status = Sample.Status.ABORTED self.non_generation_time = 0.0 self.spec_info = Sample.SpecInfo() diff --git a/scripts/models/k2mova-36b.sh b/scripts/models/k2mova-36b.sh new file mode 100644 index 00000000000..89579d8e45d --- /dev/null +++ b/scripts/models/k2mova-36b.sh @@ -0,0 +1,62 @@ +MODEL_ARGS=( + --bf16 + --num-layers 48 + --hidden-size 2560 + --num-attention-heads 32 + --group-query-attention + --num-query-groups 8 + --kv-channels 128 + --ffn-hidden-size 6144 + + # FFN MoE + --num-experts 100 + --moe-ffn-hidden-size 768 + --moe-router-topk 8 + --moe-router-score-function sigmoid + --moe-router-topk-scaling-factor 2.5 + --moe-router-enable-expert-bias + --moe-router-bias-update-rate 0.001 + --moe-router-load-balancing-type aux_loss + --moe-aux-loss-coeff 1.1111111111111112e-6 + --moe-router-dtype fp32 + --moe-shared-expert-intermediate-size 768 + --moe-grouped-gemm + --moe-token-dispatcher-type alltoall + --moe-permute-fusion + --moe-shared-expert-overlap + + # routed value experts + --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-load-balancing-type aux_loss + --mova-router-aux-loss-coeff 1.1111111111111112e-6 + --mova-num-dense-layers 3 + --mova-norm-num-groups 2 + --mova-attention-gate-function softplus + --mova-value-backend grouped_gemm + --xllm-router-compatibility + --xllm-router-gemm-partitions 2 + + --attention-output-gate + --disable-bias-linear + --swiglu + --normalization RMSNorm + --norm-epsilon 1e-6 + --apply-layernorm-1p + --position-embedding-type rope + --rotary-percent 1.0 + --rotary-base 10000000 + --rotary-interleaved + # The pinned fused packed-THD rope path drops rotary_interleaved, so MoVA + # keeps the unfused path to stay in parity with SGLang. + --no-rope-fusion + --max-position-embeddings 524288 + --untie-embeddings-and-output-weights + --init-method-std 0.02 + --vocab-size 250624 + --make-vocab-size-divisible-by 1 +) diff --git a/tests/e2e/megatron/test_mova_r3.py b/tests/e2e/megatron/test_mova_r3.py new file mode 100644 index 00000000000..71d4fdda400 --- /dev/null +++ b/tests/e2e/megatron/test_mova_r3.py @@ -0,0 +1,127 @@ +import os + +import miles.utils.external_utils.command_utils as U + +TIGHT_HOST_MEMORY = bool(int(os.environ.get("MILES_TEST_TIGHT_HOST_MEMORY", "1"))) + +MODEL_TYPE = os.environ.get("MILES_TEST_MOVA_MODEL_TYPE", "k2mova-36b") +HF_DIR = os.environ.get( + "MILES_TEST_MOVA_HF_DIR", + "/mnt/weka/shrd/k2m/yash.akhauri/k2mova36b/k2mova-36b-mid3_v3/checkpoint_0005500", +) +TORCH_DIST = os.environ.get( + "MILES_TEST_MOVA_TORCH_DIST", + "/mnt/weka/shrd/k2m/yash.akhauri/k2mova36b/megatron-k2mova-36b-mid3-v3-0005500-native-router-meta", +) +PROMPT_DATA = os.environ.get( + "MILES_TEST_MOVA_PROMPT_DATA", + "/mnt/weka/shrd/k2pta/datasets/dapo-math-17k/dapo-math-17k.jsonl", +) +RAY_PYTHONPATH = os.environ.get("MILES_TEST_MOVA_RAY_PYTHONPATH", "/root/Megatron-LM") +NUM_GPUS = 8 + + +def execute(): + for path in (HF_DIR, TORCH_DIST, PROMPT_DATA): + if not os.path.exists(path): + raise FileNotFoundError(f"MoVA R3 e2e needs pre-staged inputs; missing: {path}") + + ckpt_args = f"--hf-checkpoint {HF_DIR} --ref-load {TORCH_DIST} " + + rollout_args = ( + f"--prompt-data {PROMPT_DATA} " + "--input-key prompt " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type math " + "--num-rollout 1 " + "--rollout-batch-size 4 " + "--n-samples-per-prompt 2 " + "--rollout-max-response-len 512 " + "--rollout-temperature 1 " + "--global-batch-size 8 " + ) + + perf_args = ( + "--tensor-model-parallel-size 4 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 1 " + "--expert-model-parallel-size 4 " + "--expert-tensor-parallel-size 1 " + "--recompute-granularity full " + "--recompute-method uniform " + "--recompute-num-layers 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 2048 " + ) + + grpo_args = ( + "--advantage-estimator gspo " + f"{'' if TIGHT_HOST_MEMORY else '--use-kl-loss '}" + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--kl-coef 0.00 " + "--entropy-coef 0.00 " + "--eps-clip 4e-4 " + "--use-rollout-routing-replay " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + sglang_args = ( + "--rollout-num-gpus-per-engine 8 " + "--sglang-dtype bfloat16 " + """--sglang-json-model-override-args '{"xllm_source_router_gemm_partitions": 2}' """ + f"--sglang-mem-fraction-static {0.7 if TIGHT_HOST_MEMORY else 0.8} " + "--sglang-max-running-requests 512 " + # The image's Rust gateway predates R3 and drops `return_routed_experts` + # from the request; miles' Python router proxies the body verbatim. + "--use-miles-router " + ) + + ci_args = "--ci-test --ci-disable-kl-checker " + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--actor-num-nodes 1 " + "--actor-num-gpus-per-node 8 " + "--colocate " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__)} " + f"{perf_args} " + f"{sglang_args} " + f"{ci_args} " + f"{misc_args} " + ) + + U.execute_train( + train_args=train_args, + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + extra_env_vars={"PYTHONPATH": RAY_PYTHONPATH}, + ) + + +if __name__ == "__main__": + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + execute() diff --git a/tests/fast/backends/megatron_utils/test_replay_columns.py b/tests/fast/backends/megatron_utils/test_replay_columns.py new file mode 100644 index 00000000000..c1d29b8af7e --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_replay_columns.py @@ -0,0 +1,68 @@ +"""Unit tests for rollout-routing-replay column mapping.""" + +from types import SimpleNamespace + +from miles.backends.megatron_utils.replay_utils import get_replay_columns +from miles.utils.replay_base import RoutingReplayManager + + +def _router(manager: RoutingReplayManager) -> SimpleNamespace: + """A router carrying the replay registered for it, like TopKRouter.__init__ leaves it.""" + replay = manager.create_replay() + return SimpleNamespace(routing_replay=replay) + + +def _chunk(manager: RoutingReplayManager, layer_numbers, num_dense_layers: int) -> SimpleNamespace: + """A model chunk whose layers above the dense prefix carry both routers. + + Layers are built back to front, so the mapping must follow the layer number, not the + build order. + """ + layers = [] + for layer_number in sorted(layer_numbers, reverse=True): + layer = SimpleNamespace(layer_number=layer_number, mlp=SimpleNamespace(), self_attention=SimpleNamespace()) + if layer_number > num_dense_layers: + layer.mlp.router = _router(manager) + layer.self_attention.value_projection = SimpleNamespace(router=_router(manager)) + layers.append(layer) + layers.reverse() + config = SimpleNamespace(mova_num_dense_layers=num_dense_layers) + return SimpleNamespace(config=config, decoder=SimpleNamespace(layers=layers)) + + +def _columns(pairs): + return [column for _, column in pairs] + + +def _replays(pairs): + return [replay for replay, _ in pairs] + + +def test_columns_are_global_layer_indices(): + """Value routing is captured for MoVA layers only, so its columns start after the dense prefix.""" + manager = RoutingReplayManager() + chunks = [_chunk(manager, [1, 2, 3], 2), _chunk(manager, [4, 5, 6], 2)] + + assert _columns(get_replay_columns("ffn", chunks)) == [2, 3, 4, 5] + assert _columns(get_replay_columns("value", chunks)) == [0, 1, 2, 3] + + +def test_pairing_follows_layer_not_registration_order(): + manager = RoutingReplayManager() + chunk = _chunk(manager, range(1, 7), 2) + + assert _replays(get_replay_columns("ffn", [chunk])) == [ + layer.mlp.router.routing_replay for layer in chunk.decoder.layers[2:] + ] + assert _replays(get_replay_columns("value", [chunk])) == [ + layer.self_attention.value_projection.router.routing_replay for layer in chunk.decoder.layers[2:] + ] + + +def test_ffn_mapping_needs_no_mova_config(): + """Vanilla MoE: no dense prefix, no MoVA config fields, no value projection anywhere.""" + manager = RoutingReplayManager() + chunk = _chunk(manager, range(1, 7), num_dense_layers=0) + chunk.config = SimpleNamespace() + + assert _columns(get_replay_columns("ffn", [chunk])) == [0, 1, 2, 3, 4, 5] diff --git a/tests/fast/utils/test_types.py b/tests/fast/utils/test_types.py index e6c12b416e8..b4b5be31cf3 100644 --- a/tests/fast/utils/test_types.py +++ b/tests/fast/utils/test_types.py @@ -15,6 +15,7 @@ def _make_sample( log_probs: bool = False, loss_mask: bool = False, routed_experts: bool = False, + routed_value_experts: bool = False, ) -> Sample: """Create a Sample with the given prompt + response token IDs.""" tokens = prompt_ids + response_ids @@ -30,6 +31,9 @@ def _make_sample( if routed_experts: # shape: (num_tokens - 1, ...) s.rollout_routed_experts = numpy.zeros((len(tokens) - 1, 2, 2), dtype=numpy.int32) + if routed_value_experts: + # shape: (num_tokens - 1, num_value_layers, mova_router_topk) + s.rollout_routed_value_experts = numpy.zeros((len(tokens) - 1, 2, 2), dtype=numpy.int32) return s @@ -84,6 +88,12 @@ def test_strip_truncates_routed_experts(self, tokenizer): s.strip_last_output_tokens(2, tokenizer) assert len(s.rollout_routed_experts) == original_len - 2 + def test_strip_truncates_routed_value_experts(self, tokenizer): + s = _make_sample([1, 2], [3, 4, 5], routed_value_experts=True) + original_len = len(s.rollout_routed_value_experts) + s.strip_last_output_tokens(2, tokenizer) + assert len(s.rollout_routed_value_experts) == original_len - 2 + def test_strip_updates_response_text(self, tokenizer): s = _make_sample([1, 2], [3, 4, 5]) s.strip_last_output_tokens(1, tokenizer)