Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsnemo_automodel.components.models.step3p7

nemo_automodel.components.models.step3p7

View as Markdown

Submodules

Package Contents

Classes

NameDescription
Step3p7ConfigTop-level configuration for Step3.7 vision-language checkpoints.
Step3p7ForConditionalGenerationNative Step3.7 VLM implementation for MedPix fine-tuning with EP and PP.
Step3p7TextConfigConfiguration for the Step3.7 language backbone.
StepRoboticsVisionEncoderConfigConfiguration for the Step robotics vision encoder.

API

class nemo_automodel.components.models.step3p7.configuration_step3p7.Step3p7Config(
    vision_config: typing.Union[dict, nemo_automodel.components.models.step3p7.configuration_step3p7.StepRoboticsVisionEncoderConfig] | None = None,
    text_config: typing.Union[dict, nemo_automodel.components.models.step3p7.configuration_step3p7.Step3p7TextConfig] | None = None,
    understand_projector_stride: int = 2,
    projector_bias: bool = False,
    image_token_id: int = 151679,
    kwargs = {}
)

Bases: PretrainedConfig

Top-level configuration for Step3.7 vision-language checkpoints.

hidden_size
= text_config.hidden_size
max_position_embeddings
= text_config.max_position_embeddings
model_type
= 'step3p7'
nemo_automodel.components.models.step3p7.configuration_step3p7.Step3p7Config.to_dict()
class nemo_automodel.components.models.step3p7.model.Step3p7ForConditionalGeneration(
    config: nemo_automodel.components.models.step3p7.configuration_step3p7.Step3p7Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs: typing.Any = {}
)

Bases: HFCheckpointingMixin, Module, MoEFSDPSyncMixin

Native Step3.7 VLM implementation for MedPix fine-tuning with EP and PP.

_keep_in_fp32_modules
= ['rotary_emb']
_pp_keep_self_forward
bool = True
backend
= backend or BackendConfig()
lm_head
model
mtp
mtp_config
pad_token_id
= getattr(config.text_config, 'pad_token_id', None)
state_dict_adapter
tie_word_embeddings_support
TieSupport = TieSupport.UNTIED_ONLY
vocab_size
= config.text_config.vocab_size
nemo_automodel.components.models.step3p7.model.Step3p7ForConditionalGeneration._build_mtp_embed_inputs_from_embeds(
    inputs_embeds: torch.Tensor
) -> tuple[torch.Tensor, ...]
nemo_automodel.components.models.step3p7.model.Step3p7ForConditionalGeneration._is_pipeline_parallel_stage() -> bool
nemo_automodel.components.models.step3p7.model.Step3p7ForConditionalGeneration._make_position_ids(
    hidden: torch.Tensor,
    position_ids: torch.Tensor | None
) -> torch.Tensor
nemo_automodel.components.models.step3p7.model.Step3p7ForConditionalGeneration.customize_pipeline_stage_modules(
    module_names_per_stage: list[list[str]],
    layers_prefix: str,
    text_model: torch.nn.Module | None = None
) -> list[list[str]]
nemo_automodel.components.models.step3p7.model.Step3p7ForConditionalGeneration.forward(
    input_ids: torch.Tensor | None = None,
    mtp_embed_inputs: torch.Tensor = (),
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    cache_position: torch.Tensor | None = None,
    logits_to_keep: typing.Union[int, torch.Tensor] = 0,
    output_hidden_states: bool | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor | nemo_automodel.components.models.step3p7.model.Step3p7CausalLMOutput

Run one microbatch with direct or explicitly staged patch-based media.

Parameters:

input_ids
torch.Tensor | NoneDefaults to None

Token IDs [batch, sequence], or [tokens] for THD.

*mtp_embed_inputs
torch.TensorDefaults to ()

Optional MTP embeddings [batch, sequence, hidden] propagated by a previous pipeline stage.

position_ids
torch.Tensor | NoneDefaults to None

Scalar positions matching the token axes.

attention_mask
torch.Tensor | NoneDefaults to None

Optional [batch, sequence] attention eligibility.

padding_mask
torch.Tensor | NoneDefaults to None

Optional boolean padding flags matching token axes.

inputs_embeds
torch.Tensor | NoneDefaults to None

Optional embeddings [batch, sequence, hidden].

cache_position
torch.Tensor | NoneDefaults to None

Optional cache positions [sequence].

logits_to_keep
Union[int, torch.Tensor]Defaults to 0

Number of trailing logits, or indices [selected_tokens].

output_hidden_states
bool | NoneDefaults to None

Include final hidden states in the model output.

**kwargs
AnyDefaults to {}

Media pixels [images, channels, height, width], patch pixels [patches, channels, height, width], patch counts [images], newline flags [patches], and optional packed cumulative offsets [segments + 1].

Returns: torch.Tensor | Step3p7CausalLMOutput

Logits [batch, sequence, vocab] or a model output containing those

nemo_automodel.components.models.step3p7.model.Step3p7ForConditionalGeneration.from_config(
    config: nemo_automodel.components.models.step3p7.configuration_step3p7.Step3p7Config,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs: typing.Any = {}
)
classmethod
nemo_automodel.components.models.step3p7.model.Step3p7ForConditionalGeneration.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args: typing.Any = (),
    kwargs: typing.Any = {}
)
classmethod
nemo_automodel.components.models.step3p7.model.Step3p7ForConditionalGeneration.get_decoder()
nemo_automodel.components.models.step3p7.model.Step3p7ForConditionalGeneration.get_input_embeddings()
nemo_automodel.components.models.step3p7.model.Step3p7ForConditionalGeneration.get_output_embeddings()
nemo_automodel.components.models.step3p7.model.Step3p7ForConditionalGeneration.get_pipeline_stage_metas(
    is_first: bool,
    microbatch_size: int,
    seq_len: int,
    dtype: torch.dtype
) -> tuple[tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]]
nemo_automodel.components.models.step3p7.model.Step3p7ForConditionalGeneration.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None
nemo_automodel.components.models.step3p7.model.Step3p7ForConditionalGeneration.prepare_model_inputs_for_cp(
    batch: dict[str, typing.Any],
    num_chunks: int = 1
) -> dict[str, typing.Any]

