MACEFittingParameters

Included in QATK.MLFF

class MACEFittingParameters(model_parameters=None, dataset_parameters=None, training_parameters=None, **kwargs)

A parameter class for passing settings to training MACE models [1][2]. The parameters are used for setting up from-scratch fitting or fine-tuning of MACE models to custom data.

Parameters:

Note

Argument changes:

number_of_channels

Deprecated since 2026.03

number_of_channels is deprecated. Use MACEModelParameters(number_of_channels=…) instead

max_l_equivariance

Deprecated since 2026.03

max_l_equivariance is deprecated. Use MACEModelParameters(max_l_equivariance=…) instead

distance_cutoff

Deprecated since 2026.03

distance_cutoff is deprecated. Use MACEModelParameters(distance_cutoff=…) instead

number_of_interactions

Deprecated since 2026.03

number_of_interactions is deprecated. Use MACEModelParameters(number_of_interactions=…) instead

correlation_order

Deprecated since 2026.03

correlation_order is deprecated. Use MACEModelParameters(correlation_order=…) instead

max_ell

Deprecated since 2026.03

max_ell is deprecated. Use MACEModelParameters(max_ell=…) instead

number_of_radial_basis_functions

Deprecated since 2026.03

number_of_radial_basis_functions is deprecated. Use MACEModelParameters(number_of_radial_basis_functions=…) instead

model_type

Deprecated since 2026.03

model_type is deprecated.

scaling_type

Deprecated since 2026.03

scaling_type is deprecated.

foundation_model_path

Deprecated since 2026.03

foundation_model_path is deprecated. Use MACEModelParameters(foundation_model_path=…) instead

use_multiheads_finetuning

Deprecated since 2026.03

use_multiheads_finetuning is deprecated. Use ForceFieldDatasetParameters(replay_finetuning_settings=MACEReplayFinetuningSettings(…)) instead

pretrained_head_train_file

Deprecated since 2026.03

pretrained_head_train_file is deprecated. Use ForceFieldDatasetParameters(replay_finetuning_settings=MACEReplayFinetuningSettings(replay_data_filepath=…)) instead

number_of_samples_from_pretrained_head

Deprecated since 2026.03

number_of_samples_from_pretrained_head is deprecated. Use ForceFieldDatasetParameters(replay_finetuning_settings=MACEReplayFinetuningSettings(number_of_samples=…)) instead

training_data_filepath

Deprecated since 2026.03

training_data_filepath is deprecated.

testing_data_filepath

Deprecated since 2026.03

testing_data_filepath is deprecated.

validation_fraction

Deprecated since 2026.03

validation_fraction is deprecated. Use ForceFieldDatasetParameters(validation_fraction=…) instead

isolated_atom_energies

Deprecated since 2026.03

isolated_atom_energies is deprecated. Use ForceFieldDatasetParameters(isolated_atom_energies=…) instead

energy_key

Deprecated since 2026.03

energy_key is deprecated. Use ForceFieldDatasetParameters(energy_key=…) instead

forces_key

Deprecated since 2026.03

forces_key is deprecated. Use ForceFieldDatasetParameters(forces_key=…) instead

stress_key

Deprecated since 2026.03

stress_key is deprecated. Use ForceFieldDatasetParameters(stress_key=…) instead

energy_weight

Deprecated since 2026.03

energy_weight is deprecated. Use ForceFieldDatasetParameters(energy_weight=…) instead

forces_weight

Deprecated since 2026.03

forces_weight is deprecated. Use ForceFieldDatasetParameters(forces_weight=…) instead

stress_weight

Deprecated since 2026.03

stress_weight is deprecated. Use ForceFieldDatasetParameters(stress_weight=…) instead

compute_stress

Deprecated since 2026.03

compute_stress is deprecated. Use ForceFieldDatasetParameters(compute_stress=…) instead

loss_function_type

Deprecated since 2026.03

loss_function_type is deprecated.

error_measure

Deprecated since 2026.03

error_measure is deprecated.

mp_data_path

Deprecated since 2026.03

mp_data_path is deprecated.

mp_descriptors_path

Deprecated since 2026.03

