nemo_automodel.components.models.step5.configuration

Checkpoint configuration for Step-5 Preview.

Module Contents

Classes

Name	Description
`Step5BackendConfig`	Model-owned execution options, disabled by default.
`Step5Config`	Step-5 language backbone with per-layer RoPE and CSA attention.
`Step5SparseConfig`	CSA region selection geometry and normalization semantics.
`Step5VLConfig`	Step-5 perception encoder and language model configuration.
`Step5VisionConfig`	Perception configuration with explicit rotary semantics.

Data

ConfigValue

API

class nemo_automodel.components.models.step5.configuration.Step5BackendConfig(
    attn: nemo_automodel.components.models.common.utils.AttentionBackend = 'te' if HAVE_TE and torch.c...,
    sparse_attn: typing.Literal['generic', 'msa'] = 'generic',
    linear: typing.Literal['torch', 'te', 'quack'] = 'te' if HAVE_TE and torch.c...,
    rms_norm: typing.Literal['torch', 'torch_fp32', 'te', 'quack'] = 'torch_fp32',
    rope: typing.Literal['torch', 'quack', 'triton'] = 'torch',
    rope_fusion: bool = HAVE_TE and torch.cuda.is_a...,
    experts: typing.Literal['torch', 'te', 'gmm', 'torch_mm', 'torch_mm_mxfp8'] = 'torch_mm' if torch.cuda.is...,
    dispatcher: typing.Literal['torch', 'deepep', 'hybridep', 'uccl_ep', 'mok'] = 'deepep' if HAVE_DEEP_EP an...,
    dispatcher_num_sms: int = 20,
    dispatcher_share_token_dispatcher: bool = True,
    dispatcher_async_dispatch: bool = False,
    dispatcher_capacity_factor: float | None = None,
    dispatcher_equal_token_counts: bool = False,
    mok: nemo_automodel.components.models.common.utils.MoKBackendConfig = MoKBackendConfig(),
    dispatcher_hybridep_permute_fusion: bool = False,
    dispatcher_hybridep_compact_routing: bool = False,
    dispatcher_hybridep_num_sms_preprocessing: int | None = None,
    dispatcher_hybridep_num_blocks_permute: int | None = None,
    dispatcher_hybridep_num_blocks_unpermute: int | None = None,
    enable_deepep: bool | None = None,
    fake_balanced_gate: bool = False,
    fake_gate_noise: float = 0.0,
    enable_hf_state_dict_adapter: bool = True,
    enable_fsdp_optimizations: bool = False,
    te_fp8: nemo_automodel.components.models.common.utils.TEFp8Config | None = None,
    gate_precision: str | torch.dtype | None = None,
    compile_attn: bool = False,
    compile_situ: bool = False,
    compile_norm: bool = False,
    compile_hc: bool = False,
    shared_expert_overlap: bool = False,
    compile_router_weight: bool = False,
    benchmark_static_routing: bool = False,
    cuda_graph: nemo_automodel.components.models.common.utils.CudaGraphConfig = CudaGraphConfig(),
    csa_forward_full_blocks: bool = False,
    csa_backward_full_blocks: bool = False,
    swa_varlen: bool = False
)

Dataclass

Bases: BackendConfig

Model-owned execution options, disabled by default.

csa_backward_full_blocks

bool = False

Skip custom masks on proven complete backward blocks under the same single-document guard. May change FP32 dK/dV accumulation order. Requires attn=“fa4”.

csa_forward_full_blocks

bool = False

Skip custom masks on proven fully selected forward blocks for single-document, unpadded B1 CP1 inputs on SM90 with head dimension 192. Backward block lists stay unchanged. Requires attn=“fa4”.

swa_varlen

bool = False

Use native FA4 window attention for valid single-batch SM90 D192 query prefixes. Padding keeps its custom-mask path. Requires enabled sparse metadata; communication is unchanged.

nemo_automodel.components.models.step5.configuration.Step5BackendConfig.__post_init__() -> None

