[TP plans] Fix some incorrects TP plans

[Cyrilvallez]

What does this PR do?

As per the title. See here for the source

TODO CYRIL: models modified only apply norm on head_dim!! So should be fine without replication. But then nano_chat needs to remove useless replication, and let's check other models with norms to be sure

cc @vasqu

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: apertus, doge, nanochat

[vasqu]

SP is super slow so removed it while I was at it

[vasqu]

Only a bit weird one is Cohere, it has this use_qk_norm attribute but it defaults to false and the base models also doesn't use it. If it were to use it, we would need to rep, otherwise not. See

transformers/src/transformers/models/cohere/modular_cohere.py

Lines 139 to 146 in 352a2e0

	if self.use_qk_norm:
	# When sharding the model using Tensor Parallelism, need to be careful to use n_local_heads
	self.q_norm = CohereLayerNorm(
	hidden_size=(config.num_attention_heads, self.head_dim), eps=config.layer_norm_eps
	)
	self.k_norm = CohereLayerNorm(
	hidden_size=(config.num_key_value_heads, self.head_dim), eps=config.layer_norm_eps
	)

Not sure if we just rep either way but left it for now

[Cyrilvallez]

Thanks a lot for double checking @vasqu! For cohere, if the main checkpoints do not use the norm, let's not replicate as it's a huge performance drawback!
Cannot approve my own PR, so just gonna merge it!