Claude/great allen 51nxz7

.gitignore

@@ -180,4 +180,7 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
 
-.vscode/
+.vscode/
+.history/
+
+*.whl

configs/experiment/examples/muse_unlearn.yaml

@@ -27,7 +27,7 @@ trainer:
     bf16_full_eval: true
     logging_steps: 5
     output_dir: ${paths.output_dir}
-    logging_dir: ${trainer.args.output_dir}/logs
+    # logging_dir removed: deprecated in transformers (removed in v5.2).
     report_to: tensorboard
     ddp_find_unused_parameters: None
     gradient_checkpointing: false

configs/experiment/examples/tofu_eval.yaml

@@ -5,7 +5,7 @@ model:
     attn_implementation: flash_attention_2
     torch_dtype: bfloat16
   tokenizer_args:
-    pretrained_model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
+    pretrained_model_name_or_path: unsloth/Llama-3.2-1B-Instruct
   template_args:
     apply_chat_template: true
     system_prompt: You are a helpful assistant.

configs/hydra/default.yaml

@@ -1,9 +1,11 @@
 # https://hydra.cc/docs/configure_hydra/intro/
 
-# enable color logging
+# enable rich logging (requires `pip install rich`)
 defaults:
-  - override hydra_logging: colorlog
-  - override job_logging: colorlog
+  - override hydra_logging: rich   # Hydra's own framework logs
+  - override job_logging: rich     # your application logs (the ones you usually care about)
+  # To revert to the old style: change these to `colorlog`, or per-run e.g.
+  # `hydra/job_logging=colorlog hydra/hydra_logging=colorlog`.
 
 # output directory, generated dynamically on each run
 run:

configs/hydra/hydra_logging/rich.yaml

@@ -0,0 +1,21 @@
+# Hydra's *own* framework logs (the hydra_logging group), routed through Rich.
+# Pairs with configs/hydra/job_logging/rich.yaml (your application logs).
+# Requires: pip install rich
+version: 1
+formatters:
+  rich:
+    (): log_utils.RichNameFormatter
+    name_style: "dim cyan"   # dim, to visually separate Hydra's own logs from job logs
+handlers:
+  rich:
+    class: rich.logging.RichHandler
+    formatter: rich
+    rich_tracebacks: true
+    show_time: true
+    show_level: true
+    show_path: false
+    markup: true
+root:
+  level: INFO
+  handlers: [rich]
+disable_existing_loggers: false

configs/hydra/job_logging/rich.yaml

@@ -0,0 +1,31 @@
+# Hydra job_logging config that routes logs through rich.logging.RichHandler.
+# Select it via the defaults list (override job_logging: rich) in
+# configs/hydra/default.yaml, or per-run on the CLI: `hydra/job_logging=rich`.
+# Requires: pip install rich
+version: 1
+formatters:
+  rich:
+    # Custom formatter (src/log_utils.py): colors the logger name and escapes
+    # the message so brackets in log text aren't parsed as Rich markup.
+    (): log_utils.RichNameFormatter
+    name_style: "bold cyan"   # change to any Rich style, e.g. "magenta", "bold green"
+  plain:
+    # Plain text for the log file (no color markup written to disk).
+    format: "[%(asctime)s][%(name)s][%(levelname)s] - %(message)s"
+handlers:
+  rich:
+    class: rich.logging.RichHandler
+    formatter: rich
+    rich_tracebacks: true
+    show_time: true
+    show_level: true
+    show_path: false
+    markup: true   # required so the name's [bold cyan]..[/] tags are colored
+  file:
+    class: logging.FileHandler
+    formatter: plain
+    filename: ${hydra.runtime.output_dir}/${hydra.job.name}.log
+root:
+  level: INFO
+  handlers: [rich, file]
+disable_existing_loggers: false

configs/model/Llama-3.2-1B-Instruct-LoRA.yaml

