DIstributed branch base
Author: 3outeille
Pull request: yes

# Current state of things

- **Training has been confirming to work in TRL** thanks to @AmineDiro cf https://github.com/huggingface/transformers/pull/46126
  > I was validated end-to-end with trl.SFTTrainer on Qwen3-30B-A3B at multiple shapes (16k / 32k / 64k, 2n / 4n / 8n). Healthy training (loss bit-exact match against our prod fork at LR=0) and very competitive MFU :clap: :clap:
- **Inference works but depends on `Dtensor`** (which we don't want as it is very slow)
- DTensor-based TP
- Native FSDP2 support in transformer
- Sequence parallelism for activations and norms via per-model `_sp_plan` entries
- MoE expert parallelism + packed-weight sharding for fused `gate_up_proj` /
`grouped_mm` kernels
- `Shard-on-read` via `DtensorShardOperation`
- Resume under a different parallelism layout.
- Standard safetensors export. `model.save_pretrained(dir)` (no split files checkpoint when distributed) writes a fully-gathered
- Distributed checkpointing for both model and optimizer.

# What's missing:

- The PR has been reverted https://github.com/huggingface/transformers/pull/45028 due to:
  - missing support of EP
  - Inference on plain tensors for Continuous batching

# TODO

- Critical (blocking)
   - [x] Bring EP => `ep_router + moe_experts_ep_allreduce ParallelStyles` (work began by @AmineDiro at https://github.com/huggingface/transformers/pull/46126) 
   - Mix plain tensors & dtensors in `Colwise/Rowwise` to enable inference for continuous batching 
- Medium:
    - push checkpoint to `hf_buckets` by default 
    - [x] `MoEExpertsParallel`
       -  go from module level to param level  `("layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts": "moe_experts_allreduce",)`
- Minor:
   - [x] handling `sparse` + `dense` in the  `sp_plan` for models like `cohere2_moe`, `mellum`, `qwen3_moe`
   - [x]  default `tp_plan + ep_plan` for inference / `sp_plan + ep_plan` training
  - Add `DistributedMixin`  to isolate (`save_pretrained(distributed_checkpoint=True)` from `PreTrainedModel`
  - [x] cleaning FSDP by unifying manual/auto path

# PR

- https://github.com/huggingface/transformers/pull/46290 (param_level, ep_router, handling sparse + dense)
	- https://github.com/huggingface/transformers/pull/46292 (tp + ep_plan / sp + ep_plan)
		- https://github.com/huggingface/transformers/pull/46394
        - tp_inference_training
- https://github.com/huggingface/transformers/pull/46679
- DistributedMixin


# Split PR

- https://github.com/huggingface/transformers/pull/46705
	- https://github.com/huggingface/transformers/pull/46707
- https://github.com/huggingface/transformers/pull/46717

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_46269). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@3outeille & @ArthurZucker  should I port https://github.com/huggingface/transformers/pull/46126  or thats something already in work and just close it ? 

@AmineDiro you can change the base to this branch ! 

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=46269&sha=309729

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

run-slow: afmoe, apertus, arcee, aria, audioflamingo3, bamba, bitnet, cohere, cohere2, cohere2_moe, csm, cwm, data2vec, dbrx, deepseek_ocr2, deepseek_v2