LoadCheckpoint

steerability.algorithms.structural_control.load_checkpoint

args

LoadCheckpoint argument validation.

control

LoadCheckpoint: a structural control that installs a checkpoint as the pipeline model.

logger = logging.getLogger(__name__) module-attribute

LoadCheckpoint

Bases: StructuralControl

Installs a saved full-weights checkpoint as the pipeline model.

On the in-process backend, steer() loads the checkpoint with AutoModelForCausalLM.from_pretrained and returns it, replacing the incoming model for subsequent controls. On vLLM backends the checkpoint is served directly (Capability.SERVE_CHECKPOINT), exported as a CheckpointArtifact. The frozen form of a trained structural control (fine-tune, merge) is an instance of this control pointing at the trained checkpoint.

Source code in steerability/algorithms/structural_control/load_checkpoint/control.py

class LoadCheckpoint(StructuralControl):
    """Installs a saved full-weights checkpoint as the pipeline model.

    On the in-process backend, `steer()` loads the checkpoint with
    `AutoModelForCausalLM.from_pretrained` and returns it, replacing the incoming model for
    subsequent controls. On vLLM backends the checkpoint is served directly
    (`Capability.SERVE_CHECKPOINT`), exported as a `CheckpointArtifact`. The frozen form of a
    trained structural control (fine-tune, merge) is an instance of this control pointing at
    the trained checkpoint.
    """

    Args = LoadCheckpointArgs
    supports_batching = True

    def artifact_capability(self) -> Capability:
        """`Capability.SERVE_CHECKPOINT`; the configuration is an on-disk checkpoint."""
        return Capability.SERVE_CHECKPOINT

    def export_artifact(self) -> CheckpointArtifact:
        """The configured checkpoint directory."""
        return CheckpointArtifact(path=str(self.path))

    def steer(
        self,
        model: PreTrainedModel | None,
        tokenizer: PreTrainedTokenizerBase | None = None,
        **kwargs,
    ) -> PreTrainedModel:
        """Load the checkpoint and return it as the new pipeline model.

        Args:
            model: The incoming pipeline model (unused; the checkpoint replaces it).
            tokenizer: The pipeline tokenizer (unused).

        Returns:
            The loaded model.
        """
        logger.info("Loading checkpoint %s.", self.path)
        return AutoModelForCausalLM.from_pretrained(
            self.path,
            device_map=self.device_map,
            trust_remote_code=self.trust_remote_code,
            **self.hf_model_kwargs,
        )

RUNTIME_KWARGS_SCHEMA = [] class-attribute instance-attribute

args = self.Args.validate(*args, **kwargs) instance-attribute

enabled = True class-attribute instance-attribute

supports_batching = True class-attribute instance-attribute

__init__(*args, **kwargs)

Source code in steerability/algorithms/core/base_control.py

def __init__(self, *args, **kwargs) -> None:
    if self.Args is None:  # null / arg-free control
        if args or kwargs:
            raise TypeError(f"{type(self).__name__} accepts no constructor arguments.")
        self._configure()
        return

    self.args: BaseArgs = self.Args.validate(*args, **kwargs)

    # move fields to attributes, skipping any name the control exposes as a property
    # (e.g. CAST.condition_point); the raw value stays reachable via self.args.<name>
    for field in fields(self.args):
        if isinstance(getattr(type(self), field.name, None), property):
            continue
        setattr(self, field.name, getattr(self.args, field.name))

    self._configure()

_configure()

Post-construction hook, called after Args fields are mirrored onto the instance.

Driver presets override this to map their mirrored args onto the fields their generic base reads, so subclasses never bypass a parent __init__. Default no-op.

Source code in steerability/algorithms/core/base_control.py

def _configure(self) -> None:
    """Post-construction hook, called after `Args` fields are mirrored onto the instance.

    Driver presets override this to map their mirrored args onto the fields their generic base
    reads, so subclasses never bypass a parent `__init__`. Default no-op.
    """
    pass

artifact_capability()

Capability.SERVE_CHECKPOINT; the configuration is an on-disk checkpoint.

Source code in steerability/algorithms/structural_control/load_checkpoint/control.py

def artifact_capability(self) -> Capability:
    """`Capability.SERVE_CHECKPOINT`; the configuration is an on-disk checkpoint."""
    return Capability.SERVE_CHECKPOINT

cleanup()

Release resources allocated during steer().

