Skip to content

TRL

steerability.algorithms.structural_control.wrappers.trl

The TRL wrapper implements a variety of methods from Hugging Face's TRL library.

The current functionality spans the following methods:

  • SFT (Supervised Fine-Tuning): Standard supervised learning to fine-tune language models on demonstration data
  • DPO (Direct Preference Optimization): Trains models directly on preference data without requiring a separate reward model
  • APO (Anchored Preference Optimization): A variant of DPO that uses an anchor model to improve training stability and performance
  • PPO (Proximal Policy Optimization): Reinforcement learning against a sequence-classification reward model (and value model)
  • GRPO (Group-Relative Policy Optimization): Critic-free reinforcement learning that takes a callable reward function (no reward model and no value model); used by PRewrite for its metric-in-the-loop reward

For documentation information, please refer to the TRL page.

args

base_mixin

TRLMixin

Small shared helpers for TRL-based structural controls.

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
class TRLMixin:
    """
    Small shared helpers for TRL-based structural controls.
    """

    # populated from Args by subclasses
    base_model_name_or_path: str | None = None
    tokenizer_name_or_path: str | None = None
    hf_model_kwargs: dict[str, Any] = {}
    trust_remote_code: bool = False

    training_args: dict[str, Any] = {}
    output_dir: str | None = None
    resume_from_checkpoint: str | None = None

    use_peft: bool = False
    peft_type: Any = None
    lora_kwargs: dict[str, Any] = {}
    adapter_name: str | None = None

    merge_lora_after_train: bool = False
    merged_output_dir: str | None = None

    # resolved at runtime
    tokenizer: PreTrainedTokenizerBase | None = None
    device = None
    _resolved_base_ref: str | None = None

    def _resolve_model_tokenizer(
        self,
        model: PreTrainedModel | None,
        tokenizer: PreTrainedTokenizerBase | None,
    ) -> tuple[PreTrainedModel, PreTrainedTokenizerBase]:
        """Resolve the model and tokenizer, returning both as locals.

        Loads the model from `base_model_name_or_path` when `model` is None, and the tokenizer from
        `tokenizer_name_or_path`, the model's `name_or_path`, or `base_model_name_or_path` when
        `tokenizer` is None. Records the base model reference on `self._resolved_base_ref` (for
        `export_artifact`) and the model's device on `self.device`. The model is not stored on the
        instance; it is returned for the caller to thread.

        Returns:
            The resolved `(model, tokenizer)` pair.

        Raises:
            ValueError: If `model` is None and `base_model_name_or_path` is unset, or the tokenizer
                path cannot be resolved.
        """
        if model is None:
            if not self.base_model_name_or_path:
                raise ValueError("TRLMixin: model is None and `base_model_name_or_path` was not provided.")
            model = AutoModelForCausalLM.from_pretrained(
                self.base_model_name_or_path,
                trust_remote_code=self.trust_remote_code,
                **(self.hf_model_kwargs or {}),
            )

        self._resolved_base_ref = self.base_model_name_or_path or getattr(model, "name_or_path", None)

        if tokenizer is None:
            path = (
                self.tokenizer_name_or_path
                or getattr(model, "name_or_path", None)
                or self.base_model_name_or_path
            )
            if not path:
                raise ValueError("TRLMixin: could not resolve tokenizer path.")
            self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=self.trust_remote_code)
        else:
            self.tokenizer = tokenizer

        self.device = next(model.parameters()).device
        return model, self.tokenizer

    def _post_train_freeze(self, model: PreTrainedModel) -> PreTrainedModel:
        """Put `model` in eval mode, freeze its parameters, and return it.

        The model is passed in and returned; it is not stored on the instance.
        """
        model.eval()
        for parameter in model.parameters():
            parameter.requires_grad_(False)
        return model

    def _maybe_save_trained_artifacts(self, trainer) -> None:
        output_dir = self.training_args.get("output_dir") or self.output_dir
        if output_dir:
            trainer.save_model(output_dir)
            try:
                self.tokenizer.save_pretrained(output_dir)
            except Exception:
                pass

    def _resolved_output_dir(self) -> str | None:
        return self.training_args.get("output_dir") or self.output_dir

    def artifact_capability(self) -> Capability | None:
        """The serve capability implied by this training configuration.

        Training runs only when `train_dataset` is set, so a configuration without one produces
        no artifact. A LoRA run without merge-back saves an adapter to the output directory
        (`Capability.SERVE_LORA`); a full fine-tune saves a checkpoint there
        (`Capability.SERVE_CHECKPOINT`); a merged LoRA run yields a checkpoint only when
        `merged_output_dir` is set.
        """
        if getattr(self, "train_dataset", None) is None:
            return None
        is_lora = bool(self.use_peft) and self.peft_type == PeftType.LORA
        if is_lora and not self.merge_lora_after_train:
            return Capability.SERVE_LORA if self._resolved_output_dir() else None
        if is_lora and self.merge_lora_after_train:
            return Capability.SERVE_CHECKPOINT if self.merged_output_dir else None
        return Capability.SERVE_CHECKPOINT if self._resolved_output_dir() else None

    def export_artifact(self) -> Artifact | None:
        """The on-disk product of this configuration's `steer()`, matching
        `artifact_capability()`."""
        capability = self.artifact_capability()
        if capability is None:
            return None
        if capability == Capability.SERVE_LORA:
            base = self.base_model_name_or_path or self._resolved_base_ref or ""
            return LoRAArtifact(path=str(self._resolved_output_dir()), base_model=str(base))
        is_lora = bool(self.use_peft) and self.peft_type == PeftType.LORA
        path = self.merged_output_dir if (is_lora and self.merge_lora_after_train) else self._resolved_output_dir()
        return CheckpointArtifact(path=str(path))

    def export_state(self) -> dict[str, Any]:
        """The trained on-disk product under the `"artifact"` key, for freezing.

        Returns an empty mapping when the configuration trains nothing (`train_dataset` is
        None), so an inert wrapper's recipe is its frozen form.

        Raises:
            NotFreezableError: If the configuration trains but produces no on-disk product
                (a merged LoRA run without `merged_output_dir`).
        """
        if getattr(self, "train_dataset", None) is None:
            return {}
        artifact = self.export_artifact()
        if artifact is None:
            raise NotFreezableError(
                f"{type(self).__name__} trains but writes no on-disk product; set output_dir "
                "(or merged_output_dir for a merged LoRA run) to make the result freezable."
            )
        return {"artifact": artifact}

    def frozen_form(self, state: dict[str, Any]) -> tuple[str, dict[str, Any]]:
        """The trained artifact as a `load_lora` or `load_checkpoint` entry."""
        artifact = state["artifact"]
        if isinstance(artifact, LoRAArtifact):
            return "structural_control/load_lora", {
                "path": artifact,
                "base_model": artifact.base_model,
                "merge": False,
            }
        return "structural_control/load_checkpoint", {"path": artifact}

    def fit_identity(self) -> Any | None:
        """The training-relevant args projection, or None when nothing trains.

        Output locations (`output_dir`, `merged_output_dir`, `resume_from_checkpoint`, and
        `training_args["output_dir"]`) are excluded, since where a product is saved does not
        change what was trained.
        """
        if getattr(self, "train_dataset", None) is None:
            return None
        args = getattr(self, "args", None)
        if args is None:
            return None
        excluded = {"output_dir", "merged_output_dir", "resume_from_checkpoint"}
        payload: dict[str, Any] = {}
        for f in fields(args):
            if not f.init or f.name in excluded:
                continue
            value = getattr(args, f.name)
            if f.name == "training_args" and isinstance(value, dict):
                value = {k: v for k, v in value.items() if k not in excluded}
            payload[f.name] = value
        return payload

    def _maybe_merge_lora_in_place(self, model: PreTrainedModel) -> PreTrainedModel:
        """Optionally merge LoRA into the base weights, returning the (possibly merged) model.

        When `use_peft` and `merge_lora_after_train` are set and `model` exposes `merge_and_unload`,
        merges the adapter, refreshes `self.device` from the merged model, and saves the merged model
        and tokenizer to `merged_output_dir` when set. Returns the merged model, or `model` unchanged
        otherwise. The model is not stored on the instance.
        """
        if not (self.use_peft and self.merge_lora_after_train):
            return model

        # trainer often returns a PEFT-wrapped model; merge if possible
        if hasattr(model, "merge_and_unload"):
            model = model.merge_and_unload()
            self.device = next(model.parameters()).device

            # save if requested
            if self.merged_output_dir:
                model.save_pretrained(self.merged_output_dir)
                try:
                    self.tokenizer.save_pretrained(self.merged_output_dir)
                except Exception:
                    pass

        return model