mp_descriptors_path is deprecated.

experiment_name

Deprecated since 2026.03

experiment_name is deprecated. Use TrainingParameters(experiment_name=…) instead

batch_size

Deprecated since 2026.03

batch_size is deprecated. Use TrainingParameters(batch_size=…) instead

validation_batch_size

Deprecated since 2026.03

validation_batch_size is deprecated.

max_number_of_epochs

Deprecated since 2026.03

max_number_of_epochs is deprecated. Use TrainingParameters(max_number_of_epochs=…) instead

patience

Deprecated since 2026.03

patience is deprecated. Use TrainingParameters(patience=…) instead

device

Deprecated since 2026.03

device is deprecated. Use TrainingParameters(device=…) instead

distributed_training

Deprecated since 2026.03

distributed_training is deprecated.

random_seed

Deprecated since 2026.03

random_seed is deprecated. Use TrainingParameters(random_seed=…) instead

keep_checkpoints

Deprecated since 2026.03

keep_checkpoints is deprecated.

number_of_workers

Deprecated since 2026.03

number_of_workers is deprecated. Use TrainingParameters(number_of_workers=…) instead

default_dtype

Deprecated since 2026.03

default_dtype is deprecated. Use TrainingParameters(default_dtype=…) instead

evaluation_interval

Deprecated since 2026.03

evaluation_interval is deprecated.

learning_rate

Deprecated since 2026.03

learning_rate is deprecated. Use TrainingParameters(learning_rate=…) instead

weight_decay

Deprecated since 2026.03

weight_decay is deprecated. Use TrainingParameters(weight_decay=…) instead

restart_from_last_checkpoint

Deprecated since 2026.03

restart_from_last_checkpoint is deprecated. Use TrainingParameters(restart_from_last_checkpoint=…) instead

use_exponential_moving_average

Deprecated since 2026.03

use_exponential_moving_average is deprecated.

exponential_moving_average_decay

Deprecated since 2026.03

exponential_moving_average_decay is deprecated.

scheduler_patience

Deprecated since 2026.03

scheduler_patience is deprecated. Use TrainingParameters(scheduler_patience=…) instead

gradient_clipping_threshold

Deprecated since 2026.03

gradient_clipping_threshold is deprecated. Use TrainingParameters(gradient_clipping_threshold=…) instead

save_for_cpu

Deprecated since 2026.03

save_for_cpu is deprecated and done automatically.

log_directory

Deprecated since 2026.03

log_directory is deprecated.

results_directory

Deprecated since 2026.03

results_directory is deprecated.

checkpoint_directory

Deprecated since 2026.03

checkpoint_directory is deprecated.

additional_parameters

Deprecated since 2026.03

additional_parameters is deprecated. Use TrainingParameters(additional_parameters=…) instead

datasetParameters()
Returns:

The dataset-specific parameters.

Return type:

ForceFieldDatasetParameters

modelParameters()
Returns:

The model-specific parameters.

Return type:

MACEModelParameters

nlinfo()
Returns:

The nlinfo.

Return type:

dict

trainingParameters()
Returns:

The training-specific parameters.

Return type:

TrainingParameters

uniqueString()

Return a unique string representing the state of the object.

Usage Examples

The MACEFittingParameters class is used for choosing settings for training MACE Machine Learned Force Field (MLFF) models [1][2] using the MachineLearnedForceFieldTrainer. Once trained, the model is automatically converted into a QATK-compatible format, and a TremoloXCalculator can be set up using the trained model.

There are three main usage scenarios of the MACEFittingParameters class.

  1. Training a MACE model from scratch.

This is the most straightforward way to train a MACE model where a model is trained only on data provided by the user. The model size and hyperparameters are fully customizable allowing users to define and utilize their own optimal tradeoff between desired model accuracy and performance. Shown in the example below is a simple training script that trains a MACE model from scratch highlighting some of the most important parameters.

model_parameters = MACEModelParameters(
    # Most important model size parameters (affects accuracy and speed)
    max_l_equivariance=0,
    number_of_channels=64,
    distance_cutoff=4.0 * Angstrom,
)

