feat: add ADAPT online reweighting weighter

examples/train_lora/weighters/adapt.yaml

@@ -0,0 +1,61 @@
+### model
+model_name_or_path: meta-llama/Llama-3.1-8B
+trust_remote_code: true
+
+### method
+stage: sft
+do_train: true
+finetuning_type: lora
+lora_target: all
+lora_rank: 16
+lora_alpha: 8
+# deepspeed: examples/deepspeed/ds_z3_config.json  # choices: [ds_z0_config.json, ds_z2_config.json, ds_z3_config.json]
+
+### dataset
+dataset: alpaca_en_demo
+eval_dataset: alpaca_en_demo  # 用作 anchor/验证集，请替换为你自己的目标分布数据
+template: llama3
+cutoff_len: 4096
+# max_samples: 100000000
+overwrite_cache: true
+preprocessing_num_workers: 16
+dataloader_num_workers: 0
+# disable_shuffling: true
+seed: 42
+
+### output
+output_dir: ../dataflex_saves/Llama-3.1-8B/adapt
+logging_steps: 10
+save_steps: 100
+plot_loss: true
+save_only_model: false
+overwrite_output_dir: true
+
+### swanlab
+report_to: none  # choices: [none, wandb, tensorboard, swanlab, mlflow]
+# use_swanlab: true
+# swanlab_project: dynamic_sft
+# swanlab_run_name: your_run_name
+# swanlab_workspace: your_workspace
+# swanlab_api_key: AnLWTMijcbd4cyEfundi3
+
+### train
+per_device_train_batch_size: 2
+gradient_accumulation_steps: 16
+learning_rate: 1.0e-4
+num_train_epochs: 1.0
+lr_scheduler_type: cosine
+warmup_ratio: 0.1
+bf16: true
+ddp_timeout: 180000000
+
+### Dataflex args
+train_type: dynamic_weight   # 选择训练器类型。可选值包括：
+                          # "dynamic_select" - 动态选择训练器
+                          # "dynamic_mix" - 动态混合训练器
+                          # "dynamic_weight" - 动态加权训练器
+                          # "static" - 默认静态训练器
+components_cfg_file: src/dataflex/configs/components.yaml
+component_name: adapt  # 选择组件名称，对应 components_cfg_file 中定义的组件
+warmup_step: 100
+train_step: 500 # Total steps; overrides num_train_epochs

skills/how_to_use.md

@@ -168,7 +168,7 @@ Dynamically adjusts per-sample loss weights during backpropagation based on samp
 ```yaml
 train_type: dynamic_weight
 components_cfg_file: src/dataflex/configs/components.yaml
-component_name: loss       # choices: loss, custom
+component_name: loss       # choices: loss, adapt, custom
 warmup_step: 100
 train_step: 500            # fixed-step example; set to 0 for num_train_epochs-based multi-epoch runs
 ```
@@ -179,9 +179,12 @@ train_step: 500            # fixed-step example; set to 0 for num_train_epochs-b
 3. `warmup_step` is measured in global optimization steps and does not reset per epoch.
 4. If `train_step > 0`, total steps = `train_step`; otherwise total steps follow the standard `num_train_epochs` calculation.
 
+> `adapt` scores each sample by the cosine similarity between its representation and an anchor set, so it requires an `eval_dataset` field (the anchor/validation set you provide).
+
 **Example:**
 ```bash
 dataflex-cli train examples/train_lora/weighters/loss.yaml
+dataflex-cli train examples/train_lora/weighters/adapt.yaml
 ```
 
 ## Component Configuration (`components.yaml`)
@@ -240,6 +243,7 @@ You select which algorithm to use via `component_name` in your training YAML.
 | Algorithm | `component_name` | Category | Description |
 |-----------|-----------------|----------|-------------|
 | Loss Reweighting | `loss` | Loss-based | Strategies: `linupper`, `uniform`, `quadratic`, `extremes` |
+| ADAPT | `adapt` | Similarity-based | Online per-sample weights from embedding similarity to an anchor set (requires `eval_dataset`) |
 | Custom | `custom` | Custom | Template for user-defined weighting logic |
 
 ## Offline Preprocessing
@@ -295,7 +299,7 @@ examples/
 ├── train_lora/
 │   ├── selectors/     # LESS, NICE, Loss, Delta Loss, TSDS, NEAR, Random, Custom
 │   ├── mixers/        # DoReMi Step 2 (LoRA), Random
-│   └── weighters/     # Loss, Custom
+│   └── weighters/     # Loss, ADAPT, Custom
 ├── train_full/
 │   └── mixers/        # DoReMi Steps 1-3 (full), ODM (full)
 ├── test/              # minimal smoke-test configs

src/dataflex/configs/components.yaml

@@ -145,6 +145,14 @@ weighters:
       strategy: linupper # choices: [linupper, uniform, quadratic, extremes]
       delta: 1.0
 
+  adapt:
+    name: adapt
+    params:
+      tau: 1.0              # 温度，越小权重区分越锐利
+      refresh_interval: 50  # 每多少步用当前模型刷新一次 anchor 向量
+      anchor_batch_size: 8  # 计算句向量时的前向 batch 大小
+      clip: null            # 可选权重上限，防梯度爆炸
+
   custom:
     name: custom
     params:

