Fine tuning OC25 model

[Alexandrina-Chen]

Hi all,

I am new to fairchem, and was trying to fine-tuning OC25 model (esen-sm-conserving-all-oc25) on my own dataset. I followed the manual fine_tuning and was able to fine tune the default uma-sm-1 checking point on my own dataset.

Then I changed the uma_sm_finetune_template.yaml's base_model_name from the default uma-s-1 to esen-sm-conserving-all-oc25. Everything is exactly the same as the default setting generated by the create_uma_finetune_dataset.py except for the base_model_name. I then ran the same command fairchem -c uma_sm_finetune_template.yaml, and the following error occured:

rain Epoch 0:   0%|          | 0/200 [00:00<?, ?it/s]
INFO:root:Saved dcp checkpoint to /tmp/uma_finetune_runs/202510-1416-1153-8328/checkpoints/step_0
ERROR:root:Exception during training! On step 0
ERROR:root:Data info: AtomicData(atomic_numbers=[768], atomic_numbers_full=[768], batch=[768], batch_full=[768], cell=[2, 3, 3], cell_offsets=[0, 3], charge=[2], dataset=[2], dataset_name=[2], edge_index=[2, 0], energy=[2], fixed=[768], forces=[768, 3], natoms=[2], nedges=[2], pbc=[2, 3], pos=[768, 3], sid=[2], spin=[2], stress=[2, 9], tags=[768])
data.dataset: ['omat', 'omat']
data.sid: ['', '']
data.natoms: tensor([384, 384], device='cuda:0')
data.atomic_numbers: tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 8, 8, 8, 8, 8, 8, 8, 8,
        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 8, 8, 8, 8, 8, 8, 8, 8,
        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8],
       device='cuda:0')