Override in subclasses that allocate GPU memory or other resources during steering. Default no-op.

Source code in steerability/algorithms/core/base_control.py

def cleanup(self) -> None:
    """Release resources allocated during `steer()`.

    Override in subclasses that allocate GPU memory or other resources during steering. Default
    no-op.
    """
    pass

clone_for_call(seed=None)

A configuration-preserving shallow clone for one generation call.

The clone shares steer-time artifacts (memories, steering vectors, attached tokenizers) with the original but has its own attribute namespace, so per-call attribute mutation on the clone never races another call using the original. When seed is given and the control defines reseed(seed), the clone's client-side RNG is re-seeded.

Parameters:

Name	Type	Description	Default
seed	int | None	Optional seed forwarded to the clone's reseed().	None

Returns:

Type	Description
BaseControl	The clone.

Source code in steerability/algorithms/core/base_control.py

def clone_for_call(self, seed: int | None = None) -> "BaseControl":
    """A configuration-preserving shallow clone for one generation call.

    The clone shares steer-time artifacts (memories, steering vectors, attached tokenizers)
    with the original but has its own attribute namespace, so per-call attribute mutation on
    the clone never races another call using the original. When `seed` is given and the
    control defines `reseed(seed)`, the clone's client-side RNG is re-seeded.

    Args:
        seed: Optional seed forwarded to the clone's `reseed()`.

    Returns:
        The clone.
    """
    clone = copy.copy(self)
    if seed is not None:
        reseed = getattr(clone, "reseed", None)
        if callable(reseed):
            reseed(seed)
    return clone

export_artifact()

The configured checkpoint directory.

Source code in steerability/algorithms/structural_control/load_checkpoint/control.py

def export_artifact(self) -> CheckpointArtifact:
    """The configured checkpoint directory."""
    return CheckpointArtifact(path=str(self.path))

export_state()

Steer-time state to persist when freezing, keyed by logical name.

Values are typed artifacts (SteeringVector, Probe, ProbeSet, a Memory implementation, CheckpointArtifact, LoRAArtifact, torch.Tensor, or a Path to an on-disk product). Must be called after steer(). The default returns an empty mapping (nothing to persist).

Returns:

Type	Description
dict[str, Any]	Mapping from logical name to artifact value.

Source code in steerability/algorithms/core/base_control.py

def export_state(self) -> dict[str, Any]:
    """Steer-time state to persist when freezing, keyed by logical name.

    Values are typed artifacts (`SteeringVector`, `Probe`, `ProbeSet`, a `Memory`
    implementation, `CheckpointArtifact`, `LoRAArtifact`, `torch.Tensor`, or a `Path` to an
    on-disk product). Must be called after `steer()`. The default returns an empty mapping
    (nothing to persist).

    Returns:
        Mapping from logical name to artifact value.
    """
    return {}

fit_identity()

The object whose canonical form defines this control's fit digest, or None.

The digest detects recipe edits that invalidate frozen artifacts, so the returned object should cover exactly the inputs a re-fit would consume (training data, fit specs) and exclude inert application parameters (strengths, multipliers). The default returns None (no fits).

Returns:

Type	Description
Any | None	The fit-identity object, or None.

Source code in steerability/algorithms/core/base_control.py

def fit_identity(self) -> Any | None:
    """The object whose canonical form defines this control's fit digest, or None.

    The digest detects recipe edits that invalidate frozen artifacts, so the returned
    object should cover exactly the inputs a re-fit would consume (training data, fit
    specs) and exclude inert application parameters (strengths, multipliers). The default
    returns None (no fits).

    Returns:
        The fit-identity object, or None.
    """
    return None

frozen_form(state)

The (registry method key, constructor kwargs) of this control's frozen form.

state is the output of export_state(); artifact values inside the returned kwargs are replaced by store references at encode time. When steer_fits() is empty and state is empty, the recipe is the frozen form and the control's own key and recipe args are returned unchanged. A control whose freeze expands into several entries may return a list of (method key, kwargs) pairs instead.

Parameters:

Name	Type	Description	Default
state	dict[str, Any]	The output of export_state().	required

Returns:

Type	Description
tuple[str, dict[str, Any]]	The frozen method key and its constructor kwargs.

Raises:

Type	Description
NotFreezableError	If the control produces fits or state but declares no frozen form.

Source code in steerability/algorithms/core/base_control.py

