checkpoint.json

{
  "version": 1,
  "trainer_config": {
    "optimizer": {
      "learning_rate": 0.003,
      "encoder_lr": 0.003,
      "head_lr": 0.003,
      "weight_decay": 0.0001,
      "max_grad_norm": 1.0,
      "warmup_steps": 100,
      "scheduler": "Cosine",
      "min_lr": 1e-06,
      "optimizer_type": "AdamW"
    },
    "loss": {
      "log_score_weight": 1.0,
      "spherical_weight": 0.75,
      "rps_weight": 1.0,
      "epsilon": 1e-07
    },
    "total_steps": 1880,
    "num_epochs": 40,
    "batch_size": 64,
    "validation_split": 0.0,
    "early_stopping_patience": 3,
    "grad_accum_steps": 1,
    "log_interval": 20,
    "checkpoint_interval": 0
  },
  "optimizer_config": {
    "learning_rate": 0.003,
    "encoder_lr": 0.003,
    "head_lr": 0.003,
    "weight_decay": 0.0001,
    "max_grad_norm": 1.0,
    "warmup_steps": 100,
    "scheduler": "Cosine",
    "min_lr": 1e-06,
    "optimizer_type": "AdamW"
  },
  "model_config": {
    "encoder": {
      "vocab_size": 261,
      "hidden_size": 128,
      "num_layers": 2,
      "num_heads": 4,
      "max_seq_len": 256,
      "dropout": 0.1,
      "sliding_window": 64,
      "global_attn_every_n_layers": 2,
      "rope_theta_global": 160000.0,
      "rope_theta_local": 10000.0,
      "ffn_hidden_size": 0,
      "norm_epsilon": 1e-05,
      "activation": "Gelu"
    },
    "head": {
      "attention": {
        "input_size": 128,
        "rank": 32,
        "num_reads": 1,
        "dropout": 0.0
      },
      "hidden_size": 128,
      "num_types": 3,
      "num_heads": 4,
      "num_head_layers": 1,
      "max_options": 32,
      "dropout": 0.1
    },
    "option_embedding_size": 64
  },
  "epoch": 8,
  "batch": 0,
  "best_val_loss": 0.08570641155044238,
  "weights_path": "model.weights"
}