New example: [REDACTED]

[qgallouedec]

LoRA SFT on [REDACTED] on 8 nodes with SFTTrainer, sharded at load time by transformers' DistributedConfig.

Same shape as #6869's examples.

Needs huggingface/transformers#48204.

---

Note

Low Risk
Documentation and a self-contained example only; no changes to TRL trainers or library APIs.

Overview
Adds sft_moe_expert_parallel, a distributed-training example that runs LoRA SFT on a large MoE with SFTTrainer across 8 nodes (64 GPUs). The training script wires transformers.distributed.DistributedConfig into model_init_kwargs (tp_size=8, fsdp_size=8, expert parallel + token dispatch) so the checkpoint is sharded at load time, uses packed sequences, save_only_model=True (no resume), and documents pins on transformers ep-fsdp-2d-mesh and peft main.

A companion sft_moe_expert_parallel.slurm launches torchrun via Slurm with IPv4 rendezvous. docs/source/example_overview.md gains an index row for the new folder.

^{Reviewed by Cursor Bugbot for commit 878a8c3. Bugbot is set up for automated code reviews on this repo. Configure here.}

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-09T01:28:44.786856Z	e4a6163	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[bot-ci-comment]

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.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: e4a616384b

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[AmineDiro]

does this work out of the box with accelerate ? imo, we should reconsider the interaction between running with torchrun vs accelerate like we discussed to catch weird behavior. For example :

• setting this in training args + run with accelerate with a different config.
• run with accelerate + distributed_config
• accelerate has different FSDP or deepspeed ?

[AmineDiro]

so this depends on : huggingface/transformers#48518 ?

[albertvillanova]

Yes, this depends on transformers#48518. Verified by ancestry, ep-token-dispatch is merged into the pinned ep-fsdp-2d-mesh branch, which is where expert_parallel_dispatch is defined (absent from transformers main).

[albertvillanova]

Thanks.

Codex raised three points: 2 were fixed, and 1 declined:

• PEP 723 header vs sbatch: declined. I think the decline is factually right but leaves the actual gap open.
  • The PEP 723 header is inert. sbatch → srun torchrun never reads it, so its pins install nothing. Note the tension: examples/README.md requires the header, and #6869's two SLURM examples have no header at all. Three files, three answers. One decision would help.

Findings

1. Merge ordering. The only install path is transformers#48204, a draft umbrella branch that will disappear when its constituent PRs land. Separately, main's docs/source/distributing_training.md has no expert-parallelism or DistributedConfig prose at all: that section comes from #6869, still open. Merging this first puts an index row in main pointing at a concept documented nowhere in TRL's docs, and a launcher pattern (srun torchrun) that contradicts the only multi-node guidance main has (srun accelerate launch).
2. No environment guard (should fix). Verified on current transformers: DistributedConfig(..., expert_parallel_dispatch=True) raises TypeError: DistributedConfig.__init__() got an unexpected keyword argument 'expert_parallel_dispatch'. On 64 ranks, and after load_dataset("allenai/tulu-3-sft-mixture", split="train[:500000]") has already been paid for, since the dataset load precedes SFTConfig. The sibling already in main (sft_qwen3_8b_1m_context) sets the precedent: a 4-line check right after the imports raising a RuntimeError that names the branch to install. That also makes Codex point above harmless.
3. Exactly one checkpoint, at the last step, never in output_dir. save_strategy defaults to "steps" and save_steps to 500, equal to max_steps, so there is a single save into output_dir/checkpoint-500, and no trainer.save_model(). At the script's own configuration (grad accum 2 → 43.7 s/step) that is ~6.1 h against --time=08:00:00, before preprocessing. A wall-clock kill or preemption at step 480 leaves nothing after 64 H100-hours. The LoRA adapter "saves in seconds" per #6869's own docstring, so a real save_steps cadence is near-free insurance. Related: the measurements are quoted at 100 steps, so was the save path ever exercised? The sibling chose save_strategy="no" plus an explicit save_model, so the two examples now show two different save recipes with no stated reason.
4. dataset_num_proc unset. TRL prepares the dataset inside main_process_first() with num_proc=args.dataset_num_proc (sft_trainer.py:1453-1455), so 500k rows get tokenized and bfd-packed single-threaded on rank 0 while 63 ranks idle, on nodes where the launcher already requests --cpus-per-task=64. This is the exact failure mode the script's own comment warns about, and it eats the 8 h budget.
5. .slurm drift across the three launchers. This copy drops --node_rank "$SLURM_NODEID"; #6869's two copies still carry it (same dead flag, same reason). This copy also drops export HF_SHARD_PREFETCH=4, which both siblings set, while prescribing node-local staging instead and never mentioning the prefetch knob. Per the repo's consistency rule, propagate the node_rank removal and reconcile the two loading recipes, otherwise the three files read as drift rather than intent.
6. Router aux loss, probably unintended. router_aux_loss_coef defaults to 0.001, so aux_loss_enabled is True for any MoE and output_router_logits=True is forwarded every step (sft_trainer.py:1410-1419). With LoRA on q/k/v/o only, the router and the experts are frozen, so the load-balancing term cannot rebalance anything; it only perturbs the attention adapters. Under expert_parallel_dispatch it is also a per-rank quantity, unlike the siblings' shared-batch EP. If unintended, router_aux_loss_coef=0.0 drops the term and the router-logit memory. If intended, that is exactly what the docstring should say.

