Running TRL methods¶

The toolkit wraps several TRL trainers as structural controls. In this guide we fine-tune a small model on preference data through a SteeringPipeline. We cover supervised fine-tuning (SFT) and direct preference optimization (DPO) with LoRA adapters, anchored preference optimization (APO) as a DPO-family variant, and a full-parameter SFT run. We also show how to resume an interrupted run from a checkpoint and how to serve the trained artifact (a merged checkpoint or a LoRA adapter) on the vLLM backend.

Setup¶

If running this from a Google Colab notebook, please uncomment the following cell to install the toolkit. The following block is not necessary if running this notebook from a virtual environment where the package has already been installed.

In [1]:

# !git clone https://github.com/generative-computing/steerability.git
# %cd steerability

# !git clone https://github.com/generative-computing/steerability.git # %cd steerability

The following authentication steps may be necessary to access any gated models (after being granted access by Hugging Face). Uncomment the following if you need to log in to the Hugging Face Hub:

In [2]:

# !pip install python-dotenv
# !pip install ipywidgets
# from dotenv import load_dotenv
# import os

# load_dotenv()
# token = os.getenv("HUGGINGFACE_TOKEN")
# from huggingface_hub import login
# login(token=token)

# !pip install python-dotenv # !pip install ipywidgets # from dotenv import load_dotenv # import os # load_dotenv() # token = os.getenv("HUGGINGFACE_TOKEN") # from huggingface_hub import login # login(token=token)

Next, we import the SteeringPipeline class (used throughout) and specify the base model, in this case a small Qwen model.

In [3]:

import torch
from datasets import load_dataset
from peft import PeftType
from transformers import AutoTokenizer

from steerability.algorithms.core.steering_pipeline import SteeringPipeline


MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct" 

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

device = "cuda" if torch.cuda.is_available() else "cpu"
print("Using device:", device)

import torch from datasets import load_dataset from peft import PeftType from transformers import AutoTokenizer from steerability.algorithms.core.steering_pipeline import SteeringPipeline MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct" tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token device = "cuda" if torch.cuda.is_available() else "cpu" print("Using device:", device)

Using device: cuda

Data Preparation¶

The controls throughout this notebook are trained using a common dataset, ultrafeedback_binarized, since it contains preference data for each prompt (which is necessary for DPO-based controls). We load each of the splits below.

In [4]:

raw_train = load_dataset("HuggingFaceH4/ultrafeedback_binarized", split="train_prefs")
raw_test  = load_dataset("HuggingFaceH4/ultrafeedback_binarized", split="test_prefs")
len(raw_train), raw_train[0].keys()

raw_train = load_dataset("HuggingFaceH4/ultrafeedback_binarized", split="train_prefs") raw_test = load_dataset("HuggingFaceH4/ultrafeedback_binarized", split="test_prefs") len(raw_train), raw_train[0].keys()

Out[4]:

(61135,
 dict_keys(['prompt', 'prompt_id', 'chosen', 'rejected', 'messages', 'score_chosen', 'score_rejected']))

Different trainers expect different data formats (i.e., tensor layouts) and thus we define two helper functions, one for SFT and one for DPO, to process the data in a way that is amenable to each.

In [5]:

def sft_preprocess(example, tokenizer, max_length=1024):
    answer = example["chosen"][-1]["content"]
    text = f"Question: {example['prompt']}\n\nAnswer: {answer}"
    encoding = tokenizer(text, truncation=True, max_length=max_length)
    labels = [
        token_id if mask == 1 else -100  # label pads as -100 so they don't contribute to loss
        for token_id, mask in zip(encoding["input_ids"], encoding["attention_mask"])
    ]
    encoding["labels"] = labels
    return encoding

def dpo_filter(example, max_prompt_chars=4000):
    return {
        "prompt": example["prompt"][:max_prompt_chars],
        "chosen": example["chosen"][-1]["content"],
        "rejected": example["rejected"][-1]["content"],
    }


subset_size = 500

sft_train = raw_train.select(range(subset_size)).map(
    lambda example: sft_preprocess(example, tokenizer, max_length=1024),
    remove_columns=raw_train.column_names
)

dpo_train = raw_train.select(range(subset_size)).map(dpo_filter, remove_columns=[])
dpo_train[0].keys()

