Skip to content

Commit 6412602

Browse files
wood-bwood-b
committed
training configs
1 parent 5d9b1c2 commit 6412602

10 files changed

Lines changed: 1032 additions & 5 deletions

‎configs/uma/benchmark/cluster/h100.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ data_root_dir: /checkpoint/ocp/shared/benchmark
22
checkpoint_root_dir: /checkpoint/ocp/shared/uma_checkpoints
33
run_dir: /checkpoint/ocp/shared/uma
44
account: ocp
5-
qos: alignment_shared
5+
qos: h100_ocp_high
66
partition: learn
77
mode: SLURM
88
device: CUDA

‎configs/uma/training_release/dataset/mptrj_omat_finetune.yaml‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
omat_train:
2+
splits:
3+
train:
4+
src:
5+
- ${cluster.data_root_dir}/fm_training/alex-10M/train/
6+
- ${cluster.data_root_dir}/fm_training/mptrj_ue/train/
7+
- ${cluster.data_root_dir}/fm_training/mptrj_ue/train/
8+
- ${cluster.data_root_dir}/fm_training/mptrj_ue/train/
9+
- ${cluster.data_root_dir}/fm_training/mptrj_ue/train/
10+
- ${cluster.data_root_dir}/fm_training/mptrj_ue/train/
11+
- ${cluster.data_root_dir}/fm_training/mptrj_ue/train/
12+
- ${cluster.data_root_dir}/fm_training/mptrj_ue/train/
13+
- ${cluster.data_root_dir}/fm_training/mptrj_ue/train/
14+
format: ase_db
15+
a2g_args:
16+
r_energy: True
17+
r_forces: True
18+
r_stress: ${regress_stress}
19+
r_edges: ${cpu_graph}
20+
radius: ${cutoff_radius}
21+
max_neigh: ${max_neighbors}
22+
key_mapping:
23+
energy: omat_energy
24+
forces: ${omat_forces_key}
25+
stress: omat_stress
26+
transforms:
27+
common_transform:
28+
dataset_name: omat
29+
stress_reshape_transform:
30+
dataset_name: omat
31+
32+
omat_val:
33+
splits:
34+
val:
35+
src: ${cluster.data_root_dir}/fm_training/mptrj_ue/val/
36+
format: ase_db
37+
a2g_args:
38+
r_energy: True
39+
r_forces: True
40+
r_stress: ${regress_stress}
41+
r_edges: ${cpu_graph}
42+
radius: ${cutoff_radius}
43+
max_neigh: ${max_neighbors}
44+
key_mapping:
45+
energy: omat_energy
46+
forces: ${omat_forces_key}
47+
stress: omat_stress
48+
transforms:
49+
common_transform:
50+
dataset_name: omat
51+
stress_reshape_transform:
52+
dataset_name: omat

‎configs/uma/training_release/element_refs/mptrj_lin_refs.yaml‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
mptrj_elem_refs:
2+
- 0.0
3+
- -3.679124355316162
4+
- -1.3378056287765503
5+
- -3.7239913940429688
6+
- -4.592716693878174
7+
- -7.687317848205566
8+
- -8.41640853881836
9+
- -7.525318145751953
10+
- -8.025588989257812
11+
- -5.389685153961182
12+
- -0.03005797415971756
13+
- -2.7558460235595703
14+
- -2.7971031665802
15+
- -4.721299171447754
16+
- -7.570775508880615
17+
- -6.981732368469238
18+
- -4.996428966522217
19+
- -3.354858160018921
20+
- -0.06283307820558548
21+
- -2.60414981842041
22+
- -5.521111488342285
23+
- -7.961789608001709
24+
- -10.151139259338379
25+
- -9.776466369628906
26+
- -10.182172775268555
27+
- -9.458876609802246
28+
- -7.932404041290283
29+
- -6.091212749481201
30+
- -5.266493320465088
31+
- -3.203958749771118
32+
- -1.3486576080322266
33+
- -3.5260260105133057
34+
- -4.564589977264404
35+
- -3.898099422454834
36+
- -4.201990127563477
37+
- -2.8809189796447754
38+
- 6.9195637702941895
39+
- -2.574320077896118
40+
- -5.184962272644043
41+
- -10.202549934387207
42+
- -11.72141170501709
43+
- -12.063482284545898
44+
- -10.50037956237793
45+
- -8.85855770111084
46+
- -7.736929416656494
47+
- -6.870950222015381
48+
- -4.921371936798096
49+
- -2.1186888217926025
50+
- -0.6435719132423401
51+
- -2.779581069946289
52+
- -3.7746050357818604
53+
- -3.6241791248321533
54+
- -3.1492624282836914
55+
- -1.8959052562713623
56+
- 10.294095993041992
57+
- -2.846658229827881
58+
- -4.987889766693115
59+
- -9.249128341674805
60+
- -8.800369262695312
61+
- -8.043109893798828
62+
- -8.190547943115234
63+
- -7.944450855255127
64+
- -8.290793418884277
65+
- -13.709291458129883
66+
- -18.544340133666992
67+
- -7.6736063957214355
68+
- -7.989833354949951
69+
- -7.509627819061279
70+
- -6.765751361846924
71+
- -7.9181227684021
72+
- -3.65626859664917
73+
- -7.298524379730225
74+
- -12.700183868408203
75+
- -14.0196533203125
76+
- -12.38825798034668
77+
- -11.336587905883789
78+
- -9.53471851348877
79+
- -7.371169090270996
80+
- -5.319552421569824
81+
- -2.4043772220611572
82+
- 0.28949788212776184
83+
- -2.39648699760437
84+
- -3.7129523754119873
85+
- -3.401271343231201
86+
- 0.0
87+
- 0.0
88+
- 0.0
89+
- 0.0
90+
- 0.0
91+
- -3.553612232208252
92+
- -10.63386058807373
93+
- -12.45141887664795
94+
- -13.549932479858398
95+
- -14.93902587890625
96+
- -15.31093692779541
97+
- 0.0
98+
- 0.0
99+
- 0.0
100+
- 0.0
101+
- 0.0
102+
- 0.0
103+
- 0.0
104+
- 0.0
105+
- 0.0
106+
- 0.0
107+
- 0.0
108+
- 0.0
109+
- 0.0
110+
- 0.0
111+
- 0.0
112+
- 0.0
113+
- 0.0
114+
- 0.0
115+
- 0.0
116+
- 0.0
117+
- 0.0
118+
- 0.0
119+
- 0.0
120+
- 0.0

