Avoid TP helper imports for non-TP distributed LoRA loads

[jiqing-feng]

Summary

Fix a ModuleNotFoundError in distributed LoRA adapter loading when torch.distributed is initialized but the model is not actually tensor-parallel sharded.

Problem

set_peft_model_state_dict calls _maybe_shard_state_dict_for_tp whenever distributed is initialized. This also happens for regular DDP runs. The helper imported transformers.integrations.tensor_parallel before checking whether any LoRA layer had TP metadata, so non-TP distributed runs could fail immediately with:

ModuleNotFoundError: No module named 'transformers.integrations.tensor_parallel'

Fix

Check for real TP-sharded LoRA layers first by looking for _hf_tp_plan and _hf_device_mesh on the base layers. If none are found, return early. Only import Transformers TP helpers when there is actually something to shard.

This keeps the existing TP path intact while avoiding false-positive TP handling for DDP/non-TP workloads.

Test

Added a regression test for a LoRA model with a model-level _tp_plan but no layer-level TP metadata:

python -m pytest -o addopts="" tests/test_other.py::test_maybe_shard_state_dict_for_tp_noops_without_tp_layers -q

Result: 1 passed.

[jiqing-feng]

Hi @BenjaminBossan . Would you please review the PR? Thanks!

[BenjaminBossan]

Thank you for this PR. I agree that we can check for TP layers first before importing the TP related class, so that part of the PR looks good. What I don't quite get is the test. Under what circumstances would we expect the test to fail? Could you perhaps add a comment to explain?

[jiqing-feng]

Thank you for this PR. I agree that we can check for TP layers first before importing the TP related class, so that part of the PR looks good. What I don't quite get is the test. Under what circumstances would we expect the test to fail? Could you perhaps add a comment to explain?

Thanks for the feedback. I updated the test comment to make the intent clearer.

It is a unit test for _maybe_shard_state_dict_for_tp. It verifies that when no LoRA base layer has TP metadata (_hf_tp_plan / _hf_device_mesh), the helper is a no-op and leaves the state dict unchanged.