def sft_preprocess(example, tokenizer, max_length=1024): answer = example["chosen"][-1]["content"] text = f"Question: {example['prompt']}\n\nAnswer: {answer}" encoding = tokenizer(text, truncation=True, max_length=max_length) labels = [ token_id if mask == 1 else -100 # label pads as -100 so they don't contribute to loss for token_id, mask in zip(encoding["input_ids"], encoding["attention_mask"]) ] encoding["labels"] = labels return encoding def dpo_filter(example, max_prompt_chars=4000): return { "prompt": example["prompt"][:max_prompt_chars], "chosen": example["chosen"][-1]["content"], "rejected": example["rejected"][-1]["content"], } subset_size = 500 sft_train = raw_train.select(range(subset_size)).map( lambda example: sft_preprocess(example, tokenizer, max_length=1024), remove_columns=raw_train.column_names ) dpo_train = raw_train.select(range(subset_size)).map(dpo_filter, remove_columns=[]) dpo_train[0].keys()

Out[5]:

dict_keys(['prompt', 'prompt_id', 'chosen', 'rejected', 'messages', 'score_chosen', 'score_rejected'])

SFT control¶

We now show how to fine-tune with SFT using LoRA. We also merge the trained adapter back into the model (using the argument merge_lora_after_train). Note the argument use_peft=True to indicate that we are not running a full fine-tune (the example near the end of this notebook will illustrate a full fine-tuning run).

We load the model in float32 via hf_model_kwargs={"dtype": torch.float32}. Transformers loads a checkpoint in its saved dtype by default (Qwen2.5 ships bfloat16), and training wants float32 master weights, so we set the dtype explicitly across the training pipelines below.

In [6]:

from steerability.algorithms.structural_control.wrappers.trl.sfttrainer.control import SFT


sft = SFT(
    # data
    train_dataset=sft_train,
    eval_dataset=None, 
    # data_collator=None  # optional; if omitted and you provided labels, you're fine

    # TRL / Trainer config (forwarded into SFTConfig)
    output_dir="./tmp/sft_lora",
    max_length=1024,
    per_device_train_batch_size=4,
    num_train_epochs=1,
    learning_rate=1e-4,
    logging_steps=50,
    save_strategy="steps",
    load_best_model_at_end=False,
    report_to="none",
    seed=42,
    training_args={"save_steps": 50},

    # PEFT (LoRA)
    use_peft=True,
    peft_type=PeftType.LORA,
    r=16,
    lora_alpha=16,
    lora_dropout=0.05,
    target_modules=["q_proj", "v_proj"],
    adapter_name="sft",

    # optionally merge LoRA into base weights after training
    merge_lora_after_train=True,
    merged_output_dir="./tmp/sft_lora_merged",
)

from steerability.algorithms.structural_control.wrappers.trl.sfttrainer.control import SFT sft = SFT( # data train_dataset=sft_train, eval_dataset=None, # data_collator=None # optional; if omitted and you provided labels, you're fine # TRL / Trainer config (forwarded into SFTConfig) output_dir="./tmp/sft_lora", max_length=1024, per_device_train_batch_size=4, num_train_epochs=1, learning_rate=1e-4, logging_steps=50, save_strategy="steps", load_best_model_at_end=False, report_to="none", seed=42, training_args={"save_steps": 50}, # PEFT (LoRA) use_peft=True, peft_type=PeftType.LORA, r=16, lora_alpha=16, lora_dropout=0.05, target_modules=["q_proj", "v_proj"], adapter_name="sft", # optionally merge LoRA into base weights after training merge_lora_after_train=True, merged_output_dir="./tmp/sft_lora_merged", )

We create a steering pipeline using the above control, without a model_name_or_path since the structural control (sft) returns a model. The pipeline is then steered which invokes the training procedure.

In [7]:

sft_pipeline = SteeringPipeline(
    model_name_or_path=MODEL_NAME,
    device_map=None,
    hf_model_kwargs={"trust_remote_code": True, "dtype": torch.float32},
    controls=[sft],
)

sft_pipeline.steer()

sft_pipeline = SteeringPipeline( model_name_or_path=MODEL_NAME, device_map=None, hf_model_kwargs={"trust_remote_code": True, "dtype": torch.float32}, controls=[sft], ) sft_pipeline.steer()

[transformers] The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None, 'pad_token_id': 151643}.

[125/125 00:32, Epoch 1/1]

Step	Training Loss
50	1.660603
100	1.579905

The above SFT-trained pipeline is now ready for inference.

In [8]:

prompt = "Question: What makes the sky look blue?\n\nAnswer:"
print(sft_pipeline.generate(prompt, max_new_tokens=64, do_sample=False))

prompt = "Question: What makes the sky look blue?\n\nAnswer:" print(sft_pipeline.generate(prompt, max_new_tokens=64, do_sample=False))

 The sky looks blue because it reflects sunlight. When light hits a surface, some of it is reflected back to us as white light, while other parts are absorbed or scattered by the surface. The remaining light that is not absorbed or scattered is what we see as color.

The colors in the sky are caused by different wavelengths

DPO control¶

DPO is instantiated in a similar fashion with the primary differences being that the training data is now triples (prompt, chosen, rejected), the trainer must keep a reference policy alongside the trainable policy, and the loss is a pair-wise KL-reg. contrastive objective rather than the token-level cross entropy loss in SFT.

Note: By default, the trainer clones the base weights and freezes them. When LoRA is enabled, the wrapper automatically passes ref_model=None, letting TRL re-create a frozen reference that shares the same LoRA adapters. If you are full fine-tuning you can still supply your own ref_model via pipeline.steer(ref_model=my_frozen_model).

In [9]:

from steerability.algorithms.structural_control.wrappers.trl.dpotrainer.control import DPO


dpo = DPO(
    train_dataset=dpo_train,

    # DPO / TRL config (forwarded into DPOConfig)
    output_dir="./tmp/dpo_lora",
    per_device_train_batch_size=2,  # often smaller than SFT
    num_train_epochs=1,
    learning_rate=5e-5,
    beta=0.1,
    loss_type="sigmoid",  # baseline DPO loss
    max_length=1024,
    prompt_format="chat_prompt",
    precompute_ref_log_probs=False,  # off: avoids the noisy per-batch reference log-prob pass; enable for multi-epoch runs where the precompute is reused
    disable_dropout=True,
    logging_steps=50,
    report_to="none",
    seed=123,

    # LoRA
    use_peft=True,
    peft_type=PeftType.LORA,
    r=16,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],
    adapter_name="dpo",

    merge_lora_after_train=False,
)