# Setup dataset-related parameters
dataset_parameters = ForceFieldDatasetParameters(
    # Weights for the different parts of the loss term
    energy_weight=1.0,
    forces_weight=100.0,
    stress_weight=1.0,
    # Ratio of training data used for validation
    validation_fraction=0.2,
    # Include stress if desired and present in the training data
    compute_stress=True,
)

# Setup training process parameters
training_parameters = TrainingParameters(
    # Name of the model - should always be set
    experiment_name="mace_experiment1",
    # Training hyperparameters
    max_number_of_epochs=200,  # Number of epochs to train for
    patience=50,  # Number of epochs without any improvement that will cause model training to finalize
    batch_size=4,  # Generally, higher means quicker training, but too high batch_sizes can cause GPU memory errors
    random_seed=123,  # Can be used for experimenting with the influence on different data splits
    default_dtype=MLParameterOptions.DTYPE.FLOAT64,  # Can be adjusted for shifting tradeoff between model accuracy and training speed
    device=Automatic,  # Device can be set explicitly but Automatic is recommended to automatically use the GPU if available
)

# Combine all parameters into MACEFittingParameters
mace_fp = MACEFittingParameters(
    model_parameters=model_parameters,
    dataset_parameters=dataset_parameters,
    training_parameters=training_parameters,
)

# Setup ML model training object
machine_learned_force_field_trainer = MachineLearnedForceFieldTrainer(
    fitting_parameters=mace_fp, training_sets=training_set, calculator=calculator
)

# Run the training
machine_learned_force_field_trainer.train()


# Load the model (which includes random_seed suffix) for inference use in QuantumATK
model_path = "mace_experiment1_123.qatkpt"

potentialSet = TremoloXPotentialSet(name="my_mace_potential1")
potential = TorchXPotential(
    dtype="float64",
    # Select the device automatically at runtime.
    device=None,
    file=model_path,
)
for symbol in potential.get_symbols(model_path):
    potentialSet.addParticleType(ParticleType(symbol))
potentialSet.addPotential(potential)

calculator = TremoloXCalculator(parameters=potentialSet)

# Alternatively, in same script as training, extract ready to use calculator
# calculator = machine_learned_force_field_trainer.fittedCalculator()

# Define configuration and attach the calculator
# configuration = ...
# configuration.setCalculator(calculator)
# ...
  1. Training with Naive Finetuning.

This is the simplest way to finetune a MACE model. The model is trained on a new dataset with model weights and hyperparameters initializing from the values of the pretrained model. Any hyperparameter settings related to the neural network architecture provided in the finetuning MACEFittingParameters object will be overwritten. This ensures that the model size and architecture remain consistent with the pretrained model.

# Setup model architecture parameters with foundation model for naive finetuning
model_parameters = MACEModelParameters(
    # When using naive finetuning, specify the foundation model path
    # Model architecture parameters will be inherited from the foundation model
    foundation_model_path="/path/to/model/data/on/run/machine/my_previous_mace_model.model",
)

# Setup dataset-related parameters
dataset_parameters = ForceFieldDatasetParameters(
    # Weights for the different parts of the loss term
    energy_weight=1.0,
    forces_weight=100.0,
    stress_weight=1.0,
    # Ratio of training data used for validation
    validation_fraction=0.2,
)

# Setup training process parameters
training_parameters = TrainingParameters(
    # Name of the model - should always be set
    experiment_name="mace_experiment2",
    # Training hyperparameters - keep max_number_of_epochs low for finetuning
    max_number_of_epochs=30,
    batch_size=4,  # Generally, higher means quicker training, but too high batch_sizes can cause GPU memory errors
    random_seed=123,  # Can be used for experimenting with the influence on different data splits
    default_dtype=MLParameterOptions.DTYPE.FLOAT64,  # Can be adjusted for shifting tradeoff between model accuracy and training speed
    device=Automatic,  # Device can be set explicitly but Automatic is recommended to automatically use the GPU if available
)

# Combine all parameters into MACEFittingParameters
mace_fp = MACEFittingParameters(
    model_parameters=model_parameters,
    dataset_parameters=dataset_parameters,
    training_parameters=training_parameters,
)