@@ -0,0 +1,35 @@
+# LoRA variant of Llama-3.2-1B-Instruct.
+# Identical to the base model config, but adds a `peft_args` block which
+# `src/model/__init__.py:get_model` uses to wrap the base model with a LoRA
+# adapter (via the `peft` library). Remove/omit `peft_args` to train full weights.
+model_args:
+  pretrained_model_name_or_path: "unsloth/Llama-3.2-1B-Instruct"
+  attn_implementation: 'flash_attention_2'
+  torch_dtype: bfloat16
+tokenizer_args:
+  pretrained_model_name_or_path: "unsloth/Llama-3.2-1B-Instruct"
+peft_args:
+  # mirrors peft.LoraConfig
+  r: 8
+  lora_alpha: 32
+  lora_dropout: 0.05
+  bias: none
+  task_type: CAUSAL_LM
+  target_modules:
+    - q_proj
+    - k_proj
+    - v_proj
+    - o_proj
+    - gate_proj
+    - up_proj
+    - down_proj
+  # path: null  # set to an existing adapter dir to resume / evaluate a LoRA checkpoint
+template_args:
+  apply_chat_template: True
+  system_prompt: You are a helpful assistant.
+  system_prompt_with_special_tokens: "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nYou are a helpful assistant.<|eot_id|>"
+  user_start_tag: "<|start_header_id|>user<|end_header_id|>\n\n"
+  user_end_tag: "<|eot_id|>"
+  asst_start_tag: "<|start_header_id|>assistant<|end_header_id|>\n\n"
+  asst_end_tag: "<|eot_id|>"
+  date_string: 10 Apr 2025

configs/model/Llama-3.2-1B-Instruct.yaml

@@ -1,9 +1,9 @@
 model_args:
-  pretrained_model_name_or_path: "meta-llama/Llama-3.2-1B-Instruct"
+  pretrained_model_name_or_path: "unsloth/Llama-3.2-1B-Instruct"
   attn_implementation: 'flash_attention_2'
   torch_dtype: bfloat16
 tokenizer_args:
-  pretrained_model_name_or_path: "meta-llama/Llama-3.2-1B-Instruct"
+  pretrained_model_name_or_path: "unsloth/Llama-3.2-1B-Instruct"
 template_args:
   apply_chat_template: True
   system_prompt: You are a helpful assistant.

configs/trainer/BoundedGradDiff.yaml

@@ -0,0 +1,9 @@
+defaults:
+  - finetune          # 继承 configs/trainer/finetune.yaml 里的 HuggingFace TrainingArguments
+
+handler: BoundedGradDiff   # 必须等于 src/trainer/unlearn/bounded_grad_diff.py 里的类名
+method_args:               # 这些会作为 **kwargs 传给 BoundedGradDiff.__init__
+  gamma: 1.0               # forget 项权重 (继承自 GradDiff)
+  alpha: 1.0               # retain 项权重 (继承自 GradDiff)
+  retain_loss_type: NLL    # NLL 或 KL (继承自 GradDiff)
+  forget_loss_bound: 4.0   # 本算法新增: forget NLL 的上界 tau

configs/trainer/finetune.yaml

@@ -8,7 +8,8 @@ args:
   bf16_full_eval: True
   logging_steps: 5
   output_dir: ${paths.output_dir}
-  logging_dir: ${trainer.args.output_dir}/logs
+  # logging_dir removed: deprecated in transformers (removed in v5.2). TensorBoard
+  # logs now default under ${output_dir}/runs. Set env TENSORBOARD_LOGGING_DIR to override.
   report_to: tensorboard
   ddp_find_unused_parameters: None
   gradient_checkpointing: False

configs/unlearn.yaml

@@ -10,7 +10,7 @@ defaults:
   - _self_
 
 trainer:
-  args: 
+  args:
     remove_unused_columns: False
 
 mode: unlearn

demos/1_finetune.sh

