fix(flex-dispatcher): gate DeepEP/HybridEP on compute capability, not device name - #30
Open
XiaohanZhangCMU wants to merge 17 commits into
Open
fix(flex-dispatcher): gate DeepEP/HybridEP on compute capability, not device name#30XiaohanZhangCMU wants to merge 17 commits into
XiaohanZhangCMU wants to merge 17 commits into
Conversation
Rebase the still-relevant Baseten trainers-main Bridge fixes onto current NVIDIA-NeMo main now that most GLM-5.x support has landed upstream. Keep the Megatron-LM submodule pointed at the clean Baseten LM rebase branch so Bridge picks up the remaining LM fixes, including absorbed-MLA LoRA support, without pulling in Jack Rao's CP branch. Keep the HuggingFace remote-code compatibility shim that makes transformers materialize transitive relative imports, which is needed for Kimi-style custom model code. Keep the GLM-5.2 raw config workaround for qk_nope_head_dim/qk_rope_head_dim because some transformers configs still collapse the MLA split dimensions. Keep GLM-5.2-FP8 blockwise dequantization on HF weight load and preserve the IndexShare skip-offset default from first_k_dense_replace. Dropped older local changes that are now native upstream, including GLM-5.x Bridge registration, DSA RoPE interleave plumbing, and the Qwen3-ASR docstring workaround.
…18) maybe_enable_recompute_inputs_grad only patched TransformerBlock, but since megatron-core 0.19 HybridStack honours recompute_granularity='full' through the same reentrant tensor_parallel.checkpoint. A reentrant checkpoint records a backward node only when a tensor input requires grad, so with a frozen base model (adapter-only training at PP=1) every hybrid decoder chunk dropped out of the autograd graph and LoRA gradients were silently zero. Patch HybridStack the same way as TransformerBlock. Verified on Nemotron-3-Super-120B-A12B LoRA (TP=8, EP=8, PP=1, B200): grad_norm goes from 0.0 to healthy values and loss descends on a repeated sample.
* perf(ckpt): vectorize FP8 blockwise dequant (bit-exact, ~17x/tensor) The per-block Python loop made GLM-5.2-FP8's 800B load CPU-bound (8 workers pegged ~50 min). Two repeat_interleaves + one multiply, verified bit-exact against the loop. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(glm5): resolve raw config.json via the hub cache for HF-id launches The qk_rope_head_dim workaround (transformers GlmMoeDsaConfig collapses the qk head-dim split, corrupting kv_a_proj shapes: 704 != 576) only read <base_model>/config.json as a filesystem path, so it silently no-opped when base_model is a hub repo id and the 800B weight load failed. Resolve the raw config through hf_hub_download when the local read misses — transformers has already cached config.json by the time the bridge runs, so this works offline (HF_HUB_OFFLINE) too. Warn loudly when neither path resolves. Prod launches GLM-5.2-FP8 by HF id, so this unblocks the registry row in trainers#592. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * build: pin Megatron-LM to the packed-CP q_causal_offsets fix (LM#14 head) * build: update Megatron-LM PR pin Keep the Bridge submodule aligned with the current LM#14 head before the dependent Bridge and Trainers PRs merge. Signed-off-by: Jack Rao <jack.rao@baseten.co> * fix(glm5): fail fast without raw config The raw config preserves GLM qk dimensions, so conversion must terminate instead of continuing with invalid MLA shapes. Remove redundant comments. Signed-off-by: Jack Rao <jack.rao@baseten.co> Co-authored-by: Cursor <cursoragent@cursor.com> * build: advance Megatron-LM PR pin Keep the Bridge submodule aligned with the current LM#14 head before the dependent PRs merge. Signed-off-by: Jack Rao <jack.rao@baseten.co> * docs(glm5): clarify raw config workaround State the Transformer configuration defect and why the raw config is required. Signed-off-by: Jack Rao <jack.rao@baseten.co> Co-authored-by: Cursor <cursoragent@cursor.com> * build: repin Megatron-LM after packed-CP merge Point the Bridge submodule at the landed trainers-main commit so the packed-CP indexer fix remains fetchable after the PR branch is removed. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Signed-off-by: Jack Rao <jack.rao@baseten.co> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Cursor <cursoragent@cursor.com>
* build: pin Megatron-LM odd top-k fallback Advance Megatron-LM to the cuDNN DSA odd top-k fallback so packed CP segments avoid the frontend vector-width assertion.
…imports hook The recursive check_imports hook walks transitive relative imports by opening files on disk, but on the hub path transformers calls check_imports before downloading the module's relative imports — on a cold HF cache the walk dies with FileNotFoundError (Kimi-K2.6: modeling_deepseek.py). Guard the walk on file existence and return missing siblings by name: get_cached_module_file downloads each named module and re-enters the hook on it, so deeper imports are covered on that pass. Local-dir loads (all files present) still return the full transitive closure. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…hed_module_file Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This reverts commit 357c464.
…ports-cold-cache [baseten] fix(model): tolerate missing remote-code siblings in check_imports hook
Upstream Megatron-LM's training loop supports `--async-ckpt-use-cpu-shm` (see 3rdparty/Megatron-LM/megatron/training/config/training_config.py:535), which flips TorchDistSaveShardedStrategy.cpu_shm_mode=True. That causes the strategy to pass use_cpu_shm_for_gpu_tensors=True to FileSystemWriterAsync, staging GPU tensors into POSIX shared-memory in the training process before the async worker consumes them. The worker never has to receive CUDA IPC handles. Bridge's CheckpointConfig doesn't currently expose this knob, so users of the bridge-owned save_checkpoint path are stuck on the default CUDA IPC hand-off. That breaks on: - Sandboxed hosts where pidfd_getfd is not permitted (kernel.yama. ptrace_scope > 0) — the nvrx worker crashes with 'Operation not permitted', which is our motivating case at Baseten. - MNNVL systems where fabric handles are exhausted — the documented upstream use case. Add the field to CheckpointConfig, and pipe it through to TorchDistSaveShardedStrategy(cpu_shm_mode=...) in save_checkpoint. Use an inspect.signature guard so older megatron-core installs that don't know the kwarg fall back to the current behavior with a warning rather than crashing. Default False keeps every existing caller on today's path. Signed-off-by: Ian Korovinsky <ian.korovinsky@baseten.co>
feat(ckpt): expose async_ckpt_use_cpu_shm on CheckpointConfig
Pulls basetenlabs/Megatron-LM@a69b6b95d: the fht git source leaked into consumers' resolutions and overrode their prebuilt-wheel routing.
… device name
device_properties.name is a marketing string and does not reliably identify the
SKU. Our B300 hosts enumerate as "NVIDIA L20D", so the
name.startswith(("NVIDIA B200", "NVIDIA B300")) test skipped DeepEP on every one
of them and silently fell back to the alltoall dispatcher -- a performance
regression with nothing in the logs to explain it.
Compute capability comes from the driver and cannot be relabelled: 8 = Ampere,
9 = Hopper, 10 = Blackwell (B200 sm_100, B300 sm_103). Both call sites now share
one named constant instead of repeating the literal tuple.
Split out of the Kimi-K3 vendoring PR (#27): it is unrelated to K3, which runs
with moe_dispatcher unset and never takes the DeepEP path.
XiaohanZhangCMU
force-pushed
the
trainers-main
branch
from
August 20, 2026 05:44
20fcf2e to
1e60218
Compare
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
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.
Stack position — independent. Not part of the Kimi-K3 stack (#33 → #34 → #27) and not required by it: K3's golden trainer config uses the alltoall dispatcher, so it never takes this path. Reviewable and mergeable on its own, in any order.
Split out of #27, which is a Kimi-K3 vendoring PR. This is unrelated to K3 and is our own change, not upstream's, so it does not belong in a PR whose contract is "no original code".
The bug
DeepEP / HybridEP eligibility was gated partly on
device_properties.name:nameis a marketing string and does not reliably identify the SKU. Our B300 hosts enumerate asNVIDIA L20D, so thatstartswithtest is false on every one of them. DeepEP was skipped and the model silently fell back to the alltoall dispatcher - a throughput regression with nothing in the logs pointing at the cause, because the warning it emits reads as if the GPU genuinely were unsupported.The fix
Gate on compute capability, which comes from the driver and cannot be relabelled:
8= Ampere (A100, sm_80/86)9= Hopper (H100/H200, sm_90)10= Blackwell (B200 sm_100, B300 sm_103)The tuple was repeated at four call sites across
apply_flex_dispatcher_backendandvalidate_flex_dispatcher_backend(two of them written asnot x in [...]). It is now one named constant,_FLEX_DISPATCHER_CC_MAJORS, documented where it is defined - so the next reader does not have to ask what10means.Scope
Behaviour changes in exactly one direction: Blackwell GPUs that previously failed the name check now pass. No compute capability that used to be accepted is now rejected -
8and9were already accepted by themajortest, and any B200/B300 matching the old name check is CC 10.Not exercised on hardware here; this is a gating predicate, and the four sites are covered by the surrounding warning/raise paths. Kimi-K3 itself does not take this path - its golden trainer config leaves
moe_dispatcherunset - which is precisely why it was wrong to carry this in the K3 PR.