Skip to content

Claude/serene einstein ej ask - #5

Merged
lindsey98 merged 7 commits into
mainfrom
claude/serene-einstein-ejASK
Jun 3, 2026
Merged

Claude/serene einstein ej ask#5
lindsey98 merged 7 commits into
mainfrom
claude/serene-einstein-ejASK

Conversation

@lindsey98

Copy link
Copy Markdown
Owner

No description provided.

claude added 7 commits June 3, 2026 13:12
Read-only integrity tool: normalizes the (system prompt, clean data) content of
a training set and an eval set (handling the different field-name schemas of
datasets/sep/train_dataset.json vs datasets/SEP_dataset.json) and reports any
overlap, exiting non-zero on leakage. Guards against train/test contamination
that would inflate results. Verified on synthetic data.
- (MAIN, result-affecting) data_generation/dpo_data_loader.py: the DPO collator
  tokenized prompt and response separately, and enc() used the default
  add_special_tokens=True, prepending a BOS to every chosen/rejected response.
  Concatenated with the prompt (which already has BOS), this put a spurious BOS
  mid-sequence at the prompt/response boundary — a train/inference mismatch
  (at inference the answer follows the response delimiter with no BOS). The BOS
  term cancels in the DPO chosen-vs-rejected margin so the preference gradient
  is largely unaffected, but the conditioning of the answer tokens is not. Fix:
  add_special_tokens=False (the explicit eos string still maps to the eos id).
  NOTE: changes training data -> re-train DRIP/DPO models to fully apply.
- testing/test.py form_llm_input: the no-attack branch passed the shared source
  dict into apply_testtime_defense, which mutates in place; deepcopy it so a
  non-'none' defense run cannot contaminate later attacks on the same data.
- Harden attack.split('_') -> split('_', 1) in both data loaders.

Verified: compiles, pytest 25 passed/1 skipped, ruff F821/F811 clean.
DPOTrainerMOE.compute_loss fetched the router aux loss with
output_router_logits=False on the policy chosen/rejected forwards, so
_seq_logps_moe always returned aux=None and the 'if chosen_aux is not None and
rejected_aux is not None' guard was never true -- the load-balancing term was
silently dropped (loss = dpo_loss only). The MoE forward only computes
load_balancing_loss_func when output_router_logits=True (qwen_moe_drip.py:305).
Enable it on the two POLICY forwards; the reference forwards stay False (frozen,
no aux needed). NOTE: affects only the Qwen3-MoE variant; re-train it to apply.

Verified: compiles, pytest 25 passed/1 skipped, ruff clean.
- train_unified.py: training saved a LoRA adapter via trainer.save_model, but
  testing/test.py's default load path expects a fully merged checkpoint (and the
  eval scripts don't pass --load_as_adapter) -- so evaluating a freshly trained
  model would load base/untrained weights. Now, on the main process, unwrap the
  model and (for non-QLoRA LoRA runs) merge_and_unload() before save_pretrained,
  producing a full checkpoint the default eval path can load. QLoRA/4-bit and
  non-PEFT (AIR/full-finetune) paths save as before; merge failures fall back to
  adapter save. Adds a post-save distributed barrier.
- trainer.py: DPOTrainerAIR.create_optimizer used a mutable default arg
  (['intermediate_shifts']); switch to the None pattern (behavior unchanged).

Verified: compiles, pytest 25 passed/1 skipped, ruff clean. NOTE: the save-path
change can't be runtime-tested here (no GPU/distributed) -- smoke-test one short
training run to confirm the merged checkpoint loads in eval.
…ints

For checkpoints saved as adapters (QLoRA runs, or models trained before
train_unified started merging at save time). It reuses testing/test.py's
load_full_model(load_as_adapter=True) so loading matches evaluation exactly,
then merge_and_unload() and saves a full checkpoint (weights + tokenizer +
delimiter config) that the default eval path loads directly. Base path and
model class are auto-detected from the adapter path (overridable). Document it
in the README Training section.
The DPO chosen-side judge (tasktracker_judge_prompt2.txt) returns INCLUDED when
the accompanying text was used as data and NOT INCLUDED when it was ignored;
data_curation_drip.py correctly retries on NOT INCLUDED. The README described it
backwards ('checks whether the probe leaked... on INCLUDED it retries'). Correct
the description to match the code; no code change (the logic was right).
Resume skipped already-done requests by COUNT (requests[len(existing):]). Since
failed requests (None results) are dropped, len(existing) under-counts attempted
items, so the length offset would skip some unprocessed requests and reprocess
completed ones. Resume by content key (instruction, injected_input) against the
existing output records' (instruction, input) instead, so failures are retried
and completed items are skipped correctly.
@lindsey98
lindsey98 merged commit 488a590 into main Jun 3, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants