| 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.""" |
|
|
