[Feature][AutoDeploy]: Improved heuristics to detect shardable regions

[greg-kwasniewski1]

🚀 The feature, motivation and pitch

Current TP sharding heuristics to detect shardable regions (identify_regions_between_residuals) fails to properly distinguish layers (attention/MoE/MLP/SSM) for some models: Nemotron, LLama4, Phi.

To efficiently shard the model, we need head parallelism. To apply head parallelism, we need a more robust way to extract layers.

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

[coderabbitai]

Caution

The CodeRabbit agent's plans did not produce any file changes.