| """ |
| Automated Quality Checker & Anomaly Detector for SLM training and inference. |
| Verifies model weights, gradient norms, loss stability, checkpoint integrity, and tokenizer consistency. |
| """ |
|
|
| import os |
| import math |
| import torch |
| import torch.nn as nn |
| from typing import List, Tuple, Dict, Any, Optional |
| from slm.tokenizer.bpe import BPETokenizer |
| from slm.utils.logger import get_logger |
|
|
| logger = get_logger("slm.utils.quality") |
|
|
|
|
| class QualityChecker: |
| """Performs pre-flight, in-flight, and post-flight quality checks and anomaly detection.""" |
|
|
| @staticmethod |
| def check_tensor_sanities(tensor: torch.Tensor, name: str = "tensor") -> bool: |
| """ |
| Verifies that a tensor contains no NaN or Inf values. |
| Raises ValueError if anomalies are found. |
| """ |
| if torch.isnan(tensor).any(): |
| raise ValueError(f"Quality Check Failed: NaN detected in {name}!") |
| if torch.isinf(tensor).any(): |
| raise ValueError(f"Quality Check Failed: Inf detected in {name}!") |
| return True |
|
|
| @classmethod |
| def check_model_weights(cls, model: nn.Module) -> Dict[str, Any]: |
| """Inspects all model parameters for NaNs, Infs, and extreme values.""" |
| total_params = 0 |
| total_nan = 0 |
| total_inf = 0 |
|
|
| for name, param in model.named_parameters(): |
| if param.requires_grad: |
| total_params += param.numel() |
| if torch.isnan(param).any(): |
| total_nan += 1 |
| if torch.isinf(param).any(): |
| total_inf += 1 |
|
|
| if total_nan > 0 or total_inf > 0: |
| raise ValueError(f"Quality Check Failed: Model weights contained {total_nan} NaN tensors and {total_inf} Inf tensors!") |
|
|
| return { |
| "total_parameters": total_params, |
| "nan_count": total_nan, |
| "inf_count": total_inf, |
| "is_healthy": True |
| } |
|
|
| @classmethod |
| def check_gradient_norms(cls, model: nn.Module, max_norm_threshold: float = 100.0) -> float: |
| """ |
| Computes total gradient norm across all model parameters. |
| Raises ValueError if gradient explosion (norm > threshold) or NaN/Inf is detected. |
| """ |
| total_sq_norm = 0.0 |
| for p in model.parameters(): |
| if p.grad is not None: |
| cls.check_tensor_sanities(p.grad, name="param_gradient") |
| param_norm = p.grad.data.norm(2) |
| total_sq_norm += param_norm.item() ** 2 |
|
|
| total_norm = math.sqrt(total_sq_norm) |
|
|
| if math.isnan(total_norm) or math.isinf(total_norm): |
| raise ValueError("Quality Check Failed: Gradient norm is NaN or Inf!") |
| if total_norm > max_norm_threshold: |
| logger.warning(f"High gradient norm detected: {total_norm:.4f} > threshold {max_norm_threshold}") |
|
|
| return total_norm |
|
|
| @staticmethod |
| def check_loss_health(loss_val: float, history: Optional[List[float]] = None) -> bool: |
| """Checks if current loss value is finite and not exploding relative to history.""" |
| if math.isnan(loss_val) or math.isinf(loss_val): |
| raise ValueError(f"Quality Check Failed: Training loss is non-finite ({loss_val})!") |
|
|
| if history and len(history) >= 5: |
| avg_recent = sum(history[-5:]) / 5.0 |
| if loss_val > avg_recent * 4.0 and loss_val > 10.0: |
| logger.warning(f"Loss Spike Detected: Current loss {loss_val:.4f} is 4x recent average {avg_recent:.4f}") |
|
|
| return True |
|
|
| @staticmethod |
| def verify_tokenizer_consistency(tokenizer: BPETokenizer, test_sample: str = "LawSLM test sentence.") -> bool: |
| """Verifies tokenizer encode/decode consistency.""" |
| ids = tokenizer.encode(test_sample, add_special_tokens=False) |
| decoded = tokenizer.decode(ids, skip_special_tokens=True) |
| if decoded.strip() != test_sample.strip(): |
| logger.warning(f"Tokenizer mismatch! Input: {test_sample} != Decoded: {decoded}") |
| return False |
| return True |
|
|
| @staticmethod |
| def verify_checkpoint(filepath: str) -> bool: |
| """Verifies that a checkpoint file exists, is non-empty, and loadable by PyTorch.""" |
| if not os.path.isfile(filepath): |
| raise FileNotFoundError(f"Checkpoint file not found: {filepath}") |
| if os.path.getsize(filepath) == 0: |
| raise ValueError(f"Checkpoint file is empty: {filepath}") |
|
|
| try: |
| checkpoint = torch.load(filepath, map_location="cpu", weights_only=False) |
| if not isinstance(checkpoint, dict) or "model_state_dict" not in checkpoint: |
| raise ValueError(f"Invalid checkpoint format in {filepath}: missing 'model_state_dict'") |
| except Exception as e: |
| raise ValueError(f"Corrupted checkpoint file {filepath}: {str(e)}") |
|
|
| return True |
|
|