adapter_name = None class-attribute instance-attribute
base_model_name_or_path = None class-attribute instance-attribute
device = None class-attribute instance-attribute
hf_model_kwargs = {} class-attribute instance-attribute
lora_kwargs = {} class-attribute instance-attribute
merge_lora_after_train = False class-attribute instance-attribute
merged_output_dir = None class-attribute instance-attribute
output_dir = None class-attribute instance-attribute
peft_type = None class-attribute instance-attribute
resume_from_checkpoint = None class-attribute instance-attribute
tokenizer = None class-attribute instance-attribute
tokenizer_name_or_path = None class-attribute instance-attribute
training_args = {} class-attribute instance-attribute
trust_remote_code = False class-attribute instance-attribute
use_peft = False class-attribute instance-attribute
artifact_capability()

The serve capability implied by this training configuration.

Training runs only when train_dataset is set, so a configuration without one produces no artifact. A LoRA run without merge-back saves an adapter to the output directory (Capability.SERVE_LORA); a full fine-tune saves a checkpoint there (Capability.SERVE_CHECKPOINT); a merged LoRA run yields a checkpoint only when merged_output_dir is set.

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
def artifact_capability(self) -> Capability | None:
    """The serve capability implied by this training configuration.

    Training runs only when `train_dataset` is set, so a configuration without one produces
    no artifact. A LoRA run without merge-back saves an adapter to the output directory
    (`Capability.SERVE_LORA`); a full fine-tune saves a checkpoint there
    (`Capability.SERVE_CHECKPOINT`); a merged LoRA run yields a checkpoint only when
    `merged_output_dir` is set.
    """
    if getattr(self, "train_dataset", None) is None:
        return None
    is_lora = bool(self.use_peft) and self.peft_type == PeftType.LORA
    if is_lora and not self.merge_lora_after_train:
        return Capability.SERVE_LORA if self._resolved_output_dir() else None
    if is_lora and self.merge_lora_after_train:
        return Capability.SERVE_CHECKPOINT if self.merged_output_dir else None
    return Capability.SERVE_CHECKPOINT if self._resolved_output_dir() else None
export_artifact()

The on-disk product of this configuration's steer(), matching artifact_capability().

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
162
163
164
165
166
167
168
169
170
171
172
173
def export_artifact(self) -> Artifact | None:
    """The on-disk product of this configuration's `steer()`, matching
    `artifact_capability()`."""
    capability = self.artifact_capability()
    if capability is None:
        return None
    if capability == Capability.SERVE_LORA:
        base = self.base_model_name_or_path or self._resolved_base_ref or ""
        return LoRAArtifact(path=str(self._resolved_output_dir()), base_model=str(base))
    is_lora = bool(self.use_peft) and self.peft_type == PeftType.LORA
    path = self.merged_output_dir if (is_lora and self.merge_lora_after_train) else self._resolved_output_dir()
    return CheckpointArtifact(path=str(path))
export_state()

The trained on-disk product under the "artifact" key, for freezing.

Returns an empty mapping when the configuration trains nothing (train_dataset is None), so an inert wrapper's recipe is its frozen form.

Raises:

Type Description
NotFreezableError

If the configuration trains but produces no on-disk product (a merged LoRA run without merged_output_dir).
Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def export_state(self) -> dict[str, Any]:
    """The trained on-disk product under the `"artifact"` key, for freezing.

    Returns an empty mapping when the configuration trains nothing (`train_dataset` is
    None), so an inert wrapper's recipe is its frozen form.

    Raises:
        NotFreezableError: If the configuration trains but produces no on-disk product
            (a merged LoRA run without `merged_output_dir`).
    """
    if getattr(self, "train_dataset", None) is None:
        return {}
    artifact = self.export_artifact()
    if artifact is None:
        raise NotFreezableError(
            f"{type(self).__name__} trains but writes no on-disk product; set output_dir "
            "(or merged_output_dir for a merged LoRA run) to make the result freezable."
        )
    return {"artifact": artifact}
fit_identity()

The training-relevant args projection, or None when nothing trains.

Output locations (output_dir, merged_output_dir, resume_from_checkpoint, and training_args["output_dir"]) are excluded, since where a product is saved does not change what was trained.

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
def fit_identity(self) -> Any | None:
    """The training-relevant args projection, or None when nothing trains.

    Output locations (`output_dir`, `merged_output_dir`, `resume_from_checkpoint`, and
    `training_args["output_dir"]`) are excluded, since where a product is saved does not
    change what was trained.
    """
    if getattr(self, "train_dataset", None) is None:
        return None
    args = getattr(self, "args", None)
    if args is None:
        return None
    excluded = {"output_dir", "merged_output_dir", "resume_from_checkpoint"}
    payload: dict[str, Any] = {}
    for f in fields(args):
        if not f.init or f.name in excluded:
            continue
        value = getattr(args, f.name)
        if f.name == "training_args" and isinstance(value, dict):
            value = {k: v for k, v in value.items() if k not in excluded}
        payload[f.name] = value
    return payload
frozen_form(state)

The trained artifact as a load_lora or load_checkpoint entry.

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
195
196
197
198
199
200
201
202
203
204
def frozen_form(self, state: dict[str, Any]) -> tuple[str, dict[str, Any]]:
    """The trained artifact as a `load_lora` or `load_checkpoint` entry."""
    artifact = state["artifact"]
    if isinstance(artifact, LoRAArtifact):
        return "structural_control/load_lora", {
            "path": artifact,
            "base_model": artifact.base_model,
            "merge": False,
        }
    return "structural_control/load_checkpoint", {"path": artifact}

resolve_config_kwargs(config_cls, training_args)

The kwargs to construct config_cls from training_args.

Every key must be a dataclass field of config_cls; None values are dropped so the config's own default applies. The toolkit forwards training_args verbatim to TRL, so a key the config does not declare would otherwise be passed to a constructor that does not accept it.

Parameters:

Name Type Description Default
config_cls type

A TRL config dataclass (DPOConfig, SFTConfig, GRPOConfig, PPOConfig).

 required
training_args dict[str, Any]

The composed training arguments.

 required

Returns:

Type Description
dict[str, Any]

The accepted kwargs, in training_args order, without None values.

Raises:

Type Description
ValueError

