huggingface / huggingface/peft
TP + LoRA is broken on transformers 5.16: add_tensor_parallel_hooks_to_module and EmbeddingParallel were removed
- Dominant language
- Python
- Stars
- 21.7k
- Forks
- 2.5k
- Avg merge
- 4d 12h
- Merged PRs (30d)
- 59
Description
### System Info
```
peft main @ 9c16ee66
transformers 5.16.1
```
`peft` declares no upper bound on `transformers`, so `pip install -U` lands here.
### Who can help?
Tensor parallel + LoRA was added in #3079 and #3096.
### Reproduction
transformers 5.16.0 turned `transformers.integrations.tensor_parallel` into a
backward-compatibility shim over the new `transformers.distributed.tensor_parallel`, and the
shim's re-export list does not carry everything the old module had. Two names `peft` imports
are gone:
```
transformers 5.16.1
ImportError: cannot import name 'add_tensor_parallel_hooks_to_module' from 'transformers.integrations.tensor_parallel'
ImportError: cannot import name 'EmbeddingParallel' from 'transformers.integrations.tensor_parallel'
```
- `src/peft/tuners/lora/model.py:318` — `add_tensor_parallel_hooks_to_module`
- `src/peft/utils/save_and_load.py:365` — `EmbeddingParallel`
Both are reached only when the base layer carries `_hf_tp_plan`, so this hits TP-sharded models
with LoRA: `inject_adapter` at load time, and `get_peft_model_state_dict` at save time.
Neither name moved — they are absent from the canonical module too:
```python
>>> import transformers.distributed.tensor_parallel as t
>>> hasattr(t, "add_tensor_parallel_hooks_to_module"), hasattr(t, "EmbeddingParallel")
(False, False)
```
They were present through 5.15.1 and removed in 5.16.0:
| transformers | `add_tensor_parallel_hooks_to_module` | `EmbeddingParallel` |
| --- | --- | --- |
| 5.4.0 … 5.15.1 | yes | yes |
| 5.16.0, 5.16.1 | no | no |
### Expected behavior
Three things changed together, and the last two are the reason I am opening an issue rather
than sending a patch — a mechanical rename would be wrong here.
**1. `add_tensor_parallel_hooks_to_module` has no direct replacement.** Its 5.15 body was
```python
tp_layer = ALL_PARALLEL_STYLES[current_module_plan]
tp_layer.validate_module(module, device_mesh, layer_name)
tp_layer.prepare_module_tp(module, device_mesh, config=model.config)
module._hf_tp_plan = current_module_plan
module._hf_device_mesh = device_mesh
```
but on 5.16 `TensorParallelLayer` no longer has `prepare_module_tp` or `validate_module`. The
surface is now `install_forward`, `shard_param`, `transform_inputs_pre_forward`,
`transform_output_post_forward`, `context_around_forward`.
**2. `EmbeddingParallel` was folded into `RowwiseParallel`, which breaks an `isinstance` chain
silently.** `save_and_load.py:421-425` dispatches on the class:
```python
if isinstance(tp_layer, ColwiseParallel): key = ...lora_B...
elif isinstance(tp_layer, RowwiseParallel): key = ...lora_A...
elif isinstance(tp_layer, EmbeddingParallel): # shards base_layer.weight and lora_embedding_A
```
and the style registry changed under it:
```
5.15.1 ALL_PARALLEL_STYLES["embedding_rowwise"] = EmbeddingParallel(embedding_dim_sharding=0)
5.16.1 ALL_PARALLEL_STYLES["embedding_rowwise"] -> RowwiseParallel
```
So just dropping the third branch would send `embedding_rowwise` into the `RowwiseParallel`
branch and save `lora_A` instead of sharding `base_layer.weight` and `lora_embedding_A` — wrong
output rather than an error. Whatever the fix is, that dispatch has to key off the plan name.
**3. The `embedding_colwise` style key is gone.** `lora/model.py:357` registers it for
`lora_embedding_A` so the save-time gather happens on the right dimension, and
`save_and_load.py:415` does `ALL_PARALLEL_STYLES[tp_plan]`. On 5.16 that key does not exist:
```python
>>> "embedding_colwise" in list(t.ALL_PARALLEL_STYLES.keys())
False
```
so even past the two ImportErrors the embedding path would `KeyError`.
Happy to do the port if you can say which direction you want for (1) — whether `peft` should
call `install_forward` itself or whether this belongs back in transformers as a helper. I have
multi-GPU hardware to verify a real TP + LoRA save/load round-trip on, which is what I would
want before touching (2) and (3).
Contributor guide
Research direction
Start by reproducing the transformers 5.16.1 failures at src/peft/tuners/lora/model.py:318 and src/peft/utils/save_and_load.py:365. Read the tensor-parallel integration around those imports and the dispatch at save_and_load.py:415-425, then compare the 5.15 and 5.16 APIs and style keys. Done means TP-sharded LoRA injection and save/load work without ImportError, KeyError, or incorrect embedding sharding, verified with a multi-GPU round-trip.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100