@@ -0,0 +1,46 @@
+#!/bin/bash
+# =============================================================================
+# Demo 1: 直接（全参数）微调一个模型
+# -----------------------------------------------------------------------------
+# 入口: src/train.py  (mode=train)
+# 通过 Hydra 把以下配置组合在一起:
+#   experiment=finetune/tofu/default
+#     -> configs/experiment/finetune/tofu/default.yaml
+#        - model:   Llama-3.2-1B-Instruct
+#        - trainer: finetune  (handler=FinetuneTrainer, configs/trainer/finetune.yaml)
+#        - data:    TOFU_QA_full (locuslab/TOFU 的 "full" split)
+#        - eval:    tofu  (训练过程中按 epoch 评测)
+#
+# 训练完的权重会保存到 paths.output_dir，默认 = saves/<mode>/<task_name>
+# 即: saves/finetune/demo_finetune_full
+# =============================================================================
+set -e
+cd "$(dirname "$0")/.."   # 切到仓库根目录
+
+# 共享集群必看: 只暴露一张【空闲】GPU, 否则 HF Trainer 会在所有可见卡上启用
+# DataParallel, 往被别人占满的卡复制模型而 CUDA OOM。先 `nvidia-smi` 选空闲卡,
+# 或运行时 `CUDA_VISIBLE_DEVICES=3 bash demos/1_finetune.sh` 覆盖。
+export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0}
+
+MODEL=Llama-3.2-1B-Instruct
+
+python src/train.py --config-name=train.yaml \
+  experiment=finetune/tofu/default \
+  model=${MODEL} \
+  task_name=demo_finetune_full \
+  trainer.args.num_train_epochs=5 \
+  trainer.args.per_device_train_batch_size=4 \
+  trainer.args.gradient_accumulation_steps=8 \
+  trainer.args.learning_rate=1e-5
+
+# 命令行上任意 key 都能覆盖 yaml，比如:
+#   trainer.args.num_train_epochs=10
+#   data/datasets@data.train=TOFU_QA_retain \
+#   data.train.TOFU_QA_retain.args.hf_args.name=retain90   # 改成训练 retain 模型
+#
+# 多卡训练把 `python src/train.py` 换成:
+#   accelerate launch --config_file configs/accelerate/default_config.yaml src/train.py ...
+
+
+python src/train.py --config-name=unlearn.yaml experiment=unlearn/tofu/default \
+  forget_split=forget10 retain_split=retain90 trainer=GradAscent task_name=SAMPLE_UNLEARN

demos/2_eval.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+# =============================================================================
+# Demo 2: 评测一个模型 (TOFU benchmark)
+# -----------------------------------------------------------------------------
+# 入口: src/eval.py  (mode=eval)
+#   1) get_model()      根据 model 配置加载模型+tokenizer
+#   2) get_evaluators() 根据 eval=tofu 构造 TOFUEvaluator
+#   3) evaluator.evaluate() 逐个跑 configs/eval/tofu.yaml 里 default 列出的指标
+#
+# 结果文件 (写入 paths.output_dir = saves/eval/<task_name>):
+#   TOFU_EVAL.json     每条样本的细粒度分数 (value_by_index)
+#   TOFU_SUMMARY.json  每个指标的聚合值 (agg_value)
+#
+# retain_logs_path: 指向 "retain 参照模型" 的 EVAL.json。
+#   forget_quality / privleak 等指标需要它来和参照模型做对比。
+#   需要先 `python setup_data.py --eval` 下载官方参照日志，或自己评一个 retain 模型。
+# =============================================================================
+set -e
+cd "$(dirname "$0")/.."
+
+# 共享集群必看: 只暴露一张【空闲】GPU 评测。先 `nvidia-smi` 选空闲卡,
+# 或运行时 `CUDA_VISIBLE_DEVICES=3 bash demos/2_eval.sh` 覆盖。
+export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0}
+
+MODEL=Llama-3.2-1B-Instruct
+
+python src/eval.py --config-name=eval.yaml \
+  experiment=eval/tofu/default \
+  model=${MODEL} \
+  model.model_args.pretrained_model_name_or_path=open-unlearning/tofu_${MODEL}_full \
+  forget_split=forget10 \
+  holdout_split=holdout10 \
+  retain_logs_path=saves/eval/tofu_${MODEL}_retain90/TOFU_EVAL.json \
+  task_name=demo_eval
+
+# 想评测 Demo 1 自己微调出来的模型, 把上面这行换成本地路径:
+#   model.model_args.pretrained_model_name_or_path=saves/finetune/demo_finetune_full
+#
+# 没有 retain 参照日志时, 去掉 retain_logs_path 即可 (forget_quality 会是 None)。

