File size: 2,569 Bytes
5c43f61 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """
Training configuration for Vortex models.
Covers both 7B and 13B variants with hardware-specific optimizations.
"""
import torch
TRAINING_CONFIG = {
# Training hyperparameters
"learning_rate": 3e-4,
"weight_decay": 0.1,
"beta1": 0.9,
"beta2": 0.95,
"clip_grad_norm": 1.0,
# Batch sizing
"global_batch_size": 512, # tokens per batch
"micro_batch_size": 8, # per GPU
"gradient_accumulation_steps": 4,
# Training schedule
"max_steps": 100000,
"warmup_steps": 2000,
"save_interval": 5000,
"eval_interval": 1000,
"log_interval": 100,
# Mixed precision
"use_amp": True,
"amp_dtype": torch.bfloat16,
# Optimizer
"optimizer": "AdamW",
"use_fused": True, # fused AdamW if available
# Curriculum learning stages (as fractions of max_steps)
"curriculum_stages": [
{"name": "foundation", "start": 0.0, "end": 0.2}, # 0-20%
{"name": "domain", "start": 0.2, "end": 0.5}, # 20-50%
{"name": "reasoning", "start": 0.5, "end": 0.8}, # 50-80%
{"name": "integration", "start": 0.8, "end": 1.0}, # 80-100%
],
# Loss weights (science-aware loss)
"loss_weights": {
"lm_loss": 1.0,
"equation_loss": 0.3,
"domain_loss": 0.1,
"citation_loss": 0.1,
"numerical_loss": 0.2,
},
# Checkpointing
"checkpoint_dir": "checkpoints",
"save_optimizer_state": True,
"save_scheduler_state": True,
# Logging
"log_dir": "logs",
"use_wandb": False,
"wandb_project": "vortex-scientific",
# Data loading
"num_workers": 8,
"prefetch_factor": 2,
"pin_memory": True,
# Device configuration
"device": "cuda", # or "mps" for Apple Silicon
"use_mps": False,
# Quantization (for 13B on 8GB VRAM)
"quantization": None, # None, "int8", "int4"
}
# Hardware-specific overrides
TRAINING_CONFIG_7B_CUDA = TRAINING_CONFIG.copy()
TRAINING_CONFIG_7B_CUDA.update({
"device": "cuda",
"quantization": None,
"micro_batch_size": 8,
})
TRAINING_CONFIG_13B_CUDA = TRAINING_CONFIG.copy()
TRAINING_CONFIG_13B_CUDA.update({
"device": "cuda",
"quantization": "int8", # 13B needs INT8 on 8GB
"micro_batch_size": 4,
})
TRAINING_CONFIG_MPS = TRAINING_CONFIG.copy()
TRAINING_CONFIG_MPS.update({
"device": "mps",
"use_mps": True,
"use_amp": False, # MPS doesn't support bfloat16 AMP well
"micro_batch_size": 4,
})
|