train_args.json

{
  "model": "mistralai/Mistral-7B-v0.3",
  "dataset": "narrativeqa",
  "stage": 3,
  "epochs": 1,
  "batch_size": 1,
  "grad_accum": 1,
  "lr": 0.0001,
  "output_dir": "checkpoints/stage3_ert_local_fresh",
  "lora_r": 16,
  "lora_alpha": 32,
  "max_samples": 128,
  "resume_from_checkpoint": null,
  "bf16": false,
  "load_in_4bit": true,
  "device": null,
  "weight_alignment": 0.1,
  "weight_robustness": 0.0,
  "max_length": 256,
  "lambda_init_stage2": 0.1,
  "ert_budgets": [
    0.25,
    0.5,
    0.75
  ],
  "gen_check_interval": 500,
  "gen_check_prompts": 5
}