【Task.31】feat(sft): add pair-aware DPO training - #257
Open
A-Words wants to merge 17 commits into
Open
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
11 tasks
A-Words
changed the title
A-Words
marked this pull request as ready for review
August 11, 2026 05:55
Member
|
Nit: Please avoid committing data JSON files directly to this PR. Consider adding a download link instead. |
Contributor
Author
|
Thanks for the suggestion. I removed the checked-in UltraFeedback manifest from this PR, updated the DPO documentation to use the reproducibility evidence bundle, and updated the evidence bundle linked in the PR description to include the manifest. |
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
2 participants
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 在现有离线 SFT 链路中加入 pair-aware Direct Preference Optimization（DPO）能力。
SeqlenBalancedSampler、DP 分配、动态 batching 和 micro-batch packing 全链路保持 preference pair 原子性。本 PR 不增加 Reward Model scalar head。Reward Modeling 和最终的共享偏好评测器位于后续 PR #258。
Why
Relax 当前主要覆盖在线 RL 工作流。Task 31 要求在不复制训练框架的前提下，复用 SFT trainer 建立标准的离线偏好优化链路。
实现遵循 RFC #208 和Task 31 要求。
Relates to #208.
Follow-up: #258.
How
每个 TransferQueue row 表示一个完整 preference pair，并通过
custom_meta携带 chosen/rejected 的合并 sequence cost。固定版本的SeqlenBalancedSampler因而可以在不拆分 pair 的前提下形成等大小 DP 分组。动态 packer 先分配完整 pair，仅在构造模型输入时展平 chosen/rejected branch。DPO loss 消费显式的 pair/branch identity，使用 completion-only mask，按全局 pair 数归一化，同时支持标准 reference 和 reference-free 公式。空 mask、非有限值、prompt prefix 不一致、截断后两侧相同、超大 pair，以及不一致的 global pair count 均 fail closed。
标准 DPO fresh 训练将初始 actor 快照为 frozen reference；resume 时则从固定 HF repository/revision 执行 model-only reference 重建。Canonical parameter、源文件 manifest、optimizer master/state 和固定 log-prob probe digest 可防止 reference 切换静默修改 actor 或 optimizer。验证后的 identity 会写入 checkpoint 旁的
relax_dpo_reference.json。Testing
pre-commit run --all-filespassespytest tests/)在最终堆叠集成 head
b57a1818f2f80d28d0ad650b314eb4d1c82df3aa的标准容器中执行：固定 TransferQueue 兼容性测试覆盖 pair-row 原子性、等大小 DP 分组、row-aligned
total_lengthscost 消费，以及无丢失、无重复的完整分配。数据准备脚本的测试覆盖 manifest：4,096 个 train pair、512 个互不重叠的 eval pair、确定性选择、拒绝原因计数和派生文件 hash。GPU 验证固定使用相同的模型/data revision 和 global batch size 32 pairs：
SeqlenBalancedSampler with dp_size=2，两张 A100 均有实际活动。Type of Change
Screenshots / Logs
本 PR 附带 DPO DP=1/DP=2 训练曲线和脱敏后的可复现性证据包。证据包包含：
公开附件不包含 checkpoint、optimizer state、模型权重、完整数据集、凭据或机器专属标识。
测量环境使用
ghcr.io/redai-infra/relaxrl:dev-20260723-8cc1e8fd；证据采集时从 GHCR 解析的 manifest digest 为sha256:8dc39af377a570e6cd7ec88c8b7fcd44c1eb820111e9d2069f1c7c3024b2ea23。Transformer Engine 无法导入顶层flash_attn_interface，并明确回退到 native attention。因此 DP 对比描述的是相同 native-attention 环境，不是 FlashAttention 性能结果。task31-pr1-dpo-evidence-public-v2.tar.gz