from steerability.algorithms.structural_control.wrappers.trl.dpotrainer.control import DPO dpo = DPO( train_dataset=dpo_train, # DPO / TRL config (forwarded into DPOConfig) output_dir="./tmp/dpo_lora", per_device_train_batch_size=2, # often smaller than SFT num_train_epochs=1, learning_rate=5e-5, beta=0.1, loss_type="sigmoid", # baseline DPO loss max_length=1024, prompt_format="chat_prompt", precompute_ref_log_probs=False, # off: avoids the noisy per-batch reference log-prob pass; enable for multi-epoch runs where the precompute is reused disable_dropout=True, logging_steps=50, report_to="none", seed=123, # LoRA use_peft=True, peft_type=PeftType.LORA, r=16, lora_alpha=16, target_modules=["q_proj", "v_proj"], adapter_name="dpo", merge_lora_after_train=False, )

As before, we create the pipeline using the control, steer the pipeline, and run inference on the steered pipeline.

In [10]:

dpo_pipeline = SteeringPipeline(
    model_name_or_path=MODEL_NAME,
    hf_model_kwargs={"trust_remote_code": True, "dtype": torch.float32},
    controls=[dpo]
)
dpo_pipeline.steer()

dpo_pipeline = SteeringPipeline( model_name_or_path=MODEL_NAME, hf_model_kwargs={"trust_remote_code": True, "dtype": torch.float32}, controls=[dpo] ) dpo_pipeline.steer()

[transformers] The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None, 'pad_token_id': 151643}.

[248/248 01:15, Epoch 1/1]

Step	Training Loss
50	0.682237
100	0.647819
150	0.716059
200	0.651778

In [11]:

print(dpo_pipeline.generate(
    messages=[{"role": "user", "content": "Is it ever helpful to be blunt with feedback?"}],
    max_new_tokens=150,
    do_sample=False,
))

print(dpo_pipeline.generate( messages=[{"role": "user", "content": "Is it ever helpful to be blunt with feedback?"}], max_new_tokens=150, do_sample=False, ))

