torchtune.utils

Checkpointing

FullModelHFCheckpointer	Checkpointer which reads and writes "full-model" checkpoints in HF's format.
FullModelMetaCheckpointer	Checkpointer which reads and writes "full-model" checkpoints in Meta's format.

Distributed

init_distributed	Initialize torch.distributed.
get_world_size_and_rank	Function that gets the current world size (aka total number of ranks) and rank number of the current trainer.

Mixed Precision

precision.get_autocast	Intelligently determines, based on the dtype if mixed precision training is supported and returns the builtin torch.autocast if applicable.
precision.get_gradient_scaler	Returns a gradient scaler for mixed-precision training.
precision.get_dtype	Get the torch.dtype corresponding to the given precision string.
precision.list_dtypes	Return a list of supported dtypes for finetuning.

Memory Management

memory.set_activation_checkpointing

Utility to setup activation checkpointing and wrap the model for checkpointing.

Metric Logging

metric_logging.WandBLogger	Logger for use w/ Weights and Biases application (https://wandb.ai/).
metric_logging.TensorBoardLogger	Logger for use w/ PyTorch's implementation of TensorBoard (https://pytorch.org/docs/stable/tensorboard.html).
metric_logging.StdoutLogger	Logger to standard output.
metric_logging.DiskLogger	Logger to disk.

Data

collate.padded_collate

Pad a batch of sequences to the longest sequence length in the batch, and convert integer lists to tensors.

Miscellaneous

argparse.TuneRecipeArgumentParser	TuneRecipeArgParser is a helpful utility subclass of the argparse ArgumentParser that adds a builtin argument "config".
logging.get_logger	Get a logger with a stream handler.
get_device	Function that takes or device or device string, verifies it's correct and availabe given the machine and distributed settings, and returns a torch.device.
seed.set_seed	Function that sets seed for pseudo-random number generators across commonly used libraries.