diff --git a/auto_antislop_config.yaml b/auto_antislop_config.yaml index 22095c8..049e505 100644 --- a/auto_antislop_config.yaml +++ b/auto_antislop_config.yaml @@ -259,7 +259,8 @@ finetune_shuffle_seed: 666 # 0 = off; 0.9 strongly downsamples overrepresented rule violations # (this is useful because the raw generated dataset is typically very skewed) ftpo_sample_rejected_regularisation_strength: 0.8 -ftpo_sample_chosen_regularisation_strength: 0.2 +# 0 = off; positive values trim globally overrepresented chosen-token slots +ftpo_sample_chosen_regularisation_strength: 0.0 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list # ── FTPO-specific hyper-parameters ───────────────────────────────────────── @@ -276,4 +277,4 @@ ftpo_tau_mse_target: 0.5 # Grace bandwidth (logits) before the above MSE l # MSE loss term 2: stronger mse term applied to remaining (non-target) vocab ftpo_lambda_mse: 0.4 -ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off" \ No newline at end of file +ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off" diff --git a/configs/GLM-4-32B-0414.yaml b/configs/GLM-4-32B-0414.yaml index 4bd5041..321923c 100644 --- a/configs/GLM-4-32B-0414.yaml +++ b/configs/GLM-4-32B-0414.yaml @@ -265,7 +265,8 @@ finetune_shuffle_seed: 666 # 0 = off; 0.9 strongly downsamples overrepresented rule violations # (this is useful because the raw generated dataset is typically very skewed) ftpo_sample_rejected_regularisation_strength: 0.8 -ftpo_sample_chosen_regularisation_strength: 0.2 +# 0 = off; positive values trim globally overrepresented chosen-token slots +ftpo_sample_chosen_regularisation_strength: 0.0 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list @@ -283,4 +284,4 @@ ftpo_tau_mse_target: 0.5 # Grace bandwidth (logits) before the above MSE l # MSE loss term 2: stronger mse term applied to remaining (non-target) vocab ftpo_lambda_mse: 0.4 -ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off" \ No newline at end of file +ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off" diff --git a/configs/Mistral-Small-3.2-24B-Instruct-2506.yaml b/configs/Mistral-Small-3.2-24B-Instruct-2506.yaml index 4245c38..00beb0a 100644 --- a/configs/Mistral-Small-3.2-24B-Instruct-2506.yaml +++ b/configs/Mistral-Small-3.2-24B-Instruct-2506.yaml @@ -265,7 +265,8 @@ finetune_shuffle_seed: 666 # 0 = off; 0.9 strongly downsamples overrepresented rule violations # (this is useful because the raw generated dataset is typically very skewed) ftpo_sample_rejected_regularisation_strength: 0.7 -ftpo_sample_chosen_regularisation_strength: 0.2 +# 0 = off; positive values trim globally overrepresented chosen-token slots +ftpo_sample_chosen_regularisation_strength: 0.0 ftpo_sample_min_chosen_tokens: 3 # filter out ftpo samples that have fewer than this number in the chosen tokens list @@ -283,4 +284,4 @@ ftpo_tau_mse_target: 0.5 # Grace bandwidth (logits) before the above MSE l # MSE loss term 2: stronger mse term applied to remaining (non-target) vocab ftpo_lambda_mse: 0.4 -ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off" \ No newline at end of file +ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off" diff --git a/configs/default.yaml b/configs/default.yaml index 8382c18..7de41f5 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -259,7 +259,8 @@ finetune_shuffle_seed: 666 # 0 = off; 0.9 strongly downsamples overrepresented rule violations # (this is useful because the raw generated dataset is typically very skewed) ftpo_sample_rejected_regularisation_strength: 0.8 -ftpo_sample_chosen_regularisation_strength: 0.2 +# 0 = off; positive values trim globally overrepresented chosen-token slots +ftpo_sample_chosen_regularisation_strength: 0.0 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list # ── FTPO-specific hyper-parameters ───────────────────────────────────────── @@ -276,4 +277,4 @@ ftpo_tau_mse_target: 0.5 # Grace bandwidth (logits) before the above MSE l # MSE loss term 2: stronger mse term applied to remaining (non-target) vocab ftpo_lambda_mse: 0.4 -ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off" \ No newline at end of file +ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off" diff --git a/configs/gemma-3-12b-it.yaml b/configs/gemma-3-12b-it.yaml index 437ff42..677ba0e 100644 --- a/configs/gemma-3-12b-it.yaml +++ b/configs/gemma-3-12b-it.yaml @@ -258,7 +258,8 @@ finetune_shuffle_seed: 666 # 0 = off; 0.9 strongly downsamples overrepresented rule violations # (this is useful because the raw generated dataset is typically very skewed) ftpo_sample_rejected_regularisation_strength: 0.8 -ftpo_sample_chosen_regularisation_strength: 0.2 +# 0 = off; positive values trim globally overrepresented chosen-token slots +ftpo_sample_chosen_regularisation_strength: 0.0 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list @@ -276,4 +277,4 @@ ftpo_tau_mse_target: 0.5 # Grace bandwidth (logits) before the above MSE l # MSE loss term 2: stronger mse term applied to remaining (non-target) vocab ftpo_lambda_mse: 0.4 -ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off" \ No newline at end of file +ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off" diff --git a/configs/gemma-3-27b-it.yaml b/configs/gemma-3-27b-it.yaml index 01654ba..5d05a3e 100644 --- a/configs/gemma-3-27b-it.yaml +++ b/configs/gemma-3-27b-it.yaml @@ -258,7 +258,8 @@ finetune_shuffle_seed: 666 # 0 = off; 0.9 strongly downsamples overrepresented rule violations # (this is useful because the raw generated dataset is typically very skewed) ftpo_sample_rejected_regularisation_strength: 0.6 -ftpo_sample_chosen_regularisation_strength: 0.2 +# 0 = off; positive values trim globally overrepresented chosen-token slots +ftpo_sample_chosen_regularisation_strength: 0.0 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list diff --git a/configs/gemma-3-4b-it.yaml b/configs/gemma-3-4b-it.yaml index cdaf076..9dca09e 100644 --- a/configs/gemma-3-4b-it.yaml +++ b/configs/gemma-3-4b-it.yaml @@ -258,7 +258,8 @@ finetune_shuffle_seed: 666 # 0 = off; 0.9 strongly downsamples overrepresented rule violations # (this is useful because the raw generated dataset is typically very skewed) ftpo_sample_rejected_regularisation_strength: 0.8 -ftpo_sample_chosen_regularisation_strength: 0.2 +# 0 = off; positive values trim globally overrepresented chosen-token slots +ftpo_sample_chosen_regularisation_strength: 0.0 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list @@ -276,4 +277,4 @@ ftpo_tau_mse_target: 0.5 # Grace bandwidth (logits) before the above MSE l # MSE loss term 2: stronger mse term applied to remaining (non-target) vocab ftpo_lambda_mse: 0.4 -ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off" \ No newline at end of file +ftpo_clip_epsilon_logits: 2 # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off" diff --git a/utils/dataset_helpers.py b/utils/dataset_helpers.py index 00d731e..ca3767b 100644 --- a/utils/dataset_helpers.py +++ b/utils/dataset_helpers.py @@ -16,6 +16,82 @@ logger = logging.getLogger(__name__) _WATCH = [" nodded", " leaned"] +def _chosen_target_quotas( + chosen_counts: Counter[str], + strength: float, +) -> dict[str, int]: + """Return per-token occurrence caps for chosen-token regularisation.""" + if not chosen_counts or strength <= 0: + return dict(chosen_counts) + + # Prevent the largest outliers from dominating the distribution before + # applying the smoother median-threshold regularisation. + if len(chosen_counts) >= 10: + cap_value = sorted(chosen_counts.values(), reverse=True)[9] + capped = { + token: min(count, cap_value) + for token, count in chosen_counts.items() + } + else: + capped = dict(chosen_counts) + + median = float(np.median(list(capped.values()))) + return { + token: int(round( + count + if count <= median + else count * (median / count) ** strength + )) + for token, count in capped.items() + } + + +def _trim_chosen_to_quotas( + rows: list[dict], + quotas: dict[str, int], + rng: np.random.Generator, +) -> list[dict]: + """Uniformly retain chosen-token occurrences up to their global quotas.""" + counts = Counter( + token + for row in rows + for token in (row.get("multi_chosen_decoded") or []) + ) + if all(quotas.get(token, count) >= count for token, count in counts.items()): + return [dict(row) for row in rows] + + # Randomise both row and slot traversal so quota allocation does not + # systematically favour early examples or a token's probability rank. + seen: Counter[str] = Counter() + trimmed: list[dict | None] = [None] * len(rows) + for row_idx_raw in rng.permutation(len(rows)): + row_idx = int(row_idx_raw) + row = rows[row_idx] + decoded = row.get("multi_chosen_decoded") or [] + keep: set[int] = set() + for slot_idx_raw in rng.permutation(len(decoded)): + slot_idx = int(slot_idx_raw) + token = decoded[slot_idx] + if seen[token] < max(0, quotas.get(token, counts[token])): + keep.add(slot_idx) + seen[token] += 1 + + new_row = dict(row) + new_row["multi_chosen_decoded"] = [ + token for slot_idx, token in enumerate(decoded) if slot_idx in keep + ] + + raw = row.get("multi_chosen_raw") + if isinstance(raw, list) and len(raw) == len(decoded): + new_row["multi_chosen_raw"] = [ + token for slot_idx, token in enumerate(raw) if slot_idx in keep + ] + + trimmed[row_idx] = new_row + + return [row for row in trimmed if row is not None] + + def load_ftpo_multi_dataset( path: Path, tokenizer, @@ -108,24 +184,10 @@ def load_ftpo_multi_dataset( _log_top(chosen_cts_orig, "ORIGINAL CHOSEN TOKENS") - # Trim the peak: cap top tokens to match the 10th highest count - if len(chosen_cts_orig) >= 10: - top_counts = sorted(chosen_cts_orig.values(), reverse=True) - cap_value = top_counts[9] # 10th highest count - chosen_cts_capped = Counter() - for tok, cnt in chosen_cts_orig.items(): - chosen_cts_capped[tok] = min(cnt, cap_value) - else: - chosen_cts_capped = chosen_cts_orig.copy() - - # Now calculate regularization on the capped distribution - med_chosen = float(np.median(list(chosen_cts_capped.values()))) - w_chosen = {tok: 1.0 if c <= med_chosen - else (med_chosen / c) ** chosen_reg_strength - for tok, c in chosen_cts_capped.items()} - - tgt_chosen = {tok: int(round(c * w_chosen.get(tok, 1.0))) - for tok, c in chosen_cts_capped.items()} + tgt_chosen = _chosen_target_quotas( + chosen_cts_orig, + chosen_reg_strength, + ) # Log the target quotas quota_items = sorted(tgt_chosen.items(), key=lambda x: x[1], reverse=True)[:20] @@ -137,7 +199,13 @@ def load_ftpo_multi_dataset( _WATCH[1], tgt_chosen.get(_WATCH[1], 0), chosen_cts_orig.get(_WATCH[1], 0), ) - _log_top(Counter(r["rejected_decoded"] for r in rows), "POST-CHOSEN") + rows = _trim_chosen_to_quotas(rows, tgt_chosen, rng) + chosen_cts_trimmed = Counter( + token + for row in rows + for token in (row["multi_chosen_decoded"] or []) + ) + _log_top(chosen_cts_trimmed, "POST-CHOSEN") # ──────────────────────────────────────────────────────────────── # 3️⃣ Apply min_chosen_tokens row filter @@ -206,6 +274,13 @@ def load_ftpo_multi_dataset( rows = selected + final_chosen_counts = Counter( + token + for row in rows + for token in (row["multi_chosen_decoded"] or []) + ) + _log_top(final_chosen_counts, "FINAL CHOSEN TOKENS") + # ── Dump the final row subset exactly as it was read (no tokenisation) ── if experiment_run_dir is not None: ts = datetime.now(timezone.utc).astimezone()\ diff --git a/utils/test_dataset_helpers.py b/utils/test_dataset_helpers.py new file mode 100644 index 0000000..ee28377 --- /dev/null +++ b/utils/test_dataset_helpers.py @@ -0,0 +1,154 @@ +import json +import tempfile +import unittest +from collections import Counter +from pathlib import Path + +import numpy as np + +from utils.dataset_helpers import ( + _chosen_target_quotas, + _trim_chosen_to_quotas, + load_ftpo_multi_dataset, +) + + +def _counts(rows): + return Counter( + token + for row in rows + for token in row["multi_chosen_decoded"] + ) + + +class _Tokenized: + def __init__(self, input_ids): + self.input_ids = input_ids + + +class _FakeTokenizer: + truncation_side = "right" + + def __call__(self, text, **_kwargs): + def encode(value): + if value.startswith("context-"): + return [1, 2] + return [100 + sum(value.encode("utf-8"))] + + if isinstance(text, list): + return _Tokenized([encode(value) for value in text]) + return _Tokenized(encode(text)) + + +class ChosenRegularisationTests(unittest.TestCase): + def test_zero_strength_disables_all_trimming(self): + counts = Counter({f"token-{i}": 20 - i for i in range(12)}) + + self.assertEqual(_chosen_target_quotas(counts, 0), dict(counts)) + + def test_positive_strength_caps_and_regularises_outliers(self): + counts = Counter({f"token-{i}": 100 - 5 * i for i in range(12)}) + + quotas = _chosen_target_quotas(counts, 0.2) + + tenth_highest = sorted(counts.values(), reverse=True)[9] + self.assertLess(quotas["token-0"], counts["token-0"]) + self.assertLessEqual(quotas["token-0"], tenth_highest) + self.assertEqual(quotas["token-11"], counts["token-11"]) + + def test_trimming_enforces_quotas_and_keeps_raw_fields_aligned(self): + rows = [ + { + "multi_chosen_decoded": [" common", " rare-a", " common"], + "multi_chosen_raw": ["raw-common-1", "raw-rare-a", "raw-common-2"], + }, + { + "multi_chosen_decoded": [" common", " rare-b"], + "multi_chosen_raw": ["raw-common-3", "raw-rare-b"], + }, + ] + quotas = {" common": 2, " rare-a": 1, " rare-b": 1} + + trimmed = _trim_chosen_to_quotas( + rows, quotas, np.random.default_rng(3407) + ) + + self.assertEqual(_counts(trimmed), Counter(quotas)) + for row in trimmed: + self.assertEqual( + len(row["multi_chosen_decoded"]), + len(row["multi_chosen_raw"]), + ) + # The helper does not mutate the source rows. + self.assertEqual(rows[0]["multi_chosen_decoded"].count(" common"), 2) + + def test_trimming_is_reproducible(self): + rows = [ + { + "multi_chosen_decoded": [" common", f" unique-{i}"], + "multi_chosen_raw": [" common", f" unique-{i}"], + } + for i in range(20) + ] + quotas = {" common": 5, **{f" unique-{i}": 1 for i in range(20)}} + + first = _trim_chosen_to_quotas( + rows, quotas, np.random.default_rng(123) + ) + second = _trim_chosen_to_quotas( + rows, quotas, np.random.default_rng(123) + ) + + self.assertEqual(first, second) + self.assertEqual(_counts(first)[" common"], 5) + + def test_minimum_filter_applies_after_trimming(self): + rows = [ + {"multi_chosen_decoded": [" common", " keep"]}, + {"multi_chosen_decoded": [" common", " other"]}, + ] + + trimmed = _trim_chosen_to_quotas( + rows, + {" common": 1, " keep": 1, " other": 1}, + np.random.default_rng(7), + ) + surviving = [ + row for row in trimmed if len(row["multi_chosen_decoded"]) >= 2 + ] + + self.assertEqual(len(surviving), 1) + + def test_loader_filters_rows_after_applying_chosen_quotas(self): + rows = [ + { + "context_with_chat_template": f"context-{i}", + "rejected_decoded": " reject", + "multi_chosen_decoded": [" common", f" unique-{i}"], + "multi_chosen_raw": ["raw-common", f"raw-unique-{i}"], + } + for i in range(4) + ] + + with tempfile.TemporaryDirectory() as tmp_dir: + path = Path(tmp_dir) / "ftpo.jsonl" + with path.open("w", encoding="utf-8") as handle: + for row in rows: + handle.write(json.dumps(row) + "\n") + + dataset = load_ftpo_multi_dataset( + path, + _FakeTokenizer(), + chosen_reg_strength=1.0, + min_chosen_tokens=2, + num_proc=1, + ) + + # " common" is trimmed from four occurrences to one, so only its + # containing row still meets the two-chosen-token minimum. + self.assertEqual(len(dataset), 1) + self.assertEqual(len(dataset[0]["chosen_ids"]), 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/trainer_dataloaders.py b/utils/trainer_dataloaders.py index b83ea5d..b7c1c83 100644 --- a/utils/trainer_dataloaders.py +++ b/utils/trainer_dataloaders.py @@ -110,7 +110,7 @@ def load_and_prepare_dataset(config: dict, experiment_run_dir: Path, tokenizer: # balance *rejected* tokens rejected_reg_strength = config.get("ftpo_sample_rejected_regularisation_strength", 0.8), # balance *chosen* tokens - chosen_reg_strength = config.get("ftpo_sample_chosen_regularisation_strength", 0.2), + chosen_reg_strength = config.get("ftpo_sample_chosen_regularisation_strength", 0.0), # hard floor on |chosen| min_chosen_tokens = config.get("ftpo_sample_min_chosen_tokens", 3), # overall training-set cap (used for per-token quotas too) @@ -158,7 +158,7 @@ def load_and_prepare_dataset(config: dict, experiment_run_dir: Path, tokenizer: experiment_run_dir = experiment_run_dir, max_seq_len = max_seq_length, rejected_reg_strength = config.get("ftpo_sample_rejected_regularisation_strength", 0.8), - chosen_reg_strength = config.get("ftpo_sample_chosen_regularisation_strength", 0.2), + chosen_reg_strength = config.get("ftpo_sample_chosen_regularisation_strength", 0.0), min_chosen_tokens = config.get("ftpo_sample_min_chosen_tokens", 3), max_train_examples = config.get("finetune_max_train_examples"), ) @@ -247,7 +247,7 @@ def load_and_prepare_dataset(config: dict, experiment_run_dir: Path, tokenizer: experiment_run_dir = experiment_run_dir, max_seq_len = max_seq_length, rejected_reg_strength = config.get("ftpo_sample_rejected_regularisation_strength", 0.8), - chosen_reg_strength = config.get("ftpo_sample_chosen_regularisation_strength", 0.2), + chosen_reg_strength = config.get("ftpo_sample_chosen_regularisation_strength", 0.0), min_chosen_tokens = config.get("ftpo_sample_min_chosen_tokens", 3), max_train_examples = config.get("finetune_max_train_examples"), ) @@ -299,7 +299,7 @@ def load_and_prepare_dataset(config: dict, experiment_run_dir: Path, tokenizer: experiment_run_dir = experiment_run_dir, max_seq_len = max_seq_length, rejected_reg_strength = config.get("ftpo_sample_rejected_regularisation_strength", 0.8), - chosen_reg_strength = config.get("ftpo_sample_chosen_regularisation_strength", 0.2), + chosen_reg_strength = config.get("ftpo_sample_chosen_regularisation_strength", 0.0), min_chosen_tokens = config.get("ftpo_sample_min_chosen_tokens", 3), max_train_examples = config.get("finetune_max_train_examples"), ) @@ -345,4 +345,4 @@ def load_and_prepare_dataset(config: dict, experiment_run_dir: Path, tokenizer: logger.error(f"Unknown finetune_mode '{mode}'. Use 'dpo' or 'ftpo'.") return None - return dpo_dataset_hf \ No newline at end of file + return dpo_dataset_hf