As an AI language model, I don't have personal preferences or emotions like humans do, but I can provide some insights based on general principles of communication and feedback.

Being blunt with feedback is generally considered beneficial in several ways:

1. **Clarification**: Blunt feedback helps clarify the issue at hand, making it easier for others to understand what went wrong and how to improve.

2. **Empowerment**: When feedback is clear and specific, it empowers individuals to take ownership of their performance and make necessary adjustments.

3. **Motivation**: Clear feedback can motivate employees to work harder and improve their skills, leading to better outcomes.

4. **Encouragement**: Being upfront about issues can encourage open dialogue and collaboration among team members,

APO control¶

APO lives in the same trainer family as DPO and uses the same DPOTrainer class (it is activated by simply choosing a different loss_type). In contrast to DPO that pushes the policy away from the reference (by a relative KL-scaled margin), APO pushes the policy toward a fixed "anchor" score. Generally, APO keeps the policy closer to the reference for the same beta, reducing the risk of over-optimization.

In [12]:

from steerability.algorithms.structural_control.wrappers.trl.apotrainer.control import APO


apo = APO(
    # data
    train_dataset=dpo_train,

    # APO / TRL config 
    output_dir="./tmp/apo_lora",
    per_device_train_batch_size=2,
    num_train_epochs=1,
    learning_rate=5e-5,
    beta=0.1,
    loss_type="apo_zero",     # APO-specific loss
    max_length=1024,
    prompt_format="chat_prompt",
    precompute_ref_log_probs=False,  # inherited default is True (APOArgs subclasses DPOArgs); off for the same reason as the DPO cell
    logging_steps=50,
    report_to="none",
    seed=99,

    # LoRA
    use_peft=True,
    peft_type=PeftType.LORA,
    r=16,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],
    adapter_name="apo",
    
    merge_lora_after_train=False,
)

from steerability.algorithms.structural_control.wrappers.trl.apotrainer.control import APO apo = APO( # data train_dataset=dpo_train, # APO / TRL config output_dir="./tmp/apo_lora", per_device_train_batch_size=2, num_train_epochs=1, learning_rate=5e-5, beta=0.1, loss_type="apo_zero", # APO-specific loss max_length=1024, prompt_format="chat_prompt", precompute_ref_log_probs=False, # inherited default is True (APOArgs subclasses DPOArgs); off for the same reason as the DPO cell logging_steps=50, report_to="none", seed=99, # LoRA use_peft=True, peft_type=PeftType.LORA, r=16, lora_alpha=16, target_modules=["q_proj", "v_proj"], adapter_name="apo", merge_lora_after_train=False, )

Steering and inference proceeds as before.

In [13]:

apo_pipeline = SteeringPipeline(
    model_name_or_path=MODEL_NAME,
    hf_model_kwargs={"trust_remote_code": True, "dtype": torch.float32},
    controls=[apo]
)
apo_pipeline.steer()

apo_pipeline = SteeringPipeline( model_name_or_path=MODEL_NAME, hf_model_kwargs={"trust_remote_code": True, "dtype": torch.float32}, controls=[apo] ) apo_pipeline.steer()

[transformers] The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None, 'pad_token_id': 151643}.

[248/248 01:17, Epoch 1/1]

Step	Training Loss
50	0.980300
100	0.982084
150	0.961600
200	0.984106

In [14]:

print(apo_pipeline.generate(
    messages=[{"role": "user", "content": "Explain why kindness can be strategic."}],
    max_new_tokens=64,
    do_sample=False,
))

print(apo_pipeline.generate( messages=[{"role": "user", "content": "Explain why kindness can be strategic."}], max_new_tokens=64, do_sample=False, ))

Kindness is often considered a strategic tool because it can have significant positive impacts on individuals and organizations alike. Here are several reasons why kindness can be strategically valuable:

1. **Building Trust**: Kindness fosters trust among people. When individuals feel valued and respected, they are more likely to share their thoughts, feelings,

Full-parameter SFT¶

Lastly, to run a full-weight fine-tune set use_peft=False, drop the LoRA arguments, and usually shrink the batch size (because every parameter now receives gradients).

Note: Full fine-tuning can be 10-20 times more memory-intensive than LoRA.

In [15]:

