Repository navigation
Conversation
|
[For maintainers] Suggested jobs to run (before merge) run-slow: qwen3_moe |
|
View the CircleCI Test Summary for this PR: https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=46126&sha=e8de65 |
|
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. |
ArthurZucker
left a comment
ArthurZucker
left a comment
There was a problem hiding this comment.
Ty, let's not push for wrong patterns
| "moe_experts_ep_allreduce": MoEExpertsParallel( | ||
| output_layouts=Replicate(), | ||
| shard_plan={ | ||
| "gate_up_proj": Shard(0), | ||
| "down_proj": Shard(0), | ||
| }, | ||
| ), |
There was a problem hiding this comment.
No no no no nonoooooo this is the worst pattern that was ever introduced in the lib, sorry you have to "pay the price" but EP should be back to how we did it: gate_up_proj gets "gemm" and down_proj as well, te module gets ep_allreduce!
There was a problem hiding this comment.
no worries 🤣 ! Ok, I was just following the exact implementation pattern:
Just sharding on the dim 0 for the experts. If I go back to :
"layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
"layers.*.mlp.experts.down_proj": "grouped_gemm",
"layers.*.mlp.experts": "moe_experts_allreduce",There is some implementation logic that uses the new TensorParallelStyle that needs to be rewritten, I guess
| local_grads = [_replicate_dtensor(g).to_local() if isinstance(g, DTensor) else g for g in grads] | ||
| total_norm = torch.nn.utils.get_total_norm(local_grads, norm_type=norm_type) | ||
| torch.nn.utils.clip_grads_with_norm_(grads, max_norm=max_norm, total_norm=total_norm) | ||
| """Mesh-aware grad-norm clip with O(1) extra memory per parameter.""" |
There was a problem hiding this comment.
remove unrelated if unrelated
There was a problem hiding this comment.
Oh yes I created I different PR : #46113 for the grad clipping issue
| "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", | ||
| "layers.*.mlp.experts.down_proj": "grouped_gemm", | ||
| "layers.*.mlp.experts": "moe_experts_allreduce", |
There was a problem hiding this comment.
this is what we want
| # deepspeed and accelerate flags covering both trainer args and accelerate launcher. | ||
| # Treated as exclusive with `is_native_distributed`: if transformers already manages | ||
| # distribution, the accelerate plugin (e.g. left over from `accelerate launch --use_fsdp`) | ||
| # is not the responsible party — skip its code paths to avoid operating on a model | ||
| # accelerate didn't wrap. | ||
| _ds = getattr(self.accelerator.state, "deepspeed_plugin", None) is not None | ||
| _fsdp = getattr(self.accelerator.state, "fsdp_plugin", None) is not None | ||
| if self.is_native_distributed and (_ds or _fsdp): | ||
| logger.warning_once( | ||
| "`distributed_config` is set on the model AND accelerate has a " | ||
| f"{'deepspeed' if _ds else 'fsdp'}_plugin in state. Trusting `distributed_config` " | ||
| "(transformers-native) and ignoring the accelerate plugin." | ||
| ) | ||
| self.is_deepspeed_enabled = _ds and not self.is_native_distributed | ||
| self.is_fsdp_enabled = _fsdp and not self.is_native_distributed |
There was a problem hiding this comment.
not super clean IMO but in trainer soooo
There was a problem hiding this comment.
Hah, agreed on cleanliness, but let me defend this 😆. I just tried to localize variables that clearly indicate which mode we are running, so we can do stuff like self.is_deepspeed_enabled or self.is_fsdp_enabled or self.is_native_distributed. Because we can be running in fsdp from accelerate of Zero3 from native and maybe have passed the DistributedConfig, we just want to avoid overlap and be explicit in the condition.
But happy to change this if you have a better idea!
Builds on #45028 (
refactor-tp-dtensor) so that Expert Parallelism actually runs on any MoE model that has abase_ep_plan:I was validated end-to-end with
trl.SFTTraineronQwen3-30B-A3Bat 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 👏🏼 👏🏼What does this PR do?
ep_router+moe_experts_ep_allreduceParallelStyles (distributed/tensor_parallel.py).RouterParallelmasks non-local expert scores.MoEExpertsParallelshards on the experts dim.model._tp_plan→model.tp_plansoenable_expert_parallel=Trueactually selects_ep_plan(the property override atmodeling_utils.py:1459).Trainer._clip_grad_normroutes to the mesh-aware streaming clip (trainer.py) when on the native path. Stockaccelerator.clip_grad_norm_across cross-mesh DTensor grads crashes on the firstoptimizer.step.Trainer.create_optimizercallsmaybe_disable_foreach_and_fused_for_mixed_dtensor_groupshelper; Adam fused/foreach rejects mixed Tensor/DTensor groups underfsdp_size=1.Trainer.compute_lossdividesloss_scalebytp_sizefor theDistributedConfigpath (trainer.py). Companion to Trainer.compute_loss: fix loss over-counting under TP and EP-as-TP #45994 (which only handlesparallelism_config.tp_size).Trainer.is_native_distributedflag, exclusive withis_fsdp_enabled/is_deepspeed_enabled@3outeille tell me what you think. This is made to make clear when we are running in native vs accelerate accross branches.base_model_ep_plansimplified to usemoe_experts_ep_allreduce: @ArthurZucker I wonder if we need to have this ep plan dynamically set by the user and not part of the model definition 🤔Note:
This PR depends on :
_clip_grad_normroute depends on the streaming O(1)-per-param clip from clip_grad_norm avoid O(global_params) replication #46113. Without it, every 30B+ EP+FSDPoptimizer.stepallocates a ~60 GB transient (currentclip_grad_normreplicates every DTensor grad). So we should land clip_grad_norm avoid O(global_params) replication #46113 first.Who can review?
@3outeille @ArthurZucker @SunMarc