wehub-resource-sync
69 lines
3.1 KiB
Python
69 lines
3.1 KiB
Python
# --- Configuration ---
|
|
|
|
import torch
|
|
|
|
# Define vocabulary size and transformer configuration (3 Billion)
|
|
VOCAB_SIZE = 50304 # Number of unique tokens in the vocabulary
|
|
CONTEXT_LENGTH = 512 # Maximum sequence length for the model
|
|
N_EMBED = 2048 # Dimension of the embedding space
|
|
N_HEAD = 16 # Number of attention heads in each transformer block
|
|
N_BLOCKS = 64 # Number of transformer blocks in the model
|
|
|
|
# Paths to training and development datasets
|
|
TRAIN_PATH = "data/train/pile_train.h5" # File path for the training dataset
|
|
DEV_PATH = "data/val/pile_dev.h5" # File path for the validation dataset
|
|
|
|
# Transformer training parameters
|
|
T_BATCH_SIZE = 32 # Number of samples per training batch
|
|
T_CONTEXT_LENGTH = 16 # Context length for training batches
|
|
T_TRAIN_STEPS = 200000 # Total number of training steps
|
|
T_EVAL_STEPS = 1000 # Frequency (in steps) to perform evaluation
|
|
T_EVAL_ITERS = 250 # Number of iterations to evaluate the model
|
|
T_LR_DECAY_STEP = 50000 # Step at which to decay the learning rate
|
|
T_LR = 5e-4 # Initial learning rate for training
|
|
T_LR_DECAYED = 5e-5 # Learning rate after decay
|
|
T_OUT_PATH = "models/transformer_B.pt" # Path to save the trained model
|
|
T_CHECKPOINT_STEPS = 0 # Save periodic checkpoints every N steps (0 disables)
|
|
T_KEEP_LAST_CHECKPOINTS = 3 # Number of periodic checkpoints to keep (0 keeps all)
|
|
T_CHECKPOINT_DIR = None # Optional checkpoint directory override
|
|
|
|
# Memory-optimisation knobs (all OFF by default => unchanged behaviour/numerics).
|
|
# These let large configs fit in less VRAM; enable them on the CLI or here. See issue #5.
|
|
USE_AMP = False # bf16/fp16 autocast (CUDA only; ignored on CPU)
|
|
AMP_DTYPE = "bf16" # "bf16" (no GradScaler) or "fp16" (GradScaler)
|
|
USE_GRADIENT_CHECKPOINTING = False # recompute block activations in backward to save VRAM
|
|
GRAD_ACCUM_STEPS = 1 # micro-batches per optimizer step (effective batch x N)
|
|
REPORT_MEMORY_BUDGET = False # print a rough VRAM budget before training (CUDA only)
|
|
|
|
# Device configuration
|
|
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
|
|
# Store all configurations in a dictionary for easy access and modification
|
|
default_config = {
|
|
'vocab_size': VOCAB_SIZE,
|
|
'context_length': CONTEXT_LENGTH,
|
|
'n_embed': N_EMBED,
|
|
'n_head': N_HEAD,
|
|
'n_blocks': N_BLOCKS,
|
|
'train_path': TRAIN_PATH,
|
|
'dev_path': DEV_PATH,
|
|
't_batch_size': T_BATCH_SIZE,
|
|
't_context_length': T_CONTEXT_LENGTH,
|
|
't_train_steps': T_TRAIN_STEPS,
|
|
't_eval_steps': T_EVAL_STEPS,
|
|
't_eval_iters': T_EVAL_ITERS,
|
|
't_lr_decay_step': T_LR_DECAY_STEP,
|
|
't_lr': T_LR,
|
|
't_lr_decayed': T_LR_DECAYED,
|
|
't_out_path': T_OUT_PATH,
|
|
't_checkpoint_steps': T_CHECKPOINT_STEPS,
|
|
't_keep_last_checkpoints': T_KEEP_LAST_CHECKPOINTS,
|
|
't_checkpoint_dir': T_CHECKPOINT_DIR,
|
|
'use_amp': USE_AMP,
|
|
'amp_dtype': AMP_DTYPE,
|
|
'use_gradient_checkpointing': USE_GRADIENT_CHECKPOINTING,
|
|
'grad_accum_steps': GRAD_ACCUM_STEPS,
|
|
'report_memory_budget': REPORT_MEMORY_BUDGET,
|
|
'device': DEVICE,
|
|
}
|