Upload config.yaml with huggingface_hub

config.yaml

@@ -0,0 +1,78 @@
+# ELECTRA Corrective MLM — full-scale 1B pretraining
+# Generator corrupts input with plausible-but-guaranteed-wrong tokens,
+# discriminator predicts original tokens at ALL positions.
+#
+# Usage:
+#   /cache/ssahn_lab/envs/voxfaplm/bin/python -m torch.distributed.run \
+#       --nproc_per_node=8 experiments/train_multinode.py \
+#       --config configs/pretrain/pretrain_electra.yaml
+
+model:
+  name: "ELECTRAProteinModel"
+  max_residues: 512
+  embed_dim: 2560
+  encoder_depth: 33
+  encoder_heads: 40
+  decoder_dim: 2560
+  decoder_depth: 3
+  decoder_heads: 40
+  # ELECTRA-specific
+  generator_depth: 8            # ~27% of disc encoder
+  generator_decoder_depth: 2
+  disc_loss_weight: 50.0         # corrective CE ≈ same scale as generator CE
+  generator_lr_multiplier: 1.0
+  generator_temperature: 2.0
+  relax_temperature_scaling: null   # null = 1/sqrt(head_dim) further relaxes attention temperature
+  # Stability: QK-Norm (Dehghani 2023, ViT-22B; also Qwen3) prevents attention
+  use_qk_norm: false
+
+# Training configuration
+training:
+  learning_rate: 1.0e-3
+  weight_decay: 0.01
+  # Qwen/Llama/PaLM use β2=0.95 at scale (default 0.999 causes loss spikes).
+  adam_beta1: 0.9
+  adam_beta2: 0.95
+  # Warm-starting from a pre-QK-Norm ckpt with calibrated q_norm/k_norm:
+  # longer warmup lets the optimizer slide the calibrated weights toward the
+  # standard RMSNorm range without shocking the learned representation.
+  warmup_steps: 5000
+  batch_size: 32
+  num_workers: 8
+  max_steps: 500000
+  total_training_steps: 500000
+  accumulate_grad_batches: 1
+  val_check_interval: 100
+  test_interval: 100
+  gradient_clip_val: 1.0
+# Data configuration
+data:
+  name: "ELECTRADataModule"
+  lmdb_dir: "/cache/ssahn_lab/lmdb/atlas_pdb_all_CHI"
+  pdb_dir: "/cache"
+  struct_format: "fullatom"
+  max_length: 512
+  bb_vocab_size: 512
+  fa_vocab_size: 512
+  mask_prob: 0.6
+  random_cropping: true
+  # Sequence packing: concat variable-length samples into one super-sequence
+  # (bsz=1) with xformers BlockDiagonalMask. Keeps flash attention and
+  # eliminates padding waste. Only affects training — probing/eval uses
+  # padded path unchanged.
+  pack_sequences: false
+  # dynamic_batch_tokens: 65536  # opt-in: variable-size batches that pack up
+  #                                to this token budget. Disables compile.
+
+# Lightning configuration
+lightning:
+  gpus: 8
+  nodes: 2
+  precision: "bf16-mixed"
+  compile: true
+  checkpoint_dir: "/cache/ssahn_lab/pretrain_model/electra_3B/checkpoints"
+  log_dir: "/cache/ssahn_lab/pretrain_model/electra_3B/logs"
+  wandb_dir: "/cache/ssahn_lab/pretrain_model/electra_3B/wandb"
+  wandb_project: "PLM_submission"  
+  wandb_name: "electra_3B"
+  save_every_n_steps: 10000