diff --git a/src/lerobot/policies/g05/modeling_g05.py b/src/lerobot/policies/g05/modeling_g05.py index 47697f5db..b53d6020b 100644 --- a/src/lerobot/policies/g05/modeling_g05.py +++ b/src/lerobot/policies/g05/modeling_g05.py @@ -178,6 +178,17 @@ class G05Policy(PreTrainedPolicy): apply_fp32_params = getattr(self.backend, "apply_fp32_params", None) if callable(apply_fp32_params): apply_fp32_params() + if self.config.predict_cot: + # The author Qwen3.5 final norm is an FP32 island and its fused CE + # kernel disables autocast, so the tied output projection must be + # FP32 as well. Otherwise CoT training reaches FLCE with FP32 hidden + # states and a BF16 weight and fails before computing text loss. + model = getattr(self.backend, "model", None) + vlm = getattr(model, "vlm", None) + output_proj = getattr(vlm, "output_proj", None) + weight = getattr(output_proj, "weight", None) + if isinstance(weight, nn.Parameter): + weight.data = weight.data.float() def to(self, *args, **kwargs) -> G05Policy: """Apply author inference precision and move the ActionCodec sidecar.""" diff --git a/tests/policies/g05/test_g05.py b/tests/policies/g05/test_g05.py index 693a5dbf6..0c9173f63 100644 --- a/tests/policies/g05/test_g05.py +++ b/tests/policies/g05/test_g05.py @@ -648,6 +648,28 @@ def test_author_inference_precision_preserves_declared_fp32_parameters(): assert backend.precision_weight.dtype is torch.float32 +def test_system2_precision_keeps_tied_lm_head_compatible_with_fp32_final_norm(): + class CoTPrecisionBackend(TinyG05Backend): + def __init__(self): + super().__init__() + self.model = nn.Module() + self.model.vlm = nn.Module() + self.model.vlm.input_proj = nn.Embedding(8, 4) + self.model.vlm.output_proj = nn.Linear(4, 8, bias=False) + self.model.vlm.output_proj.weight = self.model.vlm.input_proj.weight + + def apply_fp32_params(self): + pass + + backend = CoTPrecisionBackend() + policy = G05Policy(_config(predict_cot=True, runtime_system="system2"), backend=backend) + + policy._apply_author_inference_precision() + + assert backend.model.vlm.output_proj.weight.dtype is torch.float32 + assert backend.model.vlm.input_proj.weight is backend.model.vlm.output_proj.weight + + def test_batch_two_preserves_each_raw_task_and_every_camera_slot(): backend = TinyG05Backend() policy = G05Policy(_config(), backend=backend)