feat: support generic freeze-config selectors

docs/guides/configuration.mdx

@@ -63,6 +63,34 @@ Only load configuration files and `_target_` values from trusted sources. A YAML
 
 The `distributed:` section is **not** instantiated using `_target_`. Recipes parse it with a fixed schema. Use `strategy: fsdp2`, `strategy: ddp`, or `strategy: megatron_fsdp`. You can also configure parallelism sizes, such as `dp_size`, `tp_size`, and `pp_size`, and strategy-specific options. When pipeline parallelism is enabled (`pp_size > 1`), add a `pipeline:` subsection with options such as `pp_schedule`, `pp_microbatch_size`, and `layers_per_stage`. For examples, see the [Pipeline Parallelism with AutoPipeline](/development/pipeline-parallelism) guide and the recipe configs.
 
+## Select Trainable Modules
+
+The LLM and VLM fine-tuning recipes accept a `freeze_config:` section. Use `freeze_modules` and `unfreeze_modules` with typed selectors to control which modules are trainable:
+
+```yaml
+freeze_config:
+  freeze_modules:
+  - path: vision_tower        # exact canonical module path
+  - path: audio_tower         # repeat the key to select more exact module paths
+  - glob: "*.speech_encoder"  # case-sensitive shell-style glob on the full module path
+  unfreeze_modules:
+  - path: multi_modal_projector
+```
+
+Each selector is a mapping with exactly one key — `path` matches one module by its exact fully qualified name, while `glob` matches module names with `fnmatch`-style wildcards (`*` crosses `.` separators, so `*_proj` matches projection modules at any depth). List entries are additive: repeat `path` or `glob` entries to select several modules, and a single entry combining both keys is rejected. Matching is recursive: a selected module's entire subtree is frozen or unfrozen. Bare strings are rejected; unknown options and selectors that match no parameters raise an error before training starts.
+
+Trainability is resolved in a fixed order. Full fine-tuning (everything trainable) or PEFT (LoRA trainable, base frozen) establishes the baseline. `freeze_modules` selectors then freeze their modules, and `unfreeze_modules` selectors unfreeze theirs, winning on overlap — this makes combinations such as LoRA plus a fully trainable multimodal projector a two-line configuration. Freezing only controls `requires_grad`; it never changes a parameter's storage dtype. The policy is validated on the complete model, re-resolved after tensor/expert parallel and activation-checkpointing surgery immediately before DDP/FSDP construction, and resolved once more after checkpoint loading before optimizer construction.
+
+The modality-specific booleans `freeze_vision_tower`, `freeze_audio_tower`, `freeze_language_model`, and `freeze_video_embedder` remain supported and keep their established attribute and substring matching, applied before `unfreeze_modules`:
+
+```yaml
+freeze_config:
+  freeze_vision_tower: true
+  freeze_audio_tower: true
+```
+
+Legacy-only configurations retain the implicit `freeze_vision_tower: true` default. A configuration that declares `freeze_modules` or `unfreeze_modules` uses explicit-selector semantics and does not implicitly freeze vision modules; add a vision selector (for example, `glob: "*vision*"`) when such a configuration should also freeze the vision tower.
+
 ## Prewarm One-Time CUDA Initialization
 
 {/* docs-review-start: mamba-ssd-prewarm */}

docs/guides/llm/sequence-classification.mdx

@@ -66,6 +66,7 @@ distributed:
   tp_size: 1
   cp_size: 1
   sequence_parallel: false
+  autocast_dtype: bfloat16
 
 peft:
   _target_: nemo_automodel.components._peft.lora.PeftConfig
@@ -76,6 +77,10 @@ peft:
   alpha: 16
   dropout: 0.1
 
+freeze_config:
+  unfreeze_modules:
+  - glob: "*classifier"
+
 dataset:
   _target_: nemo_automodel.components.datasets.llm.seq_cls.GLUE_MRPC
   split: train
@@ -110,6 +115,8 @@ optimizer:
 
 - `target_modules`: glob to select linear layers (e.g., `"*.proj"`).
 - `dim` (rank), `alpha`, `dropout`: tune per model/compute budget. Values `dim=8, alpha=16, dropout=0.1` are a good starting point for RoBERTa.
+- `freeze_config.unfreeze_modules`: keeps the classification head fully trainable while PEFT freezes other non-LoRA parameters.
+- `distributed.autocast_dtype`: runs the forward pass in the selected compute dtype while trainable parameters retain their configured storage dtype, including when single-rank FSDP is skipped.
 - The recipe automatically applies the adapters; no additional code changes are required.
 
 ## Running on Multiple GPUs