full_sft = SFT(
    train_dataset=sft_train,
    use_peft=False,  # full FT
    output_dir="./tmp/sft_full",
    per_device_train_batch_size=1,
    num_train_epochs=1,
    learning_rate=5e-6,
    logging_steps=50,
    report_to="none",
    seed=7,
)
full_pipeline = SteeringPipeline(
    model_name_or_path=MODEL_NAME,
    hf_model_kwargs={"trust_remote_code": True, "dtype": torch.float32},
    controls=[full_sft]
)
full_pipeline.steer()

full_sft = SFT( train_dataset=sft_train, use_peft=False, # full FT output_dir="./tmp/sft_full", per_device_train_batch_size=1, num_train_epochs=1, learning_rate=5e-6, logging_steps=50, report_to="none", seed=7, ) full_pipeline = SteeringPipeline( model_name_or_path=MODEL_NAME, hf_model_kwargs={"trust_remote_code": True, "dtype": torch.float32}, controls=[full_sft] ) full_pipeline.steer()

[transformers] The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None, 'pad_token_id': 151643}.

[500/500 00:56, Epoch 1/1]

Step	Training Loss
50	1.744838
100	1.584822
150	1.637241
200	1.796518
250	1.574530
300	1.714003
350	1.542453
400	1.565740
450	1.639349
500	1.652375

The wrapper resumes an interrupted run when resume_from_checkpoint names a checkpoint directory written by an earlier run; the other training arguments should match that run.

In [16]:

resume_sft = SFT(
    # data
    train_dataset=sft_train,

    # TRL / Trainer config (forwarded into SFTConfig)
    output_dir="./tmp/sft_lora",
    resume_from_checkpoint="./tmp/sft_lora/checkpoint-100",
    max_length=1024,
    per_device_train_batch_size=4,
    num_train_epochs=1,
    learning_rate=1e-4,
    logging_steps=50,
    save_strategy="steps",
    load_best_model_at_end=False,
    report_to="none",
    seed=42,
    training_args={"save_steps": 50},

    # PEFT (LoRA)
    use_peft=True,
    peft_type=PeftType.LORA,
    r=16,
    lora_alpha=16,
    lora_dropout=0.05,
    target_modules=["q_proj", "v_proj"],
    adapter_name="sft",
)
resume_pipeline = SteeringPipeline(
    model_name_or_path=MODEL_NAME,
    hf_model_kwargs={"trust_remote_code": True, "dtype": torch.float32},
    controls=[resume_sft]
)
resume_pipeline.steer()

resume_sft = SFT( # data train_dataset=sft_train, # TRL / Trainer config (forwarded into SFTConfig) output_dir="./tmp/sft_lora", resume_from_checkpoint="./tmp/sft_lora/checkpoint-100", max_length=1024, per_device_train_batch_size=4, num_train_epochs=1, learning_rate=1e-4, logging_steps=50, save_strategy="steps", load_best_model_at_end=False, report_to="none", seed=42, training_args={"save_steps": 50}, # PEFT (LoRA) use_peft=True, peft_type=PeftType.LORA, r=16, lora_alpha=16, lora_dropout=0.05, target_modules=["q_proj", "v_proj"], adapter_name="sft", ) resume_pipeline = SteeringPipeline( model_name_or_path=MODEL_NAME, hf_model_kwargs={"trust_remote_code": True, "dtype": torch.float32}, controls=[resume_sft] ) resume_pipeline.steer()

[transformers] The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None, 'pad_token_id': 151643}.

[125/125 00:05, Epoch 1/1]

Step	Training Loss

Serving the trained artifact on vLLM¶

Structural controls train on live weights, so on an engine backend the steer phase runs on a temporary in-process model (the stage) that is freed before the engine boots. The exported artifact carries the training across to the engine. A full fine-tune or a merged LoRA run exports a checkpoint (CheckpointArtifact), which overrides the model the engine serves. A LoRA run without merging exports the adapter (LoRAArtifact) instead, which the engine attaches as a LoRA request (enable_lora is set for you). No plugin is involved since the artifact is plain weights, so any vLLM install serves it. Note that running this section requires the toolkit's vllm extra, and the vllm-serve backend works the same way against a running server.

We rerun the earlier LoRA SFT configuration with fresh output directories inside a single pipeline whose backend is the offline engine. The steer() call trains on the staged model exactly as before, and generation then runs on vLLM serving the merged checkpoint. With merge_lora_after_train=False the engine would serve the base model with the adapter attached instead.