‎configs/uma/training_release/tasks/mptrj_omat_conserving_stress.yaml‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
- _target_: fairchem.core.units.mlip_unit.mlip_unit.Task
2+
name: omat_energy
3+
level: system
4+
property: energy
5+
loss_fn:
6+
_target_: fairchem.core.modules.loss.DDPMTLoss
7+
loss_fn:
8+
_target_: fairchem.core.modules.loss.PerAtomMAELoss
9+
coefficient: ${energy_coef}
10+
out_spec:
11+
dim: [1]
12+
dtype: float32
13+
normalizer:
14+
_target_: fairchem.core.modules.normalization.normalizer.Normalizer
15+
mean: 0.0
16+
rmsd: ${normalizer_rmsd}
17+
element_references:
18+
_target_: fairchem.core.modules.normalization.element_references.ElementReferences
19+
element_references:
20+
_target_: torch.DoubleTensor
21+
_args_:
22+
- ${element_refs.mptrj_elem_refs}
23+
datasets:
24+
- omat
25+
metrics:
26+
- mae
27+
- per_atom_mae
28+
- _target_: fairchem.core.units.mlip_unit.mlip_unit.Task
29+
name: omat_forces
30+
level: atom
31+
property: forces
32+
train_on_free_atoms: True
33+
eval_on_free_atoms: True
34+
loss_fn:
35+
_target_: fairchem.core.modules.loss.DDPMTLoss
36+
loss_fn:
37+
_target_: fairchem.core.modules.loss.L2NormLoss
38+
reduction: mean
39+
coefficient: ${force_coef}
40+
out_spec:
41+
dim: [3]
42+
dtype: float32
43+
normalizer:
44+
_target_: fairchem.core.modules.normalization.normalizer.Normalizer
45+
mean: 0.0
46+
rmsd: ${normalizer_rmsd}
47+
datasets:
48+
- omat
49+
metrics:
50+
- mae
51+
- cosine_similarity
52+
- magnitude_error
53+
- _target_: fairchem.core.units.mlip_unit.mlip_unit.Task
54+
name: omat_stress
55+
level: system
56+
property: stress
57+
loss_fn:
58+
_target_: fairchem.core.modules.loss.DDPMTLoss
59+
loss_fn:
60+
_target_: fairchem.core.modules.loss.MAELoss
61+
reduction: mean
62+
coefficient: ${stress_coef}
63+
out_spec:
64+
dim: [1, 9]
65+
dtype: float32
66+
normalizer:
67+
_target_: fairchem.core.modules.normalization.normalizer.Normalizer
68+
mean: 0.0
69+
rmsd: ${normalizer_rmsd}
70+
datasets:
71+
- omat
72+
metrics:
73+
- mae

‎configs/uma/training_release/uma_sm_full_conserve_1p2.yaml‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,15 +11,15 @@ job:
1111
scheduler:
1212
mode: ${cluster.mode}
1313
ranks_per_node: ${cluster.ranks_per_node}
14-
num_nodes: 16
14+
num_nodes: 32
1515
slurm:
1616
account: ${cluster.account}
1717
qos: h200_ocp_high
1818
mem_gb: ${cluster.mem_gb}
1919
cpus_per_task: ${cluster.cpus_per_task}
2020
debug: ${cluster.debug}
2121
run_dir: ${cluster.run_dir}
22-
run_name: uma_sm_full_conserving_v1_2b_retain_bug_fix_cpp
22+
run_name: uma_sm_full_conserving_v1_2b_retain_bug_fix_cpp_speed_test
2323
logger:
2424
_target_: fairchem.core.common.logger.WandBSingletonLogger.init_wandb
2525
_partial_: true
@@ -30,9 +30,9 @@ moe_layer_type: fairchem_cpp
3030
num_moe_experts: 64
3131
max_neighbors: 300
3232
cutoff_radius: 6
33-
epochs: 2
33+
epochs: 2.5
3434
steps: null #3216000 # 140B atoms, 128 ranks, max atoms 700 (mean atoms 650)
35-
max_atoms: 700
35+
max_atoms: 350
3636
bf16: False
3737
cpu_graph: True
3838
otf_graph: False

0 commit comments

Comments
 (0)