huggingface / huggingface/peft

TP + LoRA is broken on transformers 5.16: add_tensor_parallel_hooks_to_module and EmbeddingParallel were removed

Open
#3,628 1 comment 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.