From afe30630cc02771977543b80a61b5ca7012cc03c Mon Sep 17 00:00:00 2001 From: pepijn223 Date: Thu, 4 Jun 2026 20:03:18 +0200 Subject: [PATCH] test(pi052): repair stale-name CE tests for fused linear CE MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _fast_ce/_shifted_ce were renamed to _fast_lin_ce/_shifted_lin_ce and changed from logits-based to Liger fused-linear-CE (hidden @ lm_head_weightᵀ). Update the tests via thin adapters that pass an identity lm_head_weight (so the computed logits equal the provided ones), run on CUDA (Liger is GPU-only) and skip otherwise, and loosen the allclose tolerance to absorb GPU-vs-CPU float noise on the tiny losses. Co-authored-by: Cursor --- .../pi052/test_pi052_attention_masking.py | 20 ++++++++++- .../pi052/test_pi052_fast_action_loss.py | 33 +++++++++++++++++-- 2 files changed, 49 insertions(+), 4 deletions(-) diff --git a/tests/policies/pi052/test_pi052_attention_masking.py b/tests/policies/pi052/test_pi052_attention_masking.py index 5c74c5488..c12412376 100644 --- a/tests/policies/pi052/test_pi052_attention_masking.py +++ b/tests/policies/pi052/test_pi052_attention_masking.py @@ -37,7 +37,24 @@ import torch pytest.importorskip("transformers") from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks # noqa: E402 -from lerobot.policies.pi052.modeling_pi052 import _mark_target_span_causal, _shifted_ce # noqa: E402 +from lerobot.policies.pi052.modeling_pi052 import ( # noqa: E402 + _mark_target_span_causal, + _shifted_lin_ce, +) + + +def _shifted_ce(logits, labels): + """Adapter: ``_shifted_lin_ce`` is Liger-fused (hidden @ lm_head_weightᵀ). + + An identity ``lm_head_weight`` makes the computed logits equal ``logits``. + Liger's Triton kernel is GPU-only, so inputs run on CUDA; the loss is + returned on CPU so grad still flows back to the CPU ``logits`` leaf. + """ + if not torch.cuda.is_available(): + pytest.skip("Liger fused CE requires CUDA") + vocab_size = logits.shape[-1] + eye = torch.eye(vocab_size, dtype=logits.dtype, device="cuda") + return _shifted_lin_ce(logits.cuda(), eye, labels.cuda()).cpu() # --------------------------------------------------------------------------- # A synthetic PI052 prefix layout: [images, prompt-lang, target-lang] @@ -139,6 +156,7 @@ def test_unmarked_mask_is_bidirectional_the_bug(): def test_shifted_ce_returns_zero_when_no_text_positions_are_supervised(): + pytest.importorskip("liger_kernel") logits = torch.randn(2, 4, 8, requires_grad=True) labels = torch.full((2, 4), -100, dtype=torch.long) diff --git a/tests/policies/pi052/test_pi052_fast_action_loss.py b/tests/policies/pi052/test_pi052_fast_action_loss.py index c5575d6fd..de99b25d6 100644 --- a/tests/policies/pi052/test_pi052_fast_action_loss.py +++ b/tests/policies/pi052/test_pi052_fast_action_loss.py @@ -21,8 +21,29 @@ import torch from torch.nn import functional as F pytest.importorskip("transformers") +pytest.importorskip("liger_kernel") -from lerobot.policies.pi052.modeling_pi052 import _fast_ce # noqa: E402 +from lerobot.policies.pi052.modeling_pi052 import _fast_lin_ce # noqa: E402 + + +def _fast_ce(logits, action_tokens, action_code_mask, predict_actions_t): + """Adapter: ``_fast_lin_ce`` is Liger-fused (hidden @ lm_head_weightᵀ). + + Feeding an identity ``lm_head_weight`` makes the computed logits equal the + provided ``logits``, so these regression tests exercise the masking/gating + logic exactly as before the fused-CE refactor. Liger's Triton kernel is + GPU-only, so inputs are moved to CUDA and the loss is returned on CPU + (keeping grad flowing back to the CPU ``logits`` leaf). + """ + if not torch.cuda.is_available(): + pytest.skip("Liger fused CE requires CUDA") + vocab_size = logits.shape[-1] + eye = torch.eye(vocab_size, dtype=logits.dtype, device="cuda") + predict = predict_actions_t.cuda() if predict_actions_t is not None else None + loss = _fast_lin_ce( + logits.cuda(), eye, action_tokens.cuda(), action_code_mask.cuda(), predict + ) + return loss.cpu() def test_fast_ce_supervises_only_discrete_action_codes(): @@ -46,7 +67,10 @@ def test_fast_ce_supervises_only_discrete_action_codes(): reduction="mean", ) - assert torch.allclose(loss, expected) + # Looser tolerance: the fused Triton kernel (GPU) differs from CPU eager + # F.cross_entropy at the ~1e-7 level, which exceeds the default rtol on + # these very small (~1e-4) losses. + assert torch.allclose(loss, expected, atol=1e-5, rtol=1e-3) def test_fast_ce_masks_non_action_samples(): @@ -72,7 +96,10 @@ def test_fast_ce_masks_non_action_samples(): reduction="mean", ) - assert torch.allclose(loss, expected) + # Looser tolerance: the fused Triton kernel (GPU) differs from CPU eager + # F.cross_entropy at the ~1e-7 level, which exceeds the default rtol on + # these very small (~1e-4) losses. + assert torch.allclose(loss, expected, atol=1e-5, rtol=1e-3) def test_fast_ce_returns_zero_when_no_action_code_positions_are_valid():