INFO:torchtnt.framework.fit:Exception during fit after the following progress: train progress: completed epochs: 0, completed steps: 0, completed steps in current epoch: 0. eval progress: completed epochs: 0, completed steps: 0, completed steps in current epoch: 0.:
'omat'
[rank0]: Traceback (most recent call last):
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/bin/fairchem", line 8, in <module>
[rank0]:     sys.exit(main())
[rank0]:              ^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/fairchem/core/_cli.py", line 134, in main
[rank0]:     local_launch(cfg, log_dir)
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/fairchem/core/launchers/slurm_launch.py", line 288, in local_launch
[rank0]:     runner_wrapper(cfg)
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/fairchem/core/launchers/slurm_launch.py", line 85, in runner_wrapper
[rank0]:     SlurmSPMDProgram()(config, run_type)
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/fairchem/core/launchers/slurm_launch.py", line 159, in __call__
[rank0]:     self.runner.run()
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/fairchem/core/components/train/train_runner.py", line 159, in run
[rank0]:     fit(
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/torchtnt/framework/fit.py", line 144, in fit
[rank0]:     raise e
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/torchtnt/framework/fit.py", line 133, in fit
[rank0]:     _train_impl(state, unit, callback_handler)
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/torchtnt/framework/train.py", line 148, in _train_impl
[rank0]:     _train_epoch_impl(state, train_unit, callback_handler)
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/torchtnt/framework/train.py", line 215, in _train_epoch_impl
[rank0]:     train_state._step_output = train_unit.train_step(state, step_input)
[rank0]:                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/fairchem/core/units/mlip_unit/mlip_unit.py", line 698, in train_step
[rank0]:     pred = self.model.forward(batch_on_device)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/torch/nn/parallel/distributed.py", line 1648, in forward
[rank0]:     else self._run_ddp_forward(*inputs, **kwargs)
[rank0]:          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/torch/nn/parallel/distributed.py", line 1474, in _run_ddp_forward
[rank0]:     return self.module(*inputs, **kwargs)  # type: ignore[index]
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/fairchem/core/models/base.py", line 178, in forward
[rank0]:     emb = self.backbone(data)
[rank0]:           ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/fairchem/core/common/utils.py", line 80, in cls_method
[rank0]:     return f(self, *args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/fairchem/core/models/uma/escn_md.py", line 412, in forward
[rank0]:     csd_mixed_emb = self.csd_embedding(
[rank0]:                     ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/fairchem/core/models/uma/escn_md.py", line 336, in csd_embedding
[rank0]:     dataset_emb = self.dataset_embedding(dataset)
[rank0]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/fairchem/core/models/uma/nn/embedding_dev.py", line 239, in forward
[rank0]:     emb_for_datasets = [
[rank0]:                        ^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/fairchem/core/models/uma/nn/embedding_dev.py", line 242, in <listcomp>
[rank0]:     else self.dataset_emb_dict[dataset](emb_idx)
[rank0]:          ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^
[rank0]:   File "/mnt/home/schen1/my-python-venv/fairchem/fairchem/lib/python3.11/site-packages/torch/nn/modules/container.py", line 543, in __getitem__
[rank0]:     return self._modules[key]
[rank0]:            ~~~~~~~~~~~~~^^^^^
[rank0]: KeyError: 'omat'

Train Epoch 0:   0%|          | 0/200 [00:01<?, ?it/s]

[rank0]:[W1014 16:12:12.711099894 ProcessGroupNCCL.cpp:1538] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())

I was wondering how I could fix this error. I assume it is probably related to the uma_conserving_data_task_energy_force_stress.yaml file generated by create_uma_finetune_dataset.py file?

If any additional information is needed, please kindly let me know

[Alexandrina-Chen]

I think I have figured the initial question out.

Another naive question I have is that I am facing NaN loss after the first few steps using no matter UMA small (+different tasks) or OC25 small models. I have tried reducing learning rate and it didn't help.

INFO:root:{'train/loss': 4.122553922740137, 'train/lr': 9.994504881061753e-07, 'train/step': 30, 'train/epoch': 0.15, 'train/samples_per_second(approx)': 1.7575262997010452, 'train/atoms_per_second(approx)': 674.8900990852014, 'train/num_atoms_on_rank': 768, 'train/num_samples_on_rank': 2}

Train Epoch 0:  16%|█▌        | 31/200 [00:43<03:14,  1.15s/it]
WARNING:root:Found nans while computing loss
WARNING:root:Found nans while computing loss
WARNING:root:Found nans while computing loss

Here is an example yaml script I used:

defaults:
- data: uma_conserving_data_task_energy_force
- _self_
job:
  device_type: CUDA
  scheduler:
    mode: LOCAL
    ranks_per_node: 1
    num_nodes: 1
  debug: true
  run_dir: /tmp/uma_finetune_runs/
  run_name: uma_finetune
  logger:
    _target_: fairchem.core.common.logger.WandBSingletonLogger.init_wandb
    _partial_: true
    entity: example
    project: uma_finetune
base_model_name: uma-s-1p1
max_neighbors: 300
epochs: 1
steps: null
batch_size: 2
lr: 1e-5
weight_decay: 1e-5
evaluate_every_n_steps: 100
checkpoint_every_n_steps: 1000
train_dataset:
  _target_: fairchem.core.datasets.mt_concat_dataset.create_concat_dataset
  dataset_configs:
    omc: ${data.train_dataset}
  combined_dataset_config:
    sampling:
      type: temperature
      temperature: 1.0
train_dataloader:
  _target_: fairchem.core.components.common.dataloader_builder.get_dataloader
  dataset: ${train_dataset}
  batch_sampler_fn:
    _target_: fairchem.core.common.data_parallel.BalancedBatchSampler
    _partial_: true
    batch_size: ${batch_size}
    shuffle: true
    seed: 0
  num_workers: 0
  collate_fn:
    _target_: fairchem.core.units.mlip_unit.mlip_unit.mt_collater_adapter
    tasks: ${data.tasks_list}
val_dataset:
  _target_: fairchem.core.datasets.mt_concat_dataset.create_concat_dataset
  dataset_configs:
    omc: ${data.val_dataset}
  combined_dataset_config:
    sampling:
      type: temperature
      temperature: 1.0
eval_dataloader:
  _target_: fairchem.core.components.common.dataloader_builder.get_dataloader
  dataset: ${val_dataset}
  batch_sampler_fn:
    _target_: fairchem.core.common.data_parallel.BalancedBatchSampler
    _partial_: true
    batch_size: ${batch_size}
    shuffle: false
    seed: 0
  num_workers: 0
  collate_fn:
    _target_: fairchem.core.units.mlip_unit.mlip_unit.mt_collater_adapter
    tasks: ${data.tasks_list}
runner:
  _target_: fairchem.core.components.train.train_runner.TrainEvalRunner
  train_dataloader: ${train_dataloader}
  eval_dataloader: ${eval_dataloader}
  train_eval_unit:
    _target_: fairchem.core.units.mlip_unit.mlip_unit.MLIPTrainEvalUnit
    job_config: ${job}
    tasks: ${data.tasks_list}
    model:
      _target_: fairchem.core.units.mlip_unit.mlip_unit.initialize_finetuning_model
      checkpoint_location:
        _target_: fairchem.core.calculate.pretrained_mlip.pretrained_checkpoint_path_from_name
        model_name: ${base_model_name}
      overrides:
        backbone:
          otf_graph: true
          max_neighbors: ${max_neighbors}
          regress_stress: ${data.regress_stress}
          always_use_pbc: false
        pass_through_head_outputs: ${data.pass_through_head_outputs}
      heads: ${data.heads}
    optimizer_fn:
      _target_: torch.optim.AdamW
      _partial_: true
      lr: ${lr}
      weight_decay: ${weight_decay}
    cosine_lr_scheduler_fn:
      _target_: fairchem.core.units.mlip_unit.mlip_unit._get_consine_lr_scheduler
      _partial_: true
      warmup_factor: 0.2
      warmup_epochs: 0.01
      lr_min_factor: 0.01
      epochs: ${epochs}
      steps: ${steps}
    print_every: 10
    clip_grad_norm: 100
  max_epochs: ${epochs}
  max_steps: ${steps}
  evaluate_every_n_steps: ${evaluate_every_n_steps}
  callbacks:
  - _target_: fairchem.core.components.train.train_runner.TrainCheckpointCallback
    checkpoint_every_n_steps: ${checkpoint_every_n_steps}
    max_saved_checkpoints: 5
  - _target_: torchtnt.framework.callbacks.TQDMProgressBar

My training data set contains some ice structures with energy and force computed using revPBE+D3(BJ).

Any advice is appreciated!

[misko]

Hi @Alexandrina-Chen ,
Sorry to miss this! What solved the first one? was there a mistake in configs somewhere?

For the second issue (NaNs) can you please try reverting this file to the commit, https://github.com/facebookresearch/fairchem/blob/main/src/fairchem/core/models/uma/common/rotation.py to the previous commit (4c6d01e) and try running again?

I am working on fixing this, hoping to have an update by Monday, the above might fix the issue for you

[Alexandrina-Chen]

I forgot to set base-model to esen-sm-conserving-all-oc25, after corrected the base-model and uma-task, it worked! It was my fault not carefully looking at the options.

For the second issue, I will try and let you know if it works.

Thank you for your help!