MarinSkyRL cannot freeze selected policy or critic parameters without architecture-specific code. This prevents config-only experiments such as freezing Qwen MoE router weights (*.mlp.gate.weight) during RL. The optimizer path already excludes parameters with requires_grad=False, but no general mechanism resolves parameter names and disables gradients before optimizer construction (optimizer grouping). The existing Grug query-bias setting is specific to replacement updates for one buffer and should remain separate.
Add trainer.policy.freeze_parameters and trainer.critic.freeze_parameters, defaulting to empty lists. Patterns should use glob matching against normalized named_parameters() names, preserving architecture prefixes while removing execution-wrapper prefixes. Apply freezing before optimizer construction so frozen tensors allocate no optimizer state.
Done when:
- Empty freeze lists leave existing training behavior unchanged.
- Each configured pattern must match at least one parameter; an unmatched pattern raises an error naming the role and pattern.
- Every matched parameter has
requires_grad=False before optimizer construction and is absent from optimizer parameter groups.
- Matching is stable across FSDP2, Megatron, and DeepSpeed wrapper names. A backend that cannot honor the setting rejects it instead of ignoring it.
- Startup logs record the resolved parameter names and count once per role.
- Training metrics include
policy/frozen_parameter_count and, when a critic exists, critic/frozen_parameter_count.
- CPU tests cover glob matching, wrapper-prefix normalization, unmatched-pattern failure, optimizer exclusion, unchanged frozen weights after a step, and the empty-list default.
MarinSkyRL cannot freeze selected policy or critic parameters without architecture-specific code. This prevents config-only experiments such as freezing Qwen MoE router weights (
*.mlp.gate.weight) during RL. The optimizer path already excludes parameters withrequires_grad=False, but no general mechanism resolves parameter names and disables gradients before optimizer construction (optimizer grouping). The existing Grug query-bias setting is specific to replacement updates for one buffer and should remain separate.Add
trainer.policy.freeze_parametersandtrainer.critic.freeze_parameters, defaulting to empty lists. Patterns should use glob matching against normalizednamed_parameters()names, preserving architecture prefixes while removing execution-wrapper prefixes. Apply freezing before optimizer construction so frozen tensors allocate no optimizer state.Done when:
requires_grad=Falsebefore optimizer construction and is absent from optimizer parameter groups.policy/frozen_parameter_countand, when a critic exists,critic/frozen_parameter_count.