train_args.json

{
  "model": "mistralai/Mistral-7B-v0.3",
  "dataset": "narrativeqa",
  "stage": 1,
  "epochs": 2,
  "batch_size": 4,
  "grad_accum": 8,
  "lr": 0.0001,
  "output_dir": "/workspace/checkpoints/stage1",
  "lora_r": 16,
  "lora_alpha": 32,
  "max_samples": null,
  "resume_from_checkpoint": null,
  "bf16": true,
  "device": null,
  "weight_alignment": 0.1,
  "weight_robustness": 0.0,
  "max_length": 2048,
  "lambda_init_stage2": 0.1
}