kempnerforge.training.freeze

Parameter freezing helpers.

freeze_params is the primitive: toggle requires_grad on parameters whose fully-qualified name matches any of the provided fnmatch patterns.

apply_freeze_specs consumes a list of FreezeSpec entries (e.g. from VLMConfig.freeze) and resolves each spec’s module field against a pattern map (typically DEFAULT_MODULE_PATTERNS or an arch-specific override on the config). A raw fnmatch pattern is passed through unchanged when it is not a known alias.

canonical_freeze_meta produces a stable, reorder-invariant serialization of a freeze-spec list for checkpoint metadata, so a checkpoint saved with [A, B] matches one loaded with [B, A] as long as the effective mask is identical.

effective_freeze resolves the active freeze-spec list at a given training step from a base list (always-on) and a list of FreezeStage step-boundary transitions. Used at save (records the post-transition state in metadata), at load (computes the expected metadata for the compare), and at the training-loop hook site (applies stage transitions when step reaches them).

freeze_meta_at_step composes the last two into the single expression both the save and the resume path use.

Functions

apply_freeze_specs(model, specs, pattern_map)

Apply a list of freeze specs to model.

canonical_freeze_meta(specs)

Return a sorted, deduplicated serialization of freeze specs.

effective_freeze(step, base, schedule[, ...])

Compute the active freeze-spec list at step.

freeze_meta_at_step(step, config)

Canonical freeze metadata for config at step.

freeze_params(model, patterns, *[, frozen])

Toggle requires_grad on parameters matching any fnmatch pattern.
kempnerforge.training.freeze.freeze_params(model, patterns, *, frozen=True)[source]

Toggle requires_grad on parameters matching any fnmatch pattern.

Only parameters whose current state differs from the target are flipped, so calling this twice with the same arguments is idempotent. Returns the number of elements (param.numel() summed) that were actually flipped.

Parameters:
Return type:

int

kempnerforge.training.freeze.apply_freeze_specs(model, specs, pattern_map)[source]

Apply a list of freeze specs to model.

For each spec, the module field is looked up in pattern_map; if present, its pattern list is used. Otherwise module is treated as a raw fnmatch pattern. Returns {spec.module: n_params_flipped}.

Parameters:
Return type:

dict[str, int]

kempnerforge.training.freeze.canonical_freeze_meta(specs)[source]

Return a sorted, deduplicated serialization of freeze specs.

The output is safe to JSON-encode and compare across runs: two semantically equivalent freeze-spec lists (same (module, frozen) pairs, any order or duplicates) produce byte-equal JSON.

Parameters:

specs (Iterable[FreezeSpec])

Return type:

list[dict[str, object]]

kempnerforge.training.freeze.effective_freeze(step, base, schedule, valid_modules=None)[source]

Compute the active freeze-spec list at step.

Resolution rule:

  • Start from base (build-time freeze list).

  • For each FreezeStage with start_step <= step, in ascending start_step order, override base entries on conflicting module keys (last-write-wins). Stages with start_step > step are ignored.

Module-key validation:

  • When valid_modules is provided, every spec.module referenced in base or in any applied stage must appear in the set; otherwise ValueError. This catches typos at load time rather than at the next step boundary.

Returns the list of active specs (one per module key).

Parameters:
Return type:

list[FreezeSpec]

kempnerforge.training.freeze.freeze_meta_at_step(step, config)[source]

Canonical freeze metadata for config at step.

Save and resume must derive this identically — otherwise the vlm_freeze_expected compare rejects checkpoints this code wrote — so both go through here. Passing the config’s alias set means a TOML typo raises instead of silently no-op’ing the freeze.

Parameters:
Return type:

list[dict[str, object]]