Generate and persist UMA model IDs

Author: lbluque

CLAUDE.md

@@ -332,6 +332,11 @@ configs/                 # Hydra YAML configs (datasets, tasks, backbone, optimi
   not depend on enumeration order, seeded random coordinates, or atom counts
   unless those values are part of the public contract; prefer composition,
   Miller index, shift, placement, and cell invariants that survive upgrades.
+- Do not make `HydraModel.model_id` unconditionally required. Older non-UMA
+  Hydra checkpoints are intentionally untagged, while new UMA MoE configs get
+  a generated ID. Generate it on rank zero after distributed setup, broadcast
+  it to every rank, and explicitly persist it in the resume/DCP model config;
+  setting only the model attribute does not make it part of the checkpoint.
 
 Anytime we learn something that could be beneficial in future coding sessions, automatically add it to CLAUDE.md.
 

configs/uma/benchmark/perf_check/training_inner.yaml

@@ -54,6 +54,7 @@ runner:
     tasks: ${tasks}
     model:
       _target_: fairchem.core.models.base.HydraModel
+      model_id: UMA-S-1.2
       backbone: ${backbone}
       heads: ${heads}
       pass_through_head_outputs: True

configs/uma/benchmark/perf_check/training_inner_2x.yaml

@@ -54,6 +54,7 @@ runner:
     tasks: ${tasks}
     model:
       _target_: fairchem.core.models.base.HydraModel
+      model_id: UMA-S-1.2
       backbone: ${backbone}
       heads: ${heads}
       pass_through_head_outputs: True

src/fairchem/core/models/base.py

@@ -21,6 +21,7 @@
 from fairchem.core.common.utils import (
     load_model_and_weights_from_checkpoint,
 )
