[yaoyu-33]

Summary

• Add an offline HF safetensors checkpoint repair script for near-zero input embedding rows.
• Repair damaged input rows from matching lm_head directions scaled to healthy input embedding RMS norm.
• Preserve HF checkpoint layout, write an audit manifest, and document the conversion example.

Nano 4B Check

Checked nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 model.safetensors via HTTP range reads of only:

• backbone.embeddings.weight: BF16 [131072, 3136]
• lm_head.weight: BF16 [131072, 3136]

With min_norm=1.0e-4:

• num rows that would be repaired: 79
• min input row norm before repair: 0.0
• healthy input RMS target norm: 0.7767779671842665
• corresponding lm_head source rows are healthy: min norm 0.2580839991569519, max norm 0.25814253091812134, unusable source rows 0

Repaired row IDs:
1192, 1193, 1245, 1246, 1247, 1248, 1249, 1250, 1251, 1252, 1253, 1254, 1255, 25304, 42819, 51067, 55014, 56286, 67044, 77579, 81422, 81691, 81695, 81772, 81815, 81819, 82005, 82267, 82312, 82500, 82690, 82725, 82737, 82977, 83362, 83855, 84006, 84105, 84115, 84121, 84521, 84794, 85015, 86060, 86366, 86409, 87020, 87411, 88346, 89406, 89412, 90101, 90294, 90320, 90703, 91368, 91515, 91608, 91762, 92529, 96374, 96385, 98573, 101707, 104097, 105465, 106921, 108262, 109896, 112147, 112277, 112327, 112497, 114442, 114755, 116654, 117123, 117744, 118025

Testing

• uvx ruff check examples/conversion/repair_hf_embedding_rows.py
• uvx ruff format --check examples/conversion/repair_hf_embedding_rows.py

Note: uv run pre-commit run --all-files was attempted after initializing 3rdparty/Megatron-LM, but cannot resolve on this macOS host because nvidia-resiliency-ext==0.6.0 only publishes manylinux wheels.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyu-33]

/ok to test 109a44b

[claude]

Light review - LGTM

Self-contained offline utility with clear contracts, good safety guards (max_rows abort, damaged-output-source rejection, post-repair norm verification, output-path nesting check), and reasonable README docs. No correctness or logic issues found, no typos, and no perf/recipe configs are touched.

Minor, non-blocking observations:

Test-coverage gaps. The three unit tests cover the sharded-repair, dry-run, and damaged-output-source paths, but several guarded branches are untested:

• Single-file (model.safetensors, no index) checkpoint path in _load_weight_map.
• max_rows exceeded raising ValueError.
• No damaged rows: returns a clone and still writes a manifest with target_norm=null.
• Non-finite (NaN/inf) input rows via ~torch.isfinite, not just the <= min_norm path (current fixtures only use zero/tiny rows).
• Explicit --input-embedding-name / --output-embedding-name, and the _resolve_tensor_name failure message.
• --overwrite vs FileExistsError when the output dir already exists.

README nit: the new repair_hf_embedding_rows.py section is inserted just above the numbered '1. hf_megatron_roundtrip.py' entry, so an unnumbered tool now precedes the numbered list. Harmless, but slightly odd ordering.

Suggested test cases

• No perf tests impacted.
• Unit: single-file (non-sharded) model.safetensors checkpoint repair path.
• Unit: max_rows exceeded raises ValueError.
• Unit: no damaged rows results in output written unchanged plus manifest with target_norm=null and empty repaired_row_ids.
• Unit: non-finite (NaN/inf) input embedding rows are detected and repaired.
• Unit: explicit --input-embedding-name / --output-embedding-name resolution, plus _resolve_tensor_name failure error.
• Unit: existing output dir raises FileExistsError without --overwrite, and is replaced with --overwrite.

[yaoyu-33]

/ok to test 14a6307