Merge branch 'main' into sort-paths-in-datasets

configs/omat24/all/eqV2_153M.yml

@@ -0,0 +1,141 @@
+includes:
+  - ../omat24_dataset.yml
+
+trainer: equiformerv2_forces
+
+evaluation_metrics:
+  primary_metric: energy_mae
+  metrics:
+    energy:
+      - mae
+      - per_atom_mae
+    forces:
+      - mae
+      - forcesx_mae
+      - forcesy_mae
+      - forcesz_mae
+      - cosine_similarity
+    stress:
+      - mae
+    stress_isotropic:
+      - mae
+    stress_anisotropic:
+      - mae
+
+outputs:
+  energy:
+    level: system
+    property: energy
+  forces:
+    level: atom
+    property: forces
+    train_on_free_atoms: True
+    eval_on_free_atoms: True
+  stress:
+    level: system
+    property: stress
+    decomposition:
+      stress_isotropic:
+        irrep_dim: 0
+      stress_anisotropic:
+        irrep_dim: 2
+
+loss_functions:
+  - energy:
+      fn: per_atom_mae
+      coefficient: 20
+  - forces:
+      fn: l2mae
+      coefficient: 20
+  - stress_isotropic:
+      fn: mae
+      coefficient: 5
+  - stress_anisotropic:
+      fn: mae
+      reduction: mean_all
+      coefficient: 5
+
+optim:
+  batch_size:                   8 # 8 node - Global 512
+  eval_batch_size:              12
+  num_workers:                  8
+  lr_initial:                   0.0004
+
+  optimizer: AdamW
+  optimizer_params:
+    weight_decay: 0.001
+
+  scheduler:                    LambdaLR
+  scheduler_params:
+    lambda_type:                cosine
+    warmup_factor:              0.2
+    warmup_epochs:              0.1
+    lr_min_factor:              0.01
+
+  max_epochs: 2
+  clip_grad_norm: 100
+  ema_decay: 0.999
+  eval_every: 5000
+  load_balancing: atoms
+
+model:
+  name: hydra
+  pass_through_head_outputs: True
+  otf_graph: True
+
+  backbone:
+    model: equiformer_v2_backbone
+    use_pbc: True
+    use_pbc_single: True
+    otf_graph: True
+
+    enforce_max_neighbors_strictly: False
+    max_neighbors: 20
+    max_radius: 12.0
+    max_num_elements: 96
+
+    avg_num_nodes: 31.17
+    avg_degree: 61.95
+
+    num_layers: 20
+    sphere_channels: 128
+    attn_hidden_channels: 64
+    num_heads: 8
+    attn_alpha_channels: 64
+    attn_value_channels: 16
+    ffn_hidden_channels: 128
+    norm_type: 'layer_norm_sh'
+
+    lmax_list: [ 6 ]
+    mmax_list: [ 3 ]
+    grid_resolution: 18
+
+    num_sphere_samples: 128
+
+    edge_channels: 128
+    use_atom_edge_embedding: True
+    distance_function: 'gaussian'
+    num_distance_basis: 512
+
+    attn_activation: 'silu'
+    use_s2_act_attn: False
+    ffn_activation: 'silu'
+    use_gate_act: False
+    use_grid_mlp: True
+
+    alpha_drop: 0.1
+    drop_path_rate: 0.1
+    proj_drop: 0.0
+
+    weight_init: 'uniform'
+
+  heads:
+    energy:
+      module: equiformer_v2_energy_head
+    forces:
+      module: equiformer_v2_force_head
+    stress:
+      module: rank2_symmetric_head
+      output_name: stress
+      use_source_target_embedding:   True
+      decompose:              True

configs/omat24/all/eqV2_31M.yml

@@ -0,0 +1,147 @@
+includes:
+  - ../omat24_dataset.yml
+
+trainer: equiformerv2_forces
+
+evaluation_metrics:
+  primary_metric: energy_mae
+  metrics:
+    energy:
+      - mae
+      - per_atom_mae
+    forces:
+      - mae
+      - forcesx_mae
+      - forcesy_mae
+      - forcesz_mae
+      - cosine_similarity
+    stress:
+      - mae
+    stress_isotropic:
+      - mae
+    stress_anisotropic:
+      - mae
+
+outputs:
+  energy:
+    level: system
+    property: energy
+  forces:
+    level: atom
+    property: forces
+    train_on_free_atoms: True
+    eval_on_free_atoms: True
+  stress:
+    level: system
+    property: stress
+    decomposition:
+      stress_isotropic:
+        irrep_dim: 0
+      stress_anisotropic:
+        irrep_dim: 2
+
+
+loss_functions:
+  - energy:
+      fn: per_atom_mae
+      coefficient: 20
+  - forces:
+      fn: l2mae
+      coefficient: 20
+  - stress_isotropic:
+      fn: mae
+      coefficient: 5
+  - stress_anisotropic:
+      fn: mae
+      reduction: mean_all
+      coefficient: 5
+
+
+optim:
+  batch_size:                   8      # 8 node - Global 512
+  eval_batch_size:              12
+  load_balancing: atoms
+  num_workers: 8
+  lr_initial:                   0.0006
+
+  optimizer: AdamW
+  optimizer_params:
+    weight_decay: 0.001
+  scheduler: LambdaLR
+  scheduler_params:
+    lambda_type: cosine
+    warmup_factor: 0.2
+    warmup_epochs: 0.01
+    lr_min_factor: 0.01
+
+  max_epochs: 2
+  clip_grad_norm: 100
+  ema_decay: 0.999
+  eval_every: 5000
+
+model:
+  name: hydra
+  pass_through_head_outputs: True
+  otf_graph: True
+
+  backbone:
+    model: equiformer_v2_backbone
+    use_pbc: True
+    use_pbc_single: True
+    otf_graph: True
+
+    enforce_max_neighbors_strictly: False
+
+    max_neighbors: 20
+    max_radius: 12.0
+    max_num_elements: 96
+
+    avg_num_nodes: 31.17
+    avg_degree: 61.95
+
+    num_layers: 8
+    sphere_channels: 128
+    attn_hidden_channels: 64
+    num_heads: 8
+    attn_alpha_channels: 64
+    attn_value_channels: 16
+    ffn_hidden_channels: 128
+    norm_type: 'layer_norm_sh'
+
+    lmax_list: [ 4 ]
+    mmax_list: [ 2 ]
+    grid_resolution: 18
+
+    num_sphere_samples: 128
+
+    edge_channels: 128
+    use_atom_edge_embedding: True
+    share_atom_edge_embedding: False
+    use_m_share_rad: False
+    distance_function: 'gaussian'
+    num_distance_basis: 512
+
+    attn_activation: 'silu'
+    use_s2_act_attn: False
+    use_attn_renorm: True
+    ffn_activation: 'silu'
+    use_gate_act: False
+    use_grid_mlp: True
+    use_sep_s2_act: True
+
+    alpha_drop: 0.1
+    drop_path_rate: 0.1
+    proj_drop: 0.0
+
+    weight_init: 'uniform'
+
+  heads:
+    energy:
+      module: equiformer_v2_energy_head
+    forces:
+      module: equiformer_v2_force_head
+    stress:
+      module: rank2_symmetric_head
+      output_name: stress
+      use_source_target_embedding:   True
+      decompose:                     True