fix(config): reject non-positive numeric training-config fields at load (#175)

Author: SuperMarioYL

tests/test_training_config_validation.py

@@ -1,11 +1,19 @@
-"""Regression tests for training batch-size validation at config load.
+"""Regression tests for training-config validation at config load.
 
 Guards the empty-dispatch NCCL-hang root cause behind the issue #126 surface:
 a ``training.micro_batch_size`` of 0 makes the derived ``dispatch_batch_size`` 0,
 so ``try_dispatch_batch`` no-op-dispatches (returns True while queuing nothing)
 and every rank blocks on the data-fetcher queue, surfacing as an NCCL
 all-gather timeout. The validator rejects non-positive sizes at config load,
 before any training process starts.
+
+Extends the same fail-closed-at-load idiom (PR #171) to four more numeric
+fields whose non-positive values currently fail *silently* (flat loss from
+``learning_rate=0`` / ``max_grad_norm=0``, sign-flipped gradients from
+``max_grad_norm<0``, inference-pool starvation from ``inference_batch_size=0``)
+or crash *opaquely and late* (post-init ``ZeroDivisionError`` / bare
+``total_steps`` assert from ``draft_accumulation_steps<=0``). Rejecting at
+load surfaces the misconfig before Ray/mooncake/vLLM init.
 """
 
 import pytest
@@ -57,3 +65,82 @@ def test_validate_training_batch_config_accepts_positive():
 
     config = _resolved_config(micro_batch_size=8)
     _validate_training_batch_config(config)  # must not raise
+
+
+# --- Numeric training-config fields (draft_accumulation_steps, learning_rate,
+#     max_grad_norm) and inference_batch_size — PR #171 idiom extended. ---
+
+
+def _resolved_training_config(**overrides):
+    """Build a fully-resolved config from the schema defaults + training overrides."""
+    config = OmegaConf.structured(Config)
+    for key, value in overrides.items():
+        setattr(config.training, key, value)
+    return config
+
+
+def _resolved_inference_config(inference_batch_size: int):
+    """Build a fully-resolved config from the schema defaults + an inference-batch override."""
+    config = OmegaConf.structured(Config)
+    config.inference.inference_batch_size = inference_batch_size
+    return config
+
+
+def test_load_config_rejects_zero_draft_accumulation_steps():
+    """draft_accumulation_steps=0 must fail at load: propagates to global_batch_size=0
+    and crashes post-init (ZeroDivisionError or bare total_steps assert)."""
+    base = _resolved_training_config(draft_accumulation_steps=0)
+    with pytest.raises(ValueError, match="draft_accumulation_steps"):
+        load_config(base_config=base)
+
+
+def test_load_config_rejects_negative_draft_accumulation_steps():
+    """Negative values propagate the same way and must be rejected at load."""
+    base = _resolved_training_config(draft_accumulation_steps=-2)
+    with pytest.raises(ValueError, match="draft_accumulation_steps"):
+        load_config(base_config=base)
+
+
+def test_load_config_rejects_zero_learning_rate():
+    """learning_rate=0 yields silent flat loss (AdamW+scheduler at lr=0) and an
+    untrained checkpoint; reject at load."""
+    base = _resolved_training_config(learning_rate=0)
+    with pytest.raises(ValueError, match="learning_rate"):
+        load_config(base_config=base)
+
+
+def test_load_config_rejects_negative_learning_rate():
+    """Negative learning rates hit a late scheduler assert; reject at load."""
+    base = _resolved_training_config(learning_rate=-1e-4)
+    with pytest.raises(ValueError, match="learning_rate"):
+        load_config(base_config=base)
+
+
+def test_load_config_rejects_zero_max_grad_norm():
+    """max_grad_norm=0 zeroes all grads via clip_grad_norm_ (silent flat loss);
+    reject at load."""
+    base = _resolved_training_config(max_grad_norm=0)
+    with pytest.raises(ValueError, match="max_grad_norm"):
+        load_config(base_config=base)
+
+
+def test_load_config_rejects_negative_max_grad_norm():
+    """Negative max_grad_norm sign-flips grads (silent gradient ascent); reject at load."""
+    base = _resolved_training_config(max_grad_norm=-0.5)
+    with pytest.raises(ValueError, match="max_grad_norm"):
+        load_config(base_config=base)
+
+
+def test_load_config_rejects_zero_inference_batch_size():
+    """inference_batch_size=0 starves the inference pool (silent dispatch spin) and
+    vLLM rejects max_num_seqs=0 late after init; reject at load."""
+    base = _resolved_inference_config(inference_batch_size=0)
+    with pytest.raises(ValueError, match="inference_batch_size"):
+        load_config(base_config=base)
+
+
+def test_load_config_rejects_negative_inference_batch_size():
+    """Negative inference batch sizes starve the pool the same way; reject at load."""
+    base = _resolved_inference_config(inference_batch_size=-1)
+    with pytest.raises(ValueError, match="inference_batch_size"):
+        load_config(base_config=base)

torchspec/config/inference_config.py

@@ -29,6 +29,8 @@
 from dataclasses import dataclass, field
 from typing import Any, Dict, Optional
 
+from omegaconf import DictConfig
+
 from torchspec.config.mooncake_config import MooncakeConfig
 
 
@@ -198,3 +200,12 @@ class HFInferenceConfig:
     trust_remote_code: bool = False
     aux_hidden_states_layers: Optional[list[int]] = None
     mooncake_config: Optional[MooncakeConfig] = None
+
+
+def _validate_inference_batch_config(config: DictConfig) -> None:
+    if config.inference.inference_batch_size <= 0:
+        raise ValueError(
+            f"inference_batch_size must be > 0 (got {config.inference.inference_batch_size}); "
+            f"0 causes inference-pool starvation (silent dispatch spin) and vLLM rejects "
+            f"max_num_seqs=0 late after init"
+        )

torchspec/config/train_config.py

@@ -26,7 +26,7 @@
 
 from omegaconf import DictConfig, OmegaConf
 
-from torchspec.config.inference_config import InferenceConfig
+from torchspec.config.inference_config import InferenceConfig, _validate_inference_batch_config
 from torchspec.data.utils import is_local_data_path
 from torchspec.utils.logging import logger
 
@@ -300,6 +300,31 @@ def _validate_training_batch_config(config: DictConfig) -> None:
         )
 
 
+def _validate_training_numeric_config(config: DictConfig) -> None:
+    """Reject non-positive values that would otherwise fail silently or crash late.
+
+    These fields share the fail-closed-at-load principle of #171: a misconfig that
+    produces flat loss, sign-flipped gradients, or an opaque post-init crash is
+    preferable to surfacing after expensive Ray/mooncake init.
+    """
+    if config.training.draft_accumulation_steps <= 0:
+        raise ValueError(
+            f"draft_accumulation_steps must be > 0 (got {config.training.draft_accumulation_steps}); "
+            f"<=0 propagates into global_batch_size/lr_total_steps and crashes post-init "
+            f"(ZeroDivisionError or bare total_steps assert)"
+        )
+    if config.training.learning_rate <= 0:
+        raise ValueError(
+            f"learning_rate must be > 0 (got {config.training.learning_rate}); "
+            f"0 yields silent flat loss (untrained checkpoint), <0 hits a late assert"
+        )
+    if config.training.max_grad_norm <= 0:
+        raise ValueError(
+            f"max_grad_norm must be > 0 (got {config.training.max_grad_norm}); "
+            f"0 zeroes all grads (silent flat loss), <0 sign-flips grads (silent gradient ascent)"
+        )
+
+
 def _save_config_snapshot(config: DictConfig) -> None:
     """Save the resolved config to output_dir/config.yaml if output_dir is set."""
     output_dir = OmegaConf.select(config, "output_dir", default=None)
@@ -347,6 +372,8 @@ def load_config(
     _validate_offline_config(config)
     _validate_vocab_mapping_config(config)
     _validate_training_batch_config(config)
+    _validate_training_numeric_config(config)
+    _validate_inference_batch_config(config)
 
     if save_snapshot:
         _save_config_snapshot(config)