mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 18:26:11 +00:00
Liang Su
3cec067795
* perf(pi052): optimize equivalent training paths * fix(pi052): guard FlexAttention backend selection
83 lines
3.2 KiB
Python
83 lines
3.2 KiB
Python
#!/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
|