mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-08 17:39:44 +00:00
fix(g05): match author inference precision
This commit is contained in:
@@ -168,6 +168,7 @@ class G05Config(PreTrainedConfig):
|
||||
discrete_action: bool = True
|
||||
continuous_action: bool = False
|
||||
return_continuous_action: bool = False
|
||||
model_weights_to_bf16: bool = True
|
||||
|
||||
policy_action_dim: int = 20
|
||||
policy_state_dim: int = 20
|
||||
|
||||
@@ -165,10 +165,25 @@ class G05Policy(PreTrainedPolicy):
|
||||
if callable(reset):
|
||||
reset()
|
||||
|
||||
def _apply_author_inference_precision(self) -> None:
|
||||
"""Match the released serving path's BF16 weights with declared FP32 islands."""
|
||||
|
||||
self.backend.to(dtype=torch.bfloat16)
|
||||
apply_fp32_params = getattr(self.backend, "apply_fp32_params", None)
|
||||
if callable(apply_fp32_params):
|
||||
apply_fp32_params()
|
||||
|
||||
def to(self, *args, **kwargs) -> G05Policy:
|
||||
"""Move the author ActionCodec sidecar along with the policy module."""
|
||||
"""Apply author inference precision and move the ActionCodec sidecar."""
|
||||
|
||||
result = super().to(*args, **kwargs)
|
||||
explicit_dtype = "dtype" in kwargs or any(isinstance(arg, torch.dtype | Tensor) for arg in args)
|
||||
if (
|
||||
self.config.model_weights_to_bf16
|
||||
and not explicit_dtype
|
||||
and next(self.backend.parameters()).device.type == "cuda"
|
||||
):
|
||||
self._apply_author_inference_precision()
|
||||
action_tokenizer = getattr(self.backend, "action_tokenizer", None)
|
||||
move_tokenizer = getattr(action_tokenizer, "to", None)
|
||||
if callable(move_tokenizer):
|
||||
|
||||
@@ -685,6 +685,25 @@ def test_policy_to_moves_non_module_action_tokenizer_sidecar():
|
||||
assert backend.action_tokenizer.device == next(policy.parameters()).device
|
||||
|
||||
|
||||
def test_author_inference_precision_preserves_declared_fp32_parameters():
|
||||
class MixedPrecisionBackend(TinyG05Backend):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.bulk_weight = nn.Parameter(torch.ones(2))
|
||||
self.precision_weight = nn.Parameter(torch.ones(2))
|
||||
|
||||
def apply_fp32_params(self):
|
||||
self.precision_weight.data = self.precision_weight.data.float()
|
||||
|
||||
backend = MixedPrecisionBackend()
|
||||
policy = G05Policy(_config(), backend=backend)
|
||||
|
||||
policy._apply_author_inference_precision()
|
||||
|
||||
assert backend.bulk_weight.dtype is torch.bfloat16
|
||||
assert backend.precision_weight.dtype is torch.float32
|
||||
|
||||
|
||||
def test_batch_two_preserves_each_raw_task_and_every_camera_slot():
|
||||
backend = TinyG05Backend()
|
||||
policy = G05Policy(_config(), backend=backend)
|
||||
|
||||
Reference in New Issue
Block a user