Skip to content

fix(flex-dispatcher): gate DeepEP/HybridEP on compute capability, not device name - #30

Open
XiaohanZhangCMU wants to merge 17 commits into
trainers-mainfrom
xiaohan/deepep-compute-capability-gate
Open

fix(flex-dispatcher): gate DeepEP/HybridEP on compute capability, not device name#30
XiaohanZhangCMU wants to merge 17 commits into
trainers-mainfrom
xiaohan/deepep-compute-capability-gate

Conversation

@XiaohanZhangCMU

@XiaohanZhangCMU XiaohanZhangCMU commented Aug 7, 2026

Copy link
Copy Markdown

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:

device_properties.major in [8, 9] or device_properties.name.startswith(("NVIDIA B200", "NVIDIA B300"))

name is a marketing string and does not reliably identify the SKU. Our B300 hosts enumerate as NVIDIA L20D, so that startswith test 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_backend and validate_flex_dispatcher_backend (two of them written as not 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 what 10 means.

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 - 8 and 9 were already accepted by the major test, 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_dispatcher unset - which is precisely why it was wrong to carry this in the K3 PR.

Paras Stefanopoulos and others added 17 commits July 9, 2026 11:31
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>
…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.
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.

6 participants