Component catalog

Every component you can select by name, and every config preset. This page is generated from the registries and src/toolbox/configs/ when the docs are built, so it always matches the code. toolbox list prints the same names in the terminal.

Config presets

Select with group=preset, e.g. toolbox train experiment=grpo_tiny reward=learned.

Group	Presets
advantage=	dr_grpo, group_relative, rloo
data=	decisions, gsm8k, planetoid, preference_arithmetic, prompts_arithmetic, sft_arithmetic, synthetic_graph, tiny_corpus, traces
env=	cartpole, gymnasium
experiment=	decision_rl_tiny, decision_tiny, dpo_qwen, dpo_tiny, gnn_finetune, gnn_pretrain, grpo_qwen_gsm8k, grpo_tiny, lm_pretrain_tiny, ppo_cartpole, reward_model_tiny, sft_from_traces, sft_qwen, sft_tiny
launcher=	local, slurm, torchrun
model=	actor_critic, gnn_node_classifier, gnn_pretrainer, qwen2_5_0_5b, qwen2_5_0_5b_lora, tiny_gpt, tiny_reward_model
objective=	decision_ce, decision_distill, decision_pg, dpo, graph_contrastive, grpo, ipo, lm_ce, masked_feature, node_classification, ppo_clip, reinforce, reward_bt, sft
optimizer=	adam, adamw, sgd
reward=	exact_match, learned, math_with_format
scheduler=	constant, cosine, linear
strategy=	auto, ddp, fsdp, single
tokenizer=	byte, qwen2_5
trainer=	grpo, ppo_env, preference, supervised

Registered components

Use the name as name: in a config section, or with the matching ExperimentBuilder.with_* method.

activation

Name	Summary	Source
gelu	Applies the Gaussian Error Linear Units function.
identity	A placeholder identity operator that is argument-insensitive.
relu	Applies the rectified linear unit function element-wise.
silu	Applies the Sigmoid Linear Unit (SiLU) function, element-wise.
tanh	Applies the Hyperbolic Tangent (Tanh) function element-wise.

advantage

Name	Summary	Source
batch_normalized	REINFORCE++-style: whiten rewards across the whole batch, ignoring groups.	toolbox.rl.advantage
group_relative	GRPO: (r − mean_group) / (std_group + ε). scale=False gives Dr. GRPO's unscaled variant, which avoids up-weighting near-deterministic groups.	toolbox.rl.advantage
rloo	RLOO (Ahmadian et al., 2024): baseline is the mean reward of the other samples.	toolbox.rl.advantage

attention

Name	Summary	Source
causal_self_attention	Multi-head causal self-attention with optional grouped-query attention (GQA).	toolbox.nn.layers.transformer

callback

