nemo_automodel.components.models.deepseek_v4.flops

Useful model FLOPs for DeepSeek V4 text training.

Module Contents

Functions

Name	Description
`_compressed_pairs`	Sum min(floor(position / ratio), cap) over one-based query positions.
`model_flops`	Count useful GEMM FLOPs for text-only full-parameter training without MTP.

API

nemo_automodel.components.models.deepseek_v4.flops._compressed_pairs(
    length: int,
    ratio: int,
    cap: int
) -> int

Sum min(floor(position / ratio), cap) over one-based query positions.

nemo_automodel.components.models.deepseek_v4.flops.model_flops(
    config: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config,
    gbs: int = 1,
    seq_len: int | None = None
) -> float

Count useful GEMM FLOPs for text-only full-parameter training without MTP.

Projections, active routed/shared experts, hyper-connection coefficient and residual mixing, and causal attention count forward and backward (6 per MAC). The hard-top-k indexer has no gradient and counts forward only (2 per MAC), including its separate overlapping compressor. Ratio-4 attention caps compressed keys at index_topk; other compressed layers use all causal pooled keys. Every layer also uses causal sliding-window attention.

Counts assume one unpadded document per sequence. Activation recomputation, masked-out dense work, communication, optimizer updates, pooling, softmax, norms, RoPE, and elementwise operations are excluded, as for other MFU formulas. Compressor projections run at every token, but have no backward pass when the sequence is too short to produce a complete pooled window.

Parameters:

config

DeepseekV4Config

DeepSeek V4 text configuration with MTP disabled.

gbs

intDefaults to 1

Number of sequences in the global optimizer step.

seq_len

int | NoneDefaults to None

Tokens per sequence; defaults to max_position_embeddings.

Returns: float

Model FLOPs across all devices for one optimizer step.

Raises:

ValueError: Vision, MTP, or invalid sequence/compression settings prevent this text-only accounting.