Reduce expected fallback warning noise (#2921)

## Summary
- suppress expected CPU-side pallas fallback warning (`requires TPU
backend`) while still falling back to XLA
- emit pallas fallback warnings once per unique message instead of
repeating on every call
- adapt inferred pallas `b_block_size`/`h_block_size` to
TPU-lane-aligned divisors of the actual shape when available
- reduce splash fallback warning spam by warning once per fallback
category
- make tracker no-op logging APIs (`jit_log`, `log_summary`, etc.)
silently drop events when no global tracker is configured (with one info
log)

## Validation
- `cd lib/levanter && PYTHONPATH=tests:src:. uv run --package levanter
--group test pytest
tests/kernels/test_pallas_fused_cross_entropy_loss.py -q`
- `cd lib/levanter && PYTHONPATH=tests:src:. uv run --package levanter
--group test pytest tests/test_tracker.py tests/test_skip_step.py
tests/test_weight_decay_mask.py -q`
- `cd lib/levanter && PYTHONPATH=tests:src:. uv run --package levanter
--group test pytest tests/grug/test_grugformer_fused_loss.py
tests/grug/test_grugformer_model_loss.py -q`
- `cd lib/levanter && PYTHONPATH=tests:src:. uv run --package levanter
--group test pytest tests/test_attention.py -q`

Author: dlwh

lib/levanter/src/levanter/kernels/pallas/fused_cross_entropy_loss/api.py

@@ -27,6 +27,7 @@
     "xla": linear_softmax_cross_entropy_loss_xla,
 }
 _DEFAULT_IMPLEMENTATION: tuple[Implementation, ...] = ("xla",)
+_PALLAS_FALLBACK_WARNINGS_EMITTED: set[str] = set()
 
 try:
     from .pallas_tpu import PallasUnsupportedError, linear_softmax_cross_entropy_loss_pallas
@@ -37,6 +38,19 @@
     PallasUnsupportedError = NotImplementedError  # type: ignore[assignment]
 
 
+def _warn_pallas_fallback_once(exc: Exception) -> None:
+    message = str(exc)
+    if "requires TPU backend" in message:
+        return
+    if message in _PALLAS_FALLBACK_WARNINGS_EMITTED:
+        return
+    _PALLAS_FALLBACK_WARNINGS_EMITTED.add(message)
+    warnings.warn(
+        f"Pallas fused cross-entropy unavailable, falling back to XLA: {message}",
+        RuntimeWarning,
+    )
+
+
 def _validate_inputs(x: jax.Array, labels: jax.Array, w: jax.Array) -> None:
     if x.ndim != 2:
         raise ValueError(f"x must be rank-2 [B, H], got shape {x.shape}.")