demos/3_lora_finetune.sh

@@ -0,0 +1,44 @@
+#!/bin/bash
+# =============================================================================
+# Demo 3: 用 LoRA 微调模型
+# -----------------------------------------------------------------------------
+# 注意: 原框架默认不支持 LoRA。本 demo 配套做了一处最小、非破坏性的扩展:
+#   - src/model/__init__.py: get_model() 在加载完基座模型后, 若 model 配置里有
+#     `peft_args` 块, 就调用 get_peft_lora_model() 用 peft 包套一层 LoRA adapter。
+#   - configs/model/Llama-3.2-1B-Instruct-LoRA.yaml: 在原模型配置上加了 peft_args。
+#
+# 前置依赖: pip install peft
+#
+# 训练流程其余部分和 Demo 1 完全一样 (同样走 FinetuneTrainer)，只是把
+# model 换成带 peft_args 的 LoRA 配置。HF Trainer 会自动只保存 adapter 权重。
+# 输出: saves/finetune/demo_lora_finetune
+# =============================================================================
+set -e
+cd "$(dirname "$0")/.."
+
+# 共享集群必看: 只暴露一张【空闲】GPU, 否则 HF Trainer 会在所有可见卡上启用
+# DataParallel, 往被别人占满的卡复制模型而 CUDA OOM。先 `nvidia-smi` 选空闲卡,
+# 或运行时 `CUDA_VISIBLE_DEVICES=3 bash demos/3_lora_finetune.sh` 覆盖。
+export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0}
+
+python src/train.py --config-name=train.yaml \
+  experiment=finetune/tofu/default \
+  model=Llama-3.2-1B-Instruct-LoRA \
+  task_name=demo_lora_finetune \
+  trainer.args.num_train_epochs=5 \
+  trainer.args.per_device_train_batch_size=4 \
+  trainer.args.gradient_accumulation_steps=8 \
+  trainer.args.learning_rate=1e-4    # LoRA 通常用比全参微调更大的学习率
+
+# 也可以完全在命令行里临时指定 LoRA 超参 (无需改 yaml), 例如:
+#   model=Llama-3.2-1B-Instruct \
+#   +model.peft_args.r=16 +model.peft_args.lora_alpha=32 \
+#   +model.peft_args.task_type=CAUSAL_LM \
+#   '+model.peft_args.target_modules=[q_proj,v_proj]'
+#
+# ---- 评测训练好的 LoRA adapter ----
+# 把 base 模型路径 + adapter 路径 (peft_args.path) 一起传给 eval:
+#   python src/eval.py experiment=eval/tofu/default \
+#     model=Llama-3.2-1B-Instruct-LoRA \
+#     +model.peft_args.path=saves/finetune/demo_lora_finetune \
+#     task_name=demo_lora_eval

demos/4_unlearn.sh

