nvalchemi.training.FineTuningStrategy

pydantic model nvalchemi.training.FineTuningStrategy

Training strategy for patching modules and selecting trainable parameters.

FineTuningStrategy is intended for workflows where a pretrained model is loaded first and then adapted in-place before optimizer construction. The strategy keeps the base TrainingStrategy loop, but prepends registration-time hooks derived from its convenience fields before any explicit hooks= supplied by the user:

• module_patches becomes a ModulePatchHook.

• freeze_patterns / trainable_patterns become a TrainableParameterHook.

Module patch targets are fully-qualified paths of the form "<model_key>.<module_path>.<child>", for example "main.model.readouts.1.linear". The parent path must already exist. The final child is replaced when it is an existing torch.nn.Module or added when missing. Use nvalchemi.training.create_model_spec() for module patches that must round-trip through to_spec_dict(); direct torch.nn.Module instances are supported at runtime but are rejected by serialization.

Parameter patterns are matched against fully-qualified names such as "main.model.readouts.1.linear.weight". trainable_patterns alone is an allow-list: only matching parameters remain trainable and enter optimizers. When freeze_patterns is also supplied, matching parameters are excluded first, then trainable_patterns are re-included. With the default freeze_mode="requires_grad", excluded parameters are temporarily marked requires_grad=False during run() and restored afterward. Use freeze_mode="optimizer_only" when excluded parameters should still receive gradients but must not be updated by optimizers.

Examples

Replace a readout head, train only that head, and serialize the workflow by declaring the replacement as a BaseSpec:

import torch

from nvalchemi.training import (
    EnergyMSELoss,
    FineTuningStrategy,
    ForceMSELoss,
    OptimizerConfig,
    create_model_spec,
    default_training_fn,
)

strategy = FineTuningStrategy(
    models=pretrained_model,
    module_patches={
        "main.model.readouts.1.linear": create_model_spec(
            torch.nn.Linear,
            in_features=128,
            out_features=1,
        )
    },
    trainable_patterns=("main.model.readouts.1.linear.*",),
    freeze_mode="requires_grad",
    optimizer_configs=OptimizerConfig(
        optimizer_cls=torch.optim.AdamW,
        optimizer_kwargs={"lr": 1e-4},
    ),
    training_fn=default_training_fn,
    loss_fn=EnergyMSELoss() + ForceMSELoss(normalize_by_atom_count=True),
    num_epochs=10,
    devices=[torch.device("cuda")],
)

strategy.run(train_loader)

Use optimizer-only filtering when excluded parameters should still receive gradients but must not be updated:

strategy = FineTuningStrategy(
    models=pretrained_model,
    freeze_patterns=("main.model.*",),
    trainable_patterns=("main.model.readouts.*",),
    freeze_mode="optimizer_only",
    optimizer_configs=optimizer_config,
    training_fn=default_training_fn,
    loss_fn=loss_fn,
    num_steps=1000,
)

field module_patches: dict[str, BaseSpec | Module] [Optional]

Ordered module patches applied before optimizer construction.

field freeze_patterns: tuple[str, ...] = ()

Glob patterns excluded from training. Exclusions can be re-included by trainable_patterns.

field trainable_patterns: tuple[str, ...] = ()

Glob patterns included in the trainable parameter allow-list. When no freeze_patterns are supplied, this is the complete allow-list.

field freeze_mode: Literal['requires_grad', 'optimizer_only'] = 'requires_grad'

Whether excluded parameters are temporarily frozen via requires_grad=False or only excluded from optimizers. Defaults to "requires_grad".