This section is too much of an implementation detail IMO, I'd skip it

Removed the section. The opt-in check is now described in the refit_verification.py module docstring instead.

Done in 4fd4c20.

This paragraph seems unnecessary

Removed.

Done in 4fd4c20.

what is "full hf seed"?

It meant the Hugging Face model directory the GSPO run started from. The command now uses --seed <hf-model-dir> and says that directory contains config.json, llm_backbone/ and the tokenizer.

Done in 4fd4c20.

The default format's name should be "safetensors", Lightning's DCP is the input not the output

checkpoint_format is the format of the input ckpt_path: lightning, or the nemo_rl DCP directory that NeMo RL saves. The output is always HF safetensors. I kept the name and rewrote the comment to say this.

Done in 4fd4c20.

Unnecessary comment

Removed.

Done in 4fd4c20.

are you sure that skipping device_mesh is None condition here is not breaking anything?

Yes. Without parameter_config, the no-mesh path behaves as before:

• Perception is still not cast: the cast now runs only with a device mesh or a parameter_config.
• Perception activation checkpointing is still applied. It moved after the LoRA setup, which only touches the LLM, so the order doesn't matter.
• The same early return (maybe_load_pretrained_models and the MTP training mode) follows.

The early return moved below the parameter_config block so that single-process construction, such as to_hf.py export, also applies the selected parameter dtypes before activation-checkpoint wrapping. The unchanged path is exercised by the shared fixture in tests/collections/speechlm2/test_salm_automodel.py, which calls configure_model() without a mesh.

I don't understand this docstring

Rewrote it: the loader checks every name and shape before copying anything. The first load must cover every perception tensor; later RL weight updates may send only some of them. The function now lives in weight_loading.py (see the merge thread).

Done in 4fd4c20.

Can this docstring explain what do you mean by "Refit"

The docstring now starts with the definition: in NeMo RL, a refit copies the trainer's updated weights into the vLLM generation engine after each training step. It also describes how to enable the check, which used to be in the docs.

Done in 4fd4c20.

Can we merge perception_weight_loading.py into this file?

Merged: load_perception_state is now in weight_loading.py, and perception_weight_loading.py is deleted.

Done in 4fd4c20.

what is this doing?

NeMo RL depends on nemo-lens for telemetry, which installs nemo.lens into the same nemo package. The GSPO guide installs Speech in editable mode, so without extend_path the nemo package resolves only to the Speech checkout and import nemo.lens fails. extend_path lets Python find the other installed nemo.* modules and changes nothing otherwise. tests/test_namespace_package.py covers it. I rewrote the comment to say this.

Done in 4fd4c20.

[P2] Apply the explicit parameter dtype policy before any blanket cast

Failure: With torch_dtype=bfloat16 and a ParameterConfig selecting FP32 trainable perception parameters, this casts those parameters to BF16 before apply_parameter_config casts them back to FP32. For example, a pristine weight of 0.1234567 becomes approximately 0.12353515625, stored as FP32. This affects initialization from pretrained ASR weights, which setup_speech_encoder has already loaded; without a subsequent full-model checkpoint load, the precision loss is permanent for that model instance.

Correction: Let the explicit policy cast parameters directly from their pristine values. Preserve the existing buffer conversion separately. A suggested replacement for this cast block is:

if dtype != torch.float32:
    if self.parameter_config is None:
        if device_mesh is not None:
            self.perception.to(dtype=dtype)
    else:
        for module in self.perception.modules():
            for name, buffer in module._buffers.items():
                if buffer is not None and (
                    buffer.is_floating_point() or buffer.is_complex()
                ):
                    module._buffers[name] = buffer.to(dtype=dtype)

Keep the following apply_parameter_config call before activation-checkpoint wrapping. The invariant is that explicit FP32 parameters must never pass through BF16 storage first.

Preservation and acceptance: Add a proposed configure_model unit test with a tiny mocked perception stack, a selected FP32 weight initialized to 0.1234567, a frozen BF16-selected parameter, and a floating frontend buffer. Assert exact equality for the selected FP32 value, the declared frozen dtype/trainability, and the existing frontend-buffer cast. Cover mesh and no-mesh construction. Without ParameterConfig, retain the current distributed BF16 cast and the no-mesh behavior that leaves perception storage unchanged. These checks are proposed, not executed.

[P2] Preserve CTC storage in the single-process Lightning export path

Failure: A Lightning checkpoint containing the newly registered CTC artifact takes this branch unless checkpoint_format=nemo_rl. The whole-model .to(bfloat16) rounds FP32 CTC decoder weights before HFHubMixin._save_pretrained saves them. Thus a CTC weight such as 1.000123 is exported as BF16 1.0, despite the preservation contract. The canonical saver exempts CTC tensors, but this path bypasses it.

Correction: Exclude the artifact from this dtype-only cast while retaining the existing save_pretrained serialization, including its shared-tensor handling. One suggested replacement for model.to(...) is:

