New example: [REDACTED] - #7121
qgallouedec wants to merge 18 commits into
Conversation
Fixed-config script plus its SLURM launcher: SFTTrainer with a load-time DistributedConfig (tp_size=8, fsdp_size=8, expert parallelism with token dispatch) on 64 GPUs. Defaults are the measured operating point; the constraints behind them are in the docstring.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
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. |
There was a problem hiding this comment.
💡 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".
…d rendezvous ignores
There was a problem hiding this comment.
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 torchrunnever reads it, so its pins install nothing. Note the tension:examples/README.mdrequires the header, and #6869's two SLURM examples have no header at all. Three files, three answers. One decision would help.
- The PEP 723 header is inert.
Findings
- 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).
- No environment guard (should fix). Verified on current transformers:
DistributedConfig(..., expert_parallel_dispatch=True)raisesTypeError: DistributedConfig.__init__() got an unexpected keyword argument 'expert_parallel_dispatch'. On 64 ranks, and afterload_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 aRuntimeErrorthat names the branch to install. That also makes Codex point above harmless. - Exactly one checkpoint, at the last step, never in
output_dir.save_strategydefaults to "steps" and save_steps to 500, equal to max_steps, so there is a single save intooutput_dir/checkpoint-500, and notrainer.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 chosesave_strategy="no"plus an explicit save_model, so the two examples now show two different save recipes with no stated reason. dataset_num_procunset. TRL prepares the dataset insidemain_process_first()withnum_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..slurmdrift 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 dropsexport 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.- Router aux loss, probably unintended.
router_aux_loss_coefdefaults to 0.001, soaux_loss_enabledis True for any MoE andoutput_router_logits=Trueis forwarded every step (sft_trainer.py:1410-1419). With LoRA onq/k/v/oonly, the router and the experts are frozen, so the load-balancing term cannot rebalance anything; it only perturbs the attention adapters. Underexpert_parallel_dispatchit is also a per-rank quantity, unlike the siblings' shared-batch EP. If unintended,router_aux_loss_coef=0.0drops 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-ignoreis inert underexamples/(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.
The branch it pinned, ep-fsdp-2d-mesh, has been superseded by ep-mesh (#48792), where ep_size is its own axis and token dispatch is selected by an ep_plan rule rather than by a boolean. peft 0.21.0 and accelerate 1.15.0 now carry the FSDP2 fixes the recipe relies on, so those become version pins. The layout changes with the API, tp=8 fsdp=8 becomes tp=1 fsdp=64 ep=8, so the measured numbers in the docstring are replaced by ones taken on this stack: 14.1 s/step at per-device batch 1, 131k tokens per step, 40 GB per GPU, loss 1.80 -> 0.72 over 100 steps.
…operational notes
…a find and replace
ep_size is not in a release yet, and the check sits before the 500k-row dataset load so 64 ranks fail in a second rather than after preprocessing. dataset_num_proc was unset, so those rows were tokenized and packed single-threaded on rank 0 while 63 ranks waited.
The docstring carried 14.1 s/step from the 100-step run at grad accum 1 while the script set 2. Dropping the accumulation makes the config the measured one: job at 8 nodes, 20 steps in 308 s, 15.4 s/step, 131k tokens/step, 40.2 GB peak.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 250afa2. Configure here.
|
thanks
It's because it's hip, when transformers releases, we will guard the transformers version. In the meantime:
Set to 32.
Fixed in 73cdd79, save_steps=50
Intended, keeping it. The router weights are frozen but its inputs are not, so routing still drifts under LoRA. But maybe I'm missing your point?
Yes, because it's WIP, we're working on this with the transformers team, in the meantime, to allow the team to test, we have to work with umbrella branches.
Out of scope here, #6869 is still wip.
torchrun is what accelerate launch runs underneath, see the internal thread. Nothing here needs an accelerate config, the model is sharded at load time by DistributedConfig.
Same answer, #6869 is wip, I'd ignore that branch for now.
Right. Setting gradient_accumulation_steps=1 so the docstring describes what was run: 14.1 s/step, 131k
removed. For the records it exists in other examples.
It was needed before v1.13.0 because of #7101. Now we can remove. |
| model=MODEL, | ||
| args=training_args, | ||
| train_dataset=train_dataset, | ||
| peft_config=LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj", "k_proj", "v_proj", "o_proj"]), |
There was a problem hiding this comment.
To be clear, all the experts and the router stay frozen, right ? Because the model is quite big, most of the params probably (90% are expert weights) so the LoRA adapter are on attention only which could end up to be very small.
Can this recipe work with something like :
LoraConfig(target_modules=[], target_parameters=["mlp.experts.gate_up_proj", "mlp.experts.down_proj"])🤔 ?

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 standalone example only; no changes to TRL trainers or core training paths, though the example pins bleeding-edge transformers for expert parallelism.
Overview
Adds a new
sft_redactedexample for LoRA SFT on a large MoE-style checkpoint across 8 nodes (64 H100s) usingSFTTrainerwith load-time sharding via transformers'DistributedConfig(fsdp_size=64,ep_size=8for expert parallelism with token dispatch).The training script loads allenai/tulu-3-sft-mixture (500k rows) before the model, enforces a minimum transformers dev version for
ep_size, uses packed sequences and parallel tokenization, and setssave_only_model=Truebecause optimizer checkpoints aren't supported for this sharding mode. A Slurm +torchrunlauncher wires multi-node rendezvous (IPv4 master selection).The examples index in
docs/source/example_overview.mdis updated with a row forsft_redacted. Running the example currently depends on unreleased transformers (noted in script dependencies / version guard).Reviewed by Cursor Bugbot for commit 94ef601. Bugbot is set up for automated code reviews on this repo. Configure here.