Add sound type annotations: training lifecycle and optimizers

Parent: NVIDIA-NeMo/Automodel#4607 / GitHub #4607.

Scope: Training lifecycle and optimizers

• training/: Training state, callbacks, schedules, RNG and utility interfaces; preserve checkpoint/resume and training lifecycle semantics.
• optim/: Optimizer factories/configs, parameter-group schemas, and schedulers; preserve state dicts and precision behavior.

Suggested first contribution

Pick one module below, inspect its actual callers, and propose a few missing/imprecise signatures or one boundary contract. Already-correct code needs no cosmetic rewrite. Claim the slice in this issue before starting; keep the group open until its remaining scope is covered.

• training/garbage_collection.py
• optim/optimizer.py

Required annotation-policy branch

Until draft PR #4653 is merged, base annotation work on akoumpa/ci/annotation-policy from akoumpa/Automodel. This branch adds the annotation-policy checker alongside ty; it is not yet part of upstream main. Current policy commit: ef93d5e70.

From a clean checkout, create your contribution branch with:

git fetch https://github.com/akoumpa/Automodel.git akoumpa/ci/annotation-policy
git switch -c my-annotation-slice FETCH_HEAD

For an existing contribution branch, fetch the same policy branch and merge FETCH_HEAD into your branch instead of creating a new one. Link #4653 as a dependency in your PR; after it merges, update your branch onto upstream main.

Run the shared entry point from the repository root in the configured development environment:

uv run --group linting python tools/check_types.py

ty.toml owns type-checker settings and the existing opt-in module list. tools/annotation_policy.toml owns exact annotation-slot exceptions and the legacy inventory. New object annotations fail unless a maintainer reviews a specific exception with a reason for genuinely opaque values. Do not add to or increase the legacy inventory; remove/reduce entries when fixing their annotations. Runtime object() sentinels and Literal/Annotated metadata are unaffected.

The policy scans nemo_automodel and examples, but ty still checks only its configured opt-in modules through this entry point. Also run a targeted uv run --group linting ty check <affected-path> for your slice when it falls outside that scope, and report any exclusions or missing dependencies. A passing policy check does not establish type soundness: it does not infer types or resolve cross-module aliases, and all requirements below still apply.

Annotation requirements (applies to every sub-issue)

Read CODING_GUIDELINES.md and AGENTS.md. This campaign requires real contracts, not annotation-count improvements alone.

• Trace actual callers and implementations before choosing types. Annotate parameters, returns, and relevant stored attributes; include -> None for constructors and methods that return nothing. Preserve supported inputs, defaults, return forms, and subclass compatibility.
• Use concrete types and parameterized containers. Prefer Sequence[T], Mapping[K, V], Iterable[T], and precisely typed callbacks when those are the operations the consumer needs. Use T | None only when None really is supported, and narrow it before use.
• Reuse existing typed configs and interfaces. Use small consumer-owned Protocols for genuine structural contracts, TypedDict for fixed-key records, and dataclasses for owned configuration/state. Use generics and overloads only for relationships demonstrated by real callers. Do not introduce cross-component imports or move types into shared merely to evade import boundaries.
• No Any, including nested Any, aliases to it, **kwargs: Any, or newly propagated implicit Any from untyped containers/callbacks/decorators/dependencies. No bare dict, list, Callable, or Callable[..., ...] used to erase a known signature. Parameterize generic bases as well.
• No escape hatches: do not add typing.cast, # type: ignore, # ty: ignore, # pyright: ignore, @no_type_check, blanket noqa, checker exclusions, disabled diagnostics, or fake stubs/protocols that merely assert an unverified interface. Existing suppressions are not precedent; remove them in the completed slice where feasible and explicitly track anything that remains.
• object is valid for a genuinely opaque value or an unvalidated boundary input, not as a replacement for Any or a bag that hides a known schema. Validate external data at its boundary; do not add reflection, speculative getattr/hasattr, runtime-checkable Protocol probes, or assertions solely to convince the checker. Retain real validation and supported compatibility behavior.
• For untyped third-party APIs, first use available upstream types; otherwise provide a narrow, truthful boundary adapter/stub whose contract is verified against the supported dependency. Do not spread dependency uncertainty through component interfaces. Report a blocked boundary instead of masking it.
• torch.Tensor does not encode shape, dtype, device, or distributed placement. Document relevant tensor and semantic contracts, including axis order, optional fields, aliasing, token/label alignment, and storage versus compute dtype where applicable.
• Preserve runtime behavior, public APIs, optional-dependency import safety, and tensor math. Use deferred annotations / TYPE_CHECKING imports where appropriate without making runtime annotation consumers fail. No new import-time GPU/dependency requirements. Keep behavioral fixes and broad refactors separate.

Validation and completion

1. Claim a small slice (one module or a few related functions) in the sub-issue before implementation. Inspect current main: already-correct code needs no cosmetic rewrite. Larger components require multiple small PRs; one starter PR does not complete the component.
2. Record before/after missing or imprecise signatures and remaining Any/suppression sites for that slice. Use the repository's annotation coverage reporter / Python annotation coverage CI as supporting evidence, not proof of type soundness.
3. Run uv run --group linting python tools/check_types.py from the annotation-policy branch, plus targeted ty checking and Ruff on the affected files, and appropriate existing tests. Record exact commands, dependency versions, results, and any pre-existing diagnostics. Confirm the changed files were actually checked; ty.toml currently excludes checkpoint _backports and relaxes some diagnostics. Do not weaken configuration. If a slice is excluded, supply a targeted check that includes it and report the limitation.
4. Audit actual callers and concrete implementations against the declared contracts; a green checker with untyped dependencies is insufficient. Run import checks for optional-dependency/type-only import changes. If runtime tensor/distributed behavior changes, require the corresponding real forward/backward or multi-process validation rather than mocks alone.
5. A completed slice has fully specified, truthful contracts with no new escape hatches and no new diagnostics. Keep the group issue open until remaining public APIs/config contracts and required internal helpers are covered, or explicitly document an audited no-change result. Keep the parent table current when group work is complete.