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
|
discrete_action: bool = True
|
||||||
continuous_action: bool = False
|
continuous_action: bool = False
|
||||||
return_continuous_action: bool = False
|
return_continuous_action: bool = False
|
||||||
|
model_weights_to_bf16: bool = True
|
||||||
|
|
||||||
policy_action_dim: int = 20
|
policy_action_dim: int = 20
|
||||||
policy_state_dim: int = 20
|
policy_state_dim: int = 20
|
||||||
|
|||||||
@@ -165,10 +165,25 @@ class G05Policy(PreTrainedPolicy):
|
|||||||
if callable(reset):
|
if callable(reset):
|
||||||
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:
|
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)
|
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)
|
action_tokenizer = getattr(self.backend, "action_tokenizer", None)
|
||||||
move_tokenizer = getattr(action_tokenizer, "to", None)
|
move_tokenizer = getattr(action_tokenizer, "to", None)
|
||||||
if callable(move_tokenizer):
|
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
|
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():
|
def test_batch_two_preserves_each_raw_task_and_every_camera_slot():
|
||||||
backend = TinyG05Backend()
|
backend = TinyG05Backend()
|
||||||
policy = G05Policy(_config(), backend=backend)
|
policy = G05Policy(_config(), backend=backend)
|
||||||
|
|||||||
Reference in New Issue
Block a user