Skip to content

conversion

Experimental Megatron↔HF noise-layer conversion.

Converts a trained MegatronTransformerCloak into core's TransformerCloak so the trained noise layer can be exported for plain-HuggingFace inference (the product direction). The reverse direction (HF→Megatron) and a numerical round-trip check are tracked as follow-up work (SG-3626).

The estimator body differs by construction — Megatron deep-copies an MCore decoder layer, core uses an HF decoder layer — so the block weights are converted with Megatron-Bridge's own AutoBridge mappings (QKV/gate-up unfusion + TP gather) by wrapping the estimator's layer in a minimal one-layer MCore model. The column-parallel head is gathered to a full dense nn.Linear, and the "start-small" std init shift (carried in Megatron's parameterization, whose head is bias-free) is folded into core's linear bias (core carries the shift there, with a zero parameterization shift).

This module is experimental and requires the Megatron runtime (NeMo container).

Functions:

Name Description
megatron_to_hf_noise_layer

Convert a trained Megatron cloak into a core TransformerCloak for HF inference.

megatron_to_hf_noise_layer

megatron_to_hf_noise_layer(
    megatron_cloak: MegatronTransformerCloak,
    bridge: AutoBridge,
    transformer_type: type[PreTrainedModel],
    config: PretrainedConfig,
    *,
    noise_layer_attention: Any = None
) -> TransformerCloak

Convert a trained Megatron cloak into a core TransformerCloak for HF inference.

Parameters:

Name Type Description Default

megatron_cloak

MegatronTransformerCloak

The trained Megatron noise layer.

required

bridge

AutoBridge

The AutoBridge for the base model architecture (used to convert the estimator decoder-layer weights MCore→HF).

required

transformer_type

type[PreTrainedModel]

The HF backbone class for the core estimators (e.g. Qwen3Model).

required

config

PretrainedConfig

The HF PretrainedConfig of the base model.

required

noise_layer_attention

Any

Optional attention implementation for the core estimators.

None

Returns:

Type Description
TransformerCloak

A core TransformerCloak whose std/mean fields reproduce the Megatron cloak

TransformerCloak

(to numerical tolerance) on the same input embeddings.

Raises:

Type Description
NotImplementedError

If the Megatron cloak bounds its mean (BoundedMeanParameterization); only an unbounded mean is supported.