Skip to content

fix(checkpoint_inspector): allow empty --param-to-param-group-map-json #4403

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
asolergi-nv merged 2 commits into from
Apr 24, 2026
Merged

fix(checkpoint_inspector): allow empty --param-to-param-group-map-json #4403

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c50a555
fix(checkpoint_inspector): allow empty --param-to-param-group-map-json
DAISY-gh Apr 21, 2026
47eb225
Merge branch 'main' into fix/checkpoint-inspector-empty-param-group-map
asolergi-nv Apr 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 9 additions & 4 deletions tools/checkpoint/checkpoint_inspector.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -815,8 +815,10 @@ def has_layer_index(key: str) -> bool:
@click.option(
"--param-to-param-group-map-json",
type=str,
default="{}",
help="JSON string representing the param to parameter group map."
default=None,
help="Path to a JSON file mapping parameter names to optimizer param group ids. "
"Required only if the source checkpoint has multiple optimizer param groups "
"(e.g. different LR/weight-decay per group). Leave unset for single-group checkpoints."
)
@click.option(
"--rename-mtp-keys",
Expand Down Expand Up @@ -921,8 +923,11 @@ def oom_observer(device, alloc, device_alloc, device_free):

ckpt_path = Path(input_dir)
output_dir = Path(output_dir)
with open(param_to_param_group_map_json, "r") as f:
param_to_param_group_map = json.load(f)
if param_to_param_group_map_json:
with open(param_to_param_group_map_json, "r") as f:
param_to_param_group_map = json.load(f)
else:
param_to_param_group_map = {}
_swiglu_modules = (
[m.strip() for m in swiglu_modules.split(",") if m.strip()]
if swiglu_modules is not None else None
Expand Down
Loading