Composite lm head tp plan (#49080)

* Declare the lm_head TP plan on every model that can tie it

`tie_word_embeddings` has the config shard `embed_tokens` `embedding_rowwise`,
and a tied head IS that parameter — so a head carrying no plan entry of its own
hands `F.linear` a plain input and a DTensor weight. 81 classes already declare
one; the composites that wrap them did not, because they keep their plan on
`text_config` rather than the outer config. 57 classes are affected.

Tying is a per-checkpoint flag, so the entry belongs on every class that CAN
tie, not only those whose shipped config does.

`test_tied_head_has_a_tp_plan_entry` pins it. The existing
`test_tp_plan_matches_params` checks the converse — that an entry names a real
module — so nothing caught an entry that was simply absent. It fails on 57
classes without this change.

Measured on 8xB200 with an untied head, so both arms are valid, against a
202048x4096 head (torch.compile, cudagraphs):

    tp   decode with/without      weights per rank
    2    0.674 / 0.732 ms (-8%)   2.844 / 3.615 GiB
    4    0.689 / 0.737 ms (-7%)   2.193 / 3.350 GiB
    8    0.724 / 0.802 ms (-10%)  1.867 / 3.216 GiB

Prefill is a wash (+2.8% at tp=2, -2.2% at tp=4 and 8). Eager numbers are not
reported: decode is host-bound there and the sign of the delta flips between TP
sizes.

Also drops `ColwiseParallel.validate_param`, which refused an output size
indivisible by the world size. `Shard._to_replicate_tensor` pads each shard
before the all-gather and trims after, and has since torch 2.5, the floor in
setup.py -- the neighbouring `test_colwise_uneven_local_shapes` already covered
the same thing for plain colwise. Verified at 2 and 8 ranks against an
unsharded reference, including GotOCR2's 151860 vocabulary: forward exact, and
weight and bias gradients bit-identical, with dgrad matching the even case to
the same 1 ulp of float64.

`_tp_plan` lands after `pack_image_features` in DeepseekOcr2: a class-level
assignment placed before a raise-only method makes the modular converter drop
that method from the generated file.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

* Declare the head plan once per inheritance chain

39 of the declarations the previous commit added were redundant: the modular
converter flattens a parent's class attributes into the child's generated file,
so a model inheriting from one that already declares the entry gets it without
restating it. `Llava` alone covers Aria, AyaVision (and Cohere2Vision through
it), FastVlm, GotOCR2, InternVL, MiniMaxM3, Mistral3, Ovis2, PerceptionLM and
VipLlava; `PaliGemma` covers the Gemma3/3n/4 chain; `Qwen2VL` covers
Qwen2.5-VL, Glm4v and its descendants. 18 declarations remain, one per chain.

Every generated file is byte-identical apart from two spots, both intended:
the converter-placement comment in DeepseekOcr2 is gone with the declaration
it explained, and the `Ernie4_5_VL_Moe` BC alias drops a copy it inherits from
`Ernie4_5_VLMoe` at runtime.

`Qwen3OmniMoeTalker` keeps its own entry — it was there already, and its head
is `codec_head` rather than `lm_head`.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

* Use GotOcr2's indivisible vocab in the uneven colwise-gather test

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

* Give uneven TP shards their global shape and test VLM heads end to end

`DTensor.from_local` infers the global shape as `local * world` unless told otherwise, so an
output or input that does not divide by the world size (a 151860 vocab at TP=8, an odd MLP
width) was padded by the gather and never trimmed. The colwise output and rowwise input hooks
now pass the real shape, read off the weight DTensor.

The TP mixin no longer rounds the tiny vocab up to the world size, which hid exactly this; its
backward comparison slices uneven shards at the true offset; and it reads the vocab and plans
off the text config, so GotOcr2 and GLM-4V-MoE now run the real multi-rank tests.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

* Regenerate MiniCPM-V 4.7 with the head plan it inherits from 4.6

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

* Review: take a sharded dim's full width from the module, and simplify the tied-head plan test

- An uneven last-dim shard's global width comes from the module's declared size (`out_features`,
  `in_features`, `embedding_dim`) rather than its weight's shape, which a packed or quantized weight does not
  carry (packed FP4 halves `in_features`).
- The tied-head test reads the head and embedding through `get_output_embeddings` / `get_input_embeddings`, over
  the generative classes.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

* Review: make the last-dim shard helper a `TensorParallelLayer` static method with a shorter docstring

It is shared by `ColwiseParallel`'s output and `RowwiseParallel`'s input, so it lives on their common base.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

* Review: keep this PR to the tied-head plans, move uneven sharding to its own PR

The failure this PR fixes is a tied `lm_head` without a TP-plan entry on composite models: it reads the
`embed_tokens` weight a tying config shards, so `F.linear` got a plain input and a DTensor weight. The support for
dims that do not divide by the TP size (the explicit global shape on `from_local`, dropping the colwise-gather
divisibility check, uneven backward slicing, the un-rounded test vocab and GotOcr2's TP tests) moves to its own PR.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

* TP: split a linear layer unevenly when its width does not divide by the TP size

A colwise layer whose output is gathered (an `lm_head`) had to divide by the TP size, so a model with a vocab such
as GotOcr2's could not use TP. A linear layer can now split unevenly: the last rank holds a shorter slice, and the
colwise output and rowwise input wrap their local slice with the module's declared width (`out_features`,
`in_features`; a packed or quantized weight's shape does not show it), since `DTensor.from_local` would otherwise
assume every rank holds the same width. A module without that width, such as an embedding, still has to divide.

The TP mixin keeps the tester's vocab as is and slices uneven gradient shards at their true offset; GotOcr2 joins
the multi-rank tests.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

* Keep this PR to the tied-head plans; uneven linear sharding moves to the linear/embedding TP split PR

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

* Review: drop the `is_torch_available` guard around GLM-4V MoE's `causal_lm_class`

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>

Author: 3 people

src/transformers/models/aria/modeling_aria.py

@@ -978,6 +978,7 @@ def _create_patch_attention_mask(self, pixel_mask):
 )
 class AriaForConditionalGeneration(AriaPreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_gather_output"}
 
     def __init__(self, config: AriaConfig):
         super().__init__(config)

src/transformers/models/aya_vision/modeling_aya_vision.py

@@ -297,6 +297,7 @@ def forward(
 )
 class AyaVisionForConditionalGeneration(AyaVisionPreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_gather_output"}
 
     def __init__(self, config: AyaVisionConfig):
         super().__init__(config)

src/transformers/models/cohere2_vision/modeling_cohere2_vision.py

@@ -252,6 +252,7 @@ def forward(
 )
 class Cohere2VisionForConditionalGeneration(Cohere2VisionPreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_gather_output"}
 
     def __init__(self, config: Cohere2VisionConfig):
         super().__init__(config)

src/transformers/models/cohere_compass/modeling_cohere_compass.py

@@ -1383,6 +1383,7 @@ def forward(
 )
 class CohereCompassForConditionalGeneration(CohereCompassPreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_gather_output"}
     # Reference: fix gemma3 grad acc #37208
     accepts_loss_kwargs = False
 

src/transformers/models/cosmos3_edge/modeling_cosmos3_edge.py

@@ -1160,6 +1160,7 @@ class Cosmos3EdgeCausalLMOutputWithPast(CausalLMOutputWithPast):
 @auto_docstring
 class Cosmos3EdgeForConditionalGeneration(Cosmos3EdgePreTrainedModel, GenerationMixin):
     _tied_weights_keys = {}
+    _tp_plan = {"lm_head": "colwise_gather_output"}
     config_class = Cosmos3EdgeConfig
     accepts_loss_kwargs = False
 

src/transformers/models/cosmos3_omni/modeling_cosmos3_omni.py

@@ -517,6 +517,7 @@ class Cosmos3OmniCausalLMOutputWithPast(CausalLMOutputWithPast):
 
 class Cosmos3OmniForConditionalGeneration(Cosmos3OmniPreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_gather_output"}
     # Reference: fix gemma3 grad acc #37208
     accepts_loss_kwargs = False
 

src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py

@@ -1553,6 +1553,7 @@ def forward(
 @auto_docstring
 class DeepseekOcr2ForConditionalGeneration(DeepseekOcr2PreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_gather_output"}
 
     def __init__(self, config: DeepseekOcr2Config):
         super().__init__(config)

src/transformers/models/deepseek_vl/modeling_deepseek_vl.py

@@ -233,6 +233,7 @@ class DeepseekVLForConditionalGeneration(DeepseekVLPreTrainedModel, GenerationMi
     _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
     output_modalities = ("text",)
     _can_compile_fullgraph = True
+    _tp_plan = {"lm_head": "colwise_gather_output"}
 
     def __init__(self, config: DeepseekVLConfig):
         super().__init__(config)

src/transformers/models/deepseek_vl_hybrid/modeling_deepseek_vl_hybrid.py

@@ -414,6 +414,7 @@ class DeepseekVLHybridForConditionalGeneration(DeepseekVLHybridPreTrainedModel,
     _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
     output_modalities = ("text",)
     _can_compile_fullgraph = True
+    _tp_plan = {"lm_head": "colwise_gather_output"}
 
     def __init__(self, config: DeepseekVLHybridConfig):
         super().__init__(config)

src/transformers/models/diffusion_gemma/modeling_diffusion_gemma.py

@@ -1586,6 +1586,7 @@ class DiffusionGemmaForBlockDiffusion(DiffusionGemmaPreTrainedModel, DiffusionGe
 
     _tied_weights_keys = {"lm_head.weight": "model.decoder.embed_tokens.weight"}
     generation_config_class = DiffusionGemmaGenerationConfig
+    _tp_plan = {"lm_head": "colwise_gather_output"}
 
     def __init__(self, config: DiffusionGemmaConfig):
         super().__init__(config)