examples/llm_seq_cls/glue/mrpc_roberta_lora.yaml

@@ -38,17 +38,21 @@ distributed:
   cp_size: 1
 
   sequence_parallel: false
+  autocast_dtype: bfloat16
 
 peft:
   _target_: nemo_automodel.components._peft.lora.PeftConfig
   target_modules:
   - "*.query"
   - "*.value"
-  # Note: classifier head is fully trained (not LoRA), unfrozen automatically in train_seq_cls.py
   dim: 8
   alpha: 16
   dropout: 0.1
 
+freeze_config:
+  unfreeze_modules:
+  - glob: "*classifier"
+
 dataset:
   _target_: nemo_automodel.components.datasets.llm.seq_cls.GLUE_MRPC
   split: train
@@ -71,4 +75,3 @@ optimizer:
   eps: 1e-8
   lr: 2.0e-5  # Standard learning rate for BERT/RoBERTa fine-tuning
   weight_decay: 0.01  # Crucial for stable training on small datasets
-

nemo_automodel/_transformers/infrastructure.py

@@ -24,6 +24,7 @@
 """
 
 import logging
+from collections.abc import Callable
 from contextlib import nullcontext
 from dataclasses import is_dataclass, replace
 from functools import partial
@@ -67,6 +68,7 @@
 from nemo_automodel.components.quantization.qat import QATConfig
 from nemo_automodel.components.utils.compile_utils import compile_model
 from nemo_automodel.components.utils.model_utils import (
+    FreezeConfig,
     _supports_logits_to_keep,
     apply_parameter_freezing,
     count_model_parameters,
@@ -75,6 +77,7 @@
     freeze_minimax_m3_indexer_params,
     freeze_unused_kv_sharing_params,
     init_empty_weights,
+    parse_freeze_config,
     print_trainable_parameters,
 )
 from nemo_automodel.shared.tied_weights import ensure_tied_lm_head
@@ -189,7 +192,7 @@ def _apply_runtime_compatibility_fixes(model):
 
 
 #  Sharding helpers
-def _shard_pp(autopipeline, model, loss_fn, parallelize_fn):
+def _shard_pp(autopipeline, model, loss_fn, parallelize_fn, reapply_trainability):
     trainable_params, total_params = count_model_parameters(model)
     # Store param info on autopipeline before splitting so it can be accessed later
     # This captures the full model's param counts before PP shards it across ranks
@@ -198,22 +201,25 @@ def _shard_pp(autopipeline, model, loss_fn, parallelize_fn):
     if get_world_size_safe() == 1:
         logger.info("World size is 1, skipping autopipeline.")
     else:
+        if parallelize_fn is not None:
+            parallelize_fn = partial(parallelize_fn, reapply_trainability=reapply_trainability)
         autopipeline.build(model, loss_fn=loss_fn, parallelize_fn=parallelize_fn)
         model = autopipeline
     return model
 
 
-def _shard_ep_fsdp(model, model_wrapper, parallelize_fn, mesh: MeshContext):
+def _shard_ep_fsdp(model, model_wrapper, parallelize_fn, mesh: MeshContext, reapply_trainability):
     """Apply EP + FSDP sharding (non-PP path)."""
     if parallelize_fn is not None and get_world_size_safe() > 1:
         parallelize_fn(
             model,
             world_mesh=mesh.device_mesh,
             moe_mesh=mesh.moe_mesh,
+            reapply_trainability=reapply_trainability,
             **mesh.parallelize_axis_kwargs(),
         )
     elif callable(getattr(model_wrapper, "parallelize", None)):
-        model = model_wrapper.parallelize(model)
+        model = model_wrapper.parallelize(model, reapply_trainability=reapply_trainability)
         model = (
             model[0] if isinstance(model, tuple) else model
         )  # MegatronFSDP will return (model, None) since we don't pass optimizer here
@@ -318,6 +324,7 @@ def parallelize_for_pp(
     model: torch.nn.Module,
     *,
     model_wrapper: Union[FSDP2Manager, MegatronFSDPManager, DDPManager] | None = None,
+    reapply_trainability: Callable[[torch.nn.Module], None] | None = None,
     **kwargs,
 ) -> torch.nn.Module:
     """Parallelize model for pipeline parallelism (non-MoE case).
