style(tests): reformat test_transformers_compat to unblock Lint - #33
Open
XiaohanZhangCMU wants to merge 18 commits into
Open
style(tests): reformat test_transformers_compat to unblock Lint#33XiaohanZhangCMU wants to merge 18 commits into
XiaohanZhangCMU wants to merge 18 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.
…D under CP) (#31) * build: bump Megatron-LM (explicit pad_between_seqs for tail-padded THD under CP) Picks up basetenlabs/Megatron-LM#25 (e43ae913f): TEDotProductAttention now passes the tail-inclusive pad_between_seqs answer explicitly under context parallelism, engaging TE's exact get_cu_seqlens_on_cp_rank path. Fixes the LPS-1063 nondeterministic forward/gradients and silent boundary-row mis-attention caused by TE's cu_seqlens[:-1] auto-detect (NVIDIA/TransformerEngine#3331). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * build: pin Megatron-LM to merged trainers-main HEAD (57efae08) --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Pre-existing on trainers-main and unrelated to any feature work, but it fails the Lint check in cicd-main.yml, and Lint gates the rest of that workflow -- so no PR against this fork can get a unit-test run until it is fixed. Found while trying to get unit tests to run on Megatron-Bridge #27: the manual cicd-main L0 dispatch (run 31925778616) failed at Lint on this file, which #27 never touched. The change is one line-length collapse produced by `ruff format`; no logic. Signed-off-by: Xiaohan Zhang <xiaohan.zhang@baseten.co>
This was referenced Aug 18, 2026
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
6 participants
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 — merge this first. It unblocks CI for everything else in this fork: Lint gates the rest of
cicd-main.yml, so #34 and #27 cannot reach a unit-test run until this lands. Independent of the Kimi-K3 work in every other respect, and deliberately not stacked under it so the unblock is not coupled to K3 review.tests/unit_tests/models/conversion/test_transformers_compat.pyfailsruff format --check. It is pre-existing ontrainers-mainand unrelated to any feature branch, but Lint gates the rest ofcicd-main.yml, so no PR against this fork can reach a unit-test run until it is fixed.Found while trying to get unit tests to run on #27: the manual
cicd-mainL0 dispatch (run 31925778616) failed at Lint on this file, which #27 never touches.One line-length collapse from
ruff format. No logic change.Worth landing on its own rather than inside a feature PR, so the unblock is not coupled to K3 review.