nemo_automodel.components.models.step5.model

Native Step-5 Preview text and vision-language training models.

Module Contents

Classes

Name	Description
`Step5ForCausalLM`	Step-5 causal LM with native grouped experts and model-owned CP.
`Step5ForConditionalGeneration`	Step-5 language model with the compatible Step perception encoder.
`Step5TextModel`	Embedding, decoder stack and final normalization.

Data

ModelClass

API

class nemo_automodel.components.models.step5.model.Step5ForCausalLM(
    config: nemo_automodel.components.models.step5.configuration.Step5Config | nemo_automodel.components.models.step5.configuration.Step5VLConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None
)

Bases: HFCheckpointingMixin, PreTrainedModel, MoEFSDPSyncMixin

Step-5 causal LM with native grouped experts and model-owned CP.

_keep_in_fp32_modules

= ['rotary_emb']

_keep_in_fp32_modules_strict

_no_split_modules

= ['Step5DecoderLayer']

_pp_keep_self_forward

bool = True

_pp_return_hidden_states

bool = False

_pp_return_hidden_states_supported

bool = True

backend

base_model_prefix

= 'model'

cp_mesh

DeviceMesh | None = None

lm_head

model

packed_mask_type

= 'document_ids'

state_dict_adapter

tie_word_embeddings_support

= TieSupport.UNTIED_ONLY

nemo_automodel.components.models.step5.model.Step5ForCausalLM._is_pipeline_parallel_stage() -> bool

Identify a stage after AutoPipeline removes modules owned by other stages.

nemo_automodel.components.models.step5.model.Step5ForCausalLM.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    logits_to_keep: int | torch.Tensor = 0,
    output_hidden_states: bool = False,
    labels: torch.Tensor | None = None,
    _packed_seq_ids: torch.Tensor | None = None,
    kwargs: nemo_automodel.components.models.step5.configuration.ConfigValue | torch.Tensor = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast | torch.Tensor

Evaluate all local query positions.

Parameters:

input_ids

torch.Tensor | NoneDefaults to None

Int64 tensor of shape [batch, local_sequence], or the FP32 residual tensor [batch, local_sequence, hidden] on later pipeline stages.

inputs_embeds

torch.Tensor | NoneDefaults to None

Optional tensor of shape [batch, local_sequence, hidden].

position_ids

torch.Tensor | NoneDefaults to None

Optional tensor of shape [batch, local_sequence], resetting at document starts.

attention_mask

torch.Tensor | NoneDefaults to None

Optional tensor of shape [batch, local_sequence], one for valid tokens.

padding_mask

torch.Tensor | NoneDefaults to None

Optional bool tensor of shape [batch, local_sequence], true for padding.

logits_to_keep

int | torch.TensorDefaults to 0

Last-token count or tensor of sequence indices to project.

output_hidden_states

boolDefaults to False

Return final hidden states for native fused CE.

labels

torch.Tensor | NoneDefaults to None

Optional tensor of shape [batch, local_sequence] for serial HF-style shifted CE.

_packed_seq_ids

torch.Tensor | NoneDefaults to None

Optional [batch, local_sequence] document IDs from native packing.

**kwargs

ConfigValue | torch.TensorDefaults to {}

Recipe metadata. BSHD inputs are used even for concatenated documents.

Returns: CausalLMOutputWithPast | torch.Tensor

Output with logits [batch, selected_sequence, vocab], optional final

nemo_automodel.components.models.step5.model.Step5ForCausalLM.from_config(
    config: nemo_automodel.components.models.step5.configuration.Step5Config | nemo_automodel.components.models.step5.configuration.Step5VLConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None
) -> nemo_automodel.components.models.step5.model.Step5ForCausalLM

classmethod

Construct a model through the AutoModel registry.

nemo_automodel.components.models.step5.model.Step5ForCausalLM.get_input_embeddings() -> torch.nn.Embedding

Return the language embedding module.

nemo_automodel.components.models.step5.model.Step5ForCausalLM.get_output_embeddings() -> torch.nn.Module

Return the untied language-model head.

nemo_automodel.components.models.step5.model.Step5ForCausalLM.get_pipeline_stage_metas(
    is_first: bool,
    microbatch_size: int,
    seq_len: int,
    dtype: torch.dtype
) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]

Declare stage I/O without rounding the FP32 residual at a pipeline boundary.

Parameters:

is_first

bool

Whether this stage receives token IDs.

microbatch_size

int

Number of sequences in each pipeline microbatch.

seq_len

int

Local sequence length after any context-parallel sharding.

dtype

torch.dtype

Compute dtype for the final normalized hidden states or logits.

