85 lines
3.3 KiB
Diff
85 lines
3.3 KiB
Diff
diff --git a/frontier/cc_backend/backends/vidur_cc_backend.py b/frontier/cc_backend/backends/vidur_cc_backend.py
|
|
index ca1983a..0c57f05 100644
|
|
--- a/frontier/cc_backend/backends/vidur_cc_backend.py
|
|
+++ b/frontier/cc_backend/backends/vidur_cc_backend.py
|
|
@@ -882,2 +882,21 @@ class VidurCCBackend(BaseCCBackend):
|
|
- # Fallback to analytical if not in cache
|
|
- logger.debug(f"num_tokens={num_tokens} not in cache, using analytical fallback")
|
|
+ # The precomputed lookup is capped at 100k elements, while realistic
|
|
+ # TP payloads are commonly much larger. A cache miss does not mean the
|
|
+ # measured-data model is unavailable: predict on demand and memoize the
|
|
+ # exact payload instead of silently switching model families.
|
|
+ with self._cache_lock:
|
|
+ model = self._models.get("all_reduce")
|
|
+ if model is not None:
|
|
+ features = pd.DataFrame({"num_tokens": [num_tokens]})
|
|
+ result = float(model.predict(features)[0])
|
|
+ with self._cache_lock:
|
|
+ self._predictions["all_reduce"][(num_tokens,)] = result
|
|
+ logger.debug(
|
|
+ f"predict_allreduce: data_size={data_size_bytes}, num_tokens={num_tokens}, "
|
|
+ f"result={result:.6f} ms (ML model, on-demand cache miss)"
|
|
+ )
|
|
+ return max(0.0, result)
|
|
+
|
|
+ logger.debug(
|
|
+ f"num_tokens={num_tokens} not in cache and model unavailable, "
|
|
+ "using analytical fallback"
|
|
+ )
|
|
diff --git a/tests/unit/test_vidur_cc_large_payload.py b/tests/unit/test_vidur_cc_large_payload.py
|
|
new file mode 100644
|
|
index 0000000..7e87aa7
|
|
--- /dev/null
|
|
+++ b/tests/unit/test_vidur_cc_large_payload.py
|
|
@@ -0,0 +1,50 @@
|
|
+from __future__ import annotations
|
|
+
|
|
+import threading
|
|
+import unittest
|
|
+
|
|
+import numpy as np
|
|
+
|
|
+from frontier.cc_backend.backends.vidur_cc_backend import VidurCCBackend
|
|
+
|
|
+
|
|
+class RecordingModel:
|
|
+ def __init__(self, value: float) -> None:
|
|
+ self.value = value
|
|
+ self.features = []
|
|
+
|
|
+ def predict(self, features):
|
|
+ self.features.append(features.copy())
|
|
+ return np.array([self.value])
|
|
+
|
|
+
|
|
+class VidurCCLargePayloadTest(unittest.TestCase):
|
|
+ def test_cache_miss_uses_measured_model_and_memoizes(self) -> None:
|
|
+ backend = object.__new__(VidurCCBackend)
|
|
+ backend._cache_lock = threading.RLock()
|
|
+ backend._num_devices = 2
|
|
+ backend._predictions = {"all_reduce": {(100000,): 0.1}}
|
|
+ model = RecordingModel(0.321)
|
|
+ backend._models = {"all_reduce": model}
|
|
+ backend._analytical_fallback_allreduce = lambda *_: self.fail(
|
|
+ "analytical fallback must not run when the measured model exists"
|
|
+ )
|
|
+
|
|
+ data_size_bytes = 2048 * 2048 * 2
|
|
+ expected_elements = data_size_bytes // 2
|
|
+ first = backend.predict_allreduce(data_size_bytes, num_devices=2)
|
|
+ second = backend.predict_allreduce(data_size_bytes, num_devices=2)
|
|
+
|
|
+ self.assertEqual(first, 0.321)
|
|
+ self.assertEqual(second, 0.321)
|
|
+ self.assertEqual(len(model.features), 1)
|
|
+ self.assertEqual(
|
|
+ int(model.features[0].iloc[0]["num_tokens"]), expected_elements
|
|
+ )
|
|
+ self.assertEqual(
|
|
+ backend._predictions["all_reduce"][(expected_elements,)], 0.321
|
|
+ )
|
|
+
|
|
+
|
|
+if __name__ == "__main__":
|
|
+ unittest.main()
|