If any key is not a field of config_cls. The message names the config class, the installed TRL version, and the unknown keys, with a closest-field suggestion when one is close.
Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def resolve_config_kwargs(config_cls: type, training_args: dict[str, Any]) -> dict[str, Any]:
    """The kwargs to construct `config_cls` from `training_args`.

    Every key must be a dataclass field of `config_cls`; `None` values are dropped so the config's
    own default applies. The toolkit forwards `training_args` verbatim to TRL, so a key the config
    does not declare would otherwise be passed to a constructor that does not accept it.

    Args:
        config_cls: A TRL config dataclass (`DPOConfig`, `SFTConfig`, `GRPOConfig`, `PPOConfig`).
        training_args: The composed training arguments.

    Returns:
        The accepted kwargs, in `training_args` order, without `None` values.

    Raises:
        ValueError: If any key is not a field of `config_cls`. The message names the config class,
            the installed TRL version, and the unknown keys, with a closest-field suggestion when
            one is close.
    """
    allowed = {field.name for field in fields(config_cls)}
    unknown = sorted(key for key in training_args if key not in allowed)
    if unknown:
        details = []
        for key in unknown:
            close = difflib.get_close_matches(key, allowed, n=1)
            details.append(f"{key!r} (did you mean {close[0]!r}?)" if close else repr(key))
        raise ValueError(
            f"{config_cls.__name__} (trl {trl.__version__}) does not accept training_args key(s): "
            f"{', '.join(details)}. training_args is forwarded verbatim to TRL; remove these keys or "
            f"use a field the installed {config_cls.__name__} declares."
        )
    return {key: value for key, value in training_args.items() if value is not None}

grpotrainer

args

base_mixin

GRPOTrainerMixin

Bases: TRLMixin, StructuralControl

GRPO structural control backed by TRL's GRPOTrainer.

GRPO is critic-free (group-relative advantage), so there is no reward model and no value model. The reward comes from one or more callables in reward_funcs, each called as reward_func(prompts, completions, **kwargs) and returning one float per completion. The trainer reads a text "prompt" column directly and samples num_generations completions per prompt, so the dataset is only standardized to a prompt column and is not tokenized into an input_ids column.

Source code in steerability/algorithms/structural_control/wrappers/trl/grpotrainer/base_mixin.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
class GRPOTrainerMixin(TRLMixin, StructuralControl):
    """GRPO structural control backed by TRL's `GRPOTrainer`.

    GRPO is critic-free (group-relative advantage), so there is no reward model and no value model.
    The reward comes from one or more callables in `reward_funcs`, each called as
    `reward_func(prompts, completions, **kwargs)` and returning one float per completion. The trainer
    reads a text `"prompt"` column directly and samples `num_generations` completions per prompt, so
    the dataset is only standardized to a prompt column and is not tokenized into an `input_ids`
    column.
    """

    train_dataset: Any | None = None
    eval_dataset: Any | None = None
    reward_funcs: Callable | list[Callable] | None = None

    def steer(
        self,
        model: PreTrainedModel | None,
        tokenizer: PreTrainedTokenizerBase | None = None,
        **_,
    ) -> torch.nn.Module:
        self.tokenizer = tokenizer or (getattr(model, "tokenizer", None) if model is not None else None)
        model, self.tokenizer = self._resolve_model_tokenizer(model, self.tokenizer)
        self.tokenizer = ensure_pad_token(self.tokenizer)

        if self.reward_funcs is None:
            raise ValueError(
                "GRPOTrainerMixin.steer: `reward_funcs` is required (a callable or list of callables)."
            )
        reward_funcs = self.reward_funcs if isinstance(self.reward_funcs, list) else [self.reward_funcs]

        # GRPO reads the text "prompt" column directly; standardize but do NOT tokenize to input_ids
        train_dataset = standardize_prompt_dataset(self.train_dataset) if self.train_dataset is not None else None
        eval_dataset = standardize_prompt_dataset(self.eval_dataset) if self.eval_dataset is not None else None

        config_kwargs = resolve_config_kwargs(GRPOConfig, self.training_args)
        training_config = GRPOConfig(**config_kwargs)

        peft_config = None
        if self.use_peft and self.peft_type == PeftType.LORA:
            peft_config = LoraConfig(**self.lora_kwargs)

        if train_dataset is not None:
            trainer = GRPOTrainer(
                model=model,
                reward_funcs=reward_funcs,
                args=training_config,
                train_dataset=train_dataset,
                eval_dataset=eval_dataset,
                processing_class=self.tokenizer,
                peft_config=peft_config,
            )
            trainer.train(resume_from_checkpoint=self.training_args.get("resume_from_checkpoint"))

            # recover the trained policy so it can be used for generation (GRPO has no .policy wrapper)
            trained_model = trainer.accelerator.unwrap_model(trainer.model)
            model = getattr(trained_model, "policy", trained_model)
            self._maybe_save_trained_artifacts(trainer)
            model = self._maybe_merge_lora_in_place(model)

        return self._post_train_freeze(model)
RUNTIME_KWARGS_SCHEMA = [] class-attribute instance-attribute
adapter_name = None class-attribute instance-attribute
args = self.Args.validate(*args, **kwargs) instance-attribute
base_model_name_or_path = None class-attribute instance-attribute
device = None class-attribute instance-attribute
enabled = True class-attribute instance-attribute
eval_dataset = None class-attribute instance-attribute
hf_model_kwargs = {} class-attribute instance-attribute
lora_kwargs = {} class-attribute instance-attribute
merge_lora_after_train = False class-attribute instance-attribute
merged_output_dir = None class-attribute instance-attribute
output_dir = None class-attribute instance-attribute
peft_type = None class-attribute instance-attribute
resume_from_checkpoint = None class-attribute instance-attribute
reward_funcs = None class-attribute instance-attribute
supports_batching = True class-attribute instance-attribute
tokenizer = None class-attribute instance-attribute
tokenizer_name_or_path = None class-attribute instance-attribute
train_dataset = None class-attribute instance-attribute
training_args = {} class-attribute instance-attribute
trust_remote_code = False class-attribute instance-attribute
use_peft = False class-attribute instance-attribute
artifact_capability()

The serve capability implied by this training configuration.

Training runs only when train_dataset is set, so a configuration without one produces no artifact. A LoRA run without merge-back saves an adapter to the output directory (Capability.SERVE_LORA); a full fine-tune saves a checkpoint there (Capability.SERVE_CHECKPOINT); a merged LoRA run yields a checkpoint only when merged_output_dir is set.

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
def artifact_capability(self) -> Capability | None:
    """The serve capability implied by this training configuration.

    Training runs only when `train_dataset` is set, so a configuration without one produces
    no artifact. A LoRA run without merge-back saves an adapter to the output directory
    (`Capability.SERVE_LORA`); a full fine-tune saves a checkpoint there
    (`Capability.SERVE_CHECKPOINT`); a merged LoRA run yields a checkpoint only when
    `merged_output_dir` is set.
    """
    if getattr(self, "train_dataset", None) is None:
        return None
    is_lora = bool(self.use_peft) and self.peft_type == PeftType.LORA
    if is_lora and not self.merge_lora_after_train:
        return Capability.SERVE_LORA if self._resolved_output_dir() else None
    if is_lora and self.merge_lora_after_train:
        return Capability.SERVE_CHECKPOINT if self.merged_output_dir else None
    return Capability.SERVE_CHECKPOINT if self._resolved_output_dir() else None
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 
195
196
197
198
199
200
201
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 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
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 on-disk product of this configuration's steer(), matching artifact_capability().

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
162
163
164
165
166
167
168
169
170
171
172
173
def export_artifact(self) -> Artifact | None:
    """The on-disk product of this configuration's `steer()`, matching
    `artifact_capability()`."""
    capability = self.artifact_capability()
    if capability is None:
        return None
    if capability == Capability.SERVE_LORA:
        base = self.base_model_name_or_path or self._resolved_base_ref or ""
        return LoRAArtifact(path=str(self._resolved_output_dir()), base_model=str(base))
    is_lora = bool(self.use_peft) and self.peft_type == PeftType.LORA
    path = self.merged_output_dir if (is_lora and self.merge_lora_after_train) else self._resolved_output_dir()
    return CheckpointArtifact(path=str(path))
