fix(checkpoint): restore Gemma4 Unified HF export keys (#3260)

* fix(checkpoint): restore Gemma4 Unified HF export keys

Signed-off-by: khazic <khazzz1c@gmail.com>

* fix(checkpoint): keep the Gemma4 Unified key rename off the base-model load path

The rename lived inside _maybe_adapt_state_dict_to_hf, which runs on three
paths: save, DCP resume, and base-model init load. That broke two of them.

DCP resume had no inverse: save wrote HF FQNs into the shards, load renamed
the destinations to HF FQNs, and _maybe_adapt_state_dict_from_hf left them
there. gemma4_unified is HF-native, so set_model_state_dict runs with
strict=True and the six vision weights surfaced as missing plus unexpected.

Base-model init already translates keys through the storage reader's
key_mapping, built from the model's _checkpoint_conversion_mapping. Renaming
the destinations on top of that double-transformed them and the vision
weights no longer matched the checkpoint.

Move the rename into _maybe_rename_gemma4_unified_keys, apply it explicitly
on save and on both directions of the DCP resume path, and skip it entirely
when is_init_step.

Signed-off-by: khazic <khazzz1c@gmail.com>

* refactor(checkpoint): move Gemma4 Unified key renames to a gemma4 module

Address review feedback: the gemma4_unified HF export key table and the
rename helper now live in
nemo_automodel/components/models/gemma4_unified/hf_key_renames.py, and
checkpointing.py only calls the public helper. Tests move alongside.

Signed-off-by: khazic <khazzz1c@gmail.com>

* refactor(checkpoint): handle Gemma4 Unified keys via a state dict adapter

Address review feedback: instead of gemma-specific renames in the common
checkpointer, gemma4_unified now gets a Gemma4UnifiedStateDictAdapter
attached at model init, so the existing to_hf/from_hf hooks cover save,
base-model init, and DCP resume symmetrically and the storage reader's
key_mapping is disabled to avoid double conversion. checkpointing.py is
back to being model-agnostic.

DCP checkpoints written before a model gained an adapter store the model's
own FQNs, so the load path now detects that layout and loads them with the
model names instead of the adapter's Hugging Face names.

Signed-off-by: khazic <khazzz1c@gmail.com>

* refactor(checkpoint): use str | None in the Gemma4 Unified adapter signature

Signed-off-by: khazic <khazzz1c@gmail.com>

* refactor(checkpoint): derive HF export key renames generically

Transformers can load a model under FQNs that differ from the ones its own
checkpoint publishes, so the consolidated export was writing the in-memory
names instead of the published ones.

Handle this at consolidation time rather than with a per-model state dict
adapter. The renames are derived from the load-time conversion mapping by
pushing each published FQN of the reference checkpoint through it, so they are
exact strings rather than reversed regexes, and ambiguous collapses are
rejected. DCP shards keep the model's own FQNs, which leaves base-model
initialization and resume untouched and removes the need for a compatibility
branch.

The renames are applied only to the consolidated safetensors headers and the
weight index, across the inline, async, single-rank, and offline paths. Models
whose experts Transformers merges, and NemotronH remote code, are excluded so
the export keeps matching the tensors actually written.

Signed-off-by: khazic <khazzz1c@gmail.com>

* test(tools): update the offline consolidation stub for export_key_renames

The happy-path test replaces consolidate_safetensors_files_on_every_rank
with a hand-written stub whose signature predates the export_key_renames
argument, so main() raised TypeError once the tool started forwarding the
sidecar. Extend the stub and assert the argument, and add a case covering
a checkpoint that ships an export_key_renames.json sidecar.

Signed-off-by: khazic <khazzz1c@gmail.com>

* refactor(checkpoint): update consolidated tensor metadata in place

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* refactor(checkpoint): preserve export name across metadata rebuild

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* refactor(checkpoint): use Gemma4 state dict adapter

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* refactor(gemma4): localize adapter attachment

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* refactor(gemma4): use custom unified model wrapper

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* chore(gemma4): remove unused model type constant

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

---------

Signed-off-by: khazic <khazzz1c@gmail.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
Co-authored-by: HuiyingLi <willwin.lee@gmail.com>