src/dataflex/train/weighter/__init__.py

@@ -1,3 +1,4 @@
 from .base_weighter import Weighter
 from .loss_weighter import *
-from .custom_weighter import CustomWeighter
+from .custom_weighter import CustomWeighter
+from .adapt_weighter import AdaptWeighter

src/dataflex/train/weighter/adapt_weighter.py

@@ -0,0 +1,90 @@
+from dataflex.core.registry import register_weighter
+from dataflex.utils.logging import logger
+from typing import Any, Optional, Union
+from torch import nn
+from torch.utils.data import DataLoader
+import torch
+from .base_weighter import Weighter
+
+
+@register_weighter('adapt')
+class AdaptWeighter(Weighter):
+    """
+    参考论文
+    Rethinking Data Curation in LLM Training: Online Reweighting Offers Better Generalization (ICLR2026)
+    以「训练样本与 anchor(验证)集的表征相似度」作为质量信号，经温度化 sigmoid 得到每样本的绝对权重。
+    """
+    def __init__(
+        self,
+        tau: float = 1.0,                # 温度，越小权重区分越锐利
+        refresh_interval: int = 50,      # 每多少步用当前模型刷新一次 anchor 向量
+        anchor_batch_size: int = 8,      # 计算句向量时的前向 batch 大小
+        clip: Optional[float] = None,    # 可选权重上限，防梯度爆炸
+        eps: float = 1e-8,
+        eval_dataset=None,
+        data_collator=None,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        if eval_dataset is None:
+            raise ValueError("AdaptWeighter 需要 anchor 集，请在配置中提供 eval_dataset。")
+        self.tau = float(tau)
+        self.refresh_interval = int(refresh_interval)
+        self.anchor_batch_size = int(anchor_batch_size)
+        self.clip = clip
+        self.eps = float(eps)
+        self.eval_dataset = eval_dataset
+        self.data_collator = data_collator
+        self.anchor_emb = None  # (M, H) 归一化后的 anchor 向量，每 refresh_interval 步刷新
+
+    @torch.no_grad()
+    def _embed(self, model, input_ids, attention_mask):
+        """位置加权均值池化 + L2 归一化，得到句向量 (B, H)。越靠后的 token 权重越大。"""
+        hidden = model(input_ids=input_ids, attention_mask=attention_mask,
+                       output_hidden_states=True).hidden_states[-1]  # (B, L, H)
+        mask = attention_mask.to(hidden.dtype)
+        pos = torch.arange(1, hidden.size(1) + 1, device=hidden.device, dtype=hidden.dtype) * mask
+        w = pos / pos.sum(dim=1, keepdim=True).clamp_min(self.eps)   # (B, L)
+        phi = (w.unsqueeze(-1) * hidden).sum(dim=1)                  # (B, H)
+        return phi / phi.norm(dim=-1, keepdim=True).clamp_min(self.eps)
+
+    @torch.no_grad()
+    def _refresh_anchors(self, model):
+        """用当前模型重算 anchor 集的句向量。"""
+        loader = DataLoader(self.eval_dataset, batch_size=self.anchor_batch_size,
+                            collate_fn=self.data_collator)
+        device = next(model.parameters()).device
+        embs = [self._embed(model, b["input_ids"].to(device), b["attention_mask"].to(device))
+                for b in loader]
+        self.anchor_emb = torch.cat(embs, dim=0)  # (M, H)
+
+    def get_weighted_loss(
+        self,
+        losses: torch.Tensor,
+        *,
+        ctx: Any = None,
+        model: nn.Module | None = None,
+        inputs: dict[str, Union[torch.Tensor, Any]] | None = None,
+    ) -> torch.Tensor:
+        # 兼容：标量或非张量 → 不加权
+        if (not torch.is_tensor(losses)) or losses.dim() == 0:
+            return losses
+        losses = losses.view(-1)
+
+        # 每 refresh_interval 步用当前模型刷新 anchor 向量（在线更新）
+        step = ctx.state.global_step if ctx is not None else 0
+        if self.anchor_emb is None or step % self.refresh_interval == 0:
+            self._refresh_anchors(model)
+
+        # 样本句向量 → 与 anchor 的平均余弦相似度 → 温度化 sigmoid 得到绝对权重
+        phi = self._embed(model, inputs["input_ids"], inputs["attention_mask"])  # (B, H)
+        score = (phi @ self.anchor_emb.t().to(phi.dtype)).mean(dim=1)            # (B,)
+        weights = torch.sigmoid(score / max(self.tau, self.eps))                # (B,)
+        if self.clip is not None:
+            weights = weights.clamp(max=self.clip)
+
+        if ctx is not None and ctx.args.local_rank in [-1, 0]:
+            logger.info(f"[Dataflex] ADAPT weights (first sample): {float(weights[0])}")
+
+        # 权重作为常量缩放每样本 loss（per-sample learning rate）
+        return torch.sum(weights.detach().to(losses.dtype) * losses)