[gc-anyscale]

What does this PR do?

Adds megatron_config.freeze_language_model, freeze_vision_model and freeze_vision_projection to choose which parts of a vision-language model train on Megatron. They are written to the Megatron-Bridge model provider, whose freeze() maps them to each family's modules (Qwen-VL, GLM-4.5V, Kimi K2.5-VL, Gemma-VL, Nemotron-VL, Ernie-VL). Rank 0 logs trainable / total parameters per module at startup.

megatron_config.freeze_language_model=true megatron_config.freeze_vision_model=true   # projection only
megatron_config.freeze_vision_model=true                                              # LM + projection

RL uses the same fields under trainer.policy.megatron_config.*. All default to false.

Rejected:

• any flag set with a non-Megatron strategy
• any flag with language_model_only=true
• all three set without LoRA (nothing would train)
• LoRA with a partial selection (LoRA + full-rank components is a follow-up PR); plain LoRA or LoRA with all three set is fine
• a flag set on a provider that doesn't support it (raises with the model name)

Behavior change: the flags are always written, so the config wins over provider defaults. Bridge's Qwen3VLMoEModelProvider defaults to freezing the LM and vision encoder, so default Qwen3-VL MoE runs previously trained only the projection; they now train the whole model. Other VLMs are unchanged.

Tests: 21 new CPU tests (provider writes/overrides, projection-only selection on a Qwen3VLModel mock, the log, RL and SFT validation). GPU smoke on Qwen3-VL-4B with projection-only: only visual.merger tensors changed.

Note: on Qwen3-VL, the DeepStack mergers get zero gradient under activation recompute (a Megatron-Bridge issue, independent of this PR). The docs recommend recompute off for projection training; in a 300-step ablation that took the eval from 69.4 to 77.3.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces configuration flags (freeze_language_model, freeze_vision_model, and freeze_vision_projection) to selectively freeze or train components of vision-language models (VLMs) under Megatron-Bridge, including support for training full-rank base parameters (like the vision projection) alongside LoRA adapters. It includes validation logic, checkpoint saving/loading updates, documentation, and unit tests. The review feedback highlights a potential fragility in make_keep_trainable_hooks where a shared closure list is used to track trainable parameters across model chunks, and suggests storing these names directly as attributes on the chunk modules to prevent execution-order bugs.

[gemini-code-assist]

Using a shared closure list (recorded) to track trainable parameters across model chunks is fragile and prone to bugs if chunks are wrapped or processed in an unexpected order (e.g., interleaved or parallelized by the framework). If record is called on multiple chunks before restore is called, recorded.clear() will wipe out previously recorded chunks, leading to a mismatch or silent failure where parameters are not restored to trainable.\n\nSuggestion: Store the recorded trainable parameter names directly as an attribute on each chunk module object (e.g., chunk._recorded_trainable_names). This completely decouples the state from the execution order of the hooks across different chunks, making the implementation robust and stateless.

def make_keep_trainable_hooks(\n    on_restore: Optional[Callable[[List[nn.Module]], None]] = None,\n) -> tuple[Callable, Callable]:\n    '''Pre-wrap hooks that keep the base parameters left trainable by the freeze flags.\n\n    Register record before the LoRA pre-wrap hook and restore after it.\n    record stores the names of the parameters that require grad (per model\n    chunk, after the provider's freeze()); restore sets\n    requires_grad=True on them again once PEFT has frozen the base, then\n    calls on_restore (e.g. to refresh PEFT's params_to_save). Adapter\n    parameters PEFT added are left as they are.\n    '''\n    def record(model):\n        for chunk in _as_list(model):\n            chunk._recorded_trainable_names = {\n                _canonical_name(n) for n, p in chunk.named_parameters() if p.requires_grad\n            }\n        return model\n\n    def restore(model):\n        for chunk in _as_list(model):\n            names = getattr(chunk, '_recorded_trainable_names', None)\n            if names is None:\n                raise RuntimeError('keep-trainable hooks: restore called before record on chunk')\n            for name, param in chunk.named_parameters():\n                if _canonical_name(name) in names:\n                    param.requires_grad = True\n            delattr(chunk, '_recorded_trainable_names')\n        if on_restore is not None:\n            on_restore(model)\n        return model\n\n    return record, restore

[gc-anyscale]

Addressed in cd19047: record now keys the kept parameters by identity across all chunks, so there is no per-chunk list and no ordering assumption, and restore checks every recorded parameter by identity. Tests: test_hooks_do_not_depend_on_chunk_order, test_restore_rejects_a_replaced_parameter.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 8053df9. Configure here.}

[greptile-apps]

Confidence Score: 3/5

[Medium impact] Adds selective freezing of vision-language model components during training.

Fix the LoRA restore matching and provider-default guard before merging.

Findings

1.  Selected base weights stay frozen ▶
2.  Plain LoRA trains base weights ▶
3.  Unsupported freeze flags pass ▶

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Freeze config flags] --> B[Bridge provider freezes model parts]
  B --> C[Record trainable base weights]
  C --> D[LoRA freezes base weights and adds adapters]
  D --> E[Restore selected base weights]
  E --> F[DDP and optimizer]
  F --> G[Save adapters and trainable base weights]

Loading

_{Reviews (1) · Last reviewed commit: "[docs] Note that Qwen3-VL DeepStack merg..." · Reviewed by Greptile}

