Claude/serene einstein ej ask - #5
Merged
Merged
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.