Inference-time model_config overrides with an allowlist of sampling-only keys #87

New Issue

2026-08-28T10:48:38+02:00

lars commented

2026-08-28 10:48:38 +02:00

Problem

Some config keys only affect sampling, not model construction or data preprocessing, so they could legitimately be changed at inference time on an already-trained checkpoint — without retraining. Today there is no way to do that: giant predict / giant rollout rebuild the models from ckpt["model_config"] (giant/checkpoint_io.py:191, build_models), and nothing lets a caller alter that dict.

Concrete motivating case: flipping stage2_model.n_sec.stop_sampling between "greedy" and "sample" for a rollout comparison. It is carried into a plain attribute on Stage2Autoregressive (giant/model/builders.py:141, giant/model/models.py:517) and read only at sample time (giant/sample.py:348). It touches no state_dict key and no shape.

giant/model/summary.py:143 already names this class of key explicitly.

Proposal

Add a config_overrides: dict | None parameter to load_for_inference that deep-merges dotted paths into model_cfg before build_models, validated against an explicit allowlist of inference-safe paths.

The allowlist is the essential part. Without it a typo or a shape-changing key (hidden_dim) surfaces as an opaque load_state_dict size-mismatch traceback; with it, it raises CheckpointCompatibilityError with the same wording style as the module's other failures.

Criteria for a key to be inference-safe:

does not affect module construction or tensor shapes, and
does not affect data preprocessing that the normalizers / vocab maps were fit under.

Initial allowlist candidates:

stage2_model.n_sec.stop_sampling (and n_sec.sampling, once the companion mode="head" sampling issue lands)
stage1_model.ddpm.n_steps / stage2_model.ddpm.n_steps (already surfaced as stage{1,2}_ddpm_steps on InferenceContext)
stage2_model.particle_type.other_policy (already surfaced as other_policy)
stage1_model.router.temperature / stage2_model.router.temperature (EnergyRouter.temperature)

Rejected by design: anything shape-bearing (hidden_dim, n_layers, k_max, emb_dim, n_experts), anything conditioning-mode-bearing (conditioning.*.type, particle_type.target), and all of [train].

Both predict and rollout get the capability for free from the single load_for_inference change. Preferred over the alternative of a post-load setattr in the CLI, which scatters override logic and does not generalize.

Surface

CLI: a repeatable --set key=value (or per-key flags) on giant predict / giant rollout
The giant rollout YAML sidecar should record the overrides actually applied, so the analysis pipeline can tell two rollouts of the same checkpoint apart
Overrides should show up in giant model summary output where relevant

Acceptance

load_for_inference(..., config_overrides=...) merges into model_cfg before build_models
Non-allowlisted or unknown dotted paths raise CheckpointCompatibilityError, never a load_state_dict mismatch
Value validation reuses giant/config.py's existing per-key checks where they exist
CLI surface on both predict and rollout; applied overrides recorded in the rollout YAML
Tests: each allowlisted key round-trips and actually changes behavior; a shape-bearing key is rejected up front

## Problem Some config keys only affect *sampling*, not model construction or data preprocessing, so they could legitimately be changed at inference time on an already-trained checkpoint — without retraining. Today there is no way to do that: `giant predict` / `giant rollout` rebuild the models from `ckpt["model_config"]` (`giant/checkpoint_io.py:191`, `build_models`), and nothing lets a caller alter that dict. Concrete motivating case: flipping `stage2_model.n_sec.stop_sampling` between `"greedy"` and `"sample"` for a rollout comparison. It is carried into a plain attribute on `Stage2Autoregressive` (`giant/model/builders.py:141`, `giant/model/models.py:517`) and read only at sample time (`giant/sample.py:348`). It touches no `state_dict` key and no shape. `giant/model/summary.py:143` already names this class of key explicitly. ## Proposal Add a `config_overrides: dict | None` parameter to `load_for_inference` that deep-merges dotted paths into `model_cfg` **before** `build_models`, validated against an explicit **allowlist** of inference-safe paths. The allowlist is the essential part. Without it a typo or a shape-changing key (`hidden_dim`) surfaces as an opaque `load_state_dict` size-mismatch traceback; with it, it raises `CheckpointCompatibilityError` with the same wording style as the module's other failures. Criteria for a key to be inference-safe: 1. does not affect module construction or tensor shapes, and 2. does not affect data preprocessing that the normalizers / vocab maps were fit under. Initial allowlist candidates: - `stage2_model.n_sec.stop_sampling` (and `n_sec.sampling`, once the companion `mode="head"` sampling issue lands) - `stage1_model.ddpm.n_steps` / `stage2_model.ddpm.n_steps` (already surfaced as `stage{1,2}_ddpm_steps` on `InferenceContext`) - `stage2_model.particle_type.other_policy` (already surfaced as `other_policy`) - `stage1_model.router.temperature` / `stage2_model.router.temperature` (`EnergyRouter.temperature`) Rejected by design: anything shape-bearing (`hidden_dim`, `n_layers`, `k_max`, `emb_dim`, `n_experts`), anything conditioning-mode-bearing (`conditioning.*.type`, `particle_type.target`), and all of `[train]`. Both `predict` and `rollout` get the capability for free from the single `load_for_inference` change. Preferred over the alternative of a post-load `setattr` in the CLI, which scatters override logic and does not generalize. ## Surface - CLI: a repeatable `--set key=value` (or per-key flags) on `giant predict` / `giant rollout` - The `giant rollout` YAML sidecar should record the overrides actually applied, so the analysis pipeline can tell two rollouts of the same checkpoint apart - Overrides should show up in `giant model summary` output where relevant ## Acceptance - [x] `load_for_inference(..., config_overrides=...)` merges into `model_cfg` before `build_models` - [x] Non-allowlisted or unknown dotted paths raise `CheckpointCompatibilityError`, never a `load_state_dict` mismatch - [x] Value validation reuses `giant/config.py`'s existing per-key checks where they exist - [x] CLI surface on both `predict` and `rollout`; applied overrides recorded in the rollout YAML - [x] Tests: each allowlisted key round-trips and actually changes behavior; a shape-bearing key is rejected up front

lars referenced this issue from a commit

2026-08-28 11:21:43 +02:00

Add inference-time model_config overrides with a sampling-key allowlist (gitea #87)

lars closed this issue

2026-08-28 11:45:18 +02:00

Sign in to join this conversation.

1 Participants

Notifications

Due Date

No due date set.

Dependencies

No dependencies set.

Reference: lars/giant#87