Add more test cases

Signed-off-by: Hollow Man <hollowman@opensuse.org>

Author: HollowMan6

nemo_rl/utils/weight_transfer_delta_tracker.py

@@ -94,7 +94,7 @@ def __init__(self, config: Mapping[str, Any]) -> None:
         self._baseline_ready_events: dict[str, torch.cuda.Event] = {}
 
     def _baseline_dtype(self, dtype: torch.dtype) -> torch.dtype:
-        return self.delta_dtype if dtype.is_floating_point else dtype
+        return dtype
 
     def should_prewarm_baseline(self) -> bool:
         return (
@@ -272,6 +272,7 @@ def _make_per_tensor_deltas(self, tensors: TensorBatch) -> TensorBatch:
             )
             baseline = self.baseline[name].to(
                 device=tensor.device,
+                dtype=self.delta_dtype,
                 non_blocking=tensor.is_cuda,
                 copy=True,
             )

tests/unit/models/generation/test_registry.py

@@ -15,6 +15,8 @@
 import importlib
 import sys
 
+import pytest
+
 from nemo_rl.models.generation import registry as registry_mod
 from nemo_rl.models.generation.constants import SGLANG_BACKEND, VLLM_BACKEND
 
@@ -26,3 +28,31 @@ def test_builtin_registration_does_not_import_sglang():
 
     assert set(registry.get_registered_backends()) == {VLLM_BACKEND, SGLANG_BACKEND}
     assert "nemo_rl.models.generation.sglang" not in sys.modules
+
+
+def test_custom_generation_backend_factory_is_used():
+    registry = importlib.reload(registry_mod)
+    created = object()
+
+    def custom_factory(cluster, config):
+        assert cluster == "cluster"
+        assert config == {"name": "custom"}
+        return created
+
+    registry.register_generation_backend("custom", custom_factory)
+
+    assert (
+        registry.create_generation(
+            "custom",
+            cluster="cluster",
+            config={"name": "custom"},
+        )
+        is created
+    )
+
+
+def test_unknown_generation_backend_lists_registered_backends():
+    registry = importlib.reload(registry_mod)
+
+    with pytest.raises(ValueError, match="Registered backends"):
+        registry.create_generation("missing", cluster=None, config={})

tests/unit/utils/test_weight_transfer.py

@@ -305,6 +305,21 @@ def test_empty_transfer_header_uses_single_control_broadcast():
     assert len(group.broadcasted_tensors) == 1
 
 
