Print, but only on rank 0.

Return the path to the cache directory for this user.

Gather and stack/cat values from all processes, if there are multiple processes.

Write a checkpoint to disk.

Slice a batch into chunks, and move each chunk to the specified device.

Save policy, optimizer, and scheduler state to disk, gathering from all processes and saving only on the rank 0 process.

Compute the log probabilities of the given labels under the given logits. Args: logits: Logits of the model (unnormalized). Shape: (batch

Run the given model on the given batch of inputs, concatenating the chosen and rejected inputs together. We do this to avoid doing

Disable dropout in a model.

Format a dictionary for printing.

Get an iterator over batches of data. Stops after n_epochs or n_examples, whichever comes first. Args: names: Names of datasets to use.

Compute the SFT or DPO loss and other metrics for the given batch of inputs.

Strip HTML tags from a string, except for <code> tags (which contain real code in the StackExchange answers).

Tokenize a single batch element. At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation in case th

A trainer subclass that uses PyTorch FSDP to shard the model across multiple GPUs. This trainer will shard both the policy and ref

Clip the gradient norm of the parameters of a non-FSDP policy.

Concatenate the chosen and rejected inputs into a single tensor. Args: batch: A batch of data. Must contain the keys 'chosen_input_id

Extract the anthropic prompt from a prompt and response pair.

Generate samples from the policy (and reference model, if doing DPO training) for the given batch of inputs.

Get the class of a block from a model, using the block's class name.

Returns a collate function for the given tokenizer. The collate function takes a list of examples (dicts, where values are lists of

Load the given dataset by name. Supported by default are 'shp', 'hh', and 'se'.

Load the Anthropic Helpful-Harmless dataset from Huggingface and convert it to the necessary format. The dataset is converted to a diction

Create a local directory to store outputs for this run, and return its path.

Load the StackExchange dataset from Huggingface, and return a dict of prompts and responses. See get_hh for the format. We strip the HTML

Load the Stanford Human Preferences dataset from Huggingface and convert it to the necessary format. See hh for the format. We filter preferen

Main entry point for training. Validates config, creates/initializes model(s), and kicks off worker process(es).

Compute the DPO loss for a batch of policy and reference model log probabilities. Args: policy_chosen_logps: Log probabilities of the pol

Save policy, optimizer, and scheduler state to disk.

Begin either SFT or DPO training, with periodic evaluation.

Main function for each worker process (may be only 1 for BasicTrainer/TensorParallelTrainer).

A trainer for a language model, supporting either SFT or DPO training. If multiple GPUs are present, naively splits the model a

A trainer subclass that uses TensorParallel to shard the model across multiple GPUs. Based on https://github.com/BlackSamorez/tensor_paral

Temporarily set the random seed, and then restore it when exiting the context.

Clip the gradient norm of the parameters of an FSDP policy, gathering the gradients across all GPUs.

Print the amount of GPU memory currently allocated for each GPU.

Save (unsharded) policy state to disk.

Returns True if str_a and str_b match up to spaces, False otherwise.