nemo_automodel.components.models.step5.flops

Effective Step-5 text FLOPs, including CSA’s forward-only discrete indexer.

Module Contents

Functions

Name	Description
`_document_work`	Count independent documents; each length excludes physical pack padding.
`_sum_history`	Sum complete historical regions for document-local queries zero through length-1.
`step5_flops`	Count Step-5 text training GEMMs for independent, fully supervised documents.

API

nemo_automodel.components.models.step5.flops._document_work(
    config: nemo_automodel.components.models.step5.configuration.Step5Config,
    lengths: list[int],
    supervised_tokens: int
) -> float

Count independent documents; each length excludes physical pack padding.

nemo_automodel.components.models.step5.flops._sum_history(
    length: int,
    region_size: int,
    cap: int | None
) -> int

Sum complete historical regions for document-local queries zero through length-1.

nemo_automodel.components.models.step5.flops.step5_flops(
    config: nemo_automodel.components.models.step5.configuration.Step5Config,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Count Step-5 text training GEMMs for independent, fully supervised documents.

The count includes trainable linear/attention forward and backward (6 x MACs) and the forward-only indexer (2 x MACs). Embedding lookup, norms, activations, rotary/softmax elementwise work, vision, MTP, optimizer work, kernel padding, masked indexer scores and activation recomputation are excluded. A CSA query at position t sees region_size * min(topk, floor(t / region_size)) historical tokens plus t % region_size + 1 tokens in its current region.

Parameters:

config

Step5Config

Actual complete text configuration, before pipeline layer partitioning.

gbs

intDefaults to 1

Number of independent documents.

seq_len

int | NoneDefaults to None

Tokens per document; defaults to max_position_embeddings.

Returns: float

Effective training FLOPs with all document tokens supervised.