+from fairchem.core.models.uma.compat import ensure_uma_model_id
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -328,6 +329,11 @@ def __init__(
         self._tasks = None
         self._dataset_to_tasks = None
 
+        model_config = {"backbone": backbone, "model_id": model_id}
+        generated_model_id = ensure_uma_model_id(model_config)
+        if generated_model_id is not None:
+            model_id = generated_model_id
+
         # Does this model support inference on single atom systems
         self.supports_single_atoms = supports_single_atoms
         # model_id string in form of NAME-VERSION e.g. UMA-1.2

src/fairchem/core/models/uma/compat.py

@@ -16,7 +16,12 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Literal
+import logging
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Literal, MutableMapping
+from uuid import uuid4
+
+from fairchem.core.common import distutils
 
 if TYPE_CHECKING:
     from fairchem.core.units.mlip_unit.api.inference import MLIPInferenceCheckpoint
@@ -37,7 +42,68 @@
 _UMA_BACKBONE_FQN_SUFFIX = "uma.escn_moe.eSCNMDMoeBackbone"
 
 
-def get_uma_version(model_config: dict | None) -> UmaVersion:
+def is_uma_moe_backbone_config(backbone_config: Mapping | None) -> bool:
+    """
+    Return whether a backbone config describes an UMA MoE model.
+
+    The UMA MoE backbone is also shared by models such as eSEN with
+    ``num_experts == 0``. Those models do not have model-ID-gated behavior and
+    therefore do not require a ``model_id``.
+
+    Args:
+        backbone_config: Backbone configuration to classify.
+
+    Returns:
+        Whether the configuration describes an UMA backbone with experts.
+    """
+    if not isinstance(backbone_config, Mapping):
+        return False
+
+    model = backbone_config.get("model")
+    if not isinstance(model, str) or not (
+        model == _UMA_BACKBONE_SHORT_NAME or model.endswith(_UMA_BACKBONE_FQN_SUFFIX)
+    ):
+        return False
+
+    num_experts = backbone_config.get("num_experts")
+    return isinstance(num_experts, int) and num_experts > 0
+
+
+def ensure_uma_model_id(model_config: MutableMapping) -> str | None:
+    """
+    Add a generated ID to an untagged UMA MoE model config.
+
+    Existing IDs are preserved, and non-UMA model configs are unchanged.
+
+    Args:
+        model_config: Model configuration to update in place.
+
+    Returns:
+        The existing or generated UMA model ID, or ``None`` for non-UMA models.
+    """
+    if not is_uma_moe_backbone_config(model_config.get("backbone")):
+        return None
+
+    model_id = model_config.get("model_id")
+    if isinstance(model_id, str) and model_id.strip():
+        return model_id
+
+    model_id = f"UMA-{uuid4().hex[:12]}" if distutils.is_master() else None
+    model_id_list = [model_id]
+    distutils.broadcast_object_list(model_id_list, src=0)
+    model_id = model_id_list[0]
+    if not isinstance(model_id, str):
+        raise RuntimeError("Failed to broadcast the generated UMA model_id")
+    model_config["model_id"] = model_id
+    if distutils.is_master():
+        logging.warning(
+            "No model_id was provided for an UMA MoE model. Generated model_id=%r.",
+            model_id,
+        )
+    return model_id
+
+
+def get_uma_version(model_config: Mapping | None) -> UmaVersion:
     """Classify what fix-up a checkpoint needs (see :func:`apply_uma_compat_fixups`).
 
     * ``"not_uma"`` — not a UMA MoE backbone. This includes non-UMA models and
@@ -50,20 +116,10 @@ def get_uma_version(model_config: dict | None) -> UmaVersion:
     * ``"tagged"`` — already has a ``model_id`` (UMA 1.2+ or custom) → no-op. The
       1.2 ``include_self`` rule lives in the backbone, keyed on ``model_id``.
     """
-    if not isinstance(model_config, dict):
+    if not isinstance(model_config, Mapping):
         return "not_uma"
     backbone = model_config.get("backbone", {})
-    if not isinstance(backbone, dict):
-        return "not_uma"
-    model = backbone.get("model")
-    if not isinstance(model, str) or not (
-        model == _UMA_BACKBONE_SHORT_NAME or model.endswith(_UMA_BACKBONE_FQN_SUFFIX)
-    ):
-        return "not_uma"
-
-    # UMA uses num_experts > 0; eSCNMDMoeBackbone with num_experts == 0
-    # (e.g. eSEN) is not UMA.
-    if not isinstance(backbone.get("num_experts"), int) or backbone["num_experts"] == 0:
+    if not is_uma_moe_backbone_config(backbone):
         return "not_uma"
 
     model_id = model_config.get("model_id")

src/fairchem/core/units/mlip_unit/mlip_unit.py

@@ -12,12 +12,12 @@
 import time
 from copy import deepcopy
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Optional, Sequence
+from typing import Any, Optional, Sequence
 
 import numpy as np
 import torch
 import torch.distributed.checkpoint as dcp
-from omegaconf import OmegaConf
+from omegaconf import DictConfig, OmegaConf, open_dict
 from torch.distributed.checkpoint.format_utils import dcp_to_torch_save
 from torch.distributed.checkpoint.state_dict import (
     get_model_state_dict,
@@ -62,12 +62,45 @@
     tf32_context_manager,
 )
 
-if TYPE_CHECKING:
-    from omegaconf import DictConfig
-
 # this is a config generated on the fly and can be used to resume a run for a given checkpoint
 UNIT_RESUME_CONFIG = "resume.yaml"
 
+
+def _get_train_eval_unit_config(config: dict | DictConfig):
+    """
+    Return the train/eval unit config from direct or Ray-wrapped runner config.
+
+    Args:
+        config: Canonical job configuration.
+
+    Returns:
+        The train/eval unit configuration.
+    """
+    runner_config = config["runner"]
+    if "train_eval_unit" not in runner_config:
+        runner_config = runner_config["runner_config"]
+    return runner_config["train_eval_unit"]
+
+
+def _set_model_id_in_config(config: dict | DictConfig, model_id: str | None) -> None:
+    """
+    Persist a resolved model ID in a canonical training configuration.
+
+    Args:
+        config: Canonical job configuration to update.
+        model_id: Resolved model ID, or ``None`` for models without one.
+    """
+    if model_id is None:
+        return
+
+    model_config = _get_train_eval_unit_config(config)["model"]
+    if isinstance(model_config, DictConfig):
+        with open_dict(model_config):
+            model_config["model_id"] = model_id
+    else:
+        model_config["model_id"] = model_id
+
+
 # this represents the inference only checkpoint generated at each checkpoint
 UNIT_INFERENCE_CHECKPOINT = "inference_ckpt.pt"
 
@@ -126,7 +159,7 @@ def convert_train_checkpoint_to_inference_checkpoint(
     inference_ckpt = torch.load(
         checkpoint_loc, map_location="cpu", weights_only=False
     )  # DCP model config
-    train_eval_unit_state = inference_ckpt["config"]["runner"]["train_eval_unit"]
+    train_eval_unit_state = _get_train_eval_unit_config(inference_ckpt["config"])
     unit_state = inference_ckpt["unit_state"]
     torch.save(
         MLIPInferenceCheckpoint(
@@ -535,6 +568,7 @@ def __init__(
         self.finetune_model_full_config = getattr(
             model, "finetune_model_full_config", None
         )
+        self.model_id = getattr(model, "model_id", None)
         # call optimizer function between wrapping in DDP
         # this is required for models that have a no_weight_decay function
         self.optimizer = _get_optimizer_wd(optimizer_fn, model)
@@ -888,7 +922,10 @@ def save_state(self, checkpoint_location: str) -> None:
 
         finetune_model_full_config = self.get_finetune_model_config()
         if finetune_model_full_config is not None:
-            config.runner.train_eval_unit.model = finetune_model_full_config
+            train_eval_unit_config = _get_train_eval_unit_config(config)
+            train_eval_unit_config["model"] = finetune_model_full_config
+
+        _set_model_id_in_config(config, self.model_id)
 
         OmegaConf.save(config, os.path.join(checkpoint_location, UNIT_RESUME_CONFIG))
 

tests/core/models/test_base.py

@@ -0,0 +1,36 @@
+"""
+Copyright (c) Meta Platforms, Inc. and affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+
+import pytest
+
+from fairchem.core.common.registry import registry
+from fairchem.core.models.base import HydraModel
+
+
+@pytest.mark.parametrize("model_id", [None, "", "   "])
+def test_uma_moe_hydra_model_generates_model_id(model_id, caplog, monkeypatch):
+    backbone = {
+        "model": "fairchem.core.models.uma.escn_moe.eSCNMDMoeBackbone",
+        "num_experts": 8,
+    }
+
+    class DummyBackbone:
+        def __init__(self, **kwargs):
+            pass
+
+    monkeypatch.setattr(registry, "get_model_class", lambda _: DummyBackbone)
+    with caplog.at_level(logging.WARNING):
+        model = HydraModel(backbone=backbone, heads={}, model_id=model_id)
+
+    assert re.fullmatch(r"UMA-[0-9a-f]{12}", model.model_id)
+    assert model.backbone.model_id == model.model_id
+    assert f"Generated model_id='{model.model_id}'" in caplog.text

tests/core/models/uma/test_compat.py

@@ -5,11 +5,15 @@
 from __future__ import annotations
 
 import pytest
+from omegaconf import OmegaConf
 
+from fairchem.core.common import distutils
 from fairchem.core.models.uma.compat import (
     UMA_1P1_MODEL_ID,
     apply_uma_compat_fixups,
+    ensure_uma_model_id,
     get_uma_version,
+    is_uma_moe_backbone_config,
 )
 from fairchem.core.units.mlip_unit.api.inference import MLIPInferenceCheckpoint
 
@@ -40,6 +44,48 @@ def uma_cfg(
     return cfg
 
 
+@pytest.mark.parametrize("backbone_model", [UMA_BACKBONE_FQN, UMA_BACKBONE_SHORT])
+def test_uma_moe_backbone_config(backbone_model):
+    assert is_uma_moe_backbone_config({"model": backbone_model, "num_experts": 8})
+
+
+def test_uma_moe_dict_config():
+    config = OmegaConf.create(uma_cfg(model_id="UMA-S-1.2"))
+
+    assert get_uma_version(config) == "tagged"
+    assert is_uma_moe_backbone_config(config.backbone)
+
+
+def test_existing_uma_model_id_is_preserved():
+    config = uma_cfg(model_id="UMA-S-custom")
+
+    assert ensure_uma_model_id(config) == "UMA-S-custom"
+    assert config["model_id"] == "UMA-S-custom"
+
+
+def test_generated_uma_model_id_is_broadcast(monkeypatch):
+    config = uma_cfg()
+
+    monkeypatch.setattr(distutils, "is_master", lambda: False)
+
+    def broadcast_model_id(model_id_list, src):
+        assert model_id_list == [None]
+        assert src == 0
+        model_id_list[0] = "UMA-from-rank-zero"
+
+    monkeypatch.setattr(distutils, "broadcast_object_list", broadcast_model_id)
+
+    assert ensure_uma_model_id(config) == "UMA-from-rank-zero"
+    assert config["model_id"] == "UMA-from-rank-zero"
+
+
+@pytest.mark.parametrize("num_experts", [0, -1, None])
+def test_uma_non_moe_backbone_config(num_experts):
+    assert not is_uma_moe_backbone_config(
+        {"model": UMA_BACKBONE_FQN, "num_experts": num_experts}
+    )
+
+
 # ---------------------------------------------------------------------------
 # UMA 1.0 / untagged (neither model_id nor model_version) — hard fail
 # ---------------------------------------------------------------------------

tests/core/units/mlip_unit/test_checkpoint_config.py

@@ -0,0 +1,29 @@
+"""
+Copyright (c) Meta Platforms, Inc. and affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from __future__ import annotations
+
+import pytest
+from omegaconf import OmegaConf
+
+from fairchem.core.units.mlip_unit.mlip_unit import (
+    _get_train_eval_unit_config,
+    _set_model_id_in_config,
+)
+
+
+@pytest.mark.parametrize("ray_wrapped", [False, True])
+def test_set_model_id_in_checkpoint_config(ray_wrapped):
+    train_runner = {"train_eval_unit": {"model": {"_target_": "model"}}}
+    runner = {"runner_config": train_runner} if ray_wrapped else train_runner
+    config = OmegaConf.create({"runner": runner})
+    OmegaConf.set_struct(config, True)
+
+    _set_model_id_in_config(config, "UMA-generated")
+
+    train_eval_unit_config = _get_train_eval_unit_config(config)
+    assert train_eval_unit_config.model.model_id == "UMA-generated"