export_state()

The trained on-disk product under the "artifact" key, for freezing.

Returns an empty mapping when the configuration trains nothing (train_dataset is None), so an inert wrapper's recipe is its frozen form.

Raises:

Type Description
NotFreezableError

If the configuration trains but produces no on-disk product (a merged LoRA run without merged_output_dir).
Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def export_state(self) -> dict[str, Any]:
    """The trained on-disk product under the `"artifact"` key, for freezing.

    Returns an empty mapping when the configuration trains nothing (`train_dataset` is
    None), so an inert wrapper's recipe is its frozen form.

    Raises:
        NotFreezableError: If the configuration trains but produces no on-disk product
            (a merged LoRA run without `merged_output_dir`).
    """
    if getattr(self, "train_dataset", None) is None:
        return {}
    artifact = self.export_artifact()
    if artifact is None:
        raise NotFreezableError(
            f"{type(self).__name__} trains but writes no on-disk product; set output_dir "
            "(or merged_output_dir for a merged LoRA run) to make the result freezable."
        )
    return {"artifact": artifact}
fit_identity()

The training-relevant args projection, or None when nothing trains.

Output locations (output_dir, merged_output_dir, resume_from_checkpoint, and training_args["output_dir"]) are excluded, since where a product is saved does not change what was trained.

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
def fit_identity(self) -> Any | None:
    """The training-relevant args projection, or None when nothing trains.

    Output locations (`output_dir`, `merged_output_dir`, `resume_from_checkpoint`, and
    `training_args["output_dir"]`) are excluded, since where a product is saved does not
    change what was trained.
    """
    if getattr(self, "train_dataset", None) is None:
        return None
    args = getattr(self, "args", None)
    if args is None:
        return None
    excluded = {"output_dir", "merged_output_dir", "resume_from_checkpoint"}
    payload: dict[str, Any] = {}
    for f in fields(args):
        if not f.init or f.name in excluded:
            continue
        value = getattr(args, f.name)
        if f.name == "training_args" and isinstance(value, dict):
            value = {k: v for k, v in value.items() if k not in excluded}
        payload[f.name] = value
    return payload
frozen_form(state)

The trained artifact as a load_lora or load_checkpoint entry.

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
195
196
197
198
199
200
201
202
203
204
def frozen_form(self, state: dict[str, Any]) -> tuple[str, dict[str, Any]]:
    """The trained artifact as a `load_lora` or `load_checkpoint` entry."""
    artifact = state["artifact"]
    if isinstance(artifact, LoRAArtifact):
        return "structural_control/load_lora", {
            "path": artifact,
            "base_model": artifact.base_model,
            "merge": False,
        }
    return "structural_control/load_checkpoint", {"path": artifact}
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 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
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, **_)

Required steering/preparation.

session is a SteeringSession on the steering backend, provided by the pipeline.

Source code in steerability/algorithms/structural_control/wrappers/trl/grpotrainer/base_mixin.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
def steer(
    self,
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase | None = None,
    **_,
) -> torch.nn.Module:
    self.tokenizer = tokenizer or (getattr(model, "tokenizer", None) if model is not None else None)
    model, self.tokenizer = self._resolve_model_tokenizer(model, self.tokenizer)
    self.tokenizer = ensure_pad_token(self.tokenizer)

    if self.reward_funcs is None:
        raise ValueError(
            "GRPOTrainerMixin.steer: `reward_funcs` is required (a callable or list of callables)."
        )
    reward_funcs = self.reward_funcs if isinstance(self.reward_funcs, list) else [self.reward_funcs]

    # GRPO reads the text "prompt" column directly; standardize but do NOT tokenize to input_ids
    train_dataset = standardize_prompt_dataset(self.train_dataset) if self.train_dataset is not None else None
    eval_dataset = standardize_prompt_dataset(self.eval_dataset) if self.eval_dataset is not None else None

    config_kwargs = resolve_config_kwargs(GRPOConfig, self.training_args)
    training_config = GRPOConfig(**config_kwargs)

    peft_config = None
    if self.use_peft and self.peft_type == PeftType.LORA:
        peft_config = LoraConfig(**self.lora_kwargs)

    if train_dataset is not None:
        trainer = GRPOTrainer(
            model=model,
            reward_funcs=reward_funcs,
            args=training_config,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            processing_class=self.tokenizer,
            peft_config=peft_config,
        )
        trainer.train(resume_from_checkpoint=self.training_args.get("resume_from_checkpoint"))

        # recover the trained policy so it can be used for generation (GRPO has no .policy wrapper)
        trained_model = trainer.accelerator.unwrap_model(trainer.model)
        model = getattr(trained_model, "policy", trained_model)
        self._maybe_save_trained_artifacts(trainer)
        model = self._maybe_merge_lora_in_place(model)

    return self._post_train_freeze(model)
steer_access()

ModelAccess.MODULE; training happens on live weights.

Source code in steerability/algorithms/structural_control/base.py 
112
113
114
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 
105
106
107
108
109
110
111
112
113
114
115
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 ()

control

GRPO

Bases: GRPOTrainerMixin

GRPO controller.

Source code in steerability/algorithms/structural_control/wrappers/trl/grpotrainer/control.py 
5
6
7
8
9
class GRPO(GRPOTrainerMixin):
    """
    GRPO controller.
    """
    Args = GRPOArgs
RUNTIME_KWARGS_SCHEMA = [] class-attribute instance-attribute
adapter_name = None class-attribute instance-attribute
args = self.Args.validate(*args, **kwargs) instance-attribute
base_model_name_or_path = None class-attribute instance-attribute
device = None class-attribute instance-attribute
enabled = True class-attribute instance-attribute
eval_dataset = None class-attribute instance-attribute
hf_model_kwargs = {} class-attribute instance-attribute
lora_kwargs = {} class-attribute instance-attribute
merge_lora_after_train = False class-attribute instance-attribute
merged_output_dir = None class-attribute instance-attribute
output_dir = None class-attribute instance-attribute
peft_type = None class-attribute instance-attribute
resume_from_checkpoint = None class-attribute instance-attribute
reward_funcs = None class-attribute instance-attribute
supports_batching = True class-attribute instance-attribute
tokenizer = None class-attribute instance-attribute
tokenizer_name_or_path = None class-attribute instance-attribute
train_dataset = None class-attribute instance-attribute
training_args = {} class-attribute instance-attribute
trust_remote_code = False class-attribute instance-attribute
use_peft = False class-attribute instance-attribute
artifact_capability()

The serve capability implied by this training configuration.

