Add type hints to helpers.py, hotswap.py, constants.py, integrations.py (consolidates #3448, #3452)

[RudrenduPaul]

What does this PR do?

Consolidates type-hint PRs into a single PR, per @BenjaminBossan's request on #3452:

If you're working on multiple PRs to add type hints, I would strongly prefer you could pool them into a single PR, even if multiple files are being touched.

This PR supersedes and replaces:

• Add type hints to rescale_adapter_scale and compute_loss in helpers.py #3448 — Add type hints to rescale_adapter_scale and compute_loss in helpers.py
• Add type hints to hotswap utility functions in hotswap.py #3452 — Add type hints to hotswap utility functions in hotswap.py

Both #3448 and #3452 will be closed in favor of this PR. No functional behavior is changed in any of the touched files; only type annotations are added to existing signatures (plus one pure variable rename needed to keep a new annotation type-correct, see below).

src/peft/helpers.py

• rescale_adapter_scale(model, multiplier) -> rescale_adapter_scale(model: nn.Module, multiplier: Union[float, int]) -> Iterator[None]
• MontecloraTrainerMixin.compute_loss(self, model, inputs, return_outputs=False, **kwargs) -> typed with nn.Module, dict[str, Any], bool, and a Union[torch.Tensor, tuple[torch.Tensor, Any]] return type

src/peft/utils/hotswap.py

• _update_scaling(lora_module, adapter_name, scaling=None) -> typed with LoraLayer, str, Optional[float], -> None
• hotswap_adapter_from_state_dict(...) -> added missing -> None return annotation
• hotswap_adapter(model, model_name_or_path, adapter_name, torch_device=None, **kwargs) -> typed with torch.nn.Module, str, str, Optional[str], -> None

This also folds in the fix @githubnemo requested on #3448 (make style): the Iterator import now comes from collections.abc instead of typing, per the deprecated-import ruff rule.

src/peft/utils/constants.py

Continuing the same pooled effort (precedent: #3144, merged), the following two functions were missing all parameter and return type hints:

• bloom_model_postprocess_past_key_value(past_key_values) -> bloom_model_postprocess_past_key_value(past_key_values: tuple[torch.Tensor, ...]) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]
• starcoder_model_postprocess_past_key_value(past_key_values) -> starcoder_model_postprocess_past_key_value(past_key_values: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]

Types were derived from the call site in peft_model.py (past_key_values is produced by Tensor.split(...), i.e. a tuple[torch.Tensor, ...]) and from tracing each function body.

Adding the tuple[torch.Tensor, ...] parameter annotation to bloom_model_postprocess_past_key_value surfaced a real pyright error: the function body reassigned the past_key_values parameter to the result of torch.cat(past_key_values) (a plain Tensor), which is incompatible with the new declared parameter type. Fixed by renaming that local variable to concatenated_past_key_values — a pure rename with no behavior change, same category of fix as the loftq_utils.py/integrations.py pyright cleanups in #3144.

src/peft/utils/integrations.py (new in this update)

The remaining functions in this file without full type hints, following the same precedent as #3144 (which added hints to merge_utils.py/other.py and did a small pyright fix in this same file):

• gather_params_ctx(param, modifier_rank=0, fwd_module=None) -> param: Union[nn.Parameter, Iterable[nn.Parameter]], -> Iterator[None] (it's a @contextmanager; param type derived from the ~15 real call sites across tuners/, which pass single nn.Parameters, .parameters() iterators, and lists)
• dequantize_bnb_weight(weight, state=None) -> state: Optional[Any], -> torch.Tensor
• get_layer_device_map(model) -> model: Any, -> Optional[dict[int, Union[int, str]]]
• map_cache_to_layer_device_map(model, cache) -> None -> model: Any, cache: transformers.Cache
• init_empty_weights, _init_on_device, _skip_init_on_device -> -> Iterator[None] (all @contextmanager)
• skip_init_on_device(func) -> func: Callable, -> Callable

model in get_layer_device_map/map_cache_to_layer_device_map is typed Any rather than torch.nn.Module: both functions read model.hf_device_map and model.config.num_hidden_layers, attributes that accelerate/transformers attach dynamically and that aren't part of nn.Module's stub. Typing model as torch.nn.Module there resolves those accesses through nn.Module.__getattr__'s Tensor | Module stub return type and introduces 13 new pyright errors that don't reflect real bugs (same category of false positive #3452 hit with LoraLayer vs torch.nn.Module); Any avoids that noise while still documenting the return type precisely.

Typing dequantize_bnb_weight's return as torch.Tensor also surfaces one pre-existing, real pyright finding in dequantize_module_weight (declared -> torch.nn.Parameter, but one code path returns whatever dequantize_bnb_weight returns, i.e. a plain Tensor) — that function already has complete hints from before this PR, so it's left unchanged and out of scope here, but flagging it for visibility.

Coordination / approval

• @BenjaminBossan requested consolidation into a single PR: Add type hints to hotswap utility functions in hotswap.py #3452 (comment)
• @githubnemo requested the make style fix: Add type hints to rescale_adapter_scale and compute_loss in helpers.py #3448 (review)
• The constants.py and integrations.py additions follow the same precedent and pooling request rather than opening separate PRs.

Testing

• ruff check and ruff format --check, run via the repo's pinned ruff~=0.15.12 (i.e. make quality's lint/format steps), pass on all four changed files.
• npx pyright src/peft/utils/constants.py -> 0 errors, 0 warnings after the variable rename described above.
• npx pyright src/peft/utils/integrations.py -> 27 errors both before and after this change save for the one pre-existing finding in dequantize_module_weight described above (26 -> 27); none are newly introduced by the Any-typed model parameters.
• python3 -m py_compile and an isolated module load (importlib, bypassing the package's accelerate import chain, which isn't installed in this sandbox) both pass; function signatures verified via inspect.signature.
• No runtime logic was changed (type-annotation-only diff, plus the one pure rename in constants.py), so no new unit tests were added; existing tests are unaffected by annotation-only changes.

AI disclosure

AI assistance (Claude Code) was used to prepare this PR, including this update: adding type hints to integrations.py by reading each function body and its call sites, running ruff/pyright on the changed file and diffing the error output against the pre-change baseline to confirm no new false positives, and documenting the one genuine pre-existing pyright finding the new annotations surfaced. All changes were reviewed by the submitter (@RudrenduPaul) before pushing.

[RudrenduPaul]

Hi @BenjaminBossan — pushed one more commit onto this pooled PR, adding type hints to the two functions in src/peft/utils/constants.py (bloom_model_postprocess_past_key_value and starcoder_model_postprocess_past_key_value) that were still missing them, following the same pattern as #3144. Details and pyright verification in the updated PR description. Kept it in this PR rather than opening a new one per your pooling request on #3452.

[BenjaminBossan]

Thanks for pooling the PRs to add more type annotations. I have a few comments, otherwise the PR looks good.

[BenjaminBossan]

Probably most type checkers accept this, but I'd still rather explicitly type fwd_module as being optional.

[BenjaminBossan]

Hmm, I think scaling can't ever be None, the default argument here doesn't really make sense. In all PEFT hotswapping tests, it's never None and I also don't see how that could be achieved. So let's remove the default value and type it as float.

[BenjaminBossan]

Let's type the output as dict[int, Union[int, str]] | None, which IMO makes more sense for outputs.

[BenjaminBossan]

Hmm, yielding f doesn't really make sense here, does it? Let's remove it, then the output type makes more sense.

[RudrenduPaul]

Thanks for the review, @BenjaminBossan! Addressed all 4 comments in commit 610410b:

• gather_params_ctx: fwd_module is now typed Optional[torch.nn.Module] to match its None default
• get_layer_device_map: return type now dict[int, Union[int, str]] | None
• init_empty_weights: removed the meaningless yield f (it was always None since _init_on_device itself yields nothing) — now a bare yield, consistent with Iterator[None]
• _update_scaling: scaling is never None at the one call site, so dropped Optional[float] = None in favor of a required float

[BenjaminBossan]

Thanks for the updates, LGTM.