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

[alanhuangyoo]

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:

>>> 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

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:

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:

>>> "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).

[BenjaminBossan]

Thanks for reporting. We're already working on this, no need for a PR.