Function setup

pretrain/main.py:213–332

Sets up the environment and configurations for training a model. Args: data_dir (Optional[Path]): Directory containing the training data. Defaults to None. out_dir (Optional[Path]): Directory to save the outputs. Defaults `models/{model_name}`. gpus (int): Number of

(
    data_dir: Path,
    out_dir: Optional[Path] = None,
    gpus: int = 8,
    global_batch_size=512,
    learning_rate=4e-4,
    micro_batch_size=8,
    max_step=48090
    * 2,  # changes are made to global_batch_size, or dataset, this needs to change.
    warmup_steps=2000,
    log_step_interval=10,
    eval_iters=100,
    save_step_interval=4000,
    eval_step_interval=4000,
    weight_decay=0.1,
    beta1=0.9,
    beta2=0.95,
    grad_clip=1.0,
    decay_lr=True,
    min_lr=4e-5,
    precision: _PRECISION_INPUT = None,
    tpu: bool = False,
    resume: Union[bool, Path] = False,
    model_name='1.5-Pints-2k',
    wandb_name: Optional[str] = None,
    wandb_project: Optional[str] = None,
    tokenizer_dir: Optional[Path] = None,
)

def setup(
    data_dir: Path,
    out_dir: Optional[Path] = None,
    gpus: int = 8,
    global_batch_size=512,
    learning_rate=4e-4,
    micro_batch_size=8,
    max_step=48090
    * 2,  # changes are made to global_batch_size, or dataset, this needs to change.
    warmup_steps=2000,
    log_step_interval=10,
    eval_iters=100,
    save_step_interval=4000,
    eval_step_interval=4000,
    weight_decay=0.1,
    beta1=0.9,
    beta2=0.95,
    grad_clip=1.0,
    decay_lr=True,
    min_lr=4e-5,
    precision: _PRECISION_INPUT = None,
    tpu: bool = False,
    resume: Union[bool, Path] = False,
    model_name='1.5-Pints-2k',
    wandb_name: Optional[str] = None,
    wandb_project: Optional[str] = None,
    tokenizer_dir: Optional[Path] = None,
) -> None:
    """
    Sets up the environment and configurations for training a model.

    Args:
        data_dir (Optional[Path]): Directory containing the training data. Defaults to None.
        out_dir (Optional[Path]): Directory to save the outputs. Defaults `models/{model_name}`.
        gpus (int): Number of GPUs to use. Defaults to 8. GOTCHA: Somehow naming this `devices` will be ignored as a CLI param.
        global_batch_size (int): This is the training batch size. A larger batch size accumulates more gradients and can be faster. A smaller batch size reduces parallelism, but improves stability. See Llama 3.1 paper https://arxiv.org/pdf/2407.21783 (3.4.1 Initial Pre-Training). Defaults to 512.
        learning_rate (float): This is the max learning rate for the optimizer. Defaults to 4e-4.
        micro_batch_size (int): Batch size for each device. Increase this to the memory limit of each GPU for faster training. Defaults to 8.
        max_step (int): Maximum number of training steps, so that the optimizer can schedule the learning rate. Defaults to 48090 * 2.
        warmup_steps (int): Number of steps for learning rate warmup. Defaults to 2000.
        log_step_interval (int): Interval for logging training progress. Defaults to 10.
        eval_iters (int): Number of iterations for evaluation. Defaults to 100.
        save_step_interval (int): Interval for saving the model checkpoints. Defaults to 4000.
        eval_step_interval (int): Interval for evaluation during training. Defaults to 4000.
        weight_decay (float): Weight decay for regularization. Defaults to 0.1.
        beta1 (float): Beta1 parameter for the Adam optimizer. Defaults to 0.9.
        beta2 (float): Beta2 parameter for the Adam optimizer. Defaults to 0.95.
        grad_clip (float): Gradient clipping value. Most people use 1.0. Defaults to 1.0.
        decay_lr (bool): Whether to decay the learning rate during training. Defaults to True.
        min_lr (float): Minimum learning rate if decay_lr is True. Defaults to 4e-5.
        precision (Optional[Literal['32-true', 'bf16-mixed', 'bf16-true', '16-mixed', '16-true']]): Precision type for training. Defaults to None.
        tpu (bool): Whether to use TPU for training. Defaults to False.
        resume (Union[bool, Path]): Whether to resume training from a checkpoint. If a Path is provided, training will resume from the specified checkpoint. Defaults to False.
        model_name (str): Name of the model. Used for organizing outputs and logging. Defaults to '1.5-Pints-2k'.
        wandb_name (Optional[str]): Name for the Weights and Biases run. Defaults to None.
        wandb_project (Optional[str]): Project name for Weights and Biases. Defaults to None.
        tokenizer_dir (Optional[Path]): Directory containing the tokenizer. Must match the tokenizer used for preparing the data. Defaults to None.
        checkpoint_path (Optional[Path]): Path to the checkpoint directory for continuing training. Defaults to None.