diff --git a/examples/run_grpo.py b/examples/run_grpo.py index 51b1f08698..2f8961de6d 100644 --- a/examples/run_grpo.py +++ b/examples/run_grpo.py @@ -248,10 +248,20 @@ def _make_policy(**kwargs): ) finally: shutdown_environments(task_to_env, val_task_to_env) + # Do not defer distributed worker teardown to Python finalizers. Ray may + # already be finalizing its core worker by then, and a destructor-triggered + # RPC can attempt to initialize it a second time and abort the driver. A + # failed run must still reach this path, otherwise the NIXL agents are + # only torn down once `sys.meta_path` is gone. Each teardown is guarded so + # a failure in one still lets the other run. try: policy_generation.shutdown() except Exception as error: print(f"Error shutting down generation: {error}", flush=True) + try: + policy.shutdown() + except Exception as error: + print(f"Error shutting down policy: {error}", flush=True) if __name__ == "__main__": diff --git a/infra/nrl_k8s/dynamo_mx/bench/attribute_variance.py b/infra/nrl_k8s/dynamo_mx/bench/attribute_variance.py new file mode 100644 index 0000000000..a1cc8d2b62 --- /dev/null +++ b/infra/nrl_k8s/dynamo_mx/bench/attribute_variance.py @@ -0,0 +1,195 @@ +#!/usr/bin/env python3 +"""Attribute run-to-run refit variance to a stage and a rank pattern. + +``summarize_refit_stages.py`` answers "how fast was this run". This answers "why +do two identical runs differ", which needs a different cut of the same records: +per-stage spread *across* runs, and whether the slowest rank is the same one each +time. + +That distinction decides what is worth optimising. If the spread lives in one +stage, optimise that stage. If it lives in a rank that is persistently slow, the +problem is placement, not code. If the slow rank moves between runs, it is +contention, and a code change aimed at the stage will not reproduce. + +Takes one JSONL file of ``MX_REFIT_STAGE`` payloads per run. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +from pathlib import Path + +# Every stage that contributes to accounted_s, in pipeline order. +STAGES = [ + "prepare_discover_s", + "prepare_capture_s", + "prepare_plan_s", + "prepare_alloc_s", + "prepare_register_s", + "prepare_handshake_s", + "wire_fused_s", + "install_s", +] + + +def load(path: Path, warmup: int) -> list[dict]: + rows = [ + json.loads(l) + for l in path.read_text().splitlines() + if l.strip().startswith("{") + ] + return [r for r in rows if r["step"] > warmup] + + +def fleet_critical(rows: list[dict]) -> list[tuple[int, float, int]]: + """Per step: the slowest rank's accounted_s. A refit ends with its last rank.""" + out = [] + for step in sorted({r["step"] for r in rows}): + at = [r for r in rows if r["step"] == step] + worst = max(at, key=lambda r: r["accounted_s"]) + out.append((step, worst["accounted_s"], worst["rank"])) + return out + + +def spread(vals: list[float]) -> float: + lo = min(vals) + return (max(vals) / lo) if lo > 0 else float("inf") + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("runs", nargs="+", type=Path) + ap.add_argument("--warmup-steps", type=int, default=1) + ap.add_argument("--json-out", type=Path) + args = ap.parse_args() + + runs = [load(p, args.warmup_steps) for p in args.runs] + names = [p.stem for p in args.runs] + + report: dict = {"runs": names, "warmup_steps_excluded": args.warmup_steps} + + # 1. The headline each run would have quoted. + print("== fleet-critical accounted_s per run ==") + headline = [] + crit_ranks = [] + for name, rows in zip(names, runs): + fc = fleet_critical(rows) + med = statistics.median([v for _, v, _ in fc]) + headline.append(med) + ranks = [rk for _, _, rk in fc] + crit_ranks.append(ranks) + mode = max(set(ranks), key=ranks.count) + print( + f" {name:28s} median={med:7.3f}s min={min(v for _, v, _ in fc):7.3f} " + f"max={max(v for _, v, _ in fc):7.3f} critical rank most often={mode} " + f"({ranks.count(mode)}/{len(ranks)} steps)" + ) + report["headline_medians"] = headline + report["headline_spread_x"] = spread(headline) + print(f" -> across-run spread: {spread(headline):.2f}x") + + # 2. Which stage carries the spread. + # + # Measure the stages of the rank that was fleet-critical on each step, not + # the median over all ranks. Those answer different questions and can + # disagree sharply: the body of the per-rank distribution can shift by ~2x + # while the tail that actually sets the refit duration barely moves. Only the + # critical rank's time is on the critical path, so that is the cut that says + # what to optimise. + print( + "\n== stage medians of the fleet-critical rank (the one on the critical path) ==" + ) + hdr = ( + " " + + "stage".ljust(22) + + "".join(n[:11].rjust(13) for n in names) + + " spread share" + ) + print(hdr) + + crit_stage_vals: list[dict[str, list[float]]] = [] + for rows in runs: + per_stage: dict[str, list[float]] = {s: [] for s in STAGES} + for step in sorted({r["step"] for r in rows}): + at = [r for r in rows if r["step"] == step] + worst = max(at, key=lambda r: r["accounted_s"]) + for s in STAGES: + per_stage[s].append(worst.get(s, 0.0)) + crit_stage_vals.append(per_stage) + + stage_rows = {} + totals = [sum(statistics.median(v[s]) for s in STAGES) for v in crit_stage_vals] + for s in STAGES: + meds = [statistics.median(v[s]) for v in crit_stage_vals] + share = ( + (statistics.median(meds) / statistics.median(totals) * 100) + if any(totals) + else 0.0 + ) + active = all(m > 0 for m in meds) + sp = spread(meds) if active else None + stage_rows[s] = { + "medians": meds, + "spread_x": sp, + "share_pct": share, + "active": active, + } + sp_txt = f"{sp:6.2f}x" if active else " --" + print( + " " + + s.ljust(22) + + "".join(f"{m:13.4f}" for m in meds) + + f" {sp_txt} {share:5.1f}%" + ) + report["stages_critical_rank"] = stage_rows + if any(not d["active"] for d in stage_rows.values()): + print( + " ('--' = stage is zero in the warm window; it only runs on the cold step)" + ) + + ranked = sorted( + (kv for kv in stage_rows.items() if kv[1]["active"]), + key=lambda kv: (kv[1]["spread_x"] - 1.0) * kv[1]["share_pct"], + reverse=True, + ) + print("\n stages ranked by contribution to the spread (excess spread x share):") + for s, d in ranked[:4]: + print( + f" {s:22s} spread {d['spread_x']:5.2f}x on {d['share_pct']:5.1f}% of the time" + ) + + # 3. Is the straggler the same rank every run? Placement vs contention. + print("\n== per-rank max accounted_s ==") + slowest = [] + for name, rows in zip(names, runs): + per_rank = { + rk: max(r["accounted_s"] for r in rows if r["rank"] == rk) + for rk in sorted({r["rank"] for r in rows}) + } + worst = max(per_rank, key=per_rank.get) + best = min(per_rank, key=per_rank.get) + slowest.append(worst) + print( + f" {name:28s} slowest=rank {worst} ({per_rank[worst]:.3f}s) " + f"fastest=rank {best} ({per_rank[best]:.3f}s) " + f"ratio={per_rank[worst] / per_rank[best]:.2f}x" + ) + report["slowest_rank_per_run"] = slowest + if len(set(slowest)) == 1: + print( + f" -> the SAME rank ({slowest[0]}) is slowest in every run: placement, not contention." + ) + else: + print( + f" -> the slowest rank MOVES ({slowest}): contention or scheduling, not a fixed rank." + ) + + if args.json_out: + args.json_out.write_text(json.dumps(report, indent=2) + "\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/infra/nrl_k8s/dynamo_mx/bench/configs/grpo_dense_qwen3_4b_tp2dp8_tp4dp4_mx_reshard.yaml b/infra/nrl_k8s/dynamo_mx/bench/configs/grpo_dense_qwen3_4b_tp2dp8_tp4dp4_mx_reshard.yaml new file mode 100644 index 0000000000..0e63d13a82 --- /dev/null +++ b/infra/nrl_k8s/dynamo_mx/bench/configs/grpo_dense_qwen3_4b_tp2dp8_tp4dp4_mx_reshard.yaml @@ -0,0 +1,95 @@ +# Dense counterpart of grpo_moe_qwen3_30b_tp2ep4dp2_tp4dp4_mx_reshard.yaml. +# +# Qwen3-4B-Thinking-2507 (dense) +# trainer TP2 x PP1 x DP8 (16 GPUs) +# -> 4 independent vLLM replicas, each TP4 (16 GPUs) +# +# Why dense: a dense model has no grouped-expert tensors, so it exercises the +# same TP2 -> TP4 reshard across 16 publishers and 16 receivers without also +# exercising vLLM's fused expert layout. That separates transport and lifecycle +# regressions from expert-capture regressions. +# +# The parallel geometry is deliberately identical to the MoE recipe so the +# transport shape -- rank counts, resharding factor, agent fan-out -- carries +# over. Only the model differs, and it is smaller, so these numbers bound scale +# and lifecycle rather than the MoE recipe's byte volume. + +defaults: ../../../../../examples/configs/grpo_math_1B_megatron.yaml + +grpo: + # One cold/warm-up cycle followed by ten measured warm refits. + max_num_steps: 11 + num_prompts_per_step: 8 + num_generations_per_prompt: 2 + max_rollout_turns: 1 + val_period: 0 + val_at_start: false + val_at_end: false + +policy: + model_name: Qwen/Qwen3-4B-Thinking-2507 + train_global_batch_size: 16 + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 1024 + + dtensor_cfg: + enabled: false + + megatron_cfg: + enabled: true + tensor_model_parallel_size: 2 + pipeline_model_parallel_size: 1 + context_parallel_size: 1 + sequence_parallel: true + activation_checkpointing: true + empty_unused_memory_level: 1 + env_vars: + PYTORCH_CUDA_ALLOC_CONF: expandable_segments:False + # Each node holds an address on all four RDMA rails (10.0.32/48/64/80.x), + # and NCCL was pairing its local rail-0 HCA against a peer's rail-3 + # address. Those are different subnets, so RoCE never connects and the + # QP dies with IBV_WC_RETRY_EXC_ERR. Forbid rail crossing so each ring + # stays on one rail. Set here rather than on the RayCluster because an + # env change there restarts the pods and would release eight scarce + # whole nodes. + NCCL_CROSS_NIC: "0" + + generation: + backend: vllm + max_new_tokens: 8 + temperature: 1.0 + top_p: 1.0 + refit_transport: mx_reshard + vllm_cfg: + async_engine: false + tensor_parallel_size: 4 + gpu_memory_utilization: 0.6 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: true + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 4 + +data: + max_input_seq_length: 1024 + +logger: + # Node-local so the run does not depend on a shared volume having free space. + # Copy the evidence off the pod after the run. + log_dir: /tmp/mx-reshard-dense + num_val_samples_to_print: 0 + wandb_enabled: false + tensorboard_enabled: true + mlflow_enabled: false + swanlab_enabled: false + monitor_gpus: true + +cluster: + gpus_per_node: 4 + num_nodes: 8 + +checkpointing: + enabled: false diff --git a/infra/nrl_k8s/dynamo_mx/bench/configs/grpo_moe_qwen3_30b_kl_sample.yaml b/infra/nrl_k8s/dynamo_mx/bench/configs/grpo_moe_qwen3_30b_kl_sample.yaml new file mode 100644 index 0000000000..04f2e44e1e --- /dev/null +++ b/infra/nrl_k8s/dynamo_mx/bench/configs/grpo_moe_qwen3_30b_kl_sample.yaml @@ -0,0 +1,114 @@ +# Identical to grpo_moe_qwen3_30b_tp2ep4dp2_tp4dp4_mx_reshard.yaml except +# max_new_tokens 8 -> 128, which raises the token sample 16x and changes +# nothing else, so batch sizes and the refit path are untouched. +# +# Qwen3-30B-A3B-Instruct-2507 +# trainer TP2 x PP1 x DP8 with EP4/ETP1 (16 GPUs) +# -> 4 independent vLLM replicas, each TP4 (16 GPUs) +# +# Purpose: separate a real trainer/generator mismatch from a small-sample +# artifact when gen_kl_error exceeds the 1e-3 guideline. The base recipe +# generates at most 8 prompts x 2 generations x 8 tokens = 128 tokens per step, +# and gen_kl_error is a k3 estimator over exactly those tokens, so a single +# outlier token moves it a long way. +# +# Performance runs use 1 warm-up + 10 measured refits. + +defaults: ../../../../../examples/configs/grpo_math_qwen30ba3b_megatron.yaml + +grpo: + # One cold/warm-up cycle followed by ten measured warm refits. + max_num_steps: 11 + # One generation per trainer rank keeps the batch divisible by 16. + num_prompts_per_step: 8 + num_generations_per_prompt: 2 + max_rollout_turns: 1 + val_period: 0 + val_at_start: false + val_at_end: false + +policy: + model_name: Qwen/Qwen3-30B-A3B-Instruct-2507 + train_global_batch_size: 16 + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 1024 + + dtensor_cfg: + enabled: false + + megatron_cfg: + enabled: true + tensor_model_parallel_size: 2 + pipeline_model_parallel_size: 1 + context_parallel_size: 1 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 4 + sequence_parallel: true + activation_checkpointing: true + empty_unused_memory_level: 1 + checkpoint: + # The base config defaults this to true, which makes + # GlobalState.initialize_async_checkpoint_worker fork a persistent + # multiprocessing.Manager during worker init. That fork intermittently dies + # before writing its address back, and the parent then raises EOFError from + # reader.recv() and takes the whole run down; it cost 2 of the first 5 MoE + # runs here and never reproduced on the smaller dense model. + # + # `checkpointing.enabled` is false for these benchmark recipes, so no + # checkpoint is ever written and the async worker is pure overhead. Turning + # it off removes the failure mode rather than retrying around it. It is not + # in the refit path, so no measured number depends on it. + async_save: false + env_vars: + PYTORCH_CUDA_ALLOC_CONF: expandable_segments:False + # Each node holds an address on all four RDMA rails (10.0.32/48/64/80.x), + # and NCCL can pair its local rail-0 HCA against a peer's rail-3 address. + # Those are different subnets, so RoCE never connects and the QP dies with + # IBV_WC_RETRY_EXC_ERR. Forbid rail crossing so each ring stays on one + # rail. The earlier MoE run survived without this only because its pod + # placement happened not to produce a rail-crossed pair; KAI reschedules, + # so that is luck rather than a property of the topology. + NCCL_CROSS_NIC: "0" + + generation: + backend: vllm + max_new_tokens: 128 + temperature: 1.0 + top_p: 1.0 + refit_transport: mx_reshard + vllm_cfg: + async_engine: false + tensor_parallel_size: 4 + expert_parallel_size: 1 + gpu_memory_utilization: 0.6 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: true + vllm_kwargs: + moe_backend: triton + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 4 + +data: + max_input_seq_length: 1024 + +logger: + # Node-local so the run does not depend on a shared volume having free space. + # Copy the evidence off the pod after the run. + log_dir: /tmp/mx-reshard-topoa + num_val_samples_to_print: 0 + wandb_enabled: false + tensorboard_enabled: true + mlflow_enabled: false + swanlab_enabled: false + monitor_gpus: true + +cluster: + gpus_per_node: 4 + num_nodes: 8 + +checkpointing: + enabled: false diff --git a/infra/nrl_k8s/dynamo_mx/bench/configs/grpo_moe_qwen3_30b_tp2ep4dp2_tp4dp4_mx_reshard.yaml b/infra/nrl_k8s/dynamo_mx/bench/configs/grpo_moe_qwen3_30b_tp2ep4dp2_tp4dp4_mx_reshard.yaml new file mode 100644 index 0000000000..55fce25c70 --- /dev/null +++ b/infra/nrl_k8s/dynamo_mx/bench/configs/grpo_moe_qwen3_30b_tp2ep4dp2_tp4dp4_mx_reshard.yaml @@ -0,0 +1,117 @@ +# Reference MoE reshard recipe for the ModelExpress `mx_reshard` refit path. +# +# Qwen3-30B-A3B-Instruct-2507 +# trainer TP2 x PP1 x DP8 with EP4/ETP1 (16 GPUs) +# -> 4 independent vLLM replicas, each TP4 (16 GPUs) +# +# The transport is selected by `policy.generation.refit_transport`, which gives +# the publisher/receiver pair added here. Generation is therefore +# NeMo-RL-managed vLLM rather than an external Dynamo deployment, because the +# receiver lives inside the generation worker. Selecting +# `cluster.weight_sync.method: mx` instead would route to the older +# ModelExpress RL surface that this path replaces. +# +# Both inference and training run inside the Ray cluster, so cluster num_nodes +# covers both: 4 nodes (16 GPUs) for the trainer and 4 nodes (16 GPUs) for the +# four TP4 vLLM replicas. +# +# Performance runs use 1 warm-up + 10 measured refits. + +defaults: ../../../../../examples/configs/grpo_math_qwen30ba3b_megatron.yaml + +grpo: + # One cold/warm-up cycle followed by ten measured warm refits. + max_num_steps: 11 + # One generation per trainer rank keeps the batch divisible by 16. + num_prompts_per_step: 8 + num_generations_per_prompt: 2 + max_rollout_turns: 1 + val_period: 0 + val_at_start: false + val_at_end: false + +policy: + model_name: Qwen/Qwen3-30B-A3B-Instruct-2507 + train_global_batch_size: 16 + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 1024 + + dtensor_cfg: + enabled: false + + megatron_cfg: + enabled: true + tensor_model_parallel_size: 2 + pipeline_model_parallel_size: 1 + context_parallel_size: 1 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 4 + sequence_parallel: true + activation_checkpointing: true + empty_unused_memory_level: 1 + checkpoint: + # The base config defaults this to true, which makes + # GlobalState.initialize_async_checkpoint_worker fork a persistent + # multiprocessing.Manager during worker init. That fork intermittently dies + # before writing its address back, and the parent then raises EOFError from + # reader.recv() and takes the whole run down; it cost 2 of the first 5 MoE + # runs here and never reproduced on the smaller dense model. + # + # `checkpointing.enabled` is false for these benchmark recipes, so no + # checkpoint is ever written and the async worker is pure overhead. Turning + # it off removes the failure mode rather than retrying around it. It is not + # in the refit path, so no measured number depends on it. + async_save: false + env_vars: + PYTORCH_CUDA_ALLOC_CONF: expandable_segments:False + # Each node holds an address on all four RDMA rails (10.0.32/48/64/80.x), + # and NCCL can pair its local rail-0 HCA against a peer's rail-3 address. + # Those are different subnets, so RoCE never connects and the QP dies with + # IBV_WC_RETRY_EXC_ERR. Forbid rail crossing so each ring stays on one + # rail. The earlier MoE run survived without this only because its pod + # placement happened not to produce a rail-crossed pair; KAI reschedules, + # so that is luck rather than a property of the topology. + NCCL_CROSS_NIC: "0" + + generation: + backend: vllm + max_new_tokens: 8 + temperature: 1.0 + top_p: 1.0 + refit_transport: mx_reshard + vllm_cfg: + async_engine: false + tensor_parallel_size: 4 + expert_parallel_size: 1 + gpu_memory_utilization: 0.6 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: true + vllm_kwargs: + moe_backend: triton + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 4 + +data: + max_input_seq_length: 1024 + +logger: + # Node-local so the run does not depend on a shared volume having free space. + # Copy the evidence off the pod after the run. + log_dir: /tmp/mx-reshard-topoa + num_val_samples_to_print: 0 + wandb_enabled: false + tensorboard_enabled: true + mlflow_enabled: false + swanlab_enabled: false + monitor_gpus: true + +cluster: + gpus_per_node: 4 + num_nodes: 8 + +checkpointing: + enabled: false diff --git a/infra/nrl_k8s/dynamo_mx/bench/summarize_refit_stages.py b/infra/nrl_k8s/dynamo_mx/bench/summarize_refit_stages.py new file mode 100644 index 0000000000..f559e6bd39 --- /dev/null +++ b/infra/nrl_k8s/dynamo_mx/bench/summarize_refit_stages.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python3 +"""Summarize MX_REFIT_STAGE records into the numbers the benchmark contract asks for. + +Reads one JSON object per line (the payload after the ``MX_REFIT_STAGE`` marker) +and reports, over the measured warm window only: + +* the fleet-critical latency per step, i.e. the slowest rank, because a refit is + not finished until its last rank is; +* min / median / p95 / max of that critical value across steps; +* per-rank maxima, so a single persistently slow rank is visible; +* stage attribution, to satisfy the ">=95% attributed or <=100 ms + unattributed" gate before any number is quoted. + +Throughput is deliberately reported per rank. Ranks overlap in wall clock but +are not measured in one shared window, so summing them would overstate the +fleet; the aggregate is labelled an upper bound. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import sys +from pathlib import Path + +# Per-step stages known at the time of writing. Kept only to flag records that +# carry a stage this tool has never seen; the attribution arithmetic below uses +# whatever per-step stages each record actually contains, so a stage added on the +# MX side can never silently masquerade as unattributed time. +KNOWN_STAGE_FIELDS = ( + "descriptor_build_s", + "wire_fused_s", + "install_s", + "reslice_s", +) + + +def stage_fields(record: dict) -> list[str]: + """Per-step stage keys in a record. + + ``accounted_s`` is the total being decomposed, and ``prepare_*`` covers + one-time discovery/handshake work that sits outside the per-step total, so + both are excluded. + """ + return [ + key + for key in record + if key.endswith("_s") + and key != "accounted_s" + and not key.startswith("prepare_") + ] + + +def percentile(values: list[float], fraction: float) -> float: + """Nearest-rank percentile; avoids interpolating across few samples.""" + if not values: + raise ValueError("no values") + ordered = sorted(values) + index = min(len(ordered) - 1, max(0, round(fraction * (len(ordered) - 1)))) + return ordered[index] + + +def _attribution_gate(framework_e2e_s: float | None, mx_critical_s: float) -> dict: + """Attribution measured against the cost the framework actually pays. + + The benchmark contract's ">=95% attributed or <=100 ms unattributed" gate is + about explaining a refit, so the denominator has to be the framework-visible + refit time. Measuring MX's stages against MX's own ``accounted_s`` subtotal + always yields ~100% and would pass this gate while an arbitrary amount of + per-refit cost sits outside MX entirely. + """ + if framework_e2e_s is None: + return { + "gate_pass": None, + "reason": ( + "no --framework-e2e-s supplied; attribution is unknown. MX stages " + "alone cannot establish it." + ), + } + unattributed = framework_e2e_s - mx_critical_s + attributed_pct = 100.0 * mx_critical_s / framework_e2e_s + return { + "framework_e2e_s": framework_e2e_s, + "mx_critical_s": mx_critical_s, + "unattributed_s": unattributed, + "attributed_pct": attributed_pct, + "gate_pass": unattributed <= 0.100 or attributed_pct >= 95.0, + } + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("records", type=Path) + parser.add_argument( + "--warmup-steps", + type=int, + default=1, + help="Leading steps to exclude as cold/warm-up.", + ) + parser.add_argument("--json-out", type=Path) + parser.add_argument( + "--framework-e2e-s", + type=float, + help=( + "Median warm framework-visible refit seconds " + "(timing/.../transfer_and_update_weights). Required to evaluate the " + "attribution gate, because MX's own stages sum to accounted_s by " + "construction and cannot reveal cost outside MX." + ), + ) + args = parser.parse_args() + + rows = [ + json.loads(line) + for line in args.records.read_text().splitlines() + if line.strip() + ] + if not rows: + print("no records", file=sys.stderr) + return 1 + + steps = sorted({r["step"] for r in rows}) + measured_steps = steps[args.warmup_steps :] + measured = [r for r in rows if r["step"] in measured_steps] + ranks = sorted({r["rank"] for r in rows}) + + # A refit ends when its slowest rank ends, so the fleet-critical figure is a + # per-step max over ranks -- never a mean. + critical = [] + for step in measured_steps: + per_step = [r["accounted_s"] for r in measured if r["step"] == step] + critical.append(max(per_step)) + + per_rank_max = { + rank: max(r["accounted_s"] for r in measured if r["rank"] == rank) + for rank in ranks + } + + unattributed = [ + r["accounted_s"] - sum(r[f] for f in stage_fields(r)) for r in measured + ] + worst_unattributed = max(unattributed) + novel_stages = sorted( + {f for r in measured for f in stage_fields(r)} - set(KNOWN_STAGE_FIELDS) + ) + if novel_stages: + print( + f"note: records carry stage(s) unknown to this tool: {novel_stages}; " + "counted as attributed. Add them to KNOWN_STAGE_FIELDS.", + file=sys.stderr, + ) + attribution_pct = 100.0 * ( + 1.0 - worst_unattributed / max(r["accounted_s"] for r in measured) + ) + + bytes_per_rank = {r["bytes"] for r in measured} + wire = [r["wire_fused_s"] for r in measured] + one_rank_bytes = next(iter(bytes_per_rank)) + per_rank_gbps = [8 * one_rank_bytes / w / 1e9 for w in wire] + + fallback_total = sum(r.get("fallback", 0) for r in rows) + full_pull_total = sum(r.get("full_pull_sources", 0) for r in rows) + + summary = { + "schema": "refit-summary-v1", + "ranks": len(ranks), + "steps_total": len(steps), + "warmup_steps_excluded": args.warmup_steps, + "measured_steps": measured_steps, + "bytes_per_rank": sorted(bytes_per_rank), + "fleet_critical_accounted_s": { + "min": min(critical), + "median": statistics.median(critical), + "p95": percentile(critical, 0.95), + "max": max(critical), + }, + "per_rank_max_accounted_s": { + "min": min(per_rank_max.values()), + "max": max(per_rank_max.values()), + "slowest_rank": max(per_rank_max, key=per_rank_max.get), + }, + "per_rank_wire_gbps": { + "min": min(per_rank_gbps), + "median": statistics.median(per_rank_gbps), + "max": max(per_rank_gbps), + }, + "aggregate_wire_gbps_upper_bound": len(ranks) + * statistics.median(per_rank_gbps), + # Internal consistency only. MX's per-step stages are defined to sum to + # accounted_s, so a high number here is near-tautological: it proves the + # stage fields were parsed, not that the refit's cost is understood. It is + # deliberately NOT named "attribution" and carries no gate. + "internal_stage_consistency": { + "worst_residual_s": worst_unattributed, + "consistency_pct_worst_case": attribution_pct, + }, + "stage_attribution": _attribution_gate( + args.framework_e2e_s, statistics.median(critical) + ), + "correctness_markers": { + "fallback_total": fallback_total, + "full_pull_sources_total": full_pull_total, + }, + } + + print(json.dumps(summary, indent=2)) + if args.json_out: + args.json_out.write_text(json.dumps(summary, indent=2) + "\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index d8d5f26763..cf868f30eb 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -393,12 +393,16 @@ def _validate_multimodal_dedup_capability(master_config: MasterConfig) -> None: def _needs_hf_refit_handshake( generation_backend: str, nccl_reshard_refit_enabled: bool, + mx_reshard_refit_enabled: bool, colocated_inference: bool, ) -> bool: """Whether setup must run the HF-schema prepare_refit_info handshake.""" if generation_backend == "megatron": return False - return not (nccl_reshard_refit_enabled and not colocated_inference) + return not ( + (nccl_reshard_refit_enabled or mx_reshard_refit_enabled) + and not colocated_inference + ) def shutdown_environments( @@ -1311,9 +1315,12 @@ def initialize_generation_with_policy( ) assert remote_transport is not None remote_synchronizer_cls = VllmRemoteSparseWeightSynchronizer - elif refit_transport is not None and refit_transport != "nccl_reshard": - # nccl_reshard is handled below via nccl_reshard_refit_enabled, - # not via checkpoint-engine. + elif refit_transport is not None and refit_transport not in { + "nccl_reshard", + "mx_reshard", + }: + # Native reshard transports are handled below, not via + # checkpoint-engine. checkpoint_engine_config = checkpoint_engine_refit_config(generation_config) assert checkpoint_engine_config is not None @@ -1513,12 +1520,19 @@ def init_dynamo(): nccl_reshard_refit_enabled = ( generation_config.get("refit_transport") == "nccl_reshard" ) + mx_reshard_refit_enabled = generation_config.get("refit_transport") == "mx_reshard" if nccl_reshard_refit_enabled: from nemo_rl.weight_sync.nccl_reshard_utils import ( check_nccl_reshard_refit_support, ) check_nccl_reshard_refit_support(master_config) + if mx_reshard_refit_enabled: + from nemo_rl.weight_sync.mx_reshard_weight_synchronizer import ( + check_mx_reshard_refit_support, + ) + + check_mx_reshard_refit_support(master_config) if generation_config.get("refit_transport") is not None and backend != "vllm": raise NotImplementedError( @@ -1546,6 +1560,20 @@ def init_dynamo(): t0 = time.perf_counter() policy_generation.weight_synchronizer.sync_weights() setup_timing_metrics.generation_init_load_time_s = time.perf_counter() - t0 + elif mx_reshard_refit_enabled: + t0 = time.perf_counter() + policy_generation.weight_synchronizer = create_weight_synchronizer( + policy=policy, + generation=policy_generation, + generation_backend=backend, + colocated=False, + train_cluster=train_cluster, + inference_cluster=inference_cluster, + ) + policy_generation.weight_synchronizer.init_communicator() + setup_timing_metrics.extras["vllm_mx_reshard_init_time_s"] = ( + time.perf_counter() - t0 + ) # if it is not colocated inference, initialize collective communication for update weights elif ( not colocated_inference @@ -1625,7 +1653,10 @@ def init_dynamo(): if getattr( policy_generation, "weight_synchronizer", None ) is None and _needs_hf_refit_handshake( - backend, nccl_reshard_refit_enabled, colocated_inference + backend, + nccl_reshard_refit_enabled, + mx_reshard_refit_enabled, + colocated_inference, ): state_dict_info = policy.prepare_refit_info() if policy_generation is not None: diff --git a/nemo_rl/distributed/mx_megatron_helpers.py b/nemo_rl/distributed/mx_megatron_helpers.py new file mode 100644 index 0000000000..d8100946ee --- /dev/null +++ b/nemo_rl/distributed/mx_megatron_helpers.py @@ -0,0 +1,728 @@ +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""Classify native Megatron parameters for external reshard publishers. + +The DTensor MX path uses the generic MX publisher directly because DTensor's +``Placement`` enum gives sharding info in a uniform way. Megatron-Core has no +such uniform API; sharding lives in the wrapper-class identity +(``ColumnParallelLinear``, ``RowParallelLinear``, ``VocabParallelEmbedding``, +fused QKV/MLP, MoE expert layers). + +This module: + +* Classifies every parameter into a Megatron sharding role by walking the model + graph and consulting **Megatron-Bridge's authoritative parallelism + registry** (``megatron.bridge.models.conversion.param_mapping.AutoMapping + ._MODULE_TYPE_REGISTRY``). Bridge's registry already classifies every + TE / Inference / Quant variant of column-parallel, row-parallel, and + replicated modules — using it directly rather than rolling our own + string-matching means we get correct classification of: + - ``TEColumnParallelLinear``, ``TELayerNormColumnParallelLinear``, + ``TEColumnParallelGroupedLinear``, ``InferenceLayerNormColumnParallelLinear`` + - ``TERowParallelLinear``, ``TERowParallelGroupedLinear``, + ``InferenceRowParallelLinear`` + - ``TENorm``, ``FusedLayerNorm``, ``WrappedTorchNorm``, ``L2Norm``, + ``InferenceTopKRouter``, ``LinearForLastLayer`` + …without us having to maintain a parallel list. If Bridge is not + importable, the helper falls back to string-matching against the + base class names — sufficient for mainline Megatron-Core. +* Extracts the local native shard (no allgather, no Megatron-Bridge + ``export_hf_weights`` call — the param tensor IS the local shard). +* Builds the descriptor extras consumed by ModelExpress's main-native + ``refit.reshard.megatron_aliases`` adapter. + +Limitations: + +* Fused QKV / fused gated MLP detection is currently keyed on common + Megatron name patterns (``linear_qkv``, ``linear_fc1``). Mainline + Megatron-Core uses these names; non-mainline forks may need a + ``megatron_role_overrides`` entry. +* MoE per-expert publishing classifies as ``expert_column`` / + ``expert_row``; the per-expert axis is assumed to be 0 (the leading + axis), matching ``detect_moe_expert_layout``'s convention. +""" + +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass, field +from functools import lru_cache +from typing import TYPE_CHECKING, Any, Callable, Iterator + +if TYPE_CHECKING: + import torch + +logger = logging.getLogger("nemo_rl.distributed.mx_megatron_helpers") + +QkvGeometry = tuple[int, int, int] +QkvGeometryResolver = Callable[[str, Any, Any], QkvGeometry | None] + + +# Match ModelExpress's main-native Megatron alias role vocabulary. +ROLE_QKV_COLUMN = "qkv_column" +ROLE_GATED_MLP_COLUMN = "gated_mlp_column" +ROLE_COLUMN = "column" +ROLE_ROW = "row" +ROLE_VOCAB_PARALLEL = "vocab_parallel" +ROLE_REPLICATED = "replicated" +ROLE_EXPERT_COLUMN = "expert_column" +ROLE_EXPERT_ROW = "expert_row" + +_TP_SHARDED_ROLES = frozenset( + { + ROLE_QKV_COLUMN, + ROLE_GATED_MLP_COLUMN, + ROLE_COLUMN, + ROLE_ROW, + ROLE_VOCAB_PARALLEL, + ROLE_EXPERT_COLUMN, + ROLE_EXPERT_ROW, + } +) + + +@dataclass +class MegatronRoleSpec: + """Per-parameter classification result. + + ``role`` is one of the role string constants. ``descriptor_extras`` carries + the per-tensor metadata consumed by ModelExpress's Megatron alias builder. + """ + + role: str + descriptor_extras: dict[str, str] = field(default_factory=dict) + is_expert: bool = False + expert_axis: int = 0 + owned_expert_ids: set[int] = field(default_factory=set) + + +@dataclass(frozen=True) +class MegatronTpShardGeometry: + """Global TP geometry for one rank-local native Megatron tensor.""" + + global_shape: tuple[int, ...] + shard_axis: int + local_shard_range: tuple[int, int] + + +def infer_megatron_tp_shard_geometry( + *, + local_shape: tuple[int, ...], + role: str, + tp_size: int, + tp_rank: int, + expert_tp_size: int | None = None, + expert_tp_rank: int | None = None, + descriptor_extras: dict[str, str] | None = None, +) -> MegatronTpShardGeometry | None: + """Describe the TP shard held by a native Megatron parameter. + + Expert tensors still have TP geometry when inference EP is disabled. The + previous publisher omitted it for every expert, which forced receivers to + pull each complete expert tensor and slice locally. + """ + is_expert = role in {ROLE_EXPERT_COLUMN, ROLE_EXPERT_ROW} + shard_world_size = ( + int(expert_tp_size) + if is_expert and expert_tp_size is not None + else int(tp_size) + ) + shard_rank = ( + int(expert_tp_rank) + if is_expert and expert_tp_rank is not None + else int(tp_rank) + ) + if not 0 <= shard_rank < shard_world_size: + raise ValueError( + f"shard rank {shard_rank} is outside [0, {shard_world_size}) " + f"for Megatron role {role!r}" + ) + if role == ROLE_REPLICATED or shard_world_size <= 1: + return None + extras = descriptor_extras or {} + expert_layout = extras.get("expert_layout", "grouped") + if role in { + ROLE_COLUMN, + ROLE_QKV_COLUMN, + ROLE_GATED_MLP_COLUMN, + ROLE_VOCAB_PARALLEL, + }: + shard_axis = 0 + elif role == ROLE_ROW: + shard_axis = 1 + elif role == ROLE_EXPERT_COLUMN: + shard_axis = 1 if expert_layout == "leading_axis" else 0 + elif role == ROLE_EXPERT_ROW: + shard_axis = 2 if expert_layout == "leading_axis" else 1 + else: + raise ValueError(f"unsupported Megatron TP shard role {role!r}") + if shard_axis >= len(local_shape): + raise ValueError( + f"Megatron role {role!r} requires shard axis {shard_axis}, " + f"but local shape is {local_shape}" + ) + local_extent = int(local_shape[shard_axis]) + global_shape = list(local_shape) + global_shape[shard_axis] = local_extent * shard_world_size + return MegatronTpShardGeometry( + global_shape=tuple(global_shape), + shard_axis=shard_axis, + local_shard_range=( + shard_rank * local_extent, + (shard_rank + 1) * local_extent, + ), + ) + + +# Heuristic name patterns for fused-QKV and fused-gate+up linears in +# mainline Megatron-Core. Callers can provide role overrides for forks that use +# different names. +_DEFAULT_FUSED_QKV_NAME_PATTERNS = ("linear_qkv", "qkv_proj", "fused_qkv") +_DEFAULT_FUSED_GATED_MLP_PATTERNS = ("linear_fc1", "gate_up_proj") +# Vocab / embedding name pattern. +_DEFAULT_VOCAB_NAME_PATTERNS = ( + "word_embeddings", + "embedding", + "lm_head", + "output_layer", +) + + +@lru_cache(maxsize=1) +def _bridge_module_type_registry() -> dict[str, frozenset[str]] | None: + """Return Bridge's authoritative module classifier registry, or None. + + Bridge ships a curated dict of + ``{"column": {classes...}, "row": {...}, "replicated": {...}}`` covering + every TE / Inference / Quant variant. Importing it lazily avoids a hard + dependency: when Bridge is not in the import path (e.g. in unit tests on + a CPU-only env), the caller falls back to substring matching against + the base class names, which is correct for mainline Megatron-Core. + + Cached because this is consulted once per parameter while classifying a + publish set, and re-running the import machinery and copying the registry + thousands of times per refit buys nothing: the installed Bridge cannot + change mid-process. Frozen sets keep the cached value from being mutated by + a caller. + """ + try: + from megatron.bridge.models.conversion.param_mapping import ( + AutoMapping as _AM, + ) + + return { + kind: frozenset(classes) + for kind, classes in _AM._MODULE_TYPE_REGISTRY.items() + } + except Exception: + return None + + +def _classify_module_class(mod_class_name: str) -> str | None: + """Map ``mod.__class__.__name__`` to a Megatron-Bridge parallelism kind. + + Returns one of ``"column"``, ``"row"``, ``"replicated"``, or ``None`` + if the class name doesn't match any known parallelism variant. + """ + if not mod_class_name: + return None + registry = _bridge_module_type_registry() + if registry is not None: + # Direct hit on Bridge's curated set (catches every TE / Inference / + # Quant variant by exact class name). + for kind, cls_set in registry.items(): + if mod_class_name in cls_set: + return kind + # Bridge also has a special-case for the TE-fused + # LayerNormColumnParallelLinear: classify as column. + if "LayerNormColumnParallelLinear" in mod_class_name: + return "column" + # Fallback: substring match against the base names. + if "ColumnParallel" in mod_class_name or "VocabParallelEmbedding" in mod_class_name: + return "column" + if "RowParallel" in mod_class_name: + return "row" + if any( + needle in mod_class_name + for needle in ( + "Norm", + "RMSNorm", + "L2Norm", + "TopKRouter", + "LinearForLastLayer", + "IdentityOp", + ) + ): + return "replicated" + return None + + +_PARAM_LEAF_NAMES = {"weight", "bias", "scale", "_extra_state"} + + +def _is_param_leaf(name_part: str) -> bool: + """Return True for any trailing name that's a parameter rather than a child module. + + Includes the standard ``weight``/``bias``/``scale``/``_extra_state`` + and the grouped-MoE per-expert convention ``weight0``, ``weight1``, + ``weight127``, ``bias0``, etc. Megatron-Core's TE-grouped linears + expose one ``weight`` ``nn.Parameter`` per local expert. + """ + if name_part in _PARAM_LEAF_NAMES: + return True + for base in ("weight", "bias", "scale"): + if name_part.startswith(base): + suffix = name_part[len(base) :] + if suffix and suffix.isdigit(): + return True + return False + + +def _expert_index_from_param(name_part: str) -> int | None: + """If ``name_part`` is ``weight``/``bias``/etc, return ``N``.""" + for base in ("weight", "bias", "scale"): + if name_part.startswith(base): + suffix = name_part[len(base) :] + if suffix and suffix.isdigit(): + return int(suffix) + return None + + +def canonicalize_grouped_expert_name( + name: str, descriptor_extras: dict[str, str] +) -> str: + """Replace an EP-local grouped-expert leaf index with its global ID. + + ModelExpress groups reshard sources by tensor name. Leaving every EP rank's + first local expert named ``weight0`` makes unrelated experts collide before + the Megatron translator can inspect descriptor extras. + """ + if descriptor_extras.get("expert_layout") != "grouped": + return name + local = descriptor_extras.get("local_expert_id") + global_id = descriptor_extras.get("expert_id") + if local is None or global_id is None: + raise ValueError( + f"grouped expert tensor {name!r} is missing local/global expert IDs" + ) + parent, separator, leaf = name.rpartition(".") + for prefix in ("weight", "bias", "scale"): + if leaf == f"{prefix}{local}": + global_leaf = f"{prefix}{global_id}" + return f"{parent}{separator}{global_leaf}" if separator else global_leaf + raise ValueError( + f"grouped expert tensor {name!r} does not end in an indexed " + "weight/bias/scale leaf" + ) + + +def _enclosing_module(name: str, model: "torch.nn.Module") -> "torch.nn.Module | None": + """Walk down model attributes to find the module that owns ``name``. + + ``name`` is a parameter name like + ``decoder.layers.0.self_attention.linear_qkv.weight`` or + ``decoder.layers.0.mlp.experts.linear_fc1.weight0`` for grouped-MoE + per-expert parameters. Return the parent module of the final + parameter token. + """ + parts = name.split(".") + if not parts or not _is_param_leaf(parts[-1]): + # Fall back to the deepest module — caller will get a leaf. + cur = model + for p in parts: + sub = getattr(cur, p, None) + if sub is None: + return None + cur = sub + return cur + cur: Any = model + for p in parts[:-1]: + sub = getattr(cur, p, None) + if sub is None: + return None + cur = sub + return cur + + +def resolve_qkv_geometry_from_param( + name: str, _param: Any, model: "torch.nn.Module" +) -> QkvGeometry | None: + """Read global Q/KV geometry from the live layer-local QKV module. + + Heterogeneous Megatron models attach the resolved per-layer + ``TransformerConfig`` to ``linear_qkv``. Reading the root model config would + stamp one geometry on every layer and is therefore only a compatibility + fallback for callers that cannot resolve the owning module. + """ + if not _is_fused_qkv_name(name): + return None + module = _enclosing_module(name, model) + config = getattr(module, "config", None) + if config is None: + return None + q_heads = getattr(config, "num_attention_heads", None) + if q_heads is None: + return None + kv_heads = getattr(config, "num_query_groups", None) or q_heads + head_dim = getattr(config, "kv_channels", None) + if head_dim is None: + hidden_size = getattr(config, "hidden_size", None) + if hidden_size is None or int(hidden_size) % int(q_heads): + return None + head_dim = int(hidden_size) // int(q_heads) + return int(q_heads), int(kv_heads), int(head_dim) + + +def _module_class_name(mod: "torch.nn.Module | None") -> str: + if mod is None: + return "" + return type(mod).__name__ + + +def _is_fused_qkv_name(name: str) -> bool: + return any(p in name for p in _DEFAULT_FUSED_QKV_NAME_PATTERNS) + + +def _is_fused_gated_mlp_name(name: str) -> bool: + return any(p in name for p in _DEFAULT_FUSED_GATED_MLP_PATTERNS) + + +def _is_vocab_name(name: str) -> bool: + return any(p in name for p in _DEFAULT_VOCAB_NAME_PATTERNS) + + +def _is_expert_name(name: str, *, expert_pattern: str) -> bool: + return expert_pattern in name + + +def detect_megatron_role( + name: str, + param: "torch.Tensor", + *, + model: "torch.nn.Module", + tp_size: int, + ep_size: int, + ep_rank: int, + num_local_experts: int | None = None, + num_attention_heads: int | None = None, + num_kv_heads: int | None = None, + head_dim: int | None = None, + qkv_geometry: QkvGeometry | None = None, + expert_pattern: str | None = None, + role_overrides: dict[str, str] | None = None, +) -> MegatronRoleSpec: + """Classify a Megatron parameter into one of seven roles. + + Returns the role plus per-tensor metadata for the alias builder. The + classifier is conservative: when we can't + determine sharding from the module class, we fall back to + ``ROLE_REPLICATED`` (rank 0 publishes, others skip). That's a + correctness-preserving default — replicated tensors round-trip via the + receiver's passthrough path. + + Args: + name: param name from ``model.named_parameters()`` (e.g. + ``decoder.layers.0.self_attention.linear_qkv.weight``). + param: the local shard tensor (Megatron stores native shards). + model: the root model module; used to walk attributes for the + enclosing module's class. + tp_size, ep_size, ep_rank: from ``parallel_state``. + num_attention_heads, num_kv_heads, head_dim: model-wide compatibility + fallback for ``qkv_column`` metadata. + qkv_geometry: optional per-tensor ``(global query heads, global KV + heads, head dimension)``. This takes precedence over model-wide + values and is required for heterogeneous attention. + expert_pattern: substring marker for MoE expert tensors; defaults to + ``"experts"`` and can be overridden with + ``NRL_MX_EXPERT_TENSOR_PATTERN``. + role_overrides: optional ``{param_name_substring: role}`` dict + for forcing a role on a specific tensor (escape hatch for + non-mainline Megatron forks). + """ + expert_pattern = expert_pattern or os.environ.get( + "NRL_MX_EXPERT_TENSOR_PATTERN", "experts" + ) + + # ---- 1. Explicit override wins. ---- + if role_overrides: + for needle, role in role_overrides.items(): + if needle in name: + return MegatronRoleSpec(role=role) + + # ---- 2a. Grouped-MoE per-expert tensors (one ``weight`` + # nn.Parameter per local expert, used by TE-grouped linears even + # when EP=1). The trailing param name carries the expert index. + if _is_expert_name(name, expert_pattern=expert_pattern): + leaf = name.rsplit(".", 1)[-1] if "." in name else name + expert_idx = _expert_index_from_param(leaf) + if expert_idx is not None: + # Per-expert grouped tensor. Each `weight` is one expert's + # full local shard; the receiver runs per_expert assembly. + # + # `weight` is LOCAL to this EP rank (every rank names its + # experts 0..num_local-1). Advertise the GLOBAL expert id so a + # receiver gathering across EP ranks (EP-trainer -> non-EP / lower-EP + # rollout) can place experts without collision and the EP filter can + # route by global ownership. global = ep_rank*num_local + local. + global_idx = expert_idx + if num_local_experts: + global_idx = ep_rank * int(num_local_experts) + expert_idx + mod_class = _module_class_name(_enclosing_module(name, model)) + sub_role = ( + ROLE_EXPERT_ROW if "RowParallel" in mod_class else ROLE_EXPERT_COLUMN + ) + return MegatronRoleSpec( + role=sub_role, + is_expert=True, + expert_axis=0, + owned_expert_ids={global_idx}, + descriptor_extras={ + "expert_axis": "0", + "expert_id": str(global_idx), + "local_expert_id": str(expert_idx), + "expert_layout": "grouped", + }, + ) + + # ---- 2b. EP>1 leading-axis grouped (legacy path: single .weight + # holds ep_size experts as the leading axis chunk). ---- + if ( + _is_expert_name(name, expert_pattern=expert_pattern) + and ep_size > 1 + and param.ndim >= 2 + ): + leading = param.shape[0] + if leading % ep_size == 0: + chunk = leading // ep_size + owned = set(range(ep_rank * chunk, (ep_rank + 1) * chunk)) + sub_role = ROLE_EXPERT_COLUMN + if _is_fused_gated_mlp_name(name): + # Per-expert fused gate+up: assembler treats it as + # gated_mlp_split inside the per-expert routing. + sub_role = ROLE_EXPERT_COLUMN + mod_class = _module_class_name(_enclosing_module(name, model)) + if "RowParallel" in mod_class: + sub_role = ROLE_EXPERT_ROW + return MegatronRoleSpec( + role=sub_role, + is_expert=True, + expert_axis=0, + owned_expert_ids=owned, + descriptor_extras={ + "expert_axis": "0", + "expert_layout": "leading_axis", + }, + ) + + # ---- 3. Walk to the enclosing module + classify against Bridge's + # AutoMapping._MODULE_TYPE_REGISTRY (or fall back to substring match). ---- + mod = _enclosing_module(name, model) + mod_class = _module_class_name(mod) + parallelism = _classify_module_class(mod_class) + + # ---- 4. VocabParallelEmbedding / lm_head sharded along rows. ---- + if mod_class == "VocabParallelEmbedding" or ( + _is_vocab_name(name) + and tp_size > 1 + and param.ndim >= 2 + and parallelism == "column" + ): + return MegatronRoleSpec(role=ROLE_VOCAB_PARALLEL) + + # ---- 5. Column-parallel linears (incl. all TE / Inference / Quant variants). ---- + if parallelism == "column": + if _is_fused_qkv_name(name): + extras: dict[str, str] = {"qkv_interleave": "by_head"} + if qkv_geometry is not None: + num_attention_heads, num_kv_heads, head_dim = qkv_geometry + if ( + num_attention_heads is not None + and num_kv_heads is not None + and head_dim is not None + ): + q_heads = int(num_attention_heads) + kv_heads = int(num_kv_heads) + qkv_head_dim = int(head_dim) + if ( + q_heads < 1 + or kv_heads < 1 + or qkv_head_dim < 1 + or q_heads % kv_heads + ): + raise ValueError( + f"{name}: invalid global Q/KV geometry " + f"{(q_heads, kv_heads, qkv_head_dim)}" + ) + expected_global_rows = (q_heads + 2 * kv_heads) * qkv_head_dim + actual_global_rows = int(param.shape[0]) * int(tp_size) + if expected_global_rows != actual_global_rows: + raise ValueError( + f"{name}: fused QKV rows {actual_global_rows} disagree " + f"with global head geometry {expected_global_rows}" + ) + extras.update( + { + "num_heads": str(q_heads), + "num_kv_heads": str(kv_heads), + "head_dim": str(qkv_head_dim), + } + ) + # Preserve compatibility with old MX clients only when local + # head counts are meaningful. Never publish a zero KV count. + if tp_size > 0 and q_heads % tp_size == 0 and kv_heads % tp_size == 0: + extras["num_heads_local"] = str(q_heads // tp_size) + extras["num_kv_heads_local"] = str(kv_heads // tp_size) + return MegatronRoleSpec(role=ROLE_QKV_COLUMN, descriptor_extras=extras) + if _is_fused_gated_mlp_name(name): + return MegatronRoleSpec( + role=ROLE_GATED_MLP_COLUMN, + descriptor_extras={"gated_mlp_order": "gate_then_up"}, + ) + return MegatronRoleSpec(role=ROLE_COLUMN) + + # ---- 6. Row-parallel linears. ---- + if parallelism == "row": + return MegatronRoleSpec(role=ROLE_ROW) + + # ---- 7. Replicated (LayerNorms, biases, scalars, routers, etc.). ---- + # Bridge's registry covers TENorm, FusedLayerNorm, WrappedTorchNorm, + # LayerNorm, RMSNorm, L2Norm, InferenceTopKRouter, IdentityOp, + # LinearForLastLayer, TopKRouter — anything unclassified here also + # falls into "replicated" as a safe default (rank 0 publishes; others + # skip), since misclassifying a sharded tensor as replicated would + # silently produce wrong logits while misclassifying a replicated + # tensor stays correct (just wastes one rank's publish bandwidth). + return MegatronRoleSpec(role=ROLE_REPLICATED) + + +def collect_megatron_publish_set( + model: "torch.nn.Module", + *, + tp_size: int, + pp_size: int, + pp_rank: int, + ep_size: int, + ep_rank: int, + tp_rank: int, + num_local_experts: int | None = None, + num_attention_heads: int | None = None, + num_kv_heads: int | None = None, + head_dim: int | None = None, + qkv_geometry_resolver: QkvGeometryResolver | None = None, + expert_pattern: str | None = None, + role_overrides: dict[str, str] | None = None, + target_dtype: "torch.dtype | None" = None, +) -> Iterator[tuple[str, "torch.Tensor", MegatronRoleSpec, dict[str, str]]]: + """Yield ``(name, local_shard, role_spec, full_extras)`` for the publisher. + + For each parameter: + + * Skips replicated tensors when ``tp_rank != 0``. The MX Megatron receiver + handles rank-0 replicated model tensors specially; publishing local + copies from non-zero TP ranks can make vLLM's rank-local loader treat + them as global tensors and slice past the end. + * Returns the parameter as-is — Megatron stores native shards, so + the param tensor IS the local shard. No allgather, no Bridge call. + * ``full_extras`` is the merged ``{megatron_role, tp_rank, tp_size, + pp_rank, pp_size, ep_rank, ep_size, ...}`` metadata consumed by + ``mx_reshard_publisher.build_megatron_alias_inputs``. + """ + for raw_name, param in model.named_parameters(): + if not param.is_floating_point(): + # Skip non-float buffers (rotary inv_freq, etc.); they aren't + # weight-refit material. + continue + + # `model.named_parameters()` returns names with a `module.` prefix + # when the model is wrapped (DDP-style). Two distinct uses of the + # name: + # + # 1. The model-walking classifier needs the ORIGINAL prefixed + # name to descend through `model.module.decoder.layers...` — + # stripping the prefix breaks `_enclosing_module` and every + # non-expert tensor falls to ROLE_REPLICATED. + # 2. The PUBLISHED name on the catalog has to match Bridge's + # name_map (which uses unprefixed names from + # `get_conversion_tasks`) so the receiver's name-map lookup + # finds the HF target names. + # + # Classify with `raw_name`; publish with the stripped form. + # (Bug surfaced on Qwen3-MoE-30B-A3B on 2026-06-10: the + # previous version stripped before classification and the + # receiver saw only `expert_column` / `replicated` because + # every TP-sharded role fell through to the default.) + name = raw_name + while name.startswith("module."): + # Megatron can wrap a chunk more than once (for example a local + # Float16Module around a distributed-data-parallel module). Bridge + # conversion tasks are keyed from the unwrapped module, so every + # leading wrapper component must be removed, not only the first. + name = name[len("module.") :] + + qkv_geometry = ( + qkv_geometry_resolver(raw_name, param, model) + if qkv_geometry_resolver is not None + else None + ) + spec = detect_megatron_role( + raw_name, + param, + model=model, + tp_size=tp_size, + ep_size=ep_size, + ep_rank=ep_rank, + num_local_experts=num_local_experts, + num_attention_heads=num_attention_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + qkv_geometry=qkv_geometry, + expert_pattern=expert_pattern, + role_overrides=role_overrides, + ) + if spec.is_expert: + name = canonicalize_grouped_expert_name(name, spec.descriptor_extras) + + if spec.role == ROLE_REPLICATED and tp_rank != 0: + continue + + local = param.detach() + if target_dtype is not None and local.dtype != target_dtype: + local = local.to(target_dtype, non_blocking=True) + local = local.contiguous() + + full_extras: dict[str, str] = { + "megatron_role": spec.role, + "tp_rank": str(tp_rank), + "tp_size": str(tp_size), + "pp_rank": str(pp_rank), + "pp_size": str(pp_size), + "ep_rank": str(ep_rank), + "ep_size": str(ep_size), + } + full_extras.update(spec.descriptor_extras) + + yield name, local, spec, full_extras + + +__all__ = [ + "MegatronRoleSpec", + "MegatronTpShardGeometry", + "ROLE_COLUMN", + "ROLE_EXPERT_COLUMN", + "ROLE_EXPERT_ROW", + "ROLE_GATED_MLP_COLUMN", + "ROLE_QKV_COLUMN", + "ROLE_REPLICATED", + "ROLE_ROW", + "ROLE_VOCAB_PARALLEL", + "canonicalize_grouped_expert_name", + "collect_megatron_publish_set", + "detect_megatron_role", + "infer_megatron_tp_shard_geometry", +] diff --git a/nemo_rl/distributed/mx_refit_verify.py b/nemo_rl/distributed/mx_refit_verify.py new file mode 100644 index 0000000000..a58ec736bd --- /dev/null +++ b/nemo_rl/distributed/mx_refit_verify.py @@ -0,0 +1,131 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. + +"""Parameter-equality verification for the ``mx_reshard`` refit path. + +Every other refit transport has a way to check the weights it moved: SGLang has +``check_weights(compare)`` and the sparse transports emit ``delta_verify/*``. +``mx_reshard`` had only the end-to-end logprob metrics, which conflate refit +fidelity with Megatron-vs-vLLM implementation divergence and so cannot answer +"did the transport install the right bytes" on their own. + +The check exploits a property of the first refit. A fresh run loads vLLM from the +HF checkpoint, and the trainer's Megatron weights are converted from that same +checkpoint with no optimizer step taken yet. So **the first refit should leave +every parameter it touches unchanged.** Anything it does change is refit or +conversion error, and it is named. + +From the second refit on the trainer has genuinely moved, so parameters are +*expected* to change and the same records instead answer a different question: +which parameters the refit reaches at all. A parameter that never changes across +many steps of training is one the refit is silently not updating - the failure +mode the coverage counters are meant to catch, cross-checked against the values. + +Fingerprints rather than copies: retaining a pre-refit copy of a 30B model's +shard costs ~15 GB per rank, while these are two allocation-free reductions. +""" + +from __future__ import annotations + +import json +import os + +_ENV_FLAG = "MX_REFIT_VERIFY" + + +def enabled() -> bool: + """Whether parameter-equality verification is enabled. + + Off unless asked for: the reductions are cheap but not free, and this runs on + the refit critical path. + """ + return os.environ.get(_ENV_FLAG, "0") not in ("", "0", "false", "False") + + +def fingerprint(tensor) -> tuple[int, int, int]: + """``(numel, sum of all raw bytes, sum of the high byte of each element)``. + + Computed over the raw bytes, so it is sensitive to any bit change rather than + only to changes large enough to move a float statistic. Both sums accumulate + in int64 via ``sum(dtype=...)``, which avoids materializing a promoted copy of + a tensor that can be hundreds of MB. + + The second statistic strides by element size to pick out each element's most + significant byte -- sign, exponent and the top mantissa bits for a + little-endian float -- so a change confined to the high bits cannot cancel + against an offsetting change elsewhere in the flat byte sum. + + This is a fingerprint, not a hash: collisions are possible in principle. Two + independent statistics make an accidental collision unlikely, and a *silent* + collision would have to survive both. + """ + import torch + + flat = tensor.detach().reshape(-1) + if not flat.is_contiguous(): + flat = flat.contiguous() + raw = flat.view(torch.uint8) + elsize = max(1, raw.numel() // max(1, flat.numel())) + high = raw[elsize - 1 :: elsize] if elsize > 1 else raw + return ( + int(flat.numel()), + int(raw.sum(dtype=torch.int64).item()), + int(high.sum(dtype=torch.int64).item()), + ) + + +def fingerprint_model(model) -> dict[str, tuple[int, int, int]]: + """Fingerprint every parameter, or ``{}`` if the model cannot be walked. + + Never raises: this is verification, and a verification failure must not be + able to fail the refit it is verifying. + + Returning ``{}`` disables the check for this refit, so say so rather than + going quiet. A verification tool that silently switches itself off reports + the same thing as one that is passing. + """ + try: + return {name: fingerprint(p) for name, p in model.named_parameters()} + except Exception as error: # noqa: BLE001 - see above + print( + "MX_REFIT_VERIFY skipped: could not fingerprint parameters " + f"({type(error).__name__}: {error})", + flush=True, + ) + return {} + + +def compare( + before: dict[str, tuple[int, int, int]], + after: dict[str, tuple[int, int, int]], +) -> dict: + """Which parameters the refit changed, and which it left alone.""" + shared = [name for name in after if name in before] + changed = [name for name in shared if before[name] != after[name]] + unchanged = [name for name in shared if before[name] == after[name]] + return { + "params_compared": len(shared), + "params_changed": len(changed), + "params_unchanged": len(unchanged), + "changed_sample": sorted(changed)[:10], + "unchanged_sample": sorted(unchanged)[:10], + } + + +def report(step: int, rank: int, before: dict, after: dict) -> dict: + """Emit one ``MX_REFIT_VERIFY`` record and return it. + + ``first_refit`` marks the record where "unchanged" is the passing outcome, so a + later reader does not have to reconstruct which expectation applied. + """ + record = { + "schema": "mx-refit-verify-v1", + "step": step, + "rank": rank, + "first_refit": step <= 1, + **compare(before, after), + } + print("MX_REFIT_VERIFY " + json.dumps(record), flush=True) + return record diff --git a/nemo_rl/distributed/mx_reshard_config.py b/nemo_rl/distributed/mx_reshard_config.py new file mode 100644 index 0000000000..7531942764 --- /dev/null +++ b/nemo_rl/distributed/mx_reshard_config.py @@ -0,0 +1,107 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Lightweight ModelExpress reshard endpoint configuration helpers.""" + +from __future__ import annotations + +from typing import Any + +DEFAULT_PUBLISHER_LISTEN_PORT_BASE = 5555 +DEFAULT_RECEIVER_LISTEN_PORT_BASE = 15555 +LEGACY_RECEIVER_PORT_OFFSET = 10000 +MAX_TCP_PORT = 65535 + + +def _get(config: Any, name: str, default: Any = None) -> Any: + if isinstance(config, dict): + return config.get(name, default) + return getattr(config, name, default) + + +def maybe_preinit_mx_reshard_nixl(config: Any) -> Any: + """Load NIXL/UCX before model construction when MX reshard is selected. + + On GB200, loading Megatron first can leave the later NIXL agent without + UCX's CUDA and InfiniBand components. Keep this agent alive for the worker + lifetime so the production publisher can create its listening agent after + the model and Bridge mappings exist. + """ + generation = _get(config, "generation") + if generation is None or _get(generation, "refit_transport") != "mx_reshard": + return None + + from nemo_rl.utils.checkpoint_engines.nixl import preinit_nixl_agent + + return preinit_nixl_agent() + + +def resolve_mx_reshard_listen_port_bases(config: Any) -> tuple[int, int]: + """Resolve distinct publisher/receiver bases, including the legacy fallback.""" + legacy = _get(config, "listen_port_base") + publisher = _get(config, "publisher_listen_port_base") + receiver = _get(config, "receiver_listen_port_base") + + if publisher is None: + publisher = legacy if legacy is not None else DEFAULT_PUBLISHER_LISTEN_PORT_BASE + if receiver is None: + receiver = ( + int(legacy) + LEGACY_RECEIVER_PORT_OFFSET + if legacy is not None + else DEFAULT_RECEIVER_LISTEN_PORT_BASE + ) + return int(publisher), int(receiver) + + +def resolve_mx_reshard_publisher_listen_port_base(config: Any) -> int: + """Resolve the base passed only to trainer publishers.""" + return resolve_mx_reshard_listen_port_bases(config)[0] + + +def resolve_mx_reshard_receiver_listen_port_base(config: Any) -> int: + """Resolve the base passed only to inference receivers.""" + return resolve_mx_reshard_listen_port_bases(config)[1] + + +def validate_mx_reshard_listen_port_ranges( + config: Any, + *, + train_world_size: int, + inference_world_size: int, +) -> tuple[int, int]: + """Require valid, disjoint physical-rank port ranges before actor startup.""" + publisher, receiver = resolve_mx_reshard_listen_port_bases(config) + if train_world_size <= 0 or inference_world_size <= 0: + raise ValueError( + "ModelExpress train and inference world sizes must be positive" + ) + + publisher_range = (publisher, publisher + train_world_size - 1) + receiver_range = (receiver, receiver + inference_world_size - 1) + for role, port_range in ( + ("publisher", publisher_range), + ("receiver", receiver_range), + ): + if port_range[0] <= 0 or port_range[1] > MAX_TCP_PORT: + raise ValueError( + f"ModelExpress {role} listen port range {port_range} is outside " + f"1..{MAX_TCP_PORT}" + ) + if max(publisher_range[0], receiver_range[0]) <= min( + publisher_range[1], receiver_range[1] + ): + raise ValueError( + "ModelExpress publisher and receiver listen port ranges overlap: " + f"publisher={publisher_range}, receiver={receiver_range}" + ) + return publisher, receiver + + +__all__ = [ + "DEFAULT_PUBLISHER_LISTEN_PORT_BASE", + "DEFAULT_RECEIVER_LISTEN_PORT_BASE", + "maybe_preinit_mx_reshard_nixl", + "resolve_mx_reshard_listen_port_bases", + "resolve_mx_reshard_publisher_listen_port_base", + "resolve_mx_reshard_receiver_listen_port_base", + "validate_mx_reshard_listen_port_ranges", +] diff --git a/nemo_rl/distributed/mx_reshard_publisher.py b/nemo_rl/distributed/mx_reshard_publisher.py new file mode 100644 index 0000000000..1de55056d6 --- /dev/null +++ b/nemo_rl/distributed/mx_reshard_publisher.py @@ -0,0 +1,455 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Publish a live Megatron layout through ModelExpress `main`'s reshard seam. + +`mx_megatron_helpers.collect_megatron_publish_set` classifies each native +Megatron parameter; ModelExpress accepts that classification through +``refit.reshard.megatron_aliases.build_hf_aliases``, which turns native storage +into HF-canonical shard records without copying, and publishes it through +``publish_registered_shard_table``. This module is the translation between the +two, plus the name resolution that neither side owns. + +ModelExpress owns alias construction and publication; NeMo-RL owns deriving +Megatron's native parameter ownership and resolving Megatron-Bridge names. This +module is the narrow adapter between those contracts. + +The two sides use the same role vocabulary (``qkv_column``, +``gated_mlp_column``, ``expert_column``, ...) and the same extras keys +(``head_dim``, global ``num_heads`` / ``num_kv_heads``, optional divisible +local-head counts, and ``gated_mlp_order``), so no renaming happens here. +""" + +from __future__ import annotations + +import json +import re +import time +from inspect import signature +from typing import Any, Callable, Iterable, Iterator + +from modelexpress.refit.reshard.megatron_aliases import ( + MegatronAliasInput, + build_hf_aliases, +) +from modelexpress.refit.reshard.megatron_publisher import ( + publish_registered_shard_table, +) +from modelexpress.refit.reshard.rendezvous import wrap_rendezvous_blob + +from nemo_rl.distributed.mx_megatron_helpers import ( + MegatronRoleSpec, + infer_megatron_tp_shard_geometry, +) + +# build_hf_aliases treats any placement that is not exactly "SHARD" as +# replicated, so the replicated spelling is ours to choose and only has to be +# stable. +PLACEMENT_SHARD = "SHARD" +PLACEMENT_REPLICATE = "REPLICATE" + +# HF MoE parameters carry the expert index as a path component. Megatron's +# grouped layout carries it as a leaf suffix (``linear_fc1.weight0``), which +# `canonicalize_grouped_expert_name` has already rewritten to the global id by +# the time we see it. +_HF_EXPERT_INDEX = re.compile(r"(?<=\.experts\.)\d+(?=\.)") + +HfNameResolver = Callable[[str, dict[str, str]], "tuple[str, ...]"] + + +class UnmappedMegatronTensor(KeyError): + """A published parameter has no HF target names. + + Fails the publish rather than skipping the tensor. A skipped source is not + visibly broken: the receiver simply never sees those bytes, its coverage + check reports a shortfall far from the cause, and with the coverage floor off + it reports nothing at all. + """ + + +class MxMegatronPublisher: + """Own one rank's NIXL registration, publication heartbeat, and MX client.""" + + def __init__( + self, + *, + items: list[MegatronAliasInput], + model_name: str, + server_url: str, + rank: int, + device_id: int, + listen_port: int, + metadata_endpoint: str, + ) -> None: + from modelexpress.client import MxClient + from modelexpress.nixl_transfer import NixlTransferManager + from modelexpress.refit.reshard.rendezvous import MxReshardRendezvous + + worker_id = f"nemo-rl-trainer-{rank}" + self._items = items + self._rank = rank + self._metadata_endpoint = metadata_endpoint + self._manager = NixlTransferManager( + agent_name=worker_id, + device_id=device_id, + listen_port=listen_port, + ) + self._manager.initialize() + self._client = MxClient(server_url=server_url) + try: + self._manager.register_tensors( + {item.name: item.tensor for item in self._items} + ) + self._rendezvous = MxReshardRendezvous( + self._client, + role="trainer", + rank=rank, + model_name=model_name, + worker_id=worker_id, + ) + except Exception: + self._client.close() + self._manager.shutdown() + raise + self._closed = False + + def publish(self, version: int) -> None: + if self._closed: + raise RuntimeError("ModelExpress publisher is closed") + started = time.perf_counter() + _source_id, published = publish_megatron_hf_aliases( + manager=self._manager, + rendezvous=self._rendezvous, + items=self._items, + metadata_endpoint=self._metadata_endpoint, + publisher_step=version, + ) + self._report_publish(version, time.perf_counter() - started, published) + + def _report_publish(self, version: int, elapsed_s: float, published: list) -> None: + """Emit the publish-side counterpart of the receiver's phase record. + + The weight synchronizer already times this half, but elapsed time alone + does not distinguish a publish that described more shards from one that + described the same shards more slowly. Never raises: telemetry must not + be able to fail the operation it measures. + """ + try: + shards = sum(len(tensor.shards) for tensor in published) + print( + "MX_PUBLISH_PHASE " + + json.dumps( + { + "schema": "mx-publish-phase-v1", + "step": version, + "rank": self._rank, + "publish_s": round(elapsed_s, 6), + "tensors": len(published), + "shards": shards, + "bytes": published_byte_count(published), + } + ), + flush=True, + ) + except Exception: # noqa: BLE001 - never worth a failed publish + pass + + def shutdown(self) -> None: + if self._closed: + return + self._closed = True + try: + self._rendezvous.close() + finally: + try: + self._manager.shutdown() + finally: + self._client.close() + + +def build_bridge_name_map( + conversion_tasks: Iterable[Any], +) -> dict[str, tuple[str, ...]]: + """Map this PP rank's unwrapped local Megatron names to ordered HF names. + + ``WeightConversionTask.param_name`` is the local name emitted by this + rank's ``model.named_parameters()`` after the collector strips ``module.``; + ``global_param_name`` is intentionally not the key because PP ranks can + renumber local layers. Tasks with ``param_weight is None`` belong to another + PP rank and are excluded. QKV dictionary order is semantic, not incidental. + """ + name_map: dict[str, tuple[str, ...]] = {} + for task in conversion_tasks: + if task is None or getattr(task, "param_weight", None) is None: + continue + name = str(task.param_name) + hf_param = getattr(task.mapping, "hf_param", None) + if isinstance(hf_param, str): + names = (hf_param,) + elif isinstance(hf_param, dict): + names = ( + tuple(hf_param[key] for key in ("q", "k", "v")) + if set(hf_param) == {"q", "k", "v"} + else tuple(hf_param.values()) + ) + else: + continue + previous = name_map.get(name) + if previous is not None and previous != names: + raise ValueError( + f"conflicting Bridge mappings for local Megatron name {name!r}: " + f"{previous!r} vs {names!r}" + ) + name_map[name] = names + return name_map + + +def make_bridge_resolver( + name_map: dict[str, Iterable[str]], + *, + strict: bool = True, +) -> HfNameResolver: + """Resolve a published Megatron name to HF target names via a Bridge map. + + ``name_map`` is what ``AutoBridge.get_conversion_tasks`` yields, keyed by + Megatron parameter name. Grouped-expert names need two steps of care: + + * the key may carry the **EP-local** expert index, because the Bridge + inspects one rank's module tree, while the published name carries the + **global** index that `canonicalize_grouped_expert_name` substituted. So a + miss on the global name is retried against the local one. + * the HF names that come back then describe the local expert, and the global + index has to be substituted into them, or every EP rank publishes over + rank 0's experts. + """ + resolved: dict[str, tuple[str, ...]] = { + key: tuple(value) for key, value in name_map.items() + } + + def resolve(name: str, extras: dict[str, str]) -> tuple[str, ...]: + hit = resolved.get(name) + if hit is not None: + return hit + + global_id = extras.get("expert_id") + local_id = extras.get("local_expert_id") + if global_id is not None and local_id is not None: + local_name = _swap_expert_leaf(name, global_id, local_id) + hit = resolved.get(local_name) + if hit is not None: + # The map describes local expert `local_id`; this rank owns + # global `global_id`. + return tuple( + _HF_EXPERT_INDEX.sub(str(global_id), hf_name) for hf_name in hit + ) + + if strict: + raise UnmappedMegatronTensor( + f"{name!r} has no HF target names. Nearest keys: " + f"{_nearest_keys(name, resolved)}" + ) + return () + + return resolve + + +def _swap_expert_leaf(name: str, from_id: str, to_id: str) -> str: + """Rewrite a trailing ``weight`` to ``weight``.""" + parent, separator, leaf = name.rpartition(".") + for prefix in ("weight", "bias", "scale"): + if leaf == f"{prefix}{from_id}": + swapped = f"{prefix}{to_id}" + return f"{parent}{separator}{swapped}" if separator else swapped + return name + + +def _nearest_keys(name: str, resolved: dict[str, tuple[str, ...]], limit: int = 3): + """A few same-suffix keys, so an unmapped name is diagnosable from the log.""" + leaf = name.rpartition(".")[2] + near = [key for key in resolved if key.rpartition(".")[2] == leaf] + return sorted(near)[:limit] or sorted(resolved)[:limit] + + +# Roles whose Megatron parameter fuses the gated MLP's two projections into one +# tensor, so MX has to split it and assign the halves to two HF names. +_GATED_ROLES = frozenset({"gated_mlp_column", "expert_column"}) + + +def _gated_mlp_extras( + name: str, role: str, hf_names: tuple[str, ...] +) -> dict[str, str]: + """The ``gated_mlp_order`` stamp MX requires for a fused gate/up parameter. + + MX refuses to infer this, correctly: it assigns the first half of the fused + tensor to ``hf_names[0]`` and the second to ``hf_names[1]``, so if the storage + order is actually the other way round it publishes the gate projection's bytes + under the up projection's name. Both names then receive exactly the bytes their + publisher advertised, so every digest agrees and the model is simply wrong. + + Megatron-Core stores a gated ``linear_fc1`` as ``[gate; up]`` concatenated on + the output axis -- that is the layout its SwiGLU expects when it chunks the + activation in two. So the order is known, but the *name* order is not ours: the + Bridge supplies ``hf_names``, and a mapping that listed up before gate would + make the stamp a lie. Rather than trust it, check the names look like the + order being claimed, and refuse when they do not -- an unrecognised naming + convention should stop the publish, not silently transpose a projection. + """ + if role not in _GATED_ROLES or len(hf_names) != 2: + return {} + first, second = hf_names[0].lower(), hf_names[1].lower() + if "gate" in first and "up" in second: + return {"gated_mlp_order": "gate_then_up"} + raise ValueError( + f"{name}: role {role!r} fuses a gated MLP, but its HF names " + f"{hf_names!r} do not read as (gate, up). Megatron stores this parameter " + f"as [gate; up] and MX assigns the halves positionally, so publishing " + f"under an unverified name order would transpose the two projections " + f"undetectably." + ) + + +def build_megatron_alias_inputs( + publish_set: Iterable[tuple[str, Any, MegatronRoleSpec, dict[str, str]]], + *, + resolve_hf_names: HfNameResolver, + tp_size: int, + tp_rank: int, + expert_tp_size: int | None = None, + expert_tp_rank: int | None = None, +) -> Iterator[MegatronAliasInput]: + """Translate a classified publish set into MX alias inputs. + + Geometry comes from `infer_megatron_tp_shard_geometry`, which returns None + for a tensor this rank holds whole. That covers two different situations + that MX describes the same way: + + * genuinely replicated tensors (norms, router gates); + * expert tensors under EP with expert-TP of 1. Megatron's grouped layout + gives each expert its own parameter, so EP partitions *names*, not an + axis, and each rank owns its experts entire. Fan-in across EP ranks then + happens by name in the rendezvous merge, with no shard arithmetic. + + Non-expert tensors under EP are byte-identical across every EP rank, so the + fleet publishes DP replicas of them. That is intended: the receiver's merge + deduplicates by geometry and reads each from one owner. + """ + for name, tensor, spec, extras in publish_set: + hf_names = resolve_hf_names(name, extras) + if not hf_names: + # Only reachable through a deliberately non-strict resolver. The + # strict one -- what the worker builds -- raises + # UnmappedMegatronTensor instead, because silently dropping a + # parameter surfaces far from its cause. + continue + + alias_extras = dict(extras) + alias_extras.update(_gated_mlp_extras(name, spec.role, tuple(hf_names))) + + geometry = infer_megatron_tp_shard_geometry( + local_shape=tuple(int(dim) for dim in tensor.shape), + role=spec.role, + tp_size=tp_size, + tp_rank=tp_rank, + expert_tp_size=expert_tp_size, + expert_tp_rank=expert_tp_rank, + descriptor_extras=spec.descriptor_extras, + ) + + if geometry is None: + yield MegatronAliasInput( + name=name, + tensor=tensor, + role=spec.role, + hf_names=tuple(hf_names), + global_shape=tuple(int(dim) for dim in tensor.shape), + placement_kind=PLACEMENT_REPLICATE, + shard_axis=None, + local_shard_range=None, + extras=alias_extras, + ) + continue + + yield MegatronAliasInput( + name=name, + tensor=tensor, + role=spec.role, + hf_names=tuple(hf_names), + global_shape=tuple(geometry.global_shape), + placement_kind=PLACEMENT_SHARD, + shard_axis=int(geometry.shard_axis), + local_shard_range=tuple(geometry.local_shard_range), + extras=alias_extras, + ) + + +def publish_megatron_hf_aliases( + *, + manager: Any, + rendezvous: Any, + items: list[MegatronAliasInput], + metadata_endpoint: str, + publisher_step: int, +) -> tuple[str, list]: + """Alias a registered Megatron layout as HF shards and publish it. + + Returns the source id and the published table, the latter so a caller can + report byte and shard counts without rebuilding it. + + The caller owns both the NIXL manager and the rendezvous. Publishing starts + the source's READY heartbeat, and only the rendezvous owner can stop it, so a + rendezvous created in here would leave that thread alive with no handle and + the source would go stale only at interpreter exit. + """ + if not items: + raise ValueError("no alias inputs to publish") + if publisher_step < 0: + raise ValueError("publisher_step must be non-negative") + + published = build_hf_aliases(items, agent_name=str(manager.agent_name)) + + # ModelExpress main's rendezvous payload owns the version stamp, while the + # current convenience publisher does not expose it. Use MX's encoder rather + # than duplicating its wire format. Keep the convenience path for a future + # MX version that grows the parameter. + if "publisher_step" in signature(publish_registered_shard_table).parameters: + source_id = publish_registered_shard_table( + manager=manager, + rendezvous=rendezvous, + published=published, + metadata_endpoint=metadata_endpoint, + publisher_step=publisher_step, + ) + else: + blob = wrap_rendezvous_blob( + manager.nixl_metadata, + str(manager.agent_name), + metadata_endpoint, + published, + publisher_step=publisher_step, + ) + source_id = rendezvous.publish(blob) + return source_id, published + + +def published_byte_count(published: Iterable[Any]) -> int: + """Bytes described by a published table, counting each shard once.""" + total = 0 + for tensor in published: + for shard in tensor.shards: + count = 1 + for dim in shard.shape: + count *= int(dim) + total += count * int(tensor.elsize) + return total + + +__all__ = [ + "PLACEMENT_REPLICATE", + "PLACEMENT_SHARD", + "MxMegatronPublisher", + "UnmappedMegatronTensor", + "build_bridge_name_map", + "build_megatron_alias_inputs", + "make_bridge_resolver", + "publish_megatron_hf_aliases", + "published_byte_count", +] diff --git a/nemo_rl/distributed/mx_vllm_reshard_receiver.py b/nemo_rl/distributed/mx_vllm_reshard_receiver.py new file mode 100644 index 0000000000..e7cbc9edd2 --- /dev/null +++ b/nemo_rl/distributed/mx_vllm_reshard_receiver.py @@ -0,0 +1,232 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. + +"""Compatibility adapter around ModelExpress' vLLM reshard receiver. + +MX exposes ``VllmReshardReceiver.update_weights(step)`` but two things this path +needs are missing from its public surface: a way to require a publisher step +before the receiver's first lazy discovery, and a receiver shutdown method. + +So this adapter owns the extra metadata client used for the version preflight, +and confines the one private dependency it cannot avoid +(``receiver._manager.shutdown()``) to a boundary that is checked at +construction rather than assumed at teardown. If MX grows public equivalents, +this adapter is what shrinks. +""" + +from __future__ import annotations + +import json +import time +from inspect import signature +from typing import Any, Optional + + +class MxVllmReshardReceiver: + """Version-gated lifecycle around ModelExpress' reshard receiver API.""" + + def __init__( + self, + *, + model: Any, + vllm_config: Any, + model_config: Any, + model_name: str, + server_url: str, + agent_name: str, + local_rank: int, + global_rank: int, + num_trainer_sources: int, + device: Any, + listen_port: int, + timeout: float, + ) -> None: + from modelexpress.client import MxClient + from modelexpress.engines.vllm.refit.receiver import VllmReshardReceiver + from modelexpress.refit.reshard.rendezvous import MxReshardRendezvous + + self._num_trainer_sources = num_trainer_sources + self._timeout = timeout + self._global_rank = global_rank + # Resolved on first discovery and cached: the installed MX cannot change + # under a live receiver. + self._tensorless_discovery: Optional[bool] = None + # Held for parameter-equality verification, which needs the live params + # both before and after an install. MX's receiver keeps its own reference. + self._model = model + self._client = MxClient(server_url=server_url) + self._rendezvous = MxReshardRendezvous( + self._client, + role="inference", + rank=global_rank, + model_name=model_name, + ) + try: + self._receiver = VllmReshardReceiver( + model=model, + vllm_config=vllm_config, + model_config=model_config, + model_name=model_name, + mx_server=server_url, + agent_name=agent_name, + local_rank=local_rank, + global_rank=global_rank, + num_trainer_sources=num_trainer_sources, + device=device, + listen_port=listen_port, + timeout=timeout, + ) + except Exception: + self._client.close() + raise + manager = getattr(self._receiver, "_manager", None) + if manager is None or not callable(getattr(manager, "shutdown", None)): + self._client.close() + raise RuntimeError( + "ModelExpress compatibility error: " + "VllmReshardReceiver._manager.shutdown() is unavailable. This " + "adapter needs it to release the receiver's NIXL registrations; " + "check the installed modelexpress version." + ) + self._closed = False + + def update_weights(self, version: int) -> dict[str, Any]: + """Require every visible trainer stamp before MX can install weights.""" + # Timed separately because MX_REFIT_STAGE records cover only the install + # that follows, leaving the majority of a MoE refit unattributed. The + # quorum check costs one list_sources plus a get_metadata round-trip per + # trainer rank, and each response carries that rank's whole shard table, + # so it scales with sources rather than with bytes moved. + discover_t0 = time.perf_counter() + payloads = self._discover_for_quorum() + discover_s = time.perf_counter() - discover_t0 + observed = [payload.publisher_step for payload in payloads] + if len(observed) != self._num_trainer_sources or any( + step != version for step in observed + ): + raise RuntimeError( + f"ModelExpress publisher version mismatch: requested {version}, " + f"observed {observed}" + ) + from nemo_rl.distributed import mx_refit_verify + + verifying = mx_refit_verify.enabled() + before = mx_refit_verify.fingerprint_model(self._model) if verifying else {} + + install_t0 = time.perf_counter() + result = self._receiver.update_weights(version, timeout=self._timeout) + install_s = time.perf_counter() - install_t0 + + if verifying and before: + mx_refit_verify.report( + version, + self._global_rank, + before, + mx_refit_verify.fingerprint_model(self._model), + ) + self._report_phases(version, discover_s, install_s, payloads) + return result + + def _discover_for_quorum(self) -> list: + """Discover trainer ranks for the version check only. + + This check needs one integer per rank. It does not read shard geometry -- + MX's own receiver discovers that once in ``_prepare`` and keeps it -- so + asking for the shard tables here rebuilds, every step, a table that is + identical every step. On Qwen3-30B-A3B that is 78,760 entries across 16 + ranks, and skipping the rebuild removes ~0.8 s of a ~5.2 s check. + + Omits the flag against an MX that predates it, so this does not have to + land in lockstep with the client change. + """ + kwargs: dict[str, Any] = {"timeout": self._timeout} + if self._supports_tensorless_discovery(): + kwargs["with_tensors"] = False + return self._rendezvous.discover_trainers(self._num_trainer_sources, **kwargs) + + def _supports_tensorless_discovery(self) -> bool: + """Whether the installed MX accepts ``with_tensors``. + + Asked of the signature rather than by calling and catching TypeError: a + genuine TypeError raised *inside* discovery would otherwise be swallowed + and retried as the expensive full fetch, turning a real bug into an + unexplained per-step slowdown. + """ + if self._tensorless_discovery is None: + try: + parameters = signature(self._rendezvous.discover_trainers).parameters + except (TypeError, ValueError): + # Unintrospectable callable (C extension, some mocks). Assume the + # older contract; the full fetch is slower but always correct. + self._tensorless_discovery = False + else: + self._tensorless_discovery = "with_tensors" in parameters + return self._tensorless_discovery + + def _report_phases( + self, + version: int, + discover_s: float, + install_s: float, + payloads: list, + ) -> None: + """Emit the phase split, and never let doing so break a refit. + + Telemetry must not be able to fail the operation it measures. The first + version of this read ``payload.tensors`` directly and raised + AttributeError on any payload shape that lacked it, which would abort a + refit that had already succeeded. + """ + + def entries(payload) -> int: + # Prefer the recorded count: the quorum path deliberately does not + # build the shard tables, so len(payload.tensors) would read 0 and + # hide the figure that showed this cost tracks source count. + counter = getattr(payload, "entry_count", None) + if callable(counter): + return int(counter()) + return len(getattr(payload, "tensors", ()) or ()) + + try: + tensors_seen = sum(entries(p) for p in payloads) + except TypeError: + tensors_seen = -1 + try: + print( + "MX_RECV_PHASE " + + json.dumps( + { + "schema": "mx-recv-phase-v1", + "step": version, + "rank": self._global_rank, + "discover_s": round(discover_s, 6), + "mx_update_s": round(install_s, 6), + "tensors_seen": tensors_seen, + "trainer_sources": len(payloads), + } + ), + flush=True, + ) + except Exception: # noqa: BLE001 - reporting is never worth a failed refit + pass + + def shutdown(self) -> None: + """Release the receiver-owned NIXL manager and both gRPC clients.""" + if self._closed: + return + self._closed = True + manager = self._receiver._manager + try: + manager.shutdown() + finally: + receiver_client = getattr(self._receiver, "_mx_client", None) + if receiver_client is not None and callable( + getattr(receiver_client, "close", None) + ): + receiver_client.close() + self._client.close() + + +__all__ = ["MxVllmReshardReceiver"] diff --git a/nemo_rl/models/generation/vllm/checkpoint_engine.py b/nemo_rl/models/generation/vllm/checkpoint_engine.py index 4a9a42bdba..d075df6f41 100644 --- a/nemo_rl/models/generation/vllm/checkpoint_engine.py +++ b/nemo_rl/models/generation/vllm/checkpoint_engine.py @@ -33,12 +33,17 @@ NIXL_VLLM_WORKER = "nemo_rl.models.generation.vllm.vllm_backend.NixlVllmWorker" _NIXL_CONFIG_KEY = "nemo_rl_checkpoint_engine" +_MX_RESHARD_NIXL_KEY = "nemo_rl_mx_reshard_nixl" def configure_nixl_worker(config: VllmConfig, vllm_kwargs: dict[str, Any]) -> None: """Configure vLLM's worker hook for early NIXL initialization.""" checkpoint_config = checkpoint_engine_refit_config(config) - if checkpoint_config is None or checkpoint_config["backend"] != "nixl": + checkpoint_nixl = ( + checkpoint_config is not None and checkpoint_config["backend"] == "nixl" + ) + mx_reshard = config.get("refit_transport") == "mx_reshard" + if not checkpoint_nixl and not mx_reshard: return worker_cls = vllm_kwargs.setdefault("worker_cls", NIXL_VLLM_WORKER) @@ -49,20 +54,26 @@ def configure_nixl_worker(config: VllmConfig, vllm_kwargs: dict[str, Any]) -> No ) additional_config = dict(vllm_kwargs.get("additional_config") or {}) - additional_config[_NIXL_CONFIG_KEY] = checkpoint_config + if checkpoint_nixl: + additional_config[_NIXL_CONFIG_KEY] = checkpoint_config + if mx_reshard: + additional_config[_MX_RESHARD_NIXL_KEY] = True vllm_kwargs["additional_config"] = additional_config def preinit_nixl_from_vllm_config(vllm_config: Any) -> Any: """Create the NIXL preinit agent carried by a vLLM internal worker.""" checkpoint_config = vllm_config.additional_config.get(_NIXL_CONFIG_KEY) - if checkpoint_config is None: + mx_reshard = vllm_config.additional_config.get(_MX_RESHARD_NIXL_KEY, False) + if checkpoint_config is None and not mx_reshard: return None - from nemo_rl.utils.checkpoint_engines.nixl import ( - preinit_nixl_agent, - resolve_nixl_backend_kwargs, - ) + from nemo_rl.utils.checkpoint_engines.nixl import preinit_nixl_agent + + if checkpoint_config is None: + return preinit_nixl_agent() + + from nemo_rl.utils.checkpoint_engines.nixl import resolve_nixl_backend_kwargs backend_name, backend_init_params = resolve_nixl_backend_kwargs( checkpoint_config["engine_kwargs"]["nixl"] diff --git a/nemo_rl/models/generation/vllm/config.py b/nemo_rl/models/generation/vllm/config.py index f6786d1e8e..2f2a758694 100644 --- a/nemo_rl/models/generation/vllm/config.py +++ b/nemo_rl/models/generation/vllm/config.py @@ -14,12 +14,26 @@ from typing import Annotated, Any, Literal, NotRequired, TypedDict, cast, get_args -from pydantic import BaseModel, Field, NonNegativeInt, PositiveFloat, PositiveInt +from pydantic import ( + BaseModel, + Field, + NonNegativeInt, + PositiveFloat, + PositiveInt, + model_validator, +) +from nemo_rl.distributed.mx_reshard_config import ( + DEFAULT_PUBLISHER_LISTEN_PORT_BASE, + DEFAULT_RECEIVER_LISTEN_PORT_BASE, + LEGACY_RECEIVER_PORT_OFFSET, +) from nemo_rl.models.generation.interfaces import GenerationConfig VllmRefitTransportName = Literal["s3", "zmq"] -VllmRefitSelector = Literal["vllm_s3_sparse", "vllm_zmq_sparse", "nixl", "nccl_reshard"] +VllmRefitSelector = Literal[ + "vllm_s3_sparse", "vllm_zmq_sparse", "nixl", "nccl_reshard", "mx_reshard" +] VLLM_SPARSE_REFIT_TRANSPORTS = frozenset({"vllm_s3_sparse", "vllm_zmq_sparse"}) @@ -134,6 +148,34 @@ class VllmNixlRefitConfig(BaseModel, extra="forbid"): shard_expert_weights: bool = False +class VllmMxReshardConfig(BaseModel, extra="forbid"): + server_url: str | None = None + timeout_s: PositiveFloat = 1200.0 + publisher_listen_port_base: Annotated[int, Field(gt=0, le=65535)] = ( + DEFAULT_PUBLISHER_LISTEN_PORT_BASE + ) + receiver_listen_port_base: Annotated[int, Field(gt=0, le=65535)] = ( + DEFAULT_RECEIVER_LISTEN_PORT_BASE + ) + # Deprecated compatibility input. When supplied without a role-specific + # override, the receiver is shifted to a disjoint range. + listen_port_base: Annotated[int, Field(gt=0, le=55535)] | None = None + + @model_validator(mode="before") + @classmethod + def resolve_legacy_listen_port_base(cls, value: Any) -> Any: + if not isinstance(value, dict) or value.get("listen_port_base") is None: + return value + resolved = dict(value) + legacy = int(resolved["listen_port_base"]) + resolved.setdefault("publisher_listen_port_base", legacy) + resolved.setdefault( + "receiver_listen_port_base", + legacy + LEGACY_RECEIVER_PORT_OFFSET, + ) + return resolved + + class VllmCheckpointEnginePluginConfig(BaseModel, extra="allow"): update_weights_bucket_memory_ratio: Annotated[float, Field(gt=0, lt=1)] = 0.05 release_after_refit: bool = False @@ -142,6 +184,7 @@ class VllmCheckpointEnginePluginConfig(BaseModel, extra="allow"): class VllmRefitConfig(BaseModel, extra="allow"): sparse: VllmSparseRefitConfig = Field(default_factory=VllmSparseRefitConfig) nixl: VllmNixlRefitConfig = Field(default_factory=VllmNixlRefitConfig) + mx_reshard: VllmMxReshardConfig = Field(default_factory=VllmMxReshardConfig) class VllmConfig(GenerationConfig): @@ -177,12 +220,13 @@ def normalize_vllm_refit_config(config: VllmConfig) -> VllmRefitConfig | None: if transport is None: return None if transport == "nccl_reshard": - # nccl_reshard doesn't takes refit_cfg. + # nccl_reshard doesn't take refit_cfg. return None if transport not in get_args(VllmRefitSelector) and ":" not in transport: raise ValueError( f"Unknown vLLM refit transport {transport!r}: expected null, " - "'nccl_reshard', 'vllm_s3_sparse', 'vllm_zmq_sparse', 'nixl', or a " + "'nccl_reshard', 'mx_reshard', 'vllm_s3_sparse', " + "'vllm_zmq_sparse', 'nixl', or a " "'module:ClassName' checkpoint-engine path." ) # The encoder-cache reset is implemented only on the collective/IPC and diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index de790b642d..2a71834699 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -183,6 +183,7 @@ class VllmInternalWorkerExtension: # load_mtp_weights_from_disk); refit then leaves those static weights alone. _mtp_drafter_from_disk: bool = False _sparse_delta_applier: Any = None + _mx_reshard_receiver: Any = None _nrl_named_parameters: dict[str, torch.nn.Parameter] def _get_named_parameters(self) -> dict[str, torch.nn.Parameter]: @@ -249,6 +250,56 @@ def init_collective( torch.cuda.empty_cache() self.model_update_group.init_nccl_communicator(device=self.device) + def init_mx_reshard_receiver( + self, + rank_prefix: int, + inference_world_size: int, + train_world_size: int, + model_name: str, + server_url: str, + timeout_s: float, + receiver_listen_port_base: int, + ) -> bool: + """Create the ModelExpress receiver around this live vLLM model.""" + if self._mx_reshard_receiver is not None: + return True + + from nemo_rl.distributed.mx_vllm_reshard_receiver import ( + MxVllmReshardReceiver, + ) + + local_rank = torch.cuda.current_device() + rollout_rank = resolve_rollout_rank(rank_prefix, inference_world_size) + self._mx_reshard_receiver = MxVllmReshardReceiver( + model=self.model_runner.model, + vllm_config=self.model_runner.vllm_config, + model_config=self.model_runner.model_config, + model_name=model_name, + server_url=server_url, + agent_name=f"nemo-rl-vllm-{rollout_rank}", + local_rank=local_rank, + global_rank=rollout_rank, + num_trainer_sources=train_world_size, + device=self.device, + listen_port=int(receiver_listen_port_base) + rollout_rank, + timeout=float(timeout_s), + ) + return True + + def update_weights_from_mx_reshard(self, version: int) -> bool: + """Pull and atomically install one stamped ModelExpress version.""" + if self._mx_reshard_receiver is None: + raise RuntimeError("ModelExpress receiver is not initialized") + self._mx_reshard_receiver.update_weights(version) + return True + + def shutdown_mx_reshard_receiver(self) -> bool: + """Release this vLLM rank's ModelExpress receiver resources.""" + if self._mx_reshard_receiver is not None: + self._mx_reshard_receiver.shutdown() + self._mx_reshard_receiver = None + return True + def init_nccl_reshard_comm_group( self, rank_prefix: int, @@ -1108,6 +1159,9 @@ def _receive_and_load_misc_params(self) -> None: def cleanup(self) -> None: """Shutdown and cleanup resources.""" + if self._mx_reshard_receiver is not None: + self._mx_reshard_receiver.shutdown() + self._mx_reshard_receiver = None # Close ZMQ socket and context if they exist if hasattr(self, "zmq_socket"): self.zmq_socket.close() diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index a8cb78bd13..b8f8f793f7 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -1030,6 +1030,9 @@ def finish_generation(self, *args: Any, **kwargs: Any) -> bool: def shutdown(self) -> bool: """Shut down all vLLM workers and clean up resources.""" + if getattr(self, "_shutdown_complete", False): + return True + self._shutdown_complete = True try: if self.weight_synchronizer is not None: self.weight_synchronizer.shutdown() @@ -1173,6 +1176,57 @@ def nccl_reshard_refit(self) -> list[ray.ObjectRef]: ) return futures + def init_mx_reshard_receiver( + self, *, train_world_size: int, inference_world_size: int + ) -> list[ray.ObjectRef]: + """Initialize one MX receiver on every vLLM backend rank.""" + if not self.worker_group or not self.worker_group.workers: + raise RuntimeError("Worker group is not initialized") + method_name = ( + "init_mx_reshard_receiver_async" + if self.cfg["vllm_cfg"]["async_engine"] + else "init_mx_reshard_receiver" + ) + total_workers = len(self.worker_group.workers) + workers_per_group = total_workers // self.dp_size + rank_prefix_list = list(range(0, total_workers, workers_per_group)) + return self.worker_group.run_all_workers_multiple_data( + method_name, + rank_prefix=rank_prefix_list, + run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], + common_kwargs={ + "train_world_size": train_world_size, + "inference_world_size": inference_world_size, + }, + ) + + def update_weights_from_mx_reshard(self, *, version: int) -> list[ray.ObjectRef]: + """Pull a stamped MX version on all vLLM replicas.""" + if not self.worker_group or not self.worker_group.workers: + raise RuntimeError("Worker group is not initialized") + method_name = ( + "update_weights_from_mx_reshard_async" + if self.cfg["vllm_cfg"]["async_engine"] + else "update_weights_from_mx_reshard" + ) + return self.worker_group.run_all_workers_single_data( + method_name, + version=version, + run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], + ) + + def shutdown_mx_reshard_receiver(self) -> list[ray.ObjectRef]: + """Shut down MX receivers on all vLLM replicas.""" + method_name = ( + "shutdown_mx_reshard_receiver_async" + if self.cfg["vllm_cfg"]["async_engine"] + else "shutdown_mx_reshard_receiver" + ) + return self.worker_group.run_all_workers_single_data( + method_name, + run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], + ) + def start_gpu_profiling(self) -> None: """Start GPU profiling.""" futures = self.worker_group.run_all_workers_single_data("start_gpu_profiling") diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index ed3f045274..6e6e77a1bb 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -369,12 +369,11 @@ def _load_model(self, bundle_indices, seed): ) vllm_kwargs: dict[str, Any] = copy.deepcopy(self.cfg.get("vllm_kwargs", {})) checkpoint_engine_config = checkpoint_engine_refit_config(self.cfg) - if checkpoint_engine_config is not None: - from nemo_rl.models.generation.vllm.checkpoint_engine import ( - configure_nixl_worker, - ) + from nemo_rl.models.generation.vllm.checkpoint_engine import ( + configure_nixl_worker, + ) - configure_nixl_worker(self.cfg, vllm_kwargs) + configure_nixl_worker(self.cfg, vllm_kwargs) # A speculative_config with num_speculative_tokens == 0 is the supported # way to disable speculative decoding (e.g. MTP) from a launch script @@ -1258,6 +1257,70 @@ def nccl_reshard_refit(self) -> bool: traceback.print_exc() return False + def init_mx_reshard_receiver( + self, + rank_prefix: int, + inference_world_size: int, + train_world_size: int, + ) -> bool: + """Initialize ModelExpress receivers in the vLLM backend workers.""" + from nemo_rl.distributed.mx_reshard_config import ( + resolve_mx_reshard_receiver_listen_port_base, + ) + + refit_cfg = self.cfg.get("refit_cfg") or {} + mx_cfg = ( + refit_cfg.mx_reshard + if hasattr(refit_cfg, "mx_reshard") + else refit_cfg.get("mx_reshard", {}) + ) + server_url = ( + ( + getattr(mx_cfg, "server_url", None) + if not isinstance(mx_cfg, dict) + else mx_cfg.get("server_url") + ) + or os.environ.get("MX_SERVER_URL") + or os.environ.get("MODEL_EXPRESS_URL") + or os.environ.get("MX_SERVER_ADDRESS") + ) + if not server_url: + raise ValueError( + "mx_reshard requires refit_cfg.mx_reshard.server_url, " + "MX_SERVER_URL, MODEL_EXPRESS_URL, or MX_SERVER_ADDRESS" + ) + timeout_s = ( + getattr(mx_cfg, "timeout_s", 1200.0) + if not isinstance(mx_cfg, dict) + else mx_cfg.get("timeout_s", 1200.0) + ) + receiver_listen_port_base = resolve_mx_reshard_receiver_listen_port_base(mx_cfg) + results = self.llm.collective_rpc( + "init_mx_reshard_receiver", + args=( + rank_prefix, + inference_world_size, + train_world_size, + self.model_name, + server_url, + timeout_s, + receiver_listen_port_base, + ), + ) + return bool(results) and all(result is True for result in results) + + def update_weights_from_mx_reshard(self, version: int) -> bool: + """Pull one ModelExpress version on every backend rank.""" + results = self.llm.collective_rpc( + "update_weights_from_mx_reshard", args=(version,) + ) + return bool(results) and all(result is True for result in results) + + def shutdown_mx_reshard_receiver(self) -> bool: + """Shut down ModelExpress receivers on every backend rank.""" + results = self.llm.collective_rpc("shutdown_mx_reshard_receiver", args=tuple()) + return bool(results) and all(result is True for result in results) + def reset_prefix_cache(self): """Reset the prefix cache of vLLM engine.""" assert self.llm is not None, ( diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index e791fa4a02..0f1f121a0a 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -1553,6 +1553,72 @@ async def nccl_reshard_refit_async(self) -> bool: traceback.print_exc() return False + async def init_mx_reshard_receiver_async( + self, + rank_prefix: int, + inference_world_size: int, + train_world_size: int, + ) -> bool: + """Initialize ModelExpress receivers in async vLLM backend workers.""" + from nemo_rl.distributed.mx_reshard_config import ( + resolve_mx_reshard_receiver_listen_port_base, + ) + + refit_cfg = self.cfg.get("refit_cfg") or {} + mx_cfg = ( + refit_cfg.mx_reshard + if hasattr(refit_cfg, "mx_reshard") + else refit_cfg.get("mx_reshard", {}) + ) + server_url = ( + ( + getattr(mx_cfg, "server_url", None) + if not isinstance(mx_cfg, dict) + else mx_cfg.get("server_url") + ) + or os.environ.get("MX_SERVER_URL") + or os.environ.get("MODEL_EXPRESS_URL") + or os.environ.get("MX_SERVER_ADDRESS") + ) + if not server_url: + raise ValueError( + "mx_reshard requires refit_cfg.mx_reshard.server_url, " + "MX_SERVER_URL, MODEL_EXPRESS_URL, or MX_SERVER_ADDRESS" + ) + timeout_s = ( + getattr(mx_cfg, "timeout_s", 1200.0) + if not isinstance(mx_cfg, dict) + else mx_cfg.get("timeout_s", 1200.0) + ) + receiver_listen_port_base = resolve_mx_reshard_receiver_listen_port_base(mx_cfg) + results = await self.llm.collective_rpc( + "init_mx_reshard_receiver", + args=( + rank_prefix, + inference_world_size, + train_world_size, + self.model_name, + server_url, + timeout_s, + receiver_listen_port_base, + ), + ) + return bool(results) and all(result is True for result in results) + + async def update_weights_from_mx_reshard_async(self, version: int) -> bool: + """Pull one ModelExpress version on every async backend rank.""" + results = await self.llm.collective_rpc( + "update_weights_from_mx_reshard", args=(version,) + ) + return bool(results) and all(result is True for result in results) + + async def shutdown_mx_reshard_receiver_async(self) -> bool: + """Shut down ModelExpress receivers on every async backend rank.""" + results = await self.llm.collective_rpc( + "shutdown_mx_reshard_receiver", args=tuple() + ) + return bool(results) and all(result is True for result in results) + async def reset_prefix_cache_async(self): """Async version of reset_prefix_cache.""" assert self.llm is not None, ( diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index 81438afa3b..3d0b969df1 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -1178,6 +1178,28 @@ def nccl_reshard_refit(self, kv_scales=None) -> list[ray.ObjectRef]: ) return futures + def init_mx_reshard_publisher( + self, *, train_world_size: int + ) -> list[ray.ObjectRef]: + """Initialize one ModelExpress publisher per Megatron rank.""" + return self.worker_group.run_all_workers_single_data( + "init_mx_reshard_publisher", + train_world_size=train_world_size, + ) + + def publish_mx_reshard_weights(self, *, version: int) -> list[ray.ObjectRef]: + """Publish all trainer ranks at one stamped version.""" + return self.worker_group.run_all_workers_single_data( + "publish_mx_reshard_weights", + version=version, + ) + + def shutdown_mx_reshard_publisher(self) -> list[ray.ObjectRef]: + """Stop every trainer publisher and stale its MX publication.""" + return self.worker_group.run_all_workers_single_data( + "shutdown_mx_reshard_publisher" + ) + def offload_before_refit(self) -> None: """Offload the optimizer and buffers to the CPU.""" futures = self.worker_group.run_all_workers_single_data("offload_before_refit") @@ -1242,9 +1264,18 @@ def finalize_async_save(self) -> None: def shutdown(self) -> bool: """Shut down all HF workers and clean up resources.""" + if getattr(self, "_shutdown_complete", False): + return True + self._shutdown_complete = True if not hasattr(self, "worker_group"): return True try: + if self.cfg["generation"].get("refit_transport") == "mx_reshard": + ray.get( + self.worker_group.run_all_workers_single_data( + "shutdown_mx_reshard_publisher" + ) + ) # Use the worker group's shutdown method with the worker's cleanup method return self.worker_group.shutdown(cleanup_method="shutdown") except ray.exceptions.RayActorError: diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index e1b96eeced..65cec4381e 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -451,6 +451,12 @@ def __init__( self.cfg = config self._router_replay_enabled = router_replay_enabled(config) self._nixl_preinit_agent = maybe_preinit_nixl_checkpoint_engine(config) + if self._nixl_preinit_agent is None: + from nemo_rl.distributed.mx_reshard_config import ( + maybe_preinit_mx_reshard_nixl, + ) + + self._nixl_preinit_agent = maybe_preinit_mx_reshard_nixl(config) # Set rank for non-collocated to check which ranks to broadcast from self.rank = get_rank_safe() @@ -2135,6 +2141,130 @@ def prepare_refit_info(self) -> None: return refit_param_info_hf + @torch.no_grad() + def init_mx_reshard_publisher(self, train_world_size: int) -> bool: + """Register this rank's live Megatron shards with ModelExpress.""" + del train_world_size # The rendezvous rank comes from torch.distributed. + if getattr(self, "_mx_reshard_publisher", None) is not None: + return True + + # Keep ModelExpress optional for every non-MX NeMo-RL environment. + from nemo_rl.distributed.mx_megatron_helpers import ( + collect_megatron_publish_set, + resolve_qkv_geometry_from_param, + ) + from nemo_rl.distributed.mx_reshard_config import ( + resolve_mx_reshard_publisher_listen_port_base, + ) + from nemo_rl.distributed.mx_reshard_publisher import ( + MxMegatronPublisher, + build_bridge_name_map, + build_megatron_alias_inputs, + make_bridge_resolver, + ) + from nemo_rl.distributed.virtual_cluster import _get_node_ip_local + + refit_cfg = self.cfg["generation"].get("refit_cfg") or {} + mx_cfg = ( + refit_cfg.mx_reshard + if hasattr(refit_cfg, "mx_reshard") + else refit_cfg.get("mx_reshard", {}) + ) + server_url = ( + ( + getattr(mx_cfg, "server_url", None) + if not isinstance(mx_cfg, dict) + else mx_cfg.get("server_url") + ) + or os.environ.get("MX_SERVER_URL") + or os.environ.get("MODEL_EXPRESS_URL") + or os.environ.get("MX_SERVER_ADDRESS") + ) + if not server_url: + raise ValueError( + "mx_reshard requires refit_cfg.mx_reshard.server_url, " + "MX_SERVER_URL, MODEL_EXPRESS_URL, or MX_SERVER_ADDRESS" + ) + publisher_port_base = resolve_mx_reshard_publisher_listen_port_base(mx_cfg) + + rank = torch.distributed.get_rank() + device_id = torch.cuda.current_device() + model_config = get_model_config(self.model) + tp_size = int(self.cfg["megatron_cfg"].get("tensor_model_parallel_size", 1)) + pp_size = int(self.cfg["megatron_cfg"].get("pipeline_model_parallel_size", 1)) + ep_size = int(self.cfg["megatron_cfg"].get("expert_model_parallel_size", 1)) + etp_size = int(self.cfg["megatron_cfg"].get("expert_tensor_parallel_size", 1)) + tp_rank = parallel_state.get_tensor_model_parallel_rank() + pp_rank = parallel_state.get_pipeline_model_parallel_rank() + ep_rank = parallel_state.get_expert_model_parallel_rank() + etp_rank_getter = getattr( + parallel_state, "get_expert_tensor_parallel_rank", lambda: 0 + ) + + tasks = self._build_refit_conversion_tasks() + name_map = build_bridge_name_map(tasks) + + num_heads = getattr(model_config, "num_attention_heads", None) + num_kv_heads = getattr(model_config, "num_query_groups", None) or num_heads + head_dim = getattr(model_config, "kv_channels", None) + num_experts = getattr(model_config, "num_moe_experts", None) + publish_set = list( + collect_megatron_publish_set( + self.model, + tp_size=tp_size, + pp_size=pp_size, + pp_rank=pp_rank, + ep_size=ep_size, + ep_rank=ep_rank, + tp_rank=tp_rank, + num_local_experts=( + int(num_experts) // ep_size if num_experts is not None else None + ), + num_attention_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + qkv_geometry_resolver=resolve_qkv_geometry_from_param, + ) + ) + items = list( + build_megatron_alias_inputs( + publish_set, + resolve_hf_names=make_bridge_resolver(name_map), + tp_size=tp_size, + tp_rank=tp_rank, + expert_tp_size=etp_size, + expert_tp_rank=etp_rank_getter(), + ) + ) + + self._mx_reshard_publisher = MxMegatronPublisher( + items=items, + model_name=self.cfg["model_name"], + server_url=server_url, + rank=rank, + device_id=device_id, + listen_port=publisher_port_base + rank, + metadata_endpoint=f"{_get_node_ip_local()}:{publisher_port_base + rank}", + ) + return True + + @torch.no_grad() + def publish_mx_reshard_weights(self, version: int) -> bool: + """Publish this rank's currently live parameter storage at ``version``.""" + publisher = getattr(self, "_mx_reshard_publisher", None) + if publisher is None: + raise RuntimeError("ModelExpress publisher is not initialized") + publisher.publish(version) + return True + + def shutdown_mx_reshard_publisher(self) -> bool: + """Stop the rendezvous heartbeat and release NIXL registrations.""" + publisher = getattr(self, "_mx_reshard_publisher", None) + if publisher is not None: + publisher.shutdown() + self._mx_reshard_publisher = None + return True + def _collect_mtp_metrics( self, metrics: dict[str, Any], diff --git a/nemo_rl/weight_sync/checkpoint_engine_config.py b/nemo_rl/weight_sync/checkpoint_engine_config.py index 8bf7ed3c9c..b9c7b7da43 100644 --- a/nemo_rl/weight_sync/checkpoint_engine_config.py +++ b/nemo_rl/weight_sync/checkpoint_engine_config.py @@ -37,6 +37,7 @@ def checkpoint_engine_refit_config( refit_config is None or transport is None or transport in VLLM_SPARSE_REFIT_TRANSPORTS + or transport == "mx_reshard" ): return None diff --git a/nemo_rl/weight_sync/factory.py b/nemo_rl/weight_sync/factory.py index c48ed19f5d..033436b971 100644 --- a/nemo_rl/weight_sync/factory.py +++ b/nemo_rl/weight_sync/factory.py @@ -127,6 +127,22 @@ def create_weight_synchronizer( "for non-colocated weight synchronization." ) + if generation.cfg.get("refit_transport") == "mx_reshard": + if generation_backend != VLLM_BACKEND: + raise NotImplementedError( + "mx_reshard is only supported with the vLLM generation backend." + ) + from nemo_rl.weight_sync.mx_reshard_weight_synchronizer import ( + MxReshardWeightSynchronizer, + ) + + return MxReshardWeightSynchronizer( + policy=policy, + generation=generation, + train_cluster=train_cluster, + inference_cluster=inference_cluster, + ) + if generation.cfg.get("refit_transport") == "nccl_reshard": from nemo_rl.weight_sync.nccl_reshard_weight_synchronizer import ( NcclReshardWeightSynchronizer, diff --git a/nemo_rl/weight_sync/mx_reshard_weight_synchronizer.py b/nemo_rl/weight_sync/mx_reshard_weight_synchronizer.py new file mode 100644 index 0000000000..2e4509596f --- /dev/null +++ b/nemo_rl/weight_sync/mx_reshard_weight_synchronizer.py @@ -0,0 +1,219 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. + +"""ModelExpress shard-to-shard refit for Megatron trainer to vLLM generation.""" + +from contextlib import nullcontext +from typing import Any, Optional + +import ray + +from nemo_rl.utils.timer import Timer +from nemo_rl.weight_sync.interfaces import WeightSynchronizer + + +def check_mx_reshard_refit_support(master_config: Any) -> None: + """Validate the deliberately narrow first production MX reshard path.""" + policy = master_config.policy + generation = policy.get("generation", {}) or {} + megatron = policy.get("megatron_cfg", {}) or {} + dtensor = policy.get("dtensor_cfg", {}) or {} + vllm = generation.get("vllm_cfg", {}) or {} + violations: list[str] = [] + + if generation.get("backend") != "vllm": + violations.append("policy.generation.backend must be 'vllm'.") + if generation.get("colocated", {}).get("enabled", False): + violations.append("policy.generation.colocated.enabled must be False.") + if not megatron.get("enabled", False): + violations.append("policy.megatron_cfg.enabled must be True.") + if dtensor.get("enabled", False): + violations.append("policy.dtensor_cfg.enabled must be False.") + if megatron.get("expert_tensor_parallel_size", 1) not in (None, 1): + violations.append("policy.megatron_cfg.expert_tensor_parallel_size must be 1.") + if str(vllm.get("kv_cache_dtype", "auto")).startswith("fp8"): + violations.append("mx_reshard does not support FP8 KV-cache scale sync.") + + # Query heads still have to divide across trainer TP. KV heads do not: + # Megatron slices the globally interleaved fused QKV rows and MX maps each + # rank's raw interval, so KV 0 + and int(heads) % int(query_groups) + ): + violations.append( + "Megatron num_attention_heads must be divisible by num_query_groups." + ) + + if violations: + raise ValueError( + "mx_reshard refit configuration is unsupported:\n- " + + "\n- ".join(violations) + ) + + +class MxReshardWeightSynchronizer(WeightSynchronizer): + """Serialize a stamped publish quorum before the matching receiver quorum.""" + + def __init__( + self, + policy: Any, + generation: Any, + train_cluster: Any, + inference_cluster: Any, + ) -> None: + self._policy = policy + self._generation = generation + self._train_cluster = train_cluster + self._inference_cluster = inference_cluster + self._version = 0 + self._stale = True + self._shutdown = False + + @staticmethod + def _require_all(results: Any, phase: str) -> None: + leaves: list[Any] = [] + + def flatten(value: Any) -> None: + if isinstance(value, (list, tuple)): + for item in value: + flatten(item) + else: + leaves.append(value) + + flatten(results) + if not leaves or not all(result is True for result in leaves): + raise RuntimeError( + f"ModelExpress mx_reshard {phase} failed on at least one rank: " + f"{results!r}" + ) + + def sync_weights( + self, + *, + timer: Optional[Timer] = None, + kv_scales: Optional[dict[str, float]] = None, + ) -> None: + if kv_scales is not None: + raise ValueError("mx_reshard does not support FP8 KV-scale synchronization") + + # `_version` is the last *committed* version, so a failed refit leaves it + # alone and the next attempt reuses this stamp. That is safe only because + # a failure here aborts the run: publishers stamp the shard table with + # this number, so a retry inside one process would advertise different + # bytes under a version a receiver may already have seen. Anything that + # adds a retry has to advance the counter per attempt instead. + version = self._version + 1 + timer_context = ( + timer.time("prepare_for_generation/transfer_and_update_weights") + if timer is not None + else nullcontext() + ) + + def phase(name: str): + """Split the refit into its two serialized halves. + + Publish and receive are strictly sequential, so both sit on the + critical path, but only the receive half reports MX_REFIT_STAGE + telemetry. Without this split the publish cost is invisible: on + Qwen3-30B-A3B it is roughly 8.4 s of an 11 s refit, which is far + larger than the transfer MX does report, while on a dense 4B model it + is only ~0.5 s. Attributing a refit therefore requires this timer, not + just MX's stages. + """ + if timer is None: + return nullcontext() + return timer.time(f"prepare_for_generation/mx_reshard_{name}") + + with timer_context: + # This ordering is intentional and safety-critical. A receiver must + # never discover a mixed fleet while some trainers still advertise + # the previous version. + with phase("publish"): + published = ray.get( + self._policy.publish_mx_reshard_weights(version=version) + ) + self._require_all(published, "publish") + + with phase("receive"): + pulled = ray.get( + self._generation.update_weights_from_mx_reshard(version=version) + ) + self._require_all(pulled, "receive") + + self._version = version + self._stale = False + + @property + def is_stale(self) -> bool: + return self._stale + + def mark_stale(self) -> None: + self._stale = True + + def init_communicator(self) -> None: + from nemo_rl.distributed.mx_reshard_config import ( + validate_mx_reshard_listen_port_ranges, + ) + + train_world_size = self._train_cluster.world_size() + inference_world_size = self._inference_cluster.world_size() + refit_cfg = self._generation.cfg.get("refit_cfg") or {} + mx_cfg = ( + refit_cfg.mx_reshard + if hasattr(refit_cfg, "mx_reshard") + else refit_cfg.get("mx_reshard", {}) + ) + validate_mx_reshard_listen_port_ranges( + mx_cfg, + train_world_size=train_world_size, + inference_world_size=inference_world_size, + ) + published = ray.get( + self._policy.init_mx_reshard_publisher(train_world_size=train_world_size) + ) + self._require_all(published, "publisher initialization") + receivers = ray.get( + self._generation.init_mx_reshard_receiver( + train_world_size=train_world_size, + inference_world_size=inference_world_size, + ) + ) + self._require_all(receivers, "receiver initialization") + + def shutdown(self) -> None: + if self._shutdown: + return + self._shutdown = True + errors: list[Exception] = [] + try: + # Receivers load publisher metadata and own the remote connection. + # They must disconnect before the publisher destroys its local UCX + # worker; reversing this order can abort in ucp_worker_destroy. + receiver_results = ray.get(self._generation.shutdown_mx_reshard_receiver()) + self._require_all(receiver_results, "receiver shutdown") + except Exception as error: + errors.append(error) + try: + trainer_results = ray.get(self._policy.shutdown_mx_reshard_publisher()) + self._require_all(trainer_results, "publisher shutdown") + except Exception as error: + errors.append(error) + if errors: + raise RuntimeError( + "ModelExpress mx_reshard shutdown failed: " + + "; ".join(str(error) for error in errors) + ) from errors[0] diff --git a/tests/functional/grpo_mx_reshard_refit.sh b/tests/functional/grpo_mx_reshard_refit.sh new file mode 100755 index 0000000000..a8166a7a40 --- /dev/null +++ b/tests/functional/grpo_mx_reshard_refit.sh @@ -0,0 +1,86 @@ +#!/bin/bash +# Functional smoke for ModelExpress reshard refit on one 2-GPU node: +# Megatron TP1 trainer (1 GPU) -> vLLM TP1 rollout (1 GPU), two GRPO steps. +# +# Requires the mcore + vllm extras, ModelExpress/NIXL, a reachable MX metadata +# server, and a 2-GPU allocation: +# MX_SERVER_URL=host:8001 \ +# uv run --extra mcore --extra vllm \ +# bash tests/functional/grpo_mx_reshard_refit.sh + +set -eou pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") +EXP_NAME=$(basename "$0" .sh) +EXP_DIR="${SCRIPT_DIR}/${EXP_NAME}" +LOG_DIR="${EXP_DIR}/logs" +JSON_METRICS="${EXP_DIR}/metrics.json" +RUN_LOG="${EXP_DIR}/run.log" + +export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}" +export MX_SERVER_URL="${MX_SERVER_URL:-modelexpress-server.kavin.svc.cluster.local:8001}" + +rm -rf "${EXP_DIR}" +mkdir -p "${LOG_DIR}" + +cd "${PROJECT_ROOT}" +uv run coverage run -a \ + --data-file="${PROJECT_ROOT}/tests/.coverage" \ + --source="${PROJECT_ROOT}/nemo_rl" \ + "${PROJECT_ROOT}/examples/run_grpo.py" \ + --config "${PROJECT_ROOT}/examples/configs/grpo_math_1B_megatron.yaml" \ + policy.model_name=Qwen/Qwen3-0.6B \ + grpo.num_prompts_per_step=4 \ + grpo.num_generations_per_prompt=8 \ + policy.train_global_batch_size=16 \ + policy.train_micro_batch_size=1 \ + policy.logprob_batch_size=1 \ + policy.max_total_sequence_length=512 \ + policy.megatron_cfg.enabled=true \ + policy.megatron_cfg.tensor_model_parallel_size=1 \ + policy.megatron_cfg.pipeline_model_parallel_size=1 \ + policy.dtensor_cfg.enabled=false \ + policy.generation.backend=vllm \ + policy.generation.colocated.enabled=false \ + policy.generation.colocated.resources.num_nodes=1 \ + policy.generation.colocated.resources.gpus_per_node=1 \ + policy.generation.vllm_cfg.tensor_parallel_size=1 \ + policy.generation.vllm_cfg.async_engine=true \ + ++policy.generation.refit_transport=mx_reshard \ + ++policy.generation.refit_cfg.mx_reshard.server_url="${MX_SERVER_URL}" \ + cluster.num_nodes=1 \ + cluster.gpus_per_node=2 \ + grpo.max_num_steps=2 \ + logger.tensorboard_enabled=true \ + logger.log_dir="${LOG_DIR}" \ + logger.wandb_enabled=false \ + checkpointing.enabled=false \ + "$@" \ + 2>&1 | tee "${RUN_LOG}" + +uv run tests/json_dump_tb_logs.py "${LOG_DIR}" --output_path "${JSON_METRICS}" + +# A broken refit corrupts generation weights and makes the train/gen +# importance-sampling ratio explode. +# +# The grad_norm check keeps the ratio check from going vacuous. GRPO's +# leave-one-out baseline gives zero advantage to any prompt group whose rewards +# are all equal, so a run where the model never gets a mixed group trains +# nothing, the second refit re-sends bit-identical weights, and the ratio check +# passes for a model a no-op would also satisfy. 4 prompts x 8 generations over +# two steps is sized so that outcome is rare; at the observed ~1/8 solve rate +# for this model it is on the order of 1e-4. +# +# The KL-family check uses js_divergence_error rather than gen_kl_error. Both +# measure trainer-vs-rollout logprob disagreement, but that disagreement has a +# non-zero floor set by Megatron-vs-vLLM kernel differences, and the floor grows +# with model size. Measured with a refit that provably changed no parameter, the +# gen_kl_error floor is 8.7e-4 to 1.3e-3 on Qwen3-30B-A3B, so an absolute +# gen_kl_error < 1e-3 gate is unpassable there no matter how correct the refit is. +# js_divergence_error is bounded and symmetric, and one bound holds across every +# scale measured: 1.6e-4 here on 0.6B, 1.3e-4 on 4B dense, 5.0e-4 on 30B MoE. +uv run tests/check_metrics.py "${JSON_METRICS}" \ + 'max(data["train/token_mult_prob_error"]) < 1.05' \ + 'max(data["train/js_divergence_error"]) < 1e-3' \ + 'max(data["train/grad_norm"]) > 0' diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index ffc0aebd5b..2a8fd41271 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -5147,22 +5147,29 @@ def test_train_fields_for_step(skip_prev_logprobs, expect_prev): @pytest.mark.parametrize( - "backend, nccl_reshard, colocated, expected", + "backend, nccl_reshard, mx_reshard, colocated, expected", [ # MInf refits through mcore's swap_model_weights and never touches HF # names; a revert here is silent (setup time + peak memory only), so # every megatron combination must stay False. - ("megatron", False, True, False), - ("megatron", False, False, False), - ("megatron", True, False, False), - ("megatron", True, True, False), - # vLLM keeps the handshake, except NCCL-reshard non-colocated, which - # builds its own refit info. - ("vllm", False, True, True), - ("vllm", False, False, True), - ("vllm", True, False, False), - ("vllm", True, True, True), + ("megatron", False, False, True, False), + ("megatron", False, False, False, False), + ("megatron", True, False, False, False), + ("megatron", True, False, True, False), + # vLLM keeps the handshake except for non-colocated reshard transports, + # which build their own refit metadata. + ("vllm", False, False, True, True), + ("vllm", False, False, False, True), + ("vllm", True, False, False, False), + ("vllm", True, False, True, True), + ("vllm", False, True, False, False), + ("vllm", False, True, True, True), ], ) -def test_needs_hf_refit_handshake(backend, nccl_reshard, colocated, expected): - assert _needs_hf_refit_handshake(backend, nccl_reshard, colocated) is expected +def test_needs_hf_refit_handshake( + backend, nccl_reshard, mx_reshard, colocated, expected +): + assert ( + _needs_hf_refit_handshake(backend, nccl_reshard, mx_reshard, colocated) + is expected + ) diff --git a/tests/unit/distributed/test_mx_megatron_helpers.py b/tests/unit/distributed/test_mx_megatron_helpers.py new file mode 100644 index 0000000000..767b537ab2 --- /dev/null +++ b/tests/unit/distributed/test_mx_megatron_helpers.py @@ -0,0 +1,303 @@ +from types import SimpleNamespace + +import pytest +import torch + +from nemo_rl.distributed.mx_megatron_helpers import ( + ROLE_EXPERT_COLUMN, + ROLE_EXPERT_ROW, + ROLE_QKV_COLUMN, + canonicalize_grouped_expert_name, + collect_megatron_publish_set, + detect_megatron_role, + infer_megatron_tp_shard_geometry, + resolve_qkv_geometry_from_param, +) + + +class ReplicatedOnlyModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(2)) + + +class ModuleWrapper(torch.nn.Module): + def __init__(self, module): + super().__init__() + self.module = module + + +class ColumnParallelLinear(torch.nn.Module): + def __init__(self, rows: int, q_heads: int, kv_heads: int, head_dim: int): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(rows, 16)) + self.config = SimpleNamespace( + num_attention_heads=q_heads, + num_query_groups=kv_heads, + kv_channels=head_dim, + ) + + +class SelfAttention(torch.nn.Module): + def __init__(self, rows: int, q_heads: int, kv_heads: int, head_dim: int): + super().__init__() + self.linear_qkv = ColumnParallelLinear(rows, q_heads, kv_heads, head_dim) + + +class TransformerLayer(torch.nn.Module): + def __init__(self, rows: int, q_heads: int, kv_heads: int, head_dim: int): + super().__init__() + self.self_attention = SelfAttention(rows, q_heads, kv_heads, head_dim) + + +class HeterogeneousAttentionModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.layers = torch.nn.ModuleList( + [ + TransformerLayer(1088, 64, 2, 128), + TransformerLayer(384, 32, 8, 64), + ] + ) + + +def _published_names(*, tp_rank: int) -> list[str]: + model = ReplicatedOnlyModule() + published = collect_megatron_publish_set( + model, + tp_size=2, + pp_size=1, + pp_rank=0, + ep_size=1, + ep_rank=0, + tp_rank=tp_rank, + ) + return [name for name, _, _, _ in published] + + +def test_collect_megatron_publish_set_skips_replicated_on_nonzero_tp_rank(): + assert _published_names(tp_rank=1) == [] + + +def test_collect_megatron_publish_set_publishes_replicated_on_zero_tp_rank(): + assert _published_names(tp_rank=0) == ["weight"] + + +def test_collect_megatron_publish_set_strips_all_leading_module_wrappers(): + model = ModuleWrapper(ModuleWrapper(ReplicatedOnlyModule())) + + published = collect_megatron_publish_set( + model, + tp_size=1, + pp_size=1, + pp_rank=0, + ep_size=1, + ep_rank=0, + tp_rank=0, + ) + + assert [name for name, _, _, _ in published] == ["weight"] + + +def test_grouped_expert_column_records_tp_geometry(): + geometry = infer_megatron_tp_shard_geometry( + local_shape=(4864, 2048), + role=ROLE_EXPERT_COLUMN, + tp_size=4, + tp_rank=2, + expert_tp_size=4, + expert_tp_rank=2, + descriptor_extras={"expert_layout": "grouped"}, + ) + + assert geometry is not None + assert geometry.global_shape == (19456, 2048) + assert geometry.shard_axis == 0 + assert geometry.local_shard_range == (9728, 14592) + + +def test_grouped_expert_row_records_tp_geometry(): + geometry = infer_megatron_tp_shard_geometry( + local_shape=(2048, 2432), + role=ROLE_EXPERT_ROW, + tp_size=4, + tp_rank=1, + expert_tp_size=4, + expert_tp_rank=1, + descriptor_extras={"expert_layout": "grouped"}, + ) + + assert geometry is not None + assert geometry.global_shape == (2048, 9728) + assert geometry.shard_axis == 1 + assert geometry.local_shard_range == (2432, 4864) + + +def test_leading_axis_expert_geometry_keeps_expert_axis(): + column = infer_megatron_tp_shard_geometry( + local_shape=(4, 4864, 2048), + role=ROLE_EXPERT_COLUMN, + tp_size=4, + tp_rank=3, + expert_tp_size=4, + expert_tp_rank=3, + descriptor_extras={"expert_layout": "leading_axis"}, + ) + row = infer_megatron_tp_shard_geometry( + local_shape=(4, 2048, 2432), + role=ROLE_EXPERT_ROW, + tp_size=4, + tp_rank=3, + expert_tp_size=4, + expert_tp_rank=3, + descriptor_extras={"expert_layout": "leading_axis"}, + ) + + assert column is not None and row is not None + assert column.global_shape == (4, 19456, 2048) + assert column.shard_axis == 1 + assert column.local_shard_range == (14592, 19456) + assert row.global_shape == (4, 2048, 9728) + assert row.shard_axis == 2 + assert row.local_shard_range == (7296, 9728) + + +def test_etp1_expert_is_not_mislabeled_as_trainer_tp_sharded(): + geometry = infer_megatron_tp_shard_geometry( + local_shape=(19456, 2048), + role=ROLE_EXPERT_COLUMN, + tp_size=2, + tp_rank=1, + expert_tp_size=1, + expert_tp_rank=0, + descriptor_extras={"expert_layout": "grouped"}, + ) + + assert geometry is None + + +def test_grouped_expert_publish_name_uses_global_id(): + assert ( + canonicalize_grouped_expert_name( + "decoder.layers.3.mlp.experts.linear_fc1.weight0", + { + "expert_layout": "grouped", + "local_expert_id": "0", + "expert_id": "64", + }, + ) + == "decoder.layers.3.mlp.experts.linear_fc1.weight64" + ) + + +def test_qkv_descriptor_carries_global_heads_when_kv_heads_are_below_tp(): + model = HeterogeneousAttentionModel() + name = "layers.0.self_attention.linear_qkv.weight" + + spec = detect_megatron_role( + name, + model.layers[0].self_attention.linear_qkv.weight, + model=model, + tp_size=8, + ep_size=1, + ep_rank=0, + qkv_geometry=(64, 2, 128), + ) + + assert spec.role == ROLE_QKV_COLUMN + assert spec.descriptor_extras == { + "qkv_interleave": "by_head", + "num_heads": "64", + "num_kv_heads": "2", + "head_dim": "128", + } + + +def test_divisible_qkv_descriptor_retains_legacy_local_head_fields(): + model = HeterogeneousAttentionModel() + name = "layers.1.self_attention.linear_qkv.weight" + + spec = detect_megatron_role( + name, + model.layers[1].self_attention.linear_qkv.weight, + model=model, + tp_size=8, + ep_size=1, + ep_rank=0, + qkv_geometry=(32, 8, 64), + ) + + assert spec.descriptor_extras == { + "qkv_interleave": "by_head", + "num_heads": "32", + "num_kv_heads": "8", + "head_dim": "64", + "num_heads_local": "4", + "num_kv_heads_local": "1", + } + + +def test_collect_reads_per_layer_geometry_from_the_live_qkv_module(): + model = HeterogeneousAttentionModel() + + published = list( + collect_megatron_publish_set( + model, + tp_size=8, + pp_size=1, + pp_rank=0, + ep_size=1, + ep_rank=0, + tp_rank=3, + qkv_geometry_resolver=resolve_qkv_geometry_from_param, + ) + ) + + assert len(published) == 2 + extras_by_name = {name: extras for name, _, _, extras in published} + assert ( + extras_by_name["layers.0.self_attention.linear_qkv.weight"]["num_kv_heads"] + == "2" + ) + assert ( + extras_by_name["layers.1.self_attention.linear_qkv.weight"]["num_kv_heads"] + == "8" + ) + assert ( + "num_kv_heads_local" + not in extras_by_name["layers.0.self_attention.linear_qkv.weight"] + ) + assert ( + extras_by_name["layers.1.self_attention.linear_qkv.weight"][ + "num_kv_heads_local" + ] + == "1" + ) + + +def test_malformed_qkv_geometry_fails_closed(): + model = HeterogeneousAttentionModel() + with pytest.raises(ValueError, match="invalid global Q/KV geometry"): + detect_megatron_role( + "layers.0.self_attention.linear_qkv.weight", + model.layers[0].self_attention.linear_qkv.weight, + model=model, + tp_size=8, + ep_size=1, + ep_rank=0, + qkv_geometry=(63, 2, 128), + ) + + +def test_qkv_geometry_must_match_the_fused_weight_rows(): + model = HeterogeneousAttentionModel() + with pytest.raises(ValueError, match="fused QKV rows"): + detect_megatron_role( + "layers.0.self_attention.linear_qkv.weight", + model.layers[0].self_attention.linear_qkv.weight, + model=model, + tp_size=8, + ep_size=1, + ep_rank=0, + qkv_geometry=(64, 4, 128), + ) diff --git a/tests/unit/distributed/test_mx_refit_verify.py b/tests/unit/distributed/test_mx_refit_verify.py new file mode 100644 index 0000000000..d800ef6c88 --- /dev/null +++ b/tests/unit/distributed/test_mx_refit_verify.py @@ -0,0 +1,152 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. + +"""Parameter-equality verification for mx_reshard refits. + +Runs on CPU: no GPU, no vLLM, no ModelExpress. +""" + +import json + +import torch + +from nemo_rl.distributed import mx_refit_verify + + +class _Model(torch.nn.Module): + def __init__(self, dtype=torch.bfloat16): + super().__init__() + self.a = torch.nn.Parameter(torch.arange(64, dtype=torch.float32).to(dtype)) + self.b = torch.nn.Parameter(torch.ones(8, 8, dtype=dtype)) + + +def test_a_single_flipped_bit_is_detected(): + """The whole mechanism rests on this. A float statistic can miss a low-mantissa + change; the fingerprint is taken over raw bytes so it cannot.""" + t = torch.ones(1024, dtype=torch.bfloat16) + before = mx_refit_verify.fingerprint(t) + + raw = t.view(torch.uint8) + raw[500] = raw[500].item() ^ 1 # flip the lowest mantissa bit of one element + + assert mx_refit_verify.fingerprint(t) != before + + +def test_high_byte_statistic_catches_an_offsetting_change(): + """Two changes can cancel in a flat byte sum. The second statistic covers only + each element's most significant byte, so a cancelling pair has to cancel in + both to hide.""" + t = torch.ones(64, dtype=torch.bfloat16) + raw = t.view(torch.uint8) + before = mx_refit_verify.fingerprint(t) + + # One byte up, one down, so the flat sum is unchanged by construction. The + # pair has to straddle the element boundary for this to be a real test: bf16 + # is 2 bytes little-endian, so index 11 is element 5's high byte and index 10 + # is its low byte. Two *low* bytes would cancel in both statistics and prove + # nothing. + raw[11] = raw[11].item() + 1 + raw[10] = raw[10].item() - 1 + after = mx_refit_verify.fingerprint(t) + + assert after[1] == before[1], "expected the flat byte sum to be fooled here" + assert after[2] != before[2], "the high-byte statistic should still differ" + assert after != before + + +def test_identical_tensors_fingerprint_identically(): + """Determinism, or every refit would look like it changed everything.""" + a = torch.arange(256, dtype=torch.float32).to(torch.bfloat16) + b = a.clone() + assert mx_refit_verify.fingerprint(a) == mx_refit_verify.fingerprint(b) + assert mx_refit_verify.fingerprint(a) == mx_refit_verify.fingerprint(a) + + +def test_non_contiguous_parameters_are_handled(): + """A transposed or sliced view cannot be reinterpreted as bytes directly, and + must not take the whole verification down with it.""" + t = torch.ones(8, 8, dtype=torch.bfloat16).t() + assert not t.is_contiguous() + assert mx_refit_verify.fingerprint(t)[0] == 64 + + +def test_float32_and_bfloat16_both_pick_the_high_byte(): + """The high-byte stride is derived from element size, not hardcoded to 2.""" + for dtype in (torch.bfloat16, torch.float32, torch.float16): + t = torch.ones(32, dtype=dtype) + numel, _flat, _high = mx_refit_verify.fingerprint(t) + assert numel == 32 + + +def test_compare_names_what_changed(): + model = _Model() + before = mx_refit_verify.fingerprint_model(model) + with torch.no_grad(): + model.b.add_(1.0) + result = mx_refit_verify.compare(before, mx_refit_verify.fingerprint_model(model)) + + assert result["params_compared"] == 2 + assert result["params_changed"] == 1 + assert result["changed_sample"] == ["b"] + assert result["unchanged_sample"] == ["a"] + + +def test_first_refit_should_change_nothing(capsys): + """The load-bearing case. On a fresh run vLLM is loaded from the HF checkpoint + and the trainer's weights are converted from the same checkpoint with no + optimizer step taken, so a faithful first refit is a no-op on every parameter. + """ + model = _Model() + before = mx_refit_verify.fingerprint_model(model) + + record = mx_refit_verify.report( + 1, 0, before, mx_refit_verify.fingerprint_model(model) + ) + + assert record["first_refit"] is True + assert record["params_changed"] == 0 + assert record["params_unchanged"] == 2 + emitted = json.loads(capsys.readouterr().out.split("MX_REFIT_VERIFY ", 1)[1]) + assert emitted == record + + +def test_later_refits_are_flagged_as_a_different_question(capsys): + """From step 2 the trainer has moved, so change is expected and the record + instead reports which params the refit reaches.""" + model = _Model() + before = mx_refit_verify.fingerprint_model(model) + with torch.no_grad(): + model.a.mul_(2.0) + model.b.mul_(3.0) + + record = mx_refit_verify.report( + 2, 5, before, mx_refit_verify.fingerprint_model(model) + ) + + assert record["first_refit"] is False + assert record["params_changed"] == 2 + capsys.readouterr() + + +def test_verification_never_raises_on_a_broken_model(): + """A verification failure must not be able to fail the refit it verifies.""" + + class Broken: + def named_parameters(self): + raise RuntimeError("model is not walkable here") + + assert mx_refit_verify.fingerprint_model(Broken()) == {} + + +def test_disabled_by_default(monkeypatch): + """It sits on the refit critical path, so it is opt-in.""" + monkeypatch.delenv("MX_REFIT_VERIFY", raising=False) + assert mx_refit_verify.enabled() is False + for value in ("0", "", "false", "False"): + monkeypatch.setenv("MX_REFIT_VERIFY", value) + assert mx_refit_verify.enabled() is False + for value in ("1", "true", "yes"): + monkeypatch.setenv("MX_REFIT_VERIFY", value) + assert mx_refit_verify.enabled() is True diff --git a/tests/unit/distributed/test_mx_reshard_publisher.py b/tests/unit/distributed/test_mx_reshard_publisher.py new file mode 100644 index 0000000000..160da5f53c --- /dev/null +++ b/tests/unit/distributed/test_mx_reshard_publisher.py @@ -0,0 +1,916 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Spec derivation for the main-native Megatron reshard publisher. + +Everything here runs on CPU with no MX server, no NIXL and no Megatron: the +translation from a classified publish set to MX alias inputs is pure geometry and +name resolution, which is the part that is wrong silently. A bad shard range +publishes an address that reads real but wrong bytes, and a bad HF name publishes +one expert's weights under another's. +""" + +from __future__ import annotations + +import json +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +pytest.importorskip( + "modelexpress", + reason="ModelExpress is an optional integration dependency", +) + +from nemo_rl.distributed.mx_reshard_config import ( + DEFAULT_PUBLISHER_LISTEN_PORT_BASE, + DEFAULT_RECEIVER_LISTEN_PORT_BASE, + maybe_preinit_mx_reshard_nixl, + resolve_mx_reshard_listen_port_bases, + resolve_mx_reshard_publisher_listen_port_base, + resolve_mx_reshard_receiver_listen_port_base, + validate_mx_reshard_listen_port_ranges, +) +from nemo_rl.distributed.mx_megatron_helpers import ( + ROLE_EXPERT_COLUMN, + ROLE_GATED_MLP_COLUMN, + ROLE_QKV_COLUMN, + ROLE_REPLICATED, + ROLE_ROW, + MegatronRoleSpec, +) +from nemo_rl.distributed.mx_reshard_publisher import ( + _gated_mlp_extras, + PLACEMENT_REPLICATE, + PLACEMENT_SHARD, + MxMegatronPublisher, + UnmappedMegatronTensor, + build_bridge_name_map, + build_megatron_alias_inputs, + make_bridge_resolver, + publish_megatron_hf_aliases, + published_byte_count, +) + + +def _entry(name, tensor, role, extras=None, **spec_kwargs): + """One (name, tensor, spec, full_extras) tuple as the collector yields it.""" + descriptor_extras = dict(extras or {}) + spec = MegatronRoleSpec( + role=role, descriptor_extras=descriptor_extras, **spec_kwargs + ) + # The collector merges the mesh position into the per-tensor extras; the + # alias builder only reads the geometry keys, but carry both so the fixture + # matches the real shape of the input. + full = {"megatron_role": role, "tp_size": "1", "ep_size": "8"} + full.update(descriptor_extras) + return name, tensor, spec, full + + +def test_publisher_step_is_forwarded_to_mx_rendezvous_encoder(): + entry = _entry("norm.weight", torch.zeros(16), ROLE_REPLICATED) + resolver = make_bridge_resolver({"norm.weight": ["model.norm.weight"]}) + items = list( + build_megatron_alias_inputs( + [entry], resolve_hf_names=resolver, tp_size=1, tp_rank=0 + ) + ) + manager = MagicMock(agent_name="trainer-r0", nixl_metadata=b"metadata") + rendezvous = MagicMock() + rendezvous.publish.return_value = "source-id" + + with patch( + "nemo_rl.distributed.mx_reshard_publisher.wrap_rendezvous_blob", + return_value=b"stamped", + ) as wrap: + source_id, _ = publish_megatron_hf_aliases( + manager=manager, + rendezvous=rendezvous, + items=items, + metadata_endpoint="trainer:5555", + publisher_step=17, + ) + + assert source_id == "source-id" + assert wrap.call_args.kwargs["publisher_step"] == 17 + rendezvous.publish.assert_called_once_with(b"stamped") + + +def test_bridge_map_uses_local_pp_name_and_local_tasks_only(): + local = SimpleNamespace( + param_name="decoder.layers.0.self_attention.linear_qkv.weight", + global_param_name="decoder.layers.24.self_attention.linear_qkv.weight", + param_weight=torch.zeros(1), + mapping=SimpleNamespace( + hf_param={ + "k": "model.layers.24.self_attn.k_proj.weight", + "q": "model.layers.24.self_attn.q_proj.weight", + "v": "model.layers.24.self_attn.v_proj.weight", + } + ), + ) + remote_pp = SimpleNamespace( + param_name="decoder.layers.0.input_layernorm.weight", + global_param_name="decoder.layers.0.input_layernorm.weight", + param_weight=None, + mapping=SimpleNamespace(hf_param="model.layers.0.input_layernorm.weight"), + ) + + name_map = build_bridge_name_map([local, remote_pp]) + + assert name_map == { + "decoder.layers.0.self_attention.linear_qkv.weight": ( + "model.layers.24.self_attn.q_proj.weight", + "model.layers.24.self_attn.k_proj.weight", + "model.layers.24.self_attn.v_proj.weight", + ) + } + + +def test_bridge_map_rejects_conflicting_local_names(): + def task(hf_name): + return SimpleNamespace( + param_name="decoder.layers.0.weight", + global_param_name="decoder.layers.0.weight", + param_weight=torch.zeros(1), + mapping=SimpleNamespace(hf_param=hf_name), + ) + + with pytest.raises(ValueError, match="conflicting Bridge mappings"): + build_bridge_name_map([task("model.layers.0.a"), task("model.layers.0.b")]) + + +def test_default_publisher_and_receiver_port_ranges_do_not_overlap(): + publisher, receiver = validate_mx_reshard_listen_port_ranges( + {}, + train_world_size=8, + inference_world_size=8, + ) + + assert publisher == DEFAULT_PUBLISHER_LISTEN_PORT_BASE + assert receiver == DEFAULT_RECEIVER_LISTEN_PORT_BASE + assert publisher + 7 < receiver + + +def test_legacy_port_base_falls_back_to_shifted_receiver_range(): + assert resolve_mx_reshard_listen_port_bases({"listen_port_base": 7000}) == ( + 7000, + 17000, + ) + + +def test_role_specific_port_bases_forward_to_the_correct_side(): + config = { + "publisher_listen_port_base": 19000, + "receiver_listen_port_base": 29000, + } + assert resolve_mx_reshard_publisher_listen_port_base(config) == 19000 + assert resolve_mx_reshard_receiver_listen_port_base(config) == 29000 + + +def test_mx_reshard_preinitializes_nixl_only_for_selected_transport(monkeypatch): + from nemo_rl.utils.checkpoint_engines import nixl + + calls = [] + agent = object() + monkeypatch.setattr( + nixl, + "preinit_nixl_agent", + lambda **kwargs: calls.append(kwargs) or agent, + ) + + assert maybe_preinit_mx_reshard_nixl({}) is None + assert ( + maybe_preinit_mx_reshard_nixl({"generation": {"refit_transport": "mx_reshard"}}) + is agent + ) + assert calls == [{}] + + +def test_overlapping_explicit_port_ranges_are_rejected(): + with pytest.raises(ValueError, match="listen port ranges overlap"): + validate_mx_reshard_listen_port_ranges( + { + "publisher_listen_port_base": 19000, + "receiver_listen_port_base": 19004, + }, + train_world_size=8, + inference_world_size=8, + ) + + +def test_publisher_is_constructed_with_the_expected_mx_arguments(): + item = SimpleNamespace(name="weight", tensor=torch.zeros(1)) + manager = MagicMock() + client = MagicMock() + rendezvous = MagicMock() + with ( + patch( + "modelexpress.nixl_transfer.NixlTransferManager", + return_value=manager, + ) as manager_cls, + patch("modelexpress.client.MxClient", return_value=client) as client_cls, + patch( + "modelexpress.refit.reshard.rendezvous.MxReshardRendezvous", + return_value=rendezvous, + ) as rendezvous_cls, + ): + publisher = MxMegatronPublisher( + items=[item], + model_name="Qwen/Qwen3-30B", + server_url="mx.example:8001", + rank=11, + device_id=3, + listen_port=19011, + metadata_endpoint="trainer-1:19011", + ) + + manager_cls.assert_called_once_with( + agent_name="nemo-rl-trainer-11", + device_id=3, + listen_port=19011, + ) + manager.initialize.assert_called_once() + manager.register_tensors.assert_called_once_with({"weight": item.tensor}) + client_cls.assert_called_once_with(server_url="mx.example:8001") + rendezvous_cls.assert_called_once_with( + client, + role="trainer", + rank=11, + model_name="Qwen/Qwen3-30B", + worker_id="nemo-rl-trainer-11", + ) + + publisher.shutdown() + rendezvous.close.assert_called_once() + manager.shutdown.assert_called_once() + client.close.assert_called_once() + + +def test_dp_replica_ranks_receive_unique_worker_ids(): + item = SimpleNamespace(name="weight", tensor=torch.zeros(1)) + worker_ids = [] + + def rendezvous(*_args, **kwargs): + worker_ids.append(kwargs["worker_id"]) + return MagicMock() + + with ( + patch("modelexpress.nixl_transfer.NixlTransferManager"), + patch("modelexpress.client.MxClient"), + patch( + "modelexpress.refit.reshard.rendezvous.MxReshardRendezvous", + side_effect=rendezvous, + ), + ): + first = MxMegatronPublisher( + items=[item], + model_name="model", + server_url="mx:8001", + rank=0, + device_id=0, + listen_port=19000, + metadata_endpoint="host:19000", + ) + second = MxMegatronPublisher( + items=[item], + model_name="model", + server_url="mx:8001", + rank=8, + device_id=0, + listen_port=19008, + metadata_endpoint="host:19008", + ) + + assert worker_ids == ["nemo-rl-trainer-0", "nemo-rl-trainer-8"] + first.shutdown() + second.shutdown() + + +# --- placement ------------------------------------------------------------- + + +def test_replicated_tensor_publishes_its_whole_shape(): + entry = _entry( + "decoder.layers.0.input_layernorm.weight", torch.zeros(16), ROLE_REPLICATED + ) + resolver = make_bridge_resolver( + { + "decoder.layers.0.input_layernorm.weight": [ + "model.layers.0.input_layernorm.weight" + ] + } + ) + + (item,) = build_megatron_alias_inputs( + [entry], resolve_hf_names=resolver, tp_size=1, tp_rank=0 + ) + + assert item.placement_kind == PLACEMENT_REPLICATE + assert item.global_shape == (16,) + assert item.shard_axis is None + assert item.local_shard_range is None + assert item.hf_names == ("model.layers.0.input_layernorm.weight",) + + +def test_tp_sharded_column_carries_its_global_shape_and_range(): + """TP2 rank 1 of a column-parallel weight owns the second row band.""" + entry = _entry( + "decoder.layers.0.mlp.linear_fc2.weight", torch.zeros(8, 32), ROLE_ROW + ) + resolver = make_bridge_resolver( + { + "decoder.layers.0.mlp.linear_fc2.weight": [ + "model.layers.0.mlp.down_proj.weight" + ] + } + ) + + (item,) = build_megatron_alias_inputs( + [entry], resolve_hf_names=resolver, tp_size=2, tp_rank=1 + ) + + # ROLE_ROW shards axis 1. + assert item.placement_kind == PLACEMENT_SHARD + assert item.shard_axis == 1 + assert item.global_shape == (8, 64) + assert item.local_shard_range == (32, 64) + + +def test_tp1_leaves_every_tensor_whole(): + """The Topology B publisher runs TP1/ETP1, so nothing is TP-sharded.""" + entries = [ + _entry("a.weight", torch.zeros(4, 8), ROLE_ROW), + _entry("b.weight", torch.zeros(4, 8), ROLE_GATED_MLP_COLUMN), + ] + resolver = make_bridge_resolver({"a.weight": ["hf.a"], "b.weight": ["hf.b"]}) + + items = list( + build_megatron_alias_inputs( + entries, resolve_hf_names=resolver, tp_size=1, tp_rank=0 + ) + ) + + assert [item.placement_kind for item in items] == [PLACEMENT_REPLICATE] * 2 + assert all(item.local_shard_range is None for item in items) + + +def test_expert_tp_geometry_is_taken_from_the_expert_mesh(): + """An expert tensor shards on the expert-TP mesh, not the dense TP mesh.""" + entry = _entry( + "decoder.layers.0.mlp.experts.linear_fc2.weight0", + torch.zeros(4, 16), + ROLE_EXPERT_COLUMN, + extras={"expert_layout": "grouped", "expert_id": "0", "local_expert_id": "0"}, + is_expert=True, + ) + resolver = make_bridge_resolver( + { + "decoder.layers.0.mlp.experts.linear_fc2.weight0": [ + "model.layers.0.mlp.experts.0.down_proj.weight" + ] + } + ) + + (item,) = build_megatron_alias_inputs( + [entry], + resolve_hf_names=resolver, + tp_size=4, + tp_rank=3, + expert_tp_size=2, + expert_tp_rank=1, + ) + + # Grouped expert_column shards axis 0, and the extent comes from ETP=2 not TP=4. + assert item.shard_axis == 0 + assert item.global_shape == (8, 16) + assert item.local_shard_range == (4, 8) + + +# --- expert name resolution ------------------------------------------------ + + +def test_global_expert_id_is_substituted_into_the_hf_name(): + """EP rank 1's first local expert is global expert 4 and must publish as such. + + The Bridge inspects one rank's module tree, so its map is keyed on the + EP-local leaf (`weight0`) and its HF names describe expert 0. Publishing + those names unchanged would have every EP rank overwrite rank 0's experts. + """ + resolver = make_bridge_resolver( + { + "decoder.layers.0.mlp.experts.linear_fc1.weight0": [ + "model.layers.0.mlp.experts.0.gate_proj.weight" + ] + } + ) + + names = resolver( + "decoder.layers.0.mlp.experts.linear_fc1.weight4", + {"expert_id": "4", "local_expert_id": "0"}, + ) + + assert names == ("model.layers.0.mlp.experts.4.gate_proj.weight",) + + +def test_exact_key_wins_over_the_local_expert_retry(): + """A map already keyed on the global name is used as-is.""" + resolver = make_bridge_resolver( + { + "e.linear_fc1.weight4": ["model.layers.0.mlp.experts.4.gate_proj.weight"], + "e.linear_fc1.weight0": ["model.layers.0.mlp.experts.0.gate_proj.weight"], + } + ) + + names = resolver("e.linear_fc1.weight4", {"expert_id": "4", "local_expert_id": "0"}) + + assert names == ("model.layers.0.mlp.experts.4.gate_proj.weight",) + + +def test_expert_substitution_rewrites_every_hf_name_of_a_fused_pair(): + """A fused gate/up expert tensor maps to two HF names, both needing the id.""" + resolver = make_bridge_resolver( + { + "e.linear_fc1.weight0": [ + "model.layers.0.mlp.experts.0.gate_proj.weight", + "model.layers.0.mlp.experts.0.up_proj.weight", + ] + } + ) + + names = resolver("e.linear_fc1.weight7", {"expert_id": "7", "local_expert_id": "0"}) + + assert names == ( + "model.layers.0.mlp.experts.7.gate_proj.weight", + "model.layers.0.mlp.experts.7.up_proj.weight", + ) + + +def test_substitution_only_touches_the_expert_index(): + """A layer number that equals the old expert id must not be rewritten.""" + resolver = make_bridge_resolver( + {"e.linear_fc1.weight3": ["model.layers.3.mlp.experts.3.gate_proj.weight"]} + ) + + names = resolver("e.linear_fc1.weight5", {"expert_id": "5", "local_expert_id": "3"}) + + assert names == ("model.layers.3.mlp.experts.5.gate_proj.weight",) + + +def test_unmapped_tensor_raises_instead_of_being_skipped(): + """A silently dropped source becomes a coverage shortfall far from its cause.""" + resolver = make_bridge_resolver({"known.weight": ["hf.known"]}) + + with pytest.raises(UnmappedMegatronTensor, match="unknown.weight"): + resolver("unknown.weight", {}) + + +def test_non_strict_resolver_reports_no_names_and_the_tensor_is_dropped(): + entry = _entry("unknown.weight", torch.zeros(4), ROLE_REPLICATED) + resolver = make_bridge_resolver({"known.weight": ["hf.known"]}, strict=False) + + items = list( + build_megatron_alias_inputs( + [entry], resolve_hf_names=resolver, tp_size=1, tp_rank=0 + ) + ) + + assert items == [] + + +# --- the extras contract that build_hf_aliases depends on ----------------- + + +def test_qkv_head_extras_survive_translation(): + """Global QKV geometry must reach MX unchanged when KV heads are below TP.""" + entry = _entry( + "decoder.layers.0.self_attention.linear_qkv.weight", + torch.zeros(1088, 2048), + ROLE_QKV_COLUMN, + extras={ + "qkv_interleave": "by_head", + "head_dim": "128", + "num_heads": "64", + "num_kv_heads": "2", + }, + ) + resolver = make_bridge_resolver( + { + "decoder.layers.0.self_attention.linear_qkv.weight": [ + "model.layers.0.self_attn.q_proj.weight", + "model.layers.0.self_attn.k_proj.weight", + "model.layers.0.self_attn.v_proj.weight", + ] + } + ) + + (item,) = build_megatron_alias_inputs( + [entry], resolve_hf_names=resolver, tp_size=8, tp_rank=3 + ) + + assert item.extras["head_dim"] == "128" + assert item.extras["num_heads"] == "64" + assert item.extras["num_kv_heads"] == "2" + assert "num_kv_heads_local" not in item.extras + assert item.global_shape == (8704, 2048) + assert item.local_shard_range == (3264, 4352) + assert len(item.hf_names) == 3 + + +def test_qkv_names_keep_q_k_v_order(): + """build_hf_aliases assigns hf_names[0..2] to the Q, K and V row bands, so a + reordering here transposes the projections with no error anywhere.""" + hf = [ + "model.layers.0.self_attn.q_proj.weight", + "model.layers.0.self_attn.k_proj.weight", + "model.layers.0.self_attn.v_proj.weight", + ] + resolver = make_bridge_resolver({"qkv": hf}) + + assert resolver("qkv", {}) == tuple(hf) + + +def test_gated_order_extra_survives_translation(): + entry = _entry( + "decoder.layers.0.mlp.linear_fc1.weight", + torch.zeros(64, 32), + ROLE_GATED_MLP_COLUMN, + extras={"gated_mlp_order": "gate_then_up"}, + ) + resolver = make_bridge_resolver( + { + "decoder.layers.0.mlp.linear_fc1.weight": [ + "model.layers.0.mlp.gate_proj.weight", + "model.layers.0.mlp.up_proj.weight", + ] + } + ) + + (item,) = build_megatron_alias_inputs( + [entry], resolve_hf_names=resolver, tp_size=1, tp_rank=0 + ) + + assert item.extras["gated_mlp_order"] == "gate_then_up" + + +# --- end-to-end against the real MX alias builder ------------------------- + + +def test_translated_items_are_accepted_by_build_hf_aliases(): + """The contract test: MX main's builder consumes what we produce. + + Runs the real `build_hf_aliases`, so a drift in either the role vocabulary or + the extras keys fails here rather than on a cluster. + """ + from modelexpress.refit.reshard.megatron_aliases import build_hf_aliases + + entries = [ + _entry("norm.weight", torch.zeros(16, dtype=torch.bfloat16), ROLE_REPLICATED), + _entry( + "mlp.linear_fc1.weight", + torch.zeros(64, 32, dtype=torch.bfloat16), + ROLE_GATED_MLP_COLUMN, + extras={"gated_mlp_order": "gate_then_up"}, + ), + ] + resolver = make_bridge_resolver( + { + "norm.weight": ["model.norm.weight"], + "mlp.linear_fc1.weight": [ + "model.mlp.gate_proj.weight", + "model.mlp.up_proj.weight", + ], + } + ) + + items = list( + build_megatron_alias_inputs( + entries, resolve_hf_names=resolver, tp_size=1, tp_rank=0 + ) + ) + published = build_hf_aliases(items, agent_name="trainer-r0") + + # The fused gate/up parent becomes two HF tensors, so three in total. + assert sorted(tensor.name for tensor in published) == [ + "model.mlp.gate_proj.weight", + "model.mlp.up_proj.weight", + "model.norm.weight", + ] + for tensor in published: + assert tensor.shards + for shard in tensor.shards: + assert shard.agent_name == "trainer-r0" + + +def test_publish_reports_its_shard_and_byte_counts(capsys): + """The synchronizer times the publish half but not what it described, so a + publish that slowed down cannot be told from one that grew without this.""" + item = SimpleNamespace(name="weight", tensor=torch.zeros(1)) + published = [SimpleNamespace(shards=[SimpleNamespace(shape=(8,))], elsize=2)] + with ( + patch("modelexpress.nixl_transfer.NixlTransferManager"), + patch("modelexpress.client.MxClient"), + patch("modelexpress.refit.reshard.rendezvous.MxReshardRendezvous"), + ): + publisher = MxMegatronPublisher( + items=[item], + model_name="model", + server_url="mx:8001", + rank=11, + device_id=0, + listen_port=19011, + metadata_endpoint="host:19011", + ) + with patch( + "nemo_rl.distributed.mx_reshard_publisher.publish_megatron_hf_aliases", + return_value=("source-1", published), + ): + publisher.publish(7) + + record = json.loads(capsys.readouterr().out.split("MX_PUBLISH_PHASE ", 1)[1]) + assert record["step"] == 7 + assert record["rank"] == 11 + assert record["tensors"] == 1 + assert record["shards"] == 1 + assert record["bytes"] == 8 * 2 + publisher.shutdown() + + +def test_publish_telemetry_cannot_fail_a_successful_publish(): + """Reporting is never worth a failed publish: the bytes are already out.""" + item = SimpleNamespace(name="weight", tensor=torch.zeros(1)) + with ( + patch("modelexpress.nixl_transfer.NixlTransferManager"), + patch("modelexpress.client.MxClient"), + patch("modelexpress.refit.reshard.rendezvous.MxReshardRendezvous"), + ): + publisher = MxMegatronPublisher( + items=[item], + model_name="model", + server_url="mx:8001", + rank=0, + device_id=0, + listen_port=19000, + metadata_endpoint="host:19000", + ) + # A payload whose shards cannot be counted at all. + with patch( + "nemo_rl.distributed.mx_reshard_publisher.publish_megatron_hf_aliases", + return_value=("source-1", [SimpleNamespace()]), + ): + publisher.publish(3) + + publisher.shutdown() + + +def test_published_byte_count_sums_shard_boxes(): + from modelexpress.refit.reshard.megatron_aliases import build_hf_aliases + + entry = _entry( + "norm.weight", torch.zeros(16, dtype=torch.bfloat16), ROLE_REPLICATED + ) + resolver = make_bridge_resolver({"norm.weight": ["model.norm.weight"]}) + items = list( + build_megatron_alias_inputs( + [entry], resolve_hf_names=resolver, tp_size=1, tp_rank=0 + ) + ) + + published = build_hf_aliases(items, agent_name="trainer-r0") + + assert published_byte_count(published) == 16 * 2 + + +# --- whole-model name coverage -------------------------------------------- +# +# Qwen3-30B-A3B-Instruct-2507's real geometry. The expected HF name set below is +# generated from these constants rather than shipped as a fixture, and the +# generator was checked against the checkpoint's own +# model.safetensors.index.json on 2026-08-12: 18,867 names, exact match, nothing +# missing and nothing extra. +QWEN3_30B_A3B = { + "layers": 48, + "experts": 128, + "tie_word_embeddings": False, +} + + +def _expected_hf_names( + layers: int, experts: int, tie_word_embeddings: bool +) -> set[str]: + """The tensor names the HF checkpoint actually contains.""" + names = {"model.embed_tokens.weight", "model.norm.weight"} + if not tie_word_embeddings: + names.add("lm_head.weight") + for layer in range(layers): + head = f"model.layers.{layer}" + names.update( + { + f"{head}.input_layernorm.weight", + f"{head}.post_attention_layernorm.weight", + f"{head}.mlp.gate.weight", + f"{head}.self_attn.q_proj.weight", + f"{head}.self_attn.k_proj.weight", + f"{head}.self_attn.v_proj.weight", + f"{head}.self_attn.o_proj.weight", + f"{head}.self_attn.q_norm.weight", + f"{head}.self_attn.k_norm.weight", + } + ) + for expert in range(experts): + stem = f"{head}.mlp.experts.{expert}" + names.update( + { + f"{stem}.gate_proj.weight", + f"{stem}.up_proj.weight", + f"{stem}.down_proj.weight", + } + ) + return names + + +def _bridge_map_for_rank(layers: int, local_experts: int) -> dict[str, list[str]]: + """A Bridge map as seen from ONE EP rank. + + The Bridge walks one rank's module tree, so grouped-expert keys carry the + **EP-local** leaf index and the HF names it returns describe that local + expert. Every rank's map therefore looks identical and describes experts + 0..local_experts-1 -- which is exactly why the resolver has to substitute the + global id, and why a map keyed on the global name would be the easier thing + to test and the wrong thing to test. + """ + name_map: dict[str, list[str]] = { + "embedding.word_embeddings.weight": ["model.embed_tokens.weight"], + "decoder.final_layernorm.weight": ["model.norm.weight"], + "output_layer.weight": ["lm_head.weight"], + } + for layer in range(layers): + megatron = f"decoder.layers.{layer}" + hf = f"model.layers.{layer}" + name_map[f"{megatron}.self_attention.linear_qkv.weight"] = [ + f"{hf}.self_attn.q_proj.weight", + f"{hf}.self_attn.k_proj.weight", + f"{hf}.self_attn.v_proj.weight", + ] + name_map[f"{megatron}.self_attention.linear_proj.weight"] = [ + f"{hf}.self_attn.o_proj.weight" + ] + name_map[f"{megatron}.self_attention.q_layernorm.weight"] = [ + f"{hf}.self_attn.q_norm.weight" + ] + name_map[f"{megatron}.self_attention.k_layernorm.weight"] = [ + f"{hf}.self_attn.k_norm.weight" + ] + name_map[f"{megatron}.input_layernorm.weight"] = [ + f"{hf}.input_layernorm.weight" + ] + name_map[f"{megatron}.pre_mlp_layernorm.weight"] = [ + f"{hf}.post_attention_layernorm.weight" + ] + name_map[f"{megatron}.mlp.router.weight"] = [f"{hf}.mlp.gate.weight"] + for local in range(local_experts): + name_map[f"{megatron}.mlp.experts.linear_fc1.weight{local}"] = [ + f"{hf}.mlp.experts.{local}.gate_proj.weight", + f"{hf}.mlp.experts.{local}.up_proj.weight", + ] + name_map[f"{megatron}.mlp.experts.linear_fc2.weight{local}"] = [ + f"{hf}.mlp.experts.{local}.down_proj.weight" + ] + return name_map + + +def _publish_whole_model(ep_size: int) -> dict[str, set[int]]: + """Resolve every rank's publish set; return hf name -> owning EP ranks.""" + layers = QWEN3_30B_A3B["layers"] + local_experts = QWEN3_30B_A3B["experts"] // ep_size + owners: dict[str, set[int]] = {} + + for ep_rank in range(ep_size): + resolve = make_bridge_resolver(_bridge_map_for_rank(layers, local_experts)) + entries: list[tuple[str, dict[str, str]]] = [ + (key, {}) + for key in _bridge_map_for_rank(layers, local_experts) + if ".experts.linear_fc" not in key + ] + for layer in range(layers): + megatron = f"decoder.layers.{layer}" + for local in range(local_experts): + global_id = ep_rank * local_experts + local + for fused in ("linear_fc1", "linear_fc2"): + entries.append( + ( + f"{megatron}.mlp.experts.{fused}.weight{global_id}", + { + "expert_id": str(global_id), + "local_expert_id": str(local), + "expert_layout": "grouped", + }, + ) + ) + for name, extras in entries: + for hf_name in resolve(name, extras): + owners.setdefault(hf_name, set()).add(ep_rank) + return owners + + +def test_ep8_publishes_every_checkpoint_tensor_exactly_once_per_owner(): + """Whole-model coverage: the union of 8 EP ranks is the checkpoint, exactly. + + A name the fleet never publishes is a coverage shortfall the receiver reports + far from its cause; a name it publishes that the model does not have is an + install into a buffer nothing owns. + """ + owners = _publish_whole_model(ep_size=8) + expected = _expected_hf_names(**QWEN3_30B_A3B) + + assert set(owners) - expected == set(), "published names absent from the checkpoint" + assert expected - set(owners) == set(), "checkpoint tensors nobody publishes" + assert len(owners) == 18867 + + +def test_every_expert_has_exactly_one_owner_under_ep8(): + """EP partitions experts, so two ranks claiming one expert means the global id + substitution failed and one rank is publishing another's weights.""" + owners = _publish_whole_model(ep_size=8) + + experts = {name: rank for name, rank in owners.items() if ".experts." in name} + assert len(experts) == 48 * 128 * 3 + assert {len(rank) for rank in experts.values()} == {1} + + +def test_non_expert_tensors_are_published_by_every_rank(): + """Not a defect: with TP1 the collector keeps replicated tensors on all ranks, + so the fleet offers 8 byte-identical copies and the receiver's merge picks one. + Recorded because it is the DP amplification c2 removes, and because a change + here silently changes the wire bytes of every measurement.""" + owners = _publish_whole_model(ep_size=8) + + non_expert = { + name: rank for name, rank in owners.items() if ".experts." not in name + } + assert len(non_expert) == 435 + assert {len(rank) for rank in non_expert.values()} == {8} + + # 18,432 expert offers + 435 x 8 replicated offers. + assert sum(len(rank) for rank in owners.values()) == 21912 + + +# --------------------------------------------------------- gated MLP fusion order +# MX assigns the two halves of a fused gate/up parameter to hf_names positionally +# and refuses to infer which half is which, because getting it wrong publishes the +# gate's bytes under the up projection's name with every digest agreeing. + + +def test_fused_expert_gate_up_is_stamped_gate_then_up(): + extras = _gated_mlp_extras( + "decoder.layers.0.mlp.experts.linear_fc1.weight0", + "expert_column", + ( + "model.layers.0.mlp.experts.0.gate_proj.weight", + "model.layers.0.mlp.experts.0.up_proj.weight", + ), + ) + assert extras == {"gated_mlp_order": "gate_then_up"} + + +def test_dense_gated_mlp_is_stamped_too(): + extras = _gated_mlp_extras( + "decoder.layers.0.mlp.linear_fc1.weight", + "gated_mlp_column", + ("model.layers.0.mlp.gate_proj.weight", "model.layers.0.mlp.up_proj.weight"), + ) + assert extras == {"gated_mlp_order": "gate_then_up"} + + +def test_reversed_hf_name_order_is_refused_not_guessed(): + with pytest.raises(ValueError, match="do not read as"): + _gated_mlp_extras( + "decoder.layers.0.mlp.experts.linear_fc1.weight0", + "expert_column", + ( + "model.layers.0.mlp.experts.0.up_proj.weight", + "model.layers.0.mlp.experts.0.gate_proj.weight", + ), + ) + + +def test_unfused_roles_are_not_stamped(): + # linear_fc2 maps to one HF name, so there is nothing to order; stamping it + # anyway would assert a layout claim about a tensor that has no halves. + assert ( + _gated_mlp_extras( + "decoder.layers.0.mlp.experts.linear_fc2.weight0", + "expert_row", + ("model.layers.0.mlp.experts.0.down_proj.weight",), + ) + == {} + ) + assert ( + _gated_mlp_extras( + "decoder.layers.0.self_attention.linear_qkv.weight", + "qkv_column", + ("q.weight", "k.weight", "v.weight"), + ) + == {} + ) diff --git a/tests/unit/models/generation/test_mx_reshard_receiver.py b/tests/unit/models/generation/test_mx_reshard_receiver.py new file mode 100644 index 0000000000..e8a04bed08 --- /dev/null +++ b/tests/unit/models/generation/test_mx_reshard_receiver.py @@ -0,0 +1,248 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import json +from types import SimpleNamespace +from unittest.mock import MagicMock, create_autospec, patch + +import pytest + +from nemo_rl.distributed.mx_vllm_reshard_receiver import ( + MxVllmReshardReceiver, +) + + +# The adapter picks its call by inspecting MX's signature, so a bare MagicMock +# (which reports only *args/**kwargs) cannot stand in for either contract. These +# two stubs carry the real shapes, and autospec'ing them makes a call the +# installed MX would reject fail here too. +def _discover_with_flag(count, *, timeout=None, with_tensors=True): + """MX that can skip building the shard tables.""" + + +def _discover_without_flag(count, *, timeout=None): + """MX predating the flag, which always builds them.""" + + +def _build_receiver(discover=_discover_with_flag): + manager = MagicMock() + receiver_client = MagicMock() + raw_receiver = MagicMock(_manager=manager, _mx_client=receiver_client) + client = MagicMock() + rendezvous = MagicMock() + rendezvous.discover_trainers = create_autospec(discover) + with ( + patch( + "modelexpress.engines.vllm.refit.receiver.VllmReshardReceiver", + return_value=raw_receiver, + ) as receiver_cls, + patch("modelexpress.client.MxClient", return_value=client) as client_cls, + patch( + "modelexpress.refit.reshard.rendezvous.MxReshardRendezvous", + return_value=rendezvous, + ) as rendezvous_cls, + ): + wrapper = MxVllmReshardReceiver( + model="model", + vllm_config="vllm-config", + model_config="model-config", + model_name="Qwen/Qwen3-30B", + server_url="mx.example:8001", + agent_name="nemo-rl-vllm-3", + local_rank=1, + global_rank=3, + num_trainer_sources=16, + device="cuda:1", + listen_port=19003, + timeout=77.0, + ) + return ( + wrapper, + raw_receiver, + manager, + receiver_client, + client, + rendezvous, + receiver_cls, + client_cls, + rendezvous_cls, + ) + + +def test_receiver_is_constructed_with_the_expected_mx_arguments(): + ( + _wrapper, + _raw, + _manager, + _receiver_client, + client, + _rendezvous, + receiver_cls, + client_cls, + rendezvous_cls, + ) = _build_receiver() + + client_cls.assert_called_once_with(server_url="mx.example:8001") + rendezvous_cls.assert_called_once_with( + client, + role="inference", + rank=3, + model_name="Qwen/Qwen3-30B", + ) + receiver_cls.assert_called_once_with( + model="model", + vllm_config="vllm-config", + model_config="model-config", + model_name="Qwen/Qwen3-30B", + mx_server="mx.example:8001", + agent_name="nemo-rl-vllm-3", + local_rank=1, + global_rank=3, + num_trainer_sources=16, + device="cuda:1", + listen_port=19003, + timeout=77.0, + ) + + +def test_version_mismatch_fails_before_receiver_install(): + wrapper, raw, _, _, _, rendezvous, *_ = _build_receiver() + rendezvous.discover_trainers.return_value = [ + SimpleNamespace(publisher_step=8) for _ in range(16) + ] + + with pytest.raises(RuntimeError, match="requested 9"): + wrapper.update_weights(9) + + raw.update_weights.assert_not_called() + + +def test_matching_version_is_forwarded_with_timeout(): + wrapper, raw, _, _, _, rendezvous, *_ = _build_receiver() + rendezvous.discover_trainers.return_value = [ + SimpleNamespace(publisher_step=9) for _ in range(16) + ] + raw.update_weights.return_value = {"step": 9} + + assert wrapper.update_weights(9) == {"step": 9} + rendezvous.discover_trainers.assert_called_once_with( + 16, timeout=77.0, with_tensors=False + ) + raw.update_weights.assert_called_once_with(9, timeout=77.0) + + +def test_phase_telemetry_cannot_fail_a_successful_refit(capsys): + """Reporting must never break the operation it measures. + + The phase split reads fields off the rendezvous payloads, and an earlier + version accessed ``payload.tensors`` directly. Any payload shape without that + field then raised AttributeError *after* the weights had already installed, + turning a successful refit into a failed one. + """ + wrapper, raw, _, _, _, rendezvous, *_ = _build_receiver() + # Payloads carry the version stamp but no shard table, which is exactly the + # shape that used to abort the refit. + rendezvous.discover_trainers.return_value = [ + SimpleNamespace(publisher_step=9) for _ in range(16) + ] + raw.update_weights.return_value = {"step": 9} + + assert wrapper.update_weights(9) == {"step": 9} + + record = json.loads(capsys.readouterr().out.split("MX_RECV_PHASE ", 1)[1]) + assert record["step"] == 9 + assert record["rank"] == 3 + assert record["trainer_sources"] == 16 + # Absent shard tables count as zero rather than aborting. + assert record["tensors_seen"] == 0 + assert record["discover_s"] >= 0.0 and record["mx_update_s"] >= 0.0 + + +def test_phase_telemetry_counts_shard_table_entries(capsys): + """``tensors_seen`` is what showed the quorum cost scales with source count + rather than bytes moved, so it has to actually count.""" + wrapper, raw, _, _, _, rendezvous, *_ = _build_receiver() + rendezvous.discover_trainers.return_value = [ + SimpleNamespace(publisher_step=4, tensors=tuple(range(100))) for _ in range(16) + ] + raw.update_weights.return_value = {"step": 4} + + wrapper.update_weights(4) + + record = json.loads(capsys.readouterr().out.split("MX_RECV_PHASE ", 1)[1]) + assert record["tensors_seen"] == 1600 + + +def test_shutdown_releases_manager_and_both_clients(): + wrapper, _, manager, receiver_client, client, _, *_ = _build_receiver() + + wrapper.shutdown() + wrapper.shutdown() + + manager.shutdown.assert_called_once() + receiver_client.close.assert_called_once() + client.close.assert_called_once() + + +def test_quorum_skips_shard_tables_it_does_not_need(): + """The version check needs one integer per rank, not the geometry. MX keeps the + geometry from its own one-time prepare, so fetching it per step rebuilds an + identical table every refit.""" + wrapper, raw, _, _, _, rendezvous, *_ = _build_receiver() + rendezvous.discover_trainers.return_value = [ + SimpleNamespace(publisher_step=2) for _ in range(16) + ] + raw.update_weights.return_value = {"step": 2} + + wrapper.update_weights(2) + + rendezvous.discover_trainers.assert_called_once_with( + 16, timeout=77.0, with_tensors=False + ) + + +def test_quorum_omits_the_flag_when_mx_predates_it(capsys): + """The two changes need not land in lockstep, so an MX without the flag must + still work rather than fail every refit with a TypeError. The autospec would + raise that TypeError if the adapter passed the flag anyway.""" + wrapper, raw, _, _, _, rendezvous, *_ = _build_receiver( + discover=_discover_without_flag + ) + rendezvous.discover_trainers.return_value = [ + SimpleNamespace(publisher_step=5, tensors=(1, 2)) for _ in range(16) + ] + raw.update_weights.return_value = {"step": 5} + + assert wrapper.update_weights(5) == {"step": 5} + + rendezvous.discover_trainers.assert_called_once_with(16, timeout=77.0) + record = json.loads(capsys.readouterr().out.split("MX_RECV_PHASE ", 1)[1]) + assert record["tensors_seen"] == 32 + + +def test_a_real_typeerror_from_discovery_is_not_swallowed(): + """Detecting the flag by catching TypeError also caught TypeErrors raised + *inside* discovery, silently retrying them as the slower full fetch. A bug in + MX has to surface as a bug rather than as an unexplained per-step slowdown.""" + wrapper, _, _, _, _, rendezvous, *_ = _build_receiver() + rendezvous.discover_trainers.side_effect = TypeError("boom inside MX") + + with pytest.raises(TypeError, match="boom inside MX"): + wrapper.update_weights(1) + + rendezvous.discover_trainers.assert_called_once() + + +def test_recorded_entry_count_is_preferred_over_the_built_table(capsys): + """On the quorum path the table is deliberately absent, so the count has to + come from the payload's own tally or the metric silently reads zero.""" + wrapper, raw, _, _, _, rendezvous, *_ = _build_receiver() + rendezvous.discover_trainers.return_value = [ + SimpleNamespace(publisher_step=6, tensors=[], entry_count=lambda: 4922) + for _ in range(16) + ] + raw.update_weights.return_value = {"step": 6} + + wrapper.update_weights(6) + + record = json.loads(capsys.readouterr().out.split("MX_RECV_PHASE ", 1)[1]) + assert record["tensors_seen"] == 16 * 4922 diff --git a/tests/unit/models/generation/test_vllm_nixl_worker.py b/tests/unit/models/generation/test_vllm_nixl_worker.py index 09bf0ac5d4..63038eb471 100644 --- a/tests/unit/models/generation/test_vllm_nixl_worker.py +++ b/tests/unit/models/generation/test_vllm_nixl_worker.py @@ -84,6 +84,18 @@ def test_configure_nixl_worker_uses_vllm_extension_points(): } +def test_configure_nixl_worker_enables_early_init_for_mx_reshard(): + vllm_kwargs = {"additional_config": {"existing": True}} + + configure_nixl_worker({"refit_transport": "mx_reshard"}, vllm_kwargs) + + assert vllm_kwargs["worker_cls"] == NIXL_VLLM_WORKER + assert vllm_kwargs["additional_config"] == { + "existing": True, + "nemo_rl_mx_reshard_nixl": True, + } + + def test_configure_nixl_worker_rejects_incompatible_worker_class(): with pytest.raises(ValueError, match="worker_cls to be unset"): configure_nixl_worker( @@ -121,6 +133,22 @@ def test_preinit_nixl_from_vllm_config_is_disabled_without_nixl_config(): assert preinit_nixl_from_vllm_config(config) is None +def test_preinit_nixl_from_vllm_config_uses_default_backend_for_mx(monkeypatch): + from nemo_rl.utils.checkpoint_engines import nixl + + calls = [] + agent = object() + monkeypatch.setattr( + nixl, + "preinit_nixl_agent", + lambda **kwargs: calls.append(kwargs) or agent, + ) + config = SimpleNamespace(additional_config={"nemo_rl_mx_reshard_nixl": True}) + + assert preinit_nixl_from_vllm_config(config) is agent + assert calls == [{}] + + @pytest.mark.vllm def test_nixl_worker_preinitializes_before_vllm_worker(monkeypatch): from nemo_rl.models.generation.vllm import vllm_backend diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index a79d01961a..15bc9a90bb 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -46,6 +46,38 @@ pytestmark = pytest.mark.mcore +def test_mx_publisher_wires_the_live_layer_qkv_geometry_resolver(): + source_path = ( + Path(__file__).parents[4] + / "nemo_rl/models/policy/workers/megatron_policy_worker.py" + ) + tree = ast.parse(source_path.read_text()) + method = next( + node + for class_node in tree.body + if isinstance(class_node, ast.ClassDef) + and class_node.name == "MegatronPolicyWorkerImpl" + for node in class_node.body + if isinstance(node, ast.FunctionDef) + and node.name == "init_mx_reshard_publisher" + ) + collect_call = next( + node + for node in ast.walk(method) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "collect_megatron_publish_set" + ) + resolver = next( + keyword.value + for keyword in collect_call.keywords + if keyword.arg == "qkv_geometry_resolver" + ) + + assert isinstance(resolver, ast.Name) + assert resolver.id == "resolve_qkv_geometry_from_param" + + def test_model_owned_packing_capability_is_detected(): from nemo_rl.models.policy.workers.megatron_policy_worker import ( _model_self_packs_for_cp, diff --git a/tests/unit/weight_sync/test_mx_reshard_weight_synchronizer.py b/tests/unit/weight_sync/test_mx_reshard_weight_synchronizer.py new file mode 100644 index 0000000000..a655534143 --- /dev/null +++ b/tests/unit/weight_sync/test_mx_reshard_weight_synchronizer.py @@ -0,0 +1,277 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from types import SimpleNamespace +from unittest.mock import MagicMock, call, patch + +import pytest + +from nemo_rl.weight_sync.mx_reshard_weight_synchronizer import ( + MxReshardWeightSynchronizer, + check_mx_reshard_refit_support, +) + + +def _cluster(size: int) -> MagicMock: + cluster = MagicMock() + cluster.world_size.return_value = size + return cluster + + +def _generation() -> MagicMock: + generation = MagicMock() + generation.cfg = {"refit_transport": "mx_reshard"} + return generation + + +def test_factory_selects_mx_reshard() -> None: + from nemo_rl.weight_sync.factory import create_weight_synchronizer + + sync = create_weight_synchronizer( + policy=MagicMock(), + generation=_generation(), + generation_backend="vllm", + colocated=False, + train_cluster=_cluster(8), + inference_cluster=_cluster(4), + ) + assert isinstance(sync, MxReshardWeightSynchronizer) + + +def test_config_normalizes_mx_scope() -> None: + from nemo_rl.models.generation.vllm.config import normalize_vllm_refit_config + + config = { + "refit_transport": "mx_reshard", + "refit_cfg": { + "mx_reshard": { + "server_url": "mx.example:8001", + "timeout_s": 30, + "listen_port_base": 19000, + } + }, + } + normalized = normalize_vllm_refit_config(config) # type: ignore[arg-type] + assert normalized is not None + assert normalized.mx_reshard.server_url == "mx.example:8001" + assert normalized.mx_reshard.timeout_s == 30 + assert normalized.mx_reshard.publisher_listen_port_base == 19000 + assert normalized.mx_reshard.receiver_listen_port_base == 29000 + + +def test_config_normalizes_distinct_role_specific_port_bases() -> None: + from nemo_rl.models.generation.vllm.config import normalize_vllm_refit_config + + config = { + "refit_transport": "mx_reshard", + "refit_cfg": { + "mx_reshard": { + "publisher_listen_port_base": 19000, + "receiver_listen_port_base": 29000, + } + }, + } + normalized = normalize_vllm_refit_config(config) # type: ignore[arg-type] + assert normalized is not None + assert normalized.mx_reshard.publisher_listen_port_base == 19000 + assert normalized.mx_reshard.receiver_listen_port_base == 29000 + + +def test_support_validation_rejects_colocation_and_dtensor() -> None: + master = SimpleNamespace( + policy={ + "generation": { + "backend": "vllm", + "colocated": {"enabled": True}, + }, + "megatron_cfg": {"enabled": True}, + "dtensor_cfg": {"enabled": True}, + } + ) + with pytest.raises(ValueError, match="mx_reshard refit configuration"): + check_mx_reshard_refit_support(master) + + +def _supported_gqa_master(*, heads: int = 64, query_groups: int = 2): + return SimpleNamespace( + policy={ + "generation": { + "backend": "vllm", + "colocated": {"enabled": False}, + "vllm_cfg": {"kv_cache_dtype": "auto"}, + }, + "megatron_cfg": { + "enabled": True, + "tensor_model_parallel_size": 8, + "expert_tensor_parallel_size": 1, + "num_attention_heads": heads, + "num_query_groups": query_groups, + }, + "dtensor_cfg": {"enabled": False}, + } + ) + + +def test_support_validation_accepts_kv_heads_below_tp() -> None: + check_mx_reshard_refit_support(_supported_gqa_master()) + + +def test_support_validation_rejects_query_heads_that_do_not_form_groups() -> None: + with pytest.raises(ValueError, match="divisible by num_query_groups"): + check_mx_reshard_refit_support(_supported_gqa_master(heads=63, query_groups=2)) + + +def test_result_validation_accepts_flattened_and_nested_ray_shapes() -> None: + MxReshardWeightSynchronizer._require_all(True, "scalar") + MxReshardWeightSynchronizer._require_all([True, (True, [True])], "nested") + with pytest.raises(RuntimeError, match="nested failure"): + MxReshardWeightSynchronizer._require_all( + [True, [True, False]], "nested failure" + ) + + +@patch("nemo_rl.weight_sync.mx_reshard_weight_synchronizer.ray") +def test_publish_quorum_completes_before_pull(mock_ray: MagicMock) -> None: + events: list[str] = [] + policy = MagicMock() + generation = _generation() + publish_refs = [object(), object()] + pull_refs = [object(), object()] + policy.publish_mx_reshard_weights.side_effect = lambda **_: ( + events.append("publish-called") or publish_refs + ) + generation.update_weights_from_mx_reshard.side_effect = lambda **_: ( + events.append("pull-called") or pull_refs + ) + + def get(refs): + if refs is publish_refs: + events.append("publish-complete") + return [True, True] + if refs is pull_refs: + events.append("pull-complete") + return [True, True] + raise AssertionError(refs) + + mock_ray.get.side_effect = get + sync = MxReshardWeightSynchronizer(policy, generation, _cluster(2), _cluster(2)) + sync.sync_weights() + + assert events == [ + "publish-called", + "publish-complete", + "pull-called", + "pull-complete", + ] + policy.publish_mx_reshard_weights.assert_called_once_with(version=1) + generation.update_weights_from_mx_reshard.assert_called_once_with(version=1) + + +@patch("nemo_rl.weight_sync.mx_reshard_weight_synchronizer.ray") +def test_publish_failure_prevents_pull(mock_ray: MagicMock) -> None: + policy = MagicMock() + generation = _generation() + policy.publish_mx_reshard_weights.return_value = ["publish-ref"] + mock_ray.get.return_value = [True, False] + sync = MxReshardWeightSynchronizer(policy, generation, _cluster(2), _cluster(2)) + + with pytest.raises(RuntimeError, match="publish failed"): + sync.sync_weights() + + generation.update_weights_from_mx_reshard.assert_not_called() + assert sync.is_stale + + +@patch("nemo_rl.weight_sync.mx_reshard_weight_synchronizer.ray") +def test_receiver_failure_does_not_commit_version(mock_ray: MagicMock) -> None: + policy = MagicMock() + generation = _generation() + policy.publish_mx_reshard_weights.return_value = ["publish-ref"] + generation.update_weights_from_mx_reshard.return_value = ["pull-ref"] + mock_ray.get.side_effect = [[True], [False]] + sync = MxReshardWeightSynchronizer(policy, generation, _cluster(1), _cluster(1)) + + with pytest.raises(RuntimeError, match="receive failed"): + sync.sync_weights() + + assert sync.is_stale + assert sync._version == 0 + assert generation.mock_calls[-1] == call.update_weights_from_mx_reshard(version=1) + + +@patch("nemo_rl.weight_sync.mx_reshard_weight_synchronizer.ray") +def test_init_uses_physical_train_world_size(mock_ray: MagicMock) -> None: + policy = MagicMock() + generation = _generation() + policy.init_mx_reshard_publisher.return_value = ["trainer-init"] + generation.init_mx_reshard_receiver.return_value = ["receiver-init"] + mock_ray.get.side_effect = [[True] * 16, [True] * 8] + sync = MxReshardWeightSynchronizer(policy, generation, _cluster(16), _cluster(8)) + + sync.init_communicator() + + policy.init_mx_reshard_publisher.assert_called_once_with(train_world_size=16) + generation.init_mx_reshard_receiver.assert_called_once_with( + train_world_size=16, + inference_world_size=8, + ) + + +@patch("nemo_rl.weight_sync.mx_reshard_weight_synchronizer.ray") +def test_init_rejects_overlapping_port_ranges_before_worker_rpcs( + mock_ray: MagicMock, +) -> None: + policy = MagicMock() + generation = _generation() + generation.cfg["refit_cfg"] = { + "mx_reshard": { + "publisher_listen_port_base": 19000, + "receiver_listen_port_base": 19001, + } + } + sync = MxReshardWeightSynchronizer(policy, generation, _cluster(2), _cluster(2)) + + with pytest.raises(ValueError, match="listen port ranges overlap"): + sync.init_communicator() + + policy.init_mx_reshard_publisher.assert_not_called() + generation.init_mx_reshard_receiver.assert_not_called() + mock_ray.get.assert_not_called() + + +@patch("nemo_rl.weight_sync.mx_reshard_weight_synchronizer.ray") +def test_shutdown_invokes_both_cleanup_quorums(mock_ray: MagicMock) -> None: + events: list[str] = [] + policy = MagicMock() + generation = _generation() + policy.shutdown_mx_reshard_publisher.side_effect = lambda: ( + events.append("publisher-called") or ["trainer-shutdown"] + ) + generation.shutdown_mx_reshard_receiver.side_effect = lambda: ( + events.append("receiver-called") or ["receiver-shutdown"] + ) + mock_ray.get.side_effect = [[True, True], [[True], True]] + sync = MxReshardWeightSynchronizer(policy, generation, _cluster(2), _cluster(2)) + + sync.shutdown() + sync.shutdown() + + policy.shutdown_mx_reshard_publisher.assert_called_once_with() + generation.shutdown_mx_reshard_receiver.assert_called_once_with() + assert events == ["receiver-called", "publisher-called"] + + +@patch("nemo_rl.weight_sync.mx_reshard_weight_synchronizer.ray") +def test_shutdown_attempts_publisher_cleanup_after_receiver_failure( + mock_ray: MagicMock, +) -> None: + policy = MagicMock() + generation = _generation() + policy.shutdown_mx_reshard_publisher.return_value = ["trainer-shutdown"] + generation.shutdown_mx_reshard_receiver.return_value = ["receiver-shutdown"] + mock_ray.get.side_effect = [RuntimeError("receiver cleanup failed"), [True]] + sync = MxReshardWeightSynchronizer(policy, generation, _cluster(1), _cluster(1)) + + with pytest.raises(RuntimeError, match="receiver cleanup failed"): + sync.shutdown() + + policy.shutdown_mx_reshard_publisher.assert_called_once_with()