LoRA finetuning pi0 fails: prepare_input not found

[WilliamForest233]

Hi, I attempted to finetune the π₀ model on the genmanip-demo dataset using LoRA, configured via the YAML config script.
However, training fails with the following error:

Traceback (most recent call last):
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/peft/peft_model.py", line 824, in __getattr__
    return super().__getattr__(name)  # defer to nn.Module's logic
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1931, in __getattr__
    raise AttributeError(
AttributeError: 'PeftModelForCausalLM' object has no attribute 'prepare_input'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/peft/tuners/lora/model.py", line 371, in __getattr__
    return super().__getattr__(name)  # defer to nn.Module's logic
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1931, in __getattr__
    raise AttributeError(
AttributeError: 'LoraModel' object has no attribute 'prepare_input'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/data/internmanip/scripts/train/train.py", line 366, in <module>
    main(config)
  File "/data/internmanip/scripts/train/train.py", line 241, in main
    trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
  File "/data/internmanip/internmanip/trainer/base.py", line 160, in train
    return super().train(resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/transformers/trainer.py", line 2245, in train
    return inner_training_loop(
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/transformers/trainer.py", line 2560, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/transformers/trainer.py", line 3736, in training_step
    loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
  File "/data/internmanip/internmanip/trainer/base.py", line 76, in compute_loss
    outputs = model(inputs)
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1643, in forward
    else self._run_ddp_forward(*inputs, **kwargs)
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1459, in _run_ddp_forward
    return self.module(*inputs, **kwargs)  # type: ignore[index]
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/accelerate/utils/operations.py", line 823, in forward
    return model_forward(*args, **kwargs)
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/accelerate/utils/operations.py", line 811, in __call__
    return convert_to_fp32(self.model_forward(*args, **kwargs))
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/amp/autocast_mode.py", line 44, in decorate_autocast
    return func(*args, **kwargs)
  File "/data/internmanip/internmanip/utils/peft.py", line 7, in _forward
    backbone_inputs, action_inputs = model.prepare_input(inputs)
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/peft/peft_model.py", line 828, in __getattr__
    return getattr(self.base_model, name)
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/peft/tuners/lora/model.py", line 375, in __getattr__
    return getattr(self.model, name)
  File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1931, in __getattr__
    raise AttributeError(
AttributeError: 'PI0Policy' object has no attribute 'prepare_input'
[rank0]: Traceback (most recent call last):
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/peft/peft_model.py", line 824, in __getattr__
[rank0]:     return super().__getattr__(name)  # defer to nn.Module's logic
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1931, in __getattr__
[rank0]:     raise AttributeError(
[rank0]: AttributeError: 'PeftModelForCausalLM' object has no attribute 'prepare_input'

[rank0]: During handling of the above exception, another exception occurred:

[rank0]: Traceback (most recent call last):
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/peft/tuners/lora/model.py", line 371, in __getattr__
[rank0]:     return super().__getattr__(name)  # defer to nn.Module's logic
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1931, in __getattr__
[rank0]:     raise AttributeError(
[rank0]: AttributeError: 'LoraModel' object has no attribute 'prepare_input'

[rank0]: During handling of the above exception, another exception occurred:

[rank0]: Traceback (most recent call last):
[rank0]:   File "/data/internmanip/scripts/train/train.py", line 366, in <module>
[rank0]:     main(config)
[rank0]:   File "/data/internmanip/scripts/train/train.py", line 241, in main
[rank0]:     trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
[rank0]:   File "/data/internmanip/internmanip/trainer/base.py", line 160, in train
[rank0]:     return super().train(resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/transformers/trainer.py", line 2245, in train
[rank0]:     return inner_training_loop(
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/transformers/trainer.py", line 2560, in _inner_training_loop
[rank0]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/transformers/trainer.py", line 3736, in training_step
[rank0]:     loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
[rank0]:   File "/data/internmanip/internmanip/trainer/base.py", line 76, in compute_loss
[rank0]:     outputs = model(inputs)
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1643, in forward
[rank0]:     else self._run_ddp_forward(*inputs, **kwargs)
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1459, in _run_ddp_forward
[rank0]:     return self.module(*inputs, **kwargs)  # type: ignore[index]
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/accelerate/utils/operations.py", line 823, in forward
[rank0]:     return model_forward(*args, **kwargs)
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/accelerate/utils/operations.py", line 811, in __call__
[rank0]:     return convert_to_fp32(self.model_forward(*args, **kwargs))
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/amp/autocast_mode.py", line 44, in decorate_autocast
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/data/internmanip/internmanip/utils/peft.py", line 7, in _forward
[rank0]:     backbone_inputs, action_inputs = model.prepare_input(inputs)
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/peft/peft_model.py", line 828, in __getattr__
[rank0]:     return getattr(self.base_model, name)
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/peft/tuners/lora/model.py", line 375, in __getattr__
[rank0]:     return getattr(self.model, name)
[rank0]:   File "/data/internmanip/.venv/model/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1931, in __getattr__
[rank0]:     raise AttributeError(
[rank0]: AttributeError: 'PI0Policy' object has no attribute 'prepare_input'

Is prepare_input expected to be implemented by custom models for LoRA to work?
If so, could you provide guidance or examples on how to adapt custom policy classes (PI0Policy) to be compatible with LoRA?
If not, when might LoRA support for such models be added?

[SiHhangChen]

你得自己写，他们没写这个函数，prepare_input本质上是只支持了gr00t模型的，直接把pi0的select_action函数搬到那边去就行了（当然还有点细节得自己改）

[SiHhangChen]

lora他们也没适配，大胆修改就是了

[zhouyunsong]

Hello, we haven’t yet adapted pi0 for LoRA fine-tuning, but we’ll update it later. My understanding is that the quickest way is:

Modify _forward as follows:

def _forward(inputs):
    action_head_outputs = model(inputs)
    # model.validate_data(action_head_outputs, backbone_outputs, is_training=True)
    return action_head_outputs

[SiHhangChen]

OK，thank you for your works

…

----------Reply to Message---------- On Tue, Aug 19, 2025 13:58 PM Yunsong ***@***.***> wrote: zhouyunsong left a comment (InternRobotics/InternManip#9) Hello, we haven’t yet adapted pi0 for LoRA fine-tuning, but we’ll update it later. My understanding is that the quickest way is: Modify _forward as follows: def _forward(inputs): action_head_outputs = model(inputs) # model.validate_data(action_head_outputs, backbone_outputs, is_training=True) return action_head_outputs — Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you commented.Message ID: ***@***.***>

[bbbbbbbbly]

Do you have pi0 model checkpoint finetune on the genmanip datasets? @SiHhangChen @zhouyunsong