Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsStep5nemo_automodel.components.models.step5.parallelization

nemo_automodel.components.models.step5.parallelization

View as Markdown

FSDP sidecar that keeps the Step 5 final norm in an FP32 unit.

Module Contents

Classes

NameDescription
Step5ModelParallelizerKeep the final decoder norm in its own FP32 FSDP unit.

Data

PARALLELIZER

_FP32_NORM_POLICY

API

class nemo_automodel.components.models.step5.parallelization.Step5ModelParallelizer()

Bases: ModelParallelizer

Keep the final decoder norm in its own FP32 FSDP unit.

nemo_automodel.components.models.step5.parallelization.Step5ModelParallelizer._fully_shard_module(
    module: torch.nn.Module,
    kwargs = {}
) -> torch.nn.Module

Shard a module, first giving any contained final norm an FP32 unit.

Decoder blocks already group their strict-FP32 norms by dtype. The final norm is owned by the root unit, which uses the BF16 parameter policy, so it is wrapped separately before that root.

Parameters:

module
nn.Module

Module to shard in place.

**kwargs
Defaults to {}

fully_shard keyword arguments (mesh, mp_policy, offload_policy, reshard_after_forward, ignored_params).

Returns: nn.Module

The input module with FSDP applied.

nemo_automodel.components.models.step5.parallelization.PARALLELIZER = Step5ModelParallelizer()
 
nemo_automodel.components.models.step5.parallelization._FP32_NORM_POLICY = MixedPrecisionPolicy(param_dtype=(torch.float32), reduce_dtype=(torch.float32), ...