Training runs only when train_dataset is set, so a configuration without one produces no artifact. A LoRA run without merge-back saves an adapter to the output directory (Capability.SERVE_LORA); a full fine-tune saves a checkpoint there (Capability.SERVE_CHECKPOINT); a merged LoRA run yields a checkpoint only when merged_output_dir is set.

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
def artifact_capability(self) -> Capability | None:
    """The serve capability implied by this training configuration.

    Training runs only when `train_dataset` is set, so a configuration without one produces
    no artifact. A LoRA run without merge-back saves an adapter to the output directory
    (`Capability.SERVE_LORA`); a full fine-tune saves a checkpoint there
    (`Capability.SERVE_CHECKPOINT`); a merged LoRA run yields a checkpoint only when
    `merged_output_dir` is set.
    """
    if getattr(self, "train_dataset", None) is None:
        return None
    is_lora = bool(self.use_peft) and self.peft_type == PeftType.LORA
    if is_lora and not self.merge_lora_after_train:
        return Capability.SERVE_LORA if self._resolved_output_dir() else None
    if is_lora and self.merge_lora_after_train:
        return Capability.SERVE_CHECKPOINT if self.merged_output_dir else None
    return Capability.SERVE_CHECKPOINT if self._resolved_output_dir() else None
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 
195
196
197
198
199
200
201
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 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
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 on-disk product of this configuration's steer(), matching artifact_capability().

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
162
163
164
165
166
167
168
169
170
171
172
173
def export_artifact(self) -> Artifact | None:
    """The on-disk product of this configuration's `steer()`, matching
    `artifact_capability()`."""
    capability = self.artifact_capability()
    if capability is None:
        return None
    if capability == Capability.SERVE_LORA:
        base = self.base_model_name_or_path or self._resolved_base_ref or ""
        return LoRAArtifact(path=str(self._resolved_output_dir()), base_model=str(base))
    is_lora = bool(self.use_peft) and self.peft_type == PeftType.LORA
    path = self.merged_output_dir if (is_lora and self.merge_lora_after_train) else self._resolved_output_dir()
    return CheckpointArtifact(path=str(path))
export_state()

The trained on-disk product under the "artifact" key, for freezing.

Returns an empty mapping when the configuration trains nothing (train_dataset is None), so an inert wrapper's recipe is its frozen form.

Raises:

Type Description
NotFreezableError

If the configuration trains but produces no on-disk product (a merged LoRA run without merged_output_dir).
Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def export_state(self) -> dict[str, Any]:
    """The trained on-disk product under the `"artifact"` key, for freezing.

    Returns an empty mapping when the configuration trains nothing (`train_dataset` is
    None), so an inert wrapper's recipe is its frozen form.

    Raises:
        NotFreezableError: If the configuration trains but produces no on-disk product
            (a merged LoRA run without `merged_output_dir`).
    """
    if getattr(self, "train_dataset", None) is None:
        return {}
    artifact = self.export_artifact()
    if artifact is None:
        raise NotFreezableError(
            f"{type(self).__name__} trains but writes no on-disk product; set output_dir "
            "(or merged_output_dir for a merged LoRA run) to make the result freezable."
        )
    return {"artifact": artifact}
fit_identity()

The training-relevant args projection, or None when nothing trains.

Output locations (output_dir, merged_output_dir, resume_from_checkpoint, and training_args["output_dir"]) are excluded, since where a product is saved does not change what was trained.

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
def fit_identity(self) -> Any | None:
    """The training-relevant args projection, or None when nothing trains.

    Output locations (`output_dir`, `merged_output_dir`, `resume_from_checkpoint`, and
    `training_args["output_dir"]`) are excluded, since where a product is saved does not
    change what was trained.
    """
    if getattr(self, "train_dataset", None) is None:
        return None
    args = getattr(self, "args", None)
    if args is None:
        return None
    excluded = {"output_dir", "merged_output_dir", "resume_from_checkpoint"}
    payload: dict[str, Any] = {}
    for f in fields(args):
        if not f.init or f.name in excluded:
            continue
        value = getattr(args, f.name)
        if f.name == "training_args" and isinstance(value, dict):
            value = {k: v for k, v in value.items() if k not in excluded}
        payload[f.name] = value
    return payload
frozen_form(state)

The trained artifact as a load_lora or load_checkpoint entry.

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
195
196
197
198
199
200
201
202
203
204
def frozen_form(self, state: dict[str, Any]) -> tuple[str, dict[str, Any]]:
    """The trained artifact as a `load_lora` or `load_checkpoint` entry."""
    artifact = state["artifact"]
    if isinstance(artifact, LoRAArtifact):
        return "structural_control/load_lora", {
            "path": artifact,
            "base_model": artifact.base_model,
            "merge": False,
        }
    return "structural_control/load_checkpoint", {"path": artifact}
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 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
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, **_)

Required steering/preparation.

session is a SteeringSession on the steering backend, provided by the pipeline.

Source code in steerability/algorithms/structural_control/wrappers/trl/grpotrainer/base_mixin.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
def steer(
    self,
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase | None = None,
    **_,
) -> torch.nn.Module:
    self.tokenizer = tokenizer or (getattr(model, "tokenizer", None) if model is not None else None)
    model, self.tokenizer = self._resolve_model_tokenizer(model, self.tokenizer)
    self.tokenizer = ensure_pad_token(self.tokenizer)

    if self.reward_funcs is None:
        raise ValueError(
            "GRPOTrainerMixin.steer: `reward_funcs` is required (a callable or list of callables)."
        )
    reward_funcs = self.reward_funcs if isinstance(self.reward_funcs, list) else [self.reward_funcs]

    # GRPO reads the text "prompt" column directly; standardize but do NOT tokenize to input_ids
    train_dataset = standardize_prompt_dataset(self.train_dataset) if self.train_dataset is not None else None
    eval_dataset = standardize_prompt_dataset(self.eval_dataset) if self.eval_dataset is not None else None

    config_kwargs = resolve_config_kwargs(GRPOConfig, self.training_args)
    training_config = GRPOConfig(**config_kwargs)

    peft_config = None
    if self.use_peft and self.peft_type == PeftType.LORA:
        peft_config = LoraConfig(**self.lora_kwargs)

    if train_dataset is not None:
        trainer = GRPOTrainer(
            model=model,
            reward_funcs=reward_funcs,
            args=training_config,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            processing_class=self.tokenizer,
            peft_config=peft_config,
        )
        trainer.train(resume_from_checkpoint=self.training_args.get("resume_from_checkpoint"))

        # recover the trained policy so it can be used for generation (GRPO has no .policy wrapper)
        trained_model = trainer.accelerator.unwrap_model(trainer.model)
        model = getattr(trained_model, "policy", trained_model)
        self._maybe_save_trained_artifacts(trainer)
        model = self._maybe_merge_lora_in_place(model)

    return self._post_train_freeze(model)
steer_access()

ModelAccess.MODULE; training happens on live weights.

Source code in steerability/algorithms/structural_control/base.py 
112
113
114
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 
105
106
107
108
109
110
111
112
113
114
115
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 ()

ppotrainer

args

base_mixin

PPOTrainerMixin

Bases: TRLMixin, StructuralControl

PPO structural control backed by TRL's PPOTrainer (shipped under trl.experimental in TRL 1.x).

Reward and value models are sequence-classification models. When value_model_name_or_path is not provided, the wrapper loads a fresh value model from the same path as the reward model, since TRL's PPOTrainer wraps the value model into its policy and rejects value_model=None.

The standardized prompt-only training dataset is tokenized into an input_ids column before being passed to the trainer, which reads data["input_ids"] directly.

Note: reward and value models must share the policy's tokenizer/vocabulary. The wrapper guards this with an error.