+def test_protocol_rejects_invalid_header_dtype_and_env(monkeypatch):
+    with pytest.raises(ValueError, match="header control"):
+        protocol_mod.decode_header_control(torch.tensor([99, 0, 0, 0]))
+
+    with pytest.raises(ValueError, match="header control"):
+        protocol_mod.decode_header_control(torch.tensor([1, 99, 0, 0]))
+
+    with pytest.raises(ValueError, match="Unsupported tensor dtype"):
+        protocol_mod.dtype_from_name("complex64")
+
+    monkeypatch.setenv("NRL_REFIT_BASELINE_PREWARM_CHUNK_BYTES", "not-an-int")
+    with pytest.raises(ValueError, match="Expected integer value"):
+        protocol_mod.env_int("NRL_REFIT_BASELINE_PREWARM_CHUNK_BYTES", default=1)
+
+
 def test_sparse_indices_encoder_coalesces_small_tensors(monkeypatch):
     monkeypatch.setenv("NRL_REFIT_SPARSE_ENCODE_COALESCE_BYTES", "1024")
     tensors = [
@@ -604,6 +619,42 @@ def test_packed_weight_transfer_non_floating_chunk_uses_full_update():
         assert torch.equal(receiver_state[name], tensor)
 
 
+def test_packed_weight_transfer_sparse_encoding_unavailable_uses_full_update(
+    monkeypatch,
+):
+    tracker = DeltaCompressionTracker(_tracker_config())
+    receiver_state = {}
+    counters = {"full": 0, "delta": 0}
+    initial = [("linear.weight", torch.zeros(4))]
+    updated = [("linear.weight", torch.ones(4))]
+
+    load_full, load_delta = _state_loaders(receiver_state, counters=counters)
+
+    with _cpu_transfer_env():
+        _run_transfer(
+            initial,
+            tracker=tracker,
+            load_full=load_full,
+            load_delta=load_delta,
+            delta_load_batch_size_bytes=1024,
+        )
+        monkeypatch.setattr(
+            sparse_codec_mod,
+            "encode_sparse_indices",
+            Mock(side_effect=sparse_codec_mod.SparseEncodingUnavailable),
+        )
+        _run_transfer(
+            updated,
+            tracker=tracker,
+            load_full=load_full,
+            load_delta=load_delta,
+            delta_load_batch_size_bytes=1024,
+        )
+
+    assert counters == {"full": 2, "delta": 0}
+    assert torch.equal(receiver_state["linear.weight"], updated[0][1])
+
+
 def test_consumer_flushes_async_sparse_decode_before_full_update():
     group = RecordingGroup()
     sparse_delta = torch.zeros(1024)
@@ -732,6 +783,17 @@ def test_full_sync_interval_one_does_not_create_baseline():
     assert tracker.baseline == {}
 
 
+def test_delta_tracker_rejects_invalid_config():
+    with pytest.raises(ValueError, match="Unsupported delta compression dtype"):
+        DeltaCompressionTracker({**_tracker_config(), "dtype": "float64"})
+
+    with pytest.raises(ValueError, match="full_sync_interval"):
+        DeltaCompressionTracker(_tracker_config(full_sync_interval=0))
+
+    with pytest.raises(ValueError, match="sparse_bucket_size_bytes"):
+        DeltaCompressionTracker(_tracker_config(sparse_bucket_size_bytes=0))
+
+
 def test_delta_baseline_prewarm_reuses_allocations(monkeypatch):
     monkeypatch.setenv("NRL_REFIT_BASELINE_PREWARM_CHUNK_BYTES", "32")
     tracker = DeltaCompressionTracker(_tracker_config())
@@ -748,6 +810,8 @@ def test_delta_baseline_prewarm_reuses_allocations(monkeypatch):
     assert tracker.baseline["linear.weight"].shape == torch.Size([2, 4])
     assert tracker.baseline["linear.bias"].dtype == torch.float32
     assert tracker.baseline["embed.weight"].dtype == torch.float16
+    assert tracker._baseline_entries["linear.weight"].arena.dtype == torch.float32
+    assert tracker._baseline_entries["embed.weight"].arena.dtype == torch.float16
 
     tensors = [
         (name, torch.ones(tuple(shape), dtype=dtype))
@@ -762,6 +826,43 @@ def test_delta_baseline_prewarm_reuses_allocations(monkeypatch):
     )
 
 
+def test_delta_baseline_prewarm_can_be_disabled(monkeypatch):
+    monkeypatch.setenv("NRL_REFIT_PREWARM_DELTA_BASELINE", "0")
+    tracker = DeltaCompressionTracker(_tracker_config())
+
+    with _cpu_transfer_env():
+        tracker.prewarm_baseline_from_metadata(
+            {"linear.weight": (torch.Size([4]), torch.float32)}
+        )
+
+    assert tracker.baseline == {}
+    assert tracker._baseline_entries == {}
+
+
+def test_delta_baseline_keeps_source_dtype_but_emits_delta_dtype():
+    tracker = DeltaCompressionTracker(_tracker_config())
+    initial = [
+        ("linear.weight", torch.tensor([1.0, 2.0], dtype=torch.float16)),
+    ]
+    updated = [
+        ("linear.weight", torch.tensor([1.5, 1.0], dtype=torch.float16)),
+    ]
+
+    with _cpu_transfer_env():
+        first_is_delta, _ = tracker.prepare_chunk(initial)
+        tracker.on_sync_succeeded()
+        is_delta, deltas = tracker.prepare_chunk(updated)
+
+    assert not first_is_delta
+    assert tracker.baseline["linear.weight"].dtype == torch.float16
+    assert is_delta
+    assert deltas[0][1].dtype == torch.float32
+    torch.testing.assert_close(
+        deltas[0][1],
+        torch.tensor([0.5, -1.0], dtype=torch.float32),
+    )
+
+
 def test_delta_baseline_records_contiguous_entries(monkeypatch):
     monkeypatch.setenv("NRL_REFIT_BASELINE_PREWARM_CHUNK_BYTES", "1024")
     tracker = DeltaCompressionTracker(_tracker_config())
@@ -1093,6 +1194,41 @@ def test_vllm_delta_transfer_config_rejects_colocated_mode():
         create_vllm_delta_transfer_tracker(generation_config)
 
 
+def test_vllm_delta_transfer_config_returns_none_when_disabled():
+    assert create_vllm_delta_transfer_tracker(None) is None
+    assert create_vllm_delta_transfer_tracker({}) is None
+    assert (
+        create_vllm_delta_transfer_tracker({"delta_compression": {"enabled": False}})
+        is None
+    )
+
+
+def test_vllm_delta_transfer_config_rejects_non_vllm_backend():
+    generation_config = {
+        "backend": "sglang",
+        "vllm_cfg": {"precision": "bfloat16"},
+        "colocated": {"enabled": False},
+        "delta_compression": _tracker_config(),
+    }
+
+    with pytest.raises(ValueError, match="vLLM"):
+        create_vllm_delta_transfer_tracker(generation_config)
+
+
+def test_vllm_delta_transfer_config_creates_tracker_when_enabled():
+    generation_config = {
+        "backend": "vllm",
+        "vllm_cfg": {"precision": "bfloat16"},
+        "colocated": {"enabled": False},
+        "delta_compression": _tracker_config(),
+    }
+
+    tracker = create_vllm_delta_transfer_tracker(generation_config)
+
+    assert isinstance(tracker, DeltaCompressionTracker)
+    assert tracker.delta_dtype == torch.float32
+
+
 def test_vllm_delta_transfer_config_rejects_quantized_modelopt_path():
     generation_config = {
         "backend": "vllm",
@@ -1104,3 +1240,15 @@ def test_vllm_delta_transfer_config_rejects_quantized_modelopt_path():
 
     with pytest.raises(NotImplementedError, match="ModelOpt"):
         create_vllm_delta_transfer_tracker(generation_config)
+
+
+def test_vllm_delta_transfer_config_rejects_fp8_precision():
+    generation_config = {
+        "backend": "vllm",
+        "vllm_cfg": {"precision": "fp8"},
+        "colocated": {"enabled": False},
+        "delta_compression": _tracker_config(),
+    }
+
+    with pytest.raises(NotImplementedError, match="FP8"):
+        create_vllm_delta_transfer_tracker(generation_config)