[janbernloehr]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Prevents source checkpoint metadata from overwriting weight-shard and GTP-remat arguments derived from the converter's target parallelism.

Background

MegatronCheckpointSaverBase._load_checkpoint_args copies arguments recorded in
the source checkpoint onto the target saver arguments unless their names appear
in args_to_keep.

This is appropriate for model properties such as layer counts, but not for
arguments derived from the requested target parallelism. The converter already
preserves values such as:

• tensor_model_parallel_size;
• pipeline_model_parallel_size;
• expert_model_parallel_size.

core_v0.19.0 added four related arguments that validate_args derives through
resolve_tensor_parallel_weight_shards:

• tensor_parallel_num_weight_shards;
• gtp_weight_remat_size;
• expert_tensor_parallel_num_weight_shards;
• expert_gtp_weight_remat_size.

These arguments were not added to args_to_keep.

For an HF source model represented with TP=1 and a requested conversion target
of TP=4, the saver therefore performs the following sequence:

1. argument parsing establishes the TP=4 target;
2. target validation derives a compatible weight-shard count;
3. _load_checkpoint_args replaces it with the source checkpoint's value of 1;
4. the saver runs validate_args;
5. validation rejects the source shard count against the target TP size.

The forked saver then aborts with:

ValueError: tensor_parallel_num_weight_shards (1) must be >=
tensor_model_parallel_size (4).

Changes

This PR adds the four target-derived arguments to args_to_keep:

tensor_parallel_num_weight_shards
gtp_weight_remat_size
expert_tensor_parallel_num_weight_shards
expert_gtp_weight_remat_size

The target values are consequently retained until validate_args derives their
final values from the requested target parallelism.

Other source checkpoint arguments continue to be inherited as before.

A focused unit test verifies both sides of this behavior:

• the four target-derived arguments retain their target values;
• an ordinary checkpoint argument (num_layers) is still inherited from the
  source metadata.

Reproduction

Reproduced with core_v0.19.0 and confirmed to remain present on main.

Use a source checkpoint whose recorded tensor parallelism is 1 and request a
target tensor parallel size greater than 1:

python tools/checkpoint/convert.py \
  --model-type GPT \
  --loader mixtral_hf \
  --saver mcore \
  --target-tensor-parallel-size 4 \
  --target-pipeline-parallel-size 2 \
  --target-expert-parallel-size 1 \
  --load-dir <local copy of mistralai/Mixtral-8x7B-v0.1 safetensors> \
  --save-dir <output directory> \
  --tokenizer-model <tokenizer.model>

The converter runs the loader in its main process and forks a separate saver
process. Once the independent process-group initialization failure from #6989
is fixed, the saver reaches:

MegatronCheckpointSaverBase.parse_megatron_args
  -> MegatronCheckpointSaverBase._load_checkpoint_args
  -> validate_args
  -> resolve_tensor_parallel_weight_shards

ValueError: tensor_parallel_num_weight_shards (1) must be >=
tensor_model_parallel_size (4).

The significant values are:

source tensor_model_parallel_size:         1
source tensor_parallel_num_weight_shards:  1
target tensor_model_parallel_size:         4
incorrect inherited shard count:           1
expected target-derived shard count:        4

The converter is explicitly re-sharding to
--target-tensor-parallel-size 4. It exposes no separate target GTP-remat flag,
so the source checkpoint's derived value must not override the value derived
from the target configuration.

The preceding loader failure is addressed by the companion PR:

Companion PR: <add link after creation>

Validation

Focused regression coverage

Added:

python -m pytest \
  tests/unit_tests/tools/checkpoint/test_checkpoint_saver_target_args.py

The test constructs source and target argument namespaces with deliberately
different values:

                                      source   target
tensor_parallel_num_weight_shards       1       4
gtp_weight_remat_size                    1       2
expert_tensor_parallel_num_weight_shards 1       8
expert_gtp_weight_remat_size             1       2
num_layers                              24      12

After _load_checkpoint_args:

tensor_parallel_num_weight_shards         4  (target preserved)
gtp_weight_remat_size                      2  (target preserved)
expert_tensor_parallel_num_weight_shards   8  (target preserved)
expert_gtp_weight_remat_size               2  (target preserved)
num_layers                                24  (source inherited)

Local preparation checks completed successfully:

target-argument regression test passed
python -m py_compile <all changed Python files>
git diff --check passed

The complete pytest suite was not run in the local preparation environment
because that Python environment does not contain PyTorch. CI should execute the
new unit test in the normal project test environment.

End-to-end conversion

The end-to-end validation included the companion process-group fix because the
loader otherwise aborts before the saver reaches this code.

With both changes applied, the original TP=4/PP=2/EP=1 conversion completes:

successfully saved checkpoint from iteration 1 to <output directory>
[ t 1/4 ... 4/4, gtp_remat 1/0, p 2/2 ]
Done!

The saver argument output shows the intended separation:

source / loader:
  tensor_model_parallel_size = 1
  tensor_parallel_num_weight_shards = 1

target / saver:
  tensor_model_parallel_size = 4
  tensor_parallel_num_weight_shards = 4

There is no longer an Overwriting default message for any of the four
target-derived arguments.

The generated checkpoint subsequently loads at TP=4/PP=2 and evaluation
completes successfully:

successfully loaded checkpoint ... [ t 1/4, gtp_remat 1/1, p 1/2 ]

bf16:       lm loss value 1.952644
fp8-hybrid: lm loss value 1.979680

No source model revision, target parallelism, or converter parameters were
changed to obtain the successful result.

Scope

This PR intentionally does not:

• change how resolve_tensor_parallel_weight_shards derives its values;
• introduce new converter command-line options;
• change checkpoint metadata;
• change model or training behavior;
• address process-group initialization in the loader or saver.

Issue tracking

Related to #6989.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests — manual end-to-end validation is documented above
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate review.

[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.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[dimapihtar]

/ok to test 2087449

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/35772172319