From d99d20f57312c1c28b8e4527b4f4ca5970a02d03 Mon Sep 17 00:00:00 2001 From: haonan3 Date: Fri, 31 Jul 2026 12:44:49 +0800 Subject: [PATCH] chore(scripts): stop scripts/ regrowing a test tree; move converters into datasets/ No-unenforced-test-tree policy (#99/#267): drop the seven ep_verify/ harnesses (#158) and verify_refl_kl_batching.py (#210) -- recoverable from git history -- and move the two dataset converters next to the data they produce (datasets/video_r1_260k/, datasets/daily_omni_av/), matching the datasets// convention and the recipes' default data paths. Also fix the daily-omni docstring example, which invoked the video-r1 script by mistake. scripts/ now holds only the check-recipe-targets CI guard. --- ...vert_daily_omni_dataset_format_to_unirl.py | 2 +- .../convert_video_r1_260k_to_unirl.py | 4 +- scripts/ep_verify/measure_ep_load.py | 67 ------- .../ep_verify/qwen3_moe_layout_roundtrip.py | 86 -------- scripts/ep_verify/unirl_ep_backend_real.py | 187 ------------------ .../unirl_ep_checkpoint_roundtrip.py | 130 ------------ scripts/ep_verify/unirl_ep_sync_verify.py | 110 ----------- scripts/ep_verify/unirl_ep_verify.py | 182 ----------------- scripts/ep_verify/unirl_grpo_ep_real.py | 160 --------------- scripts/verify_refl_kl_batching.py | 166 ---------------- 10 files changed, 3 insertions(+), 1091 deletions(-) rename {scripts => datasets/daily_omni_av}/convert_daily_omni_dataset_format_to_unirl.py (98%) rename {scripts => datasets/video_r1_260k}/convert_video_r1_260k_to_unirl.py (98%) delete mode 100644 scripts/ep_verify/measure_ep_load.py delete mode 100644 scripts/ep_verify/qwen3_moe_layout_roundtrip.py delete mode 100644 scripts/ep_verify/unirl_ep_backend_real.py delete mode 100644 scripts/ep_verify/unirl_ep_checkpoint_roundtrip.py delete mode 100644 scripts/ep_verify/unirl_ep_sync_verify.py delete mode 100644 scripts/ep_verify/unirl_ep_verify.py delete mode 100644 scripts/ep_verify/unirl_grpo_ep_real.py delete mode 100755 scripts/verify_refl_kl_batching.py diff --git a/scripts/convert_daily_omni_dataset_format_to_unirl.py b/datasets/daily_omni_av/convert_daily_omni_dataset_format_to_unirl.py similarity index 98% rename from scripts/convert_daily_omni_dataset_format_to_unirl.py rename to datasets/daily_omni_av/convert_daily_omni_dataset_format_to_unirl.py index 7abf3a288..32f5afe56 100755 --- a/scripts/convert_daily_omni_dataset_format_to_unirl.py +++ b/datasets/daily_omni_av/convert_daily_omni_dataset_format_to_unirl.py @@ -5,7 +5,7 @@ embedded audio track from that same file when ``use_audio_in_video=true``. Example: - python scripts/convert_video_r1_260k_to_unirl.py \ + python datasets/daily_omni_av/convert_daily_omni_dataset_format_to_unirl.py \ --train-input /path/to/daily_omni_av_train.jsonl \ --val-input /path/to/daily_omni_av_val.jsonl \ --out-dir datasets/daily_omni_av diff --git a/scripts/convert_video_r1_260k_to_unirl.py b/datasets/video_r1_260k/convert_video_r1_260k_to_unirl.py similarity index 98% rename from scripts/convert_video_r1_260k_to_unirl.py rename to datasets/video_r1_260k/convert_video_r1_260k_to_unirl.py index f54c4e323..307849dcc 100644 --- a/scripts/convert_video_r1_260k_to_unirl.py +++ b/datasets/video_r1_260k/convert_video_r1_260k_to_unirl.py @@ -10,9 +10,9 @@ Missing videos are skipped unless ``--keep-missing`` is set. Example:: - python scripts/convert_video_r1_260k_to_unirl.py \ + python datasets/video_r1_260k/convert_video_r1_260k_to_unirl.py \ --data-root /path/to/Video-R1-data \ - --out-dir /path/to/output/video_r1_260k \ + --out-dir datasets/video_r1_260k \ --sources CLEVRER,STAR,NeXT-QA,PerceptionTest \ --max-total 20000 --val-count 200 """ diff --git a/scripts/ep_verify/measure_ep_load.py b/scripts/ep_verify/measure_ep_load.py deleted file mode 100644 index 0a9ee0e90..000000000 --- a/scripts/ep_verify/measure_ep_load.py +++ /dev/null @@ -1,67 +0,0 @@ -"""Isolate the EP-loader host-RAM behavior: full-read vs per-rank sliced-read. - -Pure safetensors I/O (no GPU/dist). Run once per mode in a FRESH process so -ru_maxrss (process high-water mark) reflects only that mode. - -Usage: python measure_ep_load.py [ep_size] [ep_rank] -""" - -import glob -import os -import resource -import sys - -from safetensors import safe_open - - -def maxrss_gb() -> float: - # Linux ru_maxrss is in KiB. - return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 / 1024 - - -def is_expert(key: str) -> bool: - return "mlp.experts." in key and ("gate_up_proj" in key or "down_proj" in key) - - -def main(): - d = sys.argv[1] - mode = sys.argv[2] - ep = int(sys.argv[3]) if len(sys.argv) > 3 else 4 - ep_rank = int(sys.argv[4]) if len(sys.argv) > 4 else 0 - shards = sorted(glob.glob(os.path.join(d, "*.safetensors"))) - - read_elems = 0 - if mode == "full": - # Emulates the OLD loader: _read_safetensors_dir -> load_file copies every - # tensor into RAM and keeps the whole dict resident (real host RAM). - from safetensors.torch import load_file - - sd = {} - for s in shards: - sd.update(load_file(s, device="cpu")) - read_elems = sum(t.numel() for t in sd.values()) - # ``sd`` stays resident until the function returns, so ru_maxrss captures - # the full footprint without an extra alias. - elif mode == "sliced": - # NEW loader: read only THIS ep rank's expert block; .clone() forces a real - # materialization off the mmap, then free it immediately (peak = one block). - for s in shards: - with safe_open(s, framework="pt", device="cpu") as f: - for k in f.keys(): - sl = f.get_slice(k) - shp = sl.get_shape() - if is_expert(k) and ep > 1: - n = shp[0] // ep - blk = sl[ep_rank * n : (ep_rank + 1) * n].clone() - else: - blk = sl[:].clone() - read_elems += blk.numel() - del blk - else: - raise SystemExit("mode must be full|sliced") - - print(f"mode={mode} ep={ep} ep_rank={ep_rank} read_elems={read_elems / 1e6:.0f}M peak_host_rss={maxrss_gb():.2f}GB") - - -if __name__ == "__main__": - main() diff --git a/scripts/ep_verify/qwen3_moe_layout_roundtrip.py b/scripts/ep_verify/qwen3_moe_layout_roundtrip.py deleted file mode 100644 index 9d651cf48..000000000 --- a/scripts/ep_verify/qwen3_moe_layout_roundtrip.py +++ /dev/null @@ -1,86 +0,0 @@ -"""Synthetic bit-exact check for the shared Qwen3-MoE expert layout.""" - -from __future__ import annotations - -import torch - -from unirl.train.backend.veomni.ep.models.qwen3_moe import ( - build_local_fused_block, - iter_hf_expert_tensors, -) - - -def main() -> None: - prefix = "model.layers.0.mlp" - num_experts, local_experts, intermediate, hidden = 4, 2, 3, 2 - source = {} - for expert in range(num_experts): - base = expert * 100 - source[f"{prefix}.experts.{expert}.gate_proj.weight"] = ( - torch.arange(intermediate * hidden).reshape(intermediate, hidden) + base - ) - source[f"{prefix}.experts.{expert}.up_proj.weight"] = ( - torch.arange(intermediate * hidden).reshape(intermediate, hidden) + base + 20 - ) - source[f"{prefix}.experts.{expert}.down_proj.weight"] = ( - torch.arange(hidden * intermediate).reshape(hidden, intermediate) + base + 40 - ) - - gate_up_blocks = [] - down_blocks = [] - for ep_rank in range(num_experts // local_experts): - gate_up_blocks.append( - build_local_fused_block( - fused_param_name=f"{prefix}.experts.gate_up_proj", - expected_shape=(local_experts, 2 * intermediate, hidden), - ep_rank=ep_rank, - available_keys=set(source), - get_tensor=source.__getitem__, - ) - ) - down_blocks.append( - build_local_fused_block( - fused_param_name=f"{prefix}.experts.down_proj", - expected_shape=(local_experts, hidden, intermediate), - ep_rank=ep_rank, - available_keys=set(source), - get_tensor=source.__getitem__, - ) - ) - - recovered = dict( - iter_hf_expert_tensors( - f"{prefix}.experts.gate_up_proj", - torch.cat(gate_up_blocks), - ) - ) - recovered.update( - iter_hf_expert_tensors( - f"{prefix}.experts.down_proj", - torch.cat(down_blocks), - ) - ) - assert source.keys() == recovered.keys() - for key in source: - torch.testing.assert_close(recovered[key], source[key], rtol=0, atol=0) - - partial = dict(source) - partial.pop(f"{prefix}.experts.1.up_proj.weight") - try: - build_local_fused_block( - fused_param_name=f"{prefix}.experts.gate_up_proj", - expected_shape=(local_experts, 2 * intermediate, hidden), - ep_rank=0, - available_keys=set(partial), - get_tensor=partial.__getitem__, - ) - except RuntimeError: - pass - else: - raise AssertionError("partial per-expert checkpoint did not fail closed") - - print(f"Qwen3-MoE layout round-trip PASS ({len(recovered)}/{len(source)} tensors)", flush=True) - - -if __name__ == "__main__": - main() diff --git a/scripts/ep_verify/unirl_ep_backend_real.py b/scripts/ep_verify/unirl_ep_backend_real.py deleted file mode 100644 index b5ad6a2a1..000000000 --- a/scripts/ep_verify/unirl_ep_backend_real.py +++ /dev/null @@ -1,187 +0,0 @@ -"""REAL UniRL VeOmniBackend EP verification. - -Unlike unirl_ep_verify.py (which replicated the init call + drove the wrap/clip -functions), this drives the **actual** ``VeOmniBackend`` class end to end: - - VeOmniBackend.__init__ (init_parallel_state w/ ep -> veomni_parallelize -> - _attach_extra_parallel_param_groups -> load_trainable_weights -> - build optimizer/scheduler) - -> backend.zero_grad / loss.backward / backend.optimizer_step (EP-aware clip) - -on a VeOmni-patched Qwen3-MoE (meta-init + real stacked safetensors load). - -Launch: torchrun --nproc_per_node=8 unirl_ep_backend_real.py -""" - -import json -import os -import sys -import time - -import torch -import torch.distributed as dist - -from unirl.train.backend.base import LrSchedulerConfig, OptimizerConfig -from unirl.train.backend.veomni.backend import VeOmniBackend -from unirl.train.configs import FSDPConfig - - -class _SimpleBundle: - """Minimal meta-init bundle: a VeOmni MoE transformer + stashed weights dir. - - Satisfies the duck-typed contract VeOmniBackend uses: - ``.transformer`` (resolve_trainable_module fallback) and - ``._transformer_weights_path`` (load_trainable_weights Pattern B). - """ - - def __init__(self, transformer, weights_dir): - self.transformer = transformer - self._transformer_weights_path = weights_dir - - def prepare_for_expert_parallel(self): - if not callable(getattr(self.transformer, "get_parallel_plan", None)): - raise RuntimeError("test transformer does not expose get_parallel_plan()") - - -def build_meta_moe(config_path): - from unirl.models.types.meta_init import finalize_meta_init - from unirl.train.backend.veomni import _compat - - _compat.ensure_qwen3_moe_installed() - from veomni.arguments import OpsImplementationConfig - from veomni.models.auto import build_foundation_model - - ops = OpsImplementationConfig( - attn_implementation="flash_attention_2", - moe_implementation="fused_triton", - cross_entropy_loss_implementation="eager", - rms_norm_implementation="eager", - swiglu_mlp_implementation="eager", - rotary_pos_emb_implementation="eager", - load_balancing_loss_implementation="eager", - ) - model = build_foundation_model( - config_path=config_path, - weights_path=None, - torch_dtype="bfloat16", - init_device="meta", - ops_implementation=ops, - ) - return finalize_meta_init(model, dtype=torch.bfloat16) - - -def validate_recovered_rope(model): - """Assert that the backend restored finite, nonzero RoPE frequencies.""" - n = 0 - for name, module in model.named_modules(): - inv_freq = getattr(module, "inv_freq", None) - if inv_freq is None: - continue - local = inv_freq.to_local() if hasattr(inv_freq, "to_local") else inv_freq - if not bool(torch.isfinite(local).all()) or int(torch.count_nonzero(local)) == 0: - raise RuntimeError(f"backend did not recover RoPE inv_freq for {name!r}") - n += 1 - if n == 0: - raise RuntimeError("test model has no RoPE inv_freq buffer to validate") - return n - - -def main(): - config_dir = sys.argv[1] # dir containing config.json + stacked model.safetensors - ep_size = int(sys.argv[2]) - out_path = sys.argv[3] - steps = int(os.environ.get("STEPS", "8")) - seq_len = int(os.environ.get("SEQ", "4096")) - - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - torch.cuda.set_device(local_rank) - device = torch.device("cuda") - rank = int(os.environ.get("RANK", "0")) - - bundle = _SimpleBundle(build_meta_moe(config_dir), config_dir) - - fsdp_cfg = FSDPConfig(param_dtype="bf16", fsdp_mode="full", reshard_after_forward=True, ep_size=ep_size) - opt_cfg = OptimizerConfig(learning_rate=1e-4, adam_beta1=0.9, adam_beta2=0.95, adam_epsilon=1e-8, weight_decay=0.0) - sched_cfg = LrSchedulerConfig(type="constant", warmup_steps=0, total_steps=100) - - # ---- the REAL backend: full __init__ runs init_parallel_state(ep) + wrap + load + optimizer ---- - backend = VeOmniBackend( - bundle=bundle, - block_class_names=("Qwen3MoeDecoderLayer",), - fsdp_cfg=fsdp_cfg, - optimizer_cfg=opt_cfg, - scheduler_cfg=sched_cfg, - trainable_attr="transformer", - device=device, - rank=rank, - ) - model = backend.model - n_rope = validate_recovered_rope(model) - - from veomni.distributed.parallel_state import get_parallel_state - - ps = get_parallel_state() - ep_enabled = ps.ep_enabled - has_groups = hasattr(model, "_extra_parallel_param_groups") - if rank == 0: - print( - f"[real] ep={ep_size} ep_enabled={ep_enabled} ep_size(ps)={ps.ep_size if ep_enabled else 1} " - f"ep_param_groups={has_groups} rope_validated={n_rope} model={type(model).__name__}", - flush=True, - ) - - vocab = model.config.vocab_size - records = [] - last = None - for step in range(steps): - input_ids = torch.randint(0, min(vocab, 1024), (1, seq_len), device=device) - labels = input_ids.clone() - backend.zero_grad() - out = model(input_ids=input_ids, labels=labels) - loss = out.loss - loss.backward() - gn = backend.optimizer_step(max_grad_norm=1.0) # EP-aware clip + step + sched + ema - torch.cuda.synchronize() - now = time.perf_counter() - dt = None if last is None else now - last - last = now - peak = torch.cuda.max_memory_allocated() / 1e9 - records.append( - {"step": step, "time_s": dt, "peak_alloc_gb": peak, "loss": float(loss.detach()), "grad_norm": float(gn)} - ) - torch.cuda.reset_peak_memory_stats() - if rank == 0: - print( - f"[real] step {step} loss={float(loss.detach()):.4f} gn={float(gn):.4f} peak={peak:.2f}GB dt={dt}", - flush=True, - ) - - local_peak = max(r["peak_alloc_gb"] for r in records) - t = torch.tensor([local_peak], device=device) - dist.all_reduce(t, op=dist.ReduceOp.MAX) - global_peak = t.item() - - if rank == 0: - steady = [r["time_s"] for r in records if r["time_s"] is not None and r["step"] >= 3] - med = sorted(steady)[len(steady) // 2] if steady else None - out = { - "ep_size": ep_size, - "world": dist.get_world_size(), - "ep_enabled": bool(ep_enabled), - "ep_param_groups": bool(has_groups), - "rope_validated": n_rope, - "global_peak_alloc_gb": global_peak, - "median_step_time_s": med, - "records": records, - } - os.makedirs(os.path.dirname(out_path), exist_ok=True) - with open(out_path, "w") as f: - json.dump(out, f, indent=2) - print(f"[real] WROTE {out_path}: ep={ep_size} peak={global_peak:.2f}GB median_step={med}", flush=True) - - dist.barrier() - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/scripts/ep_verify/unirl_ep_checkpoint_roundtrip.py b/scripts/ep_verify/unirl_ep_checkpoint_roundtrip.py deleted file mode 100644 index 13b0b2b75..000000000 --- a/scripts/ep_verify/unirl_ep_checkpoint_roundtrip.py +++ /dev/null @@ -1,130 +0,0 @@ -"""Small EP checkpoint round-trip using synthetic expert tensors. - -This exercises the production single-file model + Adam-state helpers without a -large MoE checkpoint or fused kernels. It defaults to CPU/Gloo; set -``DEVICE=cuda`` for NCCL: - - torchrun --nproc_per_node=4 scripts/ep_verify/unirl_ep_checkpoint_roundtrip.py \ - 2 /tmp/unirl_ep_checkpoint_roundtrip.pt -""" - -from __future__ import annotations - -import os -import sys - -import torch -import torch.distributed as dist -from torch import nn - -from unirl.train.backend.veomni.ep.checkpoint import ( - gather_ep_model_state_dict, - gather_ep_optimizer_state_dict, - load_ep_model_state_dict, - load_ep_optimizer_state_dict, -) - - -class _TinyExpertModel(nn.Module): - def __init__(self, expert_param: nn.Parameter) -> None: - super().__init__() - self.experts = expert_param - self._extra_parallel_param_groups = { - "ep": [self.experts], - "non_extra_parallel": [], - } - - -def main() -> None: - ep_size = int(sys.argv[1]) - checkpoint_path = sys.argv[2] - local_rank = int(os.environ["LOCAL_RANK"]) - device_type = os.environ.get("DEVICE", "cpu").strip().lower() - if device_type == "cuda": - torch.cuda.set_device(local_rank) - device = torch.device("cuda", local_rank) - backend = "nccl" - elif device_type == "cpu": - device = torch.device("cpu") - backend = "gloo" - else: - raise ValueError(f"DEVICE must be 'cpu' or 'cuda', got {device_type!r}") - dist.init_process_group(backend) - - from unirl.train.backend.veomni import _compat - - _compat.ensure_installed() - from veomni.distributed.parallel_state import get_parallel_state, init_parallel_state - - world = dist.get_world_size() - if world % ep_size: - raise ValueError(f"world_size={world} must be divisible by ep_size={ep_size}") - init_parallel_state( - dp_size=world, - ulysses_size=1, - dp_mode="fsdp2", - device_type=device_type, - extra_parallel_sizes=(ep_size,), - extra_parallel_names=("ep",), - ) - ps = get_parallel_state() - ep_rank = int(ps.extra_parallel_rank("ep")) - - from torch.distributed.tensor import Replicate, Shard, distribute_tensor - - # Match VeOmni exactly: the outer EP split is already applied to this - # rank's [E/ep,H] block; the DTensor records only the inner ep_fsdp shard. - local_expert_block = torch.full((2, 4), float(ep_rank + 1), device=device) - full_ep_mesh = ps.extra_parallel_fsdp_device_mesh["ep"] - ep_fsdp_mesh = full_ep_mesh["ep_fsdp"] - assert ep_fsdp_mesh.mesh_dim_names == ("ep_fsdp",) - placements = [Replicate()] * (ep_fsdp_mesh.ndim - 1) + [Shard(1)] - expert_dtensor = distribute_tensor(local_expert_block, ep_fsdp_mesh, placements) - model = _TinyExpertModel(nn.Parameter(expert_dtensor)) - optimizer = torch.optim.AdamW(model.parameters(), lr=0.01, foreach=False) - model.experts.grad = torch.full_like(model.experts, 0.1 * (ep_rank + 1)) - optimizer.step() - - expected_param = model.experts.to_local().detach().clone() - expected_optim = { - key: (value.to_local() if hasattr(value, "to_local") else value).detach().clone() - for key, value in optimizer.state[model.experts].items() - if isinstance(value, torch.Tensor) - } - - model_state = gather_ep_model_state_dict(model) - optimizer_state = gather_ep_optimizer_state_dict(model, optimizer) - if dist.get_rank() == 0: - local_global_shape = tuple(model.experts.shape) - expected_shape = (local_global_shape[0] * ep_size, *local_global_shape[1:]) - assert tuple(model_state["experts"].shape) == expected_shape - assert tuple(optimizer_state["state"]["experts"]["exp_avg"].shape) == expected_shape - blocks = model_state["experts"].reshape(ep_size, local_global_shape[0], *local_global_shape[1:]) - assert any(not torch.equal(blocks[0], blocks[index]) for index in range(1, ep_size)) - torch.save({"model": model_state, "optimizer": optimizer_state}, checkpoint_path) - dist.barrier() - - with torch.no_grad(): - model.experts.zero_() - for value in optimizer.state[model.experts].values(): - if isinstance(value, torch.Tensor): - value.zero_() - - checkpoint = torch.load(checkpoint_path, map_location="cpu") - load_ep_model_state_dict(model, checkpoint["model"], strict=True) - load_ep_optimizer_state_dict(model, optimizer, checkpoint["optimizer"]) - - torch.testing.assert_close(model.experts.to_local(), expected_param, rtol=0, atol=0) - for key, expected in expected_optim.items(): - actual = optimizer.state[model.experts][key] - actual = actual.to_local() if hasattr(actual, "to_local") else actual - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - - dist.barrier() - if dist.get_rank() == 0: - print(f"EP checkpoint round-trip PASS (world={world}, ep={ep_size})", flush=True) - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/scripts/ep_verify/unirl_ep_sync_verify.py b/scripts/ep_verify/unirl_ep_sync_verify.py deleted file mode 100644 index 3303f680c..000000000 --- a/scripts/ep_verify/unirl_ep_sync_verify.py +++ /dev/null @@ -1,110 +0,0 @@ -"""Verify EP-aware weight sync round-trips the EP-sharded model to HF per-expert -format (no SGLang needed). - -Builds the real VeOmniBackend (ep_size>1) on a MoE checkpoint, runs the ACTUAL -EP weight walk (FullWeightSync._iter_full_tensors_ep), and checks that the emitted -HF per-expert tensors (experts.{e}.gate_proj/up_proj/down_proj) exactly match the -original checkpoint's per-expert weights (cast to the train dtype). This proves the -ep all-gather + stacked->per-expert reverse-convert is correct, which is the only -EP-specific part of pushing weights into a rollout engine. - -Launch: torchrun --nproc_per_node=8 unirl_ep_sync_verify.py -""" - -import sys -import types - -import torch -import torch.distributed as dist -from safetensors import safe_open - -from unirl.distributed.weight_sync.full.base import FullWeightSync -from unirl.models.qwen3_moe import Qwen3MoeBundle -from unirl.train.backend.base import LrSchedulerConfig, OptimizerConfig -from unirl.train.backend.veomni.backend import VeOmniBackend -from unirl.train.configs import FSDPConfig - - -def main(): - ckpt = sys.argv[1] - ep_size = int(sys.argv[2]) if len(sys.argv) > 2 else 4 - - local_rank = int(__import__("os").environ.get("LOCAL_RANK", "0")) - torch.cuda.set_device(local_rank) - device = torch.device("cuda") - rank = int(__import__("os").environ.get("RANK", "0")) - - tok = types.SimpleNamespace(pad_token_id=0, eos_token_id=0, pad_token="", eos_token="") - bundle = Qwen3MoeBundle.from_config(pretrained_model_ckpt_path=ckpt, tokenizer=tok) - backend = VeOmniBackend( - bundle=bundle, - block_class_names=("Qwen3MoeDecoderLayer",), - fsdp_cfg=FSDPConfig(param_dtype="bf16", fsdp_mode="full", reshard_after_forward=True, ep_size=ep_size), - optimizer_cfg=OptimizerConfig( - learning_rate=1e-4, adam_beta1=0.9, adam_beta2=0.95, adam_epsilon=1e-8, weight_decay=0.0 - ), - scheduler_cfg=LrSchedulerConfig(type="constant", warmup_steps=0, total_steps=10), - trainable_attr="transformer", - device=device, - rank=rank, - ) - - # Drive the REAL EP weight walk via a minimal FullWeightSync-shaped object. - fake = types.SimpleNamespace(_backend=backend, _wire_dtype=None, _name_remap={}) - emitted = {} - n_expert_keys = 0 - for name, tensor in FullWeightSync._iter_full_tensors_ep(fake): - if ".experts." in name and any( - name.endswith(s) for s in (".gate_proj.weight", ".up_proj.weight", ".down_proj.weight") - ): - n_expert_keys += 1 - if rank == 0: - emitted[name] = tensor.detach().to("cpu", torch.float32) - - if rank == 0: - f = safe_open(f"{ckpt}/model.safetensors", framework="pt", device="cpu") - ckpt_keys = set(f.keys()) - # Sample experts across layers/expert-ids; verify bit-exact (file cast to bf16, - # the train dtype the backend loaded into, then back to fp32 for compare). - import re - - layer_ids = sorted( - {int(m.group(1)) for k in ckpt_keys if (m := re.search(r"layers\.(\d+)\.mlp\.experts\.0\.", k))} - ) - E = 1 + max(int(m.group(1)) for k in ckpt_keys if (m := re.search(r"experts\.(\d+)\.gate_proj", k))) - checks, ok = 0, 0 - sample_layers = [layer_ids[0], layer_ids[-1]] if layer_ids else [] - sample_experts = sorted({0, 1, E // 2, E - 1}) - for L in sample_layers: - for e in sample_experts: - for proj in ("gate_proj", "up_proj", "down_proj"): - key = f"model.layers.{L}.mlp.experts.{e}.{proj}.weight" - if key not in emitted or key not in ckpt_keys: - print(f"[sync] MISSING {key} (emitted={key in emitted}, ckpt={key in ckpt_keys})", flush=True) - checks += 1 - continue - ref = f.get_tensor(key).to(torch.bfloat16).to(torch.float32) # match backend bf16 load - got = emitted[key] - checks += 1 - if got.shape == ref.shape and torch.equal(got, ref): - ok += 1 - else: - md = (got - ref).abs().max().item() if got.shape == ref.shape else -1 - print( - f"[sync] MISMATCH {key} shape got={tuple(got.shape)} ref={tuple(ref.shape)} maxdiff={md}", - flush=True, - ) - total_expert_keys = len(layer_ids) * E * 3 - print( - f"[sync] EP={ep_size} emitted_expert_keys={len(emitted)} (expected {total_expert_keys}); " - f"bit-exact {ok}/{checks} sampled", - flush=True, - ) - print(f"[sync] RESULT: {'PASS' if ok == checks and len(emitted) == total_expert_keys else 'FAIL'}", flush=True) - - dist.barrier() - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/scripts/ep_verify/unirl_ep_verify.py b/scripts/ep_verify/unirl_ep_verify.py deleted file mode 100644 index dc76a4ce5..000000000 --- a/scripts/ep_verify/unirl_ep_verify.py +++ /dev/null @@ -1,182 +0,0 @@ -"""Verify UniRL's VeOmni-backend EP path on a real MoE model. - -Drives the *actual* UniRL functions touched by the EP change: - * the init_parallel_state(...) call exactly as VeOmniBackend.__init__ issues it - (with extra_parallel_sizes=(ep,)) — i.e. FSDPConfig.ep_size in action; - * unirl.train.backend.veomni.wrap.veomni_parallelize (the real wrap); - * unirl.train.backend.veomni.state.clip_grad_norm (the EP-aware clip). - -Builds a VeOmni-patched Qwen3-MoE (random init on meta — no checkpoint needed), -runs fwd/bwd/clip/step, and records per-step time + per-GPU peak memory so we can -compare ep_size=1 (pure FSDP) vs ep_size>1 (expert-parallel). - -Launch: torchrun --nproc_per_node=8 unirl_ep_verify.py -""" - -import json -import os -import sys -import time - -import torch -import torch.distributed as dist - -from unirl.train.backend.veomni.state import clip_grad_norm -from unirl.train.backend.veomni.wrap import veomni_parallelize - -# --- UniRL code under test --- -from unirl.train.configs import FSDPConfig - - -def main(): - config_path = sys.argv[1] - ep_size = int(sys.argv[2]) - out_path = sys.argv[3] - sp_size = int(sys.argv[4]) if len(sys.argv) > 4 else 1 - steps = int(os.environ.get("STEPS", "8")) - seq_len = int(os.environ.get("SEQ", "4096")) - - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - torch.cuda.set_device(local_rank) - if not dist.is_initialized(): - dist.init_process_group(backend="nccl") - rank = dist.get_rank() - world = dist.get_world_size() - device = torch.device("cuda") - - # FSDPConfig is the real config object the backend consumes; ep_size is the - # field added for EP. - fsdp_cfg = FSDPConfig( - param_dtype="bf16", fsdp_mode="full", reshard_after_forward=True, sp_size=sp_size, ep_size=ep_size - ) - - # ---- EXACT init_parallel_state call from VeOmniBackend.__init__ ---- - from unirl.train.backend.veomni import _compat - - _compat.ensure_qwen3_moe_installed() - from veomni.distributed.parallel_state import get_parallel_state, init_parallel_state - - if world % fsdp_cfg.sp_size != 0: - raise ValueError("world not divisible by sp") - if fsdp_cfg.ep_size > 1 and world % fsdp_cfg.ep_size != 0: - raise ValueError("world not divisible by ep") - init_parallel_state( - dp_size=world // fsdp_cfg.sp_size, - ulysses_size=fsdp_cfg.sp_size, - extra_parallel_sizes=(fsdp_cfg.ep_size,), - extra_parallel_names=("ep",), - extra_parallel_placement_innermost=(False,), - dp_mode="fsdp2", - device_type="cuda", - ) - ps = get_parallel_state() - ep_enabled = ps.ep_enabled - if rank == 0: - print( - f"[verify] world={world} sp={sp_size} ep={ep_size} ep_enabled={ep_enabled} " - f"ep_size(ps)={ps.ep_size if ep_enabled else 1}", - flush=True, - ) - - # ---- build VeOmni MoE model on meta (random init, no checkpoint) ---- - from veomni.arguments import OpsImplementationConfig - from veomni.models.auto import build_foundation_model - - ops = OpsImplementationConfig( - attn_implementation=os.environ.get("ATTN_IMPLEMENTATION", "flash_attention_2"), - moe_implementation="fused_triton", - cross_entropy_loss_implementation="eager", - rms_norm_implementation="eager", - swiglu_mlp_implementation="eager", - rotary_pos_emb_implementation="eager", - load_balancing_loss_implementation="eager", - ) - model = build_foundation_model( - config_path=config_path, - weights_path=None, - torch_dtype="bfloat16", - init_device="meta", - ops_implementation=ops, - ) - # model must expose get_parallel_plan for EP (Shard(0) experts) - has_plan = getattr(model, "get_parallel_plan", None) is not None - if rank == 0: - print(f"[verify] model={type(model).__name__} has_parallel_plan={has_plan}", flush=True) - - # ---- the REAL UniRL wrap (forwards to parallelize_model_fsdp2 -> applies EP) ---- - veomni_parallelize( - model, - block_class_names=("Qwen3MoeDecoderLayer",), - param_dtype=fsdp_cfg.param_dtype, - reshard_after_forward=fsdp_cfg.reshard_after_forward, - ) - has_ep_groups = hasattr(model, "_extra_parallel_param_groups") - if rank == 0: - print(f"[verify] wrapped; _extra_parallel_param_groups={has_ep_groups}", flush=True) - - # Mirror UniRL's real build_optimizer EXACTLY: a single AdamW with - # foreach=False. The single-tensor (per-param) kernel steps each DTensor - # independently, so EP-sharded experts (ep_fsdp mesh) and non-EP params - # (dp_shard mesh) never get stacked across meshes — UniRL's existing - # optimizer already handles EP with no change. - optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-4, foreach=False) - - vocab = model.config.vocab_size - records = [] - last = None - for step in range(steps): - input_ids = torch.randint(0, min(vocab, 1024), (1, seq_len), device=device) - labels = input_ids.clone() - optimizer.zero_grad(set_to_none=True) - out = model(input_ids=input_ids, labels=labels) - loss = out.loss - loss.backward() - gn = clip_grad_norm(model, 1.0) - optimizer.step() - torch.cuda.synchronize() - now = time.perf_counter() - dt = None if last is None else now - last - last = now - peak = torch.cuda.max_memory_allocated() / 1e9 - records.append( - {"step": step, "time_s": dt, "peak_alloc_gb": peak, "loss": float(loss.detach()), "grad_norm": float(gn)} - ) - torch.cuda.reset_peak_memory_stats() - if rank == 0: - print( - f"[verify] step {step} loss={float(loss.detach()):.4f} gn={float(gn):.4f} peak={peak:.2f}GB dt={dt}", - flush=True, - ) - - local_peak = max(r["peak_alloc_gb"] for r in records) - t = torch.tensor([local_peak], device=device) - dist.all_reduce(t, op=dist.ReduceOp.MAX) - global_peak = t.item() - - if rank == 0: - steady = [r["time_s"] for r in records if r["time_s"] is not None and r["step"] >= 3] - med = sorted(steady)[len(steady) // 2] if steady else None - out = { - "ep_size": ep_size, - "sp_size": sp_size, - "world": world, - "ep_enabled": bool(ep_enabled), - "has_parallel_plan": bool(has_plan), - "has_ep_param_groups": bool(has_ep_groups), - "global_peak_alloc_gb": global_peak, - "median_step_time_s": med, - "records": records, - } - out_dir = os.path.dirname(out_path) - if out_dir: - os.makedirs(out_dir, exist_ok=True) - with open(out_path, "w") as f: - json.dump(out, f, indent=2) - print(f"[verify] WROTE {out_path}: ep={ep_size} peak={global_peak:.2f}GB median_step={med}", flush=True) - - dist.barrier() - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/scripts/ep_verify/unirl_grpo_ep_real.py b/scripts/ep_verify/unirl_grpo_ep_real.py deleted file mode 100644 index c24d249f9..000000000 --- a/scripts/ep_verify/unirl_grpo_ep_real.py +++ /dev/null @@ -1,160 +0,0 @@ -"""REAL GRPO training-side step on an EP MoE, through the actual UniRL stack. - -Pipeline (all real UniRL classes): - Qwen3MoeBundle (VeOmni MoE, meta-init, EP-capable) - -> VeOmniBackend(ep_size=N) # full __init__: EP shard + load + optimizer - -> Qwen3ARStage(model=bundle) # installs the replay forward - -> GRPO(stage=...).compute_loss_and_backward(conds, segment, advantages, ...) - == stage.replay (policy fwd on EP MoE) + PPO clip loss + backward - -> backend.optimizer_step(max_grad_norm) # EP-aware clip + step + sched - -Rollout + reward are synthesized (random prompts / responses / advantages) — EP -only affects the TRAINING backend, not the reward signal. old_logp is seeded -from a no-grad replay so the step-0 ratio == 1 (a clean GRPO ratio). - -Launch: torchrun --nproc_per_node=8 unirl_grpo_ep_real.py -""" - -import json -import os -import sys -import time -from types import SimpleNamespace - -import torch -import torch.distributed as dist - -from unirl.algorithms.grpo import GRPO -from unirl.models.qwen3.ar import Qwen3ARStage -from unirl.models.qwen3.conditions import Qwen3ARConditions -from unirl.models.qwen3_moe import Qwen3MoeBundle -from unirl.train.backend.base import LrSchedulerConfig, OptimizerConfig -from unirl.train.backend.veomni.backend import VeOmniBackend -from unirl.train.configs import FSDPConfig -from unirl.types.conditions import TextTokenCondition -from unirl.types.segments.text import TextSegment - - -def main(): - cfg_dir = sys.argv[1] - ep_size = int(sys.argv[2]) - out_path = sys.argv[3] - steps = int(os.environ.get("STEPS", "6")) - B = int(os.environ.get("B", "4")) # prompts (== global batch on this single DP group) - P = int(os.environ.get("P", "64")) # prompt len - R = int(os.environ.get("R", "128")) # response len - - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - torch.cuda.set_device(local_rank) - device = torch.device("cuda") - rank = int(os.environ.get("RANK", "0")) - - # Stub tokenizer: vocab-4096 toy model needs pad_id < vocab (a real Qwen - # tokenizer's pad id 151643 would index past the toy embedding). - tok = SimpleNamespace(pad_token_id=0, eos_token_id=0, pad_token="", eos_token="") - bundle = Qwen3MoeBundle.from_config(pretrained_model_ckpt_path=cfg_dir, tokenizer=tok) - - fsdp_cfg = FSDPConfig(param_dtype="bf16", fsdp_mode="full", reshard_after_forward=True, ep_size=ep_size) - backend = VeOmniBackend( - bundle=bundle, - block_class_names=("Qwen3MoeDecoderLayer",), - fsdp_cfg=fsdp_cfg, - optimizer_cfg=OptimizerConfig( - learning_rate=1e-4, adam_beta1=0.9, adam_beta2=0.95, adam_epsilon=1e-8, weight_decay=0.0 - ), - scheduler_cfg=LrSchedulerConfig(type="constant", warmup_steps=0, total_steps=100), - trainable_attr="transformer", - device=device, - rank=rank, - ) - - from veomni.distributed.parallel_state import get_parallel_state - - ps = get_parallel_state() - ep_enabled = ps.ep_enabled - - stage = Qwen3ARStage(model=bundle) - grpo = GRPO(stage=stage, conditions_cls=Qwen3ARConditions, clip_range=0.2, sampling_temperature=1.0) - - vocab = bundle.transformer.config.vocab_size - hi = min(vocab, 1024) - # Fixed synthetic rollout (same across ep sizes via seeded RNG → comparable). - g = torch.Generator(device="cpu").manual_seed(1234) - prompt_ids = torch.randint(0, hi, (B, P), generator=g).to(device) - attn = torch.ones((B, P), dtype=torch.long, device=device) - conds = {"prompt": TextTokenCondition(input_ids=prompt_ids, attention_mask=attn)} - resp = [torch.randint(0, hi, (R,), generator=g).to(device) for _ in range(B)] - - # Seed old_logp from a no-grad replay so step-0 ratio == 1 (clean GRPO). - seg0 = TextSegment.pack(tokens=resp, log_probs=[torch.zeros(R, device=device) for _ in range(B)]) - with torch.no_grad(): - old_flat = stage.replay(Qwen3ARConditions.from_dict(conds), segment=seg0, temperature=1.0) - old_lists, off = [], 0 - for _ in range(B): - old_lists.append(old_flat[off : off + R].float().cpu()) - off += R - segment = TextSegment.pack(tokens=[r.cpu() for r in resp], log_probs=old_lists) - advantages = torch.randn(B, generator=g) - - records = [] - last = None - for step in range(steps): - backend.zero_grad() - res = grpo.compute_loss_and_backward( - conditions=conds, - segment=segment, - advantages=advantages, - training_progress=0.0, - loss_scale=1.0, - ) - gn = backend.optimizer_step(max_grad_norm=1.0) if res.has_backward else float("nan") - torch.cuda.synchronize() - now = time.perf_counter() - dt = None if last is None else now - last - last = now - peak = torch.cuda.max_memory_allocated() / 1e9 - m = res.metrics - records.append( - { - "step": step, - "time_s": dt, - "peak_alloc_gb": peak, - "policy_loss": m.get("policy_loss"), - "grad_norm": float(gn), - "ratio_mean": m.get("ratio_mean"), - "logp_absdiff": m.get("rollout_replay_logp_absdiff"), - } - ) - torch.cuda.reset_peak_memory_stats() - if rank == 0: - print( - f"[grpo] step {step} loss={m.get('policy_loss'):.5f} gn={float(gn):.4f} " - f"ratio_mean={m.get('ratio_mean')} peak={peak:.2f}GB dt={dt}", - flush=True, - ) - - local_peak = max(r["peak_alloc_gb"] for r in records) - t = torch.tensor([local_peak], device=device) - dist.all_reduce(t, op=dist.ReduceOp.MAX) - if rank == 0: - steady = [r["time_s"] for r in records if r["time_s"] is not None and r["step"] >= 2] - med = sorted(steady)[len(steady) // 2] if steady else None - out = { - "ep_size": ep_size, - "ep_enabled": bool(ep_enabled), - "world": dist.get_world_size(), - "global_peak_alloc_gb": t.item(), - "median_step_time_s": med, - "records": records, - } - os.makedirs(os.path.dirname(out_path), exist_ok=True) - with open(out_path, "w") as f: - json.dump(out, f, indent=2) - print(f"[grpo] WROTE {out_path}: ep={ep_size} peak={t.item():.2f}GB median_step={med}", flush=True) - - dist.barrier() - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/scripts/verify_refl_kl_batching.py b/scripts/verify_refl_kl_batching.py deleted file mode 100755 index cd76e979b..000000000 --- a/scripts/verify_refl_kl_batching.py +++ /dev/null @@ -1,166 +0,0 @@ -#!/usr/bin/env python3 -"""Verify KL DP-batching semantics for the refl recipe (CPU, no Ray, no GPU). - -Regression verification for the P1 review finding on PR #210: the original -``REFLGenerated.kl_loss`` was a per-shard *scalar* ``shared_field`` — DP -collect kept only rank 0's KL and re-broadcast it to every actor rank, which -happened to "work" only on the verified ``batch_size == actor_dp == 8`` -topology (logs duplicated rank 0; other B/dp splits risked a hard shape -mismatch between the routed KL grad and each rank's saved scalar). - -The fix makes ``kl_loss`` a batch-aligned per-sample ``[B]`` concat field. -This script pins the invariants at the exact wire layer DP dispatch uses -(``pytree_chunk`` / ``pytree_cat`` / ``infer_batch_size``) across the -topologies called out in review: B == dp, B > dp, non-power-of-two, dp == 1, -and unequal actor/reward dp. - -Standalone by design (repo policy after #99/#267 is no unenforced test tree): -run it directly whenever the refl recipe's KL/reward wire types change:: - - python scripts/verify_refl_kl_batching.py -""" - -from __future__ import annotations - -import sys -from pathlib import Path - -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -import torch # noqa: E402 - -from experimental.refl.roles import REFLGenerated, REFLLossMetrics # noqa: E402 -from unirl.distributed.tensor.pytree import infer_batch_size, pytree_cat, pytree_chunk # noqa: E402 - -# (batch_size, dp) — B == dp (the only previously-verified shape), B > dp, -# non-power-of-two, and the degenerate dp == 1. -TOPOLOGIES = [(8, 8), (8, 2), (6, 3), (12, 4), (4, 1)] - - -def _generated(batch: int) -> REFLGenerated: - decoded = torch.arange(batch * 6, dtype=torch.float32).reshape(batch, 2, 3) - kl = torch.arange(batch, dtype=torch.float32) + 1.0 # distinct per sample - return REFLGenerated(decoded=decoded, kl_loss=kl) - - -def check_refl_generated_chunk_cat_roundtrip(batch: int, dp: int) -> None: - """Each DP shard sees exactly its own KL rows; merge restores the batch.""" - gen = _generated(batch) - assert infer_batch_size((gen,), {}) == batch - - shards = pytree_chunk(gen, dp, batch) - assert len(shards) == dp - per = batch // dp - for rank, shard in enumerate(shards): - expect = gen.kl_loss[rank * per : (rank + 1) * per] - assert torch.equal(shard.kl_loss, expect), f"rank {rank} got foreign KL rows" - assert torch.equal(shard.decoded, gen.decoded[rank * per : (rank + 1) * per]) - - merged = pytree_cat(shards) - assert torch.equal(merged.kl_loss, gen.kl_loss) - assert torch.equal(merged.decoded, gen.decoded) - - -def check_forward_backward_payload_alignment(batch: int, dp: int) -> None: - """rewards and kl_loss chunk in lockstep — the forward_backward_loss wire. - - The old scalar-shared KL made this payload rewards=[B] + kl=[1]; chunking - by the inferred batch could not keep the two aligned off the B == dp - topology. Per-sample KL makes both first-class batch columns. - """ - kwargs = { - "rewards": torch.randn(batch), - "kl_loss": torch.arange(batch, dtype=torch.float32), - } - assert infer_batch_size((), kwargs) == batch - shards = pytree_chunk(kwargs, dp, batch) - per = batch // dp - for rank, shard in enumerate(shards): - assert shard["rewards"].shape == (per,) - assert shard["kl_loss"].shape == (per,) - assert torch.equal(shard["kl_loss"], kwargs["kl_loss"][rank * per : (rank + 1) * per]) - - -def check_per_shard_backward_grad_shape(batch: int, dp: int) -> None: - """The KL grad each rank produces matches its saved generate output rows. - - Mirrors ReflActorRole.forward_backward_loss on one shard: the KL input is - a grad leaf of shape [B/dp]; after backward its .grad must be the same - shape, because GradContext routes it as out_grads onto the SAME rank's - saved kl tensor from generate_samples. With the old scalar KL this pairing - was [broadcast scalar] vs [rank-local scalar] and only lined up by luck. - """ - per = batch // dp - for _rank in range(dp): - rewards = torch.randn(per, requires_grad=True) - kl = torch.rand(per, requires_grad=True) - reward_loss = (-(rewards.to(torch.bfloat16) - 0.5) / 0.25 * 1.0).mean() - loss = reward_loss + 1.0 * kl.float().mean() - loss.backward() - assert kl.grad is not None and kl.grad.shape == kl.shape - assert rewards.grad is not None and rewards.grad.shape == rewards.shape - - -def check_unequal_actor_reward_dp_stays_aligned() -> None: - """B-length columns survive reward-dp merge → actor-dp re-chunk. - - The recipe colocates actor and reward (equal dp by construction), but the - wire contract must not depend on that: scoring merged at reward dp=2 and - re-scattered at actor dp=4 must hand every actor rank the reward/KL rows - of its own samples. - """ - batch, rdp, adp = 8, 2, 4 - rewards = torch.arange(batch, dtype=torch.float32) - reward_shards = pytree_chunk({"r": rewards}, rdp, batch) - merged = pytree_cat(reward_shards)["r"] - assert torch.equal(merged, rewards) - - kwargs = {"rewards": merged, "kl_loss": torch.arange(batch, dtype=torch.float32) * 10.0} - actor_shards = pytree_chunk(kwargs, adp, batch) - per = batch // adp - for rank, shard in enumerate(actor_shards): - assert torch.equal(shard["rewards"], rewards[rank * per : (rank + 1) * per]) - assert torch.equal(shard["kl_loss"], kwargs["kl_loss"][rank * per : (rank + 1) * per]) - - -def check_per_sample_kl_equals_legacy_scalar_mean() -> None: - """Per-sample reduction then batch-mean == the legacy global scalar mean.""" - torch.manual_seed(0) - kl_pred = torch.randn(3, 4, 2, 5, 5) - ref = torch.randn(3, 4, 2, 5, 5) - sigma = torch.tensor(0.7) - per_sample = ((kl_pred - ref) ** 2 / (2.0 * sigma**2)).flatten(1).mean(dim=1) - legacy = ((kl_pred - ref) ** 2 / (2.0 * sigma**2)).mean() - assert per_sample.shape == (3,) - assert torch.allclose(per_sample.mean(), legacy, atol=1e-6) - - -def check_loss_metrics_concat_keeps_every_shard() -> None: - """REFLLossMetrics concat lists must surface every shard's scalars.""" - shards = [ - REFLLossMetrics(loss=[0.1], reward_loss=[0.2], kl_loss=[0.3], reward_mean=[0.4]), - REFLLossMetrics(loss=[1.1], reward_loss=[1.2], kl_loss=[1.3], reward_mean=[1.4]), - ] - merged = pytree_cat(shards) - assert merged.loss == [0.1, 1.1] - assert merged.kl_loss == [0.3, 1.3] - - -def main() -> int: - for batch, dp in TOPOLOGIES: - check_refl_generated_chunk_cat_roundtrip(batch, dp) - check_forward_backward_payload_alignment(batch, dp) - check_per_shard_backward_grad_shape(batch, dp) - print(f"[ok] topology B={batch} dp={dp}: chunk/cat, payload lockstep, backward shapes") - check_unequal_actor_reward_dp_stays_aligned() - print("[ok] unequal actor/reward dp (rdp=2 → adp=4) stays row-aligned") - check_per_sample_kl_equals_legacy_scalar_mean() - print("[ok] per-sample KL batch-mean equals legacy scalar mean") - check_loss_metrics_concat_keeps_every_shard() - print("[ok] REFLLossMetrics keeps every shard's scalars") - print("verify-refl-kl-batching: ALL CHECKS PASSED") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main())