Add GLM-5/5.1/5.2 (744B MoE) LoRA RL support by yushengsu-thu · Pull Request #1559 · radixark/miles

yushengsu-thu · 2026-07-01T22:17:44Z

Results glm5.2 lora RL training:

miles:
- Add GLM-5/5.1/5.2 (744B MoE) LoRA RL support #1559
SGLang (miles)
- [LoRA] DSA indexer targets + MoE-LoRA cuda-graph and RL adapter-reload fixes sgl-project/sglang#29874
megatron-bridge:
- GLM-5/5.1/5.2 DSA + LoRA enablement (DSA option: megatron-bridge-native and glm-native) Megatron-Bridge#15
  which is split into:

To-do and limitation:

To reduce the review effort, I will submit another lora weight checker PR (since it cannot use currently base model weight checker and require adding APIs on sglang side.)
currently, we only use up_proj but did not include down_proj (moe expert down projection weight) since grad_nore will aggregate to there and cause larger logprobdiff and effect the reward. We are fixing this in issue #
memory optimization - currently the mem take too much
fp8 rollout and improve perf.

Below is CC’s summary of the dev journal and unit test correctness. You can ignore it.

Summary. End-to-end GLM-5 / 5.1 / 5.2 (744B-A40B MoE + MLA + DSA) LoRA RL on the Megatron-Bridge path (--megatron-to-hf-mode bridge): selectable DSA kernel backend, MoE-expert LoRA in two layouts, MLA LoRA targets with HF↔Megatron name mapping, per-model runners + registries + CI combination-matrix tests, and several colocate-rollout correctness fixes.

Fixes

Rollout served the base model: the refactor rollout path never tagged lora_path (c4d9d49cb); lora_B=0 init masked it at step 0.
MLA up-proj adapters silently dropped: stale q_b/kv_b exclusion emptied (_SGLANG_UNSUPPORTED_HF_TARGETS = {}).
--lora-base-cpu-backup defaults ON in the runners — OFF silently gives fake on-policy (KL ~1.0 vs ~1e-4).
DSA-indexer RoPE style (INDEXER_ROPE_NEOX_STYLE=0) and --sglang-max-lora-rank / --sglang-lora-backend triton pinned by the runners (engine-init crashes otherwise).
NVSHMEM_DISABLE_NCCL=1 defaulted (second NCCL communicator hung SGLang CUDA-graph replay).
Weight-update checker works under frozen-base LoRA (base re-ship forced when the checker is on; GLM CI passes the engine-side stacked params via --check-weight-update-skip-list).

Features

