Problem when finetuning with all heads

[SunXT-0719]

Hi fairchem team!
I'm on a research of finetuning uma-s-1p1 with all heads pretrained. I see the PR #1766 and delete the head part in my finetune config.
uma_sm_finetune_template.yaml
from:

    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}

to:

    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}

I do some basic debugs to make the programm run but the loss at step 0 is still abnormally high, which is:

INFO:root:{'train/loss': 8990.620638182878, 'train/lr': 1e-05, 'train/step': 0, 'train/epoch': 0.0, 'train/samples_per_second(approx)': 6.1664387292354395, 'train/atoms_per_second(approx)': 197.7114417561113, 'train/num_atoms_on_rank': 1026, 'train/num_samples_on_rank': 32}
/data/sunxuetin/anaconda3/envs/UMA/lib/python3.12/site-packages/torch/optim/lr_scheduler.py:332: UserWarning: To get the last learning rate computed by the scheduler, please use `get_last_lr()`.
  _warn_get_lr_called_within_step(self)
INFO:root:{'train/loss': 11867.852604975855, 'train/lr': 1e-05, 'train/step': 0, 'train/epoch': 0.0, 'train/samples_per_second(approx)': 6.168348108451524, 'train/atoms_per_second(approx)': 197.5799003488379, 'train/num_atoms_on_rank': 1025, 'train/num_samples_on_rank': 32}
/data/sunxuetin/anaconda3/envs/UMA/lib/python3.12/site-packages/torch/optim/lr_scheduler.py:332: UserWarning: To get the last learning rate computed by the scheduler, please use `get_last_lr()`.
  _warn_get_lr_called_within_step(self)
INFO:root:{'train/loss': 10284.808734129449, 'train/lr': 1e-05, 'train/step': 0, 'train/epoch': 0.0, 'train/samples_per_second(approx)': 6.165504733162675, 'train/atoms_per_second(approx)': 197.68149550702827, 'train/num_atoms_on_rank': 1026, 'train/num_samples_on_rank': 32}
/data/sunxuetin/anaconda3/envs/UMA/lib/python3.12/site-packages/torch/optim/lr_scheduler.py:332: UserWarning: To get the last learning rate computed by the scheduler, please use `get_last_lr()`.
  _warn_get_lr_called_within_step(self)
INFO:root:{'train/loss': 10899.674122548213, 'train/lr': 1e-05, 'train/step': 0, 'train/epoch': 0.0, 'train/samples_per_second(approx)': 6.131760968243467, 'train/atoms_per_second(approx)': 195.83311592327573, 'train/num_atoms_on_rank': 1022, 'train/num_samples_on_rank': 32}

which is same as the loss using re-initialized heads.

Do I need to do anything else when editing the config to make sure the heads are succefully loaded?
Thanks for your reply!

[wood-b]

Hi @SunXT-0719,

can you send your entire updated config? hard to know what going on otherwise

[SunXT-0719]

Hi @SunXT-0719,

can you send your entire updated config? hard to know what going on otherwise

Sure, here is the entire uma_sm_finetune_template.yaml

defaults:
- data: uma_conserving_data_task_energy_force
- _self_
job:
  device_type: CUDA
  scheduler:
    mode: LOCAL
    ranks_per_node: 4
    num_nodes: 1
  debug: false
  run_dir: /data/sunxuetin/tmp/uma_finetune_runs/none_sketch_finetune
  run_name: none_sketch_finetune
  logger:
    _target_: fairchem.core.common.logger.WandBSingletonLogger.init_wandb
    _partial_: true
    entity: SUNXUETING
    project: uma_finetune
base_model_name: uma-m-1p1
max_neighbors: 300
epochs: 1
steps: null
batch_size: 32
lr: 5e-5
weight_decay: 1e-3
evaluate_every_n_steps: 10000
checkpoint_every_n_steps: 2000
train_dataset:
  _target_: fairchem.core.datasets.mt_concat_dataset.create_concat_dataset
  dataset_configs:
    finetune_dataset: ${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: false
    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:
    finetune_dataset: ${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}
    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: 500
    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

[wood-b]

nothing immediate sticks out, we typically use LR 4e-4 for the finetuning but we also have much larger global batch sizes. You can play around with the inital LR. Also, maybe a good idea to finetune UMA-S instead of UMA-M until you get things figured out. If you send the data file with tasks I can also look at that. Does the loss decrease rapidly or stays very high?

[SunXT-0719]

the entire data file is:

dataset_name: omol
elem_refs:
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- -2.082805460035017
- 0.0
- 0.0
- 0.0
- -3.256404920200298
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
- 0.0
normalizer_rmsd: 1.6831237759602915
train_dataset:
  splits:
    train:
      src: /data/sunxuetin/projects/fairchem/datasets/neutral_train
  format: ase_db
  transforms:
    common_transform:
      dataset_name: ${data.dataset_name}
    stress_reshape_transform:
      dataset_name: ${data.dataset_name}
val_dataset:
  splits:
    val:
      src: /data/sunxuetin/projects/fairchem/packages/fairchem-core/src/fairchem/data/om/neutral_val
  format: ase_db
  transforms:
    common_transform:
      dataset_name: ${data.dataset_name}
    stress_reshape_transform:
      dataset_name: ${data.dataset_name}
regress_stress: false
pass_through_head_outputs: true
heads:
  efs:
    module: fairchem.core.models.uma.escn_md.MLP_EFS_Head
tasks_list:
- _target_: fairchem.core.units.mlip_unit.mlip_unit.Task
  name: energy
  level: system
  property: energy
  loss_fn:
    _target_: fairchem.core.modules.loss.DDPMTLoss
    loss_fn:
      _target_: fairchem.core.modules.loss.PerAtomMAELoss
    coefficient: 20
  out_spec:
    dim:
    - 1
    dtype: float32
  normalizer:
    _target_: fairchem.core.modules.normalization.normalizer.Normalizer
    mean: 0.0
    rmsd: ${data.normalizer_rmsd}
  element_references:
    _target_: fairchem.core.modules.normalization.element_references.ElementReferences
    element_references:
      _target_: torch.DoubleTensor
      _args_:
      - ${data.elem_refs}
  datasets:
  - ${data.dataset_name}
  metrics:
  - mae
  - per_atom_mae

the loss decrease at the begining but not very much, and it comes to about 100 after one epoch of omol25 neutral-train dataset, does it make sense?
Besides, the MAE(energy) of inference result is everything alright, which is 0.01 on omol neutral-val dataset, so I'm wondering is there any difference of heads between inference and finetuning?

[wood-b]

Did you try fine-tuning the OMol task without changing the refs/normalizer?

[SunXT-0719]

Did you try fine-tuning the OMol task without changing the refs/normalizer?

sorrt, I didn't get it, what does "changing refs/normalizer?" mean?