API: Trainer Configs

RapidFire AI’s core APIs for trainer specifications are all thin wrappers around the corresponding APIs of Hugging Face’s TRL libraries.

See also

Most trainer knobs here (learning rate, scheduler, batch size, epochs, etc.) can be modified in a Clone-Modify IC Op, with or without warm-start. See which knobs can be modified on the IC Ops page.

Note

Checkpointing and save_strategy. RapidFire AI manages checkpoints itself: it forces the wrapped Hugging Face trainer’s save_strategy to "no" and always writes a final checkpoint for every run. On top of the usual Hugging Face values, these classes accept a RapidFire AI-specific save_strategy="chunk" sentinel, which is not a valid Hugging Face SaveStrategy. It is translated to "no" for the wrapped TRL class while the original value is preserved on the config, and the worker reads it back to also write an intermediate checkpoint to disk at every chunk boundary, one folder per save. Use it when you want per-chunk checkpoints for post-hoc auditing or external inspection. Any other value, including the "epoch" used in the examples below, leaves checkpointing entirely to RapidFire AI’s own cadence.

RFSFTConfig

This is a wrapper around SFTConfig in HF TRL. The full signature and list of arguments are available on this page.

The only difference here is that the individual arguments (knobs) can be List valued or Range valued in RFSFTConfig. That is how you can specify a base set of knob combinations from which a config group can be produced. Also read the Multi-Config Specification page. Other than the multi-config specification, this class preserves all semantics of Hugging Face’s SFT trainer under the hood.

Examples:

# From the SFT tutorial notebook
RFSFTConfig(
        learning_rate=2e-4,
        lr_scheduler_type = "linear",
        per_device_train_batch_size=4,
        per_device_eval_batch_size=8,
        gradient_accumulation_steps=4,
        num_train_epochs=2,
        logging_steps=5,
        eval_strategy="steps",
        eval_steps=25,
        fp16=True,
        save_strategy="epoch",
)

# Two knobs have list values here
RFSFTConfig(
        learning_rate=List([2e-4, 1e-5]),
        lr_scheduler_type=List(["linear", "cosine"]),
        per_device_train_batch_size=4,
        gradient_accumulation_steps=4,
        per_device_eval_batch_size=8,
        num_train_epochs=2,
        logging_steps=5,
        eval_strategy="steps",
        eval_steps=25,
        fp16=True,
        save_strategy="epoch",
)

For larger LLMs that do not fit on a single GPU and need cross-GPU model partitioning (within a machine) to fit in aggregate GPU memory, specify the FSDP configuration to apply with its usual settings given a dictionary as illustrated below.

Examples:

# From the SFT FSDP Large tutorial notebook
RFSFTConfig(
        ...
        fsdp="full_shard auto_wrap",
fsdp_config={
                "backward_prefetch": "backward_pre",
                "forward_prefetch": False,
                "use_orig_params": False,
                "cpu_ram_efficient_loading": True,
                "offload_params": False,
                "sync_module_states": True,
                "limit_all_gathers": True,
                "sharding_strategy": "FULL_SHARD",
                "auto_wrap_policy": "TRANSFORMER_BASED_WRAP"
        }
)

# From the SFT FSDP Lite tutorial notebook
RFSFTConfig(
        ...
        fsdp="full_shard auto_wrap",
fsdp_config={
                "backward_prefetch": "backward_pre",
                "forward_prefetch": True,
                "use_orig_params": False,
                "cpu_ram_efficient_loading": True,
                "offload_params": True,
                "sync_module_states": True,
                "limit_all_gathers": True,
                "sharding_strategy": "FULL_SHARD",
                "auto_wrap_policy": "TRANSFORMER_BASED_WRAP"
        }
)

RFDPOConfig

This is a wrapper around DPOConfig in HF TRL. The full signature and list of arguments are available on this page.

Again, the only difference here is that the individual arguments (knobs) can be List valued or Range valued in RFDPOConfig. That is how you can specify a base set of knob combinations from which a config group can be produced. Also read the Multi-Config Specification page. Other than the multi-config specification, this class preserves all semantics of Hugging Face’s DPO trainer under the hood.

Example:

# Based on the DPO tutorial notebook; one knob has list of values
base_dpo_config = RFDPOConfig(
        model_adapter_name="default",
        ref_adapter_name="reference",
        force_use_ref_model=False,
        loss_type="sigmoid",
        beta=List([0.1,0.001]),
        max_prompt_length=1024,
        max_completion_length=1024,
        max_length=2048,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        learning_rate=5e-6,
        warmup_ratio=0.1,
        weight_decay=0,
        lr_scheduler_type="linear",
        optim="adamw_8bit",
        num_train_epochs=1,
        logging_strategy="steps",
        logging_steps=1,
        bf16=True,
        save_strategy="epoch",
)

Just like for SFT, you can specify an FSDP configuration for DPO too for larger LLMs that need cross-GPU partitioning (within a machine).

Example:

# From the DPO FSDP Lite notebook
base_dpo_config_lite = RFDPOConfig(
        ...
        fsdp="full_shard auto_wrap",
        fsdp_config={
                "backward_prefetch": "backward_pre",
                "forward_prefetch": True,
                "use_orig_params": False,
                "cpu_ram_efficient_loading": True,
                "offload_params": False,
                "sync_module_states": True,
                "min_num_params": 1000000,
                "limit_all_gathers": True,
                "sharding_strategy": "FULL_SHARD",
                "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
                "activation_checkpointing":False
        }
)

RFGRPOConfig

This is a wrapper around GRPOConfig in HF TRL. The full signature and list of arguments are available on this page.

Again, the only difference here is that the individual arguments (knobs) can be List valued or Range valued in RFGRPOConfig. That is how you can specify a base set of knob combinations from which a config group can be produced. Also read the Multi-Config Specification page. Other than the multi-config specification, this class preserves all semantics of Hugging Face’s GRPO trainer under the hood.

Example:

# Based on the GRPO tutorial notebook
RFGRPOConfig(
        learning_rate=5e-6,
        warmup_ratio=0.1,
        weight_decay=0.1,
        max_grad_norm=0.1,
        adam_beta1=0.9,
        adam_beta2=0.99,
        lr_scheduler_type = "linear",
        per_device_train_batch_size=4,
        gradient_accumulation_steps=4,
        num_generations=8,
        optim ="adamw_8bit",
        num_train_epochs=2,
        max_prompt_length=1024,
        max_completion_length=1024,
        logging_steps=2,
        eval_steps=5,
)

Note

As of this writing, out-of-the-box support for FSDP for GRPO is still in the works. Watch this space for updates.