[greptile-apps]

Selected base weights stay frozen

With freeze_vision_model=true and LM LoRA, the language model should also train at full rank. But Bridge moves targeted base weights under to_wrap: linear_qkv.weight becomes linear_qkv.to_wrap.weight. The restore hook compares names recorded before LoRA, and _canonical_name only strips module., so those weights stay frozen.

Match parameters by identity or account for LoRA wrapper names. Add a test that wraps the original linear rather than only attaching an adapter.

[gc-anyscale]

Confirmed, fixed in cd19047, with a different resolution than re-enabling the wrapped weights: a weight cannot usefully train both at full rank and through an adapter. restore now matches by identity and raises a ValueError that lists the overlapping parameters and suggests scoping target_modules (or freezing that component). This also caught a worse case than the one described: target_modules=all-linear matches Qwen-VL's merger layers, so "LoRA + full projection" silently trained the projection adapter-only. Tests use a Bridge-shaped wrapping mock (to_wrap): test_lora_wrapping_frozen_layers_restores_the_projection, test_all_linear_lora_over_a_trainable_projection_raises, test_lora_over_an_unfrozen_language_model_raises. On GPU (Qwen3-VL-4B): fails at model build as intended: ValueError: LoRA target_modules wrap 16 parameter(s) that the megatron_config.freeze_* flags leave training at full rank (the merger and the three DeepStack mergers, linear_fc1/linear_fc2 weight and bias). Before this change that configuration would have trained the projection through adapters only, with no error.

[greptile-apps]

Plain LoRA trains base weights

keep_trainable reads provider defaults, not whether the user enabled the new flags. Providers with a default freeze flag, such as the Qwen3-VL MoE provider described in this PR, therefore enable the restore hooks even when all three config flags are false. Plain LoRA can then train the remaining vision or projection weights at full rank.

The validator returns early for that configuration, so adapter-only RL sync can miss those updates too. Enable these hooks from the explicit config flags rather than provider defaults.

[gc-anyscale]

Same root cause as the Bugbot thread above, fixed in cd19047: the hooks now follow the megatron_config flags, so they are off exactly when the validator returns early.

[greptile-apps]

Unsupported freeze flags pass

RL calls validate_megatron_vlm_freeze inside validate_megatron_cfg, but only calls that function for strategy="megatron". An FSDP RL run therefore accepts these flags and silently ignores the requested freeze selection. A user requesting projection-only training gets no warning that the selection does not apply.

Call the shared freeze validator from the unconditional config checks and add an RL test using FSDP.

[gc-anyscale]

Fixed in cd19047: validate_policy_vlm_freeze runs from validate_cfg for every strategy, so an FSDP RL run with a freeze flag fails with "requires strategy=megatron". Test: test_vlm_freeze_rejected_on_a_non_megatron_rl_run.

[gc-anyscale]

this is ready for review, @SumanthRH

[xinyuangui2]

Right now, False means two different things depending on what the user did:

If at least one flag is set, the other two are written to the provider as False.
If no flag is set, nothing is written, so all three follow the Bridge provider’s own defaults.

Those defaults aren’t always False. In the pinned Bridge (c4067f5), Qwen3VLMoEModelProvider defaults freeze_language_model=True and freeze_vision_model=True (qwen3_vl_provider.py:272-273). So a default full fine-tune of Qwen3-VL-30B-A3B trains only the vision projection, with no warning, even though our config says all three are False.

Suggestion: make the rule simple. All three default to False, and we always write all three to the provider, whether the user set them or left the default. Then the megatron_config value is always the one that applies, and the provider’s defaults never leak through.

[gc-anyscale]

Agreed, and thanks for catching the MoE default. That was a real silent bug on main. Done in 7b860c3: all three default to false and are always written to any provider that declares them, so megatron_config is the only source of truth. The worker logs what it wrote and any provider default it replaced (for Qwen3-VL MoE: (provider defaults overridden: freeze_language_model=True, freeze_vision_model=True)). A provider without the fields is left alone unless a flag is set, which still raises. Test: test_apply_overrides_provider_defaults_when_no_flag_is_set. I also called it out in the description as a behavior change for default Qwen3-VL MoE runs.

[xinyuangui2]

I suggest moving LoRA + full-rank base weights into a separate PR. That’s these hooks (make_keep_trainable_hooks / register_lora_pre_wrap_hooks), lora_checkpoint_state_dict in megatron_strategy.py, and the LoRA parts of validate_megatron_vlm_freeze and the docs. It’s most of the new logic and serves a single recipe that’s mostly SFT, so it’s easier to review and test on its own. In the follow-up I’d also keep it SFT-only and reject it in RL.

For this PR, let’s reject LoRA together with freeze flags that leave part of the base model trainable (e.g. the projection). Plain LoRA and LoRA with all three flags set should still work.

[gc-anyscale]

Makes sense, done in 7b860c3. I removed the keep-trainable hooks, the LoRA checkpoint change in megatron_strategy.py, and the LoRA parts of the validator and docs. This PR now rejects LoRA with a selection that leaves part of the base model trainable, and accepts plain LoRA and LoRA with all three flags set (tests in test_config.py / test_sft_config.py). The LoRA + full-rank work is preserved on gc-anyscale:megatron-vlm-freeze-lora. I will rebase it onto this PR, make it SFT-only with RL rejected as you suggested, and open it once this lands.