Returns: tuple[torch.Tensor, ...]

Input and output metadata tuples. First-stage input IDs have shape

nemo_automodel.components.models.step5.model.Step5ForCausalLM.init_weights(
    buffer_device: torch.device | None = None
) -> None

Initialize all trainable tensors for checkpoint-free construction.

nemo_automodel.components.models.step5.model.Step5ForCausalLM.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None

Initialize after meta allocation and apply precision-aware casting.

nemo_automodel.components.models.step5.model.Step5ForCausalLM.prepare_model_inputs_for_cp(
    batch: dict[str, torch.Tensor],
    num_chunks: int = 1
) -> dict[str, nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder | torch.Tensor]

Choose the framework’s contiguous sharder.

Parameters:

batch

dict[str, torch.Tensor]

Mapping with [batch, global_sequence] token and label tensors, optional [batch, global_sequence, hidden] embeddings and image tensors.

num_chunks

intDefaults to 1

Number of pipeline microbatches; each uses the same sequence sharder.

Returns: dict[str, ContextParallelSharder | torch.Tensor]

Mapping containing the model’s ContextParallelSharder.

class nemo_automodel.components.models.step5.model.Step5ForConditionalGeneration(
    config: nemo_automodel.components.models.step5.configuration.Step5VLConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None
)

Bases: Step5ForCausalLM

Step-5 language model with the compatible Step perception encoder.

vision_model

= Step5VisionEncoder(config.vision_config)

vit_large_projector

nemo_automodel.components.models.step5.model.Step5ForConditionalGeneration._encode_images(
    pixels: torch.Tensor,
    compute_dtype: torch.dtype
) -> torch.Tensor

Encode overview or crop images.

Parameters:

pixels

torch.Tensor

Tensor of shape [images, channels, height, width].

compute_dtype

torch.dtype

Dtype of the completed language embedding projection.

Returns: torch.Tensor

Tensor of shape [images, image_tokens, hidden].

nemo_automodel.components.models.step5.model.Step5ForConditionalGeneration.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    pixel_values: torch.Tensor | None = None,
    patch_pixel_values: torch.Tensor | None = None,
    num_patches: torch.Tensor | list[int] | None = None,
    kwargs: nemo_automodel.components.models.step5.configuration.ConfigValue | torch.Tensor = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast

Inject projected crop/overview tokens and run the text backbone.

Parameters:

input_ids

torch.Tensor | NoneDefaults to None

Tensor of shape [batch, local_sequence].

inputs_embeds

torch.Tensor | NoneDefaults to None

Optional tensor of shape [batch, local_sequence, hidden].

pixel_values

torch.Tensor | NoneDefaults to None

Optional tensor of shape [images, channels, height, width].

patch_pixel_values

torch.Tensor | NoneDefaults to None

Optional tensor of shape [crops, channels, height, width].

num_patches

torch.Tensor | list[int] | NoneDefaults to None

Crop counts of shape [images], or equivalent list.

**kwargs

ConfigValue | torch.TensorDefaults to {}

Tensor and scalar arguments from Step5ForCausalLM.forward.

Returns: CausalLMOutputWithPast

Language output with logits of shape [batch, local_sequence, vocab].

nemo_automodel.components.models.step5.model.Step5ForConditionalGeneration.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None

Initialize both towers after meta allocation, then apply compute dtype.

class nemo_automodel.components.models.step5.model.Step5TextModel(
    config: nemo_automodel.components.models.step5.configuration.Step5Config,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    dtype: torch.dtype
)

Bases: Module

Embedding, decoder stack and final normalization.

csa_backward_full_blocks

csa_forward_full_blocks

embed_tokens

layers

norm

swa_varlen

nemo_automodel.components.models.step5.model.Step5TextModel.forward(
    embeddings: torch.Tensor,
    positions: torch.Tensor,
    padding_mask: torch.Tensor,
    packed_documents: torch.Tensor | None = None
) -> torch.Tensor

Evaluate the decoder stack with FP32 residuals.

Parameters:

embeddings

torch.Tensor

Tensor of shape [batch, local_sequence, hidden].

positions

torch.Tensor

Tensor of shape [batch, local_sequence].

padding_mask

torch.Tensor

Boolean tensor of shape [batch, local_sequence], true for padding.

packed_documents

torch.Tensor | NoneDefaults to None

Optional independent document IDs [batch, local_sequence], positive for real tokens and zero for padding.

Returns: torch.Tensor

Tensor of shape [batch, local_sequence, hidden]: FP32 residuals on

nemo_automodel.components.models.step5.model.ModelClass = Step5ForCausalLM