utils

Utilities for the Megatron noise-training integration.

Contains helpers shared across the Megatron noise layer, noisy model, and distillation provider:

• deepcopy_module_with_process_groups — deep-copy an MCore module while sharing (not copying) its process groups and CUDA streams, which are not copyable.
• freeze_module — disable gradients on every parameter of a module.
• should_attach_noise — decide whether the current pipeline stage should carry the noise layer (only the first pipeline stage does).

Only should_attach_noise needs the Megatron runtime; it imports megatron.core lazily so the copy/freeze helpers remain importable (and testable) without the NeMo container.

Functions:

Name	Description
deepcopy_module_with_process_groups	Deep-copy an MCore module while sharing its process groups and CUDA streams.
freeze_module	Disable gradient computation for every parameter of module in place.
should_attach_noise	Return whether the current pipeline stage should carry the noise layer.

deepcopy_module_with_process_groups

deepcopy_module_with_process_groups(
    module: Module,
) -> nn.Module

Deep-copy an MCore module while sharing its process groups and CUDA streams.

MCore modules hold references to torch.distributed.ProcessGroup and torch.cuda.Stream objects, which cannot be deep-copied. This collects those objects and seeds copy.deepcopy's memo with them so they are shared with the original rather than copied — every other tensor/parameter is copied as usual.

Parameters:

Name	Type	Description	Default
module	Module	The module to copy (e.g. an MCore transformer decoder block).	required

Returns:

Type	Description
nn.Module	A deep copy of module that shares the original's process groups and
nn.Module	CUDA streams.

freeze_module

freeze_module(module: Module) -> None

Disable gradient computation for every parameter of module in place.

Parameters:

Name	Type	Description	Default
module	Module	The module whose parameters should be frozen.	required

should_attach_noise

should_attach_noise(vp_stage: int | None) -> bool

Return whether the current pipeline stage should carry the noise layer.

The noise layer is attached at the model's embedding, so only the first pipeline stage needs it. When distributed state is not initialized (e.g. a single-process build), the answer is always True.

Pipeline parallelism is not supported by this integration; this guard exists so that, in any incidental multi-stage build, noise is attached on exactly the first stage rather than replicated.

Parameters:

Name	Type	Description	Default
vp_stage	int | None	The virtual pipeline stage index, or None when virtual pipeline parallelism is not in use.	required

Returns:

Type	Description
bool	True if the noise layer should be attached on this stage.