_tp_plan / _cp_plan: allow entries for modules that are not always present

@yiyixuxu

Follow-up to #14609. This is a discussion, not a request for a PR: I want to hear opinions on whether we
need this and which shape it should take before anyone starts on it.

Problem

_tp_plan and _cp_plan are class attributes, and both are resolved by attribute lookup:
resolve_tp_shard_specs / _resolve_tp_plan in src/diffusers/hooks/tensor_parallel.py walk each
key with getattr, and _find_submodule_by_name in src/diffusers/hooks/context_parallel.py
raises ValueError for an atom that is not a submodule. So every key must resolve on the instance,
or applying the plan fails before anything is sharded.

That means a plan can only list modules that exist on every instance of the class. #14609 hit this
with fused projections and had to leave them out:

# `attn.to_qkv` is deliberately not listed: it only exists after `fuse_projections()`, and the plan is
# resolved by attribute lookup, so an unconditional entry would break the ordinary unfused model.

The same limit applies to any module that depends on the config: a model with optional layers
(a guidance embedder, an extra projection enabled by a flag, a block variant chosen per
checkpoint) cannot have one plan that covers all of its configurations.

Options

One way is to let a plan mark some entries as optional, so they are skipped when the module is absent
and still raise for a typo in a required entry. Two possible shapes:

an explicit wrapper, e.g. "transformer_blocks.*.attn.to_qkv": Optional(PackedColwiseParallel([1, 1, 1]))
match plan keys against the modules that exist, like transformers does, and warn for an unmatched key

(1) keeps the plan strict by default; (2) is less code in each model but a typo goes unnoticed. A third
option is to leave the plans as they are and accept that a model with config-dependent layers needs
separate handling. Whatever we pick should apply to both _tp_plan and _cp_plan.

Opinions welcome, especially from anyone who has written a plan for a model with optional layers.

Nothing stops a user from calling fuse_qkv_projections() on a model that is already sharded with TP,
and it cannot work: the sharded to_q/to_k/to_v weights are DTensors, and fuse_projections
concatenates them into a new plain nn.Linear that has no TP hooks and no plan entry. We should
either raise there, like we already do for offloading, saving and LoRA on a sharded model, or let the
plan cover the fused Linear, which needs the optional entry discussed above.

Posted by Claude Code on behalf of @yiyixuxu; the text was reviewed and approved by them.

@sayakpaul

cc @sayakpaul @JingyaHuang
let me know if we need to make _tp_plan/_cp_plan work with optional layers

Thanks for opening the thread!

I think the benefits of fusion start to diminish with parallelism. The benefits start showing up with specific quantization schemes, depending on the model parameters and input sizes. But we don't support TP + quant yet. To get more ideas on this, I ran a small benchmark with H3 and TP:

Full artifact: https://claude.ai/artifact/CNbMdEuNJFuiZuoJZFkrNW

(it's possible that the implementation is suboptimal but that is a different issue IMO)

So, based on the numbers, I think instead of working on the first option or any other more intrusive changes, we should simply warn the users (i.e., the second option). WDYT?

@sayakpaul

@sayakpaul

thanks for looking into this! so it's feature is not just for fusing use case in particular - although that's what inspired it

more generally, currently we don't really know how to create a plan involve optional layers in models, i.e. if a model class support 2 different checkpoint with slightly different layout, e.g. add_k_proj/add_v_proj are None in one checkpoint but not in another etc

Ah yes, that is true. Let think about it a bit and get back to you.

Okay some further thoughts.

I think #15197 is quite related to this discussion because it's linked to #15137 where the Cosmos checkpoints have optional layers.

#15197 basically exposes a model-level _get_tp_plan() which model classes can override when there is conditional logic to be implemented. For example (from #15137):

def _get_tp_plan(self):
    tp_plan = super()._get_tp_plan().copy()
    if self.config.hidden_act == "silu":
        tp_plan.update(
            {
                "layers.*.mlp.gate_proj": "colwise",
                "layers.*.mlp_moe_gen.gate_proj": "colwise",
            }
        )
    return tp_plan

The idea is that the TP-related code would read from what is returned from _get_tp_plan() instead of just the class-level _tp_plan attribute.

I think marking with Optional like you suggested could also work. IMO, the _get_tp_plan() approach is more explicit.

@DN6

ohh i see, indeedd it's related cc @DN6 into the discussion too

I don't think _get_tp_plan, as how it is currently implemented in the PR, is the right design for this though.

it's basically to express "this layer is optional" with bunch of if ... else in an override. , i.e. compare this

_tp_plan = {
    "layers.*.mlp.up_proj": "colwise",
    "layers.*.mlp.down_proj": "rowwise",
    ...
}

def _get_tp_plan(self):
    tp_plan = super()._get_tp_plan().copy()
    if self.config.hidden_act == "silu":
        tp_plan.update(
            {
                "layers.*.mlp.gate_proj": "colwise",
                "layers.*.mlp_moe_gen.gate_proj": "colwise",
            }
        )
    return tp_plan

with something like this

_tp_plan = {
    "layers.*.mlp.up_proj": "colwise",
    "layers.*.mlp.down_proj": "rowwise",
    # only in the SiLU checkpoints
    "layers.*.mlp.gate_proj": Optional("colwise"),
    "layers.*.mlp_moe_gen.gate_proj": Optional("colwise"),
    ...
}

I think it's expressing the exact same amount of info, but the plan lives in two places instead of one so that's not very nice/ easy to read and I think better to keep it declarative like either one of options proposed in the issue

However, if we want the plan to be more customizable (not sure if there is a use case), we can allow the override from the checkpoint level (not the model level like in the PR ) -> so that would be something user config in config.json

@DN6

Then Option 1 (Optional) seems like a better choice to me. I think we can also warn for unmatched modules just for comprehensive logging, too.

Let's see what @DN6 has to say and I can open a PR soon because this blocks TP for latest Cosmos checkpoints (and I think this will be a recurring case). Cc: @yzhautouskay

Option 1 looks cleanest to me 👍🏽. Just have to make sure we have no clashes with typing.Optional (although I think we removed all of these when we bumped the min Python version)

Navigation Menu

Repository navigation

`_tp_plan` / `_cp_plan`: allow entries for modules that are not always present #15218

Problem

Options

Related

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

_tp_plan / _cp_plan: allow entries for modules that are not always present #15218

Description

Problem

Options

Related

Activity

yiyixuxu commented on Nov 9, 2026

sayakpaul commented on Nov 9, 2026

yiyixuxu commented on Nov 9, 2026

sayakpaul commented on Nov 9, 2026

sayakpaul commented on Nov 9, 2026

yiyixuxu commented on Nov 9, 2026

sayakpaul commented on Nov 10, 2026

DN6 commented on Nov 10, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

`_tp_plan` / `_cp_plan`: allow entries for modules that are not always present #15218