[REGRESSION] core_v0.19.0: HF->mcore checkpoint converter broken twice — GTP_remat process groups uninitialized, and saver inherits target-derived weight-shard args

**Describe the regression**

Two independent regressions in `tools/checkpoint` make the HF → mcore checkpoint
converter unusable for Mixtral-8x7B (and, for the second one, for any conversion
whose `--target-tensor-parallel-size > 1`) at `core_v0.19.0`
(`5be9626709af2722333bf54797c954c09edeada3`). Both work at `core_v0.18.2`
(`571370c829ca768fe37244f4e2e7f28d8accc4ab`). Both are still present on `main`
(`f2f0f7bfd88fcb1243df55275988d6af52daea35`).

**1. GTP_remat process groups are not initialized in the converter path**

`core_v0.19.0` adds the GTP_remat axis and the corresponding
`_DATA_PARALLEL_GROUP_WITH_CP_WITH_GTP_REMAT` /
`_DATA_PARALLEL_GLOBAL_RANKS_WITH_CP_WITH_GTP_REMAT` (and `_INTRA_*`) globals in
`megatron/core/parallel_state.py`. `LanguageModule.__init__`
(`megatron/core/models/common/language_module/language_module.py:51`) calls
`ProcessGroupCollection.use_mpu_process_groups()` with no arguments, so
`process_groups_config.py:319` resolves *every* entry of `pg_to_func`, four of
which reach a getter that hard-asserts on an uninitialized group (`dp`, `dp_cp`,
`dp_cp_gtp_remat`, `intra_dp_cp`).

The converter never calls `initialize_model_parallel`; it fakes single-rank
groups by assigning `megatron.core.parallel_state` globals directly (see
`_ConverterFakeProcessGroup` in `tools/checkpoint/utils.py` and its use in
`loader_*.py` / `saver_base.py`). Those assignments were not extended with the
new GTP_remat globals, so building the model in the loader now asserts.

**2. `saver_base._load_checkpoint_args` inherits target-derived weight-shard args
from the source checkpoint**

`args_to_keep` in `tools/checkpoint/saver_base.py` lists the args that must keep
their *target* values instead of being inherited from `md.checkpoint_args`. It
lists `tensor_model_parallel_size`, `pipeline_model_parallel_size`,
`expert_model_parallel_size`, … but `core_v0.19.0` added four args that
`validate_args` derives from the target parallelism via
`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` — and did not add them to that list.
`tools/checkpoint/saver_base.py` is byte-identical between `core_v0.18.2` and
`core_v0.19.0`, so the list simply did not keep up with the new args.

Result: an HF source (TP=1, shards=1) overwrites the target-derived default and
`validate_args` rejects its own arguments inside the forked saver process.

**To Reproduce**

```bash
# core_v0.19.0 (or current main), single process, no torchrun
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 \
--save-dir \
--tokenizer-model
```

Regression 1 fires first, in the loader process. With the GTP_remat globals
faked to single-rank groups, regression 2 fires next, in the forked saver
process. Regression 2 needs only `--target-tensor-parallel-size > 1`.

The conversion itself runs on CPU in a single process (`convert.py` forks the
saver), so no GPU, no `torchrun`, and no distributed init are needed to
reproduce either regression.

**Expected behavior**

The conversion completes and writes an mcore checkpoint at the requested target
parallelism, as it does at `core_v0.18.2`: the loader builds the source model,
the forked saver validates its own target-derived args, and the run ends with

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

The resulting checkpoint then loads in `pretrain_gpt.py` with
`--tensor-model-parallel-size 4 --pipeline-model-parallel-size 2`.

**Actual behavior**

At `core_v0.19.0` and on `main` nothing is written — `` is left empty
and `convert.py` exits non-zero — because the run aborts in one of two places:

1. The loader process aborts while building the source model, before any weight
is read: `AssertionError: data parallel group with CP (with GTP_remat) is not
initialized`, raised from `parallel_state.get_data_parallel_group` via
`ProcessGroupCollection.use_mpu_process_groups()`.
2. With the GTP_remat globals faked to single-rank groups so the loader gets
past that, the forked saver process aborts in its own `validate_args`:
`ValueError: tensor_parallel_num_weight_shards (1) must be >=
tensor_model_parallel_size (4).` — the source checkpoint's TP=1-derived shard
count has been inherited over the target-derived default.

Full tracebacks for both are below.

**Previous performance**

At `core_v0.18.2` the same command converts the checkpoint and prints
`Done!`; the converted checkpoint then loads in `pretrain_gpt.py` with
`--tensor-model-parallel-size 4 --pipeline-model-parallel-size 2`.

**New performance**

At `core_v0.19.0` and on `main` the conversion aborts before writing anything.

**Stack trace/logs**

Regression 1 (loader process):