In [17]:

from steerability.algorithms.core.execution import BackendSpec

sft_vllm = SFT(
    # data
    train_dataset=sft_train,

    # TRL / Trainer config (forwarded into SFTConfig)
    output_dir="./tmp/sft_lora_vllm",
    max_length=1024,
    per_device_train_batch_size=4,
    num_train_epochs=1,
    learning_rate=1e-4,
    logging_steps=50,
    report_to="none",
    seed=42,

    # PEFT (LoRA)
    use_peft=True,
    peft_type=PeftType.LORA,
    r=16,
    lora_alpha=16,
    lora_dropout=0.05,
    target_modules=["q_proj", "v_proj"],
    adapter_name="sft",

    # merge so the exported artifact is a checkpoint the engine serves directly
    merge_lora_after_train=True,
    merged_output_dir="./tmp/sft_lora_vllm_merged",
)

torch.cuda.empty_cache()

engine_spec = BackendSpec(
    kind="vllm",
    model=MODEL_NAME,
    options={
        "trust_remote_code": True,
        "engine_kwargs": {"gpu_memory_utilization": 0.35, "max_model_len": 2048, "dtype": "bfloat16"},
    },
)

with SteeringPipeline(
    controls=[sft_vllm],
    backend=engine_spec,
    hf_model_kwargs={"trust_remote_code": True, "dtype": torch.float32},
) as vllm_pipeline:
    vllm_pipeline.steer()
    prompt = "Question: What makes the sky look blue?\n\nAnswer:"
    print(vllm_pipeline.generate(prompt, max_new_tokens=64, do_sample=False))

from steerability.algorithms.core.execution import BackendSpec sft_vllm = SFT( # data train_dataset=sft_train, # TRL / Trainer config (forwarded into SFTConfig) output_dir="./tmp/sft_lora_vllm", max_length=1024, per_device_train_batch_size=4, num_train_epochs=1, learning_rate=1e-4, logging_steps=50, report_to="none", seed=42, # PEFT (LoRA) use_peft=True, peft_type=PeftType.LORA, r=16, lora_alpha=16, lora_dropout=0.05, target_modules=["q_proj", "v_proj"], adapter_name="sft", # merge so the exported artifact is a checkpoint the engine serves directly merge_lora_after_train=True, merged_output_dir="./tmp/sft_lora_vllm_merged", ) torch.cuda.empty_cache() engine_spec = BackendSpec( kind="vllm", model=MODEL_NAME, options={ "trust_remote_code": True, "engine_kwargs": {"gpu_memory_utilization": 0.35, "max_model_len": 2048, "dtype": "bfloat16"}, }, ) with SteeringPipeline( controls=[sft_vllm], backend=engine_spec, hf_model_kwargs={"trust_remote_code": True, "dtype": torch.float32}, ) as vllm_pipeline: vllm_pipeline.steer() prompt = "Question: What makes the sky look blue?\n\nAnswer:" print(vllm_pipeline.generate(prompt, max_new_tokens=64, do_sample=False))

[transformers] The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None, 'pad_token_id': 151643}.

[125/125 00:31, Epoch 1/1]

Step	Training Loss
50	1.660352
100	1.578647

(EngineCore pid=484837) 
Loading safetensors checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]

(EngineCore pid=484837) 
Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:07<00:00,  7.52s/it]
(EngineCore pid=484837) 
Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:07<00:00,  7.52s/it]
(EngineCore pid=484837) 

 The sky appears blue because it reflects sunlight. When sunlight hits a surface, it scatters into different colors, including blue. The blue color of the sky is due to the scattering of sunlight by the Earth's atmosphere, which is made up of tiny particles called dust and water vapor. These particles scatter the blue light more

The SFT step needs module access and therefore runs on the stage, i.e., a temporary in-process copy of the model. The generate phase is supported since the configuration exports a checkpoint that the engine can serve, whereas a LoRA configuration without an output_dir would be supported in process only. The staged weights are freed before the engine boots, so the trained copy and the served copy are never in memory at the same time. Exiting the with block shuts the engine down. Note that releasing the offline engine clears vLLM's distributed state for the whole process, so we run this section with no other live vLLM engine in the process.