Upload 2 files

megascale_pretrain_config.yaml

@@ -0,0 +1,173 @@
+nn:
+  model: GeometricTransformer
+  dataset: GeometricTransformerERADatasetPretrain
+  dtype: float32
+  device: cuda
+  data_in_memory: false
+  load_model: /data2/scratch/group_scratch/era/protein_era/models/esm3clm/esm3_clm.pt
+  batch_size: 64
+  num_workers: 4
+  collate_fn: geometric_transformer_era_pretrain_collate_fn
+  model_args:
+    dim_model: 1536
+    unified_transformer_args:
+      n_layers: 48
+      geom_layer_indices:
+      - 0
+      mha_layer_indices:
+      - 0
+      - 1
+      - 2
+      - 3
+      - 4
+      - 5
+      - 6
+      - 7
+      - 8
+      - 9
+      - 10
+      - 11
+      - 12
+      - 13
+      - 14
+      - 15
+      - 16
+      - 17
+      - 18
+      - 19
+      - 20
+      - 21
+      - 22
+      - 23
+      - 24
+      - 25
+      - 26
+      - 27
+      - 28
+      - 29
+      - 30
+      - 31
+      - 32
+      - 33
+      - 34
+      - 35
+      - 36
+      - 37
+      - 38
+      - 39
+      - 40
+      - 41
+      - 42
+      - 43
+      - 44
+      - 45
+      - 46
+      - 47
+      bias: false
+      mha_args:
+        num_heads: 24
+        bias: false
+        qk_layernorm: true
+      gha_args:
+        num_heads: 256
+        num_vector_messages: 1
+        mask_and_zero_frameless: true
+        bias: false
+      scaling_factor: 1.1547005383792515
+      ffn_type: swiglu
+      norm_type: layer_norm
+      expansion_ratio: 2.66666666667
+      ida_layer_indices: []
+    struc_token_info:
+      mask: 4096
+      eos: 4097
+      bos: 4098
+      pad: 4099
+      total: 5001
+      max_non_special_token: 4095
+    residue_token_info:
+      mask: 32
+      eos: 2
+      bos: 0
+      pad: 1
+      total: 33
+      max_non_special_token: null
+    sasa_token_info:
+      mask: 0
+      eos: 0
+      bos: 0
+      pad: 0
+      total: null
+      max_non_special_token: null
+    sec_struct_token_info:
+      mask: 0
+      eos: 0
+      bos: 0
+      pad: 0
+      total: null
+      max_non_special_token: null
+    res_annot_token_info:
+      mask: 0
+      eos: 0
+      bos: 0
+      pad: 0
+      total: null
+      max_non_special_token: null
+  dataset_split_args:
+    train: 1.0
+    val: 0.0
+    test: 0.0
+train:
+  lightning_model: BidirectionalModel
+  resume_training_path: lightning_logs/version_7673/checkpoints/step_step=100000.ckpt
+  lightning_model_args:
+    eval_type: era
+    beta: -10.0
+    gamma: 0
+    sampling_temperature: 1.0
+    optimizer: AdamW
+    optimizer_args:
+      lr: 1.0e-06
+      betas:
+      - 0.9
+      - 0.99
+      weight_decay: 0.01
+    lr_scheduler: null
+    interval: step
+    monitor: train/ERALoss
+    sync_dist: true
+    on_step: true
+    lr_scheduler_args: null
+  trainer_args:
+    accelerator: cuda
+    devices: 4
+    precision: 16-mixed
+    log_every_n_steps: 50
+    max_epochs: 1
+    enable_progress_bar: false
+    gradient_clip_val: 1.0
+  strategy: DDPStrategy
+  strategy_args:
+    find_unused_parameters: true
+  every_epoch_checkpoint_args:
+    filename: step_{step:02d}
+    every_n_epochs: null
+    every_n_train_steps: 50000
+    save_top_k: -1
+  best_checkpoint_args:
+    filename: best_model
+    monitor: train/ERALoss
+    mode: min
+    save_top_k: 1
+  logger:
+    loggertype: TensorBoard
+    logger_args:
+      version: null
+  seed_args:
+    seed: 42
+    workers: true
+global_args:
+  dataset_filename: alignment_dataset.h5
+  keys_to_test:
+  - nn.model
+  - nn.model_args

megascale_pretrained_model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ac49bba2f7f95a524bf0a7b3bd79bc42f93e553de588fabc6a67a8ae03f9e663
+size 16632521682