diff --git a/grail/infrastructure/delta_checkpoint.py b/grail/infrastructure/delta_checkpoint.py index 741bd52..42226bb 100644 --- a/grail/infrastructure/delta_checkpoint.py +++ b/grail/infrastructure/delta_checkpoint.py @@ -134,6 +134,8 @@ def apply_sparse_delta( # Infer dtype from base_state if not specified if target_dtype is None: + if not base_state: + raise ValueError("Cannot infer target_dtype: base_state is empty") target_dtype = next(iter(base_state.values())).dtype for name, base_tensor in base_state.items(): diff --git a/tests/unit/infrastructure/test_delta_checkpoint.py b/tests/unit/infrastructure/test_delta_checkpoint.py index bdbb3bf..c43d561 100644 --- a/tests/unit/infrastructure/test_delta_checkpoint.py +++ b/tests/unit/infrastructure/test_delta_checkpoint.py @@ -202,6 +202,16 @@ def test_dtype_conversion(self) -> None: assert result["weight"].dtype == torch.bfloat16 + def test_empty_base_state_raises_value_error(self) -> None: + """Test that empty base_state with target_dtype=None raises ValueError instead of StopIteration.""" + with pytest.raises(ValueError, match="Cannot infer target_dtype: base_state is empty"): + apply_sparse_delta({}, {}, {}, target_dtype=None) + + def test_empty_base_state_with_explicit_dtype(self) -> None: + """Test that empty base_state with explicit target_dtype returns empty dict.""" + result = apply_sparse_delta({}, {}, {}, target_dtype=torch.float32) + assert result == {} + class TestRoundTrip: """Test round-trip: compute delta, apply delta, verify identical."""