Tensor parallelism skips cross-attention blocks: apply_tensor_parallel never visits cross_attention_layers #341

Summary

apply_tensor_parallel builds a TP plan for every block in model.layers and never looks at model.cross_attention_layers. On a VLM with arch = "cross_attention" and tp > 1, the CrossAttentionBlocks stay unsharded plain modules between TP'd, sequence-parallel blocks, and no config check rejects the combination.

Where (fb23328)

• kempnerforge/distributed/tensor_parallel.py:143-160: only model.layers is parallelized; :129 turns sequence parallel on for a dense, untied, non-PP model.
• kempnerforge/distributed/parallel.py:461-462: the VLM build applies TP to the inner Transformer for every arch.
• kempnerforge/model/transformer.py:193-211 (the CA blocks) and :527-550 (h = ca(h, image_features, image_mask) after selected blocks).
• kempnerforge/config/job.py:182-199: no TP check by arch.

Why it matters

With sequence parallel on, each TP'd block hands the next a Shard(1) DTensor. A CA block that was never parallelized receives that DTensor with plain-tensor weights. From the code, this most likely fails at the first CA block's norm (kempnerforge/model/norm.py:30) with a mixed torch.Tensor/DTensor error. If some path lets it run instead, the CA weights are replicated on every TP rank with no TP-group gradient reduction, so the copies can drift apart. Either way TP does not apply to the CA blocks (on 17248f8, none of a CA block's parameters became DTensors under tp = 2). No shipped or example config combines cross-attention with tp > 1 (examples/vlm/configs/*cross_attn*.toml, vlm_7b_freeze_schedule.toml and vlm_debug_moe.toml leave tp = 1); it is reached by adding --distributed.tp to one of them.

How to see it

On 2 GPUs, build a small cross-attention VLM (for example a scaled-down examples/vlm/configs/vlm_7b_cross_attn.toml, or a tiny ModelConfig with CrossAttentionConfig(cross_attention_every_n_layers=2)) with --distributed.tp=2. (1) After build_parallel_model, list which parameters are DTensors: the CA blocks' are not. (2) Run one forward and backward to see whether it raises; if it trains, compare one CA block's weights across the two TP ranks after a few steps.

Suggested fix

Either add a TP plan for CrossAttentionBlock and iterate cross_attention_layers in apply_tensor_parallel, or reject tp > 1 with arch = "cross_attention" in JobConfig.validate until that exists. The block's attribute names differ from TransformerBlock (attn_norm, attn.{q,k,v,o}_proj, mlp_norm, mlp), so _build_block_tp_plan cannot be reused as is, and k_proj/v_proj take image_features, which is not sequence-sharded. Keep any guard specific to cross-attention: the joint-decoder arch's transformer blocks do get sharded.

Related

Open PR #284 also edits apply_tensor_parallel; #214 / #267 (norm initialization on the tensor-parallel meta-device path).