amachine.am_transformers.am_training_config
1import json 2from dataclasses import dataclass, asdict 3from pathlib import Path 4import hashlib 5from typing import Any 6from typing import Literal 7import torch 8 9@dataclass 10class TrainingConfig: 11 12 # Data 13 data: str # Training parquet file 14 metadata: str | None = None # Path to metadata JSON containing the alphabet 15 input_col : str = "symbol_index" # Column name for tokens in parquet 16 eval_data: str | None = None # Eval parquet file (optional) 17 resume_reset_data : bool = False # to start from the beginning of the data (e.g., if resuming from new dataset) 18 tokenizer_type: Literal["symbolic", "text"] = "symbolic" 19 pretokenized : bool = False 20 21 # Important to set these correctly 22 add_eos: bool = False # whether to add eos token between rows 23 shuffle_buffer: int = 0 # Reservoir size for data shuffling (0->no shuffling) 24 num_workers: int = 1 # Dataset workers (for sequential chunks, num_workers must be 1) 25 26 # Paths 27 model_dir: str = "./" # Directory containing config.json and tokenizer.json 28 output_dir: str = "./checkpoints" # Where to save model weights and trainer state 29 resume: str | None = None # Checkpoint directory to resume from 30 31 # Architecture / sequence 32 seq_len: int = 1024 # Context window length 33 attn_implementation: Literal["eager", "sdpa", "flash_attention_2", None] = "sdpa" 34 35 # Trainable parameters 36 train_layers : Literal[ "all" ] | list[int] = "all" 37 train_input_embeddings : bool = True 38 train_output_layer : bool = True 39 40 # Optimisation 41 optimizer : Literal[ "adamw", "sgdw", "sdg", "novograd" ] = "adamw" 42 momentum : float = 0.9 # only used if optimizer is sdg 43 44 batch_size: int = 64 # Per-device micro-batch size 45 grad_accum: int = 4 # Number of steps before optimizer update 46 lr: float = 2e-4 # Peak learning rate 47 min_lr_ratio: float = 0.001 # Final LR as a fraction of peak LR 48 weight_decay: float = 0.05 # L2 regularization strength 49 grad_clip: float = 1.0 # Maximum gradient norm (0.0 to disable) 50 steps: int = 3200 # Total number of training steps 51 warmup_steps: int = 200 # Linear LR warmup duration 52 53 lr_schedule: Literal["cosine", "wsd", "constant" ] = "cosine" 54 decay_steps: int | None = None # WSD only: explicit length of the final decay phase (in steps). 55 # Leave as None (recommended) to derive it from decay_ratio instead, 56 # which stays sensible if you extend `steps` across resumes. 57 # Stable phase runs from warmup_steps to (steps - decay_steps). 58 decay_ratio: float = 0.2 # WSD only, used when decay_steps is None: decay phase length as a 59 # fraction of (steps - warmup_steps). 0.1-0.2 matches common WSD practice. 60 beta1: float = 0.9 # AdamW beta1 61 beta2: float = 0.995 # AdamW beta2 62 eps: float = 1e-8 # AdamW epsilon; default is fine, revisit only if updates 63 # stall despite nonzero gradients late in a long run 64 65 # Loss spike detection 66 spike_factor: float = 4.5 # Skip optimizer step if loss > spike_factor * loss_ema 67 spike_ema_alpha: float = 0.98 # EMA decay for reference loss used in spike detection 68 spike_warmup_steps: int = 100 # Wait N steps before enabling spike detection 69 70 track_grad_norms: bool = True 71 track_weight_norms: bool = True 72 73 # Logging / saving / eval 74 log_every: int = 10 # Steps between console logging 75 save_every: int = 200 # Steps between checkpoint saves 76 eval_every: int = 200 # Run eval loop every N steps (requires eval_data) 77 eval_steps: int = 100 # Number of eval batches per eval run 78 test_prompts: list[str] | None = None 79 test_prompt_gen_len : int = 300 80 81 # System 82 seed: int = 42 # Random seed for reproducibility 83 pdtype: Literal[ "bfloat16", "float32" ] = "bfloat16" 84 wdtype: Literal[ "bfloat16", "float32" ] = "float32" 85 86 compile: bool = False # Enable torch.compile (requires PyTorch >= 2.0) 87 no_grad_ckpt: bool = False # Disable gradient checkpointing (default is False/Enabled) 88 89 # aim Tracking 90 aim_repo: str | None = None # Path to Aim repository for experiment tracking 91 experiment : str | None = None 92 experiment_params : dict[str,Any] | None = None 93 94 # Analysis 95 analysis_ema_alpha : float = 0.98 # EMA decay for loss used in for information theoretic analysis 96 baseline_entropy_rate : float | None = None # if you know the entropy rate intrinsic to the data generator 97 baseline_excess_entropy : float | None = None # if you know the excess entropy of the data generator 98 99 extra_loss_terms : list[str] | None = None 100 101 compute_effective_rank : bool = False 102 compute_channel_stats : bool = False 103 104 evaluate_basis_independence : bool = False 105 data_metadata : dict[str,Any] | None = None 106 metrics_dir : str | None = None 107 108 def __post_init__(self): 109 110 if self.warmup_steps >= self.steps: 111 raise ValueError( 112 f"warmup_steps ({self.warmup_steps}) must be less than steps ({self.steps})" 113 ) 114 115 if not (0.0 <= self.min_lr_ratio <= 1.0): 116 raise ValueError( 117 f"min_lr_ratio ({self.min_lr_ratio}) must be in [0.0, 1.0]" 118 ) 119 120 if self.lr_schedule not in ("cosine", "wsd", "constant"): 121 raise ValueError( 122 f"lr_schedule ({self.lr_schedule}) must be one of 'cosine', 'wsd', 'constant'" 123 ) 124 125 if self.lr_schedule == "wsd": 126 if self.decay_steps is not None: 127 if not (0 < self.decay_steps <= self.steps - self.warmup_steps): 128 raise ValueError( 129 f"decay_steps ({self.decay_steps}) must be in " 130 f"(0, steps - warmup_steps] = (0, {self.steps - self.warmup_steps}]" 131 ) 132 elif not (0.0 < self.decay_ratio <= 1.0): 133 raise ValueError( 134 f"decay_ratio ({self.decay_ratio}) must be in (0.0, 1.0] " 135 f"when decay_steps is not set" 136 ) 137 138 if self.eval_data is not None and self.eval_every % self.log_every != 0: 139 raise ValueError( 140 f"eval_every ({self.eval_every}) must be a multiple of " 141 f"log_every ({self.log_every}), otherwise eval results are silently dropped" 142 ) 143 144 if self.spike_factor <= 1.0: 145 raise ValueError( 146 f"spike_factor ({self.spike_factor}) must be > 1.0" 147 ) 148 149 if not (0.0 < self.spike_ema_alpha < 1.0): 150 raise ValueError( 151 f"spike_ema_alpha ({self.spike_ema_alpha}) must be in (0.0, 1.0)" 152 ) 153 154 if self.tokenizer_type == "symbolic": 155 if not self.metadata or not Path(self.metadata).exists(): 156 raise FileNotFoundError(f"Metadata file containing alphabet is required for symbolic mode: {self.metadata}") 157 158 if not Path(self.data).exists(): 159 raise FileNotFoundError(f"Training data not found: {self.data}") 160 161 if not Path(self.data).exists(): 162 raise FileNotFoundError(f"Training data not found: {self.data}") 163 164 if self.eval_data is not None and not Path(self.eval_data).exists(): 165 raise FileNotFoundError(f"Eval data not found: {self.eval_data}") 166 167 model_dir = Path(self.model_dir) 168 169 if not model_dir.exists(): 170 raise FileNotFoundError(f"Model directory not found: {self.model_dir}") 171 172 if not (model_dir / "config.json").exists(): 173 raise FileNotFoundError(f"Model config (config.json) not found in {self.model_dir}") 174 175 if not (model_dir / "tokenizer.json").exists(): 176 raise FileNotFoundError(f"tokenizer.json not found in {self.model_dir}") 177 178 if self.metadata : 179 with open(self.metadata) as f: 180 self.data_metadata = json.load( f ) 181 182 valid_attns = ("eager", "sdpa", "flash_attention_2", None) 183 if self.attn_implementation not in valid_attns: 184 raise ValueError(f"attn_implementation must be one of {valid_attns}, got '{self.attn_implementation}'") 185 186 if self.data_metadata is not None and "machine_metadata" in self.data_metadata : 187 188 # To avoid cluttering aim with extranious machine metadata, 189 # retain only name, description, and complexity 190 191 updated_metadata = { 192 "name" : self.data_metadata[ "machine_metadata" ].get( "name", "" ), 193 "description" : self.data_metadata[ "machine_metadata" ].get( "description", "" ) 194 } 195 196 if "complexity" in self.data_metadata[ "machine_metadata" ] : 197 198 if "h_mu" in self.data_metadata["machine_metadata"][ "complexity" ] : 199 self.baseline_entropy_rate = self.data_metadata["machine_metadata"][ "complexity" ][ "h_mu" ] 200 201 if "E" in self.data_metadata["machine_metadata"][ "complexity" ] : 202 self.baseline_excess_entropy = self.data_metadata["machine_metadata"][ "complexity" ][ "E" ] 203 204 updated_metadata[ "complexity" ] = self.data_metadata["machine_metadata"][ "complexity" ] 205 206 self.data_metadata[ "machine_metadata" ] = updated_metadata 207 208 print( f"shuffle_buffer: {self.shuffle_buffer}" ) 209 210 def to_dict(self) -> dict: 211 return asdict(self) 212 213 def save(self, path): 214 with open(path, "w") as f: 215 json.dump(asdict(self), f, indent=2) 216 217 def stable_hash(self) -> str: 218 """ 219 Deterministic hash of config contents. 220 Independent of key order. 221 """ 222 cfg_str = json.dumps(self.to_dict(), sort_keys=True) 223 return hashlib.sha256(cfg_str.encode("utf-8")).hexdigest()[:10] # short hash 224 225 @classmethod 226 def from_json(cls, json_path: str, **overrides): 227 with open(json_path) as f: 228 cfg = json.load(f) 229 cfg.update(overrides) 230 return cls(**cfg) 231 232 233@dataclass 234class TrainerState : 235 global_step : int = 0 236 global_tokens : int = 0 237 epoch : int = 0 238 spike_loss_ema : int | None = None # cold-starts on first step; restored from ckpt 239 analysis_loss_ema : int | None = None 240 spikes_skipped : int = 0 241 consecutive_spikes : int = 0 242 max_consecutive_spikes : int = 5 243 batches_consumed : int = 0 244 245 # Inside TrainerState 246 def process_loss(self, step_loss: float, config: TrainingConfig, is_warm: bool) -> bool: 247 248 if self.spike_loss_ema is None: 249 self.spike_loss_ema = step_loss 250 self.analysis_loss_ema = step_loss 251 252 is_spike = is_warm and (step_loss > config.spike_factor * self.spike_loss_ema) 253 254 if is_spike: 255 self.consecutive_spikes += 1 256 self.spikes_skipped += 1 257 return True 258 259 self.consecutive_spikes = 0 260 self.spike_loss_ema = (config.spike_ema_alpha * self.spike_loss_ema) + ((1 - config.spike_ema_alpha) * step_loss) 261 self.analysis_loss_ema = (config.analysis_ema_alpha * self.analysis_loss_ema) + ((1 - config.analysis_ema_alpha) * step_loss) 262 263 return False 264 265@dataclass 266class TrainerContext: 267 model: Any 268 optimizer: Any 269 device: torch.device 270 ptdtype: torch.dtype 271 use_amp: bool 272 pad_id: int 273 eos_id: int 274 training_config: TrainingConfig 275 aim_run: Any = None
@dataclass
class
TrainingConfig:
10@dataclass 11class TrainingConfig: 12 13 # Data 14 data: str # Training parquet file 15 metadata: str | None = None # Path to metadata JSON containing the alphabet 16 input_col : str = "symbol_index" # Column name for tokens in parquet 17 eval_data: str | None = None # Eval parquet file (optional) 18 resume_reset_data : bool = False # to start from the beginning of the data (e.g., if resuming from new dataset) 19 tokenizer_type: Literal["symbolic", "text"] = "symbolic" 20 pretokenized : bool = False 21 22 # Important to set these correctly 23 add_eos: bool = False # whether to add eos token between rows 24 shuffle_buffer: int = 0 # Reservoir size for data shuffling (0->no shuffling) 25 num_workers: int = 1 # Dataset workers (for sequential chunks, num_workers must be 1) 26 27 # Paths 28 model_dir: str = "./" # Directory containing config.json and tokenizer.json 29 output_dir: str = "./checkpoints" # Where to save model weights and trainer state 30 resume: str | None = None # Checkpoint directory to resume from 31 32 # Architecture / sequence 33 seq_len: int = 1024 # Context window length 34 attn_implementation: Literal["eager", "sdpa", "flash_attention_2", None] = "sdpa" 35 36 # Trainable parameters 37 train_layers : Literal[ "all" ] | list[int] = "all" 38 train_input_embeddings : bool = True 39 train_output_layer : bool = True 40 41 # Optimisation 42 optimizer : Literal[ "adamw", "sgdw", "sdg", "novograd" ] = "adamw" 43 momentum : float = 0.9 # only used if optimizer is sdg 44 45 batch_size: int = 64 # Per-device micro-batch size 46 grad_accum: int = 4 # Number of steps before optimizer update 47 lr: float = 2e-4 # Peak learning rate 48 min_lr_ratio: float = 0.001 # Final LR as a fraction of peak LR 49 weight_decay: float = 0.05 # L2 regularization strength 50 grad_clip: float = 1.0 # Maximum gradient norm (0.0 to disable) 51 steps: int = 3200 # Total number of training steps 52 warmup_steps: int = 200 # Linear LR warmup duration 53 54 lr_schedule: Literal["cosine", "wsd", "constant" ] = "cosine" 55 decay_steps: int | None = None # WSD only: explicit length of the final decay phase (in steps). 56 # Leave as None (recommended) to derive it from decay_ratio instead, 57 # which stays sensible if you extend `steps` across resumes. 58 # Stable phase runs from warmup_steps to (steps - decay_steps). 59 decay_ratio: float = 0.2 # WSD only, used when decay_steps is None: decay phase length as a 60 # fraction of (steps - warmup_steps). 0.1-0.2 matches common WSD practice. 61 beta1: float = 0.9 # AdamW beta1 62 beta2: float = 0.995 # AdamW beta2 63 eps: float = 1e-8 # AdamW epsilon; default is fine, revisit only if updates 64 # stall despite nonzero gradients late in a long run 65 66 # Loss spike detection 67 spike_factor: float = 4.5 # Skip optimizer step if loss > spike_factor * loss_ema 68 spike_ema_alpha: float = 0.98 # EMA decay for reference loss used in spike detection 69 spike_warmup_steps: int = 100 # Wait N steps before enabling spike detection 70 71 track_grad_norms: bool = True 72 track_weight_norms: bool = True 73 74 # Logging / saving / eval 75 log_every: int = 10 # Steps between console logging 76 save_every: int = 200 # Steps between checkpoint saves 77 eval_every: int = 200 # Run eval loop every N steps (requires eval_data) 78 eval_steps: int = 100 # Number of eval batches per eval run 79 test_prompts: list[str] | None = None 80 test_prompt_gen_len : int = 300 81 82 # System 83 seed: int = 42 # Random seed for reproducibility 84 pdtype: Literal[ "bfloat16", "float32" ] = "bfloat16" 85 wdtype: Literal[ "bfloat16", "float32" ] = "float32" 86 87 compile: bool = False # Enable torch.compile (requires PyTorch >= 2.0) 88 no_grad_ckpt: bool = False # Disable gradient checkpointing (default is False/Enabled) 89 90 # aim Tracking 91 aim_repo: str | None = None # Path to Aim repository for experiment tracking 92 experiment : str | None = None 93 experiment_params : dict[str,Any] | None = None 94 95 # Analysis 96 analysis_ema_alpha : float = 0.98 # EMA decay for loss used in for information theoretic analysis 97 baseline_entropy_rate : float | None = None # if you know the entropy rate intrinsic to the data generator 98 baseline_excess_entropy : float | None = None # if you know the excess entropy of the data generator 99 100 extra_loss_terms : list[str] | None = None 101 102 compute_effective_rank : bool = False 103 compute_channel_stats : bool = False 104 105 evaluate_basis_independence : bool = False 106 data_metadata : dict[str,Any] | None = None 107 metrics_dir : str | None = None 108 109 def __post_init__(self): 110 111 if self.warmup_steps >= self.steps: 112 raise ValueError( 113 f"warmup_steps ({self.warmup_steps}) must be less than steps ({self.steps})" 114 ) 115 116 if not (0.0 <= self.min_lr_ratio <= 1.0): 117 raise ValueError( 118 f"min_lr_ratio ({self.min_lr_ratio}) must be in [0.0, 1.0]" 119 ) 120 121 if self.lr_schedule not in ("cosine", "wsd", "constant"): 122 raise ValueError( 123 f"lr_schedule ({self.lr_schedule}) must be one of 'cosine', 'wsd', 'constant'" 124 ) 125 126 if self.lr_schedule == "wsd": 127 if self.decay_steps is not None: 128 if not (0 < self.decay_steps <= self.steps - self.warmup_steps): 129 raise ValueError( 130 f"decay_steps ({self.decay_steps}) must be in " 131 f"(0, steps - warmup_steps] = (0, {self.steps - self.warmup_steps}]" 132 ) 133 elif not (0.0 < self.decay_ratio <= 1.0): 134 raise ValueError( 135 f"decay_ratio ({self.decay_ratio}) must be in (0.0, 1.0] " 136 f"when decay_steps is not set" 137 ) 138 139 if self.eval_data is not None and self.eval_every % self.log_every != 0: 140 raise ValueError( 141 f"eval_every ({self.eval_every}) must be a multiple of " 142 f"log_every ({self.log_every}), otherwise eval results are silently dropped" 143 ) 144 145 if self.spike_factor <= 1.0: 146 raise ValueError( 147 f"spike_factor ({self.spike_factor}) must be > 1.0" 148 ) 149 150 if not (0.0 < self.spike_ema_alpha < 1.0): 151 raise ValueError( 152 f"spike_ema_alpha ({self.spike_ema_alpha}) must be in (0.0, 1.0)" 153 ) 154 155 if self.tokenizer_type == "symbolic": 156 if not self.metadata or not Path(self.metadata).exists(): 157 raise FileNotFoundError(f"Metadata file containing alphabet is required for symbolic mode: {self.metadata}") 158 159 if not Path(self.data).exists(): 160 raise FileNotFoundError(f"Training data not found: {self.data}") 161 162 if not Path(self.data).exists(): 163 raise FileNotFoundError(f"Training data not found: {self.data}") 164 165 if self.eval_data is not None and not Path(self.eval_data).exists(): 166 raise FileNotFoundError(f"Eval data not found: {self.eval_data}") 167 168 model_dir = Path(self.model_dir) 169 170 if not model_dir.exists(): 171 raise FileNotFoundError(f"Model directory not found: {self.model_dir}") 172 173 if not (model_dir / "config.json").exists(): 174 raise FileNotFoundError(f"Model config (config.json) not found in {self.model_dir}") 175 176 if not (model_dir / "tokenizer.json").exists(): 177 raise FileNotFoundError(f"tokenizer.json not found in {self.model_dir}") 178 179 if self.metadata : 180 with open(self.metadata) as f: 181 self.data_metadata = json.load( f ) 182 183 valid_attns = ("eager", "sdpa", "flash_attention_2", None) 184 if self.attn_implementation not in valid_attns: 185 raise ValueError(f"attn_implementation must be one of {valid_attns}, got '{self.attn_implementation}'") 186 187 if self.data_metadata is not None and "machine_metadata" in self.data_metadata : 188 189 # To avoid cluttering aim with extranious machine metadata, 190 # retain only name, description, and complexity 191 192 updated_metadata = { 193 "name" : self.data_metadata[ "machine_metadata" ].get( "name", "" ), 194 "description" : self.data_metadata[ "machine_metadata" ].get( "description", "" ) 195 } 196 197 if "complexity" in self.data_metadata[ "machine_metadata" ] : 198 199 if "h_mu" in self.data_metadata["machine_metadata"][ "complexity" ] : 200 self.baseline_entropy_rate = self.data_metadata["machine_metadata"][ "complexity" ][ "h_mu" ] 201 202 if "E" in self.data_metadata["machine_metadata"][ "complexity" ] : 203 self.baseline_excess_entropy = self.data_metadata["machine_metadata"][ "complexity" ][ "E" ] 204 205 updated_metadata[ "complexity" ] = self.data_metadata["machine_metadata"][ "complexity" ] 206 207 self.data_metadata[ "machine_metadata" ] = updated_metadata 208 209 print( f"shuffle_buffer: {self.shuffle_buffer}" ) 210 211 def to_dict(self) -> dict: 212 return asdict(self) 213 214 def save(self, path): 215 with open(path, "w") as f: 216 json.dump(asdict(self), f, indent=2) 217 218 def stable_hash(self) -> str: 219 """ 220 Deterministic hash of config contents. 221 Independent of key order. 222 """ 223 cfg_str = json.dumps(self.to_dict(), sort_keys=True) 224 return hashlib.sha256(cfg_str.encode("utf-8")).hexdigest()[:10] # short hash 225 226 @classmethod 227 def from_json(cls, json_path: str, **overrides): 228 with open(json_path) as f: 229 cfg = json.load(f) 230 cfg.update(overrides) 231 return cls(**cfg)
TrainingConfig( data: str, metadata: str | None = None, input_col: str = 'symbol_index', eval_data: str | None = None, resume_reset_data: bool = False, tokenizer_type: Literal['symbolic', 'text'] = 'symbolic', pretokenized: bool = False, add_eos: bool = False, shuffle_buffer: int = 0, num_workers: int = 1, model_dir: str = './', output_dir: str = './checkpoints', resume: str | None = None, seq_len: int = 1024, attn_implementation: Literal['eager', 'sdpa', 'flash_attention_2', None] = 'sdpa', train_layers: Union[Literal['all'], list[int]] = 'all', train_input_embeddings: bool = True, train_output_layer: bool = True, optimizer: Literal['adamw', 'sgdw', 'sdg', 'novograd'] = 'adamw', momentum: float = 0.9, batch_size: int = 64, grad_accum: int = 4, lr: float = 0.0002, min_lr_ratio: float = 0.001, weight_decay: float = 0.05, grad_clip: float = 1.0, steps: int = 3200, warmup_steps: int = 200, lr_schedule: Literal['cosine', 'wsd', 'constant'] = 'cosine', decay_steps: int | None = None, decay_ratio: float = 0.2, beta1: float = 0.9, beta2: float = 0.995, eps: float = 1e-08, spike_factor: float = 4.5, spike_ema_alpha: float = 0.98, spike_warmup_steps: int = 100, track_grad_norms: bool = True, track_weight_norms: bool = True, log_every: int = 10, save_every: int = 200, eval_every: int = 200, eval_steps: int = 100, test_prompts: list[str] | None = None, test_prompt_gen_len: int = 300, seed: int = 42, pdtype: Literal['bfloat16', 'float32'] = 'bfloat16', wdtype: Literal['bfloat16', 'float32'] = 'float32', compile: bool = False, no_grad_ckpt: bool = False, aim_repo: str | None = None, experiment: str | None = None, experiment_params: dict[str, typing.Any] | None = None, analysis_ema_alpha: float = 0.98, baseline_entropy_rate: float | None = None, baseline_excess_entropy: float | None = None, extra_loss_terms: list[str] | None = None, compute_effective_rank: bool = False, compute_channel_stats: bool = False, evaluate_basis_independence: bool = False, data_metadata: dict[str, typing.Any] | None = None, metrics_dir: str | None = None)
def
stable_hash(self) -> str:
218 def stable_hash(self) -> str: 219 """ 220 Deterministic hash of config contents. 221 Independent of key order. 222 """ 223 cfg_str = json.dumps(self.to_dict(), sort_keys=True) 224 return hashlib.sha256(cfg_str.encode("utf-8")).hexdigest()[:10] # short hash
Deterministic hash of config contents. Independent of key order.
@dataclass
class
TrainerState:
234@dataclass 235class TrainerState : 236 global_step : int = 0 237 global_tokens : int = 0 238 epoch : int = 0 239 spike_loss_ema : int | None = None # cold-starts on first step; restored from ckpt 240 analysis_loss_ema : int | None = None 241 spikes_skipped : int = 0 242 consecutive_spikes : int = 0 243 max_consecutive_spikes : int = 5 244 batches_consumed : int = 0 245 246 # Inside TrainerState 247 def process_loss(self, step_loss: float, config: TrainingConfig, is_warm: bool) -> bool: 248 249 if self.spike_loss_ema is None: 250 self.spike_loss_ema = step_loss 251 self.analysis_loss_ema = step_loss 252 253 is_spike = is_warm and (step_loss > config.spike_factor * self.spike_loss_ema) 254 255 if is_spike: 256 self.consecutive_spikes += 1 257 self.spikes_skipped += 1 258 return True 259 260 self.consecutive_spikes = 0 261 self.spike_loss_ema = (config.spike_ema_alpha * self.spike_loss_ema) + ((1 - config.spike_ema_alpha) * step_loss) 262 self.analysis_loss_ema = (config.analysis_ema_alpha * self.analysis_loss_ema) + ((1 - config.analysis_ema_alpha) * step_loss) 263 264 return False
TrainerState( global_step: int = 0, global_tokens: int = 0, epoch: int = 0, spike_loss_ema: int | None = None, analysis_loss_ema: int | None = None, spikes_skipped: int = 0, consecutive_spikes: int = 0, max_consecutive_spikes: int = 5, batches_consumed: int = 0)
247 def process_loss(self, step_loss: float, config: TrainingConfig, is_warm: bool) -> bool: 248 249 if self.spike_loss_ema is None: 250 self.spike_loss_ema = step_loss 251 self.analysis_loss_ema = step_loss 252 253 is_spike = is_warm and (step_loss > config.spike_factor * self.spike_loss_ema) 254 255 if is_spike: 256 self.consecutive_spikes += 1 257 self.spikes_skipped += 1 258 return True 259 260 self.consecutive_spikes = 0 261 self.spike_loss_ema = (config.spike_ema_alpha * self.spike_loss_ema) + ((1 - config.spike_ema_alpha) * step_loss) 262 self.analysis_loss_ema = (config.analysis_ema_alpha * self.analysis_loss_ema) + ((1 - config.analysis_ema_alpha) * step_loss) 263 264 return False
@dataclass
class
TrainerContext:
266@dataclass 267class TrainerContext: 268 model: Any 269 optimizer: Any 270 device: torch.device 271 ptdtype: torch.dtype 272 use_amp: bool 273 pad_id: int 274 eos_id: int 275 training_config: TrainingConfig 276 aim_run: Any = None
TrainerContext( model: Any, optimizer: Any, device: torch.device, ptdtype: torch.dtype, use_amp: bool, pad_id: int, eos_id: int, training_config: TrainingConfig, aim_run: Any = None)
training_config: TrainingConfig