@@ -0,0 +1,45 @@
+#!/bin/bash
+# =============================================================================
+# Demo 4: 用一个【已有的】遗忘方法做 unlearning (这里用 GradDiff)
+# -----------------------------------------------------------------------------
+# 入口: src/train.py  (mode=unlearn)
+# experiment=unlearn/tofu/default -> configs/experiment/unlearn/tofu/default.yaml
+#   - model:   待遗忘的目标模型 (默认 open-unlearning/tofu_Llama-3.2-1B-Instruct_full)
+#   - data:    unlearn -> 同时加载 forget 和 retain 两个数据集
+#              collator 把每个 batch 组织成 {"forget": {...}, "retain": {...}}
+#   - trainer: 由命令行 trainer=GradDiff 指定 (handler=GradDiff)
+#   - eval:    tofu (遗忘过程中/结束后评测)
+#
+# 遗忘方法的核心只有一个函数: trainer 的 compute_loss(model, inputs)。
+#   - GradAscent: loss = -forget_loss
+#   - GradDiff:   loss = gamma*(-forget_loss) + alpha*retain_loss
+#   - NPO/SimNPO/DPO/RMU/...: 各自不同的 compute_loss
+#
+# 输出: saves/unlearn/demo_unlearn_graddiff
+# =============================================================================
+set -e
+cd "$(dirname "$0")/.."
+
+# 共享集群必看: 只暴露一张【空闲】GPU, 否则 HF Trainer 会在所有可见卡上启用
+# DataParallel, 往被别人占满的卡复制模型而 CUDA OOM。先 `nvidia-smi` 选一张空闲卡,
+# 改下面的默认 0, 或运行时 `CUDA_VISIBLE_DEVICES=3 bash demos/4_unlearn.sh` 覆盖。
+export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0}
+
+MODEL=Llama-3.2-1B-Instruct
+
+python src/train.py --config-name=unlearn.yaml \
+  experiment=unlearn/tofu/default \
+  model=${MODEL} \
+  trainer=GradDiff \
+  trainer.method_args.gamma=1.0 \
+  trainer.method_args.alpha=1.0 \
+  trainer.method_args.retain_loss_type=NLL \
+  forget_split=forget10 \
+  retain_split=retain90 \
+  holdout_split=holdout10 \
+  retain_logs_path=saves/eval/tofu_${MODEL}_retain90/TOFU_EVAL.json \
+  task_name=demo_unlearn_graddiff
+
+# 换方法只需改 trainer= : GradAscent / NPO / SimNPO / DPO / RMU / UNDIAL / WGA / CEU ...
+# 对应方法的额外超参在 trainer.method_args.* 下覆盖, 例如 NPO:
+#   trainer=NPO trainer.method_args.beta=0.1 trainer.method_args.gamma=1.0

demos/5_custom_unlearn.sh

@@ -0,0 +1,41 @@
+#!/bin/bash
+# =============================================================================
+# Demo 5: 运行【自定义】遗忘算法 BoundedGradDiff
+# -----------------------------------------------------------------------------
+# 把自己的算法接入本项目, 只需 3 步 (本 demo 已帮你做好):
+#   1. 写 trainer 类:  src/trainer/unlearn/bounded_grad_diff.py
+#                      - 继承 GradDiff (或更底层的 UnlearnTrainer)
+#                      - 只重写 compute_loss(model, inputs, ...)
+#   2. 注册:           src/trainer/__init__.py 里 _register_trainer(BoundedGradDiff)
+#   3. 写配置:         configs/trainer/BoundedGradDiff.yaml (handler + method_args)
+#
+# 之后用法和内置方法完全一样, 只是 trainer=BoundedGradDiff:
+# 输出: saves/unlearn/demo_custom_unlearn
+# =============================================================================
+set -e
+cd "$(dirname "$0")/.."
+
+# 共享集群必看: 只暴露一张【空闲】GPU, 否则 HF Trainer 会在所有可见卡上启用
+# DataParallel, 往被别人占满的卡复制模型而 CUDA OOM。先 `nvidia-smi` 选一张空闲卡,
+# 改下面的默认 0, 或运行时 `CUDA_VISIBLE_DEVICES=3 bash demos/5_custom_unlearn.sh` 覆盖。
+export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0}
+
+MODEL=Llama-3.2-1B-Instruct
+
+python src/train.py --config-name=unlearn.yaml \
+  experiment=unlearn/tofu/default \
+  model=${MODEL} \
+  trainer=BoundedGradDiff \
+  trainer.method_args.forget_loss_bound=4.0 \
+  trainer.method_args.gamma=1.0 \
+  trainer.method_args.alpha=1.0 \
+  forget_split=forget10 \
+  retain_split=retain90 \
+  holdout_split=holdout10 \
+  retain_logs_path=saves/eval/tofu_${MODEL}_retain90/TOFU_EVAL.json \
+  task_name=demo_custom_unlearn
+
+# 调用链 (谁调用了你的 compute_loss):
+#   src/train.py -> load_trainer() 按 handler 从 TRAINER_REGISTRY 取出 BoundedGradDiff,
+#   用 trainer.args(=TrainingArguments) 和 **method_args 实例化 -> trainer.train()
+#   -> HuggingFace Trainer 训练循环每个 step 调用 compute_loss(model, inputs)。