@@ -161,19 +175,13 @@ def fused_cross_entropy_loss_and_logsumexp_penalty(
             except PallasUnsupportedError as e:
                 if explicit:
                     raise
-                warnings.warn(
-                    f"Pallas fused cross-entropy unavailable, falling back to XLA: {e}",
-                    RuntimeWarning,
-                )
+                _warn_pallas_fallback_once(e)
                 errors.append(e)
                 continue
             except NotImplementedError as e:
                 if explicit:
                     raise
-                warnings.warn(
-                    f"Pallas fused cross-entropy unavailable, falling back to XLA: {e}",
-                    RuntimeWarning,
-                )
+                _warn_pallas_fallback_once(e)
                 errors.append(e)
                 continue
         else:
@@ -193,19 +201,13 @@ def fused_cross_entropy_loss_and_logsumexp_penalty(
             except PallasUnsupportedError as e:
                 if explicit:
                     raise
-                warnings.warn(
-                    f"Pallas fused cross-entropy unavailable, falling back to XLA: {e}",
-                    RuntimeWarning,
-                )
+                _warn_pallas_fallback_once(e)
                 errors.append(e)
                 continue
             except NotImplementedError as e:
                 if explicit:
                     raise
-                warnings.warn(
-                    f"Pallas fused cross-entropy unavailable, falling back to XLA: {e}",
-                    RuntimeWarning,
-                )
+                _warn_pallas_fallback_once(e)
                 errors.append(e)
                 continue
 

lib/levanter/src/levanter/layers/attention.py

@@ -62,6 +62,14 @@ class AttentionBackend(StrEnum):
 
 
 DEFAULT_SPLASH_BLOCK_SIZE = 512
+_SPLASH_FALLBACK_WARNINGS_EMITTED: set[str] = set()
+
+
+def _warn_splash_fallback_once(message: str) -> None:
+    if message in _SPLASH_FALLBACK_WARNINGS_EMITTED:
+        return
+    _SPLASH_FALLBACK_WARNINGS_EMITTED.add(message)
+    warnings.warn(message, stacklevel=3)
 
 
 def default_attention_type() -> AttentionBackend:
@@ -1185,13 +1193,13 @@ def _try_tpu_splash_attention(
     if dropout != 0.0:
         if force_flash:
             raise NotImplementedError("Splash attention does not support dropout.")
-        warnings.warn("Splash attention does not support. Falling back to the reference implementation.")
+        _warn_splash_fallback_once("Splash attention does not support dropout. Falling back to the reference.")
         return None
 
     if bias is not None:
         if force_flash:
             raise NotImplementedError("Splash attention does not support bias.")
-        warnings.warn("Splash attention does not support bias. Falling back to the reference implementation.")
+        _warn_splash_fallback_once("Splash attention does not support bias. Falling back to the reference.")
         return None
 
     try:
@@ -1219,16 +1227,17 @@ def _try_tpu_splash_attention(
             raise
         if force_flash:
             raise ImportError("Could not import splash attention. You need to update your JAX to at least 0.7.2.")
-        warnings.warn(
+        _warn_splash_fallback_once(
             "Could not import splash attention. You need to update your JAX to at least 0.7.2. "
-            "Falling back to the reference implementation."
+            "Falling back to the reference implementation.",
         )
         return None
     except NotImplementedError as e:
         message = str(e)
         if force_flash:
             raise NotImplementedError(f"Could not use splash attention: {message}")
-        warnings.warn(f"Could not use splash attention: {message}. Falling back to the reference")
+        logger.info("Could not use splash attention. Falling back to the reference implementation: %s", message)
+        _warn_splash_fallback_once("Could not use splash attention. Falling back to the reference implementation.")
         return None
 
 

lib/levanter/src/levanter/tracker/tracker_fns.py

@@ -29,10 +29,19 @@
 
 _should_use_callback = True
 _global_tracker: Optional["Tracker"] = None
+_has_logged_missing_tracker = False
 
 LoggableValue: typing.TypeAlias = Scalar | jax.Array | str | dict | Histogram
 
 
+def _log_missing_tracker_once() -> None:
+    global _has_logged_missing_tracker
+    if _has_logged_missing_tracker:
+        return
+    _has_logged_missing_tracker = True
+    logger.info("No global tracker set; tracker logs are being dropped.")
+
+
 def log(metrics: typing.Mapping[str, LoggableValue | Any], *, step: Optional[int], commit: Optional[bool] = None):
     """
     Log metrics to the global tracker.
@@ -45,7 +54,8 @@ def log(metrics: typing.Mapping[str, LoggableValue | Any], *, step: Optional[int
     """
     global _global_tracker
     if _global_tracker is None:
-        raise RuntimeError("No global tracker set")
+        _log_missing_tracker_once()
+        return
 
     if is_inside_jit():
         # we're inside a jit, so we need to log from the host
@@ -71,7 +81,7 @@ def log_metrics(
 def _do_jit_log(metrics, *, step=None):
     try:
         if _global_tracker is None:
-            warnings.warn("No global tracker set")
+            _log_missing_tracker_once()
         else:
             _global_tracker.log(metrics, step=step, commit=False)
     except Exception:
@@ -91,7 +101,7 @@ def jit_log(metrics, *, step=None):
     We strongly recommend using the first method, as it is much more performant.
     """
     if _global_tracker is None:
-        warnings.warn("No global tracker set")
+        _log_missing_tracker_once()
         return
     if not _should_use_callback:
         # we're not using the callback, so we assume we're inside a defer_tracker_for_jit context manager
@@ -154,7 +164,7 @@ def log_summary(metrics: dict[str, Any]):
     """
     global _global_tracker
     if _global_tracker is None:
-        warnings.warn("No global tracker set")
+        _log_missing_tracker_once()
         return
 
     _global_tracker.log_summary(metrics)
@@ -169,7 +179,7 @@ def log_hyperparameters(hparams: dict[str, Any]):
     """
     global _global_tracker
     if _global_tracker is None:
-        warnings.warn("No global tracker set")
+        _log_missing_tracker_once()
         return
 
     _global_tracker.log_hyperparameters(hparams)
@@ -185,7 +195,7 @@ def log_configuration(hparams: Any, config_name: Optional[str] = None):
     """
     global _global_tracker
     if _global_tracker is None:
-        warnings.warn("No global tracker set")
+        _log_missing_tracker_once()
         return
 
     hparams_dict = hparams_to_dict(hparams)

lib/levanter/tests/kernels/test_pallas_fused_cross_entropy_loss.py

@@ -123,6 +123,52 @@ def test_fused_cross_entropy_pallas_requires_tpu():
         )
 
 
+def test_infer_block_sizes_adapts_to_supported_divisors():
+    block_sizes = infer_block_sizes(
+        b=512,
+        h=128,
+        v=4096,
+        dtype=jnp.float32,
+        device_kind="TPU v5e",
+    )
+
+    assert block_sizes.b_block_size == 512
+    assert block_sizes.h_block_size == 128
+
+
+def test_infer_block_sizes_preserves_defaults_without_128_aligned_divisors():
+    block_sizes = infer_block_sizes(
+        b=96,
+        h=64,
+        v=4096,
+        dtype=jnp.float32,
+        device_kind="TPU v5e",
+    )
+
+    assert block_sizes.b_block_size == 1024
+    assert block_sizes.h_block_size == 512
+
+
+def test_default_implementation_on_cpu_skips_expected_tpu_warning():
+    if jax.default_backend() == "tpu":
+        pytest.skip("requires non-TPU backend")
+
+    x = jnp.zeros((32, 64), dtype=jnp.float32)
+    w = jnp.zeros((64, 128), dtype=jnp.float32)
+    y = jnp.zeros((32,), dtype=jnp.int32)
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        fused_api.fused_cross_entropy_loss_and_logsumexp_penalty(
+            x,
+            y,
+            w,
+            reduction=None,
+        )
+
+    assert not any("requires TPU backend" in str(warning.message) for warning in caught)
+
+
 def test_fused_cross_entropy_default_matches_reference():
     backend = jax.default_backend()
     if backend == "tpu":

lib/levanter/tests/test_tracker.py

@@ -3,6 +3,7 @@
 
 # NOTE: Do not explicitly import wandb/other trackers here, as this will cause the tests to trivially pass.
 import dataclasses
+import warnings
 from typing import Tuple
 
 import pytest
@@ -82,3 +83,20 @@ def test_get_tracker_by_name(monkeypatch):
 
         with pytest.raises(KeyError):
             levanter.tracker.get_tracker("foo")
+
+
+def test_tracker_logging_without_global_tracker_emits_no_warning(monkeypatch):
+    import levanter.tracker.tracker_fns as tracker_fns
+
+    monkeypatch.setattr(tracker_fns, "_global_tracker", None)
+    monkeypatch.setattr(tracker_fns, "_has_logged_missing_tracker", False)
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        tracker_fns.log({"metric": 1.0}, step=0)
+        tracker_fns.jit_log({"metric": 1.0}, step=0)
+        tracker_fns.log_summary({"metric": 1.0})
+        tracker_fns.log_hyperparameters({"metric": 1.0})
+        tracker_fns.log_configuration({"metric": 1.0})
+
+    assert not caught