```
building GPT model ...
Traceback (most recent call last):
File "/workspace/megatron-lm/tools/checkpoint/convert.py", line 168, in
main()
File "/workspace/megatron-lm/tools/checkpoint/convert.py", line 160, in main
loader.load_checkpoint(queue, args)
File "/workspace/megatron-lm/tools/checkpoint/loader_mixtral_hf.py", line 349, in load_checkpoint
_load_checkpoint(queue, args)
File "/workspace/megatron-lm/tools/checkpoint/loader_mixtral_hf.py", line 289, in _load_checkpoint
model = load_checkpoint_to_model(margs)
File "/workspace/megatron-lm/tools/checkpoint/loader_mixtral_hf.py", line 142, in load_checkpoint_to_model
model = model_provider(gpt_builder, pre_process=True, post_process=True).to(args.params_dtype)
File "/workspace/megatron-lm/model_provider.py", line 59, in model_provider
return model_builder(args, pre_process, post_process, vp_stage, config=config, pg_collection=pg_collection)
File "/workspace/megatron-lm/gpt_builders.py", line 86, in gpt_builder
model = GPTModel(
File "/workspace/megatron-lm/megatron/core/models/gpt/gpt_model.py", line 125, in __init__
super().__init__(config=config, pg_collection=pg_collection)
File "/workspace/megatron-lm/megatron/core/models/common/language_module/language_module.py", line 51, in __init__
pg_collection = ProcessGroupCollection.use_mpu_process_groups()
File "/workspace/megatron-lm/megatron/core/process_groups_config.py", line 319, in use_mpu_process_groups
init_dict = {pg: pg_to_func[pg]() for pg in required_pgs}
File "/workspace/megatron-lm/megatron/core/parallel_state.py", line 1806, in get_data_parallel_group
assert group is not None, f"{description} is not initialized"
AssertionError: data parallel group with CP (with GTP_remat) is not initialized
```

Regression 2 (forked saver process, after the GTP_remat groups are faked):

```
Process Process-1:
Traceback (most recent call last):
File "/usr/lib/python3.12/multiprocessing/process.py", line 314, in _bootstrap
self.run()
File "/usr/lib/python3.12/multiprocessing/process.py", line 108, in run
self._target(*self._args, **self._kwargs)
File "/workspace/megatron-lm/tools/checkpoint/saver_core.py", line 71, in save_checkpoint
raise e
File "/workspace/megatron-lm/tools/checkpoint/saver_core.py", line 69, in save_checkpoint
saver.save()
File "/workspace/megatron-lm/tools/checkpoint/saver_base.py", line 338, in save
self.parse_megatron_args()
File "/workspace/megatron-lm/tools/checkpoint/saver_base.py", line 119, in parse_megatron_args
validate_args(margs)
File "/workspace/megatron-lm/megatron/training/arguments.py", line 416, in validate_args
resolve_tensor_parallel_weight_shards(
File "/workspace/megatron-lm/megatron/core/model_parallel_config.py", line 32, in resolve_tensor_parallel_weight_shards
raise ValueError(
ValueError: tensor_parallel_num_weight_shards (1) must be >= tensor_model_parallel_size (4).
```

**Environment**

- Megatron-LM: bad `core_v0.19.0` (`5be9626709af2722333bf54797c954c09edeada3`,
`megatron.core.__version__` reports `0.19.1`); good `core_v0.18.2`
(`571370c829ca768fe37244f4e2e7f28d8accc4ab`); also reproduced by inspection on
`main` (`f2f0f7bfd88fcb1243df55275988d6af52daea35`)
- PyTorch 2.14.0a0 (dev build), Python 3.12, Transformer Engine 2.19.0,
transformers 4.49.0
- Conversion runs on CPU in a single process (`convert.py` forks the saver); no
distributed init, so this is not hardware-specific
- Model: `mistralai/Mixtral-8x7B-v0.1` safetensors → mcore, TP=4, PP=2, EP=1

**Proposed fix**

1. Process groups — pick one:
- extend the converter's fake-group assignment in `tools/checkpoint` to cover
the new GTP_remat globals (what we do downstream today), or
- have the converter path request only the process groups it needs instead of
the full `pg_to_func` set, or
- make the GTP_remat entries in `pg_to_func` degrade like the tolerant
entries rather than hard-asserting.
2. Checkpoint args — add the four target-derived args to `args_to_keep` in
`tools/checkpoint/saver_base.py`:
`tensor_parallel_num_weight_shards`, `gtp_weight_remat_size`,
`expert_tensor_parallel_num_weight_shards`,
`expert_gtp_weight_remat_size`; or re-derive them from the target parallelism
after `_load_checkpoint_args` has copied the source args. The converter
re-shards to `--target-tensor-parallel-size` and exposes no GTP target flag,
so keeping them target-derived is the rule already applied to
`tensor_model_parallel_size` itself.

Both changes are local to `tools/checkpoint` (plus, for option 1c,
`process_groups_config.py`) and do not affect training paths.

**Additional context**

Both fixes have been validated together: with the two changes applied to
`core_v0.19.0`, `convert.py` completes (`successfully saved checkpoint … [ t 4/4,
gtp_remat 1/0, p 2/2 ] … Done!`), `pretrain_gpt.py` loads the result at TP=4/PP=2
(`successfully loaded checkpoint … [ t 1/4, gtp_remat 1/1, p 1/2 ] at iteration
1`), and evaluation reports `lm loss value: 1.952644E+00` (bf16) /
`1.979680E+00` (fp8-hybrid) on 2×4 GB200 NVL GPUs. No workload, target
parallelism, or Megatron revision was changed to obtain that result.

A caveat for whoever fixes regression 1: `convert.py` runs the loader in-process
and forks the saver, and each side installs its own `parallel_state` globals, so
the loader and the saver each need the full set — fixing only one side moves the
crash rather than removing it.

---

_This issue was drafted with assistance from the `opus` AI model._

Start by running the documented tools/checkpoint/convert.py reproduction, then trace the loader through tools/checkpoint/utils.py and loader_*.py, and the saver through tools/checkpoint/saver_base.py. Review the process-group mappings in process_groups_config.py and globals in parallel_state.py. Done means Mixtral conversion completes with TP=4/PP=2, prints Done!, and the checkpoint loads successfully.

Written by the indexing model from the issue text.