Fix ZeRO-3 deadlock: mark MoE blocks as z3 leaf modules - #7
Open
Xuan-1998 wants to merge 6 commits into
Open
Conversation
- model_utils.get_model called get_quantization_config and get_kbit_device_map without importing them, raising NameError on first use. Import both from trl. - get_tokenizer crashed with ValueError for tokenizers that have no chat template, because tokenizer.get_chat_template() raises rather than returning None; inspect tokenizer.chat_template directly so the default template is applied as intended. Also pass use_fast=True (fast_tokenizer is not a recognized argument). - load_moe_bias_states skipped nothing for modules whose bias attribute is None (e.g. nn.Linear(bias=False)); guard with an isinstance check. - is_slurm_available crashed with CalledProcessError when sinfo exists but fails; treat that as Slurm unavailable. Verified on CPU: tokenizer without a template now receives the default chat template, get_model loads a tiny OLMoE checkpoint, and MoE bias states survive a save/load roundtrip on a model that also contains bias=None Linear modules.
Two defects in the OLMoE training path: - CustomSFTTrainer gated the deferred bias update on (global_step + 1) %% gradient_accumulation_steps == 0. global_step counts optimizer steps, not microbatches, so with accumulation N the bias updated on only 1/N of optimizer steps (and fired mid- accumulation N times on those). Replace the override with a MoeBiasUpdateCallback hooked on on_optimizer_step, which fires exactly once per optimizer step regardless of accumulation. - With the default remove_unused_columns=True the Trainer drops the raw 'messages' column that DataCollatorForChatML reads at collate time, so OLMoE/Qwen training crashed with KeyError: 'messages' unless the user knew to pass --remove_unused_columns False. Force it off automatically when the ChatML collator is used. Verified with a tiny OLMoE end-to-end SFT run on CPU (4 optimizer steps, 2 MoE layers): bias now updates exactly 8 times with both gradient_accumulation_steps=2 and 3 (previously 6 of 8 with 3), and training runs with default remove_unused_columns.
Under DeepSpeed ZeRO-3 (the configuration shown in the README), sparse expert dispatch makes different ranks fetch different expert parameters in different orders, and training deadlocks at the first step (observed hanging on an 8-GPU OLMoE-1B-7B run). Apply DeepSpeed's documented fix for MoE models: mark the whole MoE block as a ZeRO-3 leaf via deepspeed.utils.set_z3_leaf_modules, so each block's parameters are gathered as a unit. No-op when DeepSpeed is absent or ZeRO-3 unused. Verified on an 8x GPU node: OLMoE-1B-7B under ZeRO-3 previously hung at step 0 for 20+ minutes; with this change training completes, with the correct once-per-optimizer-step bias cadence on every rank. CPU and DDP paths unchanged.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Builds on #4/#5 (branched from #5 so the distributed experiment could run; after those merge this shows only one commit).
Problem
Training with DeepSpeed ZeRO-3 — the configuration shown in the README usage example — deadlocks at step 0. Reproduced on an 8-GPU node with the real allenai/OLMoE-1B-7B-0924: the run hung at
0/3steps for 20+ minutes before being cancelled. Root cause is the classic MoE × ZeRO-3 interaction: sparse expert dispatch makes different ranks gather different expert parameters in different orders, so the collective all-gathers never line up.Fix
Apply DeepSpeed's documented fix for MoE models: mark the whole MoE block as a ZeRO-3 leaf via
deepspeed.utils.set_z3_leaf_modules, so each block's parameters are gathered as a unit. For DeepSeek the (trust_remote_code) class is looked up on the instantiated model. No-op when DeepSpeed is not installed or ZeRO-3 is not used.all_reducebranch fires on every update, andmoe_bias_states.jsonis saved with non-zero biases for all 16 layers.accelerate launch --multi_gpu, 8 GPUs) still passes with identical counters, and the CPU end-to-end run is unchanged (the marking is a no-op there).