--dsa-attention-backend {megatron,tilelang} (default tilelang): tilelang = fused TileLang kernels (thd, matches slime rollout numerics), megatron = portable unfused megatron-core kernels (bshd). Backend is a GLM5ModelProvider dataclass field on the MB side; runners derive --qkv-format automatically.
MoE-expert LoRA layouts: per-expert by default; --experts-shared-outer-loras (same field name end-to-end, incl. the serve-side auto-pair) opts into shared-outer. Serving uses sglang virtual-experts by default (--no-sglang-lora-use-virtual-experts selects the alignment path).
--target-modules all-linear adds the MLA projections only when the HF config is MLA (dense models stay dense); the DSA indexer (wq_b/wk/weights_proj) is opt-in by explicit list.
Runners scripts/run_glm5_{1,2}_744b_a40b_lora.py + scripts/models/*_lora.sh registries; CI: GLM-5.2 / GLM-5.1 tests run the combination matrix {shared-outer+virtual-experts, per-expert+no-virtual-experts} × {tilelang, megatron} (all must pass), plus a new gpt-oss-20b MoE-LoRA CI with both serving combinations.

Validation (8×H200)

CI combination matrix: 10/10 green (GLM-5.2 4/4, GLM-5.1 4/4, gpt-oss 2/2); Qwen2.5 dense LoRA CI (all-linear) green.
Full 744B GLM-5.2, 64 GPU: rollout→train→save to 50+ steps; post-tagging-fix abs_diff ~0.0096 flat, rollout/raw_reward 0.45→0.8-0.99.
down_proj ablation (full 744B, only difference = expert down_proj in targets): with it, abs_diff drifts to 0.058@43 (KL 0.02); without it, flat 0.0044-0.0077 (KL ~2e-4) at no reward cost — the down-proj delta lands on the residual stream and amplifies the bf16 kernel-numerics gap (matches Kimi team's finding). Recommendation: exclude expert down-proj for clean parity; use_tis / rollout-logprobs are the mitigations if its capacity is needed. wandb: https://wandb.ai/glm-zero/miles-glm5_att_unfused_moe (group no-down-proj-260702-0948).

Naming note. --dsa-attention-backend values are now megatron / tilelang (earlier reviews may say megatron-bridge[-native] / slime / glm-native); pairs with Megatron-Bridge PRs #16/#17. Runner-split parity verification: https://gist.github.com/yushengsu-thu/a23d258e7b211a86e97d63631cc09970

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.

… 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.

…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.

…g crash run_glm5_lora.py: - Set INDEXER_ROPE_NEOX_STYLE=0 + SGLANG_NSA_FORCE_MLA=1 in the rollout runtime-env. GLM-5's DSA indexer uses INTERLEAVED (not NeoX) RoPE; without it the sparse-attention top-k index is computed on wrongly-rotated q/k, so the rollout selects the wrong tokens and emits gibberish on long sequences — the root cause of zero-reward full-model rollouts (the 5-layer toy masked it). - Emit the full-model rollout sglang config for GLM-5 (dp-attention, EP/DP, --sglang-attention-backend nsa + flashmla_sparse, dp-lm-head, moe-dense-tp 1, triton moe/lora backend, max-lora-rank); remove the toy reasoning/tool parser. - Drop gate_proj/up_proj/down_proj from the rollout LoRA targets: sglang's mem_pool LoRA-B probe mis-sizes the fully-DP dense gate_up and crashes ("LoRA B output dim != base partition prefix dim"); LoRA only attention until the sglang mem_pool probe is patched. - Apply the full rollout config to all GLM-5 models incl the 5-layer toy. run_glm5_lora_multinode.sh: - Raise ulimit -n to the hard cap before ray start (a 64-GPU job spawns many workers; default soft 1024 fatally crashes the head raylet "Too many workers"). - Add ROLLOUT_GPUS_PER_ENGINE / SGLANG_MEM_FRACTION / SGLANG_LORA_BACKEND knobs and the alltoall MoE dispatcher for the full model. Verified on 8x8 H200: full GLM-5.2 LoRA rollout produces coherent, correct dapo-math solutions (reward=1).

- scripts/run_glm5_lora_multinode_full.sh (new): preset wrapper pinning the validated 8-node 64-GPU GLM-5.2 LoRA-RL e2e recipe (slime fused DSA backend, R3 off, EP32/PP1/CP1, seq 1024 for bring-up) and SELF-HEALING editable installs (pip install -e /personal/{sglang,miles,Megatron-Bridge}) so a worker pod that OOM-restarted does not import the stale image sglang (cannot import name 'ParallelismContext'). - scripts/run_glm5_lora_multinode.sh: add R3 and LORA_BASE_CPU_BACKUP knobs (default off). R3 off drops the sglang IndexerTopkCapturer host pinned buffer (~78-128 GB/rank) that host-OOM'd the colocate pod; LORA_BASE_CPU_BACKUP off drops the sglang base-weight CPU mirror (~372 GB/node). - scripts/run_glm5_lora.py: --lora-base-cpu-backup now opt-in (moved to the .sh knob); note that PYTORCH_CUDA_ALLOC_CONF=expandable_segments is incompatible with torch_memory_saver (colocate offload).

…t spot) Rollout-only raw_reward seq sweep on full 744B GLM-5.2 (slime, R3 off): seq 1024 -> 0.0 (100% truncated), 2048 -> 0.125, 4096 -> 0.25, 8192 -> 0.3125. Diminishing returns past 4096 (+0.0625 reward for 2x rollout time), so 4096 is the reward/throughput sweet spot for the long dapo-math reasoning traces. - run_glm5_lora_multinode.sh: SEQ/RESP_LEN default now task-aware -- dapo-math gets 4096/3584, other tasks keep 8192/7168. - run_glm5_lora_multinode_full.sh: default SEQ 1024->4096, RESP_LEN 512->3584 (the 1024 was a fast-bring-up placeholder that 100%-truncated dapo -> reward 0).

At current dapo reward/truncation (seq 4096: raw_reward ~0.25, ~75% truncated -> many all-zero-std groups), the check_reward_nonzero_std filter rejects most groups and resamples indefinitely, so rollout 0 never fills a batch and training never starts. Default it off in the full preset so each step proceeds to train (GRPO still learns from mixed-reward groups); re-enable once rewards are dense (longer seq / trained policy).

R3 on now adds only --use-rollout-routing-replay (MoE top-8 replay for rollout<->train on-policy parity; cheap sglang routed-experts capturer ~0.5GB/rank). We no longer add --use-rollout-indexer-replay: * It is a DEBUG aid -- only verifies rollout-vs-train DSA indexer top-k parity. The slime kernel recomputes the indexer top-k, so training does not need it. * It made sglang allocate the IndexerTopkCapturer HOST pinned buffer (max_total_num_tokens, num_layers, index_topk=2048) int32 ~= 78-128 GB/rank x8/node, which blew the ~1.78 TB colocate pod cgroup -> RolloutManager host-OOM. So R3 can stay ON (routing parity) WITHOUT the indexer host buffer. - run_glm5_lora.py: r3_args drops the slime indexer-replay branch (routing only) + why. - run_glm5_lora_multinode_full.sh: R3 default off->on (safe now); header updated. - run_glm5_lora_multinode.sh: R3 knob comment updated (routing-only).

The use_r3 field docstring still claimed the slime backend ALSO adds --use-rollout-indexer-replay; that branch was removed (indexer replay is debug-only and the source of the host-OOM buffer). Update the comment to match.

Make rollout ROUTING replay (R3, arxiv 2510.11370) the global default: it's cheap (sglang routed-experts capturer ~0.5 GB/rank, no host buffer) and gives rollout<->train on-policy MoE parity. Use BooleanOptionalAction so it stays togglable (--no-use-rollout-routing-replay). Distinct from --use-rollout-indexer-replay, which stays default-OFF (debug-only; it triggers the ~78-128 GB/rank host buffer that OOMs). - miles/utils/arguments.py: --use-rollout-routing-replay store_true/default=False -> BooleanOptionalAction/default=True (+ help clarifies vs indexer-replay). - scripts/run_glm5_lora.py: r3_args emits --use-rollout-routing-replay (use_r3) or --no-use-rollout-routing-replay (--no-use-r3) so the R3 knob still fully controls the now-default-on flag.

… LoRA) + R3 routing replay (already on) 5-layer A/B PROVED the colocate base re-ship corrupts the sglang-served base: train_rollout_kl 1.04 (--lora-base-cpu-backup OFF, base re-shipped) -> 0.0004 (ON, no re-ship, rollout==train). So --lora-base-cpu-backup is REQUIRED for correct colocate LoRA, not just a perf knob -> default it ON in the full preset. (skip_base_sync=True -> megatron stops re-shipping the base; sglang serves its own correctly-loaded base + ~372GB/node host mirror across pause/resume; host fits now that R3 is routing-only with no indexer buffer.) R3 (routing replay) stays default ON (on-policy MoE parity).

…ut adapter config The colocate adapter-sync config (target_modules_hf_for_sglang_rollout) filtered out the MLA up-projections q_b_proj / kv_b_proj via _SGLANG_UNSUPPORTED_HF_TARGETS, so sglang declared only 8 of the 10 trained modules and silently SKIPPED the shipped q_b/kv_b adapter tensors (mem_pool.py) -- the trained MLA up-proj LoRA never reached the rollout once LoRA_B became nonzero. sglang now supports them: get_hidden_dim returns real dims for q_b_proj/kv_b_proj (sglang lora/utils.py:176-187, sglang-miles-glm-dev / PR #28110) and they are in the known target set (lora/utils.py:353-354). Empty the frozenset so the colocate rollout declares the SAME modules Megatron trains. Verified on the running /personal/sglang.

…el launcher, 4-node plan - scripts/run_glm5_5layer_likefull_2node.sh: run the 5-layer prune with the FULL recipe (dapo-math/seq/recompute/mem-frac 0.70/alltoall), differing ONLY in parallelism (2 nodes; single-node TP=EP/PP1/CP1; rollout-engine 2) — fast validation of the full rollout config before the 64-GPU run. - scripts/run_glm5_lora_multinode_full_model.sh: full 78L GLM-5.2 (744B) LoRA multinode launcher (NODES=4 smoke -> NODES=8 prod). - GLM52_4node_plan.md: 4-node (32xH200) full LoRA + DAPO plan.

gsm8k is short-answer grade-school math (~256-tok answers, dense DSA since <index_topk 2048), so a 256 window is enough for a fast smoke. Make the task-conditional defaults 3-way: dapo-math -> 4096/3584 (long-CoT sweet spot), gsm8k -> 256/256, other -> 8192/7168. - run_glm5_lora_multinode.sh: SEQ/RESP_LEN conditionals add the gsm8k=256 branch. - run_glm5_lora_multinode_full.sh: TASK moved before SEQ; SEQ/RESP_LEN now task-aware (gsm8k -> 256/256, else dapo -> 4096/3584). Override either via env.

…to miles-dev-2026-06-16-merging

Per review: per-expert is now Megatron-Bridge's own default (share_expert_adapters=False), so create_lora_instance shrinks to the original opt-in form and forwards a single, same-named kwarg — experts_shared_outer_loras — when the flag is set (with an assert that the standard LoRA adapter type is in use). miles no longer references share_expert_adapters anywhere. Validated e2e on GLM-5.2_5layer LoRA RL --ci-test in BOTH layouts (weight check clean, step-1 train_rollout_kl ~1e-4 each; layout log prints shared-outer / per-expert respectively). Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>

Per review: miles overrides the sglang pass-through default to True right after add_sglang_arguments. The flag is required whenever expert projections are LoRA targets (engine-init LoRA-B dim mismatch otherwise) and is inert without MoE-expert LoRA (consumed only inside the MoE-LoRA runners; engines built against an sglang without the field drop the kwarg via the dataclasses.fields filter), so defaulting it on is safe. The now-dead validate-time auto-enable and the GLM-5 runners' explicit flag emission are removed. Validated e2e on GLM-5.2_5layer --ci-test with NO explicit flag anywhere: sglang logs 'Virtual expert computation enabled', weight check clean, step-1 train_rollout_kl ~1e-4. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>

…and_utils Per review: restore execute_train's ray job submit to main's exact form. The env-gated --no-wait/--submission-id option served only out-of-repo multinode dev launchers (nothing in the repo sets it); those launchers no longer export it either. Validated e2e on GLM-5.2_5layer LoRA RL --ci-test over the restored blocking submit path (weight check clean, step-1 train_rollout_kl ~1e-4). Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>

Per review: replace the hand-rolled config.json read (try/except with a pseudo-config fallback) with miles' existing load_hf_config helper and plain getattr gating; unreadable checkpoints now fail loudly instead of silently assuming MLA. Verified attr-vs-raw-JSON equivalence for kv_lora_rank/q_lora_rank on GLM-5.2_5layer, GLM-5.1-6layer and Qwen2.5-0.5B, plus GLM-5.2_5layer --ci-test e2e (weight check clean, step-1 train_rollout_kl ~1e-4). Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>

…CI tests; add gpt-oss-20b MoE LoRA CI Per review: the GLM-5.2/5.1 LoRA CI tests now run the MoE-expert LoRA combination matrix sequentially — {shared-outer + virtual-experts} and {per-expert + no-virtual-experts}, each on both DSA kernel backends (tilelang / megatron) — and every combination must pass. A new gpt-oss-20b MoE LoRA CI test (expert-only targets, grouped-GEMM experts, TP=4, bridge mode) covers the same two serving combinations. Adds --no-sglang-lora-use-virtual-experts so the alignment-path (non virtual-experts) serving combination is expressible now that the flag defaults on. All 10 combinations validated e2e on H200s: GLM-5.2 4/4, GLM-5.1 4/4, gpt-oss 2/2 (weight check clean where enabled, healthy step-1 on-policy metrics throughout). Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>

Per review; the resolved list is visible via the emitted train command and the adapter config anyway. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>

…at 10 Per review: the env-gated near-tie/real-divergence diagnostic moves out of the tree (kept as a local patch for debugging); it never altered the threshold semantics — per-model tuning via replay_check_max_mismatch_fraction stays the mechanism. The per-token mismatch dump cap rises from 5 to 10 examples (still bounded: an unbounded loop syncs the GPU ~15x per mismatched token and floods the log by hundreds of thousands of lines on a wholesale divergence, drowning the real error). Validated e2e on GLM-5.2_5layer --ci-test (R3 routing replay on; weight check clean, step-1 train_rollout_kl ~1e-4). Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>

yushengsu-thu · 2026-07-07T03:39:26Z

        mismatch_threshold = threshold * orig_flat.shape[0]
        mismatch_indices = is_mismatch.nonzero(as_tuple=False).squeeze(1)
-        for idx in mismatch_indices:
+        for idx in mismatch_indices[:10]:


I still modify here.
Because if r3 fail, it will stay here too longer to show the mismatch_indices. So, i just let it extract final 10 to quickly move forward and the users can also get the error msg.

Comment-only cleanup of comments introduced by this PR (verified: stripped ASTs identical before/after for every touched .py; .sh non-comment lines byte-identical). Pre-existing comments are untouched. Keeps only the non-obvious constraint notes, compressed to the terse one-to-two-line form the core files use; drops narration that restates the code. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>

…6-merging

yushengsu-thu · 2026-07-07T19:00:45Z

Lora related CI run in local:

Environment: 1 node × 8 H200 (tests use 4 GPUs, TP=EP=4); Megatron-Bridge
bridge-dev-glm-merging @ d5b93af1; sglang @ 7c9dad97cb (#29874 content); each combo is the
full rollout → train → weight-update-check → save loop; each test file executed verbatim
(python3 tests/e2e/megatron/<test>.py), exit code 0.

Test	Combination	Result	step-1 abs_diff / KL
`test_glm5_2_744b_a40b_5layer_lora_ci.py`	tilelang + shared-outer + virtual-experts	✅	0.0101 / 1.11e-4
	megatron + shared-outer + virtual-experts	✅	0.0098 / 1.06e-4
	tilelang + per-expert + no-virtual-experts	✅	0.0100 / 1.09e-4
	megatron + per-expert + no-virtual-experts	✅	~0.010 / ~1.1e-4
`test_glm5_1_744b_a40b_6layer_lora_ci.py`	same 4 combinations	✅ 4/4	~0.0104 / ~1.2e-4
`test_gpt_oss_20b_moe_lora_ci.py`	shared-outer + virtual-experts	✅	—
	per-expert + no-virtual-experts	✅	0.0111 / 4.9e-4

Weight-update checker ON for all GLM combos (--ci-test +
--check-weight-update-skip-list fused_qkv_a_proj_with_mqa indexer.):
0 × check tensor equality failed across all runs.
gpt-oss uses expert-only LoRA targets (gate/up/down, grouped-GEMM experts, TP=4, bridge mode,
lmsys/gpt-oss-20b-bf16).

These should turn green on the hosted CI once the runners pick up the rebuilt v0.5.14 image
(its Dockerfile checks out the sglang-miles branch tip, which now contains the fixes).

kailash109 · 2026-07-13T20:40:12Z

Hello, I've added a PR to sglang that enables GLM 5.2 lora fp8 rollouts (there was a bug in Lora mem pool miscalculating tp dims): sgl-project/sglang#31066

When enabling fp8 rollouts in miles for GLM 5.2 lora, logprob diff hovers around 0.05, config available here (apologies it's through modal so different syntax)

kailash109 · 2026-07-15T20:20:21Z

Hello, updated to do a full run with fp8 rollouts using GLM5.2 lora, config is here

Relevant patches to sglang:

miles/glm5_tilelang_safe_indices.patch — TileLang sparse-MLA backward produced NaN gradients: out-of-bounds loads/atomics through padded indices, plus the TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE miscompile. Added index guards + disabled merge pass.
miles/sglang_fp8_lora_fix.py — SGLang's LoRA-B buffers were incorrectly sized on quantized column-parallel layers (input-axis sharding probe is poisoned by Fp8LinearMethod, which stamps input_size_per_partition == input_size), crashing engine init at LoRA attach. Instead I size buffers from the base module's output_size_per_partition instead.
miles/sglang_tp1_shared_expert_fix.py — with SGLANG_SHARED_EXPERT_TP1=1 (required: 128×128 quant blocks can't shard the 2048-row shared expert across TP32), the replicated shared-expert output was added after the all reduce, but FlashInfer AllReduce Fusion and dp-attention's reduce-scatterv defer/replace that all-reduce — so the shared output was summed once per TP rank (32×), which corrupts every Moe layer. Instead folded it into the pre-reduction add scaled by 1/tp_size

validation run with dapo-math 17k: logprob diff ~ 0.04-0.05

reward goes up to ~0.7-0.8:

yushengsu-thu added 30 commits June 18, 2026 01:00

utils: add --dsa-attention-backend arg (GLM DSA fused/unfused selection)

06b6416

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.

backends: wire --dsa-attention-backend onto the GLM bridge provider

86edaea

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.

Merge remote-tracking branch 'origin/miles-dev-2026-06-16-merging' in…

cf9ca2f

…to miles-dev-2026-06-16-merging

yushengsu-thu commented Jul 6, 2026

View reviewed changes

Comment thread miles/backends/megatron_utils/lora_utils.py Outdated

yushengsu-thu added 7 commits July 6, 2026 16:18

docs: drop the resolved-target-modules info log

2299d88

Per review; the resolved list is visible via the emitted train command and the adapter config anyway. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>

yushengsu-thu commented Jul 7, 2026

View reviewed changes

yushengsu-thu mentioned this pull request Jul 7, 2026

GLM-5/5.1/5.2 DSA + LoRA enablement (DSA option: megatron-bridge-native and glm-native) radixark/Megatron-Bridge#15

Closed

yushengsu-thu and others added 2 commits July 6, 2026 23:50

Merge remote-tracking branch 'radixark/main' into miles-dev-2026-06-1…

4c665ee

…6-merging

Replace weight-checker comment with TODO

6f9877c

guapisolo reviewed Jul 7, 2026

View reviewed changes

Comment thread miles/ray/placement_group.py Outdated

yueming-yuan added 2 commits July 7, 2026 11:50

Remove obsolete comments

4d940be

Move backend-neutral LoRA helpers out of megatron_utils

2c4e33a

yueming-yuan approved these changes Jul 7, 2026

View reviewed changes

Remove comment in generate_endpoint_utils

aeff684

yueming-yuan added 2 commits July 7, 2026 12:02

Repoint LORA_ADAPTER_NAME imports to miles.utils.lora; fix pre-commit

e079735

Fix test imports of LORA_ADAPTER_NAME after move to miles.utils.lora

cfac946

yushengsu-thu merged commit 5bdd942 into radixark:main Jul 7, 2026
64 of 70 checks passed

This was referenced Jul 8, 2026

Move FT modules into topic folders (mechanical) #1596

Merged

Route dsa/dsv4 through the experimental attention variant module spec dispatch radixark/Megatron-LM#64

Closed

kailash109 mentioned this pull request Jul 10, 2026

Adding GLM 5.2 LoRA to modal multinode modal-labs/multinode-training-guide#91

Open

This was referenced Aug 10, 2026

docs: expand LoRA training and serving guide #2349

Merged

[lora] Native LoRA for standard Megatron-core raw mode models #1792

Open

Navigation Menu

Uh oh!

Add GLM-5/5.1/5.2 (744B MoE) LoRA RL support - #1559

Add GLM-5/5.1/5.2 (744B MoE) LoRA RL support #1559
yushengsu-thu merged 92 commits into
radixark:mainfrom
yushengsu-thu:miles-dev-2026-06-16-merging

yushengsu-thu commented Jul 1, 2026 •

edited

Loading

Uh oh!

Uh oh!

yushengsu-thu Jul 7, 2026 •

edited

Loading

Uh oh!

Uh oh!

yushengsu-thu commented Jul 7, 2026

Uh oh!

Uh oh!

kailash109 commented Jul 13, 2026

Uh oh!

kailash109 commented Jul 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Conversation

yushengsu-thu commented Jul 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Results glm5.2 lora RL training:

Below is CC’s summary of the dev journal and unit test correctness. You can ignore it.

Uh oh!

Uh oh!

yushengsu-thu Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yushengsu-thu commented Jul 7, 2026

Uh oh!

Uh oh!

kailash109 commented Jul 13, 2026

Uh oh!

kailash109 commented Jul 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

yushengsu-thu commented Jul 1, 2026 •

edited

Loading

yushengsu-thu Jul 7, 2026 •

edited

Loading