gptjtest_original.json

{
  "layers": 28,
  "d_model": 4096,
  "n_heads": 16,
  "d_head": 8,
  "n_vocab": 50400,
  "norm": "layernorm",
  "pe": "rotary",
  "pe_rotary_dims": 64,

  "seq": 2048,
  "cores_per_replica": 8,
  "per_replica_batch": 1,
  "gradient_accumulation_steps": 598,

  "warmup_steps": 1,
  "anneal_steps": 9,
  "lr": 1.7e-5,
  "end_lr": 2.4e-6,
  "weight_decay": 0.1,
  "total_steps": 10,

  "tpu_size": 8,

  "bucket": "gptjtest2",
  "model_dir": "puub",

  "train_set": "gptjtest.train.index",
  "val_set": {
    "pile": "gptjtest.val.index",
    "owt": "gptjtest.val.index"
  },

  "eval_harness_tasks": [
    "lambada",
    "piqa",
    "hellaswag",
    "winogrande",
    "mathqa",
    "pubmedqa"
  ],
  
  
  "val_batches": 21,
  "val_every": 100,
  "ckpt_every": 500,
  "keep_every": 10000,

  "name": "gptjtest",
  "wandb_project": "mesh-transformer-jax",
  "comment": ""
}