noise_layer

Megatron-native noise layer for Stained Glass noise distillation.

Provides:

• MegatronTransformerBlockEstimator — a deep-copied MCore transformer decoder block followed by a column-parallel head that maps hidden states to raw estimate values (rhos).
• MegatronTransformerCloak — the noise layer proper: independent std and mean estimators whose raw outputs are mapped to a standard-deviation field and a mean field by core's parameterizations, then combined into additive Gaussian noise on the embedding.

Design notes:

• This is a Megatron-native module (MegatronModule + ColumnParallelLinear + a deep-copied MCore decoder block); it does not subclass core's BaseNoiseLayer (whose torch-DTensor tensor-parallel backend and RNG-fork hooks conflict with Megatron TP/CP). It reuses core's parameterization math, which is pure and elementwise and therefore tensor-parallel safe.
• Unlike the flat POC (std = clamp(min + |rho|, max)), the std field is produced by CloakStandardDeviationParameterization (tanh, invertible). The mean is unbounded by default (matching core's TransformerCloak); passing mean_bound applies BoundedMeanParameterization (tanh-bounded — structurally limits mean drift under bf16 + RMSNorm-saturated gradients) instead.
• Mirrors core's TransformerCloak, which builds independent std and mean estimator bodies.

Classes:

Name	Description
MegatronTransformerBlockEstimator	A deep-copied MCore decoder block plus a column-parallel head producing raw estimates.
MegatronTransformerCloak	Megatron noise layer: independent std/mean estimators + additive Gaussian noise.

MegatronTransformerBlockEstimator

Bases: MegatronModule

A deep-copied MCore decoder block plus a column-parallel head producing raw estimates.

The block is a deep copy of a source decoder layer (so its parameters are independent and trainable), and the head is a bias-free ColumnParallelLinear that gathers its output to the full hidden dimension. The returned rhos are unconstrained; a parameterization maps them to a std or mean field.

Methods:

Name	Description
__init__	Initialize the estimator.
forward	Compute raw estimate values from hidden states.

__init__

__init__(config: Any, source_layer: Module) -> None

Initialize the estimator.

Parameters:

Name	Type	Description	Default
config	Any	The MCore transformer config of the base model.	required
source_layer	Module	The decoder layer to deep-copy as the estimator body.	required

forward

forward(
    hidden_states: Tensor,
    attention_mask: Tensor,
    **block_kwargs: Any
) -> Tensor

Compute raw estimate values from hidden states.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Input hidden states, shape (sequence, batch, hidden).	required
attention_mask	Tensor	Attention mask forwarded to the decoder block.	required
**block_kwargs	Any	Additional decoder-block keyword arguments (e.g. rotary position embeddings) forwarded verbatim.	{}

Returns:

Type	Description
Tensor	Raw estimate tensor (rhos) of shape (sequence, batch, hidden).

MegatronTransformerCloak

Bases: MegatronModule

Megatron noise layer: independent std/mean estimators + additive Gaussian noise.

Produces, from a clean embedding tensor, a per-element standard-deviation field and mean field, then additive Gaussian noise noise = z * std + mean (with z ~ N(0, I)). The forward returns (clean, noisy, std, mean) so the distillation loss has everything it needs at the injection point without a super-batch.

Methods:

Name	Description
__init__	Initialize the cloak.
forward	Compute clean/noisy embeddings and the std/mean fields.

__init__

__init__(
    config: Any,
    source_layer: Module,
    *,
    std_scale: tuple[float, float],
    std_init_shift: float,
    std_shallow: float = 1.0,
    mean_bound: float | None = None,
    mean_init_shift: float = 0.0
) -> None

Initialize the cloak.

The standard-deviation bounds (std_scale) and init shift (std_init_shift) are required — they govern the privacy floor/ceiling of the noise, so they are always chosen explicitly rather than defaulted. The mean is unbounded by default (mean_bound=None), matching core's TransformerCloak; pass a bound to constrain it.

Parameters:

Name	Type	Description	Default
config	Any	The MCore transformer config of the base model.	required
source_layer	Module	The decoder layer deep-copied by each estimator body.	required
std_scale	tuple[float, float]	Asymptotic (min, max) bounds of the std field.	required
std_init_shift	float	Shift added to the std estimator's raw output before the tanh; a negative value starts the std field near min_scale.	required
std_shallow	float	Temperature of the std parameterization's tanh.	1.0
mean_bound	float | None	Magnitude bound of the mean field ((-bound, +bound)). None (the default) leaves the mean unbounded.	None
mean_init_shift	float	Shift added to the mean estimator's raw output (only applied when mean_bound is set).	0.0

forward

forward(
    hidden_states: Tensor,
    attention_mask: Tensor,
    noise_mask: Tensor | None = None,
    **block_kwargs: Any
) -> tuple[Tensor, Tensor, Tensor, Tensor]

Compute clean/noisy embeddings and the std/mean fields.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Clean embedding tensor, shape (sequence, batch, hidden).	required
attention_mask	Tensor	Attention mask forwarded to both estimator blocks.	required
noise_mask	Tensor | None	Optional mask broadcastable to hidden_states that zeroes the additive noise at masked positions. None (the default, and the distillation path) applies noise everywhere; padding is masked downstream by the loss reduction.	None
**block_kwargs	Any	Additional decoder-block keyword arguments (e.g. rotary position embeddings) forwarded to both estimator blocks.	{}

Returns:

Type	Description
tuple[Tensor, Tensor, Tensor, Tensor]	A tuple (clean_embeddings, noisy_embeddings, stds, means).