perf(pi052): optimize flow and full-training paths (#3974)

* perf(pi052): optimize equivalent training paths

* fix(pi052): guard FlexAttention backend selection

tests/policies/pi052/test_pi052_bucketed_ce.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+from lerobot.policies.pi052.modeling_pi052 import _lin_ce_flat
+
+
+@pytest.mark.parametrize("z_loss_weight", [0.0, 1e-4])
+@pytest.mark.parametrize("rows,valid_rows", [(24, 9), (48, 25)])
+def test_bucketed_ce_matches_dense_loss_and_gradients(z_loss_weight, rows, valid_rows):
+    generator = torch.Generator().manual_seed(23)
+    hidden_size, vocab_size = 7, 19
+    hidden_ref = torch.randn(rows, hidden_size, generator=generator, dtype=torch.float64, requires_grad=True)
+    weight_ref = torch.randn(
+        vocab_size, hidden_size, generator=generator, dtype=torch.float64, requires_grad=True
+    )
+    labels = torch.full((rows,), -100, dtype=torch.long)
+    valid_indices = torch.randperm(rows, generator=generator)[:valid_rows]
+    labels[valid_indices] = torch.randint(0, vocab_size, (valid_rows,), generator=generator)
+    hidden_bucketed = hidden_ref.detach().clone().requires_grad_(True)
+    weight_bucketed = weight_ref.detach().clone().requires_grad_(True)
+
+    import lerobot.policies.pi052.modeling_pi052 as modeling_pi052
+
+    loss_ref = _lin_ce_flat(hidden_ref, weight_ref, labels, z_loss_weight=z_loss_weight)
+    old_limit = modeling_pi052._LOGITS_CE_MAX_POSITIONS
+    modeling_pi052._LOGITS_CE_MAX_POSITIONS = 16
+    try:
+        loss_bucketed = _lin_ce_flat(
+            hidden_bucketed,
+            weight_bucketed,
+            labels,
+            z_loss_weight=z_loss_weight,
+        )
+    finally:
+        modeling_pi052._LOGITS_CE_MAX_POSITIONS = old_limit
+
+    loss_ref.backward()
+    loss_bucketed.backward()
+
+    torch.testing.assert_close(loss_bucketed, loss_ref, rtol=1e-6, atol=1e-6)
+    torch.testing.assert_close(hidden_bucketed.grad, hidden_ref.grad, rtol=1e-12, atol=1e-12)
+    torch.testing.assert_close(weight_bucketed.grad, weight_ref.grad, rtol=1e-12, atol=1e-12)
+
+
+def test_bucketed_ce_all_ignored_preserves_zero_gradients():
+    hidden = torch.randn(24, 7, dtype=torch.float64, requires_grad=True)
+    weight = torch.randn(19, 7, dtype=torch.float64, requires_grad=True)
+    labels = torch.full((24,), -100, dtype=torch.long)
+
+    import lerobot.policies.pi052.modeling_pi052 as modeling_pi052
+
+    old_limit = modeling_pi052._LOGITS_CE_MAX_POSITIONS
+    modeling_pi052._LOGITS_CE_MAX_POSITIONS = 16
+    try:
+        loss = _lin_ce_flat(hidden, weight, labels)
+    finally:
+        modeling_pi052._LOGITS_CE_MAX_POSITIONS = old_limit
+    loss.backward()
+
+    assert loss.item() == 0.0
+    assert hidden.grad is not None
+    assert weight.grad is not None
+    assert torch.count_nonzero(hidden.grad) == 0
+    assert torch.count_nonzero(weight.grad) == 0

tests/policies/pi052/test_pi052_flex_attention.py

@@ -0,0 +1,65 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+import lerobot.policies.pi052.modeling_pi052 as modeling_pi052  # noqa: E402
+from lerobot.policies.pi052.configuration_pi052 import PI052Config  # noqa: E402
+
+
+def test_flex_backend_skips_non_cuda_without_initializing(monkeypatch):
+    monkeypatch.setattr(modeling_pi052, "_flex_fns", None)
+    monkeypatch.setattr(torch, "compile", lambda *args, **kwargs: pytest.fail("torch.compile was called"))
+    monkeypatch.setattr(
+        torch.cuda,
+        "get_device_properties",
+        lambda *args, **kwargs: pytest.fail("CUDA properties were queried"),
+    )
+
+    assert modeling_pi052._get_flex_fns(torch.device("cpu")) is None
+    assert modeling_pi052._get_flex_kernel_options(torch.device("cpu")) is None
+    assert modeling_pi052._flex_fns is None
+
+
+def test_flex_initialization_failure_falls_back(monkeypatch, caplog):
+    monkeypatch.setattr(modeling_pi052, "_flex_fns", None)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+
+    def fail_compile(*args, **kwargs):
+        raise RuntimeError("compile failed")
+
+    monkeypatch.setattr(torch, "compile", fail_compile)
+
+    with caplog.at_level(logging.WARNING, logger=modeling_pi052.__name__):
+        assert modeling_pi052._get_flex_fns(torch.device("cuda", 0)) is None
+
+    assert modeling_pi052._flex_fns is False
+    assert "FlexAttention unavailable" in caplog.text
+
+
+def test_flex_rejects_single_repeat_configuration():
+    with pytest.raises(ValueError, match="use_flex_attention requires flow_num_repeats > 1"):
+        PI052Config(use_flex_attention=True, flow_num_repeats=1)
+
+
+def test_flex_accepts_amortized_repeat_configuration():
+    config = PI052Config(use_flex_attention=True, flow_num_repeats=5)
+    assert config.use_flex_attention

tests/policies/pi052/test_pi052_vision_checkpointing.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import MethodType, SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+pytest.importorskip("transformers")
+
+from lerobot.policies.pi052.modeling_pi052 import PI05Pytorch
+
+
+class _MockVisionTower:
+    def __init__(self):
+        self.enable_kwargs = None
+        self.disable_calls = 0
+
+    def gradient_checkpointing_enable(self, **kwargs):
+        self.enable_kwargs = kwargs
+
+    def gradient_checkpointing_disable(self):
+        self.disable_calls += 1
+
+
+def _checkpoint_model():
+    tower = _MockVisionTower()
+    language_model = SimpleNamespace(gradient_checkpointing=False)
+    expert_model = SimpleNamespace(gradient_checkpointing=False)
+    model = SimpleNamespace(
+        gradient_checkpointing_enabled=False,
+        paligemma_with_expert=SimpleNamespace(
+            paligemma=SimpleNamespace(
+                model=SimpleNamespace(language_model=language_model, vision_tower=tower)
+            ),
+            gemma_expert=SimpleNamespace(model=expert_model),
+        ),
+    )
+    return model, tower, language_model, expert_model
+
+
+def test_gradient_checkpointing_uses_vision_tower_layer_api():
+    model, tower, language_model, expert_model = _checkpoint_model()
+
+    PI05Pytorch.gradient_checkpointing_enable(model)
+
+    assert model.gradient_checkpointing_enabled
+    assert language_model.gradient_checkpointing
+    assert expert_model.gradient_checkpointing
+    assert tower.enable_kwargs == {"gradient_checkpointing_kwargs": {"use_reentrant": False}}
+
+    PI05Pytorch.gradient_checkpointing_disable(model)
+
+    assert not model.gradient_checkpointing_enabled
+    assert not language_model.gradient_checkpointing
+    assert not expert_model.gradient_checkpointing
+    assert tower.disable_calls == 1
+
+
+def test_siglip_layers_recompute_individually():
+    from transformers.models.siglip.configuration_siglip import SiglipVisionConfig
+    from transformers.models.siglip.modeling_siglip import SiglipVisionModel
+
+    config = SiglipVisionConfig(
+        hidden_size=16,
+        intermediate_size=32,
+        num_hidden_layers=2,
+        num_attention_heads=2,
+        num_channels=3,
+        image_size=16,
+        patch_size=8,
+    )
+    tower = SiglipVisionModel(config).train()
+    tower.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+    calls = [0] * config.num_hidden_layers
+
+    for index, layer in enumerate(tower.vision_model.encoder.layers):
+        original_forward = layer.forward
+
+        def counted_forward(self, *args, _index=index, _forward=original_forward, **kwargs):
+            calls[_index] += 1
+            return _forward(*args, **kwargs)
+
+        layer.forward = MethodType(counted_forward, layer)
+
+    pixels = torch.randn(2, config.num_channels, config.image_size, config.image_size)
+    tower(pixels).last_hidden_state.sum().backward()
+
+    assert calls == [2] * config.num_hidden_layers
+
+
+def test_embed_prefix_does_not_wrap_the_whole_vision_tower_checkpoint():
+    model = PI05Pytorch.__new__(PI05Pytorch)
+    nn.Module.__init__(model)
+    model.config = SimpleNamespace()
+    model.gradient_checkpointing_enabled = True
+    model.train()
+
+    image_calls = []
+
+    def embed_image(image):
+        image_calls.append(image.shape)
+        return image[:, :1, 0, :2]
+
+    def embed_language_tokens(tokens):
+        return tokens.to(torch.float32).unsqueeze(-1).expand(*tokens.shape, 2)
+
+    model.paligemma_with_expert = SimpleNamespace(
+        embed_image=embed_image,
+        embed_language_tokens=embed_language_tokens,
+    )
+    outer_checkpoint_calls = []
+
+    def apply_checkpoint(func, value):
+        outer_checkpoint_calls.append(value.shape)
+        return func(value)
+
+    model._apply_checkpoint = apply_checkpoint
+
+    images = [torch.randn(2, 3, 4, 4), torch.randn(2, 3, 4, 4)]
+    image_masks = [torch.ones(2, dtype=torch.bool) for _ in images]
+    tokens = torch.ones(2, 3, dtype=torch.long)
+    token_masks = torch.ones_like(tokens, dtype=torch.bool)
+
+    embeddings, _, _ = model.embed_prefix(images, image_masks, tokens, token_masks)
+
+    assert image_calls == [image.shape for image in images]
+    assert outer_checkpoint_calls == [tokens.shape]
+    assert embeddings.shape == (2, 5, 2)