From 68a7b9515339e9985377f148e8b8b79204e4c005 Mon Sep 17 00:00:00 2001 From: tongruiliu <952814125@qq.com> Date: Tue, 16 Jun 2026 18:58:30 +0800 Subject: [PATCH] feat: add ADAPT online reweighting weighter --- examples/train_lora/weighters/adapt.yaml | 61 +++++++++++++ skills/how_to_use.md | 8 +- src/dataflex/configs/components.yaml | 8 ++ src/dataflex/train/weighter/__init__.py | 3 +- src/dataflex/train/weighter/adapt_weighter.py | 90 +++++++++++++++++++ 5 files changed, 167 insertions(+), 3 deletions(-) create mode 100644 examples/train_lora/weighters/adapt.yaml create mode 100644 src/dataflex/train/weighter/adapt_weighter.py diff --git a/examples/train_lora/weighters/adapt.yaml b/examples/train_lora/weighters/adapt.yaml new file mode 100644 index 0000000..ce60d30 --- /dev/null +++ b/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 diff --git a/skills/how_to_use.md b/skills/how_to_use.md index 2101ae5..efefc01 100644 --- a/skills/how_to_use.md +++ b/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 diff --git a/src/dataflex/configs/components.yaml b/src/dataflex/configs/components.yaml index e5b65a9..b20c235 100644 --- a/src/dataflex/configs/components.yaml +++ b/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: diff --git a/src/dataflex/train/weighter/__init__.py b/src/dataflex/train/weighter/__init__.py index a114322..a31c6dc 100644 --- a/src/dataflex/train/weighter/__init__.py +++ b/src/dataflex/train/weighter/__init__.py @@ -1,3 +1,4 @@ from .base_weighter import Weighter from .loss_weighter import * -from .custom_weighter import CustomWeighter \ No newline at end of file +from .custom_weighter import CustomWeighter +from .adapt_weighter import AdaptWeighter \ No newline at end of file diff --git a/src/dataflex/train/weighter/adapt_weighter.py b/src/dataflex/train/weighter/adapt_weighter.py new file mode 100644 index 0000000..8af385f --- /dev/null +++ b/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)