feat: support generic freeze-config selectors

[pstjohn]

Adds support for custom freeze_modules / unfreeze_modules in freeze_config. This version still supports the legacy freeze_vision_tower arguments

freeze_config:
  freeze_vision_tower: false
  freeze_modules:
  - encoder
  unfreeze_modules:
  - projector

this now also supports the same glob-style pattern of peft's target_modules:

freeze_config:
  freeze_modules: "vision*"

porting all these existing freeze_vision_tower args over to freeze_modules would be a pretty big refactor; for the time i've just emitted a deprecation warning letting people know about the freeze_modules option.

Alternative to #3661 to enable freezing / unfreezing modules in between model construction and optimizer creation.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[github-actions]

🌿 Preview your docs: https://nvidia-preview-preview-dcea089731f8.docs.buildwithfern.com/nemo/automodel

[jgerh]

Completed tech pubs review of docs/guides/llm/sequence-classification.mdx and provided a few copyedits.

[yuhezhang-ai]

Thanks @pstjohn

For the ci failure, I thinkfreeze_embeddings was never really used (it was ignored by the code). I would say we can remove all occurrences specifically under freeze_config from the affected configs and documentation while keeping strict rejection. Please do not remove the separately supported freeze_embeddings options in speculative-training recipe arguments.

If a user genuinely wants to freeze token embeddings, they can use an explicit model-specific selector such as:

freeze_modules:
  - path: model.language_model.embed_tokens

There are a few other items:

1. Preserve parameters intentionally frozen by the model

   The current full-fine-tuning baseline [sets every parameter to trainable](

   Automodel/nemo_automodel/_transformers/infrastructure.py

   Lines 495 to 500 in cff2a57

   	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)

   ) whenever generic selectors are used.

   This can unintentionally unfreeze model-owned constants. For example, Bagel has a [fixed positional-embedding parameter](

   Automodel/nemo_automodel/components/models/bagel/embeddings.py

   Lines 75 to 85 in cff2a57

   	class BagelGridPositionEmbedding(nn.Module):
   	"""Frozen 2D sine/cosine position table for patch and latent grids."""
   	def __init__(self, max_num_patch_per_side: int, hidden_size: int) -> None:
   	super().__init__()
   	self.max_num_patch_per_side = max_num_patch_per_side
   	self.hidden_size = hidden_size
   	self.pos_embed = nn.Parameter(
   	torch.zeros(max_num_patch_per_side**2, hidden_size),
   	requires_grad=False,
   	)

   ) constructed with requires_grad=False.

   Generic freeze/unfreeze selectors should overlay the model’s existing trainability state rather than globally resetting it. Please add a test showing that an unrelated, initially frozen parameter remains frozen.

2. Distinguish empty selector fields from legacy configuration

   These currently become indistinguishable:

   freeze_config:
     freeze_modules: []

   freeze_config: {}

   The first explicitly opts into selector-based configuration but selects nothing; the second retains legacy behavior, including the historical implicit vision freeze.

   Please preserve whether freeze_modules or unfreeze_modules was declared. Empty lists can mean “no modules selected,” or they can be rejected, but they should not silently activate legacy defaults. Values such as null, false, and 0 should also be rejected instead of being treated as empty lists.

3. Add real FSDP2 lifecycle coverage

   The two-rank DDP test is valuable. However, this PR also changes trainability timing around FSDP2, Megatron-FSDP, TP/EP, and pipeline transformations, while the new FSDP2 tests currently mock the wrapping order.

   Please add the smallest practical two-rank FSDP2 test performing forward, backward, and one optimizer step, verifying that the final selected parameters remain trainable and synchronized after parameter replacement and wrapping.

[pstjohn]

Thanks, addressed!

1. Preserve parameters intentionally frozen by the model
   Generic selectors now overlay the model’s existing requires_grad state during full fine-tuning rather than first making every parameter trainable. Added coverage showing that an unrelated model-owned parameter initialized with requires_grad=False remains frozen through parameter replacement and wrapper construction. The PEFT LoRA-only baseline remains unchanged.
   

   Automodel/tests/unit_tests/_transformers/test_infrastructure.py

   Lines 350 to 368 in dcea089

   	def test_freeze_config_rebinds_after_name_changing_replacement_under_full_finetuning():
   	"""Under full fine-tuning, selectors re-resolve on the post-surgery module hierarchy."""
   	freeze_config = {
   	"freeze_modules": [{"path": "base"}, {"glob": "ext*"}],
   	"unfreeze_modules": [{"path": "extension"}],
   	}
   	model = _run_trainability_infrastructure(_TinyTrainabilityModel(), freeze_config)
   	assert not model.base.weight.requires_grad
   	# The renamed extension parameter is re-selected through its parent module path.
   	assert model.extension[0].weight.requires_grad
   	# Parameters the policy does not select retain their model-owned trainability.
   	assert not model.model_constant.requires_grad
   	assert model.parallel_parameter.requires_grad
   	assert not model.trainability_at_wrapper_construction["base.weight"]
   	assert model.trainability_at_wrapper_construction["extension.0.weight"]
   	assert not model.trainability_at_wrapper_construction["model_constant"]
   	assert model.trainability_at_wrapper_construction["parallel_parameter"]