Name	Summary	Source
console	Prints metrics every log_every steps (rank 0 only).	toolbox.trainer.callbacks
early_stopping	Stops training when an eval metric has not improved for patience evaluations.	toolbox.trainer.callbacks
jsonl	Appends {"step", "split", **metrics} lines to <output_dir>/metrics.jsonl.	toolbox.trainer.callbacks
memory_monitor	Adds peak accelerator memory and host RSS (MB) to step metrics.	toolbox.trainer.callbacks
memory_writer	Stores RL rollouts with reward ≥ min_reward as episodic memories.	toolbox.memory
sample_logger	Writes a few RL rollouts per step to <output_dir>/samples.jsonl for inspection.	toolbox.trainer.callbacks
tensorboard	Base class; override any subset of the hooks.	toolbox.trainer.callbacks
throughput	Adds step_time_s and (when batches carry input_ids) tokens_per_s.	toolbox.trainer.callbacks
torch_profiler	Profiles a window of steps with torch.profiler and exports a Chrome/Perfetto trace (open <output_dir>/profiler/*.json in https://ui.perfetto.dev).	toolbox.trainer.callbacks
trace_rollouts	Logs every RL rollout (prompt → completion, reward) as a generation action.	toolbox.tracing.integrations
wandb	Base class; override any subset of the hooks.	toolbox.trainer.callbacks

dataset

Name	Summary	Source
decisions	Records {"state", "question", "options", "answer", ["probs"]}.	toolbox.decision.training
planetoid	Cora / CiteSeer / PubMed through PyTorch Geometric (pip install toolbox[gnn]).	toolbox.data.datasets.graph
preference	Preference pairs {"prompt", "chosen", "rejected"} for DPO / reward modelling.	toolbox.data.datasets.text
prompts	Prompts (+ reference answers) for online RL; batches are lists of strings.	toolbox.data.datasets.text
sft	Supervised fine-tuning on {"prompt": ..., "response": ...} JSONL records.	toolbox.data.datasets.text
synthetic_graph	Offline stand-in for Cora-style citation graphs (no download needed).	toolbox.data.datasets.graph
text_lm	Pack a plain-text corpus into fixed-length blocks for next-token prediction.	toolbox.data.datasets.text
traces	Build a training set from <path>/<key>.jsonl traces (see module docstring).	toolbox.tracing.datasets

engine

Name	Summary	Source
hf	Hugging Face transformers model behind the same interface as :class:TorchEngine.	toolbox.serving.engines
openai_http	Client for any OpenAI-compatible /v1/completions server.	toolbox.serving.engines
sglang	In-process SGLang Engine (RadixAttention prefix caching, continuous batching).	toolbox.serving.engines
torch	Runs any toolbox :class:~toolbox.nn.models.base.CausalLM in-process.	toolbox.serving.engines
vllm	In-process vLLM LLM (PagedAttention, continuous batching).	toolbox.serving.engines

env

Name	Summary	Source
cartpole	Classic cart-pole balancing (Barto et al., 1983), dependency-free.	toolbox.rl.envs
gymnasium	Wrap any discrete-action gymnasium environment by id.	toolbox.rl.envs

evaluator

Name	Summary	Source
decision_accuracy	Accuracy plus flip rate: how often the answer changes when options are reversed — a direct measure of position bias (lower is better).	toolbox.eval
generation_reward	Mean reward of greedy completions (e.g. accuracy on a math benchmark).	toolbox.eval

feedforward

Name	Summary	Source
gelu_mlp	Classic transformer feed-forward: W2(gelu(W1 x)) (GPT-2 / BERT).	toolbox.nn.layers.transformer
swiglu	Gated feed-forward: W2(silu(W1 x) * W3 x) (Shazeer, 2020).	toolbox.nn.layers.transformer

graph_conv

Name	Summary	Source
gat	Graph attention (Veličković et al., 2018), multi-head with concatenated outputs.	toolbox.nn.layers.graph
gcn	Graph convolution (Kipf & Welling, 2017): D^-1/2 (A + I) D^-1/2 X W + b.	toolbox.nn.layers.graph
pyg	Any torch_geometric.nn layer by class name, e.g. layer="GATv2Conv".	toolbox.nn.layers.pyg
pyg_gat	PyG GATConv (Veličković et al., 2018); heads concatenated.	toolbox.nn.layers.pyg
pyg_gatv2	PyG GATv2Conv (Brody et al., 2022).	toolbox.nn.layers.pyg
pyg_gcn	PyG GCNConv (Kipf & Welling, 2017).	toolbox.nn.layers.pyg
pyg_gin	PyG GINConv with a 2-layer MLP (Xu et al., 2019).	toolbox.nn.layers.pyg
pyg_sage	PyG SAGEConv (Hamilton et al., 2017).	toolbox.nn.layers.pyg
pyg_transformer	PyG TransformerConv (Shi et al., 2021).	toolbox.nn.layers.pyg
sage	GraphSAGE with mean aggregation (Hamilton et al., 2017).	toolbox.nn.layers.graph

launcher

Name	Summary	Source
local	Train in the current process.	toolbox.launch
slurm	Render an sbatch script from the template and submit it.	toolbox.launch
torchrun	Single-node multi-GPU training: torchrun --nproc-per-node N.	toolbox.launch

memory

Name	Summary	Source
episodic	Append-only experience store with similarity search.	toolbox.memory

mlp

Name	Summary	Source
linear	A single linear projection.	toolbox.nn.layers.common
mlp	Linear → [norm] → activation → [dropout] per hidden layer; the last layer is linear.	toolbox.nn.layers.common

model

Name	Summary	Source
actor_critic	Separate MLP policy (categorical actions) and value networks over a flat observation.	toolbox.nn.models.policy
gnn_encoder	Stack of conv → norm → activation → dropout blocks (the last conv is bare).	toolbox.nn.models.gnn
gnn_node_classifier	Encoder → activation → head, or with head=None the classic Kipf & Welling GNN.	toolbox.nn.models.gnn
gnn_pretrainer	Encoder + heads used only during self-supervised pretraining.	toolbox.nn.models.gnn
hf_causal_lm	Wrap any Hub/local causal LM, optionally with LoRA adapters.	toolbox.nn.models.hf
pyg_model	A torch_geometric.nn.models GNN (GCN, GraphSAGE, GAT, GIN ...).	toolbox.nn.models.external
tiny_gpt	Decoder-only transformer: token embedding → N pre-norm blocks → norm → LM head.	toolbox.nn.models.gpt
tiny_reward_model	A transformer backbone with a scalar head read at the last non-padding token.	toolbox.nn.models.gpt
torch_module	Any nn.Module by import path ("module.path:ClassName"), built with kwargs.	toolbox.nn.models.external

norm

Name	Summary	Source
batchnorm	BatchNorm over the feature dimension of (N, dim) inputs (e.g. node features).	toolbox.nn.layers.common
identity	No normalisation.	toolbox.nn.layers.common
layernorm	LayerNorm over the last dimension.	toolbox.nn.layers.common
rmsnorm	Root-mean-square layer norm (Zhang & Sennrich, 2019): no mean-centering, no bias.	toolbox.nn.layers.transformer

objective

Name	Summary	Source
decision_ce	loss = f(model, batch) — LM pretraining, SFT, reward modelling, GNN training.	toolbox.decision.training
decision_distill	KL(teacher ‖ model) over options; falls back to one-hot gold labels without a teacher.	toolbox.decision.training
decision_pg	REINFORCE over options with a mean-reward baseline and entropy bonus.	toolbox.decision.training
dpo	Direct Preference Optimization (Rafailov et al., 2023), with conservative label smoothing.	toolbox.objectives.preference
graph_contrastive	GRACE (Zhu et al., 2020): two stochastic views per graph, node-level InfoNCE where each node's counterpart in the other view is the positive and all other nodes are negatives.	toolbox.objectives.graph
grpo	Group Relative Policy Optimization (Shao et al., 2024) defaults: KL 0.04, per-seq mean.	toolbox.objectives.policy
ipo	Identity Preference Optimisation (Azar et al., 2023): regress the margin to 1/(2β).	toolbox.objectives.preference
lm_ce	Next-token cross-entropy. With prompt labels masked (IGNORE_INDEX) this is SFT.	toolbox.objectives.language
masked_feature	GraphMAE (Hou et al., 2022): mask node features, reconstruct them with a scaled cosine error (1 − cos(x̂, x))^γ on the masked nodes only.	toolbox.objectives.graph
node_classification	Cross-entropy over the nodes of a split: mask_key while training and eval_mask_key in eval mode, so eval/accuracy is held-out accuracy.	toolbox.objectives.graph
ppo_clip	PPO clipped surrogate (Schulman et al., 2017) with optional KL penalty to a reference.	toolbox.objectives.policy
reinforce	Vanilla policy gradient -A · log π (use with RLOO / baseline advantages).	toolbox.objectives.policy
reward_bt	Train a scalar reward model: -log σ(r(x, y_w) − r(x, y_l)).	toolbox.objectives.preference
sft	Next-token cross-entropy. With prompt labels masked (IGNORE_INDEX) this is SFT.	toolbox.objectives.language
slic	SLiC-HF style hinge loss on the calibrated margin (Zhao et al., 2023).	toolbox.objectives.preference

optimizer

Name	Summary	Source
adam	Adam with L2 weight decay (the classic GNN setting).	toolbox.trainer.optim
adamw	AdamW with weight decay applied only to matrices (not biases or norm scales).	toolbox.trainer.optim
sgd	SGD with momentum.	toolbox.trainer.optim
torch	Any torch.optim optimizer by class name, e.g. {name: torch, optimizer: RAdam}.	toolbox.trainer.optim

reward

Name	Summary	Source
exact_match	1.0 if the extracted answer equals the reference, else 0.0.	toolbox.rewards
format	1.0 if the completion matches pattern (e.g. <think>.*</think><answer>.*</answer>).	toolbox.rewards
length_penalty	-\|len(completion) - target\| / target in characters — discourages rambling.	toolbox.rewards
llm_judge	LLM-as-a-judge: asks an inference engine to grade each completion on [0, max_score].	toolbox.rewards
reward_model	Learned reward from a toolbox reward-model checkpoint (trained with reward_bt).	toolbox.rewards
weighted_sum	Composite reward: Σ wᵢ · rewardᵢ. Children are built from config specs.	toolbox.rewards

scheduler

Name	Summary	Source
constant	Linear warmup, then a constant learning rate.	toolbox.trainer.optim
cosine	Linear warmup, then cosine decay to lr * min_lr_ratio.	toolbox.trainer.optim
hf	Any Hugging Face transformers schedule (linear, cosine, polynomial, warmup_stable_decay ...) via transformers.get_scheduler.	toolbox.trainer.optim
linear	Linear warmup, then linear decay to zero.	toolbox.trainer.optim
torch	Any torch.optim.lr_scheduler class by name, e.g. {name: torch, scheduler: StepLR, step_size: 100}. max_steps is available to configs that need it (T_max).	toolbox.trainer.optim

strategy

Name	Summary	Source
auto	DDP when launched with WORLD_SIZE > 1 (torchrun / srun), otherwise single device.	toolbox.distributed
ddp	Data parallelism with torch.nn.parallel.DistributedDataParallel.	toolbox.distributed
fsdp	Shard parameters, gradients and optimizer state (FSDP) for models that don't fit on one GPU. Wraps each transformer block as its own FSDP unit when auto_wrap is on.	toolbox.distributed
single	How the model and data are distributed across processes.	toolbox.distributed

tokenizer

Name	Summary	Source
byte	UTF-8 bytes as tokens (ids 0–255) plus <pad>, <bos>, <eos>.	toolbox.data.tokenization
hf	Adapter over transformers.AutoTokenizer.	toolbox.data.tokenization

trainer

Name	Summary	Source
grpo	Online RL for language models with verifiable or learned rewards.	toolbox.trainer.trainers
ppo_env	PPO with an actor-critic on a (Gymnasium-style) environment — classic control, games, or any sequential decision problem exposed through :class:~toolbox.rl.envs.Env.	toolbox.trainer.trainers
preference	Offline preference optimisation (DPO / IPO / SLiC) against a frozen reference model.	toolbox.trainer.trainers
supervised	Minimises a :class:SupervisedObjective: LM pretraining, SFT, reward models, GNNs, decision models.	toolbox.trainer.trainers
trl	Runs a TRL trainer from toolbox components. Call :meth:fit like any toolbox trainer.	toolbox.trainer.trl_backend