Source code in steerability/algorithms/structural_control/wrappers/trl/ppotrainer/base_mixin.py 
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
class PPOTrainerMixin(TRLMixin, StructuralControl):
    """PPO structural control backed by TRL's `PPOTrainer` (shipped under `trl.experimental` in TRL 1.x).

    Reward and value models are sequence-classification models. When
    `value_model_name_or_path` is not provided, the wrapper loads a fresh value model from the same
    path as the reward model, since TRL's `PPOTrainer` wraps the value model into its policy and
    rejects `value_model=None`.

    The standardized prompt-only training dataset is tokenized into an `input_ids` column before
    being passed to the trainer, which reads `data["input_ids"]` directly.

    Note: reward and value models must share the policy's tokenizer/vocabulary. The wrapper guards
    this with an error.
    """

    train_dataset: Any | None = None
    eval_dataset: Any | None = None
    reward_model_name_or_path: str | None = None
    value_model_name_or_path: str | None = None
    max_prompt_length: int = 512

    def steer(
        self,
        model: PreTrainedModel | None,
        tokenizer: PreTrainedTokenizerBase | None = None,
        ref_model: PreTrainedModel | None = None,
        **_,
    ) -> torch.nn.Module:
        if self.training_args.get("resume_from_checkpoint"):
            raise ValueError(
                "PPO does not support resume_from_checkpoint; TRL's PPOTrainer.train() takes no checkpoint argument."
            )

        self.tokenizer = tokenizer or (getattr(model, "tokenizer", None) if model is not None else None)
        model, self.tokenizer = self._resolve_model_tokenizer(model, self.tokenizer)
        self.tokenizer = ensure_pad_token(self.tokenizer)

        # reward + value models (sequence-classification heads)
        reward_model = AutoModelForSequenceClassification.from_pretrained(
            self.reward_model_name_or_path,
            num_labels=1,
            trust_remote_code=self.trust_remote_code,
        )
        value_path = self.value_model_name_or_path or self.reward_model_name_or_path
        value_model = AutoModelForSequenceClassification.from_pretrained(
            value_path,
            num_labels=1,
            trust_remote_code=self.trust_remote_code,
        )

        self._check_scoring_vocab(reward_model, value_model)

        # standardize then tokenize so the dataset has the input_ids column PPOTrainer expects
        train_dataset = self._prepare_dataset(self.train_dataset) if self.train_dataset is not None else None
        eval_dataset = self._prepare_dataset(self.eval_dataset) if self.eval_dataset is not None else None

        config_kwargs = resolve_config_kwargs(PPOConfig, self.training_args)
        training_config = PPOConfig(**config_kwargs)

        peft_config = None
        if self.use_peft and self.peft_type == PeftType.LORA:
            peft_config = LoraConfig(**self.lora_kwargs)
            ref_model = None

        if train_dataset is not None:
            trainer = PPOTrainer(
                args=training_config,
                processing_class=self.tokenizer,
                model=model,
                ref_model=ref_model,
                reward_model=reward_model,
                value_model=value_model,
                train_dataset=train_dataset,
                eval_dataset=eval_dataset,
                peft_config=peft_config,
            )
            trainer.train()

            # recover the trained policy so it can be used for generation
            trained_model = trainer.accelerator.unwrap_model(trainer.model)
            model = getattr(trained_model, "policy", trained_model)
            self._maybe_save_trained_artifacts(trainer)
            model = self._maybe_merge_lora_in_place(model)

        return self._post_train_freeze(model)

    def _check_scoring_vocab(self, reward_model, value_model) -> None:
        """Verify the reward/value models can index every policy token id.
        """
        policy_vocab = len(self.tokenizer)
        for role, scoring_model, path in (
            ("reward", reward_model, self.reward_model_name_or_path),
            ("value", value_model, self.value_model_name_or_path or self.reward_model_name_or_path),
        ):
            scoring_vocab = getattr(scoring_model.config, "vocab_size", None)
            if scoring_vocab is not None and scoring_vocab < policy_vocab:
                raise ValueError(
                    f"PPO {role} model '{path}' has vocab_size {scoring_vocab}, smaller than the "
                    f"policy tokenizer's {policy_vocab}. TRL's PPOTrainer scores the policy's own token "
                    f"ids with one shared tokenizer, so the {role} model must share the policy's "
                    "tokenizer/vocabulary (use a reward/value model from the policy's model family)."
                )

    def _prepare_dataset(self, dataset):
        """Standardize a prompt-bearing dataset and tokenize prompts into the `input_ids` column TRL expects."""
        ds = standardize_prompt_dataset(dataset)
        max_length = self.max_prompt_length

        def tokenize(example):
            return self.tokenizer(
                example["prompt"],
                truncation=True,
                max_length=max_length,
            )

        return ds.map(tokenize, remove_columns=["prompt"])
RUNTIME_KWARGS_SCHEMA = [] class-attribute instance-attribute
adapter_name = None class-attribute instance-attribute
args = self.Args.validate(*args, **kwargs) instance-attribute
base_model_name_or_path = None class-attribute instance-attribute
device = None class-attribute instance-attribute
enabled = True class-attribute instance-attribute
eval_dataset = None class-attribute instance-attribute
hf_model_kwargs = {} class-attribute instance-attribute
lora_kwargs = {} class-attribute instance-attribute
max_prompt_length = 512 class-attribute instance-attribute
merge_lora_after_train = False class-attribute instance-attribute
merged_output_dir = None class-attribute instance-attribute
output_dir = None class-attribute instance-attribute
peft_type = None class-attribute instance-attribute
resume_from_checkpoint = None class-attribute instance-attribute
reward_model_name_or_path = None class-attribute instance-attribute
supports_batching = True class-attribute instance-attribute
tokenizer = None class-attribute instance-attribute
tokenizer_name_or_path = None class-attribute instance-attribute
train_dataset = None class-attribute instance-attribute
training_args = {} class-attribute instance-attribute
trust_remote_code = False class-attribute instance-attribute
use_peft = False class-attribute instance-attribute
value_model_name_or_path = None class-attribute instance-attribute
artifact_capability()

The serve capability implied by this training configuration.

Training runs only when train_dataset is set, so a configuration without one produces no artifact. A LoRA run without merge-back saves an adapter to the output directory (Capability.SERVE_LORA); a full fine-tune saves a checkpoint there (Capability.SERVE_CHECKPOINT); a merged LoRA run yields a checkpoint only when merged_output_dir is set.

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
def artifact_capability(self) -> Capability | None:
    """The serve capability implied by this training configuration.

    Training runs only when `train_dataset` is set, so a configuration without one produces
    no artifact. A LoRA run without merge-back saves an adapter to the output directory
    (`Capability.SERVE_LORA`); a full fine-tune saves a checkpoint there
    (`Capability.SERVE_CHECKPOINT`); a merged LoRA run yields a checkpoint only when
    `merged_output_dir` is set.
    """
    if getattr(self, "train_dataset", None) is None:
        return None
    is_lora = bool(self.use_peft) and self.peft_type == PeftType.LORA
    if is_lora and not self.merge_lora_after_train:
        return Capability.SERVE_LORA if self._resolved_output_dir() else None
    if is_lora and self.merge_lora_after_train:
        return Capability.SERVE_CHECKPOINT if self.merged_output_dir else None
    return Capability.SERVE_CHECKPOINT if self._resolved_output_dir() else None
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 
195
196
197
198
199
200
201
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 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
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 on-disk product of this configuration's steer(), matching artifact_capability().

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
162
163
164
165
166
167
168
169
170
171
172
173
def export_artifact(self) -> Artifact | None:
    """The on-disk product of this configuration's `steer()`, matching
    `artifact_capability()`."""
    capability = self.artifact_capability()
    if capability is None:
        return None
    if capability == Capability.SERVE_LORA:
        base = self.base_model_name_or_path or self._resolved_base_ref or ""
        return LoRAArtifact(path=str(self._resolved_output_dir()), base_model=str(base))
    is_lora = bool(self.use_peft) and self.peft_type == PeftType.LORA
    path = self.merged_output_dir if (is_lora and self.merge_lora_after_train) else self._resolved_output_dir()
    return CheckpointArtifact(path=str(path))