Nits

• the headline number (22.3 s/step, 23k tok/s) describes a configuration the script does not use; I checked the arithmetic and it is internally consistent (64 × 4 × 2048 / 22.3 = 23.5k tok/s), it just is not this config.
• # docstyle-ignore is inert under examples/ (the doc-builder hook is scoped to ^(trl|tests|docs/source)/).
• trl @ git+... where the 1M example pins plain "trl", and nothing here appears to need trl main.

Verified correct

Worth saying explicitly, since several claims in the docstring are load-bearing:

• The block-diagonal-mask claim is accurate: find_packed_sequence_indices plus packed_sequence_mask_function build it from position_ids when attention_mask is None.
• save_only_model=True is precisely what upstream's guard demands, and the comment paraphrases the upstream error faithfully.
• expert_parallel_dispatch's three preconditions (train_sampling_strategy == "random", sized dataset, dispatch_batches False) all hold: TRL leaves the strategy at the default and does not override _get_train_sampler.
• peft pin is right: #3578 merged 2026-08-27, v0.20.0 released 2026-07-28.
• The model_init_kwargs path (new here, the siblings call from_pretrained directly) is safe only because TRL forces device_map=None when distributed_type is MULTI_GPU, which is what torchrun gives. Running the same script single-process would inject device_map="auto" alongside distributed_config.
• report_to="trackio" with a non-JSON model_init_kwargs is safe: trackio's to_json_safe converts DistributedConfig via to_dict() and torch.bfloat16 via str(). Adjacent, not this PR: args.to_json_string() does raise TypeError: Object of type DistributedConfig is not JSON serializable, so report_to="tensorboard" would break on this and on the already-merged 1M example.

CI

Everything passes except "Tests with latest dependencies": test_grpo_trainer.py::TestGRPOTrainerVLM::test_train_vlm_with_pad_to_multiple_of, ValueError: Image features and image tokens do not match, tokens: 513, features: 512. A docs-and-examples diff cannot cause a GRPO VLM failure, and only the latest-deps job is red (dev, minimum and no-optional all pass), so it is an upstream-latest issue on a 4-commit-stale base. Rebase and re-run to clear it. mergeable_state is blocked, no human approval yet.

[albertvillanova]

Yes, this depends on transformers#48518. Verified by ancestry, ep-token-dispatch is merged into the pinned ep-fsdp-2d-mesh branch, which is where expert_parallel_dispatch is defined (absent from transformers main).