# Setup ML model training object
machine_learned_force_field_trainer = MachineLearnedForceFieldTrainer(
    fitting_parameters=mace_fp, training_sets=training_set, calculator=calculator
)

# Run the training
machine_learned_force_field_trainer.train()


# Load the model (which includes random_seed suffix) for inference use in QuantumATK
model_path = "mace_experiment2_123.qatkpt"

potentialSet = TremoloXPotentialSet(name="my_mace_potential2")
potential = TorchXPotential(
    dtype="float64",
    # Select the device automatically at runtime.
    device=None,
    file=model_path,
)
for symbol in potential.get_symbols(model_path):
    potentialSet.addParticleType(ParticleType(symbol))
potentialSet.addPotential(potential)

calculator = TremoloXCalculator(parameters=potentialSet)

# Alternatively, in same script as training, extract ready to use calculator
# calculator = machine_learned_force_field_trainer.fittedCalculator()

# Define configuration and attach the calculator
# configuration = ...
# configuration.setCalculator(calculator)
# ...
  1. Training with Replay Finetuning.

This is the most advanced way to finetune a MACE model. It is also the recommended approach for finetuning any materials project foundation model [3][4] as this usually leads to more robust and stable models. The model is trained on both a new dataset as well as (some of) the data used for training the original model in separate model heads. The model heads have different readout weights but share the remaining model weights to be optimized during the fitting. This ensures that fitting takes both data sets into account, learning from new data while retaining the knowledge of the original model as best as possible. The model weights and the hyperparameters initialize from the values of the pretrained model. Any hyperparameter settings related to the neural network architecture provided in the finetuning MACEFittingParameters object will be overwritten. This ensures that the model size and architecture remain consistent with the pretrained model.

# Setup replay finetuning settings for Materials Project foundation models
# Note: The replay data file (mp_traj_combined.xyz) must be downloaded beforehand
# from the mace-foundations repository
replay_settings = MACEReplayFinetuningSettings(
    # Path to the original training data used for the foundation model
    replay_data_filepath="/path/to/model/data/on/run/machine/mp_traj_combined.xyz",
    # Number of samples from the original model training data to use for replay finetuning
    number_of_samples=15000,
    # Method for subselecting samples from the replay data
    # Can use MLParameterOptions: MLParameterOptions.REPLAY_SUBSELECT.RANDOM or MLParameterOptions.REPLAY_SUBSELECT.FPS
    # Or string equivalents: 'random' or 'fps'
    subselect_method=MLParameterOptions.REPLAY_SUBSELECT.RANDOM,
    # Type of filtering to apply based on chemical elements
    # Can use MLParameterOptions: MLParameterOptions.REPLAY_FILTERING.NONE, .COMBINATIONS, .EXCLUSIVE, or .INCLUSIVE
    # Or string equivalents: 'none', 'combinations', 'exclusive', or 'inclusive'
    filtering_type=MLParameterOptions.REPLAY_FILTERING.COMBINATIONS,
)

# Setup model architecture parameters with foundation model for replay finetuning
model_parameters = MACEModelParameters(
    # Path to the foundation model - model architecture will be inherited
    foundation_model_path="/path/to/model/data/on/run/machine/mace-mpa-0-medium.model",
)

# Setup dataset-related parameters
dataset_parameters = ForceFieldDatasetParameters(
    # Weights for the different parts of the loss term
    energy_weight=1.0,
    forces_weight=100.0,
    stress_weight=1.0,
    # Ratio of training data used for validation
    validation_fraction=0.2,
    # Include replay finetuning settings for multihead training
    replay_finetuning_settings=replay_settings,
)

# Setup training process parameters
training_parameters = TrainingParameters(
    # Name of the model - has to be set
    experiment_name="mace_experiment3",
    # Training hyperparameters - keep max_number_of_epochs low for finetuning
    max_number_of_epochs=20,
    batch_size=4,  # Generally, higher means quicker training, but too high batch_sizes can cause GPU memory errors
    random_seed=123,  # Can be used for experimenting with the influence on different data splits
    default_dtype=MLParameterOptions.DTYPE.FLOAT64,  # Can be adjusted for shifting tradeoff between model accuracy and training speed
    device=Automatic,  # Device can be set explicitly but Automatic is recommended to automatically use the GPU if available
)

