Add Parallel Cross Entropy

[zhenglongjiepheonix]

What does this PR do?

• add parallel cross-entropy implementation
• add parallel axis propagation rules for cross-entropy
• add parallel cross-entropy replacement logic in pass
• modify tests to trigger the critical path

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

[michaelbenayoun]

LGTM!
Is it for the default backend?

[zhenglongjiepheonix]

LGTM! Is it for the default backend?

Nanotron backend will also use it because there is a label_mask input in Nanotron which is not compatible with transformers current implementation, the current sharded cross-entropy is compatible with transformers implementation, and it does not matter because cross-entropy layer does not contain any backend-specific stuffs, it does not even contain a single parameter

[zhenglongjiepheonix]

If everything is good, can we merge this, I don't have write access to the main branch now :(