2. Distinguish empty selector fields from legacy configuration
   freeze_modules and unfreeze_modules now preserve whether the field was declared.
   • Omitted fields retain legacy behavior, including the implicit vision freeze.
   • An explicitly declared empty list opts into selector semantics while selecting nothing.
   • null, false, and numeric values are rejected.

3. Add real FSDP2 lifecycle coverage
   Added a real two-rank NCCL/FSDP2 test using FSDP2Manager. It verifies that the selected parameter remains the only trainable parameter after DTensor replacement and wrapping, runs forward/backward and an optimizer step, compares the synchronized gradient and update with a reference, and confirms both ranks have identical final weights.
   

   Automodel/tests/functional_tests/hf_transformer/test_freeze_config_fsdp2.py

   Lines 95 to 195 in dcea089

   	Tensor with the input's global shape, replicated on every rank.
   	"""
   	return tensor.full_tensor() if isinstance(tensor, DTensor) else tensor
   	def _worker(rank: int, port: int) -> None:
   	"""Run one rank of the real FSDP2 trainability lifecycle regression."""
   	torch.cuda.set_device(rank)
   	dist.init_process_group(
   	"nccl",
   	init_method=f"tcp://127.0.0.1:{port}",
   	rank=rank,
   	world_size=_WORLD_SIZE,
   	)
   	try:
   	torch.manual_seed(1234)
   	model = _TinyFreezeModel().cuda(rank)
   	reference = copy.deepcopy(model)
   	reference.backbone.requires_grad_(False)
   	reference.classifier.requires_grad_(True)
   	mesh = init_device_mesh(
   	"cuda",
   	(1, _WORLD_SIZE, 1),
   	mesh_dim_names=("dp_replicate", "dp_shard_cp", "tp"),
   	)
   	config = FSDP2Config(
   	mp_policy=MixedPrecisionPolicy(
   	param_dtype=torch.float32,
   	reduce_dtype=torch.float32,
   	output_dtype=torch.float32,
   	),
   	enable_fsdp2_prefetch=False,
   	)
   	model = apply_model_infrastructure(
   	model=model,
   	is_meta_device=False,
   	device=torch.device("cuda", rank),
   	load_base_model=False,
   	model_wrapper=FSDP2Manager(config, device_mesh=mesh),
   	mesh=MeshContext.from_meshes(mesh),
   	freeze_config={
   	"freeze_modules": [{"path": "backbone"}],
   	"unfreeze_modules": [{"path": "classifier"}],
   	},
   	)
   	assert isinstance(model.backbone.layers[0].weight, DTensor)
   	assert isinstance(model.classifier.weight, DTensor)
   	assert isinstance(model.model_constant, DTensor)
   	assert not model.backbone.layers[0].weight.requires_grad
   	assert model.classifier.weight.requires_grad
   	assert not model.model_constant.requires_grad
   	assert [name for name, parameter in model.named_parameters() if parameter.requires_grad] == [
   	"classifier.weight"
   	]
   	optimizer = torch.optim.SGD((parameter for parameter in model.parameters() if parameter.requires_grad), lr=0.1)
   	reference_optimizer = torch.optim.SGD(
   	(parameter for parameter in reference.parameters() if parameter.requires_grad), lr=0.1
   	)
   	rank_inputs = torch.full((2, _FEATURES), float(rank + 1), device=rank)
   	model(rank_inputs).sum().backward()
   	reference_loss = (
   	sum(
   	reference(torch.full((2, _FEATURES), float(source_rank + 1), device=rank)).sum()
   	for source_rank in range(_WORLD_SIZE)
   	)
   	/ _WORLD_SIZE
   	)
   	reference_loss.backward()
   	assert model.backbone.layers[0].weight.grad is None
   	assert model.classifier.weight.grad is not None
   	torch.testing.assert_close(
   	_full_tensor(model.classifier.weight.grad),
   	reference.classifier.weight.grad,
   	)
   	optimizer.step()
   	reference_optimizer.step()
   	full_classifier_weight = _full_tensor(model.classifier.weight)
   	torch.testing.assert_close(full_classifier_weight, reference.classifier.weight)
   	gathered_weights = [torch.empty_like(full_classifier_weight) for _ in range(_WORLD_SIZE)]
   	dist.all_gather(gathered_weights, full_classifier_weight)
   	for gathered_weight in gathered_weights:
   	torch.testing.assert_close(gathered_weight, full_classifier_weight)
   	finally:
   	dist.destroy_process_group()
   	@pytest.mark.skipif(
   	not torch.cuda.is_available() or torch.cuda.device_count() < _WORLD_SIZE,
   	reason="requires two CUDA GPUs",
   	)
   	def test_freeze_config_survives_real_fsdp2_forward_backward_and_optimizer_step() -> None:
   	"""Selected FSDP2 parameters remain trainable and synchronized through one update."""
   	mp.spawn(_worker, args=(_free_port(),), nprocs=_WORLD_SIZE, join=True)

[yuhezhang-ai]

/ok to test dcea089