【Task.31】feat(sft): add reward modeling and preference evaluation - #258
Open
A-Words wants to merge 31 commits into
Open
【Task.31】feat(sft): add reward modeling and preference evaluation#258A-Words wants to merge 31 commits into
A-Words wants to merge 31 commits into
Conversation
# 🐛 Bug Fix ## Reject --ref-load under preference objectives - Standard DPO snapshots the frozen reference from the initialized policy and rebuilds it from --hf-checkpoint on resume; --ref-load was silently ignored on that path while also rerouting the bridge mode policy-init fallback, so validate_preference_args now fails fast when it is set (RFC redai-studio#208: --ref-load is not the v1 DPO reference source) - Drop the misleading --ref-load from the DPO recipe; behavior is unchanged because the bridge fallback already resolves to the same HF checkpoint --- # ♻️ Refactor ## Deduplicate the preference-mode predicate - Route data.py get_data_iterator and loss.py loss_function through is_preference_mode() instead of inline loss_type/sft_objective checks, keeping objective dispatch on a single source of truth --- # ✅ Tests ## Cover --ref-load rejection - Parametrize standard and reference-free DPO rejection cases in test_preference_runtime.py
# ♻️ Refactor ## Remove unreachable DPO loss fallback - Drop the dead `0 * logits.sum()` branch in dpo_loss_function: build_preference_pair_indices already raises on an empty micro-batch ## Attach stable reason codes at preference raise sites - Introduce _PreferenceRowError carrying an explicit reason_code from every normalization/split/truncation rejection; message matching in _classify_preference_error remains only as a fallback for errors raised outside the module, so reworded messages can no longer silently degrade rejections to "schema" --- # ⚡ Performance ## Merge DP pair-count collectives - Reduce _get_preference_data_iterator from three all_reduces to two by riding MIN/MAX on a single MAX all_reduce over [count, -count] ## Skip redundant same-tag model restores - _switch_model now returns early when the target tag is already active, eliminating the duplicate full-weight CPU->GPU restore after the ref-forward finally block; paths that deliberately dirty weights already clear the tag first (covered by the injected-loader-failure reference integrity test) --- # 📝 Documentation ## Note probe bitwise-determinism prerequisite - en/zh DPO guides now state that the reference probe digest assumes an identical GPU/driver/image/kernel stack on resume, and that a probe mismatch signals environment drift rather than data corruption ## Justify the checkpoint-save barrier - Comment that the post-sidecar barrier is functional (peers must not pass before rank 0 persists the identity file), and use an explicit gloo group for the rank-0 check per distributed code rules
# 🐛 Bug Fix ## Preserve standard DPO likelihood semantics - Reject non-unit or non-finite rollout temperatures before actor construction - Replace CUDA-to-Python condition checks with device-side asynchronous assertions - Consolidate finite-value validation on the final DPO logits ## Validate preference data preparation inputs - Validate message objects, roles, and string content on both preference branches - Classify non-object source rows as schema rejections --- # ✅ Tests ## Cover validation regressions - Test finite and non-unit temperature failures - Verify CUDA conditions avoid Python boolean conversion - Cover malformed chosen and rejected message schemas
# 🐛 Bug Fix ## Pin standard DPO frozen reference provenance - Resolve the reference from the declared repository and revision in the configured HF checkpoint directory - Require Hugging Face local metadata to verify the pinned snapshot without downloading during actor startup - Rebuild the frozen reference from that verified directory for both fresh starts and resumes - Fail clearly when the configured checkpoint is unavailable, unverified, or resolves elsewhere --- # 📝 Documentation ## Document pinned local model preparation - Show the fixed-revision hf download command used by the public DPO recipe --- # ✅ Tests ## Cover local reference resolution - Verify repository, revision, local directory, and local-only cache resolution - Verify missing metadata and mismatched resolved directories fail before model loading
# 🐛 Bug Fix ## Require pinned local snapshot metadata - Require an immutable 40-character commit SHA for standard DPO references - Verify Hugging Face local-dir tree metadata before accepting the configured checkpoint - Reject copied or incomplete local snapshots that lack pinned provenance metadata --- # ✅ Tests ## Cover metadata validation - Exercise the pinned tree lookup and missing-metadata rejection paths - Document the full commit-SHA requirement for the DPO recipe
# 🐛 Bug Fix ## Validate pinned Hugging Face local snapshots - Parse the per-file metadata format used by pinned huggingface_hub 1.7.2 - Require every reference file to match the configured commit revision - Verify regular Git files with Git blob SHA-1 and LFS files with SHA-256 - Stream large-file hashing to keep reference validation memory bounded --- # ✅ Tests ## Cover reference provenance and integrity - Exercise Git and LFS ETag validation - Reject missing or mismatched local metadata - Reject replaced weights even when their original mtime is restored
# 🐛 Bug Fix ## Enforce local reference completeness - Require a supported single-file weight or standard Transformers weight index - Validate non-empty index weight maps, safe shard paths, and format-specific suffixes - Reject indexes that reference missing model shards --- # ✅ Tests ## Cover missing reference weights - Accept a complete multi-shard safetensors index - Reject snapshots without a supported weight entry point - Reject a deleted shard even when its local metadata is also removed
# 🐛 Bug Fix ## Move preference scheduling to the CPU control plane - Replace CUDA count reductions and scalar readbacks with one DP-Gloo control gather - Preserve pair-row, global-denominator, and micro-batch-count agreement across ranks - Require Gloo process groups for preference objectives and enable them in the DPO recipe --- # ⚡ Performance ## Remove hot-path GPU synchronization - Eliminate all item calls and control tensors from the preference iterator - Avoid GPU-to-CPU synchronization for train, reference, and evaluation iterators --- # ✅ Tests ## Guard preference iterator contracts - Cover unequal DP pair-row rejection through the Gloo control path - Assert the iterator contains no item or all-reduce scalar readback path
# ♻️ Refactor ## Keep one reference identity path - Remove the unused helper that reads the already-backed-up ref weights - Preserve rebuild validation and backup ordering without duplicate digest work --- # ✅ Tests - Pass the full pre-commit suite - Compile the Megatron actor module
# 🐛 Bug Fix ## Enforce reward-model runtime contracts - Avoid CUDA synchronization in the margin finite check - Validate terminal masks, tokens, packed segments, and pair identity - Restore step-zero and unconditional final preference evaluation ## Harden Megatron checkpoint resume - Persist and validate reward-model objective, role, and scalar-head metadata - Reject incompatible actor, SFT, critic, and partial-state resumes - Preserve restored optimizer, scheduler, RNG, and next-step behavior --- # ⭐ Feature ## Produce reproducible preference-evaluation evidence - Freeze the 512-pair probe and canonical batch-plan hashes - Write per-pair JSONL and PCG64 paired-bootstrap summaries - Align reward-model metrics and recipe batch size with RFC values --- # ✅ Tests ## Cover reward-model loss, checkpoint, and evaluation behavior - Add focused checkpoint and loss contract suites - Test bootstrap digests, artifacts, pair reordering, and eval scheduling --- # 📝 Documentation ## Document reward-model recipes and acceptance artifacts - Describe checkpoint compatibility and required evidence outputs
# 🐛 Bug Fix ## Break the step-zero evaluation wait cycle - Treat the first preference eval chunk as the colocate actor's initial input - Preserve the baseline barrier before exposing the step-zero train partition - Keep causal SFT, RL, and resumed steps gated on their normal train partitions --- # ✅ Tests ## Cover colocate partition readiness - Verify preference step zero accepts only the first baseline eval chunk - Verify causal SFT and resumed preference steps retain existing readiness behavior
# 🐛 Bug Fix ## Keep preference outputs bound to original pairs - Restore per-sample forward outputs whenever the iterator supplies a packed index schedule - Validate packed schedules before reordering scores and log probabilities ## Preserve the complete fixed probe - Consume all 512 preference pairs with capacity-bounded partial chunks - Use actual per-rank and global batch sizes for final chunks - Reject chunk plans that cannot divide cleanly across data-parallel ranks ## Correct causal SFT prediction cadence - Evaluate prediction intervals against completed optimizer steps - Prevent baseline and pre-boundary prediction triggers --- # ✅ Tests ## Cover ordering, chunking, and scheduling regressions - Test packed output restoration independently of the dynamic-batch flag - Cover probe chunk plans across small, uneven, exact, and oversized batch sizes - Verify completed-step prediction boundaries
# 🐛 Bug Fix ## Use completed steps for every SFT post-train evaluation - Map zero-based SFT rollout IDs to completed optimizer-step values - Apply the same mapping to sync, hybrid, and fully-async training paths - Preserve zero-based rollout evaluation steps for RL and step zero for the preference baseline --- # ✅ Tests ## Cover SFT and RL evaluation-step mapping - Verify initial and interval-boundary rollout IDs for both training modes
# 🐛 Bug Fix ## Mirror Megatron checkpoint tracker semantics - Resolve release metadata to the official release checkpoint directory - Preserve iteration and checkpoint-step handling for non-release trackers - Honor an explicitly requested checkpoint step of zero - Reject invalid and negative checkpoint metadata with explicit errors ## Preserve legacy checkpoint compatibility - Detect distributed checkpoints before loading common state metadata - Delegate legacy SFT, DPO, and PPO checkpoints to the upstream Megatron loader - Reject legacy and model-only release checkpoints for reward-model resumes --- # ✅ Tests ## Cover checkpoint path and format contracts - Test release, zero, iteration, explicit step zero, checkpoint-step, invalid, and negative tracker values - Verify non-RM legacy checkpoints bypass distributed metadata and reach the upstream loader - Verify RM legacy and release checkpoints fail before checkpoint state is loaded
# 🐛 Bug Fix ## Read pair identity after TransferQueue expansion - Prefer pair-level IDs from expanded preference data during evaluation - Keep raw pair-row compatibility for direct callers and test fixtures - Reject missing pair-level identity instead of using duplicated branch IDs --- # ✅ Tests ## Cover raw and expanded preference schemas - Verify pair IDs are extracted once per pair from both data shapes - Verify expanded data removes raw IDs and preserves branch-level identity separately - Reject data that exposes only branch-level pair IDs
# 🐛 Bug Fix ## Reject unsupported reward-model exports - Fail before Serve startup when reward modeling is combined with --save-hf - Direct users to native Megatron checkpoints that preserve the scalar reward head - Keep existing Hugging Face export behavior for DPO and causal SFT --- # ✅ Tests ## Cover export configuration boundaries - Verify reward modeling rejects configured HF export paths - Verify reward modeling without HF export remains valid - Verify DPO and causal SFT continue to accept HF export paths
# 🐛 Bug Fix ## Keep the PR2 recipe compatible with its PR1 base - Enable Gloo process groups required by the preference iterator control plane - Prevent reward-model training from failing preference argument validation after the PR1 rebase
# 🐛 Bug Fix ## Keep preference eval chunk plans consistent - Validate the frozen 512-pair probe before optional artifact handling - Reject nonconforming eval sets before any TransferQueue partition is pushed - Verify preference packing preserves every expected probe pair --- # ✅ Tests ## Cover evaluation without artifact output - Reject non-512 probes when --save is unset or empty - Allow the frozen 512-pair probe when artifact output is disabled
A-Words
changed the title
11 tasks
A-Words
marked this pull request as ready for review
August 11, 2026 05:56
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
本堆叠 PR 增加 Reward Modeling，以及 Task 31 最终可复现的偏好评测链路。
Why
Task 31 同时要求离线 DPO 和基于公开 chosen/rejected 数据训练 Reward Model。PR1 建立共享偏好数据与 DPO contract；本 PR 在不增加第二套 trainer 的前提下，补充 RM 专用模型/checkpoint lifecycle 和最终 held-out 验收证据。
实现遵循 RFC #208 和Task 31 要求。
Relates to #208.
Depends on #257.
本 PR 堆叠在 #257（PR1）之上。请优先审阅 PR2-only range：
0e8c334..b57a181。PR1 合并后，本 PR 将重放到最新main，使最终 diff 仅保留 Reward Modeling 和共享偏好评测部分。How
Actor 在
--sft-objective reward_model下安装专用 scalar output layer。通过 packed THD metadata 选择每个 branch 最后一个有效 completion 位置，并产生一个 scalar score。Pairwise logistic objective 按全局 pair 数归一化，评测过程保留显式 pair ID 和 chosen/rejected identity。Checkpoint metadata 区分 causal-LM、DPO、Reward Model 和 critic/value-head 格式。加载过程双向拒绝不兼容 objective，支持 Megatron distributed checkpoint layout，在 resume 时恢复 RM head，并拒绝当前不受支持的 HF export。
共享 evaluator 在 step 0 和 final evaluation 固定使用 manifest 中相同的 512 个 pair。写出逐 pair JSONL 和 10,000-replicate FP64 PCG64 paired-bootstrap summary 前，会校验预处理 identity、pair 顺序、截断结果和 DP/micro-batch plan。
Testing
pre-commit run --all-filespassespytest tests/)在
b57a1818f2f80d28d0ad650b314eb4d1c82df3aa的标准容器中执行：Reward Model 在固定 512-pair held-out 集合上的结果：
accuracy=0.708984、lower-95=0.675781。SeqlenBalancedSampler with dp_size=2，两张 A100 均有实际活动。Type of Change
Screenshots / Logs
本 PR 附带 Reward Model DP=1/DP=2 训练曲线和脱敏后的可复现性证据包。证据包包含：
公开附件不包含 checkpoint、optimizer state、模型权重、完整数据集、凭据或机器专属标识。DPO 的完整训练证据由 PR1 提供；本 PR 的最终报告保留 DPO/RM 汇总结果用于 Task 31 验收。
测量环境使用
ghcr.io/redai-infra/relaxrl:dev-20260723-8cc1e8fd；证据采集时从 GHCR 解析的 manifest digest 为sha256:8dc39af377a570e6cd7ec88c8b7fcd44c1eb820111e9d2069f1c7c3024b2ea23。由于顶层flash_attn_interfaceimport 不可用，Transformer Engine 明确回退到 native attention。这些结果用于证明 correctness 和同环境 DP 行为，不用于声明 FlashAttention 性能。task31-pr2-rm-evidence-public.tar.gz