# Combine all parameters into MACEFittingParameters
mace_fp = MACEFittingParameters(
    model_parameters=model_parameters,
    dataset_parameters=dataset_parameters,
    training_parameters=training_parameters,
)

# Setup ML model training object
machine_learned_force_field_trainer = MachineLearnedForceFieldTrainer(
    fitting_parameters=mace_fp, training_sets=training_set, calculator=calculator
)

# Run the training
machine_learned_force_field_trainer.train()


# Load the model for evaluation use in QuantumATK
model_path = "mace_experiment3_123.qatkpt"

potentialSet = TremoloXPotentialSet(name="my_mace_potential3")
potential = TorchXPotential(
    dtype="float64",
    # Select the device automatically at runtime.
    device=None,
    file=model_path,
)
for symbol in potential.get_symbols(model_path):
    potentialSet.addParticleType(ParticleType(symbol))
potentialSet.addPotential(potential)

calculator = TremoloXCalculator(parameters=potentialSet)

# Alternatively, in same script as training, extract ready to use calculator
# calculator = machine_learned_force_field_trainer.fittedCalculator()

# Define configuration and attach the calculator
# configuration = ...
# configuration.setCalculator(calculator)
# ...

The example scripts, including a full training setup with MachineLearnedForceFieldTrainer and TremoloXCalculator for retrieving the trained models as calculators, are available for downloading:

  1. mace_example1.py

  2. mace_example2.py

  3. mace_example3.py

Notes

The MACEFittingParameters class provides an interface for training MACE MLFFs [1][2]. The framework provides a multitude of possible settings aiding the creation of accurate and performant MLFFs.

In order to use any MACE training functionality, it is assumed that at least one TrainingSet is available with precomputed energies, forces, and, if desired, stress. The calculator used for calculating the training data is similarly assumed to be either attached to the TrainingSet or to be defined in the calculator parameter of the MachineLearnedForceFieldTrainer class. When passing a DFT reference calculator, isolated atom energy reference values are automatically computed using spin polarized calculations.

Parameter Organization:

The MACEFittingParameters class is organized into three modular parameter groups for improved clarity and flexibility:

  • MACEModelParameters (MACEModelParameters): Groups all parameters related to the neural network architecture, including the number of embedding channels (number_of_channels), equivariance level (max_l_equivariance), radial cutoff distance (distance_cutoff), and the path to a foundation model for finetuning (foundation_model_path). These parameters define the model’s capacity and expressiveness. See MACEModelParameters for the complete list of available parameters.

  • ForceFieldDatasetParameters (ForceFieldDatasetParameters): Groups all parameters related to dataset handling and preprocessing, including the fraction of data for validation (validation_fraction), forces capping to remove outliers (forces_cap), and replay finetuning configuration (replay_finetuning_settings). These parameters control how training data is loaded, weighted, and processed. See ForceFieldDatasetParameters for the complete list of available parameters.

  • TrainingParameters (TrainingParameters): Groups all parameters related to the training process itself, including experiment naming (experiment_name), maximum training duration (max_number_of_epochs), early stopping criterion (patience), batch size (batch_size), reproducibility seed (random_seed), and numerical precision (default_dtype). These parameters control the optimization and training behavior. See TrainingParameters for the complete list of available parameters.

These parameter objects are passed to MACEFittingParameters to configure the complete training setup. This modular approach simplifies configuration while maintaining full control over all training aspects.

The MACEFittingParameters class is a simplified version of the complete set of MACE training parameters available in the open source package with more self-explanatory parameter namings and a reduced number of total parameters for simplifying general user interaction. MACE training parameters that have been left out of the MACEFittingParameters class, but that are available in the open source training framework, can, however, still be passed into the MACEFittingParameters object by supplying them as key-value pairs in a python dictionary to be inserted via the additional_parameters parameter of TrainingParameters.

