[Lora] Fix bridge-LoRA path silently dropping recompute args - #1593
Conversation
_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>
There was a problem hiding this comment.
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.
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>
|
I've already verified lora CI locally: Local validation on 4×H200 (TP4) — all passed (rc=0)Final PR code (both commits,
A/B validation of the first commit (
Not covered locally (delegated to PR CI): GLM After these 2 mb PRs are in the docker, |
What
_setup_lora_model_via_bridgebuilds the Megatron model from a Megatron-Bridge provider and copies a hand-picked set of megatron args onto it beforefinalize(). The recompute settings were missing from that list, so--recompute-granularity/--recompute-method/--recompute-num-layerswere 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:Runs that do not pass recompute flags are unaffected:
Noneis forwarded, which matches the provider default.