API: Multi-Config Specification
===============
The core value of RapidFire AI is in the ability to launch, compare, and dynamically
control multiple configurations (configs) in one go.
It is already common practice in AI to do hyperparameter sweeps via grid search,
random search, or AutoML heuristics to generate knob values.
RapidFire AI generalizes that notion to any type of config knobs, not just
regular hyperparameters, but also base model architectures, prompt schemes,
LoRA adapters, and optimizers (for training), as well as data chunking, embedding,
retrieval, reranking, generation, and prompt schemes (for RAG) and any other
user-defined knobs.
Knob Set Generators
-------
To create a multi-config specification, you need two things: **knob set generators** for
knob values and **config group generators** that take a config with set-valued knobs
to generate groups of full configs.
We currently support two common knob set generators: :func:`List()` for a discrete
set of values and :func:`Range()` for sampling from a continuous value interval.
.. py:function:: List(values: List[Any])
:param values: List of discrete values for a knob. We recommend using values of the same python data type within a given :func:`List()`, although this is not enforced.
:type values: List[Any]
.. py:function:: Range(start: int | float, end: int | float, dtype: str = None, log: bool = False, step: int | float = None, sample_n: int = 3, seed: int = None)
:param start: Lower bound of range interval, inclusive.
:type start: int | float
:param end: Upper bound of range interval, inclusive.
:type end: int | float
:param dtype: Data type of value to be sampled, either :code:`"int"` or :code:`"float"`. Optional; if omitted, it is inferred from :code:`start` and :code:`end`, viz., :code:`"int"` only if both of them are python ints and :code:`"float"` otherwise.
:type dtype: str, optional
:param log: Whether to sample in log space instead of uniformly; both :code:`start` and :code:`end` must then be greater than 0. Default is :code:`False`. Mutually exclusive with :code:`step`.
:type log: bool, optional
:param step: Discretization step; sampled values are then multiples of :code:`step` counting up from :code:`start`. Default is :code:`None`, i.e., the interval is treated as continuous. Mutually exclusive with :code:`log`.
:type step: int | float, optional
:param sample_n: Number of distinct values to draw when the interval has to be enumerated up front instead of sampled per run. This applies to index-affecting knobs under :func:`RFOptuna()` in evals mode — the embedding, text splitter, and prompt manager settings that determine which RAG index a config needs — because that set has to be finite and known before any query runs. It applies whether :code:`build_all_indexes` is :code:`True` or :code:`False`; the flag decides how many of those indexes get built, not whether the range is discretized. It does not change :func:`RFRandomSearch()`, which draws one value per run. Default is 3.
:type sample_n: int, optional
:param seed: Seed for this range's own generator. Default is :code:`None`. Both :func:`RFRandomSearch()` and :func:`RFOptuna()` stamp their own constructor :code:`seed` onto every :func:`Range()` in the search space, so you rarely need to set this yourself.
:type seed: int, optional
**Notes:**
By default, :func:`Range()` performs uniform sampling within the given interval.
Set :code:`log=True` for log-uniform sampling, which is usually a better fit for scale-free knobs
such as learning rate, or set :code:`step` to restrict sampling to a discrete grid of values within
the interval. Giving both :code:`log` and :code:`step` raises an error.
These variants line up one-to-one with Optuna's :code:`suggest_int()` / :code:`suggest_float()`
distributions; so, the same :func:`Range()` specification carries over to :func:`RFOptuna()` below.
Note that the return types of the knob set generators are internal to RapidFire AI and
they are usable only within the context of the config group generators below.
Config Group Generators
-----
We currently support three config group generators: :func:`RFGridSearch()` for grid search,
:func:`RFRandomSearch()` for random search, and :func:`RFOptuna()` for search driven by the popular
AutoML library Optuna, including adaptive pruning of unpromising runs.
More support for AutoML heuristics such as SHA and HyperOpt is coming soon.
Likewise for RAG/context engineering, we also plan to support the AutoML heuristic syftr.
All three take the config(s) to expand as their first argument and :code:`trainer_type` as their
second. The first accepts a single config, a plain python :code:`list` of configs, or a
:func:`List()` of configs. For :func:`run_fit()`, the configs must be :class:`RFModelConfig`
instances; for :func:`run_evals()`, they must be config dictionaries. Passing a
:code:`trainer_type` is what puts a generator in fit mode; leaving it unset puts it in evals mode.
The examples below pass it by keyword, which is also the safer habit: in 0.16.1 and earlier the
second parameter of :func:`RFGridSearch()` and :func:`RFRandomSearch()` was :code:`create_model_fn`,
so a positional :code:`"SFT"` binds to a different argument there.
.. py:function:: RFGridSearch(configs: RFModelConfig | Dict[str, Any] | list, trainer_type: str = None, num_runs: int = 1)
:param configs: The config(s) to expand, with :func:`List()` for at least one knob. An :class:`RFModelConfig` (or a list/:func:`List()` of them) for :func:`run_fit()`; a config dictionary (or a list/:func:`List()` of them) for :func:`run_evals()`.
:type configs: RFModelConfig | Dict[str, Any] | list
:param trainer_type: The fine-tuning/post-training control flow to use: :code:`"SFT"`, :code:`"DPO"`, or :code:`"GRPO"`. Skip this argument for :func:`run_evals()`.
:type trainer_type: str, optional
:param num_runs: Ignored by grid search, which always yields the full cross product of knob values. Default is 1.
:type num_runs: int, optional
.. py:function:: RFRandomSearch(configs: RFModelConfig | Dict[str, Any] | list, trainer_type: str = None, num_runs: int = 1, seed: int = 42)
:param configs: The config(s) to sample from, with :func:`List()` or :func:`Range()` for at least one knob. An :class:`RFModelConfig` (or a list/:func:`List()` of them) for :func:`run_fit()`; a config dictionary (or a list/:func:`List()` of them) for :func:`run_evals()`.
:type configs: RFModelConfig | Dict[str, Any] | list
:param trainer_type: The fine-tuning/post-training control flow to use: :code:`"SFT"`, :code:`"DPO"`, or :code:`"GRPO"`. Skip this argument for :func:`run_evals()`.
:type trainer_type: str, optional
:param num_runs: Number of runs (full combinations of knob values) to sample in total. Default is 1; so, set this explicitly.
:type num_runs: int, optional
:param seed: Seed for the sampling this generator performs. Default is 42.
:type seed: int, optional
**Notes:**
For :func:`RFGridSearch()`, each knob can have either a single value or a :func:`List()` of values.
Grid search expands only :func:`List()` valued knobs. A :func:`Range()` is rejected outright with an
error that names the offending knob path, because grid search enumerates every combination and a
continuous interval has no finite set of values. Either replace it with a :func:`List()` of the exact
values to try, or use :func:`RFRandomSearch()` or :func:`RFOptuna()`, both of which sample from a
:func:`Range()`.
For :func:`RFRandomSearch()`, each knob can have either a single value, or a :func:`List()` of values, or a
:func:`Range()` of values. The semantics of sampling are independently-identically-distributed (IID), i.e.,
we uniformly randomly pick a value from each discrete set and from each continuous set to construct the
knob combination for one run.
Then we repeat that sampling process in an IID way to accumulate :code:`num_runs` distinct combinations.
The :code:`seed` you pass to the :func:`RFRandomSearch()` constructor is what governs which configs it
samples. The :code:`seed` argument of :func:`run_fit()` / :func:`run_evals()` is ignored for those draws
and only governs the surrounding infrastructure such as dataset sharding; so, change the constructor
:code:`seed` if you want a different set of sampled configs.
In evals mode, the generator recognizes exactly three keys for the inference pipeline in your config
dictionary: :code:`"api_config"`, :code:`"vllm_config"`, and :code:`"pipeline"`. Everything else in
the dictionary is treated as a regular knob. Also read
:doc:`the Generator Configs page`.
Note that the return types of the config group generators are internal to RapidFire AI and they are usable only
within the context of :func:`run_fit()` or :func:`run_evals()` in the :class:`Experiment` class.
**Examples:**
.. code-block:: python
# Example 1: Based on SFT tutorial notebook
from rapidfireai.automl import List, RFGridSearch, RFLoraConfig, RFModelConfig, RFSFTConfig
# 2 LoRA adapter configs with different capacities
peft_configs = List([
RFLoraConfig(
r=16, lora_alpha=32, lora_dropout=0.05,
target_modules=["q_proj", "v_proj"], bias="none"
),
RFLoraConfig(
r=128, lora_alpha=256, lora_dropout=0.05,
target_modules=["q_proj","k_proj", "v_proj","o_proj"], bias="none"
)
])
# 2 base models x 2 LoRA configs = 4 combinations in total
config_set = List([
RFModelConfig(
model_name="meta-llama/Llama-3.1-8B-Instruct",
peft_config=peft_configs,
training_args=RFSFTConfig(learning_rate=2e-4, num_train_epochs=2, ...),
model_type="causal_lm",
...
),
RFModelConfig(
model_name="mistralai/Mistral-7B-Instruct-v0.3",
peft_config=peft_configs,
training_args=RFSFTConfig(learning_rate=2e-4, num_train_epochs=2, ...),
model_type="causal_lm",
...
)
])
config_group = RFGridSearch(
configs=config_set,
trainer_type="SFT"
)
# Example 2: Based on GSM8K tutorial notebook
from rapidfireai.automl import List, RFAPIModelConfig, RFGridSearch, RFPromptManager
gemini_config1 = RFAPIModelConfig(
client_config={"max_retries": 2},
endpoint_config={
"provider": "gemini",
"api_key_name": "my_gemini_key",
"api_key": GOOGLE_API_KEY,
"endpoint": {
"model": "gemini-2.5-flash-lite",
"name": "gemini-2.5-flash-lite",
},
},
model_config={"max_completion_tokens": 2048},
rpm_limit=30_000,
tpm_limit=30_000_000,
prompt_manager=fewshot_prompt_manager,
)
gemini_config2 = RFAPIModelConfig(
client_config={"max_retries": 2},
endpoint_config={
"provider": "gemini",
"api_key_name": "my_gemini_key", # reusing the same stored secret
"endpoint": {
"model": "gemini-3-flash-preview",
"name": "gemini-3-flash-preview",
},
},
model_config={
"max_completion_tokens": 4096,
"reasoning_effort": List(["medium", "low"]), # 2 different reasoning levels
},
rpm_limit=20_000,
tpm_limit=20_000_000,
prompt_manager=fewshot_prompt_manager,
)
config_set = {
"api_config": List(
[gemini_config1, gemini_config2]
),
"batch_size": batch_size,
"preprocess_fn": sample_preprocess_fn,
"compute_metrics_fn": sample_compute_metrics_fn,
...
}
config_group = RFGridSearch(config_set)
RFOptuna
^^^^^^^^
:func:`RFOptuna()` is a drop-in replacement for the above two generators that uses
`Optuna `__'s ask-and-tell API to pick knob combinations adaptively.
It supports single- and multi-objective optimization, and it can *prune* an underperforming run at a
shard boundary and spend the freed GPU slot on a fresh config, up to a total :code:`budget` of runs.
.. note::
:code:`from rapidfireai.automl import RFOptuna` always succeeds, but instantiating it raises an
:code:`ImportError` with install instructions if Optuna is not importable. Install it with
:code:`pip install "rapidfireai[optuna]"` (which pins :code:`optuna>=4.0.0`) or with plain
:code:`python -m pip install optuna`. Optuna must go into the *same* interpreter as your Jupyter
kernel; check with :code:`import sys; print(sys.executable)` in a notebook cell and restart the
kernel afterward.
.. py:function:: RFOptuna(configs: RFModelConfig | Dict[str, Any] | list, trainer_type: str = None, n_initial: int = 16, budget: int = 40, objective: str = "minimize:eval_loss", sampler: str = "tpe", pruner: str = "median", seed: int = 42, granularity: str = "shard", build_all_indexes: bool = True)
:param configs: The config template(s) that define the search space, with :func:`List()` or :func:`Range()` for at least one knob; otherwise it errors out. When you give more than one template, Optuna treats the choice of template as a categorical knob of its own.
:type configs: RFModelConfig | Dict[str, Any] | list
:param trainer_type: The fine-tuning/post-training control flow to use: :code:`"SFT"`, :code:`"DPO"`, or :code:`"GRPO"`. Skip this argument for :func:`run_evals()`.
:type trainer_type: str, optional
:param n_initial: Number of configs to launch up front. Default is 16.
:type n_initial: int, optional
:param budget: Maximum total number of runs, viz., the initial ones plus any replacements accepted after pruning. It is raised to :code:`n_initial` if you give a smaller number; set it equal to :code:`n_initial` to disable replacement altogether. Default is 40.
:type budget: int, optional
:param objective: What to optimize, given as :code:`"minimize:"` or :code:`"maximize:"`, e.g., :code:`"maximize:accuracy"`. Comma-separate two or more of these for multi-objective optimization, e.g., :code:`"maximize:rougeL,maximize:bleu"`. Default is :code:`"minimize:eval_loss"`.
:type objective: str, optional
:param sampler: Optuna sampler to use: :code:`"tpe"`, :code:`"cmaes"`, or :code:`"random"`. Default is :code:`"tpe"`.
:type sampler: str, optional
:param pruner: Pruner to use: :code:`"median"` or :code:`None` to disable pruning. :code:`"median"` selects RapidFire AI's *adapted* median pruner described below, not Optuna's stock :code:`MedianPruner`. Ignored for a multi-objective study, which always prunes by Pareto dominance. Default is :code:`"median"`.
:type pruner: str, optional
:param seed: Seed for the algorithm's own stochastic state: the Optuna sampler, every :func:`Range()` generator, and the RNG behind :func:`List()` draws. Default is 42.
:type seed: int, optional
:param granularity: When pruning decisions get evaluated during :func:`run_fit()`: :code:`"shard"` (after every shard) or :code:`"epoch"` (after every epoch). Ignored in evals mode. Default is :code:`"shard"`, which was spelled :code:`"chunk"` in 0.16.1 and earlier; passing :code:`"chunk"` now raises an error.
:type granularity: str, optional
:param build_all_indexes: Evals mode only; ignored in fit mode. When :code:`True` (the default), every RAG index the search space can reach is built up front, so any replacement Optuna suggests already has its retriever available. When :code:`False`, only the indexes the :code:`n_initial` configs need are built. Default is :code:`True`.
:type build_all_indexes: bool, optional
**Notes:**
*Seeding.* Just as with :func:`RFRandomSearch()`, the :code:`seed` given to the :func:`RFOptuna()`
constructor is the single seed governing the algorithm's own draws. The :code:`seed` argument of
:func:`run_fit()` / :func:`run_evals()` is ignored for those draws, so changing it alone will not
change which configs Optuna samples; change it on the :func:`RFOptuna()` constructor instead.
*How pruning actually works.* Optuna's built-in pruners compare a running trial against *completed*
trials, but in RapidFire AI's concurrent loop every trial stays running until the whole search
finishes, so a stock pruner would never have reference data. :code:`pruner="median"` therefore selects
an adapted median pruner that compares each trial against the median of its peers *at the same
progress point*. Concretely:
* In fit mode, a trial reports one value per shard, anchored on the cumulative number of shards it has completed, so trials with different effective batch sizes are still compared at the same point. The comparison uses the trial's best value so far, and NaN values are dropped from the peer pool.
* In evals mode, the comparison uses the trial's latest value at the current shard id, and a NaN current value prunes immediately.
* The adapted comparison applies no startup-trial threshold and no minimum step, so a run can be pruned after its very first shard. Whichever run reaches a given shard first has no peers yet and is never pruned at that point; the trace line for that case reads :code:`no_peers_at_step`.
* Because roughly half of the population sits below the median of the rest by construction, the practical effect resembles successive halving: the field thins out at each successive shard boundary.
* One :code:`[RFOptuna prune-check]` line is logged per single-objective decision with the trial, step, direction, current value, peer values, median, and reason, so you can see exactly why a run was or was not pruned. Multi-objective studies use the Pareto rule below and do not emit these lines.
:code:`pruner=None` disables single-objective pruning entirely: no run is pruned and no replacement is
created. It does *not* disable multi-objective pruning, which always applies and prunes a trial that
is Pareto-dominated by more than half of its active peers at that step.
*Budget accounting.* Only accepted suggestions count against :code:`budget`. A suggestion that is
rejected before launch (see index coverage below) is marked failed, kept out of the sampler's model,
and resampled without consuming budget. Each accepted replacement takes over the slot freed by the
run it replaces, so the number of concurrently active runs stays roughly level until the budget is
exhausted; after that, pruning simply shrinks the field.
*RAG index coverage in evals mode.* A replacement config can only run if its RAG index already exists.
Index-affecting :func:`Range()` knobs (embedding, text splitter, prompt manager) are always discretized
to :code:`Range.sample_n` values, whatever :code:`build_all_indexes` is set to, because the set of
indexes must be finite and known before any query runs. Retrieval-only :code:`search_cfg` and
:code:`reranker_cfg` ranges stay continuous. What the flag controls is how many of those indexes get
built up front.
With the default :code:`build_all_indexes=True`, RapidFire AI enumerates every index-affecting
combination the search space can reach and builds all of them before the run starts, which is the same
index count :func:`RFGridSearch()` would build for the equivalent space. If that enumeration exceeds 64
candidate contexts, it errors out rather than silently launching that many corpus builds. The limit is
checked on the enumerated candidates, before overlapping ones are collapsed by context hash, so a space
whose *distinct* index count is under the limit can still trip it. Narrow the index-affecting knobs,
lower :code:`sample_n` on them, or set :code:`build_all_indexes=False`.
With :code:`build_all_indexes=False`, only what the initial configs need is built, and a replacement
that would need an unbuilt index is rejected and resampled; that is cheaper up front but narrows the
space Optuna can explore, and after 10 failed attempts the freed slot is left unused.
Pruned runs are marked STOPPED, the same status a run you stop yourself with IC Ops ends up in.
The two are still distinguishable after the fact: a pruned run records :code:`ended_by` as
:code:`OPTUNA_PRUNED`, while a manual stop records :code:`INTERACTIVE_CONTROL`. Each replacement run
appears alongside them as it is created.
See the :code:`rf-tutorial-optuna-sft-chatqa-tiny`, :code:`rf-tutorial-optuna-rag-fiqa`, and
:code:`rf-tutorial-optuna-rag-scifact`
:doc:`tutorial notebooks` for end-to-end examples.
.. note::
**How much the sampler actually learns during a run.** The pruner was adapted to compare running
trials with each other; the sampler was not. TPE draws at random until 10 trials have *finished*,
and it only ever models finished trials — it does not read the intermediate values of runs that are
still in flight. In RapidFire AI no trial completes until the whole search finishes, so early in a
run the replacement configs are effectively random draws. Pruned trials do count as finished,
though: once 10 runs have been pruned, TPE starts modelling them, ranking each by its last reported
value. So a search with a generous :code:`budget` becomes progressively better informed, while a
short one stays close to random sampling with aggressive pruning. Feeding the sampler from runs that
are still in flight is tracked in
`issue #304 `__.
**Example:**
.. code-block:: python
# Based on the Optuna SFT tutorial notebook
from rapidfireai.automl import List, RFLoraConfig, RFModelConfig, RFOptuna, RFSFTConfig
config_template = RFModelConfig(
model_name="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
peft_config=RFLoraConfig(
r=List([8, 32]),
lora_alpha=List([16, 32, 64, 128]),
lora_dropout=0.1,
target_modules=List([
["q_proj", "v_proj"],
["q_proj", "k_proj", "v_proj", "o_proj"],
]),
bias="none",
),
training_args=RFSFTConfig(learning_rate=List([1e-3, 1e-4]), ...),
model_type="causal_lm",
...
)
# 4 configs launched up front; up to 4 more created as poor performers are pruned
config_group = RFOptuna(
configs=[config_template],
trainer_type="SFT",
n_initial=4,
budget=8,
objective="maximize:rougeL",
sampler="tpe",
pruner="median",
seed=42,
)