artifact = getattr(model, "ctc_timestamp", None)
protected = set() if artifact is None else {
    id(t) for t in (*artifact.parameters(), *artifact.buffers())
}
target_dtype = str_to_dtype(cfg.dtype)

def cast_for_export(tensor):
    if id(tensor) in protected or not (
        tensor.is_floating_point() or tensor.is_complex()
    ):
        return tensor
    return tensor.to(dtype=target_dtype)

model = model._apply(cast_for_export)

The invariant is that CTC parameters and tokenizer bytes retain their loaded storage and values, while ordinary floating weights still use the requested export dtype.

Acceptance: Add a proposed test through the non-native single-process branch using the existing CTC bundle fixture and dtype=bfloat16. Reload the output and assert exact FP32 CTC values, unchanged uint8 tokenizer bytes, identical CTC metadata, and BF16 ordinary weights. Retain the existing native/distributed saver preservation assertions and no-CTC Lightning behavior. This patch and these checks have not been executed.

[P2] Do not invent an MTP head when exporting a headless native checkpoint

Failure: Native export unconditionally sets retain_native_mtp=True. For a valid SpeechLM checkpoint whose backbone has MTP depth zero and whose SpeechLM MTP block is disabled, configure_model consequently supplies positive-depth MTP overrides. The LLM loader treats the absent head as a request to initialize a fresh one. Strict native loading then rejects the checkpoint for missing llm.mtp.* tensors. Export should retain the actual topology, including the absence of a head.

Correction: Resolve whether a head is expected before overriding the training switch. For example, in the native-format block in main, preserve the original enabled/retention switches and resolve the backbone config before calling this helper:

from transformers import AutoConfig
from nemo.collections.speechlm2.parts.pretrained import _automodel_config_mtp_depth

mtp_requested = bool((model_cfg.get("mtp") or {}).get("enabled", False))
retention_requested = bool(model_cfg.get("retain_native_mtp", False))
backbone_cfg = AutoConfig.from_pretrained(
    model_cfg["pretrained_llm"],
    trust_remote_code=model_cfg.get("trust_remote_code", False),
)
model_cfg = prepare_native_export_model_config(model_cfg)
model_cfg["retain_native_mtp"] = (
    mtp_requested or retention_requested or _automodel_config_mtp_depth(backbone_cfg) > 0
)

Do not infer optionality solely from missing checkpoint tensors: that would weaken strict coverage for a genuinely expected head.

Preservation and acceptance: Proposed tests should export a depth-zero/disabled native model successfully with no MTP tensors and disabled output metadata; retain a native positive-depth head; retain repeated-head physical depth one plus its logical depth; and retain a recipe-created head when the original SpeechLM MTP switch was enabled even though the base backbone had depth zero. Keep the existing assertion that an explicitly requested retained head which was not constructed fails export. No tests were executed.

[P2] Normalize nullable optional configuration blocks before merging

Failure: prepare_native_export_model_config({"mtp": None}) raises TypeError at **None; the same happens with automodel_backend=None. Both are valid disabled/default representations in the model constructor: MTP is normalized with or {}, and a null backend is explicitly skipped. A seed/config with either null block therefore fails before native checkpoint loading.

Correction: Use these replacements in this helper:

model_cfg["mtp"] = {**(model_cfg.get("mtp") or {}), "enabled": False}
model_cfg["automodel_backend"] = {
    **(model_cfg.get("automodel_backend") or {}),
    "enable_hf_state_dict_adapter": True,
}

Apply the same nullable-MTP normalization to the retained-head merge in _hf_export_config, since direct model configurations can retain a head while leaving the optional training block null.

Preservation and acceptance: Parameterize proposed tests over absent, null, and dictionary-valued blocks. Absent/null inputs must prepare successfully; dictionary inputs must retain backend settings and repeated-head options while changing only the intended switches. Preserve the existing assertion that preparation does not mutate the input. With an attached positive-depth retained head and mtp=None, config export must produce enabled MTP metadata rather than a TypeError. These are static deductions and proposed checks, not executed results.

[P2] Initialize the new loader state in the existing FP8 test fixture

Failure: This new unconditional attribute access breaks the existing test_plugin_perception_loader_loads_fp8_checkpoint. That test constructs the model through object.__new__, initializes nn.Module, and sets _uses_pe_encoder=True, but never runs the constructor that initializes _perception_weights_initialized. In a vLLM-enabled test environment, a valid FP8 checkpoint now raises AttributeError before loading instead of satisfying the existing regression assertions.

Correction: Add this line to that fixture before _load_perception_weights(state):

model._perception_weights_initialized = False

Production construction already initializes the flag correctly; this does not require weakening the first-load coverage check.

Preservation and acceptance: Run that test in a vLLM-enabled environment and retain its assertions that every perception name is reported loaded, FP8 weight bytes match exactly, and the diarization weights are BF16. Also retain the neighboring invalid-quantization test's ValueError expectation and the updated PE loader test's initial-complete/subsequent-partial behavior. The failure is established by source inspection; tests were not executed here.