Training-related arguments
| 6 | |
| 7 | @dataclass |
| 8 | class TrainArgs: |
| 9 | """Training-related arguments""" |
| 10 | |
| 11 | save_interval: Optional[int] = 1000 |
| 12 | """Number of optimizer steps between saving checkpoints""" |
| 13 | log_interval: int = 1 |
| 14 | """Number of iterations between logging calls""" |
| 15 | global_batch_size: int = 64 |
| 16 | """Number of samples between optimizer steps across data-parallel ranks""" |
| 17 | micro_batch_size: int = 4 |
| 18 | """Number of samples per data-parallel rank""" |
| 19 | lr_warmup_steps: int = 100 |
| 20 | """Number of iterations with learning rate warmup active""" |
| 21 | epochs: Optional[int] = None |
| 22 | """Number of epochs to train on""" |
| 23 | # TODO: `pretrain` is the only script using `max_tokens` explicitly. replace it with epoch_size*epochs? |
| 24 | max_tokens: Optional[int] = None |
| 25 | """Total number of tokens to train on""" |
| 26 | max_steps: Optional[int] = None |
| 27 | """Limits the number of optimizer steps to run""" |
| 28 | max_seq_length: Optional[int] = None |
| 29 | """Limits the length of samples""" |
| 30 | tie_embeddings: Optional[bool] = None |
| 31 | """Whether to tie the embedding weights with the language modeling head weights""" |
| 32 | |
| 33 | # Optimization args |
| 34 | |
| 35 | # The default of 1e-3 that came from litgpt caused a loss spike. Hence, lowered to 1e-5. |
| 36 | # See original litgpt code: https://github.com/Lightning-AI/litgpt/blob/64bd9eb32e7fd2bebe8ff187c6f4847b85fe16e8/litgpt/args.py#L36 |
| 37 | learning_rate: float = 1e-5 |
| 38 | weight_decay: float = 0.02 |
| 39 | beta1: float = 0.9 |
| 40 | beta2: float = 0.95 |
| 41 | max_norm: Optional[float] = None |
| 42 | # Tinyllama https://arxiv.org/pdf/2401.02385.pdf |
| 43 | min_lr: float = 4e-5 |
| 44 | |
| 45 | def gradient_accumulation_iters(self, devices: int) -> int: |
| 46 | """Number of iterations between gradient synchronizations""" |
| 47 | gradient_accumulation_iters = self.batch_size(devices) // self.micro_batch_size |
| 48 | assert gradient_accumulation_iters > 0 |
| 49 | return gradient_accumulation_iters |
| 50 | |
| 51 | def batch_size(self, devices: int) -> int: |
| 52 | """Number of samples between optimizer steps per data-parallel rank""" |
| 53 | batch_size = self.global_batch_size // devices |
| 54 | assert batch_size > 0 |
| 55 | return batch_size |
| 56 | |
| 57 | |
| 58 | @dataclass |