Skip to content

New example: [REDACTED] - #7121

Open
qgallouedec wants to merge 18 commits into
mainfrom
sft-moe-expert-parallel-8-nodes-v2
Open

qgallouedec wants to merge 18 commits into
mainfrom
sft-moe-expert-parallel-8-nodes-v2

Conversation

@qgallouedec

@qgallouedec qgallouedec commented Sep 9, 2026

Copy link
Copy Markdown
Member

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_redacted example for LoRA SFT on a large MoE-style checkpoint across 8 nodes (64 H100s) using SFTTrainer with load-time sharding via transformers' DistributedConfig (fsdp_size=64, ep_size=8 for 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 sets save_only_model=True because optimizer checkpoints aren't supported for this sharding mode. A Slurm + torchrun launcher wires multi-node rendezvous (IPv4 master selection).

The examples index in docs/source/example_overview.md is updated with a row for sft_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.

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.
@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Sep 9, 2026

Copy link
Copy Markdown

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

bot-ci-comment Bot commented Sep 9, 2026

Copy link
Copy Markdown

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 chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment thread examples/sft_moe_expert_parallel/sft_moe_expert_parallel.slurm Outdated
Comment thread examples/sft_moe_expert_parallel/sft_moe_expert_parallel.py Outdated
Comment thread examples/sft_moe_expert_parallel/sft_moe_expert_parallel.py Outdated
Comment thread examples/sft_moe_expert_parallel/sft_moe_expert_parallel.py Outdated
Comment thread examples/sft_moe_expert_parallel/sft_moe_expert_parallel.py Outdated

@albertvillanova albertvillanova left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread examples/sft_moe_expert_parallel/sft_moe_expert_parallel.py Outdated
qgallouedec and others added 14 commits September 15, 2026 12:17
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.
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.

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

Fix All in Cursor

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

Comment thread examples/sft_redacted/sft_redacted.py
@qgallouedec

Copy link
Copy Markdown
Member Author

thanks

No environment guard (should fix).

It's because it's hip, when transformers releases, we will guard the transformers version. In the meantime:
Fixed in 2cee151

dataset_num_proc unset.

Set to 32.

Exactly one checkpoint, at the last step, never in output_dir.

Fixed in 73cdd79, save_steps=50

Router aux loss, probably unintended.

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?

The only install path is transformers#48204, a draft umbrella branch that will disappear when its constituent
PRs land.

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.

Merging this first puts an index row in main pointing at a concept documented nowhere in TRL's docs

Out of scope here, #6869 is still wip.

a launcher pattern (srun torchrun) that contradicts the only multi-node guidance main has

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.

.slurm drift across the three launchers.

Same answer, #6869 is wip, I'd ignore that branch for now.

the headline number (22.3 s/step, 23k tok/s) describes a configuration the script does not use

Right. Setting gradient_accumulation_steps=1 so the docstring describes what was run: 14.1 s/step, 131k
tokens/step, 40.2 GB.

# docstyle-ignore is inert under examples/

removed. For the records it exists in other examples.

trl @ git+... where the 1M example pins plain "trl", and nothing here appears to need trl main.

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"]),

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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"])

🤔 ?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants