config.json

{
  "hidden_size": 4096,
  "num_head_layers": 1,
  "num_heads": 8,
  "use_adaln": true,
  "use_mask_embeddings": true,
  "use_block_embeddings": false,
  "learning_rate": 5e-05,
  "weight_decay": 0.01,
  "num_epochs": 3,
  "batch_size": 8,
  "gradient_accumulation_steps": 1,
  "max_grad_norm": 1.0,
  "warmup_steps": 100,
  "exploitation_loss_weight": 1.0,
  "exploration_loss_weight": 1.0,
  "gen_length": 512,
  "mask_ratio_range": [
    0.1,
    0.9
  ],
  "log_interval": 5,
  "eval_interval": 2000,
  "save_interval": 10000,
  "max_checkpoints": 2,
  "use_wandb": true,
  "wandb_project": "planner-sft-ete",
  "wandb_run_name": "sft_ete_6train_1val",
  "output_dir": "/accounts/projects/songmei/hengyuf/hengyu_scratch/hengyuf/ete-checkpoints",
  "resume_from_checkpoint": null,
  "device": "cuda:0",
  "mixed_precision": true
}