def frozen_form(self, state: dict[str, Any]) -> tuple[str, dict[str, Any]]:
    """The `(registry method key, constructor kwargs)` of this control's frozen form.

    `state` is the output of `export_state()`; artifact values inside the returned kwargs
    are replaced by store references at encode time. When `steer_fits()` is empty and
    `state` is empty, the recipe is the frozen form and the control's own key and recipe
    args are returned unchanged. A control whose freeze expands into several entries may
    return a list of `(method key, kwargs)` pairs instead.

    Args:
        state: The output of `export_state()`.

    Returns:
        The frozen method key and its constructor kwargs.

    Raises:
        NotFreezableError: If the control produces fits or state but declares no frozen
            form.
    """
    if not state and not self.steer_fits():
        from steerability.algorithms.core.registry import method_key_for

        args = getattr(self, "args", None)
        kwargs = {field.name: getattr(args, field.name) for field in fields(args) if field.init} \
            if args is not None else {}
        return method_key_for(type(self)), kwargs
    raise NotFreezableError(
        f"{type(self).__name__} produces steer-time state but declares no frozen form; "
        "override frozen_form() (and export_state()) to support freezing."
    )

requirements()

Backend requirements computed from this instance's configuration, per phase.

The generate phase requires Capability.IN_PROCESS_TORCH for in-process adoption of the returned model; when the configuration produces an on-disk artifact (artifact_capability()), serving that artifact is an alternative, so a backend advertising the matching serve capability also supports the generate phase.

Returns:

Type	Description
Requirements	The control's phase-keyed requirements.

Source code in steerability/algorithms/structural_control/base.py

def requirements(self) -> Requirements:
    """Backend requirements computed from this instance's configuration, per phase.

    The generate phase requires `Capability.IN_PROCESS_TORCH` for in-process adoption of
    the returned model; when the configuration produces an on-disk artifact
    (`artifact_capability()`), serving that artifact is an alternative, so a backend
    advertising the matching serve capability also supports the generate phase.

    Returns:
        The control's phase-keyed requirements.
    """
    generate = needs(Capability.IN_PROCESS_TORCH)
    capability = self.artifact_capability()
    if capability is not None:
        generate = any_of(
            generate,
            needs(capability, hint="serve the steer-time artifact on a vLLM backend"),
        )
    return Requirements(generate=generate)

steer(model, tokenizer=None, **kwargs)

Load the checkpoint and return it as the new pipeline model.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	The incoming pipeline model (unused; the checkpoint replaces it).	required
tokenizer	PreTrainedTokenizerBase | None	The pipeline tokenizer (unused).	None

Returns:

Type	Description
PreTrainedModel	The loaded model.

Source code in steerability/algorithms/structural_control/load_checkpoint/control.py

def steer(
    self,
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase | None = None,
    **kwargs,
) -> PreTrainedModel:
    """Load the checkpoint and return it as the new pipeline model.

    Args:
        model: The incoming pipeline model (unused; the checkpoint replaces it).
        tokenizer: The pipeline tokenizer (unused).

    Returns:
        The loaded model.
    """
    logger.info("Loading checkpoint %s.", self.path)
    return AutoModelForCausalLM.from_pretrained(
        self.path,
        device_map=self.device_map,
        trust_remote_code=self.trust_remote_code,
        **self.hf_model_kwargs,
    )

steer_access()

ModelAccess.MODULE; training happens on live weights.

Source code in steerability/algorithms/structural_control/base.py

def steer_access(self) -> ModelAccess:
    """`ModelAccess.MODULE`; training happens on live weights."""
    return ModelAccess.MODULE

steer_fits()

The fit artifacts this instance's steer step will produce, for the steer plan.

Each entry is (artifact, artifact_class), where artifact is the fit source or recipe class name and artifact_class is "direction" or "calibrated". The default is an empty tuple (no fits).

Returns:

Type	Description
tuple[tuple[str, str], ...]	The declared fit artifacts, in declaration order.

Source code in steerability/algorithms/core/base_control.py

def steer_fits(self) -> tuple[tuple[str, str], ...]:
    """The fit artifacts this instance's steer step will produce, for the steer plan.

    Each entry is `(artifact, artifact_class)`, where `artifact` is the fit source or
    recipe class name and `artifact_class` is `"direction"` or `"calibrated"`. The default
    is an empty tuple (no fits).

    Returns:
        The declared fit artifacts, in declaration order.
    """
    return ()