lora: GLM-5.1 + 5.2 (MoE+MLA+DSA) LoRA training — bridge support, launcher, registries, CI - #1373
lora: GLM-5.1 + 5.2 (MoE+MLA+DSA) LoRA training — bridge support, launcher, registries, CI#1373yushengsu-thu wants to merge 54 commits into
Conversation
GLM-5.1 / DeepSeek-V3.2 DSA indexer projections are named wq_b/wk/weights_proj in HF/SGLang but linear_wq_b/linear_wk/linear_weights_proj in the Megatron (bridge) model. Add them to _MLA_HF_TO_MEGATRON so a single --target-modules name resolves to the Megatron name for training and back to the HF name for SGLang rollout — the same mechanism already used for MLA (q_b_proj↔linear_q_up_proj).
--target-modules all-linear previously expanded to dense-only HF names (q/k/v/o/gate/up/down), which don't match MLA/DSA models. Extend it to also cover MLA (q_a_proj, kv_a_proj_with_mqa, q_b_proj, kv_b_proj) and the DSA indexer (wq_b, wk, weights_proj) so all-linear is meaningful for GLM-5.1 / DeepSeek-V3.2-family models. Names absent from a given model simply don't match (no-op); SGLang-unsupported targets are still filtered for rollout by target_modules_hf_for_sglang_rollout.
…A (GLM-5.1) The Megatron-Bridge model build for GLM-5.1 goes through megatron-core's experimental-attention dispatcher (not miles' --spec get_glm5_spec path). That dispatcher only wires "gated_delta_net" and raises for "dsa", even though megatron-core already ships a DSA builder (get_dsa_module_spec_for_backend). Monkey-patch get_experimental_attention_variant_module_spec (same pattern as deepseek_v4.get_dsv4_spec) to route "dsa" -> the existing DSA builder while the model is built, restoring the original in a finally block. The DSA builder omits metainfo, so also set metainfo["fuse_input_layernorm"]=False (MLA-based DSA keeps a separate input layernorm, like dsv4) which the variant layer builder requires.
…istry Typer launcher scripts/run_glm5_lora.py (modeled on run_deepseek_v4.py), a single-node LoRA example examples/lora/run-glm5.1-6layer-megatron-lora.sh, and the 6-layer registry entry scripts/models/glm5-744B-A40B_6layer.sh for GLM-5.1 GRPO LoRA via Megatron-Bridge. Both launchers verified e2e (rollout -> train -> save, TRAIN EXIT 0 + PEFT adapter) on jybsuper/GLM-5.1-6layer.
There was a problem hiding this comment.
Code Review
This pull request introduces support for GLM-5.1 GRPO LoRA training via the Megatron-Bridge path. It adds a single-node training example shell script, a Python training script (run_glm5_lora.py), a 6-layer pruned model configuration, and monkey-patches megatron-core's experimental-attention dispatcher to support the 'dsa' variant. Additionally, it maps DSA indexer modules for LoRA targeting. Feedback on these changes suggests adding a safety check for the DSA module spec to prevent potential crashes, utilizing a dataclass default_factory for unique run IDs, correcting a PYTHONUNBUFFERED environment variable typo, and calling super().post_init() in the subclassed arguments dataclass.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| spec = _eav_specs.get_dsa_module_spec_for_backend(config=config, backend=backend) | ||
| # get_dsa_module_spec_for_backend omits metainfo, but the experimental-variant | ||
| # layer builder reads attention.metainfo["fuse_input_layernorm"] (KeyError otherwise). | ||
| # MLA-based DSA keeps a separate (non-fused) input layernorm -- same as the | ||
| # deepseek_v4 dsv4 spec -> False (gated_delta_net uses True). | ||
| if spec.metainfo is None: | ||
| spec.metainfo = {} | ||
| spec.metainfo.setdefault("fuse_input_layernorm", False) | ||
| return spec |
There was a problem hiding this comment.
If get_dsa_module_spec_for_backend returns None (for instance, if the backend is unsupported or misconfigured), accessing spec.metainfo will raise an AttributeError. Adding a None check for spec protects against this potential crash.
| spec = _eav_specs.get_dsa_module_spec_for_backend(config=config, backend=backend) | |
| # get_dsa_module_spec_for_backend omits metainfo, but the experimental-variant | |
| # layer builder reads attention.metainfo["fuse_input_layernorm"] (KeyError otherwise). | |
| # MLA-based DSA keeps a separate (non-fused) input layernorm -- same as the | |
| # deepseek_v4 dsv4 spec -> False (gated_delta_net uses True). | |
| if spec.metainfo is None: | |
| spec.metainfo = {} | |
| spec.metainfo.setdefault("fuse_input_layernorm", False) | |
| return spec | |
| spec = _eav_specs.get_dsa_module_spec_for_backend(config=config, backend=backend) | |
| if spec is not None: | |
| # get_dsa_module_spec_for_backend omits metainfo, but the experimental-variant | |
| # layer builder reads attention.metainfo["fuse_input_layernorm"] (KeyError otherwise). | |
| # MLA-based DSA keeps a separate (non-fused) input layernorm -- same as the | |
| # deepseek_v4 dsv4 spec -> False (gated_delta_net uses True). | |
| if spec.metainfo is None: | |
| spec.metainfo = {} | |
| spec.metainfo.setdefault("fuse_input_layernorm", False) | |
| return spec |
|
|
||
| @dataclass | ||
| class ScriptArgs(U.ExecuteTrainConfig): | ||
| run_id: str = U.create_run_id() |
There was a problem hiding this comment.
In Python dataclasses, defining a field with a default function call like run_id: str = U.create_run_id() evaluates the function only once at module import time. This means all instances of ScriptArgs created during the lifetime of the process will share the same run_id. Using field(default_factory=U.create_run_id) guarantees that a unique run_id is generated for each instance.
| run_id: str = U.create_run_id() | |
| run_id: str = field(default_factory=U.create_run_id) |
|
|
||
| export GPUS_PER_NODE=${GPUS_PER_NODE:-4} | ||
| export HF_HOME=${HF_HOME:-/cluster-storage/models} | ||
| export PYTHONBUFFERED=16 |
There was a problem hiding this comment.
| def __post_init__(self): | ||
| if self.hf_checkpoint is None: | ||
| self.hf_checkpoint = ( | ||
| _DEFAULT_6LAYER_CKPT | ||
| if self.model_name == "GLM-5.1-6layer" | ||
| else f"/root/models/{self.model_name}" | ||
| ) |
There was a problem hiding this comment.
When subclassing a dataclass like U.ExecuteTrainConfig and overriding __post_init__, calling super().__post_init__() is a best practice to prevent bypassing any initialization or validation logic defined in the parent class.
| def __post_init__(self): | |
| if self.hf_checkpoint is None: | |
| self.hf_checkpoint = ( | |
| _DEFAULT_6LAYER_CKPT | |
| if self.model_name == "GLM-5.1-6layer" | |
| else f"/root/models/{self.model_name}" | |
| ) | |
| def __post_init__(self): | |
| super().__post_init__() | |
| if self.hf_checkpoint is None: | |
| self.hf_checkpoint = ( | |
| _DEFAULT_6LAYER_CKPT | |
| if self.model_name == "GLM-5.1-6layer" | |
| else f"/root/models/{self.model_name}" | |
| ) |
… monkey-patch) The GLM-5 / GLM-5.1 "dsa" experimental-attention-variant spec is now registered by the Megatron-Bridge GLM-5 bridge itself (glm5_bridge.py: provider.transformer_layer_spec = _build_glm5_dsa_block_spec, feature-detected so it is a no-op on newer megatron-core). So _setup_lora_model_via_bridge no longer needs the caller-side monkey-patch of megatron-core's experimental-attention dispatcher added in b0f45cb -- it just builds via the bridge. Requires a Megatron-Bridge that carries the dsa spec (bridge-dev-glm / the bridge PR). Verified e2e: GLM-5.1 6-layer GRPO LoRA via bridge -> Job succeeded + PEFT adapter saved, with no dsa patch in miles.
Add scripts/models/glm5.2-744B-A40B.sh (= GLM-5.1 dims, --rotary-base 8e6; converges with PR radixark#1376's registry) and glm5.2-744B-A40B_7layer.sh (7-layer prune: 3 dense + 4 MoE). Extend scripts/run_glm5_lora.py to map GLM-5.2 / GLM-5.2-7layer model-names to those registries, add an _HF_REPO download map, and default --hf-checkpoint to a local {model_dir}/{model_name} path. GLM-5.2's DSA cross-layer index sharing (index_topk_freq) is read from the HF config by the Megatron-Bridge GLM5 provider (CrossLayerDSAttention) -- no extra CLI args here. Verified e2e: run_glm5_lora.py GLM-5.2-7layer train-only (replay of a GLM-5.1 rollout dump) -> TRAIN EXIT 0 + PEFT adapter.
Add tests/e2e/megatron/test_glm5_lora_6layer_ci.py (GLM-5.1 full rollout->train smoke test) and test_glm5_2_lora_7layer_ci.py (GLM-5.2 train-only: sglang cannot serve the cross-layer rollout yet, so it replays a GLM-5.1 dump -- both toys share the GLM tokenizer + vocab 154880). Both follow the existing test_glm5_744b_a40b_4layer_ci.py pattern (register_cuda_ci + prepare/execute). Remove the ad-hoc examples/lora/run-glm5.1-6layer-megatron-lora.sh, now superseded by scripts/run_glm5_lora.py + the CI tests.
dd16ec8 to
96064c3
Compare
…sion clarity
Make the version explicit in the names of the GLM-5.1-only files I added, to
mirror the GLM-5.2 ones (glm5.2-* / glm5_2):
scripts/models/glm5-744B-A40B_6layer.sh -> glm5.1-744B-A40B_6layer.sh
tests/e2e/megatron/test_glm5_lora_6layer_ci.py -> test_glm5_1_lora_6layer_ci.py
(.py uses glm5_1, not glm5.1 — a dot is not a valid Python module name)
run_glm5_lora.py keeps its name: one launcher serves BOTH GLM-5.1 and GLM-5.2
(selected via --model-name), so it is not version-specific. The pre-existing
glm5-744B-A40B{,_4layer,_20layer}.sh registries are untouched (not added here).
Updated references: GLM-5.1-6layer -> glm5.1-744B-A40B_6layer in
run_glm5_lora.py's registry map, plus a comment in glm5.2-744B-A40B_7layer.sh.
Add --dsa-attention-backend {megatron-bridge,slime} (default megatron-bridge) under --megatron-to-hf-mode bridge; selects the GLM DSA sparse-MLA kernel backend. No effect on non-DSA models or the raw path.
model_provider sets provider.dsa_attention_backend from the arg (hasattr-guarded, DSA providers only) so Megatron-Bridge selects the fused (slime) vs unfused DSA attention backend.
…both backends run_glm5_lora.py: promote dsa_attention_backend to a first-class ScriptArgs field (default megatron-bridge, matching the arg default) instead of only reaching it via --extra-args; wire it into ckpt_args; and pick --qkv-format from the backend -- thd for slime's fused kernels, bshd for the unfused megatron-core path. Document the two backends: both run GLM-5.1 and GLM-5.2, full or LoRA; slime is training/forward-only and needs the optional tilelang dep + thd layout. utils/arguments.py: extend the --dsa-attention-backend help to flag slime's constraints (requires --qkv-format thd, training/forward-only) and that both backends cover GLM-5.1/5.2, full or LoRA.
….2_5layer model Swap the GLM-5.2 LoRA toy from jybsuper/GLM-5.2-7layer to Pinaster/GLM-5.2_5layer (3 dense + 2 MoE; computing layers 1,2,3 + skip 4,5), matching the full-FT example run_glm5_2_744b_a40b.py and its 5-layer CI so both GLM-5.2 paths use the same model. - scripts/models/glm5.2-744B-A40B_5layer.sh: add (source glm5.2-744B-A40B.sh, N_MOE_LAYERS=2) - scripts/models/glm5.2-744B-A40B.sh: adopt the canonical comment (MODEL_ARGS unchanged) - scripts/models/glm5.2-744B-A40B_7layer.sh: drop - run_glm5_lora.py: GLM-5.2_5layer in _HF_REPO / _MEGATRON_MODEL_TYPE / model_name + docstring - tests: rename test_glm5_2_lora_7layer_ci.py -> test_glm5_2_lora_5layer_ci.py (GLM-5.2_5layer)
…un_glm5_lora Make the slime (fused TileLang) backend the default for the GLM bridge path -- it is the rollout<->train-parity backend (matches slime's rollout kernels), so it is the right default for on-policy LoRA. And enable R3 (rollout routing replay, arxiv 2510.11370) by default. - arguments.py: --dsa-attention-backend default megatron-bridge -> slime (help reworded) - run_glm5_lora.py: dsa_attention_backend field default -> slime - run_glm5_lora.py: add use_r3 (default True) -> --use-rollout-routing-replay; on the slime backend also add --use-rollout-indexer-replay (only slime self-registers the indexer replay stream; the unfused megatron-bridge path has none, so it is skipped there) - run_glm5_lora.py: note that the DSA indexer only does sparse top-k (cross-layer path) when the full sequence exceeds index_topk (2048); at shorter seq it degenerates to dense
…ild path The LoRA path builds via _setup_lora_model_via_bridge (NOT model_provider's wrapped_bridge_provider), which did not propagate args.dsa_attention_backend to the provider, so SlimeMLASelfAttention fell back to the unfused (bshd) path on a thd input and crashed with 'not enough values to unpack (expected 4, got 3)'. Set provider.dsa_attention_backend before finalize and force the configured backend onto every module config (and each SlimeMLASelfAttention.config) after provide_distributed_model -- same value for both backends, so the unfused default is preserved. Document in _get_parallel_config that both backends run under the canonical TP=EP=ngpu + sequence-parallel layout (the slime fused path is SP-aware via the Megatron-Bridge slime_mla.py reconciliation).
… log-tail ray job submit (blocking) streams logs over a WebSocket; if that WebSocket drops (close 1006) the submit command fails and tears down the run even though the job is healthy. Gate a --no-wait submit behind MILES_RAY_SUBMIT_NO_WAIT (+ MILES_RAY_SUBMISSION_ID to pin the job id) so the job runs detached under Ray and can be polled via ray job status / ray job logs.
When the R3 replay consistency check (check_replay_result, --ci-test) flags mismatched tokens, MILES_R3_DIAG prints per mismatched token the router-score near-tie gap (rank topk vs topk+1) and the ratio of the replayed experts' score to the recompute's own top-k -- distinguishing a benign near-uniform-routing tie (ratio ~1) from a real divergence. Also cap the per-token warning loop at 5 to avoid flooding logs.
run_glm5_lora.py is single-node by design (hardcodes --actor-num-nodes 1 and execute_train only does a local `ray start --head`). This wrapper drives the multi-node flow without editing the launcher via three roles -- head / worker / launch: form a Ray cluster manually (head + workers), then submit with MILES_SCRIPT_EXTERNAL_RAY=1 and override --actor-num-nodes through --extra-args (argparse last-wins). TP=EP stay intra-node (NVLink); nodes are crossed with DP (PP stays 1, so the GLM-5.2 cross-layer DSA PP-split assert is a no-op). Critical: --num-gpus-per-node must equal the REAL per-node GPU count -- the rollout addr allocator uses it to map sglang engines to nodes, and a wrong value hands every engine the head node's dist_init_addr (worker-node engines then time out on a cross-node TCPStore rendezvous). It is passed both as the script flag and inside --extra-args. Env knobs: HEAD_IP/GPUS_PER_NODE/NUM_NODES/MODEL_NAME/DSA_BACKEND/NUM_ROLLOUT/ SAVE_INTERVAL, plus WANDB (on|offline|off, default on) with WANDB_API_KEY/_TEAM/ _PROJECT/_GROUP. NCCL/GLOO iface auto-detected (ip, then ifconfig fallback). Validated: GLM-5.2_5layer, 2 nodes x 4xH200, unfused, 50 steps -> Ray job SUCCEEDED.
Add a second example task beside gsm8k: --task dapo-math trains on DAPO-Math-17k
(zhuzilin/dapo-math-17k, hard long-CoT competition math) with the same boxed/SymPy
verifier (--rm-type math). Mirrors the task-dispatch pattern in run_deepseek_v4.py:
_download_dataset and the rollout_args now switch on args.task -- gsm8k keeps
{messages,label} parquet, dapo-math uses the {prompt,label} jsonl with --input-key
prompt. gsm8k flags are unchanged (existing example + CI default to it).
Also adds opt-in DAPO dynamic sampling (--dapo-dynamic-sampling +
--over-sampling-batch-size), wiring check_reward_nonzero_std to drop all-same-reward
groups. Off by default: on a model that scores 0 on every sample (the toy pruned
checkpoints) it would reject every batch and resample forever; enable on a model that
solves some problems. For dapo-math pass a longer --rollout-max-response-len (e.g.
4096); a >2048 total seq is also what makes the GLM-5.2 DSA indexer go genuinely sparse.
Validated end-to-end: --task dapo-math, GLM-5.2_5layer, 2 nodes x 4xH200, unfused,
1 step -> Ray job SUCCEEDED; rollout served real DAPO-Math problems, adapter saved.
Thread the run_glm5_lora.py task selection through the multi-node wrapper so DAPO-Math can be trained on N nodes, not just gsm8k. New env knobs on the launch role: TASK (gsm8k|dapo-math, default gsm8k), RESP_LEN (optional --rollout-max-response-len; use ~4096 for dapo-math long CoT), and DAPO_DYNAMIC_SAMPLING (on -> --dapo-dynamic-sampling; real model only -- a toy that scores 0 reward would resample forever). Header gets a DAPO-Math multi-node example. Verified the launch role emits --task dapo-math --input-key prompt --prompt-data .../dapo-math-17k.jsonl --rollout-max-response-len 1024.
6d851fb to
38777bf
Compare
The GLM-5 LoRA scripts had no eval. Add it, mirroring run_deepseek_v4.py:
gsm8k -> eval on the held-out gsm8k test split (data_dir/gsm8k/test.parquet),
n-samples 1, eval-max-response-len = train resp (512).
dapo-math -> eval on AIME-2024 (data_dir/aime-2024/aime-2024.jsonl), n-samples 8,
eval-max-response-len 4096. NB: AIME-2024 is not bundled with dapo-math-17k --
supply it (or run EVAL=off) for dapo eval; gsm8k works out of the box.
- run_glm5_lora.py: add enable_eval (default True) + eval_interval (default 5) fields; build
eval_args per task (explicit --eval-input-key/--eval-label-key since miles defaults them None);
splice {eval_args} into train_args.
- run_glm5_lora_multinode.sh: EVAL (on/off) + EVAL_INTERVAL (default 5) knobs; pass
--no-enable-eval when EVAL=off and --eval-interval through to the launcher.
- run_glm5_lora_multinode_full.sh: export EVAL / EVAL_INTERVAL presets.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MuN5ow5hHFxPze36uhgZC4
Root cause (GLM-5.2 744B eval hang in colocate): eval_rollout_single_dataset created an asyncio task for the ENTIRE eval set (e.g. 1319 gsm8k prompts) up front, so in-flight concurrency was bounded only by the large global rollout semaphore (= sglang_server_concurrency * rollout_num_gpus / rollout_num_gpus_per_engine = 512*64/32 = 1024). That burst floods a dp-attention colocate engine (dp-size 32) and deadlocks its scheduler: 'Eval gsm8k: 0/1319', 0 GPU util, nothing generated. Training rollout never hits this -- it windows submission to over_sampling_batch_size (~16 in flight). The two green colocate+eval CI tests are dense / ~1 GPU-per-engine / no dp-attention, so they never exercise the burst path. Fix: window eval like generate_rollout_async -- keep at most --eval-concurrency (default 64) generations in flight via asyncio.wait(FIRST_COMPLETED) + top-up, instead of creating all tasks at once. 64 fills both engines' 64 dp-workers once and is 16x under the 1024 flood; lower it if eval still hangs. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MuN5ow5hHFxPze36uhgZC4
This reverts commit c378269. The windowing fix removed the 1024-request flood hang, but eval still stalls under the GLM-5.2 dp-attention colocate config (sglang generates a few requests at ~0.6 tok/s then the scheduler freezes, util 0). eval-in-colocate-dp-attention needs a deeper fix (e.g. a non-dp-attention eval dispatch); reverting to keep miles core unchanged and running EVAL=off for now. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MuN5ow5hHFxPze36uhgZC4
DeepEP initializes NVSHMEM, whose built-in NCCL path opens a second NCCL communicator that collides with our own NCCL and hangs during SGLang CUDA-graph replay. Set NVSHMEM_DISABLE_NCCL=1 centrally on the rollout (server_group) and training (actor_group) actor env_vars so it covers all scripts; overridable via the env var.
e37bb14 to
ff8b4fe
Compare
…-outer grouped-expert adapter run_glm5_lora.py: keep gate_proj/up_proj/down_proj in the LoRA targets by default (MoE/dense-MLP LoRA on). Set KEEP_MOE_LORA=0 to restore the attention-only drop (needed only for sglang colocate rollout serving, which cannot yet serve MoE-expert LoRA). lora_utils.py: drive the grouped-expert adapter layout from --experts-shared-outer-loras -- unset => share_expert_adapters=False (regular per-expert MoE LoRA); set => shared-outer (SGLang PR #21466 contract).
…rtual-experts run_glm5_lora.py: when KEEP_MOE_LORA=1 (default) the two MoE-expert-LoRA flags go on TOGETHER -- --experts-shared-outer-loras (train-side shared-outer layout) and --sglang-lora-use-virtual-experts (serve-side virtual-experts path). They are two independent flags that each must be on; emitting only one breaks serving (expert gate_up LoRA-B dim mismatch). KEEP_MOE_LORA=0 -> attention-only LoRA. Dropped the separate EXPERTS_SHARED_OUTER / VIRTUAL_EXPERTS knobs (the half-enabled combos only crash). run_glm5_lora_multinode.sh: document KEEP_MOE_LORA (default 1 = MoE layer LoRA on) and export it to the driver.
…subset of layers actor_train OOM on full GLM-5.2 (744B) colocate RL is dominated by the many-layer MoE-expert grouped-GEMM backward activations. MOE_LORA_LAYERS lets you put the MoE-expert LoRA (MLP gate/up/down -> linear_fc1/linear_fc2) on only a subset of layers while attention LoRA stays on ALL layers. run_glm5_lora.py: when MOE_LORA_LAYERS is set (accepts ranges/commas, e.g. "58-77" or "60,65,70"), drop the bare gate/up/down (which match every layer) and emit Megatron-Bridge ModuleMatcher wildcard patterns "*.layers.<N>.*.linear_fc1/fc2" for only the selected layers. Empty (default) = every layer (unchanged). miles passes "*"-patterns through to MB unchanged; attention bare names still map to all-layer LoRA. run_glm5_lora_multinode.sh: add MOE_LORA_LAYERS knob + export; full GLM-5.2 (78 layers) defaults to the last 20 layers (58-77), toy/others default to all layers. Verified on GLM-5.2_5layer (MOE_LORA_LAYERS=4): megatron target_modules carries the layer-4 expert patterns; wildcard_match matches layer 4 and not layer 3; colocate LoRA RL e2e (cuda graph on) trains/serves/saves (save_model end).
…5.2) Previously the full GLM-5.2 (78L) path defaulted MOE_LORA_LAYERS to 58-77 (last 20 layers); now both the 5-layer toy and the full model default to empty = every MoE layer. Set MOE_LORA_LAYERS explicitly (e.g. 58-77) to restrict and cut actor_train backward-activation / optimizer memory at full scale.
…LORA_LAYERS subset - run_glm5_lora.py: append the CPU-Adam triple (--optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer) by default, gated by OPTIMIZER_CPU_OFFLOAD (=1 default; =0 keeps Adam on GPU). Matches the full-FT run_glm5_744b_a40b.py / deepseek / qwen3 recipes; needs --use-distributed-optimizer (already on). - run_glm5_lora_multinode.sh: add the OPTIMIZER_CPU_OFFLOAD knob (default on) + export it. - Disable the MOE_LORA_LAYERS subset->wildcard rewrite (commented out in both files); MoE-expert LoRA now always covers all MoE layers. Kept a warning shim if the env is set.
run_glm5_lora.py emitted no --sglang-lora-backend, so a bare `python run_glm5_lora.py`, run_glm5_lora_multinode_full_model.sh, and the LoRA CI smoke tests all fell back to sglang's default csgmv (only run_glm5_lora_multinode.sh pinned triton via --extra-args). csgmv is fragile on the GLM-5.2 DSA MoE-LoRA dp-attention rollout (gate_up slice miscount -> "scheduler died"). Add an overridable sglang_lora_backend field defaulting to triton and emit --sglang-lora-backend in both the full and toy sglang_args branches, so every entrypoint inherits triton.
|
lora rl: add periodic held-out eval (interval default 5) |
|
note this lora hyper |
This reverts commit c5051fe.
…launcher) Self-contained multi-node launcher: full GLM-5.2 744B, attention-only LoRA (KEEP_MOE_LORA=0), megatron-bridge backend default (recompute auto-off to avoid the bshd+cross-layer-DSA forward rejection), gsm8k, --lora-base-cpu-backup, EVAL toggle (default off; eval hangs under dp-attention colocate), save-every-step, online wandb. Eval flags passed straight to train.py so it does not depend on the reverted GLM-wrapper eval wiring. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- Rework into a single `all` orchestrator role runnable from the local rx box (auto-detect rank-0 HEAD_IP incl. dash-form node names, form Ray cluster, wait for N nodes, submit) on top of the existing head/worker/launch/stop roles. - Default DSA backend = megatron-bridge (unfused, bshd, LoRA-A natively differentiable); qkv-format + recompute follow the backend automatically. - Add pre-run cleanup: kill ray/sglang/raylet on every rank (avoids leftover raylets counted as extra nodes) and clear /personal/checkpoints (CLEAN_CKPT=0 to keep) so disk does not fill mid-run. - Set rollout-batch-size 4, n-samples-per-prompt 16, global-batch-size 64. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- Add MoE expert projections gate_proj,up_proj,down_proj to --target-modules
(mapped to Megatron linear_fc1/linear_fc2; previously attention-only).
- MoE-expert LoRA requires both flags ON together or the sglang colocate
rollout crashes ("scheduler died": expert gate_up LoRA-B dim vs base
mismatch under EP=32 / dp-attention with --sglang-moe-dense-tp-size 1):
add --experts-shared-outer-loras (train) and --sglang-lora-use-virtual-experts
(serve; not auto-set by arguments.py).
- Set --lora-rank 4 (was 16); scale --lora-alpha 32->8 to keep the 2:1 ratio.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
rank=4 fails actor init with megatron-core ensure_divisibility "4 is not divisible by 8": the LoRA rank dim is tensor-parallel-sharded across --tensor-model-parallel-size 8, and MoE normalize_moe_lora also requires dim % moe_router_topk(=8) == 0. Smallest valid rank is 8. alpha 8->16 (2:1). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…NCCL=0 Full-744B MoE-expert LoRA on all 75 MoE layers OOMs the step-0 backward (recompute is OFF on bshd): rank GPU hits ~98% and dies. Two changes: - NVSHMEM_DISABLE_NCCL=0 in runtime env: NVSHMEM (DeepEP MoE path) was using miles' own NCCL communicator and the failure surfaced as an opaque ncclUnhandledCudaError; with this it reports a clean torch.OutOfMemoryError, confirming the root cause is memory (also avoids a SGLang CUDA-graph hang). - MOE_LORA_LAST_N (default 10): restrict MoE-expert LoRA to the last N MoE layers via Megatron-Bridge ModuleMatcher wildcards (*.layers.<N>.*.linear_fc1/linear_fc2); attention LoRA stays on all layers. --target-modules is now built into TARGET_MODULES and QUOTED so the * glob does not expand. MOE_LORA_LAST_N=0 restores all-layer MoE (OOM-prone). Verified: last-10 MoE LoRA clears step 0 (peak GPU 82% vs 98% OOM), grad_norm 0.26, train_rollout_abs_diff 0.0054; rollout MoE virtual-experts LoRA serves OK. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Validated config (rank 8, MoE LoRA last 10 layers, NVSHMEM_DISABLE_NCCL=0) clears step 0 cleanly, so extend from the 20-step bring-up to a 10000-step training run. save-interval stays 20 (frequent LoRA checkpoints). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…oRA-only sync check After the first LoRA weight sync, query the rollout engine's `lora_checksum` action and verify the rollout actually holds served LoRA (lora_gpu:: buffers present), i.e. it is NOT silently serving the BASE model. This is the LoRA-safe replacement for --check-weight-update-equal on GLM e2e: lora_checksum is LoRA-only, so it skips the base-model hash that gpu_tensor_hash CUDA-IMAs on for some GLM base tensors (which is why --check-weight-update-equal's snapshot/compare/checksum path dies on GLM). - arguments.py: new --check-lora-weight-update-equal (BooleanOptionalAction, default None). Defaults ON whenever LoRA is enabled (lora_rank>0); pass --no-check-lora-weight-update-equal to disable. Non-LoRA runs leave it off. All GLM LoRA scripts (launch_glm_rl.sh, examples/lora/*.sh, run_glm5_lora*.py/.sh) parse via arguments.py, so they inherit the default automatically. - update_weight_from_tensor.py: one-shot verify hook in update_weights() (after the first LoRA sync, before the gloo barrier) + _verify_lora_served(), which on the IPC-src rank queries check_weights(action="lora_checksum") on the colocated rollout engine and counts served-GPU LoRA buffers. All ranks all-reduce the verdict so the raise is collective (a single failing src rank cannot skip the barrier and deadlock). Degrades gracefully: a query error (e.g. an sglang build without the lora_checksum action) logs a warning and skips rather than aborting training. Depends on the sglang lora_checksum action (sgl-project/sglang PR #28703, commit fe43763976). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GZ6ujQSTDHBSgWxR3tYbTP
…LoRA): LoRA-only sync check" This reverts commit f0562bd. The served-buffer check (lora_gpu:: > 0) at the 2nd weight-sync hook is a structural false positive in colocate: between the rollout and that hook the engine is offloaded (release_memory_occupation), trained on the same GPUs, and resumed, which tears down the served LoRA buffer / uid_to_buffer_id. Verified on GLM-5.2_5layer single-node: - --lora-base-cpu-backup ON -> rollout serves LoRA correctly (train_rollout_kl=1.16e-4), yet lora_gpu reads 0. - --lora-base-cpu-backup OFF -> rollout serves BASE (train_rollout_kl=1.03), lora_gpu=0. lora_gpu=0 in BOTH cases, so the check carries no serving-correctness signal at that hook. The original delivery variant also keyed off gpu==0 (always 0 at the sync point, since the GPU buffer is lazily populated), which false-failed healthy runs. train_rollout_kl is the authoritative serving-base signal (~1.0 = BASE, ~1e-4 = LoRA), and the existing _check_weight_sync_results RPC-success check already catches delivery failures (same path the base sync uses), so this flag is removed rather than patched. A correct served check would hook post-rollout / pre-offload (where uid_to_buffer_id is still live); deferred. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GZ6ujQSTDHBSgWxR3tYbTP
- rollout-batch-size 4 -> 8 (128 samples/rollout => 2 train steps/rollout at global-batch 64) - update stale ~/Downloads/miles_dev references to ~/Downloads/miles_lora (devbox_config source path, release hints, flow comments) Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… (per-expert MoE LoRA, TOY=1 smoke) - rename launch_glm_rl.sh -> launch_glm_rl_att_fused_moe.sh (attention LoRA + fused MoE-expert LoRA: --experts-shared-outer-loras + virtual-experts); rename JOB_ID/WANDB_PROJECT to match - new launch_glm_rl_att_unfused_moe.sh: per-expert MoE-expert LoRA (shared-outer OFF, virtual-experts ON — EP-correct), TOY=1 toggles the 5-layer single-node smoke config - fix phantom-worker bug: BSD `seq 1 0` emits "1 0" on macOS, spawning bogus worker ranks on 1-node runs -> C-style for loop - fix TOY not forwarded through _pod -> pods silently ran the full 78-layer config and hung on 1 node Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
Adds end-to-end support for GLM-5.1 / GLM-5.2 (MoE+MLA+DSA) GRPO LoRA training via the Megatron-Bridge (“bridge”) path, including launchers/registries and CUDA CI smoke coverage, plus supporting runtime/LoRA plumbing updates.
Changes:
- Introduce
run_glm5_lora.py+ multi-node wrappers for GLM-5.x LoRA RL runs (including backend selection and rollout/train wiring). - Add GLM-5.2 model registries (rope base 8e6) and new e2e CI tests (GLM-5.1 full e2e; GLM-5.2 train-only via replay).
- Extend miles internals for DSA backend selection, LoRA target mapping, R3 defaults/diagnostics, and Ray submit robustness.
Reviewed changes
Copilot reviewed 22 out of 22 changed files in this pull request and generated 11 comments.
Show a summary per file
| File | Description |
|---|---|
tests/e2e/megatron/test_glm5_2_lora_5layer_ci.py |
GLM-5.2 5-layer train-only CI smoke test using replayed rollout data. |
tests/e2e/megatron/test_glm5_1_lora_6layer_ci.py |
GLM-5.1 6-layer full rollout→train→save CI smoke test. |
scripts/run_glm5_lora.py |
New Typer-based single-node launcher for GLM-5.x bridge-mode LoRA RL training. |
scripts/run_glm5_lora_multinode.sh |
Unified multi-node wrapper that forms an external Ray cluster and submits jobs. |
scripts/run_glm5_lora_multinode_full.sh |
Preset wrapper for validated full-scale GLM-5.2 LoRA RL defaults. |
scripts/run_glm5_lora_multinode_full_model.sh |
Alternate multi-node full-model wrapper (parallelism/seq/recompute knobs). |
scripts/run_glm5_5layer_likefull_2node.sh |
Convenience wrapper to run the 5-layer toy using “full recipe” knobs. |
scripts/models/glm5.2-744B-A40B.sh |
GLM-5.2 full registry (rotary base 8e6) for Megatron model args. |
scripts/models/glm5.2-744B-A40B_5layer.sh |
5-layer GLM-5.2 toy registry override for CI/smoke usage. |
scripts/models/glm5.1-744B-A40B_6layer.sh |
6-layer GLM-5.1 toy registry override and spec notes. |
miles/utils/replay_base.py |
R3 replay mismatch diagnostics + reduced mismatch log spam. |
miles/utils/external_utils/command_utils.py |
Add ray job submit --no-wait/submission-id support to avoid WS-tail fragility. |
miles/utils/arguments.py |
Add --dsa-attention-backend; default-enable --use-rollout-routing-replay; expand all-linear LoRA targets. |
miles/ray/rollout/server_group.py |
Set NVSHMEM_DISABLE_NCCL for rollout actors to avoid NCCL conflicts/hangs. |
miles/ray/actor_group.py |
Set NVSHMEM_DISABLE_NCCL for training actors to avoid NCCL conflicts/hangs. |
miles/backends/megatron_utils/model_provider.py |
Thread miles --dsa-attention-backend into bridge providers when supported. |
miles/backends/megatron_utils/lora_utils.py |
Extend HF↔Megatron module-name mapping for DSA indexer + MoE LoRA layout wiring. |
miles/backends/megatron_utils/bridge_lora_helpers.py |
Ensure DSA backend choice reaches module configs in bridge LoRA model setup. |
launch_glm_rl_att_unfused_moe.sh |
Rx/devbox orchestration script for attention+MoE LoRA runs (unfused/fused selectable). |
launch_glm_rl_att_fused_moe.sh |
Rx/devbox orchestration script for fused MoE-expert LoRA runs. |
GLM52_4node_plan.md |
Operational run plan for a 4-node GLM-5.2 LoRA experiment. |
docs/platforms/nvidia.md |
Document why miles defaults NVSHMEM_DISABLE_NCCL=1 and how to opt out. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # first-dropped (rank topk+1) expert (≈0 => ill-conditioned near-tie), and the mean score of | ||
| # the REPLAYED experts vs the recompute's own top-k (ratio≈1 => benign; <<1 => real). | ||
| if os.environ.get("MILES_R3_DIAG"): | ||
| sf = scores.view(-1, scores.shape[-1]).float() |
| "--use-rollout-routing-replay", | ||
| action="store_true", | ||
| default=False, | ||
| help="The rollout routing replay technique from https://arxiv.org/abs/2510.11370", | ||
| action=argparse.BooleanOptionalAction, | ||
| default=True, | ||
| help="R3 rollout routing replay (https://arxiv.org/abs/2510.11370): replay the rollout's " | ||
| "MoE top-8 in training for rollout<->train on-policy parity. DEFAULT ON -- cheap " |
| DSA kernel backend (``--dsa-attention-backend``; bridge path only). Orthogonal to model version | ||
| and to LoRA -- BOTH backends run GLM-5.1 *and* GLM-5.2, full or LoRA: | ||
| * ``megatron-bridge`` (default): the portable unfused megatron-core DSA kernels (DSAttention / | ||
| CrossLayerDSAttention). No extra deps. Uses the ``bshd`` query layout. | ||
| * ``slime``: the vendored fused TileLang kernels (SparseMLA + lighting_indexer), matching slime's | ||
| rollout kernels for rollout<->train numerical parity (incl. R3 indexer replay). Needs the | ||
| optional ``tilelang`` dep and the ``thd`` (packed) layout. Training/forward-only (no KV cache): | ||
| it cannot serve generation -- the rollout is always served by sglang. | ||
| This launcher selects the matching ``--qkv-format`` automatically from the backend (see | ||
| ``_get_parallel_config``); just pass ``--dsa-attention-backend slime`` or leave the default. See | ||
| the Megatron-Bridge ``models/glm_moe_dsa/__init__.py`` docstring for the full backend matrix. |
| python scripts/run_glm5_lora.py prepare --model-name GLM-5.1-6layer # download model + task dataset (default gsm8k) | ||
| # default (megatron-bridge / unfused) backend: | ||
| python scripts/run_glm5_lora.py full-train --model-name GLM-5.1-6layer --num-gpus-per-node 4 | ||
| # fused slime backend (GLM-5.2 shown; works for GLM-5.1 too): | ||
| python scripts/run_glm5_lora.py full-train --model-name GLM-5.2_5layer \\ | ||
| --dsa-attention-backend slime --num-gpus-per-node 4 |
| # script is functional, not model accuracy. Uses 4 of the suite's GPUs (TP=EP=4, | ||
| # the validated single-node layout; bshd + micro-batch-size 1 for DSA). |
| # NOTE on --spec (inherited from glm5-744B-A40B.sh): in the LoRA-via-bridge path | ||
| # (--megatron-to-hf-mode bridge + --lora-rank>0) the model is built by the Megatron-Bridge | ||
| # provider + miles' "dsa" experimental-attention-variant monkey-patch (bridge_lora_helpers.py), | ||
| # NOT get_glm5_spec — args.spec is never imported/invoked there (model.py dispatch bypasses |
| pkill -9 sglang || true; sleep 3; pkill -9 miles || true; sleep 3; pkill -9 miles || true; pkill -9 redis || true; true | ||
| export RAY_ADDRESS="http://$HEAD_IP:$DASH_PORT" PYTHONUNBUFFERED=1 HF_HOME=/cluster-storage/models | ||
| RUNTIME_ENV_JSON="$(cat <<JSON | ||
| {"env_vars":{"PYTHONPATH":"$MEGATRON_PATH","CUDA_DEVICE_MAX_CONNECTIONS":"1","NCCL_NVLS_ENABLE":"1","NCCL_SOCKET_IFNAME":"$NCCL_IFNAME","GLOO_SOCKET_IFNAME":"$NCCL_IFNAME","no_proxy":"127.0.0.1,127.0.0.1","MASTER_ADDR":"127.0.0.1","MILES_EXPERIMENTAL_ROLLOUT_REFACTOR":"1","INDEXER_ROPE_NEOX_STYLE":"0","SGLANG_NSA_FORCE_MLA":"1","NVSHMEM_DISABLE_NCCL":"0"}} |
| pkill -9 sglang || true; sleep 3; pkill -9 miles || true; sleep 3; pkill -9 miles || true; pkill -9 redis || true; true | ||
| export RAY_ADDRESS="http://$HEAD_IP:$DASH_PORT" PYTHONUNBUFFERED=1 HF_HOME=/cluster-storage/models | ||
| RUNTIME_ENV_JSON="$(cat <<JSON | ||
| {"env_vars":{"PYTHONPATH":"$MEGATRON_PATH","CUDA_DEVICE_MAX_CONNECTIONS":"1","NCCL_NVLS_ENABLE":"1","NCCL_SOCKET_IFNAME":"$NCCL_IFNAME","GLOO_SOCKET_IFNAME":"$NCCL_IFNAME","no_proxy":"127.0.0.1,127.0.0.1","MASTER_ADDR":"127.0.0.1","MILES_EXPERIMENTAL_ROLLOUT_REFACTOR":"1","INDEXER_ROPE_NEOX_STYLE":"0","SGLANG_NSA_FORCE_MLA":"1","NVSHMEM_DISABLE_NCCL":"0"}} |
| ## ⚠️ Real root cause of the acquire failures (NOT image pull) | ||
| `rx devbox why` revealed: pods get **Scheduled** then stick on | ||
| `FailedMount … csi: driver name s3.radixark.io not found in the list of registered CSI drivers` | ||
| — the **`/global-s3` mount's CSI node-plugin is missing/unhealthy on some H200 nodes**. The pod | ||
| never initializes → 17-min provisioning timeout. prep-cpu ran fine on `gpu1-10-220-51-39`, so | ||
| **some nodes are healthy, some are broken** — it's node-level. The earlier "image pull" / | ||
| "RDMA" / "warm-cache via re-acquire" theories were wrong. |
| f"export no_proxy=127.0.0.1 && export PYTHONBUFFERED=16 && " | ||
| f"{cmd_megatron_model_source}" | ||
| f"""ray job submit {'' if 'RAY_ADDRESS' in os.environ else '--address="http://127.0.0.1:8265" '}""" | ||
| f"""ray job submit {'--no-wait ' if _no_wait else ''}{f'--submission-id {_sub_id} ' if _sub_id else ''}{'' if 'RAY_ADDRESS' in os.environ else '--address="http://127.0.0.1:8265" '}""" |
What
Enable GRPO LoRA training for GLM-5.1 and GLM-5.2 (MoE + MLA + DSA) via the Megatron-Bridge path (
--megatron-to-hf-mode bridge), with a launcher, model registries, and e2e CI tests. GLM-5.2 additionally has DSA cross-layer index sharing — that is handled entirely in Megatron-Bridge#13, so nothing GLM-5.2-specific is needed in miles beyond a registry (rope 8e6) and a model-name.LoRA enablement (GLM-5.1 / DSA)
cb16e28map DSA indexer module names across sglang ↔ megatron-bridge — extend_MLA_HF_TO_MEGATRONinlora_utils.pywithwq_b→linear_wq_b,wk→linear_wk,weights_proj→linear_weights_proj(mirrors the MLAq_b_proj↔linear_q_up_projmap), so one target-module name resolves on both rollout (sglang) and training (Megatron-Bridge).610746ainclude MLA + DSA indexer modules in theall-linearexpansion (arguments.py).b0f45cb→00f3dba"dsa" experimental-attention spec registration — originally a caller-side monkey-patch inbridge_lora_helpers.py(route"dsa"→ megatron-core'sget_dsa_module_spec_for_backend+ the omittedmetainfo, same try/finally pattern asdeepseek_v4.get_dsv4_spec), then moved into Megatron-Bridge (00f3dbadrops the patch;_setup_lora_model_via_bridgenow just callsprovide_distributed_model). The fix lives in Megatron-Bridge#13 (feature-detected, self-disabling). FixesValueError: Invalid experimental attention variant: dsa.Launcher + model registries
980e420GLM-5.1 launcher + 6-layer registry —scripts/run_glm5_lora.py(typerScriptArgs(ExecuteTrainConfig), modeled onrun_deepseek_v4.py) andscripts/models/glm5-744B-A40B_6layer.sh.eb4ddc2GLM-5.2 registries + launcher extension —scripts/models/glm5.2-744B-A40B.sh(= GLM-5.1 dims, only--rotary-base 8e6differs; converges with the slime-path PR Support GLM-5.2 744B-A40B #1376) andglm5.2-744B-A40B_7layer.sh(3 dense + 4 MoE).run_glm5_lora.pyis extended to GLM-5.1 and GLM-5.2 (model-name → _MEGATRON_MODEL_TYPEregistry + an_HF_REPOdownload map;--hf-checkpointdefaults to a local{model_dir}/{model_name}path — miles assertsargs.loadis an existing dir). The registry--spec "miles_plugins.models.glm5.glm5" get_glm5_specline is inert under bridge LoRA (the Megatron-Bridge provider overridestransformer_layer_spec, soget_glm5_specis never imported), so GLM-5/5.1/5.2 share one set ofscripts/models/*.shregistries across both the slime full-FT path (Support GLM-5.2 744B-A40B #1376) and this LoRA path.CI
dd16ec8e2e CI tests + drop ad-hoc example —tests/e2e/megatron/test_glm5_lora_6layer_ci.py(GLM-5.1 full rollout→train smoke test) andtest_glm5_2_lora_7layer_ci.py(GLM-5.2 train-only: sglang can't serve the 5.2 cross-layer rollout yet, so it replays a GLM-5.1 rollout dump — both toys share the GLM tokenizer + vocab 154880). Both follow the existingtest_glm5_744b_a40b_4layer_ci.pypattern (register_cuda_ci+prepare/execute). Removes the ad-hocexamples/lora/run-glm5.1-6layer-megatron-lora.sh(superseded by the launcher + CI tests).Verified (4×H200, colocate, bshd, TP4/EP4)
TRAIN EXIT 0+ valid HF PEFT adapter (indexer excluded from the default target list).run_glm5_lora.py train --model-name GLM-5.2-7layer:Job succeeded+ adapter (confirms the launcher + the newglm5.2-744B-A40B_7layer.shregistry + the inert--specend-to-end).Notes / deps
glm_moe_dsabridge build).--qkv-format bshd(+--micro-batch-size 1) is required: megatron-core's DSA core-attention needs a 4D query; the defaultthdpacking yields a 3D query → "not enough values to unpack".🤖 Generated with Claude Code