export_state()

The trained on-disk product under the "artifact" key, for freezing.

Returns an empty mapping when the configuration trains nothing (train_dataset is None), so an inert wrapper's recipe is its frozen form.

Raises:

Type Description
NotFreezableError

If the configuration trains but produces no on-disk product (a merged LoRA run without merged_output_dir).
Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def export_state(self) -> dict[str, Any]:
    """The trained on-disk product under the `"artifact"` key, for freezing.

    Returns an empty mapping when the configuration trains nothing (`train_dataset` is
    None), so an inert wrapper's recipe is its frozen form.

    Raises:
        NotFreezableError: If the configuration trains but produces no on-disk product
            (a merged LoRA run without `merged_output_dir`).
    """
    if getattr(self, "train_dataset", None) is None:
        return {}
    artifact = self.export_artifact()
    if artifact is None:
        raise NotFreezableError(
            f"{type(self).__name__} trains but writes no on-disk product; set output_dir "
            "(or merged_output_dir for a merged LoRA run) to make the result freezable."
        )
    return {"artifact": artifact}
fit_identity()

The training-relevant args projection, or None when nothing trains.

Output locations (output_dir, merged_output_dir, resume_from_checkpoint, and training_args["output_dir"]) are excluded, since where a product is saved does not change what was trained.

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
def fit_identity(self) -> Any | None:
    """The training-relevant args projection, or None when nothing trains.

    Output locations (`output_dir`, `merged_output_dir`, `resume_from_checkpoint`, and
    `training_args["output_dir"]`) are excluded, since where a product is saved does not
    change what was trained.
    """
    if getattr(self, "train_dataset", None) is None:
        return None
    args = getattr(self, "args", None)
    if args is None:
        return None
    excluded = {"output_dir", "merged_output_dir", "resume_from_checkpoint"}
    payload: dict[str, Any] = {}
    for f in fields(args):
        if not f.init or f.name in excluded:
            continue
        value = getattr(args, f.name)
        if f.name == "training_args" and isinstance(value, dict):
            value = {k: v for k, v in value.items() if k not in excluded}
        payload[f.name] = value
    return payload
frozen_form(state)

The trained artifact as a load_lora or load_checkpoint entry.

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
195
196
197
198
199
200
201
202
203
204
def frozen_form(self, state: dict[str, Any]) -> tuple[str, dict[str, Any]]:
    """The trained artifact as a `load_lora` or `load_checkpoint` entry."""
    artifact = state["artifact"]
    if isinstance(artifact, LoRAArtifact):
        return "structural_control/load_lora", {
            "path": artifact,
            "base_model": artifact.base_model,
            "merge": False,
        }
    return "structural_control/load_checkpoint", {"path": artifact}
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 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
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, ref_model=None, **_)

Required steering/preparation.

session is a SteeringSession on the steering backend, provided by the pipeline.

Source code in steerability/algorithms/structural_control/wrappers/trl/ppotrainer/base_mixin.py 
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def steer(
    self,
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase | None = None,
    ref_model: PreTrainedModel | None = None,
    **_,
) -> torch.nn.Module:
    if self.training_args.get("resume_from_checkpoint"):
        raise ValueError(
            "PPO does not support resume_from_checkpoint; TRL's PPOTrainer.train() takes no checkpoint argument."
        )

    self.tokenizer = tokenizer or (getattr(model, "tokenizer", None) if model is not None else None)
    model, self.tokenizer = self._resolve_model_tokenizer(model, self.tokenizer)
    self.tokenizer = ensure_pad_token(self.tokenizer)

    # reward + value models (sequence-classification heads)
    reward_model = AutoModelForSequenceClassification.from_pretrained(
        self.reward_model_name_or_path,
        num_labels=1,
        trust_remote_code=self.trust_remote_code,
    )
    value_path = self.value_model_name_or_path or self.reward_model_name_or_path
    value_model = AutoModelForSequenceClassification.from_pretrained(
        value_path,
        num_labels=1,
        trust_remote_code=self.trust_remote_code,
    )

    self._check_scoring_vocab(reward_model, value_model)

    # standardize then tokenize so the dataset has the input_ids column PPOTrainer expects
    train_dataset = self._prepare_dataset(self.train_dataset) if self.train_dataset is not None else None
    eval_dataset = self._prepare_dataset(self.eval_dataset) if self.eval_dataset is not None else None

    config_kwargs = resolve_config_kwargs(PPOConfig, self.training_args)
    training_config = PPOConfig(**config_kwargs)

    peft_config = None
    if self.use_peft and self.peft_type == PeftType.LORA:
        peft_config = LoraConfig(**self.lora_kwargs)
        ref_model = None

    if train_dataset is not None:
        trainer = PPOTrainer(
            args=training_config,
            processing_class=self.tokenizer,
            model=model,
            ref_model=ref_model,
            reward_model=reward_model,
            value_model=value_model,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            peft_config=peft_config,
        )
        trainer.train()

        # recover the trained policy so it can be used for generation
        trained_model = trainer.accelerator.unwrap_model(trainer.model)
        model = getattr(trained_model, "policy", trained_model)
        self._maybe_save_trained_artifacts(trainer)
        model = self._maybe_merge_lora_in_place(model)

    return self._post_train_freeze(model)
steer_access()

ModelAccess.MODULE; training happens on live weights.

Source code in steerability/algorithms/structural_control/base.py 
112
113
114
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 
105
106
107
108
109
110
111
112
113
114
115
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 ()

control

PPO

Bases: PPOTrainerMixin

PPO controller.

Source code in steerability/algorithms/structural_control/wrappers/trl/ppotrainer/control.py 
5
6
7
8
9
class PPO(PPOTrainerMixin):
    """
    PPO controller.
    """
    Args = PPOArgs
RUNTIME_KWARGS_SCHEMA = [] class-attribute instance-attribute
adapter_name = None class-attribute instance-attribute
args = self.Args.validate(*args, **kwargs) instance-attribute
base_model_name_or_path = None class-attribute instance-attribute
device = None class-attribute instance-attribute
enabled = True class-attribute instance-attribute
eval_dataset = None class-attribute instance-attribute
hf_model_kwargs = {} class-attribute instance-attribute
lora_kwargs = {} class-attribute instance-attribute
max_prompt_length = 512 class-attribute instance-attribute
merge_lora_after_train = False class-attribute instance-attribute
merged_output_dir = None class-attribute instance-attribute
output_dir = None class-attribute instance-attribute
peft_type = None class-attribute instance-attribute
resume_from_checkpoint = None class-attribute instance-attribute
reward_model_name_or_path = None class-attribute instance-attribute
supports_batching = True class-attribute instance-attribute
tokenizer = None class-attribute instance-attribute
tokenizer_name_or_path = None class-attribute instance-attribute
train_dataset = None class-attribute instance-attribute
training_args = {} class-attribute instance-attribute
trust_remote_code = False class-attribute instance-attribute
use_peft = False class-attribute instance-attribute
value_model_name_or_path = None class-attribute instance-attribute
artifact_capability()

The serve capability implied by this training configuration.

