huggingface / huggingface/diffusers

better reporting of errors when using TP

Open
#14,533 0 comments 0 reactions 1 assignee Claimed by @JingyaHuang View on GitHub
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

I think the following aren't implemented at the moment (which is fine; just flagging).

* **Sharded loading**. from_pretrained should stream shards straight to each rank's DTensor rather than materializing the full checkpoint then slicing — otherwise TP saves you nothing at load time. And check `save_pretrained` / state_dict calls `.full_tensor()` or uses DCP. I think we should at least raise when `save_pretrained()` is called in case TP is enabled?
* **LoRA loading**. for a colwise base layer, lora_A replicated + lora_B colwise; for rowwise, lora_A rowwise + lora_B replicated. If the plan doesn't cover PEFT layers, loading an adapter onto a TP model will either error or be wrong. I think we should detect if the model has `peft` layers injected and raise if TP is requested?
* **Quantization, offloading**. We should probably also raise when these are requested?

_Originally posted by @sayakpaul in https://github.com/huggingface/diffusers/pull/13718#discussion_r3662821938_

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.