6B_roto_256.json

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

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

  "warmup_steps": 3000,
  "anneal_steps": 300000,
  "lr": 1.2e-4,
  "end_lr": 1.2e-5,
  "weight_decay": 0.1,
  "total_steps": 350000,

  "tpu_size": 256,

  "bucket": "neo-models",
  "model_dir": "mesh_jax_pile_6B_rotary",

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

  "eval_harness_tasks": [
    "lambada",
    "piqa",
    "hellaswag",
    "winogrande",
    "mathqa",
    "pubmedqa"
  ],

  "val_batches": 100,
  "val_every": 500,
  "ckpt_every": 500,
  "keep_every": 10000,

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