Training runs only when train_dataset is set, so a configuration without one produces no artifact. A LoRA run without merge-back saves an adapter to the output directory (Capability.SERVE_LORA); a full fine-tune saves a checkpoint there (Capability.SERVE_CHECKPOINT); a merged LoRA run yields a checkpoint only when merged_output_dir is set.

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
def artifact_capability(self) -> Capability | None:
    """The serve capability implied by this training configuration.

    Training runs only when `train_dataset` is set, so a configuration without one produces
    no artifact. A LoRA run without merge-back saves an adapter to the output directory
    (`Capability.SERVE_LORA`); a full fine-tune saves a checkpoint there
    (`Capability.SERVE_CHECKPOINT`); a merged LoRA run yields a checkpoint only when
    `merged_output_dir` is set.
    """
    if getattr(self, "train_dataset", None) is None:
        return None
    is_lora = bool(self.use_peft) and self.peft_type == PeftType.LORA
    if is_lora and not self.merge_lora_after_train:
        return Capability.SERVE_LORA if self._resolved_output_dir() else None
    if is_lora and self.merge_lora_after_train:
        return Capability.SERVE_CHECKPOINT if self.merged_output_dir else None
    return Capability.SERVE_CHECKPOINT if self._resolved_output_dir() else None
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 
195
196
197
198
199
200
201
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 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
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 on-disk product of this configuration's steer(), matching artifact_capability().

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
162
163
164
165
166
167
168
169
170
171
172
173
def export_artifact(self) -> Artifact | None:
    """The on-disk product of this configuration's `steer()`, matching
    `artifact_capability()`."""
    capability = self.artifact_capability()
    if capability is None:
        return None
    if capability == Capability.SERVE_LORA:
        base = self.base_model_name_or_path or self._resolved_base_ref or ""
        return LoRAArtifact(path=str(self._resolved_output_dir()), base_model=str(base))
    is_lora = bool(self.use_peft) and self.peft_type == PeftType.LORA
    path = self.merged_output_dir if (is_lora and self.merge_lora_after_train) else self._resolved_output_dir()
    return CheckpointArtifact(path=str(path))
export_state()

The trained on-disk product under the "artifact" key, for freezing.

Returns an empty mapping when the configuration trains nothing (train_dataset is None), so an inert wrapper's recipe is its frozen form.

Raises:

Type Description
NotFreezableError

If the configuration trains but produces no on-disk product (a merged LoRA run without merged_output_dir).
Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def export_state(self) -> dict[str, Any]:
    """The trained on-disk product under the `"artifact"` key, for freezing.

    Returns an empty mapping when the configuration trains nothing (`train_dataset` is
    None), so an inert wrapper's recipe is its frozen form.

    Raises:
        NotFreezableError: If the configuration trains but produces no on-disk product
            (a merged LoRA run without `merged_output_dir`).
    """
    if getattr(self, "train_dataset", None) is None:
        return {}
    artifact = self.export_artifact()
    if artifact is None:
        raise NotFreezableError(
            f"{type(self).__name__} trains but writes no on-disk product; set output_dir "
            "(or merged_output_dir for a merged LoRA run) to make the result freezable."
        )
    return {"artifact": artifact}
fit_identity()

The training-relevant args projection, or None when nothing trains.

Output locations (output_dir, merged_output_dir, resume_from_checkpoint, and training_args["output_dir"]) are excluded, since where a product is saved does not change what was trained.

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
def fit_identity(self) -> Any | None:
    """The training-relevant args projection, or None when nothing trains.

    Output locations (`output_dir`, `merged_output_dir`, `resume_from_checkpoint`, and
    `training_args["output_dir"]`) are excluded, since where a product is saved does not
    change what was trained.
    """
    if getattr(self, "train_dataset", None) is None:
        return None
    args = getattr(self, "args", None)
    if args is None:
        return None
    excluded = {"output_dir", "merged_output_dir", "resume_from_checkpoint"}
    payload: dict[str, Any] = {}
    for f in fields(args):
        if not f.init or f.name in excluded:
            continue
        value = getattr(args, f.name)
        if f.name == "training_args" and isinstance(value, dict):
            value = {k: v for k, v in value.items() if k not in excluded}
        payload[f.name] = value
    return payload
frozen_form(state)

The trained artifact as a load_lora or load_checkpoint entry.

Source code in steerability/algorithms/structural_control/wrappers/trl/base_mixin.py 
195
196
197
198
199
200
201
202
203
204
def frozen_form(self, state: dict[str, Any]) -> tuple[str, dict[str, Any]]:
    """The trained artifact as a `load_lora` or `load_checkpoint` entry."""
    artifact = state["artifact"]
    if isinstance(artifact, LoRAArtifact):
        return "structural_control/load_lora", {
            "path": artifact,
            "base_model": artifact.base_model,
            "merge": False,
        }
    return "structural_control/load_checkpoint", {"path": artifact}
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 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
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, ref_model=None, **_)

Required steering/preparation.

session is a SteeringSession on the steering backend, provided by the pipeline.

Source code in steerability/algorithms/structural_control/wrappers/trl/ppotrainer/base_mixin.py 
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def steer(
    self,
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase | None = None,
    ref_model: PreTrainedModel | None = None,
    **_,
) -> torch.nn.Module:
    if self.training_args.get("resume_from_checkpoint"):
        raise ValueError(
            "PPO does not support resume_from_checkpoint; TRL's PPOTrainer.train() takes no checkpoint argument."
        )

    self.tokenizer = tokenizer or (getattr(model, "tokenizer", None) if model is not None else None)
    model, self.tokenizer = self._resolve_model_tokenizer(model, self.tokenizer)
    self.tokenizer = ensure_pad_token(self.tokenizer)

    # reward + value models (sequence-classification heads)
    reward_model = AutoModelForSequenceClassification.from_pretrained(
        self.reward_model_name_or_path,
        num_labels=1,
        trust_remote_code=self.trust_remote_code,
    )
    value_path = self.value_model_name_or_path or self.reward_model_name_or_path
    value_model = AutoModelForSequenceClassification.from_pretrained(
        value_path,
        num_labels=1,
        trust_remote_code=self.trust_remote_code,
    )

    self._check_scoring_vocab(reward_model, value_model)

    # standardize then tokenize so the dataset has the input_ids column PPOTrainer expects
    train_dataset = self._prepare_dataset(self.train_dataset) if self.train_dataset is not None else None
    eval_dataset = self._prepare_dataset(self.eval_dataset) if self.eval_dataset is not None else None

    config_kwargs = resolve_config_kwargs(PPOConfig, self.training_args)
    training_config = PPOConfig(**config_kwargs)

    peft_config = None
    if self.use_peft and self.peft_type == PeftType.LORA:
        peft_config = LoraConfig(**self.lora_kwargs)
        ref_model = None

    if train_dataset is not None:
        trainer = PPOTrainer(
            args=training_config,
            processing_class=self.tokenizer,
            model=model,
            ref_model=ref_model,
            reward_model=reward_model,
            value_model=value_model,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            peft_config=peft_config,
        )
        trainer.train()

        # recover the trained policy so it can be used for generation
        trained_model = trainer.accelerator.unwrap_model(trainer.model)
        model = getattr(trained_model, "policy", trained_model)
        self._maybe_save_trained_artifacts(trainer)
        model = self._maybe_merge_lora_in_place(model)

    return self._post_train_freeze(model)
steer_access()

ModelAccess.MODULE; training happens on live weights.

Source code in steerability/algorithms/structural_control/base.py 
112
113
114
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 
105
106
107
108
109
110
111
112
113
114
115
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 ()