Return a sharder-only CP backend; embed + splice + shard happen in forward.

Embedding and the vision multimodal scatter now run inside forward per microbatch (see the CP branch that calls get_multimodal_embeddings

  • prepare_inputs_embeds + shard_sequence_for_cp_round_robin). The returned ContextParallelSharder round-robin-shards only the no-grad aux streams (labels/position_ids/loss_mask/padding_mask) and leaves input_ids and the media inputs full-length for the forward. Step3.7 uses plain 1-D positions, so no position_ids are computed here (the aux shard injects and slices them).

Parameters:

batch
dict[str, Any]

The full-sequence batch (with input_ids [batch, sequence]); left intact.

num_chunks
intDefaults to 1

Accepted for hook-signature parity; unused (round-robin CP).

nemo_automodel.components.models.step3p7.model.Step3p7ForConditionalGeneration.set_decoder(
    decoder
)
nemo_automodel.components.models.step3p7.model.Step3p7ForConditionalGeneration.set_input_embeddings(
    value
)
nemo_automodel.components.models.step3p7.model.Step3p7ForConditionalGeneration.set_output_embeddings(
    new_embeddings
)
class nemo_automodel.components.models.step3p7.configuration_step3p7.Step3p7TextConfig(
    hidden_size: int = 4096,
    intermediate_size: int = 11264,
    num_attention_heads: int = 64,
    num_attention_groups: int = 8,
    num_hidden_layers: int = 45,
    num_nextn_predict_layers: int = 0,
    mtp_base_layer_idx: int | None = None,
    max_seq_len: int = 128000,
    vocab_size: int = 128815,
    rms_norm_eps: float = 1e-05,
    moe_intermediate_size: int = 1280,
    moe_num_experts: int = 288,
    moe_top_k: int = 8,
    rope_theta: float = 10000,
    rope_scaling: dict[str, typing.Any] | None = None,
    max_position_embeddings: int = 128000,
    share_expert_dims: int = 1280,
    share_expert_dim: int | None = None,
    head_dim: int = 128,
    norm_expert_weight: bool = True,
    layer_types: list[str] = None,
    sliding_window: int | None = None,
    pad_token_id: int = 1,
    attention_dropout: float = 0.0,
    use_head_wise_attn_gate: bool = False,
    use_moe_router_bias: bool = False,
    moe_router_activation: str = 'softmax',
    moe_router_scaling_factor: float = 1.0,
    need_fp32_gate: bool = False,
    attention_other_setting: dict[str, typing.Any] | None = None,
    swiglu_limits: list[float | None] | None = None,
    swiglu_limits_shared: list[float | None] | None = None,
    use_rope_layers: list[bool] | None = None,
    yarn_only_types: list[str] | None = None,
    moe_layers_enum: tuple[int] = (3, 4, 5, 6, 7, 8, 9, 10, 1...,
    kwargs = {}
)

Bases: PretrainedConfig

Configuration for the Step3.7 language backbone.

architectures
= ['Step3p5ForCausalLM']
model_type
= 'step3p5'
mtp_base_layer_idx
= int(mtp_base_layer_idx)
mtp_layer_types
mtp_partial_rotary_factors
mtp_rope_theta
mtp_swiglu_limits
mtp_swiglu_limits_shared
mtp_use_rope_layers
nemo_automodel.components.models.step3p7.configuration_step3p7.Step3p7TextConfig.to_dict()
class nemo_automodel.components.models.step3p7.configuration_step3p7.StepRoboticsVisionEncoderConfig(
    width = 1536,
    layers = 47,
    heads = 16,
    num_channels = 3,
    image_size = 728,
    mlp_ratio = 8960 / 1536,
    patch_size = 14,
    hidden_act = 'quick_gelu',
    layer_norm_eps = 1e-05,
    ues_cls_token = False,
    use_cls_token: bool | None = None,
    use_ln_pre = True,
    use_ln_post = False,
    use_abs_posemb = True,
    use_rope2d = True,
    ls_init_value = 0.1,
    kwargs = {}
)

Bases: PretrainedConfig

Configuration for the Step robotics vision encoder.

model_type
= 'perception_encoder'