Most of the parameters have default values corresponding to the defaults in the open source package, but some have been adapted in the MACEFittingParameters class for ease of use and more appropriate settings for most QuantumATK users. All applied default values are documented in the respective parameter class documentation.

Finetuning Approaches:

For naive finetuning, set the foundation_model_path parameter in MACEModelParameters to point to a pretrained .model file. The model will initialize with weights from the foundation model.

For replay finetuning (recommended for Materials Project foundation models), additionally provide a MACEReplayFinetuningSettings object via the replay_finetuning_settings parameter in ForceFieldDatasetParameters. The MACEReplayFinetuningSettings class configures how replay data from the foundation model’s original training set is incorporated during finetuning. The replay_subselect_method options are available as MLParameterOptions.REPLAY_SUBSELECT constants (or string values like 'random', 'fps'), and the replay_filtering_type options are available as MLParameterOptions.REPLAY_FILTERING constants (or string values like 'none', 'combinations', 'exclusive', 'inclusive'). This approach trains the model with separate heads for new and replay data, sharing most model weights while maintaining separate readout layers, which helps preserve the foundation model’s learned knowledge while adapting to new data.

The MP foundation models [3], their underlying training data [5], are available via the mace-foundations repo of ACESuit [4]. It should be noted, that the training data was calculated with other DFT codes and that total compatibility with QuantumATK-native calculators is not guaranteed. This is therefore a prime use case for replay finetuning, which supports heads using different DFT definitions. The total energies may differ, but derivative quantities such as forces and stress for similar levels of theory show signs of good agreement (for same level of theory). This means that in terms of absolute energies, pretrained models may have to somewhat relearn different energy scales, but the inter-atomic relationships learned by a foundation model are still a strong foundation for continued learning although the level/type of theory is updated. In practice, the different model heads will have separate weights for the readout layer, but otherwise share the remaining model weights.

GPU Acceleration and Model Formats:

GPU acceleration of MACE training using CUDA is supported. The default way to activate GPU acceleration is by checking the “Enable proprietary GPU acceleration” option in the Job Manager or by using the atkpython_gpu command instead of atkpython, as described in the technical notes.

When GPU acceleration is enabled on Linux machines, MACE training automatically utilizes NVIDIA cuEquivariance acceleration. This provides optimized operations during the training process, which can improve training performance on large training sets. The trained models are automatically saved in CuEquivariance-accelerated formats in addition to the standard format. For inference, e.g. during Molecular Dynamics simulations, CuEquivariance acceleration provides significant speedups whether the training has been performed with CuEquivariance acceleration or not.

An important consideration for CuEquivariance models is that the floating-point precision (data type) is ingrained into the model file and cannot be changed after creation. A CuEquivariance-accelerated MACE model is compiled for either floating-point 32 (fp32) or floating-point 64 (fp64) precision.

By default, when GPU training is enabled, multiple model versions are saved after training:

  • <experiment_name>_<random_seed>.qatkpt: The standard MACE model file in the format used during training

  • <experiment_name>_<random_seed>_cue_fp32.qatkpt: CuEquivariance-accelerated model with fp32 precision

  • <experiment_name>_<random_seed>_cue_fp64.qatkpt: CuEquivariance-accelerated model with fp64 precision

If you prefer to save only the model format used during training (to save disk space), set the save_all_available_model_formats parameter to False in the TrainingParameters class.

In addition to the .qatkpt format files that are used for inference in QuantumATK, an additional model file with the .model extension is also saved (<experiment_name>_<random_seed>.model). The .model format has special uses:

  • It can be used as input to the foundation_model_path parameter in MACEModelParameters for further finetuning of the model

  • It can be converted to any of the QuantumATK-compatible formats (.qatkpt, _cue_fp32.qatkpt, _cue_fp64.qatkpt) using the convertMACEModelToQATKFormat utility function

After training a model via the general training flow outlined above, the resulting model can be used directly with the TremoloXCalculator, as evident in the downloadable example scripts. In the Workflow Builder in Nanolab the trained model file can be loaded via the MachineLearnedForceFieldCalculator to set up a TremoloXCalculator.