noisy_model

MegatronNoisyModel: wrap an MCore GPT model and inject noise at the embedding.

MegatronNoisyModel wraps a frozen MCore GPTModel and injects a MegatronTransformerCloak via forward hooks on the embedding: a pre-hook on the base model stashes the forward context (attention mask, inference/packed-sequence params), and a forward hook on the embedding recomputes the rotary position embeddings and replaces the clean embedding output with the noisy one, caching (clean, noisy, std, mean) for the distillation loss.

wrap_models_with_noise applies the wrapper stage-aware: only the pipeline stage that owns the embedding (the first stage) receives the noise layer; other stages are frozen and returned unchanged. Pipeline parallelism is not supported by this integration, but the guard keeps any incidental multi-stage build correct.

Classes:

Name	Description
MegatronNoisyModel	Wrap an MCore GPT model and inject an embedding-level noise layer.

Functions:

Name	Description
wrap_models_with_noise	Wrap the embedding-owning model chunk(s) with MegatronNoisyModel.

MegatronNoisyModel

Bases: MegatronModule

Wrap an MCore GPT model and inject an embedding-level noise layer.

The base model is frozen; only the noise layer is trainable. During each forward, the clean embedding output is replaced by the noisy embedding, and the noise-layer intermediates are cached on the instance for the distillation loss.

Methods:

Name	Description
__getattr__	Delegate unknown attributes to the wrapped base model.
__init__	Initialize the noisy model.
forward	Run the base model (with noise injected via the embedding hook).
remove_noise_hook	Remove the registered forward hooks, restoring the base model's behavior.
set_input_tensor	Delegate the pipeline input-tensor plumbing to the base model.

__getattr__

__getattr__(name: str) -> Any

Delegate unknown attributes to the wrapped base model.

Parameters:

Name	Type	Description	Default
name	str	The attribute name.	required

Returns:

Type	Description
Any	The attribute from this module if present, otherwise from the base model.

__init__

__init__(
    base_model: MegatronModule,
    *,
    noise_layer_kwargs: Mapping[str, Any] | None = None
) -> None

Initialize the noisy model.

Parameters:

Name	Type	Description	Default
base_model	MegatronModule	The MCore GPTModel to wrap. It is frozen in place.	required
noise_layer_kwargs	Mapping[str, Any] | None	Optional keyword arguments forwarded to MegatronTransformerCloak (e.g. std_scale, mean_bound).	None

forward

forward(*args: Any, **kwargs: Any) -> Any

Run the base model (with noise injected via the embedding hook).

Parameters:

Name	Type	Description	Default
*args	Any	Positional arguments forwarded to the base model.	()
**kwargs	Any	Keyword arguments forwarded to the base model.	{}

Returns:

Type	Description
Any	The base model's output.

remove_noise_hook

remove_noise_hook() -> None

Remove the registered forward hooks, restoring the base model's behavior.

set_input_tensor

set_input_tensor(
    input_tensor: Tensor | list[Tensor],
) -> None

Delegate the pipeline input-tensor plumbing to the base model.

Parameters:

Name	Type	Description	Default
input_tensor	Tensor | list[Tensor]	The input tensor(s) from the previous pipeline stage.	required

wrap_models_with_noise

wrap_models_with_noise(
    models: Sequence[MegatronModule],
    *,
    noise_layer_kwargs: Mapping[str, Any] | None = None
) -> list[MegatronModule]

Wrap the embedding-owning model chunk(s) with MegatronNoisyModel.

Only the pipeline stage that owns the embedding (the first stage) is wrapped; other stages are frozen and returned unchanged. Pipeline parallelism is not supported, but this keeps any incidental multi-stage build correct.

Parameters:

Name	Type	Description	Default
models	Sequence[MegatronModule]	The model chunks produced by the Megatron provider.	required
noise_layer_kwargs	Mapping[str, Any] | None	Optional keyword arguments forwarded to the noise layer.	None

Returns:

Type	Description
list[MegatronModule]	The wrapped model chunks.