File size: 1,805 Bytes
7b592f7 | 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 | from dataclasses import dataclass
from typing import Optional
@dataclass
class TrainArgs:
"""Training-related arguments."""
save_interval: Optional[int] = 1000
"""Number of optimizer steps between checkpoint saves."""
log_interval: int = 1
"""Number of iterations between log lines."""
global_batch_size: int = 64
"""Total samples per optimizer step across all data-parallel ranks."""
micro_batch_size: int = 4
"""Samples per data-parallel rank per forward/backward pass."""
lr_warmup_steps: Optional[int] = 100
"""Number of warmup iterations (linear ramp from 0 to max_lr)."""
epochs: Optional[int] = None
"""Number of epochs to train."""
max_steps: Optional[int] = None
"""Hard cap on optimizer steps (overrides epochs if reached first)."""
max_seq_length: Optional[int] = None
"""Truncate samples longer than this."""
def gradient_accumulation_iters(self, devices: int, num_nodes: int = 1) -> int:
n = self.batch_size(devices, num_nodes) // self.micro_batch_size
assert n > 0
return n
def batch_size(self, devices: int, num_nodes: int = 1) -> int:
n = self.global_batch_size // (devices * num_nodes)
assert n > 0
return n
@dataclass
class EvalArgs:
"""Evaluation-related arguments."""
interval: int = 600
"""Number of optimizer steps between validation runs."""
max_new_tokens: Optional[int] = None
"""Generation length cap (only used by validation-time generation paths)."""
max_iters: int = 100
"""Max number of validation batches per run."""
initial_validation: bool = False
"""Run one validation pass before the first training step."""
final_validation: bool = True
"""Run one validation pass after training completes."""
|