class nemo_automodel.components.models.step5.configuration.Step5Config(
    sparse_config: nemo_automodel.components.models.step5.configuration.Step5SparseConfig | dict[str, nemo_automodel.components.models.step5.configuration.ConfigValue] | None = None,
    quantization_config: dict[str, nemo_automodel.components.models.step5.configuration.ConfigValue] | None = None,
    moe_layer_list: list[int] | None = None,
    fp32_residual_connection: bool = True,
    use_moe: bool = True,
    norm_dtype: str = 'float32',
    att_impl_type: str = 'GQA',
    num_sliding_attention_heads: int | None = None,
    head_dim_full_attention: int | None = None,
    head_dim_sliding_attention: int | None = None,
    num_attention_groups_full_attention: int | None = None,
    num_attention_groups_sliding_attention: int | None = None,
    kwargs: nemo_automodel.components.models.step5.configuration.ConfigValue = {}
)

Bases: Step3p7TextConfig

Step-5 language backbone with per-layer RoPE and CSA attention.

architectures

= ['Step5ForCausalLM']

base_config_key

= 'text_config'

model_type

= 'step5'

moe_layer_list

= list(moe_layer_list)

num_experts

= self.moe_num_experts

num_key_value_heads

= self.num_attention_groups

sparse_config

class nemo_automodel.components.models.step5.configuration.Step5SparseConfig(
    enabled: bool = True,
    proxy_dim: int = 256,
    sparse_indexer_rope_dim: int = 32,
    sparse_indexer_use_rope: bool = True,
    sparse_indexer_num_heads: int = 16,
    sparse_indexer_num_k_heads: int = 1,
    sparse_indexer_q_norm_type: str = 'rmsnorm',
    sparse_indexer_k_norm_type: str = 'layernorm',
    sparse_indexer_csa_z_norm_type: str = 'none',
    sparse_indexer_softmax_variant: str = 'ssmax',
    sparse_indexer_ssmax_s_granularity: str = 'q_head',
    num_provider_groups: int = 4,
    topk: int = 512,
    region_block_size: int = 8,
    compression_method: str = 'csa_block_compress',
    attention_impl: str = 'sparse_gqa',
    apply_to_layer_types: list[str] | None = None,
    kwargs: nemo_automodel.components.models.step5.configuration.ConfigValue = {}
)

Bases: PretrainedConfig

CSA region selection geometry and normalization semantics.

apply_to_layer_types

model_type

= 'step5_sparse'

class nemo_automodel.components.models.step5.configuration.Step5VLConfig(
    text_config: nemo_automodel.components.models.step5.configuration.Step5Config | dict[str, nemo_automodel.components.models.step5.configuration.ConfigValue] | None = None,
    vision_config: nemo_automodel.components.models.step5.configuration.Step5VisionConfig | dict[str, nemo_automodel.components.models.step5.configuration.ConfigValue] | None = None,
    image_token_id: int = 128001,
    projector_bias: bool = False,
    understand_projector_stride: int = 2,
    kwargs: nemo_automodel.components.models.step5.configuration.ConfigValue = {}
)

Bases: PretrainedConfig

Step-5 perception encoder and language model configuration.

hidden_size

= self.text_config.hidden_size

max_position_embeddings

= self.text_config.max_position_embeddings

model_type

= 'step5_vl'

sub_configs

text_config

vision_config

vocab_size

= self.text_config.vocab_size

class nemo_automodel.components.models.step5.configuration.Step5VisionConfig(
    rope_theta: float = 10000.0,
    rope_theta_rescale_factor: float = 1.0,
    kwargs: nemo_automodel.components.models.step5.configuration.ConfigValue = {}
)

Bases: StepRoboticsVisionEncoderConfig

Perception configuration with explicit rotary semantics.

nemo_automodel.components.models.step5.configuration.ConfigValue: TypeAlias = bool | int | float | str | torch.dtype | None | list['ConfigValue'] | dict[str, ...