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, including training hyperparameters, base model architectures, adapter knobs, optimizer knobs, or any other user-provided knobs.

Knob Set Generators

To create a multi-config specification, we need two things: knob set generators for values of knobs and config group generators that take a config with set-value knobs to generate groups of full configs.

We currently support two common knob set generators: List() for a discrete set of values and Range() for sampling from a continuous value interval.

List(values: List[Any])

Parameters:: values – List of discrete values for a knob; all values must be the same python data type.

Range(start: int | float, end: int | float, dtype: str [= "int" or "float"])

Parameters:

start (int | float) – Lower bound of range interval
end (int | float) – Upper bound of range interval
dtype (str [= "int" or "float"]) – Data type of value to be sampled

Notes:

As of this writing, Range() performs uniform sampling within the given interval. We plan to continue expanding this API and add more functionality on this front based on feedback.

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 two common config group generators: RFGridSearch() for grid search and RFRandomSearch() for random search.

More support for AutoML heuristics such as SHA and HyperOpt, as well as an integration with the popular AutoML library Optuna are coming soon.

RFGridSearch(configs: Dict[str, Any] | List[Dict[str, Any]], trainer_type: str [= "SFT" or "DPO" or "GRPO"])

Parameters:

configs (Dict[str, Any] | List[Dict[str, Any]]) – A config dictionary with List() for at least one knob; can be a list of such config dictionaries too
trainer_type (str [= "SFT" or "DPO" or "GRPO"]) – The fine-tuning/post-training control flow to use; names from HF TRL

RFRandomSearch(configs: Dict[str, Any], trainer_type: str [= "SFT" or "DPO" or "GRPO"], num_runs: int, seed: int)

Parameters:

configs (Dict[str, Any]) – A config dictionary with List() or Range() for at least one knob
trainer_type (str [= "SFT" or "DPO" or "GRPO"]) – The fine-tuning/post-training control flow to use; based on HuggingFace TRL
num_runs (int) – Number of runs (full combinations of knob values) to sample in total
seed (int, optional) – Seed for random sampling of knob values to construct combinations

Notes:

For RFGridSearch(), each knob can have either a single value or a List() of values but no knob should have Range() of values; otherwise, it will error out.

For RFRandomSearch(), each knob can have either a single value, or a List() of values, or a: 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 num_runs distinct combinations.

Note that the return types of the config group generators are internal to RapidFire AI and they are usable only within the context of run_fit() in the Experiment class.

Examples:

# Based on SFT tutorial notebook
from rapidfireai.automl import RFModelConfig, RFLoraConfig, RFSFTConfig, List

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"
        )
])

config_group=RFGridSearch(
        configs=config_grid,
        trainer_type="SFT"
)