@@ -328,6 +335,8 @@ def parallelize_for_pp(
     Args:
         model: The model to parallelize.
         model_wrapper: Distributed manager instance.
+        reapply_trainability: Callback that re-resolves the trainability policy
+            after pipeline-stage surgery and immediately before wrapping.
         **kwargs: Additional arguments (world_mesh, moe_mesh, axis names) passed by
             AutoPipeline but unused for non-MoE parallelization.
 
@@ -336,7 +345,7 @@ def parallelize_for_pp(
     """
     if model_wrapper is not None:
         if callable(getattr(model_wrapper, "parallelize", None)):
-            model = model_wrapper.parallelize(model)
+            model = model_wrapper.parallelize(model, reapply_trainability=reapply_trainability)
     return model
 
 
@@ -462,6 +471,43 @@ def _uses_thd_only_te_attention(model) -> bool:
     )
 
 
+def _apply_trainability_policy(
+    model: torch.nn.Module,
+    *,
+    peft_enabled: bool,
+    freeze_config: FreezeConfig | None,
+    strict: bool,
+) -> None:
+    """Resolve the complete trainability policy on the current module hierarchy.
+
+    Parallelization and checkpoint loading can replace modules and parameters.
+    Re-running this policy after each such surgery selects the current objects by
+    module path instead of transferring stale parameter-name state.
+
+    Args:
+        model: Model or pipeline stage whose trainability is being resolved.
+        peft_enabled: Whether the PEFT baseline should freeze non-LoRA parameters.
+        freeze_config: Optional user freeze/unfreeze policy.
+        strict: Whether every generic selector must match this model. Full-model
+            validation is strict; pipeline stages use non-strict rebinding because
+            each rank owns only part of the hierarchy.
+    """
+    if peft_enabled:
+        for name, param in model.named_parameters(remove_duplicate=False):
+            param.requires_grad_("lora_" in name)
+    elif freeze_config is not None and freeze_config.has_generic_selectors():
+        for param in model.parameters():
+            param.requires_grad_(True)
+    if freeze_config is not None:
+        apply_parameter_freezing(model, freeze_config, strict=strict)
+
+    # These are framework invariants, so they are always applied last and cannot
+    # be overridden by a user unfreeze selector.
+    freeze_unused_kv_sharing_params(model)
+    freeze_deepseek_v4_indexer_params(model)
+    freeze_minimax_m3_indexer_params(model)
+
+
 #  apply_model_infrastructure  --  the main post-init orchestration function
 def apply_model_infrastructure(
     model,
@@ -622,16 +668,21 @@ def apply_model_infrastructure(
             _maybe_adapt_state_dict_to_hf(model, model.state_dict(), quantization=False).keys()
         )
 
-    # Apply freezing before sharding
-    freeze_config = _kwargs.get("freeze_config")
-    if freeze_config is not None:
-        apply_parameter_freezing(model, freeze_config)
-
-    # Freeze dead K/V parameters in KV-shared layers (e.g. Gemma4 E2B/E4B)
-    # so the optimizer never tracks them and checkpoint save/resume stay consistent.
-    freeze_unused_kv_sharing_params(model)
-    freeze_deepseek_v4_indexer_params(model)
-    freeze_minimax_m3_indexer_params(model)
+    # Validate selectors on the complete pre-parallelization hierarchy. The
+    # same policy is rebound after model surgery and before DDP/FSDP capture.
+    freeze_config = parse_freeze_config(_kwargs.get("freeze_config"))
+    _apply_trainability_policy(
+        model,
+        peft_enabled=peft_config is not None,
+        freeze_config=freeze_config,
+        strict=True,
+    )
+    reapply_trainability = partial(
+        _apply_trainability_policy,
+        peft_enabled=peft_config is not None,
+        freeze_config=freeze_config,
+        strict=False,
+    )
 
     # NemotronOmni RADIO: opt into the fused SDPA path on ViT attention blocks.
     enable_radio_vit_fused_attn(model)
@@ -644,11 +695,11 @@ def apply_model_infrastructure(
     # Note: AutoPipeline takes care of applying PP + EP + FSDP. _shard_ep_fsdp will take care of applying EP + FSDP if no PP.
     mfsdp_param_attrs = None
     if autopipeline is not None:
-        model = _shard_pp(autopipeline, model, loss_fn, parallelize_fn)
+        model = _shard_pp(autopipeline, model, loss_fn, parallelize_fn, reapply_trainability)
         for part in model.parts:
             setattr(part, "_pre_shard_hf_state_dict_keys", pre_shard_hf_state_dict_keys)
     else:
-        model = _shard_ep_fsdp(model, model_wrapper, parallelize_fn, mesh)
+        model = _shard_ep_fsdp(model, model_wrapper, parallelize_fn, mesh, reapply_trainability)
         # Megatron-FSDP stamps load-bearing per-parameter state (owning-model back-ref,
         # tied-weight ``_is_shared`` marker, ``orig_param`` and friends) during wrapping.
         # The lm-head re-tie and post-wrap checkpoint reload below rebuild Parameter
@@ -724,14 +775,23 @@ def apply_model_infrastructure(
         checkpoint_loaded=bool(checkpoint_already_loaded or weights_already_loaded or should_load_checkpoint),
     )
 
-    # Freeze parameters after checkpoint loading and parallelization
-    # This catches params created during parallelization (e.g., GroupedExpertsTE in init_token_dispatcher)
-    if peft_config is not None:
-        models_to_freeze = model.parts if hasattr(model, "parts") else [model]
-        for mp in models_to_freeze:
-            for name, param in mp.named_parameters():
-                if "lora_" not in name and param.requires_grad:
-                    param.requires_grad_(False)
+    # Checkpoint loading can perform another round of parameter replacement, so
+    # re-resolve once more on the final model parts before optimizer construction.
+    trainability_models: list[torch.nn.Module]
+    if hasattr(model, "parts"):
+        trainability_models = list(model.parts)
+    elif isinstance(model_wrapper, (DDPManager, MegatronFSDPManager)):
+        trainability_models = [getattr(model, "module", model)]
+    else:
+        trainability_models = [model]
+    for mp in trainability_models:
+        reapply_trainability(mp)
+    if peft_config is not None or freeze_config is not None:
+        if not any(param.requires_grad for mp in trainability_models for param in mp.parameters()):
+            logger.warning(
+                "The configured trainability policy left no trainable parameters; "
+                "check freeze_config and the PEFT configuration."
+            )
 
     if autopipeline is None:
         print_trainable_parameters(model)  # Once model's been sharded
@@ -809,8 +869,10 @@ def apply_model_infrastructure(
     # module also keeps its storage-dtype params. Under fp32 master weights + bf16 compute
     # that leaves frozen fp32 tensors feeding bf16 trainable modules -> dtype-mismatch
     # matmul at the seam. Cast frozen params/buffers to the compute dtype so the whole
-    # forward runs uniformly. No-op for pure-fp32 / pure-bf16 runs and when no mp_policy
-    # is available (DDP/PP).
+    # forward runs uniformly. Freeze configuration only controls requires_grad; trainable
+    # parameters keep their storage dtype (compute dtype is owned by autocast or the
+    # distributed mixed-precision policy). No-op for pure-fp32 / pure-bf16 runs and when
+    # no mp_policy is available (DDP/PP).
     compute_dtype = getattr(getattr(model_wrapper, "mp_policy", None), "param_dtype", None)
     if compute_dtype is not None:
         for mp in model.parts if hasattr(model, "parts") else [model]:

nemo_automodel/components/distributed/ddp.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+from collections.abc import Callable
 
 import torch
 import torch.distributed as dist
@@ -90,7 +91,11 @@ def _setup_distributed(self):
         else:
             self.device = torch.device("cpu")
 
-    def parallelize(self, model):
+    def parallelize(
+        self,
+        model: torch.nn.Module,
+        reapply_trainability: Callable[[torch.nn.Module], None] | None = None,
+    ) -> torch.nn.Module:
         """
         Wraps the given model with DistributedDataParallel (DDP).
 
@@ -99,6 +104,8 @@ def parallelize(self, model):
 
         Args:
             model (torch.nn.Module): The PyTorch model to be wrapped.
+            reapply_trainability: Optional callback that re-resolves parameter
+                trainability after model surgery and before DDP construction.
 
         Returns:
             torch.nn.parallel.DistributedDataParallel: The DDP-wrapped model.
@@ -124,6 +131,8 @@ def parallelize(self, model):
                         model.gradient_checkpointing_enable()
                     else:
                         apply_submodule_checkpointing(layers, detect_kv_sharing_and_maybe_disable_cache(model))
+            if reapply_trainability is not None:
+                reapply_trainability(model)
             return model
 
         if self.activation_checkpointing:
@@ -152,4 +161,7 @@ def parallelize(self, model):
         if self.bucket_cap_mb is not None:
             ddp_kwargs["bucket_cap_mb"] = self.bucket_cap_mb
 
-        return DDP(model.to(self.device), **ddp_kwargs)
+        model = model.to(self.device)
+        if reapply_trainability is not None:
+            reapply_trainability(model)
+        return DDP(model, **ddp_kwargs)