Add joint_update_aware weighter

[tongruiliu]

概述

新增一个动态数据加权器 Joint-Update-Aware Reweighting（

本方法是 batch 级、joint-update-aware 的：一个样本的边际价值不再独立计算，而是取决于当前加权 batch 已经覆盖的方向。它把一个 pointwise效用信号 s 修正为单纯形上的 batch 相关权重，求解带熵正则的目标：

max_w   s^T w  -  (beta / 2) w^T S w  +  tau * H(w)

• s_i = <u, z_i>：样本 embedding z_i 对目标方向 u 的对齐度（pointwise 效用）
• S = Z Z^T：L2 归一化 embedding 的余弦 Gram 矩阵，作为 梯度交互矩阵的廉价代理
  （避免存每样本的全参数梯度）
• H(w)：熵正则，保证内部唯一最优

唯一内部最优是定点 w_i = softmax([s - beta·Sw] / tau)_i，用阻尼迭代求解：
w <- (1 - rho) w + rho · softmax([s - beta·Sw] / tau)。

改动文件

文件	改动
src/dataflex/train/weighter/joint_update_aware_weighter.py	新增，JointUpdateAwareWeighter
src/dataflex/train/weighter/__init__.py	导出 JointUpdateAwareWeighter
src/dataflex/configs/components.yaml	注册 joint_update_aware 组件及默认超参
examples/train_lora/weighters/joint_update_aware.yaml	新增训练示例配置

实现要点

• 接入方式：继承 Weighter，实现 get_weighted_loss，走已有的 WeightTrainer 动态加权流程
  （train_type: dynamic_weight），无需改动 trainer。
• Embedding 提取：前向一次取指定层 hidden states，按 pooling 池化为每样本句向量后 L2 归一化。
• 目标方向 u：取 eval_dataset（anchor / 目标分布集）的平均 embedding，每 target_update_step
  步用当前模型重算并缓存；eval 不可用时回退为当前 batch 的平均 embedding。
• 分布式：对变长 batch 做 padding + all_gather，在 全局 batch 上求解权重，再按 rank 切回本卡；
  最后 × world_size 抵消 DDP 的跨卡梯度平均，使加权目标尺度与单机一致。

可配置开关

核心超参：

参数	默认	含义
beta	0.1	交互 / 冗余强度
tau	0.05	熵温度（>0），越小权重越锐利
fixed_point_iters	5	阻尼定点迭代次数
damping	1.0	阻尼系数 rho ∈ (0, 1]，1.0 表示不阻尼
target_update_step	50	每多少步用 eval 集刷新一次目标向量
target_batch_size	1	计算目标向量时的前向 batch 大小
embed_normalize	true	是否对 embedding 做 L2 归一化
embed_layer	-1	取哪一层 hidden state

功能开关：

• pooling：last_token（默认）/ mean_pool（在 response token 上取均值）——两种句向量池化方式。
• objective_mode：消融用开关，控制求解时的 score 项
  • full（默认）：s - beta·Sw
  • align_only：s（只对齐，退回 pointwise）
  • diverse_only：-beta·Sw（只压冗余）
  • uniform：直接 1/B（均匀权重，对照基线）

说明：objective_mode 的后三项和 pooling 的 mean_pool 为消融 / 对照选项，

用法

llamafactory-cli train examples/train_lora/weighters/joint_update_aware.yaml

关键配置：train_type: dynamic_weight、component_name: joint_update_aware，并将 eval_dataset
设为目标分布 / anchor 数据集。

测试

• python -m py_compile 通过，components.yaml 与示例 yaml 均 YAML 解析通过。