🚀 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...
🚀 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...