checkpoints/ablations/mult/baseline/mult_default.yaml

run:
  train_model: mult
  name: rp_mult_baseline_240726
  devices:
  - 0
  - 1
  - 2
  - 3
  precision: 32
  matmul_precision: medium
  dtype: float32
  epochs: 10
  dataset_size: 20000000
  compile: false
  gradient_clip_val: 1.0
  val_check_interval: 0.2
  limit_val_batches: 4
  strategy: ddp
  ckpt_dir: null
logging:
  comet: true
  project: lego_pdgid
config:
  dl_conf:
    lds_args:
      data: <dataset_path>
      frac: 1.0
      cutoff_mev: 10
      min_particles: 0
    bs: 4096
    num_workers: 16
  mm_conf:
    in_dim: 11
    cond_scalars:
    - Density
    - Z
    - A
    - Size
    dropout: 0.05
    h_dim: 128
    n_layers: 6
    n_heads: 6
    bs: 4096
    lr: 0.001
    weight_decay: 0.01
    warmup_steps: 0
    post_emb_norm: false
    use_abs_pos_emb: true
    train_inverse: false
    model_args:
      ff_swish: true
      ff_glu: true
      attn_qk_norm: true
      rotary_pos_emb: true
      use_adaptive_rmsnorm: true
      use_adaptive_layerscale: true
      residual_attn: true
    inv_model_args:
      ff_swish: true
      ff_glu: true
      attn_qk_norm: true
      rotary_pos_emb: true
      use_adaptive_rmsnorm: true
      use_adaptive_layerscale: true
      residual_attn: true
  opt_conf:
    opt: muon
    lr: 0.01
    momentum: 0.95
    nesterov: true
    ns_steps: 5
    weight_decay: 0.007
    weight_decouple: true
    adamw_lr: 0.003
    adamw_betas:
    - 0.9
    - 0.999
    adamw_wd: 0.01
    adamw_eps: 1.0e-08
    scheduler:
      cls: warmup_cosine
      warmup_frac: 0.1
      eta_min: 1.0e-06
      interval: step
  additional:
    notes: ''