Skip to content

[Lora] Fix bridge-LoRA path silently dropping recompute args - #1593

Merged
yushengsu-thu merged 3 commits into
radixark:mainfrom
yushengsu-thu:miles-lora-mem
Jul 8, 2026
Merged

[Lora] Fix bridge-LoRA path silently dropping recompute args#1593
yushengsu-thu merged 3 commits into
radixark:mainfrom
yushengsu-thu:miles-lora-mem

Conversation

@yushengsu-thu

Copy link
Copy Markdown
Collaborator

What

_setup_lora_model_via_bridge builds the Megatron model from a Megatron-Bridge provider and copies a hand-picked set of megatron args onto it before finalize(). The recompute settings were missing from that list, so --recompute-granularity / --recompute-method / --recompute-num-layers were silently dropped on the bridge-LoRA path and every layer's activations stayed resident during backward.

This PR forwards the three recompute fields to the provider.

Measured effect

gpt-oss-20B MoE LoRA (examples/lora/run-gpt-oss-20B-megatron-moe-lora.sh, 4xH200 TP4), train peak allocator memory per GPU:

peak alloc / GPU
before (recompute args dropped) 90.9 GB
after (recompute active) 44.1 GB (-46.8 GB, -51%)

Runs that do not pass recompute flags are unaffected: None is forwarded, which matches the provider default.

_setup_lora_model_via_bridge copies a hand-picked set of megatron args
onto the Megatron-Bridge provider; the recompute settings were missing
from that list, so --recompute-granularity/--recompute-method/
--recompute-num-layers were silently ignored and full activations
stayed resident during backward. Forward the three fields before
finalize(). Runs without recompute flags are unaffected (None is
forwarded, matching the provider default).

Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
Copilot AI review requested due to automatic review settings July 8, 2026 01:53

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request updates the _setup_lora_model_via_bridge function in miles/backends/megatron_utils/bridge_lora_helpers.py to forward activation recomputation configurations (recompute_granularity, recompute_method, and recompute_num_layers) from args to the provider. The reviewer recommends also forwarding recompute_modules to ensure the activation recomputation configuration is fully applied and not silently dropped.

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.

Comment thread miles/backends/megatron_utils/bridge_lora_helpers.py
@yushengsu-thu yushengsu-thu changed the title Fix bridge-LoRA path silently dropping recompute args [Lora] Fix bridge-LoRA path silently dropping recompute args Jul 8, 2026
Mirror the full recompute family forwarded by the non-LoRA bridge path
(_apply_bridge_runtime_config). Both default to the provider defaults
(None / False), so runs without the flags are unaffected.

Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
@yushengsu-thu

yushengsu-thu commented Jul 8, 2026

Copy link
Copy Markdown
Collaborator Author

I've already verified lora CI locally:

Local validation on 4×H200 (TP4) — all passed (rc=0)

Final PR code (both commits, 68eebf254):

# Test Config Result
1 tests/e2e/megatron/test_gpt_oss_20b_moe_lora_ci.py — combo shared-outer + virtual-experts gpt-oss-20b-bf16, TP4/EP1, triton LoRA backend, --recompute-granularity full --recompute-method uniform --recompute-num-layers 1 ✅ PASS
2 scripts/run_glm5_2_744b_a40b_lora.py — GLM-5.2_5layer, tilelang, shared-outer + virtual-experts TP4/EP4, --ci-test --ci-disable-logprobs-checker --disable-weights-backuper --check-weight-update-skip-list fused_qkv_a_proj_with_mqa indexer. + recompute full ✅ PASS

A/B validation of the first commit (2fa21676e):

# Test Notes Result
3–4 gpt-oss, same combo — baseline (recompute args dropped) vs. fixed Device-level train-phase peak: 79.1 → 72.3 GB/GPU ✅ ×2
5–6 GLM-5.2_5layer, same config — baseline vs. fixed Smoke + no regression ✅ ×2
7–10 Both configs re-run with --record-memory-history snapshots Allocator-level backward peak (gpt-oss): 7.58 → 0.73 GB/GPU (−90%) ✅ ×4

Not covered locally (delegated to PR CI): GLM megatron DSA backend combos, per-expert + no-virtual-experts combos, GLM-5.1 6-layer, and the dense Qwen LoRA CI — included in the full matrix triggered by the run-ci-lora / run-ci-model-scripts / run-ci-megatron labels.

After these 2 mb PRs are in the docker,

@yushengsu-thu
yushengsu-thu merged commit 8dccb85 into radixark:main Jul 8, 2026
12 checks passed
mouad-hpc pushed a commit to Osmosis-AI/miles that referenced this pull request Jul 14, 2026
mouad-hpc pushed a commit to Osmosis-AI/miles that referenced this pull request Jul 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run-ci-megatron run-ci-model-scripts Run model script smoke tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants